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