diff --git a/Cargo.lock b/Cargo.lock index 03ed0bba0..6cc195af8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10615,6 +10615,8 @@ dependencies = [ "rustfs-s3select-api", "rustfs-test-utils", "s3s", + "serde", + "serde_json", "tokio", "tracing", ] diff --git a/crates/s3select-api/src/input_stream.rs b/crates/s3select-api/src/input_stream.rs index ff2b501dd..06942d4b4 100644 --- a/crates/s3select-api/src/input_stream.rs +++ b/crates/s3select-api/src/input_stream.rs @@ -326,13 +326,13 @@ fn invalid_gzip_header_error() -> io::Error { io::Error::new(io::ErrorKind::InvalidData, CompressionFormat::Gzip.invalid_header_error()) } -struct BlockingChannelReader { +pub(crate) struct BlockingChannelReader { receiver: mpsc::Receiver>, current: Bytes, } impl BlockingChannelReader { - fn new(receiver: mpsc::Receiver>) -> Self { + pub(crate) fn new(receiver: mpsc::Receiver>) -> Self { Self { receiver, current: Bytes::new(), @@ -639,7 +639,7 @@ fn error_chain_contains(error: &(dyn StdError + 'static)) find_error_source::(error).is_some() } -fn find_error_source<'a, T: StdError + 'static>(error: &'a (dyn StdError + 'static)) -> Option<&'a T> { +pub(crate) fn find_error_source<'a, T: StdError + 'static>(error: &'a (dyn StdError + 'static)) -> Option<&'a T> { let mut current = Some(error); for _ in 0..MAX_ERROR_SOURCE_DEPTH { let Some(error) = current else { diff --git a/crates/s3select-api/src/object_store.rs b/crates/s3select-api/src/object_store.rs index c135000aa..499436333 100644 --- a/crates/s3select-api/src/object_store.rs +++ b/crates/s3select-api/src/object_store.rs @@ -16,15 +16,11 @@ 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, + input_stream::{CompressionFormat, compressed_input_reader, compressed_input_stream, processed_bytes_limit}, query::{ ast::{JsonPathSegment, JsonSource}, parser::RustFsDialect, - session::{QueryExecutionGuard, QueryExecutionTracker}, + session::QueryExecutionTracker, }, resolve_select_object_store_handle, select_is_err_bucket_not_found, select_is_err_object_not_found, select_is_err_version_not_found, @@ -33,8 +29,7 @@ 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}, + execution::memory_pool::{MemoryPool, UnboundedMemoryPool}, object_store::{ Attributes, CopyOptions, Error as o_Error, GetOptions, GetRange, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result, path::Path, @@ -51,7 +46,6 @@ 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, @@ -63,15 +57,19 @@ use s3s::{ }; use std::collections::VecDeque; use std::ops::Range; +use std::sync::Arc; #[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; +mod json_document; + +use json_document::{compressed_json_document_ndjson_stream, json_document_ndjson_stream}; + fn select_default_read_buffer_size_u64() -> u64 { u64::try_from(SELECT_DEFAULT_READ_BUFFER_SIZE).unwrap_or(u64::MAX) } @@ -85,29 +83,18 @@ fn compression_format(input: &InputSerialization) -> Result 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)?; + // Size validation is eager; input and row production remain lazy. let stream = json_document_ndjson_stream( reader.stream, original_size, @@ -932,7 +907,7 @@ impl ObjectStore for EcObjectStore { 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(); @@ -1271,660 +1246,6 @@ impl ScanRangeState { } } -/// 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, - memory_pool: Arc, - query_tracker: Option, -) -> futures_core::stream::BoxStream<'static, Result> { - 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, - query_tracker: Option, -) -> futures_core::stream::BoxStream<'static, Result> { - 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

( - stream: SelectInputReader, - read_mode: JsonDocumentReadMode, - json_source: JsonSource, - memory_pool: Arc, - query_tracker: Option, - parser: P, -) -> futures_core::stream::BoxStream<'static, Result> -where - P: FnOnce(Vec, JsonSource, Arc) -> std::io::Result> + Send + 'static, -{ - AsyncTryStream::::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, - 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, -} - -struct JsonDocumentParseTask

{ - parser: P, - all_bytes: Vec, - json_source: JsonSource, - task_resources: JsonDocumentTaskResources, -} - -struct JsonDocumentCancellation { - cancelled: Arc, - queued: Arc>>, -} - -impl JsonDocumentCancellation { - fn new(cancelled: Arc, queued: Arc>>) -> Self { - Self { cancelled, queued } - } -} - -impl Drop for JsonDocumentCancellation { - 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 { - 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 { - 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 { - if let Some(select_error) = error.get_ref().and_then(|source| source.downcast_ref::()) { - 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>, -} - -enum QueryExecutionGuardState { - Pending(Option), - Started, - Cancelled, -} - -impl PendingQueryExecutionGuard { - fn new(query_tracker: Option) -> Self { - Self { - state: Arc::new(Mutex::new(QueryExecutionGuardState::Pending(query_tracker))), - } - } - - fn task_state(&self) -> Arc> { - Arc::clone(&self.state) - } - - fn start(state: &Mutex) -> std::io::Result> { - 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> { - 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> { - 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 = 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> { - // 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, - name: &str, - quoted: bool, - cancelled: &AtomicBool, -) -> std::io::Result> { - 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, - name: &str, - quoted: bool, - checkpoint: &mut impl FnMut() -> std::io::Result<()>, -) -> std::io::Result> { - 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 { - 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 { - 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 { - 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, - 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 { - 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 { - 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( stream: S, input_metrics: Arc, @@ -1965,21 +1286,6 @@ where }) } -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", @@ -1993,14 +1299,13 @@ fn incomplete_object_stream_error(remaining: impl std::fmt::Display) -> o_Error #[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, + EcObjectStore, EcObjectStoreBuildError, 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, + escaped_json_object_with_size, find_delimiter, flatten_json_document_to_ndjson, http_range_spec_from_get_range, + json_document_ndjson_stream, 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}; @@ -2011,7 +1316,7 @@ mod test { use bytes::Bytes; use datafusion::{ common::DataFusionError, - execution::memory_pool::{GreedyMemoryPool, MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation}, + execution::memory_pool::{GreedyMemoryPool, MemoryLimit, MemoryPool, MemoryReservation}, execution::{config::SessionConfig, context::SessionContext}, object_store::{self, GetOptions, GetRange, GetResultPayload, ObjectStore as _, path::Path}, physical_plan::ExecutionPlanProperties, @@ -2037,6 +1342,8 @@ mod test { use tokio::{io::AsyncReadExt, sync::Semaphore}; + const JSON_DOCUMENT_TEST_MEMORY_BYTES: usize = 8 * 1024 * 1024; + #[derive(Debug)] struct AdmissionObservingMemoryPool { inner: GreedyMemoryPool, @@ -3771,9 +3078,14 @@ mod test { .expect("put JSON input metrics fixture"); let snapshot = prepare_test_snapshot(BUCKET, object).await; let input_metrics = Arc::new(SelectInputMetrics::default()); + let memory_limit = if json_type == JSONType::DOCUMENT { + JSON_DOCUMENT_TEST_MEMORY_BYTES + } else { + 1024 * 1024 + }; let store = EcObjectStore::build_with_snapshot( json_input(BUCKET, object, json_type), - Arc::new(GreedyMemoryPool::new(1024 * 1024)), + Arc::new(GreedyMemoryPool::new(memory_limit)), None, Arc::clone(&input_metrics), snapshot, @@ -3858,8 +3170,7 @@ mod test { #[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 = Arc::new(GreedyMemoryPool::new(required - 1)); + let memory_pool: Arc = Arc::new(GreedyMemoryPool::new(1)); let mut output = json_document_ndjson_stream( Box::new(std::io::Cursor::new(input.clone())), input.len() as u64, @@ -3867,7 +3178,8 @@ mod test { Arc::new(SelectInputMetrics::default()), memory_pool, None, - ); + ) + .expect("build constrained JSON document stream"); let err = output .next() @@ -3903,7 +3215,7 @@ mod test { let output = compressed_json_document_ndjson_stream( reader, JsonSource::default(), - Arc::new(GreedyMemoryPool::new(1024 * 1024)), + Arc::new(GreedyMemoryPool::new(JSON_DOCUMENT_TEST_MEMORY_BYTES)), None, ) .try_collect::>() @@ -3937,16 +3249,16 @@ mod test { let mut output = compressed_json_document_ndjson_stream( reader, JsonSource::default(), - Arc::new(GreedyMemoryPool::new(1024 * 1024)), + Arc::new(GreedyMemoryPool::new(64 * 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() + while let Some(result) = output.next().await { + if let Err(source) = result { + assert!(output.next().await.is_none()); + return QueryError::from(DataFusionError::ObjectStore(Box::new(source))).select_error(); + } + } + panic!("compressed JSON DOCUMENT decoding must fail"); } #[tokio::test] @@ -3975,86 +3287,44 @@ mod test { } } - #[test] - fn compressed_json_document_buffer_grows_amortized_and_reserves_capacity() { - let memory_pool: Arc = 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; + #[tokio::test] + async fn json_document_input_record_limit_applies_after_decompression() { + let at_limit = escaped_json_object_with_size(MAX_SELECT_RECORD_BYTES); + let over_limit = escaped_json_object_with_size(MAX_SELECT_RECORD_BYTES + 1); - 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); + for format in [CompressionFormat::Gzip, CompressionFormat::Bzip2] { + let compressed = encode_compressed_fixture(format, &at_limit).await; + 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()), + u64::MAX, + None, + ); + compressed_json_document_ndjson_stream( + reader, + JsonSource::default(), + Arc::new(GreedyMemoryPool::new(64 * 1024 * 1024)), + None, + ) + .try_collect::>() + .await + .expect("one-megabyte compressed input record should be accepted"); + + let compressed = encode_compressed_fixture(format, &over_limit).await; assert_eq!( - reservation.size(), - super::json_document_memory_reservation_bytes(buffer.capacity(), &json_source) - .expect("test reservation should fit") + compressed_json_document_select_error(format, compressed, u64::MAX).await, + SelectError::OverMaxRecordSize ); } - - 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 = 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::(), - Some(DataFusionError::ResourcesExhausted(_)) - )); - - let memory_pool = Arc::new(GreedyMemoryPool::new(required)); - let output: Vec = 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 memory_pool = Arc::new(GreedyMemoryPool::new(JSON_DOCUMENT_TEST_MEMORY_BYTES)); let input_metrics = Arc::new(SelectInputMetrics::default()); let output: Vec = json_document_ndjson_stream( Box::new(std::io::Cursor::new(input.clone())), @@ -4064,6 +3334,7 @@ mod test { memory_pool.clone(), None, ) + .expect("build JSON document stream") .try_collect() .await .expect("JSON conversion should fit the pool"); @@ -4076,11 +3347,11 @@ mod test { } #[tokio::test] - async fn test_json_document_stream_rejects_early_eof() { + async fn test_json_document_stream_reports_early_eof_after_completed_rows() { 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 = Arc::new(GreedyMemoryPool::new(4 * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER)); + let memory_pool: Arc = Arc::new(GreedyMemoryPool::new(JSON_DOCUMENT_TEST_MEMORY_BYTES)); let mut output = json_document_ndjson_stream( Box::new(std::io::Cursor::new(input)), 4, @@ -4088,7 +3359,15 @@ mod test { Arc::clone(&input_metrics), memory_pool, None, - ); + ) + .expect("build short JSON document stream"); + + let row = output + .next() + .await + .expect("completed row") + .expect("completed rows may precede a terminal stream error"); + assert_eq!(row, Bytes::from_static(b"{}\n")); let err = output .next() @@ -4109,8 +3388,7 @@ mod test { #[tokio::test] async fn malformed_json_document_stream_has_typed_select_error() { let input = b"{bad".to_vec(); - let memory_pool: Arc = - Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER)); + let memory_pool: Arc = Arc::new(GreedyMemoryPool::new(JSON_DOCUMENT_TEST_MEMORY_BYTES)); let mut output = json_document_ndjson_stream( Box::new(std::io::Cursor::new(input.clone())), input.len() as u64, @@ -4118,7 +3396,8 @@ mod test { Arc::new(SelectInputMetrics::default()), memory_pool, None, - ); + ) + .expect("build malformed JSON document stream"); let source = output .next() @@ -4131,34 +3410,6 @@ mod test { 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 = - 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(); @@ -4213,7 +3464,7 @@ mod test { } #[test] - fn test_json_document_cancelled_queued_parse_releases_before_dequeue() { + fn test_json_document_parser_does_not_use_tokio_blocking_pool() { let runtime = tokio::runtime::Builder::new_multi_thread() .worker_threads(2) .max_blocking_threads(1) @@ -4243,10 +3494,8 @@ mod test { 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, reservation_released) = + AdmissionObservingMemoryPool::new(JSON_DOCUMENT_TEST_MEMORY_BYTES, Arc::clone(&admission)); let memory_pool = Arc::new(memory_pool); let query_memory_pool: Arc = memory_pool.clone(); let mut output = json_document_ndjson_stream( @@ -4256,243 +3505,41 @@ mod test { Arc::new(SelectInputMetrics::default()), query_memory_pool, Some(query_tracker), - ); + ) + .expect("build JSON document stream"); - { - let next = output.next(); - futures::pin_mut!(next); - assert!(futures::poll!(next.as_mut()).is_pending()); - } + let row = tokio::time::timeout(std::time::Duration::from_secs(1), output.next()) + .await + .expect("dedicated parser must not wait for Tokio's blocking pool") + .expect("JSON document should produce one row") + .expect("JSON document row should be valid"); + assert_eq!(row, Bytes::from_static(b"{}\n")); + assert!( + tokio::time::timeout(std::time::Duration::from_secs(1), output.next()) + .await + .expect("parser should finish while the blocking pool is occupied") + .is_none() + ); drop(output); + tokio::pin!(reservation_released); assert!( - tokio::time::timeout(std::time::Duration::from_millis(100), reservation_released) + tokio::time::timeout(std::time::Duration::from_secs(1), &mut reservation_released) .await - .expect("queued parse cancellation should release its memory immediately") + .expect("completed parser should release memory") .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()) + tokio::time::timeout(std::time::Duration::from_secs(1), 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 = - 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::().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 = - 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::().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 = - 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("completed parser should release query admission") .expect("query admission should remain open"); drop(recovered_permit); + + release_blocking_tx.send(()).expect("release blocking worker"); + blocker.await.expect("blocking worker should finish"); assert_eq!(admission.available_permits(), 1); }); } @@ -4583,14 +3630,18 @@ mod test { #[test] fn json_document_logical_record_enforces_one_megabyte_limit() { - const OBJECT_OVERHEAD: usize = br#"{"v":""}"#.len(); + const SCALAR_WRAPPER_BYTES: usize = br#"{"":0}"#.len(); + let input = br#"{"v":0}"#; + let path = vec![source_key("v")]; - 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 alias = "a".repeat(MAX_SELECT_RECORD_BYTES - SCALAR_WRAPPER_BYTES); + let source = JsonSource::new(path.clone(), Some(alias)); + let line = super::parse_json_document_to_lines(input, &source).expect("one-megabyte logical record should be accepted"); + assert_eq!(line[0].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"); + let alias = "a".repeat(MAX_SELECT_RECORD_BYTES + 1 - SCALAR_WRAPPER_BYTES); + let source = JsonSource::new(path.clone(), Some(alias)); + let error = super::parse_json_document_to_lines(input, &source).expect_err("oversized logical record must fail"); assert_eq!(error.kind(), std::io::ErrorKind::InvalidData); assert!(error.get_ref().is_some_and(|source| { source @@ -4598,8 +3649,10 @@ mod test { .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"); + let escaped_alias = "\0".repeat(MAX_SELECT_RECORD_BYTES / 6); + let source = JsonSource::new(path, Some(escaped_alias)); + let error = super::parse_json_document_to_lines(input, &source) + .expect_err("escaped output must be bounded while it is serialized"); assert!(error.get_ref().is_some_and(|source| { source .downcast_ref::() @@ -4905,8 +3958,7 @@ mod test { ]; for (case, input, path) in cases { - let memory_pool: Arc = - Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER)); + let memory_pool: Arc = Arc::new(GreedyMemoryPool::new(JSON_DOCUMENT_TEST_MEMORY_BYTES)); let mut output = json_document_ndjson_stream( Box::new(std::io::Cursor::new(input.to_vec())), input.len() as u64, @@ -4914,7 +3966,8 @@ mod test { Arc::new(SelectInputMetrics::default()), memory_pool, None, - ); + ) + .expect("build source-path JSON document stream"); let source = output .next() .await diff --git a/crates/s3select-api/src/object_store/json_document.rs b/crates/s3select-api/src/object_store/json_document.rs new file mode 100644 index 000000000..b8936ff5e --- /dev/null +++ b/crates/s3select-api/src/object_store/json_document.rs @@ -0,0 +1,1968 @@ +// 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 super::MAX_JSON_DOCUMENT_BYTES; +use crate::{ + QueryError, SelectError, SelectInputMetrics, + input_stream::{ + BlockingChannelReader, MAX_SELECT_RECORD_BYTES, SELECT_DECODE_CHUNK_BYTES, SelectInputReader, find_error_source, + }, + metrics::SelectInputMetricsRecorder, + query::{ + ast::{JsonPathSegment, JsonSource}, + session::{QueryExecutionGuard, QueryExecutionTracker}, + }, +}; +use bytes::Bytes; +use datafusion::{ + common::runtime::SpawnedTask, + execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}, + object_store::{Error as ObjectStoreError, Result as ObjectStoreResult}, +}; +use futures::{StreamExt as _, TryStreamExt as _, stream}; +use futures_core::Stream; +use std::{ + fmt, + future::Future as _, + io, + pin::Pin, + sync::{ + Arc, LazyLock, + atomic::{AtomicBool, Ordering}, + }, + task::{Context, Poll}, +}; +use tokio::{ + io::AsyncReadExt, + sync::{Semaphore, mpsc, oneshot}, +}; +use tokio_util::io::ReaderStream; + +const JSON_DOCUMENT_INPUT_CHANNEL_CAPACITY: usize = 2; +const JSON_DOCUMENT_ROW_CHANNEL_CAPACITY: usize = 1; +const JSON_DOCUMENT_OUTPUT_BATCH_BYTES: usize = SELECT_DECODE_CHUNK_BYTES; +const JSON_DOCUMENT_MAX_DEPTH: usize = 100; +const JSON_DOCUMENT_RAW_MEMORY_RESERVATION_MULTIPLIER: usize = 2; +const JSON_DOCUMENT_DOM_MEMORY_RESERVATION_MULTIPLIER: usize = 32; +const JSON_SCALAR_COLUMN_MEMORY_RESERVATION_MULTIPLIER: usize = 2; +const JSON_DOCUMENT_MIN_PARSER_THREADS: usize = 4; +const JSON_DOCUMENT_MAX_PARSER_THREADS: usize = 32; +pub(super) const JSON_CANCELLATION_CHECK_BYTES: usize = 64 * 1024; +pub(super) const JSON_CANCELLATION_CHECK_KEYS: usize = 1024; + +static JSON_DOCUMENT_PARSER_THREADS: LazyLock = LazyLock::new(|| { + std::thread::available_parallelism().map_or(JSON_DOCUMENT_MIN_PARSER_THREADS, |parallelism| { + parallelism + .get() + .saturating_mul(2) + .clamp(JSON_DOCUMENT_MIN_PARSER_THREADS, JSON_DOCUMENT_MAX_PARSER_THREADS) + }) +}); +static JSON_DOCUMENT_PARSER_ADMISSION: LazyLock = LazyLock::new(|| Semaphore::new(*JSON_DOCUMENT_PARSER_THREADS)); + +pub(super) fn json_document_ndjson_stream( + stream: SelectInputReader, + original_size: u64, + json_source: JsonSource, + input_metrics: Arc, + memory_pool: Arc, + query_tracker: Option, +) -> ObjectStoreResult> + Send + 'static>>> { + validate_json_document_size(original_size)?; + Ok(json_document_stream( + stream, + JsonDocumentReadMode::Exact { + original_size, + input_metrics: input_metrics.recorder(), + }, + json_source, + memory_pool, + query_tracker, + )) +} + +pub(super) fn compressed_json_document_ndjson_stream( + stream: SelectInputReader, + json_source: JsonSource, + memory_pool: Arc, + query_tracker: Option, +) -> Pin> + Send + 'static>> { + json_document_stream(stream, JsonDocumentReadMode::Bounded, json_source, memory_pool, query_tracker) +} + +enum JsonDocumentReadMode { + Exact { + original_size: u64, + input_metrics: SelectInputMetricsRecorder, + }, + Bounded, +} + +fn json_document_stream( + stream: SelectInputReader, + read_mode: JsonDocumentReadMode, + json_source: JsonSource, + memory_pool: Arc, + query_tracker: Option, +) -> Pin> + Send + 'static>> { + let config = JsonDocumentPipelineConfig { + stream, + read_mode, + json_source, + memory_pool, + query_tracker, + }; + Box::pin(stream::once(JsonDocumentPipeline::start(config)).try_flatten()) +} + +struct JsonDocumentPipelineConfig { + stream: SelectInputReader, + read_mode: JsonDocumentReadMode, + json_source: JsonSource, + memory_pool: Arc, + query_tracker: Option, +} + +struct JsonDocumentPipeline { + receiver: mpsc::Receiver, + read_task: Option>, + parse_result: oneshot::Receiver>, + done: bool, + cancellation: Arc, + _resources: Arc, +} + +struct JsonDocumentPipelineResources { + _channel_reservation: MemoryReservation, + _query_guard: Option, +} + +impl JsonDocumentPipeline { + async fn start(config: JsonDocumentPipelineConfig) -> ObjectStoreResult { + let parser_permit = JSON_DOCUMENT_PARSER_ADMISSION + .acquire() + .await + .map_err(|_| json_document_worker_store_error("parser admission"))?; + let JsonDocumentPipelineConfig { + stream, + read_mode, + json_source, + memory_pool, + query_tracker, + } = config; + let query_guard = match query_tracker.as_ref() { + Some(query_tracker) => Some(query_tracker.query_guard().ok_or_else(json_document_cancelled_store_error)?), + None => None, + }; + let channel_reservation = MemoryConsumer::new("S3 Select JSON document channels").register(&memory_pool); + channel_reservation + .try_resize(json_document_channel_reservation_bytes()) + .map_err(resource_exhausted_store_error)?; + let resources = Arc::new(JsonDocumentPipelineResources { + _channel_reservation: channel_reservation, + _query_guard: query_guard, + }); + + let (input_tx, input_rx) = mpsc::channel(JSON_DOCUMENT_INPUT_CHANNEL_CAPACITY); + let (row_tx, row_rx) = mpsc::channel(JSON_DOCUMENT_ROW_CHANNEL_CAPACITY); + let cancellation = Arc::new(AtomicBool::new(false)); + + let parse_cancellation = Arc::clone(&cancellation); + let parser_memory_pool = Arc::clone(&memory_pool); + let parser_resources = Arc::clone(&resources); + let (parse_result_tx, parse_result_rx) = oneshot::channel(); + std::thread::Builder::new() + .name("s3select-json-document".to_string()) + .spawn(move || { + let _parser_permit = parser_permit; + let _resources = parser_resources; + let reservation = MemoryConsumer::new("S3 Select JSON document record").register(&parser_memory_pool); + let result = parse_json_document( + BlockingChannelReader::new(input_rx), + &json_source, + &reservation, + parse_cancellation.as_ref(), + |rows| row_tx.blocking_send(rows).map_err(|_| json_document_cancelled_error()), + ) + .map_err(classify_json_document_error); + let _ = parse_result_tx.send(result); + }) + .map_err(|error| json_document_worker_spawn_store_error("parser", error))?; + + // The parser waits on bounded channels, so it must not occupy Tokio's + // blocking pool. Query admission bounds the dedicated worker threads. + let read_cancellation = Arc::clone(&cancellation); + let read_resources = Arc::clone(&resources); + let read_task = SpawnedTask::spawn(async move { + let _resources = read_resources; + forward_json_document_input(stream, read_mode, input_tx, read_cancellation).await; + }); + + Ok(Self { + receiver: row_rx, + read_task: Some(read_task), + parse_result: parse_result_rx, + done: false, + cancellation, + _resources: resources, + }) + } + + fn poll_terminal(&mut self, context: &mut Context<'_>) -> Poll>> { + if let Some(read_task) = self.read_task.as_mut() { + match Pin::new(read_task).poll(context) { + Poll::Ready(Ok(())) => self.read_task = None, + Poll::Ready(Err(_)) => { + self.done = true; + return Poll::Ready(Some(Err(json_document_worker_store_error("reader")))); + } + Poll::Pending => return Poll::Pending, + } + } + match Pin::new(&mut self.parse_result).poll(context) { + Poll::Ready(Ok(Ok(()))) => { + self.done = true; + Poll::Ready(None) + } + Poll::Ready(Ok(Err(error))) => { + self.done = true; + Poll::Ready(Some(Err(error))) + } + Poll::Ready(Err(_)) => { + self.done = true; + Poll::Ready(Some(Err(json_document_worker_store_error("parser")))) + } + Poll::Pending => Poll::Pending, + } + } +} + +impl Stream for JsonDocumentPipeline { + type Item = ObjectStoreResult; + + fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + if self.done { + return Poll::Ready(None); + } + match self.receiver.poll_recv(context) { + Poll::Ready(Some(bytes)) => Poll::Ready(Some(Ok(bytes))), + Poll::Ready(None) => self.poll_terminal(context), + Poll::Pending => Poll::Pending, + } + } +} + +impl Drop for JsonDocumentPipeline { + fn drop(&mut self) { + self.cancellation.store(true, Ordering::Release); + } +} + +async fn forward_json_document_input( + stream: SelectInputReader, + read_mode: JsonDocumentReadMode, + input_tx: mpsc::Sender>, + cancellation: Arc, +) { + let (stream, expected_size, input_metrics): (SelectInputReader, Option, Option) = + match read_mode { + JsonDocumentReadMode::Exact { + original_size, + input_metrics, + } => (Box::new(stream.take(original_size)), Some(original_size), Some(input_metrics)), + JsonDocumentReadMode::Bounded => (stream, None, None), + }; + let mut stream = ReaderStream::with_capacity(stream, SELECT_DECODE_CHUNK_BYTES); + let mut remaining = expected_size; + let mut total = 0_u64; + + loop { + if cancellation.load(Ordering::Acquire) || input_tx.is_closed() { + return; + } + let bytes = tokio::select! { + biased; + _ = input_tx.closed() => return, + result = stream.next() => match result { + Some(Ok(bytes)) => bytes, + Some(Err(error)) => { + send_input_error(&input_tx, error).await; + return; + } + None => { + if let Some(remaining) = remaining + && remaining > 0 + { + send_input_error( + &input_tx, + io::Error::new( + io::ErrorKind::UnexpectedEof, + format!("object stream ended with {remaining} bytes remaining"), + ), + ) + .await; + } + return; + } + }, + }; + + if let Some(input_metrics) = input_metrics.as_ref() { + input_metrics.record_uncompressed(bytes.len()); + } + let read = match u64::try_from(bytes.len()) { + Ok(read) => read, + Err(_) => { + send_input_error(&input_tx, resource_exhausted_io_error()).await; + return; + } + }; + total = match total.checked_add(read) { + Some(total) => total, + None => { + send_input_error(&input_tx, resource_exhausted_io_error()).await; + return; + } + }; + if total > MAX_JSON_DOCUMENT_BYTES { + send_input_error(&input_tx, resource_exhausted_io_error()).await; + return; + } + if let Some(value) = remaining.as_mut() { + let Some(new_remaining) = value.checked_sub(read) else { + send_input_error(&input_tx, resource_exhausted_io_error()).await; + return; + }; + *value = new_remaining; + } + let sent = tokio::select! { + biased; + _ = input_tx.closed() => false, + result = input_tx.send(Ok(bytes)) => result.is_ok(), + }; + if !sent { + return; + } + } +} + +async fn send_input_error(input_tx: &mpsc::Sender>, error: io::Error) { + let _ = input_tx.send(Err(error)).await; +} + +fn parse_json_document( + reader: R, + json_source: &JsonSource, + reservation: &MemoryReservation, + cancellation: &AtomicBool, + emit: F, +) -> io::Result<()> +where + R: io::Read, + F: FnMut(Bytes) -> io::Result<()>, +{ + let reader = io::BufReader::with_capacity(SELECT_DECODE_CHUNK_BYTES, reader); + let mut decoder = JsonDocumentDecoder::new(json_source, reservation, cancellation, emit); + match parse_json_document_reader(reader, &mut decoder, cancellation) { + Ok(()) => Ok(()), + Err(error) => { + decoder.flush_output()?; + Err(error) + } + } +} + +fn parse_json_document_reader( + mut reader: R, + decoder: &mut JsonDocumentDecoder<'_, F>, + cancellation: &AtomicBool, +) -> io::Result<()> +where + R: io::BufRead, + F: FnMut(Bytes) -> io::Result<()>, +{ + loop { + ensure_json_parse_active(cancellation)?; + let available = reader.fill_buf()?; + if available.is_empty() { + return decoder.finish(); + } + let consumed = available.len(); + for byte in available.iter().copied() { + decoder.push(byte)?; + } + reader.consume(consumed); + } +} + +struct JsonDocumentDecoder<'a, F> { + source: &'a JsonSource, + reservation: &'a MemoryReservation, + cancellation: &'a AtomicBool, + emit: F, + state: JsonRootState, + record: Vec, + stack: Vec, + in_string: bool, + escape_next: bool, + array_after_comma: bool, + array_value_terminated: bool, + array_index: usize, + output: Vec, + emitted_rows: usize, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum JsonRootState { + Start, + Array(RootArrayPlan), + Single, + Done, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum RootArrayPlan { + Each { path_start: usize }, + Index { index: usize, path_start: usize }, +} + +impl<'a, F> JsonDocumentDecoder<'a, F> +where + F: FnMut(Bytes) -> io::Result<()>, +{ + fn new(source: &'a JsonSource, reservation: &'a MemoryReservation, cancellation: &'a AtomicBool, emit: F) -> Self { + Self { + source, + reservation, + cancellation, + emit, + state: JsonRootState::Start, + record: Vec::new(), + stack: Vec::with_capacity(JSON_DOCUMENT_MAX_DEPTH), + in_string: false, + escape_next: false, + array_after_comma: false, + array_value_terminated: false, + array_index: 0, + output: Vec::new(), + emitted_rows: 0, + } + } + + fn push(&mut self, byte: u8) -> io::Result<()> { + match self.state { + JsonRootState::Start => self.push_start(byte), + JsonRootState::Array(plan) => self.push_array(byte, plan), + JsonRootState::Single => self.push_single(byte), + JsonRootState::Done => { + if is_json_whitespace(byte) { + Ok(()) + } else { + Err(json_document_parse_error(JsonDocumentParseError::Malformed)) + } + } + } + } + + fn push_start(&mut self, byte: u8) -> io::Result<()> { + if is_json_whitespace(byte) { + return Ok(()); + } + if byte == b'[' { + let plan = root_array_plan(self.source.path())?; + self.stack.push(b']'); + self.state = JsonRootState::Array(plan); + return Ok(()); + } + + self.state = JsonRootState::Single; + self.push_single(byte) + } + + fn push_single(&mut self, byte: u8) -> io::Result<()> { + if is_json_whitespace(byte) && self.stack.is_empty() && !self.in_string && record_has_value(&self.record) { + return self.emit_single_record(false); + } + self.reserve_record_byte()?; + self.record.push(byte); + if self.observe_nested_byte(byte)? { + self.emit_single_record(false)?; + } + Ok(()) + } + + fn push_array(&mut self, byte: u8, plan: RootArrayPlan) -> io::Result<()> { + if self.in_string { + self.reserve_record_byte()?; + self.record.push(byte); + if self.escape_next { + self.escape_next = false; + } else if byte == b'\\' { + self.escape_next = true; + } else if byte == b'"' { + self.in_string = false; + } + return Ok(()); + } + + if self.stack.len() == 1 { + if is_json_whitespace(byte) { + self.array_value_terminated |= record_has_value(&self.record); + return Ok(()); + } + if self.array_value_terminated && !matches!(byte, b',' | b']') { + return Err(json_document_parse_error(JsonDocumentParseError::Malformed)); + } + } + + match byte { + b'"' => { + self.reserve_record_byte()?; + self.record.push(byte); + self.in_string = true; + self.array_after_comma = false; + Ok(()) + } + b'{' | b'[' => { + self.reserve_record_byte()?; + self.record.push(byte); + self.push_container(byte)?; + self.array_after_comma = false; + Ok(()) + } + b'}' => { + if self.stack.last().copied() != Some(b'}') { + return Err(json_document_parse_error(JsonDocumentParseError::Malformed)); + } + self.stack.pop(); + self.reserve_record_byte()?; + self.record.push(byte); + self.array_after_comma = false; + Ok(()) + } + b']' if self.stack.len() == 1 => self.finish_array(plan), + b']' => { + if self.stack.last().copied() != Some(b']') { + return Err(json_document_parse_error(JsonDocumentParseError::Malformed)); + } + self.stack.pop(); + self.reserve_record_byte()?; + self.record.push(byte); + self.array_after_comma = false; + Ok(()) + } + b',' if self.stack.len() == 1 => { + if !record_has_value(&self.record) { + return Err(json_document_parse_error(JsonDocumentParseError::Malformed)); + } + self.emit_array_record(plan)?; + self.array_after_comma = true; + Ok(()) + } + _ => { + if !is_json_whitespace(byte) { + self.array_after_comma = false; + } + self.reserve_record_byte()?; + self.record.push(byte); + Ok(()) + } + } + } + + fn observe_nested_byte(&mut self, byte: u8) -> io::Result { + if self.in_string { + if self.escape_next { + self.escape_next = false; + } else if byte == b'\\' { + self.escape_next = true; + } else if byte == b'"' { + self.in_string = false; + } + return Ok(!self.in_string && self.stack.is_empty()); + } + + match byte { + b'"' => self.in_string = true, + b'{' | b'[' => self.push_container(byte)?, + b'}' | b']' => { + let expected = if byte == b'}' { b'}' } else { b']' }; + if self.stack.last().copied() != Some(expected) { + return Err(json_document_parse_error(JsonDocumentParseError::Malformed)); + } + self.stack.pop(); + } + _ => {} + } + Ok(!self.in_string && self.stack.is_empty() && matches!(byte, b'}' | b']')) + } + + fn push_container(&mut self, opening: u8) -> io::Result<()> { + if self.stack.len() >= JSON_DOCUMENT_MAX_DEPTH { + return Err(json_document_parse_error(JsonDocumentParseError::DepthLimit { + limit: JSON_DOCUMENT_MAX_DEPTH, + })); + } + self.stack.push(if opening == b'{' { b'}' } else { b']' }); + Ok(()) + } + + fn finish_array(&mut self, plan: RootArrayPlan) -> io::Result<()> { + if record_has_value(&self.record) { + self.emit_array_record(plan)?; + } else if self.array_after_comma { + return Err(json_document_parse_error(JsonDocumentParseError::Malformed)); + } + self.stack.clear(); + self.record = Vec::new(); + self.reservation.free(); + self.state = JsonRootState::Done; + Ok(()) + } + + fn emit_array_record(&mut self, plan: RootArrayPlan) -> io::Result<()> { + ensure_json_parse_active(self.cancellation)?; + let path = match plan { + RootArrayPlan::Each { path_start } => Some(&self.source.path()[path_start..]), + RootArrayPlan::Index { index, path_start } if index == self.array_index => Some(&self.source.path()[path_start..]), + RootArrayPlan::Index { .. } => None, + }; + if let Some(path) = path { + self.process_record(path, false, false)?; + } else { + validate_json_value(&self.record, self.cancellation, false)?; + } + self.array_index = self.array_index.checked_add(1).ok_or_else(resource_exhausted_io_error)?; + self.array_value_terminated = false; + self.reset_record_buffer()?; + Ok(()) + } + + fn finish(&mut self) -> io::Result<()> { + ensure_json_parse_active(self.cancellation)?; + match self.state { + JsonRootState::Start => Err(json_document_parse_error(JsonDocumentParseError::Malformed)), + JsonRootState::Array(_) => { + if record_has_value(&self.record) { + validate_json_value(&self.record, self.cancellation, true)?; + } + Err(json_document_parse_error(JsonDocumentParseError::Truncated)) + } + JsonRootState::Done => self.flush_output(), + JsonRootState::Single => self.emit_single_record(true).and_then(|()| self.flush_output()), + } + } + + fn emit_single_record(&mut self, eof: bool) -> io::Result<()> { + let source_path = self.source.path(); + let path = source_path + .strip_prefix(&[JsonPathSegment::ArrayWildcard]) + .unwrap_or(source_path); + let implicitly_expand_arrays = matches!(source_path, [] | [JsonPathSegment::Key { .. }]); + self.process_record(path, implicitly_expand_arrays, eof)?; + self.record = Vec::new(); + self.reservation.free(); + self.state = JsonRootState::Done; + Ok(()) + } + + fn process_record(&mut self, path: &[JsonPathSegment], implicitly_expand_arrays: bool, eof: bool) -> io::Result<()> { + if path.is_empty() { + validate_json_value(&self.record, self.cancellation, eof)?; + return self.append_raw_value(is_json_object(&self.record)); + } + + resize_json_document_dom_reservation(self.reservation, self.record.capacity(), self.source)?; + let value = deserialize_json_value(&self.record, self.cancellation, eof)?; + let defer_output_flush = !implicitly_expand_arrays + && !path + .iter() + .any(|segment| matches!(segment, JsonPathSegment::ArrayWildcard | JsonPathSegment::ObjectWildcard)); + let emitted_rows = self.emitted_rows; + let result = visit_json_path(value, path, self.cancellation, &mut |value| { + self.emit_value(value, implicitly_expand_arrays, defer_output_flush) + }); + let shrink_result = resize_json_document_record_reservation(self.reservation, self.record.capacity(), self.source); + result?; + shrink_result?; + if defer_output_flush + && self.emitted_rows > emitted_rows + && (self.emitted_rows == 1 || self.output.len() >= JSON_DOCUMENT_OUTPUT_BATCH_BYTES) + { + self.flush_output()?; + } + Ok(()) + } + + fn emit_value( + &mut self, + value: serde_json::Value, + implicitly_expand_arrays: bool, + defer_output_flush: bool, + ) -> io::Result<()> { + if implicitly_expand_arrays && let serde_json::Value::Array(array) = value { + for item in array { + ensure_json_parse_active(self.cancellation)?; + self.append_json_value(item, defer_output_flush)?; + } + return Ok(()); + } + self.append_json_value(value, defer_output_flush) + } + + fn append_json_value(&mut self, value: serde_json::Value, defer_output_flush: bool) -> io::Result<()> { + let scalar_column = (!value.is_object()).then(|| scalar_column(self.source)); + let record_start = self.output.len(); + let (serialize_result, limit_exceeded) = { + let mut writer = CancellableJsonWriter { + inner: &mut self.output, + record_start, + cancellation: self.cancellation, + bytes_since_check: 0, + limit_exceeded: false, + }; + let serialize_result = match scalar_column { + Some(column) => { + io::Write::write_all(&mut writer, b"{").is_ok() + && serde_json::to_writer(&mut writer, column).is_ok() + && io::Write::write_all(&mut writer, b":").is_ok() + && serde_json::to_writer(&mut writer, &value).is_ok() + && io::Write::write_all(&mut writer, b"}").is_ok() + } + None => serde_json::to_writer(&mut writer, &value).is_ok(), + }; + (serialize_result, writer.limit_exceeded) + }; + if limit_exceeded { + self.output.truncate(record_start); + return Err(over_max_record_size_error()); + } + if !serialize_result { + self.output.truncate(record_start); + ensure_json_parse_active(self.cancellation)?; + return Err(json_document_parse_error(JsonDocumentParseError::Malformed)); + } + if let Err(error) = ensure_json_parse_active(self.cancellation) { + self.output.truncate(record_start); + return Err(error); + } + drop(value); + self.output.push(b'\n'); + self.complete_output_row(!defer_output_flush) + } + + fn append_raw_value(&mut self, object: bool) -> io::Result<()> { + let record = trim_json_whitespace(&self.record); + if record.len() > MAX_SELECT_RECORD_BYTES { + return Err(over_max_record_size_error()); + } + let scalar_column = (!object).then(|| scalar_column(self.source)); + let wrapper_bytes = scalar_column + .map_or(Ok(0), |column| { + column + .len() + .checked_mul(6) + .and_then(|bytes| bytes.checked_add(5)) + .ok_or_else(resource_exhausted_io_error) + })? + .min(MAX_SELECT_RECORD_BYTES); + let estimated_record_bytes = record + .len() + .checked_add(wrapper_bytes) + .ok_or_else(resource_exhausted_io_error)? + .min(MAX_SELECT_RECORD_BYTES); + let required = self + .output + .len() + .checked_add(estimated_record_bytes) + .and_then(|length| length.checked_add(1)) + .ok_or_else(resource_exhausted_io_error)?; + self.output + .try_reserve(required.saturating_sub(self.output.len())) + .map_err(|_| resource_exhausted_io_error())?; + let record_start = self.output.len(); + if let Some(column) = scalar_column { + let (serialize_result, limit_exceeded) = { + let mut writer = CancellableJsonWriter { + inner: &mut self.output, + record_start, + cancellation: self.cancellation, + bytes_since_check: 0, + limit_exceeded: false, + }; + let serialize_result = io::Write::write_all(&mut writer, b"{").is_ok() + && serde_json::to_writer(&mut writer, column).is_ok() + && io::Write::write_all(&mut writer, b":").is_ok(); + (serialize_result, writer.limit_exceeded) + }; + if limit_exceeded { + self.output.truncate(record_start); + return Err(over_max_record_size_error()); + } + if !serialize_result { + self.output.truncate(record_start); + ensure_json_parse_active(self.cancellation)?; + return Err(json_document_parse_error(JsonDocumentParseError::Malformed)); + } + } + let mut in_string = false; + let mut escape_next = false; + for chunk in record.chunks(JSON_CANCELLATION_CHECK_BYTES) { + if let Err(error) = ensure_json_parse_active(self.cancellation) { + self.output.truncate(record_start); + return Err(error); + } + for byte in chunk.iter().copied() { + let emit_byte = in_string || byte == b'"' || !is_json_whitespace(byte); + if emit_byte && self.output.len().saturating_sub(record_start) >= MAX_SELECT_RECORD_BYTES { + self.output.truncate(record_start); + return Err(over_max_record_size_error()); + } + if in_string { + self.output.push(byte); + if escape_next { + escape_next = false; + } else if byte == b'\\' { + escape_next = true; + } else if byte == b'"' { + in_string = false; + } + } else if byte == b'"' { + in_string = true; + self.output.push(byte); + } else if !is_json_whitespace(byte) { + self.output.push(byte); + } + } + } + if !object { + if self.output.len().saturating_sub(record_start) >= MAX_SELECT_RECORD_BYTES { + self.output.truncate(record_start); + return Err(over_max_record_size_error()); + } + self.output.push(b'}'); + } + self.output.push(b'\n'); + self.complete_output_row(true) + } + + fn complete_output_row(&mut self, flush_allowed: bool) -> io::Result<()> { + self.emitted_rows = self.emitted_rows.checked_add(1).ok_or_else(resource_exhausted_io_error)?; + if flush_allowed && (self.emitted_rows == 1 || self.output.len() >= JSON_DOCUMENT_OUTPUT_BATCH_BYTES) { + self.flush_output()?; + } + Ok(()) + } + + fn flush_output(&mut self) -> io::Result<()> { + if self.output.is_empty() { + return Ok(()); + } + (self.emit)(Bytes::from(std::mem::take(&mut self.output))) + } + + fn reset_record_buffer(&mut self) -> io::Result<()> { + if self.record.capacity() > SELECT_DECODE_CHUNK_BYTES.saturating_mul(2) { + self.record = Vec::new(); + self.reservation.free(); + } else { + self.record.clear(); + resize_json_document_record_reservation(self.reservation, self.record.capacity(), self.source)?; + } + Ok(()) + } + + fn reserve_record_byte(&mut self) -> io::Result<()> { + let required = self.record.len().checked_add(1).ok_or_else(resource_exhausted_io_error)?; + if required > MAX_SELECT_RECORD_BYTES { + return Err(over_max_record_size_error()); + } + if required <= self.record.capacity() { + return Ok(()); + } + let target_capacity = required + .checked_next_power_of_two() + .ok_or_else(resource_exhausted_io_error)? + .min(MAX_SELECT_RECORD_BYTES); + if target_capacity < required { + return Err(resource_exhausted_io_error()); + } + resize_json_document_record_reservation(self.reservation, target_capacity, self.source)?; + self.record + .try_reserve_exact(target_capacity.saturating_sub(self.record.len())) + .map_err(|_| resource_exhausted_io_error())?; + resize_json_document_record_reservation(self.reservation, self.record.capacity(), self.source) + } +} + +fn deserialize_json_value(bytes: &[u8], cancellation: &AtomicBool, eof: bool) -> io::Result { + let bytes = trim_json_whitespace(bytes); + if bytes.len() <= JSON_CANCELLATION_CHECK_BYTES { + return serde_json::from_slice(bytes).map_err(|error| classify_serde_json_error(error, cancellation, eof)); + } + let reader = io::BufReader::with_capacity( + JSON_CANCELLATION_CHECK_BYTES, + CancellableSliceReader { + inner: io::Cursor::new(bytes), + cancellation, + }, + ); + serde_json::from_reader(reader).map_err(|error| classify_serde_json_error(error, cancellation, eof)) +} + +fn validate_json_value(bytes: &[u8], cancellation: &AtomicBool, eof: bool) -> io::Result<()> { + ensure_json_parse_active(cancellation)?; + let bytes = trim_json_whitespace(bytes); + if bytes.len() <= JSON_CANCELLATION_CHECK_BYTES { + serde_json::from_slice::<&serde_json::value::RawValue>(bytes) + .map(drop) + .map_err(|error| classify_serde_json_error(error, cancellation, eof))?; + } else { + let reader = io::BufReader::with_capacity( + JSON_CANCELLATION_CHECK_BYTES, + CancellableSliceReader { + inner: io::Cursor::new(bytes), + cancellation, + }, + ); + serde_json::from_reader::<_, Box>(reader) + .map(drop) + .map_err(|error| classify_serde_json_error(error, cancellation, eof))?; + } + ensure_json_parse_active(cancellation) +} + +fn classify_serde_json_error(error: serde_json::Error, cancellation: &AtomicBool, eof: bool) -> io::Error { + match ensure_json_parse_active(cancellation) { + Err(cancelled) => cancelled, + Ok(()) => json_document_parse_error(if eof && error.is_eof() { + JsonDocumentParseError::Truncated + } else { + JsonDocumentParseError::Malformed + }), + } +} + +struct CancellableSliceReader<'a> { + inner: io::Cursor<&'a [u8]>, + cancellation: &'a AtomicBool, +} + +impl io::Read for CancellableSliceReader<'_> { + fn read(&mut self, buffer: &mut [u8]) -> io::Result { + ensure_json_parse_active(self.cancellation)?; + io::Read::read(&mut self.inner, buffer) + } +} + +fn root_array_plan(path: &[JsonPathSegment]) -> io::Result { + if let Some(rest) = path.strip_prefix(&[JsonPathSegment::ArrayWildcard]) { + return Ok(match rest.first() { + Some(JsonPathSegment::ArrayWildcard) => RootArrayPlan::Each { path_start: 2 }, + Some(JsonPathSegment::Index(index)) => RootArrayPlan::Index { + index: *index, + path_start: 2, + }, + Some(JsonPathSegment::Key { .. } | JsonPathSegment::ObjectWildcard) | None => RootArrayPlan::Each { path_start: 1 }, + }); + } + Ok(match path.first() { + None => RootArrayPlan::Each { path_start: 0 }, + Some(JsonPathSegment::Index(index)) => RootArrayPlan::Index { + index: *index, + path_start: 1, + }, + Some(JsonPathSegment::Key { .. } | JsonPathSegment::ArrayWildcard | JsonPathSegment::ObjectWildcard) => { + return Err(invalid_json_source_path("JSON source path segment does not match the input value")); + } + }) +} + +fn visit_json_path( + root: serde_json::Value, + path: &[JsonPathSegment], + cancellation: &AtomicBool, + emit: &mut impl FnMut(serde_json::Value) -> io::Result<()>, +) -> io::Result<()> { + if path.is_empty() { + ensure_json_parse_active(cancellation)?; + return emit(root); + } + let mut values = vec![root]; + for segment in path { + ensure_json_parse_active(cancellation)?; + let mut expanded = Vec::new(); + for value in values { + ensure_json_parse_active(cancellation)?; + 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, cancellation)? { + 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(array)) => expanded.extend(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; + if values.is_empty() { + return Ok(()); + } + } + for value in values { + ensure_json_parse_active(cancellation)?; + emit(value)?; + } + Ok(()) +} + +fn remove_json_source_key( + object: &mut serde_json::Map, + name: &str, + quoted: bool, + cancellation: &AtomicBool, +) -> io::Result> { + let mut checkpoint = || ensure_json_parse_active(cancellation); + remove_json_source_key_with_checkpoint(object, name, quoted, &mut checkpoint) +} + +pub(super) fn remove_json_source_key_with_checkpoint( + object: &mut serde_json::Map, + name: &str, + quoted: bool, + checkpoint: &mut impl FnMut() -> io::Result<()>, +) -> io::Result> { + 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(io::Error::new(io::ErrorKind::InvalidData, SelectError::AmbiguousFieldName)); + } + matched = Some(key.clone()); + } + } + Ok(matched.and_then(|key| object.remove(&key))) +} + +pub(super) fn json_key_eq_ignore_ascii_case_with_checkpoint( + key: &str, + expected: &str, + checkpoint: &mut impl FnMut() -> io::Result<()>, +) -> io::Result { + 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) +} + +fn invalid_json_source_path(message: &'static str) -> io::Error { + json_document_parse_error(JsonDocumentParseError::SourcePathMismatch(message)) +} + +fn scalar_column(source: &JsonSource) -> &str { + source.scalar_column().unwrap_or_else(|| match source.path().last() { + Some(JsonPathSegment::Key { name, .. }) => name, + Some(JsonPathSegment::Index(_) | JsonPathSegment::ArrayWildcard | JsonPathSegment::ObjectWildcard) | None => "_1", + }) +} + +struct CancellableJsonWriter<'a> { + inner: &'a mut Vec, + record_start: usize, + cancellation: &'a AtomicBool, + bytes_since_check: usize, + limit_exceeded: bool, +} + +impl io::Write for CancellableJsonWriter<'_> { + fn write(&mut self, buffer: &[u8]) -> io::Result { + let Some(new_len) = self.inner.len().checked_add(buffer.len()) else { + self.limit_exceeded = true; + return Err(io::Error::new(io::ErrorKind::InvalidData, SelectError::OverMaxRecordSize)); + }; + if new_len.saturating_sub(self.record_start) > MAX_SELECT_RECORD_BYTES { + self.limit_exceeded = true; + return Err(io::Error::new(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.cancellation)?; + self.bytes_since_check %= JSON_CANCELLATION_CHECK_BYTES; + } + self.inner.extend_from_slice(buffer); + Ok(buffer.len()) + } + + fn flush(&mut self) -> io::Result<()> { + ensure_json_parse_active(self.cancellation)?; + self.bytes_since_check = 0; + Ok(()) + } +} + +fn is_json_whitespace(byte: u8) -> bool { + matches!(byte, b' ' | b'\t' | b'\r' | b'\n') +} + +fn trim_json_whitespace(bytes: &[u8]) -> &[u8] { + let start = bytes + .iter() + .position(|byte| !is_json_whitespace(*byte)) + .unwrap_or(bytes.len()); + let end = bytes + .iter() + .rposition(|byte| !is_json_whitespace(*byte)) + .map_or(start, |index| index + 1); + &bytes[start..end] +} + +fn record_has_value(bytes: &[u8]) -> bool { + bytes.iter().any(|byte| !is_json_whitespace(*byte)) +} + +fn is_json_object(bytes: &[u8]) -> bool { + trim_json_whitespace(bytes).first() == Some(&b'{') +} + +fn ensure_json_parse_active(cancellation: &AtomicBool) -> io::Result<()> { + if cancellation.load(Ordering::Acquire) { + Err(json_document_cancelled_error()) + } else { + Ok(()) + } +} + +fn json_document_channel_reservation_bytes() -> usize { + // Include queued chunks, a blocked send, and both endpoint read buffers. + let input_chunks = JSON_DOCUMENT_INPUT_CHANNEL_CAPACITY.saturating_add(3); + let input_bytes = input_chunks.saturating_mul(SELECT_DECODE_CHUNK_BYTES); + // Include queued output, the consumer's current chunk, and the producer's + // current or blocked chunk after Vec growth. + let output_rows = JSON_DOCUMENT_ROW_CHANNEL_CAPACITY.saturating_add(2); + let row_capacity = MAX_SELECT_RECORD_BYTES + .saturating_add(1) + .checked_next_power_of_two() + .unwrap_or(usize::MAX); + let row_bytes = output_rows.saturating_mul(row_capacity); + input_bytes.saturating_add(row_bytes) +} + +pub(super) fn json_document_memory_reservation_bytes( + record_capacity: usize, + json_source: &JsonSource, +) -> datafusion::common::Result { + json_document_record_reservation_bytes(record_capacity, json_source, JSON_DOCUMENT_RAW_MEMORY_RESERVATION_MULTIPLIER) +} + +fn json_document_dom_memory_reservation_bytes( + record_capacity: usize, + json_source: &JsonSource, +) -> datafusion::common::Result { + json_document_record_reservation_bytes(record_capacity, json_source, JSON_DOCUMENT_DOM_MEMORY_RESERVATION_MULTIPLIER) +} + +fn json_document_record_reservation_bytes( + record_capacity: usize, + json_source: &JsonSource, + multiplier: usize, +) -> datafusion::common::Result { + let record_bytes = record_capacity.checked_mul(multiplier).ok_or_else(|| { + datafusion::common::DataFusionError::ResourcesExhausted(format!( + "JSON DOCUMENT memory reservation overflow for {record_capacity} input bytes" + )) + })?; + let scalar_bytes = scalar_column(json_source) + .len() + .checked_mul(JSON_SCALAR_COLUMN_MEMORY_RESERVATION_MULTIPLIER) + .ok_or_else(|| { + datafusion::common::DataFusionError::ResourcesExhausted(format!( + "JSON DOCUMENT scalar column reservation overflow for {record_capacity} input bytes" + )) + })?; + record_bytes.checked_add(scalar_bytes).ok_or_else(|| { + datafusion::common::DataFusionError::ResourcesExhausted(format!( + "JSON DOCUMENT memory reservation overflow for {record_capacity} input bytes" + )) + }) +} + +fn resize_json_document_record_reservation( + reservation: &MemoryReservation, + record_capacity: usize, + json_source: &JsonSource, +) -> io::Result<()> { + let bytes = + json_document_memory_reservation_bytes(record_capacity, json_source).map_err(|_| resource_exhausted_io_error())?; + reservation.try_resize(bytes).map_err(|_| resource_exhausted_io_error()) +} + +fn resize_json_document_dom_reservation( + reservation: &MemoryReservation, + record_capacity: usize, + json_source: &JsonSource, +) -> io::Result<()> { + let bytes = + json_document_dom_memory_reservation_bytes(record_capacity, json_source).map_err(|_| resource_exhausted_io_error())?; + reservation.try_resize(bytes).map_err(|_| resource_exhausted_io_error()) +} + +pub(super) fn validate_json_document_size(original_size: u64) -> ObjectStoreResult<()> { + if original_size <= MAX_JSON_DOCUMENT_BYTES { + Ok(()) + } else { + Err(resource_exhausted_store_error(datafusion::common::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 resource_exhausted_io_error() -> io::Error { + io::Error::new(io::ErrorKind::OutOfMemory, SelectError::ResourceExhausted) +} + +fn over_max_record_size_error() -> io::Error { + io::Error::new(io::ErrorKind::InvalidData, SelectError::OverMaxRecordSize) +} + +fn json_document_cancelled_error() -> io::Error { + io::Error::new(io::ErrorKind::Interrupted, SelectError::Canceled) +} + +fn json_document_cancelled_store_error() -> ObjectStoreError { + ObjectStoreError::Generic { + store: "EcObjectStore", + source: Box::new(QueryError::Cancel), + } +} + +fn json_document_worker_store_error(worker: &'static str) -> ObjectStoreError { + ObjectStoreError::Generic { + store: "EcObjectStore", + source: Box::new(io::Error::other(format!("JSON DOCUMENT {worker} worker terminated unexpectedly"))), + } +} + +fn json_document_worker_spawn_store_error(worker: &'static str, source: io::Error) -> ObjectStoreError { + ObjectStoreError::Generic { + store: "EcObjectStore", + source: Box::new(io::Error::new( + source.kind(), + format!("failed to start JSON DOCUMENT {worker} worker: {source}"), + )), + } +} + +fn resource_exhausted_store_error(source: datafusion::common::DataFusionError) -> ObjectStoreError { + ObjectStoreError::Generic { + store: "EcObjectStore", + source: Box::new(source), + } +} + +#[derive(Debug)] +enum JsonDocumentParseError { + Malformed, + Truncated, + DepthLimit { limit: usize }, + SourcePathMismatch(&'static str), +} + +impl fmt::Display for JsonDocumentParseError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Malformed => formatter.write_str("malformed JSON document"), + Self::Truncated => formatter.write_str("truncated JSON document"), + Self::DepthLimit { limit } => write!(formatter, "JSON document exceeds the maximum nesting depth of {limit}"), + Self::SourcePathMismatch(message) => formatter.write_str(message), + } + } +} + +impl std::error::Error for JsonDocumentParseError {} + +fn json_document_parse_error(source: JsonDocumentParseError) -> io::Error { + let kind = if matches!(source, JsonDocumentParseError::Truncated) { + io::ErrorKind::UnexpectedEof + } else { + io::ErrorKind::InvalidData + }; + io::Error::new(kind, source) +} + +fn classify_json_document_error(error: io::Error) -> ObjectStoreError { + let source: Box = if let Some(select_error) = find_error_source::(&error) { + Box::new(select_error.clone()) + } else if find_error_source::(&error).is_some() { + Box::new(SelectError::JsonParsingError) + } else { + Box::new(error) + }; + ObjectStoreError::Generic { + store: "EcObjectStore", + source, + } +} + +#[cfg(test)] +pub(super) fn parse_json_document_to_lines(bytes: &[u8], json_source: &JsonSource) -> io::Result> { + let memory_pool: Arc = Arc::new(datafusion::execution::memory_pool::UnboundedMemoryPool::default()); + let reservation = MemoryConsumer::new("JSON document parser test").register(&memory_pool); + let mut rows = Vec::new(); + parse_json_document(io::Cursor::new(bytes), json_source, &reservation, &AtomicBool::new(false), |batch| { + split_test_rows(batch, &mut rows); + Ok(()) + })?; + Ok(rows) +} + +#[cfg(test)] +fn split_test_rows(mut batch: Bytes, rows: &mut Vec) { + while let Some(newline) = batch.iter().position(|byte| *byte == b'\n') { + rows.push(batch.split_to(newline + 1)); + } + assert!(batch.is_empty(), "parser batches must end on a row boundary"); +} + +#[cfg(test)] +pub(super) fn flatten_json_document_to_ndjson(bytes: &[u8], json_source_path: &[JsonPathSegment]) -> io::Result { + let rows = parse_json_document_to_lines(bytes, &JsonSource::from_path(json_source_path.to_vec()))?; + let total = rows.iter().map(Bytes::len).sum(); + let mut output = Vec::with_capacity(total); + for row in rows { + output.extend_from_slice(&row); + } + Ok(Bytes::from(output)) +} + +#[cfg(test)] +pub(super) fn escaped_json_object_with_size(size: usize) -> Vec { + const PREFIX: &[u8] = br#"{"value":""#; + const SUFFIX: &[u8] = br#""}"#; + const ESCAPED_A: &[u8] = br#"\u0061"#; + let payload = size + .checked_sub(PREFIX.len() + SUFFIX.len()) + .expect("fixture size must fit an object"); + let mut input = Vec::with_capacity(size); + input.extend_from_slice(PREFIX); + for _ in 0..payload / ESCAPED_A.len() { + input.extend_from_slice(ESCAPED_A); + } + input.extend(std::iter::repeat_n(b'a', payload % ESCAPED_A.len())); + input.extend_from_slice(SUFFIX); + assert_eq!(input.len(), size); + input +} + +#[cfg(test)] +mod tests { + use super::*; + use tokio::io::AsyncWriteExt as _; + + struct ByteAtATime(R); + + impl io::Read for ByteAtATime { + fn read(&mut self, buffer: &mut [u8]) -> io::Result { + if buffer.is_empty() { + return Ok(0); + } + self.0.read(&mut buffer[..1]) + } + } + + fn parse_reader(reader: impl io::Read) -> io::Result> { + let memory_pool: Arc = Arc::new(datafusion::execution::memory_pool::UnboundedMemoryPool::default()); + let reservation = MemoryConsumer::new("JSON document chunk test").register(&memory_pool); + let mut rows = Vec::new(); + parse_json_document(reader, &JsonSource::default(), &reservation, &AtomicBool::new(false), |batch| { + split_test_rows(batch, &mut rows); + Ok(()) + })?; + Ok(rows) + } + + #[test] + fn root_array_plan_preserves_existing_source_semantics() { + assert_eq!(root_array_plan(&[]).expect("default source"), RootArrayPlan::Each { path_start: 0 }); + assert_eq!( + root_array_plan(&[JsonPathSegment::ArrayWildcard]).expect("array wildcard source"), + RootArrayPlan::Each { path_start: 1 } + ); + assert_eq!( + root_array_plan(&[JsonPathSegment::ArrayWildcard, JsonPathSegment::ArrayWildcard]) + .expect("double array wildcard source"), + RootArrayPlan::Each { path_start: 2 } + ); + assert_eq!( + root_array_plan(&[JsonPathSegment::Index(2)]).expect("array index source"), + RootArrayPlan::Index { index: 2, path_start: 1 } + ); + assert!( + root_array_plan(&[JsonPathSegment::Key { + name: "x".to_string(), + quoted: false, + }]) + .is_err() + ); + } + + #[tokio::test] + async fn first_array_row_is_available_before_input_finishes() { + let (mut writer, reader) = tokio::io::duplex(128); + let stream = json_document_ndjson_stream( + Box::new(reader), + 19, + JsonSource::default(), + Arc::new(SelectInputMetrics::default()), + Arc::new(datafusion::execution::memory_pool::GreedyMemoryPool::new(8 * 1024 * 1024)), + None, + ) + .expect("build lazy JSON document stream"); + let write_task = tokio::spawn(async move { + writer.write_all(b"[{\"id\":1},").await.expect("write first row"); + tokio::time::sleep(std::time::Duration::from_secs(30)).await; + }); + tokio::pin!(stream); + + let first = tokio::time::timeout(std::time::Duration::from_secs(1), stream.try_next()) + .await + .expect("first row should not wait for the document tail") + .expect("stream should remain valid") + .expect("first row should be present"); + assert_eq!(first, Bytes::from_static(b"{\"id\":1}\n")); + + write_task.abort(); + } + + #[tokio::test] + async fn root_object_is_available_when_its_value_finishes() { + let (mut writer, reader) = tokio::io::duplex(128); + let stream = json_document_ndjson_stream( + Box::new(reader), + 16, + JsonSource::default(), + Arc::new(SelectInputMetrics::default()), + Arc::new(datafusion::execution::memory_pool::GreedyMemoryPool::new(8 * 1024 * 1024)), + None, + ) + .expect("build lazy JSON document stream"); + let write_task = tokio::spawn(async move { + writer.write_all(b"{\"id\":1}").await.expect("write root object"); + futures::future::pending::<()>().await; + }); + tokio::pin!(stream); + + let first = tokio::time::timeout(std::time::Duration::from_secs(1), stream.try_next()) + .await + .expect("root object should not wait for trailing input") + .expect("stream should remain valid") + .expect("root object should be present"); + assert_eq!(first, Bytes::from_static(b"{\"id\":1}\n")); + + write_task.abort(); + } + + #[tokio::test] + async fn malformed_tail_is_reported_after_completed_rows() { + let input = br#"[{"id":1},{"id":"#.to_vec(); + let input_len = u64::try_from(input.len()).expect("fixture size should fit in u64"); + let mut stream = json_document_ndjson_stream( + Box::new(io::Cursor::new(input)), + input_len, + JsonSource::default(), + Arc::new(SelectInputMetrics::default()), + Arc::new(datafusion::execution::memory_pool::GreedyMemoryPool::new(8 * 1024 * 1024)), + None, + ) + .expect("build lazy JSON document stream"); + + assert_eq!( + stream.try_next().await.expect("read completed prefix row"), + Some(Bytes::from_static(b"{\"id\":1}\n")) + ); + let error = stream + .try_next() + .await + .expect_err("malformed tail must terminate the stream with an error"); + assert_eq!( + QueryError::from(datafusion::common::DataFusionError::ObjectStore(Box::new(error))).select_error(), + SelectError::JsonParsingError + ); + assert!(stream.try_next().await.expect("error must terminate the stream").is_none()); + } + + #[test] + fn chunk_boundaries_preserve_escapes_and_utf8() { + let input = "[\"\\\"\u{1f980}\",{\"name\":\"中\"}]".as_bytes(); + let rows = parse_reader(ByteAtATime(io::Cursor::new(input))).expect("parse split-sensitive JSON"); + assert_eq!(rows.len(), 2); + assert_eq!( + serde_json::from_slice::(&rows[0]).expect("first row"), + serde_json::json!({"_1": "\"🦀"}) + ); + assert_eq!( + serde_json::from_slice::(&rows[1]).expect("second row"), + serde_json::json!({"name": "中"}) + ); + } + + #[test] + fn multiline_objects_are_compacted_to_one_row() { + let root = parse_json_document_to_lines(b"{\n \"id\": 1,\n \"text\": \"a b\"\n}", &JsonSource::default()) + .expect("parse multiline root object"); + assert_eq!(root, vec![Bytes::from_static(b"{\"id\":1,\"text\":\"a b\"}\n")]); + + let array = parse_json_document_to_lines(b"[{\n \"id\": 1\n}, {\n \"id\": 2\n}]", &JsonSource::default()) + .expect("parse multiline array objects"); + assert_eq!(array, vec![Bytes::from_static(b"{\"id\":1}\n"), Bytes::from_static(b"{\"id\":2}\n")]); + + let nested_array = + parse_json_document_to_lines(b"[[\n 1, 2\n]]", &JsonSource::default()).expect("parse multiline nested array"); + assert_eq!(nested_array, vec![Bytes::from_static(b"{\"_1\":[1,2]}\n")]); + } + + #[test] + fn scalar_root_yields_one_record() { + let rows = parse_json_document_to_lines(b" 42 ", &JsonSource::default()).expect("parse scalar root"); + assert_eq!(rows, vec![Bytes::from_static(b"{\"_1\":42}\n")]); + } + + #[test] + fn raw_scalar_output_limit_has_an_exact_boundary() { + const SCALAR_WRAPPER_BYTES: usize = br#"{"":0}"#.len(); + let alias = "a".repeat(MAX_SELECT_RECORD_BYTES - SCALAR_WRAPPER_BYTES); + let source = JsonSource::new(Vec::new(), Some(alias)); + let rows = parse_json_document_to_lines(b"0", &source).expect("maximum scalar output record should parse"); + assert_eq!(rows[0].len(), MAX_SELECT_RECORD_BYTES + 1); + + let alias = "a".repeat(MAX_SELECT_RECORD_BYTES + 1 - SCALAR_WRAPPER_BYTES); + let source = JsonSource::new(Vec::new(), Some(alias)); + let error = + parse_json_document_to_lines(b"0", &source).expect_err("scalar output above one megabyte must fail while streaming"); + assert!(find_error_source::(&error).is_some_and(|error| error == &SelectError::OverMaxRecordSize)); + } + + #[test] + fn depth_limit_is_distinct_before_protocol_mapping() { + let mut at_limit = vec![b'['; JSON_DOCUMENT_MAX_DEPTH]; + at_limit.extend(std::iter::repeat_n(b']', JSON_DOCUMENT_MAX_DEPTH)); + parse_json_document_to_lines(&at_limit, &JsonSource::default()) + .expect("document at the nesting-depth limit should parse"); + + let mut input = vec![b'['; JSON_DOCUMENT_MAX_DEPTH + 1]; + input.extend(std::iter::repeat_n(b']', JSON_DOCUMENT_MAX_DEPTH + 1)); + let error = parse_json_document_to_lines(&input, &JsonSource::default()).expect_err("depth limit should fail"); + assert!(find_error_source::(&error).is_some_and(|error| { + matches!( + error, + JsonDocumentParseError::DepthLimit { + limit: JSON_DOCUMENT_MAX_DEPTH + } + ) + })); + } + + #[test] + fn source_path_length_is_not_a_document_depth_limit() { + let path = (0..JSON_DOCUMENT_MAX_DEPTH * 100) + .map(|index| JsonPathSegment::Key { + name: format!("missing_{index}"), + quoted: false, + }) + .collect(); + let rows = parse_json_document_to_lines(b"{}", &JsonSource::new(path, None)) + .expect("a shallow document must not inherit the source-path depth"); + assert!(rows.is_empty()); + } + + #[test] + fn json_whitespace_is_limited_to_the_rfc_tokens() { + for input in [b"\x0b{}".as_slice(), b"{}\x0c".as_slice(), b"[\x0b]".as_slice()] { + let error = parse_json_document_to_lines(input, &JsonSource::default()) + .expect_err("vertical tab and form feed are not JSON whitespace"); + assert!(matches!( + find_error_source::(&error), + Some(JsonDocumentParseError::Malformed) + )); + } + } + + #[test] + fn root_array_separator_whitespace_is_not_a_record() { + let mut empty = Vec::with_capacity(MAX_SELECT_RECORD_BYTES + 3); + empty.push(b'['); + empty.extend(std::iter::repeat_n(b' ', MAX_SELECT_RECORD_BYTES + 1)); + empty.push(b']'); + assert!( + parse_json_document_to_lines(&empty, &JsonSource::default()) + .expect("parse empty array") + .is_empty() + ); + + let mut input = Vec::with_capacity(MAX_SELECT_RECORD_BYTES + 4); + input.push(b'['); + input.extend(escaped_json_object_with_size(MAX_SELECT_RECORD_BYTES)); + input.extend_from_slice(b" ]"); + assert_eq!( + parse_json_document_to_lines(&input, &JsonSource::default()) + .expect("separator whitespace must not increase the record size") + .len(), + 1 + ); + + let error = parse_json_document_to_lines(b"[1 2]", &JsonSource::default()) + .expect_err("separator whitespace must not merge adjacent values"); + assert!(matches!( + find_error_source::(&error), + Some(JsonDocumentParseError::Malformed) + )); + } + + #[test] + fn malformed_open_record_is_not_classified_as_truncated() { + let error = parse_json_document_to_lines(b"[!", &JsonSource::default()) + .expect_err("an invalid token is malformed, even when the root array is unfinished"); + assert!(matches!( + find_error_source::(&error), + Some(JsonDocumentParseError::Malformed) + )); + } + + #[test] + fn raw_input_record_limit_has_exact_boundary() { + let at_limit = escaped_json_object_with_size(MAX_SELECT_RECORD_BYTES); + parse_json_document_to_lines(&at_limit, &JsonSource::default()).expect("one-megabyte input record should be accepted"); + + let over_limit = escaped_json_object_with_size(MAX_SELECT_RECORD_BYTES + 1); + let error = parse_json_document_to_lines(&over_limit, &JsonSource::default()) + .expect_err("input record above one megabyte must fail before deserialization"); + assert!(find_error_source::(&error).is_some_and(|error| error == &SelectError::OverMaxRecordSize)); + } + + #[test] + fn truncated_input_is_distinct_before_protocol_mapping() { + let error = + parse_json_document_to_lines(br#"[{"id":1}"#, &JsonSource::default()).expect_err("truncated document should fail"); + assert!(matches!( + find_error_source::(&error), + Some(JsonDocumentParseError::Truncated) + )); + } + + #[tokio::test] + async fn large_root_array_uses_bounded_query_memory() { + const RECORDS: usize = 20_000; + let mut input = Vec::with_capacity(RECORDS * 9 + 1); + input.push(b'['); + for index in 0..RECORDS { + if index > 0 { + input.push(b','); + } + input.extend_from_slice(br#"{"id":0}"#); + } + input.push(b']'); + let input_len = u64::try_from(input.len()).expect("fixture size should fit in u64"); + + let memory_pool = Arc::new(datafusion::execution::memory_pool::GreedyMemoryPool::new(8 * 1024 * 1024)); + let mut stream = json_document_ndjson_stream( + Box::new(io::Cursor::new(input)), + input_len, + JsonSource::default(), + Arc::new(SelectInputMetrics::default()), + memory_pool.clone(), + None, + ) + .expect("build lazy JSON document stream"); + let mut rows = 0; + let mut batches = 0; + while let Some(batch) = stream.try_next().await.expect("stream large root array") { + rows += batch.iter().filter(|byte| **byte == b'\n').count(); + batches += 1; + } + + assert_eq!(rows, RECORDS); + assert!(batches < 100, "small records should cross the worker boundary in bounded batches"); + assert_eq!(memory_pool.reserved(), 0); + } + + #[tokio::test] + async fn legal_large_record_fits_the_default_query_memory_limit() { + let value = "x".repeat(600 * 1024); + let input = format!(r#"{{"value":"{value}"}}"#).into_bytes(); + let memory_pool = Arc::new(datafusion::execution::memory_pool::GreedyMemoryPool::new(64 * 1024 * 1024)); + let output = json_document_ndjson_stream( + Box::new(io::Cursor::new(input.clone())), + u64::try_from(input.len()).expect("fixture size should fit in u64"), + JsonSource::default(), + Arc::new(SelectInputMetrics::default()), + memory_pool.clone(), + None, + ) + .expect("build large-record JSON document stream") + .try_collect::>() + .await + .expect("a legal record below one megabyte should fit the default pool") + .concat(); + + assert_eq!(output.len(), input.len() + 1); + assert_eq!(memory_pool.reserved(), 0); + } + + #[test] + fn non_expanding_path_releases_dom_memory_before_blocking_emit() { + let input = format!("[{{\"payload\":\"{}\"}}]", "x".repeat(512 * 1024)); + let source = JsonSource::from_path(vec![ + JsonPathSegment::ArrayWildcard, + JsonPathSegment::Key { + name: "payload".to_string(), + quoted: false, + }, + ]); + let memory_pool = Arc::new(datafusion::execution::memory_pool::GreedyMemoryPool::new(64 * 1024 * 1024)); + let query_memory_pool: Arc = memory_pool.clone(); + let reservation = MemoryConsumer::new("JSON DOM emit overlap test").register(&query_memory_pool); + let mut observed_reservation = None; + + parse_json_document(io::Cursor::new(input), &source, &reservation, &AtomicBool::new(false), |batch| { + observed_reservation = Some(memory_pool.reserved()); + assert!(batch.ends_with(b"\n")); + Ok(()) + }) + .expect("non-expanding source path should stream"); + + assert!( + observed_reservation.is_some_and(|bytes| bytes < 8 * 1024 * 1024), + "DOM reservation must shrink before a potentially blocking output send" + ); + assert_eq!(memory_pool.reserved(), 0); + } + + #[tokio::test] + async fn parser_worker_failure_is_a_terminal_internal_error() { + let (row_tx, row_rx) = mpsc::channel(1); + drop(row_tx); + let (parse_result_tx, parse_result_rx) = oneshot::channel(); + drop(parse_result_tx); + let memory_pool: Arc = Arc::new(datafusion::execution::memory_pool::UnboundedMemoryPool::default()); + let reservation = MemoryConsumer::new("JSON worker failure test").register(&memory_pool); + let resources = Arc::new(JsonDocumentPipelineResources { + _channel_reservation: reservation, + _query_guard: None, + }); + let mut pipeline = JsonDocumentPipeline { + receiver: row_rx, + read_task: Some(SpawnedTask::spawn(async {})), + parse_result: parse_result_rx, + done: false, + cancellation: Arc::new(AtomicBool::new(false)), + _resources: resources, + }; + + let error = pipeline + .try_next() + .await + .expect_err("a failed parser task must not become successful EOF"); + assert_eq!( + QueryError::from(datafusion::common::DataFusionError::ObjectStore(Box::new(error))).select_error(), + SelectError::InternalError + ); + assert!( + pipeline + .try_next() + .await + .expect("terminal error should close the stream") + .is_none() + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn trackerless_parser_waits_for_global_admission() { + let permits = tokio::time::timeout( + std::time::Duration::from_secs(5), + JSON_DOCUMENT_PARSER_ADMISSION + .acquire_many(u32::try_from(*JSON_DOCUMENT_PARSER_THREADS).expect("parser thread limit should fit in u32")), + ) + .await + .expect("other parser tests should release global admission") + .expect("global parser admission should remain open"); + let memory_pool = Arc::new(datafusion::execution::memory_pool::GreedyMemoryPool::new(8 * 1024 * 1024)); + let mut stream = json_document_ndjson_stream( + Box::new(io::Cursor::new(b"{}")), + 2, + JsonSource::default(), + Arc::new(SelectInputMetrics::default()), + memory_pool.clone(), + None, + ) + .expect("build trackerless JSON document stream"); + + { + let next = stream.try_next(); + futures::pin_mut!(next); + assert!(futures::poll!(next.as_mut()).is_pending()); + } + assert_eq!(memory_pool.reserved(), 0, "waiting streams must not reserve parser memory"); + + drop(permits); + assert_eq!( + tokio::time::timeout(std::time::Duration::from_secs(1), stream.try_next()) + .await + .expect("parser should start after admission is released") + .expect("trackerless parser should remain valid"), + Some(Bytes::from_static(b"{}\n")) + ); + drop(stream); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while memory_pool.reserved() != 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("trackerless parser should release query memory"); + } + + #[test] + fn cancelled_row_build_keeps_only_completed_output() { + let memory_pool: Arc = Arc::new(datafusion::execution::memory_pool::UnboundedMemoryPool::default()); + let reservation = MemoryConsumer::new("JSON cancellation output test").register(&memory_pool); + let cancellation = AtomicBool::new(true); + let source = JsonSource::default(); + let mut decoder = JsonDocumentDecoder::new(&source, &reservation, &cancellation, |_batch| Ok(())); + decoder.output.extend_from_slice(b"{\"done\":true}\n"); + let completed = decoder.output.clone(); + + decoder.record.extend_from_slice(b"0"); + let error = decoder + .append_raw_value(false) + .expect_err("raw scalar construction must observe cancellation"); + assert_eq!(error.kind(), io::ErrorKind::Interrupted); + assert_eq!(decoder.output, completed); + + let error = decoder + .append_json_value(serde_json::json!(1), false) + .expect_err("DOM scalar construction must observe cancellation"); + assert_eq!(error.kind(), io::ErrorKind::Interrupted); + assert_eq!(decoder.output, completed); + } + + #[tokio::test] + async fn scalar_alias_expansion_is_reserved_per_streamed_row() { + const EXPECTED_SCALAR_ALIAS_MULTIPLIER: usize = 2; + assert_eq!(JSON_SCALAR_COLUMN_MEMORY_RESERVATION_MULTIPLIER, EXPECTED_SCALAR_ALIAS_MULTIPLIER); + let alias = "alias".repeat(40_000); + let source = JsonSource::new(vec![JsonPathSegment::ArrayWildcard], Some(alias.clone())); + let alias_reservation = alias + .len() + .checked_mul(EXPECTED_SCALAR_ALIAS_MULTIPLIER) + .expect("alias reservation should fit"); + let insufficient = json_document_channel_reservation_bytes() + alias_reservation - 1; + let mut stream = json_document_ndjson_stream( + Box::new(io::Cursor::new(b"[0]".to_vec())), + 3, + source.clone(), + Arc::new(SelectInputMetrics::default()), + Arc::new(datafusion::execution::memory_pool::GreedyMemoryPool::new(insufficient)), + None, + ) + .expect("build constrained JSON document stream"); + let error = stream + .try_next() + .await + .expect_err("unreserved scalar alias expansion must fail"); + assert_eq!( + QueryError::from(datafusion::common::DataFusionError::ObjectStore(Box::new(error))).select_error(), + SelectError::ResourceExhausted + ); + + let memory_pool = Arc::new(datafusion::execution::memory_pool::GreedyMemoryPool::new(insufficient + 1024 * 1024)); + let rows = json_document_ndjson_stream( + Box::new(io::Cursor::new(b"[0]".to_vec())), + 3, + source, + Arc::new(SelectInputMetrics::default()), + memory_pool.clone(), + None, + ) + .expect("build JSON document stream") + .try_collect::>() + .await + .expect("reserved scalar alias expansion should stream"); + assert_eq!(rows.len(), 1); + assert_eq!(memory_pool.reserved(), 0); + } + + #[tokio::test] + async fn dropping_active_stream_releases_parser_memory() { + let (mut writer, reader) = tokio::io::duplex(128); + let memory_pool = Arc::new(datafusion::execution::memory_pool::GreedyMemoryPool::new(8 * 1024 * 1024)); + let mut stream = json_document_ndjson_stream( + Box::new(reader), + 19, + JsonSource::default(), + Arc::new(SelectInputMetrics::default()), + memory_pool.clone(), + None, + ) + .expect("build active JSON document stream"); + let writer_task = tokio::spawn(async move { + writer.write_all(b"[{\"id\":1},").await.expect("write first row"); + futures::future::pending::<()>().await; + }); + + assert!(stream.try_next().await.expect("read first row").is_some()); + drop(stream); + writer_task.abort(); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while memory_pool.reserved() != 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("active parser cancellation should release memory"); + } +} diff --git a/crates/s3select-query/Cargo.toml b/crates/s3select-query/Cargo.toml index dfec9520b..f40152682 100644 --- a/crates/s3select-query/Cargo.toml +++ b/crates/s3select-query/Cargo.toml @@ -51,6 +51,8 @@ derive_builder = { workspace = true } futures = { workspace = true } parking_lot = { workspace = true } s3s = { workspace = true, features = ["minio"] } +serde.workspace = true +serde_json.workspace = true tokio = { workspace = true, features = ["fs", "rt-multi-thread", "sync", "time"] } tracing = { workspace = true } diff --git a/crates/s3select-query/src/dispatcher/json_document_schema.rs b/crates/s3select-query/src/dispatcher/json_document_schema.rs new file mode 100644 index 000000000..d61d9ff0b --- /dev/null +++ b/crates/s3select-query/src/dispatcher/json_document_schema.rs @@ -0,0 +1,494 @@ +// 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 datafusion::{ + arrow::{ + datatypes::{DataType, Field, Schema}, + error::ArrowError, + }, + common::{DataFusionError, Result as DFResult}, +}; +use serde::de::{self, DeserializeSeed, MapAccess, SeqAccess, Visitor}; +use std::{ + collections::HashMap, + fmt, + io::{self, Cursor}, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, +}; + +const CANCELLATION_CHECK_ITEMS: usize = 1024; +const READER_BUFFER_BYTES: usize = 64 * 1024; + +pub(super) fn infer_schema( + bytes: &[u8], + expected_records: usize, + cancellation: &AtomicBool, + max_schema_bytes: usize, +) -> DFResult<(Schema, usize, usize)> { + let mut context = InferenceContext { + cancellation, + items_since_check: 0, + }; + let mut fields = InferredObject::default(); + let mut records = 0; + + for row in bytes.split_inclusive(|byte| *byte == b'\n') { + if row.is_empty() { + continue; + } + context.check_now()?; + let reader = io::BufReader::with_capacity( + READER_BUFFER_BYTES, + CancellableReader { + inner: Cursor::new(row), + cancellation, + }, + ); + let mut deserializer = serde_json::Deserializer::from_reader(reader); + let inferred = InferSeed { context: &mut context } + .deserialize(&mut deserializer) + .map_err(|error| schema_parse_error(error, cancellation))?; + deserializer.end().map_err(|error| schema_parse_error(error, cancellation))?; + let InferredType::Object(record) = inferred else { + return Err(DataFusionError::Execution( + "JSON DOCUMENT schema inference expected an object row".to_string(), + )); + }; + if record.retained_size() > max_schema_bytes { + return Err(schema_complexity_error()); + } + fields.merge(record, &mut context)?; + if fields.retained_size() > max_schema_bytes { + return Err(schema_complexity_error()); + } + records += 1; + } + + if records != expected_records { + return Err(DataFusionError::Execution(format!( + "JSON DOCUMENT schema prefix contained {expected_records} records but inference read {records}" + ))); + } + let schema = Schema::new(fields.into_fields(&mut context)?); + let schema_bytes = std::mem::size_of::() + .checked_add(schema.fields().size()) + .ok_or_else(schema_complexity_error)?; + if schema_bytes > max_schema_bytes { + return Err(schema_complexity_error()); + } + Ok((schema, records, schema_bytes)) +} + +fn schema_parse_error(error: serde_json::Error, cancellation: &AtomicBool) -> DataFusionError { + if cancellation.load(Ordering::Acquire) { + DataFusionError::Execution("JSON DOCUMENT schema inference canceled".to_string()) + } else { + schema_json_error(format!("JSON DOCUMENT schema inference failed: {error}")) + } +} + +fn schema_json_error(message: impl Into) -> DataFusionError { + DataFusionError::ArrowError(Box::new(ArrowError::JsonError(message.into())), None) +} + +fn schema_complexity_error() -> DataFusionError { + DataFusionError::ResourcesExhausted("JSON DOCUMENT schema exceeds the inference complexity limit".to_string()) +} + +struct CancellableReader<'a> { + inner: Cursor<&'a [u8]>, + cancellation: &'a AtomicBool, +} + +impl io::Read for CancellableReader<'_> { + fn read(&mut self, buffer: &mut [u8]) -> io::Result { + if self.cancellation.load(Ordering::Acquire) { + return Err(io::Error::new(io::ErrorKind::Interrupted, "JSON DOCUMENT schema inference canceled")); + } + io::Read::read(&mut self.inner, buffer) + } +} + +struct InferenceContext<'a> { + cancellation: &'a AtomicBool, + items_since_check: usize, +} + +impl InferenceContext<'_> { + fn checkpoint(&mut self) -> Result<(), E> { + self.items_since_check += 1; + if self.items_since_check < CANCELLATION_CHECK_ITEMS { + return Ok(()); + } + self.items_since_check = 0; + if self.cancellation.load(Ordering::Acquire) { + Err(E::custom("JSON DOCUMENT schema inference canceled")) + } else { + Ok(()) + } + } + + fn check_now(&self) -> DFResult<()> { + if self.cancellation.load(Ordering::Acquire) { + Err(DataFusionError::Execution("JSON DOCUMENT schema inference canceled".to_string())) + } else { + Ok(()) + } + } + + fn checkpoint_df(&mut self) -> DFResult<()> { + self.items_since_check += 1; + if self.items_since_check < CANCELLATION_CHECK_ITEMS { + return Ok(()); + } + self.items_since_check = 0; + self.check_now() + } +} + +struct InferSeed<'a, 'b> { + context: &'a mut InferenceContext<'b>, +} + +impl<'de> DeserializeSeed<'de> for InferSeed<'_, '_> { + type Value = InferredType; + + fn deserialize(self, deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + deserializer.deserialize_any(InferVisitor { context: self.context }) + } +} + +struct InferVisitor<'a, 'b> { + context: &'a mut InferenceContext<'b>, +} + +impl<'de> Visitor<'de> for InferVisitor<'_, '_> { + type Value = InferredType; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("a JSON value") + } + + fn visit_unit(self) -> Result { + Ok(InferredType::Any) + } + + fn visit_none(self) -> Result { + Ok(InferredType::Any) + } + + fn visit_bool(self, _value: bool) -> Result { + Ok(InferredType::Scalar(ScalarKinds::BOOLEAN)) + } + + fn visit_i64(self, _value: i64) -> Result { + Ok(InferredType::Scalar(ScalarKinds::INTEGER)) + } + + fn visit_u64(self, value: u64) -> Result { + Ok(InferredType::Scalar(if i64::try_from(value).is_ok() { + ScalarKinds::INTEGER + } else { + ScalarKinds::FLOAT + })) + } + + fn visit_f64(self, _value: f64) -> Result { + Ok(InferredType::Scalar(ScalarKinds::FLOAT)) + } + + fn visit_str(self, _value: &str) -> Result { + Ok(InferredType::Scalar(ScalarKinds::STRING)) + } + + fn visit_string(self, _value: String) -> Result { + Ok(InferredType::Scalar(ScalarKinds::STRING)) + } + + fn visit_seq(self, mut sequence: A) -> Result + where + A: SeqAccess<'de>, + { + let mut inner = InferredType::Any; + let mut first = true; + while let Some(value) = sequence.next_element_seed(InferSeed { context: self.context })? { + self.context.checkpoint()?; + if first && matches!(value, InferredType::Any) { + inner = InferredType::Scalar(ScalarKinds::default()); + } else { + inner.merge_array_element(value, self.context).map_err(de::Error::custom)?; + } + first = false; + } + Ok(InferredType::Array(Box::new(inner))) + } + + fn visit_map(self, mut map: A) -> Result + where + A: MapAccess<'de>, + { + let mut object = InferredObject::default(); + while let Some(key) = map.next_key::()? { + self.context.checkpoint()?; + let value = map.next_value_seed(InferSeed { context: self.context })?; + object.insert_record_field(key, value); + } + Ok(InferredType::Object(object)) + } +} + +#[derive(Clone, Copy, Debug, Default)] +struct ScalarKinds(u8); + +impl ScalarKinds { + const BOOLEAN: Self = Self(1 << 0); + const INTEGER: Self = Self(1 << 1); + const FLOAT: Self = Self(1 << 2); + const STRING: Self = Self(1 << 3); + + fn merge(&mut self, other: Self) { + self.0 |= other.0; + } + + fn data_type(self) -> DataType { + match self.0 { + value if value == Self::BOOLEAN.0 => DataType::Boolean, + value if value == Self::INTEGER.0 => DataType::Int64, + value if value == Self::FLOAT.0 || value == (Self::INTEGER.0 | Self::FLOAT.0) => DataType::Float64, + value if value == Self::STRING.0 => DataType::Utf8, + _ => DataType::Utf8, + } + } +} + +#[derive(Debug)] +enum InferredType { + Scalar(ScalarKinds), + Array(Box), + Object(InferredObject), + Any, +} + +impl InferredType { + fn merge(&mut self, other: Self, context: &mut InferenceContext<'_>) -> DFResult<()> { + match (self, other) { + (Self::Array(current), Self::Array(other)) => current.merge(*other, context), + (Self::Scalar(current), Self::Scalar(other)) => { + current.merge(other); + Ok(()) + } + (Self::Object(current), Self::Object(other)) => current.merge(other, context), + (current @ Self::Any, value) => { + *current = value; + Ok(()) + } + (_, Self::Any) => Ok(()), + (Self::Array(inner), scalar @ Self::Scalar(_)) => inner.merge(scalar, context), + (current @ Self::Scalar(_), Self::Array(mut inner)) => { + let scalar = std::mem::replace(current, Self::Any); + inner.merge(scalar, context)?; + *current = Self::Array(inner); + Ok(()) + } + (current, other) => Err(schema_json_error(format!( + "incompatible JSON types during schema inference: {} and {}", + current.kind(), + other.kind() + ))), + } + } + + fn merge_array_element(&mut self, other: Self, context: &mut InferenceContext<'_>) -> DFResult<()> { + match (self, other) { + (current @ Self::Any, value) => { + *current = value; + Ok(()) + } + (_, Self::Any) => Ok(()), + (Self::Scalar(current), Self::Scalar(other)) => { + current.merge(other); + Ok(()) + } + (Self::Object(current), Self::Object(other)) => current.merge(other, context), + (Self::Array(current), Self::Array(other)) => current.merge(*other, context), + (current, other) => Err(schema_json_error(format!( + "incompatible JSON array elements during schema inference: {} and {}", + current.kind(), + other.kind() + ))), + } + } + + fn kind(&self) -> &'static str { + match self { + Self::Scalar(_) => "scalar", + Self::Array(_) => "array", + Self::Object(_) => "object", + Self::Any => "null", + } + } + + fn into_data_type(self, context: &mut InferenceContext<'_>) -> DFResult { + context.check_now()?; + match self { + Self::Scalar(kinds) => Ok(kinds.data_type()), + Self::Array(inner) => Ok(DataType::List(Arc::new(Field::new_list_field(inner.into_data_type(context)?, true)))), + Self::Object(object) => Ok(DataType::Struct(object.into_fields(context)?.into())), + Self::Any => Ok(DataType::Null), + } + } + + fn retained_size(&self) -> usize { + std::mem::size_of::() + + match self { + Self::Array(inner) => inner.retained_size(), + Self::Object(object) => object.retained_size(), + Self::Scalar(_) | Self::Any => 0, + } + } +} + +#[derive(Debug, Default)] +struct InferredObject { + indexes: HashMap, usize>, + fields: Vec<(Arc, InferredType)>, +} + +impl InferredObject { + fn insert_record_field(&mut self, key: String, value: InferredType) { + if let Some(index) = self.indexes.get(key.as_str()).copied() { + self.fields[index].1 = value; + return; + } + let key: Arc = Arc::from(key); + let index = self.fields.len(); + self.indexes.insert(Arc::clone(&key), index); + self.fields.push((key, value)); + } + + fn merge(&mut self, other: Self, context: &mut InferenceContext<'_>) -> DFResult<()> { + for (key, value) in other.fields { + context.checkpoint_df()?; + if let Some(index) = self.indexes.get(key.as_ref()).copied() { + self.fields[index].1.merge(value, context)?; + } else { + let index = self.fields.len(); + self.indexes.insert(Arc::clone(&key), index); + self.fields.push((key, value)); + } + } + Ok(()) + } + + fn into_fields(self, context: &mut InferenceContext<'_>) -> DFResult> { + self.fields + .into_iter() + .map(|(name, inferred)| { + context.check_now()?; + Ok(Field::new(name.as_ref(), inferred.into_data_type(context)?, true)) + }) + .collect() + } + + fn retained_size(&self) -> usize { + let index_entry = std::mem::size_of::<(Arc, usize)>() + 16; + let mut bytes = std::mem::size_of::() + .saturating_add(self.indexes.capacity().saturating_mul(index_entry)) + .saturating_add( + self.fields + .capacity() + .saturating_mul(std::mem::size_of::<(Arc, InferredType)>()), + ); + for (name, inferred) in &self.fields { + bytes = bytes.saturating_add(name.len()).saturating_add(inferred.retained_size()); + } + bytes + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::json::reader::infer_json_schema as arrow_infer_json_schema; + + const TEST_SCHEMA_LIMIT: usize = 4 * 1024 * 1024; + + fn assert_matches_arrow(input: &[u8]) { + let records = input.iter().filter(|byte| **byte == b'\n').count(); + let (expected, expected_records) = arrow_infer_json_schema(io::BufReader::new(input), Some(records)) + .expect("Arrow should infer the compatibility fixture"); + let (actual, actual_records, _) = infer_schema(input, records, &AtomicBool::new(false), TEST_SCHEMA_LIMIT) + .expect("streaming inference should accept the compatibility fixture"); + assert_eq!(actual_records, expected_records); + assert_eq!(actual, expected); + } + + #[test] + fn inference_matches_arrow_for_nested_and_coerced_types() { + assert_matches_arrow( + br#"{"a":1,"values":[1,2],"nested":{"enabled":true},"nullable":null} +{"a":1.5,"values":3,"nested":{"name":"ok"},"nullable":"set"} +"#, + ); + assert_matches_arrow( + br#"{"matrix":[[1,2],[3]],"objects":[{"id":1},{"name":"two"}],"empty":[]} +{"matrix":[[4.5]],"objects":[],"empty":[null]} +"#, + ); + } + + #[test] + fn inference_matches_arrow_for_null_first_arrays_and_duplicate_keys() { + assert_matches_arrow( + br#"{"values":[null,1,2],"duplicate":1,"duplicate":"last"} +"#, + ); + + let input = br#"{"values":[null,{"id":1}]} +"#; + let records = 1; + assert!(arrow_infer_json_schema(io::BufReader::new(input.as_slice()), Some(records)).is_err()); + assert!(infer_schema(input, records, &AtomicBool::new(false), TEST_SCHEMA_LIMIT).is_err()); + } + + #[test] + fn inference_checkpoints_observe_cancellation() { + let cancellation = AtomicBool::new(true); + let mut context = InferenceContext { + cancellation: &cancellation, + items_since_check: CANCELLATION_CHECK_ITEMS - 1, + }; + let error = context + .checkpoint::() + .expect_err("the item checkpoint must observe cancellation"); + assert!(error.to_string().contains("schema inference canceled")); + } + + #[test] + fn incompatible_types_preserve_arrow_json_error_classification() { + let error = infer_schema(b"{\"a\":{}}\n{\"a\":1}\n", 2, &AtomicBool::new(false), TEST_SCHEMA_LIMIT) + .expect_err("object and scalar fields must remain incompatible"); + + assert!(matches!( + error, + DataFusionError::ArrowError(source, None) if matches!(*source, ArrowError::JsonError(_)) + )); + } +} diff --git a/crates/s3select-query/src/dispatcher/json_document_table.rs b/crates/s3select-query/src/dispatcher/json_document_table.rs new file mode 100644 index 000000000..767c5ddfb --- /dev/null +++ b/crates/s3select-query/src/dispatcher/json_document_table.rs @@ -0,0 +1,1284 @@ +// 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 async_trait::async_trait; +use bytes::{Buf as _, Bytes}; +use datafusion::{ + arrow::{datatypes::SchemaRef, json::ReaderBuilder, record_batch::RecordBatch}, + catalog::Session, + common::{DataFusionError, Result as DFResult, project_schema, runtime::SpawnedTask}, + datasource::TableProvider, + execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}, + logical_expr::{Expr, TableProviderFilterPushDown, TableType}, + object_store::{Error as ObjectStoreError, ObjectStoreExt as _, path::Path}, + physical_plan::{ + ExecutionPlan, RecordBatchStream, SendableRecordBatchStream, + stream::RecordBatchStreamAdapter, + streaming::{PartitionStream, StreamingTableExec}, + }, +}; +use futures::{Stream, StreamExt as _, stream}; +use parking_lot::Mutex; +use rustfs_s3select_api::QueryResult; +use std::{ + fmt, + pin::Pin, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, + task::{Context, Poll}, +}; + +use super::json_document_schema::infer_schema; + +const SCHEMA_INFERENCE_MAX_RECORDS: usize = 1000; +const SCHEMA_INFERENCE_MAX_BYTES: usize = 4 * 1024 * 1024; +const SCHEMA_INFERENCE_MAX_SCHEMA_BYTES: usize = 4 * 1024 * 1024; +const SCHEMA_INFERENCE_WORK_MULTIPLIER: usize = 16; +// Once the byte target is reached, finish only the current protocol-sized +// logical record before flushing. The multiplier covers Arrow's tape, offsets, +// structural positions, output buffers, and Vec growth for dense minified arrays. +const JSON_DECODE_BATCH_BYTES: usize = 1024 * 1024; +const JSON_DECODE_MEMORY_MULTIPLIER: usize = 32; +const JSON_DECODE_POLL_BYTES: usize = 64 * 1024; +const JSON_DECODER_TAPE_TARGET_BYTES: usize = 4 * 1024 * 1024; +const JSON_DECODER_PER_FIELD_METADATA_BYTES: usize = 256; + +pub(super) struct JsonDocumentTable { + schema: SchemaRef, + source: Arc>>, +} + +impl fmt::Debug for JsonDocumentTable { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("JsonDocumentTable") + .field("schema", &self.schema) + .finish_non_exhaustive() + } +} + +impl JsonDocumentTable { + pub(super) async fn try_new(state: &dyn Session, bucket: &str, key: &str) -> QueryResult> { + let object_store_url = datafusion::execution::object_store::ObjectStoreUrl::parse(format!("s3://{bucket}"))?; + let object_store = state.runtime_env().object_store(&object_store_url)?; + let result = object_store.get(&Path::from(key)).await.map_err(DataFusionError::from)?; + let memory_pool = Arc::clone(&state.runtime_env().memory_pool); + let (schema, source) = prepare_source(result.into_stream(), memory_pool).await?; + Ok(Arc::new(Self { + schema, + source: Arc::new(Mutex::new(Some(source))), + })) + } +} + +#[async_trait] +impl TableProvider for JsonDocumentTable { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + state: &dyn Session, + projection: Option<&Vec>, + filters: &[Expr], + limit: Option, + ) -> DFResult> { + let scan_limit = filters.is_empty().then_some(limit).flatten(); + let default_batch_size = usize::from(state.config_options().execution.batch_size); + let batch_size = scan_limit + .filter(|limit| *limit > 0) + .map_or(default_batch_size, |limit| limit.min(default_batch_size)); + let projected_schema = project_schema(&self.schema, projection)?; + let partition: Arc = Arc::new(JsonDocumentPartition { + schema: Arc::clone(&projected_schema), + source: Arc::clone(&self.source), + batch_size, + }); + let plan = StreamingTableExec::try_new(projected_schema, vec![partition], None, [], false, scan_limit)?; + Ok(Arc::new(plan)) + } + + fn supports_filters_pushdown(&self, filters: &[&Expr]) -> DFResult> { + Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()]) + } +} + +#[derive(Debug)] +struct JsonDocumentPartition { + schema: SchemaRef, + source: Arc>>, + batch_size: usize, +} + +impl PartitionStream for JsonDocumentPartition { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + fn execute(&self, _context: Arc) -> SendableRecordBatchStream { + let source = self.source.lock().take(); + let Some(source) = source else { + return error_stream( + Arc::clone(&self.schema), + DataFusionError::Execution("JSON DOCUMENT input stream was consumed more than once".to_string()), + ); + }; + match JsonDocumentBatchStream::try_new(Arc::clone(&self.schema), source, self.batch_size) { + Ok(stream) => Box::pin(stream), + Err(error) => error_stream(Arc::clone(&self.schema), error), + } + } +} + +fn error_stream(schema: SchemaRef, error: DataFusionError) -> SendableRecordBatchStream { + Box::pin(RecordBatchStreamAdapter::new(schema, stream::once(async move { Err(error) }))) +} + +async fn prepare_source( + stream: futures::stream::BoxStream<'static, Result>, + memory_pool: Arc, +) -> DFResult<(SchemaRef, JsonDocumentSource)> { + let reservation = MemoryConsumer::new("S3 Select JSON schema inference").register(&memory_pool); + let mut input = JsonDocumentInput::new(stream); + let mut sample = Vec::new(); + let mut records = 0; + let mut largest_record = 0; + + while records < SCHEMA_INFERENCE_MAX_RECORDS { + let Some(row) = input.next_row().await? else { + break; + }; + let required = sample.len().checked_add(row.len()).ok_or_else(schema_prefix_resource_error)?; + if required > SCHEMA_INFERENCE_MAX_BYTES { + return Err(schema_prefix_too_large_error()); + } + reserve_schema_prefix(&mut sample, required, &reservation)?; + largest_record = largest_record.max(row.len()); + sample.extend_from_slice(&row); + records += 1; + } + + reservation.try_resize(schema_inference_reservation_bytes(sample.capacity(), largest_record)?)?; + let cancellation = Arc::new(AtomicBool::new(false)); + let _cancel_on_drop = SchemaInferenceCancellation(Arc::clone(&cancellation)); + let inference_task = SpawnedTask::spawn_blocking(move || { + let result = infer_schema(&sample, records, cancellation.as_ref(), SCHEMA_INFERENCE_MAX_SCHEMA_BYTES); + (sample, reservation, input, result) + }); + let (sample, reservation, input, inference) = inference_task + .join() + .await + .map_err(|error| DataFusionError::Execution(format!("JSON DOCUMENT schema inference worker failed: {error}")))?; + let (schema, inferred_records, schema_bytes) = inference?; + if inferred_records != records { + return Err(DataFusionError::Execution(format!( + "JSON DOCUMENT schema prefix contained {records} records but Arrow decoded {inferred_records}" + ))); + } + let replay_capacity = sample.capacity(); + let retained_bytes = replay_capacity + .checked_add(schema_bytes) + .ok_or_else(schema_prefix_resource_error)?; + reservation.try_resize(retained_bytes)?; + let replay = (!sample.is_empty()).then(|| Bytes::from(sample)); + Ok(( + Arc::new(schema), + JsonDocumentSource { + replay, + replay_reservation: replay_capacity, + input, + reservation, + memory_pool, + }, + )) +} + +struct SchemaInferenceCancellation(Arc); + +impl Drop for SchemaInferenceCancellation { + fn drop(&mut self) { + self.0.store(true, Ordering::Release); + } +} + +fn schema_inference_reservation_bytes(sample_capacity: usize, largest_record: usize) -> DFResult { + largest_record + .checked_mul(SCHEMA_INFERENCE_WORK_MULTIPLIER) + .and_then(|work| work.checked_add(sample_capacity)) + .and_then(|bytes| bytes.checked_add(SCHEMA_INFERENCE_MAX_SCHEMA_BYTES)) + .ok_or_else(schema_prefix_resource_error) +} + +fn reserve_schema_prefix(buffer: &mut Vec, required: usize, reservation: &MemoryReservation) -> DFResult<()> { + if required <= buffer.capacity() { + return Ok(()); + } + let target = required + .checked_next_power_of_two() + .ok_or_else(schema_prefix_resource_error)?; + reservation.try_resize(target)?; + buffer + .try_reserve_exact(target.saturating_sub(buffer.len())) + .map_err(|_| schema_prefix_resource_error())?; + reservation.try_resize(buffer.capacity()) +} + +fn schema_prefix_resource_error() -> DataFusionError { + DataFusionError::ResourcesExhausted("JSON DOCUMENT schema prefix exceeds the query memory limit".to_string()) +} + +fn schema_prefix_too_large_error() -> DataFusionError { + DataFusionError::ResourcesExhausted(format!( + "JSON DOCUMENT schema inference exceeds the bounded {SCHEMA_INFERENCE_MAX_BYTES}-byte prefix" + )) +} + +/// The object-store adapter emits chunks containing one or more complete +/// newline-terminated rows. A row split across chunks violates that boundary. +struct JsonDocumentInput { + stream: futures::stream::BoxStream<'static, Result>, + current: Bytes, + finished: bool, +} + +impl fmt::Debug for JsonDocumentInput { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("JsonDocumentInput") + .field("current_bytes", &self.current.len()) + .field("finished", &self.finished) + .finish_non_exhaustive() + } +} + +impl JsonDocumentInput { + fn new(stream: futures::stream::BoxStream<'static, Result>) -> Self { + Self { + stream, + current: Bytes::new(), + finished: false, + } + } + + async fn next_row(&mut self) -> DFResult> { + loop { + if let Some(newline) = self.current.iter().position(|byte| *byte == b'\n') { + return Ok(Some(self.current.split_to(newline + 1))); + } + if !self.current.is_empty() { + return Err(DataFusionError::Execution( + "JSON DOCUMENT row stream produced a non-newline-terminated chunk".to_string(), + )); + } + if self.finished { + return Ok(None); + } + match self.stream.next().await { + Some(Ok(bytes)) => self.current = bytes, + Some(Err(error)) => return Err(DataFusionError::ObjectStore(Box::new(error))), + None => self.finished = true, + } + } + } + + fn poll_next_bytes(&mut self, context: &mut Context<'_>) -> Poll>> { + if !self.current.is_empty() { + return Poll::Ready(Some(Ok(std::mem::take(&mut self.current)))); + } + if self.finished { + return Poll::Ready(None); + } + match self.stream.as_mut().poll_next(context) { + Poll::Ready(Some(Ok(bytes))) => Poll::Ready(Some(Ok(bytes))), + Poll::Ready(Some(Err(error))) => Poll::Ready(Some(Err(DataFusionError::ObjectStore(Box::new(error))))), + Poll::Ready(None) => { + self.finished = true; + Poll::Ready(None) + } + Poll::Pending => Poll::Pending, + } + } +} + +struct JsonDocumentSource { + replay: Option, + replay_reservation: usize, + input: JsonDocumentInput, + reservation: MemoryReservation, + memory_pool: Arc, +} + +impl fmt::Debug for JsonDocumentSource { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("JsonDocumentSource") + .field("replay_bytes", &self.replay.as_ref().map_or(0, Bytes::len)) + .field("input", &self.input) + .finish_non_exhaustive() + } +} + +impl JsonDocumentSource { + fn poll_next(&mut self, context: &mut Context<'_>) -> Poll>> { + if let Some(replay) = self.replay.take() { + return Poll::Ready(Some(Ok((replay, self.replay_reservation)))); + } + self.input + .poll_next_bytes(context) + .map(|item| item.map(|result| result.map(|bytes| (bytes, 0)))) + } +} + +struct JsonDocumentBatchStream { + schema: SchemaRef, + current: Bytes, + ready: Option, + flush_task: Option>, + decoder_base_reservation: usize, + batch_size: usize, + batch_input_bytes: usize, + row_bytes_since_newline: usize, + row_probe_bytes: usize, + row_boundary_verified: bool, + release_decoder_work: bool, + current_reservation: usize, + done: bool, +} + +struct JsonDocumentDecoderReady { + source: JsonDocumentSource, + decoder: datafusion::arrow::json::reader::Decoder, + reservation: MemoryReservation, +} + +struct JsonDocumentFlushOutput { + ready: JsonDocumentDecoderReady, + result: DFResult>, + final_flush: bool, +} + +impl JsonDocumentBatchStream { + fn try_new(schema: SchemaRef, source: JsonDocumentSource, batch_size: usize) -> DFResult { + let (batch_size, decoder_base_reservation) = json_decoder_layout(&schema, batch_size.max(1))?; + let reservation = MemoryConsumer::new("S3 Select JSON decoder").register(&source.memory_pool); + reservation.try_resize(decoder_base_reservation)?; + let decoder = ReaderBuilder::new(Arc::clone(&schema)) + .with_batch_size(batch_size) + .build_decoder()?; + Ok(Self { + schema, + current: Bytes::new(), + ready: Some(JsonDocumentDecoderReady { + source, + decoder, + reservation, + }), + flush_task: None, + decoder_base_reservation, + batch_size, + batch_input_bytes: 0, + row_bytes_since_newline: 0, + row_probe_bytes: 0, + row_boundary_verified: false, + release_decoder_work: false, + current_reservation: 0, + done: false, + }) + } + + fn release_current(&mut self) -> DFResult<()> { + self.current = Bytes::new(); + self.row_bytes_since_newline = 0; + self.row_probe_bytes = 0; + self.row_boundary_verified = false; + let reservation = std::mem::take(&mut self.current_reservation); + if reservation > 0 + && let Some(ready) = self.ready.as_ref() + { + ready.source.reservation.try_shrink(reservation)?; + } + Ok(()) + } + + fn reserve_decoder_input(&mut self, bytes: usize) -> DFResult<()> { + let input_bytes = self.batch_input_bytes.checked_add(bytes).ok_or_else(decoder_resource_error)?; + let work = input_bytes + .checked_mul(JSON_DECODE_MEMORY_MULTIPLIER) + .ok_or_else(decoder_resource_error)?; + let reservation = self + .decoder_base_reservation + .checked_add(work) + .ok_or_else(decoder_resource_error)?; + self.ready + .as_ref() + .ok_or_else(decoder_worker_error)? + .reservation + .try_resize(reservation) + } + + fn release_decoder_work(&mut self) -> DFResult<()> { + if let Some(ready) = self.ready.as_ref() { + ready.reservation.try_resize(self.decoder_base_reservation)?; + } + Ok(()) + } + + fn start_flush(&mut self, context: &mut Context<'_>, final_flush: bool) -> Poll>> { + let Some(mut ready) = self.ready.take() else { + return self.fail(decoder_worker_error()); + }; + self.flush_task = Some(SpawnedTask::spawn_blocking(move || { + let result = ready.decoder.flush().map_err(DataFusionError::from); + JsonDocumentFlushOutput { + ready, + result, + final_flush, + } + })); + self.poll_flush(context) + } + + fn poll_flush(&mut self, context: &mut Context<'_>) -> Poll>> { + let poll = match self.flush_task.as_mut() { + Some(task) => Pin::new(task).poll(context), + None => return self.fail(decoder_worker_error()), + }; + let Poll::Ready(joined) = poll else { + return Poll::Pending; + }; + self.flush_task = None; + let output = match joined { + Ok(output) => output, + Err(_) => return self.fail(decoder_worker_error()), + }; + let final_flush = output.final_flush; + self.ready = Some(output.ready); + match output.result { + Ok(Some(batch)) => { + self.batch_input_bytes = 0; + self.release_decoder_work = true; + self.done = final_flush; + Poll::Ready(Some(Ok(batch))) + } + Ok(None) if final_flush => { + self.batch_input_bytes = 0; + if let Err(error) = self.release_decoder_work() { + return self.fail(error); + } + self.done = true; + Poll::Ready(None) + } + Ok(None) => self.fail(DataFusionError::Execution( + "Arrow JSON decoder reached a flush boundary without producing a batch".to_string(), + )), + Err(error) => self.fail(error), + } + } + + fn fail(&mut self, mut error: DataFusionError) -> Poll>> { + self.done = true; + if let Err(release_error) = self.release_current() { + error = release_error; + } + if let Err(release_error) = self.release_decoder_work() { + error = release_error; + } + Poll::Ready(Some(Err(error))) + } +} + +impl RecordBatchStream for JsonDocumentBatchStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +impl Stream for JsonDocumentBatchStream { + type Item = DFResult; + + fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + if self.flush_task.is_some() { + return self.poll_flush(context); + } + if self.release_decoder_work { + if let Err(error) = self.release_decoder_work() { + return self.fail(error); + } + self.release_decoder_work = false; + } + if self.done { + return Poll::Ready(None); + } + let mut poll_budget = JSON_DECODE_POLL_BYTES; + loop { + let Some(ready) = self.ready.as_ref() else { + return self.fail(decoder_worker_error()); + }; + if ready.decoder.len() >= self.batch_size + || (self.batch_input_bytes >= JSON_DECODE_BATCH_BYTES + && !ready.decoder.has_partial_record() + && !ready.decoder.is_empty()) + { + return self.start_flush(context, false); + } + if poll_budget == 0 { + context.waker().wake_by_ref(); + return Poll::Pending; + } + if self.current.is_empty() { + if let Err(error) = self.release_current() { + return self.fail(error); + } + let source_poll = match self.ready.as_mut() { + Some(ready) => ready.source.poll_next(context), + None => return self.fail(decoder_worker_error()), + }; + match source_poll { + Poll::Ready(Some(Ok((bytes, reservation)))) => { + self.current = bytes; + self.current_reservation = reservation; + } + Poll::Ready(Some(Err(error))) => return self.fail(error), + Poll::Ready(None) => return self.start_flush(context, true), + Poll::Pending => return Poll::Pending, + } + } + + let batch_budget = JSON_DECODE_BATCH_BYTES.saturating_sub(self.batch_input_bytes); + let decoder_has_complete_rows = self + .ready + .as_ref() + .is_some_and(|ready| !ready.decoder.has_partial_record() && !ready.decoder.is_empty()) + && self.row_bytes_since_newline == 0; + if decoder_has_complete_rows && self.current.len() > batch_budget && !self.row_boundary_verified { + let probe_start = self.row_probe_bytes; + let probe_end = probe_start.saturating_add(poll_budget).min(batch_budget); + if let Some(newline) = self.current[probe_start..probe_end].iter().position(|byte| *byte == b'\n') { + poll_budget = poll_budget.saturating_sub(newline + 1); + self.row_probe_bytes = 0; + self.row_boundary_verified = true; + } else if probe_end == batch_budget { + self.row_probe_bytes = 0; + self.row_boundary_verified = false; + return self.start_flush(context, false); + } else { + self.row_probe_bytes = probe_end; + context.waker().wake_by_ref(); + return Poll::Pending; + } + } + if poll_budget == 0 { + context.waker().wake_by_ref(); + return Poll::Pending; + } + let mut available = self.current.len().min(poll_budget); + if batch_budget > 0 { + available = available.min(batch_budget); + } + if self.row_bytes_since_newline > 0 { + if let Some(newline) = self.current[..available].iter().position(|byte| *byte == b'\n') { + available = newline + 1; + } + } else if decoder_has_complete_rows + && let Some(newline) = self.current[..available].iter().rposition(|byte| *byte == b'\n') + { + available = newline + 1; + } + if let Err(error) = self.reserve_decoder_input(available) { + return self.fail(error); + } + let decode_result = { + let stream = self.as_mut().get_mut(); + match stream.ready.as_mut() { + Some(ready) => ready.decoder.decode(&stream.current[..available]), + None => return stream.fail(decoder_worker_error()), + } + }; + let decoded = match decode_result { + Ok(decoded) => decoded, + Err(error) => return self.fail(DataFusionError::from(error)), + }; + if decoded == 0 { + return self.fail(DataFusionError::Execution("Arrow JSON decoder made no progress".to_string())); + } + self.row_bytes_since_newline = match self.current[..decoded].iter().rposition(|byte| *byte == b'\n') { + Some(newline) => decoded - newline - 1, + None => match self.row_bytes_since_newline.checked_add(decoded) { + Some(bytes) => bytes, + None => return self.fail(decoder_resource_error()), + }, + }; + self.current.advance(decoded); + self.row_boundary_verified = false; + self.batch_input_bytes = match self.batch_input_bytes.checked_add(decoded) { + Some(bytes) => bytes, + None => return self.fail(decoder_resource_error()), + }; + poll_budget = poll_budget.saturating_sub(decoded); + } + } +} + +fn json_decoder_layout(schema: &SchemaRef, requested_batch_size: usize) -> DFResult<(usize, usize)> { + let fields = schema.flattened_fields().len(); + let field_slots = fields.checked_mul(2).ok_or_else(decoder_resource_error)?; + let offset_bytes_per_row = field_slots + .checked_mul(std::mem::size_of::()) + .ok_or_else(decoder_resource_error)?; + let tape_bytes_per_row = field_slots + .checked_add(2) + .and_then(|elements| elements.checked_mul(std::mem::size_of::())) + .ok_or_else(decoder_resource_error)?; + let bytes_per_row = offset_bytes_per_row + .checked_add(tape_bytes_per_row) + .ok_or_else(decoder_resource_error)?; + let metadata_bytes = fields + .checked_mul(JSON_DECODER_PER_FIELD_METADATA_BYTES) + .and_then(|bytes| bytes.checked_add(field_slots.checked_mul(std::mem::size_of::())?)) + .and_then(|bytes| bytes.checked_add(10 * std::mem::size_of::())) + .ok_or_else(decoder_resource_error)?; + let tape_budget = JSON_DECODER_TAPE_TARGET_BYTES.saturating_sub(metadata_bytes); + let bounded_batch_size = requested_batch_size.min((tape_budget / bytes_per_row).max(1)); + let tape_bytes = bytes_per_row + .checked_mul(bounded_batch_size) + .and_then(|bytes| bytes.checked_add(std::mem::size_of::() + std::mem::size_of::())) + .ok_or_else(decoder_resource_error)?; + let reservation = metadata_bytes.checked_add(tape_bytes).ok_or_else(decoder_resource_error)?; + Ok((bounded_batch_size, reservation)) +} + +fn decoder_resource_error() -> DataFusionError { + DataFusionError::ResourcesExhausted("JSON DOCUMENT decoder exceeds the query memory limit".to_string()) +} + +fn decoder_worker_error() -> DataFusionError { + DataFusionError::Execution("JSON DOCUMENT decoder worker terminated unexpectedly".to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::{ + arrow::{array::Int64Array, datatypes::Schema}, + object_store::memory::InMemory, + prelude::SessionContext, + }; + use futures::TryStreamExt as _; + + #[tokio::test] + async fn schema_prefix_is_replayed_without_reopening_the_object() { + let store = Arc::new(InMemory::new()); + let path = Path::from("input.json"); + store + .put(&path, Bytes::from_static(b"{\"id\":1}\n{\"id\":2}\n").into()) + .await + .expect("put fixture"); + let context = SessionContext::new(); + let url = datafusion::execution::object_store::ObjectStoreUrl::parse("s3://bucket").expect("store URL"); + context.register_object_store(url.as_ref(), store); + let provider = JsonDocumentTable::try_new(&context.state(), "bucket", "input.json") + .await + .expect("prepare streaming provider"); + context.register_table("S3Object", provider).expect("register table"); + + let batches = context + .sql("SELECT id FROM S3Object ORDER BY id") + .await + .expect("plan query") + .collect() + .await + .expect("execute query"); + let values: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .expect("id column") + .values() + .iter() + .copied() + }) + .collect(); + assert_eq!(values, vec![1, 2]); + + let error = context + .sql("SELECT id FROM S3Object") + .await + .expect("plan second scan") + .collect() + .await + .expect_err("JSON DOCUMENT source must be consumed exactly once"); + assert!(error.to_string().contains("consumed more than once")); + } + + #[tokio::test] + async fn schema_prefix_fails_closed_at_its_byte_limit() { + const ROWS: usize = 6; + const PAYLOAD_BYTES: usize = 700 * 1024; + let rows = (0..ROWS) + .map(|id| { + Ok::<_, ObjectStoreError>(Bytes::from(format!("{{\"id\":{id},\"payload\":\"{}\"}}\n", "x".repeat(PAYLOAD_BYTES)))) + }) + .collect::>(); + let memory_pool: Arc = Arc::new(datafusion::execution::memory_pool::UnboundedMemoryPool::default()); + let error = prepare_source(stream::iter(rows).boxed(), memory_pool) + .await + .expect_err("an incomplete schema prefix must not silently drop later fields"); + assert!(matches!(error, DataFusionError::ResourcesExhausted(_))); + } + + #[tokio::test] + async fn empty_row_stream_produces_an_empty_query_result() { + let memory_pool: Arc = Arc::new(datafusion::execution::memory_pool::UnboundedMemoryPool::default()); + let (schema, source) = prepare_source(stream::empty().boxed(), memory_pool) + .await + .expect("prepare empty JSON DOCUMENT stream"); + assert!(schema.fields().is_empty()); + + let provider = Arc::new(JsonDocumentTable { + schema, + source: Arc::new(Mutex::new(Some(source))), + }); + let context = SessionContext::new(); + context.register_table("S3Object", provider).expect("register empty table"); + let batches = context + .sql("SELECT COUNT(*) FROM S3Object") + .await + .expect("plan empty query") + .collect() + .await + .expect("execute empty query"); + let count = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .expect("count column") + .value(0); + assert_eq!(count, 0); + } + + #[tokio::test] + async fn object_row_with_an_empty_schema_is_counted() { + let input = stream::once(async { Ok::<_, ObjectStoreError>(Bytes::from_static(b"{}\n")) }).boxed(); + let memory_pool: Arc = Arc::new(datafusion::execution::memory_pool::UnboundedMemoryPool::default()); + let (schema, source) = prepare_source(input, memory_pool) + .await + .expect("prepare empty-object JSON DOCUMENT stream"); + assert!(schema.fields().is_empty()); + + let provider = Arc::new(JsonDocumentTable { + schema, + source: Arc::new(Mutex::new(Some(source))), + }); + let context = SessionContext::new(); + context.register_table("S3Object", provider).expect("register table"); + let batches = context + .sql("SELECT COUNT(*) FROM S3Object") + .await + .expect("plan empty-object query") + .collect() + .await + .expect("execute empty-object query"); + let count = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .expect("count column") + .value(0); + assert_eq!(count, 1); + } + + #[tokio::test] + async fn limit_drops_a_stalled_tail_after_the_bounded_schema_prefix() { + const EXPECTED_SCHEMA_PREFIX_RECORDS: usize = 1000; + assert_eq!(SCHEMA_INFERENCE_MAX_RECORDS, EXPECTED_SCHEMA_PREFIX_RECORDS); + let rows = (0..EXPECTED_SCHEMA_PREFIX_RECORDS).map(|id| Ok(Bytes::from(format!("{{\"id\":{id}}}\n")))); + let input = stream::iter(rows).chain(stream::pending()).boxed(); + let memory_pool = Arc::new(datafusion::execution::memory_pool::GreedyMemoryPool::new(16 * 1024 * 1024)); + let (schema, source) = prepare_source(input, memory_pool.clone()) + .await + .expect("prepare bounded schema prefix"); + let provider = Arc::new(JsonDocumentTable { + schema, + source: Arc::new(Mutex::new(Some(source))), + }); + let context = SessionContext::new(); + context + .register_table("S3Object", provider) + .expect("register streaming table"); + + let batches = tokio::time::timeout(std::time::Duration::from_secs(1), async { + context + .sql("SELECT id FROM S3Object LIMIT 1") + .await + .expect("plan limited query") + .collect() + .await + }) + .await + .expect("LIMIT should not wait for the unread object tail") + .expect("execute limited query"); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!(memory_pool.reserved(), 0); + } + + #[tokio::test] + async fn row_chunks_must_end_on_a_record_boundary() { + let chunks = stream::iter([ + Ok::<_, ObjectStoreError>(Bytes::from_static(b"{\"id\":")), + Ok(Bytes::from_static(b"1}\n")), + ]) + .boxed(); + let mut input = JsonDocumentInput::new(chunks); + + let error = input + .next_row() + .await + .expect_err("a producer must not split one JSON row across chunks"); + assert!(error.to_string().contains("non-newline-terminated chunk")); + } + + #[test] + fn schema_inference_has_no_fixed_field_count_limit() { + const FIELDS: usize = 4097; + let cancellation = AtomicBool::new(false); + let mut input = json_object_with_unique_keys(FIELDS); + input.push(b'\n'); + + let (schema, records, _) = infer_schema(&input, 1, &cancellation, SCHEMA_INFERENCE_MAX_SCHEMA_BYTES) + .expect("a protocol-valid wide object should be governed by memory, not a field-count constant"); + assert_eq!(records, 1); + assert_eq!(schema.fields().len(), FIELDS); + } + + #[tokio::test] + async fn schema_prefix_respects_the_query_memory_pool() { + let memory_pool = Arc::new(datafusion::execution::memory_pool::GreedyMemoryPool::new(1)); + let input = stream::once(async { Ok::<_, ObjectStoreError>(Bytes::from_static(b"{\"id\":1}\n")) }).boxed(); + + let error = prepare_source(input, memory_pool.clone()) + .await + .expect_err("schema prefix allocation must use the query memory pool"); + assert!(matches!(error, DataFusionError::ResourcesExhausted(_))); + assert_eq!(memory_pool.reserved(), 0); + } + + #[test] + fn schema_inference_cancellation_reaches_cpu_boundaries() { + let cancellation = Arc::new(AtomicBool::new(false)); + drop(SchemaInferenceCancellation(Arc::clone(&cancellation))); + assert!(cancellation.load(Ordering::Acquire)); + + let error = infer_schema(b"{}\n", 1, cancellation.as_ref(), SCHEMA_INFERENCE_MAX_SCHEMA_BYTES) + .expect_err("schema inference must observe cancellation"); + assert!(error.to_string().contains("schema inference canceled")); + } + + #[test] + fn queued_schema_inference_retains_query_admission() { + 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(tokio::sync::Semaphore::new(1)); + let query_guard = Arc::new( + Arc::clone(&admission) + .acquire_owned() + .await + .expect("query admission should be available"), + ); + let input = stream::iter([Ok::<_, ObjectStoreError>(Bytes::from_static(b"{\"id\":1}\n"))]) + .map(move |row| { + let _query_guard = &query_guard; + row + }) + .boxed(); + let memory_pool: Arc = Arc::new(datafusion::execution::memory_pool::UnboundedMemoryPool::default()); + { + let inference = prepare_source(input, memory_pool); + futures::pin_mut!(inference); + assert!(futures::poll!(inference.as_mut()).is_pending()); + } + + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), Arc::clone(&admission).acquire_owned(),) + .await + .is_err(), + "queued schema inference must retain query admission through teardown" + ); + + release_blocking_tx.send(()).expect("release blocking worker"); + blocker.await.expect("blocking worker should finish"); + let recovered = tokio::time::timeout(std::time::Duration::from_secs(1), Arc::clone(&admission).acquire_owned()) + .await + .expect("schema worker teardown should release query admission") + .expect("query admission should remain open"); + drop(recovered); + }); + } + + #[tokio::test] + async fn projection_is_applied_before_json_value_decoding() { + let rows = stream::iter( + (0..SCHEMA_INFERENCE_MAX_RECORDS) + .map(|id| Ok::<_, ObjectStoreError>(Bytes::from(format!("{{\"id\":{id},\"payload\":\"ok\"}}\n")))), + ) + .chain(stream::once(async { + Ok(Bytes::from_static(b"{\"id\":1000,\"payload\":{\"shape\":\"incompatible\"}}\n")) + })) + .boxed(); + let memory_pool: Arc = Arc::new(datafusion::execution::memory_pool::UnboundedMemoryPool::default()); + let (schema, source) = prepare_source(rows, memory_pool).await.expect("prepare projection fixture"); + let provider = Arc::new(JsonDocumentTable { + schema, + source: Arc::new(Mutex::new(Some(source))), + }); + let context = SessionContext::new(); + context.register_table("S3Object", provider).expect("register table"); + + let batches = context + .sql("SELECT id FROM S3Object") + .await + .expect("plan projected query") + .collect() + .await + .expect("unprojected type changes must not affect selected columns"); + let values: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .expect("id column") + .values() + .iter() + .copied() + }) + .collect(); + assert_eq!(values.len(), SCHEMA_INFERENCE_MAX_RECORDS + 1); + assert_eq!(values.last(), Some(&1000)); + } + + #[tokio::test] + async fn filter_columns_remain_available_to_the_physical_filter() { + let rows = stream::iter([ + Ok::<_, ObjectStoreError>(Bytes::from_static(b"{\"id\":1,\"payload\":\"keep\"}\n")), + Ok(Bytes::from_static(b"{\"id\":2,\"payload\":\"drop\"}\n")), + ]) + .boxed(); + let memory_pool: Arc = Arc::new(datafusion::execution::memory_pool::UnboundedMemoryPool::default()); + let (schema, source) = prepare_source(rows, memory_pool).await.expect("prepare filter fixture"); + let provider = Arc::new(JsonDocumentTable { + schema, + source: Arc::new(Mutex::new(Some(source))), + }); + let context = SessionContext::new(); + context.register_table("S3Object", provider).expect("register table"); + + let batches = context + .sql("SELECT id FROM S3Object WHERE payload = 'keep'") + .await + .expect("plan filtered query") + .collect() + .await + .expect("execute filtered query"); + let values: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .expect("id column") + .values() + .iter() + .copied() + }) + .collect(); + assert_eq!(values, vec![1]); + } + + #[tokio::test] + async fn decoder_flushes_large_input_on_a_byte_boundary() { + use datafusion::arrow::datatypes::{DataType, Field}; + + const ROWS: usize = 10; + const PAYLOAD_BYTES: usize = 512 * 1024; + let rows = (0..ROWS) + .map(|id| { + Ok::<_, ObjectStoreError>(Bytes::from(format!("{{\"id\":{id},\"payload\":\"{}\"}}\n", "x".repeat(PAYLOAD_BYTES)))) + }) + .collect::>(); + let memory_pool: Arc = Arc::new(datafusion::execution::memory_pool::UnboundedMemoryPool::default()); + let source = JsonDocumentSource { + replay: None, + replay_reservation: 0, + input: JsonDocumentInput::new(stream::iter(rows).boxed()), + reservation: MemoryConsumer::new("byte-batch source").register(&memory_pool), + memory_pool, + }; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, true), + Field::new("payload", DataType::Utf8, true), + ])); + let batches = JsonDocumentBatchStream::try_new(schema, source, 1024) + .expect("build JSON decoder") + .try_collect::>() + .await + .expect("decode large input"); + + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), ROWS); + assert!(batches.len() >= 2, "the decoder must not retain more than one byte-bounded batch"); + } + + #[tokio::test] + async fn decoder_flushes_before_a_large_row_exceeds_the_remaining_budget() { + use datafusion::arrow::datatypes::{DataType, Field}; + + let row = |total_bytes: usize| { + const PREFIX: &str = "{\"payload\":\""; + const SUFFIX: &str = "\"}\n"; + let payload_bytes = total_bytes + .checked_sub(PREFIX.len() + SUFFIX.len()) + .expect("row fixture must fit its JSON wrapper"); + Bytes::from(format!("{PREFIX}{}{SUFFIX}", "x".repeat(payload_bytes))) + }; + let rows = stream::iter([ + Ok::<_, ObjectStoreError>(row(JSON_DECODE_BATCH_BYTES - JSON_DECODE_POLL_BYTES + 1)), + Ok(row(JSON_DECODE_BATCH_BYTES)), + ]) + .boxed(); + let memory_pool = Arc::new(datafusion::execution::memory_pool::GreedyMemoryPool::new(40 * 1024 * 1024)); + let query_memory_pool: Arc = memory_pool.clone(); + let source = JsonDocumentSource { + replay: None, + replay_reservation: 0, + input: JsonDocumentInput::new(rows), + reservation: MemoryConsumer::new("large-row boundary source").register(&query_memory_pool), + memory_pool: query_memory_pool, + }; + let schema = Arc::new(Schema::new(vec![Field::new("payload", DataType::Utf8, true)])); + + let batches = JsonDocumentBatchStream::try_new(schema, source, 1024) + .expect("build JSON decoder") + .try_collect::>() + .await + .expect("a legal large row must start in a fresh byte-bounded batch"); + + assert_eq!(batches.len(), 2); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + assert_eq!(memory_pool.reserved(), 0); + } + + #[tokio::test] + async fn decoder_yields_after_its_per_poll_byte_budget() { + const EXPECTED_POLL_BYTES: usize = 64 * 1024; + assert_eq!(JSON_DECODE_POLL_BYTES, EXPECTED_POLL_BYTES); + let row = Bytes::from(format!("{{\"payload\":\"{}\"}}\n", "x".repeat(256 * 1024))); + let memory_pool: Arc = Arc::new(datafusion::execution::memory_pool::UnboundedMemoryPool::default()); + let (schema, source) = prepare_source(stream::once(async move { Ok::<_, ObjectStoreError>(row) }).boxed(), memory_pool) + .await + .expect("prepare cooperative decoder fixture"); + let mut decoder = JsonDocumentBatchStream::try_new(schema, source, 1024).expect("build JSON decoder"); + let waker = futures::task::noop_waker(); + let mut context = Context::from_waker(&waker); + + assert!( + Pin::new(&mut decoder).poll_next(&mut context).is_pending(), + "one poll must not decode an entire large row" + ); + let batches = decoder + .try_collect::>() + .await + .expect("finish decoding after cooperative yield"); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + } + + #[tokio::test] + async fn dense_nested_values_are_charged_before_arrow_allocation() { + use datafusion::arrow::datatypes::{DataType, Field}; + + let replay = Bytes::from(format!("{{\"values\":[{}0]}}\n", "0,".repeat(32 * 1024))); + let schema = Arc::new(Schema::new(vec![Field::new( + "values", + DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true))), + true, + )])); + let (_, decoder_base_reservation) = json_decoder_layout(&schema, 1024).expect("calculate decoder reservation"); + let pool_size = replay + .len() + .checked_mul(JSON_DECODE_MEMORY_MULTIPLIER - 1) + .and_then(|work| work.checked_add(replay.len())) + .and_then(|bytes| bytes.checked_add(decoder_base_reservation)) + .and_then(|bytes| bytes.checked_sub(1)) + .expect("decoder memory fixture should fit"); + let memory_pool = Arc::new(datafusion::execution::memory_pool::GreedyMemoryPool::new(pool_size)); + let query_memory_pool: Arc = memory_pool.clone(); + let reservation = MemoryConsumer::new("JSON decoder memory test").register(&query_memory_pool); + reservation + .try_resize(replay.len()) + .expect("reserve the retained replay buffer"); + let source = JsonDocumentSource { + replay_reservation: replay.len(), + replay: Some(replay), + input: JsonDocumentInput::new(stream::empty().boxed()), + reservation, + memory_pool: query_memory_pool, + }; + let mut decoder = JsonDocumentBatchStream::try_new(schema, source, 1024).expect("build JSON decoder"); + + let error = decoder + .next() + .await + .expect("decoder must emit a memory error") + .expect_err("dense nested values must be charged before Arrow allocates their tape"); + assert!(matches!(error, DataFusionError::ResourcesExhausted(_))); + drop(decoder); + assert_eq!(memory_pool.reserved(), 0); + } + + #[test] + fn wide_schema_bounds_decoder_preallocation() { + use datafusion::arrow::datatypes::{DataType, Field}; + + const FIELDS: usize = 1000; + const REQUESTED_BATCH_SIZE: usize = 8192; + let schema = Arc::new(Schema::new( + (0..FIELDS) + .map(|index| Field::new(format!("field_{index}"), DataType::Utf8, true)) + .collect::>(), + )); + let (batch_size, reservation) = json_decoder_layout(&schema, REQUESTED_BATCH_SIZE).expect("calculate decoder layout"); + + assert!(batch_size < REQUESTED_BATCH_SIZE); + assert!(reservation <= JSON_DECODER_TAPE_TARGET_BYTES + FIELDS * JSON_DECODER_PER_FIELD_METADATA_BYTES); + let memory_pool = Arc::new(datafusion::execution::memory_pool::GreedyMemoryPool::new(reservation)); + let query_memory_pool: Arc = memory_pool.clone(); + let source = JsonDocumentSource { + replay: None, + replay_reservation: 0, + input: JsonDocumentInput::new(stream::empty().boxed()), + reservation: MemoryConsumer::new("wide schema source").register(&query_memory_pool), + memory_pool: query_memory_pool, + }; + let decoder = JsonDocumentBatchStream::try_new(schema, source, REQUESTED_BATCH_SIZE) + .expect("bounded decoder should fit its registered reservation"); + assert_eq!(decoder.batch_size, batch_size); + assert_eq!(memory_pool.reserved(), reservation); + drop(decoder); + assert_eq!(memory_pool.reserved(), 0); + } + + #[test] + fn queued_decoder_flush_retains_upstream_query_admission() { + use datafusion::arrow::datatypes::{DataType, Field}; + + 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(tokio::sync::Semaphore::new(1)); + let query_guard = Arc::new( + Arc::clone(&admission) + .acquire_owned() + .await + .expect("query admission should be available"), + ); + let input = stream::iter([Ok::<_, ObjectStoreError>(Bytes::from_static(b"{\"id\":1}\n"))]) + .map(move |row| { + let _query_guard = &query_guard; + row + }) + .boxed(); + let memory_pool: Arc = Arc::new(datafusion::execution::memory_pool::UnboundedMemoryPool::default()); + let source = JsonDocumentSource { + replay: None, + replay_reservation: 0, + input: JsonDocumentInput::new(input), + reservation: MemoryConsumer::new("queued decoder source").register(&memory_pool), + memory_pool, + }; + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, true)])); + { + let mut decoder = JsonDocumentBatchStream::try_new(schema, source, 1).expect("build JSON decoder"); + let waker = futures::task::noop_waker(); + let mut context = Context::from_waker(&waker); + assert!(Pin::new(&mut decoder).poll_next(&mut context).is_pending()); + } + + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), Arc::clone(&admission).acquire_owned(),) + .await + .is_err(), + "queued decoder teardown must retain upstream query admission" + ); + + release_blocking_tx.send(()).expect("release blocking worker"); + blocker.await.expect("blocking worker should finish"); + let recovered = tokio::time::timeout(std::time::Duration::from_secs(1), Arc::clone(&admission).acquire_owned()) + .await + .expect("decoder teardown should release query admission") + .expect("query admission should remain open"); + drop(recovered); + }); + } + + fn json_object_with_unique_keys(keys: usize) -> Vec { + use std::fmt::Write as _; + + let mut object = String::from("{"); + for index in 0..keys { + if index > 0 { + object.push(','); + } + write!(&mut object, "\"field_{index}\":0").expect("write key fixture"); + } + object.push('}'); + object.into_bytes() + } +} diff --git a/crates/s3select-query/src/dispatcher/manager.rs b/crates/s3select-query/src/dispatcher/manager.rs index 0f69949ff..05a00e3e5 100644 --- a/crates/s3select-query/src/dispatcher/manager.rs +++ b/crates/s3select-query/src/dispatcher/manager.rs @@ -61,7 +61,7 @@ use tokio::{ }; use crate::{ - dispatcher::parquet_table::ParquetSelectTable, + dispatcher::{json_document_table::JsonDocumentTable, parquet_table::ParquetSelectTable}, execution::factory::QueryExecutionFactoryRef, instance::{DEFAULT_MAX_CONCURRENT_QUERIES, DEFAULT_QUERY_TIMEOUT_SECS}, metadata::{ContextProviderExtension, MetadataProvider, TableHandleProviderRef, base_table::BaseTableProvider}, @@ -181,7 +181,9 @@ impl QueryDispatcher for SimpleQueryDispatcher { Ok(logical_plan) }) .await?; - query_state_machine.query.input_metrics().reset(); + if !is_json_document_input(&self.input) { + query_state_machine.query.input_metrics().reset(); + } if !query_tracker.mark_planned(&self.query_execution_owner) { drop(logical_plan); return Err(self.query_tracker_error(&query_tracker)); @@ -415,6 +417,17 @@ impl SimpleQueryDispatcher { return Ok(metadata_provider); } + if is_json_document_input(&self.input) { + let provider = JsonDocumentTable::try_new(session.inner(), &self.input.bucket, &self.input.key).await?; + let current_session_table_provider = self.build_table_handle_provider()?; + return Ok(MetadataProvider::new( + provider, + current_session_table_provider, + self.func_manager.clone(), + session.clone(), + )); + } + let path = format!("s3://{}/{}", self.input.bucket, self.input.key); let table_path = ListingTableUrl::parse(path)?; let compressed_input = self @@ -542,19 +555,24 @@ impl SimpleQueryDispatcher { } } +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() == JSONType::DOCUMENT) +} + fn validate_json_source_path_input(input: &SelectObjectContentInput, source_path: &[JsonPathSegment]) -> QueryResult<()> { if source_path.is_empty() { return Ok(()); } - let Some(json) = input.request.input_serialization.json.as_ref() else { + if input.request.input_serialization.json.is_none() { return Err(SelectError::DataSourcePathUnsupported.into()); - }; - if !source_path_requires_expansion(source_path) - || json - .type_ - .as_ref() - .is_some_and(|json_type| json_type.as_str() == JSONType::DOCUMENT) - { + } + if !source_path_requires_expansion(source_path) || is_json_document_input(input) { return Ok(()); } Err(SelectError::DataSourcePathUnsupported.into()) @@ -1803,6 +1821,63 @@ mod tests { } } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn json_document_schema_prefix_is_counted_once() { + const DOCUMENT: &[u8] = br#"["one","two"]"#; + let mut input = test_input(); + input.bucket = "s3select-json-document-metrics".to_string(); + input.key = "input.json".to_string(); + input.request.expression = "SELECT _1 FROM S3Object[*]".to_string(); + input.request.input_serialization = InputSerialization { + json: Some(JSONInput { + type_: Some(JSONType::from_static(JSONType::DOCUMENT)), + }), + ..Default::default() + }; + input.request.output_serialization = OutputSerialization { + json: Some(JSONOutput::default()), + ..Default::default() + }; + let input = Arc::new(input); + let env = snapshot_test_env().await; + env.make_bucket(&input.bucket, false).await; + env.put_object_bytes(&input.bucket, &input.key, DOCUMENT.to_vec()).await; + let snapshot = env.prepare_select_object_snapshot(&input.bucket, &input.key).await; + let dispatcher = production_dispatcher(Arc::clone(&input)); + let query = Query::new_with_snapshot( + QueryContext { + input: Arc::clone(&input), + }, + input.request.expression.clone(), + snapshot, + ); + let state_machine = dispatcher + .build_query_state_machine(query) + .await + .expect("build tracked JSON DOCUMENT query"); + let input_metrics = Arc::clone(state_machine.query.input_metrics()); + + let logical_plan = dispatcher + .build_logical_plan(Arc::clone(&state_machine)) + .await + .expect("JSON DOCUMENT query should plan") + .expect("SELECT should produce a logical plan"); + let expected_bytes = u64::try_from(DOCUMENT.len()).expect("fixture size should fit in u64"); + assert_eq!(input_metrics.snapshot().bytes_scanned, expected_bytes); + assert_eq!(input_metrics.snapshot().bytes_processed, expected_bytes); + + let values = collect_utf8_output( + dispatcher + .execute_logical_plan(logical_plan, state_machine) + .await + .expect("JSON DOCUMENT query should execute"), + ) + .await; + assert_eq!(values, ["one", "two"]); + assert_eq!(input_metrics.snapshot().bytes_scanned, expected_bytes); + assert_eq!(input_metrics.snapshot().bytes_processed, expected_bytes); + } + #[tokio::test] async fn staged_query_rejects_unbound_session() { let admission = Arc::new(Semaphore::new(1)); diff --git a/crates/s3select-query/src/dispatcher/mod.rs b/crates/s3select-query/src/dispatcher/mod.rs index 83ecdb005..4b9d89930 100644 --- a/crates/s3select-query/src/dispatcher/mod.rs +++ b/crates/s3select-query/src/dispatcher/mod.rs @@ -12,5 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod json_document_schema; +mod json_document_table; pub mod manager; mod parquet_table;