Ensure we query one and only one partition per request (spiceai#9416)

* Ensure we query one and only one partition per request

* fix merge

* clippy

* clip

* clippy again??

* fix

* fix get_partitions_legacy

* clippy

* clippy

Entire-Checkpoint: 236c651af457

* clippy

* fix comments

* fix merge

* PR comments

* clippy

* docs

* Use in_memory partition manager by default + timeout update

* Fix table creation

* Fix cayenne_catalog::tests::test_create_table_detects

---------

Co-authored-by: Luke Kim <80174+lukekim@users.noreply.github.com>
Co-authored-by: Viktor Yershov <viktor@spice.ai>

Author: 3 people

crates/cayenne/src/cayenne_catalog.rs

@@ -2023,7 +2023,7 @@ mod tests {
         assert!(metadata.primary_key.is_empty());
         assert_eq!(metadata.table_id, table_id_1);
 
-        // Try to recreate with a primary key change — should return ChangedConfiguration error
+        // Now try to create with a primary key change — should get ChangedConfiguration error
         let options_changed = CreateTableOptions {
             table_name: "test_table".to_string(),
             schema: Arc::clone(&schema),

crates/object_store_occ/src/state.rs

@@ -77,6 +77,16 @@ pub struct ObjectState<T> {
     _marker: PhantomData<T>,
 }
 
+impl<T> std::fmt::Debug for ObjectState<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ObjectState")
+            .field("prefix", &self.prefix)
+            .field("store", &"Arc<dyn ObjectStore>")
+            .field("cache_size", &self.cache.read().len())
+            .finish_non_exhaustive()
+    }
+}
+
 impl<T> ObjectState<T>
 where
     T: Serialize + DeserializeOwned + Clone + Send + Sync,

crates/runtime-datafusion/src/analyzer_rule/partitioned_table_scan_rewrite.rs

@@ -39,7 +39,7 @@ pub trait TablePartitionProvider: Send + Sync + Debug {
     fn get_partitions(
         &self,
         table: &TableReference,
-        schema: SchemaRef,
+        schema: &SchemaRef,
     ) -> Vec<(Arc<dyn TableProvider>, Vec<Expr>)>;
 
     /// Whether partitioning should be applied to the given table.
@@ -104,7 +104,7 @@ impl AnalyzerRule for PartitionedTableScanRewrite {
 
             let providers = self
                 .partition_provider
-                .get_partitions(&scan.table_name, scan.source.schema());
+                .get_partitions(&scan.table_name, &scan.source.schema());
 
             tracing::debug!(
                 "PartitionedTableScanRewrite: {} partitions for '{}' table.",

crates/runtime/src/builder.rs

@@ -16,7 +16,9 @@ limitations under the License.
 
 use crate::cluster::DistributedNode;
 use crate::cluster::ExecutorRegistry;
+use crate::cluster::PartitionManager;
 use crate::cluster::ResolvedClusterConfig;
+use crate::cluster::partition;
 use crate::config::ClusterRole;
 use crate::config::Config;
 use crate::datafusion::udf::register_udfs;
@@ -225,19 +227,63 @@ impl RuntimeBuilder {
 
         // Create resource monitor early so it can be passed to DataFusion
         let resource_monitor = crate::resource_monitor::ResourceMonitor::new();
+        let secrets = Arc::new(RwLock::new(Self::load_secrets(self.app.as_ref()).await));
 
-        let distributed = match self
+        let distributed: Option<DistributedNode> = match self
             .resolved_cluster_config
             .as_ref()
             .and_then(ResolvedClusterConfig::effective_role)
         {
-            Some(ClusterRole::Scheduler) => Some(DistributedNode::Scheduler {
-                peers: Arc::new(RwLock::new(HashMap::new())),
-                // Initialized later when scheduler registry starts
-                job_executor: Arc::new(RwLock::new(None)),
-                executor_registry: Arc::new(ExecutorRegistry::new()),
-                partition_manager: Arc::new(RwLock::new(None)),
-            }),
+            Some(ClusterRole::Scheduler) => {
+                if let Some(scheduler_config) = self
+                    .app
+                    .as_ref()
+                    .and_then(|app| app.runtime.scheduler.clone())
+                {
+                    match partition::build_partition_metadata_store(
+                        io_runtime.clone(),
+                        Arc::clone(&secrets),
+                        &scheduler_config,
+                    )
+                    .await
+                    {
+                        Ok(store) => {
+                            let partition_manager = Arc::new(PartitionManager::new(store));
+
+                            Some(DistributedNode::Scheduler {
+                                peers: Arc::new(RwLock::new(HashMap::new())),
+                                // Initialized later when scheduler registry starts
+                                job_executor: Arc::new(RwLock::new(None)),
+                                executor_registry: Arc::new(ExecutorRegistry::new(Arc::clone(
+                                    &partition_manager,
+                                ))),
+                                partition_manager,
+                            })
+                        }
+                        Err(e) => {
+                            tracing::error!(
+                                "Failed to initialize partition metadata store for scheduler: {e}"
+                            );
+                            None
+                        }
+                    }
+                } else {
+                    tracing::warn!(
+                        "'--role scheduler' was specified but no `runtime.scheduler` field was found in spicepod.yaml. Using in-memory partition store."
+                    );
+                    let partition_manager = Arc::new(PartitionManager::new(Arc::new(
+                        object_store::memory::InMemory::new(),
+                    )));
+                    Some(DistributedNode::Scheduler {
+                        peers: Arc::new(RwLock::new(HashMap::new())),
+                        job_executor: Arc::new(RwLock::new(None)),
+                        executor_registry: Arc::new(ExecutorRegistry::new(Arc::clone(
+                            &partition_manager,
+                        ))),
+                        partition_manager,
+                    })
+                }
+            }
             Some(ClusterRole::Executor) => Some(DistributedNode::Executor {
                 partition_assignments: Arc::new(RwLock::new(HashMap::new())),
             }),
@@ -297,8 +343,6 @@ impl RuntimeBuilder {
             None
         };
 
-        let secrets = Self::load_secrets(self.app.as_ref()).await;
-
         let evals = self
             .app
             .as_ref()
@@ -318,7 +362,7 @@ impl RuntimeBuilder {
             tools: Arc::new(RwLock::new(HashMap::new())),
             tool_factories: Arc::new(Mutex::new(HashMap::new())),
             pods_watcher: Arc::new(RwLock::new(self.pods_watcher)),
-            secrets: Arc::new(RwLock::new(secrets)),
+            secrets,
             spaced_tracer: Arc::new(tracers::SpacedTracer::new(Duration::from_secs(15))),
             autoload_extensions: Arc::new(self.autoload_extensions),
             extensions: Arc::new(RwLock::new(HashMap::new())),