adapter: Frontend peek sequencing -- statistics_oracle

Author: ggevay

src/adapter/src/coord/sequencer.rs

@@ -12,7 +12,7 @@
 
 //! Logic for executing a planned SQL query.
 
-use std::collections::BTreeSet;
+use std::collections::{BTreeMap, BTreeSet};
 use std::str::FromStr;
 use std::sync::Arc;
 
@@ -40,17 +40,21 @@ use mz_sql::plan::{
 };
 use mz_sql::rbac;
 use mz_sql::session::metadata::SessionMetadata;
+use mz_sql::session::vars;
 use mz_sql::session::vars::SessionVars;
 use mz_sql_parser::ast::{Raw, Statement};
 use mz_storage_client::client::TableData;
 use mz_storage_client::storage_collections::StorageCollections;
 use mz_storage_types::connections::inline::IntoInlineConnection;
+use mz_storage_types::controller::StorageError;
 use mz_storage_types::stats::RelationPartStats;
 use mz_transform::dataflow::DataflowMetainfo;
 use mz_transform::notice::{OptimizerNoticeApi, OptimizerNoticeKind, RawOptimizerNotice};
+use mz_transform::{EmptyStatisticsOracle, StatisticsOracle};
 use timely::progress::Antichain;
+use timely::progress::Timestamp as TimelyTimestamp;
 use tokio::sync::oneshot;
-use tracing::{Instrument, Level, Span, event};
+use tracing::{Instrument, Level, Span, event, warn};
 
 use crate::ExecuteContext;
 use crate::catalog::{Catalog, CatalogState};
@@ -1205,3 +1209,88 @@ pub(crate) async fn explain_plan_inner(
 
     Ok(rows)
 }
+
+/// Creates a statistics oracle for query optimization.
+///
+/// This is a free-standing function that can be called from both the old peek sequencing
+/// and the new frontend peek sequencing.
+pub(crate) async fn statistics_oracle(
+    session: &Session,
+    source_ids: &BTreeSet<GlobalId>,
+    query_as_of: &Antichain<Timestamp>,
+    is_oneshot: bool,
+    system_config: &vars::SystemVars,
+    storage_collections: &dyn StorageCollections<Timestamp = Timestamp>,
+) -> Result<Box<dyn StatisticsOracle>, AdapterError> {
+    if !session.vars().enable_session_cardinality_estimates() {
+        return Ok(Box::new(EmptyStatisticsOracle));
+    }
+
+    let timeout = if is_oneshot {
+        // TODO(mgree): ideally, we would shorten the timeout even more if we think the query could take the fast path
+        system_config.optimizer_oneshot_stats_timeout()
+    } else {
+        system_config.optimizer_stats_timeout()
+    };
+
+    let cached_stats = mz_ore::future::timeout(
+        timeout,
+        CachedStatisticsOracle::new(source_ids, query_as_of, storage_collections),
+    )
+    .await;
+
+    match cached_stats {
+        Ok(stats) => Ok(Box::new(stats)),
+        Err(mz_ore::future::TimeoutError::DeadlineElapsed) => {
+            warn!(
+                is_oneshot = is_oneshot,
+                "optimizer statistics collection timed out after {}ms",
+                timeout.as_millis()
+            );
+
+            Ok(Box::new(EmptyStatisticsOracle))
+        }
+        Err(mz_ore::future::TimeoutError::Inner(e)) => Err(AdapterError::Storage(e)),
+    }
+}
+
+#[derive(Debug)]
+struct CachedStatisticsOracle {
+    cache: BTreeMap<GlobalId, usize>,
+}
+
+impl CachedStatisticsOracle {
+    pub async fn new<T: TimelyTimestamp>(
+        ids: &BTreeSet<GlobalId>,
+        as_of: &Antichain<T>,
+        storage_collections: &dyn mz_storage_client::storage_collections::StorageCollections<Timestamp = T>,
+    ) -> Result<Self, StorageError<T>> {
+        let mut cache = BTreeMap::new();
+
+        for id in ids {
+            let stats = storage_collections.snapshot_stats(*id, as_of.clone()).await;
+
+            match stats {
+                Ok(stats) => {
+                    cache.insert(*id, stats.num_updates);
+                }
+                Err(StorageError::IdentifierMissing(id)) => {
+                    ::tracing::debug!("no statistics for {id}")
+                }
+                Err(e) => return Err(e),
+            }
+        }
+
+        Ok(Self { cache })
+    }
+}
+
+impl StatisticsOracle for CachedStatisticsOracle {
+    fn cardinality_estimate(&self, id: GlobalId) -> Option<usize> {
+        self.cache.get(&id).map(|estimate| *estimate)
+    }
+
+    fn as_map(&self) -> BTreeMap<GlobalId, usize> {
+        self.cache.clone()
+    }
+}

src/adapter/src/coord/sequencer/inner.rs

@@ -86,11 +86,9 @@ use mz_storage_client::controller::{CollectionDescription, DataSource, ExportDes
 use mz_storage_types::AlterCompatible;
 use mz_storage_types::connections::inline::IntoInlineConnection;
 use mz_storage_types::controller::StorageError;
-use mz_transform::EmptyStatisticsOracle;
 use mz_transform::dataflow::DataflowMetainfo;
 use smallvec::SmallVec;
 use timely::progress::Antichain;
-use timely::progress::Timestamp as TimelyTimestamp;
 use tokio::sync::{oneshot, watch};
 use tracing::{Instrument, Span, info, warn};
 
@@ -5124,93 +5122,23 @@ impl Coordinator {
 
         Ok(ExecuteResponse::AlteredObject(ObjectType::Table))
     }
-}
-
-#[derive(Debug)]
-struct CachedStatisticsOracle {
-    cache: BTreeMap<GlobalId, usize>,
-}
-
-impl CachedStatisticsOracle {
-    pub async fn new<T: TimelyTimestamp>(
-        ids: &BTreeSet<GlobalId>,
-        as_of: &Antichain<T>,
-        storage_collections: &dyn mz_storage_client::storage_collections::StorageCollections<Timestamp = T>,
-    ) -> Result<Self, StorageError<T>> {
-        let mut cache = BTreeMap::new();
-
-        for id in ids {
-            let stats = storage_collections.snapshot_stats(*id, as_of.clone()).await;
 
-            match stats {
-                Ok(stats) => {
-                    cache.insert(*id, stats.num_updates);
-                }
-                Err(StorageError::IdentifierMissing(id)) => {
-                    ::tracing::debug!("no statistics for {id}")
-                }
-                Err(e) => return Err(e),
-            }
-        }
-
-        Ok(Self { cache })
-    }
-}
-
-impl mz_transform::StatisticsOracle for CachedStatisticsOracle {
-    fn cardinality_estimate(&self, id: GlobalId) -> Option<usize> {
-        self.cache.get(&id).map(|estimate| *estimate)
-    }
-
-    fn as_map(&self) -> BTreeMap<GlobalId, usize> {
-        self.cache.clone()
-    }
-}
-
-impl Coordinator {
     pub(super) async fn statistics_oracle(
         &self,
         session: &Session,
         source_ids: &BTreeSet<GlobalId>,
         query_as_of: &Antichain<Timestamp>,
         is_oneshot: bool,
     ) -> Result<Box<dyn mz_transform::StatisticsOracle>, AdapterError> {
-        if !session.vars().enable_session_cardinality_estimates() {
-            return Ok(Box::new(EmptyStatisticsOracle));
-        }
-
-        let timeout = if is_oneshot {
-            // TODO(mgree): ideally, we would shorten the timeout even more if we think the query could take the fast path
-            self.catalog()
-                .system_config()
-                .optimizer_oneshot_stats_timeout()
-        } else {
-            self.catalog().system_config().optimizer_stats_timeout()
-        };
-
-        let cached_stats = mz_ore::future::timeout(
-            timeout,
-            CachedStatisticsOracle::new(
-                source_ids,
-                query_as_of,
-                self.controller.storage_collections.as_ref(),
-            ),
+        super::statistics_oracle(
+            session,
+            source_ids,
+            query_as_of,
+            is_oneshot,
+            self.catalog().system_config(),
+            self.controller.storage_collections.as_ref(),
         )
-        .await;
-
-        match cached_stats {
-            Ok(stats) => Ok(Box::new(stats)),
-            Err(mz_ore::future::TimeoutError::DeadlineElapsed) => {
-                warn!(
-                    is_oneshot = is_oneshot,
-                    "optimizer statistics collection timed out after {}ms",
-                    timeout.as_millis()
-                );
-
-                Ok(Box::new(EmptyStatisticsOracle))
-            }
-            Err(mz_ore::future::TimeoutError::Inner(e)) => Err(AdapterError::Storage(e)),
-        }
+        .await
     }
 }
 

src/adapter/src/frontend_peek.rs

@@ -36,7 +36,7 @@ use tracing_opentelemetry::OpenTelemetrySpanExt;
 use crate::catalog::CatalogState;
 use crate::command::Command;
 use crate::coord::peek::PeekPlan;
-use crate::coord::sequencer::eval_copy_to_uri;
+use crate::coord::sequencer::{eval_copy_to_uri, statistics_oracle};
 use crate::coord::timestamp_selection::TimestampDetermination;
 use crate::coord::{Coordinator, CopyToContext, ExplainContext, ExplainPlanContext, TargetCluster};
 use crate::explain::insights::PlanInsightsContext;
@@ -515,14 +515,16 @@ impl PeekClient {
 
         // # From peek_optimize
 
-        if session.vars().enable_session_cardinality_estimates() {
-            debug!(
-                "Bailing out from try_frontend_peek_inner, because of enable_session_cardinality_estimates"
-            );
-            return Ok(None);
-        }
-        // TODO(peek-seq): wire up statistics
-        let stats = Box::new(EmptyStatisticsOracle);
+        let stats = statistics_oracle(
+            session,
+            &source_ids,
+            &determination.timestamp_context.antichain(),
+            true,
+            catalog.system_config(),
+            &*self.storage_collections,
+        )
+        .await
+        .unwrap_or_else(|_| Box::new(EmptyStatisticsOracle));
 
         // Generate data structures that can be moved to another task where we will perform possibly
         // expensive optimizations.

src/storage-client/src/storage_collections.rs

@@ -88,7 +88,7 @@ mod metrics;
 /// - Hands out [ReadHolds](ReadHold) that prevent a collection's since from
 /// advancing while it needs to be read at a specific time.
 #[async_trait]
-pub trait StorageCollections: Debug {
+pub trait StorageCollections: Debug + Sync {
     type Timestamp: TimelyTimestamp;
 
     /// On boot, reconcile this [StorageCollections] with outside state. We get