gsadmin/rustfs
1
0
Fork 0
mirror of https://github.com/rustfs/rustfs.git synced 2026-09-01 09:48:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(s3select): expand typed JSON source paths (#6864)

Browse Source
This commit is contained in:
GatewayJ
2026-08-30 14:46:36 +08:00
committed by GitHub
parent d6f9a7c462
commit 4932af080b
10 changed files with 2068 additions and 275 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/s3select-api/src/object_store.rs
+882 -199
View File
File diff suppressed because it is too large Load Diff
crates/s3select-api/src/query/ast.rs
+37
View File
@@ -13,6 +13,43 @@
// limitations under the License.
use datafusion::sql::sqlparser::ast::Statement;
use std::sync::Arc;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum JsonPathSegment {
Key { name: String, quoted: bool },
Index(usize),
ArrayWildcard,
ObjectWildcard,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct JsonSource {
path: Arc<[JsonPathSegment]>,
scalar_column: Option<String>,
}
impl JsonSource {
pub fn new(path: Vec<JsonPathSegment>, scalar_column: Option<String>) -> Self {
Self {
path: path.into(),
scalar_column,
}
}
#[cfg(test)]
pub(crate) fn from_path(path: Vec<JsonPathSegment>) -> Self {
Self::new(path, None)
}
pub fn path(&self) -> &[JsonPathSegment] {
&self.path
}
pub fn scalar_column(&self) -> Option<&str> {
self.scalar_column.as_deref()
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ExtStatement {
crates/s3select-api/src/query/execution.rs
+18
View File
@@ -29,6 +29,7 @@ use tracing::debug;
use crate::{QueryError, QueryResult};
use super::Query;
use super::ast::ExtStatement;
use super::logical_planner::Plan;
use super::session::{QueryExecutionTracker, SessionCtx};
@@ -172,6 +173,7 @@ pub struct QueryStateMachine {
pub session: SessionCtx,
pub query: Query,
prepared_statement: Option<ExtStatement>,
query_tracker: Option<QueryExecutionTracker>,
state: RwLock<QueryState>,
start: Instant,
@@ -196,6 +198,7 @@ impl QueryStateMachine {
Self {
session,
query,
prepared_statement: None,
query_tracker: None,
state: RwLock::new(QueryState::ACCEPTING),
start: Instant::now(),
@@ -211,6 +214,21 @@ impl QueryStateMachine {
Ok(state_machine)
}
pub fn begin_tracked_prepared(
query: Query,
session: SessionCtx,
query_tracker: QueryExecutionTracker,
prepared_statement: ExtStatement,
) -> QueryResult<Self> {
let mut state_machine = Self::begin_tracked(query, session, query_tracker)?;
state_machine.prepared_statement = Some(prepared_statement);
Ok(state_machine)
}
pub fn prepared_statement(&self) -> Option<&ExtStatement> {
self.prepared_statement.as_ref()
}
pub fn query_tracker(&self) -> Option<&QueryExecutionTracker> {
self.query_tracker.as_ref()
}
crates/s3select-api/src/query/parser.rs
+4
View File
@@ -34,6 +34,10 @@ impl Dialect for RustFsDialect {
fn supports_group_by_expr(&self) -> bool {
true
}
fn supports_partiql(&self) -> bool {
true
}
}
pub trait Parser {
crates/s3select-api/src/query/session.rs
+240 -16
View File
@@ -12,9 +12,11 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use crate::query::{Context, Query};
use crate::{QueryError, QueryResult, object_store::EcObjectStore};
use crate::{SelectInputMetrics, SelectObjectSnapshot};
use crate::query::{Context, Query, ast::JsonSource};
use crate::{
QueryError, QueryResult, SelectInputMetrics, SelectObjectSnapshot,
object_store::{EcObjectStore, is_json_document_input, legacy_json_source_from_input},
};
use datafusion::{
arrow::{
array::{Int32Array, StringArray},
@@ -314,8 +316,15 @@ impl SessionCtxFactory {
}
pub async fn create_session_ctx(&self, context: &Context) -> QueryResult<SessionCtx> {
self.create_session_ctx_inner(context, None, None, None, DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES)
.await
self.create_session_ctx_inner(
context,
None,
legacy_json_source_from_input(&context.input),
None,
None,
DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES,
)
.await
}
pub async fn create_session_ctx_with_tracker_and_memory_limit(
@@ -324,8 +333,15 @@ impl SessionCtxFactory {
query_tracker: QueryExecutionTracker,
memory_limit_bytes: usize,
) -> QueryResult<SessionCtx> {
self.create_session_ctx_inner(context, None, Some(query_tracker), None, memory_limit_bytes)
.await
self.create_session_ctx_inner(
context,
None,
legacy_json_source_from_input(&context.input),
Some(query_tracker),
None,
memory_limit_bytes,
)
.await
}
pub async fn create_session_ctx_with_snapshot_and_tracker_and_memory_limit(
@@ -335,8 +351,33 @@ impl SessionCtxFactory {
query_tracker: QueryExecutionTracker,
memory_limit_bytes: usize,
) -> QueryResult<SessionCtx> {
self.create_session_ctx_inner(context, Some(snapshot), Some(query_tracker), None, memory_limit_bytes)
.await
self.create_session_ctx_inner(
context,
Some(snapshot),
legacy_json_source_from_input(&context.input),
Some(query_tracker),
None,
memory_limit_bytes,
)
.await
}
pub async fn create_session_ctx_for_query_with_source_and_tracker_and_memory_limit(
&self,
query: &Query,
source: JsonSource,
query_tracker: QueryExecutionTracker,
memory_limit_bytes: usize,
) -> QueryResult<SessionCtx> {
self.create_session_ctx_inner(
query.context(),
query.snapshot().cloned(),
source,
Some(query_tracker),
Some(Arc::clone(query.input_metrics())),
memory_limit_bytes,
)
.await
}
pub async fn create_session_ctx_for_query_with_tracker_and_memory_limit(
@@ -348,6 +389,7 @@ impl SessionCtxFactory {
self.create_session_ctx_inner(
query.context(),
query.snapshot().cloned(),
legacy_json_source_from_input(&query.context().input),
Some(query_tracker),
Some(Arc::clone(query.input_metrics())),
memory_limit_bytes,
@@ -359,12 +401,13 @@ impl SessionCtxFactory {
&self,
context: &Context,
snapshot: Option<Arc<SelectObjectSnapshot>>,
source: JsonSource,
query_tracker: Option<QueryExecutionTracker>,
input_metrics: Option<Arc<SelectInputMetrics>>,
memory_limit_bytes: usize,
) -> QueryResult<SessionCtx> {
let df_session_ctx = self
.build_df_session_context(context, snapshot, query_tracker.clone(), input_metrics, memory_limit_bytes)
.build_df_session_context(context, snapshot, source, query_tracker.clone(), input_metrics, memory_limit_bytes)
.await?;
Ok(SessionCtx {
@@ -378,6 +421,7 @@ impl SessionCtxFactory {
&self,
context: &Context,
snapshot: Option<Arc<SelectObjectSnapshot>>,
source: JsonSource,
query_tracker: Option<QueryExecutionTracker>,
input_metrics: Option<Arc<SelectInputMetrics>>,
memory_limit_bytes: usize,
@@ -401,10 +445,12 @@ impl SessionCtxFactory {
.is_some_and(|delimiter| delimiter.len() == 2 && delimiter.as_bytes() != b"\r\n");
let scan_range_requires_single_file_scan =
context.input.request.scan_range.is_some() && context.input.request.input_serialization.parquet.is_none();
let json_document_requires_single_file_scan = is_json_document_input(&context.input);
let metered_input_requires_single_file_scan =
input_metrics.is_some() && context.input.request.input_serialization.parquet.is_none();
let config = if custom_two_byte_record_delimiter
|| scan_range_requires_single_file_scan
|| json_document_requires_single_file_scan
|| metered_input_requires_single_file_scan
{
config.with_repartition_file_scans(false)
@@ -463,14 +509,21 @@ impl SessionCtxFactory {
} else {
let input_metrics = input_metrics.unwrap_or_else(|| Arc::new(SelectInputMetrics::default()));
let store: EcObjectStore = match query_tracker {
Some(query_tracker) => EcObjectStore::new_with_query_tracker(
Some(query_tracker) => EcObjectStore::new_with_query_tracker_and_source(
context.input.clone(),
memory_pool,
query_tracker,
input_metrics,
snapshot,
source,
),
None => EcObjectStore::new_with_memory_pool_and_source(
context.input.clone(),
memory_pool,
input_metrics,
snapshot,
source,
),
None => EcObjectStore::new_with_memory_pool(context.input.clone(), memory_pool, input_metrics, snapshot),
}
.map_err(|err| QueryError::Datafusion {
source: Box::new(DataFusionError::External(Box::new(err))),
@@ -543,15 +596,15 @@ mod tests {
use crate::storage_api::object_store::ObjectIO as _;
use datafusion::{
datasource::{
file_format::csv::CsvFormat,
file_format::{csv::CsvFormat, json::JsonFormat},
listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl},
},
execution::memory_pool::MemoryLimit,
};
use http::HeaderMap;
use s3s::dto::{
CSVInput, CSVOutput, ExpressionType, InputSerialization, JSONInput, OutputSerialization, ParquetInput, ScanRange,
SelectObjectContentInput, SelectObjectContentRequest,
CSVInput, CSVOutput, ExpressionType, InputSerialization, JSONInput, JSONType, OutputSerialization, ParquetInput,
ScanRange, SelectObjectContentInput, SelectObjectContentRequest,
};
use std::io::Write as _;
@@ -592,6 +645,103 @@ mod tests {
)
}
async fn test_query_tracker() -> QueryExecutionTracker {
let permit = Arc::new(tokio::sync::Semaphore::new(1))
.acquire_owned()
.await
.expect("query permit should be available");
QueryExecutionTracker::new(
&QueryExecutionOwner::new(),
Arc::new(permit),
Instant::now() + std::time::Duration::from_secs(300),
300,
)
}
async fn assert_legacy_json_column(bucket: &str, expression: &str, document: &[u8], column_name: &str, expected: &[&str]) {
const OBJECT: &str = "input.json";
let env = crate::storage_api::select_test_ecstore_env().await;
let mut context = test_context();
{
let input = Arc::make_mut(&mut context.input);
input.bucket = bucket.to_string();
input.key = OBJECT.to_string();
input.request.expression = expression.to_string();
input.request.input_serialization.csv = None;
input.request.input_serialization.json = Some(JSONInput {
type_: Some(JSONType::from_static(JSONType::DOCUMENT)),
});
}
env.make_bucket(bucket, false).await;
env.put_object_bytes(bucket, OBJECT, document.to_vec()).await;
let factory = SessionCtxFactory::new(false);
let lazy_session = factory
.create_session_ctx(&context)
.await
.expect("legacy lazy session should preserve the JSON source");
let tracked_lazy_session = factory
.create_session_ctx_with_tracker_and_memory_limit(
&context,
test_query_tracker().await,
DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES,
)
.await
.expect("legacy tracked lazy session should preserve the JSON source");
let snapshot = prepare_test_snapshot(&context).await;
let snapshot_session = factory
.create_session_ctx_with_snapshot_and_tracker_and_memory_limit(
&context,
snapshot,
test_query_tracker().await,
DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES,
)
.await
.expect("legacy snapshot session should preserve the JSON source");
for (kind, session) in [
("lazy", lazy_session),
("tracked lazy", tracked_lazy_session),
("tracked snapshot", snapshot_session),
] {
let table_path = ListingTableUrl::parse(format!("s3://{bucket}/{OBJECT}")).expect("parse JSON table URL");
let listing_options = ListingOptions::new(Arc::new(JsonFormat::default())).with_file_extension(".json");
let schema = listing_options
.infer_schema(session.inner(), &table_path)
.await
.expect("infer expanded JSON schema");
let table = ListingTable::try_new(
ListingTableConfig::new(table_path)
.with_listing_options(listing_options)
.with_schema(schema),
)
.expect("build expanded JSON table");
let query_context = SessionContext::new_with_state(session.inner().clone());
query_context
.register_table("legacy_input", Arc::new(table))
.expect("register expanded JSON table");
let batches = query_context
.sql(&format!("SELECT {column_name} FROM legacy_input"))
.await
.expect("plan expanded JSON query")
.collect()
.await
.expect("execute expanded JSON query");
let mut values = Vec::new();
for batch in batches {
let column = batch
.column(0)
.as_any()
.downcast_ref::<StringArray>()
.expect("expanded column should be Utf8");
for row in 0..batch.num_rows() {
values.push(column.value(row).to_string());
}
}
assert_eq!(values, expected, "{kind} legacy constructor");
}
}
#[test]
fn session_factory_fields_remain_source_compatible() {
let factory = SessionCtxFactory {
@@ -629,6 +779,7 @@ mod tests {
.create_session_ctx_inner(
context,
None,
JsonSource::default(),
None,
Some(Arc::new(SelectInputMetrics::default())),
DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES,
@@ -680,6 +831,40 @@ mod tests {
assert!(!session.inner().config().options().optimizer.repartition_file_scans);
}
#[tokio::test]
async fn json_lines_without_scan_range_keeps_file_repartitioning() {
let mut context = test_context();
let request = &mut Arc::make_mut(&mut context.input).request;
request.input_serialization.csv = None;
request.input_serialization.json = Some(JSONInput::default());
let session = SessionCtxFactory::new(true)
.with_target_partitions(2)
.create_session_ctx(&context)
.await
.expect("JSON LINES session should be created");
assert!(session.inner().config().options().optimizer.repartition_file_scans);
}
#[tokio::test]
async fn json_document_disables_file_repartitioning() {
let mut context = test_context();
let request = &mut Arc::make_mut(&mut context.input).request;
request.input_serialization.csv = None;
request.input_serialization.json = Some(JSONInput {
type_: Some(JSONType::from_static(JSONType::DOCUMENT)),
});
let session = SessionCtxFactory::new(true)
.with_target_partitions(2)
.create_session_ctx(&context)
.await
.expect("JSON DOCUMENT session should be created");
assert!(!session.inner().config().options().optimizer.repartition_file_scans);
}
#[tokio::test]
async fn csv_scan_range_disables_file_repartitioning() {
let mut context = test_context();
@@ -755,7 +940,7 @@ mod tests {
async fn session_factory_applies_memory_limit() {
let factory = SessionCtxFactory::new(true);
let session = factory
.create_session_ctx_inner(&test_context(), None, None, None, 1024)
.create_session_ctx_inner(&test_context(), None, JsonSource::default(), None, None, 1024)
.await
.expect("session should be created with a bounded memory pool");
@@ -803,6 +988,45 @@ mod tests {
assert!(session.is_bound_to(&tracker));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial_test::serial]
async fn legacy_session_factory_preserves_single_key_json_source() {
assert_legacy_json_column(
"s3select-legacy-session-json-source",
"SELECT e.name FROM S3Object.employees AS e",
br#"{"employees":[{"name":"Alice"}]}"#,
"name",
&["Alice"],
)
.await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial_test::serial]
async fn legacy_session_factory_preserves_implicit_root_alias() {
assert_legacy_json_column(
"s3select-legacy-session-root-alias",
"SELECT S3Object FROM S3Object",
br#"["one","two"]"#,
"s3object",
&["one", "two"],
)
.await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial_test::serial]
async fn legacy_session_factory_preserves_quoted_root_alias() {
assert_legacy_json_column(
"s3select-legacy-session-quoted-root-alias",
"SELECT \"V\" FROM S3Object AS \"V\"",
br#"["one","two"]"#,
"\"V\"",
&["one", "two"],
)
.await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial_test::serial]
async fn session_factory_propagates_query_guard_to_ec_store() {