Skip to main content

mz_persist_client/
fetch.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Fetching batches of data from persist's backing store
11
12use std::fmt::{self, Debug};
13use std::marker::PhantomData;
14use std::sync::Arc;
15use std::time::{Duration, Instant};
16
17use anyhow::anyhow;
18use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, Int64Array};
19use arrow::compute::FilterBuilder;
20use differential_dataflow::difference::Monoid;
21use differential_dataflow::lattice::Lattice;
22use differential_dataflow::trace::Description;
23use itertools::EitherOrBoth;
24use mz_dyncfg::{Config, ConfigSet, ConfigValHandle};
25use mz_ore::bytes::SegmentedBytes;
26use mz_ore::cast::CastFrom;
27use mz_ore::{soft_assert_or_log, soft_panic_no_log, soft_panic_or_log};
28use mz_persist::indexed::columnar::arrow::{realloc_any, realloc_array};
29use mz_persist::indexed::columnar::{ColumnarRecords, ColumnarRecordsStructuredExt};
30use mz_persist::indexed::encoding::{BlobTraceBatchPart, BlobTraceUpdates};
31use mz_persist::location::{Blob, SeqNo};
32use mz_persist::metrics::ColumnarMetrics;
33use mz_persist_types::arrow::ArrayOrd;
34use mz_persist_types::columnar::{ColumnDecoder, Schema, data_type};
35use mz_persist_types::part::Codec64Mut;
36use mz_persist_types::schema::backward_compatible;
37use mz_persist_types::stats::PartStats;
38use mz_persist_types::{Codec, Codec64};
39use mz_proto::RustType;
40use serde::{Deserialize, Serialize};
41use timely::PartialOrder;
42use timely::progress::frontier::AntichainRef;
43use timely::progress::{Antichain, Timestamp};
44use tracing::{Instrument, debug, debug_span, trace_span};
45
46use crate::ShardId;
47use crate::cfg::PersistConfig;
48use crate::error::InvalidUsage;
49use crate::internal::apply::Applier;
50use crate::internal::encoding::{LazyInlineBatchPart, LazyPartStats, LazyProto, Schemas};
51use crate::internal::machine::retry_external;
52use crate::internal::metrics::{Metrics, MetricsPermits, ReadMetrics, ShardMetrics};
53use crate::internal::paths::BlobKey;
54use crate::internal::state::{
55    BatchPart, HollowBatchPart, ProtoHollowBatchPart, ProtoInlineBatchPart,
56};
57use crate::read::LeasedReaderId;
58use crate::schema::{PartMigration, SchemaCache};
59
60pub(crate) const FETCH_SEMAPHORE_COST_ADJUSTMENT: Config<f64> = Config::new(
61    "persist_fetch_semaphore_cost_adjustment",
62    // We use `encoded_size_bytes` as the number of permits, but the parsed size
63    // is larger than the encoded one, so adjust it. This default value is from
64    // eyeballing graphs in experiments that were run on tpch loadgen data.
65    1.2,
66    "\
67    An adjustment multiplied by encoded_size_bytes to approximate an upper \
68    bound on the size in lgalloc, which includes the decoded version.",
69);
70
71pub(crate) const FETCH_SEMAPHORE_PERMIT_ADJUSTMENT: Config<f64> = Config::new(
72    "persist_fetch_semaphore_permit_adjustment",
73    1.0,
74    "\
75    A limit on the number of outstanding persist bytes being fetched and \
76    parsed, expressed as a multiplier of the process's memory limit. This data \
77    all spills to lgalloc, so values > 1.0 are safe. Only applied to cc \
78    replicas.",
79);
80
81pub(crate) const PART_DECODE_FORMAT: Config<&'static str> = Config::new(
82    "persist_part_decode_format",
83    PartDecodeFormat::default().as_str(),
84    "\
85    Format we'll use to decode a Persist Part, either 'row', \
86    'row_with_validate', or 'arrow' (Materialize).",
87);
88
89pub(crate) const OPTIMIZE_IGNORED_DATA_FETCH: Config<bool> = Config::new(
90    "persist_optimize_ignored_data_fetch",
91    true,
92    "CYA to allow opt-out of a performance optimization to skip fetching ignored data",
93);
94
95pub(crate) const VALIDATE_PART_BOUNDS_ON_READ: Config<bool> = Config::new(
96    "persist_validate_part_bounds_on_read",
97    false,
98    "Validate the part lower <= the batch lower and the part upper <= batch upper,\
99    for the batch containing that part",
100);
101
102#[derive(Debug, Clone)]
103pub(crate) struct FetchConfig {
104    pub(crate) validate_bounds_on_read: bool,
105}
106
107impl FetchConfig {
108    pub fn from_persist_config(cfg: &PersistConfig) -> Self {
109        Self {
110            validate_bounds_on_read: VALIDATE_PART_BOUNDS_ON_READ.get(cfg),
111        }
112    }
113}
114
115#[derive(Debug, Clone)]
116pub(crate) struct BatchFetcherConfig {
117    pub(crate) part_decode_format: ConfigValHandle<String>,
118    pub(crate) fetch_config: FetchConfig,
119}
120
121impl BatchFetcherConfig {
122    pub fn new(value: &PersistConfig) -> Self {
123        Self {
124            part_decode_format: PART_DECODE_FORMAT.handle(value),
125            fetch_config: FetchConfig::from_persist_config(value),
126        }
127    }
128
129    pub fn part_decode_format(&self) -> PartDecodeFormat {
130        PartDecodeFormat::from_str(self.part_decode_format.get().as_str())
131    }
132}
133
134/// Capable of fetching [`LeasedBatchPart`] while not holding any capabilities.
135#[derive(Debug)]
136pub struct BatchFetcher<K, V, T, D>
137where
138    T: Timestamp + Lattice + Codec64,
139    // These are only here so we can use them in the auto-expiring `Drop` impl.
140    K: Debug + Codec,
141    V: Debug + Codec,
142    D: Monoid + Codec64 + Send + Sync,
143{
144    pub(crate) cfg: BatchFetcherConfig,
145    pub(crate) blob: Arc<dyn Blob>,
146    pub(crate) metrics: Arc<Metrics>,
147    pub(crate) shard_metrics: Arc<ShardMetrics>,
148    pub(crate) shard_id: ShardId,
149    pub(crate) read_schemas: Schemas<K, V>,
150    pub(crate) schema_cache: SchemaCache<K, V, T, D>,
151    pub(crate) is_transient: bool,
152
153    // Ensures that `BatchFetcher` is of the same type as the `ReadHandle` it's
154    // derived from.
155    pub(crate) _phantom: PhantomData<fn() -> (K, V, T, D)>,
156}
157
158// Hand-written (rather than derived) so cloning does not require `K: Clone`
159// etc.: every field is an `Arc` or independently `Clone`. The `schema_cache`
160// clone shares the schema-lookup maps and applier, so clones reuse cached
161// schema fetches and only duplicate a small per-clone migration memo. Used to
162// run several `fetch_leased_part` calls concurrently, each on its own clone.
163impl<K, V, T, D> Clone for BatchFetcher<K, V, T, D>
164where
165    T: Timestamp + Lattice + Codec64,
166    K: Debug + Codec,
167    V: Debug + Codec,
168    D: Monoid + Codec64 + Send + Sync,
169{
170    fn clone(&self) -> Self {
171        Self {
172            cfg: self.cfg.clone(),
173            blob: Arc::clone(&self.blob),
174            metrics: Arc::clone(&self.metrics),
175            shard_metrics: Arc::clone(&self.shard_metrics),
176            shard_id: self.shard_id.clone(),
177            read_schemas: self.read_schemas.clone(),
178            schema_cache: self.schema_cache.clone(),
179            is_transient: self.is_transient,
180            _phantom: PhantomData,
181        }
182    }
183}
184
185impl<K, V, T, D> BatchFetcher<K, V, T, D>
186where
187    K: Debug + Codec,
188    V: Debug + Codec,
189    T: Timestamp + Lattice + Codec64 + Sync,
190    D: Monoid + Codec64 + Send + Sync,
191{
192    /// Trade in an exchange-able [LeasedBatchPart] for the data it represents.
193    ///
194    /// Note to check the `LeasedBatchPart` documentation for how to handle the
195    /// returned value.
196    pub async fn fetch_leased_part(
197        &mut self,
198        part: ExchangeableBatchPart<T>,
199    ) -> Result<Result<FetchedBlob<K, V, T, D>, BlobKey>, InvalidUsage<T>> {
200        let ExchangeableBatchPart {
201            shard_id,
202            encoded_size_bytes: _,
203            desc,
204            filter,
205            filter_pushdown_audit,
206            part,
207            reader_id: _,
208        } = part;
209        let part: BatchPart<T> = part.decode_to().expect("valid part");
210        if shard_id != self.shard_id {
211            return Err(InvalidUsage::BatchNotFromThisShard {
212                batch_shard: shard_id,
213                handle_shard: self.shard_id.clone(),
214            });
215        }
216
217        let migration =
218            PartMigration::new(&part, self.read_schemas.clone(), &mut self.schema_cache)
219                .await
220                .unwrap_or_else(|read_schemas| {
221                    panic!(
222                        "could not decode part {:?} with schema: {:?}",
223                        part.schema_id(),
224                        read_schemas
225                    )
226                });
227
228        let (buf, fetch_permit) = match &part {
229            BatchPart::Hollow(x) => {
230                let fetch_permit = self
231                    .metrics
232                    .semaphore
233                    .acquire_fetch_permits(x.encoded_size_bytes)
234                    .await;
235                let read_metrics = if self.is_transient {
236                    &self.metrics.read.unindexed
237                } else {
238                    &self.metrics.read.batch_fetcher
239                };
240                let buf = fetch_batch_part_blob(
241                    &shard_id,
242                    self.blob.as_ref(),
243                    &self.metrics,
244                    &self.shard_metrics,
245                    read_metrics,
246                    x,
247                )
248                .await;
249                let buf = match buf {
250                    Ok(buf) => buf,
251                    Err(key) => return Ok(Err(key)),
252                };
253                let buf = FetchedBlobBuf::Hollow {
254                    buf,
255                    part: x.clone(),
256                };
257                (buf, Some(Arc::new(fetch_permit)))
258            }
259            BatchPart::Inline {
260                updates,
261                ts_rewrite,
262                ..
263            } => {
264                let buf = FetchedBlobBuf::Inline {
265                    desc: desc.clone(),
266                    updates: updates.clone(),
267                    ts_rewrite: ts_rewrite.clone(),
268                };
269                (buf, None)
270            }
271        };
272        let fetched_blob = FetchedBlob {
273            metrics: Arc::clone(&self.metrics),
274            read_metrics: self.metrics.read.batch_fetcher.clone(),
275            buf,
276            registered_desc: desc.clone(),
277            migration,
278            filter: filter.clone(),
279            filter_pushdown_audit,
280            structured_part_audit: self.cfg.part_decode_format(),
281            fetch_permit,
282            _phantom: PhantomData,
283            fetch_config: self.cfg.fetch_config.clone(),
284        };
285        Ok(Ok(fetched_blob))
286    }
287
288    /// Diagnoses a missing-blob fetch failure for a part leased by the given
289    /// reader. See the free function `missing_blob_diagnostics`.
290    pub async fn missing_blob_diagnostics(&self, reader_id: &LeasedReaderId) -> String {
291        missing_blob_diagnostics(self.schema_cache.applier(), reader_id).await
292    }
293}
294
295/// Diagnoses a missing-blob fetch failure: refreshes the shard state and
296/// reports whether the reader that leased the part is still present in it.
297///
298/// A missing blob means garbage collection deleted a blob that the part's
299/// lease (a seqno hold in shard state) should have protected. If the reader
300/// has been expired out of state, the lease was lost: this can happen when the
301/// process fails to heartbeat the reader for longer than the lease duration,
302/// e.g. because the machine went to sleep, was starved of CPU or memory, or
303/// was partitioned from consensus. If the reader is still present, the hold
304/// did not protect the blob, which points at a GC or lease-tracking bug.
305pub(crate) async fn missing_blob_diagnostics<K, V, T, D>(
306    applier: &Applier<K, V, T, D>,
307    reader_id: &LeasedReaderId,
308) -> String
309where
310    K: Debug + Codec,
311    V: Debug + Codec,
312    T: Timestamp + Lattice + Codec64 + Sync,
313    D: Monoid + Codec64,
314{
315    // Refreshing state talks to consensus; this runs on an already-fatal path
316    // and a partition from consensus may be the very reason the lease was
317    // lost, so don't let the diagnosis block the restart indefinitely.
318    let refresh = applier.fetch_and_update_state(None);
319    if tokio::time::timeout(Duration::from_secs(30), refresh)
320        .await
321        .is_err()
322    {
323        return format!(
324            "reader {reader_id}: could not refresh state within 30s to diagnose the lease; \
325             partitioned from consensus?"
326        );
327    }
328    match applier.reader_lease(reader_id.clone()) {
329        Some(lease_state) => format!(
330            "reader {reader_id} is still present in state ({lease_state:?}); \
331             a missing blob despite a live lease indicates a GC bug"
332        ),
333        None => format!(
334            "reader {reader_id} has been expired out of state; \
335             the process likely failed to heartbeat it within the lease duration \
336             (machine sleep, CPU/memory starvation, or a partition from consensus?)"
337        ),
338    }
339}
340
341#[derive(Debug, Clone, Serialize, Deserialize)]
342pub(crate) enum FetchBatchFilter<T> {
343    Snapshot {
344        as_of: Antichain<T>,
345    },
346    Listen {
347        as_of: Antichain<T>,
348        lower: Antichain<T>,
349    },
350    Compaction {
351        since: Antichain<T>,
352    },
353}
354
355impl<T: Timestamp + Lattice> FetchBatchFilter<T> {
356    pub(crate) fn filter_ts(&self, t: &mut T) -> bool {
357        match self {
358            FetchBatchFilter::Snapshot { as_of } => {
359                // This time is covered by a listen
360                if as_of.less_than(t) {
361                    return false;
362                }
363                t.advance_by(as_of.borrow());
364                true
365            }
366            FetchBatchFilter::Listen { as_of, lower } => {
367                // This time is covered by a snapshot
368                if !as_of.less_than(t) {
369                    return false;
370                }
371
372                // Because of compaction, the next batch we get might also
373                // contain updates we've already emitted. For example, we
374                // emitted `[1, 2)` and then compaction combined that batch with
375                // a `[2, 3)` batch into a new `[1, 3)` batch. If this happens,
376                // we just need to filter out anything < the frontier. This
377                // frontier was the upper of the last batch (and thus exclusive)
378                // so for the == case, we still emit.
379                if !lower.less_equal(t) {
380                    return false;
381                }
382                true
383            }
384            FetchBatchFilter::Compaction { since } => {
385                t.advance_by(since.borrow());
386                true
387            }
388        }
389    }
390}
391
392/// Trade in an exchange-able [LeasedBatchPart] for the data it represents.
393///
394/// Note to check the `LeasedBatchPart` documentation for how to handle the
395/// returned value.
396pub(crate) async fn fetch_leased_part<K, V, T, D>(
397    cfg: &PersistConfig,
398    part: &LeasedBatchPart<T>,
399    blob: &dyn Blob,
400    metrics: Arc<Metrics>,
401    read_metrics: &ReadMetrics,
402    shard_metrics: &ShardMetrics,
403    reader_id: &LeasedReaderId,
404    read_schemas: Schemas<K, V>,
405    schema_cache: &mut SchemaCache<K, V, T, D>,
406) -> FetchedPart<K, V, T, D>
407where
408    K: Debug + Codec,
409    V: Debug + Codec,
410    T: Timestamp + Lattice + Codec64 + Sync,
411    D: Monoid + Codec64 + Send + Sync,
412{
413    let fetch_config = FetchConfig::from_persist_config(cfg);
414    let encoded_part = match EncodedPart::fetch(
415        &fetch_config,
416        &part.shard_id,
417        blob,
418        &metrics,
419        shard_metrics,
420        read_metrics,
421        &part.desc,
422        &part.part,
423    )
424    .await
425    {
426        Ok(x) => x,
427        Err(blob_key) => {
428            // Ideally, readers should never encounter a missing blob. They place a seqno
429            // hold as they consume their snapshot/listen, preventing any blobs they need
430            // from being deleted by garbage collection, and all blob implementations are
431            // linearizable so there should be no possibility of stale reads.
432            //
433            // If we do have a bug and a reader does encounter a missing blob, the state
434            // cannot be recovered, and our best option is to panic and retry the whole
435            // process.
436            let diagnostics = missing_blob_diagnostics(schema_cache.applier(), reader_id).await;
437            panic!("could not fetch batch part {}: {}", blob_key, diagnostics)
438        }
439    };
440    let part_cfg = BatchFetcherConfig::new(cfg);
441    let migration = PartMigration::new(&part.part, read_schemas, schema_cache)
442        .await
443        .unwrap_or_else(|read_schemas| {
444            panic!(
445                "could not decode part {:?} with schema: {:?}",
446                part.part.schema_id(),
447                read_schemas
448            )
449        });
450    FetchedPart::new(
451        metrics,
452        encoded_part,
453        migration,
454        part.filter.clone(),
455        part.filter_pushdown_audit,
456        part_cfg.part_decode_format(),
457        part.part.stats(),
458    )
459}
460
461pub(crate) async fn fetch_batch_part_blob<T>(
462    shard_id: &ShardId,
463    blob: &dyn Blob,
464    metrics: &Metrics,
465    shard_metrics: &ShardMetrics,
466    read_metrics: &ReadMetrics,
467    part: &HollowBatchPart<T>,
468) -> Result<SegmentedBytes, BlobKey> {
469    let now = Instant::now();
470    let get_span = debug_span!("fetch_batch::get");
471    let blob_key = part.key.complete(shard_id);
472    let value = retry_external(&metrics.retries.external.fetch_batch_get, || async {
473        shard_metrics.blob_gets.inc();
474        // Name the blob in the error. A GET stuck retrying forever surfaces only in the retry log
475        // (see `retry_external`), which prints the error, so without this the log cannot say which
476        // blob, and thus which shard, is wedged.
477        blob.get(&blob_key)
478            .await
479            .map_err(|err| err.context(format!("blob {blob_key}")))
480    })
481    .instrument(get_span.clone())
482    .await
483    .ok_or(blob_key)?;
484
485    drop(get_span);
486
487    read_metrics.part_count.inc();
488    read_metrics.part_bytes.inc_by(u64::cast_from(value.len()));
489    read_metrics.seconds.inc_by(now.elapsed().as_secs_f64());
490
491    Ok(value)
492}
493
494pub(crate) fn decode_batch_part_blob<T>(
495    cfg: &FetchConfig,
496    metrics: &Metrics,
497    read_metrics: &ReadMetrics,
498    registered_desc: Description<T>,
499    part: &HollowBatchPart<T>,
500    buf: &SegmentedBytes,
501) -> EncodedPart<T>
502where
503    T: Timestamp + Lattice + Codec64,
504{
505    trace_span!("fetch_batch::decode").in_scope(|| {
506        let parsed = metrics
507            .codecs
508            .batch
509            .decode(|| BlobTraceBatchPart::decode(buf, &metrics.columnar))
510            .map_err(|err| anyhow!("couldn't decode batch at key {}: {}", part.key, err))
511            // We received a State that we couldn't decode. This could happen if
512            // persist messes up backward/forward compatibility, if the durable
513            // data was corrupted, or if operations messes up deployment. In any
514            // case, fail loudly.
515            .expect("internal error: invalid encoded state");
516        read_metrics
517            .part_goodbytes
518            .inc_by(u64::cast_from(parsed.updates.goodbytes()));
519        EncodedPart::from_hollow(cfg, read_metrics.clone(), registered_desc, part, parsed)
520    })
521}
522
523pub(crate) async fn fetch_batch_part<T>(
524    cfg: &FetchConfig,
525    shard_id: &ShardId,
526    blob: &dyn Blob,
527    metrics: &Metrics,
528    shard_metrics: &ShardMetrics,
529    read_metrics: &ReadMetrics,
530    registered_desc: &Description<T>,
531    part: &HollowBatchPart<T>,
532) -> Result<EncodedPart<T>, BlobKey>
533where
534    T: Timestamp + Lattice + Codec64,
535{
536    let buf =
537        fetch_batch_part_blob(shard_id, blob, metrics, shard_metrics, read_metrics, part).await?;
538    let part = decode_batch_part_blob(
539        cfg,
540        metrics,
541        read_metrics,
542        registered_desc.clone(),
543        part,
544        &buf,
545    );
546    Ok(part)
547}
548
549/// This represents the lease of a seqno. It's generally paired with some external state,
550/// like a hollow part: holding this lease indicates that we may still want to fetch that part,
551/// and should hold back GC to keep it around.
552///
553/// Generally the state and lease are bundled together, as in [LeasedBatchPart]... but sometimes
554/// it's necessary to handle them separately, so this struct is exposed as well. Handle with care.
555#[derive(Clone, Debug)]
556pub struct Lease(Arc<SeqNo>);
557
558impl Lease {
559    /// Creates a new [Lease] that holds the given [SeqNo].
560    pub fn new(seqno: SeqNo) -> Self {
561        Self(Arc::new(seqno))
562    }
563
564    /// Returns the inner [SeqNo] of this [Lease].
565    pub fn seqno(&self) -> SeqNo {
566        *self.0
567    }
568
569    /// Returns the number of live copies of this lease, including this one.
570    pub fn count(&self) -> usize {
571        Arc::strong_count(&self.0)
572    }
573}
574
575/// A token representing one fetch-able batch part.
576///
577/// It is tradeable via `crate::fetch::fetch_batch` for the resulting data
578/// stored in the part.
579///
580/// # Exchange
581///
582/// You can exchange `LeasedBatchPart`:
583/// - If `leased_seqno.is_none()`
584/// - By converting it to [`ExchangeableBatchPart`] through
585///   `Self::into_exchangeable_part`. [`ExchangeableBatchPart`] is exchangeable,
586///   including over the network.
587///
588/// n.b. `Self::into_exchangeable_part` is known to be equivalent to
589/// `SerdeLeasedBatchPart::from(self)`, but we want the additional warning message to
590/// be visible and sufficiently scary.
591///
592/// # Panics
593/// `LeasedBatchPart` panics when dropped unless a very strict set of invariants are
594/// held:
595///
596/// `LeasedBatchPart` may only be dropped if it:
597/// - Does not have a leased `SeqNo (i.e. `self.leased_seqno.is_none()`)
598///
599/// In any other circumstance, dropping `LeasedBatchPart` panics.
600#[derive(Debug)]
601pub struct LeasedBatchPart<T> {
602    pub(crate) metrics: Arc<Metrics>,
603    pub(crate) shard_id: ShardId,
604    pub(crate) filter: FetchBatchFilter<T>,
605    pub(crate) desc: Description<T>,
606    pub(crate) part: BatchPart<T>,
607    /// The lease that prevents this part from being GCed. Code should ensure that this lease
608    /// lives as long as the part is needed.
609    pub(crate) lease: Lease,
610    /// The id of the reader that leased this part, for diagnostics: if the
611    /// blob backing the part goes missing, knowing whether this reader is
612    /// still present in state distinguishes a lost lease from a GC bug.
613    pub(crate) reader_id: LeasedReaderId,
614    pub(crate) filter_pushdown_audit: bool,
615}
616
617impl<T> LeasedBatchPart<T>
618where
619    T: Timestamp + Codec64,
620{
621    /// Takes `self` into a [`ExchangeableBatchPart`], which allows `self` to be
622    /// exchanged (potentially across the network).
623    ///
624    /// !!!WARNING!!!
625    ///
626    /// This method also returns the [Lease] associated with the given part, since
627    /// that can't travel across process boundaries. The caller is responsible for
628    /// ensuring that the lease is held for as long as the batch part may be in use:
629    /// dropping it too early may cause a fetch to fail.
630    pub(crate) fn into_exchangeable_part(self) -> (ExchangeableBatchPart<T>, Lease) {
631        // If `x` has a lease, we've effectively transferred it to `r`.
632        let lease = self.lease.clone();
633        let part = ExchangeableBatchPart {
634            shard_id: self.shard_id,
635            encoded_size_bytes: self.part.encoded_size_bytes(),
636            desc: self.desc.clone(),
637            filter: self.filter.clone(),
638            part: LazyProto::from(&self.part.into_proto()),
639            reader_id: self.reader_id.clone(),
640            filter_pushdown_audit: self.filter_pushdown_audit,
641        };
642        (part, lease)
643    }
644
645    /// The encoded size of this part in bytes
646    pub fn encoded_size_bytes(&self) -> usize {
647        self.part.encoded_size_bytes()
648    }
649
650    /// The filter has indicated we don't need this part, we can verify the
651    /// ongoing end-to-end correctness of corner cases via "audit". This means
652    /// we fetch the part like normal and if the MFP keeps anything from it,
653    /// then something has gone horribly wrong.
654    pub fn request_filter_pushdown_audit(&mut self) {
655        self.filter_pushdown_audit = true;
656    }
657
658    /// Returns the pushdown stats for this part.
659    ///
660    /// Stats written by a newer version may not decode; those return `None`,
661    /// the same as a part that carries no stats.
662    pub fn stats(&self) -> Option<PartStats> {
663        self.part.stats().and_then(|x| x.try_decode().ok())
664    }
665
666    /// Apply any relevant projection pushdown optimizations, assuming that the data in the part
667    /// is equivalent to the provided key and value.
668    pub fn maybe_optimize(&mut self, cfg: &ConfigSet, key: ArrayRef, val: ArrayRef) {
669        assert_eq!(key.len(), 1, "expect a single-row key array");
670        assert_eq!(val.len(), 1, "expect a single-row val array");
671        let as_of = match &self.filter {
672            FetchBatchFilter::Snapshot { as_of } => as_of,
673            FetchBatchFilter::Listen { .. } | FetchBatchFilter::Compaction { .. } => return,
674        };
675        if !OPTIMIZE_IGNORED_DATA_FETCH.get(cfg) {
676            return;
677        }
678        let (diffs_sum, _stats) = match &self.part {
679            BatchPart::Hollow(x) => (x.diffs_sum, x.stats.as_ref()),
680            BatchPart::Inline { .. } => return,
681        };
682        debug!(
683            "try_optimize_ignored_data_fetch diffs_sum={:?} as_of={:?} lower={:?} upper={:?}",
684            // This is only used for debugging, so hack to assume that D is i64.
685            diffs_sum.map(i64::decode),
686            as_of.elements(),
687            self.desc.lower().elements(),
688            self.desc.upper().elements()
689        );
690        let as_of = match &as_of.elements() {
691            &[as_of] => as_of,
692            _ => return,
693        };
694        // NOTE: `diffs_sum` sums every row physically in the blob, while
695        // reads truncate rows outside the registered desc. Substituting it is
696        // sound only while no writer registers a batch with tighter bounds
697        // than the blob holds (none does today, and rewritten batches prove
698        // it), which nothing here can re-check without fetching the blob.
699        let eligible = self.desc.upper().less_equal(as_of) && self.desc.since().less_equal(as_of);
700        if !eligible {
701            return;
702        }
703        let Some(diffs_sum) = diffs_sum else {
704            return;
705        };
706
707        debug!(
708            "try_optimize_ignored_data_fetch faked {:?} diffs at ts {:?} skipping fetch of {} bytes",
709            // This is only used for debugging, so hack to assume that D is i64.
710            i64::decode(diffs_sum),
711            as_of,
712            self.part.encoded_size_bytes(),
713        );
714        self.metrics.pushdown.parts_faked_count.inc();
715        self.metrics
716            .pushdown
717            .parts_faked_bytes
718            .inc_by(u64::cast_from(self.part.encoded_size_bytes()));
719        let timestamps = {
720            let mut col = Codec64Mut::with_capacity(1);
721            col.push(as_of);
722            col.finish()
723        };
724        let diffs = {
725            let mut col = Codec64Mut::with_capacity(1);
726            col.push_raw(diffs_sum);
727            col.finish()
728        };
729        let updates = BlobTraceUpdates::Structured {
730            key_values: ColumnarRecordsStructuredExt { key, val },
731            timestamps,
732            diffs,
733        };
734        let faked_data = LazyInlineBatchPart::from(&ProtoInlineBatchPart {
735            desc: Some(self.desc.into_proto()),
736            index: 0,
737            updates: Some(updates.into_proto()),
738        });
739        self.part = BatchPart::Inline {
740            updates: faked_data,
741            ts_rewrite: None,
742            schema_id: None,
743            deprecated_schema_id: None,
744        };
745    }
746}
747
748impl<T> Drop for LeasedBatchPart<T> {
749    /// For details, see [`LeasedBatchPart`].
750    fn drop(&mut self) {
751        self.metrics.lease.dropped_part.inc()
752    }
753}
754
755/// A [Blob] object that has been fetched, but not at all decoded.
756///
757/// In contrast to [FetchedPart], this representation hasn't yet done parquet
758/// decoding.
759#[derive(Debug)]
760pub struct FetchedBlob<K: Codec, V: Codec, T, D> {
761    metrics: Arc<Metrics>,
762    read_metrics: ReadMetrics,
763    buf: FetchedBlobBuf<T>,
764    registered_desc: Description<T>,
765    migration: PartMigration<K, V>,
766    filter: FetchBatchFilter<T>,
767    filter_pushdown_audit: bool,
768    structured_part_audit: PartDecodeFormat,
769    fetch_permit: Option<Arc<MetricsPermits>>,
770    fetch_config: FetchConfig,
771    _phantom: PhantomData<fn() -> D>,
772}
773
774#[derive(Debug, Clone)]
775enum FetchedBlobBuf<T> {
776    Hollow {
777        buf: SegmentedBytes,
778        part: HollowBatchPart<T>,
779    },
780    Inline {
781        desc: Description<T>,
782        updates: LazyInlineBatchPart,
783        ts_rewrite: Option<Antichain<T>>,
784    },
785}
786
787impl<K: Codec, V: Codec, T: Clone, D> Clone for FetchedBlob<K, V, T, D> {
788    fn clone(&self) -> Self {
789        Self {
790            metrics: Arc::clone(&self.metrics),
791            read_metrics: self.read_metrics.clone(),
792            buf: self.buf.clone(),
793            registered_desc: self.registered_desc.clone(),
794            migration: self.migration.clone(),
795            filter: self.filter.clone(),
796            filter_pushdown_audit: self.filter_pushdown_audit.clone(),
797            fetch_permit: self.fetch_permit.clone(),
798            structured_part_audit: self.structured_part_audit.clone(),
799            fetch_config: self.fetch_config.clone(),
800            _phantom: self._phantom.clone(),
801        }
802    }
803}
804
805/// [FetchedPart] but with an accompanying permit from the fetch mem/disk
806/// semaphore.
807pub struct ShardSourcePart<K: Codec, V: Codec, T, D> {
808    /// The underlying [FetchedPart].
809    pub part: FetchedPart<K, V, T, D>,
810    fetch_permit: Option<Arc<MetricsPermits>>,
811}
812
813impl<K, V, T: Debug, D: Debug> Debug for ShardSourcePart<K, V, T, D>
814where
815    K: Codec + Debug,
816    <K as Codec>::Storage: Debug,
817    V: Codec + Debug,
818    <V as Codec>::Storage: Debug,
819{
820    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
821        let ShardSourcePart { part, fetch_permit } = self;
822        f.debug_struct("ShardSourcePart")
823            .field("part", part)
824            .field("fetch_permit", fetch_permit)
825            .finish()
826    }
827}
828
829impl<K: Codec, V: Codec, T: Timestamp + Lattice + Codec64, D> FetchedBlob<K, V, T, D> {
830    /// Partially decodes this blob into a [FetchedPart].
831    pub fn parse(&self) -> ShardSourcePart<K, V, T, D> {
832        self.parse_internal(&self.fetch_config)
833    }
834
835    /// Partially decodes this blob into a [FetchedPart].
836    pub(crate) fn parse_internal(&self, cfg: &FetchConfig) -> ShardSourcePart<K, V, T, D> {
837        let (part, stats) = match &self.buf {
838            FetchedBlobBuf::Hollow { buf, part } => {
839                let parsed = decode_batch_part_blob(
840                    cfg,
841                    &self.metrics,
842                    &self.read_metrics,
843                    self.registered_desc.clone(),
844                    part,
845                    buf,
846                );
847                (parsed, part.stats.as_ref())
848            }
849            FetchedBlobBuf::Inline {
850                desc,
851                updates,
852                ts_rewrite,
853            } => {
854                let parsed = EncodedPart::from_inline(
855                    cfg,
856                    &self.metrics,
857                    self.read_metrics.clone(),
858                    desc.clone(),
859                    updates,
860                    ts_rewrite.as_ref(),
861                );
862                (parsed, None)
863            }
864        };
865        let part = FetchedPart::new(
866            Arc::clone(&self.metrics),
867            part,
868            self.migration.clone(),
869            self.filter.clone(),
870            self.filter_pushdown_audit,
871            self.structured_part_audit,
872            stats,
873        );
874        ShardSourcePart {
875            part,
876            fetch_permit: self.fetch_permit.clone(),
877        }
878    }
879
880    /// Decodes and returns the pushdown stats for this part, if known.
881    ///
882    /// Stats written by a newer version may not decode; those return `None`,
883    /// the same as a part that carries no stats.
884    pub fn stats(&self) -> Option<PartStats> {
885        match &self.buf {
886            FetchedBlobBuf::Hollow { part, .. } => {
887                part.stats.as_ref().and_then(|x| x.try_decode().ok())
888            }
889            FetchedBlobBuf::Inline { .. } => None,
890        }
891    }
892}
893
894/// A [Blob] object that has been fetched, but not yet fully decoded.
895///
896/// In contrast to [FetchedBlob], this representation has already done parquet
897/// decoding.
898#[derive(Debug)]
899pub struct FetchedPart<K: Codec, V: Codec, T, D> {
900    metrics: Arc<Metrics>,
901    ts_filter: FetchBatchFilter<T>,
902    // If migration is Either, then the columnar one will have already been
903    // applied here on the structured data only.
904    part: EitherOrBoth<
905        ColumnarRecords,
906        (
907            <K::Schema as Schema<K>>::Decoder,
908            <V::Schema as Schema<V>>::Decoder,
909        ),
910    >,
911    timestamps: Int64Array,
912    diffs: Int64Array,
913    migration: PartMigration<K, V>,
914    filter_pushdown_audit: Option<LazyPartStats>,
915    peek_stash: Option<((K, V), T, D)>,
916    part_cursor: usize,
917    key_storage: Option<K::Storage>,
918    val_storage: Option<V::Storage>,
919
920    _phantom: PhantomData<fn() -> D>,
921}
922
923impl<K: Codec, V: Codec, T: Timestamp + Lattice + Codec64, D> FetchedPart<K, V, T, D> {
924    pub(crate) fn new(
925        metrics: Arc<Metrics>,
926        part: EncodedPart<T>,
927        migration: PartMigration<K, V>,
928        ts_filter: FetchBatchFilter<T>,
929        filter_pushdown_audit: bool,
930        part_decode_format: PartDecodeFormat,
931        stats: Option<&LazyPartStats>,
932    ) -> Self {
933        let part_len = u64::cast_from(part.part.updates.len());
934        match &migration {
935            PartMigration::SameSchema { .. } => metrics.schema.migration_count_same.inc(),
936            PartMigration::Schemaless { .. } => {
937                metrics.schema.migration_count_codec.inc();
938                metrics.schema.migration_len_legacy_codec.inc_by(part_len);
939            }
940            PartMigration::Either { .. } => {
941                metrics.schema.migration_count_either.inc();
942                match part_decode_format {
943                    PartDecodeFormat::Row {
944                        validate_structured: false,
945                    } => metrics.schema.migration_len_either_codec.inc_by(part_len),
946                    PartDecodeFormat::Row {
947                        validate_structured: true,
948                    } => {
949                        metrics.schema.migration_len_either_codec.inc_by(part_len);
950                        metrics.schema.migration_len_either_arrow.inc_by(part_len);
951                    }
952                    PartDecodeFormat::Arrow => {
953                        metrics.schema.migration_len_either_arrow.inc_by(part_len)
954                    }
955                }
956            }
957        }
958
959        let filter_pushdown_audit = if filter_pushdown_audit {
960            stats.cloned()
961        } else {
962            None
963        };
964
965        let downcast_structured = |structured: ColumnarRecordsStructuredExt,
966                                   structured_only: bool| {
967            let key_size_before = ArrayOrd::new(&structured.key).goodbytes();
968
969            let structured = match &migration {
970                PartMigration::SameSchema { .. } => structured,
971                PartMigration::Schemaless { read } if structured_only => {
972                    // We don't know the source schema, but we do know the source datatype; migrate it directly.
973                    let start = Instant::now();
974                    let read_key = data_type::<K>(&*read.key).ok()?;
975                    let read_val = data_type::<V>(&*read.val).ok()?;
976                    let key_migration = backward_compatible(structured.key.data_type(), &read_key)?;
977                    let val_migration = backward_compatible(structured.val.data_type(), &read_val)?;
978                    let key = key_migration.migrate(structured.key);
979                    let val = val_migration.migrate(structured.val);
980                    metrics
981                        .schema
982                        .migration_migrate_seconds
983                        .inc_by(start.elapsed().as_secs_f64());
984                    ColumnarRecordsStructuredExt { key, val }
985                }
986                PartMigration::Schemaless { .. } => return None,
987                PartMigration::Either {
988                    write: _,
989                    read: _,
990                    key_migration,
991                    val_migration,
992                } => {
993                    let start = Instant::now();
994                    let key = key_migration.migrate(structured.key);
995                    let val = val_migration.migrate(structured.val);
996                    metrics
997                        .schema
998                        .migration_migrate_seconds
999                        .inc_by(start.elapsed().as_secs_f64());
1000                    ColumnarRecordsStructuredExt { key, val }
1001                }
1002            };
1003
1004            let read_schema = migration.codec_read();
1005            let key = K::Schema::decoder_any(&*read_schema.key, &*structured.key);
1006            let val = V::Schema::decoder_any(&*read_schema.val, &*structured.val);
1007
1008            match &key {
1009                Ok(key_decoder) => {
1010                    let key_size_after = key_decoder.goodbytes();
1011                    let key_diff = key_size_before.saturating_sub(key_size_after);
1012                    metrics
1013                        .pushdown
1014                        .parts_projection_trimmed_bytes
1015                        .inc_by(u64::cast_from(key_diff));
1016                }
1017                Err(e) => {
1018                    soft_panic_or_log!("failed to create decoder: {e:#?}");
1019                }
1020            }
1021
1022            Some((key.ok()?, val.ok()?))
1023        };
1024
1025        let updates = part.normalize(&metrics.columnar);
1026        let timestamps = updates.timestamps().clone();
1027        let diffs = updates.diffs().clone();
1028        let part = match updates {
1029            // If only one encoding is available, decode via that encoding.
1030            BlobTraceUpdates::Row(records) => EitherOrBoth::Left(records),
1031            BlobTraceUpdates::Structured { key_values, .. } => EitherOrBoth::Right(
1032                // The structured-only data format was added after schema ids were recorded everywhere,
1033                // so we expect this data to be present.
1034                downcast_structured(key_values, true).expect("valid schemas for structured data"),
1035            ),
1036            // If both are available, respect the specified part decode format.
1037            BlobTraceUpdates::Both(records, ext) => match part_decode_format {
1038                PartDecodeFormat::Row {
1039                    validate_structured: false,
1040                } => EitherOrBoth::Left(records),
1041                PartDecodeFormat::Row {
1042                    validate_structured: true,
1043                } => match downcast_structured(ext, false) {
1044                    Some(decoders) => EitherOrBoth::Both(records, decoders),
1045                    None => EitherOrBoth::Left(records),
1046                },
1047                PartDecodeFormat::Arrow => match downcast_structured(ext, false) {
1048                    Some(decoders) => EitherOrBoth::Right(decoders),
1049                    None => EitherOrBoth::Left(records),
1050                },
1051            },
1052        };
1053
1054        FetchedPart {
1055            metrics,
1056            ts_filter,
1057            part,
1058            peek_stash: None,
1059            timestamps,
1060            diffs,
1061            migration,
1062            filter_pushdown_audit,
1063            part_cursor: 0,
1064            key_storage: None,
1065            val_storage: None,
1066            _phantom: PhantomData,
1067        }
1068    }
1069
1070    /// Returns Some if this part was only fetched as part of a filter pushdown
1071    /// audit. See [LeasedBatchPart::request_filter_pushdown_audit].
1072    ///
1073    /// If set, the value in the Option is for debugging and should be included
1074    /// in any error messages.
1075    pub fn is_filter_pushdown_audit(&self) -> Option<impl std::fmt::Debug + use<K, V, T, D>> {
1076        self.filter_pushdown_audit.clone()
1077    }
1078}
1079
1080/// A [Blob] object that has been fetched, but has no associated decoding
1081/// logic.
1082#[derive(Debug)]
1083pub(crate) struct EncodedPart<T> {
1084    metrics: ReadMetrics,
1085    registered_desc: Description<T>,
1086    part: BlobTraceBatchPart<T>,
1087    needs_truncation: bool,
1088    ts_rewrite: Option<Antichain<T>>,
1089}
1090
1091impl<K, V, T, D> FetchedPart<K, V, T, D>
1092where
1093    K: Debug + Codec,
1094    V: Debug + Codec,
1095    T: Timestamp + Lattice + Codec64,
1096    D: Monoid + Codec64 + Send + Sync,
1097{
1098    /// [Self::next] but optionally providing a `K` and `V` for alloc reuse.
1099    ///
1100    /// When `result_override` is specified, return it instead of decoding data.
1101    /// This is used when we know the decoded result will be ignored.
1102    pub fn next_with_storage(
1103        &mut self,
1104        key: &mut Option<K>,
1105        val: &mut Option<V>,
1106    ) -> Option<((K, V), T, D)> {
1107        let mut consolidated = self.peek_stash.take();
1108        loop {
1109            // Fetch and decode the next tuple in the sequence. (Or break if there is none.)
1110            let next = if self.part_cursor < self.timestamps.len() {
1111                let next_idx = self.part_cursor;
1112                self.part_cursor += 1;
1113                // These `to_le_bytes` calls were previously encapsulated by `ColumnarRecords`.
1114                // TODO(structured): re-encapsulate these once we've finished the structured migration.
1115                let mut t = T::decode(self.timestamps.values()[next_idx].to_le_bytes());
1116                if !self.ts_filter.filter_ts(&mut t) {
1117                    continue;
1118                }
1119                let d = D::decode(self.diffs.values()[next_idx].to_le_bytes());
1120                if d.is_zero() {
1121                    continue;
1122                }
1123                let kv = self.decode_kv(next_idx, key, val);
1124                (kv, t, d)
1125            } else {
1126                break;
1127            };
1128
1129            // Attempt to consolidate in the next tuple, stashing it if that's not possible.
1130            if let Some((kv, t, d)) = &mut consolidated {
1131                let (kv_next, t_next, d_next) = &next;
1132                if kv == kv_next && t == t_next {
1133                    d.plus_equals(d_next);
1134                    if d.is_zero() {
1135                        consolidated = None;
1136                    }
1137                } else {
1138                    self.peek_stash = Some(next);
1139                    break;
1140                }
1141            } else {
1142                consolidated = Some(next);
1143            }
1144        }
1145
1146        let (kv, t, d) = consolidated?;
1147
1148        Some((kv, t, d))
1149    }
1150
1151    fn decode_kv(&mut self, index: usize, key: &mut Option<K>, val: &mut Option<V>) -> (K, V) {
1152        let decoded = self
1153            .part
1154            .as_ref()
1155            .map_left(|codec| {
1156                let ((ck, cv), _, _) = codec.get(index).expect("valid index");
1157                let (k, v) = Self::decode_codec(
1158                    &*self.metrics,
1159                    self.migration.codec_read(),
1160                    ck,
1161                    cv,
1162                    key,
1163                    val,
1164                    &mut self.key_storage,
1165                    &mut self.val_storage,
1166                );
1167                (k.expect("valid legacy key"), v.expect("valid legacy value"))
1168            })
1169            .map_right(|(structured_key, structured_val)| {
1170                self.decode_structured(index, structured_key, structured_val, key, val)
1171            });
1172
1173        match decoded {
1174            EitherOrBoth::Both((k, v), (k_s, v_s)) => {
1175                // Purposefully do not trace to prevent blowing up Sentry.
1176                let is_valid = self
1177                    .metrics
1178                    .columnar
1179                    .arrow()
1180                    .key()
1181                    .report_valid(|| k_s == k);
1182                if !is_valid {
1183                    soft_panic_no_log!("structured key did not match, {k_s:?} != {k:?}");
1184                }
1185                // Purposefully do not trace to prevent blowing up Sentry.
1186                let is_valid = self
1187                    .metrics
1188                    .columnar
1189                    .arrow()
1190                    .val()
1191                    .report_valid(|| v_s == v);
1192                if !is_valid {
1193                    soft_panic_no_log!("structured val did not match, {v_s:?} != {v:?}");
1194                }
1195
1196                (k, v)
1197            }
1198            EitherOrBoth::Left(kv) => kv,
1199            EitherOrBoth::Right(kv) => kv,
1200        }
1201    }
1202
1203    fn decode_codec(
1204        metrics: &Metrics,
1205        read_schemas: &Schemas<K, V>,
1206        key_buf: &[u8],
1207        val_buf: &[u8],
1208        key: &mut Option<K>,
1209        val: &mut Option<V>,
1210        key_storage: &mut Option<K::Storage>,
1211        val_storage: &mut Option<V::Storage>,
1212    ) -> (Result<K, String>, Result<V, String>) {
1213        let k = metrics.codecs.key.decode(|| match key.take() {
1214            Some(mut key) => {
1215                match K::decode_from(&mut key, key_buf, key_storage, &read_schemas.key) {
1216                    Ok(()) => Ok(key),
1217                    Err(err) => Err(err),
1218                }
1219            }
1220            None => K::decode(key_buf, &read_schemas.key),
1221        });
1222        let v = metrics.codecs.val.decode(|| match val.take() {
1223            Some(mut val) => {
1224                match V::decode_from(&mut val, val_buf, val_storage, &read_schemas.val) {
1225                    Ok(()) => Ok(val),
1226                    Err(err) => Err(err),
1227                }
1228            }
1229            None => V::decode(val_buf, &read_schemas.val),
1230        });
1231        (k, v)
1232    }
1233
1234    fn decode_structured(
1235        &self,
1236        idx: usize,
1237        keys: &<K::Schema as Schema<K>>::Decoder,
1238        vals: &<V::Schema as Schema<V>>::Decoder,
1239        key: &mut Option<K>,
1240        val: &mut Option<V>,
1241    ) -> (K, V) {
1242        let mut key = key.take().unwrap_or_default();
1243        keys.decode(idx, &mut key);
1244
1245        let mut val = val.take().unwrap_or_default();
1246        vals.decode(idx, &mut val);
1247
1248        (key, val)
1249    }
1250}
1251
1252impl<K, V, T, D> Iterator for FetchedPart<K, V, T, D>
1253where
1254    K: Debug + Codec,
1255    V: Debug + Codec,
1256    T: Timestamp + Lattice + Codec64,
1257    D: Monoid + Codec64 + Send + Sync,
1258{
1259    type Item = ((K, V), T, D);
1260
1261    fn next(&mut self) -> Option<Self::Item> {
1262        self.next_with_storage(&mut None, &mut None)
1263    }
1264
1265    fn size_hint(&self) -> (usize, Option<usize>) {
1266        // We don't know in advance how restrictive the filter will be.
1267        let max_len = self.timestamps.len();
1268        (0, Some(max_len))
1269    }
1270}
1271
1272impl<T> EncodedPart<T>
1273where
1274    T: Timestamp + Lattice + Codec64,
1275{
1276    pub async fn fetch(
1277        cfg: &FetchConfig,
1278        shard_id: &ShardId,
1279        blob: &dyn Blob,
1280        metrics: &Metrics,
1281        shard_metrics: &ShardMetrics,
1282        read_metrics: &ReadMetrics,
1283        registered_desc: &Description<T>,
1284        part: &BatchPart<T>,
1285    ) -> Result<Self, BlobKey> {
1286        match part {
1287            BatchPart::Hollow(x) => {
1288                fetch_batch_part(
1289                    cfg,
1290                    shard_id,
1291                    blob,
1292                    metrics,
1293                    shard_metrics,
1294                    read_metrics,
1295                    registered_desc,
1296                    x,
1297                )
1298                .await
1299            }
1300            BatchPart::Inline {
1301                updates,
1302                ts_rewrite,
1303                ..
1304            } => Ok(EncodedPart::from_inline(
1305                cfg,
1306                metrics,
1307                read_metrics.clone(),
1308                registered_desc.clone(),
1309                updates,
1310                ts_rewrite.as_ref(),
1311            )),
1312        }
1313    }
1314
1315    pub(crate) fn from_inline(
1316        cfg: &FetchConfig,
1317        metrics: &Metrics,
1318        read_metrics: ReadMetrics,
1319        desc: Description<T>,
1320        x: &LazyInlineBatchPart,
1321        ts_rewrite: Option<&Antichain<T>>,
1322    ) -> Self {
1323        let parsed = x.decode(&metrics.columnar).expect("valid inline part");
1324        Self::new(cfg, read_metrics, desc, "inline", ts_rewrite, parsed)
1325    }
1326
1327    pub(crate) fn from_hollow(
1328        cfg: &FetchConfig,
1329        metrics: ReadMetrics,
1330        registered_desc: Description<T>,
1331        part: &HollowBatchPart<T>,
1332        parsed: BlobTraceBatchPart<T>,
1333    ) -> Self {
1334        Self::new(
1335            cfg,
1336            metrics,
1337            registered_desc,
1338            &part.key.0,
1339            part.ts_rewrite.as_ref(),
1340            parsed,
1341        )
1342    }
1343
1344    pub(crate) fn new(
1345        cfg: &FetchConfig,
1346        metrics: ReadMetrics,
1347        registered_desc: Description<T>,
1348        printable_name: &str,
1349        ts_rewrite: Option<&Antichain<T>>,
1350        parsed: BlobTraceBatchPart<T>,
1351    ) -> Self {
1352        // There are two types of batches in persist:
1353        // - Batches written by a persist user (either directly or indirectly
1354        //   via BatchBuilder). These always have a since of the minimum
1355        //   timestamp and may be registered in persist state with a tighter set
1356        //   of bounds than are inline in the batch (truncation). To read one of
1357        //   these batches, all data physically in the batch but outside of the
1358        //   truncated bounds must be ignored. Not every user batch is
1359        //   truncated.
1360        // - Batches written by compaction. These always have an inline desc
1361        //   lower and upper that matches the registered desc lower and upper,
1362        //   and a since that is less than or equal to the registered desc.
1363        //   The inline since may be less than the registered desc since,
1364        //   this is because of incremental compaction, where we might rewrite
1365        //   certain runs in a batch but not others.
1366        let inline_desc = &parsed.desc;
1367        let needs_truncation = inline_desc.lower() != registered_desc.lower()
1368            || inline_desc.upper() != registered_desc.upper();
1369        if needs_truncation {
1370            if cfg.validate_bounds_on_read {
1371                soft_assert_or_log!(
1372                    PartialOrder::less_equal(inline_desc.lower(), registered_desc.lower()),
1373                    "key={} inline={:?} registered={:?}",
1374                    printable_name,
1375                    inline_desc,
1376                    registered_desc
1377                );
1378
1379                if ts_rewrite.is_none() {
1380                    // The ts rewrite feature allows us to advance the registered
1381                    // upper of a batch that's already been staged (the inline
1382                    // upper), so if it's been used, then there's no useful
1383                    // invariant that we can assert here.
1384                    soft_assert_or_log!(
1385                        PartialOrder::less_equal(registered_desc.upper(), inline_desc.upper()),
1386                        "key={} inline={:?} registered={:?}",
1387                        printable_name,
1388                        inline_desc,
1389                        registered_desc
1390                    );
1391                }
1392            }
1393            // As mentioned above, batches that needs truncation will always have a
1394            // since of the minimum timestamp. Technically we could truncate any
1395            // batch where the since is less_than the output_desc's lower, but we're
1396            // strict here so we don't get any surprises.
1397            assert_eq!(
1398                inline_desc.since(),
1399                &Antichain::from_elem(T::minimum()),
1400                "key={} inline={:?} registered={:?}",
1401                printable_name,
1402                inline_desc,
1403                registered_desc
1404            );
1405        } else {
1406            soft_assert_or_log!(
1407                PartialOrder::less_equal(inline_desc.since(), registered_desc.since()),
1408                "key={} inline={:?} registered={:?}",
1409                printable_name,
1410                inline_desc,
1411                registered_desc
1412            );
1413            assert_eq!(
1414                inline_desc.lower(),
1415                registered_desc.lower(),
1416                "key={} inline={:?} registered={:?}",
1417                printable_name,
1418                inline_desc,
1419                registered_desc
1420            );
1421            assert_eq!(
1422                inline_desc.upper(),
1423                registered_desc.upper(),
1424                "key={} inline={:?} registered={:?}",
1425                printable_name,
1426                inline_desc,
1427                registered_desc
1428            );
1429        }
1430
1431        EncodedPart {
1432            metrics,
1433            registered_desc,
1434            part: parsed,
1435            needs_truncation,
1436            ts_rewrite: ts_rewrite.cloned(),
1437        }
1438    }
1439
1440    pub(crate) fn maybe_unconsolidated(&self) -> bool {
1441        // At time of writing, only user parts may be unconsolidated, and they are always
1442        // written with a since of [T::minimum()].
1443        self.part.desc.since().borrow() == AntichainRef::new(&[T::minimum()])
1444    }
1445
1446    pub(crate) fn updates(&self) -> &BlobTraceUpdates {
1447        &self.part.updates
1448    }
1449
1450    /// Returns the updates with all truncation / timestamp rewriting applied.
1451    pub(crate) fn normalize(&self, metrics: &ColumnarMetrics) -> BlobTraceUpdates {
1452        let updates = self.part.updates.clone();
1453        if !self.needs_truncation && self.ts_rewrite.is_none() {
1454            return updates;
1455        }
1456
1457        let mut codec = updates
1458            .records()
1459            .map(|r| (r.keys().clone(), r.vals().clone()));
1460        let mut structured = updates.structured().cloned();
1461        let mut timestamps = updates.timestamps().clone();
1462        let mut diffs = updates.diffs().clone();
1463
1464        if let Some(rewrite) = self.ts_rewrite.as_ref() {
1465            timestamps = arrow::compute::unary(&timestamps, |i: i64| {
1466                let mut t = T::decode(i.to_le_bytes());
1467                t.advance_by(rewrite.borrow());
1468                i64::from_le_bytes(T::encode(&t))
1469            });
1470        }
1471
1472        let reallocated = if self.needs_truncation {
1473            let filter = BooleanArray::from_unary(&timestamps, |i| {
1474                let t = T::decode(i.to_le_bytes());
1475                let truncate_t = {
1476                    !self.registered_desc.lower().less_equal(&t)
1477                        || self.registered_desc.upper().less_equal(&t)
1478                };
1479                !truncate_t
1480            });
1481            if filter.false_count() == 0 {
1482                // If we're not filtering anything in practice, skip filtering and reallocating.
1483                false
1484            } else {
1485                let filter = FilterBuilder::new(&filter).optimize().build();
1486                let do_filter = |array: &dyn Array| filter.filter(array).expect("valid filter len");
1487                if let Some((keys, vals)) = codec {
1488                    codec = Some((
1489                        realloc_array(do_filter(&keys).as_binary(), metrics),
1490                        realloc_array(do_filter(&vals).as_binary(), metrics),
1491                    ));
1492                }
1493                if let Some(ext) = structured {
1494                    structured = Some(ColumnarRecordsStructuredExt {
1495                        key: realloc_any(do_filter(&*ext.key), metrics),
1496                        val: realloc_any(do_filter(&*ext.val), metrics),
1497                    });
1498                }
1499                timestamps = realloc_array(do_filter(&timestamps).as_primitive(), metrics);
1500                diffs = realloc_array(do_filter(&diffs).as_primitive(), metrics);
1501                true
1502            }
1503        } else {
1504            false
1505        };
1506
1507        if self.ts_rewrite.is_some() && !reallocated {
1508            timestamps = realloc_array(&timestamps, metrics);
1509        }
1510
1511        if self.ts_rewrite.is_some() {
1512            self.metrics
1513                .ts_rewrite
1514                .inc_by(u64::cast_from(timestamps.len()));
1515        }
1516
1517        match (codec, structured) {
1518            (Some((key, value)), None) => {
1519                BlobTraceUpdates::Row(ColumnarRecords::new(key, value, timestamps, diffs))
1520            }
1521            (Some((key, value)), Some(ext)) => {
1522                BlobTraceUpdates::Both(ColumnarRecords::new(key, value, timestamps, diffs), ext)
1523            }
1524            (None, Some(ext)) => BlobTraceUpdates::Structured {
1525                key_values: ext,
1526                timestamps,
1527                diffs,
1528            },
1529            (None, None) => unreachable!(),
1530        }
1531    }
1532}
1533
1534/// This represents the serde encoding for [`LeasedBatchPart`]. We expose the struct
1535/// itself (unlike other encodable structs) to attempt to provide stricter drop
1536/// semantics on `LeasedBatchPart`, i.e. `SerdeLeasedBatchPart` is exchangeable
1537/// (including over the network), where `LeasedBatchPart` is not.
1538///
1539/// For more details see documentation and comments on:
1540/// - [`LeasedBatchPart`]
1541/// - `From<SerdeLeasedBatchPart>` for `LeasedBatchPart<T>`
1542#[derive(Debug, Serialize, Deserialize, Clone)]
1543pub struct ExchangeableBatchPart<T> {
1544    shard_id: ShardId,
1545    // Duplicated with the one serialized in the proto for use in backpressure.
1546    encoded_size_bytes: usize,
1547    desc: Description<T>,
1548    filter: FetchBatchFilter<T>,
1549    part: LazyProto<ProtoHollowBatchPart>,
1550    /// The id of the reader that leased this part. See the corresponding field
1551    /// on [LeasedBatchPart].
1552    reader_id: LeasedReaderId,
1553    filter_pushdown_audit: bool,
1554}
1555
1556impl<T> ExchangeableBatchPart<T> {
1557    /// Returns the encoded size of the given part.
1558    pub fn encoded_size_bytes(&self) -> usize {
1559        self.encoded_size_bytes
1560    }
1561
1562    /// Returns the id of the reader that leased this part.
1563    pub fn reader_id(&self) -> &LeasedReaderId {
1564        &self.reader_id
1565    }
1566}
1567
1568/// Format we'll use when decoding a [`Part`].
1569///
1570/// [`Part`]: mz_persist_types::part::Part
1571#[derive(Debug, Copy, Clone)]
1572pub enum PartDecodeFormat {
1573    /// Decode from opaque `Codec` data.
1574    Row {
1575        /// Will also decode the structured data, and validate it matches.
1576        validate_structured: bool,
1577    },
1578    /// Decode from arrow data
1579    Arrow,
1580}
1581
1582impl PartDecodeFormat {
1583    /// Returns a default value for [`PartDecodeFormat`].
1584    pub const fn default() -> Self {
1585        PartDecodeFormat::Arrow
1586    }
1587
1588    /// Parses a [`PartDecodeFormat`] from the provided string, falling back to the default if the
1589    /// provided value is unrecognized.
1590    pub fn from_str(s: &str) -> Self {
1591        match s {
1592            "row" => PartDecodeFormat::Row {
1593                validate_structured: false,
1594            },
1595            "row_with_validate" => PartDecodeFormat::Row {
1596                validate_structured: true,
1597            },
1598            "arrow" => PartDecodeFormat::Arrow,
1599            x => {
1600                let default = PartDecodeFormat::default();
1601                soft_panic_or_log!("Invalid part decode format: '{x}', falling back to {default}");
1602                default
1603            }
1604        }
1605    }
1606
1607    /// Returns a string representation of [`PartDecodeFormat`].
1608    pub const fn as_str(&self) -> &'static str {
1609        match self {
1610            PartDecodeFormat::Row {
1611                validate_structured: false,
1612            } => "row",
1613            PartDecodeFormat::Row {
1614                validate_structured: true,
1615            } => "row_with_validate",
1616            PartDecodeFormat::Arrow => "arrow",
1617        }
1618    }
1619}
1620
1621impl fmt::Display for PartDecodeFormat {
1622    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1623        f.write_str(self.as_str())
1624    }
1625}
1626
1627#[mz_ore::test]
1628fn client_exchange_data() {
1629    // The whole point of SerdeLeasedBatchPart is that it can be exchanged
1630    // between timely workers, including over the network. Enforce then that it
1631    // implements ExchangeData.
1632    fn is_exchange_data<T: timely::ExchangeData>() {}
1633    is_exchange_data::<ExchangeableBatchPart<u64>>();
1634    is_exchange_data::<ExchangeableBatchPart<u64>>();
1635}