feat(s3select): report uncompressed input byte metrics (#6865)

crates/s3select-api/src/query/dispatcher.rs

@@ -25,6 +25,8 @@ use super::{
     session::QueryAdmission,
 };
 
+pub type DispatchedQuery = (Query, Output);
+
 #[async_trait]
 pub trait QueryDispatcher: Send + Sync {
     // fn create_query_id(&self) -> QueryId;
@@ -41,6 +43,18 @@ pub trait QueryDispatcher: Send + Sync {
         self.execute_query(query).await
     }
 
+    async fn dispatch_query(&self, query: &Query) -> QueryResult<DispatchedQuery> {
+        let execution_query = query.for_execution();
+        let output = self.execute_query(&execution_query).await?;
+        Ok((execution_query, output))
+    }
+
+    async fn dispatch_query_admitted(&self, query: &Query, admission: QueryAdmission) -> QueryResult<DispatchedQuery> {
+        let execution_query = query.for_execution();
+        let output = self.execute_query_admitted(&execution_query, admission).await?;
+        Ok((execution_query, output))
+    }
+
     async fn build_logical_plan(&self, query_state_machine: Arc<QueryStateMachine>) -> QueryResult<Option<Plan>>;
 
     async fn execute_logical_plan(&self, logical_plan: Plan, query_state_machine: Arc<QueryStateMachine>) -> QueryResult<Output>;
@@ -53,3 +67,155 @@ pub trait QueryDispatcher: Send + Sync {
 
     // fn cancel_query(&self, id: &QueryId);
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::query::test_query;
+    use parking_lot::Mutex;
+
+    #[derive(Default)]
+    struct DefaultDispatchDispatcher {
+        executed_metrics: Mutex<Vec<Arc<crate::SelectInputMetrics>>>,
+    }
+
+    #[async_trait]
+    impl QueryDispatcher for DefaultDispatchDispatcher {
+        async fn execute_query(&self, query: &Query) -> QueryResult<Output> {
+            self.executed_metrics.lock().push(Arc::clone(query.input_metrics()));
+            Ok(Output::Nil(()))
+        }
+
+        async fn build_logical_plan(&self, _query_state_machine: Arc<QueryStateMachine>) -> QueryResult<Option<Plan>> {
+            unreachable!("default dispatch test does not plan queries")
+        }
+
+        async fn execute_logical_plan(
+            &self,
+            _logical_plan: Plan,
+            _query_state_machine: Arc<QueryStateMachine>,
+        ) -> QueryResult<Output> {
+            unreachable!("default dispatch test does not execute plans")
+        }
+
+        async fn build_query_state_machine(&self, _query: Query) -> QueryResult<Arc<QueryStateMachine>> {
+            unreachable!("default dispatch test does not build state machines")
+        }
+    }
+
+    #[derive(Default)]
+    struct DistinctAdmittedDispatcher {
+        plain_metrics: Mutex<Vec<Arc<crate::SelectInputMetrics>>>,
+        admitted_metrics: Mutex<Vec<Arc<crate::SelectInputMetrics>>>,
+        fail_plain: bool,
+        fail_admitted: bool,
+    }
+
+    #[async_trait]
+    impl QueryDispatcher for DistinctAdmittedDispatcher {
+        async fn execute_query(&self, query: &Query) -> QueryResult<Output> {
+            self.plain_metrics.lock().push(Arc::clone(query.input_metrics()));
+            if self.fail_plain {
+                Err(crate::QueryError::Cancel)
+            } else {
+                Ok(Output::Nil(()))
+            }
+        }
+
+        async fn execute_query_admitted(&self, query: &Query, _admission: QueryAdmission) -> QueryResult<Output> {
+            self.admitted_metrics.lock().push(Arc::clone(query.input_metrics()));
+            if self.fail_admitted {
+                Err(crate::QueryError::Cancel)
+            } else {
+                Ok(Output::Nil(()))
+            }
+        }
+
+        async fn build_logical_plan(&self, _query_state_machine: Arc<QueryStateMachine>) -> QueryResult<Option<Plan>> {
+            unreachable!("dispatch routing test does not plan queries")
+        }
+
+        async fn execute_logical_plan(
+            &self,
+            _logical_plan: Plan,
+            _query_state_machine: Arc<QueryStateMachine>,
+        ) -> QueryResult<Output> {
+            unreachable!("dispatch routing test does not execute plans")
+        }
+
+        async fn build_query_state_machine(&self, _query: Query) -> QueryResult<Arc<QueryStateMachine>> {
+            unreachable!("dispatch routing test does not build state machines")
+        }
+    }
+
+    #[tokio::test]
+    async fn plain_dispatch_propagates_override_errors() {
+        let dispatcher = DistinctAdmittedDispatcher {
+            fail_plain: true,
+            ..Default::default()
+        };
+
+        let error = match dispatcher.dispatch_query(&test_query()).await {
+            Err(error) => error,
+            Ok(_) => panic!("plain override error should propagate"),
+        };
+        assert!(matches!(error, crate::QueryError::Cancel));
+        assert_eq!(dispatcher.plain_metrics.lock().len(), 1);
+        assert!(dispatcher.admitted_metrics.lock().is_empty());
+    }
+
+    #[tokio::test]
+    async fn default_dispatch_methods_use_distinct_execution_metrics() {
+        let dispatcher = DefaultDispatchDispatcher::default();
+        let query = test_query();
+
+        let (first, _) = dispatcher
+            .dispatch_query(&query)
+            .await
+            .expect("first dispatch should execute");
+        let (second, _) = dispatcher
+            .dispatch_query(&query)
+            .await
+            .expect("second dispatch should execute");
+        let (admitted, _) = dispatcher
+            .dispatch_query_admitted(&query, QueryAdmission::unmanaged())
+            .await
+            .expect("admitted dispatch should execute");
+        let executed_metrics = dispatcher.executed_metrics.lock();
+
+        assert!(!Arc::ptr_eq(first.input_metrics(), second.input_metrics()));
+        assert!(!Arc::ptr_eq(first.input_metrics(), admitted.input_metrics()));
+        assert!(Arc::ptr_eq(first.input_metrics(), &executed_metrics[0]));
+        assert!(Arc::ptr_eq(second.input_metrics(), &executed_metrics[1]));
+        assert!(Arc::ptr_eq(admitted.input_metrics(), &executed_metrics[2]));
+    }
+
+    #[tokio::test]
+    async fn admitted_dispatch_uses_the_admitted_override_and_propagates_errors() {
+        let dispatcher = DistinctAdmittedDispatcher::default();
+        let query = test_query();
+
+        let (dispatched, _) = dispatcher
+            .dispatch_query_admitted(&query, QueryAdmission::unmanaged())
+            .await
+            .expect("admitted dispatch should execute through its override");
+        assert!(dispatcher.plain_metrics.lock().is_empty());
+        {
+            let admitted_metrics = dispatcher.admitted_metrics.lock();
+            assert_eq!(admitted_metrics.len(), 1);
+            assert!(Arc::ptr_eq(dispatched.input_metrics(), &admitted_metrics[0]));
+        }
+
+        let failing = DistinctAdmittedDispatcher {
+            fail_admitted: true,
+            ..Default::default()
+        };
+        let error = match failing.dispatch_query_admitted(&query, QueryAdmission::unmanaged()).await {
+            Err(error) => error,
+            Ok(_) => panic!("admitted override error should propagate"),
+        };
+        assert!(matches!(error, crate::QueryError::Cancel));
+        assert!(failing.plain_metrics.lock().is_empty());
+        assert_eq!(failing.admitted_metrics.lock().len(), 1);
+    }
+}

crates/s3select-api/src/query/mod.rs

@@ -15,7 +15,7 @@
 use s3s::dto::SelectObjectContentInput;
 use std::sync::Arc;
 
-use crate::SelectObjectSnapshot;
+use crate::{SelectInputMetrics, SelectObjectSnapshot};
 
 pub mod analyzer;
 pub mod ast;
@@ -40,6 +40,7 @@ pub struct Query {
     context: Context,
     content: String,
     snapshot: Option<Arc<SelectObjectSnapshot>>,
+    input_metrics: Arc<SelectInputMetrics>,
 }
 
 impl Query {
@@ -49,6 +50,7 @@ impl Query {
             context,
             content,
             snapshot: None,
+            input_metrics: Arc::new(SelectInputMetrics::default()),
         }
     }
 
@@ -58,6 +60,7 @@ impl Query {
             context,
             content,
             snapshot: Some(snapshot),
+            input_metrics: Arc::new(SelectInputMetrics::default()),
         }
     }
 
@@ -72,4 +75,46 @@ impl Query {
     pub fn snapshot(&self) -> Option<&Arc<SelectObjectSnapshot>> {
         self.snapshot.as_ref()
     }
+
+    pub fn input_metrics(&self) -> &Arc<SelectInputMetrics> {
+        &self.input_metrics
+    }
+
+    pub fn for_execution(&self) -> Self {
+        Self {
+            context: self.context.clone(),
+            content: self.content.clone(),
+            snapshot: self.snapshot.clone(),
+            input_metrics: Arc::new(SelectInputMetrics::default()),
+        }
+    }
+}
+
+#[cfg(test)]
+fn test_query() -> Query {
+    use s3s::dto::{CSVInput, CSVOutput, ExpressionType, InputSerialization, OutputSerialization, SelectObjectContentRequest};
+
+    let input = SelectObjectContentInput {
+        bucket: "bucket".to_string(),
+        expected_bucket_owner: None,
+        key: "input.csv".to_string(),
+        sse_customer_algorithm: None,
+        sse_customer_key: None,
+        sse_customer_key_md5: None,
+        request: SelectObjectContentRequest {
+            expression: "SELECT * FROM S3Object".to_string(),
+            expression_type: ExpressionType::from_static(ExpressionType::SQL),
+            input_serialization: InputSerialization {
+                csv: Some(CSVInput::default()),
+                ..Default::default()
+            },
+            output_serialization: OutputSerialization {
+                csv: Some(CSVOutput::default()),
+                ..Default::default()
+            },
+            request_progress: None,
+            scan_range: None,
+        },
+    };
+    Query::new(Context { input: Arc::new(input) }, "SELECT * FROM S3Object".to_string())
 }

crates/s3select-api/src/query/session.rs

@@ -12,9 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::SelectObjectSnapshot;
-use crate::query::Context;
+use crate::query::{Context, Query};
 use crate::{QueryError, QueryResult, object_store::EcObjectStore};
+use crate::{SelectInputMetrics, SelectObjectSnapshot};
 use datafusion::{
     arrow::{
         array::{Int32Array, StringArray},
@@ -314,7 +314,7 @@ impl SessionCtxFactory {
     }
 
     pub async fn create_session_ctx(&self, context: &Context) -> QueryResult<SessionCtx> {
-        self.create_session_ctx_inner(context, None, None, DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES)
+        self.create_session_ctx_inner(context, None, None, None, DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES)
             .await
     }
 
@@ -324,7 +324,7 @@ impl SessionCtxFactory {
         query_tracker: QueryExecutionTracker,
         memory_limit_bytes: usize,
     ) -> QueryResult<SessionCtx> {
-        self.create_session_ctx_inner(context, None, Some(query_tracker), memory_limit_bytes)
+        self.create_session_ctx_inner(context, None, Some(query_tracker), None, memory_limit_bytes)
             .await
     }
 
@@ -335,19 +335,36 @@ impl SessionCtxFactory {
         query_tracker: QueryExecutionTracker,
         memory_limit_bytes: usize,
     ) -> QueryResult<SessionCtx> {
-        self.create_session_ctx_inner(context, Some(snapshot), Some(query_tracker), memory_limit_bytes)
+        self.create_session_ctx_inner(context, Some(snapshot), Some(query_tracker), None, memory_limit_bytes)
             .await
     }
 
+    pub async fn create_session_ctx_for_query_with_tracker_and_memory_limit(
+        &self,
+        query: &Query,
+        query_tracker: QueryExecutionTracker,
+        memory_limit_bytes: usize,
+    ) -> QueryResult<SessionCtx> {
+        self.create_session_ctx_inner(
+            query.context(),
+            query.snapshot().cloned(),
+            Some(query_tracker),
+            Some(Arc::clone(query.input_metrics())),
+            memory_limit_bytes,
+        )
+        .await
+    }
+
     async fn create_session_ctx_inner(
         &self,
         context: &Context,
         snapshot: Option<Arc<SelectObjectSnapshot>>,
         query_tracker: Option<QueryExecutionTracker>,
+        input_metrics: Option<Arc<SelectInputMetrics>>,
         memory_limit_bytes: usize,
     ) -> QueryResult<SessionCtx> {
         let df_session_ctx = self
-            .build_df_session_context(context, snapshot, query_tracker.clone(), memory_limit_bytes)
+            .build_df_session_context(context, snapshot, query_tracker.clone(), input_metrics, memory_limit_bytes)
             .await?;
 
         Ok(SessionCtx {
@@ -362,6 +379,7 @@ impl SessionCtxFactory {
         context: &Context,
         snapshot: Option<Arc<SelectObjectSnapshot>>,
         query_tracker: Option<QueryExecutionTracker>,
+        input_metrics: Option<Arc<SelectInputMetrics>>,
         memory_limit_bytes: usize,
     ) -> QueryResult<SessionContext> {
         let path = format!("s3://{}", context.input.bucket);
@@ -383,7 +401,12 @@ impl SessionCtxFactory {
             .is_some_and(|delimiter| delimiter.len() == 2 && delimiter.as_bytes() != b"\r\n");
         let scan_range_requires_single_file_scan =
             context.input.request.scan_range.is_some() && context.input.request.input_serialization.parquet.is_none();
-        let config = if custom_two_byte_record_delimiter || scan_range_requires_single_file_scan {
+        let metered_input_requires_single_file_scan =
+            input_metrics.is_some() && context.input.request.input_serialization.parquet.is_none();
+        let config = if custom_two_byte_record_delimiter
+            || scan_range_requires_single_file_scan
+            || metered_input_requires_single_file_scan
+        {
             config.with_repartition_file_scans(false)
         } else {
             config
@@ -438,11 +461,16 @@ impl SessionCtxFactory {
 
             df_session_state.with_object_store(&store_url, store).build()
         } else {
+            let input_metrics = input_metrics.unwrap_or_else(|| Arc::new(SelectInputMetrics::default()));
             let store: EcObjectStore = match query_tracker {
-                Some(query_tracker) => {
-                    EcObjectStore::new_with_query_tracker(context.input.clone(), memory_pool, query_tracker, snapshot)
-                }
-                None => EcObjectStore::new_with_memory_pool(context.input.clone(), memory_pool, snapshot),
+                Some(query_tracker) => EcObjectStore::new_with_query_tracker(
+                    context.input.clone(),
+                    memory_pool,
+                    query_tracker,
+                    input_metrics,
+                    snapshot,
+                ),
+                None => EcObjectStore::new_with_memory_pool(context.input.clone(), memory_pool, input_metrics, snapshot),
             }
             .map_err(|err| QueryError::Datafusion {
                 source: Box::new(DataFusionError::External(Box::new(err))),
@@ -587,6 +615,31 @@ mod tests {
         assert!(session.inner().config().options().optimizer.repartition_file_scans);
     }
 
+    #[tokio::test]
+    async fn metered_csv_and_json_inputs_disable_file_repartitioning() {
+        let factory = SessionCtxFactory::new(true).with_target_partitions(3);
+        let csv_context = test_context();
+        let mut json_context = test_context();
+        let json_request = &mut Arc::make_mut(&mut json_context.input).request;
+        json_request.input_serialization.csv = None;
+        json_request.input_serialization.json = Some(JSONInput::default());
+
+        for context in [&csv_context, &json_context] {
+            let session = factory
+                .create_session_ctx_inner(
+                    context,
+                    None,
+                    None,
+                    Some(Arc::new(SelectInputMetrics::default())),
+                    DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES,
+                )
+                .await
+                .expect("metered session should be created");
+
+            assert!(!session.inner().config().options().optimizer.repartition_file_scans);
+        }
+    }
+
     #[tokio::test]
     async fn parquet_scan_range_keeps_file_repartitioning() {
         let mut context = test_context();
@@ -702,7 +755,7 @@ mod tests {
     async fn session_factory_applies_memory_limit() {
         let factory = SessionCtxFactory::new(true);
         let session = factory
-            .create_session_ctx_inner(&test_context(), None, None, 1024)
+            .create_session_ctx_inner(&test_context(), None, None, None, 1024)
             .await
             .expect("session should be created with a bounded memory pool");