fix(datafusion): address review findings

- fix: CAST unwrapping in classify_filter — reuse predicate::column_name
  so time-range predicates are correctly classified as Inexact and
  passed to scan(); previously CAST-wrapped filters were silently dropped
- fix: declare parquet sort order (metric_name, timestamp_secs ASC) on
  FileScanConfig so DataFusion avoids redundant sort operators
- fix: get_opts now respects GetOptions.range — dispatches to get_slice
  for Bounded/Suffix ranges instead of always downloading the full file
- fix: to_object_store_error propagates file path on NotFound
- fix: register_for_worker made a no-op; lazy scan-path registration is
  sufficient and avoids O(indexes) metastore RPCs per worker task;
  removes stale comment claiming a non-existent object-store cache
- fix: extract is_index_not_found helper, removing duplicated downcast
  block from try_consume_read_rel and create_default_table_provider
- fix: sort before dedup in QuickwitSchemaProvider::table_names
- fix: empty searcher pool returns Ok(vec![]) for local execution fallback
- fix: remove dead builder methods with_udf_batch, with_codec_applier,
  with_physical_optimizer_rule from DataSourceContributions
- feat: add tracing spans to execute_substrait and execute_sql
- feat: wire 4 GiB memory limit on DataFusionSessionBuilder in serve
- refactor: extract stream_to_receiver helper in gRPC handler

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

Author: alexanderbianchi

quickwit/Cargo.lock

@@ -40,7 +40,7 @@ datafusion-datasource = "52"
 datafusion-sql = "52"
 datafusion-physical-plan = "52"
 datafusion-datasource-parquet = "52"
-datafusion-distributed = { git = "https://github.com/datafusion-contrib/datafusion-distributed" }
+datafusion-distributed = { git = "https://github.com/datafusion-contrib/datafusion-distributed", rev = "0f2c8be3e148b0bd5c7f17b23f2df8bb1201d5fb" }
 object_store = "0.12"
 
 [dev-dependencies]

quickwit/quickwit-datafusion/Cargo.toml

@@ -88,7 +88,10 @@ impl SchemaProvider for QuickwitSchemaProvider {
                         names.append(&mut source_names);
                     }
                 }
-                // Deduplicate in case multiple sources claim the same name.
+                // Sort then deduplicate in case multiple sources claim the same name.
+                // `dedup()` only removes consecutive duplicates, so sorting first
+                // is required to remove all duplicates.
+                names.sort();
                 names.dedup();
                 names
             })

quickwit/quickwit-datafusion/src/catalog.rs

@@ -121,38 +121,12 @@ impl Default for DataSourceContributions {
 }
 
 impl DataSourceContributions {
-    /// Add a physical optimizer rule.
-    pub fn with_physical_optimizer_rule(
-        mut self,
-        rule: Arc<dyn PhysicalOptimizerRule + Send + Sync>,
-    ) -> Self {
-        self.physical_optimizer_rules.push(rule);
-        self
-    }
-
     /// Add a scalar UDF.
     pub fn with_udf(mut self, udf: Arc<ScalarUDF>) -> Self {
         self.udfs.push(udf);
         self
     }
 
-    /// Add multiple scalar UDFs at once.
-    pub fn with_udf_batch(mut self, udfs: impl IntoIterator<Item = Arc<ScalarUDF>>) -> Self {
-        self.udfs.extend(udfs);
-        self
-    }
-
-    /// Add a codec / builder-extension callback.
-    ///
-    /// Logs uses this to call `.with_distributed_user_codec(TantivyCodec)`.
-    pub fn with_codec_applier(
-        mut self,
-        f: impl FnOnce(SessionStateBuilder) -> SessionStateBuilder + Send + Sync + 'static,
-    ) -> Self {
-        self.codec_appliers.push(Box::new(f));
-        self
-    }
-
     pub(crate) fn udf_names(&self) -> Vec<String> {
         self.udfs.iter().map(|udf| udf.name().to_string()).collect()
     }

quickwit/quickwit-datafusion/src/data_source.rs

@@ -23,6 +23,7 @@ pub(crate) mod substrait;
 pub(crate) mod task_estimator;
 pub(crate) mod worker;
 
+pub use datafusion::execution::SendableRecordBatchStream;
 pub use resolver::QuickwitWorkerResolver;
 pub use service::DataFusionService;
 pub use session::DataFusionSessionBuilder;

quickwit/quickwit-datafusion/src/lib.rs

@@ -51,9 +51,11 @@ impl WorkerResolver for QuickwitWorkerResolver {
     fn get_urls(&self) -> Result<Vec<Url>, DataFusionError> {
         let addrs: Vec<SocketAddr> = self.searcher_pool.keys();
         if addrs.is_empty() {
-            return Err(DataFusionError::Execution(
-                "no searcher nodes available in the cluster".to_string(),
-            ));
+            // Empty pool means no searcher workers are registered (e.g. single-node
+            // local execution).  Return an empty list so the distributed optimizer
+            // sees zero workers and falls back to local execution rather than
+            // treating it as a hard error.
+            return Ok(vec![]);
         }
         let scheme = if self.use_tls { "https" } else { "http" };
         addrs

quickwit/quickwit-datafusion/src/resolver.rs

@@ -77,10 +77,12 @@ impl DataFusionService {
     /// decodes the plan, and returns a streaming `RecordBatch` iterator.
     /// The caller decides whether to collect, send via gRPC, or pipe to Arrow
     /// Flight — no materialization happens inside this method.
+    #[tracing::instrument(skip(self, plan_bytes), fields(plan_bytes_len = plan_bytes.len()))]
     pub async fn execute_substrait(
         &self,
         plan_bytes: &[u8],
     ) -> DFResult<SendableRecordBatchStream> {
+        tracing::info!(plan_bytes_len = plan_bytes.len(), "executing substrait plan");
         use datafusion_substrait::substrait::proto::Plan;
         use prost::Message;
 
@@ -128,7 +130,9 @@ impl DataFusionService {
     ///
     /// Returns an error if `sql` is empty after splitting, or if any statement
     /// fails to parse or execute.
+    #[tracing::instrument(skip(self, sql), fields(sql_len = sql.len()))]
     pub async fn execute_sql(&self, sql: &str) -> DFResult<SendableRecordBatchStream> {
+        tracing::info!(sql_len = sql.len(), "executing SQL query");
         let ctx = self.builder.build_session()?;
 
         // Split on `;` and discard empty fragments (trailing `;` etc.).

quickwit/quickwit-datafusion/src/service.rs

@@ -36,12 +36,11 @@
 //! variant), use `with_worker_resolver(resolver)` to supply any type that
 //! implements `datafusion_distributed::WorkerResolver`.
 //!
-//! ## Result materialization
+//! ## Result streaming
 //!
-//! `execute_substrait` collects all result batches into memory before returning.
-//! For large rollup queries this is unsuitable for production use.  A streaming
-//! variant is deferred; A downstream caller can wrap this via its own gRPC handler.
-//! Use `with_memory_limit()` to bound memory usage until streaming is in place.
+//! `execute_substrait` returns a `SendableRecordBatchStream` backed by the
+//! DataFusion streaming executor — no intermediate materialisation occurs.
+//! Use `with_memory_limit()` to bound query memory for large rollup queries.
 
 use std::collections::HashSet;
 use std::sync::Arc;
@@ -193,31 +192,31 @@ impl DataFusionSessionBuilder {
         Ok(())
     }
 
-    /// Execute a Substrait plan (protobuf bytes) and return the results.
+    /// Execute a Substrait plan (protobuf bytes) and return a streaming result.
     ///
     /// Builds a fresh session, converts the plan via `QuickwitSubstraitConsumer`,
-    /// and collects all results into memory.  See the module-level doc on
-    /// materialization limits.
+    /// and returns a `SendableRecordBatchStream` that the caller can consume
+    /// incrementally — no intermediate materialization occurs here.
     pub async fn execute_substrait(
         &self,
         plan_bytes: &[u8],
-    ) -> DFResult<Vec<arrow::array::RecordBatch>> {
+    ) -> DFResult<datafusion::physical_plan::SendableRecordBatchStream> {
         use datafusion_substrait::substrait::proto::Plan;
         use prost::Message;
 
         let plan = Plan::decode(plan_bytes)
             .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?;
 
         let ctx = self.build_session()?;
-        crate::substrait::execute_substrait_plan(&plan, &ctx, &self.sources).await
+        crate::substrait::execute_substrait_plan_streaming(&plan, &ctx, &self.sources).await
     }
 
     /// Build a `SessionContext` backed by the shared `RuntimeEnv`.
     ///
     /// Does NOT call `check_invariants()` — callers should invoke that once at
     /// startup, not on every query.
     pub fn build_session(&self) -> DFResult<SessionContext> {
-        let mut config = SessionConfig::new().with_target_partitions(1);
+        let mut config = SessionConfig::new();
         config.options_mut().catalog.default_catalog = "quickwit".to_string();
         config.options_mut().catalog.default_schema = "public".to_string();
         config.options_mut().catalog.information_schema = true;
@@ -231,23 +230,27 @@ impl DataFusionSessionBuilder {
             // at startup (via init) or lazily (via scan) are globally visible.
             .with_runtime_env(Arc::clone(&self.runtime));
 
+        // Accumulate contributions from all sources and apply them first, so that
+        // source-specific physical optimizer rules (e.g. tantivy pushdown) run
+        // before the distributed rule inspects the fully-optimized plan.
+        let mut combined = crate::data_source::DataSourceContributions::default();
+        for source in &self.sources {
+            combined.merge(source.contributions());
+        }
+        builder = combined.apply_to_builder(builder);
+
         if let Some(resolver) = &self.worker_resolver {
             // Clone the Arc so ownership passes into the distributed extension.
             // `Arc<dyn WorkerResolver>` implements `WorkerResolver` via deref,
             // so the forwarding wrapper is not needed.
+            // DistributedPhysicalOptimizerRule is added LAST so it sees the
+            // fully source-optimized plan.
             builder = builder
                 .with_distributed_worker_resolver(ArcWorkerResolver(Arc::clone(resolver)))
                 .with_distributed_task_estimator(QuickwitTaskEstimator)
                 .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule));
         }
 
-        // Accumulate contributions from all sources and apply them at once.
-        let mut combined = crate::data_source::DataSourceContributions::default();
-        for source in &self.sources {
-            combined.merge(source.contributions());
-        }
-        builder = combined.apply_to_builder(builder);
-
         let mut state = builder.build();
 
         for source in &self.sources {

quickwit/quickwit-datafusion/src/session.rs

@@ -98,20 +98,9 @@ impl MetricsSplitProvider for MetastoreSplitProvider {
 
 /// Convert a DataFusion `MetricsSplitQuery` to a metastore `ListMetricsSplitsQuery`.
 ///
-/// Note: The OSS parquet column names are bare (service, env, etc.) but the
-/// metastore `ListMetricsSplitsQuery` still uses the `tag_service`, `tag_env`
-/// field names — this is just the metastore's internal naming convention.
-///
-/// # Tag field pushdown limitation
-///
-/// `ListMetricsSplitsQuery` accepts at most one value per tag field
-/// (`Option<String>`). When a DataFusion `IN (...)` predicate produces
-/// multiple candidate values for a tag column, the metastore cannot express
-/// the full filter, so **no metastore-level pruning is applied for that
-/// dimension** — the value is left as `None`. The parquet-level filter
-/// (applied after the split is opened) will still enforce the predicate
-/// correctly. Only single-value equalities (`WHERE service = 'web'`) or
-/// single-element IN lists are pushed down to the metastore.
+/// Only metric name and time range are forwarded — the only dimensions the
+/// metastore reliably populates today. Tag-based pruning will be wired once
+/// the zonemap/bloom-filter mechanism is in place.
 fn to_metastore_query(index_uid: &IndexUid, query: &MetricsSplitQuery) -> ListMetricsSplitsQuery {
     let mut metastore_query = ListMetricsSplitsQuery::for_index(index_uid.clone());
 
@@ -127,27 +116,8 @@ fn to_metastore_query(index_uid: &IndexUid, query: &MetricsSplitQuery) -> ListMe
         metastore_query.time_range_end = Some(end as i64);
     }
 
-    // Push down a tag filter to the metastore only when there is exactly one
-    // candidate value. Multi-value IN lists cannot be expressed as a single
-    // `Option<String>` on `ListMetricsSplitsQuery`; passing only the first
-    // value would silently skip splits that match the other values, producing
-    // incorrect (incomplete) results. For multi-value lists we pass `None`
-    // (no metastore pruning) and rely on the parquet-level filter instead.
-    metastore_query.tag_service = single_value(query.tag_service.as_deref());
-    metastore_query.tag_env = single_value(query.tag_env.as_deref());
-    metastore_query.tag_datacenter = single_value(query.tag_datacenter.as_deref());
-    metastore_query.tag_region = single_value(query.tag_region.as_deref());
-    metastore_query.tag_host = single_value(query.tag_host.as_deref());
 
     metastore_query
 }
 
-/// Returns the single element of `values` as `Some(value)`, or `None` if
-/// `values` is absent, empty, or contains more than one element.
-fn single_value(values: Option<&[String]>) -> Option<String> {
-    match values {
-        Some([single]) => Some(single.clone()),
-        _ => None,
-    }
-}