feat(s3-select): stream JSON document input (#6980)
This commit is contained in:
GatewayJ
Generated
+2
@@ -10615,6 +10615,8 @@ dependencies = [
|
||||
"rustfs-s3select-api",
|
||||
"rustfs-test-utils",
|
||||
"s3s",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
@@ -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<io::Result<Bytes>>,
|
||||
current: Bytes,
|
||||
}
|
||||
|
||||
impl BlockingChannelReader {
|
||||
fn new(receiver: mpsc::Receiver<io::Result<Bytes>>) -> Self {
|
||||
pub(crate) fn new(receiver: mpsc::Receiver<io::Result<Bytes>>) -> Self {
|
||||
Self {
|
||||
receiver,
|
||||
current: Bytes::new(),
|
||||
@@ -639,7 +639,7 @@ fn error_chain_contains<T: StdError + 'static>(error: &(dyn StdError + 'static))
|
||||
find_error_source::<T>(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 {
|
||||
|
||||
+131
-1078
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -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 }
|
||||
|
||||
|
||||
@@ -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::<Schema>()
|
||||
.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<String>) -> 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<usize> {
|
||||
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<E: de::Error>(&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<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
|
||||
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<E>(self) -> Result<Self::Value, E> {
|
||||
Ok(InferredType::Any)
|
||||
}
|
||||
|
||||
fn visit_none<E>(self) -> Result<Self::Value, E> {
|
||||
Ok(InferredType::Any)
|
||||
}
|
||||
|
||||
fn visit_bool<E>(self, _value: bool) -> Result<Self::Value, E> {
|
||||
Ok(InferredType::Scalar(ScalarKinds::BOOLEAN))
|
||||
}
|
||||
|
||||
fn visit_i64<E>(self, _value: i64) -> Result<Self::Value, E> {
|
||||
Ok(InferredType::Scalar(ScalarKinds::INTEGER))
|
||||
}
|
||||
|
||||
fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E> {
|
||||
Ok(InferredType::Scalar(if i64::try_from(value).is_ok() {
|
||||
ScalarKinds::INTEGER
|
||||
} else {
|
||||
ScalarKinds::FLOAT
|
||||
}))
|
||||
}
|
||||
|
||||
fn visit_f64<E>(self, _value: f64) -> Result<Self::Value, E> {
|
||||
Ok(InferredType::Scalar(ScalarKinds::FLOAT))
|
||||
}
|
||||
|
||||
fn visit_str<E>(self, _value: &str) -> Result<Self::Value, E> {
|
||||
Ok(InferredType::Scalar(ScalarKinds::STRING))
|
||||
}
|
||||
|
||||
fn visit_string<E>(self, _value: String) -> Result<Self::Value, E> {
|
||||
Ok(InferredType::Scalar(ScalarKinds::STRING))
|
||||
}
|
||||
|
||||
fn visit_seq<A>(self, mut sequence: A) -> Result<Self::Value, A::Error>
|
||||
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<A>(self, mut map: A) -> Result<Self::Value, A::Error>
|
||||
where
|
||||
A: MapAccess<'de>,
|
||||
{
|
||||
let mut object = InferredObject::default();
|
||||
while let Some(key) = map.next_key::<String>()? {
|
||||
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<InferredType>),
|
||||
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<DataType> {
|
||||
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::<Self>()
|
||||
+ 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<Arc<str>, usize>,
|
||||
fields: Vec<(Arc<str>, 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<str> = 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<Vec<Field>> {
|
||||
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<str>, usize)>() + 16;
|
||||
let mut bytes = std::mem::size_of::<Self>()
|
||||
.saturating_add(self.indexes.capacity().saturating_mul(index_entry))
|
||||
.saturating_add(
|
||||
self.fields
|
||||
.capacity()
|
||||
.saturating_mul(std::mem::size_of::<(Arc<str>, 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::<serde_json::Error>()
|
||||
.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(_))
|
||||
));
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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?;
|
||||
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));
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user