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