From 67a403757d3cafb7f8aca001909edd59fea99b1f Mon Sep 17 00:00:00 2001 From: Lanqing Yang Date: Wed, 16 Sep 2026 03:46:11 +0000 Subject: [PATCH 1/4] Perf/http sql limit materialization (#9148) perf(servers): limit HTTP SQL rows before materialization Signed-off-by: lyang24 --- src/servers/benches/to_http_output.rs | 250 ++++++++++++++++++- src/servers/src/http/handler.rs | 331 +++++++++++++++++++++++++- typos.toml | 3 + 3 files changed, 570 insertions(+), 14 deletions(-) diff --git a/src/servers/benches/to_http_output.rs b/src/servers/benches/to_http_output.rs index 2f2708a2636..f3a9149083a 100644 --- a/src/servers/benches/to_http_output.rs +++ b/src/servers/benches/to_http_output.rs @@ -13,15 +13,34 @@ // limitations under the License. use std::sync::Arc; -use std::time::Instant; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::{Duration, Instant}; -use arrow::array::StringArray; +use arrow::array::{ArrayRef, StringArray, UInt64Array}; use arrow_schema::{DataType, Field, Schema}; -use common_recordbatch::RecordBatch; -use criterion::{Criterion, criterion_group, criterion_main}; +use async_trait::async_trait; +use axum::extract::{Query, State}; +use axum::response::IntoResponse; +use axum::{Extension, Form}; +use common_query::Output; +use common_recordbatch::{RecordBatch, RecordBatchStreamWrapper}; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group}; +use datafusion_expr::LogicalPlan; use datatypes::schema::SchemaRef; use datatypes::vectors::StringVector; -use servers::http::HttpRecordsOutput; +use futures::stream; +use query::parser::PromQuery; +use query::query_engine::DescribeResult; +use servers::error::Result; +use servers::http::handler::{SqlQuery, sql}; +use servers::http::{ApiState, HttpRecordsOutput}; +use servers::query_handler::sql::SqlQueryHandler; +use session::context::{QueryContext, QueryContextRef}; +use sql::statements::statement::Statement; + +#[cfg(not(windows))] +#[global_allocator] +static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; fn mock_schema() -> SchemaRef { let mut fields = Vec::with_capacity(10); @@ -70,12 +89,227 @@ fn bench_convert_record_batch_to_http_output(c: &mut Criterion) { }); } +struct StreamingSqlHandler { + rows: usize, + schema: SchemaRef, + consumed: Arc, +} + +impl StreamingSqlHandler { + fn new(rows: usize) -> Self { + let fields = std::iter::once(Field::new("id", DataType::UInt64, false)) + .chain((1..10).map(|i| Field::new(format!("field{i}"), DataType::Utf8, false))); + let schema = Arc::new(Schema::new(fields.collect::>())); + Self { + rows, + schema: Arc::new(schema.try_into().unwrap()), + consumed: Arc::new(AtomicUsize::new(0)), + } + } +} + +#[async_trait] +impl SqlQueryHandler for StreamingSqlHandler { + async fn do_query(&self, _: &str, _: QueryContextRef) -> Vec> { + let schema = self.schema.clone(); + let consumed = self.consumed.clone(); + let rows = self.rows; + consumed.store(0, Ordering::Relaxed); + // Allocate independent buffers lazily, as a query stream would. Prebuilding + // all batches or sharing one batch's buffers would hide retention costs. + let batches = (0..rows).step_by(4096).map(move |offset| { + let end = (offset + 4096).min(rows); + let mut columns: Vec = vec![Arc::new(UInt64Array::from_iter_values( + offset as u64..end as u64, + ))]; + columns.extend((1..10).map(|_| { + Arc::new(StringArray::from_iter_values(std::iter::repeat_n( + "0123456789abcdef0123456789abcdef", + end - offset, + ))) as ArrayRef + })); + let batch = + arrow::record_batch::RecordBatch::try_new(schema.arrow_schema().clone(), columns) + .unwrap(); + consumed.fetch_add(end - offset, Ordering::Relaxed); + Ok(RecordBatch::from_df_record_batch(schema.clone(), batch)) + }); + vec![Ok(Output::new_with_stream(Box::pin( + RecordBatchStreamWrapper::new(self.schema.clone(), stream::iter(batches)), + )))] + } + + async fn do_analyze_stream_query(&self, _: &str, _: QueryContextRef) -> Result { + unimplemented!() + } + + async fn do_exec_plan( + &self, + _: LogicalPlan, + _: Option, + _: QueryContextRef, + ) -> Result { + unimplemented!() + } + + async fn do_promql_query(&self, _: &PromQuery, _: QueryContextRef) -> Vec> { + unimplemented!() + } + + async fn do_describe( + &self, + _: Statement, + _: QueryContextRef, + ) -> Result> { + unimplemented!() + } + + async fn is_valid_schema(&self, _: &str, _: &str) -> Result { + Ok(true) + } +} + +#[derive(Clone, Copy)] +enum LimitPlacement { + Late, + Early, +} + +async fn limited_response( + handler: Arc, + limit: usize, + placement: LimitPlacement, +) -> bytes::Bytes { + let response = sql( + State(ApiState { + sql_handler: handler.clone(), + }), + Query(SqlQuery { + sql: Some("select * from benchmark_input".to_string()), + limit: matches!(placement, LimitPlacement::Early).then_some(limit), + ..Default::default() + }), + Extension(QueryContext::with_db_name(None)), + Form(SqlQuery::default()), + ) + .await; + // Reconstruct the old path with the existing post-conversion truncation API. + let response = match placement { + LimitPlacement::Late => response.with_limit(limit), + LimitPlacement::Early => response, + }; + let response = response.with_execution_time(0).into_response(); + assert!(response.status().is_success()); + assert_eq!(handler.consumed.load(Ordering::Relaxed), handler.rows); + axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap() +} + +fn bench_http_sql_limit(c: &mut Criterion) { + let runtime = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + let mut group = c.benchmark_group("http_sql_limit"); + group.sample_size(10); + group.warm_up_time(Duration::from_millis(500)); + group.measurement_time(Duration::from_secs(2)); + for (rows, limit) in [ + (10_000, 100), + (100_000, 100), + (1_000_000, 100), + (1_000_000, 0), + (100_000, 100_000), + ] { + let handler = Arc::new(StreamingSqlHandler::new(rows)); + { + let late = runtime.block_on(limited_response( + handler.clone(), + limit, + LimitPlacement::Late, + )); + let early = runtime.block_on(limited_response( + handler.clone(), + limit, + LimitPlacement::Early, + )); + assert_eq!(late, early); + let json: serde_json::Value = serde_json::from_slice(&early).unwrap(); + assert_eq!( + json["output"][0]["records"]["rows"] + .as_array() + .unwrap() + .len(), + rows.min(limit) + ); + } + group.throughput(Throughput::Elements(rows as u64)); + for (name, placement) in [ + ("late", LimitPlacement::Late), + ("early", LimitPlacement::Early), + ] { + group.bench_function( + BenchmarkId::new(name, format!("rows_{rows}_limit_{limit}")), + |b| { + b.iter(|| { + runtime.block_on(limited_response(handler.clone(), limit, placement)) + }); + }, + ); + } + } + group.finish(); +} + #[cfg(not(windows))] criterion_group! { name = benches; config = Criterion::default().with_profiler(pprof::criterion::PProfProfiler::new(101, pprof::criterion::Output::Flamegraph(None))); - targets = bench_convert_record_batch_to_http_output + targets = bench_convert_record_batch_to_http_output, bench_http_sql_limit } #[cfg(windows)] -criterion_group!(benches, bench_convert_record_batch_to_http_output); -criterion_main!(benches); +criterion_group!( + benches, + bench_convert_record_batch_to_http_output, + bench_http_sql_limit +); + +fn main() { + // Run each memory case in a fresh process, e.g. with `/usr/bin/time -v`: + // GREPTIME_HTTP_LIMIT_MEMORY_CASE=late,1000000,100 + if let Ok(case) = std::env::var("GREPTIME_HTTP_LIMIT_MEMORY_CASE") { + let parts: Vec<_> = case.split(',').collect(); + assert_eq!(parts.len(), 3, "expected late|early,rows,limit"); + let placement = match parts[0] { + "late" => LimitPlacement::Late, + "early" => LimitPlacement::Early, + _ => panic!("expected late or early"), + }; + let rows = parts[1].parse().unwrap(); + let limit = parts[2].parse().unwrap(); + let handler = Arc::new(StreamingSqlHandler::new(rows)); + let runtime = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + #[cfg(not(windows))] + let allocated = tikv_jemalloc_ctl::thread::allocatedp::read().unwrap(); + #[cfg(not(windows))] + let before = allocated.get(); + let start = Instant::now(); + let response = runtime.block_on(limited_response(handler, limit, placement)); + let elapsed = start.elapsed(); + // Memory mode writes measurements to stdout for external benchmark tools. + #[allow(clippy::print_stdout)] + { + #[cfg(not(windows))] + println!("allocated_bytes={}", allocated.get() - before); + println!( + "case={case} elapsed={elapsed:?} response_bytes={}", + response.len() + ); + } + return; + } + benches(); + Criterion::default().configure_from_args().final_summary(); +} diff --git a/src/servers/src/http/handler.rs b/src/servers/src/http/handler.rs index ee2c3fd4144..bd31ee35ace 100644 --- a/src/servers/src/http/handler.rs +++ b/src/servers/src/http/handler.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::HashMap; +use std::future; use std::panic::AssertUnwindSafe; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; @@ -28,7 +29,8 @@ use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_plugins::GREPTIME_EXEC_WRITE_COST; use common_query::{Output, OutputData}; -use common_recordbatch::util; +use common_recordbatch::error::Result as RecordBatchResult; +use common_recordbatch::{RecordBatch, RecordBatchStreamWrapper, RecordBatches, util}; use common_telemetry::tracing; use datafusion::physical_plan::ExecutionPlan; use futures::{FutureExt, StreamExt}; @@ -42,7 +44,7 @@ use sql::parser::{ParseOptions, ParserContext}; use sql::statements::statement::Statement; use tokio::sync::{Notify, watch}; -use crate::error::{FailedToParseQuerySnafu, InvalidQuerySnafu, Result}; +use crate::error::{CollectRecordbatchSnafu, FailedToParseQuerySnafu, InvalidQuerySnafu, Result}; use crate::http::header::collect_plan_metrics; use crate::http::result::arrow_result::ArrowResponse; use crate::http::result::csv_result::CsvResponse; @@ -192,7 +194,25 @@ pub async fn sql( Ok(outputs) => outputs, }; - let mut resp = match format { + let outputs = match query_params.limit { + Some(limit) + if matches!( + format, + ResponseFormat::Csv(..) + | ResponseFormat::Table + | ResponseFormat::GreptimedbV1 + | ResponseFormat::Json + ) => + { + outputs + .into_iter() + .map(|output| output.and_then(|output| limit_output_rows(output, limit))) + .collect() + } + _ => outputs, + }; + + let resp = match format { ResponseFormat::Arrow => { ArrowResponse::from_output(outputs, query_params.compression).await } @@ -206,12 +226,57 @@ pub async fn sql( ResponseFormat::Null => NullResponse::from_output(outputs).await, }; - if let Some(limit) = query_params.limit { - resp = resp.with_limit(limit); - } resp.with_execution_time(start.elapsed().as_millis() as u64) } +/// Limits each statement's response before collecting batches and converting rows. +fn limit_output_rows(output: Output, limit: usize) -> Result { + let mut remaining = limit; + let data = match output.data { + OutputData::AffectedRows(rows) => OutputData::AffectedRows(rows), + OutputData::RecordBatches(batches) => { + let schema = batches.schema(); + let batches = batches + .into_iter() + .filter_map(|batch| take_response_rows(batch, &mut remaining).transpose()) + .collect::>>() + .context(CollectRecordbatchSnafu)?; + OutputData::RecordBatches( + RecordBatches::try_new(schema, batches).context(CollectRecordbatchSnafu)?, + ) + } + OutputData::Stream(stream) => { + let schema = stream.schema(); + // This is a response limit. Drain the input even after reaching it so + // execution errors and terminal query metrics are still observed. + let stream = stream.filter_map(move |batch| { + future::ready( + batch + .and_then(|batch| take_response_rows(batch, &mut remaining)) + .transpose(), + ) + }); + OutputData::Stream(Box::pin(RecordBatchStreamWrapper::new(schema, stream))) + } + }; + Ok(Output::new(data, output.meta)) +} + +fn take_response_rows( + batch: RecordBatch, + remaining: &mut usize, +) -> RecordBatchResult> { + let num_rows = batch.num_rows().min(*remaining); + *remaining -= num_rows; + if num_rows == 0 { + Ok(None) + } else if num_rows == batch.num_rows() { + Ok(Some(batch)) + } else { + batch.slice(0, num_rows).map(Some) + } +} + /// Handler to stream partial `EXPLAIN ANALYZE VERBOSE` metrics as SSE. /// /// This endpoint is POST-only SSE, so browser `EventSource` does @@ -925,7 +990,261 @@ pub async fn index() -> axum::response::Html { #[cfg(test)] mod tests { + use std::sync::Mutex; + use std::sync::atomic::AtomicUsize; + + use async_trait::async_trait; + use common_query::OutputMeta; + use datafusion_expr::LogicalPlan; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::{ColumnSchema, Schema}; + use datatypes::vectors::{BinaryVector, UInt32Vector, VectorRef}; + use futures::stream; + use query::query_engine::DescribeResult; + use super::*; + use crate::query_handler::sql::SqlQueryHandler; + + struct TestSqlQueryHandler { + outputs: Mutex>>, + } + + #[async_trait] + impl SqlQueryHandler for TestSqlQueryHandler { + async fn do_query(&self, _: &str, _: QueryContextRef) -> Vec> { + std::mem::take(&mut *self.outputs.lock().unwrap()) + } + + async fn do_analyze_stream_query(&self, _: &str, _: QueryContextRef) -> Result { + unimplemented!() + } + + async fn do_exec_plan( + &self, + _: LogicalPlan, + _: Option, + _: QueryContextRef, + ) -> Result { + unimplemented!() + } + + async fn do_promql_query(&self, _: &PromQuery, _: QueryContextRef) -> Vec> { + unimplemented!() + } + + async fn do_describe( + &self, + _: Statement, + _: QueryContextRef, + ) -> Result> { + unimplemented!() + } + + async fn is_valid_schema(&self, _: &str, _: &str) -> Result { + Ok(true) + } + } + + async fn sql_response( + outputs: Vec>, + format: &str, + limit: Option, + ) -> HttpResponse { + sql( + State(ApiState { + sql_handler: Arc::new(TestSqlQueryHandler { + outputs: Mutex::new(outputs), + }), + }), + Query(SqlQuery { + sql: Some("select number from numbers".to_string()), + format: Some(format.to_string()), + limit, + ..Default::default() + }), + Extension(QueryContext::with_db_name(None)), + Form(SqlQuery::default()), + ) + .await + } + + fn number_batches() -> RecordBatches { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "number", + ConcreteDataType::uint32_datatype(), + false, + )])); + let batches = [vec![], vec![0, 1], vec![], vec![2, 3, 4], vec![]] + .into_iter() + .map(|values| { + let columns: Vec = vec![Arc::new(UInt32Vector::from_slice(values))]; + RecordBatch::new(schema.clone(), columns).unwrap() + }) + .collect(); + RecordBatches::try_new(schema, batches).unwrap() + } + + fn number_output(streaming: bool) -> Output { + let batches = number_batches(); + if streaming { + Output::new_with_stream(batches.as_stream()) + } else { + Output::new_with_record_batches(batches) + } + } + + #[tokio::test] + async fn test_sql_response_limit_across_batches_and_formats() { + for format in [ + "greptimedb_v1", + "json", + "csv", + "csvWithNames", + "csvWithNamesAndTypes", + "table", + ] { + for streaming in [false, true] { + for limit in [ + None, + Some(0), + Some(1), + Some(2), + Some(3), + Some(5), + Some(10), + Some(usize::MAX), + ] { + let response = + sql_response(vec![Ok(number_output(streaming))], format, limit).await; + let output = match &response { + HttpResponse::GreptimedbV1(response) => response.output(), + HttpResponse::Json(response) => response.output(), + HttpResponse::Csv(response) => response.output(), + HttpResponse::Table(response) => response.output(), + _ => panic!("unexpected response: {response:?}"), + }; + let GreptimeQueryOutput::Records(records) = &output[0] else { + panic!("expected records"); + }; + let expected_rows: Vec<_> = (0..limit.unwrap_or(5).min(5)) + .map(|number| vec![Value::from(number)]) + .collect(); + assert_eq!(records.rows(), &expected_rows); + assert_eq!(records.total_rows, expected_rows.len()); + assert_eq!(records.schema.column_schemas[0].name, "number"); + } + } + } + } + + #[tokio::test] + async fn test_sql_response_limit_is_per_statement_and_preserves_write_cost() { + let affected = Output::new( + OutputData::AffectedRows(7), + OutputMeta { + cost: 42, + ..Default::default() + }, + ); + let response = sql_response( + vec![ + Ok(number_output(true)), + Ok(affected), + Ok(number_output(false)), + ], + "greptimedb_v1", + Some(1), + ) + .await; + let HttpResponse::GreptimedbV1(response) = response else { + panic!("expected greptimedb response"); + }; + assert_eq!(response.output().len(), 3); + for index in [0, 2] { + let GreptimeQueryOutput::Records(records) = &response.output()[index] else { + panic!("expected records"); + }; + assert_eq!(records.rows(), &vec![vec![Value::from(0)]]); + } + assert!(matches!( + response.output()[1], + GreptimeQueryOutput::AffectedRows(7) + )); + assert_eq!( + response.resp_metrics[GREPTIME_EXEC_WRITE_COST], + Value::from(42) + ); + } + + #[tokio::test] + async fn test_sql_response_limit_does_not_affect_other_formats() { + for format in ["arrow", "influxdb_v1", "null"] { + let unlimited = sql_response(vec![Ok(number_output(true))], format, None).await; + let limited = sql_response(vec![Ok(number_output(true))], format, Some(0)).await; + assert_eq!( + serde_json::to_value(unlimited.with_execution_time(0)).unwrap(), + serde_json::to_value(limited.with_execution_time(0)).unwrap(), + ); + } + } + + #[tokio::test] + async fn test_sql_response_limit_drains_input_and_propagates_late_errors() { + for limit in [0, 1, 3] { + for late_error in [false, true] { + let batches = number_batches(); + let schema = batches.schema(); + let mut input: Vec<_> = batches.into_iter().map(Ok).collect(); + if late_error { + input.push( + common_recordbatch::error::CreateRecordBatchesSnafu { + reason: "late stream error", + } + .fail(), + ); + } + let expected_batches = input.len(); + let observed = Arc::new(AtomicUsize::new(0)); + let counter = observed.clone(); + let stream = stream::iter(input).inspect(move |_| { + counter.fetch_add(1, Ordering::Relaxed); + }); + let output = Output::new_with_stream(Box::pin(RecordBatchStreamWrapper::new( + schema, stream, + ))); + let response = sql_response(vec![Ok(output)], "greptimedb_v1", Some(limit)).await; + assert_eq!(observed.load(Ordering::Relaxed), expected_batches); + assert_eq!(matches!(response, HttpResponse::Error(_)), late_error); + } + } + } + + #[tokio::test] + async fn test_sql_response_limit_skips_conversion_of_discarded_rows() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "payload", + ConcreteDataType::json_datatype(), + false, + )])); + let columns: Vec = vec![Arc::new(BinaryVector::from(vec![ + datatypes::types::parse_string_to_jsonb(r#"{"ok":1}"#).unwrap(), + b"invalid jsonb".to_vec(), + ]))]; + let batch = RecordBatch::new(schema.clone(), columns).unwrap(); + for streaming in [false, true] { + for limit in [0, 1, 2] { + let batches = RecordBatches::try_new(schema.clone(), vec![batch.clone()]).unwrap(); + let output = if streaming { + Output::new_with_stream(batches.as_stream()) + } else { + Output::new_with_record_batches(batches) + }; + let response = sql_response(vec![Ok(output)], "greptimedb_v1", Some(limit)).await; + // Only retained rows should be decoded into HTTP response values. + assert_eq!(matches!(response, HttpResponse::Error(_)), limit == 2); + } + } + } #[test] fn test_final_analyze_event_uses_error_event_for_conversion_error() { diff --git a/typos.toml b/typos.toml index bb14c40bab4..44dcda61687 100644 --- a/typos.toml +++ b/typos.toml @@ -12,6 +12,9 @@ consts = "consts" # We need to extend the word list to avoid typos in our codebase. Exceded = "Exceded" +[default.extend-identifiers] +allocatedp = "allocatedp" + [files] extend-exclude = [ "corrupted", From 1c01b541ec617e580d4c8a563bccf56886e39957 Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Wed, 16 Sep 2026 04:44:32 +0000 Subject: [PATCH 2/4] refactor(json2): JSON2 parquet projection and schema alignment (#9137) * refactor(mito2): make JSON schema alignment targets explicit Signed-off-by: fys * fix: cr * fix: cargo fmt --------- Signed-off-by: fys --- src/datatypes/src/vectors/json/array.rs | 451 +------------- .../{json_align/stream.rs => json_align.rs} | 505 ++++++++++++---- src/mito2/src/sst/parquet/json_align/mod.rs | 24 - src/mito2/src/sst/parquet/read_columns.rs | 568 +++++++----------- src/mito2/src/sst/parquet/reader.rs | 39 +- 5 files changed, 639 insertions(+), 948 deletions(-) rename src/mito2/src/sst/parquet/{json_align/stream.rs => json_align.rs} (53%) delete mode 100644 src/mito2/src/sst/parquet/json_align/mod.rs diff --git a/src/datatypes/src/vectors/json/array.rs b/src/datatypes/src/vectors/json/array.rs index 153d6b490ef..e8b75dee259 100644 --- a/src/datatypes/src/vectors/json/array.rs +++ b/src/datatypes/src/vectors/json/array.rs @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::cmp::Ordering; use std::sync::Arc; use arrow::compute::{can_cast_types, cast}; @@ -21,13 +20,12 @@ use arrow_array::types::{ Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type, UInt64Type, }; -use arrow_array::{Array, ArrayRef, GenericListArray, ListArray, StructArray, new_null_array}; -use arrow_schema::{DataType, Field, FieldRef}; -use common_telemetry::trace; +use arrow_array::{Array, ArrayRef, GenericListArray, StructArray, new_null_array}; +use arrow_schema::{DataType, Field}; use serde_json::Value; use snafu::{OptionExt, ResultExt}; -use crate::arrow_array::{MutableBinaryArray, binary_array_value, string_array_value}; +use crate::arrow_array::{binary_array_value, string_array_value}; use crate::data_type::ConcreteDataType; use crate::error::{ AlignJsonArraySnafu, ArrowComputeSnafu, InvalidJsonSnafu, InvalidJsonbSnafu, Result, @@ -200,171 +198,6 @@ impl JsonArray<'_> { Ok(values) } - /// Normalizes a JSON2 array to the wider `expect` data type without losing - /// information. - /// - /// This is mainly used for write/flush-time JSON2 schema alignment: - /// - fields missing from the source are filled with typed null arrays; - /// - fields present in the source must also exist in `expect`; - /// - fields present in both are widened recursively when their types differ. - /// - /// Narrowing conversions and any other conversions that may lose information - /// are rejected. - pub fn widen_to(&self, expect: &DataType) -> Result { - let data_type = self.inner.data_type(); - - if data_type == expect { - return Ok(self.inner.clone()); - } - - trace!( - "Try aligning JSON array {} to data type {}", - data_type, expect - ); - - let struct_array = self.inner.as_struct_opt().context(AlignJsonArraySnafu { - reason: "expect struct array", - })?; - let array_fields = struct_array.fields(); - let array_columns = struct_array.columns(); - let DataType::Struct(expect_fields) = expect else { - return AlignJsonArraySnafu { - reason: "expect struct datatype", - } - .fail(); - }; - let mut aligned = Vec::with_capacity(expect_fields.len()); - - // Compare the fields in the JSON array and the to-be-aligned schema, amending with null - // arrays on the way. It's very important to note that fields in the JSON array and those - // in the JSON type are both **SORTED**, which can be guaranteed because the fields in the - // JSON type implementation are sorted. - debug_assert!(expect_fields.iter().map(|f| f.name()).is_sorted()); - debug_assert!(array_fields.iter().map(|f| f.name()).is_sorted()); - - let mut i = 0; // point to the expect fields - let mut j = 0; // point to the array fields - while i < expect_fields.len() && j < array_fields.len() { - let expect_field = &expect_fields[i]; - let array_field = &array_fields[j]; - match expect_field.name().cmp(array_field.name()) { - Ordering::Equal => { - if expect_field.data_type() == array_field.data_type() { - aligned.push(array_columns[j].clone()); - } else { - let expect_type = expect_field.data_type(); - let array_type = array_field.data_type(); - let array = match (expect_type, array_type) { - (DataType::Struct(_), DataType::Struct(_)) => { - JsonArray::from(&array_columns[j]).widen_to(expect_type)? - } - (DataType::List(expect_item), DataType::List(array_item)) => { - let list_array = array_columns[j].as_list::(); - widen_list(list_array, array_item, expect_item)? - } - _ => JsonArray::from(&array_columns[j]).widen_scalar_to(expect_type)?, - }; - aligned.push(array); - } - i += 1; - j += 1; - } - Ordering::Less => { - aligned.push(new_null_array(expect_field.data_type(), struct_array.len())); - i += 1; - } - Ordering::Greater => { - return AlignJsonArraySnafu { - reason: format!( - "source field {} does not exist in target schema", - array_field.name() - ), - } - .fail(); - } - } - } - if j < array_fields.len() { - return AlignJsonArraySnafu { - reason: format!( - "source field {} does not exist in target schema", - array_fields[j].name() - ), - } - .fail(); - } - if i < expect_fields.len() { - for field in &expect_fields[i..] { - aligned.push(new_null_array(field.data_type(), struct_array.len())); - } - } - - let json_array = StructArray::try_new_with_length( - expect_fields.clone(), - aligned, - struct_array.nulls().cloned(), - struct_array.len(), - ) - .map_err(|e| { - AlignJsonArraySnafu { - reason: e.to_string(), - } - .build() - })?; - Ok(Arc::new(json_array)) - } - - /// Widens an array to the merged JSON2 physical type without losing information. - /// - /// Supported conversions: - /// - identical types are returned unchanged; - /// - null arrays become typed null arrays; - /// - concrete JSON values are encoded as JSONB when the target type is binary. - /// - /// All other conversions are rejected. - fn widen_scalar_to(&self, to_type: &DataType) -> Result { - let from_type = self.inner.data_type(); - if from_type == to_type { - return Ok(self.inner.clone()); - } - - if from_type == &DataType::Null { - return Ok(new_null_array(to_type, self.inner.len())); - } - - if !from_type.is_binary() && to_type.is_binary() { - return self.encode_variant(); - } - - AlignJsonArraySnafu { - reason: format!("unable to widen {from_type} to {to_type}"), - } - .fail() - } - - fn encode_variant(&self) -> Result { - let len = self.inner.len(); - let mut encoded = Vec::with_capacity(len); - let mut total_bytes = 0; - - for i in 0..len { - let value = self.try_get_value(i)?; - if value.is_null() { - encoded.push(None); - } else { - let bytes = encode_serde_json_as_jsonb(value); - total_bytes += bytes.len(); - encoded.push(Some(bytes)); - } - } - - let mut builder = MutableBinaryArray::with_capacity(len, total_bytes); - for value in encoded { - builder.append_option(value); - } - Ok(Arc::new(builder.finish())) - } - /// Projects this JSON array to `target` for query evaluation. /// /// Unlike [`Self::widen_to`], projection tolerates lossy conversions: @@ -567,28 +400,6 @@ fn project_json_value_to_type(value: Value, to_type: &ConcreteDataType) -> Resul Ok(to_type.try_cast(value).unwrap_or(GreptimeValue::Null)) } -fn widen_list(list_array: &ListArray, actual: &FieldRef, expected: &FieldRef) -> Result { - let item_aligned = match (actual.data_type(), expected.data_type()) { - (DataType::Struct(_), DataType::Struct(_)) => { - JsonArray::from(list_array.values()).widen_to(expected.data_type())? - } - (DataType::List(actual), DataType::List(expected)) => { - let list_array = list_array.values().as_list::(); - widen_list(list_array, actual, expected)? - } - _ => JsonArray::from(list_array.values()).widen_scalar_to(expected.data_type())?, - }; - Ok(Arc::new( - GenericListArray::::try_new( - expected.clone(), - list_array.offsets().clone(), - item_aligned, - list_array.nulls().cloned(), - ) - .context(ArrowComputeSnafu)?, - )) -} - impl<'a> From<&'a ArrayRef> for JsonArray<'a> { fn from(inner: &'a ArrayRef) -> Self { Self { inner } @@ -740,262 +551,6 @@ mod test { Ok(()) } - #[test] - fn test_widen_null_to_any_type() -> Result<()> { - let nulls = new_null_array(&DataType::Null, 2); - let target_types = [ - DataType::Boolean, - DataType::UInt64, - DataType::Utf8View, - DataType::Binary, - DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true))), - DataType::Struct(Fields::from(vec![Field::new( - "value", - DataType::Int64, - true, - )])), - ]; - - for target_type in target_types { - let widened = JsonArray::from(&nulls).widen_scalar_to(&target_type)?; - assert_eq!(&target_type, widened.data_type()); - assert_eq!(2, widened.len()); - assert_eq!(2, widened.null_count()); - } - - Ok(()) - } - - #[test] - fn test_widen_non_null_to_utf8_view_fails() { - let bools: ArrayRef = Arc::new(BooleanArray::from(vec![true])); - let err = JsonArray::from(&bools) - .widen_scalar_to(&DataType::Utf8View) - .unwrap_err(); - - assert_eq!( - "Failed to align JSON array, reason: unable to widen Boolean to Utf8View", - err.to_string() - ); - } - - #[test] - fn test_widen_variant_to_non_binary_fails() { - let value = jsonb::parse_value(b"true").unwrap().to_vec(); - let variants: ArrayRef = Arc::new(BinaryArray::from(vec![value.as_slice()])); - let err = JsonArray::from(&variants) - .widen_scalar_to(&DataType::Boolean) - .unwrap_err(); - - assert_eq!( - "Failed to align JSON array, reason: unable to widen Binary to Boolean", - err.to_string() - ); - } - - #[test] - fn test_widen_between_number_types_fails() { - let values: ArrayRef = Arc::new(UInt64Array::from(vec![1])); - let err = JsonArray::from(&values) - .widen_scalar_to(&DataType::Int64) - .unwrap_err(); - - assert_eq!( - "Failed to align JSON array, reason: unable to widen UInt64 to Int64", - err.to_string() - ); - } - - #[test] - fn test_widen_numbers_to_variant_preserves_values() -> Result<()> { - let cases: [(ArrayRef, Value); 3] = [ - (Arc::new(UInt64Array::from(vec![u64::MAX])), json!(u64::MAX)), - (Arc::new(Int64Array::from(vec![i64::MIN])), json!(i64::MIN)), - (Arc::new(Float64Array::from(vec![1.25])), json!(1.25)), - ]; - - for (values, expected) in cases { - let widened = JsonArray::from(&values).widen_scalar_to(&DataType::Binary)?; - assert_eq!(&DataType::Binary, widened.data_type()); - assert_eq!(expected, JsonArray::from(&widened).try_get_value(0)?); - } - - Ok(()) - } - - #[test] - fn test_align_json_array() -> Result<()> { - struct TestCase { - json_array: ArrayRef, - schema_type: DataType, - expected: std::result::Result, - } - - impl TestCase { - fn new( - json_array: StructArray, - schema_type: Fields, - expected: std::result::Result, String>, - ) -> Self { - Self { - json_array: Arc::new(json_array), - schema_type: DataType::Struct(schema_type.clone()), - expected: expected - .map(|x| Arc::new(StructArray::new(schema_type, x, None)) as ArrayRef), - } - } - - fn test(self) -> Result<()> { - let result = JsonArray::from(&self.json_array).widen_to(&self.schema_type); - match (result, self.expected) { - (Ok(json_array), Ok(expected)) => assert_eq!(&json_array, &expected), - (Ok(json_array), Err(e)) => { - panic!("expecting error {e} but actually get: {json_array:?}") - } - (Err(e), Err(expected)) => assert_eq!(e.to_string(), expected), - (Err(e), Ok(_)) => return Err(e), - } - Ok(()) - } - } - - // Test empty json array can be aligned with a complex json type. - TestCase::new( - StructArray::new_empty_fields(2, None), - Fields::from(vec![ - Field::new("int", DataType::Int64, true), - Field::new_struct( - "nested", - vec![Field::new("bool", DataType::Boolean, true)], - true, - ), - Field::new("string", DataType::Utf8, true), - ]), - Ok(vec![ - Arc::new(Int64Array::new_null(2)) as ArrayRef, - Arc::new(StructArray::new_null( - Fields::from(vec![Arc::new(Field::new("bool", DataType::Boolean, true))]), - 2, - )), - Arc::new(StringArray::new_null(2)), - ]), - ) - .test()?; - - // Test simple json array alignment. - TestCase::new( - StructArray::from(vec![( - Arc::new(Field::new("float", DataType::Float64, true)), - Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])) as ArrayRef, - )]), - Fields::from(vec![ - Field::new("float", DataType::Float64, true), - Field::new("string", DataType::Utf8, true), - ]), - Ok(vec![ - Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])) as ArrayRef, - Arc::new(StringArray::new_null(3)), - ]), - ) - .test()?; - - // Test complex json array alignment. - TestCase::new( - StructArray::from(vec![ - ( - Arc::new(Field::new_list( - "list", - Field::new_list_field(DataType::Int64, true), - true, - )), - Arc::new(ListArray::from_iter_primitive::(vec![ - Some(vec![Some(1)]), - None, - Some(vec![Some(2), Some(3)]), - ])) as ArrayRef, - ), - ( - Arc::new(Field::new_struct( - "nested", - vec![Field::new("int", DataType::Int64, true)], - true, - )), - Arc::new(StructArray::from(vec![( - Arc::new(Field::new("int", DataType::Int64, true)), - Arc::new(Int64Array::from(vec![-1, -2, -3])) as ArrayRef, - )])), - ), - ( - Arc::new(Field::new("string", DataType::Utf8, true)), - Arc::new(StringArray::from(vec!["a", "b", "c"])), - ), - ]), - Fields::from(vec![ - Field::new("bool", DataType::Boolean, true), - Field::new_list("list", Field::new_list_field(DataType::Int64, true), true), - Field::new_struct( - "nested", - vec![ - Field::new("float", DataType::Float64, true), - Field::new("int", DataType::Int64, true), - ], - true, - ), - Field::new("string", DataType::Utf8, true), - ]), - Ok(vec![ - Arc::new(BooleanArray::new_null(3)) as ArrayRef, - Arc::new(ListArray::from_iter_primitive::(vec![ - Some(vec![Some(1)]), - None, - Some(vec![Some(2), Some(3)]), - ])), - Arc::new(StructArray::from(vec![ - ( - Arc::new(Field::new("float", DataType::Float64, true)), - Arc::new(Float64Array::new_null(3)) as ArrayRef, - ), - ( - Arc::new(Field::new("int", DataType::Int64, true)), - Arc::new(Int64Array::from(vec![-1, -2, -3])), - ), - ])), - Arc::new(StringArray::from(vec!["a", "b", "c"])), - ]), - ) - .test()?; - - // Source fields that do not exist in the target schema must not be discarded. - TestCase::new( - StructArray::from(vec![( - Arc::new(Field::new("a", DataType::Boolean, true)), - Arc::new(BooleanArray::from(vec![true])) as ArrayRef, - )]), - Fields::from(vec![Field::new("b", DataType::Boolean, true)]), - Err( - "Failed to align JSON array, reason: source field a does not exist in target schema" - .to_string(), - ), - ) - .test()?; - - // Trailing source fields must also be rejected after all target fields are processed. - TestCase::new( - StructArray::from(vec![( - Arc::new(Field::new("b", DataType::Boolean, true)), - Arc::new(BooleanArray::from(vec![true])) as ArrayRef, - )]), - Fields::from(vec![Field::new("a", DataType::Boolean, true)]), - Err( - "Failed to align JSON array, reason: source field b does not exist in target schema" - .to_string(), - ), - ) - .test()?; - - Ok(()) - } - #[test] fn test_align_variant_to_struct() -> Result<()> { let encode = |json: &[u8]| jsonb::parse_value(json).unwrap().to_vec(); diff --git a/src/mito2/src/sst/parquet/json_align/stream.rs b/src/mito2/src/sst/parquet/json_align.rs similarity index 53% rename from src/mito2/src/sst/parquet/json_align/stream.rs rename to src/mito2/src/sst/parquet/json_align.rs index a474fd781d0..8af88f6ac0f 100644 --- a/src/mito2/src/sst/parquet/json_align/stream.rs +++ b/src/mito2/src/sst/parquet/json_align.rs @@ -19,12 +19,15 @@ use std::task::{Context, Poll}; use datafusion_common::cast_column; use datafusion_common::format::DEFAULT_CAST_OPTIONS; use datatypes::arrow::array::{ArrayRef, new_null_array}; -use datatypes::arrow::datatypes::{DataType, Field, FieldRef, SchemaRef}; +use datatypes::arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::extension::json::{JsonMetadata, is_json2_extension_type}; use datatypes::json::JsonSettings; use datatypes::vectors::json::array::JsonArray; +use datatypes::vectors::json::json2_physical_data_type; use futures::Stream; +use futures::stream::BoxStream; +use serde_json::from_str; use snafu::{ResultExt, ensure}; use crate::error::{ @@ -32,32 +35,44 @@ use crate::error::{ }; use crate::sst::parquet::Json2TargetLayout; +pub(crate) type ProjectedRecordBatchStream = BoxStream<'static, Result>; + +/// Specifies how JSON columns in a record batch are aligned. #[derive(Debug)] -struct Json2RewriteSettings { - logical_settings: JsonSettings, - target_layout: JsonSettings, +pub(crate) enum AlignMode { + /// Aligns JSON columns to the logical fields in the output schema. + AlignToSchema, + /// Rewrites JSON columns to physical layouts, typically for compaction. + Rewrite { + /// Target layouts keyed by root column name, not nested field path. + /// + /// Only listed columns are rewritten. Other existing arrays are reused + /// unchanged and must already match their output field types. + /// An empty map therefore only fills missing roots. + columns: HashMap, + }, } -/// Aligns projected batches to the expected output schema for nested projections. +/// Alignment mode with parsed rewrite metadata and validated target layouts. +#[derive(Debug)] +enum ResolvedAlignMode { + AlignToSchema, + Rewrite { + columns: HashMap, + }, +} + +/// Adapts Parquet record batches to the output schema expected by the reader. /// -/// Background -/// ---------- -/// Nested projection may ask parquet to read leaves under a root column. If none -/// of the requested leaves exists in the current parquet file, parquet decoding -/// omits the whole root from the physical [`RecordBatch`]. +/// Nested projection can return only part of a JSON2 column, or omit its root +/// entirely when no requested leaves are read. This stream restores missing +/// roots with null arrays of the expected types. /// -/// In addition, after nested-path filtering, returned struct arrays may contain -/// only a subset of fields. The current output schema is not pruned by nested -/// paths, so physical struct fields can be a subset of the expected struct -/// fields, and their nested schema can differ from the expected output schema. -/// -/// To keep projected batches schema-consistent before entering upper readers: -/// - Root-column presence alignment restores missing projected root columns by -/// inserting root-level null arrays. -/// - Nested struct alignment aligns struct arrays to the expected nested field -/// layout. +/// Existing JSON2 columns are aligned to the logical schema inferred from +/// type hints or rewritten to the specified JSON2 physical layout, as selected by +/// [`AlignMode`]. #[derive(derive_more::Debug)] -pub struct NestedSchemaAligner { +pub struct JsonSchemaAligner { #[debug(skip)] inner: S, /// Output schema expected by the upper reader. @@ -70,79 +85,55 @@ pub struct NestedSchemaAligner { /// Whether all projected roots are present and the stream can pass batches /// through. all_roots_present: bool, - /// JSON2 columns that require semantic source-to-target layout rewriting. - json2_rewrite_targets: HashMap, + /// Alignment mode with parsed and validated rewrite settings. + mode: ResolvedAlignMode, /// The cache for whether incoming batches already match output schema. is_schema_matched: Option, } -impl NestedSchemaAligner +impl JsonSchemaAligner where S: Stream>, { - pub fn new( + /// Creates an aligner with a shared output schema and an explicit operation. + /// Parses rewrite metadata once and validates layouts against output field types. + pub(crate) fn new( inner: S, projected_root_presence: Vec, output_schema: SchemaRef, - ) -> Result> { + mode: AlignMode, + ) -> Result> { ensure!( projected_root_presence.len() == output_schema.fields().len(), UnexpectedSnafu { reason: format!( - "NestedSchemaAligner projected root presence len {} does not match output schema columns {}", + "JsonSchemaAligner projected root presence len {} does not match output schema columns {}", projected_root_presence.len(), output_schema.fields().len() ), } ); + let mode = resolve_align_mode(mode, output_schema.as_ref())?; + let expected_input_col_num = projected_root_presence .iter() .filter(|matched| **matched) .count(); let all_roots_present = projected_root_presence.iter().all(|&m| m); - Ok(NestedSchemaAligner { + Ok(JsonSchemaAligner { inner, output_schema, projected_root_presence, expected_input_col_num, all_roots_present, - json2_rewrite_targets: HashMap::new(), + mode, is_schema_matched: None, }) } - - /// Sets JSON2 columns that must be rewritten into the output field layout. - pub(crate) fn with_json2_rewrite_targets( - mut self, - targets: &HashMap, - ) -> Result { - self.json2_rewrite_targets = targets - .iter() - .map(|(name, layout)| { - let metadata = serde_json::from_str::(&layout.extension_metadata) - .map_err(|e| { - UnexpectedSnafu { - reason: format!( - "invalid JSON2 extension metadata for column '{name}': {e}" - ), - } - .build() - })?; - Ok(( - name.clone(), - Json2RewriteSettings { - logical_settings: metadata.into_json_settings(), - target_layout: layout.target_layout.clone(), - }, - )) - }) - .collect::>()?; - Ok(self) - } } -impl Stream for NestedSchemaAligner +impl Stream for JsonSchemaAligner where S: Stream> + Unpin, { @@ -153,7 +144,8 @@ where match Pin::new(&mut this.inner).poll_next(cx) { Poll::Ready(Some(Ok(rb))) => { - let is_schema_matched = this.all_roots_present + let is_schema_matched = matches!(this.mode, ResolvedAlignMode::AlignToSchema) + && this.all_roots_present && *this .is_schema_matched .get_or_insert_with(|| rb.schema() == this.output_schema); @@ -166,7 +158,7 @@ where &this.output_schema, &this.projected_root_presence, this.expected_input_col_num, - &this.json2_rewrite_targets, + &this.mode, ))) } } @@ -182,13 +174,13 @@ fn align_projected_batch( output_schema: &SchemaRef, projected_root_presence: &[bool], expected_input_col_num: usize, - json2_rewrite_targets: &HashMap, + mode: &ResolvedAlignMode, ) -> Result { ensure!( rb.columns().len() == expected_input_col_num, UnexpectedSnafu { reason: format!( - "NestedSchemaAligner expected {} input columns but got {}", + "JsonSchemaAligner expected {} input columns but got {}", expected_input_col_num, rb.columns().len() ), @@ -205,12 +197,16 @@ fn align_projected_batch( continue; } - cols.push(align_array( - rb.column(idx), - input_schema.field(idx), - field, - json2_rewrite_targets.get(field.name()), - )?); + let array = match mode { + ResolvedAlignMode::AlignToSchema => { + align_array(rb.column(idx), input_schema.field(idx), field)? + } + ResolvedAlignMode::Rewrite { columns } => match columns.get(field.name()) { + Some(settings) => rewrite_array(rb.column(idx), input_schema.field(idx), settings)?, + None => rb.column(idx).clone(), + }, + }; + cols.push(array); idx += 1; } @@ -218,36 +214,111 @@ fn align_projected_batch( } fn align_array( - array: &ArrayRef, - source: &Field, - field: &FieldRef, - rewrite_settings: Option<&Json2RewriteSettings>, + source_array: &ArrayRef, + source_field: &Field, + target_field: &FieldRef, ) -> Result { - if let Some(settings) = rewrite_settings { - return JsonArray::from(array) - .rewrite_to_v2(source, &settings.logical_settings, &settings.target_layout) - .context(DataTypeMismatchSnafu); - } - if array.data_type() == field.data_type() { - return Ok(array.clone()); + if source_array.data_type() == target_field.data_type() { + return Ok(source_array.clone()); } - if is_json2_extension_type(field) { - if is_json2_extension_type(source) { - return JsonArray::from(array) - .project_to_v2(source, field.data_type()) + if is_json2_extension_type(target_field) { + if is_json2_extension_type(source_field) { + return JsonArray::from(source_array) + .project_to_v2(source_field, target_field.data_type()) .context(DataTypeMismatchSnafu); } - return JsonArray::from(array) - .project_to(field.data_type()) + return JsonArray::from(source_array) + .project_to(target_field.data_type()) .context(DataTypeMismatchSnafu); } - if !matches!(field.data_type(), DataType::Struct(_)) { - return Ok(array.clone()); + if !matches!(target_field.data_type(), DataType::Struct(_)) { + return Ok(source_array.clone()); } - cast_column(array, field.data_type(), &DEFAULT_CAST_OPTIONS).context(CastColumnSnafu) + cast_column( + source_array, + target_field.data_type(), + &DEFAULT_CAST_OPTIONS, + ) + .context(CastColumnSnafu) +} + +fn rewrite_array( + source_array: &ArrayRef, + source_field: &Field, + settings: &RewriteSettings, +) -> Result { + JsonArray::from(source_array) + .rewrite_to_v2( + source_field, + &settings.logical_settings, + &settings.target_layout, + ) + .context(DataTypeMismatchSnafu) +} + +/// Resolved settings for rewriting one JSON2 column to a target physical layout. +/// +/// Created from [`Json2TargetLayout`] when resolving the alignment mode, so +/// extension metadata is parsed once and reused across batches. +#[derive(Debug)] +struct RewriteSettings { + /// Logical settings parsed from extension metadata and applied to JSON values + /// before encoding them into the target layout. + logical_settings: JsonSettings, + /// Settings defining the physical Arrow layout of the rewritten column. + target_layout: JsonSettings, +} + +impl TryFrom<&Json2TargetLayout> for RewriteSettings { + type Error = crate::error::Error; + + fn try_from(layout: &Json2TargetLayout) -> Result { + let metadata = from_str::(&layout.extension_metadata).map_err(|e| { + UnexpectedSnafu { + reason: format!("invalid JSON2 extension metadata: {e}"), + } + .build() + })?; + Ok(Self { + logical_settings: metadata.into_json_settings(), + target_layout: layout.target_layout.clone(), + }) + } +} + +/// Parses rewrite metadata and validates target layouts against the output schema. +fn resolve_align_mode(mode: AlignMode, output_schema: &Schema) -> Result { + let AlignMode::Rewrite { columns } = mode else { + return Ok(ResolvedAlignMode::AlignToSchema); + }; + + let mut rewrite_columns = HashMap::with_capacity(columns.len()); + for (name, layout) in columns { + let settings = RewriteSettings::try_from(&layout)?; + let field = output_schema.field_with_name(&name).map_err(|_| { + UnexpectedSnafu { + reason: format!("JSON2 rewrite column '{name}' is missing from output schema"), + } + .build() + })?; + ensure!( + is_json2_extension_type(field) + && field.data_type() == &json2_physical_data_type(&settings.target_layout), + UnexpectedSnafu { + reason: format!( + "JSON2 rewrite layout for column '{name}' does not match output field" + ), + } + ); + rewrite_columns.insert(name, settings); + } + + Ok(ResolvedAlignMode::Rewrite { + columns: rewrite_columns, + }) } #[cfg(test)] @@ -279,14 +350,21 @@ mod tests { target_layout: target_layout.clone(), }, )]); - let aligner = NestedSchemaAligner::new( + let aligner = JsonSchemaAligner::new( stream::empty::>(), - vec![], - schema(Vec::::new()), - )? - .with_json2_rewrite_targets(&rewrite_targets)?; - - let settings = &aligner.json2_rewrite_targets["j"]; + vec![false], + schema([ + Field::new("j", json2_physical_data_type(&target_layout), true) + .with_extension_type(Json2ExtensionType::default()), + ]), + AlignMode::Rewrite { + columns: rewrite_targets, + }, + )?; + let ResolvedAlignMode::Rewrite { columns } = &aligner.mode else { + panic!("expected rewrite mode"); + }; + let settings = &columns["j"]; assert_eq!(logical_settings, settings.logical_settings); assert_eq!(target_layout, settings.target_layout); Ok(()) @@ -305,8 +383,13 @@ mod tests { .unwrap(); let stream = stream::iter([Ok(input.clone())]); - let mut aligner = - NestedSchemaAligner::new(stream, vec![true, true], output_schema.clone()).unwrap(); + let mut aligner = JsonSchemaAligner::new( + stream, + vec![true, true], + output_schema.clone(), + AlignMode::AlignToSchema, + ) + .unwrap(); let output = aligner.next().await.unwrap().unwrap(); assert_eq!(input, output); @@ -324,9 +407,13 @@ mod tests { let input = RecordBatch::try_new(input_schema, vec![int_array([10, 20])]).unwrap(); let stream = stream::iter([Ok(input)]); - let mut aligner = - NestedSchemaAligner::new(stream, vec![true, false, false], output_schema.clone()) - .unwrap(); + let mut aligner = JsonSchemaAligner::new( + stream, + vec![true, false, false], + output_schema.clone(), + AlignMode::AlignToSchema, + ) + .unwrap(); let output = aligner.next().await.unwrap().unwrap(); assert_eq!(output_schema, output.schema()); @@ -362,8 +449,13 @@ mod tests { let input = RecordBatch::try_new(input_schema, vec![int_array([10, 20])]).unwrap(); let stream = stream::iter([Ok(input)]); - let mut aligner = - NestedSchemaAligner::new(stream, vec![true, false], output_schema.clone()).unwrap(); + let mut aligner = JsonSchemaAligner::new( + stream, + vec![true, false], + output_schema.clone(), + AlignMode::AlignToSchema, + ) + .unwrap(); let output = aligner.next().await.unwrap().unwrap(); assert_eq!(output_schema, output.schema()); @@ -377,8 +469,13 @@ mod tests { let output_schema = schema([Field::new("a", DataType::Int64, true)]); let stream = stream::iter([]); - let err = match NestedSchemaAligner::new(stream, vec![true, false], output_schema) { - Ok(_) => panic!("NestedSchemaAligner should reject projection length mismatch"), + let err = match JsonSchemaAligner::new( + stream, + vec![true, false], + output_schema, + AlignMode::AlignToSchema, + ) { + Ok(_) => panic!("JsonSchemaAligner should reject projection length mismatch"), Err(err) => err, }; @@ -399,8 +496,13 @@ mod tests { let input = RecordBatch::try_new(input_schema, vec![int_array([1, 2])]).unwrap(); let stream = stream::iter([Ok(input)]); - let mut aligner = - NestedSchemaAligner::new(stream, vec![true, true, false], output_schema).unwrap(); + let mut aligner = JsonSchemaAligner::new( + stream, + vec![true, true, false], + output_schema, + AlignMode::AlignToSchema, + ) + .unwrap(); let err = aligner.next().await.unwrap().unwrap_err(); assert!( @@ -410,7 +512,7 @@ mod tests { } #[tokio::test] - async fn test_nested_schema_aligner_aligns_struct_field() { + async fn test_json_schema_aligner_aligns_struct_field() { let output_schema = schema([Field::new( "nested", DataType::Struct(Fields::from(vec![ @@ -432,9 +534,13 @@ mod tests { ) .unwrap(); - let mut aligner = - NestedSchemaAligner::new(stream::iter([Ok(input)]), vec![true], output_schema.clone()) - .unwrap(); + let mut aligner = JsonSchemaAligner::new( + stream::iter([Ok(input)]), + vec![true], + output_schema.clone(), + AlignMode::AlignToSchema, + ) + .unwrap(); let output = aligner.next().await.unwrap().unwrap(); assert_eq!(output_schema, output.schema()); @@ -448,7 +554,7 @@ mod tests { } #[tokio::test] - async fn test_nested_schema_aligner_decodes_variant_to_struct() { + async fn test_json_schema_aligner_decodes_variant_to_struct() { let source_values = [ Some(parse_string_to_jsonb("1").unwrap()), Some(parse_string_to_jsonb(r#"{"b":2}"#).unwrap()), @@ -481,9 +587,13 @@ mod tests { true, ) .with_extension_type(Json2ExtensionType::default())]); - let mut aligner = - NestedSchemaAligner::new(stream::iter([Ok(input)]), vec![true], output_schema.clone()) - .unwrap(); + let mut aligner = JsonSchemaAligner::new( + stream::iter([Ok(input)]), + vec![true], + output_schema.clone(), + AlignMode::AlignToSchema, + ) + .unwrap(); let output = aligner.next().await.unwrap().unwrap(); assert_eq!(output_schema, output.schema()); @@ -516,7 +626,7 @@ mod tests { } #[tokio::test] - async fn test_nested_schema_aligner_preserves_struct_siblings() { + async fn test_json_schema_aligner_preserves_struct_siblings() { let source_values = [ Some(parse_string_to_jsonb(r#"{"x":1}"#).unwrap()), Some(parse_string_to_jsonb(r#"{"x":2}"#).unwrap()), @@ -570,9 +680,13 @@ mod tests { true, ) .with_extension_type(Json2ExtensionType::default())]); - let mut aligner = - NestedSchemaAligner::new(stream::iter([Ok(input)]), vec![true], output_schema.clone()) - .unwrap(); + let mut aligner = JsonSchemaAligner::new( + stream::iter([Ok(input)]), + vec![true], + output_schema.clone(), + AlignMode::AlignToSchema, + ) + .unwrap(); let output = aligner.next().await.unwrap().unwrap(); assert_eq!(output_schema, output.schema()); @@ -605,6 +719,159 @@ mod tests { ); } + #[tokio::test] + async fn test_rewrite_multiple_columns_and_fill_missing_roots() { + let logical_settings = JsonSettings::default(); + let target_layout = JsonSettings::try_new(vec![], Some(0)).unwrap(); + let target_type = json2_physical_data_type(&target_layout); + let output_schema = schema([ + Field::new("j", target_type.clone(), true) + .with_extension_type(Json2ExtensionType::default()), + Field::new("missing", target_type.clone(), true) + .with_extension_type(Json2ExtensionType::default()), + Field::new("k", target_type.clone(), true) + .with_extension_type(Json2ExtensionType::default()), + Field::new("a", DataType::Int64, true), + ]); + let values = [Some(parse_string_to_jsonb(r#"{"x":1}"#).unwrap()), None]; + let source = Arc::new(BinaryArray::from_iter( + values.iter().map(|value| value.as_deref()), + )) as ArrayRef; + let source_field = Field::new("j", DataType::Binary, true) + .with_extension_type(Json2ExtensionType::default()); + let expected = JsonArray::from(&source) + .rewrite_to_v2(&source_field, &logical_settings, &target_layout) + .unwrap(); + let input = RecordBatch::try_new( + schema([ + source_field, + Field::new("k", DataType::Binary, true) + .with_extension_type(Json2ExtensionType::default()), + Field::new("a", DataType::Int64, true), + ]), + vec![source.clone(), source, int_array([10, 20])], + ) + .unwrap(); + let columns = ["j", "missing", "k"] + .into_iter() + .map(|name| { + ( + name.to_string(), + Json2TargetLayout { + extension_metadata: serde_json::to_string(&JsonMetadata::new( + logical_settings.clone(), + )) + .unwrap(), + target_layout: target_layout.clone(), + }, + ) + }) + .collect(); + let mut aligner = JsonSchemaAligner::new( + stream::iter([Ok(input)]), + vec![true, false, true, true], + output_schema.clone(), + AlignMode::Rewrite { columns }, + ) + .unwrap(); + let output = aligner.next().await.unwrap().unwrap(); + assert_eq!(output_schema, output.schema()); + assert_eq!(expected.as_ref(), output.column(0).as_ref()); + assert_eq!(expected.as_ref(), output.column(2).as_ref()); + assert_eq!(&target_type, output.column(1).data_type()); + assert_eq!(2, output.column(1).null_count()); + assert_eq!(int_array([10, 20]).as_ref(), output.column(3).as_ref()); + } + + #[test] + fn test_rewrite_rejects_mismatched_output_layout() { + let columns = HashMap::from([( + "j".to_string(), + Json2TargetLayout { + extension_metadata: serde_json::to_string(&JsonMetadata::new( + JsonSettings::default(), + )) + .unwrap(), + target_layout: JsonSettings::try_new(vec![], Some(0)).unwrap(), + }, + )]); + let result = JsonSchemaAligner::new( + stream::empty::>(), + vec![false], + schema([Field::new("j", DataType::Binary, true) + .with_extension_type(Json2ExtensionType::default())]), + AlignMode::Rewrite { columns }, + ); + assert!( + result + .unwrap_err() + .to_string() + .contains("does not match output field") + ); + } + + #[tokio::test] + async fn test_empty_rewrite_only_fills_missing_roots() { + let source = int_array([10, 20]); + let input = RecordBatch::try_new( + schema([Field::new("a", DataType::Int64, true)]), + vec![source.clone()], + ) + .unwrap(); + let output_schema = schema([ + Field::new("missing", DataType::Utf8, true), + Field::new("a", DataType::Int64, true), + ]); + let mut aligner = JsonSchemaAligner::new( + stream::iter([Ok(input)]), + vec![false, true], + output_schema.clone(), + AlignMode::Rewrite { + columns: HashMap::new(), + }, + ) + .unwrap(); + let output = aligner.next().await.unwrap().unwrap(); + assert_eq!(output_schema, output.schema()); + assert_eq!(2, output.num_rows()); + assert_eq!(2, output.column(0).null_count()); + assert!(Arc::ptr_eq(&source, output.column(1))); + } + + #[tokio::test] + async fn test_empty_rewrite_does_not_align_existing_struct() { + let source = Arc::new(StructArray::from(vec![( + Arc::new(Field::new("x", DataType::Int64, true)), + int_array([1, 2]), + )])) as ArrayRef; + let input = RecordBatch::try_new( + schema([Field::new("j", source.data_type().clone(), true)]), + vec![source], + ) + .unwrap(); + let output_schema = schema([Field::new( + "j", + DataType::Struct(Fields::from(vec![ + Field::new("x", DataType::Int64, true), + Field::new("y", DataType::Utf8, true), + ])), + true, + )]); + let mut aligner = JsonSchemaAligner::new( + stream::iter([Ok(input)]), + vec![true], + output_schema, + AlignMode::Rewrite { + columns: HashMap::new(), + }, + ) + .unwrap(); + assert!(matches!( + aligner.next().await.unwrap(), + Err(crate::error::Error::NewRecordBatch { .. }) + )); + } + fn schema(fields: impl IntoIterator) -> SchemaRef { Arc::new(Schema::new(fields.into_iter().collect::>())) } diff --git a/src/mito2/src/sst/parquet/json_align/mod.rs b/src/mito2/src/sst/parquet/json_align/mod.rs deleted file mode 100644 index 54c97bdbb89..00000000000 --- a/src/mito2/src/sst/parquet/json_align/mod.rs +++ /dev/null @@ -1,24 +0,0 @@ -// Copyright 2023 Greptime 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 datatypes::arrow::record_batch::RecordBatch; -use futures::stream::BoxStream; - -mod stream; - -pub(crate) use stream::NestedSchemaAligner; - -use crate::error::Result; - -pub(crate) type ProjectedRecordBatchStream = BoxStream<'static, Result>; diff --git a/src/mito2/src/sst/parquet/read_columns.rs b/src/mito2/src/sst/parquet/read_columns.rs index 58c7d0d2f8d..07614375dcf 100644 --- a/src/mito2/src/sst/parquet/read_columns.rs +++ b/src/mito2/src/sst/parquet/read_columns.rs @@ -12,13 +12,17 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; +use std::ops::Range; -use datatypes::extension::json::JSON2_REMAINDER_FIELD_NAME; +use datatypes::arrow::datatypes::Schema as ArrowSchema; +use datatypes::extension::json::{JSON2_REMAINDER_FIELD_NAME, is_json2_extension_type}; use parquet::arrow::ProjectionMask; use parquet::basic::{ConvertedType, Type as PhysicalType}; use parquet::schema::types::{ColumnDescriptor, SchemaDescriptor}; +use crate::error::Result as MitoResult; + /// A nested field access path inside one parquet root column. pub type ParquetNestedPath = Vec; @@ -147,6 +151,32 @@ impl ParquetReadColumn { } } +/// Nested leaf selection semantics. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum NestedSelectionPolicy { + /// Also read the nearest JSONB Variant ancestor, or the remainder when it may + /// contain an unresolved path or children of a materialized object. + /// + /// For example, if `j.cold` has no matching field or Variant ancestor, read the + /// remainder because it may contain `j.cold`. When requesting a materialized + /// object such as `j.commit`, also read the remainder because it may contain + /// unmaterialized children of `j.commit`. + Json2, +} + +impl NestedSelectionPolicy { + /// Selects all leaves needed for one JSON2 root, including fallback data. + fn select_leaves( + self, + schema: &SchemaDescriptor, + leaf_range: Range, + col: &ParquetReadColumn, + selected: &mut HashSet, + ) { + select_json2_leaves(schema, leaf_range, col, selected); + } +} + /// Projection plan built for a parquet file. #[derive(Clone)] pub struct ProjectionMaskPlan { @@ -173,6 +203,9 @@ pub struct ProjectionMaskPlan { /// file. It is used to resolve requested nested paths to actual leaf /// column indices. /// +/// `source_schema` is the Arrow schema of the current parquet file. It is used +/// only for nested projections to identify JSON2 root fields. +/// /// See [`ProjectionMaskPlan`] for the returned value. /// /// For example, if the query requests `j.a` and `k`, but the current @@ -183,18 +216,20 @@ pub struct ProjectionMaskPlan { pub(crate) fn build_projection_plan( parquet_read_cols: &ParquetReadColumns, parquet_schema_desc: &SchemaDescriptor, -) -> ProjectionMaskPlan { + source_schema: &ArrowSchema, +) -> MitoResult { if !parquet_read_cols.has_nested() { let mask = ProjectionMask::roots(parquet_schema_desc, parquet_read_cols.root_indices_iter()); - return ProjectionMaskPlan { + + return Ok(ProjectionMaskPlan { mask, projected_root_presence: vec![true; parquet_read_cols.columns().len()], - }; + }); } let (matched_leaves, matched_roots) = - build_parquet_leaves_indices(parquet_schema_desc, parquet_read_cols); + build_parquet_leaves_indices(parquet_schema_desc, parquet_read_cols, source_schema)?; let projected_root_presence = parquet_read_cols .columns() @@ -203,10 +238,11 @@ pub(crate) fn build_projection_plan( .collect(); let mask = ProjectionMask::leaves(parquet_schema_desc, matched_leaves); - ProjectionMaskPlan { + + Ok(ProjectionMaskPlan { mask, projected_root_presence, - } + }) } /// Builds parquet leaf-column indices for reading a parquet file. @@ -217,89 +253,114 @@ pub(crate) fn build_projection_plan( fn build_parquet_leaves_indices( parquet_schema_desc: &SchemaDescriptor, projection: &ParquetReadColumns, -) -> (Vec, HashSet) { - let mut map = HashMap::with_capacity(projection.cols.len()); - for col in &projection.cols { - map.insert(col.root_index, col); - } + source_schema: &ArrowSchema, +) -> MitoResult<(Vec, HashSet)> { + let root_leaf_ranges = group_requested_leaf_ranges(parquet_schema_desc, projection); let mut matched_leaves = HashSet::new(); - let mut matched_roots = HashSet::with_capacity(projection.cols.len()); + let mut matched_roots = HashSet::with_capacity(projection.columns().len()); - // Tracks whether each requested nested path matched by prefix without fallback. - let mut prefix_matched = HashMap::>::new(); - for col in &projection.cols { - prefix_matched.insert(col.root_index, vec![false; col.nested_paths.len()]); - } + for col in projection.columns() { + let before = matched_leaves.len(); - // First select root/nested prefix matches. - for (leaf_idx, leaf_col) in parquet_schema_desc.columns().iter().enumerate() { - let root_idx = parquet_schema_desc.get_column_root_idx(leaf_idx); - let Some(col) = map.get(&root_idx) else { - continue; - }; - if col.nested_paths.is_empty() { - matched_leaves.insert(leaf_idx); - matched_roots.insert(root_idx); - continue; + let leaf_range = root_leaf_ranges[col.root_index()].clone(); + + if col.nested_paths().is_empty() { + matched_leaves.extend(leaf_range); + } else if is_json2_extension_type(&source_schema.fields()[col.root_index()]) { + NestedSelectionPolicy::Json2.select_leaves( + parquet_schema_desc, + leaf_range, + col, + &mut matched_leaves, + ); + } else { + select_prefix_leaves(parquet_schema_desc, leaf_range, col, &mut matched_leaves); } - let leaf_path = leaf_col.path().parts(); - let mut matched = false; - for (path_idx, _) in col - .nested_paths - .iter() - .enumerate() - .filter(|(_, nested_path)| leaf_path.starts_with(nested_path)) - { - prefix_matched.get_mut(&root_idx).unwrap()[path_idx] = true; - matched = true; - } - - if matched { - matched_leaves.insert(leaf_idx); - matched_roots.insert(root_idx); - } - } - - // Then include v2 remainder leaves or fallback prefix misses to their nearest variant parent. - // TODO(fys): Gate fallback planning on the root being JSON2. A raw Binary - // leaf is a JSONB variant only under a JSON2 root; plain struct Binary - // children should not enter this fallback path. - for col in &projection.cols { - let path_matches = &prefix_matched[&col.root_index]; - let mut needs_remainder = false; - for (matched, nested_path) in path_matches.iter().zip(&col.nested_paths) { - if *matched { - if !needs_remainder { - needs_remainder = - path_points_to_struct(parquet_schema_desc, col.root_index, nested_path); - } - continue; - } - - if let Some(leaf_idx) = - find_nearest_variant_parent(parquet_schema_desc, col.root_index, nested_path) - { - matched_leaves.insert(leaf_idx); - matched_roots.insert(col.root_index); - } else { - needs_remainder = true; - } - } - - if needs_remainder { - let remainder_leaves = find_remainder_leaves(parquet_schema_desc, col.root_index); - if !remainder_leaves.is_empty() { - matched_leaves.extend(remainder_leaves); - matched_roots.insert(col.root_index); - } + // Requested roots are unique, and leaves belong to exactly one root. + if matched_leaves.len() > before { + matched_roots.insert(col.root_index()); } } let mut matched_leaves = matched_leaves.into_iter().collect::>(); matched_leaves.sort_unstable(); - (matched_leaves, matched_roots) + + Ok((matched_leaves, matched_roots)) +} + +/// Groups file-level leaf ranges for requested roots. +fn group_requested_leaf_ranges( + schema: &SchemaDescriptor, + projection: &ParquetReadColumns, +) -> Vec> { + let root_count = schema.root_schema().get_fields().len(); + let mut requested = vec![false; root_count]; + for col in projection.columns() { + requested[col.root_index()] = true; + } + + let mut ranges = vec![0..0; root_count]; + let mut leaf_idx = 0; + while leaf_idx < schema.num_columns() { + let root_idx = schema.get_column_root_idx(leaf_idx); + let start = leaf_idx; + while leaf_idx < schema.num_columns() && schema.get_column_root_idx(leaf_idx) == root_idx { + leaf_idx += 1; + } + if requested[root_idx] { + ranges[root_idx] = start..leaf_idx; + } + } + ranges +} + +/// V2 can additionally store missing paths and object children in the remainder. +fn select_json2_leaves( + schema: &SchemaDescriptor, + leaf_range: Range, + col: &ParquetReadColumn, + selected: &mut HashSet, +) { + let prefix_matched = select_prefix_leaves(schema, leaf_range, col, selected); + let mut needs_remainder = false; + for (matched, path) in prefix_matched.iter().zip(&col.nested_paths) { + if *matched { + needs_remainder |= path_points_to_struct(schema, col.root_index, path); + } else if let Some(idx) = find_nearest_variant_parent(schema, col.root_index, path) { + selected.insert(idx); + } else { + needs_remainder = true; + } + } + if needs_remainder { + selected.extend(find_remainder_leaves(schema, col.root_index)); + } +} + +/// Selects prefix matches and returns a match flag for each requested nested path. +fn select_prefix_leaves( + schema: &SchemaDescriptor, + leaf_range: Range, + col: &ParquetReadColumn, + selected: &mut HashSet, +) -> Vec { + let mut prefix_matched = vec![false; col.nested_paths().len()]; + for leaf_idx in leaf_range { + let leaf_path = schema.columns()[leaf_idx].path().parts(); + let mut matched_leaf = false; + for (path, matched) in col.nested_paths().iter().zip(&mut prefix_matched) { + if leaf_path.starts_with(path) { + *matched = true; + matched_leaf = true; + } + } + if matched_leaf { + selected.insert(leaf_idx); + } + } + prefix_matched } /// Returns whether a nested path points to an explicitly materialized object. @@ -386,7 +447,10 @@ fn is_variant_leaf(leaf_col: &ColumnDescriptor) -> bool { mod tests { use std::sync::Arc; - use parquet::basic::{ConvertedType, LogicalType, Repetition, VariantType}; + use datatypes::arrow::datatypes::{DataType, Field, Fields, Schema as ArrowSchema}; + use datatypes::extension::json::{Json2ExtensionType, JsonMetadata}; + use datatypes::json::JsonSettings; + use parquet::basic::{LogicalType, Repetition, VariantType}; use parquet::errors::ParquetError; use parquet::schema::types::Type; @@ -397,7 +461,7 @@ mod tests { let parquet_schema_desc = build_test_nested_parquet_schema(); let projection = ParquetReadColumns::from_deduped_root_indices([0, 1]); - let plan = build_projection_plan(&projection, &parquet_schema_desc); + let plan = build_projection_plan(&projection, &parquet_schema_desc, &[None; 2]); assert_eq!(vec![true, true], plan.projected_root_presence); assert_eq!( @@ -412,8 +476,11 @@ mod tests { let projection = ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0)]); - let (matched_leaves, matched_roots) = - build_parquet_leaves_indices(&parquet_schema_desc, &projection); + let (matched_leaves, matched_roots) = build_parquet_leaves_indices_with_policies( + &parquet_schema_desc, + &projection, + &[None; 2], + ); assert_eq!(vec![0, 1, 2], matched_leaves); assert_eq!(HashSet::from([0]), matched_roots); } @@ -428,8 +495,11 @@ mod tests { ParquetReadColumn::new(1), ]); - let (matched_leaves, matched_roots) = - build_parquet_leaves_indices(&parquet_schema_desc, &projection); + let (matched_leaves, matched_roots) = build_parquet_leaves_indices_with_policies( + &parquet_schema_desc, + &projection, + &[None; 2], + ); assert_eq!(vec![1, 2, 3], matched_leaves); assert_eq!(HashSet::from([0, 1]), matched_roots); } @@ -443,8 +513,11 @@ mod tests { .with_nested_paths(vec![vec!["j".to_string(), "b".to_string()]]), ]); - let (matched_leaves, matched_roots) = - build_parquet_leaves_indices(&parquet_schema_desc, &projection); + let (matched_leaves, matched_roots) = build_parquet_leaves_indices_with_policies( + &parquet_schema_desc, + &projection, + &[None; 2], + ); assert_eq!(vec![1, 2], matched_leaves); assert_eq!(HashSet::from([0]), matched_roots); } @@ -460,8 +533,11 @@ mod tests { let read_column = ParquetReadColumn::new(0).with_nested_paths(nested_paths); let projection = ParquetReadColumns::from_deduped(vec![read_column]); - let (matched_leaves, matched_roots) = - build_parquet_leaves_indices(&parquet_schema_desc, &projection); + let (matched_leaves, matched_roots) = build_parquet_leaves_indices_with_policies( + &parquet_schema_desc, + &projection, + &[None; 2], + ); assert_eq!(vec![1, 2], matched_leaves); assert_eq!(HashSet::from([0]), matched_roots); } @@ -475,8 +551,11 @@ mod tests { vec![vec!["j".to_string(), "b".to_string(), "c".to_string()]], )]); - let (matched_leaves, matched_roots) = - build_parquet_leaves_indices(&parquet_schema_desc, &projection); + let (matched_leaves, matched_roots) = build_parquet_leaves_indices_with_policies( + &parquet_schema_desc, + &projection, + &[None; 2], + ); assert_eq!(vec![1], matched_leaves); assert_eq!(HashSet::from([0]), matched_roots); } @@ -491,7 +570,7 @@ mod tests { ParquetReadColumn::new(1), ]); - let plan = build_projection_plan(&projection, &parquet_schema_desc); + let plan = build_projection_plan(&projection, &parquet_schema_desc, &[None; 2]); assert_eq!(vec![false, true], plan.projected_root_presence); assert_eq!( @@ -511,7 +590,8 @@ mod tests { ], )]); - let plan = build_projection_plan(&projection, &parquet); + let plan = + build_projection_plan(&projection, &parquet, &[Some(NestedSelectionPolicy::Json2)]); assert_eq!(vec![true], plan.projected_root_presence); assert_eq!(ProjectionMask::leaves(&parquet, [0, 1]), plan.mask); @@ -526,7 +606,8 @@ mod tests { .with_nested_paths(vec![vec!["j".to_string(), "hot".to_string()]]), ]); - let plan = build_projection_plan(&projection, &parquet); + let plan = + build_projection_plan(&projection, &parquet, &[Some(NestedSelectionPolicy::Json2)]); assert_eq!(vec![true], plan.projected_root_presence); assert_eq!(ProjectionMask::leaves(&parquet, [3]), plan.mask); @@ -541,7 +622,8 @@ mod tests { .with_nested_paths(vec![vec!["j".to_string(), "commit".to_string()]]), ]); - let plan = build_projection_plan(&projection, &parquet); + let plan = + build_projection_plan(&projection, &parquet, &[Some(NestedSelectionPolicy::Json2)]); assert_eq!(vec![true], plan.projected_root_presence); assert_eq!(ProjectionMask::leaves(&parquet, [0, 1, 2]), plan.mask); @@ -563,7 +645,8 @@ mod tests { ]], )]); - let plan = build_projection_plan(&projection, &parquet); + let plan = + build_projection_plan(&projection, &parquet, &[Some(NestedSelectionPolicy::Json2)]); assert_eq!(vec![true], plan.projected_root_presence); assert_eq!(ProjectionMask::leaves(&parquet, [4]), plan.mask); @@ -582,8 +665,11 @@ mod tests { ], )]); - let (matched_leaves, matched_roots) = - build_parquet_leaves_indices(&parquet_schema_desc, &projection); + let (matched_leaves, matched_roots) = build_parquet_leaves_indices_with_policies( + &parquet_schema_desc, + &projection, + &[None; 2], + ); assert_eq!(vec![0, 2], matched_leaves); assert_eq!(HashSet::from([0]), matched_roots); } @@ -614,137 +700,6 @@ mod tests { assert!(col.nested_paths().is_empty()); } - #[test] - fn test_fallback_to_nearest_variant_parent() { - let parquet_schema_desc = build_test_variant_parent_schema(); - let projection = - ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths( - vec![vec!["j".to_string(), "a".to_string(), "x".to_string()]], - )]); - - let plan = build_projection_plan(&projection, &parquet_schema_desc); - - assert_eq!(vec![true], plan.projected_root_presence); - assert_eq!( - ProjectionMask::leaves(&parquet_schema_desc, vec![0]), - plan.mask - ); - } - - #[test] - fn test_prefix_match_prevents_variant_parent_fallback() { - let parquet_schema_desc = build_test_variant_parent_schema(); - let projection = - ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths( - vec![vec!["j".to_string(), "b".to_string(), "x".to_string()]], - )]); - - let plan = build_projection_plan(&projection, &parquet_schema_desc); - - assert_eq!(vec![true], plan.projected_root_presence); - assert_eq!( - ProjectionMask::leaves(&parquet_schema_desc, vec![1]), - plan.mask - ); - } - - #[test] - fn test_mixed_prefix_and_fallback_paths() { - let parquet_schema_desc = build_test_variant_parent_schema(); - let nested_paths = vec![ - vec!["j".to_string(), "a".to_string(), "x".to_string()], - vec!["j".to_string(), "b".to_string(), "x".to_string()], - ]; - let projection = ParquetReadColumns::from_deduped(vec![ - ParquetReadColumn::new(0).with_nested_paths(nested_paths), - ]); - - let plan = build_projection_plan(&projection, &parquet_schema_desc); - - assert_eq!(vec![true], plan.projected_root_presence); - assert_eq!( - ProjectionMask::leaves(&parquet_schema_desc, vec![0, 1]), - plan.mask - ); - } - - #[test] - fn test_fallback_selects_multiple_variant_parents() { - let parquet_schema_desc = build_test_two_variant_parents_schema(); - let nested_paths = vec![ - vec!["j".to_string(), "a".to_string(), "y".to_string()], - vec!["j".to_string(), "b".to_string(), "d".to_string()], - vec!["j".to_string(), "a".to_string(), "x".to_string()], - ]; - let projection = ParquetReadColumns::from_deduped(vec![ - ParquetReadColumn::new(0).with_nested_paths(nested_paths), - ]); - - let plan = build_projection_plan(&projection, &parquet_schema_desc); - - assert_eq!(vec![true], plan.projected_root_presence); - assert_eq!( - ProjectionMask::leaves(&parquet_schema_desc, vec![0, 1]), - plan.mask - ); - } - - #[test] - fn test_nested_paths_fallback_to_variant_parent_by_default() { - let parquet_schema_desc = build_test_variant_parent_schema(); - let projection = - ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths( - vec![vec!["j".to_string(), "a".to_string(), "x".to_string()]], - )]); - - let plan = build_projection_plan(&projection, &parquet_schema_desc); - - assert_eq!(vec![true], plan.projected_root_presence); - assert_eq!( - ProjectionMask::leaves(&parquet_schema_desc, vec![0]), - plan.mask - ); - } - - #[test] - fn test_non_variant_parent_does_not_fallback() { - let parquet_schema_desc = build_test_nested_parquet_schema(); - let projection = - ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths( - vec![vec!["j".to_string(), "a".to_string(), "x".to_string()]], - )]); - - let plan = build_projection_plan(&projection, &parquet_schema_desc); - - assert_eq!(vec![false], plan.projected_root_presence); - } - - #[test] - fn test_utf8_parent_does_not_fallback() { - let parquet_schema_desc = build_test_utf8_parent_schema(); - let projection = - ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths( - vec![vec!["j".to_string(), "a".to_string(), "x".to_string()]], - )]); - - let plan = build_projection_plan(&projection, &parquet_schema_desc); - - assert_eq!(vec![false], plan.projected_root_presence); - } - - #[test] - fn test_root_variant_does_not_fallback() { - let parquet_schema_desc = build_test_root_variant_schema(); - let projection = ParquetReadColumns::from_deduped(vec![ - ParquetReadColumn::new(0) - .with_nested_paths(vec![vec!["j".to_string(), "a".to_string()]]), - ]); - - let plan = build_projection_plan(&projection, &parquet_schema_desc); - - assert_eq!(vec![false], plan.projected_root_presence); - } - // Test schema: // schema // |- j @@ -859,127 +814,44 @@ mod tests { ))) } - // Test schema: - // schema - // `- j - // |- a: BYTE_ARRAY - // `- b - // `- x: INT64 - fn build_test_variant_parent_schema() -> SchemaDescriptor { - let leaf_a = Arc::new( - Type::primitive_type_builder("a", parquet::basic::Type::BYTE_ARRAY) - .with_repetition(Repetition::REQUIRED) - .build() - .unwrap(), - ); - let leaf_x = Arc::new( - Type::primitive_type_builder("x", parquet::basic::Type::INT64) - .with_repetition(Repetition::REQUIRED) - .build() - .unwrap(), - ); - let group_b = Arc::new( - Type::group_type_builder("b") - .with_repetition(Repetition::REQUIRED) - .with_fields(vec![leaf_x]) - .build() - .unwrap(), - ); - let root_j = Arc::new( - Type::group_type_builder("j") - .with_repetition(Repetition::REQUIRED) - .with_fields(vec![leaf_a, group_b]) - .build() - .unwrap(), - ); - let schema = Arc::new( - Type::group_type_builder("schema") - .with_fields(vec![root_j]) - .build() - .unwrap(), - ); - - SchemaDescriptor::new(schema) + fn source_schema(policies: &[Option]) -> ArrowSchema { + ArrowSchema::new( + policies + .iter() + .enumerate() + .map(|(index, policy)| { + let field = Field::new( + format!("root_{index}"), + DataType::Struct(Fields::empty()), + true, + ); + match policy { + None => field, + Some(NestedSelectionPolicy::Json2) => { + field.with_extension_type(Json2ExtensionType::new(Arc::new( + JsonMetadata::new(JsonSettings::default()), + ))) + } + } + }) + .collect::>(), + ) } - // Test schema: - // schema - // `- j - // |- a: BYTE_ARRAY - // `- b: BYTE_ARRAY - fn build_test_two_variant_parents_schema() -> SchemaDescriptor { - let leaf_a = Arc::new( - Type::primitive_type_builder("a", parquet::basic::Type::BYTE_ARRAY) - .with_repetition(Repetition::REQUIRED) - .build() - .unwrap(), - ); - let leaf_b = Arc::new( - Type::primitive_type_builder("b", parquet::basic::Type::BYTE_ARRAY) - .with_repetition(Repetition::REQUIRED) - .build() - .unwrap(), - ); - let root_j = Arc::new( - Type::group_type_builder("j") - .with_repetition(Repetition::REQUIRED) - .with_fields(vec![leaf_a, leaf_b]) - .build() - .unwrap(), - ); - let schema = Arc::new( - Type::group_type_builder("schema") - .with_fields(vec![root_j]) - .build() - .unwrap(), - ); - - SchemaDescriptor::new(schema) + fn build_projection_plan( + cols: &ParquetReadColumns, + parquet_schema: &SchemaDescriptor, + policies: &[Option], + ) -> ProjectionMaskPlan { + super::build_projection_plan(cols, parquet_schema, &source_schema(policies)).unwrap() } - fn build_test_utf8_parent_schema() -> SchemaDescriptor { - let leaf_a = Arc::new( - Type::primitive_type_builder("a", parquet::basic::Type::BYTE_ARRAY) - .with_repetition(Repetition::REQUIRED) - .with_logical_type(Some(LogicalType::String)) - .with_converted_type(ConvertedType::UTF8) - .build() - .unwrap(), - ); - let root_j = Arc::new( - Type::group_type_builder("j") - .with_repetition(Repetition::REQUIRED) - .with_fields(vec![leaf_a]) - .build() - .unwrap(), - ); - let schema = Arc::new( - Type::group_type_builder("schema") - .with_fields(vec![root_j]) - .build() - .unwrap(), - ); - - SchemaDescriptor::new(schema) - } - - // Test schema: - // schema - // `- j: BYTE_ARRAY - fn build_test_root_variant_schema() -> SchemaDescriptor { - let root_j = Arc::new( - Type::primitive_type_builder("j", parquet::basic::Type::BYTE_ARRAY) - .with_repetition(Repetition::REQUIRED) - .build() - .unwrap(), - ); - let schema = Arc::new( - Type::group_type_builder("schema") - .with_fields(vec![root_j]) - .build() - .unwrap(), - ); - - SchemaDescriptor::new(schema) + fn build_parquet_leaves_indices_with_policies( + parquet_schema: &SchemaDescriptor, + projection: &ParquetReadColumns, + policies: &[Option], + ) -> (Vec, HashSet) { + super::build_parquet_leaves_indices(parquet_schema, projection, &source_schema(policies)) + .unwrap() } } diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 59ef12d9e1e..e3232c189d6 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -85,7 +85,7 @@ use crate::sst::parquet::file_range::{ }; use crate::sst::parquet::flat_format::{FlatReadFormat, primary_key_column_index}; use crate::sst::parquet::format::{INTERNAL_COLUMN_NUM, need_override_sequence}; -use crate::sst::parquet::json_align::{NestedSchemaAligner, ProjectedRecordBatchStream}; +use crate::sst::parquet::json_align::{AlignMode, JsonSchemaAligner, ProjectedRecordBatchStream}; use crate::sst::parquet::metadata::MetadataLoader; use crate::sst::parquet::prefilter::{ PrefilterContextBuilder, build_reader_filter_plan, execute_prefilter, @@ -522,13 +522,14 @@ impl ParquetReaderBuilder { let file_metadata = parquet_meta.file_metadata(); let parquet_schema_desc = file_metadata.schema_descr(); - let file_schema = + let file_schema = Arc::new( parquet_to_arrow_schema(parquet_schema_desc, file_metadata.key_value_metadata()) - .context(ParquetToArrowSchemaSnafu { file: &file_path })?; + .context(ParquetToArrowSchemaSnafu { file: &file_path })?, + ); let mut read_format = FlatReadFormat::new( region_meta.clone(), read_cols, - Some(Arc::new(file_schema)), + Some(file_schema.clone()), &file_path, skip_auto_convert, )?; @@ -565,7 +566,8 @@ impl ParquetReaderBuilder { // Computes the projection mask. let parquet_read_cols = read_format.parquet_read_columns(); - let projection_plan = build_projection_plan(parquet_read_cols, parquet_schema_desc); + let projection_plan = + build_projection_plan(parquet_read_cols, parquet_schema_desc, &file_schema)?; let has_nested_projection = parquet_read_cols.has_nested(); let selection = self .row_groups_to_read(&read_format, &parquet_meta, &mut metrics.filter_metrics) @@ -2070,12 +2072,20 @@ impl RowGroupReaderBuilder { return Ok(stream); } - Ok(NestedSchemaAligner::new( + let mode = if self.json2_rewrite_targets.is_empty() { + AlignMode::AlignToSchema + } else { + AlignMode::Rewrite { + columns: self.json2_rewrite_targets.clone(), + } + }; + + Ok(JsonSchemaAligner::new( stream, self.projection.projected_root_presence.clone(), self.output_schema.clone(), + mode, )? - .with_json2_rewrite_targets(&self.json2_rewrite_targets)? .boxed()) } @@ -2693,7 +2703,17 @@ mod tests { let output_schema = read_format.arrow_schema().clone(); let parquet_schema = parquet_meta.file_metadata().schema_descr(); - let projection = build_projection_plan(read_format.parquet_read_columns(), parquet_schema); + let source_schema = parquet_to_arrow_schema( + parquet_schema, + parquet_meta.file_metadata().key_value_metadata(), + ) + .unwrap(); + let projection = build_projection_plan( + read_format.parquet_read_columns(), + parquet_schema, + &source_schema, + ) + .unwrap(); let arrow_metadata = ArrowReaderMetadata::try_new(parquet_meta.clone(), ArrowReaderOptions::new()).unwrap(); ( @@ -2989,7 +3009,8 @@ mod tests { ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths( vec![vec!["j".to_string(), "a".to_string(), "x".to_string()]], )]); - let projection_plan = build_projection_plan(&projection, parquet_schema); + let projection_plan = + build_projection_plan(&projection, parquet_schema, batch.schema_ref()).unwrap(); assert_eq!(vec![true], projection_plan.projected_root_presence); assert_eq!( projection_plan.mask, From 7aacc947b30deb79aff8fd18bfe8a9eef942bedf Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Wed, 16 Sep 2026 06:15:52 +0000 Subject: [PATCH 3/4] perf(mito2): lazily extract sparse primary key index values (#9176) Signed-off-by: Lei, HUANG --- src/index/src/bloom_filter/creator.rs | 65 +++++ src/mito-codec/src/index.rs | 45 ++- src/mito-codec/src/row_converter/sparse.rs | 3 + .../src/row_converter/sparse/checked.rs | 244 ++++++++++++++++ src/mito2/Cargo.toml | 5 + src/mito2/benches/bench_index_update.rs | 276 ++++++++++++++++++ src/mito2/src/sst/index.rs | 63 +--- .../src/sst/index/bloom_filter/creator.rs | 79 +++-- .../src/sst/index/inverted_index/creator.rs | 79 +++-- src/mito2/src/sst/index/primary_key.rs | 108 +++++++ src/mito2/src/sst/index/sparse_test.rs | 198 +++++++++++++ src/mito2/src/test_util/bench_util.rs | 2 + 12 files changed, 1020 insertions(+), 147 deletions(-) create mode 100644 src/mito-codec/src/row_converter/sparse/checked.rs create mode 100644 src/mito2/benches/bench_index_update.rs create mode 100644 src/mito2/src/sst/index/primary_key.rs create mode 100644 src/mito2/src/sst/index/sparse_test.rs diff --git a/src/index/src/bloom_filter/creator.rs b/src/index/src/bloom_filter/creator.rs index 5a9b5a0b36a..8bbd4cf1c13 100644 --- a/src/index/src/bloom_filter/creator.rs +++ b/src/index/src/bloom_filter/creator.rs @@ -149,6 +149,35 @@ impl BloomFilterCreator { Ok(()) } + /// Adds `nrows` copies of a single borrowed value (or null), copying it only when + /// it is new to a segment. Row counts advance for nulls as well. + pub async fn push_n_row_elem(&mut self, mut nrows: usize, elem: Option<&[u8]>) -> Result<()> { + while nrows > 0 { + let rows_to_seg_end = + self.rows_per_segment - (self.accumulated_row_count % self.rows_per_segment); + let rows_to_push = nrows.min(rows_to_seg_end); + nrows -= rows_to_push; + self.accumulated_row_count += rows_to_push; + + if let Some(elem) = elem + && !self.cur_seg_distinct_elems.contains(elem) + { + self.cur_seg_distinct_elems.insert(elem.to_vec()); + self.cur_seg_distinct_elems_mem_usage += elem.len(); + self.global_memory_usage + .fetch_add(elem.len(), Ordering::Relaxed); + } + if self + .accumulated_row_count + .is_multiple_of(self.rows_per_segment) + { + self.finalize_segment().await?; + self.finalized_row_count = self.accumulated_row_count; + } + } + Ok(()) + } + /// Adds a row of elements to the bloom filter. If the number of accumulated rows /// reaches `rows_per_segment`, it finalizes the current segment. pub async fn push_row_elems(&mut self, elems: impl IntoIterator) -> Result<()> { @@ -420,6 +449,42 @@ mod tests { } } + #[tokio::test] + async fn borrowed_single_value_matches_owned_rows_across_segments() { + let make_creator = || { + BloomFilterCreator::new( + 3, + 0.01, + Arc::new(MockExternalTempFileProvider::new()), + Arc::new(AtomicUsize::new(0)), + None, + ) + }; + let mut borrowed = make_creator(); + let mut owned = make_creator(); + // Zero rows, nulls, empty values, duplicates and runs crossing segment boundaries. + for (rows, elem) in [ + (0, Some(b"ignored".as_slice())), + (1, None), + (5, Some(b"".as_slice())), + (1, Some(b"".as_slice())), + (8, Some(b"label".as_slice())), + (1, None), + ] { + borrowed.push_n_row_elem(rows, elem).await.unwrap(); + owned + .push_n_row_elems(rows, elem.map(<[u8]>::to_vec)) + .await + .unwrap(); + assert_eq!(borrowed.memory_usage(), owned.memory_usage()); + } + let mut borrowed_blob = Cursor::new(Vec::new()); + let mut owned_blob = Cursor::new(Vec::new()); + borrowed.finish(&mut borrowed_blob).await.unwrap(); + owned.finish(&mut owned_blob).await.unwrap(); + assert_eq!(borrowed_blob.into_inner(), owned_blob.into_inner()); + } + #[tokio::test] async fn test_final_seg_all_null() { let mut writer = Cursor::new(Vec::new()); diff --git a/src/mito-codec/src/index.rs b/src/mito-codec/src/index.rs index 539a3adeddb..e0f2c819afe 100644 --- a/src/mito-codec/src/index.rs +++ b/src/mito-codec/src/index.rs @@ -24,13 +24,56 @@ use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::ColumnMetadata; use store_api::storage::ColumnId; -use crate::error::{FieldTypeMismatchSnafu, IndexEncodeNullSnafu, Result}; +use crate::error::{ + FieldTypeMismatchSnafu, IndexEncodeNullSnafu, InvalidSparsePrimaryKeySnafu, Result, +}; +use crate::row_converter::sparse::{ + RESERVED_COLUMN_ID_TABLE_ID, RESERVED_COLUMN_ID_TSID, SparsePrimaryKeyView, +}; use crate::row_converter::{PrimaryKeyCodec, SortField, build_primary_key_codec_with_fields}; /// Encodes index values according to their data types for sorting and storage use. pub struct IndexValueCodec; impl IndexValueCodec { + /// Extracts one sparse PK column in index format, without constructing a Value. + /// Numeric reserved fields borrow the PK; strings are unchunked into the reusable buffer. + /// Missing and null labels return None, whereas an empty string returns Some(&[]). + pub fn encode_sparse_value<'a, 'pk: 'a>( + pk: &mut SparsePrimaryKeyView<'pk, '_>, + column_id: ColumnId, + buffer: &'a mut Vec, + ) -> Result> { + let Some(encoded) = pk.encoded_value(column_id)? else { + return Ok(None); + }; + if encoded[0] == 0 { + return Ok(None); + } + if matches!( + column_id, + RESERVED_COLUMN_ID_TABLE_ID | RESERVED_COLUMN_ID_TSID + ) { + return Ok(Some(encoded)); + } + + buffer.clear(); + // The view has checked the Option marker, bytes marker and every chunk length. + // Reserving once avoids growing the buffer for each 8-byte chunk. + buffer.reserve(encoded.len() - 2); + for chunk in encoded[2..].chunks_exact(9) { + let len = usize::from(chunk[8]).min(8); + buffer.extend_from_slice(&chunk[..len]); + } + std::str::from_utf8(buffer).map_err(|_| { + InvalidSparsePrimaryKeySnafu { + reason: "label is not valid UTF-8", + } + .build() + })?; + Ok(Some(buffer.as_slice())) + } + /// Serializes a non-null `ValueRef` using the data type defined in `SortField` and writes /// the result into a buffer. /// diff --git a/src/mito-codec/src/row_converter/sparse.rs b/src/mito-codec/src/row_converter/sparse.rs index c70293886f8..b2b3e0e3778 100644 --- a/src/mito-codec/src/row_converter/sparse.rs +++ b/src/mito-codec/src/row_converter/sparse.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod checked; + use std::collections::{HashMap, HashSet}; use std::sync::Arc; @@ -34,6 +36,7 @@ use crate::error::{ use crate::key_values::KeyValue; use crate::primary_key_filter::SparsePrimaryKeyFilter; use crate::row_converter::dense::SortField; +pub use crate::row_converter::sparse::checked::SparsePrimaryKeyView; use crate::row_converter::{CompositeValues, PrimaryKeyCodec, PrimaryKeyFilter}; /// A codec for sparse key of metrics. diff --git a/src/mito-codec/src/row_converter/sparse/checked.rs b/src/mito-codec/src/row_converter/sparse/checked.rs new file mode 100644 index 00000000000..63666f91c4c --- /dev/null +++ b/src/mito-codec/src/row_converter/sparse/checked.rs @@ -0,0 +1,244 @@ +// Copyright 2023 Greptime 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 snafu::{OptionExt, ensure}; +use store_api::storage::ColumnId; + +use crate::error::{InvalidSparsePrimaryKeySnafu, Result}; +use crate::row_converter::sparse::{ + COLUMN_ID_ENCODE_SIZE, RESERVED_COLUMN_ID_TABLE_ID, RESERVED_COLUMN_ID_TSID, + SparseOffsetsCache, TABLE_ID_VALUE_OFFSET, TAGS_START_OFFSET, TSID_VALUE_OFFSET, +}; + +/// A borrowing view of one sparse primary key, with checked, lazy field lookup. +/// +/// The scratch cache is reset on construction and retains its capacity between keys. +/// Ordinary labels can be skipped without schema lookup: their encoding is always String. +pub struct SparsePrimaryKeyView<'a, 'b> { + pk: &'a [u8], + cache: &'b mut SparseOffsetsCache, +} + +impl<'a, 'b> SparsePrimaryKeyView<'a, 'b> { + /// Validates the fixed prefix and binds fresh offset discovery to this key. + pub fn new(pk: &'a [u8], cache: &'b mut SparseOffsetsCache) -> Result { + ensure!( + pk.len() >= TAGS_START_OFFSET + && pk[..COLUMN_ID_ENCODE_SIZE] == RESERVED_COLUMN_ID_TABLE_ID.to_be_bytes() + && pk[TSID_VALUE_OFFSET - COLUMN_ID_ENCODE_SIZE..TSID_VALUE_OFFSET] + == RESERVED_COLUMN_ID_TSID.to_be_bytes() + && pk[TABLE_ID_VALUE_OFFSET] == 1 + && pk[TSID_VALUE_OFFSET] == 1, + InvalidSparsePrimaryKeySnafu { + reason: "invalid table_id/tsid prefix" + } + ); + cache.clear(); + Ok(Self { pk, cache }) + } + + /// Returns the encoded value, including its null marker, or None for an absent label. + /// The result borrows the key, not the scratch cache. + pub fn encoded_value(&mut self, column_id: ColumnId) -> Result> { + match column_id { + RESERVED_COLUMN_ID_TABLE_ID => { + return Ok(Some( + &self.pk[TABLE_ID_VALUE_OFFSET..TSID_VALUE_OFFSET - COLUMN_ID_ENCODE_SIZE], + )); + } + RESERVED_COLUMN_ID_TSID => { + return Ok(Some(&self.pk[TSID_VALUE_OFFSET..TAGS_START_OFFSET])); + } + _ => {} + } + if let Some(offset) = self.cache.get(column_id) { + return encoded_label(&self.pk[offset..]).map(Some); + } + while !self.cache.finished && self.cache.cursor < self.pk.len() { + let bytes = &self.pk[self.cache.cursor..]; + let id = bytes.first_chunk::().context( + InvalidSparsePrimaryKeySnafu { + reason: "truncated label column id", + }, + )?; + let id = u32::from_be_bytes(*id); + let offset = self.cache.cursor + COLUMN_ID_ENCODE_SIZE; + let value = encoded_label(&self.pk[offset..])?; + self.cache.insert(id, offset); + self.cache.cursor = offset + value.len(); + if id == column_id { + return Ok(Some(value)); + } + } + self.cache.finished = true; + Ok(None) + } +} + +/// Finds the end of an Option without allocating or reading its payload. +/// Unlike memcomparable's skip_bytes, all advances are checked for truncated input. +fn encoded_label(bytes: &[u8]) -> Result<&[u8]> { + match bytes.first() { + Some(0) => return Ok(&bytes[..1]), + Some(1) => {} + _ => { + return InvalidSparsePrimaryKeySnafu { + reason: "invalid label null marker", + } + .fail(); + } + } + match bytes.get(1) { + Some(0) => return Ok(&bytes[..2]), + Some(1) => {} + _ => { + return InvalidSparsePrimaryKeySnafu { + reason: "invalid label bytes marker", + } + .fail(); + } + } + let mut end = 2; + loop { + let chunk = bytes + .get(end..end + 9) + .context(InvalidSparsePrimaryKeySnafu { + reason: "truncated label chunk", + })?; + end += 9; + match chunk[8] { + 1..=8 => return Ok(&bytes[..end]), + 9 => {} + _ => { + return InvalidSparsePrimaryKeySnafu { + reason: "invalid label chunk length", + } + .fail(); + } + } + } +} + +#[cfg(test)] +mod tests { + use datatypes::data_type::ConcreteDataType; + use datatypes::value::ValueRef; + + use super::*; + use crate::index::IndexValueCodec; + use crate::row_converter::{PrimaryKeyCodec, SortField, SparsePrimaryKeyCodec}; + + #[test] + fn index_bytes_match_full_decode_across_chunks_and_lookup_orders() { + let codec = SparsePrimaryKeyCodec::schemaless(); + let mut cache = SparseOffsetsCache::new(); + let mut buffer = Vec::new(); + // Reusing the cache across different keys must not reuse their field offsets. + for len in [0, 1, 7, 8, 9, 16, 24, 65] { + let label = "x".repeat(len); + let values = [ + (RESERVED_COLUMN_ID_TABLE_ID, ValueRef::UInt32(42)), + (RESERVED_COLUMN_ID_TSID, ValueRef::UInt64(u64::MAX)), + (90, ValueRef::String(&label)), + (7, ValueRef::String("")), + (1, ValueRef::Null), + (3, ValueRef::String("last中文")), + ]; + let mut pk = Vec::new(); + codec.encode_value_refs(&values, &mut pk).unwrap(); + // Exercise the offset cache's overflow storage as well as its inline entries. + codec + .encode_raw_tag_value((100..140).map(|id| (id, label.as_bytes())), &mut pk) + .unwrap(); + // Writers normally omit null labels; existing encoded nulls must remain null too. + pk.extend_from_slice(&2_u32.to_be_bytes()); + pk.push(0); + let decoded = codec.decode(&pk).unwrap().into_sparse(); + for order in [ + [ + RESERVED_COLUMN_ID_TABLE_ID, + RESERVED_COLUMN_ID_TSID, + 90, + 7, + 1, + 3, + ], + [ + 3, + 1, + 7, + 90, + RESERVED_COLUMN_ID_TSID, + RESERVED_COLUMN_ID_TABLE_ID, + ], + ] { + let mut view = SparsePrimaryKeyView::new(&pk, &mut cache).unwrap(); + for id in order.into_iter().chain([139, 100, 132, 2]) { + let expected = decoded.get(&id).filter(|v| !v.is_null()).map(|value| { + let field = SortField::new(match id { + RESERVED_COLUMN_ID_TABLE_ID => ConcreteDataType::uint32_datatype(), + RESERVED_COLUMN_ID_TSID => ConcreteDataType::uint64_datatype(), + _ => ConcreteDataType::string_datatype(), + }); + let mut bytes = Vec::new(); + IndexValueCodec::encode_nonnull_value( + value.as_value_ref(), + &field, + &mut bytes, + ) + .unwrap(); + bytes + }); + let actual = + IndexValueCodec::encode_sparse_value(&mut view, id, &mut buffer).unwrap(); + assert_eq!(actual, expected.as_deref(), "len={len}, column={id}"); + } + assert!(view.encoded_value(999).unwrap().is_none()); + } + } + } + + #[test] + fn malformed_sparse_values_return_errors() { + let codec = SparsePrimaryKeyCodec::schemaless(); + let mut pk = Vec::new(); + codec.encode_internal(1, 2, &mut pk).unwrap(); + codec + .encode_raw_tag_value([(1, b"0123456789".as_slice())].into_iter(), &mut pk) + .unwrap(); + let mut cache = SparseOffsetsCache::new(); + for end in 0..pk.len() { + // A complete reserved prefix is a valid key with no labels. + if end == TAGS_START_OFFSET { + continue; + } + let result = SparsePrimaryKeyView::new(&pk[..end], &mut cache) + .and_then(|mut view| view.encoded_value(1)); + assert!(result.is_err(), "truncation at {end}"); + } + for (offset, value) in [(0, 0), (4, 0), (13, 2), (26, 2), (27, 2), (36, 0), (45, 10)] { + let mut invalid = pk.clone(); + invalid[offset] = value; + assert!( + SparsePrimaryKeyView::new(&invalid, &mut cache) + .and_then(|mut view| view.encoded_value(1)) + .is_err() + ); + } + let mut invalid_utf8 = pk; + invalid_utf8[28] = 0xff; + let mut view = SparsePrimaryKeyView::new(&invalid_utf8, &mut cache).unwrap(); + assert!(IndexValueCodec::encode_sparse_value(&mut view, 1, &mut Vec::new()).is_err()); + } +} diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index 100a6fb1023..ccc3fd6f176 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -145,5 +145,10 @@ required-features = ["test"] name = "bench_wal_encode" harness = false +[[bench]] +name = "bench_index_update" +harness = false +required-features = ["testing"] + [package.metadata.cargo-udeps.ignore] normal = ["aquamarine"] diff --git a/src/mito2/benches/bench_index_update.rs b/src/mito2/benches/bench_index_update.rs new file mode 100644 index 00000000000..eba89cf60ff --- /dev/null +++ b/src/mito2/benches/bench_index_update.rs @@ -0,0 +1,276 @@ +// Copyright 2023 Greptime 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. + +//! Measures real index updates, excluding input generation, creator setup and cleanup. +//! +//! Run on the baseline revision with `--save-baseline before`, then on the candidate +//! with `--baseline before` (arguments after `--` in the command below): +//! `CARGO_PROFILE_BENCH_DEBUG=0 cargo bench -p mito2 --features testing --bench bench_index_update`. + +use std::collections::HashSet; +use std::num::NonZeroUsize; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use api::v1::SemanticType; +use common_test_util::temp_dir::create_temp_dir; +use criterion::{Criterion, Throughput, criterion_group, criterion_main}; +use datatypes::arrow::array::{ + ArrayRef, BinaryDictionaryBuilder, TimestampMillisecondArray, UInt8Array, UInt64Array, +}; +use datatypes::arrow::datatypes::UInt32Type; +use datatypes::arrow::record_batch::RecordBatch; +use datatypes::data_type::ConcreteDataType; +use datatypes::schema::{ColumnSchema, SkippingIndexOptions}; +use mito_codec::row_converter::SparsePrimaryKeyCodec; +use mito2::sst::index::intermediate::IntermediateManager; +use mito2::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; +use mito2::test_util::bench_util::{BloomFilterIndexer, InvertedIndexer}; +use store_api::codec::PrimaryKeyEncoding; +use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef}; +use store_api::storage::consts::ReservedColumnId; +use store_api::storage::{ColumnId, FileId, RegionId}; + +const ROWS: usize = 4096; +const SEGMENT_ROWS: usize = 1024; + +struct Shape { + name: &'static str, + tags: u32, + indexed_tags: u32, + inverted: bool, + bloom: bool, + rows_per_key: usize, +} + +fn input(shape: &Shape) -> (RegionMetadataRef, HashSet, RecordBatch) { + let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1)); + let mut pk = vec![ReservedColumnId::table_id(), ReservedColumnId::tsid()]; + let mut inverted_columns = HashSet::new(); + let mut add_tag = |id, name: String, data_type, indexed| { + let mut schema = ColumnSchema::new(name, data_type, true).with_inverted_index(false); + if indexed && shape.inverted { + inverted_columns.insert(id); + } + if indexed && shape.bloom { + schema = schema + .with_skipping_options(SkippingIndexOptions { + granularity: SEGMENT_ROWS as _, + ..Default::default() + }) + .unwrap(); + } + builder.push_column_metadata(ColumnMetadata { + column_schema: schema, + semantic_type: SemanticType::Tag, + column_id: id, + }); + }; + add_tag( + pk[0], + "__table_id".into(), + ConcreteDataType::uint32_datatype(), + shape.indexed_tags == 0, + ); + add_tag( + pk[1], + "__tsid".into(), + ConcreteDataType::uint64_datatype(), + false, + ); + for id in 0..shape.tags { + // A single indexed label is deliberately last in the key. + add_tag( + id, + format!("tag_{id:03}"), + ConcreteDataType::string_datatype(), + id >= shape.tags - shape.indexed_tags, + ); + pk.push(id); + } + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: shape.tags, + }) + .primary_key(pk) + .primary_key_encoding(PrimaryKeyEncoding::Sparse); + let metadata = Arc::new(builder.build().unwrap()); + let codec = SparsePrimaryKeyCodec::schemaless(); + let mut keys = BinaryDictionaryBuilder::::new(); + for series in 0..ROWS / shape.rows_per_key { + let mut key = Vec::new(); + codec + .encode_internal((series / 128) as u32, series as u64, &mut key) + .unwrap(); + let labels: Vec<_> = (0..shape.tags) + .map(|id| (id, format!("tag-{id:03}-value-{series:010}"))) + .collect(); + codec + .encode_raw_tag_value( + labels.iter().map(|(id, value)| (*id, value.as_bytes())), + &mut key, + ) + .unwrap(); + for _ in 0..shape.rows_per_key { + keys.append(&key).unwrap(); + } + } + let schema = to_flat_sst_arrow_schema( + &metadata, + &FlatSchemaOptions::from_encoding(PrimaryKeyEncoding::Sparse), + ); + let columns: Vec = vec![ + Arc::new(TimestampMillisecondArray::from_iter_values(0..ROWS as i64)), + Arc::new(keys.finish()), + Arc::new(UInt64Array::from(vec![1; ROWS])), + Arc::new(UInt8Array::from(vec![1; ROWS])), + ]; + let batch = RecordBatch::try_new(schema, columns).unwrap(); + (metadata, inverted_columns, batch) +} + +fn bench_index_update(c: &mut Criterion) { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let dir = create_temp_dir("bench_index_update"); + let intermediate = runtime + .block_on(IntermediateManager::init_fs(dir.path().to_str().unwrap())) + .unwrap(); + let mut group = c.benchmark_group("sparse_index_update"); + group.sample_size(30); + group.warm_up_time(Duration::from_secs(1)); + group.measurement_time(Duration::from_secs(3)); + group.throughput(Throughput::Elements(ROWS as u64)); + for shape in [ + Shape { + name: "table_id_bloom_40tags", + tags: 40, + indexed_tags: 0, + inverted: false, + bloom: true, + rows_per_key: 1, + }, + Shape { + name: "inverted_one_40tags", + tags: 40, + indexed_tags: 1, + inverted: true, + bloom: false, + rows_per_key: 1, + }, + Shape { + name: "bloom_one_40tags", + tags: 40, + indexed_tags: 1, + inverted: false, + bloom: true, + rows_per_key: 1, + }, + Shape { + name: "both_one_40tags", + tags: 40, + indexed_tags: 1, + inverted: true, + bloom: true, + rows_per_key: 1, + }, + Shape { + name: "both_all_40tags", + tags: 40, + indexed_tags: 40, + inverted: true, + bloom: true, + rows_per_key: 1, + }, + Shape { + name: "both_all_40tags_8rpk", + tags: 40, + indexed_tags: 40, + inverted: true, + bloom: true, + rows_per_key: 8, + }, + Shape { + name: "both_one_10tags", + tags: 10, + indexed_tags: 1, + inverted: true, + bloom: true, + rows_per_key: 1, + }, + Shape { + name: "both_all_10tags", + tags: 10, + indexed_tags: 10, + inverted: true, + bloom: true, + rows_per_key: 1, + }, + ] { + let (metadata, inverted_columns, batch) = input(&shape); + group.bench_function(shape.name, |b| { + b.iter_custom(|iterations| { + runtime.block_on(async { + let mut elapsed = Duration::ZERO; + for _ in 0..iterations { + let file_id = FileId::random(); + let mut inverted = shape.inverted.then(|| { + InvertedIndexer::new( + file_id, + &metadata, + intermediate.clone(), + None, + NonZeroUsize::new(SEGMENT_ROWS).unwrap(), + inverted_columns.clone(), + ) + }); + let mut bloom = if shape.bloom { + BloomFilterIndexer::new(file_id, &metadata, intermediate.clone(), None) + .unwrap() + } else { + None + }; + let start = Instant::now(); + if let Some(indexer) = &mut inverted { + indexer.update_flat(&batch).await.unwrap(); + } + if let Some(indexer) = &mut bloom { + indexer.update_flat(&batch).await.unwrap(); + } + elapsed += start.elapsed(); + if let Some(indexer) = &mut inverted { + indexer.abort().await.unwrap(); + } + if let Some(indexer) = &mut bloom { + indexer.abort().await.unwrap(); + } + } + elapsed + }) + }); + }); + } + group.finish(); +} + +criterion_group!(benches, bench_index_update); +criterion_main!(benches); diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 268b8387397..c8fcd6fbc7e 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -17,7 +17,10 @@ pub(crate) mod fulltext_index; mod indexer; pub mod intermediate; pub(crate) mod inverted_index; +mod primary_key; pub mod puffin_manager; +#[cfg(test)] +mod sparse_test; mod statistics; pub(crate) mod store; #[cfg(feature = "vector_index")] @@ -30,14 +33,11 @@ use std::sync::Arc; use bloom_filter::creator::BloomFilterIndexer; use common_telemetry::{debug, error, info, warn}; -use datatypes::arrow::array::BinaryArray; use datatypes::arrow::record_batch::RecordBatch; -use mito_codec::index::IndexValuesCodec; -use mito_codec::row_converter::CompositeValues; use object_store::ObjectStore; use puffin_manager::SstPuffinManager; use smallvec::{SmallVec, smallvec}; -use snafu::{OptionExt, ResultExt}; +use snafu::ResultExt; use statistics::{ByteCount, RowCount}; use store_api::metadata::RegionMetadataRef; use store_api::storage::{ColumnId, FileId, RegionId}; @@ -54,8 +54,8 @@ use crate::cache::{CacheManagerRef, CacheStrategy}; use crate::config::VectorIndexConfig; use crate::config::{BloomFilterConfig, FulltextIndexConfig, InvertedIndexConfig}; use crate::error::{ - BuildIndexAsyncSnafu, DecodeSnafu, Error, InvalidRecordBatchSnafu, RegionClosedSnafu, - RegionDroppedSnafu, RegionTruncatedSnafu, Result, + BuildIndexAsyncSnafu, Error, RegionClosedSnafu, RegionDroppedSnafu, RegionTruncatedSnafu, + Result, }; use crate::metrics::{ INDEX_ARTIFACT_CLEANUP_FAILURE_TOTAL, INDEX_CREATE_MEMORY_USAGE, INDEX_PUBLICATION_STALE_TOTAL, @@ -79,8 +79,6 @@ use crate::sst::index::fulltext_index::creator::FulltextIndexer; use crate::sst::index::intermediate::IntermediateManager; use crate::sst::index::inverted_index::creator::InvertedIndexer; use crate::sst::parquet::SstInfo; -use crate::sst::parquet::flat_format::primary_key_column_index; -use crate::sst::parquet::format::PrimaryKeyArray; use crate::worker::WorkerListener; pub(crate) const TYPE_INVERTED_INDEX: &str = "inverted_index"; @@ -1450,55 +1448,6 @@ impl IndexBuildScheduler { } } -/// Decodes primary keys from a flat format RecordBatch. -/// Returns a list of (decoded_pk_value, count) tuples where count is the number of occurrences. -pub(crate) fn decode_primary_keys_with_counts( - batch: &RecordBatch, - codec: &IndexValuesCodec, -) -> Result> { - let primary_key_index = primary_key_column_index(batch.num_columns()); - let pk_dict_array = batch - .column(primary_key_index) - .as_any() - .downcast_ref::() - .context(InvalidRecordBatchSnafu { - reason: "Primary key column is not a dictionary array", - })?; - let pk_values_array = pk_dict_array - .values() - .as_any() - .downcast_ref::() - .context(InvalidRecordBatchSnafu { - reason: "Primary key values are not binary array", - })?; - let keys = pk_dict_array.keys(); - - // Decodes primary keys and count consecutive occurrences - let mut result: Vec<(CompositeValues, usize)> = Vec::new(); - let mut prev_key: Option = None; - - let pk_indices = keys.values(); - for ¤t_key in pk_indices.iter().take(keys.len()) { - // Checks if current key is the same as previous key - if let Some(prev) = prev_key - && prev == current_key - { - // Safety: We already have a key in the result vector. - result.last_mut().unwrap().1 += 1; - continue; - } - - // New key, decodes it. - let pk_bytes = pk_values_array.value(current_key as usize); - let decoded_value = codec.decoder().decode(pk_bytes).context(DecodeSnafu)?; - - result.push((decoded_value, 1)); - prev_key = Some(current_key); - } - - Ok(result) -} - #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/src/mito2/src/sst/index/bloom_filter/creator.rs b/src/mito2/src/sst/index/bloom_filter/creator.rs index 7828cdb9fdd..cfdd751bf7f 100644 --- a/src/mito2/src/sst/index/bloom_filter/creator.rs +++ b/src/mito2/src/sst/index/bloom_filter/creator.rs @@ -24,8 +24,10 @@ use datatypes::vectors::Helper; use index::bloom_filter::creator::BloomFilterCreator; use index::target::IndexTarget; use mito_codec::index::{IndexValueCodec, IndexValuesCodec}; -use mito_codec::row_converter::{CompositeValues, SortField}; +use mito_codec::row_converter::sparse::SparsePrimaryKeyView; +use mito_codec::row_converter::{SortField, SparseOffsetsCache}; use puffin::puffin_manager::{PuffinWriter, PutOptions}; +use smallvec::SmallVec; use snafu::{ResultExt, ensure}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::RegionMetadataRef; @@ -33,17 +35,18 @@ use store_api::storage::{ColumnId, FileId}; use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; use crate::error::{ - BiErrorsSnafu, BloomFilterFinishSnafu, EncodeSnafu, IndexOptionsSnafu, + BiErrorsSnafu, BloomFilterFinishSnafu, DecodeSnafu, EncodeSnafu, IndexOptionsSnafu, OperateAbortedIndexSnafu, PuffinAddBlobSnafu, PushBloomFilterValueSnafu, Result, }; use crate::read::Batch; +use crate::sst::index::TYPE_BLOOM_FILTER_INDEX; use crate::sst::index::bloom_filter::INDEX_BLOB_TYPE; use crate::sst::index::intermediate::{ IntermediateLocation, IntermediateManager, TempFileProvider, }; +use crate::sst::index::primary_key::PrimaryKeyRuns; use crate::sst::index::puffin_manager::SstPuffinWriter; use crate::sst::index::statistics::{ByteCount, RowCount, Statistics}; -use crate::sst::index::{TYPE_BLOOM_FILTER_INDEX, decode_primary_keys_with_counts}; /// The buffer size for the pipe used to send index data to the puffin blob. const PIPE_BUFFER_SIZE_FOR_SENDING_BLOB: usize = 8192; @@ -58,6 +61,9 @@ pub struct BloomFilterIndexer { /// Codec for decoding primary keys. codec: IndexValuesCodec, + /// Scratch storage for extracting indexed tags from sparse primary keys. + pk_offsets: SparseOffsetsCache, + value_buf: Vec, /// Whether the indexing process has been aborted. aborted: bool, @@ -124,6 +130,8 @@ impl BloomFilterIndexer { creators, temp_file_provider, codec, + pk_offsets: SparseOffsetsCache::new(), + value_buf: Vec::new(), aborted: false, stats: Statistics::new(TYPE_BLOOM_FILTER_INDEX), global_memory_usage, @@ -185,7 +193,7 @@ impl BloomFilterIndexer { /// Returns the number of rows and bytes written. /// /// TODO(zhongzc): duplicate with `mito2::sst::index::inverted_index::creator::InvertedIndexCreator` - pub async fn finish( + pub(crate) async fn finish( &mut self, puffin_writer: &mut SstPuffinWriter, ) -> Result<(RowCount, ByteCount)> { @@ -292,7 +300,8 @@ impl BloomFilterIndexer { guard.inc_row_count(n); let is_sparse = self.metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse; - let mut decoded_pks: Option> = None; + let mut sparse_columns: SmallVec<[(ColumnId, &mut BloomFilterCreator); 8]> = + SmallVec::new(); for (col_id, creator) in &mut self.creators { // Safety: `creators` are created from the metadata so it won't be None. @@ -321,45 +330,8 @@ impl BloomFilterIndexer { .context(PushBloomFilterValueSnafu)?; } } else if is_sparse && column_meta.semantic_type == SemanticType::Tag { - // Column not found in batch, tries to decode from primary keys for sparse encoding. - if decoded_pks.is_none() { - decoded_pks = Some(decode_primary_keys_with_counts(batch, &self.codec)?); - } - - let pk_values_with_counts = decoded_pks.as_ref().unwrap(); - let Some(col_info) = self.codec.pk_col_info(*col_id) else { - debug!( - "Column {} not found in primary key during building bloom filter index", - column_name - ); - continue; - }; - let pk_index = col_info.idx; - let field = &col_info.field; - for (decoded, count) in pk_values_with_counts { - let value = match decoded { - CompositeValues::Dense(dense) => dense.get(pk_index).map(|v| &v.1), - CompositeValues::Sparse(sparse) => sparse.get(col_id), - }; - - let elems = value - .filter(|v| !v.is_null()) - .map(|v| { - let mut buf = vec![]; - IndexValueCodec::encode_nonnull_value( - v.as_value_ref(), - field, - &mut buf, - ) - .context(EncodeSnafu)?; - Ok(buf) - }) - .transpose()?; - - creator - .push_n_row_elems(*count, elems) - .await - .context(PushBloomFilterValueSnafu)?; + if self.codec.pk_col_info(*col_id).is_some() { + sparse_columns.push((*col_id, creator)); } } else { debug!( @@ -369,6 +341,25 @@ impl BloomFilterIndexer { } } + if !sparse_columns.is_empty() { + for (pk, count) in PrimaryKeyRuns::try_new(batch)? { + let mut view = + SparsePrimaryKeyView::new(pk, &mut self.pk_offsets).context(DecodeSnafu)?; + for (col_id, creator) in &mut sparse_columns { + let value = IndexValueCodec::encode_sparse_value( + &mut view, + *col_id, + &mut self.value_buf, + ) + .context(DecodeSnafu)?; + creator + .push_n_row_elem(count, value) + .await + .context(PushBloomFilterValueSnafu)?; + } + } + } + Ok(()) } diff --git a/src/mito2/src/sst/index/inverted_index/creator.rs b/src/mito2/src/sst/index/inverted_index/creator.rs index 9cad6fa84ee..0243016fc7a 100644 --- a/src/mito2/src/sst/index/inverted_index/creator.rs +++ b/src/mito2/src/sst/index/inverted_index/creator.rs @@ -27,8 +27,10 @@ use index::inverted_index::create::sort_create::SortIndexCreator; use index::inverted_index::format::writer::InvertedIndexBlobWriter; use index::target::IndexTarget; use mito_codec::index::{IndexValueCodec, IndexValuesCodec}; -use mito_codec::row_converter::{CompositeValues, SortField}; +use mito_codec::row_converter::sparse::SparsePrimaryKeyView; +use mito_codec::row_converter::{SortField, SparseOffsetsCache}; use puffin::puffin_manager::{PuffinWriter, PutOptions}; +use smallvec::SmallVec; use snafu::{ResultExt, ensure}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::RegionMetadataRef; @@ -37,17 +39,18 @@ use tokio::io::duplex; use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; use crate::error::{ - BiErrorsSnafu, EncodeSnafu, IndexFinishSnafu, OperateAbortedIndexSnafu, PuffinAddBlobSnafu, - PushIndexValueSnafu, Result, + BiErrorsSnafu, DecodeSnafu, EncodeSnafu, IndexFinishSnafu, OperateAbortedIndexSnafu, + PuffinAddBlobSnafu, PushIndexValueSnafu, Result, }; use crate::read::Batch; +use crate::sst::index::TYPE_INVERTED_INDEX; use crate::sst::index::intermediate::{ IntermediateLocation, IntermediateManager, TempFileProvider, }; use crate::sst::index::inverted_index::INDEX_BLOB_TYPE; +use crate::sst::index::primary_key::PrimaryKeyRuns; use crate::sst::index::puffin_manager::SstPuffinWriter; use crate::sst::index::statistics::{ByteCount, RowCount, Statistics}; -use crate::sst::index::{TYPE_INVERTED_INDEX, decode_primary_keys_with_counts}; /// The minimum memory usage threshold for one column. const MIN_MEMORY_USAGE_THRESHOLD_PER_COLUMN: usize = 1024 * 1024; // 1MB @@ -66,6 +69,8 @@ pub struct InvertedIndexer { codec: IndexValuesCodec, /// Reusable buffer for encoding index values. value_buf: Vec, + /// Scratch offsets shared by indexed tags of one sparse primary key. + pk_offsets: SparseOffsetsCache, /// Statistics of index creation. stats: Statistics, @@ -124,6 +129,7 @@ impl InvertedIndexer { index_creator, temp_file_provider, value_buf: vec![], + pk_offsets: SparseOffsetsCache::new(), stats: Statistics::new(TYPE_INVERTED_INDEX), aborted: false, memory_usage, @@ -173,7 +179,7 @@ impl InvertedIndexer { guard.inc_row_count(batch.num_rows()); let is_sparse = self.metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse; - let mut decoded_pks: Option> = None; + let mut sparse_columns: SmallVec<[(ColumnId, &str); 8]> = SmallVec::new(); for (col_id, target_key) in &self.indexed_column_ids { let Some(column_meta) = self.metadata.column_by_id(*col_id) else { @@ -213,45 +219,8 @@ impl InvertedIndexer { } } } else if is_sparse && column_meta.semantic_type == SemanticType::Tag { - // Column not found in batch, tries to decode from primary keys for sparse encoding. - if decoded_pks.is_none() { - decoded_pks = Some(decode_primary_keys_with_counts(batch, &self.codec)?); - } - - let pk_values_with_counts = decoded_pks.as_ref().unwrap(); - let Some(col_info) = self.codec.pk_col_info(*col_id) else { - debug!( - "Column {} not found in primary key during building bloom filter index", - column_name - ); - continue; - }; - let pk_index = col_info.idx; - let field = &col_info.field; - for (decoded, count) in pk_values_with_counts { - let value = match decoded { - CompositeValues::Dense(dense) => dense.get(pk_index).map(|v| &v.1), - CompositeValues::Sparse(sparse) => sparse.get(col_id), - }; - - let elem = value - .filter(|v| !v.is_null()) - .map(|v| { - self.value_buf.clear(); - IndexValueCodec::encode_nonnull_value( - v.as_value_ref(), - field, - &mut self.value_buf, - ) - .context(EncodeSnafu)?; - Ok(self.value_buf.as_slice()) - }) - .transpose()?; - - self.index_creator - .push_with_name_n(target_key, elem, *count) - .await - .context(PushIndexValueSnafu)?; + if self.codec.pk_col_info(*col_id).is_some() { + sparse_columns.push((*col_id, target_key)); } } else { debug!( @@ -261,12 +230,32 @@ impl InvertedIndexer { } } + if !sparse_columns.is_empty() { + for (pk, count) in PrimaryKeyRuns::try_new(batch)? { + let mut view = + SparsePrimaryKeyView::new(pk, &mut self.pk_offsets).context(DecodeSnafu)?; + // Visit all needed tags before moving to the next PK so offset discovery is shared. + for &(col_id, target_key) in &sparse_columns { + let value = IndexValueCodec::encode_sparse_value( + &mut view, + col_id, + &mut self.value_buf, + ) + .context(DecodeSnafu)?; + self.index_creator + .push_with_name_n(target_key, value, count) + .await + .context(PushIndexValueSnafu)?; + } + } + } + Ok(()) } /// Finishes index creation and cleans up garbage. /// Returns the number of rows and bytes written. - pub async fn finish( + pub(crate) async fn finish( &mut self, puffin_writer: &mut SstPuffinWriter, ) -> Result<(RowCount, ByteCount)> { diff --git a/src/mito2/src/sst/index/primary_key.rs b/src/mito2/src/sst/index/primary_key.rs new file mode 100644 index 00000000000..4724ae44e6f --- /dev/null +++ b/src/mito2/src/sst/index/primary_key.rs @@ -0,0 +1,108 @@ +// Copyright 2023 Greptime 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 datatypes::arrow::array::{Array, BinaryArray}; +use datatypes::arrow::record_batch::RecordBatch; +use snafu::{OptionExt, ensure}; + +use crate::error::{InvalidRecordBatchSnafu, Result}; +use crate::sst::parquet::flat_format::primary_key_column_index; +use crate::sst::parquet::format::PrimaryKeyArray; + +/// Iterates consecutive dictionary-key runs, preserving their row order without decoding. +pub(crate) struct PrimaryKeyRuns<'a> { + keys: &'a [u32], + values: &'a BinaryArray, +} + +impl<'a> PrimaryKeyRuns<'a> { + pub(crate) fn try_new(batch: &'a RecordBatch) -> Result { + let pk = batch + .column(primary_key_column_index(batch.num_columns())) + .as_any() + .downcast_ref::() + .context(InvalidRecordBatchSnafu { + reason: "Primary key column is not a dictionary array", + })?; + let values = pk.values().as_any().downcast_ref::().context( + InvalidRecordBatchSnafu { + reason: "Primary key values are not binary array", + }, + )?; + ensure!( + pk.null_count() == 0 && values.null_count() == 0, + InvalidRecordBatchSnafu { + reason: "Primary keys must not be null" + } + ); + Ok(Self { + keys: pk.keys().values(), + values, + }) + } +} + +impl<'a> Iterator for PrimaryKeyRuns<'a> { + type Item = (&'a [u8], usize); + + fn next(&mut self) -> Option { + let &key = self.keys.first()?; + let count = self + .keys + .iter() + .take_while(|&¤t| current == key) + .count(); + self.keys = &self.keys[count..]; + Some((self.values.value(key as usize), count)) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datatypes::arrow::array::{ArrayRef, BinaryDictionaryBuilder, UInt8Array}; + use datatypes::arrow::datatypes::UInt32Type; + + use super::*; + + #[test] + fn sliced_runs_preserve_order_and_nonconsecutive_keys() { + let mut keys = BinaryDictionaryBuilder::::new(); + for key in ["a", "a", "b", "b", "b", "a", "c"] { + keys.append(key).unwrap(); + } + let batch = RecordBatch::try_from_iter([ + ("pk", Arc::new(keys.finish()) as ArrayRef), + ("seq", Arc::new(UInt8Array::from(vec![0; 7])) as ArrayRef), + ("op", Arc::new(UInt8Array::from(vec![0; 7])) as ArrayRef), + ]) + .unwrap(); + let slice = batch.slice(1, 5); + assert_eq!( + PrimaryKeyRuns::try_new(&slice).unwrap().collect::>(), + vec![ + (b"a".as_slice(), 1), + (b"b".as_slice(), 3), + (b"a".as_slice(), 1) + ] + ); + assert!( + PrimaryKeyRuns::try_new(&batch.slice(0, 0)) + .unwrap() + .next() + .is_none() + ); + } +} diff --git a/src/mito2/src/sst/index/sparse_test.rs b/src/mito2/src/sst/index/sparse_test.rs new file mode 100644 index 00000000000..15d26bd1bc7 --- /dev/null +++ b/src/mito2/src/sst/index/sparse_test.rs @@ -0,0 +1,198 @@ +// Copyright 2023 Greptime 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 std::collections::BTreeMap; +use std::sync::Arc; + +use api::v1::SemanticType; +use datatypes::arrow::array::{ + ArrayRef, BinaryDictionaryBuilder, TimestampMillisecondArray, UInt8Array, UInt64Array, +}; +use datatypes::arrow::datatypes::UInt32Type; +use datatypes::arrow::record_batch::RecordBatch; +use datatypes::schema::SkippingIndexOptions; +use datatypes::value::ValueRef; +use index::bloom_filter::reader::{BloomFilterReader, BloomFilterReaderImpl}; +use index::inverted_index::format::reader::{InvertedIndexBlobReader, InvertedIndexReader}; +use mito_codec::row_converter::{PrimaryKeyCodec, SparsePrimaryKeyCodec}; +use object_store::ObjectStore; +use object_store::services::Memory; +use prost::Message; +use puffin::puffin_manager::{PuffinManager, PuffinReader}; +use store_api::codec::PrimaryKeyEncoding; +use store_api::metadata::{RegionMetadataBuilder, RegionMetadataRef}; +use store_api::storage::FileId; +use store_api::storage::consts::ReservedColumnId; + +use super::{IndexBuildType, IndexerBuilder, IndexerBuilderImpl}; +use crate::region::options::IndexOptions; +use crate::sst::file::{RegionFileId, RegionIndexId}; +use crate::sst::index::bloom_filter::creator::tests::TestPathProvider; +use crate::sst::index::intermediate::IntermediateManager; +use crate::sst::index::puffin_manager::PuffinManagerFactory; +use crate::sst::parquet::flat_format::FlatReadFormat; +use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema}; +use crate::test_util::sst_util::sst_region_metadata_with_encoding; + +fn sparse_input() -> (RegionMetadataRef, RecordBatch) { + let mut metadata = sst_region_metadata_with_encoding(PrimaryKeyEncoding::Sparse); + for column in &mut metadata.column_metadatas { + if column.semantic_type == SemanticType::Tag { + column.column_schema.set_inverted_index(true); + column + .column_schema + .set_skipping_options(&SkippingIndexOptions { + granularity: 3, + ..Default::default() + }) + .unwrap(); + } + } + let metadata = Arc::new( + RegionMetadataBuilder::from_existing(metadata) + .build() + .unwrap(), + ); + let codec = SparsePrimaryKeyCodec::new(&metadata); + let mut keys = BinaryDictionaryBuilder::::new(); + for (series, (first, second, count)) in [ + (Some("中文12345678"), Some(""), 5), + (None, Some("abcdefgh"), 2), + (Some(""), None, 4), + ] + .into_iter() + .enumerate() + { + let mut key = Vec::new(); + codec + .encode_value_refs( + &[ + (ReservedColumnId::table_id(), ValueRef::UInt32(42)), + (ReservedColumnId::tsid(), ValueRef::UInt64(series as u64)), + (0, first.map_or(ValueRef::Null, ValueRef::String)), + (1, second.map_or(ValueRef::Null, ValueRef::String)), + ], + &mut key, + ) + .unwrap(); + for _ in 0..count { + keys.append(&key).unwrap(); + } + } + let schema = to_flat_sst_arrow_schema( + &metadata, + &FlatSchemaOptions::from_encoding(PrimaryKeyEncoding::Sparse), + ); + let columns: Vec = vec![ + Arc::new(UInt64Array::from(vec![1; 11])), + Arc::new(TimestampMillisecondArray::from_iter_values(0..11)), + Arc::new(keys.finish()), + Arc::new(UInt64Array::from(vec![1; 11])), + Arc::new(UInt8Array::from(vec![1; 11])), + ]; + (metadata, RecordBatch::try_new(schema, columns).unwrap()) +} + +#[tokio::test] +async fn sparse_and_materialized_tags_produce_identical_indexes() { + let (metadata, sparse) = sparse_input(); + // Use the independent whole-PK read conversion as the oracle for both index types. + let materialized = FlatReadFormat::new_with_all_columns(metadata.clone()) + .convert_batch(sparse.clone(), None) + .unwrap(); + let projection = materialized + .schema() + .fields() + .iter() + .enumerate() + .filter_map(|(i, field)| { + (field.name() == "tag_0" || sparse.column_by_name(field.name()).is_some()).then_some(i) + }) + .collect::>(); + let mixed = materialized.project(&projection).unwrap(); + + let (dir, factory) = PuffinManagerFactory::new_for_test_async("sparse_index_bytes").await; + let puffin = factory.build( + ObjectStore::new(Memory::default()).unwrap(), + TestPathProvider, + ); + let mut options = IndexOptions::default(); + options.inverted_index.segment_row_count = 3; + let builder = IndexerBuilderImpl { + build_type: IndexBuildType::Flush, + metadata: metadata.clone(), + puffin_manager: puffin.clone(), + write_cache_enabled: false, + intermediate_manager: IntermediateManager::init_fs(dir.path().to_str().unwrap()) + .await + .unwrap(), + index_options: options, + inverted_index_config: Default::default(), + fulltext_index_config: Default::default(), + bloom_filter_index_config: Default::default(), + #[cfg(feature = "vector_index")] + vector_index_config: Default::default(), + }; + let mut expected = None; + for batch in [materialized, sparse, mixed] { + let file = RegionFileId::new(metadata.region_id, FileId::random()); + let mut indexer = builder.build(file, 0, None).await; + // Split a repeated PK across batches; runs also cross 3-row segment boundaries. + indexer.update_flat(&batch.slice(0, 2)).await; + indexer.update_flat(&batch.slice(2, 9)).await; + let output = indexer.finish().await; + assert_eq!(output.inverted_index.row_count, 11); + assert_eq!(output.bloom_filter.row_count, 11); + assert_eq!(output.inverted_index.columns.len(), 4); + assert_eq!(output.bloom_filter.columns.len(), 4); + + let reader = puffin.reader(&RegionIndexId::new(file, 0)).await.unwrap(); + let blob = reader.blob("greptime-inverted-index-v1").await.unwrap(); + let inverted = InvertedIndexBlobReader::new(blob.reader().await.unwrap()); + let metas = inverted.metadata(None).await.unwrap(); + let mut index_bytes = BTreeMap::new(); + // Columns can be emitted in a different order; compare each column's FST and bitmaps. + for (name, meta) in &metas.metas { + index_bytes.insert( + format!("inverted/{name}"), + inverted + .range_read(meta.base_offset, meta.inverted_index_size as u32, None) + .await + .unwrap(), + ); + } + for id in &metadata.primary_key { + let blob = reader + .blob(&format!("greptime-bloom-filter-v1-{id}")) + .await + .unwrap(); + let bloom = BloomFilterReaderImpl::new(blob.reader().await.unwrap()); + let meta = bloom.metadata(None).await.unwrap(); + assert_eq!(meta.row_count, 11); + assert_eq!(meta.segment_count, 4); + let mut bytes = bloom + .range_read(0, meta.bloom_filter_size as u32, None) + .await + .unwrap() + .to_vec(); + bytes.extend_from_slice(&meta.encode_to_vec()); + index_bytes.insert(format!("bloom/{id}"), bytes); + } + if let Some(expected) = &expected { + assert_eq!(&index_bytes, expected); + } else { + expected = Some(index_bytes); + } + } +} diff --git a/src/mito2/src/test_util/bench_util.rs b/src/mito2/src/test_util/bench_util.rs index 8f182e41578..1fa1a53a439 100644 --- a/src/mito2/src/test_util/bench_util.rs +++ b/src/mito2/src/test_util/bench_util.rs @@ -33,6 +33,8 @@ use store_api::storage::RegionId; use table::predicate::Predicate; use crate::memtable::KeyValues; +pub use crate::sst::index::bloom_filter::creator::BloomFilterIndexer; +pub use crate::sst::index::inverted_index::creator::InvertedIndexer; use crate::test_util::memtable_util::region_metadata_to_row_schema; pub struct Host { From 8af3a04ed7ec225e9b3fa41bda8fd6459f77179f Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 16 Sep 2026 06:28:15 +0000 Subject: [PATCH 4/4] fix: preserve structured query errors through distributed execution (#9161) * fix: preserve structured query errors through distributed execution Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test: update SQL expectations for preserved query error codes Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/common/query/src/error.rs | 148 +++++++- src/common/recordbatch/src/error.rs | 124 ++++++- src/flow/src/batching_mode/task/ckpt.rs | 45 --- src/flow/src/batching_mode/task/test.rs | 214 ++--------- src/query/src/datafusion.rs | 347 +++++++++++++++++- src/query/src/dist_plan/merge_scan.rs | 117 ++++-- src/query/src/error.rs | 15 + src/servers/src/error.rs | 36 +- tests-integration/src/grpc/flight.rs | 8 +- .../optimizer/filter_push_down.result | 2 +- .../standalone/common/order/limit.result | 18 +- .../optimizer/filter_push_down.result | 2 +- 12 files changed, 791 insertions(+), 285 deletions(-) diff --git a/src/common/query/src/error.rs b/src/common/query/src/error.rs index dd2b29adf3a..4c53be86b6e 100644 --- a/src/common/query/src/error.rs +++ b/src/common/query/src/error.rs @@ -270,18 +270,158 @@ pub fn datafusion_status_code( e: &DataFusionError, default_status: Option, ) -> StatusCode { - match e { + let mut error = e; + loop { + error = match error { + DataFusionError::Shared(inner) => inner, + DataFusionError::Context(_, inner) | DataFusionError::Diagnostic(_, inner) => inner, + _ => break, + }; + } + + match error { DataFusionError::Internal(_) => StatusCode::Internal, DataFusionError::NotImplemented(_) => StatusCode::Unsupported, DataFusionError::Plan(_) => StatusCode::PlanQuery, - DataFusionError::External(e) => { - if let Some(ext) = (*e).downcast_ref::() { + DataFusionError::External(error) => { + if let Some(ext) = (*error).downcast_ref::() { + ext.status_code() + } else if let Some(ext) = (*error).downcast_ref::() { ext.status_code() } else { default_status.unwrap_or(StatusCode::EngineExecuteQuery) } } - DataFusionError::Diagnostic(_, e) => datafusion_status_code::(e, default_status), _ => default_status.unwrap_or(StatusCode::EngineExecuteQuery), } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_error::ext::PlainError; + + use super::*; + + #[test] + fn test_datafusion_status_code_preserves_external_errors_through_wrappers() { + let boxed_error = |status| { + DataFusionError::External(Box::new(BoxedError::new(PlainError::new( + "neutral error".to_string(), + status, + )))) + }; + + for status in [ + StatusCode::RequestOutdated, + StatusCode::Unknown, + StatusCode::Unsupported, + ] { + let errors = [ + boxed_error(status), + DataFusionError::Shared(Arc::new(boxed_error(status))), + DataFusionError::Context("context".to_string(), Box::new(boxed_error(status))), + DataFusionError::Diagnostic( + Box::new(datafusion_common::Diagnostic::new_error("diagnostic", None)), + Box::new(boxed_error(status)), + ), + DataFusionError::Shared(Arc::new(DataFusionError::Context( + "context".to_string(), + Box::new(DataFusionError::Diagnostic( + Box::new(datafusion_common::Diagnostic::new_error("diagnostic", None)), + Box::new(boxed_error(status)), + )), + ))), + ]; + + for error in errors { + assert_eq!(datafusion_status_code::(&error, None), status); + assert_eq!( + datafusion_status_code::(&error, Some(StatusCode::PlanQuery)), + status + ); + } + } + + let direct_error = DataFusionError::External(Box::new(Error::DynFilterPayloadTooLarge { + payload_size_bytes: 2, + max_payload_bytes: 1, + location: Location::default(), + })); + assert_eq!( + datafusion_status_code::(&direct_error, Some(StatusCode::Internal)), + StatusCode::PlanQuery + ); + + let boundary_error = + DataFusionError::Shared(Arc::new(DataFusionError::External(Box::new( + BoxedError::new(common_recordbatch::error::Error::PhysicalExpr { + error: DataFusionError::NotImplemented("inner error".to_string()), + location: Location::default(), + }), + )))); + assert_eq!( + datafusion_status_code::(&boundary_error, Some(StatusCode::PlanQuery)), + StatusCode::Internal + ); + } + + #[test] + fn test_datafusion_status_code_uses_default_for_untyped_errors() { + let wrap = |error| { + DataFusionError::Shared(Arc::new(DataFusionError::Context( + "context".to_string(), + Box::new(DataFusionError::Diagnostic( + Box::new(datafusion_common::Diagnostic::new_error("diagnostic", None)), + Box::new(error), + )), + ))) + }; + let errors = || { + [ + ( + DataFusionError::External(Box::new(std::io::Error::other("neutral io error"))), + StatusCode::EngineExecuteQuery, + StatusCode::Unknown, + ), + ( + DataFusionError::Internal("neutral internal error".to_string()), + StatusCode::Internal, + StatusCode::Internal, + ), + ( + DataFusionError::NotImplemented("neutral not implemented error".to_string()), + StatusCode::Unsupported, + StatusCode::Unsupported, + ), + ( + DataFusionError::Plan("neutral plan error".to_string()), + StatusCode::PlanQuery, + StatusCode::PlanQuery, + ), + ( + DataFusionError::External(Box::new(DataFusionError::Internal( + "inner error".to_string(), + ))), + StatusCode::EngineExecuteQuery, + StatusCode::Unknown, + ), + ] + }; + + for (error, none_expected, default_expected) in + errors() + .into_iter() + .chain(errors().map(|(error, none_expected, default_expected)| { + (wrap(error), none_expected, default_expected) + })) + { + assert_eq!(datafusion_status_code::(&error, None), none_expected); + assert_eq!( + datafusion_status_code::(&error, Some(StatusCode::Unknown)), + default_expected + ); + } + } +} diff --git a/src/common/recordbatch/src/error.rs b/src/common/recordbatch/src/error.rs index 469320c19cf..bddb8ae8032 100644 --- a/src/common/recordbatch/src/error.rs +++ b/src/common/recordbatch/src/error.rs @@ -205,7 +205,26 @@ impl ErrorExt for Error { | Error::PhysicalExpr { .. } | Error::RecordBatchSliceIndexOverflow { .. } => StatusCode::Internal, - Error::PollStream { .. } => StatusCode::EngineExecuteQuery, + Error::PollStream { error, .. } => { + let mut error = error; + loop { + error = match error { + datafusion::error::DataFusionError::Shared(inner) => inner, + datafusion::error::DataFusionError::Context(_, inner) + | datafusion::error::DataFusionError::Diagnostic(_, inner) => inner, + _ => break, + }; + } + + match error { + datafusion::error::DataFusionError::External(source) => source + .downcast_ref::() + .map_or(StatusCode::EngineExecuteQuery, |source| { + source.status_code() + }), + _ => StatusCode::EngineExecuteQuery, + } + } Error::ArrowCompute { .. } => StatusCode::IllegalState, @@ -242,3 +261,106 @@ impl ErrorExt for Error { } } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_error::ext::PlainError; + use datafusion::error::DataFusionError; + + use super::*; + + #[test] + fn poll_stream_status_code_preserves_boxed_error_through_wrappers() { + let boxed_error = |status| { + DataFusionError::External(Box::new(BoxedError::new(PlainError::new( + "neutral error".to_string(), + status, + )))) + }; + + for status in [ + StatusCode::RequestOutdated, + StatusCode::Unknown, + StatusCode::Unsupported, + ] { + let errors = [ + boxed_error(status), + DataFusionError::Shared(Arc::new(boxed_error(status))), + DataFusionError::Context("context".to_string(), Box::new(boxed_error(status))), + DataFusionError::Diagnostic( + Box::new(datafusion::common::Diagnostic::new_error( + "diagnostic", + None, + )), + Box::new(boxed_error(status)), + ), + DataFusionError::Shared(Arc::new(DataFusionError::Context( + "context".to_string(), + Box::new(DataFusionError::Diagnostic( + Box::new(datafusion::common::Diagnostic::new_error( + "diagnostic", + None, + )), + Box::new(boxed_error(status)), + )), + ))), + ]; + + for error in errors { + let error = Error::PollStream { + error, + location: Location::default(), + }; + assert_eq!(error.status_code(), status); + } + } + + let error = Error::PollStream { + error: DataFusionError::Shared(Arc::new(DataFusionError::External(Box::new( + BoxedError::new(Error::PhysicalExpr { + error: DataFusionError::NotImplemented("inner error".to_string()), + location: Location::default(), + }), + )))), + location: Location::default(), + }; + assert_eq!(error.status_code(), StatusCode::Internal); + } + + #[test] + fn poll_stream_status_code_defaults_for_other_datafusion_errors() { + let wrap = |error| { + DataFusionError::Shared(Arc::new(DataFusionError::Context( + "context".to_string(), + Box::new(DataFusionError::Diagnostic( + Box::new(datafusion::common::Diagnostic::new_error( + "diagnostic", + None, + )), + Box::new(error), + )), + ))) + }; + let errors = || { + [ + DataFusionError::External(Box::new(std::io::Error::other("neutral io error"))), + DataFusionError::Internal("neutral internal error".to_string()), + DataFusionError::NotImplemented("neutral not implemented error".to_string()), + DataFusionError::Plan("neutral plan error".to_string()), + DataFusionError::External(Box::new(DataFusionError::Internal( + "inner error".to_string(), + ))), + ] + }; + + for error in errors().into_iter().chain(errors().map(wrap)) { + let error = Error::PollStream { + error, + location: Location::default(), + }; + assert_eq!(error.status_code(), StatusCode::EngineExecuteQuery); + } + } +} diff --git a/src/flow/src/batching_mode/task/ckpt.rs b/src/flow/src/batching_mode/task/ckpt.rs index 960e541c73c..c736a0fdbdf 100644 --- a/src/flow/src/batching_mode/task/ckpt.rs +++ b/src/flow/src/batching_mode/task/ckpt.rs @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::error::Error as StdError; use std::time::Duration; use client::OutputWithMetrics; @@ -31,42 +30,6 @@ use crate::metrics::{ }; use crate::{Error, FlowId}; -/// Liveness guard: when a fenced repair query fails with a wrapped error whose -/// text indicates a stale snapshot fence (even when `StatusCode::RequestOutdated` -/// was lost through client layers), classify it as `SnapshotFenceExpired` to -/// break the retry loop and force a rebind of the fence high `H`. -/// -/// Long-term the structured `StatusCode` / retry hint path should be preserved -/// end-to-end; this text fallback is a narrow safety measure. -fn matches_stale_snapshot_fence_text(err: &Error) -> bool { - let markers = [ - "STALE_SNAPSHOT_FENCE", - "REBIND_SNAPSHOT_FENCE", - "snapshot upper bound stale", - ]; - // Check the top-level error Display and Debug. - let debug_str = format!("{:?}", err); - let display_str = err.to_string(); - for marker in &markers { - if debug_str.contains(marker) || display_str.contains(marker) { - return true; - } - } - // Walk the error source chain. - let mut source = err.source(); - while let Some(s) = source { - let debug_str = format!("{:?}", s); - let display_str = s.to_string(); - for marker in &markers { - if debug_str.contains(marker) || display_str.contains(marker) { - return true; - } - } - source = s.source(); - } - false -} - impl BatchingTask { /// Classify execution errors into checkpoint fallback reasons. A stale /// snapshot fence is special only for fenced repair chunks. @@ -80,14 +43,6 @@ impl BatchingTask { } else { FlowQueryFallbackReason::StaleCursor } - } else if matches!(coverage, QueryCoverage::FencedRepairChunk { .. }) - && matches_stale_snapshot_fence_text(err) - { - // Narrow text-based fallback for wrapped errors where the - // structured StatusCode::RequestOutdated was lost through - // frontend/client layers. Without this fenced repair will - // retry the same stale `given_seq` every refresh tick. - FlowQueryFallbackReason::SnapshotFenceExpired } else if matches!(coverage, QueryCoverage::IncrementalDelta) { FlowQueryFallbackReason::IncrementalQueryFailure } else { diff --git a/src/flow/src/batching_mode/task/test.rs b/src/flow/src/batching_mode/task/test.rs index 24cdbf883e0..acb872d8fea 100644 --- a/src/flow/src/batching_mode/task/test.rs +++ b/src/flow/src/batching_mode/task/test.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::sync::Arc; use catalog::RegisterTableRequest; use catalog::memory::MemoryCatalogManager; @@ -391,50 +392,6 @@ fn flow_error_with_status(status_code: StatusCode) -> Error { .unwrap_err() } -/// Test-only error that carries a non-RequestOutdated status code but -/// displays a stale-snapshot-fence marker string, simulating the real-world -/// scenario where the structured status code is lost through frontend/client -/// wrapping layers. -#[derive(Debug)] -struct StaleFenceTextError { - code: StatusCode, - message: String, -} - -impl std::fmt::Display for StaleFenceTextError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}", self.message) - } -} - -impl std::error::Error for StaleFenceTextError {} - -impl common_error::ext::ErrorExt for StaleFenceTextError { - fn status_code(&self) -> StatusCode { - self.code - } - fn as_any(&self) -> &dyn std::any::Any { - self - } -} - -impl common_error::ext::StackError for StaleFenceTextError { - fn debug_fmt(&self, _: usize, _: &mut Vec) {} - fn next(&self) -> Option<&dyn common_error::ext::StackError> { - None - } -} - -fn flow_error_with_code_and_text(code: StatusCode, text: &str) -> Error { - let inner = StaleFenceTextError { - code, - message: text.to_string(), - }; - Err::<(), _>(BoxedError::new(inner)) - .context(crate::error::ExternalSnafu) - .unwrap_err() -} - fn dirty_range(start: i64, end: i64) -> DirtyTimeWindows { let mut dirty = DirtyTimeWindows::default(); dirty.add_window( @@ -1447,78 +1404,6 @@ fn test_query_failure_reason_distinguishes_fenced_repair_stale_fence() { ); } -/// Wrapped errors carrying stale snapshot fence marker text in their -/// Display/Debug chain should be classified as `SnapshotFenceExpired` on -/// fenced repair coverage, even when the structured `StatusCode::RequestOutdated` -/// was lost through client layering. This prevents an infinite retry loop -/// where the fenced chunk re-sends the same stale `given_seq` every tick. -#[test] -fn test_query_failure_reason_text_fallback_stale_snapshot_fence() { - let high = BTreeMap::new(); - let fenced = QueryCoverage::FencedRepairChunk { high: high.clone() }; - - // STALE_SNAPSHOT_FENCE marker with a non-RequestOutdated status code - let err = flow_error_with_code_and_text( - StatusCode::Internal, - "gRPC error: STALE_SNAPSHOT_FENCE: snapshot upper bound stale, region: 1024/0", - ); - assert_eq!( - BatchingTask::query_failure_reason(&err, &fenced), - FlowQueryFallbackReason::SnapshotFenceExpired - ); - - // REBIND_SNAPSHOT_FENCE marker - let err = flow_error_with_code_and_text( - StatusCode::Internal, - "STALE_SNAPSHOT_FENCE ... retry_hint: REBIND_SNAPSHOT_FENCE", - ); - assert_eq!( - BatchingTask::query_failure_reason(&err, &fenced), - FlowQueryFallbackReason::SnapshotFenceExpired - ); - - // snapshot upper bound stale marker (the natural-language fragment) - let err = flow_error_with_code_and_text( - StatusCode::Internal, - "query failed: snapshot upper bound stale, consider rebinding", - ); - assert_eq!( - BatchingTask::query_failure_reason(&err, &fenced), - FlowQueryFallbackReason::SnapshotFenceExpired - ); - - // Fenced coverage with a generic wrapped error (no stale-fence marker) → - // still QueryFailure - let generic_err = - flow_error_with_code_and_text(StatusCode::Internal, "some transient network error"); - assert_eq!( - BatchingTask::query_failure_reason(&generic_err, &fenced), - FlowQueryFallbackReason::QueryFailure - ); - - // Non-fenced incremental coverage with stale-fence marker text must NOT - // classify as SnapshotFenceExpired; it should remain IncrementalQueryFailure. - let err = flow_error_with_code_and_text( - StatusCode::Internal, - "STALE_SNAPSHOT_FENCE blob in unexpected context", - ); - assert_eq!( - BatchingTask::query_failure_reason(&err, &QueryCoverage::IncrementalDelta), - FlowQueryFallbackReason::IncrementalQueryFailure - ); - - // Existing RequestOutdated behavior is unchanged. - let outdated_err = flow_error_with_status(StatusCode::RequestOutdated); - assert_eq!( - BatchingTask::query_failure_reason(&outdated_err, &fenced), - FlowQueryFallbackReason::SnapshotFenceExpired - ); - assert_eq!( - BatchingTask::query_failure_reason(&outdated_err, &QueryCoverage::IncrementalDelta), - FlowQueryFallbackReason::StaleCursor - ); -} - #[test] fn test_fenced_repair_coverage_produces_snapshot_seq_map_for_distributed_metadata_path() { // Covers the metadata boundary between QueryCoverage and the @@ -1563,11 +1448,28 @@ async fn test_fenced_repair_stale_fence_next_plan_is_scoped_base_repair() { { let mut state = task.state.write().unwrap(); + let error = Err::<(), _>(BoxedError::new( + common_recordbatch::error::Error::PollStream { + error: datafusion::error::DataFusionError::Shared(Arc::new( + datafusion::error::DataFusionError::External(Box::new(BoxedError::new( + MockError::new(StatusCode::RequestOutdated), + ))), + )), + location: snafu::Location::default(), + }, + )) + .context(crate::error::ExternalSnafu) + .unwrap_err(); + let reason = BatchingTask::query_failure_reason( + &error, + &QueryCoverage::FencedRepairChunk { high: high.clone() }, + ); + assert_eq!(reason, FlowQueryFallbackReason::SnapshotFenceExpired); let decision = BatchingTask::apply_query_failure_to_state( &mut state, std::time::Duration::from_millis(1), &QueryCoverage::FencedRepairChunk { high }, - FlowQueryFallbackReason::SnapshotFenceExpired, + reason, ); assert_eq!( decision, @@ -1577,9 +1479,11 @@ async fn test_fenced_repair_stale_fence_next_plan_is_scoped_base_repair() { }) ); assert!(state.pending_fenced_repair().is_none()); + assert_eq!(state.dirty_time_windows.len(), 1); // Simulate the outer execution failure restore for the in-flight chunk. state.restore_scoped_windows(&filter); + assert_eq!(state.dirty_time_windows.len(), 2); } let plan = task @@ -1666,82 +1570,6 @@ fn test_fenced_repair_transient_non_stale_failure_retries_same_high() { ); } -/// When `query_failure_reason` classifies a wrapped error as -/// `SnapshotFenceExpired` via the text-marker fallback (not via -/// `StatusCode::RequestOutdated`), the state machine must still -/// abandon the fenced repair and produce a `ScopedBaseRepair` plan -/// next, exactly like the structured-code path. -#[tokio::test] -async fn test_text_fallback_stale_fence_produces_scoped_base_repair() { - let TestTaskParts { - task, - query_engine, - .. - } = new_time_window_test_task_with_query( - "SELECT number, date_bin(INTERVAL '5 second', ts) AS time_window FROM numbers_with_ts GROUP BY time_window, number", - ) - .await; - let high = BTreeMap::from([(1_u64, 10_u64), (2_u64, 20_u64)]); - let filter = { - let mut state = task.state.write().unwrap(); - state - .dirty_time_windows - .add_window(Timestamp::new_second(10), Some(Timestamp::new_second(15))); - state - .dirty_time_windows - .add_window(Timestamp::new_second(100), Some(Timestamp::new_second(105))); - state.start_fenced_repair(high.clone()).unwrap(); - next_fenced_repair_filter(&mut state, 1) - }; - - // Construct a wrapped error that hits the text fallback (non-RequestOutdated - // status code with STALE_SNAPSHOT_FENCE marker text). - let err = flow_error_with_code_and_text( - StatusCode::Internal, - "STALE_SNAPSHOT_FENCE: snapshot upper bound stale, retry_hint: REBIND_SNAPSHOT_FENCE", - ); - let coverage = QueryCoverage::FencedRepairChunk { high }; - let reason = BatchingTask::query_failure_reason(&err, &coverage); - assert_eq!(reason, FlowQueryFallbackReason::SnapshotFenceExpired); - - { - let mut state = task.state.write().unwrap(); - let decision = BatchingTask::apply_query_failure_to_state( - &mut state, - std::time::Duration::from_millis(1), - &coverage, - reason, - ); - assert_eq!( - decision, - Some(FlowCheckpointDecision::FallbackToFullSnapshot { - previous_mode: CheckpointMode::FullSnapshot, - reason: FlowQueryFallbackReason::SnapshotFenceExpired, - }) - ); - assert!(state.pending_fenced_repair().is_none()); - - // Simulate the outer execution failure restore for the in-flight chunk. - state.restore_scoped_windows(&filter); - } - - let plan = task - .gen_query_with_time_window( - query_engine, - &aggregate_time_window_sink_schema(), - &[], - false, - Some(1), - ) - .await - .unwrap() - .expect("text-fallback stale fence should restore dirty windows for a fresh scoped repair"); - assert!( - matches!(plan.coverage, QueryCoverage::ScopedBaseRepair), - "next plan after text-fallback stale fence should be ScopedBaseRepair" - ); -} - #[test] fn test_checkpoint_decision_labels_are_stable() { let advance = FlowCheckpointDecision::AdvancedIncremental { diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 63b834bed55..7df99f059ca 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -814,17 +814,20 @@ mod tests { use api::v1::SemanticType; use arrow::array::{ArrayRef, UInt64Array}; use arrow_schema::SortOptions; + use async_trait::async_trait; use catalog::RegisterTableRequest; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, NUMBERS_TABLE_ID}; use common_error::ext::BoxedError; - use common_recordbatch::{EmptyRecordBatchStream, SendableRecordBatchStream, util}; + use common_recordbatch::{ + EmptyRecordBatchStream, RecordBatch, SendableRecordBatchStream, util, + }; use datafusion::physical_plan::display::{DisplayAs, DisplayFormatType}; use datafusion::physical_plan::expressions::PhysicalSortExpr; use datafusion::physical_plan::joins::{HashJoinExec, JoinOn, PartitionMode}; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion::physical_plan::{ExecutionPlan, PhysicalExpr}; use datafusion::prelude::{col, lit}; - use datafusion_common::{JoinType, NullEquality}; + use datafusion_common::{JoinType, NullEquality, ScalarValue}; use datafusion_physical_expr::expressions::Column; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, SchemaRef}; @@ -835,12 +838,13 @@ mod tests { PartitionRange, PrepareRequest, QueryScanContext, RegionScanner, ScannerProperties, }; use store_api::storage::{RegionId, ScanRequest}; + use table::metadata::{TableInfoBuilder, TableMetaBuilder}; use table::table::numbers::{NUMBERS_TABLE_NAME, NumbersTable}; use table::table::scan::RegionScanExec; use super::*; use crate::options::QueryOptions; - use crate::parser::QueryLanguageParser; + use crate::parser::{QueryLanguageParser, QueryStatement}; use crate::part_sort::PartSortExec; use crate::query_engine::{QueryEngineFactory, QueryEngineRef}; @@ -1422,4 +1426,341 @@ mod tests { assert!(right_update_calls.load(Ordering::Relaxed) > 0); assert!(right_last_filter_len.load(Ordering::Relaxed) > 0); } + #[derive(Default)] + struct RecordingMutationHandler { + inserts: std::sync::Mutex>, + } + + #[async_trait] + impl common_function::handlers::TableMutationHandler for RecordingMutationHandler { + async fn insert( + &self, + request: table::requests::InsertRequest, + _ctx: session::context::QueryContextRef, + ) -> common_query::error::Result { + self.inserts.lock().unwrap().push(request); + Ok(common_query::Output::new_with_affected_rows(1)) + } + + async fn delete( + &self, + _request: table::requests::DeleteRequest, + _ctx: session::context::QueryContextRef, + ) -> common_query::error::Result { + unimplemented!("unexpected delete") + } + + async fn flush( + &self, + _request: table::requests::FlushTableRequest, + _ctx: session::context::QueryContextRef, + ) -> common_query::error::Result { + unimplemented!("unexpected flush") + } + + async fn compact( + &self, + _request: table::requests::CompactTableRequest, + _ctx: session::context::QueryContextRef, + ) -> common_query::error::Result { + unimplemented!("unexpected compact") + } + + async fn build_index( + &self, + _request: table::requests::BuildIndexTableRequest, + _ctx: session::context::QueryContextRef, + ) -> common_query::error::Result { + unimplemented!("unexpected build_index") + } + + async fn flush_region( + &self, + _region_id: store_api::storage::RegionId, + _ctx: session::context::QueryContextRef, + ) -> common_query::error::Result { + unimplemented!("unexpected flush_region") + } + + async fn compact_region( + &self, + _region_id: store_api::storage::RegionId, + _ctx: session::context::QueryContextRef, + ) -> common_query::error::Result { + unimplemented!("unexpected compact_region") + } + + async fn discard_unflushed_data( + &self, + _region_id: store_api::storage::RegionId, + _ctx: session::context::QueryContextRef, + ) -> common_query::error::Result { + unimplemented!("unexpected discard_unflushed_data") + } + + async fn discard_unflushed_data_by_table( + &self, + _table_name: table::table_name::TableName, + _ctx: session::context::QueryContextRef, + ) -> common_query::error::Result { + unimplemented!("unexpected discard_unflushed_data_by_table") + } + } + + fn native_schema() -> Arc { + Arc::new(Schema::new(vec![ + ColumnSchema::new("dim", ConcreteDataType::date_datatype(), true), + ColumnSchema::new("amount", ConcreteDataType::decimal128_datatype(30, 2), true), + ColumnSchema::new( + "elapsed", + ConcreteDataType::duration_millisecond_datatype(), + true, + ), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + "updated_at", + ConcreteDataType::timestamp_millisecond_datatype(), + true, + ), + ColumnSchema::new("marker", ConcreteDataType::uint8_datatype(), true), + ColumnSchema::new("payload", ConcreteDataType::binary_datatype(), true), + ColumnSchema::new("epoch", ConcreteDataType::uint64_datatype(), true), + ])) + } + + fn register_native_tables(catalog: &catalog::memory::MemoryCatalogManager) { + let schema = native_schema(); + let meta = TableMetaBuilder::empty() + .schema(schema.clone()) + .primary_key_indices(vec![]) + .value_indices((0..schema.num_columns()).collect()) + .next_column_id(8) + .build() + .unwrap(); + let info = TableInfoBuilder::default() + .name("native_regression") + .table_id(9001) + .table_version(0) + .meta(meta) + .build() + .unwrap(); + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "native_regression".to_string(), + table_id: 9001, + table: table::test_util::EmptyTable::from_table_info(&info), + }) + .unwrap(); + + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new("dim", ConcreteDataType::date_datatype(), true), + ColumnSchema::new("amount", ConcreteDataType::decimal128_datatype(20, 2), true), + ColumnSchema::new( + "elapsed", + ConcreteDataType::duration_millisecond_datatype(), + true, + ), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ])); + let rows = RecordBatch::new( + schema, + vec![ + Arc::new(datatypes::vectors::DateVector::from_slice([0, 2])) as VectorRef, + Arc::new( + datatypes::vectors::Decimal128Vector::from_slice([10000, 20000]) + .with_precision_and_scale(20, 2) + .unwrap(), + ) as VectorRef, + Arc::new(datatypes::vectors::DurationMillisecondVector::from_values( + [10, 20], + )) as VectorRef, + Arc::new(datatypes::vectors::TimestampMillisecondVector::from_slice( + [1, 2], + )) as VectorRef, + ], + ) + .unwrap(); + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "native_aggregate".to_string(), + table_id: 9002, + table: table::test_util::MemTable::table("native_aggregate", rows), + }) + .unwrap(); + } + + async fn run_sql(engine: &QueryEngineRef, sql: &str) -> Vec { + let stmt = QueryLanguageParser::parse_sql(sql, &QueryContext::arc()).unwrap(); + let plan = engine + .planner() + .plan(&stmt, QueryContext::arc()) + .await + .unwrap(); + match engine + .execute(plan, QueryContext::arc()) + .await + .unwrap() + .data + { + OutputData::Stream(stream) => util::collect(stream).await.unwrap(), + _ => unreachable!(), + } + } + + #[tokio::test] + async fn test_native_executor_null_insert_and_aggregates() { + let catalog = catalog::memory::new_memory_catalog_manager().unwrap(); + register_native_tables(&catalog); + let handler = Arc::new(RecordingMutationHandler::default()); + let engine = QueryEngineFactory::new( + catalog, + None, + Some(handler.clone()), + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + + // The CASTs force Insert::can_extract_values false. This records the + // QueryEngine DML path selected by operator/src/statement/dml.rs. + let insert_sql = "INSERT INTO native_regression (dim, amount, elapsed, ts, updated_at, marker, payload, epoch) VALUES (NULL, NULL, NULL, CAST(-62135596799999 AS TIMESTAMP(3)), CAST(NULL AS TIMESTAMP(3)), CAST(1 AS UInt8), X'0102', CAST(1 AS UInt64))"; + let stmt = QueryLanguageParser::parse_sql(insert_sql, &QueryContext::arc()).unwrap(); + let QueryStatement::Sql(sql::statements::statement::Statement::Insert(insert)) = &stmt + else { + unreachable!() + }; + assert!(!insert.can_extract_values()); + let plan = engine + .planner() + .plan(&stmt, QueryContext::arc()) + .await + .unwrap(); + assert!(matches!( + engine + .execute(plan, QueryContext::arc()) + .await + .unwrap() + .data, + OutputData::AffectedRows(1) + )); + + let request = handler.inserts.lock().unwrap().pop().unwrap(); + for (name, data_type) in [ + ("dim", ConcreteDataType::date_datatype()), + ("amount", ConcreteDataType::decimal128_datatype(30, 2)), + ("elapsed", ConcreteDataType::duration_millisecond_datatype()), + ( + "updated_at", + ConcreteDataType::timestamp_millisecond_datatype(), + ), + ] { + let vector = request.columns_values.get(name).unwrap(); + assert_eq!(vector.len(), 1, "{name}"); + assert_eq!(vector.data_type(), data_type, "{name}"); + assert!(vector.is_null(0), "{name}"); + } + assert!(!request.columns_values["ts"].is_null(0)); + for (name, data_type) in [ + ("marker", ConcreteDataType::uint8_datatype()), + ("payload", ConcreteDataType::binary_datatype()), + ("epoch", ConcreteDataType::uint64_datatype()), + ] { + assert_eq!( + request.columns_values[name].data_type(), + data_type, + "{name}" + ); + } + + let batches = run_sql( + &engine, + "SELECT MIN(dim), MAX(dim), SUM(amount), SUM(elapsed) FROM native_aggregate", + ) + .await; + let batch = &batches[0]; + assert_eq!(batch.num_rows(), 1); + for (index, (data_type, value)) in [ + ( + ConcreteDataType::date_datatype(), + ScalarValue::Date32(Some(0)), + ), + ( + ConcreteDataType::date_datatype(), + ScalarValue::Date32(Some(2)), + ), + ( + ConcreteDataType::decimal128_datatype(30, 2), + ScalarValue::Decimal128(Some(30000), 30, 2), + ), + ( + ConcreteDataType::duration_millisecond_datatype(), + ScalarValue::DurationMillisecond(Some(30)), + ), + ] + .into_iter() + .enumerate() + { + assert_eq!(batch.schema.column_schemas()[index].data_type, data_type); + assert_eq!( + ScalarValue::try_from_array(batch.column(index).as_ref(), 0).unwrap(), + value + ); + } + + let batches = run_sql( + &engine, + "SELECT aggregate.amount_sum + aggregate.amount_sum AS amount_add, \ + aggregate.elapsed_sum + aggregate.elapsed_sum AS elapsed_add, \ + CASE WHEN aggregate.min_dim < CAST('1970-01-02' AS DATE) THEN true ELSE false END AS min_before, \ + CASE WHEN aggregate.max_dim > CAST('1970-01-01' AS DATE) THEN true ELSE false END AS max_after \ + FROM (SELECT MIN(dim) AS min_dim, MAX(dim) AS max_dim, SUM(amount) AS amount_sum, SUM(elapsed) AS elapsed_sum \ + FROM native_aggregate) AS aggregate", + ) + .await; + let batch = &batches[0]; + assert_eq!(batch.num_rows(), 1); + for (index, (data_type, value)) in [ + ( + ConcreteDataType::decimal128_datatype(31, 2), + ScalarValue::Decimal128(Some(60000), 31, 2), + ), + ( + ConcreteDataType::duration_millisecond_datatype(), + ScalarValue::DurationMillisecond(Some(60)), + ), + ( + ConcreteDataType::boolean_datatype(), + ScalarValue::Boolean(Some(true)), + ), + ( + ConcreteDataType::boolean_datatype(), + ScalarValue::Boolean(Some(true)), + ), + ] + .into_iter() + .enumerate() + { + assert_eq!(batch.schema.column_schemas()[index].data_type, data_type); + assert_eq!( + ScalarValue::try_from_array(batch.column(index).as_ref(), 0).unwrap(), + value + ); + } + } } diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index e7e78bfec61..2efdd7b1fa2 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -24,6 +24,7 @@ use arrow_schema::{ }; use async_stream::stream; use common_catalog::parse_catalog_and_schema_from_db_string; +use common_error::ext::BoxedError; use common_plugins::GREPTIME_EXEC_READ_COST; use common_query::request::QueryRequest; use common_recordbatch::adapter::{RecordBatchMetrics, region_scan_output_bytes}; @@ -746,7 +747,7 @@ impl MergeScanExec { } let mut stream = do_get_result.map_err(|e| { MERGE_SCAN_ERRORS_TOTAL.inc(); - DataFusionError::External(Box::new(e)) + DataFusionError::External(Box::new(BoxedError::new(e))) })?; if let Some(subscriber_rollback) = subscriber_rollback.as_mut() { @@ -815,7 +816,8 @@ impl MergeScanExec { let poll_elapsed = poll_timer.elapsed(); poll_duration += poll_elapsed; - let batch = batch.map_err(|e| DataFusionError::External(Box::new(e)))?; + let batch = batch + .map_err(|e| DataFusionError::External(Box::new(BoxedError::new(e))))?; let df_batch = batch.into_df_record_batch(); if !Arc::ptr_eq(&advertised_schema, df_batch.schema_ref()) { validate_remote_schema( @@ -1423,6 +1425,8 @@ mod tests { use arrow_schema::{DataType as TestArrowDataType, Field, TimeUnit}; use async_trait::async_trait; use common_base::Plugins; + use common_error::ext::{ErrorExt, PlainError}; + use common_error::status_code::StatusCode; use common_meta::peer::Peer; use common_query::request::{ INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, InitialDynFilterRegs, @@ -1434,6 +1438,7 @@ mod tests { use datafusion::config::ConfigOptions; use datafusion::execution::SessionStateBuilder; use datafusion::physical_plan::filter_pushdown::ChildFilterPushdownResult; + use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::{StatisticsArgs, StatisticsContext}; use datafusion_common::TableReference; use datafusion_expr::{LogicalPlanBuilder, col, lit}; @@ -1448,6 +1453,7 @@ mod tests { use session::ReadPreference; use session::context::QueryContext; use session::query_id::QueryId; + use snafu::IntoError; use table::table::scan::REGION_SCAN_EXEC_NAME; use table::table_name::TableName; use tokio::sync::{Notify, oneshot}; @@ -1830,7 +1836,7 @@ mod tests { } #[tokio::test] - async fn failed_do_get_rolls_back_new_subscriber_without_starting_fanout() { + async fn failed_do_get_preserves_status_code_and_rolls_back_subscriber() { let handler = Arc::new(FailingRegionQueryHandler::default()); let query_ctx = QueryContext::arc(); let state = Arc::new(QueryEngineState::new( @@ -1886,8 +1892,12 @@ mod tests { ) .unwrap(); - let mut stream = exec.to_stream(task_ctx, 0).unwrap(); - assert!(stream.next().await.unwrap().is_err()); + let mut stream = common_recordbatch::adapter::RecordBatchStreamAdapter::try_new( + exec.to_stream(task_ctx, 0).unwrap(), + ) + .unwrap(); + let error = stream.next().await.unwrap().unwrap_err(); + assert_eq!(error.status_code(), StatusCode::RequestOutdated); assert_eq!(handler.do_get_calls.load(Ordering::SeqCst), 1); assert!(handler.saw_subscriber.load(Ordering::SeqCst)); @@ -1897,6 +1907,46 @@ mod tests { assert!(!entries[0].fanout_started_for_test()); } + #[tokio::test] + async fn repartitioned_merge_scan_later_stream_error_preserves_status_code() { + let region_id = RegionId::new(1024, 1); + let handler = Arc::new(TestRegionQueryHandler::with_responses(vec![( + region_id, + int64_schema(&["a", "b"]), + vec![Err(common_recordbatch::error::ExternalSnafu.into_error( + BoxedError::new(PlainError::new( + "neutral stream error".to_string(), + StatusCode::RequestOutdated, + )), + ))], + )])); + let merge_scan = Arc::new(merge_scan_exec_with_handler( + vec![region_id], + expected_int64_schema(), + handler, + 1, + )); + let repartition = + RepartitionExec::try_new(merge_scan, Partitioning::RoundRobinBatch(2)).unwrap(); + assert_eq!( + repartition + .properties() + .output_partitioning() + .partition_count(), + 2 + ); + + let mut stream = common_recordbatch::adapter::RecordBatchStreamAdapter::try_new( + repartition + .execute(0, Arc::new(TaskContext::default())) + .unwrap(), + ) + .unwrap(); + + let error = stream.next().await.unwrap().unwrap_err(); + assert_eq!(error.status_code(), StatusCode::RequestOutdated); + } + #[tokio::test] async fn aborting_pending_do_get_poll_rolls_back_subscriber_without_starting_fanout() { let handler = Arc::new(PendingDoGetHandler::default()); @@ -2097,10 +2147,9 @@ mod tests { assert_eq!(registry_manager.registry_count(), 0); } - #[derive(Clone)] struct TestRegionResponse { advertised_schema: Arc, - batches: Vec, + batches: Vec>, } #[derive(Default)] @@ -2117,7 +2166,7 @@ mod tests { region_id, TestRegionResponse { advertised_schema: batch.schema.clone(), - batches: vec![batch], + batches: vec![Ok(batch)], }, ) }) @@ -2126,7 +2175,13 @@ mod tests { } fn with_responses( - responses: impl IntoIterator, Vec)>, + responses: impl IntoIterator< + Item = ( + RegionId, + Arc, + Vec>, + ), + >, ) -> Self { let responses = responses .into_iter() @@ -2146,19 +2201,17 @@ mod tests { struct TestRecordBatchStream { schema: Arc, - batches: Vec, - index: usize, + batches: Vec>, } impl Stream for TestRecordBatchStream { type Item = common_recordbatch::error::Result; fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - if let Some(batch) = self.batches.get(self.index).cloned() { - self.index += 1; - Poll::Ready(Some(Ok(batch))) - } else { + if self.batches.is_empty() { Poll::Ready(None) + } else { + Poll::Ready(Some(self.batches.remove(0))) } } } @@ -2219,10 +2272,13 @@ mod tests { }), Ordering::SeqCst, ); - crate::error::UnimplementedSnafu { - operation: "test do_get failure", - } - .fail() + Err(crate::error::Error::QueryExecution { + source: BoxedError::new(PlainError::new( + "neutral do_get error".to_string(), + StatusCode::RequestOutdated, + )), + location: snafu::Location::default(), + }) } async fn handle_remote_dyn_filter_update( @@ -2540,8 +2596,19 @@ mod tests { .expect("test handler needs a response for every requested region"); Ok(Box::pin(TestRecordBatchStream { schema: response.advertised_schema.clone(), - batches: response.batches.clone(), - index: 0, + batches: response + .batches + .iter() + .map(|batch| match batch { + Ok(batch) => Ok(batch.clone()), + Err(error) => Err(common_recordbatch::error::ExternalSnafu.into_error( + BoxedError::new(PlainError::new( + error.to_string(), + error.status_code(), + )), + )), + }) + .collect(), })) } @@ -3102,7 +3169,7 @@ mod tests { Arc::new(TestRegionQueryHandler::with_responses(vec![( RegionId::new(1024, 1), remote_schema, - vec![batch], + vec![Ok(batch)], )])), 1, )) @@ -3139,7 +3206,7 @@ mod tests { Arc::new(TestRegionQueryHandler::with_responses(vec![( region_id, advertised_schema, - vec![batch], + vec![Ok(batch)], )])), 1, ); @@ -3172,7 +3239,7 @@ mod tests { Arc::new(TestRegionQueryHandler::with_responses(vec![( region_id, advertised_schema, - vec![batch], + vec![Ok(batch)], )])), 1, ); @@ -3207,7 +3274,7 @@ mod tests { Arc::new(TestRegionQueryHandler::with_responses(vec![( region_id, advertised_schema, - vec![batch], + vec![Ok(batch)], )])), 1, )) diff --git a/src/query/src/error.rs b/src/query/src/error.rs index fdfeaafde97..a2946415570 100644 --- a/src/query/src/error.rs +++ b/src/query/src/error.rs @@ -498,8 +498,23 @@ impl From for DataFusionError { #[cfg(test)] mod tests { + use common_error::ext::PlainError; + use super::*; + #[test] + fn test_datafusion_external_boxed_error_status_code() { + let error = Error::DataFusion { + error: DataFusionError::External(Box::new(BoxedError::new(PlainError::new( + "neutral error".to_string(), + StatusCode::RequestOutdated, + )))), + location: Location::default(), + }; + + assert_eq!(error.status_code(), StatusCode::RequestOutdated); + } + #[test] fn test_build_backend_delegates_error_metadata() { let source = common_datasource::error::LocalFileAccessDisabledSnafu { diff --git a/src/servers/src/error.rs b/src/servers/src/error.rs index fa5957759c7..fe1d226a10c 100644 --- a/src/servers/src/error.rs +++ b/src/servers/src/error.rs @@ -745,7 +745,7 @@ impl ErrorExt for Error { #[cfg(not(windows))] UpdateJemallocMetrics { .. } => StatusCode::Internal, - CollectRecordbatch { .. } => StatusCode::EngineExecuteQuery, + CollectRecordbatch { source, .. } => source.status_code(), ExecuteQuery { source, .. } | ExecutePlan { source, .. } @@ -986,3 +986,37 @@ pub fn status_code_to_http_status(status_code: &StatusCode) -> HttpStatusCode { | StatusCode::EngineExecuteQuery => HttpStatusCode::INTERNAL_SERVER_ERROR, } } + +#[cfg(test)] +mod tests { + use common_error::GREPTIME_DB_HEADER_ERROR_CODE; + use common_error::ext::PlainError; + + use super::*; + + #[test] + fn collect_recordbatch_preserves_poll_stream_status_in_tonic_status() { + let error = Error::CollectRecordbatch { + source: common_recordbatch::error::Error::PollStream { + error: DataFusionError::External(Box::new(BoxedError::new(PlainError::new( + "neutral error".to_string(), + StatusCode::RequestOutdated, + )))), + location: Location::default(), + }, + location: Location::default(), + }; + + let status: tonic::Status = error.into(); + assert_eq!(status.code(), tonic::Code::InvalidArgument); + assert_eq!( + status + .metadata() + .get(GREPTIME_DB_HEADER_ERROR_CODE) + .unwrap() + .to_str() + .unwrap(), + (StatusCode::RequestOutdated as u32).to_string() + ); + } +} diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index 2486dddee05..fb27a273516 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -33,6 +33,8 @@ mod test { use client::region::RegionRequester; use client::{Client, Database}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use common_error::ext::ErrorExt; + use common_error::status_code::StatusCode; use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_grpc::flight::do_put::{DoPutMetadata, DoPutResponse}; use common_grpc::flight::{FlightDecoder, FlightEncoder, FlightMessage}; @@ -784,12 +786,14 @@ mod test { let mut stale_fence_error = None; while let Some(batch) = stream.next().await { if let Err(err) = batch { - stale_fence_error = Some(format!("{err:?}")); + stale_fence_error = Some(err); break; } } - let err_msg = stale_fence_error.expect("expected stale snapshot fence rejection"); + let stale_fence_error = stale_fence_error.expect("expected stale snapshot fence rejection"); + assert_eq!(stale_fence_error.status_code(), StatusCode::RequestOutdated); + let err_msg = format!("{stale_fence_error:?}"); assert!( err_msg.contains("STALE_SNAPSHOT_FENCE") || err_msg.contains("RequestOutdated") diff --git a/tests/cases/distributed/optimizer/filter_push_down.result b/tests/cases/distributed/optimizer/filter_push_down.result index 83c438de4e9..c4eab6b1e09 100644 --- a/tests/cases/distributed/optimizer/filter_push_down.result +++ b/tests/cases/distributed/optimizer/filter_push_down.result @@ -157,7 +157,7 @@ SELECT * FROM integers i1 WHERE NOT EXISTS(SELECT i FROM integers WHERE i=i1.i) SELECT i1.i,i2.i FROM integers i1, integers i2 WHERE i1.i=(SELECT i FROM integers WHERE i1.i=i) AND i1.i=i2.i ORDER BY i1.i; -Error: 3001(EngineExecuteQuery), Error during planning: Correlated scalar subquery must be aggregated to return at most one row +Error: 3000(PlanQuery), Error during planning: Correlated scalar subquery must be aggregated to return at most one row SELECT * FROM (SELECT i1.i AS a, i2.i AS b FROM integers i1, integers i2) a1 WHERE a=b ORDER BY 1; diff --git a/tests/cases/standalone/common/order/limit.result b/tests/cases/standalone/common/order/limit.result index 558f8f6a0fa..b69f44e0465 100644 --- a/tests/cases/standalone/common/order/limit.result +++ b/tests/cases/standalone/common/order/limit.result @@ -25,7 +25,7 @@ SELECT b FROM test ORDER BY b LIMIT 2 OFFSET 0; SELECT a FROM test LIMIT 1.25; -Error: 3001(EngineExecuteQuery), Error during planning: Expected LIMIT to be an integer or null, but got Float64 +Error: 3000(PlanQuery), Error during planning: Expected LIMIT to be an integer or null, but got Float64 SELECT a FROM test LIMIT 2-1; @@ -49,7 +49,7 @@ Error: 1001(Unsupported), This feature is not implemented: Unsupported LIMIT exp SELECT a FROM test LIMIT row_number() OVER (); -Error: 3001(EngineExecuteQuery), This feature is not implemented: Unsupported LIMIT expression: Some(Cast(Cast { expr: WindowFunction(WindowFunction { fun: WindowUDF(WindowUDF { inner: RowNumber { signature: Signature { type_signature: Nullary, volatility: Immutable, parameter_names: None } } }), params: WindowFunctionParams { args: [], partition_by: [], order_by: [], window_frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }, filter: None, null_treatment: None, distinct: false } }), field: Field { name: "", data_type: Int64, nullable: true } })) +Error: 1001(Unsupported), This feature is not implemented: Unsupported LIMIT expression: Some(Cast(Cast { expr: WindowFunction(WindowFunction { fun: WindowUDF(WindowUDF { inner: RowNumber { signature: Signature { type_signature: Nullary, volatility: Immutable, parameter_names: None } } }), params: WindowFunctionParams { args: [], partition_by: [], order_by: [], window_frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }, filter: None, null_treatment: None, distinct: false } }), field: Field { name: "", data_type: Int64, nullable: true } })) CREATE TABLE test2 (a STRING, ts TIMESTAMP TIME INDEX); @@ -69,7 +69,7 @@ SELECT * FROM test2 LIMIT 3; select 1 limit date '1992-01-01'; -Error: 3001(EngineExecuteQuery), Error during planning: Expected LIMIT to be an integer or null, but got Date32 +Error: 3000(PlanQuery), Error during planning: Expected LIMIT to be an integer or null, but got Date32 CREATE TABLE integers(i TIMESTAMP TIME INDEX); @@ -102,23 +102,23 @@ SELECT * FROM integers LIMIT 4; SELECT * FROM integers as int LIMIT (SELECT MIN(integers.i) FROM integers); -Error: 3001(EngineExecuteQuery), Error during planning: Expected LIMIT to be an integer or null, but got Timestamp(ms) +Error: 3000(PlanQuery), Error during planning: Expected LIMIT to be an integer or null, but got Timestamp(ms) SELECT * FROM integers as int OFFSET (SELECT MIN(integers.i) FROM integers); -Error: 3001(EngineExecuteQuery), Error during planning: Expected OFFSET to be an integer or null, but got Timestamp(ms) +Error: 3000(PlanQuery), Error during planning: Expected OFFSET to be an integer or null, but got Timestamp(ms) SELECT * FROM integers as int LIMIT (SELECT MAX(integers.i) FROM integers) OFFSET (SELECT MIN(integers.i) FROM integers); -Error: 3001(EngineExecuteQuery), Error during planning: Expected LIMIT to be an integer or null, but got Timestamp(ms) +Error: 3000(PlanQuery), Error during planning: Expected LIMIT to be an integer or null, but got Timestamp(ms) SELECT * FROM integers as int LIMIT (SELECT max(integers.i) FROM integers where i > 5); -Error: 3001(EngineExecuteQuery), Error during planning: Expected LIMIT to be an integer or null, but got Timestamp(ms) +Error: 3000(PlanQuery), Error during planning: Expected LIMIT to be an integer or null, but got Timestamp(ms) SELECT * FROM integers as int LIMIT (SELECT max(integers.i) FROM integers where i > 5); -Error: 3001(EngineExecuteQuery), Error during planning: Expected LIMIT to be an integer or null, but got Timestamp(ms) +Error: 3000(PlanQuery), Error during planning: Expected LIMIT to be an integer or null, but got Timestamp(ms) SELECT * FROM integers as int LIMIT (SELECT NULL); @@ -130,7 +130,7 @@ Error: 1001(Unsupported), This feature is not implemented: Unsupported LIMIT exp SELECT * FROM integers as int LIMIT (SELECT 'ab'); -Error: 3001(EngineExecuteQuery), Error during planning: Expected LIMIT to be an integer or null, but got Utf8 +Error: 3000(PlanQuery), Error during planning: Expected LIMIT to be an integer or null, but got Utf8 DROP TABLE integers; diff --git a/tests/cases/standalone/optimizer/filter_push_down.result b/tests/cases/standalone/optimizer/filter_push_down.result index 6705bb3c83f..242cebd1167 100644 --- a/tests/cases/standalone/optimizer/filter_push_down.result +++ b/tests/cases/standalone/optimizer/filter_push_down.result @@ -139,7 +139,7 @@ SELECT * FROM integers i1 WHERE NOT EXISTS(SELECT i FROM integers WHERE i=i1.i) SELECT i1.i,i2.i FROM integers i1, integers i2 WHERE i1.i=(SELECT i FROM integers WHERE i1.i=i) AND i1.i=i2.i ORDER BY i1.i; -Error: 3001(EngineExecuteQuery), Error during planning: Correlated scalar subquery must be aggregated to return at most one row +Error: 3000(PlanQuery), Error during planning: Correlated scalar subquery must be aggregated to return at most one row SELECT * FROM (SELECT i1.i AS a, i2.i AS b FROM integers i1, integers i2) a1 WHERE a=b ORDER BY 1;