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> {
663 self.part.stats().and_then(|x| x.try_decode().ok())
664 }
665
666 pub fn maybe_optimize(&mut self, cfg: &ConfigSet, key: ArrayRef, val: ArrayRef) {
669 assert_eq!(key.len(), 1, "expect a single-row key array");
670 assert_eq!(val.len(), 1, "expect a single-row val array");
671 let as_of = match &self.filter {
672 FetchBatchFilter::Snapshot { as_of } => as_of,
673 FetchBatchFilter::Listen { .. } | FetchBatchFilter::Compaction { .. } => return,
674 };
675 if !OPTIMIZE_IGNORED_DATA_FETCH.get(cfg) {
676 return;
677 }
678 let (diffs_sum, _stats) = match &self.part {
679 BatchPart::Hollow(x) => (x.diffs_sum, x.stats.as_ref()),
680 BatchPart::Inline { .. } => return,
681 };
682 debug!(
683 "try_optimize_ignored_data_fetch diffs_sum={:?} as_of={:?} lower={:?} upper={:?}",
684 diffs_sum.map(i64::decode),
686 as_of.elements(),
687 self.desc.lower().elements(),
688 self.desc.upper().elements()
689 );
690 let as_of = match &as_of.elements() {
691 &[as_of] => as_of,
692 _ => return,
693 };
694 let eligible = self.desc.upper().less_equal(as_of) && self.desc.since().less_equal(as_of);
700 if !eligible {
701 return;
702 }
703 let Some(diffs_sum) = diffs_sum else {
704 return;
705 };
706
707 debug!(
708 "try_optimize_ignored_data_fetch faked {:?} diffs at ts {:?} skipping fetch of {} bytes",
709 i64::decode(diffs_sum),
711 as_of,
712 self.part.encoded_size_bytes(),
713 );
714 self.metrics.pushdown.parts_faked_count.inc();
715 self.metrics
716 .pushdown
717 .parts_faked_bytes
718 .inc_by(u64::cast_from(self.part.encoded_size_bytes()));
719 let timestamps = {
720 let mut col = Codec64Mut::with_capacity(1);
721 col.push(as_of);
722 col.finish()
723 };
724 let diffs = {
725 let mut col = Codec64Mut::with_capacity(1);
726 col.push_raw(diffs_sum);
727 col.finish()
728 };
729 let updates = BlobTraceUpdates::Structured {
730 key_values: ColumnarRecordsStructuredExt { key, val },
731 timestamps,
732 diffs,
733 };
734 let faked_data = LazyInlineBatchPart::from(&ProtoInlineBatchPart {
735 desc: Some(self.desc.into_proto()),
736 index: 0,
737 updates: Some(updates.into_proto()),
738 });
739 self.part = BatchPart::Inline {
740 updates: faked_data,
741 ts_rewrite: None,
742 schema_id: None,
743 deprecated_schema_id: None,
744 };
745 }
746}
747
748impl<T> Drop for LeasedBatchPart<T> {
749 fn drop(&mut self) {
751 self.metrics.lease.dropped_part.inc()
752 }
753}
754
755#[derive(Debug)]
760pub struct FetchedBlob<K: Codec, V: Codec, T, D> {
761 metrics: Arc<Metrics>,
762 read_metrics: ReadMetrics,
763 buf: FetchedBlobBuf<T>,
764 registered_desc: Description<T>,
765 migration: PartMigration<K, V>,
766 filter: FetchBatchFilter<T>,
767 filter_pushdown_audit: bool,
768 structured_part_audit: PartDecodeFormat,
769 fetch_permit: Option<Arc<MetricsPermits>>,
770 fetch_config: FetchConfig,
771 _phantom: PhantomData<fn() -> D>,
772}
773
774#[derive(Debug, Clone)]
775enum FetchedBlobBuf<T> {
776 Hollow {
777 buf: SegmentedBytes,
778 part: HollowBatchPart<T>,
779 },
780 Inline {
781 desc: Description<T>,
782 updates: LazyInlineBatchPart,
783 ts_rewrite: Option<Antichain<T>>,
784 },
785}
786
787impl<K: Codec, V: Codec, T: Clone, D> Clone for FetchedBlob<K, V, T, D> {
788 fn clone(&self) -> Self {
789 Self {
790 metrics: Arc::clone(&self.metrics),
791 read_metrics: self.read_metrics.clone(),
792 buf: self.buf.clone(),
793 registered_desc: self.registered_desc.clone(),
794 migration: self.migration.clone(),
795 filter: self.filter.clone(),
796 filter_pushdown_audit: self.filter_pushdown_audit.clone(),
797 fetch_permit: self.fetch_permit.clone(),
798 structured_part_audit: self.structured_part_audit.clone(),
799 fetch_config: self.fetch_config.clone(),
800 _phantom: self._phantom.clone(),
801 }
802 }
803}
804
805pub struct ShardSourcePart<K: Codec, V: Codec, T, D> {
808 pub part: FetchedPart<K, V, T, D>,
810 fetch_permit: Option<Arc<MetricsPermits>>,
811}
812
813impl<K, V, T: Debug, D: Debug> Debug for ShardSourcePart<K, V, T, D>
814where
815 K: Codec + Debug,
816 <K as Codec>::Storage: Debug,
817 V: Codec + Debug,
818 <V as Codec>::Storage: Debug,
819{
820 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
821 let ShardSourcePart { part, fetch_permit } = self;
822 f.debug_struct("ShardSourcePart")
823 .field("part", part)
824 .field("fetch_permit", fetch_permit)
825 .finish()
826 }
827}
828
829impl<K: Codec, V: Codec, T: Timestamp + Lattice + Codec64, D> FetchedBlob<K, V, T, D> {
830 pub fn parse(&self) -> ShardSourcePart<K, V, T, D> {
832 self.parse_internal(&self.fetch_config)
833 }
834
835 pub(crate) fn parse_internal(&self, cfg: &FetchConfig) -> ShardSourcePart<K, V, T, D> {
837 let (part, stats) = match &self.buf {
838 FetchedBlobBuf::Hollow { buf, part } => {
839 let parsed = decode_batch_part_blob(
840 cfg,
841 &self.metrics,
842 &self.read_metrics,
843 self.registered_desc.clone(),
844 part,
845 buf,
846 );
847 (parsed, part.stats.as_ref())
848 }
849 FetchedBlobBuf::Inline {
850 desc,
851 updates,
852 ts_rewrite,
853 } => {
854 let parsed = EncodedPart::from_inline(
855 cfg,
856 &self.metrics,
857 self.read_metrics.clone(),
858 desc.clone(),
859 updates,
860 ts_rewrite.as_ref(),
861 );
862 (parsed, None)
863 }
864 };
865 let part = FetchedPart::new(
866 Arc::clone(&self.metrics),
867 part,
868 self.migration.clone(),
869 self.filter.clone(),
870 self.filter_pushdown_audit,
871 self.structured_part_audit,
872 stats,
873 );
874 ShardSourcePart {
875 part,
876 fetch_permit: self.fetch_permit.clone(),
877 }
878 }
879
880 pub fn stats(&self) -> Option<PartStats> {
885 match &self.buf {
886 FetchedBlobBuf::Hollow { part, .. } => {
887 part.stats.as_ref().and_then(|x| x.try_decode().ok())
888 }
889 FetchedBlobBuf::Inline { .. } => None,
890 }
891 }
892}
893
894#[derive(Debug)]
899pub struct FetchedPart<K: Codec, V: Codec, T, D> {
900 metrics: Arc<Metrics>,
901 ts_filter: FetchBatchFilter<T>,
902 part: EitherOrBoth<
905 ColumnarRecords,
906 (
907 <K::Schema as Schema<K>>::Decoder,
908 <V::Schema as Schema<V>>::Decoder,
909 ),
910 >,
911 timestamps: Int64Array,
912 diffs: Int64Array,
913 migration: PartMigration<K, V>,
914 filter_pushdown_audit: Option<LazyPartStats>,
915 peek_stash: Option<((K, V), T, D)>,
916 part_cursor: usize,
917 key_storage: Option<K::Storage>,
918 val_storage: Option<V::Storage>,
919
920 _phantom: PhantomData<fn() -> D>,
921}
922
923impl<K: Codec, V: Codec, T: Timestamp + Lattice + Codec64, D> FetchedPart<K, V, T, D> {
924 pub(crate) fn new(
925 metrics: Arc<Metrics>,
926 part: EncodedPart<T>,
927 migration: PartMigration<K, V>,
928 ts_filter: FetchBatchFilter<T>,
929 filter_pushdown_audit: bool,
930 part_decode_format: PartDecodeFormat,
931 stats: Option<&LazyPartStats>,
932 ) -> Self {
933 let part_len = u64::cast_from(part.part.updates.len());
934 match &migration {
935 PartMigration::SameSchema { .. } => metrics.schema.migration_count_same.inc(),
936 PartMigration::Schemaless { .. } => {
937 metrics.schema.migration_count_codec.inc();
938 metrics.schema.migration_len_legacy_codec.inc_by(part_len);
939 }
940 PartMigration::Either { .. } => {
941 metrics.schema.migration_count_either.inc();
942 match part_decode_format {
943 PartDecodeFormat::Row {
944 validate_structured: false,
945 } => metrics.schema.migration_len_either_codec.inc_by(part_len),
946 PartDecodeFormat::Row {
947 validate_structured: true,
948 } => {
949 metrics.schema.migration_len_either_codec.inc_by(part_len);
950 metrics.schema.migration_len_either_arrow.inc_by(part_len);
951 }
952 PartDecodeFormat::Arrow => {
953 metrics.schema.migration_len_either_arrow.inc_by(part_len)
954 }
955 }
956 }
957 }
958
959 let filter_pushdown_audit = if filter_pushdown_audit {
960 stats.cloned()
961 } else {
962 None
963 };
964
965 let downcast_structured = |structured: ColumnarRecordsStructuredExt,
966 structured_only: bool| {
967 let key_size_before = ArrayOrd::new(&structured.key).goodbytes();
968
969 let structured = match &migration {
970 PartMigration::SameSchema { .. } => structured,
971 PartMigration::Schemaless { read } if structured_only => {
972 let start = Instant::now();
974 let read_key = data_type::<K>(&*read.key).ok()?;
975 let read_val = data_type::<V>(&*read.val).ok()?;
976 let key_migration = backward_compatible(structured.key.data_type(), &read_key)?;
977 let val_migration = backward_compatible(structured.val.data_type(), &read_val)?;
978 let key = key_migration.migrate(structured.key);
979 let val = val_migration.migrate(structured.val);
980 metrics
981 .schema
982 .migration_migrate_seconds
983 .inc_by(start.elapsed().as_secs_f64());
984 ColumnarRecordsStructuredExt { key, val }
985 }
986 PartMigration::Schemaless { .. } => return None,
987 PartMigration::Either {
988 write: _,
989 read: _,
990 key_migration,
991 val_migration,
992 } => {
993 let start = Instant::now();
994 let key = key_migration.migrate(structured.key);
995 let val = val_migration.migrate(structured.val);
996 metrics
997 .schema
998 .migration_migrate_seconds
999 .inc_by(start.elapsed().as_secs_f64());
1000 ColumnarRecordsStructuredExt { key, val }
1001 }
1002 };
1003
1004 let read_schema = migration.codec_read();
1005 let key = K::Schema::decoder_any(&*read_schema.key, &*structured.key);
1006 let val = V::Schema::decoder_any(&*read_schema.val, &*structured.val);
1007
1008 match &key {
1009 Ok(key_decoder) => {
1010 let key_size_after = key_decoder.goodbytes();
1011 let key_diff = key_size_before.saturating_sub(key_size_after);
1012 metrics
1013 .pushdown
1014 .parts_projection_trimmed_bytes
1015 .inc_by(u64::cast_from(key_diff));
1016 }
1017 Err(e) => {
1018 soft_panic_or_log!("failed to create decoder: {e:#?}");
1019 }
1020 }
1021
1022 Some((key.ok()?, val.ok()?))
1023 };
1024
1025 let updates = part.normalize(&metrics.columnar);
1026 let timestamps = updates.timestamps().clone();
1027 let diffs = updates.diffs().clone();
1028 let part = match updates {
1029 BlobTraceUpdates::Row(records) => EitherOrBoth::Left(records),
1031 BlobTraceUpdates::Structured { key_values, .. } => EitherOrBoth::Right(
1032 downcast_structured(key_values, true).expect("valid schemas for structured data"),
1035 ),
1036 BlobTraceUpdates::Both(records, ext) => match part_decode_format {
1038 PartDecodeFormat::Row {
1039 validate_structured: false,
1040 } => EitherOrBoth::Left(records),
1041 PartDecodeFormat::Row {
1042 validate_structured: true,
1043 } => match downcast_structured(ext, false) {
1044 Some(decoders) => EitherOrBoth::Both(records, decoders),
1045 None => EitherOrBoth::Left(records),
1046 },
1047 PartDecodeFormat::Arrow => match downcast_structured(ext, false) {
1048 Some(decoders) => EitherOrBoth::Right(decoders),
1049 None => EitherOrBoth::Left(records),
1050 },
1051 },
1052 };
1053
1054 FetchedPart {
1055 metrics,
1056 ts_filter,
1057 part,
1058 peek_stash: None,
1059 timestamps,
1060 diffs,
1061 migration,
1062 filter_pushdown_audit,
1063 part_cursor: 0,
1064 key_storage: None,
1065 val_storage: None,
1066 _phantom: PhantomData,
1067 }
1068 }
1069
1070 pub fn is_filter_pushdown_audit(&self) -> Option<impl std::fmt::Debug + use<K, V, T, D>> {
1076 self.filter_pushdown_audit.clone()
1077 }
1078}
1079
1080#[derive(Debug)]
1083pub(crate) struct EncodedPart<T> {
1084 metrics: ReadMetrics,
1085 registered_desc: Description<T>,
1086 part: BlobTraceBatchPart<T>,
1087 needs_truncation: bool,
1088 ts_rewrite: Option<Antichain<T>>,
1089}
1090
1091impl<K, V, T, D> FetchedPart<K, V, T, D>
1092where
1093 K: Debug + Codec,
1094 V: Debug + Codec,
1095 T: Timestamp + Lattice + Codec64,
1096 D: Monoid + Codec64 + Send + Sync,
1097{
1098 pub fn next_with_storage(
1103 &mut self,
1104 key: &mut Option<K>,
1105 val: &mut Option<V>,
1106 ) -> Option<((K, V), T, D)> {
1107 let mut consolidated = self.peek_stash.take();
1108 loop {
1109 let next = if self.part_cursor < self.timestamps.len() {
1111 let next_idx = self.part_cursor;
1112 self.part_cursor += 1;
1113 let mut t = T::decode(self.timestamps.values()[next_idx].to_le_bytes());
1116 if !self.ts_filter.filter_ts(&mut t) {
1117 continue;
1118 }
1119 let d = D::decode(self.diffs.values()[next_idx].to_le_bytes());
1120 if d.is_zero() {
1121 continue;
1122 }
1123 let kv = self.decode_kv(next_idx, key, val);
1124 (kv, t, d)
1125 } else {
1126 break;
1127 };
1128
1129 if let Some((kv, t, d)) = &mut consolidated {
1131 let (kv_next, t_next, d_next) = &next;
1132 if kv == kv_next && t == t_next {
1133 d.plus_equals(d_next);
1134 if d.is_zero() {
1135 consolidated = None;
1136 }
1137 } else {
1138 self.peek_stash = Some(next);
1139 break;
1140 }
1141 } else {
1142 consolidated = Some(next);
1143 }
1144 }
1145
1146 let (kv, t, d) = consolidated?;
1147
1148 Some((kv, t, d))
1149 }
1150
1151 fn decode_kv(&mut self, index: usize, key: &mut Option<K>, val: &mut Option<V>) -> (K, V) {
1152 let decoded = self
1153 .part
1154 .as_ref()
1155 .map_left(|codec| {
1156 let ((ck, cv), _, _) = codec.get(index).expect("valid index");
1157 let (k, v) = Self::decode_codec(
1158 &*self.metrics,
1159 self.migration.codec_read(),
1160 ck,
1161 cv,
1162 key,
1163 val,
1164 &mut self.key_storage,
1165 &mut self.val_storage,
1166 );
1167 (k.expect("valid legacy key"), v.expect("valid legacy value"))
1168 })
1169 .map_right(|(structured_key, structured_val)| {
1170 self.decode_structured(index, structured_key, structured_val, key, val)
1171 });
1172
1173 match decoded {
1174 EitherOrBoth::Both((k, v), (k_s, v_s)) => {
1175 let is_valid = self
1177 .metrics
1178 .columnar
1179 .arrow()
1180 .key()
1181 .report_valid(|| k_s == k);
1182 if !is_valid {
1183 soft_panic_no_log!("structured key did not match, {k_s:?} != {k:?}");
1184 }
1185 let is_valid = self
1187 .metrics
1188 .columnar
1189 .arrow()
1190 .val()
1191 .report_valid(|| v_s == v);
1192 if !is_valid {
1193 soft_panic_no_log!("structured val did not match, {v_s:?} != {v:?}");
1194 }
1195
1196 (k, v)
1197 }
1198 EitherOrBoth::Left(kv) => kv,
1199 EitherOrBoth::Right(kv) => kv,
1200 }
1201 }
1202
1203 fn decode_codec(
1204 metrics: &Metrics,
1205 read_schemas: &Schemas<K, V>,
1206 key_buf: &[u8],
1207 val_buf: &[u8],
1208 key: &mut Option<K>,
1209 val: &mut Option<V>,
1210 key_storage: &mut Option<K::Storage>,
1211 val_storage: &mut Option<V::Storage>,
1212 ) -> (Result<K, String>, Result<V, String>) {
1213 let k = metrics.codecs.key.decode(|| match key.take() {
1214 Some(mut key) => {
1215 match K::decode_from(&mut key, key_buf, key_storage, &read_schemas.key) {
1216 Ok(()) => Ok(key),
1217 Err(err) => Err(err),
1218 }
1219 }
1220 None => K::decode(key_buf, &read_schemas.key),
1221 });
1222 let v = metrics.codecs.val.decode(|| match val.take() {
1223 Some(mut val) => {
1224 match V::decode_from(&mut val, val_buf, val_storage, &read_schemas.val) {
1225 Ok(()) => Ok(val),
1226 Err(err) => Err(err),
1227 }
1228 }
1229 None => V::decode(val_buf, &read_schemas.val),
1230 });
1231 (k, v)
1232 }
1233
1234 fn decode_structured(
1235 &self,
1236 idx: usize,
1237 keys: &<K::Schema as Schema<K>>::Decoder,
1238 vals: &<V::Schema as Schema<V>>::Decoder,
1239 key: &mut Option<K>,
1240 val: &mut Option<V>,
1241 ) -> (K, V) {
1242 let mut key = key.take().unwrap_or_default();
1243 keys.decode(idx, &mut key);
1244
1245 let mut val = val.take().unwrap_or_default();
1246 vals.decode(idx, &mut val);
1247
1248 (key, val)
1249 }
1250}
1251
1252impl<K, V, T, D> Iterator for FetchedPart<K, V, T, D>
1253where
1254 K: Debug + Codec,
1255 V: Debug + Codec,
1256 T: Timestamp + Lattice + Codec64,
1257 D: Monoid + Codec64 + Send + Sync,
1258{
1259 type Item = ((K, V), T, D);
1260
1261 fn next(&mut self) -> Option<Self::Item> {
1262 self.next_with_storage(&mut None, &mut None)
1263 }
1264
1265 fn size_hint(&self) -> (usize, Option<usize>) {
1266 let max_len = self.timestamps.len();
1268 (0, Some(max_len))
1269 }
1270}
1271
1272impl<T> EncodedPart<T>
1273where
1274 T: Timestamp + Lattice + Codec64,
1275{
1276 pub async fn fetch(
1277 cfg: &FetchConfig,
1278 shard_id: &ShardId,
1279 blob: &dyn Blob,
1280 metrics: &Metrics,
1281 shard_metrics: &ShardMetrics,
1282 read_metrics: &ReadMetrics,
1283 registered_desc: &Description<T>,
1284 part: &BatchPart<T>,
1285 ) -> Result<Self, BlobKey> {
1286 match part {
1287 BatchPart::Hollow(x) => {
1288 fetch_batch_part(
1289 cfg,
1290 shard_id,
1291 blob,
1292 metrics,
1293 shard_metrics,
1294 read_metrics,
1295 registered_desc,
1296 x,
1297 )
1298 .await
1299 }
1300 BatchPart::Inline {
1301 updates,
1302 ts_rewrite,
1303 ..
1304 } => Ok(EncodedPart::from_inline(
1305 cfg,
1306 metrics,
1307 read_metrics.clone(),
1308 registered_desc.clone(),
1309 updates,
1310 ts_rewrite.as_ref(),
1311 )),
1312 }
1313 }
1314
1315 pub(crate) fn from_inline(
1316 cfg: &FetchConfig,
1317 metrics: &Metrics,
1318 read_metrics: ReadMetrics,
1319 desc: Description<T>,
1320 x: &LazyInlineBatchPart,
1321 ts_rewrite: Option<&Antichain<T>>,
1322 ) -> Self {
1323 let parsed = x.decode(&metrics.columnar).expect("valid inline part");
1324 Self::new(cfg, read_metrics, desc, "inline", ts_rewrite, parsed)
1325 }
1326
1327 pub(crate) fn from_hollow(
1328 cfg: &FetchConfig,
1329 metrics: ReadMetrics,
1330 registered_desc: Description<T>,
1331 part: &HollowBatchPart<T>,
1332 parsed: BlobTraceBatchPart<T>,
1333 ) -> Self {
1334 Self::new(
1335 cfg,
1336 metrics,
1337 registered_desc,
1338 &part.key.0,
1339 part.ts_rewrite.as_ref(),
1340 parsed,
1341 )
1342 }
1343
1344 pub(crate) fn new(
1345 cfg: &FetchConfig,
1346 metrics: ReadMetrics,
1347 registered_desc: Description<T>,
1348 printable_name: &str,
1349 ts_rewrite: Option<&Antichain<T>>,
1350 parsed: BlobTraceBatchPart<T>,
1351 ) -> Self {
1352 let inline_desc = &parsed.desc;
1367 let needs_truncation = inline_desc.lower() != registered_desc.lower()
1368 || inline_desc.upper() != registered_desc.upper();
1369 if needs_truncation {
1370 if cfg.validate_bounds_on_read {
1371 soft_assert_or_log!(
1372 PartialOrder::less_equal(inline_desc.lower(), registered_desc.lower()),
1373 "key={} inline={:?} registered={:?}",
1374 printable_name,
1375 inline_desc,
1376 registered_desc
1377 );
1378
1379 if ts_rewrite.is_none() {
1380 soft_assert_or_log!(
1385 PartialOrder::less_equal(registered_desc.upper(), inline_desc.upper()),
1386 "key={} inline={:?} registered={:?}",
1387 printable_name,
1388 inline_desc,
1389 registered_desc
1390 );
1391 }
1392 }
1393 assert_eq!(
1398 inline_desc.since(),
1399 &Antichain::from_elem(T::minimum()),
1400 "key={} inline={:?} registered={:?}",
1401 printable_name,
1402 inline_desc,
1403 registered_desc
1404 );
1405 } else {
1406 soft_assert_or_log!(
1407 PartialOrder::less_equal(inline_desc.since(), registered_desc.since()),
1408 "key={} inline={:?} registered={:?}",
1409 printable_name,
1410 inline_desc,
1411 registered_desc
1412 );
1413 assert_eq!(
1414 inline_desc.lower(),
1415 registered_desc.lower(),
1416 "key={} inline={:?} registered={:?}",
1417 printable_name,
1418 inline_desc,
1419 registered_desc
1420 );
1421 assert_eq!(
1422 inline_desc.upper(),
1423 registered_desc.upper(),
1424 "key={} inline={:?} registered={:?}",
1425 printable_name,
1426 inline_desc,
1427 registered_desc
1428 );
1429 }
1430
1431 EncodedPart {
1432 metrics,
1433 registered_desc,
1434 part: parsed,
1435 needs_truncation,
1436 ts_rewrite: ts_rewrite.cloned(),
1437 }
1438 }
1439
1440 pub(crate) fn maybe_unconsolidated(&self) -> bool {
1441 self.part.desc.since().borrow() == AntichainRef::new(&[T::minimum()])
1444 }
1445
1446 pub(crate) fn updates(&self) -> &BlobTraceUpdates {
1447 &self.part.updates
1448 }
1449
1450 pub(crate) fn normalize(&self, metrics: &ColumnarMetrics) -> BlobTraceUpdates {
1452 let updates = self.part.updates.clone();
1453 if !self.needs_truncation && self.ts_rewrite.is_none() {
1454 return updates;
1455 }
1456
1457 let mut codec = updates
1458 .records()
1459 .map(|r| (r.keys().clone(), r.vals().clone()));
1460 let mut structured = updates.structured().cloned();
1461 let mut timestamps = updates.timestamps().clone();
1462 let mut diffs = updates.diffs().clone();
1463
1464 if let Some(rewrite) = self.ts_rewrite.as_ref() {
1465 timestamps = arrow::compute::unary(×tamps, |i: i64| {
1466 let mut t = T::decode(i.to_le_bytes());
1467 t.advance_by(rewrite.borrow());
1468 i64::from_le_bytes(T::encode(&t))
1469 });
1470 }
1471
1472 let reallocated = if self.needs_truncation {
1473 let filter = BooleanArray::from_unary(×tamps, |i| {
1474 let t = T::decode(i.to_le_bytes());
1475 let truncate_t = {
1476 !self.registered_desc.lower().less_equal(&t)
1477 || self.registered_desc.upper().less_equal(&t)
1478 };
1479 !truncate_t
1480 });
1481 if filter.false_count() == 0 {
1482 false
1484 } else {
1485 let filter = FilterBuilder::new(&filter).optimize().build();
1486 let do_filter = |array: &dyn Array| filter.filter(array).expect("valid filter len");
1487 if let Some((keys, vals)) = codec {
1488 codec = Some((
1489 realloc_array(do_filter(&keys).as_binary(), metrics),
1490 realloc_array(do_filter(&vals).as_binary(), metrics),
1491 ));
1492 }
1493 if let Some(ext) = structured {
1494 structured = Some(ColumnarRecordsStructuredExt {
1495 key: realloc_any(do_filter(&*ext.key), metrics),
1496 val: realloc_any(do_filter(&*ext.val), metrics),
1497 });
1498 }
1499 timestamps = realloc_array(do_filter(×tamps).as_primitive(), metrics);
1500 diffs = realloc_array(do_filter(&diffs).as_primitive(), metrics);
1501 true
1502 }
1503 } else {
1504 false
1505 };
1506
1507 if self.ts_rewrite.is_some() && !reallocated {
1508 timestamps = realloc_array(×tamps, metrics);
1509 }
1510
1511 if self.ts_rewrite.is_some() {
1512 self.metrics
1513 .ts_rewrite
1514 .inc_by(u64::cast_from(timestamps.len()));
1515 }
1516
1517 match (codec, structured) {
1518 (Some((key, value)), None) => {
1519 BlobTraceUpdates::Row(ColumnarRecords::new(key, value, timestamps, diffs))
1520 }
1521 (Some((key, value)), Some(ext)) => {
1522 BlobTraceUpdates::Both(ColumnarRecords::new(key, value, timestamps, diffs), ext)
1523 }
1524 (None, Some(ext)) => BlobTraceUpdates::Structured {
1525 key_values: ext,
1526 timestamps,
1527 diffs,
1528 },
1529 (None, None) => unreachable!(),
1530 }
1531 }
1532}
1533
1534#[derive(Debug, Serialize, Deserialize, Clone)]
1543pub struct ExchangeableBatchPart<T> {
1544 shard_id: ShardId,
1545 encoded_size_bytes: usize,
1547 desc: Description<T>,
1548 filter: FetchBatchFilter<T>,
1549 part: LazyProto<ProtoHollowBatchPart>,
1550 reader_id: LeasedReaderId,
1553 filter_pushdown_audit: bool,
1554}
1555
1556impl<T> ExchangeableBatchPart<T> {
1557 pub fn encoded_size_bytes(&self) -> usize {
1559 self.encoded_size_bytes
1560 }
1561
1562 pub fn reader_id(&self) -> &LeasedReaderId {
1564 &self.reader_id
1565 }
1566}
1567
1568#[derive(Debug, Copy, Clone)]
1572pub enum PartDecodeFormat {
1573 Row {
1575 validate_structured: bool,
1577 },
1578 Arrow,
1580}
1581
1582impl PartDecodeFormat {
1583 pub const fn default() -> Self {
1585 PartDecodeFormat::Arrow
1586 }
1587
1588 pub fn from_str(s: &str) -> Self {
1591 match s {
1592 "row" => PartDecodeFormat::Row {
1593 validate_structured: false,
1594 },
1595 "row_with_validate" => PartDecodeFormat::Row {
1596 validate_structured: true,
1597 },
1598 "arrow" => PartDecodeFormat::Arrow,
1599 x => {
1600 let default = PartDecodeFormat::default();
1601 soft_panic_or_log!("Invalid part decode format: '{x}', falling back to {default}");
1602 default
1603 }
1604 }
1605 }
1606
1607 pub const fn as_str(&self) -> &'static str {
1609 match self {
1610 PartDecodeFormat::Row {
1611 validate_structured: false,
1612 } => "row",
1613 PartDecodeFormat::Row {
1614 validate_structured: true,
1615 } => "row_with_validate",
1616 PartDecodeFormat::Arrow => "arrow",
1617 }
1618 }
1619}
1620
1621impl fmt::Display for PartDecodeFormat {
1622 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1623 f.write_str(self.as_str())
1624 }
1625}
1626
1627#[mz_ore::test]
1628fn client_exchange_data() {
1629 fn is_exchange_data<T: timely::ExchangeData>() {}
1633 is_exchange_data::<ExchangeableBatchPart<u64>>();
1634 is_exchange_data::<ExchangeableBatchPart<u64>>();
1635}