1use std::fmt::{self, Debug};
13use std::marker::PhantomData;
14use std::sync::Arc;
15use std::time::{Duration, Instant};
16
17use anyhow::anyhow;
18use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, Int64Array};
19use arrow::compute::FilterBuilder;
20use differential_dataflow::difference::Monoid;
21use differential_dataflow::lattice::Lattice;
22use differential_dataflow::trace::Description;
23use itertools::EitherOrBoth;
24use mz_dyncfg::{Config, ConfigSet, ConfigValHandle};
25use mz_ore::bytes::SegmentedBytes;
26use mz_ore::cast::CastFrom;
27use mz_ore::{soft_assert_or_log, soft_panic_no_log, soft_panic_or_log};
28use mz_persist::indexed::columnar::arrow::{realloc_any, realloc_array};
29use mz_persist::indexed::columnar::{ColumnarRecords, ColumnarRecordsStructuredExt};
30use mz_persist::indexed::encoding::{BlobTraceBatchPart, BlobTraceUpdates};
31use mz_persist::location::{Blob, SeqNo};
32use mz_persist::metrics::ColumnarMetrics;
33use mz_persist_types::arrow::ArrayOrd;
34use mz_persist_types::columnar::{ColumnDecoder, Schema, data_type};
35use mz_persist_types::part::Codec64Mut;
36use mz_persist_types::schema::backward_compatible;
37use mz_persist_types::stats::PartStats;
38use mz_persist_types::{Codec, Codec64};
39use mz_proto::RustType;
40use serde::{Deserialize, Serialize};
41use timely::PartialOrder;
42use timely::progress::frontier::AntichainRef;
43use timely::progress::{Antichain, Timestamp};
44use tracing::{Instrument, debug, debug_span, trace_span};
45
46use crate::ShardId;
47use crate::cfg::PersistConfig;
48use crate::error::InvalidUsage;
49use crate::internal::apply::Applier;
50use crate::internal::encoding::{LazyInlineBatchPart, LazyPartStats, LazyProto, Schemas};
51use crate::internal::machine::retry_external;
52use crate::internal::metrics::{Metrics, MetricsPermits, ReadMetrics, ShardMetrics};
53use crate::internal::paths::BlobKey;
54use crate::internal::state::{
55 BatchPart, HollowBatchPart, ProtoHollowBatchPart, ProtoInlineBatchPart,
56};
57use crate::read::LeasedReaderId;
58use crate::schema::{PartMigration, SchemaCache};
59
60pub(crate) const FETCH_SEMAPHORE_COST_ADJUSTMENT: Config<f64> = Config::new(
61 "persist_fetch_semaphore_cost_adjustment",
62 1.2,
66 "\
67 An adjustment multiplied by encoded_size_bytes to approximate an upper \
68 bound on the size in lgalloc, which includes the decoded version.",
69);
70
71pub(crate) const FETCH_SEMAPHORE_PERMIT_ADJUSTMENT: Config<f64> = Config::new(
72 "persist_fetch_semaphore_permit_adjustment",
73 1.0,
74 "\
75 A limit on the number of outstanding persist bytes being fetched and \
76 parsed, expressed as a multiplier of the process's memory limit. This data \
77 all spills to lgalloc, so values > 1.0 are safe. Only applied to cc \
78 replicas.",
79);
80
81pub(crate) const PART_DECODE_FORMAT: Config<&'static str> = Config::new(
82 "persist_part_decode_format",
83 PartDecodeFormat::default().as_str(),
84 "\
85 Format we'll use to decode a Persist Part, either 'row', \
86 'row_with_validate', or 'arrow' (Materialize).",
87);
88
89pub(crate) const OPTIMIZE_IGNORED_DATA_FETCH: Config<bool> = Config::new(
90 "persist_optimize_ignored_data_fetch",
91 true,
92 "CYA to allow opt-out of a performance optimization to skip fetching ignored data",
93);
94
95pub(crate) const VALIDATE_PART_BOUNDS_ON_READ: Config<bool> = Config::new(
96 "persist_validate_part_bounds_on_read",
97 false,
98 "Validate the part lower <= the batch lower and the part upper <= batch upper,\
99 for the batch containing that part",
100);
101
102#[derive(Debug, Clone)]
103pub(crate) struct FetchConfig {
104 pub(crate) validate_bounds_on_read: bool,
105}
106
107impl FetchConfig {
108 pub fn from_persist_config(cfg: &PersistConfig) -> Self {
109 Self {
110 validate_bounds_on_read: VALIDATE_PART_BOUNDS_ON_READ.get(cfg),
111 }
112 }
113}
114
115#[derive(Debug, Clone)]
116pub(crate) struct BatchFetcherConfig {
117 pub(crate) part_decode_format: ConfigValHandle<String>,
118 pub(crate) fetch_config: FetchConfig,
119}
120
121impl BatchFetcherConfig {
122 pub fn new(value: &PersistConfig) -> Self {
123 Self {
124 part_decode_format: PART_DECODE_FORMAT.handle(value),
125 fetch_config: FetchConfig::from_persist_config(value),
126 }
127 }
128
129 pub fn part_decode_format(&self) -> PartDecodeFormat {
130 PartDecodeFormat::from_str(self.part_decode_format.get().as_str())
131 }
132}
133
134#[derive(Debug)]
136pub struct BatchFetcher<K, V, T, D>
137where
138 T: Timestamp + Lattice + Codec64,
139 K: Debug + Codec,
141 V: Debug + Codec,
142 D: Monoid + Codec64 + Send + Sync,
143{
144 pub(crate) cfg: BatchFetcherConfig,
145 pub(crate) blob: Arc<dyn Blob>,
146 pub(crate) metrics: Arc<Metrics>,
147 pub(crate) shard_metrics: Arc<ShardMetrics>,
148 pub(crate) shard_id: ShardId,
149 pub(crate) read_schemas: Schemas<K, V>,
150 pub(crate) schema_cache: SchemaCache<K, V, T, D>,
151 pub(crate) is_transient: bool,
152
153 pub(crate) _phantom: PhantomData<fn() -> (K, V, T, D)>,
156}
157
158impl<K, V, T, D> Clone for BatchFetcher<K, V, T, D>
164where
165 T: Timestamp + Lattice + Codec64,
166 K: Debug + Codec,
167 V: Debug + Codec,
168 D: Monoid + Codec64 + Send + Sync,
169{
170 fn clone(&self) -> Self {
171 Self {
172 cfg: self.cfg.clone(),
173 blob: Arc::clone(&self.blob),
174 metrics: Arc::clone(&self.metrics),
175 shard_metrics: Arc::clone(&self.shard_metrics),
176 shard_id: self.shard_id.clone(),
177 read_schemas: self.read_schemas.clone(),
178 schema_cache: self.schema_cache.clone(),
179 is_transient: self.is_transient,
180 _phantom: PhantomData,
181 }
182 }
183}
184
185impl<K, V, T, D> BatchFetcher<K, V, T, D>
186where
187 K: Debug + Codec,
188 V: Debug + Codec,
189 T: Timestamp + Lattice + Codec64 + Sync,
190 D: Monoid + Codec64 + Send + Sync,
191{
192 pub async fn fetch_leased_part(
197 &mut self,
198 part: ExchangeableBatchPart<T>,
199 ) -> Result<Result<FetchedBlob<K, V, T, D>, BlobKey>, InvalidUsage<T>> {
200 let ExchangeableBatchPart {
201 shard_id,
202 encoded_size_bytes: _,
203 desc,
204 filter,
205 filter_pushdown_audit,
206 part,
207 reader_id: _,
208 } = part;
209 let part: BatchPart<T> = part.decode_to().expect("valid part");
210 if shard_id != self.shard_id {
211 return Err(InvalidUsage::BatchNotFromThisShard {
212 batch_shard: shard_id,
213 handle_shard: self.shard_id.clone(),
214 });
215 }
216
217 let migration =
218 PartMigration::new(&part, self.read_schemas.clone(), &mut self.schema_cache)
219 .await
220 .unwrap_or_else(|read_schemas| {
221 panic!(
222 "could not decode part {:?} with schema: {:?}",
223 part.schema_id(),
224 read_schemas
225 )
226 });
227
228 let (buf, fetch_permit) = match &part {
229 BatchPart::Hollow(x) => {
230 let fetch_permit = self
231 .metrics
232 .semaphore
233 .acquire_fetch_permits(x.encoded_size_bytes)
234 .await;
235 let read_metrics = if self.is_transient {
236 &self.metrics.read.unindexed
237 } else {
238 &self.metrics.read.batch_fetcher
239 };
240 let buf = fetch_batch_part_blob(
241 &shard_id,
242 self.blob.as_ref(),
243 &self.metrics,
244 &self.shard_metrics,
245 read_metrics,
246 x,
247 )
248 .await;
249 let buf = match buf {
250 Ok(buf) => buf,
251 Err(key) => return Ok(Err(key)),
252 };
253 let buf = FetchedBlobBuf::Hollow {
254 buf,
255 part: x.clone(),
256 };
257 (buf, Some(Arc::new(fetch_permit)))
258 }
259 BatchPart::Inline {
260 updates,
261 ts_rewrite,
262 ..
263 } => {
264 let buf = FetchedBlobBuf::Inline {
265 desc: desc.clone(),
266 updates: updates.clone(),
267 ts_rewrite: ts_rewrite.clone(),
268 };
269 (buf, None)
270 }
271 };
272 let fetched_blob = FetchedBlob {
273 metrics: Arc::clone(&self.metrics),
274 read_metrics: self.metrics.read.batch_fetcher.clone(),
275 buf,
276 registered_desc: desc.clone(),
277 migration,
278 filter: filter.clone(),
279 filter_pushdown_audit,
280 structured_part_audit: self.cfg.part_decode_format(),
281 fetch_permit,
282 _phantom: PhantomData,
283 fetch_config: self.cfg.fetch_config.clone(),
284 };
285 Ok(Ok(fetched_blob))
286 }
287
288 pub async fn missing_blob_diagnostics(&self, reader_id: &LeasedReaderId) -> String {
291 missing_blob_diagnostics(self.schema_cache.applier(), reader_id).await
292 }
293}
294
295pub(crate) async fn missing_blob_diagnostics<K, V, T, D>(
306 applier: &Applier<K, V, T, D>,
307 reader_id: &LeasedReaderId,
308) -> String
309where
310 K: Debug + Codec,
311 V: Debug + Codec,
312 T: Timestamp + Lattice + Codec64 + Sync,
313 D: Monoid + Codec64,
314{
315 let refresh = applier.fetch_and_update_state(None);
319 if tokio::time::timeout(Duration::from_secs(30), refresh)
320 .await
321 .is_err()
322 {
323 return format!(
324 "reader {reader_id}: could not refresh state within 30s to diagnose the lease; \
325 partitioned from consensus?"
326 );
327 }
328 match applier.reader_lease(reader_id.clone()) {
329 Some(lease_state) => format!(
330 "reader {reader_id} is still present in state ({lease_state:?}); \
331 a missing blob despite a live lease indicates a GC bug"
332 ),
333 None => format!(
334 "reader {reader_id} has been expired out of state; \
335 the process likely failed to heartbeat it within the lease duration \
336 (machine sleep, CPU/memory starvation, or a partition from consensus?)"
337 ),
338 }
339}
340
341#[derive(Debug, Clone, Serialize, Deserialize)]
342pub(crate) enum FetchBatchFilter<T> {
343 Snapshot {
344 as_of: Antichain<T>,
345 },
346 Listen {
347 as_of: Antichain<T>,
348 lower: Antichain<T>,
349 },
350 Compaction {
351 since: Antichain<T>,
352 },
353}
354
355impl<T: Timestamp + Lattice> FetchBatchFilter<T> {
356 pub(crate) fn filter_ts(&self, t: &mut T) -> bool {
357 match self {
358 FetchBatchFilter::Snapshot { as_of } => {
359 if as_of.less_than(t) {
361 return false;
362 }
363 t.advance_by(as_of.borrow());
364 true
365 }
366 FetchBatchFilter::Listen { as_of, lower } => {
367 if !as_of.less_than(t) {
369 return false;
370 }
371
372 if !lower.less_equal(t) {
380 return false;
381 }
382 true
383 }
384 FetchBatchFilter::Compaction { since } => {
385 t.advance_by(since.borrow());
386 true
387 }
388 }
389 }
390}
391
392pub(crate) async fn fetch_leased_part<K, V, T, D>(
397 cfg: &PersistConfig,
398 part: &LeasedBatchPart<T>,
399 blob: &dyn Blob,
400 metrics: Arc<Metrics>,
401 read_metrics: &ReadMetrics,
402 shard_metrics: &ShardMetrics,
403 reader_id: &LeasedReaderId,
404 read_schemas: Schemas<K, V>,
405 schema_cache: &mut SchemaCache<K, V, T, D>,
406) -> FetchedPart<K, V, T, D>
407where
408 K: Debug + Codec,
409 V: Debug + Codec,
410 T: Timestamp + Lattice + Codec64 + Sync,
411 D: Monoid + Codec64 + Send + Sync,
412{
413 let fetch_config = FetchConfig::from_persist_config(cfg);
414 let encoded_part = match EncodedPart::fetch(
415 &fetch_config,
416 &part.shard_id,
417 blob,
418 &metrics,
419 shard_metrics,
420 read_metrics,
421 &part.desc,
422 &part.part,
423 )
424 .await
425 {
426 Ok(x) => x,
427 Err(blob_key) => {
428 let diagnostics = missing_blob_diagnostics(schema_cache.applier(), reader_id).await;
437 panic!("could not fetch batch part {}: {}", blob_key, diagnostics)
438 }
439 };
440 let part_cfg = BatchFetcherConfig::new(cfg);
441 let migration = PartMigration::new(&part.part, read_schemas, schema_cache)
442 .await
443 .unwrap_or_else(|read_schemas| {
444 panic!(
445 "could not decode part {:?} with schema: {:?}",
446 part.part.schema_id(),
447 read_schemas
448 )
449 });
450 FetchedPart::new(
451 metrics,
452 encoded_part,
453 migration,
454 part.filter.clone(),
455 part.filter_pushdown_audit,
456 part_cfg.part_decode_format(),
457 part.part.stats(),
458 )
459}
460
461pub(crate) async fn fetch_batch_part_blob<T>(
462 shard_id: &ShardId,
463 blob: &dyn Blob,
464 metrics: &Metrics,
465 shard_metrics: &ShardMetrics,
466 read_metrics: &ReadMetrics,
467 part: &HollowBatchPart<T>,
468) -> Result<SegmentedBytes, BlobKey> {
469 let now = Instant::now();
470 let get_span = debug_span!("fetch_batch::get");
471 let blob_key = part.key.complete(shard_id);
472 let value = retry_external(&metrics.retries.external.fetch_batch_get, || async {
473 shard_metrics.blob_gets.inc();
474 blob.get(&blob_key)
478 .await
479 .map_err(|err| err.context(format!("blob {blob_key}")))
480 })
481 .instrument(get_span.clone())
482 .await
483 .ok_or(blob_key)?;
484
485 drop(get_span);
486
487 read_metrics.part_count.inc();
488 read_metrics.part_bytes.inc_by(u64::cast_from(value.len()));
489 read_metrics.seconds.inc_by(now.elapsed().as_secs_f64());
490
491 Ok(value)
492}
493
494pub(crate) fn decode_batch_part_blob<T>(
495 cfg: &FetchConfig,
496 metrics: &Metrics,
497 read_metrics: &ReadMetrics,
498 registered_desc: Description<T>,
499 part: &HollowBatchPart<T>,
500 buf: &SegmentedBytes,
501) -> EncodedPart<T>
502where
503 T: Timestamp + Lattice + Codec64,
504{
505 trace_span!("fetch_batch::decode").in_scope(|| {
506 let parsed = metrics
507 .codecs
508 .batch
509 .decode(|| BlobTraceBatchPart::decode(buf, &metrics.columnar))
510 .map_err(|err| anyhow!("couldn't decode batch at key {}: {}", part.key, err))
511 .expect("internal error: invalid encoded state");
516 read_metrics
517 .part_goodbytes
518 .inc_by(u64::cast_from(parsed.updates.goodbytes()));
519 EncodedPart::from_hollow(cfg, read_metrics.clone(), registered_desc, part, parsed)
520 })
521}
522
523pub(crate) async fn fetch_batch_part<T>(
524 cfg: &FetchConfig,
525 shard_id: &ShardId,
526 blob: &dyn Blob,
527 metrics: &Metrics,
528 shard_metrics: &ShardMetrics,
529 read_metrics: &ReadMetrics,
530 registered_desc: &Description<T>,
531 part: &HollowBatchPart<T>,
532) -> Result<EncodedPart<T>, BlobKey>
533where
534 T: Timestamp + Lattice + Codec64,
535{
536 let buf =
537 fetch_batch_part_blob(shard_id, blob, metrics, shard_metrics, read_metrics, part).await?;
538 let part = decode_batch_part_blob(
539 cfg,
540 metrics,
541 read_metrics,
542 registered_desc.clone(),
543 part,
544 &buf,
545 );
546 Ok(part)
547}
548
549#[derive(Clone, Debug)]
556pub struct Lease(Arc<SeqNo>);
557
558impl Lease {
559 pub fn new(seqno: SeqNo) -> Self {
561 Self(Arc::new(seqno))
562 }
563
564 pub fn seqno(&self) -> SeqNo {
566 *self.0
567 }
568
569 pub fn count(&self) -> usize {
571 Arc::strong_count(&self.0)
572 }
573}
574
575#[derive(Debug)]
601pub struct LeasedBatchPart<T> {
602 pub(crate) metrics: Arc<Metrics>,
603 pub(crate) shard_id: ShardId,
604 pub(crate) filter: FetchBatchFilter<T>,
605 pub(crate) desc: Description<T>,
606 pub(crate) part: BatchPart<T>,
607 pub(crate) lease: Lease,
610 pub(crate) reader_id: LeasedReaderId,
614 pub(crate) filter_pushdown_audit: bool,
615}
616
617impl<T> LeasedBatchPart<T>
618where
619 T: Timestamp + Codec64,
620{
621 pub(crate) fn into_exchangeable_part(self) -> (ExchangeableBatchPart<T>, Lease) {
631 let lease = self.lease.clone();
633 let part = ExchangeableBatchPart {
634 shard_id: self.shard_id,
635 encoded_size_bytes: self.part.encoded_size_bytes(),
636 desc: self.desc.clone(),
637 filter: self.filter.clone(),
638 part: LazyProto::from(&self.part.into_proto()),
639 reader_id: self.reader_id.clone(),
640 filter_pushdown_audit: self.filter_pushdown_audit,
641 };
642 (part, lease)
643 }
644
645 pub fn encoded_size_bytes(&self) -> usize {
647 self.part.encoded_size_bytes()
648 }
649
650 pub fn request_filter_pushdown_audit(&mut self) {
655 self.filter_pushdown_audit = true;
656 }
657
658 pub fn stats(&self) -> Option<PartStats> {
660 self.part.stats().map(|x| x.decode())
661 }
662
663 pub fn maybe_optimize(&mut self, cfg: &ConfigSet, key: ArrayRef, val: ArrayRef) {
666 assert_eq!(key.len(), 1, "expect a single-row key array");
667 assert_eq!(val.len(), 1, "expect a single-row val array");
668 let as_of = match &self.filter {
669 FetchBatchFilter::Snapshot { as_of } => as_of,
670 FetchBatchFilter::Listen { .. } | FetchBatchFilter::Compaction { .. } => return,
671 };
672 if !OPTIMIZE_IGNORED_DATA_FETCH.get(cfg) {
673 return;
674 }
675 let (diffs_sum, _stats) = match &self.part {
676 BatchPart::Hollow(x) => (x.diffs_sum, x.stats.as_ref()),
677 BatchPart::Inline { .. } => return,
678 };
679 debug!(
680 "try_optimize_ignored_data_fetch diffs_sum={:?} as_of={:?} lower={:?} upper={:?}",
681 diffs_sum.map(i64::decode),
683 as_of.elements(),
684 self.desc.lower().elements(),
685 self.desc.upper().elements()
686 );
687 let as_of = match &as_of.elements() {
688 &[as_of] => as_of,
689 _ => return,
690 };
691 let eligible = self.desc.upper().less_equal(as_of) && self.desc.since().less_equal(as_of);
692 if !eligible {
693 return;
694 }
695 let Some(diffs_sum) = diffs_sum else {
696 return;
697 };
698
699 debug!(
700 "try_optimize_ignored_data_fetch faked {:?} diffs at ts {:?} skipping fetch of {} bytes",
701 i64::decode(diffs_sum),
703 as_of,
704 self.part.encoded_size_bytes(),
705 );
706 self.metrics.pushdown.parts_faked_count.inc();
707 self.metrics
708 .pushdown
709 .parts_faked_bytes
710 .inc_by(u64::cast_from(self.part.encoded_size_bytes()));
711 let timestamps = {
712 let mut col = Codec64Mut::with_capacity(1);
713 col.push(as_of);
714 col.finish()
715 };
716 let diffs = {
717 let mut col = Codec64Mut::with_capacity(1);
718 col.push_raw(diffs_sum);
719 col.finish()
720 };
721 let updates = BlobTraceUpdates::Structured {
722 key_values: ColumnarRecordsStructuredExt { key, val },
723 timestamps,
724 diffs,
725 };
726 let faked_data = LazyInlineBatchPart::from(&ProtoInlineBatchPart {
727 desc: Some(self.desc.into_proto()),
728 index: 0,
729 updates: Some(updates.into_proto()),
730 });
731 self.part = BatchPart::Inline {
732 updates: faked_data,
733 ts_rewrite: None,
734 schema_id: None,
735 deprecated_schema_id: None,
736 };
737 }
738}
739
740impl<T> Drop for LeasedBatchPart<T> {
741 fn drop(&mut self) {
743 self.metrics.lease.dropped_part.inc()
744 }
745}
746
747#[derive(Debug)]
752pub struct FetchedBlob<K: Codec, V: Codec, T, D> {
753 metrics: Arc<Metrics>,
754 read_metrics: ReadMetrics,
755 buf: FetchedBlobBuf<T>,
756 registered_desc: Description<T>,
757 migration: PartMigration<K, V>,
758 filter: FetchBatchFilter<T>,
759 filter_pushdown_audit: bool,
760 structured_part_audit: PartDecodeFormat,
761 fetch_permit: Option<Arc<MetricsPermits>>,
762 fetch_config: FetchConfig,
763 _phantom: PhantomData<fn() -> D>,
764}
765
766#[derive(Debug, Clone)]
767enum FetchedBlobBuf<T> {
768 Hollow {
769 buf: SegmentedBytes,
770 part: HollowBatchPart<T>,
771 },
772 Inline {
773 desc: Description<T>,
774 updates: LazyInlineBatchPart,
775 ts_rewrite: Option<Antichain<T>>,
776 },
777}
778
779impl<K: Codec, V: Codec, T: Clone, D> Clone for FetchedBlob<K, V, T, D> {
780 fn clone(&self) -> Self {
781 Self {
782 metrics: Arc::clone(&self.metrics),
783 read_metrics: self.read_metrics.clone(),
784 buf: self.buf.clone(),
785 registered_desc: self.registered_desc.clone(),
786 migration: self.migration.clone(),
787 filter: self.filter.clone(),
788 filter_pushdown_audit: self.filter_pushdown_audit.clone(),
789 fetch_permit: self.fetch_permit.clone(),
790 structured_part_audit: self.structured_part_audit.clone(),
791 fetch_config: self.fetch_config.clone(),
792 _phantom: self._phantom.clone(),
793 }
794 }
795}
796
797pub struct ShardSourcePart<K: Codec, V: Codec, T, D> {
800 pub part: FetchedPart<K, V, T, D>,
802 fetch_permit: Option<Arc<MetricsPermits>>,
803}
804
805impl<K, V, T: Debug, D: Debug> Debug for ShardSourcePart<K, V, T, D>
806where
807 K: Codec + Debug,
808 <K as Codec>::Storage: Debug,
809 V: Codec + Debug,
810 <V as Codec>::Storage: Debug,
811{
812 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
813 let ShardSourcePart { part, fetch_permit } = self;
814 f.debug_struct("ShardSourcePart")
815 .field("part", part)
816 .field("fetch_permit", fetch_permit)
817 .finish()
818 }
819}
820
821impl<K: Codec, V: Codec, T: Timestamp + Lattice + Codec64, D> FetchedBlob<K, V, T, D> {
822 pub fn parse(&self) -> ShardSourcePart<K, V, T, D> {
824 self.parse_internal(&self.fetch_config)
825 }
826
827 pub(crate) fn parse_internal(&self, cfg: &FetchConfig) -> ShardSourcePart<K, V, T, D> {
829 let (part, stats) = match &self.buf {
830 FetchedBlobBuf::Hollow { buf, part } => {
831 let parsed = decode_batch_part_blob(
832 cfg,
833 &self.metrics,
834 &self.read_metrics,
835 self.registered_desc.clone(),
836 part,
837 buf,
838 );
839 (parsed, part.stats.as_ref())
840 }
841 FetchedBlobBuf::Inline {
842 desc,
843 updates,
844 ts_rewrite,
845 } => {
846 let parsed = EncodedPart::from_inline(
847 cfg,
848 &self.metrics,
849 self.read_metrics.clone(),
850 desc.clone(),
851 updates,
852 ts_rewrite.as_ref(),
853 );
854 (parsed, None)
855 }
856 };
857 let part = FetchedPart::new(
858 Arc::clone(&self.metrics),
859 part,
860 self.migration.clone(),
861 self.filter.clone(),
862 self.filter_pushdown_audit,
863 self.structured_part_audit,
864 stats,
865 );
866 ShardSourcePart {
867 part,
868 fetch_permit: self.fetch_permit.clone(),
869 }
870 }
871
872 pub fn stats(&self) -> Option<PartStats> {
874 match &self.buf {
875 FetchedBlobBuf::Hollow { part, .. } => part.stats.as_ref().map(|x| x.decode()),
876 FetchedBlobBuf::Inline { .. } => None,
877 }
878 }
879}
880
881#[derive(Debug)]
886pub struct FetchedPart<K: Codec, V: Codec, T, D> {
887 metrics: Arc<Metrics>,
888 ts_filter: FetchBatchFilter<T>,
889 part: EitherOrBoth<
892 ColumnarRecords,
893 (
894 <K::Schema as Schema<K>>::Decoder,
895 <V::Schema as Schema<V>>::Decoder,
896 ),
897 >,
898 timestamps: Int64Array,
899 diffs: Int64Array,
900 migration: PartMigration<K, V>,
901 filter_pushdown_audit: Option<LazyPartStats>,
902 peek_stash: Option<((K, V), T, D)>,
903 part_cursor: usize,
904 key_storage: Option<K::Storage>,
905 val_storage: Option<V::Storage>,
906
907 _phantom: PhantomData<fn() -> D>,
908}
909
910impl<K: Codec, V: Codec, T: Timestamp + Lattice + Codec64, D> FetchedPart<K, V, T, D> {
911 pub(crate) fn new(
912 metrics: Arc<Metrics>,
913 part: EncodedPart<T>,
914 migration: PartMigration<K, V>,
915 ts_filter: FetchBatchFilter<T>,
916 filter_pushdown_audit: bool,
917 part_decode_format: PartDecodeFormat,
918 stats: Option<&LazyPartStats>,
919 ) -> Self {
920 let part_len = u64::cast_from(part.part.updates.len());
921 match &migration {
922 PartMigration::SameSchema { .. } => metrics.schema.migration_count_same.inc(),
923 PartMigration::Schemaless { .. } => {
924 metrics.schema.migration_count_codec.inc();
925 metrics.schema.migration_len_legacy_codec.inc_by(part_len);
926 }
927 PartMigration::Either { .. } => {
928 metrics.schema.migration_count_either.inc();
929 match part_decode_format {
930 PartDecodeFormat::Row {
931 validate_structured: false,
932 } => metrics.schema.migration_len_either_codec.inc_by(part_len),
933 PartDecodeFormat::Row {
934 validate_structured: true,
935 } => {
936 metrics.schema.migration_len_either_codec.inc_by(part_len);
937 metrics.schema.migration_len_either_arrow.inc_by(part_len);
938 }
939 PartDecodeFormat::Arrow => {
940 metrics.schema.migration_len_either_arrow.inc_by(part_len)
941 }
942 }
943 }
944 }
945
946 let filter_pushdown_audit = if filter_pushdown_audit {
947 stats.cloned()
948 } else {
949 None
950 };
951
952 let downcast_structured = |structured: ColumnarRecordsStructuredExt,
953 structured_only: bool| {
954 let key_size_before = ArrayOrd::new(&structured.key).goodbytes();
955
956 let structured = match &migration {
957 PartMigration::SameSchema { .. } => structured,
958 PartMigration::Schemaless { read } if structured_only => {
959 let start = Instant::now();
961 let read_key = data_type::<K>(&*read.key).ok()?;
962 let read_val = data_type::<V>(&*read.val).ok()?;
963 let key_migration = backward_compatible(structured.key.data_type(), &read_key)?;
964 let val_migration = backward_compatible(structured.val.data_type(), &read_val)?;
965 let key = key_migration.migrate(structured.key);
966 let val = val_migration.migrate(structured.val);
967 metrics
968 .schema
969 .migration_migrate_seconds
970 .inc_by(start.elapsed().as_secs_f64());
971 ColumnarRecordsStructuredExt { key, val }
972 }
973 PartMigration::Schemaless { .. } => return None,
974 PartMigration::Either {
975 write: _,
976 read: _,
977 key_migration,
978 val_migration,
979 } => {
980 let start = Instant::now();
981 let key = key_migration.migrate(structured.key);
982 let val = val_migration.migrate(structured.val);
983 metrics
984 .schema
985 .migration_migrate_seconds
986 .inc_by(start.elapsed().as_secs_f64());
987 ColumnarRecordsStructuredExt { key, val }
988 }
989 };
990
991 let read_schema = migration.codec_read();
992 let key = K::Schema::decoder_any(&*read_schema.key, &*structured.key);
993 let val = V::Schema::decoder_any(&*read_schema.val, &*structured.val);
994
995 match &key {
996 Ok(key_decoder) => {
997 let key_size_after = key_decoder.goodbytes();
998 let key_diff = key_size_before.saturating_sub(key_size_after);
999 metrics
1000 .pushdown
1001 .parts_projection_trimmed_bytes
1002 .inc_by(u64::cast_from(key_diff));
1003 }
1004 Err(e) => {
1005 soft_panic_or_log!("failed to create decoder: {e:#?}");
1006 }
1007 }
1008
1009 Some((key.ok()?, val.ok()?))
1010 };
1011
1012 let updates = part.normalize(&metrics.columnar);
1013 let timestamps = updates.timestamps().clone();
1014 let diffs = updates.diffs().clone();
1015 let part = match updates {
1016 BlobTraceUpdates::Row(records) => EitherOrBoth::Left(records),
1018 BlobTraceUpdates::Structured { key_values, .. } => EitherOrBoth::Right(
1019 downcast_structured(key_values, true).expect("valid schemas for structured data"),
1022 ),
1023 BlobTraceUpdates::Both(records, ext) => match part_decode_format {
1025 PartDecodeFormat::Row {
1026 validate_structured: false,
1027 } => EitherOrBoth::Left(records),
1028 PartDecodeFormat::Row {
1029 validate_structured: true,
1030 } => match downcast_structured(ext, false) {
1031 Some(decoders) => EitherOrBoth::Both(records, decoders),
1032 None => EitherOrBoth::Left(records),
1033 },
1034 PartDecodeFormat::Arrow => match downcast_structured(ext, false) {
1035 Some(decoders) => EitherOrBoth::Right(decoders),
1036 None => EitherOrBoth::Left(records),
1037 },
1038 },
1039 };
1040
1041 FetchedPart {
1042 metrics,
1043 ts_filter,
1044 part,
1045 peek_stash: None,
1046 timestamps,
1047 diffs,
1048 migration,
1049 filter_pushdown_audit,
1050 part_cursor: 0,
1051 key_storage: None,
1052 val_storage: None,
1053 _phantom: PhantomData,
1054 }
1055 }
1056
1057 pub fn is_filter_pushdown_audit(&self) -> Option<impl std::fmt::Debug + use<K, V, T, D>> {
1063 self.filter_pushdown_audit.clone()
1064 }
1065}
1066
1067#[derive(Debug)]
1070pub(crate) struct EncodedPart<T> {
1071 metrics: ReadMetrics,
1072 registered_desc: Description<T>,
1073 part: BlobTraceBatchPart<T>,
1074 needs_truncation: bool,
1075 ts_rewrite: Option<Antichain<T>>,
1076}
1077
1078impl<K, V, T, D> FetchedPart<K, V, T, D>
1079where
1080 K: Debug + Codec,
1081 V: Debug + Codec,
1082 T: Timestamp + Lattice + Codec64,
1083 D: Monoid + Codec64 + Send + Sync,
1084{
1085 pub fn next_with_storage(
1090 &mut self,
1091 key: &mut Option<K>,
1092 val: &mut Option<V>,
1093 ) -> Option<((K, V), T, D)> {
1094 let mut consolidated = self.peek_stash.take();
1095 loop {
1096 let next = if self.part_cursor < self.timestamps.len() {
1098 let next_idx = self.part_cursor;
1099 self.part_cursor += 1;
1100 let mut t = T::decode(self.timestamps.values()[next_idx].to_le_bytes());
1103 if !self.ts_filter.filter_ts(&mut t) {
1104 continue;
1105 }
1106 let d = D::decode(self.diffs.values()[next_idx].to_le_bytes());
1107 if d.is_zero() {
1108 continue;
1109 }
1110 let kv = self.decode_kv(next_idx, key, val);
1111 (kv, t, d)
1112 } else {
1113 break;
1114 };
1115
1116 if let Some((kv, t, d)) = &mut consolidated {
1118 let (kv_next, t_next, d_next) = &next;
1119 if kv == kv_next && t == t_next {
1120 d.plus_equals(d_next);
1121 if d.is_zero() {
1122 consolidated = None;
1123 }
1124 } else {
1125 self.peek_stash = Some(next);
1126 break;
1127 }
1128 } else {
1129 consolidated = Some(next);
1130 }
1131 }
1132
1133 let (kv, t, d) = consolidated?;
1134
1135 Some((kv, t, d))
1136 }
1137
1138 fn decode_kv(&mut self, index: usize, key: &mut Option<K>, val: &mut Option<V>) -> (K, V) {
1139 let decoded = self
1140 .part
1141 .as_ref()
1142 .map_left(|codec| {
1143 let ((ck, cv), _, _) = codec.get(index).expect("valid index");
1144 let (k, v) = Self::decode_codec(
1145 &*self.metrics,
1146 self.migration.codec_read(),
1147 ck,
1148 cv,
1149 key,
1150 val,
1151 &mut self.key_storage,
1152 &mut self.val_storage,
1153 );
1154 (k.expect("valid legacy key"), v.expect("valid legacy value"))
1155 })
1156 .map_right(|(structured_key, structured_val)| {
1157 self.decode_structured(index, structured_key, structured_val, key, val)
1158 });
1159
1160 match decoded {
1161 EitherOrBoth::Both((k, v), (k_s, v_s)) => {
1162 let is_valid = self
1164 .metrics
1165 .columnar
1166 .arrow()
1167 .key()
1168 .report_valid(|| k_s == k);
1169 if !is_valid {
1170 soft_panic_no_log!("structured key did not match, {k_s:?} != {k:?}");
1171 }
1172 let is_valid = self
1174 .metrics
1175 .columnar
1176 .arrow()
1177 .val()
1178 .report_valid(|| v_s == v);
1179 if !is_valid {
1180 soft_panic_no_log!("structured val did not match, {v_s:?} != {v:?}");
1181 }
1182
1183 (k, v)
1184 }
1185 EitherOrBoth::Left(kv) => kv,
1186 EitherOrBoth::Right(kv) => kv,
1187 }
1188 }
1189
1190 fn decode_codec(
1191 metrics: &Metrics,
1192 read_schemas: &Schemas<K, V>,
1193 key_buf: &[u8],
1194 val_buf: &[u8],
1195 key: &mut Option<K>,
1196 val: &mut Option<V>,
1197 key_storage: &mut Option<K::Storage>,
1198 val_storage: &mut Option<V::Storage>,
1199 ) -> (Result<K, String>, Result<V, String>) {
1200 let k = metrics.codecs.key.decode(|| match key.take() {
1201 Some(mut key) => {
1202 match K::decode_from(&mut key, key_buf, key_storage, &read_schemas.key) {
1203 Ok(()) => Ok(key),
1204 Err(err) => Err(err),
1205 }
1206 }
1207 None => K::decode(key_buf, &read_schemas.key),
1208 });
1209 let v = metrics.codecs.val.decode(|| match val.take() {
1210 Some(mut val) => {
1211 match V::decode_from(&mut val, val_buf, val_storage, &read_schemas.val) {
1212 Ok(()) => Ok(val),
1213 Err(err) => Err(err),
1214 }
1215 }
1216 None => V::decode(val_buf, &read_schemas.val),
1217 });
1218 (k, v)
1219 }
1220
1221 fn decode_structured(
1222 &self,
1223 idx: usize,
1224 keys: &<K::Schema as Schema<K>>::Decoder,
1225 vals: &<V::Schema as Schema<V>>::Decoder,
1226 key: &mut Option<K>,
1227 val: &mut Option<V>,
1228 ) -> (K, V) {
1229 let mut key = key.take().unwrap_or_default();
1230 keys.decode(idx, &mut key);
1231
1232 let mut val = val.take().unwrap_or_default();
1233 vals.decode(idx, &mut val);
1234
1235 (key, val)
1236 }
1237}
1238
1239impl<K, V, T, D> Iterator for FetchedPart<K, V, T, D>
1240where
1241 K: Debug + Codec,
1242 V: Debug + Codec,
1243 T: Timestamp + Lattice + Codec64,
1244 D: Monoid + Codec64 + Send + Sync,
1245{
1246 type Item = ((K, V), T, D);
1247
1248 fn next(&mut self) -> Option<Self::Item> {
1249 self.next_with_storage(&mut None, &mut None)
1250 }
1251
1252 fn size_hint(&self) -> (usize, Option<usize>) {
1253 let max_len = self.timestamps.len();
1255 (0, Some(max_len))
1256 }
1257}
1258
1259impl<T> EncodedPart<T>
1260where
1261 T: Timestamp + Lattice + Codec64,
1262{
1263 pub async fn fetch(
1264 cfg: &FetchConfig,
1265 shard_id: &ShardId,
1266 blob: &dyn Blob,
1267 metrics: &Metrics,
1268 shard_metrics: &ShardMetrics,
1269 read_metrics: &ReadMetrics,
1270 registered_desc: &Description<T>,
1271 part: &BatchPart<T>,
1272 ) -> Result<Self, BlobKey> {
1273 match part {
1274 BatchPart::Hollow(x) => {
1275 fetch_batch_part(
1276 cfg,
1277 shard_id,
1278 blob,
1279 metrics,
1280 shard_metrics,
1281 read_metrics,
1282 registered_desc,
1283 x,
1284 )
1285 .await
1286 }
1287 BatchPart::Inline {
1288 updates,
1289 ts_rewrite,
1290 ..
1291 } => Ok(EncodedPart::from_inline(
1292 cfg,
1293 metrics,
1294 read_metrics.clone(),
1295 registered_desc.clone(),
1296 updates,
1297 ts_rewrite.as_ref(),
1298 )),
1299 }
1300 }
1301
1302 pub(crate) fn from_inline(
1303 cfg: &FetchConfig,
1304 metrics: &Metrics,
1305 read_metrics: ReadMetrics,
1306 desc: Description<T>,
1307 x: &LazyInlineBatchPart,
1308 ts_rewrite: Option<&Antichain<T>>,
1309 ) -> Self {
1310 let parsed = x.decode(&metrics.columnar).expect("valid inline part");
1311 Self::new(cfg, read_metrics, desc, "inline", ts_rewrite, parsed)
1312 }
1313
1314 pub(crate) fn from_hollow(
1315 cfg: &FetchConfig,
1316 metrics: ReadMetrics,
1317 registered_desc: Description<T>,
1318 part: &HollowBatchPart<T>,
1319 parsed: BlobTraceBatchPart<T>,
1320 ) -> Self {
1321 Self::new(
1322 cfg,
1323 metrics,
1324 registered_desc,
1325 &part.key.0,
1326 part.ts_rewrite.as_ref(),
1327 parsed,
1328 )
1329 }
1330
1331 pub(crate) fn new(
1332 cfg: &FetchConfig,
1333 metrics: ReadMetrics,
1334 registered_desc: Description<T>,
1335 printable_name: &str,
1336 ts_rewrite: Option<&Antichain<T>>,
1337 parsed: BlobTraceBatchPart<T>,
1338 ) -> Self {
1339 let inline_desc = &parsed.desc;
1354 let needs_truncation = inline_desc.lower() != registered_desc.lower()
1355 || inline_desc.upper() != registered_desc.upper();
1356 if needs_truncation {
1357 if cfg.validate_bounds_on_read {
1358 soft_assert_or_log!(
1359 PartialOrder::less_equal(inline_desc.lower(), registered_desc.lower()),
1360 "key={} inline={:?} registered={:?}",
1361 printable_name,
1362 inline_desc,
1363 registered_desc
1364 );
1365
1366 if ts_rewrite.is_none() {
1367 soft_assert_or_log!(
1372 PartialOrder::less_equal(registered_desc.upper(), inline_desc.upper()),
1373 "key={} inline={:?} registered={:?}",
1374 printable_name,
1375 inline_desc,
1376 registered_desc
1377 );
1378 }
1379 }
1380 assert_eq!(
1385 inline_desc.since(),
1386 &Antichain::from_elem(T::minimum()),
1387 "key={} inline={:?} registered={:?}",
1388 printable_name,
1389 inline_desc,
1390 registered_desc
1391 );
1392 } else {
1393 soft_assert_or_log!(
1394 PartialOrder::less_equal(inline_desc.since(), registered_desc.since()),
1395 "key={} inline={:?} registered={:?}",
1396 printable_name,
1397 inline_desc,
1398 registered_desc
1399 );
1400 assert_eq!(
1401 inline_desc.lower(),
1402 registered_desc.lower(),
1403 "key={} inline={:?} registered={:?}",
1404 printable_name,
1405 inline_desc,
1406 registered_desc
1407 );
1408 assert_eq!(
1409 inline_desc.upper(),
1410 registered_desc.upper(),
1411 "key={} inline={:?} registered={:?}",
1412 printable_name,
1413 inline_desc,
1414 registered_desc
1415 );
1416 }
1417
1418 EncodedPart {
1419 metrics,
1420 registered_desc,
1421 part: parsed,
1422 needs_truncation,
1423 ts_rewrite: ts_rewrite.cloned(),
1424 }
1425 }
1426
1427 pub(crate) fn maybe_unconsolidated(&self) -> bool {
1428 self.part.desc.since().borrow() == AntichainRef::new(&[T::minimum()])
1431 }
1432
1433 pub(crate) fn updates(&self) -> &BlobTraceUpdates {
1434 &self.part.updates
1435 }
1436
1437 pub(crate) fn normalize(&self, metrics: &ColumnarMetrics) -> BlobTraceUpdates {
1439 let updates = self.part.updates.clone();
1440 if !self.needs_truncation && self.ts_rewrite.is_none() {
1441 return updates;
1442 }
1443
1444 let mut codec = updates
1445 .records()
1446 .map(|r| (r.keys().clone(), r.vals().clone()));
1447 let mut structured = updates.structured().cloned();
1448 let mut timestamps = updates.timestamps().clone();
1449 let mut diffs = updates.diffs().clone();
1450
1451 if let Some(rewrite) = self.ts_rewrite.as_ref() {
1452 timestamps = arrow::compute::unary(×tamps, |i: i64| {
1453 let mut t = T::decode(i.to_le_bytes());
1454 t.advance_by(rewrite.borrow());
1455 i64::from_le_bytes(T::encode(&t))
1456 });
1457 }
1458
1459 let reallocated = if self.needs_truncation {
1460 let filter = BooleanArray::from_unary(×tamps, |i| {
1461 let t = T::decode(i.to_le_bytes());
1462 let truncate_t = {
1463 !self.registered_desc.lower().less_equal(&t)
1464 || self.registered_desc.upper().less_equal(&t)
1465 };
1466 !truncate_t
1467 });
1468 if filter.false_count() == 0 {
1469 false
1471 } else {
1472 let filter = FilterBuilder::new(&filter).optimize().build();
1473 let do_filter = |array: &dyn Array| filter.filter(array).expect("valid filter len");
1474 if let Some((keys, vals)) = codec {
1475 codec = Some((
1476 realloc_array(do_filter(&keys).as_binary(), metrics),
1477 realloc_array(do_filter(&vals).as_binary(), metrics),
1478 ));
1479 }
1480 if let Some(ext) = structured {
1481 structured = Some(ColumnarRecordsStructuredExt {
1482 key: realloc_any(do_filter(&*ext.key), metrics),
1483 val: realloc_any(do_filter(&*ext.val), metrics),
1484 });
1485 }
1486 timestamps = realloc_array(do_filter(×tamps).as_primitive(), metrics);
1487 diffs = realloc_array(do_filter(&diffs).as_primitive(), metrics);
1488 true
1489 }
1490 } else {
1491 false
1492 };
1493
1494 if self.ts_rewrite.is_some() && !reallocated {
1495 timestamps = realloc_array(×tamps, metrics);
1496 }
1497
1498 if self.ts_rewrite.is_some() {
1499 self.metrics
1500 .ts_rewrite
1501 .inc_by(u64::cast_from(timestamps.len()));
1502 }
1503
1504 match (codec, structured) {
1505 (Some((key, value)), None) => {
1506 BlobTraceUpdates::Row(ColumnarRecords::new(key, value, timestamps, diffs))
1507 }
1508 (Some((key, value)), Some(ext)) => {
1509 BlobTraceUpdates::Both(ColumnarRecords::new(key, value, timestamps, diffs), ext)
1510 }
1511 (None, Some(ext)) => BlobTraceUpdates::Structured {
1512 key_values: ext,
1513 timestamps,
1514 diffs,
1515 },
1516 (None, None) => unreachable!(),
1517 }
1518 }
1519}
1520
1521#[derive(Debug, Serialize, Deserialize, Clone)]
1530pub struct ExchangeableBatchPart<T> {
1531 shard_id: ShardId,
1532 encoded_size_bytes: usize,
1534 desc: Description<T>,
1535 filter: FetchBatchFilter<T>,
1536 part: LazyProto<ProtoHollowBatchPart>,
1537 reader_id: LeasedReaderId,
1540 filter_pushdown_audit: bool,
1541}
1542
1543impl<T> ExchangeableBatchPart<T> {
1544 pub fn encoded_size_bytes(&self) -> usize {
1546 self.encoded_size_bytes
1547 }
1548
1549 pub fn reader_id(&self) -> &LeasedReaderId {
1551 &self.reader_id
1552 }
1553}
1554
1555#[derive(Debug, Copy, Clone)]
1559pub enum PartDecodeFormat {
1560 Row {
1562 validate_structured: bool,
1564 },
1565 Arrow,
1567}
1568
1569impl PartDecodeFormat {
1570 pub const fn default() -> Self {
1572 PartDecodeFormat::Arrow
1573 }
1574
1575 pub fn from_str(s: &str) -> Self {
1578 match s {
1579 "row" => PartDecodeFormat::Row {
1580 validate_structured: false,
1581 },
1582 "row_with_validate" => PartDecodeFormat::Row {
1583 validate_structured: true,
1584 },
1585 "arrow" => PartDecodeFormat::Arrow,
1586 x => {
1587 let default = PartDecodeFormat::default();
1588 soft_panic_or_log!("Invalid part decode format: '{x}', falling back to {default}");
1589 default
1590 }
1591 }
1592 }
1593
1594 pub const fn as_str(&self) -> &'static str {
1596 match self {
1597 PartDecodeFormat::Row {
1598 validate_structured: false,
1599 } => "row",
1600 PartDecodeFormat::Row {
1601 validate_structured: true,
1602 } => "row_with_validate",
1603 PartDecodeFormat::Arrow => "arrow",
1604 }
1605 }
1606}
1607
1608impl fmt::Display for PartDecodeFormat {
1609 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1610 f.write_str(self.as_str())
1611 }
1612}
1613
1614#[mz_ore::test]
1615fn client_exchange_data() {
1616 fn is_exchange_data<T: timely::ExchangeData>() {}
1620 is_exchange_data::<ExchangeableBatchPart<u64>>();
1621 is_exchange_data::<ExchangeableBatchPart<u64>>();
1622}