inner method for read_batches_consolidated

Author: aljoscha

src/persist-client/src/lib.rs

@@ -21,6 +21,7 @@ use std::sync::Arc;
 
 use differential_dataflow::difference::Semigroup;
 use differential_dataflow::lattice::Lattice;
+use itertools::Itertools;
 use mz_build_info::{BuildInfo, build_info};
 use mz_dyncfg::ConfigSet;
 use mz_ore::{instrument, soft_assert_or_log};
@@ -32,22 +33,20 @@ use semver::Version;
 use timely::progress::{Antichain, Timestamp};
 
 use crate::async_runtime::IsolatedRuntime;
-use crate::batch::{BATCH_DELETE_ENABLED, BLOB_TARGET_SIZE, Batch, BatchBuilder, ProtoBatch};
+use crate::batch::{BATCH_DELETE_ENABLED, Batch, BatchBuilder, ProtoBatch};
 use crate::cache::{PersistClientCache, StateCache};
-use crate::cfg::{COMPACTION_MEMORY_BOUND_BYTES, PersistConfig};
+use crate::cfg::PersistConfig;
 use crate::critical::{CriticalReaderId, SinceHandle};
 use crate::error::InvalidUsage;
-use crate::fetch::{BatchFetcher, BatchFetcherConfig, FetchBatchFilter, Lease};
+use crate::fetch::{BatchFetcher, BatchFetcherConfig, Lease};
 use crate::internal::compact::Compactor;
 use crate::internal::encoding::parse_id;
 use crate::internal::gc::GarbageCollector;
 use crate::internal::machine::{Machine, retry_external};
 use crate::internal::state_versions::StateVersions;
-use crate::iter::{Consolidator, StructuredSort};
 use crate::metrics::Metrics;
 use crate::read::{
-    Cursor, CursorConsolidator, LazyPartStats, LeasedReaderId, READER_LEASE_DURATION, ReadHandle,
-    Since,
+    Cursor, LazyPartStats, LeasedReaderId, READER_LEASE_DURATION, ReadHandle, Since,
 };
 use crate::rpc::PubSubSender;
 use crate::schema::CaESchema;
@@ -643,57 +642,31 @@ impl PersistClient {
         should_fetch_part: impl for<'a> Fn(Option<&'a LazyPartStats>) -> bool,
     ) -> Result<Cursor<K, V, T, D>, Since<T>>
     where
-        K: Debug + Codec,
-        V: Debug + Codec,
+        K: Debug + Codec + Ord,
+        V: Debug + Codec + Ord,
         T: Timestamp + Lattice + Codec64 + Sync,
         D: Semigroup + Ord + Codec64 + Send + Sync,
     {
-        let context = format!("{}[as_of={:?}]", shard_id, as_of.elements());
-        let filter = FetchBatchFilter::Snapshot {
-            as_of: as_of.clone(),
-        };
-
+        // WIP!
+        let lease = Lease::new(SeqNo(0));
         let shard_metrics = self.metrics.shards.shard(&shard_id, "peek_stash");
 
-        let consolidator = {
-            let mut consolidator = Consolidator::new(
-                context,
-                shard_id,
-                StructuredSort::new(read_schemas.clone()),
-                Arc::clone(&self.blob),
-                Arc::clone(&self.metrics),
-                Arc::clone(&shard_metrics),
-                self.metrics.read.snapshot.clone(),
-                filter,
-                COMPACTION_MEMORY_BOUND_BYTES.get(&self.cfg),
-            );
-            for batch in batches {
-                for (meta, run) in batch.batch.runs() {
-                    consolidator.enqueue_run(
-                        &batch.batch.desc,
-                        meta,
-                        run.into_iter()
-                            .filter(|p| should_fetch_part(p.stats()))
-                            .cloned(),
-                    );
-                }
-            }
-            CursorConsolidator::Structured {
-                consolidator,
-                // This default may end up consolidating more records than previously
-                // for cases like fast-path peeks, where only the first few entries are used.
-                // If this is a noticeable performance impact, thread the max-len in from the caller.
-                max_len: self.cfg.compaction_yield_after_n_updates,
-                max_bytes: BLOB_TARGET_SIZE.get(&self.cfg).max(1),
-            }
-        };
+        let hollow_batches = batches.iter().map(|b| b.batch.clone()).collect_vec();
 
-        Ok(Cursor {
-            consolidator,
-            // WIP: What to do about this?
-            _lease: Lease::new(SeqNo(0)),
+        ReadHandle::read_batches_consolidated(
+            &self.cfg,
+            Arc::clone(&self.metrics),
+            shard_metrics,
+            self.metrics.read.snapshot.clone(),
+            Arc::clone(&self.blob),
+            lease,
+            shard_id,
+            as_of,
             read_schemas,
-        })
+            &hollow_batches,
+            should_fetch_part,
+        )
+        .await
     }
 
     /// Returns the requested schema, if known at the current state.

src/persist-client/src/read.rs

@@ -43,7 +43,7 @@ use crate::cfg::{COMPACTION_MEMORY_BOUND_BYTES, RetryParameters};
 use crate::fetch::{FetchBatchFilter, FetchedPart, Lease, LeasedBatchPart, fetch_leased_part};
 use crate::internal::encoding::Schemas;
 use crate::internal::machine::{ExpireFn, Machine};
-use crate::internal::metrics::Metrics;
+use crate::internal::metrics::{Metrics, ReadMetrics, ShardMetrics};
 use crate::internal::state::{BatchPart, HollowBatch};
 use crate::internal::watch::StateWatch;
 use crate::iter::{Consolidator, StructuredSort};
@@ -985,24 +985,53 @@ where
         should_fetch_part: impl for<'a> Fn(Option<&'a LazyPartStats>) -> bool,
     ) -> Result<Cursor<K, V, T, D>, Since<T>> {
         let batches = self.machine.snapshot(&as_of).await?;
+        let lease = self.lease_seqno();
+
+        Self::read_batches_consolidated(
+            &self.cfg,
+            Arc::clone(&self.metrics),
+            Arc::clone(&self.machine.applier.shard_metrics),
+            self.metrics.read.snapshot.clone(),
+            Arc::clone(&self.blob),
+            lease,
+            self.shard_id(),
+            as_of,
+            self.read_schemas.clone(),
+            &batches,
+            should_fetch_part,
+        )
+        .await
+    }
 
-        let context = format!("{}[as_of={:?}]", self.shard_id(), as_of.elements());
+    pub(crate) fn read_batches_consolidated(
+        persist_cfg: &PersistConfig,
+        metrics: Arc<Metrics>,
+        shard_metrics: Arc<ShardMetrics>,
+        read_metrics: ReadMetrics,
+        blob: Arc<dyn Blob>,
+        lease: Lease,
+        shard_id: ShardId,
+        as_of: Antichain<T>,
+        schemas: Schemas<K, V>,
+        batches: &[HollowBatch<T>],
+        should_fetch_part: impl for<'a> Fn(Option<&'a LazyPartStats>) -> bool,
+    ) -> Result<Cursor<K, V, T, D>, Since<T>> {
+        let context = format!("{}[as_of={:?}]", shard_id, as_of.elements());
         let filter = FetchBatchFilter::Snapshot {
             as_of: as_of.clone(),
         };
-        let lease = self.lease_seqno();
 
         let consolidator = {
             let mut consolidator = Consolidator::new(
                 context,
-                self.shard_id(),
-                StructuredSort::new(self.read_schemas.clone()),
-                Arc::clone(&self.blob),
-                Arc::clone(&self.metrics),
-                Arc::clone(&self.machine.applier.shard_metrics),
-                self.metrics.read.snapshot.clone(),
+                shard_id,
+                StructuredSort::new(schemas.clone()),
+                blob,
+                metrics,
+                shard_metrics,
+                read_metrics,
                 filter,
-                COMPACTION_MEMORY_BOUND_BYTES.get(&self.cfg),
+                COMPACTION_MEMORY_BOUND_BYTES.get(persist_cfg),
             );
             for batch in batches {
                 for (meta, run) in batch.runs() {
@@ -1020,15 +1049,15 @@ where
                 // This default may end up consolidating more records than previously
                 // for cases like fast-path peeks, where only the first few entries are used.
                 // If this is a noticeable performance impact, thread the max-len in from the caller.
-                max_len: self.cfg.compaction_yield_after_n_updates,
-                max_bytes: BLOB_TARGET_SIZE.get(&self.cfg).max(1),
+                max_len: persist_cfg.compaction_yield_after_n_updates,
+                max_bytes: BLOB_TARGET_SIZE.get(persist_cfg).max(1),
             }
         };
 
         Ok(Cursor {
             consolidator,
             _lease: lease,
-            read_schemas: self.read_schemas.clone(),
+            read_schemas: schemas,
         })
     }