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

mz_persist_client/
fetch.rs
