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    pub fn stats(&self) -> Option<PartStats> {
660        self.part.stats().map(|x| x.decode())
661    }
662
663    /// Apply any relevant projection pushdown optimizations, assuming that the data in the part
664    /// is equivalent to the provided key and value.
665    pub fn maybe_optimize(&mut self, cfg: &ConfigSet, key: ArrayRef, val: ArrayRef) {
666        assert_eq!(key.len(), 1, "expect a single-row key array");
667        assert_eq!(val.len(), 1, "expect a single-row val array");
668        let as_of = match &self.filter {
669            FetchBatchFilter::Snapshot { as_of } => as_of,
670            FetchBatchFilter::Listen { .. } | FetchBatchFilter::Compaction { .. } => return,
671        };
672        if !OPTIMIZE_IGNORED_DATA_FETCH.get(cfg) {
673            return;
674        }
675        let (diffs_sum, _stats) = match &self.part {
676            BatchPart::Hollow(x) => (x.diffs_sum, x.stats.as_ref()),
677            BatchPart::Inline { .. } => return,
678        };
679        debug!(
680            "try_optimize_ignored_data_fetch diffs_sum={:?} as_of={:?} lower={:?} upper={:?}",
681            // This is only used for debugging, so hack to assume that D is i64.
682            diffs_sum.map(i64::decode),
683            as_of.elements(),
684            self.desc.lower().elements(),
685            self.desc.upper().elements()
686        );
687        let as_of = match &as_of.elements() {
688            &[as_of] => as_of,
689            _ => return,
690        };
691        let eligible = self.desc.upper().less_equal(as_of) && self.desc.since().less_equal(as_of);
692        if !eligible {
693            return;
694        }
695        let Some(diffs_sum) = diffs_sum else {
696            return;
697        };
698
699        debug!(
700            "try_optimize_ignored_data_fetch faked {:?} diffs at ts {:?} skipping fetch of {} bytes",
701            // This is only used for debugging, so hack to assume that D is i64.
702            i64::decode(diffs_sum),
703            as_of,
704            self.part.encoded_size_bytes(),
705        );
706        self.metrics.pushdown.parts_faked_count.inc();
707        self.metrics
708            .pushdown
709            .parts_faked_bytes
710            .inc_by(u64::cast_from(self.part.encoded_size_bytes()));
711        let timestamps = {
712            let mut col = Codec64Mut::with_capacity(1);
713            col.push(as_of);
714            col.finish()
715        };
716        let diffs = {
717            let mut col = Codec64Mut::with_capacity(1);
718            col.push_raw(diffs_sum);
719            col.finish()
720        };
721        let updates = BlobTraceUpdates::Structured {
722            key_values: ColumnarRecordsStructuredExt { key, val },
723            timestamps,
724            diffs,
725        };
726        let faked_data = LazyInlineBatchPart::from(&ProtoInlineBatchPart {
727            desc: Some(self.desc.into_proto()),
728            index: 0,
729            updates: Some(updates.into_proto()),
730        });
731        self.part = BatchPart::Inline {
732            updates: faked_data,
733            ts_rewrite: None,
734            schema_id: None,
735            deprecated_schema_id: None,
736        };
737    }
738}
739
740impl<T> Drop for LeasedBatchPart<T> {
741    /// For details, see [`LeasedBatchPart`].
742    fn drop(&mut self) {
743        self.metrics.lease.dropped_part.inc()
744    }
745}
746
747/// A [Blob] object that has been fetched, but not at all decoded.
748///
749/// In contrast to [FetchedPart], this representation hasn't yet done parquet
750/// decoding.
751#[derive(Debug)]
752pub struct FetchedBlob<K: Codec, V: Codec, T, D> {
753    metrics: Arc<Metrics>,
754    read_metrics: ReadMetrics,
755    buf: FetchedBlobBuf<T>,
756    registered_desc: Description<T>,
757    migration: PartMigration<K, V>,
758    filter: FetchBatchFilter<T>,
759    filter_pushdown_audit: bool,
760    structured_part_audit: PartDecodeFormat,
761    fetch_permit: Option<Arc<MetricsPermits>>,
762    fetch_config: FetchConfig,
763    _phantom: PhantomData<fn() -> D>,
764}
765
766#[derive(Debug, Clone)]
767enum FetchedBlobBuf<T> {
768    Hollow {
769        buf: SegmentedBytes,
770        part: HollowBatchPart<T>,
771    },
772    Inline {
773        desc: Description<T>,
774        updates: LazyInlineBatchPart,
775        ts_rewrite: Option<Antichain<T>>,
776    },
777}
778
779impl<K: Codec, V: Codec, T: Clone, D> Clone for FetchedBlob<K, V, T, D> {
780    fn clone(&self) -> Self {
781        Self {
782            metrics: Arc::clone(&self.metrics),
783            read_metrics: self.read_metrics.clone(),
784            buf: self.buf.clone(),
785            registered_desc: self.registered_desc.clone(),
786            migration: self.migration.clone(),
787            filter: self.filter.clone(),
788            filter_pushdown_audit: self.filter_pushdown_audit.clone(),
789            fetch_permit: self.fetch_permit.clone(),
790            structured_part_audit: self.structured_part_audit.clone(),
791            fetch_config: self.fetch_config.clone(),
792            _phantom: self._phantom.clone(),
793        }
794    }
795}
796
797/// [FetchedPart] but with an accompanying permit from the fetch mem/disk
798/// semaphore.
799pub struct ShardSourcePart<K: Codec, V: Codec, T, D> {
800    /// The underlying [FetchedPart].
801    pub part: FetchedPart<K, V, T, D>,
802    fetch_permit: Option<Arc<MetricsPermits>>,
803}
804
805impl<K, V, T: Debug, D: Debug> Debug for ShardSourcePart<K, V, T, D>
806where
807    K: Codec + Debug,
808    <K as Codec>::Storage: Debug,
809    V: Codec + Debug,
810    <V as Codec>::Storage: Debug,
811{
812    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
813        let ShardSourcePart { part, fetch_permit } = self;
814        f.debug_struct("ShardSourcePart")
815            .field("part", part)
816            .field("fetch_permit", fetch_permit)
817            .finish()
818    }
819}
820
821impl<K: Codec, V: Codec, T: Timestamp + Lattice + Codec64, D> FetchedBlob<K, V, T, D> {
822    /// Partially decodes this blob into a [FetchedPart].
823    pub fn parse(&self) -> ShardSourcePart<K, V, T, D> {
824        self.parse_internal(&self.fetch_config)
825    }
826
827    /// Partially decodes this blob into a [FetchedPart].
828    pub(crate) fn parse_internal(&self, cfg: &FetchConfig) -> ShardSourcePart<K, V, T, D> {
829        let (part, stats) = match &self.buf {
830            FetchedBlobBuf::Hollow { buf, part } => {
831                let parsed = decode_batch_part_blob(
832                    cfg,
833                    &self.metrics,
834                    &self.read_metrics,
835                    self.registered_desc.clone(),
836                    part,
837                    buf,
838                );
839                (parsed, part.stats.as_ref())
840            }
841            FetchedBlobBuf::Inline {
842                desc,
843                updates,
844                ts_rewrite,
845            } => {
846                let parsed = EncodedPart::from_inline(
847                    cfg,
848                    &self.metrics,
849                    self.read_metrics.clone(),
850                    desc.clone(),
851                    updates,
852                    ts_rewrite.as_ref(),
853                );
854                (parsed, None)
855            }
856        };
857        let part = FetchedPart::new(
858            Arc::clone(&self.metrics),
859            part,
860            self.migration.clone(),
861            self.filter.clone(),
862            self.filter_pushdown_audit,
863            self.structured_part_audit,
864            stats,
865        );
866        ShardSourcePart {
867            part,
868            fetch_permit: self.fetch_permit.clone(),
869        }
870    }
871
872    /// Decodes and returns the pushdown stats for this part, if known.
873    pub fn stats(&self) -> Option<PartStats> {
874        match &self.buf {
875            FetchedBlobBuf::Hollow { part, .. } => part.stats.as_ref().map(|x| x.decode()),
876            FetchedBlobBuf::Inline { .. } => None,
877        }
878    }
879}
880
881/// A [Blob] object that has been fetched, but not yet fully decoded.
882///
883/// In contrast to [FetchedBlob], this representation has already done parquet
884/// decoding.
885#[derive(Debug)]
886pub struct FetchedPart<K: Codec, V: Codec, T, D> {
887    metrics: Arc<Metrics>,
888    ts_filter: FetchBatchFilter<T>,
889    // If migration is Either, then the columnar one will have already been
890    // applied here on the structured data only.
891    part: EitherOrBoth<
892        ColumnarRecords,
893        (
894            <K::Schema as Schema<K>>::Decoder,
895            <V::Schema as Schema<V>>::Decoder,
896        ),
897    >,
898    timestamps: Int64Array,
899    diffs: Int64Array,
900    migration: PartMigration<K, V>,
901    filter_pushdown_audit: Option<LazyPartStats>,
902    peek_stash: Option<((K, V), T, D)>,
903    part_cursor: usize,
904    key_storage: Option<K::Storage>,
905    val_storage: Option<V::Storage>,
906
907    _phantom: PhantomData<fn() -> D>,
908}
909
910impl<K: Codec, V: Codec, T: Timestamp + Lattice + Codec64, D> FetchedPart<K, V, T, D> {
911    pub(crate) fn new(
912        metrics: Arc<Metrics>,
913        part: EncodedPart<T>,
914        migration: PartMigration<K, V>,
915        ts_filter: FetchBatchFilter<T>,
916        filter_pushdown_audit: bool,
917        part_decode_format: PartDecodeFormat,
918        stats: Option<&LazyPartStats>,
919    ) -> Self {
920        let part_len = u64::cast_from(part.part.updates.len());
921        match &migration {
922            PartMigration::SameSchema { .. } => metrics.schema.migration_count_same.inc(),
923            PartMigration::Schemaless { .. } => {
924                metrics.schema.migration_count_codec.inc();
925                metrics.schema.migration_len_legacy_codec.inc_by(part_len);
926            }
927            PartMigration::Either { .. } => {
928                metrics.schema.migration_count_either.inc();
929                match part_decode_format {
930                    PartDecodeFormat::Row {
931                        validate_structured: false,
932                    } => metrics.schema.migration_len_either_codec.inc_by(part_len),
933                    PartDecodeFormat::Row {
934                        validate_structured: true,
935                    } => {
936                        metrics.schema.migration_len_either_codec.inc_by(part_len);
937                        metrics.schema.migration_len_either_arrow.inc_by(part_len);
938                    }
939                    PartDecodeFormat::Arrow => {
940                        metrics.schema.migration_len_either_arrow.inc_by(part_len)
941                    }
942                }
943            }
944        }
945
946        let filter_pushdown_audit = if filter_pushdown_audit {
947            stats.cloned()
948        } else {
949            None
950        };
951
952        let downcast_structured = |structured: ColumnarRecordsStructuredExt,
953                                   structured_only: bool| {
954            let key_size_before = ArrayOrd::new(&structured.key).goodbytes();
955
956            let structured = match &migration {
957                PartMigration::SameSchema { .. } => structured,
958                PartMigration::Schemaless { read } if structured_only => {
959                    // We don't know the source schema, but we do know the source datatype; migrate it directly.
960                    let start = Instant::now();
961                    let read_key = data_type::<K>(&*read.key).ok()?;
962                    let read_val = data_type::<V>(&*read.val).ok()?;
963                    let key_migration = backward_compatible(structured.key.data_type(), &read_key)?;
964                    let val_migration = backward_compatible(structured.val.data_type(), &read_val)?;
965                    let key = key_migration.migrate(structured.key);
966                    let val = val_migration.migrate(structured.val);
967                    metrics
968                        .schema
969                        .migration_migrate_seconds
970                        .inc_by(start.elapsed().as_secs_f64());
971                    ColumnarRecordsStructuredExt { key, val }
972                }
973                PartMigration::Schemaless { .. } => return None,
974                PartMigration::Either {
975                    write: _,
976                    read: _,
977                    key_migration,
978                    val_migration,
979                } => {
980                    let start = Instant::now();
981                    let key = key_migration.migrate(structured.key);
982                    let val = val_migration.migrate(structured.val);
983                    metrics
984                        .schema
985                        .migration_migrate_seconds
986                        .inc_by(start.elapsed().as_secs_f64());
987                    ColumnarRecordsStructuredExt { key, val }
988                }
989            };
990
991            let read_schema = migration.codec_read();
992            let key = K::Schema::decoder_any(&*read_schema.key, &*structured.key);
993            let val = V::Schema::decoder_any(&*read_schema.val, &*structured.val);
994
995            match &key {
996                Ok(key_decoder) => {
997                    let key_size_after = key_decoder.goodbytes();
998                    let key_diff = key_size_before.saturating_sub(key_size_after);
999                    metrics
1000                        .pushdown
1001                        .parts_projection_trimmed_bytes
1002                        .inc_by(u64::cast_from(key_diff));
1003                }
1004                Err(e) => {
1005                    soft_panic_or_log!("failed to create decoder: {e:#?}");
1006                }
1007            }
1008
1009            Some((key.ok()?, val.ok()?))
1010        };
1011
1012        let updates = part.normalize(&metrics.columnar);
1013        let timestamps = updates.timestamps().clone();
1014        let diffs = updates.diffs().clone();
1015        let part = match updates {
1016            // If only one encoding is available, decode via that encoding.
1017            BlobTraceUpdates::Row(records) => EitherOrBoth::Left(records),
1018            BlobTraceUpdates::Structured { key_values, .. } => EitherOrBoth::Right(
1019                // The structured-only data format was added after schema ids were recorded everywhere,
1020                // so we expect this data to be present.
1021                downcast_structured(key_values, true).expect("valid schemas for structured data"),
1022            ),
1023            // If both are available, respect the specified part decode format.
1024            BlobTraceUpdates::Both(records, ext) => match part_decode_format {
1025                PartDecodeFormat::Row {
1026                    validate_structured: false,
1027                } => EitherOrBoth::Left(records),
1028                PartDecodeFormat::Row {
1029                    validate_structured: true,
1030                } => match downcast_structured(ext, false) {
1031                    Some(decoders) => EitherOrBoth::Both(records, decoders),
1032                    None => EitherOrBoth::Left(records),
1033                },
1034                PartDecodeFormat::Arrow => match downcast_structured(ext, false) {
1035                    Some(decoders) => EitherOrBoth::Right(decoders),
1036                    None => EitherOrBoth::Left(records),
1037                },
1038            },
1039        };
1040
1041        FetchedPart {
1042            metrics,
1043            ts_filter,
1044            part,
1045            peek_stash: None,
1046            timestamps,
1047            diffs,
1048            migration,
1049            filter_pushdown_audit,
1050            part_cursor: 0,
1051            key_storage: None,
1052            val_storage: None,
1053            _phantom: PhantomData,
1054        }
1055    }
1056
1057    /// Returns Some if this part was only fetched as part of a filter pushdown
1058    /// audit. See [LeasedBatchPart::request_filter_pushdown_audit].
1059    ///
1060    /// If set, the value in the Option is for debugging and should be included
1061    /// in any error messages.
1062    pub fn is_filter_pushdown_audit(&self) -> Option<impl std::fmt::Debug + use<K, V, T, D>> {
1063        self.filter_pushdown_audit.clone()
1064    }
1065}
1066
1067/// A [Blob] object that has been fetched, but has no associated decoding
1068/// logic.
1069#[derive(Debug)]
1070pub(crate) struct EncodedPart<T> {
1071    metrics: ReadMetrics,
1072    registered_desc: Description<T>,
1073    part: BlobTraceBatchPart<T>,
1074    needs_truncation: bool,
1075    ts_rewrite: Option<Antichain<T>>,
1076}
1077
1078impl<K, V, T, D> FetchedPart<K, V, T, D>
1079where
1080    K: Debug + Codec,
1081    V: Debug + Codec,
1082    T: Timestamp + Lattice + Codec64,
1083    D: Monoid + Codec64 + Send + Sync,
1084{
1085    /// [Self::next] but optionally providing a `K` and `V` for alloc reuse.
1086    ///
1087    /// When `result_override` is specified, return it instead of decoding data.
1088    /// This is used when we know the decoded result will be ignored.
1089    pub fn next_with_storage(
1090        &mut self,
1091        key: &mut Option<K>,
1092        val: &mut Option<V>,
1093    ) -> Option<((K, V), T, D)> {
1094        let mut consolidated = self.peek_stash.take();
1095        loop {
1096            // Fetch and decode the next tuple in the sequence. (Or break if there is none.)
1097            let next = if self.part_cursor < self.timestamps.len() {
1098                let next_idx = self.part_cursor;
1099                self.part_cursor += 1;
1100                // These `to_le_bytes` calls were previously encapsulated by `ColumnarRecords`.
1101                // TODO(structured): re-encapsulate these once we've finished the structured migration.
1102                let mut t = T::decode(self.timestamps.values()[next_idx].to_le_bytes());
1103                if !self.ts_filter.filter_ts(&mut t) {
1104                    continue;
1105                }
1106                let d = D::decode(self.diffs.values()[next_idx].to_le_bytes());
1107                if d.is_zero() {
1108                    continue;
1109                }
1110                let kv = self.decode_kv(next_idx, key, val);
1111                (kv, t, d)
1112            } else {
1113                break;
1114            };
1115
1116            // Attempt to consolidate in the next tuple, stashing it if that's not possible.
1117            if let Some((kv, t, d)) = &mut consolidated {
1118                let (kv_next, t_next, d_next) = &next;
1119                if kv == kv_next && t == t_next {
1120                    d.plus_equals(d_next);
1121                    if d.is_zero() {
1122                        consolidated = None;
1123                    }
1124                } else {
1125                    self.peek_stash = Some(next);
1126                    break;
1127                }
1128            } else {
1129                consolidated = Some(next);
1130            }
1131        }
1132
1133        let (kv, t, d) = consolidated?;
1134
1135        Some((kv, t, d))
1136    }
1137
1138    fn decode_kv(&mut self, index: usize, key: &mut Option<K>, val: &mut Option<V>) -> (K, V) {
1139        let decoded = self
1140            .part
1141            .as_ref()
1142            .map_left(|codec| {
1143                let ((ck, cv), _, _) = codec.get(index).expect("valid index");
1144                let (k, v) = Self::decode_codec(
1145                    &*self.metrics,
1146                    self.migration.codec_read(),
1147                    ck,
1148                    cv,
1149                    key,
1150                    val,
1151                    &mut self.key_storage,
1152                    &mut self.val_storage,
1153                );
1154                (k.expect("valid legacy key"), v.expect("valid legacy value"))
1155            })
1156            .map_right(|(structured_key, structured_val)| {
1157                self.decode_structured(index, structured_key, structured_val, key, val)
1158            });
1159
1160        match decoded {
1161            EitherOrBoth::Both((k, v), (k_s, v_s)) => {
1162                // Purposefully do not trace to prevent blowing up Sentry.
1163                let is_valid = self
1164                    .metrics
1165                    .columnar
1166                    .arrow()
1167                    .key()
1168                    .report_valid(|| k_s == k);
1169                if !is_valid {
1170                    soft_panic_no_log!("structured key did not match, {k_s:?} != {k:?}");
1171                }
1172                // Purposefully do not trace to prevent blowing up Sentry.
1173                let is_valid = self
1174                    .metrics
1175                    .columnar
1176                    .arrow()
1177                    .val()
1178                    .report_valid(|| v_s == v);
1179                if !is_valid {
1180                    soft_panic_no_log!("structured val did not match, {v_s:?} != {v:?}");
1181                }
1182
1183                (k, v)
1184            }
1185            EitherOrBoth::Left(kv) => kv,
1186            EitherOrBoth::Right(kv) => kv,
1187        }
1188    }
1189
1190    fn decode_codec(
1191        metrics: &Metrics,
1192        read_schemas: &Schemas<K, V>,
1193        key_buf: &[u8],
1194        val_buf: &[u8],
1195        key: &mut Option<K>,
1196        val: &mut Option<V>,
1197        key_storage: &mut Option<K::Storage>,
1198        val_storage: &mut Option<V::Storage>,
1199    ) -> (Result<K, String>, Result<V, String>) {
1200        let k = metrics.codecs.key.decode(|| match key.take() {
1201            Some(mut key) => {
1202                match K::decode_from(&mut key, key_buf, key_storage, &read_schemas.key) {
1203                    Ok(()) => Ok(key),
1204                    Err(err) => Err(err),
1205                }
1206            }
1207            None => K::decode(key_buf, &read_schemas.key),
1208        });
1209        let v = metrics.codecs.val.decode(|| match val.take() {
1210            Some(mut val) => {
1211                match V::decode_from(&mut val, val_buf, val_storage, &read_schemas.val) {
1212                    Ok(()) => Ok(val),
1213                    Err(err) => Err(err),
1214                }
1215            }
1216            None => V::decode(val_buf, &read_schemas.val),
1217        });
1218        (k, v)
1219    }
1220
1221    fn decode_structured(
1222        &self,
1223        idx: usize,
1224        keys: &<K::Schema as Schema<K>>::Decoder,
1225        vals: &<V::Schema as Schema<V>>::Decoder,
1226        key: &mut Option<K>,
1227        val: &mut Option<V>,
1228    ) -> (K, V) {
1229        let mut key = key.take().unwrap_or_default();
1230        keys.decode(idx, &mut key);
1231
1232        let mut val = val.take().unwrap_or_default();
1233        vals.decode(idx, &mut val);
1234
1235        (key, val)
1236    }
1237}
1238
1239impl<K, V, T, D> Iterator for FetchedPart<K, V, T, D>
1240where
1241    K: Debug + Codec,
1242    V: Debug + Codec,
1243    T: Timestamp + Lattice + Codec64,
1244    D: Monoid + Codec64 + Send + Sync,
1245{
1246    type Item = ((K, V), T, D);
1247
1248    fn next(&mut self) -> Option<Self::Item> {
1249        self.next_with_storage(&mut None, &mut None)
1250    }
1251
1252    fn size_hint(&self) -> (usize, Option<usize>) {
1253        // We don't know in advance how restrictive the filter will be.
1254        let max_len = self.timestamps.len();
1255        (0, Some(max_len))
1256    }
1257}
1258
1259impl<T> EncodedPart<T>
1260where
1261    T: Timestamp + Lattice + Codec64,
1262{
1263    pub async fn fetch(
1264        cfg: &FetchConfig,
1265        shard_id: &ShardId,
1266        blob: &dyn Blob,
1267        metrics: &Metrics,
1268        shard_metrics: &ShardMetrics,
1269        read_metrics: &ReadMetrics,
1270        registered_desc: &Description<T>,
1271        part: &BatchPart<T>,
1272    ) -> Result<Self, BlobKey> {
1273        match part {
1274            BatchPart::Hollow(x) => {
1275                fetch_batch_part(
1276                    cfg,
1277                    shard_id,
1278                    blob,
1279                    metrics,
1280                    shard_metrics,
1281                    read_metrics,
1282                    registered_desc,
1283                    x,
1284                )
1285                .await
1286            }
1287            BatchPart::Inline {
1288                updates,
1289                ts_rewrite,
1290                ..
1291            } => Ok(EncodedPart::from_inline(
1292                cfg,
1293                metrics,
1294                read_metrics.clone(),
1295                registered_desc.clone(),
1296                updates,
1297                ts_rewrite.as_ref(),
1298            )),
1299        }
1300    }
1301
1302    pub(crate) fn from_inline(
1303        cfg: &FetchConfig,
1304        metrics: &Metrics,
1305        read_metrics: ReadMetrics,
1306        desc: Description<T>,
1307        x: &LazyInlineBatchPart,
1308        ts_rewrite: Option<&Antichain<T>>,
1309    ) -> Self {
1310        let parsed = x.decode(&metrics.columnar).expect("valid inline part");
1311        Self::new(cfg, read_metrics, desc, "inline", ts_rewrite, parsed)
1312    }
1313
1314    pub(crate) fn from_hollow(
1315        cfg: &FetchConfig,
1316        metrics: ReadMetrics,
1317        registered_desc: Description<T>,
1318        part: &HollowBatchPart<T>,
1319        parsed: BlobTraceBatchPart<T>,
1320    ) -> Self {
1321        Self::new(
1322            cfg,
1323            metrics,
1324            registered_desc,
1325            &part.key.0,
1326            part.ts_rewrite.as_ref(),
1327            parsed,
1328        )
1329    }
1330
1331    pub(crate) fn new(
1332        cfg: &FetchConfig,
1333        metrics: ReadMetrics,
1334        registered_desc: Description<T>,
1335        printable_name: &str,
1336        ts_rewrite: Option<&Antichain<T>>,
1337        parsed: BlobTraceBatchPart<T>,
1338    ) -> Self {
1339        // There are two types of batches in persist:
1340        // - Batches written by a persist user (either directly or indirectly
1341        //   via BatchBuilder). These always have a since of the minimum
1342        //   timestamp and may be registered in persist state with a tighter set
1343        //   of bounds than are inline in the batch (truncation). To read one of
1344        //   these batches, all data physically in the batch but outside of the
1345        //   truncated bounds must be ignored. Not every user batch is
1346        //   truncated.
1347        // - Batches written by compaction. These always have an inline desc
1348        //   lower and upper that matches the registered desc lower and upper,
1349        //   and a since that is less than or equal to the registered desc.
1350        //   The inline since may be less than the registered desc since,
1351        //   this is because of incremental compaction, where we might rewrite
1352        //   certain runs in a batch but not others.
1353        let inline_desc = &parsed.desc;
1354        let needs_truncation = inline_desc.lower() != registered_desc.lower()
1355            || inline_desc.upper() != registered_desc.upper();
1356        if needs_truncation {
1357            if cfg.validate_bounds_on_read {
1358                soft_assert_or_log!(
1359                    PartialOrder::less_equal(inline_desc.lower(), registered_desc.lower()),
1360                    "key={} inline={:?} registered={:?}",
1361                    printable_name,
1362                    inline_desc,
1363                    registered_desc
1364                );
1365
1366                if ts_rewrite.is_none() {
1367                    // The ts rewrite feature allows us to advance the registered
1368                    // upper of a batch that's already been staged (the inline
1369                    // upper), so if it's been used, then there's no useful
1370                    // invariant that we can assert here.
1371                    soft_assert_or_log!(
1372                        PartialOrder::less_equal(registered_desc.upper(), inline_desc.upper()),
1373                        "key={} inline={:?} registered={:?}",
1374                        printable_name,
1375                        inline_desc,
1376                        registered_desc
1377                    );
1378                }
1379            }
1380            // As mentioned above, batches that needs truncation will always have a
1381            // since of the minimum timestamp. Technically we could truncate any
1382            // batch where the since is less_than the output_desc's lower, but we're
1383            // strict here so we don't get any surprises.
1384            assert_eq!(
1385                inline_desc.since(),
1386                &Antichain::from_elem(T::minimum()),
1387                "key={} inline={:?} registered={:?}",
1388                printable_name,
1389                inline_desc,
1390                registered_desc
1391            );
1392        } else {
1393            soft_assert_or_log!(
1394                PartialOrder::less_equal(inline_desc.since(), registered_desc.since()),
1395                "key={} inline={:?} registered={:?}",
1396                printable_name,
1397                inline_desc,
1398                registered_desc
1399            );
1400            assert_eq!(
1401                inline_desc.lower(),
1402                registered_desc.lower(),
1403                "key={} inline={:?} registered={:?}",
1404                printable_name,
1405                inline_desc,
1406                registered_desc
1407            );
1408            assert_eq!(
1409                inline_desc.upper(),
1410                registered_desc.upper(),
1411                "key={} inline={:?} registered={:?}",
1412                printable_name,
1413                inline_desc,
1414                registered_desc
1415            );
1416        }
1417
1418        EncodedPart {
1419            metrics,
1420            registered_desc,
1421            part: parsed,
1422            needs_truncation,
1423            ts_rewrite: ts_rewrite.cloned(),
1424        }
1425    }
1426
1427    pub(crate) fn maybe_unconsolidated(&self) -> bool {
1428        // At time of writing, only user parts may be unconsolidated, and they are always
1429        // written with a since of [T::minimum()].
1430        self.part.desc.since().borrow() == AntichainRef::new(&[T::minimum()])
1431    }
1432
1433    pub(crate) fn updates(&self) -> &BlobTraceUpdates {
1434        &self.part.updates
1435    }
1436
1437    /// Returns the updates with all truncation / timestamp rewriting applied.
1438    pub(crate) fn normalize(&self, metrics: &ColumnarMetrics) -> BlobTraceUpdates {
1439        let updates = self.part.updates.clone();
1440        if !self.needs_truncation && self.ts_rewrite.is_none() {
1441            return updates;
1442        }
1443
1444        let mut codec = updates
1445            .records()
1446            .map(|r| (r.keys().clone(), r.vals().clone()));
1447        let mut structured = updates.structured().cloned();
1448        let mut timestamps = updates.timestamps().clone();
1449        let mut diffs = updates.diffs().clone();
1450
1451        if let Some(rewrite) = self.ts_rewrite.as_ref() {
1452            timestamps = arrow::compute::unary(&timestamps, |i: i64| {
1453                let mut t = T::decode(i.to_le_bytes());
1454                t.advance_by(rewrite.borrow());
1455                i64::from_le_bytes(T::encode(&t))
1456            });
1457        }
1458
1459        let reallocated = if self.needs_truncation {
1460            let filter = BooleanArray::from_unary(&timestamps, |i| {
1461                let t = T::decode(i.to_le_bytes());
1462                let truncate_t = {
1463                    !self.registered_desc.lower().less_equal(&t)
1464                        || self.registered_desc.upper().less_equal(&t)
1465                };
1466                !truncate_t
1467            });
1468            if filter.false_count() == 0 {
1469                // If we're not filtering anything in practice, skip filtering and reallocating.
1470                false
1471            } else {
1472                let filter = FilterBuilder::new(&filter).optimize().build();
1473                let do_filter = |array: &dyn Array| filter.filter(array).expect("valid filter len");
1474                if let Some((keys, vals)) = codec {
1475                    codec = Some((
1476                        realloc_array(do_filter(&keys).as_binary(), metrics),
1477                        realloc_array(do_filter(&vals).as_binary(), metrics),
1478                    ));
1479                }
1480                if let Some(ext) = structured {
1481                    structured = Some(ColumnarRecordsStructuredExt {
1482                        key: realloc_any(do_filter(&*ext.key), metrics),
1483                        val: realloc_any(do_filter(&*ext.val), metrics),
1484                    });
1485                }
1486                timestamps = realloc_array(do_filter(&timestamps).as_primitive(), metrics);
1487                diffs = realloc_array(do_filter(&diffs).as_primitive(), metrics);
1488                true
1489            }
1490        } else {
1491            false
1492        };
1493
1494        if self.ts_rewrite.is_some() && !reallocated {
1495            timestamps = realloc_array(&timestamps, metrics);
1496        }
1497
1498        if self.ts_rewrite.is_some() {
1499            self.metrics
1500                .ts_rewrite
1501                .inc_by(u64::cast_from(timestamps.len()));
1502        }
1503
1504        match (codec, structured) {
1505            (Some((key, value)), None) => {
1506                BlobTraceUpdates::Row(ColumnarRecords::new(key, value, timestamps, diffs))
1507            }
1508            (Some((key, value)), Some(ext)) => {
1509                BlobTraceUpdates::Both(ColumnarRecords::new(key, value, timestamps, diffs), ext)
1510            }
1511            (None, Some(ext)) => BlobTraceUpdates::Structured {
1512                key_values: ext,
1513                timestamps,
1514                diffs,
1515            },
1516            (None, None) => unreachable!(),
1517        }
1518    }
1519}
1520
1521/// This represents the serde encoding for [`LeasedBatchPart`]. We expose the struct
1522/// itself (unlike other encodable structs) to attempt to provide stricter drop
1523/// semantics on `LeasedBatchPart`, i.e. `SerdeLeasedBatchPart` is exchangeable
1524/// (including over the network), where `LeasedBatchPart` is not.
1525///
1526/// For more details see documentation and comments on:
1527/// - [`LeasedBatchPart`]
1528/// - `From<SerdeLeasedBatchPart>` for `LeasedBatchPart<T>`
1529#[derive(Debug, Serialize, Deserialize, Clone)]
1530pub struct ExchangeableBatchPart<T> {
1531    shard_id: ShardId,
1532    // Duplicated with the one serialized in the proto for use in backpressure.
1533    encoded_size_bytes: usize,
1534    desc: Description<T>,
1535    filter: FetchBatchFilter<T>,
1536    part: LazyProto<ProtoHollowBatchPart>,
1537    /// The id of the reader that leased this part. See the corresponding field
1538    /// on [LeasedBatchPart].
1539    reader_id: LeasedReaderId,
1540    filter_pushdown_audit: bool,
1541}
1542
1543impl<T> ExchangeableBatchPart<T> {
1544    /// Returns the encoded size of the given part.
1545    pub fn encoded_size_bytes(&self) -> usize {
1546        self.encoded_size_bytes
1547    }
1548
1549    /// Returns the id of the reader that leased this part.
1550    pub fn reader_id(&self) -> &LeasedReaderId {
1551        &self.reader_id
1552    }
1553}
1554
1555/// Format we'll use when decoding a [`Part`].
1556///
1557/// [`Part`]: mz_persist_types::part::Part
1558#[derive(Debug, Copy, Clone)]
1559pub enum PartDecodeFormat {
1560    /// Decode from opaque `Codec` data.
1561    Row {
1562        /// Will also decode the structured data, and validate it matches.
1563        validate_structured: bool,
1564    },
1565    /// Decode from arrow data
1566    Arrow,
1567}
1568
1569impl PartDecodeFormat {
1570    /// Returns a default value for [`PartDecodeFormat`].
1571    pub const fn default() -> Self {
1572        PartDecodeFormat::Arrow
1573    }
1574
1575    /// Parses a [`PartDecodeFormat`] from the provided string, falling back to the default if the
1576    /// provided value is unrecognized.
1577    pub fn from_str(s: &str) -> Self {
1578        match s {
1579            "row" => PartDecodeFormat::Row {
1580                validate_structured: false,
1581            },
1582            "row_with_validate" => PartDecodeFormat::Row {
1583                validate_structured: true,
1584            },
1585            "arrow" => PartDecodeFormat::Arrow,
1586            x => {
1587                let default = PartDecodeFormat::default();
1588                soft_panic_or_log!("Invalid part decode format: '{x}', falling back to {default}");
1589                default
1590            }
1591        }
1592    }
1593
1594    /// Returns a string representation of [`PartDecodeFormat`].
1595    pub const fn as_str(&self) -> &'static str {
1596        match self {
1597            PartDecodeFormat::Row {
1598                validate_structured: false,
1599            } => "row",
1600            PartDecodeFormat::Row {
1601                validate_structured: true,
1602            } => "row_with_validate",
1603            PartDecodeFormat::Arrow => "arrow",
1604        }
1605    }
1606}
1607
1608impl fmt::Display for PartDecodeFormat {
1609    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1610        f.write_str(self.as_str())
1611    }
1612}
1613
1614#[mz_ore::test]
1615fn client_exchange_data() {
1616    // The whole point of SerdeLeasedBatchPart is that it can be exchanged
1617    // between timely workers, including over the network. Enforce then that it
1618    // implements ExchangeData.
1619    fn is_exchange_data<T: timely::ExchangeData>() {}
1620    is_exchange_data::<ExchangeableBatchPart<u64>>();
1621    is_exchange_data::<ExchangeableBatchPart<u64>>();
1622}