rustfs/crates/s3select-api/src/object_store.rs

// Copyright 2024 RustFS Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::{
    PrepareSelectObjectSnapshotError, QueryError, SELECT_DEFAULT_READ_BUFFER_SIZE, SelectError, SelectGetObjectReader,
    SelectInputMetrics, SelectObjectOptions, SelectObjectSnapshot, SelectObjectSnapshotReadError, SelectStorageError,
    SelectStore, SnapshotConsistencyError,
    input_stream::{
        CompressionFormat, MAX_SELECT_RECORD_BYTES, SELECT_DECODE_CHUNK_BYTES, SelectInputReader, compressed_input_reader,
        compressed_input_stream, input_io_error, processed_bytes_limit,
    },
    metrics::SelectInputMetricsRecorder,
    query::{
        ast::{JsonPathSegment, JsonSource},
        parser::RustFsDialect,
        session::{QueryExecutionGuard, QueryExecutionTracker},
    },
    resolve_select_object_store_handle, select_is_err_bucket_not_found, select_is_err_object_not_found,
    select_is_err_version_not_found,
};
use async_trait::async_trait;
use bytes::Bytes;
use chrono::{DateTime, Utc};
use datafusion::{
    common::{DataFusionError, runtime::SpawnedTask},
    execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation, UnboundedMemoryPool},
    object_store::{
        Attributes, CopyOptions, Error as o_Error, GetOptions, GetRange, GetResult, GetResultPayload, ListResult,
        MultipartUpload, ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result, path::Path,
    },
    sql::{
        planner::IdentNormalizer,
        sqlparser::{
            ast::{Expr, Ident, JsonPathElem, ObjectNamePart, SetExpr, Statement, TableFactor},
            parser::Parser as SqlParser,
        },
    },
};
use futures::pin_mut;
use futures::{Stream, StreamExt, TryStreamExt, future::ready, stream};
use futures_core::stream::BoxStream;
use http::{HeaderMap, HeaderValue, header::HeaderName};
use parking_lot::Mutex;
use rustfs_common::DEFAULT_DELIMITER;
use s3s::header::{
    X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_ALGORITHM, X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY,
    X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY_MD5,
};
use s3s::{
    S3Error, S3ErrorCode, S3Result,
    dto::{CompressionType, InputSerialization, ScanRange, SelectObjectContentInput},
};
use std::collections::VecDeque;
use std::ops::Range;
#[cfg(test)]
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, atomic::AtomicBool};
use tokio::{io::AsyncReadExt, sync::OnceCell};
use tokio_util::io::ReaderStream;
use transform_stream::AsyncTryStream;

use crate::storage_api::object_store::HTTPRangeSpec;

fn select_default_read_buffer_size_u64() -> u64 {
    u64::try_from(SELECT_DEFAULT_READ_BUFFER_SIZE).unwrap_or(u64::MAX)
}

fn compression_format(input: &InputSerialization) -> Result<Option<CompressionFormat>, SelectError> {
    match input.compression_type.as_ref().map(|value| value.as_str()) {
        None | Some(CompressionType::NONE) => Ok(None),
        Some(CompressionType::GZIP) => Ok(Some(CompressionFormat::Gzip)),
        Some(CompressionType::BZIP2) => Ok(Some(CompressionFormat::Bzip2)),
        Some(_) => Err(SelectError::InvalidCompressionFormat),
    }
}

/// Maximum allowed object size for JSON DOCUMENT mode.
///
/// JSON DOCUMENT format requires loading the entire file into memory for DOM
/// parsing, so memory consumption grows linearly with file size.  Objects
/// larger than this threshold are rejected with an error rather than risking
/// an OOM condition.
///
/// To process larger JSON files, convert the input to **JSON LINES** (NDJSON,
/// `type = LINES`), which supports line-by-line streaming with no memory
/// size limit.
///
/// Default: 128 MiB.  This matches the AWS S3 Select limit for JSON DOCUMENT
/// inputs. The query memory pool also applies: RustFS reserves 64 times the
/// input size for parsing and output. With the default 64 MiB query memory
/// limit, JSON DOCUMENT inputs larger than 1 MiB are rejected; scalar source
/// aliases reserve additional space for their maximum per-row expansion.
/// Raise `RUSTFS_S3SELECT_MEMORY_LIMIT_BYTES` to process larger inputs, up to
/// this hard cap.
pub const MAX_JSON_DOCUMENT_BYTES: u64 = 128 * 1024 * 1024;
const JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER: usize = 64;
const JSON_SCALAR_COLUMN_MEMORY_RESERVATION_MULTIPLIER: usize = 14;
const JSON_CANCELLATION_CHECK_BYTES: usize = 64 * 1024;
const JSON_CANCELLATION_CHECK_KEYS: usize = 1024;
pub const INVALID_SCAN_RANGE_MESSAGE: &str =
    "The value of a parameter in ScanRange element is invalid. Check the service API documentation and try again.";
const NORMALIZED_RECORD_DELIMITER: &[u8] = b"\r\n";
const NORMALIZED_FIELD_DELIMITER: &[u8] = &[DEFAULT_DELIMITER];

/// Returns true for the MinIO-compatible full-scan range marker.
pub fn is_noop_scan_range(scan_range: &ScanRange) -> bool {
    scan_range.start == Some(0) && scan_range.end.is_none()
}

pub struct EcObjectStore {
    input: Arc<SelectObjectContentInput>,
    need_convert: bool,
    delimiter: String,
    /// True when the JSON input type is DOCUMENT (multi-line formatted JSON).
    /// In that case the raw bytes are buffered and flattened to NDJSON before
    /// being handed to DataFusion's Arrow JSON reader.
    is_json_document: bool,
    /// JSON source path produced by the SQL compatibility analyzer.
    json_source: JsonSource,
    input_metrics: Arc<SelectInputMetrics>,
    memory_pool: Arc<dyn MemoryPool>,
    query_tracker: Option<QueryExecutionTracker>,
    store: Option<Arc<SelectStore>>,
    snapshot: OnceCell<Arc<SelectObjectSnapshot>>,
    #[cfg(test)]
    reader_open_count: Arc<AtomicUsize>,
}

#[derive(Debug, thiserror::Error)]
pub(crate) enum EcObjectStoreBuildError {
    #[error("ec store not inited")]
    StoreUnavailable,
    #[error("SelectObjectContent snapshot consistency failure: {0}")]
    Snapshot(#[source] SnapshotConsistencyError),
}

#[derive(Debug, thiserror::Error)]
pub(crate) enum SelectObjectStoreError {
    #[error("SelectObjectContent bucket does not exist")]
    BucketNotFound {
        #[source]
        source: SelectStorageError,
    },
    #[error("SelectObjectContent object does not exist")]
    ObjectNotFound {
        #[source]
        source: SelectStorageError,
    },
    #[error("SelectObjectContent storage failure")]
    Storage {
        #[source]
        source: SelectStorageError,
    },
    #[error("SelectObjectContent ScanRange is invalid")]
    InvalidScanRange,
}

impl SelectObjectStoreError {
    pub(crate) fn select_error(&self) -> SelectError {
        match self {
            Self::BucketNotFound { .. } => SelectError::BucketNotFound,
            Self::ObjectNotFound { .. } => SelectError::ObjectNotFound,
            Self::InvalidScanRange => SelectError::InvalidScanRange,
            Self::Storage { .. } => SelectError::InternalError,
        }
    }
}

#[derive(Clone, Copy, Debug)]
pub struct SelectScanRange {
    start: u64,
    end: u64,
}

impl SelectScanRange {
    pub const fn new(start: u64, end: u64) -> Self {
        Self { start, end }
    }

    pub const fn start(&self) -> u64 {
        self.start
    }

    pub const fn end(&self) -> u64 {
        self.end
    }
}

#[derive(Clone, Copy, Debug)]
pub struct InvalidScanRange;

impl EcObjectStore {
    pub fn new(input: Arc<SelectObjectContentInput>) -> S3Result<Self> {
        let source = legacy_json_source_from_input(&input);
        Self::build_lazy(
            input,
            Arc::new(UnboundedMemoryPool::default()),
            None,
            Arc::new(SelectInputMetrics::default()),
            source,
        )
        .map_err(map_build_error_to_s3)
    }

    pub fn new_with_snapshot(input: Arc<SelectObjectContentInput>, snapshot: Arc<SelectObjectSnapshot>) -> S3Result<Self> {
        let source = legacy_json_source_from_input(&input);
        Self::build_with_snapshot(
            input,
            Arc::new(UnboundedMemoryPool::default()),
            None,
            Arc::new(SelectInputMetrics::default()),
            snapshot,
            source,
        )
        .map_err(map_build_error_to_s3)
    }

    pub(crate) fn new_with_memory_pool_and_source(
        input: Arc<SelectObjectContentInput>,
        memory_pool: Arc<dyn MemoryPool>,
        input_metrics: Arc<SelectInputMetrics>,
        snapshot: Option<Arc<SelectObjectSnapshot>>,
        source: JsonSource,
    ) -> std::result::Result<Self, EcObjectStoreBuildError> {
        match snapshot {
            Some(snapshot) => Self::build_with_snapshot(input, memory_pool, None, input_metrics, snapshot, source),
            None => Self::build_lazy(input, memory_pool, None, input_metrics, source),
        }
    }

    pub(crate) fn new_with_query_tracker_and_source(
        input: Arc<SelectObjectContentInput>,
        memory_pool: Arc<dyn MemoryPool>,
        query_tracker: QueryExecutionTracker,
        input_metrics: Arc<SelectInputMetrics>,
        snapshot: Option<Arc<SelectObjectSnapshot>>,
        source: JsonSource,
    ) -> std::result::Result<Self, EcObjectStoreBuildError> {
        match snapshot {
            Some(snapshot) => Self::build_with_snapshot(input, memory_pool, Some(query_tracker), input_metrics, snapshot, source),
            None => Self::build_lazy(input, memory_pool, Some(query_tracker), input_metrics, source),
        }
    }

    fn build_lazy(
        input: Arc<SelectObjectContentInput>,
        memory_pool: Arc<dyn MemoryPool>,
        query_tracker: Option<QueryExecutionTracker>,
        input_metrics: Arc<SelectInputMetrics>,
        source: JsonSource,
    ) -> std::result::Result<Self, EcObjectStoreBuildError> {
        let store = resolve_select_object_store_handle().ok_or(EcObjectStoreBuildError::StoreUnavailable)?;
        Ok(Self::build(input, memory_pool, query_tracker, input_metrics, Some(store), None, source))
    }

    fn build_with_snapshot(
        input: Arc<SelectObjectContentInput>,
        memory_pool: Arc<dyn MemoryPool>,
        query_tracker: Option<QueryExecutionTracker>,
        input_metrics: Arc<SelectInputMetrics>,
        snapshot: Arc<SelectObjectSnapshot>,
        source: JsonSource,
    ) -> std::result::Result<Self, EcObjectStoreBuildError> {
        if !snapshot.is_for(&input.bucket, &input.key) {
            return Err(EcObjectStoreBuildError::Snapshot(SnapshotConsistencyError::ObjectChanged));
        }
        Ok(Self::build(
            input,
            memory_pool,
            query_tracker,
            input_metrics,
            None,
            Some(snapshot),
            source,
        ))
    }

    fn build(
        input: Arc<SelectObjectContentInput>,
        memory_pool: Arc<dyn MemoryPool>,
        query_tracker: Option<QueryExecutionTracker>,
        input_metrics: Arc<SelectInputMetrics>,
        store: Option<Arc<SelectStore>>,
        snapshot: Option<Arc<SelectObjectSnapshot>>,
        source: JsonSource,
    ) -> Self {
        let (need_convert, delimiter) = if let Some(csv) = input.request.input_serialization.csv.as_ref() {
            if let Some(delimiter) = csv.field_delimiter.as_ref() {
                if delimiter.len() > 1 {
                    (true, delimiter.to_owned())
                } else {
                    (false, String::new())
                }
            } else {
                (false, String::new())
            }
        } else {
            (false, String::new())
        };

        // Detect JSON DOCUMENT type: the entire file is a single (possibly
        // multi-line) JSON object/array, NOT newline-delimited JSON.
        let is_json_document = is_json_document_input(&input);

        Self {
            input,
            need_convert,
            delimiter,
            is_json_document,
            json_source: source,
            input_metrics,
            memory_pool,
            query_tracker,
            store,
            snapshot: match snapshot {
                Some(snapshot) => OnceCell::new_with(Some(snapshot)),
                None => OnceCell::new(),
            },
            #[cfg(test)]
            reader_open_count: Arc::new(AtomicUsize::new(0)),
        }
    }

    fn scan_range(&self, object_size: u64) -> Result<Option<SelectScanRange>> {
        let Some(scan_range) = self.input.request.scan_range.as_ref() else {
            return Ok(None);
        };
        if is_noop_scan_range(scan_range) {
            return Ok(None);
        }
        scan_range_from_bounds(scan_range.start, scan_range.end, object_size)
    }

    fn record_delimiter(&self) -> Vec<u8> {
        self.input
            .request
            .input_serialization
            .csv
            .as_ref()
            .and_then(|csv| csv.record_delimiter.as_ref())
            .map(|delimiter| delimiter.as_bytes().to_vec())
            .unwrap_or_else(|| b"\n".to_vec())
    }

    fn record_delimiter_for_conversion(&self) -> Option<Vec<u8>> {
        let delimiter = self.record_delimiter();
        (self.need_convert || (delimiter.len() == 2 && delimiter != NORMALIZED_RECORD_DELIMITER)).then_some(delimiter)
    }

    fn csv_has_header(&self) -> bool {
        self.input
            .request
            .input_serialization
            .csv
            .as_ref()
            .and_then(|csv| csv.file_header_info.as_ref())
            .is_some_and(|info| matches!(info.as_str(), "USE" | "IGNORE"))
    }

    async fn snapshot(&self, version: Option<&str>) -> Result<&Arc<SelectObjectSnapshot>> {
        let snapshot = self
            .snapshot
            .get_or_try_init(|| async {
                let store = self.store.as_ref().ok_or_else(|| o_Error::Generic {
                    store: "EcObjectStore",
                    source: "prepared snapshot is unavailable".into(),
                })?;
                let opts = SelectObjectOptions {
                    version_id: version.map(|version| {
                        let version = version.trim();
                        if version.eq_ignore_ascii_case("null") {
                            uuid::Uuid::nil().to_string()
                        } else {
                            version.to_owned()
                        }
                    }),
                    ..Default::default()
                };
                let snapshot = store
                    .prepare_select_object_snapshot(&self.input.bucket, &self.input.key, &select_read_headers(&self.input), &opts)
                    .await
                    .map_err(|err| map_prepare_snapshot_error(&self.input.bucket, &self.input.key, err))?;
                Ok::<_, o_Error>(Arc::new(snapshot))
            })
            .await?;
        if let Some(requested) = version
            && !snapshot.matches_version(requested)
        {
            return Err(o_Error::Generic {
                store: "EcObjectStore",
                source: "prepared snapshot is pinned to a different object version".into(),
            });
        }
        Ok(snapshot)
    }

    async fn object_reader(&self, range: Option<HTTPRangeSpec>) -> Result<SelectGetObjectReader> {
        #[cfg(test)]
        self.reader_open_count.fetch_add(1, Ordering::Relaxed);
        self.snapshot(None)
            .await?
            .open_reader(range)
            .await
            .map_err(|err| snapshot_read_error(&self.input.bucket, &self.input.key, err))
    }

    async fn read_raw_range(&self, range: Range<u64>) -> Result<Bytes> {
        if range.is_empty() {
            return Ok(Bytes::new());
        }
        let snapshot = self.snapshot(None).await?;
        let reader = self.object_reader(Some(http_range_spec_from_range(range.clone()))).await?;
        let object_size = snapshot.logical_size();
        let resolved_range = GetRange::Bounded(range)
            .as_range(object_size)
            .map_err(|err| o_Error::Generic {
                store: "EcObjectStore",
                source: Box::new(err),
            })?;
        let expected_size = usize::try_from(resolved_range.end - resolved_range.start).map_err(|err| o_Error::Generic {
            store: "EcObjectStore",
            source: Box::new(err),
        })?;
        let mut reader = reader.stream;
        let mut bytes = Vec::new();
        reader.read_to_end(&mut bytes).await.map_err(|err| o_Error::Generic {
            store: "EcObjectStore",
            source: Box::new(err),
        })?;
        if bytes.len() < expected_size {
            return Err(incomplete_object_stream_error(expected_size - bytes.len()));
        }
        bytes.truncate(expected_size);
        Ok(Bytes::from(bytes))
    }

    async fn read_header_record(&self, object_size: u64, delimiter: &[u8]) -> Result<Bytes> {
        if object_size == 0 {
            return Ok(Bytes::new());
        }

        let mut end = select_default_read_buffer_size_u64().min(object_size);
        loop {
            let bytes = self.read_raw_range(0..end).await?;
            if let Some(pos) = find_delimiter(&bytes, delimiter) {
                return Ok(bytes.slice(0..pos + delimiter.len()));
            }
            if end == object_size {
                return Ok(bytes);
            }
            end = end.saturating_mul(2).min(object_size);
        }
    }

    async fn scan_range_read_start(&self, scan_range: SelectScanRange, delimiter: &[u8]) -> Result<u64> {
        let delimiter_len = u64::try_from(delimiter.len()).unwrap_or(u64::MAX);
        let fallback_start = scan_range.start().saturating_sub(delimiter_len);
        if delimiter.len() != 2 || delimiter[0] != delimiter[1] || scan_range.start() == 0 {
            return Ok(fallback_start);
        }

        let context_start = scan_range.start().saturating_sub(select_default_read_buffer_size_u64());
        let context = self.read_raw_range(context_start..scan_range.start()).await?;
        let suffix_len = context.iter().rev().take_while(|byte| **byte == delimiter[0]).count();
        if suffix_len == context.len() && context_start > 0 {
            return Err(o_Error::Generic {
                store: "EcObjectStore",
                source: "self-overlapping CSV record delimiter exceeds the bounded ScanRange context".into(),
            });
        }
        if suffix_len == 0 {
            return Ok(fallback_start);
        }
        Ok(scan_range.start().saturating_sub(u64::from(suffix_len % 2 != 0)))
    }
}

pub(crate) fn is_json_document_input(input: &SelectObjectContentInput) -> bool {
    input
        .request
        .input_serialization
        .json
        .as_ref()
        .and_then(|json| json.type_.as_ref())
        .is_some_and(|json_type| json_type.as_str() == "DOCUMENT")
}

/// Preserves the pre-typed-path single-key behavior of public legacy constructors.
pub(crate) fn legacy_json_source_from_input(input: &SelectObjectContentInput) -> JsonSource {
    if !is_json_document_input(input) {
        return JsonSource::default();
    }
    let Ok(mut statements) = SqlParser::parse_sql(&RustFsDialect, &input.request.expression) else {
        return JsonSource::default();
    };
    if statements.len() != 1 {
        return JsonSource::default();
    }
    let Some(Statement::Query(query)) = statements.pop() else {
        return JsonSource::default();
    };
    let SetExpr::Select(select) = query.body.as_ref() else {
        return JsonSource::default();
    };
    let [table] = select.from.as_slice() else {
        return JsonSource::default();
    };
    let TableFactor::Table {
        name, alias, json_path, ..
    } = &table.relation
    else {
        return JsonSource::default();
    };
    let Some(ObjectNamePart::Identifier(table_name)) = name.0.first() else {
        return JsonSource::default();
    };
    if name.0.len() > 2 {
        return JsonSource::default();
    }
    let is_s3_object = if table_name.quote_style.is_some() {
        table_name.value == "S3Object"
    } else {
        table_name.value.eq_ignore_ascii_case("S3Object")
    };
    if !is_s3_object {
        return JsonSource::default();
    }
    let path = match (name.0.get(1), json_path.as_ref()) {
        (Some(ObjectNamePart::Identifier(sub_path)), None) => vec![JsonPathSegment::Key {
            name: sub_path.value.clone(),
            quoted: sub_path.quote_style.is_some(),
        }],
        (None, None) => Vec::new(),
        (None, Some(json_path)) if matches!(json_path.path.as_slice(), [JsonPathElem::Bracket { key: Expr::Wildcard(_) }]) => {
            vec![JsonPathSegment::ArrayWildcard]
        }
        _ => return JsonSource::default(),
    };
    let scalar_column = alias
        .as_ref()
        .map(|alias| IdentNormalizer::default().normalize(alias.name.clone()))
        .or_else(|| match path.as_slice() {
            [] => Some(IdentNormalizer::default().normalize(table_name.clone())),
            [JsonPathSegment::Key { name, quoted }] => {
                let alias = if *quoted {
                    Ident::with_quote('"', name)
                } else {
                    Ident::new(name)
                };
                Some(IdentNormalizer::default().normalize(alias))
            }
            [JsonPathSegment::ArrayWildcard] => Some(IdentNormalizer::default().normalize(Ident::new("_1"))),
            _ => None,
        });
    JsonSource::new(path, scalar_column)
}

impl std::fmt::Debug for EcObjectStore {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("EcObjectStore")
            .field("bucket", &self.input.bucket)
            .field("object", &self.input.key)
            .field("need_convert", &self.need_convert)
            .field("is_json_document", &self.is_json_document)
            .field("json_source", &self.json_source)
            .finish_non_exhaustive()
    }
}

impl std::fmt::Display for EcObjectStore {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("EcObjectStore")
    }
}

fn unsupported_store_error(op: &str) -> o_Error {
    o_Error::Generic {
        store: "s3select-api",
        source: Box::new(std::io::Error::new(
            std::io::ErrorKind::Unsupported,
            format!("operation {op} is not supported in EcObjectStore"),
        )),
    }
}

fn insert_header(headers: &mut HeaderMap, name: HeaderName, value: Option<&str>) {
    if let Some(value) = value
        && let Ok(value) = HeaderValue::from_str(value)
    {
        headers.insert(name, value);
    }
}

fn select_read_headers(input: &SelectObjectContentInput) -> HeaderMap {
    let mut headers = HeaderMap::new();
    insert_header(
        &mut headers,
        X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_ALGORITHM,
        input.sse_customer_algorithm.as_deref(),
    );
    insert_header(&mut headers, X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY, input.sse_customer_key.as_deref());
    insert_header(
        &mut headers,
        X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY_MD5,
        input.sse_customer_key_md5.as_deref(),
    );
    headers
}

fn http_range_spec_from_get_range(range: &GetRange) -> HTTPRangeSpec {
    match range {
        GetRange::Bounded(range) => http_range_spec_from_range(range.clone()),
        GetRange::Offset(start) => HTTPRangeSpec {
            is_suffix_length: false,
            start: *start as i64,
            end: -1,
        },
        GetRange::Suffix(length) => HTTPRangeSpec {
            is_suffix_length: true,
            start: *length as i64,
            end: -1,
        },
    }
}

fn http_range_spec_from_range(range: Range<u64>) -> HTTPRangeSpec {
    HTTPRangeSpec {
        is_suffix_length: false,
        start: range.start as i64,
        end: range.end.saturating_sub(1) as i64,
    }
}

fn http_range_spec_from_start(start: u64) -> HTTPRangeSpec {
    HTTPRangeSpec {
        is_suffix_length: false,
        start: start as i64,
        end: -1,
    }
}

fn find_delimiter(bytes: &[u8], delimiter: &[u8]) -> Option<usize> {
    if delimiter.is_empty() {
        return None;
    }
    bytes.windows(delimiter.len()).position(|window| window == delimiter)
}

fn map_prepare_snapshot_error(bucket: &str, object: &str, err: PrepareSelectObjectSnapshotError) -> o_Error {
    match err {
        PrepareSelectObjectSnapshotError::Storage(err) => map_storage_error(bucket, object, err),
        err => o_Error::Generic {
            store: "EcObjectStore",
            source: Box::new(err),
        },
    }
}

fn map_build_error_to_s3(error: EcObjectStoreBuildError) -> S3Error {
    let mut s3_error = S3Error::with_message(S3ErrorCode::InternalError, SelectError::InternalError.to_string());
    s3_error.set_source(Box::new(error));
    s3_error
}

fn snapshot_read_error(bucket: &str, object: &str, err: SelectObjectSnapshotReadError) -> o_Error {
    match err {
        SelectObjectSnapshotReadError::Storage(err) => map_storage_error(bucket, object, err),
        err => o_Error::Generic {
            store: "EcObjectStore",
            source: Box::new(err),
        },
    }
}

fn map_storage_error(bucket: &str, object: &str, err: SelectStorageError) -> o_Error {
    if select_is_err_bucket_not_found(&err) {
        return o_Error::NotFound {
            path: format!("{bucket}/{object}"),
            source: Box::new(SelectObjectStoreError::BucketNotFound { source: err }),
        };
    }
    if select_is_err_object_not_found(&err) || select_is_err_version_not_found(&err) {
        return o_Error::NotFound {
            path: format!("{bucket}/{object}"),
            source: Box::new(SelectObjectStoreError::ObjectNotFound { source: err }),
        };
    }
    o_Error::Generic {
        store: "EcObjectStore",
        source: Box::new(SelectObjectStoreError::Storage { source: err }),
    }
}

fn snapshot_last_modified(snapshot: &SelectObjectSnapshot) -> Result<DateTime<Utc>> {
    let mod_time = snapshot.object_info().mod_time.ok_or_else(|| o_Error::Generic {
        store: "EcObjectStore",
        source: std::io::Error::new(std::io::ErrorKind::InvalidData, "snapshot metadata has no modification time").into(),
    })?;
    DateTime::<Utc>::from_timestamp(mod_time.unix_timestamp(), mod_time.nanosecond()).ok_or_else(|| o_Error::Generic {
        store: "EcObjectStore",
        source: std::io::Error::new(std::io::ErrorKind::InvalidData, "snapshot modification time is out of range").into(),
    })
}

pub fn scan_range_from_bounds(start: Option<i64>, end: Option<i64>, object_size: u64) -> Result<Option<SelectScanRange>> {
    parse_scan_range_from_bounds(start, end, object_size).map_err(|_| invalid_scan_range_store_error())
}

pub fn validate_scan_range_bounds(
    start: Option<i64>,
    end: Option<i64>,
    object_size: u64,
) -> std::result::Result<(), InvalidScanRange> {
    parse_scan_range_from_bounds(start, end, object_size).map(|_| ())
}

fn parse_scan_range_from_bounds(
    start: Option<i64>,
    end: Option<i64>,
    object_size: u64,
) -> std::result::Result<Option<SelectScanRange>, InvalidScanRange> {
    if start.is_none() && end.is_none() {
        return Ok(None);
    }
    if start.is_some_and(|value| value < 0) || end.is_some_and(|value| value < 0) {
        return Err(InvalidScanRange);
    }
    if let (Some(start), Some(end)) = (start, end)
        && start > end
    {
        return Err(InvalidScanRange);
    }
    if let Some(start) = start {
        let start = start as u64;
        if object_size == 0 {
            if start > 0 {
                return Err(InvalidScanRange);
            }
            return Ok(Some(SelectScanRange::new(0, 0)));
        }
        if start >= object_size {
            return Err(InvalidScanRange);
        }
    }
    if object_size == 0 {
        return Ok(Some(SelectScanRange::new(0, 0)));
    }

    let last_byte = object_size - 1;
    let (start, end) = match (start, end) {
        (Some(start), Some(end)) => (start as u64, (end as u64).min(last_byte)),
        (Some(start), None) => (start as u64, last_byte),
        (None, Some(suffix_len)) => {
            let suffix_len = suffix_len as u64;
            (object_size.saturating_sub(suffix_len), last_byte)
        }
        (None, None) => return Ok(None),
    };
    Ok(Some(SelectScanRange::new(start, end)))
}

fn invalid_scan_range_store_error() -> o_Error {
    o_Error::Generic {
        store: "EcObjectStore",
        source: Box::new(SelectObjectStoreError::InvalidScanRange),
    }
}

#[async_trait]
impl ObjectStore for EcObjectStore {
    async fn put_opts(&self, _location: &Path, _payload: PutPayload, _opts: PutOptions) -> Result<PutResult> {
        Err(unsupported_store_error("put_opts"))
    }

    async fn put_multipart_opts(&self, _location: &Path, _opts: PutMultipartOptions) -> Result<Box<dyn MultipartUpload>> {
        Err(unsupported_store_error("put_multipart_opts"))
    }

    async fn get_opts(&self, location: &Path, options: GetOptions) -> Result<GetResult> {
        // SelectObjectContent has no version-id input. For direct ObjectStore
        // compatibility, a version supplied on the first operation defines
        // this instance's immutable snapshot; later operations reuse it.
        let snapshot = self.snapshot(options.version.as_deref()).await?;
        let original_size = snapshot.logical_size();
        let compression = compression_format(&self.input.request.input_serialization).map_err(|source| o_Error::Generic {
            store: "EcObjectStore",
            source: Box::new(source),
        })?;
        let has_effective_request_range = self
            .input
            .request
            .scan_range
            .as_ref()
            .is_some_and(|scan_range| !is_noop_scan_range(scan_range));
        if compression.is_some() && (options.range.is_some() || has_effective_request_range) {
            return Err(o_Error::Generic {
                store: "EcObjectStore",
                source: Box::new(SelectError::UnsupportedScanRangeInput),
            });
        }
        let object_info = snapshot.object_info();
        let meta = ObjectMeta {
            location: location.clone(),
            last_modified: snapshot_last_modified(snapshot)?,
            size: original_size,
            e_tag: object_info.etag.clone(),
            version: object_info.version_id.map(|version| version.to_string()),
        };
        let result_range = match options.range.as_ref() {
            Some(range) => range.as_range(original_size).map_err(|err| o_Error::Generic {
                store: "EcObjectStore",
                source: Box::new(err),
            })?,
            None => 0..original_size,
        };
        if options.head {
            return Ok(GetResult {
                payload: GetResultPayload::Stream(stream::empty().boxed()),
                meta,
                range: result_range,
                attributes: Attributes::default(),
            });
        }

        let record_delimiter = self.record_delimiter_for_conversion();
        let needs_scan_context = options.range.is_none() && has_effective_request_range;
        let scan_context = if needs_scan_context {
            if let Some(scan_range) = self.scan_range(original_size)? {
                let delimiter = self.record_delimiter();
                let read_start = self.scan_range_read_start(scan_range, &delimiter).await?;
                Some((scan_range, read_start))
            } else {
                None
            }
        } else {
            None
        };

        let range = options.range.as_ref().map(http_range_spec_from_get_range);
        let reader = if let Some((_, read_start)) = scan_context.as_ref() {
            let range = (original_size > 0).then(|| http_range_spec_from_start(*read_start));
            self.object_reader(range).await?
        } else {
            self.object_reader(range).await?
        };

        let meter_input = self.input.request.input_serialization.parquet.is_none();
        let payload = if let Some(compression) = compression {
            let max_processed_bytes = processed_bytes_limit();
            let query_guard = match self.query_tracker.as_ref() {
                Some(query_tracker) => Some(query_tracker.query_guard().ok_or_else(|| o_Error::Generic {
                    store: "EcObjectStore",
                    source: Box::new(QueryError::Cancel),
                })?),
                None => None,
            };
            if self.is_json_document {
                let reader = compressed_input_reader(
                    reader.stream,
                    original_size,
                    compression,
                    Arc::clone(&self.input_metrics),
                    max_processed_bytes,
                    query_guard,
                );
                let stream = compressed_json_document_ndjson_stream(
                    reader,
                    self.json_source.clone(),
                    Arc::clone(&self.memory_pool),
                    self.query_tracker.clone(),
                );
                GetResultPayload::Stream(stream)
            } else {
                let input_record_delimiter = if self.input.request.input_serialization.csv.is_some() {
                    self.record_delimiter()
                } else {
                    b"\n".to_vec()
                };
                let stream = compressed_input_stream(
                    reader.stream,
                    original_size,
                    compression,
                    Arc::clone(&self.input_metrics),
                    input_record_delimiter,
                    max_processed_bytes,
                    query_guard,
                )?;
                let stream =
                    convert_csv_delimiter_stream(stream, record_delimiter, self.need_convert.then(|| self.delimiter.clone()));
                GetResultPayload::Stream(stream)
            }
        } else if options.range.is_some() {
            let size = usize::try_from(result_range.end - result_range.start).map_err(|err| o_Error::Generic {
                store: "EcObjectStore",
                source: Box::new(err),
            })?;
            let stream = bytes_stream(ReaderStream::with_capacity(reader.stream, SELECT_DEFAULT_READ_BUFFER_SIZE), size);
            if meter_input {
                GetResultPayload::Stream(meter_uncompressed_input_stream(stream, Arc::clone(&self.input_metrics)).boxed())
            } else {
                GetResultPayload::Stream(stream.boxed())
            }
        } else if self.is_json_document {
            // JSON DOCUMENT mode: gate on object size before doing any I/O.
            //
            // Small files (<= MAX_JSON_DOCUMENT_BYTES): build a lazy stream
            // that defers all I/O and JSON parsing until DataFusion first
            // polls it.  Parsing runs inside spawn_blocking so the async
            // runtime thread is never blocked.
            //
            // Large files (> MAX_JSON_DOCUMENT_BYTES): return an error
            // immediately.  JSON DOCUMENT relies on serde_json DOM parsing
            // which must load the whole file into memory; rejecting oversized
            // files upfront is safer than risking OOM.  Users should convert
            // their data to JSON LINES (NDJSON) format for large files.
            validate_json_document_size(original_size)?;
            let stream = json_document_ndjson_stream(
                reader.stream,
                original_size,
                self.json_source.clone(),
                Arc::clone(&self.input_metrics),
                Arc::clone(&self.memory_pool),
                self.query_tracker.clone(),
            );
            GetResultPayload::Stream(stream)
        } else if let Some((scan_range, read_start)) = scan_context {
            let delimiter = self.record_delimiter();
            let include_header = self.csv_has_header();
            let header = if include_header && read_start > 0 {
                let input_metrics = self.input_metrics.recorder();
                let header = self.read_header_record(original_size, &delimiter).await?;
                input_metrics.record_uncompressed(header.len());
                Some(header)
            } else {
                None
            };
            let stream = scan_range_stream(
                meter_uncompressed_input_stream(
                    ReaderStream::with_capacity(reader.stream, SELECT_DEFAULT_READ_BUFFER_SIZE),
                    Arc::clone(&self.input_metrics),
                ),
                delimiter,
                scan_range,
                include_header && header.is_none(),
                read_start,
                original_size,
            )
            .boxed();
            let stream = if let Some(header) = header {
                stream::once(ready(Ok(header))).chain(stream).boxed()
            } else {
                stream
            };
            let stream =
                convert_csv_delimiter_stream(stream, record_delimiter, self.need_convert.then(|| self.delimiter.clone()));
            GetResultPayload::Stream(stream)
        } else {
            let stream_size = usize::try_from(original_size).map_err(|err| o_Error::Generic {
                store: "EcObjectStore",
                source: Box::new(err),
            })?;
            let stream = bytes_stream(ReaderStream::with_capacity(reader.stream, SELECT_DEFAULT_READ_BUFFER_SIZE), stream_size);
            if meter_input {
                let stream = meter_uncompressed_input_stream(stream, Arc::clone(&self.input_metrics));
                let stream =
                    convert_csv_delimiter_stream(stream, record_delimiter, self.need_convert.then(|| self.delimiter.clone()));
                GetResultPayload::Stream(stream)
            } else {
                GetResultPayload::Stream(stream.boxed())
            }
        };

        Ok(GetResult {
            payload,
            meta,
            range: result_range,
            attributes: Attributes::default(),
        })
    }

    async fn get_ranges(&self, _location: &Path, ranges: &[Range<u64>]) -> Result<Vec<Bytes>> {
        let mut out = Vec::with_capacity(ranges.len());
        for range in ranges {
            out.push(self.read_raw_range(range.clone()).await?);
        }
        Ok(out)
    }

    fn delete_stream(&self, _locations: BoxStream<'static, Result<Path>>) -> BoxStream<'static, Result<Path>> {
        stream::once(ready(Err(unsupported_store_error("delete_stream")))).boxed()
    }

    fn list(&self, _prefix: Option<&Path>) -> BoxStream<'static, Result<ObjectMeta>> {
        stream::once(ready(Err(unsupported_store_error("list")))).boxed()
    }

    async fn list_with_delimiter(&self, _prefix: Option<&Path>) -> Result<ListResult> {
        Err(unsupported_store_error("list_with_delimiter"))
    }

    async fn copy_opts(&self, _from: &Path, _to: &Path, _options: CopyOptions) -> Result<()> {
        Err(unsupported_store_error("copy_opts"))
    }
}

struct CsvDelimiterConverter {
    record_delimiter: Option<Vec<u8>>,
    field_delimiter: Option<Vec<u8>>,
    carry: Vec<u8>,
}

impl CsvDelimiterConverter {
    fn new(record_delimiter: Option<Vec<u8>>, field_delimiter: Option<Vec<u8>>) -> Self {
        Self {
            record_delimiter: record_delimiter.filter(|delimiter| !delimiter.is_empty()),
            field_delimiter: field_delimiter.filter(|delimiter| !delimiter.is_empty()),
            carry: Vec::new(),
        }
    }

    fn max_delimiter_len(&self) -> usize {
        self.record_delimiter
            .as_ref()
            .into_iter()
            .chain(self.field_delimiter.as_ref())
            .map(Vec::len)
            .max()
            .unwrap_or(1)
    }

    fn convert_prefix(&self, bytes: &[u8], end: usize) -> (Vec<u8>, usize) {
        let mut converted = Vec::with_capacity(bytes.len());
        let mut pos = 0;
        while pos < end {
            let record_match = self
                .record_delimiter
                .as_ref()
                .filter(|delimiter| bytes[pos..].starts_with(delimiter));
            let field_match = self
                .field_delimiter
                .as_ref()
                .filter(|delimiter| bytes[pos..].starts_with(delimiter));
            if let Some(delimiter) = field_match
                && record_match.is_none_or(|record_delimiter| delimiter.len() > record_delimiter.len())
            {
                converted.extend_from_slice(NORMALIZED_FIELD_DELIMITER);
                pos += delimiter.len();
            } else if let Some(delimiter) = record_match {
                if delimiter.len() == 2 && delimiter != NORMALIZED_RECORD_DELIMITER {
                    converted.extend_from_slice(NORMALIZED_RECORD_DELIMITER);
                } else {
                    converted.extend_from_slice(delimiter);
                }
                pos += delimiter.len();
            } else {
                converted.push(bytes[pos]);
                pos += 1;
            }
        }
        (converted, pos)
    }

    fn convert_chunk(&mut self, chunk: &[u8]) -> Vec<u8> {
        let mut combined = Vec::with_capacity(self.carry.len() + chunk.len());
        combined.extend_from_slice(&self.carry);
        combined.extend_from_slice(chunk);

        let safe_end = combined.len().saturating_sub(self.max_delimiter_len().saturating_sub(1));
        let (converted, pos) = self.convert_prefix(&combined, safe_end);
        self.carry.clear();
        self.carry.extend_from_slice(&combined[pos..]);
        converted
    }

    fn finish(&mut self) -> Vec<u8> {
        let (converted, _) = self.convert_prefix(&self.carry, self.carry.len());
        self.carry.clear();
        converted
    }
}

fn convert_delimiter_stream<S>(
    stream: S,
    record_delimiter: Option<Vec<u8>>,
    field_delimiter: Option<Vec<u8>>,
) -> BoxStream<'static, Result<Bytes>>
where
    S: Stream<Item = Result<Bytes>> + Send + 'static,
{
    AsyncTryStream::<Bytes, o_Error, _>::new(|mut y| async move {
        let mut converter = CsvDelimiterConverter::new(record_delimiter, field_delimiter);
        pin_mut!(stream);
        while let Some(result) = stream.next().await {
            let bytes = result?;
            let converted = converter.convert_chunk(&bytes);
            if !converted.is_empty() {
                y.yield_ok(Bytes::from(converted)).await;
            }
        }
        let converted = converter.finish();
        if !converted.is_empty() {
            y.yield_ok(Bytes::from(converted)).await;
        }
        Ok(())
    })
    .boxed()
}

#[cfg(test)]
fn convert_record_delimiter_stream<S>(stream: S, delimiter: Vec<u8>) -> BoxStream<'static, Result<Bytes>>
where
    S: Stream<Item = Result<Bytes>> + Send + 'static,
{
    // DataFusion's CSV reader treats CRLF as a record terminator.
    convert_delimiter_stream(stream, Some(delimiter), None)
}

#[cfg(test)]
fn convert_field_delimiter_stream<S>(stream: S, delimiter: String) -> BoxStream<'static, Result<Bytes>>
where
    S: Stream<Item = Result<Bytes>> + Send + 'static,
{
    convert_delimiter_stream(stream, None, Some(delimiter.into_bytes()))
}

fn convert_csv_delimiter_stream<S>(
    stream: S,
    record_delimiter: Option<Vec<u8>>,
    field_delimiter: Option<String>,
) -> BoxStream<'static, Result<Bytes>>
where
    S: Stream<Item = Result<Bytes>> + Send + 'static,
{
    match (record_delimiter, field_delimiter) {
        (None, None) => stream.boxed(),
        (record, field) => convert_delimiter_stream(stream, record, field.map(String::into_bytes)),
    }
}

struct ScanRangeState<S> {
    stream: S,
    delimiter: Vec<u8>,
    range: SelectScanRange,
    include_header: bool,
    offset: u64,
    record_start: u64,
    record: Vec<u8>,
    pending: VecDeque<Bytes>,
    done: bool,
    expected_end: u64,
}

fn scan_range_stream<S>(
    stream: S,
    delimiter: Vec<u8>,
    range: SelectScanRange,
    include_header: bool,
    base_offset: u64,
    expected_end: u64,
) -> BoxStream<'static, Result<Bytes>>
where
    S: Stream<Item = std::io::Result<Bytes>> + Send + Unpin + 'static,
{
    let state = ScanRangeState {
        stream,
        delimiter,
        range,
        include_header,
        offset: base_offset,
        record_start: base_offset,
        record: Vec::new(),
        pending: VecDeque::new(),
        done: false,
        expected_end,
    };

    stream::unfold(state, |mut state| async move {
        loop {
            if let Some(bytes) = state.pending.pop_front() {
                return Some((Ok(bytes), state));
            }
            if state.done {
                return None;
            }
            match state.stream.next().await {
                Some(Ok(bytes)) => state.push_chunk(&bytes),
                Some(Err(err)) => {
                    state.done = true;
                    return Some((
                        Err(o_Error::Generic {
                            store: "EcObjectStore",
                            source: Box::new(err),
                        }),
                        state,
                    ));
                }
                None => {
                    if state.offset < state.expected_end {
                        state.done = true;
                        return Some((Err(incomplete_object_stream_error(state.expected_end - state.offset)), state));
                    }
                    state.finish_pending_record();
                    state.done = true;
                }
            }
        }
    })
    .boxed()
}

impl<S> ScanRangeState<S> {
    fn push_chunk(&mut self, bytes: &[u8]) {
        if bytes.is_empty() || self.done {
            return;
        }
        if self.record.is_empty() {
            self.record_start = self.offset;
        }
        let search_start = self.record.len().saturating_sub(self.delimiter.len().saturating_sub(1));
        self.record.extend_from_slice(bytes);
        self.offset = self.offset.saturating_add(bytes.len() as u64);

        let mut search_start = search_start;
        while let Some(pos) = find_delimiter(&self.record[search_start..], &self.delimiter) {
            let record_end = search_start + pos + self.delimiter.len();
            self.finish_record(record_end);
            if self.done {
                break;
            }
            search_start = 0;
        }
    }

    fn finish_record(&mut self, record_end: usize) {
        let record = self.record.drain(..record_end).collect::<Vec<_>>();
        let record_start = self.record_start;
        self.record_start = self.record_start.saturating_add(record_end as u64);
        self.push_record(record, record_start);
    }

    fn finish_pending_record(&mut self) {
        if self.record.is_empty() {
            return;
        }
        let record = std::mem::take(&mut self.record);
        let record_start = self.record_start;
        self.push_record(record, record_start);
    }

    fn push_record(&mut self, record: Vec<u8>, record_start: u64) {
        let include_header = self.include_header && record_start == 0;
        let include_record = record_start >= self.range.start() && record_start <= self.range.end();
        if include_header || include_record {
            self.pending.push_back(Bytes::from(record));
        } else {
            if record_start > self.range.end() {
                self.done = true;
            }
        }
    }
}

/// Build a lazy NDJSON stream from a JSON DOCUMENT reader.
///
/// `get_opts` calls this and returns immediately – no I/O is performed until
/// DataFusion begins polling the returned stream.  The pipeline is:
///
/// 1. **Read** – the object bytes are read asynchronously from `stream` only
///    when the returned stream is first polled.
/// 2. **Parse** – JSON deserialization runs inside
///    `tokio::task::spawn_blocking` so the async runtime is never blocked by
///    CPU-bound work, even for very large documents.
/// 3. **Yield** – each NDJSON line (one per array element, or one line for a
///    scalar/object root) is yielded as a separate [`Bytes`] chunk, so
///    DataFusion can pipeline row processing as lines arrive.
fn json_document_ndjson_stream(
    stream: SelectInputReader,
    original_size: u64,
    json_source: JsonSource,
    input_metrics: Arc<SelectInputMetrics>,
    memory_pool: Arc<dyn MemoryPool>,
    query_tracker: Option<QueryExecutionTracker>,
) -> futures_core::stream::BoxStream<'static, Result<Bytes>> {
    json_document_ndjson_stream_with_parser(
        stream,
        JsonDocumentReadMode::Exact {
            original_size,
            input_metrics: input_metrics.recorder(),
        },
        json_source,
        memory_pool,
        query_tracker,
        |all_bytes, json_source, cancellation| {
            parse_json_document_to_lines_cancellable(&all_bytes, &json_source, cancellation.as_ref())
        },
    )
}

fn compressed_json_document_ndjson_stream(
    stream: SelectInputReader,
    json_source: JsonSource,
    memory_pool: Arc<dyn MemoryPool>,
    query_tracker: Option<QueryExecutionTracker>,
) -> futures_core::stream::BoxStream<'static, Result<Bytes>> {
    json_document_ndjson_stream_with_parser(
        stream,
        JsonDocumentReadMode::Bounded,
        json_source,
        memory_pool,
        query_tracker,
        |all_bytes, json_source, cancellation| {
            parse_json_document_to_lines_cancellable(&all_bytes, &json_source, cancellation.as_ref())
        },
    )
}

enum JsonDocumentReadMode {
    Exact {
        original_size: u64,
        input_metrics: SelectInputMetricsRecorder,
    },
    Bounded,
}

fn json_document_ndjson_stream_with_parser<P>(
    stream: SelectInputReader,
    read_mode: JsonDocumentReadMode,
    json_source: JsonSource,
    memory_pool: Arc<dyn MemoryPool>,
    query_tracker: Option<QueryExecutionTracker>,
    parser: P,
) -> futures_core::stream::BoxStream<'static, Result<Bytes>>
where
    P: FnOnce(Vec<u8>, JsonSource, Arc<AtomicBool>) -> std::io::Result<Vec<Bytes>> + Send + 'static,
{
    AsyncTryStream::<Bytes, o_Error, _>::new(|mut y| async move {
        let reservation = MemoryConsumer::new("S3 Select JSON document").register(&memory_pool);

        // ── 1. Read phase (lazy: only runs when the stream is polled) ────
        pin_mut!(stream);
        let all_bytes = match read_mode {
            JsonDocumentReadMode::Exact {
                original_size,
                input_metrics,
            } => {
                let buffer_capacity = usize::try_from(original_size).map_err(|_| o_Error::Generic {
                    store: "EcObjectStore",
                    source: Box::new(DataFusionError::ResourcesExhausted(format!(
                        "JSON DOCUMENT input size {original_size} does not fit in memory"
                    ))),
                })?;
                resize_json_document_reservation(&reservation, buffer_capacity, &json_source)?;
                let mut all_bytes = Vec::with_capacity(buffer_capacity);
                let read_result = stream.take(original_size).read_to_end(&mut all_bytes).await;
                input_metrics.record_uncompressed(all_bytes.len());
                read_result.map_err(input_io_error)?;
                if all_bytes.len() != buffer_capacity {
                    return Err(incomplete_object_stream_error(buffer_capacity - all_bytes.len()));
                }
                all_bytes
            }
            JsonDocumentReadMode::Bounded => {
                let mut all_bytes = Vec::new();
                let mut buffer = vec![0; SELECT_DECODE_CHUNK_BYTES];
                loop {
                    let read = stream.read(&mut buffer).await.map_err(input_io_error)?;
                    if read == 0 {
                        break;
                    }
                    let new_len = all_bytes.len().checked_add(read).ok_or_else(|| o_Error::Generic {
                        store: "EcObjectStore",
                        source: Box::new(json_document_memory_reservation_overflow(all_bytes.len())),
                    })?;
                    let new_len_u64 = u64::try_from(new_len).map_err(|_| o_Error::Generic {
                        store: "EcObjectStore",
                        source: Box::new(DataFusionError::ResourcesExhausted(format!(
                            "JSON DOCUMENT input size {new_len} does not fit in the object size type"
                        ))),
                    })?;
                    validate_json_document_size(new_len_u64)?;
                    grow_json_document_buffer(&mut all_bytes, new_len, &reservation, &json_source)?;
                    all_bytes.extend_from_slice(&buffer[..read]);
                }
                all_bytes
            }
        };

        // ── 2. Parse phase (blocking thread pool, non-blocking runtime) ──
        let queued_query_guard = match query_tracker.as_ref() {
            Some(query_tracker) => Some(query_tracker.query_guard().ok_or_else(|| o_Error::Generic {
                store: "EcObjectStore",
                source: Box::new(json_document_parse_interrupted_error()),
            })?),
            None => None,
        };
        let task_resources = JsonDocumentTaskResources {
            _reservation: reservation,
            query_guard: queued_query_guard,
        };
        let pending_query_guard = PendingQueryExecutionGuard::new(query_tracker);
        let task_query_guard = pending_query_guard.task_state();
        let cancellation = Arc::new(AtomicBool::new(false));
        let queued_task = Arc::new(Mutex::new(Some(JsonDocumentParseTask {
            parser,
            all_bytes,
            json_source,
            task_resources,
        })));
        let _cancel_on_drop = JsonDocumentCancellation::new(Arc::clone(&cancellation), Arc::clone(&queued_task));
        let (lines, _task_resources) = SpawnedTask::spawn_blocking(move || {
            let JsonDocumentParseTask {
                parser,
                all_bytes,
                json_source,
                mut task_resources,
            } = queued_task.lock().take().ok_or_else(json_document_parse_interrupted_error)?;
            let query_guard = PendingQueryExecutionGuard::start(&task_query_guard)?;
            task_resources.query_guard = query_guard;
            parser(all_bytes, json_source, cancellation).map(|lines| (lines, task_resources))
        })
        .await
        .map_err(|e| o_Error::Generic {
            store: "EcObjectStore",
            source: e.to_string().into(),
        })?
        .map_err(|e| o_Error::Generic {
            store: "EcObjectStore",
            source: classify_json_document_parse_error(e),
        })?;

        // ── 3. Yield phase (one Bytes per NDJSON line) ───────────────────
        for line in lines {
            y.yield_ok(line).await;
        }
        Ok(())
    })
    .boxed()
}

fn grow_json_document_buffer(
    buffer: &mut Vec<u8>,
    required_len: usize,
    reservation: &MemoryReservation,
    json_source: &JsonSource,
) -> Result<()> {
    if required_len <= buffer.capacity() {
        return Ok(());
    }

    let max_capacity = usize::try_from(MAX_JSON_DOCUMENT_BYTES).map_err(|_| o_Error::Generic {
        store: "EcObjectStore",
        source: Box::new(DataFusionError::ResourcesExhausted(
            "JSON DOCUMENT size limit does not fit in memory".to_string(),
        )),
    })?;
    let target_capacity = required_len
        .checked_next_power_of_two()
        .ok_or_else(|| o_Error::Generic {
            store: "EcObjectStore",
            source: Box::new(DataFusionError::ResourcesExhausted(format!(
                "JSON DOCUMENT buffer capacity overflow at {required_len} bytes"
            ))),
        })?
        .min(max_capacity);
    if target_capacity < required_len {
        return Err(o_Error::Generic {
            store: "EcObjectStore",
            source: Box::new(DataFusionError::ResourcesExhausted(format!(
                "JSON DOCUMENT input size {required_len} exceeds the maximum buffer capacity"
            ))),
        });
    }

    resize_json_document_reservation(reservation, target_capacity, json_source)?;
    buffer
        .try_reserve_exact(target_capacity - buffer.len())
        .map_err(|_| o_Error::Generic {
            store: "EcObjectStore",
            source: Box::new(DataFusionError::ResourcesExhausted(format!(
                "JSON DOCUMENT input buffer allocation failed at {target_capacity} bytes"
            ))),
        })?;
    resize_json_document_reservation(reservation, buffer.capacity(), json_source)
}

fn resize_json_document_reservation(reservation: &MemoryReservation, input_bytes: usize, json_source: &JsonSource) -> Result<()> {
    let reservation_bytes =
        json_document_memory_reservation_bytes(input_bytes, json_source).map_err(|source| o_Error::Generic {
            store: "EcObjectStore",
            source: Box::new(source),
        })?;
    reservation.try_resize(reservation_bytes).map_err(|source| o_Error::Generic {
        store: "EcObjectStore",
        source: Box::new(source),
    })
}

struct JsonDocumentTaskResources {
    // Struct fields drop in declaration order, so admission covers the reservation through teardown.
    _reservation: MemoryReservation,
    query_guard: Option<QueryExecutionGuard>,
}

struct JsonDocumentParseTask<P> {
    parser: P,
    all_bytes: Vec<u8>,
    json_source: JsonSource,
    task_resources: JsonDocumentTaskResources,
}

struct JsonDocumentCancellation<T> {
    cancelled: Arc<AtomicBool>,
    queued: Arc<Mutex<Option<T>>>,
}

impl<T> JsonDocumentCancellation<T> {
    fn new(cancelled: Arc<AtomicBool>, queued: Arc<Mutex<Option<T>>>) -> Self {
        Self { cancelled, queued }
    }
}

impl<T> Drop for JsonDocumentCancellation<T> {
    fn drop(&mut self) {
        self.cancelled.store(true, std::sync::atomic::Ordering::Release);
        let queued = self.queued.lock().take();
        drop(queued);
    }
}

struct CancellableJsonReader<'a> {
    inner: std::io::Cursor<&'a [u8]>,
    cancelled: &'a AtomicBool,
}

impl std::io::Read for CancellableJsonReader<'_> {
    fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
        ensure_json_parse_active(self.cancelled)?;
        std::io::Read::read(&mut self.inner, buffer)
    }
}

fn ensure_json_parse_active(cancelled: &AtomicBool) -> std::io::Result<()> {
    if cancelled.load(std::sync::atomic::Ordering::Acquire) {
        Err(std::io::Error::new(std::io::ErrorKind::Interrupted, SelectError::Canceled))
    } else {
        Ok(())
    }
}

fn json_document_memory_reservation_bytes(input_bytes: usize, json_source: &JsonSource) -> datafusion::common::Result<usize> {
    let base = input_bytes
        .checked_mul(JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER)
        .ok_or_else(|| json_document_memory_reservation_overflow(input_bytes))?;
    let scalar_column = json_source.scalar_column().unwrap_or_default();
    // A scalar row holds one key copy and its JSON encoding. One source byte
    // can expand to six escaped bytes, and the serializer buffer can grow to
    // twice its final length.
    let scalar_column_per_row = scalar_column
        .len()
        .checked_mul(JSON_SCALAR_COLUMN_MEMORY_RESERVATION_MULTIPLIER)
        .ok_or_else(|| json_document_memory_reservation_overflow(input_bytes))?;
    let scalar_column_max = scalar_column_per_row
        .checked_mul(input_bytes)
        .ok_or_else(|| json_document_memory_reservation_overflow(input_bytes))?;
    base.checked_add(scalar_column_max)
        .ok_or_else(|| json_document_memory_reservation_overflow(input_bytes))
}

fn json_document_memory_reservation_overflow(input_bytes: usize) -> DataFusionError {
    DataFusionError::ResourcesExhausted(format!("JSON DOCUMENT memory reservation overflow for {input_bytes} input bytes"))
}

fn classify_json_document_parse_error(error: std::io::Error) -> Box<dyn std::error::Error + Send + Sync> {
    if let Some(select_error) = error.get_ref().and_then(|source| source.downcast_ref::<SelectError>()) {
        Box::new(select_error.clone())
    } else if error.kind() == std::io::ErrorKind::InvalidData {
        Box::new(SelectError::JsonParsingError)
    } else {
        Box::new(error)
    }
}

struct PendingQueryExecutionGuard {
    state: Arc<Mutex<QueryExecutionGuardState>>,
}

enum QueryExecutionGuardState {
    Pending(Option<QueryExecutionTracker>),
    Started,
    Cancelled,
}

impl PendingQueryExecutionGuard {
    fn new(query_tracker: Option<QueryExecutionTracker>) -> Self {
        Self {
            state: Arc::new(Mutex::new(QueryExecutionGuardState::Pending(query_tracker))),
        }
    }

    fn task_state(&self) -> Arc<Mutex<QueryExecutionGuardState>> {
        Arc::clone(&self.state)
    }

    fn start(state: &Mutex<QueryExecutionGuardState>) -> std::io::Result<Option<QueryExecutionGuard>> {
        let mut state = state.lock();
        match std::mem::replace(&mut *state, QueryExecutionGuardState::Started) {
            QueryExecutionGuardState::Pending(None) => Ok(None),
            QueryExecutionGuardState::Pending(Some(query_tracker)) => query_tracker
                .query_guard()
                .map(Some)
                .ok_or_else(json_document_parse_interrupted_error),
            QueryExecutionGuardState::Cancelled => {
                *state = QueryExecutionGuardState::Cancelled;
                Err(json_document_parse_interrupted_error())
            }
            QueryExecutionGuardState::Started => {
                *state = QueryExecutionGuardState::Started;
                Err(std::io::Error::other("JSON DOCUMENT parse started more than once"))
            }
        }
    }
}

fn json_document_parse_interrupted_error() -> std::io::Error {
    std::io::Error::new(std::io::ErrorKind::Interrupted, "JSON DOCUMENT parse was cancelled before it started")
}

impl Drop for PendingQueryExecutionGuard {
    fn drop(&mut self) {
        let query_guard = {
            let mut state = self.state.lock();
            match std::mem::replace(&mut *state, QueryExecutionGuardState::Cancelled) {
                QueryExecutionGuardState::Pending(query_guard) => query_guard,
                QueryExecutionGuardState::Started => {
                    *state = QueryExecutionGuardState::Started;
                    None
                }
                QueryExecutionGuardState::Cancelled => None,
            }
        };
        drop(query_guard);
    }
}

/// Parse a JSON DOCUMENT (a single JSON value, possibly multi-line) into a
/// list of NDJSON lines – one [`Bytes`] per record.
///
/// `json_source` is produced from the SQL AST and expands nested source
/// arrays before DataFusion infers the table schema.
///
/// - A JSON array → one line per element.
/// - A JSON object or scalar root → one line.
#[cfg(test)]
fn parse_json_document_to_lines(bytes: &[u8], json_source: &JsonSource) -> std::io::Result<Vec<Bytes>> {
    parse_json_document_to_lines_cancellable(bytes, json_source, &AtomicBool::new(false))
}

fn parse_json_document_to_lines_cancellable(
    bytes: &[u8],
    json_source: &JsonSource,
    cancelled: &AtomicBool,
) -> std::io::Result<Vec<Bytes>> {
    let reader = std::io::BufReader::with_capacity(
        JSON_CANCELLATION_CHECK_BYTES,
        CancellableJsonReader {
            inner: std::io::Cursor::new(bytes),
            cancelled,
        },
    );
    let root: serde_json::Value = match serde_json::from_reader(reader) {
        Ok(root) => root,
        Err(error) => {
            ensure_json_parse_active(cancelled)?;
            return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, error));
        }
    };
    ensure_json_parse_active(cancelled)?;
    let json_source_path = json_source.path();
    let values = expand_json_source(root, json_source_path, cancelled)?;
    // Preserve the two pre-path-AST forms that flattened arrays implicitly.
    // Explicit indexes and wildcards already identify the intended records
    // and must not flatten an array-valued result a second time.
    let implicitly_expand_arrays = matches!(json_source_path, [] | [JsonPathSegment::Key { .. }]);
    let scalar_column = json_source.scalar_column().unwrap_or_else(|| match json_source_path.last() {
        Some(JsonPathSegment::Key { name, .. }) => name,
        Some(JsonPathSegment::Index(_) | JsonPathSegment::ArrayWildcard | JsonPathSegment::ObjectWildcard) | None => "_1",
    });
    let mut lines: Vec<Bytes> = Vec::new();
    for value in values {
        ensure_json_parse_active(cancelled)?;
        match value {
            serde_json::Value::Array(array) if implicitly_expand_arrays => {
                for item in array {
                    ensure_json_parse_active(cancelled)?;
                    lines.push(json_value_to_line_cancellable(item, scalar_column, cancelled)?);
                }
            }
            other => lines.push(json_value_to_line_cancellable(other, scalar_column, cancelled)?),
        }
    }
    Ok(lines)
}

fn expand_json_source(
    root: serde_json::Value,
    json_source_path: &[JsonPathSegment],
    cancelled: &AtomicBool,
) -> std::io::Result<Vec<serde_json::Value>> {
    // S3Object[*] identifies the input record stream. JSON DOCUMENT already
    // presents the root value as that stream, so the leading marker is not a
    // lookup against the root object.
    let (path, mut values) = match json_source_path.strip_prefix(&[JsonPathSegment::ArrayWildcard]) {
        // Preserve RustFS's existing S3Object[*] root-array expansion while
        // also allowing the AWS canonical S3Object[*][*] form.
        Some(path) => match (root, path.first()) {
            (root @ serde_json::Value::Array(_), Some(JsonPathSegment::ArrayWildcard | JsonPathSegment::Index(_))) => {
                (path, vec![root])
            }
            (serde_json::Value::Array(array), _) => (path, array),
            (root, _) => (path, vec![root]),
        },
        None => (json_source_path, vec![root]),
    };

    for segment in path {
        ensure_json_parse_active(cancelled)?;
        let mut expanded = Vec::new();
        for value in values {
            ensure_json_parse_active(cancelled)?;
            match (segment, value) {
                (JsonPathSegment::Key { name, quoted }, serde_json::Value::Object(mut object)) => {
                    if let Some(value) = remove_json_source_key(&mut object, name, *quoted, cancelled)? {
                        expanded.push(value);
                    }
                }
                (JsonPathSegment::Index(index), serde_json::Value::Array(array)) => {
                    if let Some(value) = array.into_iter().nth(*index) {
                        expanded.push(value);
                    }
                }
                (JsonPathSegment::ArrayWildcard, serde_json::Value::Array(mut array)) => {
                    if expanded.is_empty() {
                        expanded = array;
                    } else {
                        expanded.append(&mut array);
                    }
                }
                (JsonPathSegment::ObjectWildcard, serde_json::Value::Object(object)) => {
                    expanded.extend(object.into_values());
                }
                (JsonPathSegment::Key { .. }, _)
                | (JsonPathSegment::Index(_), _)
                | (JsonPathSegment::ArrayWildcard, _)
                | (JsonPathSegment::ObjectWildcard, _) => {
                    return Err(invalid_json_source_path("JSON source path segment does not match the input value"));
                }
            }
        }
        values = expanded;
    }

    Ok(values)
}

fn remove_json_source_key(
    object: &mut serde_json::Map<String, serde_json::Value>,
    name: &str,
    quoted: bool,
    cancelled: &AtomicBool,
) -> std::io::Result<Option<serde_json::Value>> {
    let mut checkpoint = || ensure_json_parse_active(cancelled);
    remove_json_source_key_with_checkpoint(object, name, quoted, &mut checkpoint)
}

fn remove_json_source_key_with_checkpoint(
    object: &mut serde_json::Map<String, serde_json::Value>,
    name: &str,
    quoted: bool,
    checkpoint: &mut impl FnMut() -> std::io::Result<()>,
) -> std::io::Result<Option<serde_json::Value>> {
    if quoted {
        return Ok(object.remove(name));
    }

    let mut matched = None;
    for (index, key) in object.keys().enumerate() {
        if index % JSON_CANCELLATION_CHECK_KEYS == 0 {
            checkpoint()?;
        }
        if json_key_eq_ignore_ascii_case_with_checkpoint(key, name, checkpoint)? {
            if matched.is_some() {
                return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, SelectError::AmbiguousFieldName));
            }
            matched = Some(key.clone());
        }
    }
    Ok(matched.and_then(|key| object.remove(&key)))
}

fn json_key_eq_ignore_ascii_case_with_checkpoint(
    key: &str,
    expected: &str,
    checkpoint: &mut impl FnMut() -> std::io::Result<()>,
) -> std::io::Result<bool> {
    if key.len() != expected.len() {
        return Ok(false);
    }

    for (key_chunk, expected_chunk) in key
        .as_bytes()
        .chunks(JSON_CANCELLATION_CHECK_BYTES)
        .zip(expected.as_bytes().chunks(JSON_CANCELLATION_CHECK_BYTES))
    {
        checkpoint()?;
        if !key_chunk.eq_ignore_ascii_case(expected_chunk) {
            return Ok(false);
        }
    }
    Ok(true)
}

#[cfg(test)]
fn json_value_to_line(value: serde_json::Value, scalar_column: &str) -> std::io::Result<Bytes> {
    json_value_to_line_cancellable(value, scalar_column, &AtomicBool::new(false))
}

fn json_value_to_line_cancellable(
    value: serde_json::Value,
    scalar_column: &str,
    cancelled: &AtomicBool,
) -> std::io::Result<Bytes> {
    let value = match value {
        value @ serde_json::Value::Object(_) => value,
        value => {
            let mut row = serde_json::Map::new();
            row.insert(scalar_column.to_string(), value);
            serde_json::Value::Object(row)
        }
    };
    let mut line = Vec::new();
    let (serialize_result, limit_exceeded) = {
        let mut writer = CancellableJsonWriter {
            inner: &mut line,
            cancelled,
            bytes_since_check: 0,
            limit_exceeded: false,
        };
        let serialize_result = serde_json::to_writer(&mut writer, &value);
        (serialize_result, writer.limit_exceeded)
    };
    if limit_exceeded {
        return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, SelectError::OverMaxRecordSize));
    }
    if let Err(error) = serialize_result {
        ensure_json_parse_active(cancelled)?;
        return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, error));
    }
    ensure_json_parse_active(cancelled)?;
    line.push(b'\n');
    Ok(Bytes::from(line))
}

struct CancellableJsonWriter<'a> {
    inner: &'a mut Vec<u8>,
    cancelled: &'a AtomicBool,
    bytes_since_check: usize,
    limit_exceeded: bool,
}

impl std::io::Write for CancellableJsonWriter<'_> {
    fn write(&mut self, buffer: &[u8]) -> std::io::Result<usize> {
        let Some(new_len) = self.inner.len().checked_add(buffer.len()) else {
            self.limit_exceeded = true;
            return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, SelectError::OverMaxRecordSize));
        };
        if new_len > MAX_SELECT_RECORD_BYTES {
            self.limit_exceeded = true;
            return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, SelectError::OverMaxRecordSize));
        }
        self.bytes_since_check = self.bytes_since_check.saturating_add(buffer.len());
        if self.bytes_since_check >= JSON_CANCELLATION_CHECK_BYTES {
            ensure_json_parse_active(self.cancelled)?;
            self.bytes_since_check %= JSON_CANCELLATION_CHECK_BYTES;
        }
        self.inner.extend_from_slice(buffer);
        Ok(buffer.len())
    }

    fn flush(&mut self) -> std::io::Result<()> {
        ensure_json_parse_active(self.cancelled)?;
        self.bytes_since_check = 0;
        Ok(())
    }
}

fn invalid_json_source_path(message: &'static str) -> std::io::Error {
    std::io::Error::new(std::io::ErrorKind::InvalidData, message)
}

/// Convert a JSON DOCUMENT to a single concatenated NDJSON [`Bytes`] blob.
///
/// This is a convenience wrapper around [`parse_json_document_to_lines`] used
/// by the unit tests.  Production code uses `json_document_ndjson_stream`
/// instead, which streams lines lazily without constructing this intermediate
/// blob.
#[cfg(test)]
fn flatten_json_document_to_ndjson(bytes: &[u8], json_source_path: &[JsonPathSegment]) -> std::io::Result<Bytes> {
    let lines = parse_json_document_to_lines(bytes, &JsonSource::from_path(json_source_path.to_vec()))?;
    let total = lines.iter().map(|b| b.len()).sum();
    let mut output = Vec::with_capacity(total);
    for line in lines {
        output.extend_from_slice(&line);
    }
    Ok(Bytes::from(output))
}

fn meter_uncompressed_input_stream<S, E>(
    stream: S,
    input_metrics: Arc<SelectInputMetrics>,
) -> impl Stream<Item = std::result::Result<Bytes, E>> + Send + 'static
where
    S: Stream<Item = std::result::Result<Bytes, E>> + Send + 'static,
    E: Send + 'static,
{
    let input_metrics = input_metrics.recorder();
    stream.inspect_ok(move |bytes| input_metrics.record_uncompressed(bytes.len()))
}

pub fn bytes_stream<S>(stream: S, content_length: usize) -> impl Stream<Item = Result<Bytes>> + Send + 'static
where
    S: Stream<Item = Result<Bytes, std::io::Error>> + Send + 'static,
{
    AsyncTryStream::<Bytes, o_Error, _>::new(|mut y| async move {
        pin_mut!(stream);
        let mut remaining: usize = content_length;
        while remaining > 0 {
            let Some(result) = stream.next().await else {
                break;
            };
            let mut bytes = result.map_err(|e| o_Error::Generic {
                store: "",
                source: Box::new(e),
            })?;
            if bytes.len() > remaining {
                bytes.truncate(remaining);
            }
            remaining -= bytes.len();
            y.yield_ok(bytes).await;
        }
        if remaining > 0 {
            return Err(incomplete_object_stream_error(remaining));
        }
        Ok(())
    })
}

fn validate_json_document_size(original_size: u64) -> Result<()> {
    if original_size <= MAX_JSON_DOCUMENT_BYTES {
        return Ok(());
    }

    Err(o_Error::Generic {
        store: "EcObjectStore",
        source: Box::new(DataFusionError::ResourcesExhausted(format!(
            "JSON DOCUMENT object is {original_size} bytes, which exceeds the maximum allowed size of \
             {MAX_JSON_DOCUMENT_BYTES} bytes ({} MiB). Convert the input to JSON LINES (NDJSON) to process large files.",
            MAX_JSON_DOCUMENT_BYTES / (1024 * 1024)
        ))),
    })
}

fn incomplete_object_stream_error(remaining: impl std::fmt::Display) -> o_Error {
    o_Error::Generic {
        store: "EcObjectStore",
        source: Box::new(std::io::Error::new(
            std::io::ErrorKind::UnexpectedEof,
            format!("object stream ended with {remaining} bytes remaining"),
        )),
    }
}

#[cfg(test)]
mod test {
    use super::{
        EcObjectStore, EcObjectStoreBuildError, JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER, JsonDocumentReadMode, OnceCell,
        SELECT_DEFAULT_READ_BUFFER_SIZE, SelectObjectOptions, SelectObjectSnapshot, SelectScanRange, SnapshotConsistencyError,
        bytes_stream, compressed_json_document_ndjson_stream, convert_csv_delimiter_stream, convert_field_delimiter_stream,
        convert_record_delimiter_stream, find_delimiter, flatten_json_document_to_ndjson, grow_json_document_buffer,
        http_range_spec_from_get_range, json_document_ndjson_stream, json_document_ndjson_stream_with_parser,
        json_key_eq_ignore_ascii_case_with_checkpoint, legacy_json_source_from_input, map_storage_error,
        meter_uncompressed_input_stream, remove_json_source_key_with_checkpoint, scan_range_from_bounds, scan_range_stream,
        select_read_headers, snapshot_last_modified, validate_json_document_size,
    };
    use crate::input_stream::{CompressionFormat, MAX_SELECT_RECORD_BYTES, compressed_input_reader, encode_compressed_fixture};
    use crate::query::ast::{JsonPathSegment, JsonSource};
    use crate::query::session::{QueryExecutionGuard, QueryExecutionOwner, QueryExecutionTracker};
    use crate::storage_api::SelectPutObjReader;
    use crate::storage_api::object_store::ObjectIO as _;
    use crate::{QueryError, SelectError, SelectInputMetrics, SelectStorageError};
    use bytes::Bytes;
    use datafusion::{
        common::DataFusionError,
        execution::memory_pool::{GreedyMemoryPool, MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation},
        execution::{config::SessionConfig, context::SessionContext},
        object_store::{self, GetOptions, GetRange, GetResultPayload, ObjectStore as _, path::Path},
        physical_plan::ExecutionPlanProperties,
        prelude::CsvReadOptions,
    };
    use futures::{StreamExt, TryStreamExt, stream};
    use http::HeaderMap;
    use parking_lot::Mutex;
    use rustfs_test_utils::PutObjectCommitBarrier;
    use s3s::S3ErrorCode;
    use s3s::dto::{
        CSVInput, CSVOutput, CompressionType, ExpressionType, FileHeaderInfo, InputSerialization, JSONInput, JSONOutput,
        JSONType, OutputSerialization, ScanRange, SelectObjectContentInput, SelectObjectContentRequest,
    };
    use s3s::header::{
        X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_ALGORITHM, X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY,
        X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY_MD5,
    };
    use std::sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    };

    use tokio::{io::AsyncReadExt, sync::Semaphore};

    #[derive(Debug)]
    struct AdmissionObservingMemoryPool {
        inner: GreedyMemoryPool,
        admission: Arc<Semaphore>,
        reservation_release: Mutex<Option<tokio::sync::oneshot::Sender<bool>>>,
    }

    impl AdmissionObservingMemoryPool {
        fn new(pool_size: usize, admission: Arc<Semaphore>) -> (Self, tokio::sync::oneshot::Receiver<bool>) {
            let (reservation_release, reservation_released) = tokio::sync::oneshot::channel();
            (
                Self {
                    inner: GreedyMemoryPool::new(pool_size),
                    admission,
                    reservation_release: Mutex::new(Some(reservation_release)),
                },
                reservation_released,
            )
        }
    }

    impl std::fmt::Display for AdmissionObservingMemoryPool {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            std::fmt::Display::fmt(&self.inner, f)
        }
    }

    impl MemoryPool for AdmissionObservingMemoryPool {
        fn name(&self) -> &str {
            self.inner.name()
        }

        fn grow(&self, reservation: &MemoryReservation, additional: usize) {
            self.inner.grow(reservation, additional);
        }

        fn shrink(&self, reservation: &MemoryReservation, shrink: usize) {
            self.inner.shrink(reservation, shrink);
            if self.inner.reserved() == 0
                && let Some(reservation_release) = self.reservation_release.lock().take()
            {
                let _ = reservation_release.send(self.admission.available_permits() == 0);
            }
        }

        fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> datafusion::common::Result<()> {
            self.inner.try_grow(reservation, additional)
        }

        fn reserved(&self) -> usize {
            self.inner.reserved()
        }

        fn memory_limit(&self) -> MemoryLimit {
            self.inner.memory_limit()
        }
    }

    fn source_key(name: &str) -> JsonPathSegment {
        JsonPathSegment::Key {
            name: name.to_string(),
            quoted: false,
        }
    }

    fn quoted_source_key(name: &str) -> JsonPathSegment {
        JsonPathSegment::Key {
            name: name.to_string(),
            quoted: true,
        }
    }

    fn csv_input(bucket: &str, object: &str) -> Arc<SelectObjectContentInput> {
        Arc::new(SelectObjectContentInput {
            bucket: bucket.to_string(),
            expected_bucket_owner: None,
            key: object.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,
            },
        })
    }

    fn json_document_input(bucket: &str, object: &str, expression: &str) -> Arc<SelectObjectContentInput> {
        Arc::new(SelectObjectContentInput {
            bucket: bucket.to_string(),
            expected_bucket_owner: None,
            key: object.to_string(),
            sse_customer_algorithm: None,
            sse_customer_key: None,
            sse_customer_key_md5: None,
            request: SelectObjectContentRequest {
                expression: expression.to_string(),
                expression_type: ExpressionType::from_static(ExpressionType::SQL),
                input_serialization: InputSerialization {
                    json: Some(JSONInput {
                        type_: Some(JSONType::from_static(JSONType::DOCUMENT)),
                    }),
                    ..Default::default()
                },
                output_serialization: OutputSerialization {
                    csv: Some(CSVOutput::default()),
                    ..Default::default()
                },
                request_progress: None,
                scan_range: None,
            },
        })
    }

    #[test]
    fn legacy_source_adapter_normalizes_scalar_bindings() {
        let input = json_document_input("bucket", "input.json", "SELECT V FROM S3Object AS V");
        let implicit = json_document_input("bucket", "input.json", "SELECT S3Object FROM S3Object");
        let wildcard = json_document_input("bucket", "input.json", "SELECT _1 FROM S3Object[*]");
        let key = json_document_input("bucket", "input.json", "SELECT * FROM S3Object.LongKey");
        let quoted_key = json_document_input("bucket", "input.json", "SELECT * FROM S3Object.\"LongKey\"");

        assert_eq!(legacy_json_source_from_input(&input), JsonSource::new(Vec::new(), Some("v".to_string())));
        assert_eq!(
            legacy_json_source_from_input(&implicit),
            JsonSource::new(Vec::new(), Some("s3object".to_string()))
        );
        assert_eq!(
            legacy_json_source_from_input(&wildcard),
            JsonSource::new(vec![JsonPathSegment::ArrayWildcard], Some("_1".to_string()))
        );
        assert_eq!(legacy_json_source_from_input(&key).scalar_column(), Some("longkey"));
        assert_eq!(legacy_json_source_from_input(&quoted_key).scalar_column(), Some("LongKey"));
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[serial_test::serial]
    async fn legacy_public_constructors_preserve_single_key_json_source() {
        const BUCKET: &str = "s3select-legacy-json-source";
        const OBJECT: &str = "input.json";

        let env = crate::storage_api::select_test_ecstore_env().await;
        env.make_bucket(BUCKET, false).await;
        env.put_object_bytes(BUCKET, OBJECT, br#"{"employees":[{"name":"Alice"}]}"#.to_vec())
            .await;
        let input = json_document_input(BUCKET, OBJECT, "SELECT e.name FROM S3Object.employees AS e");
        let snapshot = prepare_test_snapshot(BUCKET, OBJECT).await;

        let lazy = EcObjectStore::new(Arc::clone(&input)).expect("legacy lazy constructor should resolve the global store");
        let pinned =
            EcObjectStore::new_with_snapshot(input, snapshot).expect("legacy pinned constructor should accept the snapshot");
        let expected = JsonSource::new(vec![source_key("employees")], Some("e".to_string()));

        assert_eq!(lazy.json_source, expected);
        assert_eq!(pinned.json_source, expected);
    }

    fn json_input(bucket: &str, object: &str, json_type: &'static str) -> Arc<SelectObjectContentInput> {
        let mut input = (*csv_input(bucket, object)).clone();
        input.request.input_serialization = InputSerialization {
            json: Some(JSONInput {
                type_: Some(JSONType::from_static(json_type)),
            }),
            ..Default::default()
        };
        input.request.output_serialization = OutputSerialization {
            json: Some(JSONOutput::default()),
            ..Default::default()
        };
        Arc::new(input)
    }

    #[test]
    fn lazy_snapshot_headers_preserve_ssec_context() {
        let mut input = (*csv_input("bucket", "object.csv")).clone();
        input.sse_customer_algorithm = Some("AES256".to_string());
        input.sse_customer_key = Some("customer-key".to_string());
        input.sse_customer_key_md5 = Some("customer-key-md5".to_string());

        let headers = select_read_headers(&input);

        assert_eq!(
            headers
                .get(X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_ALGORITHM)
                .and_then(|value| value.to_str().ok()),
            Some("AES256")
        );
        assert_eq!(
            headers
                .get(X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY)
                .and_then(|value| value.to_str().ok()),
            Some("customer-key")
        );
        assert_eq!(
            headers
                .get(X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY_MD5)
                .and_then(|value| value.to_str().ok()),
            Some("customer-key-md5")
        );
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[serial_test::serial]
    async fn legacy_constructor_retries_snapshot_after_not_found() {
        let env = crate::storage_api::select_test_ecstore_env().await;
        let bucket = "s3select-lazy-snapshot-retry";
        let object = "input.csv";
        env.make_bucket(bucket, false).await;
        let store = EcObjectStore::new(csv_input(bucket, object)).expect("legacy constructor should resolve the global store");

        let error = store
            .get_opts(
                &Path::from(object),
                GetOptions {
                    head: true,
                    ..Default::default()
                },
            )
            .await
            .expect_err("missing object should remain a typed not-found error");
        assert!(matches!(error, object_store::Error::NotFound { .. }));

        env.put_object_bytes(bucket, object, b"id,name\n1,Alice\n".to_vec()).await;
        let result = store
            .get_opts(
                &Path::from(object),
                GetOptions {
                    head: true,
                    ..Default::default()
                },
            )
            .await
            .expect("failed snapshot initialization must not be cached");
        assert_eq!(result.meta.size, 16);
    }

    #[tokio::test]
    #[serial_test::serial]
    async fn legacy_constructor_maps_missing_bucket_to_not_found() {
        let _env = crate::storage_api::select_test_ecstore_env().await;
        let bucket = "s3select-lazy-snapshot-missing-bucket";
        let object = "input.csv";
        let store = EcObjectStore::new(csv_input(bucket, object)).expect("legacy constructor should resolve the global store");

        let error = store
            .get_opts(
                &Path::from(object),
                GetOptions {
                    head: true,
                    ..Default::default()
                },
            )
            .await
            .expect_err("missing bucket should remain a typed not-found error");

        assert!(matches!(error, object_store::Error::NotFound { .. }));
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[serial_test::serial]
    async fn legacy_constructor_reuses_head_snapshot_for_body() {
        const BUCKET: &str = "s3select-lazy-snapshot-head-body";
        const OBJECT: &str = "input.csv";
        const OLD_DATA: &[u8] = b"id,name\n1,old\n";
        const NEW_DATA: &[u8] = b"id,name\n1,new\n";

        let env = crate::storage_api::select_test_ecstore_env().await;
        env.make_bucket(BUCKET, false).await;
        env.put_object_bytes(BUCKET, OBJECT, OLD_DATA.to_vec()).await;
        let store = EcObjectStore::new(csv_input(BUCKET, OBJECT)).expect("legacy constructor should resolve the global store");
        let head = store
            .get_opts(
                &Path::from(OBJECT),
                GetOptions {
                    head: true,
                    ..Default::default()
                },
            )
            .await
            .expect("HEAD should lazily prepare the snapshot");
        assert_eq!(head.meta.size, u64::try_from(OLD_DATA.len()).expect("fixture length should fit in u64"));

        let commit_barrier = PutObjectCommitBarrier::before_namespace(BUCKET, OBJECT);
        let writer = tokio::spawn(async move {
            env.put_object_bytes(BUCKET, OBJECT, NEW_DATA.to_vec()).await;
        });
        commit_barrier.wait_until_paused().await;
        commit_barrier.release_and_wait_until_namespace_pending().await;
        assert!(!writer.is_finished(), "overwrite must wait for the lazy snapshot");

        let result = store
            .get_opts(&Path::from(OBJECT), GetOptions::default())
            .await
            .expect("body should reuse the HEAD snapshot");
        let GetResultPayload::Stream(stream) = result.payload else {
            panic!("expected streaming snapshot body");
        };
        let bytes = stream
            .try_collect::<Vec<_>>()
            .await
            .expect("collect lazy snapshot body")
            .concat();
        assert_eq!(bytes, OLD_DATA);
        assert_eq!(store.reader_open_count.load(Ordering::Relaxed), 1);

        drop(store);
        tokio::time::timeout(std::time::Duration::from_secs(5), writer)
            .await
            .expect("overwrite should finish after the lazy snapshot is released")
            .expect("overwrite task should join");
        assert_eq!(read_current_object(BUCKET, OBJECT).await, NEW_DATA);
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[serial_test::serial]
    async fn legacy_constructor_first_version_pins_later_reads() {
        const BUCKET: &str = "s3select-lazy-snapshot-version";
        const OBJECT: &str = "input.csv";
        const OLD_DATA: &[u8] = b"old-marker\n";
        const NEW_DATA: &[u8] = b"new-poison-value\n";

        let env = crate::storage_api::select_test_ecstore_env().await;
        env.make_bucket(BUCKET, true).await;
        let versioned_opts = SelectObjectOptions {
            versioned: true,
            ..Default::default()
        };
        let mut old_reader = SelectPutObjReader::from_vec(OLD_DATA.to_vec());
        let old_info = env
            .ecstore
            .put_object(BUCKET, OBJECT, &mut old_reader, &versioned_opts)
            .await
            .expect("put old version fixture");
        let old_version = old_info
            .version_id
            .expect("versioned PUT should return a version ID")
            .to_string();
        let mut new_reader = SelectPutObjReader::from_vec(NEW_DATA.to_vec());
        let new_info = env
            .ecstore
            .put_object(BUCKET, OBJECT, &mut new_reader, &versioned_opts)
            .await
            .expect("put latest version poison fixture");
        let new_version = new_info
            .version_id
            .expect("versioned PUT should return a version ID")
            .to_string();

        let store = EcObjectStore::new(csv_input(BUCKET, OBJECT)).expect("legacy constructor should resolve the global store");
        let head = store
            .get_opts(
                &Path::from(OBJECT),
                GetOptions {
                    head: true,
                    version: Some(old_version.to_uppercase()),
                    ..Default::default()
                },
            )
            .await
            .expect("first HEAD should bind the requested old version");
        assert_eq!(head.meta.version.as_deref(), Some(old_version.as_str()));

        let mismatch = store
            .get_opts(
                &Path::from(OBJECT),
                GetOptions {
                    head: true,
                    version: Some(new_version),
                    ..Default::default()
                },
            )
            .await
            .expect_err("an explicit different version must not reuse the pinned snapshot");
        assert!(mismatch.to_string().contains("different object version"));

        let range = store
            .get_opts(
                &Path::from(OBJECT),
                GetOptions {
                    range: Some(GetRange::Bounded(0..3)),
                    ..Default::default()
                },
            )
            .await
            .expect("later range should reuse the old-version snapshot");
        let GetResultPayload::Stream(range_stream) = range.payload else {
            panic!("expected ranged snapshot stream");
        };
        assert_eq!(
            range_stream
                .try_collect::<Vec<_>>()
                .await
                .expect("collect old-version range")
                .concat(),
            b"old"
        );

        let body = store
            .get_opts(&Path::from(OBJECT), GetOptions::default())
            .await
            .expect("later body should reuse the old-version snapshot");
        let GetResultPayload::Stream(body_stream) = body.payload else {
            panic!("expected full snapshot stream");
        };
        assert_eq!(
            body_stream
                .try_collect::<Vec<_>>()
                .await
                .expect("collect old-version body")
                .concat(),
            OLD_DATA
        );
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[serial_test::serial]
    async fn legacy_constructor_normalizes_null_version_before_snapshot_prepare() {
        const BUCKET: &str = "s3select-lazy-snapshot-null-version";
        const OBJECT: &str = "input.csv";
        const DATA: &[u8] = b"null-version-marker\n";

        let env = crate::storage_api::select_test_ecstore_env().await;
        env.make_bucket(BUCKET, false).await;
        env.put_object_bytes(BUCKET, OBJECT, DATA.to_vec()).await;
        let store = EcObjectStore::new(csv_input(BUCKET, OBJECT)).expect("legacy constructor should resolve the global store");

        let head = store
            .get_opts(
                &Path::from(OBJECT),
                GetOptions {
                    head: true,
                    version: Some("NULL".to_string()),
                    ..Default::default()
                },
            )
            .await
            .expect("null version should prepare an unversioned snapshot");
        assert!(head.meta.version.is_none());

        let body = store
            .get_opts(
                &Path::from(OBJECT),
                GetOptions {
                    version: Some(uuid::Uuid::nil().to_string()),
                    ..Default::default()
                },
            )
            .await
            .expect("nil UUID should match the pinned null-version snapshot");
        let GetResultPayload::Stream(body_stream) = body.payload else {
            panic!("expected null-version snapshot stream");
        };
        assert_eq!(
            body_stream
                .try_collect::<Vec<_>>()
                .await
                .expect("collect null-version body")
                .concat(),
            DATA
        );
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[serial_test::serial]
    async fn prepared_snapshot_rejects_a_different_query_object() {
        const BUCKET: &str = "s3select-snapshot-identity";
        const OBJECT: &str = "source.csv";

        let env = crate::storage_api::select_test_ecstore_env().await;
        env.make_bucket(BUCKET, false).await;
        env.put_object_bytes(BUCKET, OBJECT, b"source-marker\n".to_vec()).await;
        let snapshot = Arc::new(
            env.ecstore
                .prepare_select_object_snapshot(BUCKET, OBJECT, &HeaderMap::new(), &Default::default())
                .await
                .expect("prepare source snapshot"),
        );

        let error = EcObjectStore::new_with_snapshot(csv_input(BUCKET, "different.csv"), snapshot)
            .expect_err("a snapshot must remain bound to its source object");

        assert_eq!(error.code(), &S3ErrorCode::InternalError);
        assert!(error.source().is_some_and(|source| {
            source
                .downcast_ref::<EcObjectStoreBuildError>()
                .is_some_and(|error| matches!(error, EcObjectStoreBuildError::Snapshot(SnapshotConsistencyError::ObjectChanged)))
        }));
    }

    async fn prepare_test_snapshot(bucket: &str, object: &str) -> Arc<SelectObjectSnapshot> {
        let env = crate::storage_api::select_test_ecstore_env().await;
        Arc::new(
            env.ecstore
                .prepare_select_object_snapshot(bucket, object, &HeaderMap::new(), &Default::default())
                .await
                .expect("prepare SelectObjectContent snapshot"),
        )
    }

    fn scan_range_csv_store(
        bucket: &str,
        object: &str,
        snapshot: Arc<SelectObjectSnapshot>,
        record_delimiter: &str,
        file_header_info: Option<FileHeaderInfo>,
        start: i64,
        end: i64,
    ) -> EcObjectStore {
        EcObjectStore::new_with_snapshot(
            Arc::new(SelectObjectContentInput {
                bucket: bucket.to_string(),
                expected_bucket_owner: None,
                key: object.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 {
                            record_delimiter: Some(record_delimiter.to_string()),
                            file_header_info,
                            ..Default::default()
                        }),
                        ..Default::default()
                    },
                    output_serialization: OutputSerialization {
                        csv: Some(CSVOutput::default()),
                        ..Default::default()
                    },
                    request_progress: None,
                    scan_range: Some(ScanRange {
                        start: Some(start),
                        end: Some(end),
                    }),
                },
            }),
            snapshot,
        )
        .expect("snapshot should match SelectObjectContent input")
    }

    async fn read_current_object(bucket: &str, object: &str) -> Vec<u8> {
        let snapshot = prepare_test_snapshot(bucket, object).await;
        let mut reader = snapshot.open_reader(None).await.expect("current object reader should open");
        let mut bytes = Vec::new();
        reader
            .stream
            .read_to_end(&mut bytes)
            .await
            .expect("current object should be readable");
        bytes
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[serial_test::serial]
    async fn head_uses_snapshot_metadata_without_opening_body() {
        let env = crate::storage_api::select_test_ecstore_env().await;
        let bucket = "s3select-snapshot-head";
        let object = "input.csv";
        env.make_bucket(bucket, false).await;
        let mut reader = SelectPutObjReader::from_vec(b"id,name\n1,Alice\n".to_vec());
        env.ecstore
            .put_object(bucket, object, &mut reader, &Default::default())
            .await
            .expect("put HEAD fixture");

        let snapshot = prepare_test_snapshot(bucket, object).await;
        let expected_modified = snapshot_last_modified(&snapshot).expect("snapshot modification time");
        let expected_size = snapshot.logical_size();
        let expected_etag = snapshot.object_info().etag.clone();
        let expected_version = snapshot.object_info().version_id.map(|version| version.to_string());
        let input = Arc::new(SelectObjectContentInput {
            bucket: bucket.to_string(),
            expected_bucket_owner: None,
            key: object.to_string(),
            sse_customer_algorithm: Some("secret-algorithm".to_string()),
            sse_customer_key: Some("secret-customer-key".to_string()),
            sse_customer_key_md5: Some("secret-customer-key-md5".to_string()),
            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,
            },
        });
        let store = EcObjectStore::new_with_snapshot(input, snapshot).expect("snapshot should match SelectObjectContent input");
        let debug = format!("{store:?}");
        assert!(!debug.contains("secret-customer-key"));

        let result = store
            .get_opts(
                &Path::from(object),
                GetOptions {
                    head: true,
                    ..Default::default()
                },
            )
            .await
            .expect("HEAD from snapshot metadata");

        assert_eq!(result.meta.last_modified, expected_modified);
        assert_eq!(result.meta.size, expected_size);
        assert_eq!(result.meta.e_tag, expected_etag);
        assert_eq!(result.meta.version, expected_version);
        assert_eq!(store.reader_open_count.load(Ordering::Relaxed), 0);
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[serial_test::serial]
    async fn snapshot_keeps_csv_header_and_body_on_one_generation_during_overwrite() {
        const BUCKET: &str = "s3select-snapshot-header-body-race";
        const OBJECT: &str = "input.csv";
        const OLD_DATA: &[u8] = b"old_header,value\nskip_old,0\nold_body,1\n";
        const NEW_DATA: &[u8] = b"new_header,value\nskip_new,0\nnew_body,1\n";

        let env = crate::storage_api::select_test_ecstore_env().await;
        env.make_bucket(BUCKET, false).await;
        let mut reader = SelectPutObjReader::from_vec(OLD_DATA.to_vec());
        env.ecstore
            .put_object(BUCKET, OBJECT, &mut reader, &Default::default())
            .await
            .expect("put old CSV header/body fixture");

        let selected_start = i64::try_from(b"old_header,value\nskip_old,0\n".len()).expect("fixture offset should fit in i64");
        let snapshot = prepare_test_snapshot(BUCKET, OBJECT).await;
        let store = scan_range_csv_store(
            BUCKET,
            OBJECT,
            snapshot,
            "\n",
            Some(FileHeaderInfo::from_static(FileHeaderInfo::USE)),
            selected_start,
            selected_start,
        );
        let commit_barrier = PutObjectCommitBarrier::before_namespace(BUCKET, OBJECT);
        let writer = tokio::spawn(async move {
            env.put_object_bytes(BUCKET, OBJECT, NEW_DATA.to_vec()).await;
        });
        commit_barrier.wait_until_paused().await;
        commit_barrier.release_and_wait_until_namespace_pending().await;
        assert!(
            !writer.is_finished(),
            "overwrite must remain blocked while the SelectObjectContent snapshot is alive"
        );

        let result = store
            .get_opts(&Path::from(OBJECT), GetOptions::default())
            .await
            .expect("read CSV header and body from one snapshot");
        let GetResultPayload::Stream(stream) = result.payload else {
            panic!("expected streaming CSV header/body payload");
        };
        let bytes = stream
            .try_collect::<Vec<_>>()
            .await
            .expect("collect CSV header/body snapshot")
            .concat();

        assert_eq!(bytes, b"old_header,value\nold_body,1\n");
        assert_eq!(store.reader_open_count.load(Ordering::Relaxed), 2);
        assert!(!writer.is_finished(), "overwrite must remain blocked after both snapshot readers finish");

        drop(store);
        tokio::time::timeout(std::time::Duration::from_secs(5), writer)
            .await
            .expect("overwrite should finish after the snapshot is released")
            .expect("overwrite task should join");
        assert_eq!(read_current_object(BUCKET, OBJECT).await, NEW_DATA);
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[serial_test::serial]
    async fn snapshot_keeps_scan_range_context_and_main_reader_on_one_generation_during_overwrite() {
        const BUCKET: &str = "s3select-snapshot-scan-context-race";
        const OBJECT: &str = "input.csv";
        const OLD_DATA: &[u8] = b"111aaa222aa333aa";
        const NEW_DATA: &[u8] = b"999aaa888aa777aa";

        let env = crate::storage_api::select_test_ecstore_env().await;
        env.make_bucket(BUCKET, false).await;
        let mut reader = SelectPutObjReader::from_vec(OLD_DATA.to_vec());
        env.ecstore
            .put_object(BUCKET, OBJECT, &mut reader, &Default::default())
            .await
            .expect("put old ScanRange context fixture");

        let snapshot = prepare_test_snapshot(BUCKET, OBJECT).await;
        let store = scan_range_csv_store(BUCKET, OBJECT, snapshot, "aa", None, 4, 5);
        let commit_barrier = PutObjectCommitBarrier::before_namespace(BUCKET, OBJECT);
        let writer = tokio::spawn(async move {
            env.put_object_bytes(BUCKET, OBJECT, NEW_DATA.to_vec()).await;
        });
        commit_barrier.wait_until_paused().await;
        commit_barrier.release_and_wait_until_namespace_pending().await;
        assert!(
            !writer.is_finished(),
            "overwrite must remain blocked while the SelectObjectContent snapshot is alive"
        );

        let result = store
            .get_opts(&Path::from(OBJECT), GetOptions::default())
            .await
            .expect("read ScanRange context and main body from one snapshot");
        let GetResultPayload::Stream(stream) = result.payload else {
            panic!("expected streaming ScanRange context payload");
        };
        let bytes = stream
            .try_collect::<Vec<_>>()
            .await
            .expect("collect ScanRange context snapshot")
            .concat();

        assert_eq!(bytes, b"a222\r\n");
        assert_eq!(store.reader_open_count.load(Ordering::Relaxed), 2);
        assert!(
            !writer.is_finished(),
            "overwrite must remain blocked after context and main readers finish"
        );

        drop(store);
        tokio::time::timeout(std::time::Duration::from_secs(5), writer)
            .await
            .expect("overwrite should finish after the snapshot is released")
            .expect("overwrite task should join");
        assert_eq!(read_current_object(BUCKET, OBJECT).await, NEW_DATA);
    }

    #[tokio::test]
    async fn test_scan_range_stream_keeps_header_and_selected_record() {
        let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"h1,h2\n1,a\n2,b\n3,c\n"))]);
        let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(10, 11), true, 0, 18);
        let mut output = Vec::new();
        while let Some(bytes) = stream.next().await {
            output.extend_from_slice(&bytes.unwrap());
        }
        assert_eq!(output, b"h1,h2\n2,b\n");
    }

    #[tokio::test]
    async fn test_scan_range_stream_skips_record_when_start_is_in_middle() {
        let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"1,a\n2,b\n3,c\n"))]);
        let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(2, 7), false, 0, 12);
        let mut output = Vec::new();
        while let Some(bytes) = stream.next().await {
            output.extend_from_slice(&bytes.unwrap());
        }
        assert_eq!(output, b"2,b\n");
    }

    #[tokio::test]
    async fn test_scan_range_stream_keeps_record_when_end_is_in_middle() {
        let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"1,a\n2,b\n3,c\n"))]);
        let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(0, 5), false, 0, 12);
        let mut output = Vec::new();
        while let Some(bytes) = stream.next().await {
            output.extend_from_slice(&bytes.unwrap());
        }
        assert_eq!(output, b"1,a\n2,b\n");
    }

    #[tokio::test]
    async fn test_scan_range_stream_uses_base_offset_for_range_reader() {
        let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"\n2,b\n3,c\n"))]);
        let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(4, 7), false, 3, 12);
        let mut output = Vec::new();
        while let Some(bytes) = stream.next().await {
            output.extend_from_slice(&bytes.unwrap());
        }
        assert_eq!(output, b"2,b\n");
    }

    #[tokio::test]
    async fn test_scan_range_stream_handles_delimiter_split_across_chunks() {
        let chunks = stream::iter(vec![
            Ok::<_, std::io::Error>(Bytes::from_static(b"h1,h2\r")),
            Ok::<_, std::io::Error>(Bytes::from_static(b"\n1,a\r\n2,b\r")),
            Ok::<_, std::io::Error>(Bytes::from_static(b"\n3,c\r\n")),
        ]);
        let mut stream = scan_range_stream(chunks, b"\r\n".to_vec(), SelectScanRange::new(12, 14), true, 0, 22);
        let mut output = Vec::new();
        while let Some(bytes) = stream.next().await {
            output.extend_from_slice(&bytes.unwrap());
        }
        assert_eq!(output, b"h1,h2\r\n2,b\r\n");
    }

    #[tokio::test]
    async fn test_scan_range_stream_converts_custom_delimiter_split_across_chunks() {
        let chunks = stream::iter(vec![
            Ok::<_, std::io::Error>(Bytes::from_static(b"h1,h2^")),
            Ok::<_, std::io::Error>(Bytes::from_static(b"Y1,a^Y2,b^")),
            Ok::<_, std::io::Error>(Bytes::from_static(b"Y3,c^Y")),
        ]);
        let stream = scan_range_stream(chunks, b"^Y".to_vec(), SelectScanRange::new(12, 14), true, 0, 22);
        let mut stream = convert_record_delimiter_stream(stream, b"^Y".to_vec());
        let mut output = Vec::new();
        while let Some(bytes) = stream.next().await {
            output.extend_from_slice(&bytes.expect("custom-delimiter ScanRange chunk should be valid"));
        }
        assert_eq!(output, b"h1,h2\r\n2,b\r\n");
    }

    #[tokio::test]
    async fn test_record_delimiter_conversion_carries_partial_delimiter() {
        let chunks = stream::iter(vec![
            Ok::<_, object_store::Error>(Bytes::from_static(b"a,1^")),
            Ok::<_, object_store::Error>(Bytes::from_static(b"Yb,2^Y")),
        ]);
        let output = convert_record_delimiter_stream(chunks, b"^Y".to_vec())
            .try_collect::<Vec<_>>()
            .await
            .expect("convert record delimiter")
            .concat();
        assert_eq!(output, b"a,1\r\nb,2\r\n");
    }

    #[tokio::test]
    async fn test_record_delimiter_conversion_preserves_overlapping_match_order() {
        let chunks = stream::iter(vec![
            Ok::<_, object_store::Error>(Bytes::from_static(b"a")),
            Ok::<_, object_store::Error>(Bytes::from_static(b"aa")),
            Ok::<_, object_store::Error>(Bytes::from_static(b"a")),
        ]);
        let output = convert_record_delimiter_stream(chunks, b"aa".to_vec())
            .try_collect::<Vec<_>>()
            .await
            .expect("convert overlapping record delimiter")
            .concat();
        assert_eq!(output, b"\r\n\r\n");
    }

    #[tokio::test]
    async fn test_record_and_field_delimiter_conversion_order() {
        let chunks = stream::iter(vec![
            Ok::<_, object_store::Error>(Bytes::from_static(b"a\r")),
            Ok::<_, object_store::Error>(Bytes::from_static(b"\n1^")),
            Ok::<_, object_store::Error>(Bytes::from_static(b"Yb\r\n2^Y")),
        ]);
        let output = convert_csv_delimiter_stream(chunks, Some(b"^Y".to_vec()), Some("\r\n".to_string()))
            .try_collect::<Vec<_>>()
            .await
            .expect("convert record and field delimiters")
            .concat();
        assert_eq!(output, b"a,1\r\nb,2\r\n");
    }

    #[tokio::test]
    async fn test_record_delimiter_takes_precedence_when_delimiters_match() {
        let chunks = stream::iter(vec![
            Ok::<_, object_store::Error>(Bytes::from_static(b"a^")),
            Ok::<_, object_store::Error>(Bytes::from_static(b"Yb^Y")),
        ]);
        let output = convert_csv_delimiter_stream(chunks, Some(b"^Y".to_vec()), Some("^Y".to_string()))
            .try_collect::<Vec<_>>()
            .await
            .expect("convert matching record and field delimiters")
            .concat();
        assert_eq!(output, b"a\r\nb\r\n");

        let chunks = stream::iter(vec![Ok::<_, object_store::Error>(Bytes::from_static(b"a\r\nb\r\n"))]);
        let output = convert_csv_delimiter_stream(chunks, Some(b"\r\n".to_vec()), Some("\r\n".to_string()))
            .try_collect::<Vec<_>>()
            .await
            .expect("preserve matching normalized record delimiter")
            .concat();
        assert_eq!(output, b"a\r\nb\r\n");
    }

    #[tokio::test]
    async fn test_longer_field_delimiter_takes_precedence_over_record_prefix() {
        let chunks = stream::iter(vec![
            Ok::<_, object_store::Error>(Bytes::from_static(b"a^")),
            Ok::<_, object_store::Error>(Bytes::from_static(b"YQb^Y")),
        ]);
        let output = convert_csv_delimiter_stream(chunks, Some(b"^Y".to_vec()), Some("^YQ".to_string()))
            .try_collect::<Vec<_>>()
            .await
            .expect("convert record delimiter that prefixes field delimiter")
            .concat();
        assert_eq!(output, b"a,b\r\n");

        let chunks = stream::iter(vec![Ok::<_, object_store::Error>(Bytes::from_static(b"a\nXb\nc\nXd\n"))]);
        let output = convert_csv_delimiter_stream(chunks, Some(b"\n".to_vec()), Some("\nX".to_string()))
            .try_collect::<Vec<_>>()
            .await
            .expect("preserve longer field delimiter with native record prefix")
            .concat();
        assert_eq!(output, b"a,b\nc,d\n");
    }

    #[tokio::test]
    async fn test_scan_range_stream_rejects_early_eof() {
        let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"1,a\n"))]);
        let mut output = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(0, 7), false, 0, 8);

        assert_eq!(output.next().await.expect("first stream item").expect("first record"), b"1,a\n"[..]);
        let err = output
            .next()
            .await
            .expect("early EOF error")
            .expect_err("short ScanRange stream must fail");
        let object_store::Error::Generic { source, .. } = err else {
            panic!("expected generic object store error");
        };
        let source = source.downcast_ref::<std::io::Error>().expect("I/O error source");
        assert_eq!(source.kind(), std::io::ErrorKind::UnexpectedEof);
        assert!(source.to_string().contains("4 bytes remaining"));
        assert!(output.next().await.is_none());
    }

    #[test]
    fn test_find_delimiter_handles_multi_byte_delimiter() {
        assert_eq!(find_delimiter(b"one\r\ntwo", b"\r\n"), Some(3));
        assert_eq!(find_delimiter(b"one\ntwo", b"\r\n"), None);
    }

    #[test]
    fn test_scan_range_end_only_uses_aws_suffix_semantics() {
        let range = scan_range_from_bounds(None, Some(35), 100).unwrap().unwrap();
        assert_eq!(range.start(), 65);
        assert_eq!(range.end(), 99);
    }

    #[test]
    fn test_scan_range_start_after_object_is_rejected_before_reader() {
        let err = scan_range_from_bounds(Some(100), None, 100).unwrap_err();
        assert!(err.to_string().contains("ScanRange"));
    }

    #[test]
    fn test_scan_range_start_after_end_is_rejected() {
        let err = scan_range_from_bounds(Some(20), Some(10), 100).unwrap_err();
        assert!(err.to_string().contains("ScanRange"));
    }

    #[test]
    fn test_get_range_conversion_for_parquet_bounded_ranges() {
        let range = http_range_spec_from_get_range(&GetRange::Bounded(10..20));
        assert!(!range.is_suffix_length);
        assert_eq!(range.start, 10);
        assert_eq!(range.end, 19);
    }

    #[tokio::test]
    async fn test_scan_range_output_can_convert_field_delimiter() {
        let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"a&&1\nb&&2\n"))]);
        let stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(0, 10), false, 0, 10);
        let mut stream = convert_field_delimiter_stream(stream, "&&".to_string());
        let mut output = Vec::new();
        while let Some(bytes) = stream.next().await {
            output.extend_from_slice(&bytes.unwrap());
        }
        assert_eq!(output, b"a,1\nb,2\n");
    }

    #[tokio::test]
    async fn delimiter_conversion_keeps_uncompressed_metrics_equal() {
        let input = Bytes::from_static(b"a&&1\nb&&2\n");
        let input_metrics = Arc::new(SelectInputMetrics::default());
        let stream = stream::iter([Ok::<_, object_store::Error>(input.clone())]);
        let stream = meter_uncompressed_input_stream(stream, Arc::clone(&input_metrics));
        let output = convert_field_delimiter_stream(stream, "&&".to_string())
            .try_collect::<Vec<_>>()
            .await
            .expect("delimiter conversion should succeed")
            .concat();

        assert_eq!(output, b"a,1\nb,2\n");
        let input_len = u64::try_from(input.len()).expect("fixture length should fit in u64");
        assert_eq!(input_metrics.snapshot().bytes_scanned, input_len);
        assert_eq!(input_metrics.snapshot().bytes_processed, input_len);
    }

    #[tokio::test]
    async fn test_field_delimiter_stream_converts_delimiter_split_across_chunks() {
        let chunks = stream::iter(vec![
            Ok::<_, object_store::Error>(Bytes::from_static(b"a&")),
            Ok::<_, object_store::Error>(Bytes::from_static(b"&1\nb&&2\n")),
        ]);
        let mut stream = convert_field_delimiter_stream(chunks, "&&".to_string());
        let mut output = Vec::new();
        while let Some(bytes) = stream.next().await {
            output.extend_from_slice(&bytes.unwrap());
        }
        assert_eq!(output, b"a,1\nb,2\n");
    }

    #[tokio::test]
    async fn test_field_delimiter_stream_converts_delimiter_at_stream_end() {
        let chunks = stream::iter(vec![
            Ok::<_, object_store::Error>(Bytes::from_static(b"a&")),
            Ok::<_, object_store::Error>(Bytes::from_static(b"&")),
        ]);
        let mut stream = convert_field_delimiter_stream(chunks, "&&".to_string());
        let mut output = Vec::new();
        while let Some(bytes) = stream.next().await {
            output.extend_from_slice(&bytes.unwrap());
        }
        assert_eq!(output, b"a,");
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[serial_test::serial]
    async fn test_scan_range_self_overlapping_delimiter_retains_record_context() {
        let env = crate::storage_api::select_test_ecstore_env().await;
        let bucket = "s3select-scan-range-record-context";
        let object = "input.csv";
        env.make_bucket(bucket, false).await;
        let mut reader = SelectPutObjReader::from_vec(b"111aaa222aa333aa".to_vec());
        env.ecstore
            .put_object(bucket, object, &mut reader, &Default::default())
            .await
            .expect("put self-overlapping delimiter ScanRange fixture");

        let make_store = |start, end, file_header_info, snapshot| EcObjectStore {
            input: Arc::new(SelectObjectContentInput {
                bucket: bucket.to_string(),
                expected_bucket_owner: None,
                key: object.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 {
                            record_delimiter: Some("aa".to_string()),
                            file_header_info,
                            ..Default::default()
                        }),
                        ..Default::default()
                    },
                    output_serialization: OutputSerialization {
                        csv: Some(CSVOutput::default()),
                        ..Default::default()
                    },
                    request_progress: None,
                    scan_range: Some(ScanRange {
                        start: Some(start),
                        end: Some(end),
                    }),
                },
            }),
            need_convert: false,
            delimiter: String::new(),
            is_json_document: false,
            json_source: JsonSource::default(),
            input_metrics: Arc::new(SelectInputMetrics::default()),
            memory_pool: Arc::new(GreedyMemoryPool::new(1024)),
            query_tracker: None,
            store: None,
            snapshot: OnceCell::new_with(Some(snapshot)),
            reader_open_count: Arc::new(AtomicUsize::new(0)),
        };

        let store = make_store(6, 6, None, prepare_test_snapshot(bucket, object).await);
        let result = store
            .get_opts(&Path::from(object), GetOptions::default())
            .await
            .expect("read ScanRange starting inside record data");
        let GetResultPayload::Stream(stream) = result.payload else {
            panic!("expected streaming ScanRange payload");
        };
        let chunks: Vec<Bytes> = stream.try_collect().await.expect("collect record-data ScanRange output");
        assert!(chunks.concat().is_empty());
        drop(store);

        let store = make_store(4, 5, None, prepare_test_snapshot(bucket, object).await);
        let result = store
            .get_opts(&Path::from(object), GetOptions::default())
            .await
            .expect("read ScanRange starting inside overlapping delimiter");
        let GetResultPayload::Stream(stream) = result.payload else {
            panic!("expected streaming overlapping-delimiter payload");
        };
        let chunks: Vec<Bytes> = stream
            .try_collect()
            .await
            .expect("collect overlapping-delimiter ScanRange output");
        assert_eq!(chunks.concat(), b"a222\r\n");
        drop(store);

        let mut reader = SelectPutObjReader::from_vec(b"111aa222aa333aa".to_vec());
        env.ecstore
            .put_object(bucket, object, &mut reader, &Default::default())
            .await
            .expect("put exact review ScanRange fixture");
        let result = make_store(6, 6, None, prepare_test_snapshot(bucket, object).await)
            .get_opts(&Path::from(object), GetOptions::default())
            .await
            .expect("read exact review ScanRange fixture");
        let GetResultPayload::Stream(stream) = result.payload else {
            panic!("expected streaming exact review ScanRange payload");
        };
        let chunks: Vec<Bytes> = stream.try_collect().await.expect("collect exact review ScanRange output");
        assert!(chunks.concat().is_empty());

        let result = make_store(5, 5, None, prepare_test_snapshot(bucket, object).await)
            .get_opts(&Path::from(object), GetOptions::default())
            .await
            .expect("read ScanRange starting after an even delimiter run");
        let GetResultPayload::Stream(stream) = result.payload else {
            panic!("expected streaming even-run ScanRange payload");
        };
        let chunks: Vec<Bytes> = stream.try_collect().await.expect("collect even-run ScanRange output");
        assert_eq!(chunks.concat(), b"222\r\n");

        let mut reader = SelectPutObjReader::from_vec(b"h1aav1aav2aa".to_vec());
        env.ecstore
            .put_object(bucket, object, &mut reader, &Default::default())
            .await
            .expect("put ScanRange header snapshot fixture");
        let result = make_store(
            8,
            8,
            Some(FileHeaderInfo::from_static(FileHeaderInfo::USE)),
            prepare_test_snapshot(bucket, object).await,
        )
        .get_opts(&Path::from(object), GetOptions::default())
        .await
        .expect("read ScanRange with a separate header read");
        let GetResultPayload::Stream(stream) = result.payload else {
            panic!("expected streaming ScanRange header payload");
        };
        let chunks: Vec<Bytes> = stream.try_collect().await.expect("collect ScanRange header output");
        assert_eq!(chunks.concat(), b"h1\r\nv2\r\n");

        let mut reader = SelectPutObjReader::from_vec(b"aaa222aa".to_vec());
        env.ecstore
            .put_object(bucket, object, &mut reader, &Default::default())
            .await
            .expect("put object-start delimiter context fixture");
        let result = make_store(3, 3, None, prepare_test_snapshot(bucket, object).await)
            .get_opts(&Path::from(object), GetOptions::default())
            .await
            .expect("read delimiter context that reaches the object start");
        let GetResultPayload::Stream(stream) = result.payload else {
            panic!("expected streaming object-start context payload");
        };
        let chunks: Vec<Bytes> = stream.try_collect().await.expect("collect object-start context output");
        assert!(chunks.concat().is_empty());

        let run_start = SELECT_DEFAULT_READ_BUFFER_SIZE + 7;
        let mut large_fixture = vec![b'b'; run_start];
        large_fixture.extend_from_slice(b"aaa222aa");
        let mut reader = SelectPutObjReader::from_vec(large_fixture);
        env.ecstore
            .put_object(bucket, object, &mut reader, &Default::default())
            .await
            .expect("put large self-overlapping delimiter ScanRange fixture");
        let scan_start = i64::try_from(run_start + 3).expect("fixture offset should fit in i64");
        let result = make_store(scan_start, scan_start, None, prepare_test_snapshot(bucket, object).await)
            .get_opts(&Path::from(object), GetOptions::default())
            .await
            .expect("read large ScanRange with bounded delimiter context");
        let GetResultPayload::Stream(stream) = result.payload else {
            panic!("expected streaming large ScanRange payload");
        };
        let chunks: Vec<Bytes> = stream.try_collect().await.expect("collect large ScanRange output");
        assert!(chunks.concat().is_empty());

        let mut oversized_run = vec![b'b'];
        oversized_run.resize(SELECT_DEFAULT_READ_BUFFER_SIZE + 2, b'a');
        let scan_start = i64::try_from(oversized_run.len()).expect("fixture offset should fit in i64");
        oversized_run.extend_from_slice(b"222aa");
        let mut reader = SelectPutObjReader::from_vec(oversized_run);
        env.ecstore
            .put_object(bucket, object, &mut reader, &Default::default())
            .await
            .expect("put oversized delimiter context fixture");
        let err = make_store(scan_start, scan_start, None, prepare_test_snapshot(bucket, object).await)
            .get_opts(&Path::from(object), GetOptions::default())
            .await
            .expect_err("oversized self-overlapping delimiter context must fail closed");
        assert!(err.to_string().contains("bounded ScanRange context"));
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[serial_test::serial]
    async fn test_self_overlapping_record_delimiter_uses_single_full_file_partition() {
        const TARGET_PARTITIONS: usize = 4;

        let mut input_bytes = Vec::with_capacity(SELECT_DEFAULT_READ_BUFFER_SIZE + 8);
        input_bytes.extend_from_slice(b"0,");
        input_bytes.resize(SELECT_DEFAULT_READ_BUFFER_SIZE - 1, b'b');
        input_bytes.extend_from_slice(b"aaaX,caa");
        assert!(input_bytes.len() > 1024 * 1024);
        assert_eq!(
            &input_bytes[SELECT_DEFAULT_READ_BUFFER_SIZE - 1..SELECT_DEFAULT_READ_BUFFER_SIZE + 2],
            b"aaa"
        );

        let env = crate::storage_api::select_test_ecstore_env().await;
        let bucket = "s3select-record-delimiter";
        let object = "input.csv";
        env.make_bucket(bucket, false).await;
        let mut reader = SelectPutObjReader::from_vec(input_bytes);
        env.ecstore
            .put_object(bucket, object, &mut reader, &Default::default())
            .await
            .expect("put multi-byte record-delimited test object");

        let input = Arc::new(SelectObjectContentInput {
            bucket: bucket.to_string(),
            expected_bucket_owner: None,
            key: object.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 {
                        record_delimiter: Some("aa".to_string()),
                        ..Default::default()
                    }),
                    ..Default::default()
                },
                output_serialization: OutputSerialization {
                    csv: Some(CSVOutput::default()),
                    ..Default::default()
                },
                request_progress: None,
                scan_range: None,
            },
        });
        let snapshot = prepare_test_snapshot(bucket, object).await;
        let store = Arc::new(EcObjectStore {
            input,
            need_convert: false,
            delimiter: String::new(),
            is_json_document: false,
            json_source: JsonSource::default(),
            input_metrics: Arc::new(SelectInputMetrics::default()),
            memory_pool: Arc::new(GreedyMemoryPool::new(32 * 1024 * 1024)),
            query_tracker: None,
            store: None,
            snapshot: OnceCell::new_with(Some(snapshot)),
            reader_open_count: Arc::new(AtomicUsize::new(0)),
        });

        let config = SessionConfig::new()
            .with_repartition_file_scans(false)
            .with_repartition_file_min_size(0)
            .with_target_partitions(TARGET_PARTITIONS);
        let context = SessionContext::new_with_config(config);
        let store_url = url::Url::parse(&format!("s3://{bucket}")).expect("valid object store URL");
        context.runtime_env().register_object_store(&store_url, store);
        context
            .register_csv("records", &format!("s3://{bucket}/{object}"), CsvReadOptions::new().has_header(false))
            .await
            .expect("register partitioned CSV");

        let scan_plan = context
            .sql("SELECT * FROM records")
            .await
            .expect("plan partitioned CSV")
            .create_physical_plan()
            .await
            .expect("create partitioned CSV physical plan");
        assert_eq!(scan_plan.output_partitioning().partition_count(), 1);

        let batches = context
            .sql("SELECT column_1 FROM records")
            .await
            .expect("plan exact-result query")
            .collect()
            .await
            .expect("query self-overlapping record-delimited CSV");
        let values = batches
            .iter()
            .flat_map(|batch| {
                let column = batch
                    .column(0)
                    .as_any()
                    .downcast_ref::<datafusion::arrow::array::StringArray>()
                    .expect("mixed first column should be Utf8");
                column.iter().map(|value| value.map(str::to_string)).collect::<Vec<_>>()
            })
            .collect::<Vec<_>>();
        assert_eq!(values, vec![Some("0".to_string()), Some("aX".to_string())]);
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[serial_test::serial]
    async fn test_get_opts_validates_raw_length_before_delimiter_conversion() {
        let env = crate::storage_api::select_test_ecstore_env().await;
        let bucket = "s3select-multi-byte-delimiter";
        let object = "input.csv";
        let input_bytes = b"a\r\n1^Y";
        env.make_bucket(bucket, false).await;
        let mut reader = SelectPutObjReader::from_vec(input_bytes.to_vec());
        env.ecstore
            .put_object(bucket, object, &mut reader, &Default::default())
            .await
            .expect("put multi-byte-delimited test object");

        let input = Arc::new(SelectObjectContentInput {
            bucket: bucket.to_string(),
            expected_bucket_owner: None,
            key: object.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 {
                        field_delimiter: Some("\r\n".to_string()),
                        record_delimiter: Some("^Y".to_string()),
                        ..Default::default()
                    }),
                    ..Default::default()
                },
                output_serialization: OutputSerialization {
                    csv: Some(CSVOutput::default()),
                    ..Default::default()
                },
                request_progress: None,
                scan_range: None,
            },
        });
        let snapshot = prepare_test_snapshot(bucket, object).await;
        let input_metrics = Arc::new(SelectInputMetrics::default());
        let store = super::EcObjectStore {
            input,
            need_convert: true,
            delimiter: "\r\n".to_string(),
            is_json_document: false,
            json_source: JsonSource::default(),
            input_metrics: Arc::clone(&input_metrics),
            memory_pool: Arc::new(GreedyMemoryPool::new(1024)),
            query_tracker: None,
            store: None,
            snapshot: OnceCell::new_with(Some(snapshot)),
            reader_open_count: Arc::new(AtomicUsize::new(0)),
        };

        let result = store
            .get_opts(&Path::from(object), GetOptions::default())
            .await
            .expect("read multi-byte-delimited test object");
        let GetResultPayload::Stream(stream) = result.payload else {
            panic!("expected streaming object payload");
        };
        let chunks: Vec<Bytes> = stream.try_collect().await.expect("collect converted object stream");

        assert_eq!(chunks.concat(), b"a,1\r\n");
        let input_len = u64::try_from(input_bytes.len()).expect("fixture length should fit in u64");
        assert_eq!(input_metrics.snapshot().bytes_scanned, input_len);
        assert_eq!(input_metrics.snapshot().bytes_processed, input_len);

        let requested_range = 3..10;
        let ranges = store
            .get_ranges(&Path::from(object), std::slice::from_ref(&requested_range))
            .await
            .expect("bounded range past EOF should return the object remainder");
        assert_eq!(ranges, vec![Bytes::from_static(b"1^Y")]);
    }

    #[tokio::test]
    async fn test_bytes_stream_stops_at_content_length() {
        let poll_count = Arc::new(AtomicUsize::new(0));
        let stream_poll_count = Arc::clone(&poll_count);
        let source = stream::unfold(0, move |index| {
            let stream_poll_count = Arc::clone(&stream_poll_count);
            async move {
                stream_poll_count.fetch_add(1, Ordering::SeqCst);
                let bytes = match index {
                    0 => Bytes::from_static(b"abcd"),
                    1 => Bytes::from_static(b"efgh"),
                    _ => return None,
                };
                Some((Ok::<_, std::io::Error>(bytes), index + 1))
            }
        });

        let chunks: Vec<Bytes> = bytes_stream(source, 4).try_collect().await.unwrap();

        assert_eq!(chunks, vec![Bytes::from_static(b"abcd")]);
        assert_eq!(poll_count.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn metered_stream_counts_only_polled_chunks() {
        let poll_count = Arc::new(AtomicUsize::new(0));
        let source_poll_count = Arc::clone(&poll_count);
        let source = stream::unfold(0, move |index| {
            let source_poll_count = Arc::clone(&source_poll_count);
            async move {
                source_poll_count.fetch_add(1, Ordering::SeqCst);
                let bytes = match index {
                    0 => Bytes::from_static(b"abcd"),
                    1 => Bytes::from_static(b"efgh"),
                    _ => return None,
                };
                Some((Ok::<_, std::io::Error>(bytes), index + 1))
            }
        });
        let input_metrics = Arc::new(SelectInputMetrics::default());
        let mut metered = Box::pin(meter_uncompressed_input_stream(source, Arc::clone(&input_metrics)));

        assert_eq!(metered.next().await.expect("first chunk").expect("valid chunk"), b"abcd"[..]);
        drop(metered);

        assert_eq!(input_metrics.snapshot().bytes_scanned, 4);
        assert_eq!(input_metrics.snapshot().bytes_processed, 4);
        assert_eq!(poll_count.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn test_bytes_stream_rejects_early_eof() {
        let source = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"ab"))]);
        let output = bytes_stream(source, 4);
        futures::pin_mut!(output);

        assert_eq!(output.next().await.expect("first stream item").expect("first chunk"), b"ab"[..]);
        let err = output
            .next()
            .await
            .expect("early EOF error")
            .expect_err("short stream must fail");
        let object_store::Error::Generic { store, source } = err else {
            panic!("expected generic object store error");
        };
        assert_eq!(store, "EcObjectStore");
        let source = source.downcast_ref::<std::io::Error>().expect("I/O error source");
        assert_eq!(source.kind(), std::io::ErrorKind::UnexpectedEof);
        assert!(source.to_string().contains("2 bytes remaining"));
        assert!(output.next().await.is_none());
    }

    #[tokio::test]
    async fn full_and_range_object_streams_record_input_metrics() {
        const BUCKET: &str = "s3select-input-metrics";
        const OBJECT: &str = "input.csv";
        const DATA: &[u8] = b"id,name\n1,a\n";

        let env = crate::storage_api::select_test_ecstore_env().await;
        env.make_bucket(BUCKET, false).await;
        let mut reader = SelectPutObjReader::from_vec(DATA.to_vec());
        env.ecstore
            .put_object(BUCKET, OBJECT, &mut reader, &Default::default())
            .await
            .expect("put input metrics fixture");
        let snapshot = prepare_test_snapshot(BUCKET, OBJECT).await;
        let input_metrics = Arc::new(SelectInputMetrics::default());
        let store = EcObjectStore::build_with_snapshot(
            csv_input(BUCKET, OBJECT),
            Arc::new(GreedyMemoryPool::new(1024 * 1024)),
            None,
            Arc::clone(&input_metrics),
            Arc::clone(&snapshot),
            JsonSource::default(),
        )
        .expect("build metrics-aware object store");

        let result = store
            .get_opts(&Path::from(OBJECT), GetOptions::default())
            .await
            .expect("open full object stream");
        let GetResultPayload::Stream(stream) = result.payload else {
            panic!("expected streaming object payload");
        };
        let body = stream
            .try_collect::<Vec<_>>()
            .await
            .expect("read full object stream")
            .concat();
        assert_eq!(body, DATA);
        let data_len = u64::try_from(DATA.len()).expect("fixture length should fit in u64");
        assert_eq!(input_metrics.snapshot().bytes_scanned, data_len);
        assert_eq!(input_metrics.snapshot().bytes_processed, data_len);

        input_metrics.reset();

        let result = store
            .get_opts(
                &Path::from(OBJECT),
                GetOptions {
                    range: Some(GetRange::Bounded(0..2)),
                    ..Default::default()
                },
            )
            .await
            .expect("open schema-style range stream");
        let GetResultPayload::Stream(stream) = result.payload else {
            panic!("expected range stream payload");
        };
        let range = stream
            .try_collect::<Vec<_>>()
            .await
            .expect("read schema-style range")
            .concat();
        assert_eq!(range, b"id"[..]);
        assert_eq!(input_metrics.snapshot().bytes_scanned, 2);
        assert_eq!(input_metrics.snapshot().bytes_processed, 2);
    }

    #[tokio::test]
    async fn compressed_object_uses_one_full_stream_and_rejects_internal_ranges() {
        const BUCKET: &str = "s3select-compressed-object";
        const OBJECT: &str = "input.csv";
        const DATA: &[u8] = b"id,name\n1,Alice\n2,Bob\n";

        let compressed = encode_compressed_fixture(CompressionFormat::Gzip, DATA).await;
        let env = crate::storage_api::select_test_ecstore_env().await;
        env.make_bucket(BUCKET, false).await;
        let mut reader = SelectPutObjReader::from_vec(compressed.clone());
        env.ecstore
            .put_object(BUCKET, OBJECT, &mut reader, &Default::default())
            .await
            .expect("put compressed CSV fixture");
        let snapshot = prepare_test_snapshot(BUCKET, OBJECT).await;
        let mut input = (*csv_input(BUCKET, OBJECT)).clone();
        input.request.input_serialization.compression_type = Some(CompressionType::from_static(CompressionType::GZIP));
        input.request.scan_range = Some(ScanRange {
            start: Some(0),
            end: None,
        });
        let input_metrics = Arc::new(SelectInputMetrics::default());
        let store = EcObjectStore::build_with_snapshot(
            Arc::new(input),
            Arc::new(GreedyMemoryPool::new(1024 * 1024)),
            None,
            Arc::clone(&input_metrics),
            Arc::clone(&snapshot),
            JsonSource::default(),
        )
        .expect("build compressed object store");

        let error = store
            .get_opts(
                &Path::from(OBJECT),
                GetOptions {
                    range: Some(GetRange::Bounded(0..1)),
                    ..Default::default()
                },
            )
            .await
            .expect_err("compressed input must reject DataFusion byte ranges");
        assert_eq!(
            QueryError::from(DataFusionError::ObjectStore(Box::new(error))).select_error(),
            SelectError::UnsupportedScanRangeInput
        );
        assert_eq!(store.reader_open_count.load(Ordering::SeqCst), 0);

        let mut scan_input = (*csv_input(BUCKET, OBJECT)).clone();
        scan_input.request.input_serialization.compression_type = Some(CompressionType::from_static(CompressionType::GZIP));
        scan_input.request.scan_range = Some(ScanRange {
            start: Some(1),
            end: None,
        });
        let scan_store = EcObjectStore::build_with_snapshot(
            Arc::new(scan_input),
            Arc::new(GreedyMemoryPool::new(1024 * 1024)),
            None,
            Arc::new(SelectInputMetrics::default()),
            snapshot,
            JsonSource::default(),
        )
        .expect("build compressed ScanRange object store");
        let error = scan_store
            .get_opts(&Path::from(OBJECT), GetOptions::default())
            .await
            .expect_err("compressed request ScanRange must fail before object I/O");
        assert_eq!(
            QueryError::from(DataFusionError::ObjectStore(Box::new(error))).select_error(),
            SelectError::UnsupportedScanRangeInput
        );
        assert_eq!(scan_store.reader_open_count.load(Ordering::SeqCst), 0);

        let result = store
            .get_opts(&Path::from(OBJECT), GetOptions::default())
            .await
            .expect("open full compressed object stream");
        let GetResultPayload::Stream(stream) = result.payload else {
            panic!("expected compressed stream payload");
        };
        let body = stream
            .try_collect::<Vec<_>>()
            .await
            .expect("decode compressed object")
            .concat();

        assert_eq!(body, DATA);
        assert_eq!(store.reader_open_count.load(Ordering::SeqCst), 1);
        assert_eq!(
            input_metrics.snapshot().bytes_scanned,
            u64::try_from(compressed.len()).expect("compressed fixture length should fit in u64")
        );
        assert_eq!(
            input_metrics.snapshot().bytes_processed,
            u64::try_from(DATA.len()).expect("input fixture length should fit in u64")
        );
    }

    #[tokio::test]
    async fn compressed_stream_throughput_is_independent_of_query_memory_pool() {
        const BUCKET: &str = "s3select-compressed-throughput";
        const OBJECT: &str = "input.csv.gz";

        let data = b"a\n".repeat(1024 * 1024);
        let compressed = encode_compressed_fixture(CompressionFormat::Gzip, &data).await;
        let env = crate::storage_api::select_test_ecstore_env().await;
        env.make_bucket(BUCKET, false).await;
        let mut reader = SelectPutObjReader::from_vec(compressed);
        env.ecstore
            .put_object(BUCKET, OBJECT, &mut reader, &Default::default())
            .await
            .expect("put compressed throughput fixture");
        let snapshot = prepare_test_snapshot(BUCKET, OBJECT).await;
        let mut input = (*csv_input(BUCKET, OBJECT)).clone();
        input.request.input_serialization.compression_type = Some(CompressionType::from_static(CompressionType::GZIP));
        let metrics = Arc::new(SelectInputMetrics::default());
        let store = EcObjectStore::build_with_snapshot(
            Arc::new(input),
            Arc::new(GreedyMemoryPool::new(1)),
            None,
            Arc::clone(&metrics),
            snapshot,
            JsonSource::default(),
        )
        .expect("build compressed throughput store");

        let result = store
            .get_opts(&Path::from(OBJECT), GetOptions::default())
            .await
            .expect("open compressed throughput stream");
        let GetResultPayload::Stream(stream) = result.payload else {
            panic!("expected compressed stream payload");
        };
        let decoded = stream
            .try_collect::<Vec<_>>()
            .await
            .expect("streamed decoded bytes should not consume the query memory pool")
            .concat();

        assert_eq!(decoded, data);
        assert_eq!(
            metrics.snapshot().bytes_processed,
            u64::try_from(data.len()).expect("fixture length should fit in u64")
        );
    }

    #[tokio::test]
    async fn dropping_real_object_stream_counts_only_consumed_bytes() {
        const BUCKET: &str = "s3select-partial-input-metrics";
        const OBJECT: &str = "large.csv";

        let data = vec![b'x'; SELECT_DEFAULT_READ_BUFFER_SIZE * 3];
        let env = crate::storage_api::select_test_ecstore_env().await;
        env.make_bucket(BUCKET, false).await;
        let mut reader = SelectPutObjReader::from_vec(data.clone());
        env.ecstore
            .put_object(BUCKET, OBJECT, &mut reader, &Default::default())
            .await
            .expect("put partial input metrics fixture");
        let snapshot = prepare_test_snapshot(BUCKET, OBJECT).await;
        let input_metrics = Arc::new(SelectInputMetrics::default());
        let store = EcObjectStore::build_with_snapshot(
            csv_input(BUCKET, OBJECT),
            Arc::new(GreedyMemoryPool::new(1024 * 1024)),
            None,
            Arc::clone(&input_metrics),
            snapshot,
            JsonSource::default(),
        )
        .expect("build metrics-aware object store");

        let result = store
            .get_opts(&Path::from(OBJECT), GetOptions::default())
            .await
            .expect("open partial object stream");
        let GetResultPayload::Stream(mut stream) = result.payload else {
            panic!("expected streaming object payload");
        };
        let first = stream
            .next()
            .await
            .expect("first object chunk")
            .expect("first object chunk should be valid");
        drop(stream);

        assert!(first.len() < data.len(), "fixture must span multiple reader chunks");
        let consumed = u64::try_from(first.len()).expect("chunk length should fit in u64");
        assert_eq!(input_metrics.snapshot().bytes_scanned, consumed);
        assert_eq!(input_metrics.snapshot().bytes_processed, consumed);
    }

    #[tokio::test]
    async fn json_object_streams_record_input_metrics() {
        const BUCKET: &str = "s3select-json-input-metrics";
        const LINES_OBJECT: &str = "input.jsonl";
        const LINES_DATA: &[u8] = b"{\"id\":1}\n{\"id\":2}\n";
        const DOCUMENT_OBJECT: &str = "input.json";
        const DOCUMENT_DATA: &[u8] = b"[{\"id\":1},{\"id\":2}]";

        let env = crate::storage_api::select_test_ecstore_env().await;
        env.make_bucket(BUCKET, false).await;
        for (object, data, json_type) in [
            (LINES_OBJECT, LINES_DATA, JSONType::LINES),
            (DOCUMENT_OBJECT, DOCUMENT_DATA, JSONType::DOCUMENT),
        ] {
            let mut reader = SelectPutObjReader::from_vec(data.to_vec());
            env.ecstore
                .put_object(BUCKET, object, &mut reader, &Default::default())
                .await
                .expect("put JSON input metrics fixture");
            let snapshot = prepare_test_snapshot(BUCKET, object).await;
            let input_metrics = Arc::new(SelectInputMetrics::default());
            let store = EcObjectStore::build_with_snapshot(
                json_input(BUCKET, object, json_type),
                Arc::new(GreedyMemoryPool::new(1024 * 1024)),
                None,
                Arc::clone(&input_metrics),
                snapshot,
                JsonSource::default(),
            )
            .expect("build metrics-aware JSON object store");

            let result = store
                .get_opts(&Path::from(object), GetOptions::default())
                .await
                .expect("open JSON object stream");
            let GetResultPayload::Stream(stream) = result.payload else {
                panic!("expected streaming JSON payload");
            };
            stream.try_collect::<Vec<_>>().await.expect("read JSON object stream");

            let input_len = u64::try_from(data.len()).expect("fixture length should fit in u64");
            assert_eq!(input_metrics.snapshot().bytes_scanned, input_len, "JSON type {json_type}");
            assert_eq!(input_metrics.snapshot().bytes_processed, input_len, "JSON type {json_type}");
        }
    }

    #[tokio::test]
    async fn scan_range_metrics_include_header_and_raw_range_once() {
        const BUCKET: &str = "s3select-scan-range-input-metrics";
        const OBJECT: &str = "input.csv";
        const DATA: &[u8] = b"h1,h2\nr1,a\nr2,b\n";
        const HEADER_LEN: usize = b"h1,h2\n".len();
        const RECORD_START: usize = b"h1,h2\nr1,a\n".len();
        const READ_START: usize = RECORD_START - 1;

        let env = crate::storage_api::select_test_ecstore_env().await;
        env.make_bucket(BUCKET, false).await;
        let mut reader = SelectPutObjReader::from_vec(DATA.to_vec());
        env.ecstore
            .put_object(BUCKET, OBJECT, &mut reader, &Default::default())
            .await
            .expect("put ScanRange input metrics fixture");
        let snapshot = prepare_test_snapshot(BUCKET, OBJECT).await;
        let mut input = (*csv_input(BUCKET, OBJECT)).clone();
        input
            .request
            .input_serialization
            .csv
            .as_mut()
            .expect("CSV input")
            .file_header_info = Some(FileHeaderInfo::from_static(FileHeaderInfo::USE));
        input.request.scan_range = Some(ScanRange {
            start: Some(i64::try_from(RECORD_START).expect("fixture offset should fit in i64")),
            end: Some(i64::try_from(RECORD_START).expect("fixture offset should fit in i64")),
        });
        let input_metrics = Arc::new(SelectInputMetrics::default());
        let store = EcObjectStore::build_with_snapshot(
            Arc::new(input),
            Arc::new(GreedyMemoryPool::new(1024 * 1024)),
            None,
            Arc::clone(&input_metrics),
            snapshot,
            JsonSource::default(),
        )
        .expect("build ScanRange metrics-aware object store");

        let result = store
            .get_opts(&Path::from(OBJECT), GetOptions::default())
            .await
            .expect("open ScanRange object stream");
        let GetResultPayload::Stream(stream) = result.payload else {
            panic!("expected streaming ScanRange payload");
        };
        let body = stream
            .try_collect::<Vec<_>>()
            .await
            .expect("read ScanRange object stream")
            .concat();

        assert_eq!(body, b"h1,h2\nr2,b\n");
        let expected_input = u64::try_from(HEADER_LEN + DATA.len() - READ_START).expect("fixture length should fit in u64");
        assert_eq!(input_metrics.snapshot().bytes_scanned, expected_input);
        assert_eq!(input_metrics.snapshot().bytes_processed, expected_input);
    }

    #[tokio::test]
    async fn test_json_document_stream_respects_query_memory_pool() {
        let input = b"{}".to_vec();
        let required = input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER;
        let memory_pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(required - 1));
        let mut output = json_document_ndjson_stream(
            Box::new(std::io::Cursor::new(input.clone())),
            input.len() as u64,
            JsonSource::default(),
            Arc::new(SelectInputMetrics::default()),
            memory_pool,
            None,
        );

        let err = output
            .next()
            .await
            .expect("memory error")
            .expect_err("reservation should exceed the pool");
        let object_store::Error::Generic { source, .. } = err else {
            panic!("expected generic object store error");
        };
        assert!(matches!(
            source.downcast_ref::<DataFusionError>(),
            Some(DataFusionError::ResourcesExhausted(_))
        ));
    }

    #[tokio::test]
    async fn compressed_json_document_uses_decoded_size_and_metric_boundaries() {
        const INPUT: &[u8] = br#"[{"id":1},{"id":2}]"#;
        const EXPECTED: &[u8] = b"{\"id\":1}\n{\"id\":2}\n";

        for format in [CompressionFormat::Gzip, CompressionFormat::Bzip2] {
            let compressed = encode_compressed_fixture(format, INPUT).await;
            let compressed_len = u64::try_from(compressed.len()).expect("compressed fixture length should fit in u64");
            let input_metrics = Arc::new(SelectInputMetrics::default());
            let reader = compressed_input_reader(
                Box::new(std::io::Cursor::new(compressed)),
                compressed_len,
                format,
                Arc::clone(&input_metrics),
                u64::MAX,
                None,
            );
            let output = compressed_json_document_ndjson_stream(
                reader,
                JsonSource::default(),
                Arc::new(GreedyMemoryPool::new(1024 * 1024)),
                None,
            )
            .try_collect::<Vec<_>>()
            .await
            .expect("compressed JSON DOCUMENT should decode and parse")
            .concat();

            assert_eq!(output, EXPECTED);
            assert_eq!(input_metrics.snapshot().bytes_scanned, compressed_len);
            assert_eq!(
                input_metrics.snapshot().bytes_processed,
                u64::try_from(INPUT.len()).expect("JSON fixture length should fit in u64")
            );
        }
    }

    async fn compressed_json_document_select_error(
        format: CompressionFormat,
        compressed: Vec<u8>,
        max_processed_bytes: u64,
    ) -> SelectError {
        let compressed_len = u64::try_from(compressed.len()).expect("compressed fixture length should fit in u64");
        let reader = compressed_input_reader(
            Box::new(std::io::Cursor::new(compressed)),
            compressed_len,
            format,
            Arc::new(SelectInputMetrics::default()),
            max_processed_bytes,
            None,
        );
        let mut output = compressed_json_document_ndjson_stream(
            reader,
            JsonSource::default(),
            Arc::new(GreedyMemoryPool::new(1024 * 1024)),
            None,
        );
        let source = output
            .next()
            .await
            .expect("decoder failure should produce one stream error")
            .expect_err("compressed JSON DOCUMENT decoding must fail");
        assert!(output.next().await.is_none());
        QueryError::from(DataFusionError::ObjectStore(Box::new(source))).select_error()
    }

    #[tokio::test]
    async fn compressed_json_document_preserves_decoder_select_errors() {
        const INPUT: &[u8] = br#"[{"id":1}]"#;

        for (format, compression) in [(CompressionFormat::Gzip, "GZIP"), (CompressionFormat::Bzip2, "BZIP2")] {
            assert_eq!(
                compressed_json_document_select_error(format, b"not compressed".to_vec(), u64::MAX).await,
                SelectError::InvalidCompressionFormatForObject { compression }
            );

            let mut truncated = encode_compressed_fixture(format, INPUT).await;
            truncated.truncate(truncated.len() - 1);
            assert_eq!(
                compressed_json_document_select_error(format, truncated, u64::MAX).await,
                SelectError::TruncatedInput
            );

            let compressed = encode_compressed_fixture(format, INPUT).await;
            let max_processed_bytes = u64::try_from(INPUT.len() - 1).expect("fixture length should fit in u64");
            assert_eq!(
                compressed_json_document_select_error(format, compressed, max_processed_bytes).await,
                SelectError::ResourceExhausted
            );
        }
    }

    #[test]
    fn compressed_json_document_buffer_grows_amortized_and_reserves_capacity() {
        let memory_pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1024 * 1024));
        let reservation = MemoryConsumer::new("compressed JSON document test").register(&memory_pool);
        let json_source = JsonSource::default();
        let mut buffer = Vec::new();
        let mut capacity_growths = 0;

        for _ in 0..1025 {
            let old_capacity = buffer.capacity();
            let required_len = buffer.len() + 1;
            grow_json_document_buffer(&mut buffer, required_len, &reservation, &json_source)
                .expect("bounded JSON buffer should grow");
            if buffer.capacity() != old_capacity {
                capacity_growths += 1;
            }
            buffer.push(0);
            assert_eq!(
                reservation.size(),
                super::json_document_memory_reservation_bytes(buffer.capacity(), &json_source)
                    .expect("test reservation should fit")
            );
        }

        assert!(capacity_growths <= 12, "power-of-two growth should stay logarithmic");
    }

    #[tokio::test]
    async fn scalar_alias_expansion_is_in_the_query_memory_reservation() {
        let input = b"[0,0]".to_vec();
        let alias = "alias".repeat(128);
        let source = JsonSource::new(vec![JsonPathSegment::ArrayWildcard], Some(alias.clone()));
        // Keep this threshold independent from the production helper so a
        // smaller scalar-alias multiplier cannot make the test self-validate.
        let required = input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER + alias.len() * 14 * input.len();
        assert!(required > input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER);
        let memory_pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(required - 1));
        let mut output = json_document_ndjson_stream(
            Box::new(std::io::Cursor::new(input.clone())),
            input.len() as u64,
            source,
            Arc::new(SelectInputMetrics::default()),
            memory_pool,
            None,
        );

        let err = output
            .next()
            .await
            .expect("memory error")
            .expect_err("scalar alias expansion must be reserved before parsing");
        let object_store::Error::Generic { source, .. } = err else {
            panic!("expected generic object store error");
        };
        assert!(matches!(
            source.downcast_ref::<DataFusionError>(),
            Some(DataFusionError::ResourcesExhausted(_))
        ));

        let memory_pool = Arc::new(GreedyMemoryPool::new(required));
        let output: Vec<Bytes> = json_document_ndjson_stream(
            Box::new(std::io::Cursor::new(input.clone())),
            input.len() as u64,
            JsonSource::new(vec![JsonPathSegment::ArrayWildcard], Some(alias)),
            Arc::new(SelectInputMetrics::default()),
            memory_pool.clone(),
            None,
        )
        .try_collect()
        .await
        .expect("scalar alias expansion should fit the exact reservation");
        assert_eq!(output.len(), 2);
        assert_eq!(memory_pool.reserved(), 0);
    }

    #[tokio::test]
    async fn test_json_document_stream_releases_memory_reservation() {
        let input = b"[1,2]".to_vec();
        let required = input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER;
        let memory_pool = Arc::new(GreedyMemoryPool::new(required));
        let input_metrics = Arc::new(SelectInputMetrics::default());
        let output: Vec<Bytes> = json_document_ndjson_stream(
            Box::new(std::io::Cursor::new(input.clone())),
            input.len() as u64,
            JsonSource::default(),
            Arc::clone(&input_metrics),
            memory_pool.clone(),
            None,
        )
        .try_collect()
        .await
        .expect("JSON conversion should fit the pool");
        let input_len = u64::try_from(input.len()).expect("fixture length should fit in u64");
        assert_eq!(input_metrics.snapshot().bytes_scanned, input_len);
        assert_eq!(input_metrics.snapshot().bytes_processed, input_len);

        assert_eq!(output, vec![Bytes::from_static(b"{\"_1\":1}\n"), Bytes::from_static(b"{\"_1\":2}\n")]);
        assert_eq!(memory_pool.reserved(), 0);
    }

    #[tokio::test]
    async fn test_json_document_stream_rejects_early_eof() {
        let input = b"{}".to_vec();
        let input_len = u64::try_from(input.len()).expect("fixture length should fit in u64");
        let input_metrics = Arc::new(SelectInputMetrics::default());
        let memory_pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(4 * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER));
        let mut output = json_document_ndjson_stream(
            Box::new(std::io::Cursor::new(input)),
            4,
            JsonSource::default(),
            Arc::clone(&input_metrics),
            memory_pool,
            None,
        );

        let err = output
            .next()
            .await
            .expect("early EOF error")
            .expect_err("short JSON document must fail");
        let object_store::Error::Generic { source, .. } = err else {
            panic!("expected generic object store error");
        };
        let source = source.downcast_ref::<std::io::Error>().expect("I/O error source");
        assert_eq!(source.kind(), std::io::ErrorKind::UnexpectedEof);
        assert!(source.to_string().contains("2 bytes remaining"));
        assert_eq!(input_metrics.snapshot().bytes_scanned, input_len);
        assert_eq!(input_metrics.snapshot().bytes_processed, input_len);
        assert!(output.next().await.is_none());
    }

    #[tokio::test]
    async fn malformed_json_document_stream_has_typed_select_error() {
        let input = b"{bad".to_vec();
        let memory_pool: Arc<dyn MemoryPool> =
            Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER));
        let mut output = json_document_ndjson_stream(
            Box::new(std::io::Cursor::new(input.clone())),
            input.len() as u64,
            JsonSource::default(),
            Arc::new(SelectInputMetrics::default()),
            memory_pool,
            None,
        );

        let source = output
            .next()
            .await
            .expect("malformed JSON should produce one stream error")
            .expect_err("malformed JSON DOCUMENT must fail");
        let error = QueryError::from(DataFusionError::ObjectStore(Box::new(source)));

        assert_eq!(error.select_error(), SelectError::JsonParsingError);
        assert!(output.next().await.is_none());
    }

    #[tokio::test]
    async fn json_document_stream_preserves_typed_parser_error() {
        let input = b"{}".to_vec();
        let memory_pool: Arc<dyn MemoryPool> =
            Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER));
        let mut output = json_document_ndjson_stream_with_parser(
            Box::new(std::io::Cursor::new(input.clone())),
            JsonDocumentReadMode::Exact {
                original_size: input.len() as u64,
                input_metrics: SelectInputMetrics::default().recorder(),
            },
            JsonSource::default(),
            memory_pool,
            None,
            |_, _, _| Err(std::io::Error::new(std::io::ErrorKind::InvalidData, SelectError::AmbiguousFieldName)),
        );

        let source = output
            .next()
            .await
            .expect("typed parser failure should produce one stream error")
            .expect_err("typed parser failure must fail the stream");
        let error = QueryError::from(DataFusionError::ObjectStore(Box::new(source)));

        assert_eq!(error.select_error(), SelectError::AmbiguousFieldName);
        assert!(output.next().await.is_none());
    }

    #[test]
    fn storage_error_mapper_preserves_protocol_classification() {
        let classify = |source| QueryError::from(DataFusionError::ObjectStore(Box::new(source))).select_error();

        assert_eq!(
            classify(map_storage_error(
                "private-bucket",
                "private-object",
                SelectStorageError::BucketNotFound("private-bucket".to_string()),
            )),
            SelectError::BucketNotFound
        );
        assert_eq!(
            classify(map_storage_error(
                "private-bucket",
                "private-object",
                SelectStorageError::ObjectNotFound("private-bucket".to_string(), "private-object".to_string()),
            )),
            SelectError::ObjectNotFound
        );
        assert_eq!(
            classify(map_storage_error("private-bucket", "private-object", SelectStorageError::LessData)),
            SelectError::InternalError
        );
        assert_eq!(
            classify(scan_range_from_bounds(Some(10), None, 10).expect_err("out-of-bounds range must fail")),
            SelectError::InvalidScanRange
        );
        let parquet_source = map_storage_error(
            "private-bucket",
            "private-object",
            SelectStorageError::ObjectNotFound("private-bucket".to_string(), "private-object".to_string()),
        );
        let parquet_error = QueryError::from(DataFusionError::ParquetError(Box::new(
            datafusion::parquet::errors::ParquetError::External(Box::new(parquet_source)),
        )));
        assert_eq!(parquet_error.select_error(), SelectError::ObjectNotFound);
    }

    #[test]
    fn test_json_document_size_error_is_resource_exhausted() {
        assert!(validate_json_document_size(super::MAX_JSON_DOCUMENT_BYTES).is_ok());

        let err = validate_json_document_size(super::MAX_JSON_DOCUMENT_BYTES + 1).expect_err("oversized JSON document must fail");
        let object_store::Error::Generic { source, .. } = err else {
            panic!("expected generic object store error");
        };
        assert!(matches!(
            source.downcast_ref::<DataFusionError>(),
            Some(DataFusionError::ResourcesExhausted(_))
        ));
    }

    #[test]
    fn test_json_document_cancelled_queued_parse_releases_before_dequeue() {
        let runtime = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2)
            .max_blocking_threads(1)
            .enable_all()
            .build()
            .expect("build test runtime");

        runtime.block_on(async {
            let (blocking_started_tx, blocking_started_rx) = tokio::sync::oneshot::channel();
            let (release_blocking_tx, release_blocking_rx) = std::sync::mpsc::channel();
            let blocker = tokio::task::spawn_blocking(move || {
                let _ = blocking_started_tx.send(());
                release_blocking_rx.recv().expect("release blocking worker");
            });
            blocking_started_rx.await.expect("blocking worker should start");

            let admission = Arc::new(Semaphore::new(1));
            let permit = Arc::clone(&admission)
                .acquire_owned()
                .await
                .expect("query permit should be available");
            let query_guard: QueryExecutionGuard = Arc::new(permit);
            let query_tracker = QueryExecutionTracker::new(
                &QueryExecutionOwner::new(),
                query_guard,
                tokio::time::Instant::now() + std::time::Duration::from_secs(30),
                30,
            );
            let input = b"{}".to_vec();
            let (memory_pool, reservation_released) = AdmissionObservingMemoryPool::new(
                input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER,
                Arc::clone(&admission),
            );
            let memory_pool = Arc::new(memory_pool);
            let query_memory_pool: Arc<dyn MemoryPool> = memory_pool.clone();
            let mut output = json_document_ndjson_stream(
                Box::new(std::io::Cursor::new(input.clone())),
                input.len() as u64,
                JsonSource::default(),
                Arc::new(SelectInputMetrics::default()),
                query_memory_pool,
                Some(query_tracker),
            );

            {
                let next = output.next();
                futures::pin_mut!(next);
                assert!(futures::poll!(next.as_mut()).is_pending());
            }
            drop(output);

            assert!(
                tokio::time::timeout(std::time::Duration::from_millis(100), reservation_released)
                    .await
                    .expect("queued parse cancellation should release its memory immediately")
                    .expect("memory reservation release observer should remain open"),
                "query admission must cover the memory reservation through teardown"
            );
            assert_eq!(memory_pool.reserved(), 0);
            let recovered_permit =
                tokio::time::timeout(std::time::Duration::from_millis(100), Arc::clone(&admission).acquire_owned())
                    .await
                    .expect("queued parse cancellation should release admission before a worker is available")
                    .expect("query admission should remain open");

            release_blocking_tx.send(()).expect("release blocking worker");
            blocker.await.expect("blocking worker should finish");
            drop(recovered_permit);
            assert_eq!(admission.available_permits(), 1);
        });
    }

    #[test]
    fn test_json_document_expired_queued_parse_does_not_start() {
        let runtime = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2)
            .max_blocking_threads(1)
            .enable_all()
            .build()
            .expect("build test runtime");

        runtime.block_on(async {
            let (blocking_started_tx, blocking_started_rx) = tokio::sync::oneshot::channel();
            let (release_blocking_tx, release_blocking_rx) = std::sync::mpsc::channel();
            let blocker = tokio::task::spawn_blocking(move || {
                let _ = blocking_started_tx.send(());
                release_blocking_rx.recv().expect("release blocking worker");
            });
            blocking_started_rx.await.expect("blocking worker should start");

            let admission = Arc::new(Semaphore::new(1));
            let permit = Arc::clone(&admission)
                .acquire_owned()
                .await
                .expect("query permit should be available");
            let owner = QueryExecutionOwner::new();
            let query_tracker = QueryExecutionTracker::new(
                &owner,
                Arc::new(permit),
                tokio::time::Instant::now() + std::time::Duration::from_secs(30),
                30,
            );
            let input = b"{}".to_vec();
            let memory_pool: Arc<dyn MemoryPool> =
                Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER));
            let parser_started = Arc::new(std::sync::atomic::AtomicBool::new(false));
            let parser_started_in_task = Arc::clone(&parser_started);
            let mut output = json_document_ndjson_stream_with_parser(
                Box::new(std::io::Cursor::new(input.clone())),
                JsonDocumentReadMode::Exact {
                    original_size: input.len() as u64,
                    input_metrics: SelectInputMetrics::default().recorder(),
                },
                JsonSource::default(),
                Arc::clone(&memory_pool),
                Some(query_tracker.clone()),
                move |_, _, _| {
                    parser_started_in_task.store(true, std::sync::atomic::Ordering::SeqCst);
                    Ok(vec![Bytes::from_static(b"{}\n")])
                },
            );

            {
                let next = output.next();
                futures::pin_mut!(next);
                assert!(futures::poll!(next.as_mut()).is_pending());
            }
            query_tracker.expire(&owner);
            assert_eq!(admission.available_permits(), 0);
            assert!(memory_pool.reserved() > 0);
            release_blocking_tx.send(()).expect("release blocking worker");
            blocker.await.expect("blocking worker should finish");

            let err = tokio::time::timeout(std::time::Duration::from_secs(5), output.next())
                .await
                .expect("queued parser should resume")
                .expect("queued parser should return an error")
                .expect_err("expired queued parser must not run");
            let object_store::Error::Generic { source, .. } = err else {
                panic!("expected generic object store error");
            };
            let source = source.downcast_ref::<std::io::Error>().expect("I/O error source");
            assert_eq!(source.kind(), std::io::ErrorKind::Interrupted);
            assert!(!parser_started.load(std::sync::atomic::Ordering::SeqCst));
            assert_eq!(memory_pool.reserved(), 0);
            assert_eq!(admission.available_permits(), 1);
        });
    }

    #[test]
    fn test_json_document_expired_before_enqueue_releases_resources_without_blocking() {
        let runtime = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2)
            .max_blocking_threads(1)
            .enable_all()
            .build()
            .expect("build test runtime");

        runtime.block_on(async {
            let (blocking_started_tx, blocking_started_rx) = tokio::sync::oneshot::channel();
            let (release_blocking_tx, release_blocking_rx) = std::sync::mpsc::channel();
            let blocker = tokio::task::spawn_blocking(move || {
                let _ = blocking_started_tx.send(());
                release_blocking_rx.recv().expect("release blocking worker");
            });
            blocking_started_rx.await.expect("blocking worker should start");

            let admission = Arc::new(Semaphore::new(1));
            let permit = Arc::clone(&admission)
                .acquire_owned()
                .await
                .expect("query permit should be available");
            let owner = QueryExecutionOwner::new();
            let query_tracker = QueryExecutionTracker::new(
                &owner,
                Arc::new(permit),
                tokio::time::Instant::now() + std::time::Duration::from_secs(30),
                30,
            );
            query_tracker.expire(&owner);

            let input = b"{}".to_vec();
            let memory_pool: Arc<dyn MemoryPool> =
                Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER));
            let parser_started = Arc::new(std::sync::atomic::AtomicBool::new(false));
            let parser_started_in_task = Arc::clone(&parser_started);
            let mut output = json_document_ndjson_stream_with_parser(
                Box::new(std::io::Cursor::new(input.clone())),
                JsonDocumentReadMode::Exact {
                    original_size: input.len() as u64,
                    input_metrics: SelectInputMetrics::default().recorder(),
                },
                JsonSource::default(),
                Arc::clone(&memory_pool),
                Some(query_tracker),
                move |_, _, _| {
                    parser_started_in_task.store(true, std::sync::atomic::Ordering::SeqCst);
                    Ok(vec![Bytes::from_static(b"{}\n")])
                },
            );

            let err = tokio::time::timeout(std::time::Duration::from_millis(100), output.next())
                .await
                .expect("expired parse must fail before entering the saturated blocking queue")
                .expect("expired parse should return an error")
                .expect_err("expired parse must not run");
            let object_store::Error::Generic { source, .. } = err else {
                panic!("expected generic object store error");
            };
            let source = source.downcast_ref::<std::io::Error>().expect("I/O error source");
            assert_eq!(source.kind(), std::io::ErrorKind::Interrupted);
            assert!(!parser_started.load(std::sync::atomic::Ordering::SeqCst));
            assert_eq!(memory_pool.reserved(), 0);
            assert_eq!(admission.available_permits(), 1);

            release_blocking_tx.send(()).expect("release blocking worker");
            blocker.await.expect("blocking worker should finish");
        });
    }

    #[test]
    fn test_json_document_started_parse_cancels_and_releases_query_guard() {
        let runtime = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2)
            .max_blocking_threads(1)
            .enable_all()
            .build()
            .expect("build test runtime");

        runtime.block_on(async {
            let admission = Arc::new(Semaphore::new(1));
            let permit = Arc::clone(&admission)
                .acquire_owned()
                .await
                .expect("query permit should be available");
            let query_guard: QueryExecutionGuard = Arc::new(permit);
            let query_tracker = QueryExecutionTracker::new(
                &QueryExecutionOwner::new(),
                query_guard,
                tokio::time::Instant::now() + std::time::Duration::from_secs(30),
                30,
            );
            let input = b"{}".to_vec();
            let memory_pool: Arc<dyn MemoryPool> =
                Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER));
            let (parse_started_tx, parse_started_rx) = tokio::sync::oneshot::channel();
            let mut output = json_document_ndjson_stream_with_parser(
                Box::new(std::io::Cursor::new(input.clone())),
                JsonDocumentReadMode::Exact {
                    original_size: input.len() as u64,
                    input_metrics: SelectInputMetrics::default().recorder(),
                },
                JsonSource::default(),
                memory_pool,
                Some(query_tracker),
                move |_, _, cancellation| {
                    let _ = parse_started_tx.send(());
                    while !cancellation.load(std::sync::atomic::Ordering::Acquire) {
                        std::thread::yield_now();
                    }
                    Err(std::io::Error::new(std::io::ErrorKind::Interrupted, SelectError::Canceled))
                },
            );

            {
                let next = output.next();
                futures::pin_mut!(next);
                assert!(futures::poll!(next.as_mut()).is_pending());
            }
            parse_started_rx.await.expect("JSON parser should start");
            assert!(Arc::clone(&admission).try_acquire_owned().is_err());
            drop(output);

            let recovered_permit =
                tokio::time::timeout(std::time::Duration::from_secs(5), Arc::clone(&admission).acquire_owned())
                    .await
                    .expect("cancelled JSON parse should release the query guard without an external unblock")
                    .expect("query admission should remain open");
            drop(recovered_permit);
            assert_eq!(admission.available_permits(), 1);
        });
    }

    /// A JSON array is split into one NDJSON line per element.
    #[test]
    fn test_flatten_array_produces_one_line_per_element() {
        let input = br#"[{"id":1,"name":"Alice"},{"id":2,"name":"Bob"}]"#;
        let result = flatten_json_document_to_ndjson(input, &[]).expect("should succeed");
        let text = std::str::from_utf8(&result).unwrap();
        let lines: Vec<&str> = text.lines().collect();
        assert_eq!(lines.len(), 2);
        // Each line must be valid JSON
        for line in &lines {
            serde_json::from_str::<serde_json::Value>(line).expect("each line must be valid JSON");
        }
        // Spot-check field values
        let first: serde_json::Value = serde_json::from_str(lines[0]).unwrap();
        assert_eq!(first["id"], 1);
        assert_eq!(first["name"], "Alice");
    }

    /// A single JSON object emits exactly one NDJSON line.
    #[test]
    fn test_flatten_single_object_produces_one_line() {
        let input = br#"{"id":42,"value":"hello world"}"#;
        let result = flatten_json_document_to_ndjson(input, &[]).expect("should succeed");
        let text = std::str::from_utf8(&result).unwrap();
        let lines: Vec<&str> = text.lines().collect();
        assert_eq!(lines.len(), 1);
        let parsed: serde_json::Value = serde_json::from_str(lines[0]).unwrap();
        assert_eq!(parsed["id"], 42);
        assert_eq!(parsed["value"], "hello world");
    }

    /// An empty JSON array produces empty output (zero bytes).
    #[test]
    fn test_flatten_empty_array_produces_no_output() {
        let input = b"[]";
        let result = flatten_json_document_to_ndjson(input, &[]).expect("should succeed");
        assert!(result.is_empty(), "empty array should yield zero bytes");
    }

    /// A multi-line (pretty-printed) JSON document is flattened correctly.
    #[test]
    fn test_flatten_pretty_printed_document() {
        let input = b"[\n  {\"a\": 1},\n  {\"a\": 2},\n  {\"a\": 3}\n]";
        let result = flatten_json_document_to_ndjson(input, &[]).expect("should succeed");
        let text = std::str::from_utf8(&result).unwrap();
        assert_eq!(text.lines().count(), 3);
    }

    /// Nested objects inside array elements are preserved as compact single-line JSON.
    #[test]
    fn test_flatten_array_with_nested_objects() {
        let input = br#"[{"outer":{"inner":99}},{"outer":{"inner":100}}]"#;
        let result = flatten_json_document_to_ndjson(input, &[]).expect("should succeed");
        let text = std::str::from_utf8(&result).unwrap();
        let lines: Vec<&str> = text.lines().collect();
        assert_eq!(lines.len(), 2);
        // Each line must not contain a newline mid-value
        for line in &lines {
            assert!(!line.is_empty());
            let v: serde_json::Value = serde_json::from_str(line).unwrap();
            assert!(v["outer"]["inner"].as_i64().unwrap() >= 99);
        }
    }

    /// Each output line ends with exactly one newline (no blank lines between records).
    #[test]
    fn test_flatten_output_ends_with_newline_per_record() {
        let input = br#"[{"x":1},{"x":2}]"#;
        let result = flatten_json_document_to_ndjson(input, &[]).expect("should succeed");
        let text = std::str::from_utf8(&result).unwrap();
        // Exactly 2 newlines for 2 records
        assert_eq!(text.chars().filter(|&c| c == '\n').count(), 2);
        // No leading blank line
        assert!(!text.starts_with('\n'));
    }

    /// Invalid JSON returns an `InvalidData` IO error.
    #[test]
    fn test_flatten_invalid_json_returns_error() {
        let input = b"{ not valid json }";
        let err = flatten_json_document_to_ndjson(input, &[]).expect_err("should fail on invalid JSON");
        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
    }

    #[test]
    fn json_document_logical_record_enforces_one_megabyte_limit() {
        const OBJECT_OVERHEAD: usize = br#"{"v":""}"#.len();

        let at_limit = serde_json::json!({"v": "x".repeat(MAX_SELECT_RECORD_BYTES - OBJECT_OVERHEAD)});
        let line = super::json_value_to_line(at_limit, "_1").expect("one-megabyte logical record should be accepted");
        assert_eq!(line.len(), MAX_SELECT_RECORD_BYTES + 1);

        let over_limit = serde_json::json!({"v": "x".repeat(MAX_SELECT_RECORD_BYTES + 1 - OBJECT_OVERHEAD)});
        let error = super::json_value_to_line(over_limit, "_1").expect_err("oversized logical record must fail");
        assert_eq!(error.kind(), std::io::ErrorKind::InvalidData);
        assert!(error.get_ref().is_some_and(|source| {
            source
                .downcast_ref::<SelectError>()
                .is_some_and(|error| error == &SelectError::OverMaxRecordSize)
        }));

        let escaped = serde_json::json!({"v": "\0".repeat(MAX_SELECT_RECORD_BYTES / 6)});
        let error = super::json_value_to_line(escaped, "_1").expect_err("escaped output must be bounded while it is serialized");
        assert!(error.get_ref().is_some_and(|source| {
            source
                .downcast_ref::<SelectError>()
                .is_some_and(|error| error == &SelectError::OverMaxRecordSize)
        }));
    }

    /// Completely empty input returns an error (not valid JSON).
    #[test]
    fn test_flatten_empty_input_returns_error() {
        let err = flatten_json_document_to_ndjson(b"", &[]).expect_err("empty bytes are not valid JSON");
        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
    }

    #[test]
    fn key_path_expands_final_array_for_legacy_queries() {
        let input = br#"{"employees":[{"id":1,"name":"Alice","salary":75000},{"id":2,"name":"Bob","salary":65000}]}"#;
        let result = flatten_json_document_to_ndjson(input, &[source_key("employees")]).expect("key path should succeed");
        let text = std::str::from_utf8(&result).unwrap();
        let lines: Vec<&str> = text.lines().collect();
        assert_eq!(lines.len(), 2, "each employee should be its own NDJSON line");
        let first: serde_json::Value = serde_json::from_str(lines[0]).unwrap();
        assert_eq!(first["name"], "Alice");
        assert_eq!(first["salary"], 75000);
        let second: serde_json::Value = serde_json::from_str(lines[1]).unwrap();
        assert_eq!(second["name"], "Bob");
    }

    #[test]
    fn nested_array_wildcards_expand_in_source_order() {
        let input = br#"{"departments":[{"employees":[{"id":1},{"id":2}]},{"employees":[{"id":3}]}]}"#;
        let path = [
            JsonPathSegment::ArrayWildcard,
            source_key("departments"),
            JsonPathSegment::ArrayWildcard,
            source_key("employees"),
            JsonPathSegment::ArrayWildcard,
        ];

        let result = flatten_json_document_to_ndjson(input, &path).expect("nested wildcard path should succeed");
        let ids: Vec<i64> = std::str::from_utf8(&result)
            .unwrap()
            .lines()
            .map(|line| {
                serde_json::from_str::<serde_json::Value>(line).unwrap()["id"]
                    .as_i64()
                    .unwrap()
            })
            .collect();

        assert_eq!(ids, vec![1, 2, 3]);
    }

    #[test]
    fn explicit_array_wildcard_does_not_expand_nested_array_records_twice() {
        let input = br#"{"groups":[[1,2],[3,4]]}"#;
        let source = JsonSource::new(vec![source_key("groups"), JsonPathSegment::ArrayWildcard], Some("g".to_string()));

        let result = super::parse_json_document_to_lines(input, &source).expect("explicit wildcard path should succeed");

        assert_eq!(
            result,
            vec![Bytes::from_static(b"{\"g\":[1,2]}\n"), Bytes::from_static(b"{\"g\":[3,4]}\n")]
        );
    }

    #[test]
    fn leading_array_wildcard_expands_a_root_array_before_nested_keys() {
        let input = br#"[{"employees":[{"id":1}]},{"employees":[{"id":2}]}]"#;
        let path = [
            JsonPathSegment::ArrayWildcard,
            source_key("employees"),
            JsonPathSegment::ArrayWildcard,
        ];

        let result = flatten_json_document_to_ndjson(input, &path).expect("root array path should succeed");
        let ids: Vec<i64> = std::str::from_utf8(&result)
            .unwrap()
            .lines()
            .map(|line| {
                serde_json::from_str::<serde_json::Value>(line).unwrap()["id"]
                    .as_i64()
                    .unwrap()
            })
            .collect();

        assert_eq!(ids, vec![1, 2]);
    }

    #[test]
    fn root_array_index_selects_one_record() {
        let input = br#"[{"id":1},{"id":2}]"#;

        let result = flatten_json_document_to_ndjson(input, &[JsonPathSegment::Index(1)]).expect("array index should succeed");

        assert_eq!(result, Bytes::from_static(b"{\"id\":2}\n"));
    }

    #[test]
    fn canonical_root_array_index_selects_one_record() {
        let input = br#"[{"id":1},{"id":2}]"#;
        let path = [JsonPathSegment::ArrayWildcard, JsonPathSegment::Index(0)];

        let result = flatten_json_document_to_ndjson(input, &path).expect("canonical array index should succeed");

        assert_eq!(result, Bytes::from_static(b"{\"id\":1}\n"));
    }

    #[test]
    fn out_of_range_root_array_indexes_produce_no_records() {
        let input = br#"[{"id":1}]"#;

        for path in [
            vec![JsonPathSegment::Index(1)],
            vec![JsonPathSegment::ArrayWildcard, JsonPathSegment::Index(1)],
        ] {
            let result = flatten_json_document_to_ndjson(input, &path).expect("out-of-range index should not fail");
            assert!(result.is_empty());
        }
    }

    #[test]
    fn canonical_key_path_does_not_expand_an_array_without_a_wildcard() {
        let input = br#"{"rules":[{"id":1},{"id":2}]}"#;
        let source = JsonSource::new(vec![JsonPathSegment::ArrayWildcard, source_key("rules")], Some("r".to_string()));

        let result = super::parse_json_document_to_lines(input, &source).expect("array-valued source path should remain one row");

        assert_eq!(result, vec![Bytes::from_static(b"{\"r\":[{\"id\":1},{\"id\":2}]}\n")]);
    }

    #[test]
    fn object_wildcard_expands_values_and_allows_continuation() {
        let input = br#"{"groups":{"first":{"id":1},"second":{"id":2}}}"#;
        let path = [
            JsonPathSegment::ArrayWildcard,
            source_key("groups"),
            JsonPathSegment::ObjectWildcard,
            source_key("id"),
        ];

        let result = flatten_json_document_to_ndjson(input, &path).expect("object wildcard should succeed");
        let values: Vec<serde_json::Value> = std::str::from_utf8(&result)
            .unwrap()
            .lines()
            .map(|line| serde_json::from_str(line).unwrap())
            .collect();

        assert_eq!(values, vec![serde_json::json!({"id": 1}), serde_json::json!({"id": 2})]);
    }

    #[test]
    fn canonical_double_wildcard_expands_a_root_array() {
        let input = br#"[{"id":1},{"id":2}]"#;
        let path = [JsonPathSegment::ArrayWildcard, JsonPathSegment::ArrayWildcard];

        let result = flatten_json_document_to_ndjson(input, &path).expect("root wildcard should succeed");

        assert_eq!(result, Bytes::from_static(b"{\"id\":1}\n{\"id\":2}\n"));
    }

    #[test]
    fn terminal_scalar_path_uses_the_terminal_key_as_its_column() {
        let input = br#"{"rules":[{"id":"one"},{"id":"two"}]}"#;
        let path = [
            JsonPathSegment::ArrayWildcard,
            source_key("rules"),
            JsonPathSegment::ArrayWildcard,
            source_key("id"),
        ];

        let result = flatten_json_document_to_ndjson(input, &path).expect("scalar source path should succeed");

        assert_eq!(result, Bytes::from_static(b"{\"id\":\"one\"}\n{\"id\":\"two\"}\n"));
    }

    #[test]
    fn terminal_scalar_path_uses_the_explicit_source_alias() {
        let input = br#"{"rules":[{"id":"one"},{"id":"two"}]}"#;
        let source = JsonSource::new(
            vec![
                JsonPathSegment::ArrayWildcard,
                source_key("rules"),
                JsonPathSegment::ArrayWildcard,
                source_key("id"),
            ],
            Some("v".to_string()),
        );

        let result =
            super::parse_json_document_to_lines(input, &source).expect("explicit scalar source alias should be preserved");

        assert_eq!(
            result,
            vec![
                Bytes::from_static(b"{\"v\":\"one\"}\n"),
                Bytes::from_static(b"{\"v\":\"two\"}\n")
            ]
        );
    }

    #[test]
    fn source_keys_follow_s3_case_sensitivity_rules() {
        let input = br#"{"Employees":[{"id":1}]}"#;

        let unquoted = flatten_json_document_to_ndjson(input, &[source_key("employees")])
            .expect("unquoted source key should be case insensitive");
        let quoted_exact = flatten_json_document_to_ndjson(input, &[quoted_source_key("Employees")])
            .expect("exact quoted source key should match");
        let quoted = flatten_json_document_to_ndjson(input, &[quoted_source_key("employees")])
            .expect("missing quoted source key should not fail");

        assert_eq!(unquoted, Bytes::from_static(b"{\"id\":1}\n"));
        assert_eq!(quoted_exact, Bytes::from_static(b"{\"id\":1}\n"));
        assert!(quoted.is_empty());
    }

    #[test]
    fn ambiguous_unquoted_source_key_is_rejected() {
        let input = br#"{"Employees":[],"employees":[]}"#;

        let error =
            flatten_json_document_to_ndjson(input, &[source_key("EMPLOYEES")]).expect_err("ambiguous source key should fail");

        assert!(
            error
                .get_ref()
                .and_then(|source| source.downcast_ref::<SelectError>())
                .is_some_and(|error| *error == SelectError::AmbiguousFieldName)
        );
    }

    #[test]
    fn unquoted_source_key_scan_honors_cancellation() {
        let mut object = serde_json::Map::new();
        for index in 0..=super::JSON_CANCELLATION_CHECK_KEYS {
            object.insert(format!("field-{index}"), serde_json::Value::Null);
        }
        let mut checkpoints = 0;
        let mut cancel_on_second_checkpoint = || {
            checkpoints += 1;
            if checkpoints == 2 {
                Err(std::io::Error::new(std::io::ErrorKind::Interrupted, SelectError::Canceled))
            } else {
                Ok(())
            }
        };

        let error = remove_json_source_key_with_checkpoint(&mut object, "x", false, &mut cancel_on_second_checkpoint)
            .expect_err("key scan should stop at its second cancellation checkpoint");

        assert_eq!(error.kind(), std::io::ErrorKind::Interrupted);
        assert_eq!(checkpoints, 2);
        assert!(
            error
                .get_ref()
                .and_then(|source| source.downcast_ref::<SelectError>())
                .is_some_and(|error| *error == SelectError::Canceled)
        );
    }

    #[test]
    fn long_json_key_comparison_honors_cancellation() {
        let key = "x".repeat(super::JSON_CANCELLATION_CHECK_BYTES * 2);
        let mut checkpoints = 0;
        let mut cancel_on_second_checkpoint = || {
            checkpoints += 1;
            if checkpoints == 2 {
                Err(std::io::Error::new(std::io::ErrorKind::Interrupted, SelectError::Canceled))
            } else {
                Ok(())
            }
        };

        let error = json_key_eq_ignore_ascii_case_with_checkpoint(&key, &key, &mut cancel_on_second_checkpoint)
            .expect_err("long-key comparison should stop at its second cancellation checkpoint");

        assert_eq!(error.kind(), std::io::ErrorKind::Interrupted);
        assert_eq!(checkpoints, 2);
    }

    #[test]
    fn missing_source_key_produces_no_records() {
        let input = br#"{"employees":[]}"#;

        let result = flatten_json_document_to_ndjson(input, &[source_key("nonexistent")])
            .expect("missing source key should not fail the query");

        assert!(result.is_empty());
    }

    #[tokio::test]
    async fn source_path_type_mismatches_are_json_parsing_errors() {
        let cases: [(&str, &[u8], Vec<JsonPathSegment>); 4] = [
            ("key on array", b"[1]", vec![source_key("id")]),
            ("index on object", b"{}", vec![JsonPathSegment::Index(0)]),
            (
                "array wildcard on object",
                b"{}",
                vec![JsonPathSegment::ArrayWildcard, JsonPathSegment::ArrayWildcard],
            ),
            ("object wildcard on array", b"[1]", vec![JsonPathSegment::ObjectWildcard]),
        ];

        for (case, input, path) in cases {
            let memory_pool: Arc<dyn MemoryPool> =
                Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER));
            let mut output = json_document_ndjson_stream(
                Box::new(std::io::Cursor::new(input.to_vec())),
                input.len() as u64,
                JsonSource::from_path(path),
                Arc::new(SelectInputMetrics::default()),
                memory_pool,
                None,
            );
            let source = output
                .next()
                .await
                .unwrap_or_else(|| panic!("{case} should produce one stream error"))
                .unwrap_err();
            let error = QueryError::from(DataFusionError::ObjectStore(Box::new(source)));

            assert_eq!(error.select_error(), SelectError::JsonParsingError, "{case}");
            assert!(output.next().await.is_none(), "{case}");
        }
    }
}