fix(s3select): case-insensitive enums, LINES key order, path names (#9021)

* fix(s3select): match request enum values case-insensitively

ExpressionType, CompressionType, CSV FileHeaderInfo, JSON Type and CSV
QuoteFields were compared exactly, so lowercase values such as
`lines`, `document` or `gzip` were rejected. MinIO accepts any case.

Match them ASCII case-insensitively in request validation and rewrite
them to the canonical uppercase spelling, so every later comparison
keeps working unchanged. Unknown values keep their existing errors.

* fix(s3select): keep source key order for JSON LINES input

Arrow's JSON schema inference walks each record as a serde_json::Map,
which sorts keys unless serde_json's process-wide preserve_order
feature is on. JSON LINES `SELECT *` therefore returned keys sorted,
and CSV output columns were permuted.

Infer the LINES schema from the same object bytes and record limit as
DataFusion's JsonFormat, keep Arrow's types, and reorder fields
(nested structs included) by first-seen key order. Enabling
preserve_order globally was rejected: it changes map ordering for the
whole binary.

Drop the test-only preserve_order dependency so the crate's tests run
with production features; the DOCUMENT inference reference now
compares against Arrow's types in source order.

* fix(s3select): name nested path projections after their last key

`SELECT s.obj0.id` returned the key `s.obj0[id]`, DataFusion's display
name for the field access. AWS names an unaliased path projection after
its last key, e.g. `{"id": ...}`.

Alias such projections to their last key when that name is unique among
the output and ORDER BY column names. Queries with a wildcard keep
today's names, since the expanded columns could collide.

* fix(s3select): preserve names beside computed projections

* fix(s3select): preserve names for ambiguous path projections

* fix(s3select): name path projections from the planned schema

Path projections that the SQL rewrite leaves unaliased are now named
after their last key once DataFusion has planned the query. A path takes
its key only if no other output column keeps or receives that name. A
colliding rename is dropped and the check repeats, so the final names
are unique whenever DataFusion's own names are.

The SQL rewrite is unchanged, so GROUP BY can still refer to the names
it assigns and every query that planned before still plans.

Review cases:
- s.obj0.id, s.obj1.id, s.obj2."s.obj0[id]" keeps the three generated
  names: the third key would duplicate the first retained name.
- s.obj."lower(s.id)", LOWER(s.id), in either order, keeps the generated
  path name: its key is the computed column's name.

Paths beside computed columns, beside duplicated keys, and in
parentheses now take their key, as AWS S3 Select names them. Paths stay
unnamed when an ORDER BY expression is not projected, as before.

---------

Co-authored-by: Hauser <housemecn@gmail.com>
Co-authored-by: overtrue <anzhengchao@gmail.com>
Co-authored-by: GatewayJ <835269233@qq.com>
This commit is contained in:
Damien Degois
2026-10-30 16:19:18 +01:00
committed by GitHub
parent 7e536843ea
commit 0d2bf69e2d
7 changed files with 1095 additions and 38 deletions
-2
View File
@@ -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
@@ -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);
@@ -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<Schema, ArrowError> {
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<Struct>` shape.
#[derive(Default)]
struct KeyOrder {
positions: HashMap<String, usize>,
children: Vec<KeyOrder>,
}
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<FieldRef> = 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<D: de::Deserializer<'de>>(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<A: MapAccess<'de>>(self, mut map: A) -> Result<(), A::Error> {
while let Some(key) = map.next_key::<String>()? {
map.next_value_seed(self.child(key))?;
}
Ok(())
}
fn visit_seq<A: SeqAccess<'de>>(self, mut seq: A) -> Result<(), A::Error> {
while seq.next_element_seed(&mut *self)?.is_some() {}
Ok(())
}
fn visit_bool<E: de::Error>(self, _: bool) -> Result<(), E> {
Ok(())
}
fn visit_i64<E: de::Error>(self, _: i64) -> Result<(), E> {
Ok(())
}
fn visit_u64<E: de::Error>(self, _: u64) -> Result<(), E> {
Ok(())
}
fn visit_f64<E: de::Error>(self, _: f64) -> Result<(), E> {
Ok(())
}
fn visit_str<E: de::Error>(self, _: &str) -> Result<(), E> {
Ok(())
}
fn visit_unit<E: de::Error>(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());
}
}
+451 -3
View File
@@ -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<SchemaRef> {
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<String>, Vec<Vec<String>>)> {
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::<std::collections::HashSet<_>>().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::<usize>(), 1);
let batch = batches.iter().find(|batch| batch.num_rows() == 1).expect("one result row");
let key = batch
.column(0)
.as_any()
.downcast_ref::<datafusion::arrow::array::Int64Array>()
.expect("integer key");
let lower = batch
.column(1)
.as_any()
.downcast_ref::<datafusion::arrow::array::StringArray>()
.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::<usize>(), 1);
let batch = batches.iter().find(|batch| batch.num_rows() == 1).expect("one result row");
let key = batch
.column(0)
.as_any()
.downcast_ref::<datafusion::arrow::array::Int64Array>()
.expect("integer key");
assert_eq!(key.value(0), 7);
}
}
@@ -14,5 +14,6 @@
mod json_document_schema;
mod json_document_table;
mod json_lines_schema;
pub mod manager;
mod parquet_table;
+288 -7
View File
@@ -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<Plan> {
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<Option<String>> = 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::<String, usize>::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<String> {
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<Option<String>> {
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<String>]) -> DFResult<LogicalPlan> {
let schema = plan.schema();
if keys.len() != schema.fields().len() {
return Ok(plan);
}
let mut renames: Vec<Option<&str>> = 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");
+165 -25
View File
@@ -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();