rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-12 15:02:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'split/9035-exact-reads' into split/9035-state-merge

Browse Source 
Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
2026-09-16 15:14:27 +08:00
parent ab3071ef23 b0be39fadc
commit 55ebb1df03
20 changed files with 2229 additions and 1109 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/datatypes/src/vectors/json/array.rs
+3 -448
View File
@@ -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<ArrayRef> {
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::<i32>();
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<ArrayRef> {
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<ArrayRef> {
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<ArrayRef> {
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::<i32>();
widen_list(list_array, actual, expected)?
}
_ => JsonArray::from(list_array.values()).widen_scalar_to(expected.data_type())?,
};
Ok(Arc::new(
GenericListArray::<i32>::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<ArrayRef, String>,
}
impl TestCase {
fn new(
json_array: StructArray,
schema_type: Fields,
expected: std::result::Result<Vec<ArrayRef>, 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::<Int64Type, _, _>(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::<Int64Type, _, _>(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();
src/index/src/bloom_filter/creator.rs
+65
View File
@@ -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<Item = Bytes>) -> 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());
src/mito-codec/src/index.rs
+44 -1
View File
@@ -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<u8>,
) -> Result<Option<&'a [u8]>> {
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.
///
src/mito-codec/src/row_converter/sparse.rs
+3
View File
@@ -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.
src/mito-codec/src/row_converter/sparse/checked.rs
+244
View File
@@ -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<Self> {
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<Option<&'a [u8]>> {
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::<COLUMN_ID_ENCODE_SIZE>().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<String> 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());
}
}
src/mito2/Cargo.toml
+5
View File
@@ -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"]
src/mito2/benches/bench_index_update.rs
+276
View File
@@ -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<ColumnId>, 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::<UInt32Type>::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<ArrayRef> = 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);
src/mito2/src/sst/index.rs
+6 -57
View File
@@ -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<Vec<(CompositeValues, usize)>> {
let primary_key_index = primary_key_column_index(batch.num_columns());
let pk_dict_array = batch
.column(primary_key_index)
.as_any()
.downcast_ref::<PrimaryKeyArray>()
.context(InvalidRecordBatchSnafu {
reason: "Primary key column is not a dictionary array",
})?;
let pk_values_array = pk_dict_array
.values()
.as_any()
.downcast_ref::<BinaryArray>()
.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<u32> = None;
let pk_indices = keys.values();
for &current_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;
src/mito2/src/sst/index/bloom_filter/creator.rs
+35 -44
View File
@@ -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<u8>,
/// 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<Vec<(CompositeValues, usize)>> = 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(())
}
src/mito2/src/sst/index/inverted_index/creator.rs
+34 -45
View File
@@ -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<u8>,
/// 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<Vec<(CompositeValues, usize)>> = 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)> {
src/mito2/src/sst/index/primary_key.rs
+108
View File
@@ -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<Self> {
let pk = batch
.column(primary_key_column_index(batch.num_columns()))
.as_any()
.downcast_ref::<PrimaryKeyArray>()
.context(InvalidRecordBatchSnafu {
reason: "Primary key column is not a dictionary array",
})?;
let values = pk.values().as_any().downcast_ref::<BinaryArray>().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<Self::Item> {
let &key = self.keys.first()?;
let count = self
.keys
.iter()
.take_while(|&&current| 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::<UInt32Type>::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<_>>(),
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()
);
}
}
src/mito2/src/sst/index/sparse_test.rs
+198
View File
@@ -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::<UInt32Type>::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<ArrayRef> = 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::<Vec<_>>();
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);
}
}
}
src/mito2/src/sst/parquet/json_align/stream.rs → src/mito2/src/sst/parquet/json_align.rs
+386 -119
View File
@@ -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<RecordBatch>>;
/// 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<String, Json2TargetLayout>,
},
}
/// 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<String, RewriteSettings>,
},
}
/// 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<S> {
pub struct JsonSchemaAligner<S> {
#[debug(skip)]
inner: S,
/// Output schema expected by the upper reader.
@@ -70,79 +85,55 @@ pub struct NestedSchemaAligner<S> {
/// 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<String, Json2RewriteSettings>,
/// Alignment mode with parsed and validated rewrite settings.
mode: ResolvedAlignMode,
/// The cache for whether incoming batches already match output schema.
is_schema_matched: Option<bool>,
}
impl<S> NestedSchemaAligner<S>
impl<S> JsonSchemaAligner<S>
where
S: Stream<Item = Result<RecordBatch>>,
{
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<bool>,
output_schema: SchemaRef,
) -> Result<NestedSchemaAligner<S>> {
mode: AlignMode,
) -> Result<JsonSchemaAligner<S>> {
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<String, Json2TargetLayout>,
) -> Result<Self> {
self.json2_rewrite_targets = targets
.iter()
.map(|(name, layout)| {
let metadata = serde_json::from_str::<JsonMetadata>(&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::<Result<_>>()?;
Ok(self)
}
}
impl<S> Stream for NestedSchemaAligner<S>
impl<S> Stream for JsonSchemaAligner<S>
where
S: Stream<Item = Result<RecordBatch>> + 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<String, Json2RewriteSettings>,
mode: &ResolvedAlignMode,
) -> Result<RecordBatch> {
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<ArrayRef> {
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<ArrayRef> {
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<Self> {
let metadata = from_str::<JsonMetadata>(&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<ResolvedAlignMode> {
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::<Result<RecordBatch>>(),
vec![],
schema(Vec::<Field>::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::<Result<RecordBatch>>(),
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<Item = Field>) -> SchemaRef {
Arc::new(Schema::new(fields.into_iter().collect::<Vec<_>>()))
}
src/mito2/src/sst/parquet/json_align/mod.rs
-24
View File
@@ -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<RecordBatch>>;
src/mito2/src/sst/parquet/read_columns.rs
+220 -348
View File
@@ -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<String>;
@@ -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<usize>,
col: &ParquetReadColumn,
selected: &mut HashSet<usize>,
) {
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<ProjectionMaskPlan> {
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<usize>, HashSet<usize>) {
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<usize>, HashSet<usize>)> {
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::<usize, Vec<bool>>::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::<Vec<_>>();
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<Range<usize>> {
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<usize>,
col: &ParquetReadColumn,
selected: &mut HashSet<usize>,
) {
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<usize>,
col: &ParquetReadColumn,
selected: &mut HashSet<usize>,
) -> Vec<bool> {
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<NestedSelectionPolicy>]) -> 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::<Vec<_>>(),
)
}
// 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<NestedSelectionPolicy>],
) -> 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<NestedSelectionPolicy>],
) -> (Vec<usize>, HashSet<usize>) {
super::build_parquet_leaves_indices(parquet_schema, projection, &source_schema(policies))
.unwrap()
}
}
src/mito2/src/sst/parquet/reader.rs
+30 -9
View File
@@ -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,
src/mito2/src/test_util/bench_util.rs
+2
View File
@@ -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 {
src/servers/benches/to_http_output.rs
+242 -8
View File
@@ -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<AtomicUsize>,
}
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::<Vec<_>>()));
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<Result<Output>> {
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<ArrayRef> = 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<Output> {
unimplemented!()
}
async fn do_exec_plan(
&self,
_: LogicalPlan,
_: Option<Statement>,
_: QueryContextRef,
) -> Result<Output> {
unimplemented!()
}
async fn do_promql_query(&self, _: &PromQuery, _: QueryContextRef) -> Vec<Result<Output>> {
unimplemented!()
}
async fn do_describe(
&self,
_: Statement,
_: QueryContextRef,
) -> Result<Option<DescribeResult>> {
unimplemented!()
}
async fn is_valid_schema(&self, _: &str, _: &str) -> Result<bool> {
Ok(true)
}
}
#[derive(Clone, Copy)]
enum LimitPlacement {
Late,
Early,
}
async fn limited_response(
handler: Arc<StreamingSqlHandler>,
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 <bench executable>
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();
}
src/servers/src/http/handler.rs
+325 -6
View File
@@ -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<Output> {
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::<RecordBatchResult<Vec<_>>>()
.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<Option<RecordBatch>> {
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<String> {
#[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<Vec<Result<Output>>>,
}
#[async_trait]
impl SqlQueryHandler for TestSqlQueryHandler {
async fn do_query(&self, _: &str, _: QueryContextRef) -> Vec<Result<Output>> {
std::mem::take(&mut *self.outputs.lock().unwrap())
}
async fn do_analyze_stream_query(&self, _: &str, _: QueryContextRef) -> Result<Output> {
unimplemented!()
}
async fn do_exec_plan(
&self,
_: LogicalPlan,
_: Option<Statement>,
_: QueryContextRef,
) -> Result<Output> {
unimplemented!()
}
async fn do_promql_query(&self, _: &PromQuery, _: QueryContextRef) -> Vec<Result<Output>> {
unimplemented!()
}
async fn do_describe(
&self,
_: Statement,
_: QueryContextRef,
) -> Result<Option<DescribeResult>> {
unimplemented!()
}
async fn is_valid_schema(&self, _: &str, _: &str) -> Result<bool> {
Ok(true)
}
}
async fn sql_response(
outputs: Vec<Result<Output>>,
format: &str,
limit: Option<usize>,
) -> 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<VectorRef> = 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<VectorRef> = 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() {
typos.toml
+3
View File
@@ -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",