fix(select): admit legacy and linked LZ4 working memory

Cargo.lock

@@ -10759,6 +10759,7 @@ dependencies = [
  "arc-swap",
  "async-compression",
  "async-trait",
+ "brotli 9.1.1",
  "bytes",
  "chrono",
  "crc-fast",
@@ -10768,6 +10769,8 @@ dependencies = [
  "futures-core",
  "hotpath",
  "http 1.5.0",
+ "lz4",
+ "lz4_flex",
  "memchr",
  "metrics",
  "parking_lot 0.13.0",

Cargo.toml

@@ -296,6 +296,7 @@ local-ip-address = "0.6.13"
 log = "0.4"
 memmap2 = "0.9.12"
 lz4 = "1.28.1"
+lz4_flex = { version = "=0.14.0", default-features = false, features = ["std", "frame"] }
 matchit = "0.9.2"
 md-5 = "0.11.0"
 mime_guess = "2.0.5"

crates/s3select-api/Cargo.toml

@@ -73,6 +73,9 @@ rustfs-storage-api.workspace = true
 futures = { workspace = true }
 futures-core = { workspace = true }
 flate2.workspace = true
+brotli.workspace = true
+lz4.workspace = true
+lz4_flex.workspace = true
 http.workspace = true
 memchr.workspace = true
 s3s = { workspace = true, features = ["minio"] }

crates/s3select-api/src/lib.rs

@@ -27,6 +27,7 @@ mod csv_input;
 mod input_stream;
 mod metrics;
 pub mod object_store;
+pub mod parquet_budget;
 pub mod query;
 pub mod server;
 mod storage_api;
@@ -218,6 +219,14 @@ fn classify_select_error_source(err: &(dyn StdError + 'static)) -> Option<Select
     if let Some(error) = err.downcast_ref::<SelectError>() {
         return Some(error.clone());
     }
+    if let Some(error) = err.downcast_ref::<parquet_budget::ParquetBudgetError>() {
+        return Some(match error {
+            parquet_budget::ParquetBudgetError::Exceeded => SelectError::ResourceExhausted,
+            parquet_budget::ParquetBudgetError::Invalid(_) | parquet_budget::ParquetBudgetError::Codec(_) => {
+                SelectError::ParquetParsingError
+            }
+        });
+    }
     if let Some(error) = err.downcast_ref::<object_store::SelectObjectStoreError>() {
         return Some(error.select_error());
     }
@@ -332,9 +341,32 @@ impl Display for ResolvedTable {
 #[cfg(test)]
 mod tests {
     use super::*;
+
     use datafusion::common::DataFusionError;
     use datafusion::sql::sqlparser::parser::ParserError;
 
+    #[test]
+    fn parquet_admission_errors_preserve_select_error_codes() {
+        for (source, expected) in [
+            (parquet_budget::ParquetBudgetError::Exceeded, SelectError::ResourceExhausted),
+            (
+                parquet_budget::ParquetBudgetError::Invalid("malformed footer".into()),
+                SelectError::ParquetParsingError,
+            ),
+            (
+                parquet_budget::ParquetBudgetError::Codec(std::io::Error::from(std::io::ErrorKind::InvalidData)),
+                SelectError::ParquetParsingError,
+            ),
+        ] {
+            let error = QueryError::from(DataFusionError::External(Box::new(source)));
+            assert_eq!(error.select_error(), expected);
+        }
+        let nested = QueryError::from(DataFusionError::ParquetError(Box::new(ParquetError::External(Box::new(
+            parquet_budget::ParquetBudgetError::Exceeded,
+        )))));
+        assert_eq!(nested.select_error(), SelectError::ResourceExhausted);
+    }
+
     #[test]
     fn test_query_error_display() {
         let err = QueryError::NotImplemented {

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

@@ -424,7 +424,8 @@ impl SimpleQueryDispatcher {
         tracker: &QueryExecutionTracker,
     ) -> QueryResult<MetadataProvider> {
         if self.input.request.input_serialization.parquet.is_some() {
-            let provider = ParquetSelectTable::try_new(session.inner(), self.input.as_ref()).await?;
+            let query_guard = tracker.query_guard().ok_or(QueryError::Cancel)?;
+            let provider = ParquetSelectTable::try_new(session.inner(), self.input.as_ref(), Some(query_guard)).await?;
             let current_session_table_provider = self.build_table_handle_provider()?;
             let metadata_provider =
                 MetadataProvider::new(provider, current_session_table_provider, self.func_manager.clone(), session.clone());

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

@@ -27,7 +27,10 @@ use datafusion::{
     datasource::{
         TableProvider,
         listing::{ListingTableUrl, PartitionedFile},
-        physical_plan::{FileScanConfigBuilder, ParquetSource, parquet::ParquetAccessPlan},
+        physical_plan::{
+            FileScanConfigBuilder, ParquetSource,
+            parquet::{ParquetAccessPlan, ParquetFileReaderFactory},
+        },
         source::DataSourceExec,
     },
     execution::{memory_pool::MemoryLimit, object_store::ObjectStoreUrl},
@@ -40,6 +43,7 @@ use datafusion::{
     },
     physical_plan::{
         ExecutionPlan, SendableRecordBatchStream,
+        metrics::ExecutionPlanMetricsSet,
         stream::RecordBatchStreamAdapter,
         streaming::{PartitionStream, StreamingTableExec},
     },
@@ -48,6 +52,8 @@ use futures::{FutureExt, TryFutureExt, TryStreamExt, future::BoxFuture, stream};
 use rustfs_s3select_api::{
     QueryResult,
     object_store::{SelectScanRange, scan_range_from_bounds},
+    parquet_budget::{ParquetBudgetError, page_decode_budget, validate_footer},
+    query::session::QueryExecutionGuard,
 };
 use s3s::dto::SelectObjectContentInput;
 
@@ -59,6 +65,7 @@ pub struct ParquetSelectTable {
     object_size: u64,
     access_plan: Option<Arc<ParquetAccessPlan>>,
     expanded_row_budget: usize,
+    reader_factory: Arc<BudgetParquetReaderFactory>,
 }
 
 impl fmt::Debug for ParquetSelectTable {
@@ -72,11 +79,30 @@ impl fmt::Debug for ParquetSelectTable {
     }
 }
 
+#[derive(Debug)]
 struct ObjectStoreParquetReader {
     store: Arc<dyn ObjectStore>,
     path: Path,
     file_size: u64,
     metadata_read_limit: usize,
+    metadata_decode_limit: usize,
+    metadata: Option<Arc<ParquetMetaData>>,
+    query_guard: Option<QueryExecutionGuard>,
+}
+
+async fn validate_in_worker<T: Send + 'static>(
+    query_guard: Option<QueryExecutionGuard>,
+    operation: impl FnOnce() -> Result<T, ParquetBudgetError> + Send + 'static,
+) -> ParquetResult<T> {
+    tokio::task::spawn_blocking(move || {
+        // Cancellation cannot release admission while the decoder preflight
+        // still owns input buffers or codec allocations.
+        let _query_guard = query_guard;
+        operation()
+    })
+    .await
+    .map_err(|error| ParquetError::External(Box::new(error)))?
+    .map_err(|error| ParquetError::External(Box::new(error)))
 }
 
 impl AsyncFileReader for ObjectStoreParquetReader {
@@ -94,7 +120,16 @@ impl AsyncFileReader for ObjectStoreParquetReader {
         if total.is_none_or(|total| total > self.metadata_read_limit as u64) {
             return async { Err(ParquetError::External(Box::new(parquet_budget_error()))) }.boxed();
         }
-        async move { self.store.get_ranges(&self.path, &ranges).await.map_err(parquet_store_error) }.boxed()
+        async move {
+            // Read requested ranges directly so sparse offsets cannot amplify
+            // a small request into a large coalesced allocation.
+            let mut bytes = Vec::with_capacity(ranges.len());
+            for range in ranges {
+                bytes.push(self.store.get_range(&self.path, range).await.map_err(parquet_store_error)?);
+            }
+            Ok(bytes)
+        }
+        .boxed()
     }
 
     fn get_metadata<'a>(
@@ -102,6 +137,27 @@ impl AsyncFileReader for ObjectStoreParquetReader {
         options: Option<&'a ArrowReaderOptions>,
     ) -> BoxFuture<'a, ParquetResult<Arc<ParquetMetaData>>> {
         async move {
+            if let Some(metadata) = self.metadata.as_ref() {
+                return Ok(Arc::clone(metadata));
+            }
+            let tail_start = self
+                .file_size
+                .checked_sub(8)
+                .ok_or_else(|| ParquetError::EOF("Parquet footer is truncated".into()))?;
+            let tail = self.get_bytes(tail_start..self.file_size).await?;
+            let tail: &[u8; 8] = tail
+                .as_ref()
+                .try_into()
+                .map_err(|_| ParquetError::EOF("Parquet footer is truncated".into()))?;
+            let footer = datafusion::parquet::file::metadata::FooterTail::try_new(tail)?;
+            let length =
+                u64::try_from(footer.metadata_length()).map_err(|_| ParquetError::External(Box::new(parquet_budget_error())))?;
+            let start = tail_start
+                .checked_sub(length)
+                .ok_or_else(|| ParquetError::EOF("Parquet metadata is truncated".into()))?;
+            let bytes = self.get_bytes(start..tail_start).await?;
+            let limit = self.metadata_decode_limit;
+            validate_in_worker(self.query_guard.clone(), move || validate_footer(&bytes, limit)).await?;
             let metadata_options = options.map(|options| options.metadata_options().clone());
             let mut metadata_reader = ParquetMetaDataReader::new().with_metadata_options(metadata_options);
 
@@ -119,8 +175,84 @@ impl AsyncFileReader for ObjectStoreParquetReader {
     }
 }
 
+#[derive(Debug)]
+struct BudgetParquetReaderFactory {
+    store: Arc<dyn ObjectStore>,
+    path: Path,
+    file_size: u64,
+    metadata: Arc<ParquetMetaData>,
+    read_limit: usize,
+}
+
+impl ParquetFileReaderFactory for BudgetParquetReaderFactory {
+    fn create_reader(
+        &self,
+        _partition_index: usize,
+        file: PartitionedFile,
+        _metadata_size_hint: Option<usize>,
+        _metrics: &ExecutionPlanMetricsSet,
+    ) -> DFResult<Box<dyn AsyncFileReader + Send>> {
+        if file.object_meta.location != self.path || file.object_meta.size != self.file_size {
+            return Err(DataFusionError::Execution("Parquet reader input differs from its admitted object".into()));
+        }
+        Ok(Box::new(ObjectStoreParquetReader {
+            store: Arc::clone(&self.store),
+            path: self.path.clone(),
+            file_size: self.file_size,
+            metadata_read_limit: self.read_limit,
+            metadata_decode_limit: self.read_limit,
+            metadata: Some(Arc::clone(&self.metadata)),
+            query_guard: None,
+        }))
+    }
+}
+
+async fn validate_parquet_pages(
+    store: &Arc<dyn ObjectStore>,
+    path: &Path,
+    file_size: u64,
+    metadata: &ParquetMetaData,
+    plan: Option<&ParquetAccessPlan>,
+    limit: usize,
+    query_guard: &Option<QueryExecutionGuard>,
+) -> DFResult<()> {
+    for (index, group) in metadata.row_groups().iter().enumerate() {
+        if plan.is_some_and(|plan| !plan.should_scan(index)) {
+            continue;
+        }
+        let mut resident = 0usize;
+        for column in group.columns() {
+            let start = u64::try_from(column.dictionary_page_offset().unwrap_or(column.data_page_offset()))
+                .map_err(|_| parquet_budget_error())?;
+            let length = u64::try_from(column.compressed_size()).map_err(|_| parquet_budget_error())?;
+            let end = start
+                .checked_add(length)
+                .filter(|end| *end <= file_size)
+                .ok_or_else(parquet_budget_error)?;
+            if length > u64::try_from(limit).map_err(|_| parquet_budget_error())? {
+                return Err(parquet_budget_error());
+            }
+            let bytes = store.get_range(path, start..end).await?;
+            let declared = usize::try_from(column.uncompressed_size()).map_err(|_| parquet_budget_error())?;
+            let codec = column.compression();
+            let required = validate_in_worker(query_guard.clone(), move || page_decode_budget(&bytes, declared, limit, codec))
+                .await
+                .map_err(DataFusionError::from)?;
+            resident = resident
+                .checked_add(required)
+                .filter(|total| *total <= limit)
+                .ok_or_else(parquet_budget_error)?;
+        }
+    }
+    Ok(())
+}
+
 impl ParquetSelectTable {
-    pub async fn try_new(state: &dyn Session, input: &SelectObjectContentInput) -> QueryResult<Arc<dyn TableProvider>> {
+    pub async fn try_new(
+        state: &dyn Session,
+        input: &SelectObjectContentInput,
+        query_guard: Option<QueryExecutionGuard>,
+    ) -> QueryResult<Arc<dyn TableProvider>> {
         let table_path = ListingTableUrl::parse(format!("s3://{}/{}", input.bucket, input.key))?;
         let object_store_url = table_path.object_store();
         let object_location = Path::from(input.key.clone());
@@ -136,9 +268,12 @@ impl ParquetSelectTable {
         };
         let reader = ObjectStoreParquetReader {
             store: Arc::clone(&store),
-            path: object_location,
+            path: object_location.clone(),
             file_size: object_meta.size,
             metadata_read_limit: (memory_budget / 16).max(1),
+            metadata_decode_limit: (memory_budget / 4).max(1),
+            metadata: None,
+            query_guard: query_guard.clone(),
         };
         let builder = ParquetRecordBatchStreamBuilder::new(reader)
             .await
@@ -153,11 +288,26 @@ impl ParquetSelectTable {
             builder.schema().metadata().clone(),
         ));
         let metadata = Arc::clone(builder.metadata());
-        // This is admission for valid Parquet declarations, not a replacement
-        // for limits inside the footer parser or native compression codecs.
         let expanded_row_budget = (memory_budget / 16).max(1);
         let access_plan = parquet_access_plan(input, object_meta.size, metadata.as_ref())?;
         admit_parquet_decode(metadata.as_ref(), access_plan.as_deref(), memory_budget / 4)?;
+        validate_parquet_pages(
+            &store,
+            &object_location,
+            object_meta.size,
+            &metadata,
+            access_plan.as_deref(),
+            memory_budget / 4,
+            &query_guard,
+        )
+        .await?;
+        let reader_factory = Arc::new(BudgetParquetReaderFactory {
+            store,
+            path: object_location,
+            file_size: object_meta.size,
+            metadata,
+            read_limit: (memory_budget / 4).max(1),
+        });
 
         Ok(Arc::new(Self {
             schema,
@@ -166,6 +316,7 @@ impl ParquetSelectTable {
             object_size: object_meta.size,
             access_plan,
             expanded_row_budget,
+            reader_factory,
         }))
     }
 
@@ -202,7 +353,12 @@ impl TableProvider for ParquetSelectTable {
             ));
         }
         let scan_limit = if filters.is_empty() { limit } else { None };
-        let file_source = Arc::new(ParquetSource::new(Arc::clone(&self.schema)));
+        let file_source = Arc::new(
+            ParquetSource::new(Arc::clone(&self.schema))
+                .with_parquet_file_reader_factory(self.reader_factory.clone())
+                .with_enable_page_index(false)
+                .with_bloom_filter_on_read(false),
+        );
         let config = FileScanConfigBuilder::new(self.object_store_url.clone(), file_source)
             .with_file(self.partitioned_file())
             .with_projection_indices(projection.cloned())?
@@ -439,6 +595,39 @@ pub(super) fn dictionary_parquet_for_test(value_bytes: usize, rows: usize) -> Ve
 #[cfg(test)]
 mod tests {
     use super::*;
+
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn parquet_cancelled_preflight_retains_admission_until_worker_finishes() {
+        let admission = Arc::new(tokio::sync::Semaphore::new(1));
+        let permit = Arc::clone(&admission).acquire_owned().await.expect("initial admission");
+        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
+        let (release_tx, release_rx) = std::sync::mpsc::channel();
+        let task = tokio::spawn(validate_in_worker(Some(Arc::new(permit)), move || {
+            started_tx.send(()).expect("report running worker");
+            release_rx
+                .recv_timeout(std::time::Duration::from_secs(5))
+                .expect("worker release");
+            Ok(())
+        }));
+        tokio::time::timeout(std::time::Duration::from_secs(5), started_rx)
+            .await
+            .expect("worker starts before deadline")
+            .expect("worker started");
+        task.abort();
+        assert!(task.await.expect_err("cancelled request task").is_cancelled());
+        assert!(Arc::clone(&admission).try_acquire_owned().is_err(), "running worker still owns admission");
+        release_tx.send(()).expect("release worker");
+        tokio::time::timeout(std::time::Duration::from_secs(5), async {
+            loop {
+                if Arc::clone(&admission).try_acquire_owned().is_ok() {
+                    break;
+                }
+                tokio::task::yield_now().await;
+            }
+        })
+        .await
+        .expect("worker releases admission after completion");
+    }
     use datafusion::{
         arrow::{
             array::Int32Array,
@@ -519,7 +708,7 @@ mod tests {
         let context = parquet_session(store);
         let state = context.state();
 
-        let error = match ParquetSelectTable::try_new(&state, &parquet_input("missing.parquet")).await {
+        let error = match ParquetSelectTable::try_new(&state, &parquet_input("missing.parquet"), None).await {
             Ok(_) => panic!("missing parquet object must fail"),
             Err(error) => error,
         };
@@ -538,7 +727,7 @@ mod tests {
         let context = parquet_session(store);
         let state = context.state();
 
-        let error = match ParquetSelectTable::try_new(&state, &parquet_input(object.as_ref())).await {
+        let error = match ParquetSelectTable::try_new(&state, &parquet_input(object.as_ref()), None).await {
             Ok(_) => panic!("corrupt parquet metadata must fail"),
             Err(error) => error,
         };
@@ -581,7 +770,7 @@ mod tests {
             .put(&Path::from("dictionary.parquet"), Bytes::from(bytes).into())
             .await
             .expect("fixture");
-        let table = ParquetSelectTable::try_new(&context.state(), &parquet_input("dictionary.parquet"))
+        let table = ParquetSelectTable::try_new(&context.state(), &parquet_input("dictionary.parquet"), None)
             .await
             .expect("valid dictionary metadata");
         context.register_table("source", table).expect("table");
@@ -752,7 +941,7 @@ mod tests {
             start: Some(starts[0]),
             end: Some(starts[0]),
         });
-        let table = ParquetSelectTable::try_new(&context.state(), &input)
+        let table = ParquetSelectTable::try_new(&context.state(), &input, None)
             .await
             .expect("an unselected large row group must not reject the small range");
         context.register_table("source", table).expect("table");
@@ -770,7 +959,7 @@ mod tests {
             start: Some(starts[1]),
             end: Some(starts[1]),
         });
-        let error = match ParquetSelectTable::try_new(&context.state(), &input).await {
+        let error = match ParquetSelectTable::try_new(&context.state(), &input, None).await {
             Ok(_) => panic!("the selected oversized row group must still be rejected before decoding"),
             Err(error) => error,
         };
@@ -806,13 +995,209 @@ mod tests {
             .put(&Path::from("large.parquet"), Bytes::from(bytes).into())
             .await
             .expect("fixture");
-        let error = match ParquetSelectTable::try_new(&context.state(), &parquet_input("large.parquet")).await {
+        let error = match ParquetSelectTable::try_new(&context.state(), &parquet_input("large.parquet"), None).await {
             Ok(_) => panic!("valid declared decoded size must be admitted before page decode"),
             Err(error) => error,
         };
         assert!(error.to_string().contains("query memory budget"));
     }
 
+    #[tokio::test]
+    async fn parquet_footer_collection_budget_is_checked_before_metadata_decode() {
+        // A root schema with no leaves and 50 empty, valid row groups. The
+        // encoded footer fits the read budget; its reservations do not.
+        let mut footer = vec![0x15, 2, 0x19, 0x1c, 0x48, 4];
+        footer.extend(b"root");
+        footer.extend([0x15, 0, 0, 0x16, 0, 0x19, 0xfc, 50]);
+        for _ in 0..50 {
+            footer.extend([0x19, 0x0c, 0x16, 0, 0x16, 0, 0]);
+        }
+        footer.push(0);
+        let decoded = ParquetMetaDataReader::decode_metadata(&footer).expect("valid compact footer fixture");
+        assert_eq!(decoded.num_row_groups(), 50);
+        let mut bytes = b"PAR1".to_vec();
+        bytes.extend(&footer);
+        bytes.extend(u32::try_from(footer.len()).expect("small footer").to_le_bytes());
+        bytes.extend(b"PAR1");
+        let store = Arc::new(InMemory::new());
+        let context = parquet_session_with_memory(store.clone(), 8192);
+        store
+            .put(&Path::from("footer.parquet"), Bytes::from(bytes).into())
+            .await
+            .expect("fixture");
+        let error = match ParquetSelectTable::try_new(&context.state(), &parquet_input("footer.parquet"), None).await {
+            Ok(_) => panic!("footer collection reservations exceed the query budget"),
+            Err(error) => error,
+        };
+        assert_eq!(error.select_error(), SelectError::ResourceExhausted);
+    }
+
+    #[tokio::test]
+    async fn parquet_forged_footer_cannot_hide_oversized_compressed_pages() {
+        use datafusion::parquet::file::metadata::ParquetMetaDataWriter;
+
+        let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Utf8, false)]));
+        let batch = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![Arc::new(datafusion::arrow::array::StringArray::from(vec![
+                "x".repeat(32 * 1024 * 1024),
+            ]))],
+        )
+        .expect("large compressible value");
+        let properties = datafusion::parquet::file::properties::WriterProperties::builder()
+            .set_dictionary_page_size_limit(64 * 1024 * 1024)
+            .set_compression(datafusion::parquet::basic::Compression::ZSTD(Default::default()))
+            .build();
+        let mut encoded = Vec::new();
+        let mut writer = ArrowWriter::try_new(&mut encoded, schema, Some(properties)).expect("writer");
+        writer.write(&batch).expect("compressed value");
+        writer.close().expect("footer");
+        let original = Bytes::from(encoded);
+        let builder = ParquetRecordBatchReaderBuilder::try_new(original.clone()).expect("valid compressed fixture");
+        let metadata = builder.metadata();
+        let groups = metadata
+            .row_groups()
+            .iter()
+            .map(|group| {
+                let columns = group
+                    .columns()
+                    .iter()
+                    .cloned()
+                    .map(|column| {
+                        column
+                            .into_builder()
+                            .set_total_uncompressed_size(64)
+                            .build()
+                            .expect("forged column")
+                    })
+                    .collect();
+                group
+                    .clone()
+                    .into_builder()
+                    .set_column_metadata(columns)
+                    .build()
+                    .expect("forged group")
+            })
+            .collect();
+        let forged = metadata.as_ref().clone().into_builder().set_row_groups(groups).build();
+        let footer_size = u32::from_le_bytes(
+            original[original.len() - 8..original.len() - 4]
+                .try_into()
+                .expect("footer size"),
+        );
+        let mut bytes = original[..original.len() - 8 - usize::try_from(footer_size).expect("footer length")].to_vec();
+        ParquetMetaDataWriter::new(&mut bytes, &forged)
+            .finish()
+            .expect("forged footer");
+        assert!(bytes.len() < 256 * 1024, "compressed input remains small");
+        admit_parquet_decode(&forged, None, 4 * 1024 * 1024).expect("footer-only admission misses the page size");
+
+        let store = Arc::new(InMemory::new());
+        let context = parquet_session_with_memory(store.clone(), 16 * 1024 * 1024);
+        store
+            .put(&Path::from("forged.parquet"), Bytes::from(bytes).into())
+            .await
+            .expect("fixture");
+        let result = ParquetSelectTable::try_new(&context.state(), &parquet_input("forged.parquet"), None).await;
+        let error = match result {
+            Ok(_) => panic!("actual query admission must reject oversized page buffers before decoding"),
+            Err(error) => error,
+        };
+        assert_eq!(error.select_error(), SelectError::ResourceExhausted);
+    }
+
+    #[tokio::test]
+    async fn parquet_legacy_lz4_checks_native_scratch_before_query_decode() {
+        use datafusion::parquet::{basic::Compression, file::metadata::ParquetMetaDataWriter};
+        let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)]));
+        let batch =
+            RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(datafusion::arrow::array::Int32Array::from(vec![42]))])
+                .expect("one integer row");
+        let properties = datafusion::parquet::file::properties::WriterProperties::builder()
+            .set_dictionary_enabled(false)
+            .build();
+        let mut original = Vec::new();
+        let mut writer = ArrowWriter::try_new(&mut original, schema, Some(properties)).expect("writer");
+        writer.write(&batch).expect("row");
+        writer.close().expect("footer");
+        let metadata = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(original))
+            .expect("metadata")
+            .metadata()
+            .clone();
+        // A valid legacy frame encoding the four little-endian bytes for 42.
+        let payload = [2, 33, 76, 24, 5, 0, 0, 0, 0x40, 42, 0, 0, 0];
+        let mut page = vec![0x15, 0, 0x15, 8, 0x15, 26, 0x2c, 0x15, 2, 0x15, 0, 0x15, 6, 0x15, 6, 0, 0];
+        page.extend(payload);
+        let column = metadata
+            .row_group(0)
+            .column(0)
+            .clone()
+            .into_builder()
+            .set_compression(Compression::LZ4)
+            .set_data_page_offset(4)
+            .set_total_compressed_size(i64::try_from(page.len()).expect("small page"))
+            .build()
+            .expect("legacy column");
+        let group = metadata
+            .row_group(0)
+            .clone()
+            .into_builder()
+            .set_column_metadata(vec![column])
+            .build()
+            .expect("row group");
+        let metadata = metadata.as_ref().clone().into_builder().set_row_groups(vec![group]).build();
+        let mut encoded = b"PAR1".to_vec();
+        encoded.extend(page);
+        ParquetMetaDataWriter::new(&mut encoded, &metadata)
+            .finish()
+            .expect("legacy footer");
+        assert!(encoded.len() < 4096);
+        let mut native = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(encoded.clone()))
+            .expect("legacy metadata")
+            .build()
+            .expect("native reader");
+        let decoded = native.next().expect("one batch").expect("legacy fallback decode");
+        assert_eq!(
+            decoded
+                .column(0)
+                .as_primitive::<datafusion::arrow::datatypes::Int32Type>()
+                .value(0),
+            42
+        );
+        let store = Arc::new(InMemory::new());
+        store
+            .put(&Path::from("legacy.parquet"), Bytes::from(encoded).into())
+            .await
+            .expect("fixture");
+        let tight = parquet_session_with_memory(store.clone(), 2 * 1024 * 1024);
+        let result = ParquetSelectTable::try_new(&tight.state(), &parquet_input("legacy.parquet"), None).await;
+        let error = match result {
+            Ok(_) => panic!("legacy decoder scratch exceeds the query budget"),
+            Err(error) => error,
+        };
+        assert_eq!(error.select_error(), SelectError::ResourceExhausted);
+        let roomy = parquet_session_with_memory(store, 128 * 1024 * 1024);
+        let table = ParquetSelectTable::try_new(&roomy.state(), &parquet_input("legacy.parquet"), None)
+            .await
+            .expect("admitted legacy file");
+        roomy.register_table("source", table).expect("table");
+        let rows = roomy
+            .sql("SELECT v FROM source")
+            .await
+            .expect("plan")
+            .collect()
+            .await
+            .expect("query");
+        assert_eq!(rows.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
+        assert_eq!(
+            rows[0]
+                .column(0)
+                .as_primitive::<datafusion::arrow::datatypes::Int32Type>()
+                .value(0),
+            42
+        );
+    }
+
     #[tokio::test]
     async fn parquet_nested_views_decode_through_the_actual_file_schema() {
         use datafusion::arrow::array::{ListArray, StringArray};

rustfs/src/table_catalog/maintenance/compaction.rs

@@ -21,8 +21,7 @@ use datafusion::arrow::datatypes::{DataType, Field, Schema};
 use datafusion::parquet::arrow::{
     add_encoded_arrow_schema_to_metadata, arrow_reader::ArrowReaderOptions, arrow_writer::ArrowWriterOptions,
 };
-use datafusion::parquet::file::FOOTER_SIZE;
-use datafusion::parquet::file::metadata::{ColumnChunkMetaData, FooterTail, RowGroupMetaData};
+use datafusion::parquet::file::metadata::ColumnChunkMetaData;
 use datafusion::parquet::file::properties::{EnabledStatistics, WriterProperties};
 use std::io::Write;
 use std::sync::OnceLock;
@@ -74,279 +73,54 @@ fn add_budget(total: &mut usize, amount: usize, limit: usize) -> TableCatalogSto
     Ok(())
 }
 
-/// Sizes read from one Parquet page header (Thrift compact encoding).
-#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
-struct PageHeaderSizes {
-    uncompressed: usize,
-    compressed: usize,
-}
-
+#[cfg(test)]
 const THRIFT_COMPACT_STOP: u8 = 0;
+#[cfg(test)]
 const THRIFT_COMPACT_BOOL_TRUE: u8 = 1;
-const THRIFT_COMPACT_BOOL_FALSE: u8 = 2;
-const THRIFT_COMPACT_BYTE: u8 = 3;
+#[cfg(test)]
 const THRIFT_COMPACT_I16: u8 = 4;
+#[cfg(test)]
 const THRIFT_COMPACT_I32: u8 = 5;
+#[cfg(test)]
 const THRIFT_COMPACT_I64: u8 = 6;
-const THRIFT_COMPACT_DOUBLE: u8 = 7;
+#[cfg(test)]
 const THRIFT_COMPACT_BINARY: u8 = 8;
+#[cfg(test)]
 const THRIFT_COMPACT_LIST: u8 = 9;
-const THRIFT_COMPACT_SET: u8 = 10;
+#[cfg(test)]
 const THRIFT_COMPACT_MAP: u8 = 11;
+#[cfg(test)]
 const THRIFT_COMPACT_STRUCT: u8 = 12;
-const THRIFT_MAX_NESTING: usize = 16;
-const PAGE_HEADER_FIELD_UNCOMPRESSED_SIZE: i16 = 2;
-const PAGE_HEADER_FIELD_COMPRESSED_SIZE: i16 = 3;
-const FILE_METADATA_FIELD_ROW_GROUPS: i16 = 4;
 
-/// Minimal Thrift compact-protocol cursor: enough to read a `PageHeader`'s
-/// size fields and skip everything else so the next header can be located.
-struct ThriftCompactCursor<'a> {
-    bytes: &'a [u8],
-    pos: usize,
-}
-
-impl ThriftCompactCursor<'_> {
-    fn error() -> TableCatalogStoreError {
-        TableCatalogStoreError::Invalid("compaction parquet thrift metadata is malformed".to_owned())
-    }
-
-    fn byte(&mut self) -> TableCatalogStoreResult<u8> {
-        let byte = *self.bytes.get(self.pos).ok_or_else(Self::error)?;
-        self.pos += 1;
-        Ok(byte)
-    }
-
-    fn skip_bytes(&mut self, count: usize) -> TableCatalogStoreResult<()> {
-        self.pos = self
-            .pos
-            .checked_add(count)
-            .filter(|end| *end <= self.bytes.len())
-            .ok_or_else(Self::error)?;
-        Ok(())
-    }
-
-    fn varint(&mut self) -> TableCatalogStoreResult<u64> {
-        let mut value = 0_u64;
-        for shift in (0..64).step_by(7) {
-            let byte = self.byte()?;
-            value |= u64::from(byte & 0x7f) << shift;
-            if byte & 0x80 == 0 {
-                return Ok(value);
-            }
-        }
-        Err(Self::error())
-    }
-
-    fn zigzag(&mut self) -> TableCatalogStoreResult<i64> {
-        let raw = self.varint()?;
-        Ok(((raw >> 1) as i64) ^ -((raw & 1) as i64))
-    }
-
-    /// Every collection element occupies at least one byte, so a declared
-    /// count beyond the remaining bytes is malformed. This bounds the walk to
-    /// the chunk's size instead of the count a hostile header declares.
-    fn collection_count(&self, count: u64) -> TableCatalogStoreResult<usize> {
-        usize::try_from(count)
-            .ok()
-            .filter(|count| *count <= self.bytes.len() - self.pos)
-            .ok_or_else(Self::error)
-    }
-
-    fn list_header(&mut self) -> TableCatalogStoreResult<(u8, usize)> {
-        let header = self.byte()?;
-        let count = match header >> 4 {
-            15 => self.varint()?,
-            short => u64::from(short),
-        };
-        Ok((header & 0x0f, self.collection_count(count)?))
-    }
-
-    fn field_id(&mut self, last: i16, header: u8) -> TableCatalogStoreResult<i16> {
-        match header >> 4 {
-            0 => i16::try_from(self.zigzag()?).map_err(|_| Self::error()),
-            delta => last.checked_add(i16::from(delta)).ok_or_else(Self::error),
-        }
-    }
-
-    /// parquet reads the page size fields by id as a zigzag varint whatever
-    /// the wire kind says. Follow the id; a kind that is not an integer is
-    /// rejected rather than skipped, so the walker never records a smaller
-    /// size than the decoder will allocate.
-    fn size_field(&mut self, kind: u8) -> TableCatalogStoreResult<usize> {
-        match kind {
-            THRIFT_COMPACT_I16 | THRIFT_COMPACT_I32 | THRIFT_COMPACT_I64 => {
-                usize::try_from(self.zigzag()?).map_err(|_| Self::error())
-            }
-            _ => Err(Self::error()),
-        }
-    }
-
-    fn skip_value(&mut self, kind: u8, depth: usize) -> TableCatalogStoreResult<()> {
-        if depth > THRIFT_MAX_NESTING {
-            return Err(Self::error());
-        }
-        match kind {
-            // A struct field's bool value lives in its header byte.
-            THRIFT_COMPACT_BOOL_TRUE | THRIFT_COMPACT_BOOL_FALSE => Ok(()),
-            THRIFT_COMPACT_BYTE => self.skip_bytes(1),
-            THRIFT_COMPACT_I16 | THRIFT_COMPACT_I32 | THRIFT_COMPACT_I64 => self.varint().map(drop),
-            THRIFT_COMPACT_DOUBLE => self.skip_bytes(8),
-            THRIFT_COMPACT_BINARY => {
-                let len = usize::try_from(self.varint()?).map_err(|_| Self::error())?;
-                self.skip_bytes(len)
-            }
-            THRIFT_COMPACT_LIST | THRIFT_COMPACT_SET => {
-                let (element, count) = self.list_header()?;
-                for _ in 0..count {
-                    self.skip_element(element, depth + 1)?;
-                }
-                Ok(())
-            }
-            THRIFT_COMPACT_MAP => {
-                let count = self.varint()?;
-                let count = self.collection_count(count)?;
-                if count == 0 {
-                    return Ok(());
-                }
-                let kinds = self.byte()?;
-                for _ in 0..count {
-                    self.skip_element(kinds >> 4, depth + 1)?;
-                    self.skip_element(kinds & 0x0f, depth + 1)?;
-                }
-                Ok(())
-            }
-            THRIFT_COMPACT_STRUCT => self.skip_struct(depth + 1),
-            _ => Err(Self::error()),
-        }
-    }
-
-    /// A collection element carries its own bytes, but parquet skips a bool
-    /// element without consuming any. Either width would let the walker and
-    /// the decoder disagree on where the next header starts; no Parquet
-    /// struct holds a bool collection, so fail closed instead.
-    fn skip_element(&mut self, kind: u8, depth: usize) -> TableCatalogStoreResult<()> {
-        match kind {
-            THRIFT_COMPACT_BOOL_TRUE | THRIFT_COMPACT_BOOL_FALSE => Err(Self::error()),
-            kind => self.skip_value(kind, depth),
-        }
-    }
-
-    fn skip_struct(&mut self, depth: usize) -> TableCatalogStoreResult<()> {
-        let mut field_id = 0_i16;
-        loop {
-            let header = self.byte()?;
-            if header == THRIFT_COMPACT_STOP {
-                return Ok(());
-            }
-            field_id = self.field_id(field_id, header)?;
-            self.skip_value(header & 0x0f, depth)?;
-        }
-    }
-
-    /// Read a `PageHeader` struct, returning its sizes and leaving the cursor at
-    /// the first page byte.
-    fn page_header(&mut self) -> TableCatalogStoreResult<PageHeaderSizes> {
-        let mut sizes = PageHeaderSizes::default();
-        let mut field_id = 0_i16;
-        loop {
-            let header = self.byte()?;
-            if header == THRIFT_COMPACT_STOP {
-                return Ok(sizes);
-            }
-            let kind = header & 0x0f;
-            field_id = self.field_id(field_id, header)?;
-            match field_id {
-                PAGE_HEADER_FIELD_UNCOMPRESSED_SIZE => sizes.uncompressed = self.size_field(kind)?,
-                PAGE_HEADER_FIELD_COMPRESSED_SIZE => sizes.compressed = self.size_field(kind)?,
-                _ => self.skip_value(kind, 1)?,
-            }
-        }
-    }
-
-    /// Walk the footer's `FileMetaData`. The pinned parquet reserves the
-    /// declared `row_groups` capacity before reading one element; charge that
-    /// reservation to the budget so a small footer cannot request gigabytes.
-    fn file_metadata(&mut self, limit: usize) -> TableCatalogStoreResult<()> {
-        let mut field_id = 0_i16;
-        let mut reserved = 0_usize;
-        loop {
-            let header = self.byte()?;
-            if header == THRIFT_COMPACT_STOP {
-                return Ok(());
-            }
-            let kind = header & 0x0f;
-            field_id = self.field_id(field_id, header)?;
-            if field_id == FILE_METADATA_FIELD_ROW_GROUPS && kind == THRIFT_COMPACT_LIST {
-                let (element, count) = self.list_header()?;
-                if element != THRIFT_COMPACT_STRUCT {
-                    return Err(Self::error());
-                }
-                let reservation = count.checked_mul(size_of::<RowGroupMetaData>()).ok_or_else(Self::error)?;
-                add_budget(&mut reserved, reservation, limit)?;
-                for _ in 0..count {
-                    self.skip_struct(2)?;
-                }
-            } else {
-                self.skip_value(kind, 0)?;
-            }
-        }
+fn parquet_admission_error(error: rustfs_s3select_api::parquet_budget::ParquetBudgetError) -> TableCatalogStoreError {
+    match error {
+        rustfs_s3select_api::parquet_budget::ParquetBudgetError::Exceeded => budget_error(),
+        error => TableCatalogStoreError::Invalid(format!("compaction {error}")),
     }
 }
 
-/// Reject footer collection counts the decoder would allocate for before it
-/// reads a single element. Bytes after the struct are ignored, as parquet does.
 fn validate_footer_metadata(data: &Bytes, limit: usize) -> TableCatalogStoreResult<()> {
-    let error = || TableCatalogStoreError::Invalid("compaction parquet footer is malformed".to_owned());
-    let tail_start = data.len().checked_sub(FOOTER_SIZE).ok_or_else(error)?;
-    let tail: &[u8; FOOTER_SIZE] = data[tail_start..].try_into().map_err(|_| error())?;
-    let tail = FooterTail::try_new(tail).map_err(|_| error())?;
-    let metadata_start = tail_start.checked_sub(tail.metadata_length()).ok_or_else(error)?;
-    let mut cursor = ThriftCompactCursor {
-        bytes: &data[metadata_start..tail_start],
-        pos: 0,
-    };
-    cursor.file_metadata(limit)
+    rustfs_s3select_api::parquet_budget::validate_file_footer(data, limit).map_err(parquet_admission_error)
 }
 
-/// Sum the page headers of one column chunk (`[start, end)` in the file) and
-/// hold every page and the chunk total to `limit`. Pages are never decompressed.
-fn validate_page_headers_in_range(chunk: &[u8], chunk_declared: usize, limit: usize) -> TableCatalogStoreResult<usize> {
-    let mut cursor = ThriftCompactCursor { bytes: chunk, pos: 0 };
-    let mut total = 0_usize;
-    while cursor.pos < chunk.len() {
-        let sizes = cursor.page_header()?;
-        add_budget(&mut total, sizes.uncompressed, limit)?;
-        if total > chunk_declared.max(limit.min(chunk_declared.saturating_mul(2))) {
-            return Err(TableCatalogStoreError::Invalid(
-                "compaction parquet page headers exceed the footer's declared chunk size".to_owned(),
-            ));
-        }
-        cursor.skip_bytes(sizes.compressed)?;
-    }
-    Ok(total)
+#[cfg(test)]
+fn validate_page_headers_in_range(chunk: &[u8], declared: usize, limit: usize) -> TableCatalogStoreResult<usize> {
+    rustfs_s3select_api::parquet_budget::validate_page_headers(chunk, declared, limit).map_err(parquet_admission_error)
 }
 
 fn validate_column_chunk_page_headers(
     data: &Bytes,
     column: &ColumnChunkMetaData,
-    chunk_declared: usize,
+    declared: usize,
     limit: usize,
     location: &str,
-) -> TableCatalogStoreResult<()> {
-    // parquet's `byte_range`: the chunk starts at the dictionary page when the
-    // footer declares one, wherever that is. Walk exactly the bytes it decodes.
-    let start = column.dictionary_page_offset().unwrap_or(column.data_page_offset());
-    let start = usize::try_from(start).map_err(|_| budget_error())?;
-    let length = usize::try_from(column.compressed_size()).map_err(|_| budget_error())?;
-    let end = start.checked_add(length).filter(|end| *end <= data.len()).ok_or_else(|| {
-        TableCatalogStoreError::Invalid(format!("compaction parquet {location} column chunk lies outside the file"))
-    })?;
-    validate_page_headers_in_range(&data[start..end], chunk_declared, limit).map_err(|err| match err {
-        TableCatalogStoreError::Invalid(message) => TableCatalogStoreError::Invalid(format!("{location}: {message}")),
-        other => other,
-    })?;
-    Ok(())
+) -> TableCatalogStoreResult<usize> {
+    rustfs_s3select_api::parquet_budget::validate_column_chunk_page_headers(data, column, declared, limit).map_err(|error| {
+        match parquet_admission_error(error) {
+            TableCatalogStoreError::Invalid(message) => TableCatalogStoreError::Invalid(format!("{location}: {message}")),
+            error => error,
+        }
+    })
 }
 
 fn view_field(field: &Field, depth: usize, fields: &mut usize) -> TableCatalogStoreResult<Arc<Field>> {
@@ -542,11 +316,11 @@ fn compact_parquet_with_budgets(
             let mut columns = 0;
             for column in group.columns() {
                 let chunk_declared = usize::try_from(column.uncompressed_size()).map_err(|_| budget_error())?;
-                add_budget(&mut columns, chunk_declared, limit)?;
                 // The footer is one claim; the page headers the decoder actually
-                // sizes its buffers from are another. Walk them without
-                // decompressing and hold both to the budget.
-                validate_column_chunk_page_headers(&data, column, chunk_declared, limit, &location)?;
+                // sizes its buffers from are another. Check the page headers
+                // and codec expansion before admitting the decoder.
+                let resident = validate_column_chunk_page_headers(&data, column, chunk_declared, limit, &location)?;
+                add_budget(&mut columns, chunk_declared.max(resident), limit)?;
             }
             let group_bytes = usize::try_from(group.total_byte_size()).map_err(|_| budget_error())?;
             add_budget(&mut declared, group_bytes.max(columns), limit)?;
@@ -850,9 +624,9 @@ mod tests {
     #[test]
     fn compaction_page_header_walker_bounds_collection_counts_by_remaining_bytes() {
         let limit = 64 * 1024 * 1024;
-        // Unknown field 4: LIST<BOOL> with long-form count 2^63. Parquet's skip
+        // Unknown field 10: LIST<BOOL> with long-form count 2^63. Parquet's skip
         // would spin; the walker must reject it before iterating.
-        let mut list = vec![(1 << 4) | THRIFT_COMPACT_LIST, 0xf0 | THRIFT_COMPACT_BOOL_TRUE];
+        let mut list = vec![(7 << 4) | THRIFT_COMPACT_LIST, 0xf0 | THRIFT_COMPACT_BOOL_TRUE];
         list.extend(varint(1 << 63));
         let started = std::time::Instant::now();
         let error = validate_page_headers_in_range(&page_chunk(compact_i32_field(1, 8), &list), 4096, limit)
@@ -864,15 +638,15 @@ mod tests {
         assert!(started.elapsed() < std::time::Duration::from_secs(1));
         // A bool collection of any size is rejected: parquet consumes no bytes
         // per element, so either width would desynchronize the two parsers.
-        let list = [(1 << 4) | THRIFT_COMPACT_LIST, (2 << 4) | THRIFT_COMPACT_BOOL_TRUE, 1, 1];
+        let list = [(7 << 4) | THRIFT_COMPACT_LIST, (2 << 4) | THRIFT_COMPACT_BOOL_TRUE, 1, 1];
         assert!(validate_page_headers_in_range(&page_chunk(compact_i32_field(1, 8), &list), 4096, limit).is_err());
         // MAP<BOOL, BOOL> with a huge count is rejected the same way.
-        let mut map = vec![(1 << 4) | THRIFT_COMPACT_MAP];
+        let mut map = vec![(7 << 4) | THRIFT_COMPACT_MAP];
         map.extend(varint(1 << 62));
         map.push((THRIFT_COMPACT_BOOL_TRUE << 4) | THRIFT_COMPACT_BOOL_TRUE);
         assert!(validate_page_headers_in_range(&page_chunk(compact_i32_field(1, 8), &map), 4096, limit).is_err());
         // Collections whose elements carry bytes still skip when the count fits.
-        let mut list = vec![(1 << 4) | THRIFT_COMPACT_LIST, (2 << 4) | THRIFT_COMPACT_I32];
+        let mut list = vec![(7 << 4) | THRIFT_COMPACT_LIST, (2 << 4) | THRIFT_COMPACT_I32];
         list.extend(compact_i32_field(0, 7)[1..].iter());
         list.extend(compact_i32_field(0, -7)[1..].iter());
         assert_eq!(
@@ -880,7 +654,7 @@ mod tests {
             8
         );
         // A LIST<I32> whose count exceeds the chunk is malformed, not truncated.
-        let list = [(1 << 4) | THRIFT_COMPACT_LIST, 0xf0 | THRIFT_COMPACT_I32, 0x7f];
+        let list = [(7 << 4) | THRIFT_COMPACT_LIST, 0xf0 | THRIFT_COMPACT_I32, 0x7f];
         assert!(validate_page_headers_in_range(&page_chunk(compact_i32_field(1, 8), &list), 4096, limit).is_err());
     }
 
@@ -956,6 +730,35 @@ mod tests {
         }
     }
 
+    #[test]
+    fn compaction_aggregates_column_decode_working_sets() {
+        let columns =
+            (0..4).map(|index| (format!("value_{index}"), Arc::new(StringArray::from(vec!["v"; 1000])) as Arc<dyn Array>));
+        let batch = RecordBatch::try_from_iter(columns).expect("multi-column batch");
+        let properties = WriterProperties::builder()
+            .set_compression(Compression::ZSTD(ZstdLevel::default()))
+            .build();
+        let mut writer = ArrowWriter::try_new(Vec::new(), batch.schema(), Some(properties)).expect("Parquet writer");
+        writer.write(&batch).expect("multi-column data");
+        let data = Bytes::from(writer.into_inner().expect("Parquet footer"));
+        let limit = 100 * 1024;
+        let reader = ParquetRecordBatchReaderBuilder::try_new(data.clone()).expect("valid fixture metadata");
+        let mut working_set = 0;
+        for column in reader.metadata().row_group(0).columns() {
+            let declared = usize::try_from(column.uncompressed_size()).expect("positive size");
+            working_set += validate_column_chunk_page_headers(&data, column, declared, limit, "columns.parquet")
+                .expect("each column fits individually");
+        }
+        assert!(working_set > limit, "the aggregate exceeds the decode budget");
+        let error = compact_parquet_with_limits(vec![("columns.parquet".to_owned(), data.to_vec())], limit, limit)
+            .err()
+            .expect("aggregate column allocations must be admitted before decoding");
+        assert!(
+            matches!(&error, TableCatalogStoreError::Invalid(message) if message.contains("budget")),
+            "{error:?}"
+        );
+    }
+
     #[test]
     fn compaction_rejects_compressed_expansion_before_decoding_pages() {
         let file = parquet_strings(&"x".repeat(64 * 1024), 32, false);