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