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