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/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/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, 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 { 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",