feat(s3select): support compressed CSV and JSON input (#6915)

crates/s3select-api/Cargo.toml

@@ -60,21 +60,26 @@ hotpath-cpu = [
 [dependencies]
 hotpath.workspace = true
 metrics = { workspace = true }
+async-compression = { workspace = true, features = ["tokio", "gzip", "bzip2"] }
 async-trait.workspace = true
+arc-swap.workspace = true
 bytes = { workspace = true, features = ["serde"] }
 chrono = { workspace = true, features = ["serde"] }
+crc-fast.workspace = true
 rustfs-common.workspace = true
 datafusion = { workspace = true, default-features = false, features = ["parquet", "recursive_protection", "sql"] }
 rustfs-ecstore.workspace = true
 rustfs-storage-api.workspace = true
 futures = { workspace = true }
 futures-core = { workspace = true }
+flate2.workspace = true
 http.workspace = true
 s3s = { workspace = true, features = ["minio"] }
 serde_json = { workspace = true, features = ["raw_value"] }
 thiserror = { workspace = true }
 parking_lot.workspace = true
 tokio = { workspace = true, features = ["fs", "rt-multi-thread"] }
+tokio-stream.workspace = true
 tokio-util = { workspace = true, features = ["io", "compat"] }
 tracing.workspace = true
 uuid.workspace = true

crates/s3select-api/src/lib.rs

@@ -23,6 +23,7 @@ use datafusion::{
 use std::{error::Error as StdError, fmt::Display};
 use thiserror::Error;
 
+mod input_stream;
 mod metrics;
 pub mod object_store;
 pub mod query;
@@ -79,6 +80,9 @@ pub enum SelectError {
     #[error("The file is not in a supported compression format. Only GZIP and BZIP2 are supported.")]
     InvalidCompressionFormat,
 
+    #[error("{compression} is not applicable to the queried object. Please correct the request and try again.")]
+    InvalidCompressionFormatForObject { compression: &'static str },
+
     #[error("The data source type is not valid. Only CSV, JSON, and Parquet are supported.")]
     InvalidDataSource,
 
@@ -87,6 +91,9 @@ pub enum SelectError {
     )]
     TruncatedInput,
 
+    #[error("Scan range queries are not supported on this type of object.")]
+    UnsupportedScanRangeInput,
+
     #[error("An error occurred while parsing the CSV file. Check the file and try again.")]
     CsvParsingError,
 
@@ -96,6 +103,9 @@ pub enum SelectError {
     #[error("An error occurred while parsing the Parquet file. Check the file and try again.")]
     ParquetParsingError,
 
+    #[error("The length of a record in the input or result is greater than the maxCharsPerRecord limit of 1 MB.")]
+    OverMaxRecordSize,
+
     #[error("{message}")]
     ParseSelectFailure { message: String },
 

crates/s3select-api/src/metrics.rs

@@ -12,7 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::sync::atomic::{AtomicU64, Ordering};
+use arc_swap::ArcSwap;
+use std::sync::{
+    Arc,
+    atomic::{AtomicU64, Ordering},
+};
 
 #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
 pub struct SelectInputMetricsSnapshot {
@@ -20,33 +24,72 @@ pub struct SelectInputMetricsSnapshot {
     pub bytes_processed: u64,
 }
 
-#[derive(Debug, Default)]
+#[derive(Debug)]
 pub struct SelectInputMetrics {
+    active: ArcSwap<SelectInputMetricBank>,
+}
+
+#[derive(Debug, Default)]
+struct SelectInputMetricBank {
     uncompressed_bytes: AtomicU64,
+    compressed_bytes_scanned: AtomicU64,
+    compressed_bytes_processed: AtomicU64,
+}
+
+#[derive(Clone, Debug)]
+pub(crate) struct SelectInputMetricsRecorder {
+    bank: Arc<SelectInputMetricBank>,
+}
+
+impl Default for SelectInputMetrics {
+    fn default() -> Self {
+        Self {
+            active: ArcSwap::from_pointee(SelectInputMetricBank::default()),
+        }
+    }
 }
 
 impl SelectInputMetrics {
     pub fn snapshot(&self) -> SelectInputMetricsSnapshot {
-        let uncompressed_bytes = self.uncompressed_bytes.load(Ordering::Relaxed);
+        let bank = self.active.load();
+        let uncompressed_bytes = bank.uncompressed_bytes.load(Ordering::Relaxed);
         SelectInputMetricsSnapshot {
-            bytes_scanned: uncompressed_bytes,
-            bytes_processed: uncompressed_bytes,
+            bytes_scanned: uncompressed_bytes.saturating_add(bank.compressed_bytes_scanned.load(Ordering::Relaxed)),
+            bytes_processed: uncompressed_bytes.saturating_add(bank.compressed_bytes_processed.load(Ordering::Relaxed)),
         }
     }
 
-    pub(crate) fn record_uncompressed(&self, bytes: usize) {
-        let increment = u64::try_from(bytes).unwrap_or(u64::MAX);
-        let _ = self
-            .uncompressed_bytes
-            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| Some(current.saturating_add(increment)));
+    pub(crate) fn recorder(&self) -> SelectInputMetricsRecorder {
+        SelectInputMetricsRecorder {
+            bank: self.active.load_full(),
+        }
     }
 
-    /// Clears planner-only reads before query execution begins.
+    /// Publishes a fresh bank so late planner writes remain isolated.
     pub fn reset(&self) {
-        self.uncompressed_bytes.store(0, Ordering::Relaxed);
+        self.active.store(Arc::new(SelectInputMetricBank::default()));
     }
 }
 
+impl SelectInputMetricsRecorder {
+    pub(crate) fn record_uncompressed(&self, bytes: usize) {
+        saturating_add(&self.bank.uncompressed_bytes, bytes);
+    }
+
+    pub(crate) fn record_scanned(&self, bytes: usize) {
+        saturating_add(&self.bank.compressed_bytes_scanned, bytes);
+    }
+
+    pub(crate) fn record_processed(&self, bytes: usize) {
+        saturating_add(&self.bank.compressed_bytes_processed, bytes);
+    }
+}
+
+fn saturating_add(counter: &AtomicU64, bytes: usize) {
+    let increment = u64::try_from(bytes).unwrap_or(u64::MAX);
+    let _ = counter.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| Some(current.saturating_add(increment)));
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -54,7 +97,7 @@ mod tests {
     #[test]
     fn records_uncompressed_input_at_both_boundaries() {
         let metrics = SelectInputMetrics::default();
-        metrics.record_uncompressed(7);
+        metrics.recorder().record_uncompressed(7);
 
         assert_eq!(
             metrics.snapshot(),
@@ -68,21 +111,51 @@ mod tests {
     #[test]
     fn counters_saturate_instead_of_wrapping() {
         let metrics = SelectInputMetrics::default();
-        metrics.uncompressed_bytes.store(u64::MAX - 1, Ordering::Relaxed);
+        metrics
+            .active
+            .load()
+            .uncompressed_bytes
+            .store(u64::MAX - 1, Ordering::Relaxed);
 
-        metrics.record_uncompressed(2);
+        metrics.recorder().record_uncompressed(2);
 
         assert_eq!(metrics.snapshot().bytes_scanned, u64::MAX);
         assert_eq!(metrics.snapshot().bytes_processed, u64::MAX);
     }
 
+    #[test]
+    fn compressed_boundaries_are_counted_independently() {
+        let metrics = SelectInputMetrics::default();
+        let recorder = metrics.recorder();
+        recorder.record_scanned(39);
+        recorder.record_processed(19);
+
+        assert_eq!(
+            metrics.snapshot(),
+            SelectInputMetricsSnapshot {
+                bytes_scanned: 39,
+                bytes_processed: 19,
+            }
+        );
+    }
+
     #[test]
     fn reset_clears_schema_inference_bytes() {
         let metrics = SelectInputMetrics::default();
-        metrics.record_uncompressed(9);
+        let planning = metrics.recorder();
+        planning.record_uncompressed(9);
 
         metrics.reset();
+        planning.record_uncompressed(5);
+        let execution = metrics.recorder();
+        execution.record_uncompressed(3);
 
-        assert_eq!(metrics.snapshot(), SelectInputMetricsSnapshot::default());
+        assert_eq!(
+            metrics.snapshot(),
+            SelectInputMetricsSnapshot {
+                bytes_scanned: 3,
+                bytes_processed: 3,
+            }
+        );
     }
 }

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

@@ -30,6 +30,7 @@ use datafusion::{
     prelude::SessionContext,
 };
 use parking_lot::Mutex;
+use s3s::dto::CompressionType;
 use std::sync::{
     Arc, Weak,
     atomic::{AtomicU8, Ordering},
@@ -446,11 +447,19 @@ impl SessionCtxFactory {
         let scan_range_requires_single_file_scan =
             context.input.request.scan_range.is_some() && context.input.request.input_serialization.parquet.is_none();
         let json_document_requires_single_file_scan = is_json_document_input(&context.input);
+        let compressed_input_requires_single_file_scan = context
+            .input
+            .request
+            .input_serialization
+            .compression_type
+            .as_ref()
+            .is_some_and(|compression| compression.as_str() != CompressionType::NONE);
         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
             || json_document_requires_single_file_scan
+            || compressed_input_requires_single_file_scan
             || metered_input_requires_single_file_scan
         {
             config.with_repartition_file_scans(false)
@@ -847,6 +856,21 @@ mod tests {
         assert!(session.inner().config().options().optimizer.repartition_file_scans);
     }
 
+    #[tokio::test]
+    async fn compressed_input_disables_file_repartitioning_without_metrics() {
+        let mut context = test_context();
+        Arc::make_mut(&mut context.input).request.input_serialization.compression_type =
+            Some(CompressionType::from_static(CompressionType::GZIP));
+
+        let session = SessionCtxFactory::new(true)
+            .with_target_partitions(2)
+            .create_session_ctx(&context)
+            .await
+            .expect("compressed session should be created");
+
+        assert!(!session.inner().config().options().optimizer.repartition_file_scans);
+    }
+
     #[tokio::test]
     async fn json_document_disables_file_repartitioning() {
         let mut context = test_context();