diff --git a/crates/s3select-query/Cargo.toml b/crates/s3select-query/Cargo.toml index 3fbf7778c..b786f2197 100644 --- a/crates/s3select-query/Cargo.toml +++ b/crates/s3select-query/Cargo.toml @@ -57,8 +57,6 @@ tracing = { workspace = true } [dev-dependencies] rustfs-test-utils = { workspace = true, features = ["put-object-commit-barrier"] } -# Match the streaming inference's field order in Arrow's JSON reference. -serde_json = { workspace = true, features = ["preserve_order"] } [lib] doctest = false diff --git a/crates/s3select-query/src/dispatcher/json_document_schema.rs b/crates/s3select-query/src/dispatcher/json_document_schema.rs index 1f8c33c46..0e73b2364 100644 --- a/crates/s3select-query/src/dispatcher/json_document_schema.rs +++ b/crates/s3select-query/src/dispatcher/json_document_schema.rs @@ -426,14 +426,18 @@ impl InferredObject { #[cfg(test)] mod tests { use super::*; + use crate::dispatcher::json_lines_schema; 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]) -> Schema { let records = input.iter().filter(|byte| **byte == b'\n').count(); - let (expected, expected_records) = arrow_infer_json_schema(io::BufReader::new(input), Some(records)) + let (_, expected_records) = arrow_infer_json_schema(io::BufReader::new(input), Some(records)) .expect("Arrow should infer the compatibility fixture"); + // Arrow's types in source key order, which is what JSON LINES uses. + let expected = json_lines_schema::infer_schema(input, records) + .expect("JSON LINES inference should accept 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); diff --git a/crates/s3select-query/src/dispatcher/json_lines_schema.rs b/crates/s3select-query/src/dispatcher/json_lines_schema.rs new file mode 100644 index 000000000..806969ce4 --- /dev/null +++ b/crates/s3select-query/src/dispatcher/json_lines_schema.rs @@ -0,0 +1,185 @@ +// 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, FieldRef, Fields, Schema}, + error::ArrowError, + json::reader::{ValueIter, infer_json_schema_from_iterator}, +}; +use serde::de::{self, DeserializeSeed, MapAccess, SeqAccess, Visitor}; +use std::{collections::HashMap, fmt, sync::Arc}; + +/// Infers a JSON LINES schema with Arrow's types and the source's first-seen +/// key order. +/// +/// Arrow walks each record as a `serde_json::Map`, which sorts keys unless the +/// process-wide `serde_json/preserve_order` feature is enabled, so `SELECT *` +/// would otherwise reorder fields and CSV output columns. +pub(super) fn infer_schema(bytes: &[u8], max_records: usize) -> Result { + let schema = infer_json_schema_from_iterator(ValueIter::new(bytes, Some(max_records)))?; + let mut order = KeyOrder::default(); + // Same records as `ValueIter`: one value per line, blank lines skipped. + // Arrow has already rejected invalid UTF-8 and JSON within this prefix. + let records = bytes + .split(|byte| *byte == b'\n') + .filter_map(|line| std::str::from_utf8(line).ok()) + .map(str::trim) + .filter(|line| !line.is_empty()) + .take(max_records); + for record in records { + let mut deserializer = serde_json::Deserializer::from_str(record); + (&mut order) + .deserialize(&mut deserializer) + .map_err(|error| ArrowError::JsonError(format!("Not valid JSON: {error}")))?; + } + Ok(Schema::new_with_metadata(order.reorder(schema.fields()), schema.metadata().clone())) +} + +/// First-seen key order of a JSON object, merged across records. Array +/// elements share their parent's node, matching Arrow's `List` shape. +#[derive(Default)] +struct KeyOrder { + positions: HashMap, + children: Vec, +} + +impl KeyOrder { + fn child(&mut self, key: String) -> &mut KeyOrder { + let next = self.children.len(); + let position = *self.positions.entry(key).or_insert(next); + if position == next { + self.children.push(KeyOrder::default()); + } + &mut self.children[position] + } + + fn reorder(&self, fields: &Fields) -> Fields { + let mut fields: Vec = fields + .iter() + .map(|field| match self.positions.get(field.name()) { + Some(&position) => { + let data_type = self.children[position].reorder_type(field.data_type()); + Arc::new(field.as_ref().clone().with_data_type(data_type)) + } + None => Arc::clone(field), + }) + .collect(); + fields.sort_by_key(|field| self.positions.get(field.name()).copied().unwrap_or(usize::MAX)); + fields.into() + } + + fn reorder_type(&self, data_type: &DataType) -> DataType { + match data_type { + DataType::Struct(fields) => DataType::Struct(self.reorder(fields)), + DataType::List(item) => { + DataType::List(Arc::new(item.as_ref().clone().with_data_type(self.reorder_type(item.data_type())))) + } + other => other.clone(), + } + } +} + +impl<'de> DeserializeSeed<'de> for &mut KeyOrder { + type Value = (); + + fn deserialize>(self, deserializer: D) -> Result<(), D::Error> { + deserializer.deserialize_any(self) + } +} + +impl<'de> Visitor<'de> for &mut KeyOrder { + type Value = (); + + fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + formatter.write_str("a JSON value") + } + + fn visit_map>(self, mut map: A) -> Result<(), A::Error> { + while let Some(key) = map.next_key::()? { + map.next_value_seed(self.child(key))?; + } + Ok(()) + } + + fn visit_seq>(self, mut seq: A) -> Result<(), A::Error> { + while seq.next_element_seed(&mut *self)?.is_some() {} + Ok(()) + } + + fn visit_bool(self, _: bool) -> Result<(), E> { + Ok(()) + } + + fn visit_i64(self, _: i64) -> Result<(), E> { + Ok(()) + } + + fn visit_u64(self, _: u64) -> Result<(), E> { + Ok(()) + } + + fn visit_f64(self, _: f64) -> Result<(), E> { + Ok(()) + } + + fn visit_str(self, _: &str) -> Result<(), E> { + Ok(()) + } + + fn visit_unit(self) -> Result<(), E> { + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::infer_schema; + use datafusion::arrow::datatypes::{DataType, Fields}; + + fn names(fields: &Fields) -> Vec<&str> { + fields.iter().map(|field| field.name().as_str()).collect() + } + + #[test] + fn keeps_first_seen_key_order_at_every_level() { + let input = b"{\"z\":1,\"a\":{\"y\":[{\"k2\":1,\"k1\":2}],\"b\":2}}\n\n{\"m\":null,\"a\":{\"c\":3}}\n"; + let schema = infer_schema(input, 1000).expect("schema"); + + assert_eq!(names(schema.fields()), ["z", "a", "m"]); + let DataType::Struct(a) = schema.field(1).data_type() else { + panic!("a should be a struct"); + }; + assert_eq!(names(a), ["y", "b", "c"]); + let DataType::List(item) = a[0].data_type() else { + panic!("y should be a list"); + }; + let DataType::Struct(y) = item.data_type() else { + panic!("y items should be structs"); + }; + assert_eq!(names(y), ["k2", "k1"]); + } + + #[test] + fn ignores_records_after_the_inference_prefix_like_arrow() { + let schema = infer_schema(b"{\"b\":1,\"a\":2}\nnot json\n", 1).expect("schema"); + + assert_eq!(names(schema.fields()), ["b", "a"]); + } + + #[test] + fn keeps_arrow_errors_for_invalid_records() { + assert!(infer_schema(b"{\"a\":1}\n[1]\n", 1000).is_err()); + assert!(infer_schema(b"{\"a\":\n", 1000).is_err()); + } +} diff --git a/crates/s3select-query/src/dispatcher/manager.rs b/crates/s3select-query/src/dispatcher/manager.rs index 06e21393c..d6d0db5e6 100644 --- a/crates/s3select-query/src/dispatcher/manager.rs +++ b/crates/s3select-query/src/dispatcher/manager.rs @@ -28,11 +28,12 @@ use datafusion::{ record_batch::RecordBatch, }, datasource::{ - file_format::{csv::CsvFormat, json::JsonFormat}, + file_format::{DEFAULT_SCHEMA_INFER_MAX_RECORD, csv::CsvFormat, json::JsonFormat}, listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl}, }, error::Result as DFResult, execution::{RecordBatchStream, SendableRecordBatchStream}, + object_store::ObjectStoreExt as _, sql::sqlparser::parser::ParserError, }; use futures::Stream; @@ -61,7 +62,7 @@ use tokio::{ }; use crate::{ - dispatcher::{json_document_table::JsonDocumentTable, parquet_table::ParquetSelectTable}, + dispatcher::{json_document_table::JsonDocumentTable, json_lines_schema, parquet_table::ParquetSelectTable}, execution::factory::QueryExecutionFactoryRef, instance::{DEFAULT_MAX_CONCURRENT_QUERIES, DEFAULT_QUERY_TIMEOUT_SECS}, metadata::{ContextProviderExtension, MetadataProvider, TableHandleProviderRef, base_table::BaseTableProvider}, @@ -508,7 +509,11 @@ impl SimpleQueryDispatcher { return Err(SelectError::InvalidDataSource.into()); }; - let resolve_schema = listing_options.infer_schema(session.inner(), &table_path).await?; + let resolve_schema = if self.input.request.input_serialization.json.is_some() { + infer_json_lines_schema(session, &table_path).await? + } else { + listing_options.infer_schema(session.inner(), &table_path).await? + }; let config = if need_rename_volume_name { let mut new_fields = Vec::new(); for (i, field) in resolve_schema.fields().iter().enumerate() { @@ -549,6 +554,14 @@ impl SimpleQueryDispatcher { } } +/// Reads the object like DataFusion's `JsonFormat::infer_schema` does, but +/// keeps the source key order (see `json_lines_schema::infer_schema`). +async fn infer_json_lines_schema(session: &SessionCtx, table_path: &ListingTableUrl) -> DFResult { + let store = session.inner().runtime_env().object_store(table_path.object_store())?; + let bytes = store.get(table_path.prefix()).await?.bytes().await?; + Ok(Arc::new(json_lines_schema::infer_schema(&bytes, DEFAULT_SCHEMA_INFER_MAX_RECORD)?)) +} + fn is_json_document_input(input: &SelectObjectContentInput) -> bool { input .request @@ -856,6 +869,7 @@ mod tests { array::{Int32Array, StringArray}, datatypes::{DataType, Field, Schema, SchemaRef}, record_batch::RecordBatch, + util::display::array_value_to_string, }, common::DataFusionError, execution::object_store::ObjectStoreUrl, @@ -1727,6 +1741,334 @@ mod tests { } } + async fn json_lines_plan_schema(data: &'static [u8], sql: &str) -> SchemaRef { + let mut input = test_input(); + input.key = "records.json".to_owned(); + input.request.input_serialization.csv = None; + input.request.input_serialization.json = Some(JSONInput { + type_: Some(JSONType::from_static(JSONType::LINES)), + }); + let input = Arc::new(input); + let optimizer = Arc::new(CascadeOptimizerBuilder::default().build()); + let dispatcher = test_dispatcher_for_input( + Arc::clone(&input), + Arc::new(Semaphore::new(1)), + Duration::from_secs(30), + Arc::new(SqlQueryExecutionFactory::new(optimizer, Arc::new(LocalScheduler {}))), + ); + let query = Query::new(QueryContext { input }, sql.to_owned()); + let machine = dispatcher.build_query_state_machine(query).await.expect("build query state"); + let store_url = ObjectStoreUrl::parse("s3://test-bucket").expect("test store URL"); + let store = machine + .session + .inner() + .runtime_env() + .object_store(&store_url) + .expect("test store"); + store + .put(&Path::from("records.json"), data.into()) + .await + .expect("write JSON LINES object"); + let Plan::Query(plan) = dispatcher + .build_logical_plan(machine) + .await + .expect("plan JSON LINES query") + .expect("select plan"); + Arc::new(plan.df_plan.schema().as_arrow().clone()) + } + + fn field_names(fields: &datafusion::arrow::datatypes::Fields) -> Vec<&str> { + fields.iter().map(|field| field.name().as_str()).collect() + } + + fn struct_fields(data_type: &DataType) -> &datafusion::arrow::datatypes::Fields { + match data_type { + DataType::Struct(fields) => fields, + DataType::List(item) => struct_fields(item.data_type()), + other => panic!("expected a struct or list of structs, got {other}"), + } + } + + #[tokio::test] + async fn json_lines_select_star_keeps_source_field_order() { + let schema = json_lines_plan_schema( + b"{\"z\":1,\"a b\":{\"y\":true,\"B\":[{\"k2\":1,\"k1\":2}]},\"A\":\"x\"}\n\n \r\n{\"A\":\"y\",\"$c\":2,\"a b\":{\"x\":1}}\r\n", + "SELECT * FROM S3Object", + ) + .await; + + assert_eq!(field_names(schema.fields()), ["z", "a b", "A", "$c"]); + let nested = struct_fields(schema.field_with_name("a b").expect("nested field").data_type()); + assert_eq!(field_names(nested), ["y", "B", "x"]); + let listed = struct_fields(nested.find("B").expect("list field").1.data_type()); + assert_eq!(field_names(listed), ["k2", "k1"]); + } + + #[tokio::test] + async fn path_projections_are_named_after_their_last_key() { + const DATA: &[u8] = b"{\"obj0\":{\"id\":\"id0\",\"inner\":{\"Key\":1}}}\n"; + for (sql, expected) in [ + ("SELECT s.obj0.id, s.obj0.inner.\"Key\" FROM S3Object s", ["id", "Key"].as_slice()), + ("SELECT obj0.id FROM S3Object", &["id"]), + ("SELECT s.obj0.id FROM S3Object s ORDER BY s.obj0.id", &["id"]), + ("SELECT s.obj0.id, COUNT(*) FROM S3Object s GROUP BY s.obj0.id", &["id", "count(*)"]), + ] { + let schema = json_lines_plan_schema(DATA, sql).await; + assert_eq!(field_names(schema.fields()), expected, "{sql}"); + } + } + + #[tokio::test] + async fn path_projection_names_do_not_clash_with_other_columns() { + const DATA: &[u8] = b"{\"id\":\"top\",\"obj0\":{\"id\":\"id0\"}}\n"; + for sql in ["SELECT *, s.obj0.id FROM S3Object s", "SELECT s.*, s.obj0.id FROM S3Object s"] { + let schema = json_lines_plan_schema(DATA, sql).await; + assert_eq!(field_names(schema.fields()), ["id", "obj0", "s.obj0[id]"], "{sql}"); + } + for sql in [ + "SELECT s.obj0.id FROM S3Object s ORDER BY s.id", + "SELECT s.obj0.id FROM S3Object s ORDER BY id", + "SELECT s.obj0.id FROM S3Object s ORDER BY upper(s.id)", + ] { + let schema = json_lines_plan_schema(DATA, sql).await; + assert_eq!(field_names(schema.fields()), ["s.obj0[id]"], "{sql}"); + } + let schema = json_lines_plan_schema(DATA, "SELECT s.obj0.id, (s.id) FROM S3Object s").await; + assert_eq!(field_names(schema.fields()), ["s.obj0[id]", "id"]); + } + + #[tokio::test] + async fn path_aliases_do_not_collide_with_retained_path_names() { + const DATA: &[u8] = br#"{"id":"top","obj0":{"id":"a"},"obj1":{"id":"b"},"obj2":{"s.obj0[id]":"c","obj0":"d"}}"#; + for (sql, expected) in [ + ( + "SELECT s.obj0.id, s.obj1.id, s.obj2.\"s.obj0[id]\" FROM S3Object s", + vec!["s.obj0[id]", "s.obj1[id]", "s.obj2[s.obj0[id]]"], + ), + ( + "SELECT s.obj2.\"s.obj0[id]\", s.obj1.id, s.obj0.id FROM S3Object s", + vec!["s.obj2[s.obj0[id]]", "s.obj1[id]", "s.obj0[id]"], + ), + ("SELECT (s.obj0.id), s.obj2.\"s.obj0[id]\" FROM S3Object s", vec!["id", "s.obj0[id]"]), + ( + "SELECT s.obj0.id, s.obj2.\"s.obj0[id]\" FROM S3Object s ORDER BY s.id", + vec!["s.obj0[id]", "s.obj2[s.obj0[id]]"], + ), + ("SELECT s.obj2.obj0 FROM S3Object s ORDER BY s.obj0.id", vec!["s.obj2[obj0]"]), + ("SELECT s.obj2.obj0 FROM S3Object s ORDER BY (s.obj0.id)", vec!["s.obj2[obj0]"]), + ] { + let schema = json_lines_plan_schema(DATA, sql).await; + assert_eq!(field_names(schema.fields()), expected, "{sql}"); + } + } + + const PATH_NAMES_DATA: &[u8] = br#"{"id":"A","x":"B","n":1,"obj":{"lower(s.id)":7,"count(*)":3,"id":"x","ID":"X","s.n":2,"s.n + Int64(1)":3},"obj0":{"id":"a"},"obj1":{"id":"b"},"obj2":{"s.obj0[id]":"c","obj0":"d"},"a":{"id":1},"b":{"id":2}} +"#; + + /// Plans and runs `sql` over a stored JSON LINES object through the + /// production dispatcher. Returns the planned output names and the rows. + async fn production_json_lines_query(bucket: &str, key: &str, sql: &str) -> QueryResult<(Vec, Vec>)> { + let mut input = (*json_snapshot_input()).clone(); + input.bucket = bucket.to_owned(); + input.key = key.to_owned(); + input.request.expression = sql.to_owned(); + let input = Arc::new(input); + let snapshot = snapshot_test_env().await.prepare_select_object_snapshot(bucket, key).await; + let dispatcher = production_dispatcher(Arc::clone(&input)); + let query = Query::new_with_snapshot(QueryContext { input }, sql.to_owned(), snapshot); + let machine = dispatcher.build_query_state_machine(query).await?; + let plan = dispatcher + .build_logical_plan(Arc::clone(&machine)) + .await? + .expect("SELECT should produce a logical plan"); + let Plan::Query(query_plan) = &plan; + let names = query_plan + .df_plan + .schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect(); + let Output::StreamData(stream) = dispatcher.execute_logical_plan(plan, machine).await? else { + panic!("SELECT should return rows"); + }; + let batches: Vec<_> = stream.try_collect().await?; + let rows = batches + .iter() + .flat_map(|batch| { + (0..batch.num_rows()).map(|row| { + batch + .columns() + .iter() + .map(|column| array_value_to_string(column, row).expect("display a result value")) + .collect() + }) + }) + .collect(); + Ok((names, rows)) + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn production_path_alias_does_not_duplicate_a_retained_path_name() { + const BUCKET: &str = "s3select-retained-path-name"; + const KEY: &str = "retained.jsonl"; + let env = snapshot_test_env().await; + env.make_bucket(BUCKET, false).await; + env.put_object_bytes( + BUCKET, + KEY, + br#"{"obj0":{"id":"a"},"obj1":{"id":"b"},"obj2":{"s.obj0[id]":"c"}} +"# + .to_vec(), + ) + .await; + let (names, rows) = + production_json_lines_query(BUCKET, KEY, "SELECT s.obj0.id, s.obj1.id, s.obj2.\"s.obj0[id]\" FROM S3Object s") + .await + .expect("the query plans and runs"); + assert_eq!(names, ["s.obj0[id]", "s.obj1[id]", "s.obj2[s.obj0[id]]"]); + assert_eq!(rows, [["a", "b", "c"]]); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn production_path_projections_take_their_key_only_when_it_is_a_unique_output_name() { + const BUCKET: &str = "s3select-path-projection-names"; + const KEY: &str = "names.jsonl"; + let env = snapshot_test_env().await; + env.make_bucket(BUCKET, false).await; + env.put_object_bytes(BUCKET, KEY, PATH_NAMES_DATA.to_vec()).await; + let cases: &[(&str, &[&str])] = &[ + ( + "SELECT s.obj.\"lower(s.id)\", LOWER(s.id) FROM S3Object s", + &["s.obj[lower(s.id)]", "lower(s.id)"], + ), + ( + "SELECT LOWER(s.id), s.obj.\"lower(s.id)\" FROM S3Object s", + &["lower(s.id)", "s.obj[lower(s.id)]"], + ), + ( + "SELECT s.obj.\"count(*)\", COUNT(*) FROM S3Object s GROUP BY s.obj.\"count(*)\"", + &["s.obj[count(*)]", "count(*)"], + ), + ("SELECT s.obj0.id, COUNT(*) FROM S3Object s GROUP BY s.obj0.id", &["id", "count(*)"]), + ("SELECT s.obj0.id, LOWER(s.id) FROM S3Object s", &["id", "lower(s.id)"]), + ("SELECT s.obj0.id, s.x AS id FROM S3Object s", &["s.obj0[id]", "id"]), + ("SELECT s.x AS id, s.obj0.id FROM S3Object s", &["id", "s.obj0[id]"]), + ( + "SELECT *, s.obj0.id FROM S3Object s", + &["id", "x", "n", "obj", "obj0", "obj1", "obj2", "a", "b", "s.obj0[id]"], + ), + ("SELECT s.obj.ID, s.id FROM S3Object s", &["s.obj[id]", "id"]), + ("SELECT s.obj.\"ID\", s.id FROM S3Object s", &["ID", "id"]), + ("SELECT s.obj.\"s.n\", CAST(s.n AS INT) FROM S3Object s", &["s.obj[s.n]", "s.n"]), + ( + "SELECT s.n + 1, s.obj.\"s.n + Int64(1)\" FROM S3Object s", + &["s.n + Int64(1)", "s.obj[s.n + Int64(1)]"], + ), + ("SELECT s.obj0.id FROM S3Object s ORDER BY s.id", &["s.obj0[id]"]), + ("SELECT s.obj0.id FROM S3Object s ORDER BY s.obj0.id", &["id"]), + ("SELECT (s.obj0.id) FROM S3Object s", &["id"]), + ("SELECT s.obj0.id, (s.id) FROM S3Object s", &["s.obj0[id]", "id"]), + ("SELECT s.a.id, s.b.id FROM S3Object s", &["s.a[id]", "s.b[id]"]), + ("SELECT s.a.id, s.b.id, s.obj2.obj0 FROM S3Object s", &["s.a[id]", "s.b[id]", "obj0"]), + ("SELECT s.obj2.\"s.obj0[id]\", s.obj0.id FROM S3Object s", &["s.obj0[id]", "id"]), + ( + "SELECT s.obj2.\"s.obj0[id]\", s.obj1.id, s.obj0.id FROM S3Object s", + &["s.obj2[s.obj0[id]]", "s.obj1[id]", "s.obj0[id]"], + ), + ( + "SELECT s.obj2.\"s.obj0[id]\", s.obj0.id, LOWER(s.id) FROM S3Object s", + &["s.obj0[id]", "id", "lower(s.id)"], + ), + ]; + for (sql, expected) in cases { + let (names, _) = production_json_lines_query(BUCKET, KEY, sql) + .await + .unwrap_or_else(|error| panic!("{sql}: {error}")); + assert_eq!(names, *expected, "{sql}"); + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn production_path_projection_naming_keeps_queries_planning() { + const BUCKET: &str = "s3select-path-projection-planning"; + const KEY: &str = "planning.jsonl"; + let env = snapshot_test_env().await; + env.make_bucket(BUCKET, false).await; + env.put_object_bytes(BUCKET, KEY, PATH_NAMES_DATA.to_vec()).await; + env.put_object_bytes(BUCKET, "no-top-level-id.jsonl", br#"{"obj0":{"id":"a"},"obj1":{"id":"b"}}"#.to_vec()) + .await; + let mut failures = Vec::new(); + // GROUP BY resolves a SELECT alias when no input column has its name. + for sql in [ + "SELECT s.obj0.id FROM S3Object s GROUP BY id", + "SELECT s.obj0.id, s.obj1.id AS k FROM S3Object s GROUP BY id, k", + ] { + if let Err(error) = production_json_lines_query(BUCKET, "no-top-level-id.jsonl", sql).await { + failures.push(format!("{sql}: {error}")); + } + } + // Each query also plans when every path projection keeps DataFusion's name. + for sql in [ + "SELECT s.obj0.id FROM S3Object s", + "SELECT obj0.id FROM S3Object", + "SELECT s.obj0.id FROM S3Object s WHERE s.obj0.id = 'a'", + "SELECT s.obj0.id FROM S3Object s WHERE s.obj0.id IN ('a', 'b') LIMIT 1", + "SELECT s.obj0.id FROM S3Object s LIMIT 0", + "SELECT s.obj0.id FROM S3Object s ORDER BY s.obj0.id DESC LIMIT 1", + "SELECT s.obj0.id FROM S3Object s ORDER BY id", + "SELECT s.obj0.id FROM S3Object s ORDER BY s.id", + "SELECT s.obj0.id FROM S3Object s ORDER BY upper(s.id)", + "SELECT s.obj0.id AS k FROM S3Object s ORDER BY k", + "SELECT s.obj0.id, s.obj2.obj0 FROM S3Object s ORDER BY s.obj0.id", + "SELECT s.obj2.obj0 FROM S3Object s ORDER BY (s.obj0.id)", + "SELECT s.obj0.id, s.id FROM S3Object s ORDER BY 2", + "SELECT s.obj0.id FROM S3Object s GROUP BY s.obj0.id", + "SELECT s.obj0.id, COUNT(*) FROM S3Object s GROUP BY s.obj0.id", + "SELECT s.obj0.id, COUNT(*) AS id FROM S3Object s GROUP BY s.obj0.id", + "SELECT COUNT(*), s.obj0.id FROM S3Object s GROUP BY s.obj0.id ORDER BY COUNT(*)", + "SELECT COUNT(s.obj0.id) FROM S3Object s", + "SELECT MAX(s.obj0.id) FROM S3Object s", + "SELECT UPPER(s.obj0.id), s.obj0.id FROM S3Object s", + "SELECT s.obj0.id, s.obj0.id = 'a' FROM S3Object s", + "SELECT s.obj0.id, s.obj0.id AS x FROM S3Object s", + "SELECT s.obj0.id, s.id FROM S3Object s", + "SELECT s.obj0.id, (s.id) FROM S3Object s", + "SELECT (s.obj0.id), s.obj2.\"s.obj0[id]\" FROM S3Object s", + "SELECT s.obj0.id AS id, s.a.id FROM S3Object s", + "SELECT s.obj0.id, s.x AS id FROM S3Object s", + "SELECT s.obj0.id AS \"s.a[id]\", s.b.id FROM S3Object s", + "SELECT s.obj.ID, s.id FROM S3Object s", + "SELECT s.obj.\"ID\", LOWER(s.id) FROM S3Object s", + "SELECT s.obj.id, s.obj0.id FROM S3Object s", + "SELECT s.a.id, s.b.id FROM S3Object s", + "SELECT s.a.id, s.b.id, s.obj2.obj0 FROM S3Object s", + "SELECT *, s.obj0.id FROM S3Object s", + "SELECT s.*, s.obj0.id FROM S3Object s", + "SELECT s.obj.\"lower(s.id)\", LOWER(s.id) FROM S3Object s", + "SELECT LOWER(s.id), s.obj.\"lower(s.id)\" FROM S3Object s", + "SELECT s.obj.\"count(*)\", COUNT(*) FROM S3Object s GROUP BY s.obj.\"count(*)\"", + "SELECT s.obj.\"s.n\", CAST(s.n AS INT) FROM S3Object s", + "SELECT s.obj.\"s.n + Int64(1)\", s.n + 1 FROM S3Object s", + "SELECT s.obj0.id, CAST(s.n AS INT) FROM S3Object s", + "SELECT s.obj0.id, s.obj1.id, s.obj2.\"s.obj0[id]\" FROM S3Object s", + "SELECT s.obj2.\"s.obj0[id]\", s.obj1.id, s.obj0.id FROM S3Object s", + "SELECT s.obj2.\"s.obj0[id]\", s.obj0.id FROM S3Object s", + "SELECT s.obj2.\"s.obj0[id]\", s.obj0.id, LOWER(s.id) FROM S3Object s", + "SELECT (s.obj0.id), (s.obj1.id) FROM S3Object s", + "SELECT s.obj.\"ID\", s.obj0.id FROM S3Object s", + ] { + match production_json_lines_query(BUCKET, KEY, sql).await { + Ok((names, _)) if names.iter().collect::>().len() != names.len() => { + failures.push(format!("{sql}: duplicate names {names:?}")); + } + Ok(_) => {} + Err(error) => failures.push(format!("{sql}: {error}")), + } + } + assert!(failures.is_empty(), "{failures:#?}"); + } + #[tokio::test] async fn csv_query_uses_custom_record_delimiter_across_file_partitions() { const ROW_COUNT: usize = 200_000; @@ -2841,4 +3183,110 @@ mod tests { assert_eq!(admission.available_permits(), 1); assert!(output.next().await.is_none()); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn path_alias_does_not_collide_with_computed_output_name() { + let mut input = (*json_snapshot_input()).clone(); + input.bucket = "s3select-review-name-collision".to_string(); + input.key = "collision.jsonl".to_string(); + input.request.expression = "SELECT s.obj.\"lower(s.id)\", LOWER(s.id) FROM S3Object s".to_string(); + 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, b"{\"id\":\"A\",\"obj\":{\"lower(s.id)\":7}}\n".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 machine = dispatcher.build_query_state_machine(query).await.expect("production session"); + let result = dispatcher.build_logical_plan(Arc::clone(&machine)).await; + if let Err(error) = &result { + panic!( + "real production planner rejected a previously valid query: {error}; S3 classification={:?}", + error.s3_select_policy_error() + ); + } + let plan = result.expect("valid production query").expect("query plan"); + let Plan::Query(query_plan) = &plan; + let schema = query_plan.df_plan.schema(); + assert_eq!(schema.fields().len(), 2); + assert_ne!(schema.field(0).name(), schema.field(1).name()); + let output = dispatcher + .execute_logical_plan(plan, machine) + .await + .expect("execute production query"); + let Output::StreamData(stream) = output else { + panic!("query must return rows"); + }; + let batches: Vec<_> = stream.try_collect().await.expect("read production result"); + assert_eq!(batches.iter().map(|batch| batch.num_rows()).sum::(), 1); + let batch = batches.iter().find(|batch| batch.num_rows() == 1).expect("one result row"); + let key = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("integer key"); + let lower = batch + .column(1) + .as_any() + .downcast_ref::() + .expect("lowercase string"); + assert_eq!(key.value(0), 7); + assert_eq!(lower.value(0), "a"); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn path_alias_does_not_collide_with_computed_sort_name() { + let mut input = (*json_snapshot_input()).clone(); + input.bucket = "s3select-sort-name-collision".to_string(); + input.key = "collision.jsonl".to_string(); + input.request.expression = "SELECT s.obj.\"lower(s.id)\" FROM S3Object s ORDER BY LOWER(s.id)".to_string(); + 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, b"{\"id\":\"A\",\"obj\":{\"lower(s.id)\":7}}\n".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 machine = dispatcher.build_query_state_machine(query).await.expect("production session"); + let result = dispatcher.build_logical_plan(Arc::clone(&machine)).await; + if let Err(error) = &result { + panic!( + "real production planner rejected a previously valid query: {error}; S3 classification={:?}", + error.s3_select_policy_error() + ); + } + let plan = result.expect("valid production query").expect("query plan"); + let Plan::Query(query_plan) = &plan; + let schema = query_plan.df_plan.schema(); + assert_eq!(schema.fields().len(), 1); + let output = dispatcher + .execute_logical_plan(plan, machine) + .await + .expect("execute production query"); + let Output::StreamData(stream) = output else { + panic!("query must return rows"); + }; + let batches: Vec<_> = stream.try_collect().await.expect("read production result"); + assert_eq!(batches.iter().map(|batch| batch.num_rows()).sum::(), 1); + let batch = batches.iter().find(|batch| batch.num_rows() == 1).expect("one result row"); + let key = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("integer key"); + assert_eq!(key.value(0), 7); + } } diff --git a/crates/s3select-query/src/dispatcher/mod.rs b/crates/s3select-query/src/dispatcher/mod.rs index 4b9d89930..04121b55d 100644 --- a/crates/s3select-query/src/dispatcher/mod.rs +++ b/crates/s3select-query/src/dispatcher/mod.rs @@ -14,5 +14,6 @@ mod json_document_schema; mod json_document_table; +mod json_lines_schema; pub mod manager; mod parquet_table; diff --git a/crates/s3select-query/src/sql/planner.rs b/crates/s3select-query/src/sql/planner.rs index 57feec67c..cd0dbbc13 100644 --- a/crates/s3select-query/src/sql/planner.rs +++ b/crates/s3select-query/src/sql/planner.rs @@ -12,14 +12,19 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::{convert::Infallible, ops::ControlFlow}; +use std::{collections::HashMap, convert::Infallible, ops::ControlFlow}; use async_trait::async_trait; -use datafusion::sql::{ - planner::{IdentNormalizer, SqlToRel}, - sqlparser::ast::{ - AccessExpr, Expr, GroupByExpr, Ident, JsonPath, JsonPathElem, ObjectNamePart, OrderByKind, Query, Select, SelectFlavor, - SetExpr, Statement, Subscript, TableAlias, TableFactor, Value, Visit, VisitMut, Visitor, VisitorMut, +use datafusion::{ + common::{Column, Result as DFResult}, + logical_expr::{Expr as LogicalExpr, LogicalPlan, LogicalPlanBuilder}, + sql::{ + planner::{IdentNormalizer, SqlToRel}, + sqlparser::ast::{ + AccessExpr, Expr, GroupByExpr, Ident, JsonPath, JsonPathElem, ObjectNamePart, OrderBy, OrderByKind, Query, Select, + SelectFlavor, SelectItem, SetExpr, Statement, Subscript, TableAlias, TableFactor, Value, Visit, VisitMut, Visitor, + VisitorMut, + }, }, }; use rustfs_s3select_api::{ @@ -75,7 +80,9 @@ impl<'a, S: ContextProviderExtension + Send + Sync + 'a> SqlPlanner<'a, S> { } async fn df_prepared_sql_to_plan(&self, stmt: Statement, _session: &SessionCtx) -> QueryResult { + let keys = path_projection_keys(&stmt); let df_plan = self.df_planner.sql_statement_to_plan(stmt).map_err(classify_planner_error)?; + let df_plan = name_path_projections(df_plan, &keys)?; Ok(Plan::Query(QueryPlan { df_plan, is_tag_scan: false, @@ -150,9 +157,203 @@ pub(crate) fn prepare_s3_select_statement(statement: &mut Statement) -> QueryRes }; let mut normalizer = PartiQlSubscriptNormalizer; let _ = VisitMut::visit(query, &mut normalizer); + if let SetExpr::Select(select) = query.body.as_mut() { + alias_path_projections(select, query.order_by.as_ref()); + } Ok(source) } +/// Names an unaliased path projection after its last key, as AWS does +/// (`s.obj0.id` -> `id`), instead of DataFusion's `s.obj0[id]`. Names that +/// could collide with another output name keep DataFusion's name. Wildcard +/// columns, computed expressions, and retained paths have unknown output names +/// here, so they disable aliasing. Unprojected ORDER BY expressions also disable +/// aliasing because DataFusion adds their input columns to the projection. +/// GROUP BY can refer to these aliases; [`name_path_projections`] names the +/// paths left here once the planned output names are known. +fn alias_path_projections(select: &mut Select, order_by: Option<&OrderBy>) { + if has_wildcard(select) || sorts_by_unprojected_expr(select, order_by) { + return; + } + let Some(qualifier) = source_qualifier(select) else { + return; + }; + let normalizer = IdentNormalizer::default(); + let names: Vec> = select + .projection + .iter() + .map(|item| { + let ident = match item { + SelectItem::UnnamedExpr(Expr::Nested(_)) => None, + SelectItem::UnnamedExpr(expr) => column_name(expr), + SelectItem::ExprWithAlias { alias, .. } => Some(alias), + _ => None, + }; + ident.map(|ident| normalizer.normalize(ident.clone())) + }) + .collect(); + if names.iter().any(Option::is_none) { + return; + } + let mut counts = HashMap::::new(); + for name in names.iter().flatten() { + *counts.entry(name.clone()).or_default() += 1; + } + // Retained paths have generated names that cannot be inferred from their last key. + // Keep the whole projection unchanged so a sibling alias cannot collide with them. + if names.iter().flatten().any(|name| counts.get(name) != Some(&1)) { + return; + } + for item in &mut select.projection { + let SelectItem::UnnamedExpr(expr) = item else { + continue; + }; + let Some(alias) = path_key(expr, &qualifier).cloned() else { + continue; + }; + let expr = std::mem::replace(expr, Expr::Identifier(Ident::new(""))); + *item = SelectItem::ExprWithAlias { expr, alias }; + } +} + +/// The output name DataFusion gives a column reference: its last part. +fn column_name(expr: &Expr) -> Option<&Ident> { + match expr { + Expr::Identifier(ident) => Some(ident), + Expr::CompoundIdentifier(parts) => parts.last(), + Expr::Nested(inner) => column_name(inner), + _ => None, + } +} + +fn has_wildcard(select: &Select) -> bool { + select + .projection + .iter() + .any(|item| matches!(item, SelectItem::Wildcard(_) | SelectItem::QualifiedWildcard(..))) +} + +fn sorts_by_unprojected_expr(select: &Select, order_by: Option<&OrderBy>) -> bool { + let Some(OrderBy { + kind: OrderByKind::Expressions(sort_exprs), + .. + }) = order_by + else { + return false; + }; + sort_exprs.iter().any(|sort_expr| { + !select + .projection + .iter() + .any(|item| matches!(item, SelectItem::UnnamedExpr(projected) if *projected == sort_expr.expr)) + }) +} + +/// The normalized table alias, or table name, that may prefix a path. +fn source_qualifier(select: &Select) -> Option { + let normalizer = IdentNormalizer::default(); + match &select.from.first()?.relation { + TableFactor::Table { alias: Some(alias), .. } => Some(normalizer.normalize(alias.name.clone())), + TableFactor::Table { name, .. } => match name.0.first() { + Some(ObjectNamePart::Identifier(table_name)) => Some(normalizer.normalize(table_name.clone())), + _ => None, + }, + _ => None, + } +} + +/// The last key of a path of at least two keys after the source qualifier +/// (`s.obj0.id` -> `id`), the name AWS gives the projection. +fn path_key<'a>(expr: &'a Expr, qualifier: &str) -> Option<&'a Ident> { + let parts = match expr { + Expr::CompoundIdentifier(parts) => parts, + Expr::Nested(inner) => return path_key(inner, qualifier), + _ => return None, + }; + let path = match parts.split_first() { + Some((first, rest)) if IdentNormalizer::default().normalize(first.clone()) == qualifier => rest, + _ => parts.as_slice(), + }; + match path { + [_, _, ..] => path.last(), + _ => None, + } +} + +/// The normalized key of each path projection left unaliased by +/// [`alias_path_projections`]. Empty when a wildcard hides the output +/// positions, or when an ORDER BY expression is not projected. +fn path_projection_keys(statement: &Statement) -> Vec> { + let Statement::Query(query) = statement else { + return Vec::new(); + }; + let SetExpr::Select(select) = query.body.as_ref() else { + return Vec::new(); + }; + if has_wildcard(select) || sorts_by_unprojected_expr(select, query.order_by.as_ref()) { + return Vec::new(); + } + let Some(qualifier) = source_qualifier(select) else { + return Vec::new(); + }; + let normalizer = IdentNormalizer::default(); + select + .projection + .iter() + .map(|item| match item { + SelectItem::UnnamedExpr(expr) => path_key(expr, &qualifier).map(|key| normalizer.normalize(key.clone())), + _ => None, + }) + .collect() +} + +/// Renames each path projection to its key from [`path_projection_keys`] +/// unless another output column keeps or receives that name. Names come from +/// the planned schema, so computed, aliased and retained columns all count. +fn name_path_projections(plan: LogicalPlan, keys: &[Option]) -> DFResult { + let schema = plan.schema(); + if keys.len() != schema.fields().len() { + return Ok(plan); + } + let mut renames: Vec> = keys + .iter() + .zip(schema.fields()) + .map(|(key, field)| key.as_deref().filter(|key| key != field.name())) + .collect(); + // Dropping a rename brings back the generated name, which can collide with another rename. + loop { + let mut counts = HashMap::<&str, usize>::new(); + for (rename, field) in renames.iter().zip(schema.fields()) { + *counts.entry(rename.unwrap_or(field.name())).or_default() += 1; + } + let mut dropped = false; + for rename in &mut renames { + if rename.is_some_and(|name| counts[name] > 1) { + *rename = None; + dropped = true; + } + } + if !dropped { + break; + } + } + if renames.iter().all(Option::is_none) { + return Ok(plan); + } + let columns: Vec<_> = schema + .iter() + .zip(renames) + .map(|(column, rename)| { + let column = LogicalExpr::Column(Column::from(column)); + match rename { + Some(name) => column.alias(name), + None => column, + } + }) + .collect(); + LogicalPlanBuilder::from(plan).project(columns)?.build() +} + struct PartiQlSubscriptNormalizer; impl VisitorMut for PartiQlSubscriptNormalizer { @@ -374,7 +575,7 @@ impl Visitor for SubqueryDetector { #[cfg(test)] mod tests { - use super::prepare_s3_select_statement; + use super::{path_projection_keys, prepare_s3_select_statement}; use crate::sql::parser::ExtParser; use datafusion::sql::sqlparser::ast::{AccessExpr, Expr, Statement, Visit, Visitor}; use rustfs_s3select_api::{ @@ -596,6 +797,86 @@ mod tests { } } + fn prepared(sql: &str) -> String { + let mut statement = parse_statement(sql); + prepare_s3_select_statement(&mut statement).expect("statement should be supported"); + statement.to_string() + } + + #[test] + fn unaliased_path_projections_are_named_after_their_last_key() { + for (sql, expected) in [ + ("SELECT s.obj0.id FROM S3Object AS s", "SELECT s.obj0.id AS id FROM S3Object AS s"), + ( + "SELECT s.a.b.c, s.x.\"Y\" FROM S3Object AS s", + "SELECT s.a.b.c AS c, s.x.\"Y\" AS \"Y\" FROM S3Object AS s", + ), + ("SELECT obj0.id FROM S3Object", "SELECT obj0.id AS id FROM S3Object"), + ( + "SELECT s.obj0.id FROM S3Object AS s WHERE s.obj0.id = 'a' ORDER BY s.obj0.id", + "SELECT s.obj0.id AS id FROM S3Object AS s WHERE s.obj0.id = 'a' ORDER BY s.obj0.id", + ), + ] { + assert_eq!(prepared(sql), expected, "{sql}"); + } + } + + #[test] + fn projections_keep_their_names_when_a_key_alias_is_absent_or_ambiguous() { + for sql in [ + "SELECT s.id FROM S3Object AS s", + "SELECT S3Object.id FROM S3Object", + "SELECT s.tags[1] FROM S3Object AS s", + "SELECT s.obj0.id AS x FROM S3Object AS s", + "SELECT s.obj0.id, s.obj1.id FROM S3Object AS s", + "SELECT s.obj0.ID, s.obj1.id FROM S3Object AS s", + "SELECT s.obj0.id, s.id FROM S3Object AS s", + "SELECT s.obj0.id, s.name AS id FROM S3Object AS s", + "SELECT s.obj0.\"Id\", s.obj1.\"Id\" FROM S3Object AS s", + "SELECT *, s.obj0.id FROM S3Object AS s", + "SELECT s.*, s.obj0.id FROM S3Object AS s", + "SELECT s.obj0.id FROM S3Object AS s ORDER BY s.id", + "SELECT s.obj0.id FROM S3Object AS s ORDER BY upper(id)", + "SELECT s.obj0.id, (s.id) FROM S3Object AS s", + "SELECT s.obj.\"lower(s.id)\", LOWER(s.id) FROM S3Object AS s", + "SELECT s.obj.\"count(*)\", COUNT(*) FROM S3Object AS s GROUP BY s.obj.\"count(*)\"", + "SELECT s.obj0.id, COUNT(*) FROM S3Object AS s GROUP BY s.obj0.id", + ] { + assert_eq!(prepared(sql), sql); + } + } + + #[test] + fn path_projection_keys_cover_paths_left_unaliased() { + let keys = |sql: &str| { + let mut statement = parse_statement(sql); + prepare_s3_select_statement(&mut statement).expect("statement should be supported"); + path_projection_keys(&statement) + }; + let key = |key: &str| Some(key.to_owned()); + for (sql, expected) in [ + ("SELECT s.obj0.id FROM S3Object AS s", vec![None]), + ("SELECT s.obj0.id, s.obj1.id FROM S3Object AS s", vec![key("id"), key("id")]), + ( + "SELECT s.a.b.c, s.x.\"Y\", s.obj0.ID, lower(s.id) FROM S3Object AS s", + vec![key("c"), key("Y"), key("id"), None], + ), + ("SELECT obj0.id, COUNT(*) FROM S3Object GROUP BY obj0.id", vec![key("id"), None]), + ("SELECT (s.obj0.id), ((s.obj1.id)) FROM S3Object AS s", vec![key("id"), key("id")]), + ( + "SELECT s.id, s.tags[1], s.obj0.id AS x, (s.id) FROM S3Object AS s", + vec![None, None, None, None], + ), + ("SELECT S3Object.id, lower(id) FROM S3Object", vec![None, None]), + ("SELECT *, s.obj0.id FROM S3Object AS s", vec![]), + ("SELECT s.*, s.obj0.id FROM S3Object AS s", vec![]), + ("SELECT s.obj0.id FROM S3Object AS s ORDER BY s.id", vec![]), + ("SELECT s.obj0.id FROM S3Object AS s ORDER BY upper(id)", vec![]), + ] { + assert_eq!(keys(sql), expected, "{sql}"); + } + } + #[test] fn accepts_group_by_and_order_by() { let statement = parse_statement("SELECT department, COUNT(*) FROM S3Object GROUP BY department ORDER BY department"); diff --git a/rustfs/src/app/select_object.rs b/rustfs/src/app/select_object.rs index 9a51272f5..3ac691a26 100644 --- a/rustfs/src/app/select_object.rs +++ b/rustfs/src/app/select_object.rs @@ -600,9 +600,10 @@ fn validate_select_request(headers: &http::HeaderMap, input: &mut SelectObjectCo if input.request.expression.len() > MAX_SELECT_EXPRESSION_BYTES { return Err(S3Error::new(S3ErrorCode::ExpressionTooLong)); } - if input.request.expression_type.as_str() != ExpressionType::SQL { + let Some(expression_type) = canonical_enum(input.request.expression_type.as_str(), &[ExpressionType::SQL]) else { return Err(S3Error::new(S3ErrorCode::InvalidExpressionType)); - } + }; + input.request.expression_type = ExpressionType::from_static(expression_type); normalize_input_serialization(&mut input.request.input_serialization)?; let compressed_input = is_compressed_input(&input.request.input_serialization); @@ -649,6 +650,15 @@ fn is_compressed_input(input: &InputSerialization) -> bool { .is_some_and(|compression| compression.as_str() != CompressionType::NONE) } +/// Matches an S3 enum value in any ASCII case, as MinIO does, and returns the +/// canonical spelling that downstream code compares against. +fn canonical_enum(value: &str, allowed: &[&'static str]) -> Option<&'static str> { + allowed + .iter() + .copied() + .find(|candidate| candidate.eq_ignore_ascii_case(value)) +} + fn normalize_input_serialization(input: &mut InputSerialization) -> S3Result<()> { let format_count = usize::from(input.csv.is_some()) + usize::from(input.json.is_some()) + usize::from(input.parquet.is_some()); @@ -659,19 +669,20 @@ fn normalize_input_serialization(input: &mut InputSerialization) -> S3Result<()> return Err(S3Error::new(S3ErrorCode::ObjectSerializationConflict)); } - if let Some(compression) = input.compression_type.as_ref() { - match compression.as_str() { - CompressionType::NONE => {} - CompressionType::GZIP | CompressionType::BZIP2 => { - if input.parquet.is_some() { - return Err(S3Error::with_message( - S3ErrorCode::InvalidRequestParameter, - INVALID_REQUEST_PARAMETER_MESSAGE, - )); - } - } - _ => return Err(map_select_error_to_s3(&SelectError::InvalidCompressionFormat)), + if let Some(compression) = input.compression_type.as_mut() { + let Some(canonical) = canonical_enum( + compression.as_str(), + &[CompressionType::NONE, CompressionType::GZIP, CompressionType::BZIP2], + ) else { + return Err(map_select_error_to_s3(&SelectError::InvalidCompressionFormat)); + }; + if canonical != CompressionType::NONE && input.parquet.is_some() { + return Err(S3Error::with_message( + S3ErrorCode::InvalidRequestParameter, + INVALID_REQUEST_PARAMETER_MESSAGE, + )); } + *compression = CompressionType::from_static(canonical); } input .compression_type @@ -681,15 +692,16 @@ fn normalize_input_serialization(input: &mut InputSerialization) -> S3Result<()> let file_header_info = csv .file_header_info .get_or_insert_with(|| FileHeaderInfo::from_static(FileHeaderInfo::NONE)); - if !matches!( + let Some(canonical) = canonical_enum( file_header_info.as_str(), - FileHeaderInfo::NONE | FileHeaderInfo::USE | FileHeaderInfo::IGNORE - ) { + &[FileHeaderInfo::NONE, FileHeaderInfo::USE, FileHeaderInfo::IGNORE], + ) else { return Err(S3Error::with_message( S3ErrorCode::InvalidFileHeaderInfo, "The FileHeaderInfo value is not valid. Only NONE, USE, and IGNORE are supported.", )); - } + }; + *file_header_info = FileHeaderInfo::from_static(canonical); validate_single_byte(csv.comments.as_deref(), S3ErrorCode::InvalidRequestParameter)?; validate_single_character(csv.quote_character.as_deref())?; validate_single_character(csv.quote_escape_character.as_deref())?; @@ -699,9 +711,10 @@ fn normalize_input_serialization(input: &mut InputSerialization) -> S3Result<()> if let Some(json) = input.json.as_mut() { let json_type = json.type_.get_or_insert_with(|| JSONType::from_static(JSONType::LINES)); - if !matches!(json_type.as_str(), JSONType::DOCUMENT | JSONType::LINES) { + let Some(canonical) = canonical_enum(json_type.as_str(), &[JSONType::DOCUMENT, JSONType::LINES]) else { return Err(S3Error::new(S3ErrorCode::InvalidJsonType)); - } + }; + *json_type = JSONType::from_static(canonical); } Ok(()) @@ -716,15 +729,16 @@ fn normalize_output_serialization(output: &mut OutputSerialization) -> S3Result< return Err(S3Error::new(S3ErrorCode::ObjectSerializationConflict)); } - if let Some(csv) = output.csv.as_ref() { + if let Some(csv) = output.csv.as_mut() { validate_single_byte(csv.field_delimiter.as_deref(), S3ErrorCode::InvalidRequestParameter)?; validate_single_byte(csv.quote_character.as_deref(), S3ErrorCode::InvalidRequestParameter)?; validate_single_byte(csv.quote_escape_character.as_deref(), S3ErrorCode::InvalidRequestParameter)?; validate_output_record_delimiter(csv.record_delimiter.as_deref())?; - if let Some(quote_fields) = csv.quote_fields.as_ref() - && !matches!(quote_fields.as_str(), QuoteFields::ALWAYS | QuoteFields::ASNEEDED) - { - return Err(S3Error::new(S3ErrorCode::InvalidQuoteFields)); + if let Some(quote_fields) = csv.quote_fields.as_mut() { + let Some(canonical) = canonical_enum(quote_fields.as_str(), &[QuoteFields::ALWAYS, QuoteFields::ASNEEDED]) else { + return Err(S3Error::new(S3ErrorCode::InvalidQuoteFields)); + }; + *quote_fields = QuoteFields::from_static(canonical); } return Ok(SelectOutputFormat::Csv(csv.clone())); } @@ -3639,6 +3653,132 @@ mod tests { ); } + fn json_input(json_type: &str) -> SelectObjectContentInput { + let mut input = base_input(); + input.request.input_serialization.csv = None; + input.request.input_serialization.json = Some(JSONInput { + type_: Some(JSONType::from(json_type.to_string())), + }); + input + } + + #[test] + fn validate_matches_enum_values_case_insensitively_and_canonicalizes_them() { + for (value, expected) in [ + ("gzip", CompressionType::GZIP), + ("bZip2", CompressionType::BZIP2), + ("none", CompressionType::NONE), + ] { + let mut input = base_input(); + input.request.input_serialization.compression_type = Some(CompressionType::from(value.to_string())); + validate_select_request(&HeaderMap::new(), &mut input).expect("CompressionType should match any case"); + assert_eq!( + input + .request + .input_serialization + .compression_type + .as_ref() + .map(|v| v.as_str()), + Some(expected) + ); + } + + for (value, expected) in [("lines", JSONType::LINES), ("Document", JSONType::DOCUMENT)] { + let mut input = json_input(value); + validate_select_request(&HeaderMap::new(), &mut input).expect("JSON Type should match any case"); + let json_type = input + .request + .input_serialization + .json + .as_ref() + .and_then(|json| json.type_.as_ref()); + assert_eq!(json_type.map(|v| v.as_str()), Some(expected)); + } + let mut input = json_input("document"); + validate_select_request(&HeaderMap::new(), &mut input).expect("lowercase DOCUMENT should be accepted"); + assert!(is_json_document(input.request.input_serialization.json.as_ref().expect("JSON input"))); + + for (value, expected) in [ + ("use", FileHeaderInfo::USE), + ("Ignore", FileHeaderInfo::IGNORE), + ("none", FileHeaderInfo::NONE), + ] { + let mut input = base_input(); + input + .request + .input_serialization + .csv + .as_mut() + .expect("base input should use CSV") + .file_header_info = Some(FileHeaderInfo::from(value.to_string())); + validate_select_request(&HeaderMap::new(), &mut input).expect("FileHeaderInfo should match any case"); + let csv = input.request.input_serialization.csv.as_ref().expect("CSV input"); + assert_eq!(csv.file_header_info.as_ref().map(|v| v.as_str()), Some(expected)); + } + + for (value, expected) in [("always", QuoteFields::ALWAYS), ("AsNeeded", QuoteFields::ASNEEDED)] { + let mut input = base_input(); + input + .request + .output_serialization + .csv + .as_mut() + .expect("base output should use CSV") + .quote_fields = Some(QuoteFields::from(value.to_string())); + let validation = validate_select_request(&HeaderMap::new(), &mut input).expect("QuoteFields should match any case"); + let SelectOutputFormat::Csv(csv) = validation.output_format else { + panic!("CSV output expected"); + }; + assert_eq!(csv.quote_fields.as_ref().map(|v| v.as_str()), Some(expected)); + } + + let mut input = base_input(); + input.request.expression_type = ExpressionType::from("sql".to_string()); + validate_select_request(&HeaderMap::new(), &mut input).expect("ExpressionType should match any case"); + assert_eq!(input.request.expression_type.as_str(), ExpressionType::SQL); + } + + #[test] + fn validate_still_rejects_unknown_enum_values_in_any_case() { + for value in ["gzipx", "snappy", " GZIP"] { + let mut input = base_input(); + input.request.input_serialization.compression_type = Some(CompressionType::from(value.to_string())); + let error = validate_select_request(&HeaderMap::new(), &mut input).expect_err("unknown compression must fail"); + assert_eq!(error.code(), &S3ErrorCode::InvalidCompressionFormat, "{value}"); + } + for value in ["ndjson", "lines "] { + let error = + validate_select_request(&HeaderMap::new(), &mut json_input(value)).expect_err("unknown JSON type must fail"); + assert_eq!(error.code(), &S3ErrorCode::InvalidJsonType, "{value}"); + } + let mut input = base_input(); + input + .request + .input_serialization + .csv + .as_mut() + .expect("base input should use CSV") + .file_header_info = Some(FileHeaderInfo::from("used".to_string())); + let error = validate_select_request(&HeaderMap::new(), &mut input).expect_err("unknown header mode must fail"); + assert_eq!(error.code(), &S3ErrorCode::InvalidFileHeaderInfo); + + let mut input = base_input(); + input + .request + .output_serialization + .csv + .as_mut() + .expect("base output should use CSV") + .quote_fields = Some(QuoteFields::from("never".to_string())); + let error = validate_select_request(&HeaderMap::new(), &mut input).expect_err("unknown QuoteFields must fail"); + assert_eq!(error.code(), &S3ErrorCode::InvalidQuoteFields); + + let mut input = base_input(); + input.request.expression_type = ExpressionType::from("sqlite".to_string()); + let error = validate_select_request(&HeaderMap::new(), &mut input).expect_err("unknown ExpressionType must fail"); + assert_eq!(error.code(), &S3ErrorCode::InvalidExpressionType); + } + #[test] fn validate_accepts_two_byte_csv_input_record_delimiter() { let mut input = base_input();