rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-09-11 07:52:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(json2): push down json2 type hints to parquet reads (#8833)

Browse Source 
* refactor(mito): push down json2 type hints to parquet reads

Signed-off-by: fys <fengys1996@gmail.com>

* refactor(mito): share json2 target types with arc

Signed-off-by: fys <fengys1996@gmail.com>

* refactor(mito): derive json2 output schema from target types

Signed-off-by: fys <fengys1996@gmail.com>

* refactor(mito): simplify read columns construction

Signed-off-by: fys <fengys1996@gmail.com>

* fix: cargo check

Signed-off-by: fys <fengys1996@gmail.com>

* fix(mito): reject JSON hints for non-JSON2 read columns

Signed-off-by: fys <fengys1996@gmail.com>

* fix: do not pushdown json type hint of non-json2-col

Signed-off-by: fys <fengys1996@gmail.com>

* fix: unit test

Signed-off-by: fys <fengys1996@gmail.com>

* refactor(query): simplify JSON type hint application

Signed-off-by: fys <fengys1996@gmail.com>

* refactor: clean up JSON2 type hint handling

Signed-off-by: fys <fengys1996@gmail.com>

* refactor(mito2): keep JSON2 hints with flat read format

Signed-off-by: fys <fengys1996@gmail.com>

* fix(mito): use raw parquet projection for output schema

Signed-off-by: fys <fengys1996@gmail.com>

* fix(query): note JSON2 hint scope limitation

Signed-off-by: fys <fengys1996@gmail.com>

* refactor(mito): store JSON target types as native types

Signed-off-by: fys <fengys1996@gmail.com>

* test(json2): cover join hint qualifier limitation

Signed-off-by: fys <fengys1996@gmail.com>

* refactor(mito): remove JSON2 fallback from compat cast

Signed-off-by: fys <fengys1996@gmail.com>

* docs(mito): document ReadColumns ordering contract

Signed-off-by: fys <fengys1996@gmail.com>

* fix: cargo clippy

Signed-off-by: fys <fengys1996@gmail.com>

---------

Signed-off-by: fys <fengys1996@gmail.com>
This commit is contained in:
fys
2026-08-14 06:21:17 +00:00
committed by GitHub
parent b1263fc65f
commit 6538db6d61
23 changed files with 729 additions and 731 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cmd/src/datanode/parquetbench.rs
+1 -1
View File
@@ -554,7 +554,7 @@ async fn run_flat_prune_iteration(
let reader_builder = ParquetReaderBuilder::new(table_dir, path_type, file_handle, object_store)
.expected_metadata(Some(region_meta))
.cache(CacheStrategy::Disabled)
.projection(projection.map(ReadColumns::from_deduped_column_ids));
.projection(projection.map(ReadColumns::new));
let mut reader_metrics = ReaderMetrics::default();
let start = Instant::now();
let mut stats = IterationStats::default();
src/mito2/src/cache.rs
+16 -4
View File
@@ -49,7 +49,7 @@ use puffin::puffin_manager::cache::{PuffinMetadataCache, PuffinMetadataCacheRef}
use smallvec::SmallVec;
use snafu::{OptionExt, ResultExt};
use store_api::metadata::{RegionMetadata, RegionMetadataRef};
use store_api::storage::{ConcreteDataType, FileId, RegionId, TimeSeriesRowSelector};
use store_api::storage::{ColumnId, ConcreteDataType, FileId, RegionId, TimeSeriesRowSelector};
use crate::cache::cache_size::parquet_meta_size;
use crate::cache::file_cache::{FileType, IndexKey};
@@ -65,6 +65,7 @@ use crate::memtable::record_batch_estimated_size;
use crate::metrics::{CACHE_BYTES, CACHE_EVICTION, CACHE_HIT, CACHE_MISS};
use crate::read::Batch;
use crate::read::range_cache::{RangeScanCacheKey, RangeScanCacheValue};
use crate::read::read_columns::JsonTargetTypes;
use crate::sst::file::{RegionFileId, RegionIndexId};
use crate::sst::parquet::PARQUET_METADATA_KEY;
use crate::sst::parquet::read_columns::ParquetReadColumns;
@@ -2025,6 +2026,11 @@ pub struct SelectorResultValue {
pub result: SelectorResult,
/// The read columns of rows.
pub read_cols: ParquetReadColumns,
/// JSON2 target types used by flat-format reads.
///
/// JSON2 projection is query-driven; the same parquet columns can produce
/// different cached batches under different type hints.
pub json_target_types: JsonTargetTypes,
}
impl SelectorResultValue {
@@ -2033,6 +2039,7 @@ impl SelectorResultValue {
SelectorResultValue {
result: SelectorResult::PrimaryKey(result),
read_cols,
json_target_types: Arc::default(),
}
}
@@ -2040,21 +2047,26 @@ impl SelectorResultValue {
pub fn new_flat(
result: Vec<RecordBatch>,
read_cols: ParquetReadColumns,
json_target_types: JsonTargetTypes,
) -> SelectorResultValue {
SelectorResultValue {
result: SelectorResult::Flat(result),
read_cols,
json_target_types,
}
}
/// Returns memory used by the value (estimated).
fn estimated_size(&self) -> usize {
match &self.result {
let result_size: usize = match &self.result {
SelectorResult::PrimaryKey(batches) => {
batches.iter().map(|batch| batch.memory_size()).sum()
}
SelectorResult::Flat(batches) => batches.iter().map(record_batch_estimated_size).sum(),
}
};
result_size
+ self.json_target_types.len()
* (mem::size_of::<ColumnId>() + mem::size_of::<ConcreteDataType>())
}
}
@@ -2752,7 +2764,7 @@ mod tests {
region_id: RegionId::new(1, 1),
row_groups: vec![(FileId::random(), 0)],
scan: ScanRequestFingerprintBuilder {
read_columns: ReadColumns::from_deduped_column_ids(std::iter::empty()),
read_columns: ReadColumns::new(std::iter::empty()),
read_column_types: vec![],
filters: vec!["tag_0 = 1".to_string()],
time_filters: vec![],
src/mito2/src/compaction/reader.rs
+23 -10
View File
@@ -12,7 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::HashMap;
use std::collections::{BTreeMap, HashMap};
use std::sync::Arc;
use common_time::Timestamp;
@@ -118,15 +118,28 @@ impl CompactionSstReaderBuilder<'_> {
};
let projection = (0..self.metadata.column_metadatas.len()).collect();
let read_columns = ReadColumns::from_deduped_column_ids(
self.metadata.column_metadatas.iter().map(|x| x.column_id),
);
let mapper = FlatProjectionMapper::new_with_read_columns(
&self.metadata,
projection,
read_columns,
json_type_hint.as_ref(),
)?;
let read_column_ids = self
.metadata
.column_metadatas
.iter()
.map(|x| x.column_id)
.collect::<Vec<_>>();
let json_target_types = json_type_hint
.as_ref()
.map(|hint| {
hint.iter()
.filter_map(|(col_name, json_type)| {
self.metadata
.column_by_name(col_name)
.map(|col| (col.column_id, json_type.clone()))
})
.collect::<BTreeMap<_, _>>()
})
.unwrap_or_default();
let read_columns =
ReadColumns::new(read_column_ids).with_json_target_types(json_target_types);
let mapper =
FlatProjectionMapper::new_with_read_columns(&self.metadata, projection, read_columns)?;
let mut scan_input = ScanInput::new(self.sst_layer, mapper)
.with_files(self.inputs.to_vec())
src/mito2/src/engine/scan_test.rs
+25 -25
View File
@@ -38,7 +38,7 @@ use store_api::storage::{RegionId, ScanRequest, TimeSeriesDistribution};
use crate::config::MitoConfig;
use crate::error::Error;
use crate::read::read_columns::{ReadColumn, ReadColumns};
use crate::read::read_columns::ReadColumns;
use crate::read::scan_region::Scanner;
use crate::test_util;
use crate::test_util::{CreateRequestBuilder, TestEnv};
@@ -108,18 +108,19 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul
};
assert_eq!(
seq_scan.input().read_cols,
ReadColumns::from_deduped_column_ids([1, 0])
ReadColumns::new([1, 0])
.with_json_target_types(BTreeMap::from([(1, JsonNativeType::Variant)]))
);
let stream = scanner.scan().await?;
let batches = RecordBatches::try_collect(stream).await?;
let expected = r#"
+------------------------------------------+-------+
| field_0 | tag_0 |
+------------------------------------------+-------+
| {a: {x: 10, y: ignored-a}, b: ignored-b} | tag-1 |
| {a: {x: 20, y: ignored-c}, b: ignored-d} | tag-2 |
+------------------------------------------+-------+
+------------------------------------------------+-------+
| field_0 | tag_0 |
+------------------------------------------------+-------+
| {"a":{"x":10,"y":"ignored-a"},"b":"ignored-b"} | tag-1 |
| {"a":{"x":20,"y":"ignored-c"},"b":"ignored-d"} | tag-2 |
+------------------------------------------------+-------+
"#;
assert_eq!(batches.pretty_print()?, expected.trim());
@@ -149,19 +150,8 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul
// whole JSON2 struct. tag_0 is still read as a normal root column.
assert_eq!(
seq_scan.input().read_cols,
ReadColumns {
cols: vec![
ReadColumn::new(
1,
vec![vec![
"field_0".to_string(),
"a".to_string(),
"x".to_string()
]]
),
ReadColumn::new(0, vec![]),
]
}
ReadColumns::new([1, 0])
.with_json_target_types(BTreeMap::from([(1, json_type_hint["field_0"].clone())]))
);
// The scanner should still return a valid RecordBatch in the requested logical projection.
@@ -189,10 +179,7 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul
let Scanner::Seq(seq_scan) = &scanner else {
unreachable!();
};
assert_eq!(
seq_scan.input().read_cols,
ReadColumns::from_deduped_column_ids([0])
);
assert_eq!(seq_scan.input().read_cols, ReadColumns::new([0]));
let stream = scanner.scan().await?;
let batches = RecordBatches::try_collect(stream).await?;
@@ -205,6 +192,19 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul
+-------+
"#;
assert_eq!(batches.pretty_print()?, expected.trim());
// JSON type hints are derived from the whole query plan, so they may include
// columns unrelated to this scan. Mito2 should ignore hints for unknown
// columns (`other_json`) instead of failing.
let request = ScanRequest {
projection: Some(vec![0]),
json_type_hint: HashMap::from([("other_json".to_string(), JsonNativeType::String)]),
..Default::default()
};
let scanner = engine.scanner(region_id, request).await?;
let stream = scanner.scan().await?;
let batches = RecordBatches::try_collect(stream).await?;
assert_eq!(batches.pretty_print()?, expected.trim());
Ok(())
}
src/mito2/src/memtable/bulk/context.rs
+2 -2
View File
@@ -72,9 +72,9 @@ impl BulkIterContext {
let codec = build_primary_key_codec(&region_metadata);
let read_cols = if let Some(col_ids) = projection {
ReadColumns::from_deduped_column_ids(col_ids.iter().copied())
ReadColumns::new(col_ids.iter().copied())
} else {
ReadColumns::from_deduped_column_ids(
ReadColumns::new(
region_metadata
.column_metadatas
.iter()
src/mito2/src/read/compat.rs
+16 -34
View File
@@ -25,11 +25,9 @@ use datatypes::arrow::compute::{TakeOptions, take};
use datatypes::arrow::datatypes::{FieldRef, Schema, SchemaRef};
use datatypes::arrow::record_batch::RecordBatch;
use datatypes::data_type::ConcreteDataType;
use datatypes::extension::json::is_json2_extension_type;
use datatypes::prelude::DataType;
use datatypes::value::Value;
use datatypes::vectors::VectorRef;
use datatypes::vectors::json::array::JsonArray;
use mito_codec::row_converter::{
CompositeValues, PrimaryKeyCodec, SortField, build_primary_key_codec,
build_primary_key_codec_with_fields,
@@ -40,8 +38,8 @@ use store_api::metadata::{RegionMetadata, RegionMetadataRef};
use store_api::storage::ColumnId;
use crate::error::{
CompatReaderSnafu, ComputeArrowSnafu, ConvertValueSnafu, CreateDefaultSnafu, DecodeSnafu,
EncodeSnafu, NewRecordBatchSnafu, Result, UnexpectedSnafu, UnsupportedOperationSnafu,
CompatReaderSnafu, ComputeArrowSnafu, CreateDefaultSnafu, DecodeSnafu, EncodeSnafu,
NewRecordBatchSnafu, Result, UnexpectedSnafu, UnsupportedOperationSnafu,
};
use crate::read::flat_projection::{FlatProjectionMapper, flat_projected_columns};
use crate::sst::parquet::flat_format::{FlatReadFormat, primary_key_column_index};
@@ -99,20 +97,12 @@ impl FlatCompatBatch {
let actual = read_format.metadata();
let format_projection = read_format.format_projection();
let mut actual_schema = flat_projected_columns(actual, format_projection);
if read_format
.arrow_schema()
.fields()
.iter()
.any(is_json2_extension_type)
{
for field in read_format.arrow_schema().fields() {
if is_json2_extension_type(field)
&& let Some(column_id) =
actual.column_by_name(field.name()).map(|x| x.column_id)
&& let Some(i) = actual_schema.iter().position(|x| x.0 == column_id)
{
actual_schema[i].1 = ConcreteDataType::from_arrow_type(field.data_type());
}
for (column_id, target_type) in read_format.json_target_types().iter() {
if let Some(i) = actual_schema
.iter()
.position(|(actual_column_id, _)| actual_column_id == column_id)
{
actual_schema[i].1 = ConcreteDataType::json2(target_type.clone());
}
}
@@ -270,16 +260,9 @@ impl FlatCompatBatch {
let old_column = batch.column(*pos);
if let Some(ty) = cast_type {
let casted = if let Some(json_type) = ty.as_json()
&& json_type.is_json2()
{
JsonArray::from(old_column)
.project_to(&json_type.as_arrow_type())
.context(ConvertValueSnafu)?
} else {
let casted =
datatypes::arrow::compute::cast(old_column, &ty.as_arrow_type())
.context(ComputeArrowSnafu)?
};
.context(ComputeArrowSnafu)?;
Ok(casted)
} else {
Ok(old_column.clone())
@@ -777,7 +760,7 @@ mod tests {
let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap();
let read_format = FlatReadFormat::new(
actual_metadata.clone(),
ReadColumns::from_deduped_column_ids([0, 1, 2, 3]),
ReadColumns::new([0, 1, 2, 3]),
None,
"test",
false,
@@ -864,13 +847,12 @@ mod tests {
let mapper = FlatProjectionMapper::new_with_read_columns(
&expected_metadata,
vec![1, 2],
ReadColumns::from_deduped_column_ids([1, 2, 3]),
None,
ReadColumns::new([1, 2, 3]),
)
.unwrap();
let read_format = FlatReadFormat::new(
actual_metadata.clone(),
ReadColumns::from_deduped_column_ids([1, 2, 3]),
ReadColumns::new([1, 2, 3]),
None,
"test",
false,
@@ -959,7 +941,7 @@ mod tests {
let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap();
let read_format = FlatReadFormat::new(
actual_metadata.clone(),
ReadColumns::from_deduped_column_ids([0, 1, 2, 3]),
ReadColumns::new([0, 1, 2, 3]),
None,
"test",
false,
@@ -1041,7 +1023,7 @@ mod tests {
let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap();
let read_format = FlatReadFormat::new(
actual_metadata,
ReadColumns::from_deduped_column_ids([0, 1, 2]),
ReadColumns::new([0, 1, 2]),
None,
"test",
false,
@@ -1109,7 +1091,7 @@ mod tests {
let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap();
let read_format = FlatReadFormat::new(
actual_metadata.clone(),
ReadColumns::from_deduped_column_ids([0, 2, 3]),
ReadColumns::new([0, 2, 3]),
None,
"test",
true,
src/mito2/src/read/flat_projection.rs
+32 -51
View File
@@ -14,7 +14,6 @@
//! Utilities for projection on flat format.
use std::collections::HashMap;
use std::sync::Arc;
use api::v1::SemanticType;
@@ -27,7 +26,7 @@ use datatypes::arrow::array::Array;
use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field};
use datatypes::extension::json::is_json2_extension_type;
use datatypes::prelude::{ConcreteDataType, DataType};
use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
use datatypes::schema::{Schema, SchemaRef};
use datatypes::types::JsonType;
use datatypes::types::json_type::JsonNativeType;
use datatypes::value::Value;
@@ -86,8 +85,8 @@ impl FlatProjectionMapper {
) -> Result<Self> {
let projection: Vec<_> = projection.into_iter().collect();
let read_column_ids = read_column_ids_from_projection(metadata, &projection)?;
let read_cols = ReadColumns::from_deduped_column_ids(read_column_ids);
Self::new_with_read_columns(metadata, projection, read_cols, None)
let read_cols = ReadColumns::new(read_column_ids);
Self::new_with_read_columns(metadata, projection, read_cols)
}
/// Returns a new mapper with output projection and explicit read columns.
@@ -95,7 +94,6 @@ impl FlatProjectionMapper {
metadata: &RegionMetadataRef,
projection: Vec<usize>,
read_cols: ReadColumns,
json_type_hint: Option<&HashMap<String, JsonNativeType>>,
) -> Result<Self> {
// If the original projection is empty.
let is_empty_projection = projection.is_empty();
@@ -115,7 +113,11 @@ impl FlatProjectionMapper {
output_col_ids.push(col.column_id);
let mut schema = col.column_schema.clone();
maybe_concretize_json2_datatype(&mut schema, json_type_hint);
if let Some(data_type) =
json2_read_datatype(col.column_id, &schema.data_type, &read_cols)
{
schema.data_type = data_type;
}
col_schemas.push(schema);
}
@@ -124,6 +126,7 @@ impl FlatProjectionMapper {
// TODO(yingwen): Support different flat schema options.
let format_projection = FormatProjection::compute_format_projection(
metadata,
&id_to_index,
// All columns with internal columns.
metadata.column_metadatas.len() + 3,
@@ -133,21 +136,8 @@ impl FlatProjectionMapper {
let mut batch_schema = flat_projected_columns(metadata, &format_projection);
for (column_id, data_type) in batch_schema.iter_mut() {
if let Some(json_type) = data_type.as_json()
&& json_type.is_json2()
{
if let Some(concretized) = metadata
.column_by_id(*column_id)
.and_then(|metadata| {
json_type_hint.and_then(|x| x.get(&metadata.column_schema.name).cloned())
})
.map(ConcreteDataType::json2)
{
*data_type = concretized;
} else if is_empty_json2_type(json_type) {
// see `merge_scan::maybe_amend_json2_field`
*data_type = ConcreteDataType::json2(JsonNativeType::Variant);
}
if let Some(updated) = json2_read_datatype(*column_id, data_type, &read_cols) {
*data_type = updated;
}
}
@@ -401,24 +391,25 @@ impl FlatProjectionMapper {
}
}
fn maybe_concretize_json2_datatype(
schema: &mut ColumnSchema,
json_type_hint: Option<&HashMap<String, JsonNativeType>>,
) {
if let Some(json_type) = schema.data_type.as_json()
&& json_type.is_json2()
{
if let Some(concretized) = json_type_hint
.and_then(|x| x.get(&schema.name))
.cloned()
.map(ConcreteDataType::json2)
{
schema.data_type = concretized;
} else if is_empty_json2_type(json_type) {
// see `merge_scan::maybe_amend_json2_field`
schema.data_type = ConcreteDataType::json2(JsonNativeType::Variant);
}
fn json2_read_datatype(
column_id: ColumnId,
data_type: &ConcreteDataType,
read_cols: &ReadColumns,
) -> Option<ConcreteDataType> {
let json_type = data_type.as_json()?;
if !json_type.is_json2() {
return None;
}
if let Some(concretized) = read_cols.json_target_type(column_id).cloned() {
return Some(ConcreteDataType::json2(concretized));
}
if is_empty_json2_type(json_type) {
return Some(ConcreteDataType::json2(JsonNativeType::Variant));
}
None
}
fn is_empty_json2_type(json_type: &JsonType) -> bool {
@@ -547,9 +538,8 @@ impl CompactionProjectionMapper {
.collect::<Vec<_>>();
let read_col_ids = metadata.column_metadatas.iter().map(|col| col.column_id);
let read_cols = ReadColumns::from_deduped_column_ids(read_col_ids);
let mapper =
FlatProjectionMapper::new_with_read_columns(metadata, projection, read_cols, None)?;
let read_cols = ReadColumns::new(read_col_ids);
let mapper = FlatProjectionMapper::new_with_read_columns(metadata, projection, read_cols)?;
let assembler = DfBatchAssembler::new(mapper.output_schema());
Ok(Self { mapper, assembler })
@@ -617,7 +607,6 @@ impl DfBatchAssembler {
#[cfg(test)]
mod tests {
use datatypes::schema::ColumnSchema;
use datatypes::types::json_type::JsonObjectType;
use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
use store_api::storage::RegionId;
@@ -687,18 +676,10 @@ mod tests {
#[test]
fn test_json_type_hint_does_not_concretize_legacy_json() {
let metadata = metadata_with_legacy_json();
let hint = HashMap::from([(
"j".to_string(),
JsonNativeType::Object(JsonObjectType::from([(
"a".to_string(),
JsonNativeType::i64(),
)])),
)]);
let mapper = FlatProjectionMapper::new_with_read_columns(
&metadata,
vec![0, 1],
ReadColumns::from_deduped_column_ids([0, 1]),
Some(&hint),
ReadColumns::new([0, 1]),
)
.unwrap();
src/mito2/src/read/last_row.rs
+12 -3
View File
@@ -29,6 +29,7 @@ use crate::cache::{
selector_result_cache_hit, selector_result_cache_miss,
};
use crate::error::{ComputeArrowSnafu, Result};
use crate::read::read_columns::JsonTargetTypes;
use crate::read::{
Batch, BatchReader, BoxedBatchReader, BoxedRecordBatchStream, timestamp_array_to_i64_slice,
};
@@ -137,6 +138,7 @@ impl FlatRowGroupLastRowCachedReader {
row_group_idx: usize,
cache_strategy: CacheStrategy,
read_cols: &ParquetReadColumns,
json_target_types: JsonTargetTypes,
reader: FlatRowGroupReader,
) -> Self {
let key = SelectorResultKey {
@@ -148,13 +150,14 @@ impl FlatRowGroupLastRowCachedReader {
if let Some(value) = cache_strategy.get_selector_result(&key) {
let is_flat = matches!(&value.result, SelectorResult::Flat(_));
let schema_matches = value.read_cols == *read_cols;
if is_flat && schema_matches {
let json_target_types_matches = value.json_target_types == json_target_types;
if is_flat && schema_matches && json_target_types_matches {
Self::new_hit(value)
} else {
Self::new_miss(key, read_cols, reader, cache_strategy)
Self::new_miss(key, read_cols, json_target_types, reader, cache_strategy)
}
} else {
Self::new_miss(key, read_cols, reader, cache_strategy)
Self::new_miss(key, read_cols, json_target_types, reader, cache_strategy)
}
}
@@ -174,6 +177,7 @@ impl FlatRowGroupLastRowCachedReader {
fn new_miss(
key: SelectorResultKey,
read_cols: &ParquetReadColumns,
json_target_types: JsonTargetTypes,
reader: FlatRowGroupReader,
cache_strategy: CacheStrategy,
) -> Self {
@@ -181,6 +185,7 @@ impl FlatRowGroupLastRowCachedReader {
Self::Miss(FlatRowGroupLastRowReader::new(
key,
read_cols.clone(),
json_target_types,
reader,
cache_strategy,
))
@@ -260,6 +265,7 @@ pub(crate) struct FlatRowGroupLastRowReader {
yielded_batches: Vec<RecordBatch>,
cache_strategy: CacheStrategy,
read_cols: ParquetReadColumns,
json_target_types: JsonTargetTypes,
/// Accumulates small selector-output batches before concatenating.
pending: BatchBuffer,
}
@@ -268,6 +274,7 @@ impl FlatRowGroupLastRowReader {
fn new(
key: SelectorResultKey,
read_cols: ParquetReadColumns,
json_target_types: JsonTargetTypes,
reader: FlatRowGroupReader,
cache_strategy: CacheStrategy,
) -> Self {
@@ -278,6 +285,7 @@ impl FlatRowGroupLastRowReader {
yielded_batches: vec![],
cache_strategy,
read_cols,
json_target_types,
pending: BatchBuffer::new(),
}
}
@@ -326,6 +334,7 @@ impl FlatRowGroupLastRowReader {
let value = Arc::new(SelectorResultValue::new_flat(
batches,
self.read_cols.clone(),
self.json_target_types.clone(),
));
self.cache_strategy.put_selector_result(self.key, value);
}
src/mito2/src/read/projection.rs
+1 -2
View File
@@ -295,8 +295,7 @@ mod tests {
let mapper = FlatProjectionMapper::new_with_read_columns(
&metadata,
vec![4, 1],
ReadColumns::from_deduped_column_ids([4, 1, 3]),
None,
ReadColumns::new([4, 1, 3]),
)
.unwrap();
assert_eq!(&[4, 1, 3], mapper.read_columns().column_ids().as_slice());
src/mito2/src/read/range_cache.rs
+2 -2
View File
@@ -852,7 +852,7 @@ pub fn bench_cache_flat_range_stream(
let cache_strategy = CacheStrategy::EnableAll(cache_manager);
let fingerprint = ScanRequestFingerprintBuilder {
read_columns: ReadColumns::from_deduped_column_ids(std::iter::empty()),
read_columns: ReadColumns::new(std::iter::empty()),
read_column_types: vec![],
filters: vec![],
time_filters: vec![],
@@ -916,7 +916,7 @@ mod tests {
filter_deleted: bool,
partition_expr_version: u64,
) -> ScanRequestFingerprint {
let read_columns = ReadColumns::from_deduped_column_ids([1, 2]);
let read_columns = ReadColumns::new([1, 2]);
ScanRequestFingerprintBuilder {
read_columns,
read_column_types: vec![None, None],
src/mito2/src/read/read_columns.rs
+39 -86
View File
@@ -12,119 +12,72 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::BTreeMap;
use std::mem;
use std::sync::Arc;
use store_api::storage::{ColumnId, NestedPath};
use datatypes::types::json_type::JsonNativeType;
use store_api::storage::ColumnId;
pub(crate) type JsonTargetTypes = Arc<BTreeMap<ColumnId, JsonNativeType>>;
/// Logical columns to read from a region.
///
/// Read columns describe which logical columns and nested fields should be read
/// from storage. Each read column is identified by its [`ColumnId`],
/// which represents the root column in the storage schema.
///
/// Nested fields under the column are specified by [`NestedPath`] entries.
/// Each path includes the root column name as its first element.
///
/// For example, assume column id `9` corresponds to a root column named `j`
/// with nested fields:
///
/// ```text
/// j
/// ├── a
/// └── b
/// └── c
/// ```
///
/// The following SQL:
///
/// SELECT j.a, j.b.c FROM t
///
/// may produce read columns like:
///
/// ```text
/// ReadColumn {
/// column_id: 9,
/// nested_paths: [
/// ["j", "a"],
/// ["j", "b", "c"],
/// ]
/// }
/// ```
///
/// If `nested_paths` is empty, the whole column will be read.
/// Read columns describe which logical root columns should be read from storage.
/// JSON2 columns can carry query-time target types that are later translated to
/// physical nested parquet paths by the parquet reader.
#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
pub struct ReadColumns {
pub cols: Vec<ReadColumn>,
pub col_ids: Vec<ColumnId>,
json_target_types: JsonTargetTypes,
}
impl ReadColumns {
pub fn from_deduped_column_ids<I>(column_ids: I) -> Self
/// Creates read columns from logical column ids.
///
/// This preserves the input order and duplicate ids.
pub fn new<I>(col_ids: I) -> Self
where
I: IntoIterator<Item = ColumnId>,
{
let cols = column_ids
.into_iter()
.map(|col_id| ReadColumn::new(col_id, vec![]))
.collect();
ReadColumns { cols }
Self {
col_ids: col_ids.into_iter().collect(),
json_target_types: Arc::default(),
}
}
pub fn with_json_target_types(
mut self,
json_target_types: BTreeMap<ColumnId, JsonNativeType>,
) -> Self {
self.json_target_types = Arc::new(json_target_types);
self
}
pub fn is_empty(&self) -> bool {
self.cols.is_empty()
self.col_ids.is_empty()
}
pub fn column_ids_iter(&self) -> impl Iterator<Item = ColumnId> + '_ {
self.cols.iter().map(|column| column.column_id)
self.col_ids.iter().copied()
}
pub fn column_ids(&self) -> Vec<ColumnId> {
self.column_ids_iter().collect()
}
pub fn columns(&self) -> &[ReadColumn] {
&self.cols
pub fn json_target_types(&self) -> &JsonTargetTypes {
&self.json_target_types
}
pub fn json_target_type(&self, column_id: ColumnId) -> Option<&JsonNativeType> {
self.json_target_types.get(&column_id)
}
pub fn estimated_size(&self) -> usize {
self.cols.capacity() * mem::size_of::<ReadColumn>()
+ self
.cols
.iter()
.map(ReadColumn::estimated_size)
.sum::<usize>()
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ReadColumn {
pub column_id: ColumnId,
/// Nested field paths under this column.
/// Empty means reading the whole column.
pub nested_paths: Vec<NestedPath>,
}
impl ReadColumn {
pub fn new(column_id: ColumnId, nested_paths: Vec<NestedPath>) -> Self {
Self {
column_id,
nested_paths,
}
}
pub fn nested_paths(&self) -> &[NestedPath] {
&self.nested_paths
}
pub fn estimated_size(&self) -> usize {
mem::size_of::<ColumnId>()
+ self.nested_paths.capacity() * mem::size_of::<NestedPath>()
+ self
.nested_paths
.iter()
.map(|path| {
path.capacity() * mem::size_of::<String>()
+ path.iter().map(|node| node.capacity()).sum::<usize>()
})
.sum::<usize>()
self.col_ids.capacity() * mem::size_of::<ColumnId>()
+ self.col_ids.len() * mem::size_of::<ColumnId>()
+ self.json_target_types.len()
* (mem::size_of::<ColumnId>() + mem::size_of::<JsonNativeType>())
}
}
src/mito2/src/read/scan_region.rs
+215 -112
View File
@@ -14,7 +14,7 @@
//! Scans a region according to the scan request.
use std::collections::HashSet;
use std::collections::{BTreeMap, HashSet};
use std::fmt;
use std::num::NonZeroU64;
use std::sync::Arc;
@@ -35,18 +35,18 @@ use datafusion_expr::Expr;
use datafusion_expr::utils::expr_to_columns;
use datatypes::arrow::array::{ArrayRef, BooleanArray, UInt64Array};
use datatypes::extension::json::is_json2_extension_type;
use datatypes::prelude::ConcreteDataType;
use datatypes::types::json_type::JsonNativeType;
use datatypes::value::timestamp_to_scalar_value;
use futures::StreamExt;
use itertools::Itertools;
use partition::expr::PartitionExpr;
use smallvec::SmallVec;
use snafu::{OptionExt, ResultExt};
use snafu::{OptionExt, ResultExt, ensure};
use store_api::metadata::{RegionMetadata, RegionMetadataRef};
use store_api::region_engine::{PartitionRange, RegionScannerRef};
use store_api::storage::{
ColumnId, NestedPath, RegionId, ScanRequest, SequenceNumber, SequenceRange,
TimeSeriesDistribution, TimeSeriesRowSelector,
ColumnId, RegionId, ScanRequest, SequenceNumber, SequenceRange, TimeSeriesDistribution,
TimeSeriesRowSelector,
};
use table::predicate::{Predicate, build_time_range_predicate, extract_time_range_from_expr};
use tokio::sync::{Semaphore, mpsc};
@@ -64,7 +64,7 @@ use crate::read::compat::{self, FlatCompatBatch};
use crate::read::flat_projection::FlatProjectionMapper;
use crate::read::range::{FileRangeBuilder, MemRangeBuilder, RangeMeta, RowGroupIndex};
use crate::read::range_cache::{ScanRequestFingerprint, implied_time_range_from_exprs};
use crate::read::read_columns::{ReadColumn, ReadColumns};
use crate::read::read_columns::ReadColumns;
use crate::read::seq_scan::SeqScan;
use crate::read::series_scan::SeriesScan;
use crate::read::stream::ScanBatchStream;
@@ -420,23 +420,7 @@ impl ScanRegion {
let read_col_ids =
self.build_read_col_ids(self.request.projection.as_deref(), &predicate)?;
// Narrow JSON2 columns to avoid reading unnecessary nested fields.
//
// `read_col_ids` selects the root columns required by the projection and predicates,
// while nested projection is currently only applied to JSON2 columns, whose type hints
// further narrow them to the requested nested fields.
let has_structured_json = metadata
.schema
.arrow_schema()
.fields()
.iter()
.any(is_json2_extension_type);
let read_cols = if has_structured_json {
self.read_columns_with_json_type_hint(&read_col_ids)
} else {
ReadColumns::from_deduped_column_ids(read_col_ids.iter().copied())
};
let read_cols = self.build_read_columns(&read_col_ids)?;
// The mapper always computes projected column ids as the schema of SSTs may change.
let projection = self
@@ -444,23 +428,7 @@ impl ScanRegion {
.projection
.clone()
.unwrap_or_else(|| (0..metadata.column_metadatas.len()).collect());
let json_type_hint = has_structured_json
.then_some(&self.request.json_type_hint)
.inspect(|json_type_hint| {
debug!(
"Concretized JSON type: {{{}}}",
json_type_hint
.iter()
.map(|(k, v)| format!("{}: {}", k, v))
.join(", ")
);
});
let mapper = FlatProjectionMapper::new_with_read_columns(
metadata,
projection,
read_cols,
json_type_hint,
)?;
let mapper = FlatProjectionMapper::new_with_read_columns(metadata, projection, read_cols)?;
let mapper = if self.request.preserve_pk_dictionary_encoding {
mapper.with_pk_dictionary_encoding()
} else {
@@ -609,7 +577,8 @@ impl ScanRegion {
Ok(input)
}
/// Builds the ordered root column ids required by the pushed-down projection and predicate.
/// Builds the deduplicated root column ids required by the projection and
/// predicate.
fn build_read_col_ids(
&self,
projection: Option<&[usize]>,
@@ -688,27 +657,50 @@ impl ScanRegion {
Ok(read_col_ids)
}
/// Builds read columns with nested paths derived from JSON type hints.
fn read_columns_with_json_type_hint(&self, col_ids: &[ColumnId]) -> ReadColumns {
let cols = col_ids
/// Builds logical read columns and attaches JSON2 target types when needed.
///
/// The behavior about JSON2 is as follows:
/// - JSON2 columns without hints use Variant to read the whole column.
/// - Hints targeting non-JSON2 read columns are rejected.
fn build_read_columns(&self, col_ids: &[ColumnId]) -> Result<ReadColumns> {
let metadata = &self.version.metadata;
let json_type_hint = &self.request.json_type_hint;
let has_json2 = metadata
.schema
.arrow_schema()
.fields()
.iter()
.map(|&col_id| {
let nested_paths = self
.version
.metadata
.column_by_id(col_id)
.and_then(|column| {
let col_name = &column.column_schema.name;
self.request
.json_type_hint
.get(col_name)
.map(|json_type| json_nested_paths(col_name, json_type))
})
.unwrap_or_default();
ReadColumn::new(col_id, nested_paths)
})
.collect();
ReadColumns { cols }
.any(is_json2_extension_type);
if !has_json2 && json_type_hint.is_empty() {
return Ok(ReadColumns::new(col_ids.iter().copied()));
}
let mut json_target_types = BTreeMap::new();
for &col_id in col_ids {
let Some(col) = metadata.column_by_id(col_id) else {
continue;
};
let col_name = &col.column_schema.name;
let hint = json_type_hint.get(col_name);
if !col.column_schema.data_type.is_json2() {
ensure!(
hint.is_none(),
InvalidRequestSnafu {
region_id: metadata.region_id,
reason: format!(
"JSON type hint targets non-JSON2 column {} (id: {}, type: {})",
col_name, col_id, col.column_schema.data_type
),
}
);
continue;
}
let target_type = hint.cloned().unwrap_or(JsonNativeType::Variant);
json_target_types.insert(col_id, target_type);
}
Ok(ReadColumns::new(col_ids.iter().copied()).with_json_target_types(json_target_types))
}
fn region_id(&self) -> RegionId {
@@ -1583,30 +1575,6 @@ fn pre_filter_mode(append_mode: bool, merge_mode: MergeMode) -> PreFilterMode {
}
}
fn json_nested_paths(column_name: &str, json_type: &JsonNativeType) -> Vec<NestedPath> {
let mut paths = Vec::new();
let mut current = vec![column_name.to_string()];
collect_json_nested_paths(json_type, &mut current, &mut paths);
paths
}
fn collect_json_nested_paths(
json_type: &JsonNativeType,
current: &mut NestedPath,
paths: &mut Vec<NestedPath>,
) {
match json_type {
JsonNativeType::Object(fields) if !fields.is_empty() => {
for (field, child) in fields {
current.push(field.clone());
collect_json_nested_paths(child, current, paths);
current.pop();
}
}
_ => paths.push(current.clone()),
}
}
/// Output of [build_scan_fingerprint]: the cache fingerprint plus the derived
/// implied time range used to decide whether the cache key can drop the time
/// predicates for a given partition (see `build_range_cache_key`).
@@ -1703,9 +1671,15 @@ pub(crate) fn build_scan_fingerprint(input: &ScanInput) -> Option<ScanFingerprin
read_column_types: read_columns
.column_ids_iter()
.map(|id| {
metadata
.column_by_id(id)
.map(|col| col.column_schema.data_type.clone())
read_columns
.json_target_type(id)
.cloned()
.map(ConcreteDataType::json2)
.or_else(|| {
metadata
.column_by_id(id)
.map(|col| col.column_schema.data_type.clone())
})
})
.collect(),
read_columns,
@@ -2096,6 +2070,7 @@ impl PredicateGroup {
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use std::sync::Arc;
use common_time::timestamp::{TimeUnit, Timestamp};
@@ -2225,30 +2200,6 @@ mod tests {
assert_eq!(256, input.batch_size());
}
#[test]
fn test_fill_json_nested_paths_from_hint() -> Result<()> {
let hint = JsonNativeType::Object(JsonObjectType::from([
("a".to_string(), JsonNativeType::i64()),
(
"b".to_string(),
JsonNativeType::Object(JsonObjectType::from([(
"c".to_string(),
JsonNativeType::String,
)])),
),
]));
fn nested_path(parts: &[&str]) -> NestedPath {
parts.iter().map(|part| part.to_string()).collect()
}
assert_eq!(
json_nested_paths("j", &hint),
vec![nested_path(&["j", "a"]), nested_path(&["j", "b", "c"])]
);
Ok(())
}
#[tokio::test]
async fn test_build_scan_fingerprint_for_eligible_scan() {
let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false));
@@ -2370,6 +2321,158 @@ mod tests {
assert_ne!(0, metadata.partition_expr_version);
}
#[tokio::test]
async fn test_build_scan_fingerprint_uses_json_target_types() {
let mut builder = RegionMetadataBuilder::new(RegionId::new(123, 456));
builder
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new(
"k0".to_string(),
ConcreteDataType::string_datatype(),
false,
),
semantic_type: SemanticType::Tag,
column_id: 0,
})
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new(
"ts".to_string(),
ConcreteDataType::timestamp_millisecond_datatype(),
false,
),
semantic_type: SemanticType::Timestamp,
column_id: 1,
})
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new(
"j".to_string(),
ConcreteDataType::json2(JsonNativeType::Variant),
true,
),
semantic_type: SemanticType::Field,
column_id: 2,
})
.primary_key(vec![0]);
let metadata = Arc::new(builder.build().unwrap());
let make_input = |target_type| async {
let env = SchedulerEnv::new().await;
let read_cols = ReadColumns::new([0, 1, 2])
.with_json_target_types(BTreeMap::from([(2, target_type)]));
let mapper =
FlatProjectionMapper::new_with_read_columns(&metadata, vec![0, 1, 2], read_cols)
.unwrap();
let predicate =
PredicateGroup::new(metadata.as_ref(), &[col("k0").eq(lit("foo"))]).unwrap();
let file = FileHandle::new(
FileMeta::default(),
Arc::new(crate::sst::file_purger::NoopFilePurger),
);
ScanInput::new(env.access_layer.clone(), mapper)
.with_predicate(predicate)
.with_cache(CacheStrategy::EnableAll(Arc::new(
CacheManager::builder()
.range_result_cache_size(1024)
.build(),
)))
.with_files(vec![file])
};
let int_target = JsonNativeType::i64();
let string_target = JsonNativeType::String;
let int_fingerprint = build_scan_fingerprint(&make_input(int_target.clone()).await)
.unwrap()
.fingerprint;
let string_fingerprint = build_scan_fingerprint(&make_input(string_target).await)
.unwrap()
.fingerprint;
assert_ne!(int_fingerprint, string_fingerprint);
assert_eq!(
Some(&Some(ConcreteDataType::json2(int_target))),
int_fingerprint.read_column_types().get(2)
);
}
#[tokio::test]
async fn test_scan_input_rejects_json_type_hint_for_non_json2_column() {
let mut builder = RegionMetadataBuilder::new(RegionId::new(123, 456));
builder
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new(
"k0".to_string(),
ConcreteDataType::string_datatype(),
false,
),
semantic_type: SemanticType::Tag,
column_id: 0,
})
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new(
"ts".to_string(),
ConcreteDataType::timestamp_millisecond_datatype(),
false,
),
semantic_type: SemanticType::Timestamp,
column_id: 1,
})
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new(
"j".to_string(),
ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::from([(
"a".to_string(),
JsonNativeType::i64(),
)]))),
true,
),
semantic_type: SemanticType::Field,
column_id: 2,
})
.push_column_metadata(ColumnMetadata {
column_schema: ColumnSchema::new(
"v0".to_string(),
ConcreteDataType::int64_datatype(),
true,
),
semantic_type: SemanticType::Field,
column_id: 3,
})
.primary_key(vec![0]);
let metadata = Arc::new(builder.build().unwrap());
let mutable = Arc::new(crate::memtable::time_partition::TimePartitions::new(
metadata.clone(),
Arc::new(crate::test_util::memtable_util::EmptyMemtableBuilder::default()),
0,
None,
));
let version = Arc::new(
crate::region::version::VersionBuilder::new(metadata.clone(), mutable).build(),
);
let env = SchedulerEnv::new().await;
let request = ScanRequest {
projection: Some(vec![0, 1, 2, 3]),
json_type_hint: std::collections::HashMap::from([(
"v0".to_string(),
JsonNativeType::i64(),
)]),
..Default::default()
};
let err = ScanRegion::new(
version,
env.access_layer.clone(),
request,
CacheStrategy::Disabled,
)
.scan_input()
.await;
let Err(err) = err else {
panic!("scan input should reject JSON type hint for non-JSON2 column");
};
assert!(err.to_string().contains("non-JSON2 column v0"));
}
#[test]
fn test_update_dyn_filters_with_empty_base_predicates() {
let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false));
src/mito2/src/sst/parquet/file_range.rs
+3 -6
View File
@@ -235,6 +235,7 @@ impl FileRange {
self.row_group_idx,
cache_strategy,
self.context.read_format().parquet_read_columns(),
self.context.read_format().json_target_types().clone(),
flat_row_group_reader,
);
FlatPruneReader::new_with_last_row_reader(self.context.clone(), reader, skip_fields)
@@ -827,9 +828,7 @@ mod tests {
let read_format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids(
metadata.column_metadatas.iter().map(|c| c.column_id),
),
ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)),
None,
"test",
true,
@@ -895,9 +894,7 @@ mod tests {
let metadata: RegionMetadataRef = Arc::new(sst_region_metadata());
let read_format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids(
metadata.column_metadatas.iter().map(|c| c.column_id),
),
ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)),
None,
"test",
true,
src/mito2/src/sst/parquet/flat_format.rs
+60 -19
View File
@@ -51,12 +51,11 @@ use crate::error::{
ComputeArrowSnafu, DecodeSnafu, InvalidParquetSnafu, InvalidRecordBatchSnafu,
NewRecordBatchSnafu, Result,
};
use crate::read::read_columns::ReadColumns;
use crate::read::read_columns::{JsonTargetTypes, ReadColumns};
use crate::sst::parquet::format::{
FIXED_POS_COLUMN_NUM, FormatProjection, INTERNAL_COLUMN_NUM, PrimaryKeyArray,
PrimaryKeyReadFormat, StatValues, column_null_counts, column_values,
};
use crate::sst::parquet::json_align::align_schema_by_nested_paths;
use crate::sst::parquet::read_columns::ParquetReadColumns;
use crate::sst::{
FlatSchemaOptions, flat_sst_arrow_schema_column_num, tag_maybe_to_dictionary_field,
@@ -181,6 +180,8 @@ pub(crate) fn field_column_start(metadata: &RegionMetadata, num_columns: usize)
pub struct FlatReadFormat {
/// Sequence number to override the sequence read from the SST.
override_sequence: Option<SequenceNumber>,
/// Logical columns requested by this read.
read_cols: ReadColumns,
/// Parquet format adapter.
parquet_adapter: ParquetAdapter,
/// Output schema to wrap binary `__primary_key` back to a dictionary; `None` disables wrapping.
@@ -210,22 +211,25 @@ impl FlatReadFormat {
// Only skip auto convert when the primary key encoding is sparse.
ParquetAdapter::PrimaryKeyToFlat(ParquetPrimaryKeyToFlat::new(
metadata,
read_cols,
read_cols.clone(),
skip_auto_convert,
))
} else {
ParquetAdapter::PrimaryKeyToFlat(ParquetPrimaryKeyToFlat::new(
metadata, read_cols, false,
metadata,
read_cols.clone(),
false,
))
}
} else {
let file_schema = file_schema
.unwrap_or_else(|| to_flat_sst_arrow_schema(&metadata, &Default::default()));
ParquetAdapter::Flat(ParquetFlat::new(metadata, read_cols, file_schema))
ParquetAdapter::Flat(ParquetFlat::new(metadata, read_cols.clone(), file_schema))
};
Ok(FlatReadFormat {
override_sequence: None,
read_cols,
parquet_adapter,
pk_dict_wrap_schema: None,
})
@@ -294,22 +298,51 @@ impl FlatReadFormat {
}
}
/// Gets the projected output schema produced by parquet reading.
/// Gets the projected output schema expected by the scan.
pub(crate) fn output_arrow_schema(&self) -> Result<SchemaRef> {
let read_columns = self.parquet_read_columns();
let projection = read_columns.root_indices();
let projection = self.parquet_read_columns().root_indices();
let mut schema = self
.arrow_schema()
.project(projection)
.context(ComputeArrowSnafu)?;
if read_columns.has_nested() {
debug_assert_eq!(schema.fields().len(), read_columns.columns().len());
let nested_paths = read_columns.columns().iter().map(|x| x.nested_paths());
align_schema_by_nested_paths(&mut schema, nested_paths);
let mut fields = schema.fields().iter().cloned().collect::<Vec<_>>();
for (column_id, target_type) in self.json_target_types().iter() {
let Some(index) = self.parquet_projected_index_by_id(*column_id) else {
continue;
};
let Some(field) = schema.fields().get(index) else {
continue;
};
fields[index] = Arc::new(
field
.as_ref()
.clone()
.with_data_type(ConcreteDataType::json2(target_type.clone()).as_arrow_type()),
);
}
schema.fields = fields.into();
Ok(Arc::new(schema))
}
/// Index of a column in the projected schema produced directly by parquet
/// reading, before any primary-key-to-flat conversion.
fn parquet_projected_index_by_id(&self, column_id: ColumnId) -> Option<usize> {
match &self.parquet_adapter {
ParquetAdapter::Flat(p) => p
.format_projection
.column_id_to_projected_index
.get(&column_id)
.copied(),
// `format_projection` addresses the post-conversion flat batch here.
// This helper needs the raw primary-key projection used by parquet reading.
ParquetAdapter::PrimaryKeyToFlat(p) => p
.format
.field_id_to_projected_index()
.get(&column_id)
.copied(),
}
}
/// Gets the metadata of the SST.
pub(crate) fn metadata(&self) -> &RegionMetadataRef {
match &self.parquet_adapter {
@@ -326,6 +359,11 @@ impl FlatReadFormat {
}
}
/// Gets JSON2 target types keyed by column id.
pub(crate) fn json_target_types(&self) -> &JsonTargetTypes {
self.read_cols.json_target_types()
}
/// Gets the projection in the flat format.
///
/// When `skip_auto_convert` is enabled (primary-key format read), this returns the
@@ -491,6 +529,7 @@ impl ParquetPrimaryKeyToFlat {
} else {
// Computes the format projection for the new format.
let format_projection = FormatProjection::compute_format_projection(
&metadata,
&id_to_index,
sst_column_num,
read_cols.clone(),
@@ -540,8 +579,12 @@ impl ParquetFlat {
let id_to_index = sst_column_id_indices(&metadata);
let sst_column_num =
flat_sst_arrow_schema_column_num(&metadata, &FlatSchemaOptions::default());
let format_projection =
FormatProjection::compute_format_projection(&id_to_index, sst_column_num, read_cols);
let format_projection = FormatProjection::compute_format_projection(
&metadata,
&id_to_index,
sst_column_num,
read_cols,
);
Self {
metadata,
@@ -841,9 +884,7 @@ impl FlatReadFormat {
pub fn new_with_all_columns(metadata: RegionMetadataRef) -> FlatReadFormat {
Self::new(
Arc::clone(&metadata),
ReadColumns::from_deduped_column_ids(
metadata.column_metadatas.iter().map(|c| c.column_id),
),
ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)),
None,
"test",
false,
@@ -965,7 +1006,7 @@ mod tests {
.collect();
let mut read_format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids(column_ids),
ReadColumns::new(column_ids),
None,
"test",
false,
@@ -1045,7 +1086,7 @@ mod tests {
let metadata = Arc::new(build_metadata(1, 2, PrimaryKeyEncoding::Dense));
let read_format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids([0_u32, 2_u32]),
ReadColumns::new([0_u32, 2_u32]),
None,
"test",
false,
src/mito2/src/sst/parquet/format.rs
+85 -69
View File
@@ -39,6 +39,7 @@ use datatypes::arrow::array::{
use datatypes::arrow::datatypes::{SchemaRef, UInt32Type};
use datatypes::arrow::record_batch::RecordBatch;
use datatypes::prelude::DataType;
use datatypes::types::json_type::JsonNativeType;
use datatypes::vectors::Helper;
use mito_codec::row_converter::{
CompositeValues, PrimaryKeyCodec, SortField, build_primary_key_codec_with_fields,
@@ -47,12 +48,12 @@ use parquet::file::metadata::{ParquetMetaData, RowGroupMetaData};
use parquet::file::statistics::Statistics;
use snafu::{OptionExt, ResultExt, ensure};
use store_api::metadata::{ColumnMetadata, RegionMetadataRef};
use store_api::storage::{ColumnId, SequenceNumber};
use store_api::storage::{ColumnId, NestedPath, SequenceNumber};
use crate::error::{
ConvertVectorSnafu, DecodeSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, Result,
};
use crate::read::read_columns::ReadColumns;
use crate::read::read_columns::{JsonTargetTypes, ReadColumns};
use crate::read::{Batch, BatchBuilder, BatchColumn};
use crate::sst::file::{FileMeta, FileTimeRange};
use crate::sst::parquet::read_columns::{ParquetReadColumn, ParquetReadColumns};
@@ -233,6 +234,7 @@ impl PrimaryKeyReadFormat {
let arrow_schema = to_sst_arrow_schema(&metadata);
let format_projection = FormatProjection::compute_format_projection(
&metadata,
&field_id_to_index,
arrow_schema.fields.len(),
read_cols,
@@ -592,18 +594,21 @@ impl FormatProjection {
///
/// `id_to_index` is a mapping from column id to the index of the column in the SST.
pub(crate) fn compute_format_projection(
metadata: &RegionMetadataRef,
id_to_index: &HashMap<ColumnId, usize>,
sst_column_num: usize,
cols: ReadColumns,
) -> Self {
let json_target_types = cols.json_target_types().clone();
let mut projected_columns: Vec<_> = cols
.cols
.col_ids
.into_iter()
.filter_map(|col| {
id_to_index
.get(&col.column_id)
.copied()
.map(|index_of_sst| (col.column_id, index_of_sst, col.nested_paths))
.filter_map(|col_id| {
id_to_index.get(&col_id).copied().map(|index_of_sst| {
let nested_paths =
json_target_nested_paths(metadata, &json_target_types, col_id);
(col_id, index_of_sst, nested_paths)
})
})
.collect();
// Sorts columns by their indices in the SST. SST uses a bitmap for projection.
@@ -681,6 +686,45 @@ impl FormatProjection {
}
}
fn json_target_nested_paths(
metadata: &RegionMetadataRef,
json_target_types: &JsonTargetTypes,
column_id: ColumnId,
) -> Vec<NestedPath> {
let Some(target_type) = json_target_types.get(&column_id) else {
return Vec::new();
};
let Some(column) = metadata.column_by_id(column_id) else {
return Vec::new();
};
json_nested_paths(&column.column_schema.name, target_type)
}
fn json_nested_paths(column_name: &str, json_type: &JsonNativeType) -> Vec<NestedPath> {
let mut paths = Vec::new();
let mut current = vec![column_name.to_string()];
collect_json_nested_paths(json_type, &mut current, &mut paths);
paths
}
fn collect_json_nested_paths(
json_type: &JsonNativeType,
current: &mut NestedPath,
paths: &mut Vec<NestedPath>,
) {
match json_type {
JsonNativeType::Object(fields) if !fields.is_empty() => {
for (field, child) in fields {
current.push(field.clone());
collect_json_nested_paths(child, current, paths);
current.pop();
}
}
_ => paths.push(current.clone()),
}
}
/// Values of column statistics of the SST.
///
/// It also distinguishes the case that a column is not found and
@@ -710,9 +754,7 @@ impl PrimaryKeyReadFormat {
pub fn new_with_all_columns(metadata: RegionMetadataRef) -> PrimaryKeyReadFormat {
Self::new(
Arc::clone(&metadata),
ReadColumns::from_deduped_column_ids(
metadata.column_metadatas.iter().map(|c| c.column_id),
),
ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)),
)
}
}
@@ -818,6 +860,7 @@ pub(crate) fn need_override_sequence(parquet_meta: &ParquetMetaData) -> bool {
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use std::sync::Arc;
use api::v1::OpType;
@@ -840,7 +883,6 @@ mod tests {
use super::*;
use crate::error::InvalidMetadataSnafu;
use crate::read::read_columns::ReadColumn;
use crate::sst::parquet::flat_format::{
FlatReadFormat, FlatWriteFormat, sequence_column_index, sst_column_id_indices,
};
@@ -986,29 +1028,25 @@ mod tests {
fn test_projection_indices() {
let metadata = build_test_region_metadata();
// Only read tag1
let read_format =
PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::from_deduped_column_ids([3]));
let read_format = PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::new([3]));
assert_eq!(
&[2, 3, 4, 5],
read_format.parquet_read_columns().root_indices()
);
// Only read field1
let read_format =
PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::from_deduped_column_ids([4]));
let read_format = PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::new([4]));
assert_eq!(
&[0, 2, 3, 4, 5],
read_format.parquet_read_columns().root_indices()
);
// Only read ts
let read_format =
PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::from_deduped_column_ids([5]));
let read_format = PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::new([5]));
assert_eq!(
&[2, 3, 4, 5],
read_format.parquet_read_columns().root_indices()
);
// Read field0, tag0, ts
let read_format =
PrimaryKeyReadFormat::new(metadata, ReadColumns::from_deduped_column_ids([2, 1, 5]));
let read_format = PrimaryKeyReadFormat::new(metadata, ReadColumns::new([2, 1, 5]));
assert_eq!(
&[1, 2, 3, 4, 5],
read_format.parquet_read_columns().root_indices()
@@ -1049,14 +1087,16 @@ mod tests {
let metadata = Arc::new(builder.build().context(InvalidMetadataSnafu)?);
let column_id_to_parquet_index = sst_column_id_indices(&metadata);
let projection = FormatProjection::compute_format_projection(
&metadata,
&column_id_to_parquet_index,
metadata.column_metadatas.len() + FIXED_POS_COLUMN_NUM,
ReadColumns {
cols: vec![ReadColumn::new(
4,
vec![vec!["j".to_string(), "a".to_string()]],
)],
},
ReadColumns::new([4]).with_json_target_types(BTreeMap::from([(
4,
JsonNativeType::Object(JsonObjectType::from([(
"a".to_string(),
JsonNativeType::i64(),
)])),
)])),
);
let columns = projection.parquet_read_cols.columns();
@@ -1111,8 +1151,7 @@ mod tests {
.iter()
.map(|col| col.column_id)
.collect();
let read_format =
PrimaryKeyReadFormat::new(metadata, ReadColumns::from_deduped_column_ids(column_ids));
let read_format = PrimaryKeyReadFormat::new(metadata, ReadColumns::new(column_ids));
assert_eq!(arrow_schema, *read_format.arrow_schema());
let record_batch = RecordBatch::new_empty(arrow_schema);
@@ -1131,8 +1170,7 @@ mod tests {
.iter()
.map(|col| col.column_id)
.collect();
let read_format =
PrimaryKeyReadFormat::new(metadata, ReadColumns::from_deduped_column_ids(column_ids));
let read_format = PrimaryKeyReadFormat::new(metadata, ReadColumns::new(column_ids));
let columns: Vec<ArrayRef> = vec![
Arc::new(Int64Array::from(vec![1, 1, 10, 10])), // field1
@@ -1160,9 +1198,7 @@ mod tests {
let metadata = build_test_region_metadata();
let read_format = PrimaryKeyReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids(
metadata.column_metadatas.iter().map(|c| c.column_id),
),
ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)),
);
let columns: Vec<ArrayRef> = vec![
@@ -1330,56 +1366,36 @@ mod tests {
// The projection includes all "fixed position" columns: ts(4), __primary_key(5), __sequence(6), __op_type(7)
// Only read tag1 (column_id=3, index=1) + fixed columns
let read_format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids([3]),
None,
"test",
false,
)
.unwrap();
let read_format =
FlatReadFormat::new(metadata.clone(), ReadColumns::new([3]), None, "test", false)
.unwrap();
assert_eq!(
&[1, 4, 5, 6, 7],
read_format.parquet_read_columns().root_indices()
);
// Only read field1 (column_id=4, index=2) + fixed columns
let read_format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids([4]),
None,
"test",
false,
)
.unwrap();
let read_format =
FlatReadFormat::new(metadata.clone(), ReadColumns::new([4]), None, "test", false)
.unwrap();
assert_eq!(
&[2, 4, 5, 6, 7],
read_format.parquet_read_columns().root_indices()
);
// Only read ts (column_id=5, index=4) + fixed columns (ts is already included in fixed)
let read_format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids([5]),
None,
"test",
false,
)
.unwrap();
let read_format =
FlatReadFormat::new(metadata.clone(), ReadColumns::new([5]), None, "test", false)
.unwrap();
assert_eq!(
&[4, 5, 6, 7],
read_format.parquet_read_columns().root_indices()
);
// Read field0(column_id=2, index=3), tag0(column_id=1, index=0), ts(column_id=5, index=4) + fixed columns
let read_format = FlatReadFormat::new(
metadata,
ReadColumns::from_deduped_column_ids([2, 1, 5]),
None,
"test",
false,
)
.unwrap();
let read_format =
FlatReadFormat::new(metadata, ReadColumns::new([2, 1, 5]), None, "test", false)
.unwrap();
assert_eq!(
&[0, 3, 4, 5, 6, 7],
read_format.parquet_read_columns().root_indices()
@@ -1391,7 +1407,7 @@ mod tests {
let metadata = build_test_region_metadata();
let mut format = FlatReadFormat::new(
metadata,
ReadColumns::from_deduped_column_ids(std::iter::once(1)), // Just read tag0
ReadColumns::new(std::iter::once(1)), // Just read tag0
Some(build_test_flat_sst_schema()),
"test",
false,
@@ -1608,7 +1624,7 @@ mod tests {
.collect();
let format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids(column_ids),
ReadColumns::new(column_ids),
Some(build_test_arrow_schema()),
"test",
false,
@@ -1674,7 +1690,7 @@ mod tests {
.collect();
let format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids(column_ids.clone()),
ReadColumns::new(column_ids.clone()),
None,
"test",
false,
@@ -1744,7 +1760,7 @@ mod tests {
let format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids(column_ids),
ReadColumns::new(column_ids),
None,
"test",
true,
src/mito2/src/sst/parquet/json_align/mod.rs
-2
View File
@@ -15,10 +15,8 @@
use datatypes::arrow::record_batch::RecordBatch;
use futures::stream::BoxStream;
mod schema;
mod stream;
pub(crate) use schema::align_schema_by_nested_paths;
pub(crate) use stream::NestedSchemaAligner;
use crate::error::Result;
src/mito2/src/sst/parquet/json_align/schema.rs
-250
View File
@@ -1,250 +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 std::collections::BTreeMap;
use std::sync::Arc;
use arrow_schema::{DataType as ArrowDataType, FieldRef};
use datatypes::arrow::datatypes::Schema;
use datatypes::extension::json::is_json2_extension_type;
use store_api::storage::NestedPath;
/// Aligns nested struct fields according to the requested nested paths.
///
/// For each root field:
/// - An empty path list keeps the whole field unchanged.
/// - Non-JSON root fields ignore nested paths and keep the whole field unchanged.
/// - JSON2 root fields are rebuilt from `nested_paths`.
/// - Existing schema fields are preserved only when they are the requested leaf.
/// - Requested paths missing from the schema are synthesized with JSONB (`Binary`)
/// leaves.
///
/// For example, if the schema has `j: struct<a: struct<x: Int64, y: Utf8>>`
/// and `nested_paths` requests `j.a.x` and `j.a.z`, the result is
/// `j: struct<a: struct<x: Int64, z: Binary>>`.
pub(crate) fn align_schema_by_nested_paths<'a, I>(schema: &mut Schema, nested_paths: I)
where
I: IntoIterator<Item = &'a [NestedPath]>,
{
let fields = schema
.fields
.into_iter()
.zip(nested_paths)
.map(|(field, paths)| {
if !paths.is_empty() && is_json2_extension_type(field) {
let child_paths = paths
.iter()
.map(|path| {
if path.first().is_some_and(|root| root == field.name()) {
&path[1..]
} else {
path
}
})
.collect::<Vec<_>>();
rebuild_field_by_nested_paths(field, &child_paths)
} else {
field.clone()
}
})
.collect::<Vec<_>>();
schema.fields = fields.into()
}
fn rebuild_field_by_nested_paths(field: &FieldRef, nested_paths: &[&[String]]) -> FieldRef {
if nested_paths.iter().any(|path| path.is_empty()) {
return field.clone();
};
let fields = group_child_paths(nested_paths)
.into_iter()
.map(|(name, paths)| {
let existing = find_struct_child(field, &name);
build_field_from_nested_paths(&name, existing, &paths)
})
.collect::<Vec<_>>();
Arc::new(
field
.as_ref()
.clone()
.with_data_type(ArrowDataType::Struct(fields.into())),
)
}
fn group_child_paths<'a>(nested_paths: &[&'a [String]]) -> BTreeMap<String, Vec<&'a [String]>> {
let mut child_paths = BTreeMap::<String, Vec<&'a [String]>>::new();
for path in nested_paths {
let Some((name, remaining)) = path.split_first() else {
continue;
};
child_paths.entry(name.clone()).or_default().push(remaining);
}
child_paths
}
fn find_struct_child<'a>(field: &'a FieldRef, name: &str) -> Option<&'a FieldRef> {
let ArrowDataType::Struct(fields) = field.data_type() else {
return None;
};
fields.iter().find(|field| field.name() == name)
}
fn build_field_from_nested_paths(
name: &str,
existing: Option<&FieldRef>,
nested_paths: &[&[String]],
) -> FieldRef {
if nested_paths.iter().any(|path| path.is_empty()) {
return existing.cloned().unwrap_or_else(|| new_jsonb_field(name));
}
let fields = group_child_paths(nested_paths)
.into_iter()
.map(|(name, paths)| {
let existing_child = existing.and_then(|field| find_struct_child(field, &name));
build_field_from_nested_paths(&name, existing_child, &paths)
})
.collect::<Vec<_>>();
let field = existing
.map(|field| field.as_ref().clone())
.unwrap_or_else(|| arrow_schema::Field::new(name, ArrowDataType::Binary, true));
Arc::new(field.with_data_type(ArrowDataType::Struct(fields.into())))
}
fn new_jsonb_field(name: &str) -> FieldRef {
Arc::new(arrow_schema::Field::new(name, ArrowDataType::Binary, true))
}
#[cfg(test)]
mod tests {
use arrow_schema::Field;
use datatypes::extension::json::Json2ExtensionType;
use super::*;
#[test]
fn test_align_schema_by_nested_paths() {
fn new_field(name: &str, data_type: ArrowDataType) -> FieldRef {
Arc::new(Field::new(name, data_type, true))
}
fn struct_field(name: &str, fields: impl IntoIterator<Item = FieldRef>) -> FieldRef {
new_field(name, ArrowDataType::Struct(fields.into_iter().collect()))
}
fn json_struct_field(name: &str, fields: impl IntoIterator<Item = FieldRef>) -> FieldRef {
Arc::new(
Field::new(
name,
ArrowDataType::Struct(fields.into_iter().collect()),
true,
)
.with_extension_type(Json2ExtensionType::default()),
)
}
let mut schema = Schema::new([
json_struct_field(
"j",
[
struct_field(
"a",
[
new_field("x", ArrowDataType::Int64),
new_field("y", ArrowDataType::Utf8),
struct_field(
"z",
[
new_field("q", ArrowDataType::Boolean),
new_field("r", ArrowDataType::Float64),
],
),
],
),
new_field("b", ArrowDataType::Utf8),
struct_field(
"c",
vec![
new_field("d", ArrowDataType::Int64),
new_field("e", ArrowDataType::Utf8),
],
),
],
),
new_field("tag", ArrowDataType::Utf8),
struct_field(
"k",
[
new_field("k_0", ArrowDataType::Int64),
new_field("k_1", ArrowDataType::Utf8),
],
),
]);
let nested_paths = [
vec![
["j", "a", "x"].iter().map(|x| x.to_string()).collect(),
["j", "a", "z", "q"].iter().map(|x| x.to_string()).collect(),
["j", "a", "m"].iter().map(|x| x.to_string()).collect(),
["j", "b", "x"].iter().map(|x| x.to_string()).collect(),
["j", "c"].iter().map(|x| x.to_string()).collect(),
["j", "d", "e"].iter().map(|x| x.to_string()).collect(),
],
vec![["tag", "ignored"].iter().map(|x| x.to_string()).collect()],
vec![],
];
align_schema_by_nested_paths(
&mut schema,
nested_paths.iter().map(|paths| paths.as_slice()),
);
let expected = Schema::new([
json_struct_field(
"j",
[
struct_field(
"a",
[
new_field("m", ArrowDataType::Binary),
new_field("x", ArrowDataType::Int64),
struct_field("z", vec![new_field("q", ArrowDataType::Boolean)]),
],
),
struct_field("b", [new_field("x", ArrowDataType::Binary)]),
struct_field(
"c",
[
new_field("d", ArrowDataType::Int64),
new_field("e", ArrowDataType::Utf8),
],
),
struct_field("d", [new_field("e", ArrowDataType::Binary)]),
],
),
new_field("tag", ArrowDataType::Utf8),
struct_field(
"k",
[
new_field("k_0", ArrowDataType::Int64),
new_field("k_1", ArrowDataType::Utf8),
],
),
]);
assert_eq!(schema, expected);
}
}
src/mito2/src/sst/parquet/prefilter.rs
+8 -22
View File
@@ -1457,9 +1457,7 @@ mod tests {
Arc::new(sst_region_metadata_with_encoding(PrimaryKeyEncoding::Dense));
let read_format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids(
metadata.column_metadatas.iter().map(|c| c.column_id),
),
ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)),
None,
"test",
false,
@@ -1495,9 +1493,7 @@ mod tests {
));
let legacy_read_format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids(
metadata.column_metadatas.iter().map(|c| c.column_id),
),
ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)),
None,
"memtable",
false,
@@ -1524,9 +1520,7 @@ mod tests {
let metadata: RegionMetadataRef = Arc::new(sst_region_metadata());
let raw_pk_read_format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids(
metadata.column_metadatas.iter().map(|c| c.column_id),
),
ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)),
None,
"memtable",
true,
@@ -1560,9 +1554,7 @@ mod tests {
let metadata: RegionMetadataRef = Arc::new(sst_region_metadata());
let full_read_format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids(
metadata.column_metadatas.iter().map(|c| c.column_id),
),
ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)),
None,
"test",
true,
@@ -1594,7 +1586,7 @@ mod tests {
let ts = metric_metadata.time_index_column().column_id;
let projected_read_format = FlatReadFormat::new(
metric_metadata.clone(),
ReadColumns::from_deduped_column_ids([field_0, ts]),
ReadColumns::new([field_0, ts]),
None,
"test",
true,
@@ -1646,9 +1638,7 @@ mod tests {
));
let read_format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids(
metadata.column_metadatas.iter().map(|c| c.column_id),
),
ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)),
None,
"test",
false,
@@ -1686,9 +1676,7 @@ mod tests {
let metadata: RegionMetadataRef = Arc::new(sst_region_metadata());
let read_format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids(
metadata.column_metadatas.iter().map(|c| c.column_id),
),
ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)),
None,
"test",
true,
@@ -1734,9 +1722,7 @@ mod tests {
Arc::new(sst_region_metadata_with_encoding(PrimaryKeyEncoding::Dense));
let read_format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids(
metadata.column_metadatas.iter().map(|c| c.column_id),
),
ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)),
None,
"test",
false,
src/mito2/src/sst/parquet/reader.rs
+5 -11
View File
@@ -451,7 +451,7 @@ impl ParquetReaderBuilder {
} else {
let expected_meta = self.expected_metadata.as_ref().unwrap_or(&region_meta);
// Lists all column ids to read, we always use the expected metadata if possible.
ReadColumns::from_deduped_column_ids(
ReadColumns::new(
expected_meta
.column_metadatas
.iter()
@@ -2472,9 +2472,7 @@ mod tests {
let metadata = Arc::new(sst_region_metadata());
let format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids(
metadata.column_metadatas.iter().map(|c| c.column_id),
),
ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)),
None,
"test",
true,
@@ -2679,7 +2677,7 @@ mod tests {
let region_metadata: RegionMetadataRef = Arc::new(sst_region_metadata());
let read_format = FlatReadFormat::new(
region_metadata.clone(),
ReadColumns::from_deduped_column_ids(
ReadColumns::new(
region_metadata
.column_metadatas
.iter()
@@ -2946,9 +2944,7 @@ mod tests {
let expected_metadata = expected_metadata_with_reused_tag_name(metadata.as_ref());
let read_format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids(
metadata.column_metadatas.iter().map(|c| c.column_id),
),
ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)),
None,
"test",
true,
@@ -2970,9 +2966,7 @@ mod tests {
let metadata: RegionMetadataRef = Arc::new(sst_region_metadata());
let read_format = FlatReadFormat::new(
metadata.clone(),
ReadColumns::from_deduped_column_ids(
metadata.column_metadatas.iter().map(|c| c.column_id),
),
ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)),
None,
"test",
true,
src/query/src/optimizer/json_type_concretize.rs
+71 -20
View File
@@ -16,12 +16,14 @@ use std::collections::HashMap;
use arrow_schema::DataType;
use common_function::scalars::json::json_get::JsonGetWithType;
use datafusion::datasource::DefaultTableSource;
use datafusion::datasource::{DefaultTableSource, TableProvider};
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
use datafusion_common::{Result, plan_datafusion_err, plan_err};
use datafusion_expr::{Expr, LogicalPlan};
use datafusion_optimizer::{OptimizerConfig, OptimizerRule};
use datatypes::extension::json::is_json2_extension_type;
use datatypes::types::json_type::{JsonNativeType, JsonObjectType};
use table::table::adapter::DfTableProviderAdapter;
use crate::dummy_catalog::DummyTableProvider;
@@ -56,22 +58,56 @@ impl OptimizerRule for JsonTypeConcretizeRule {
return Ok(Transformed::no(plan));
};
let Some(adapter) = source
.table_provider
.as_any()
.downcast_ref::<DummyTableProvider>()
else {
return Ok(Transformed::no(plan));
};
adapter.with_json_type_hint(json_types.clone());
Ok(Transformed::yes(plan))
if apply_json_type_hint(source.table_provider.as_ref(), &json_types) {
Ok(Transformed::yes(plan))
} else {
Ok(Transformed::no(plan))
}
}
_ => Ok(Transformed::no(plan)),
})
}
}
// FIXME: `json_types` is keyed only by unqualified column name. In joins with
// same-named JSON2 columns, a hint deduced from one scan can be applied to
// another scan. Carry the originating relation/scan when deducing hints.
/// Applies JSON type hints to providers that can carry scan request hints.
///
/// Returns `true` if at least one JSON2 hint is retained and written to the provider.
fn apply_json_type_hint(
provider: &dyn TableProvider,
json_types: &HashMap<String, JsonNativeType>,
) -> bool {
let schema = provider.schema();
let json_types = json_types
.iter()
.filter(|(column, _)| {
schema
.fields()
.iter()
.any(|field| field.name() == *column && is_json2_extension_type(field))
})
.map(|(column, json_type)| (column.clone(), json_type.clone()))
.collect::<HashMap<_, _>>();
if json_types.is_empty() {
return false;
}
if let Some(adapter) = provider.as_any().downcast_ref::<DummyTableProvider>() {
adapter.with_json_type_hint(json_types);
return true;
}
if let Some(adapter) = provider.as_any().downcast_ref::<DfTableProviderAdapter>() {
adapter.with_json_type_hint(json_types);
return true;
}
false
}
fn deduce_json_types(plan: &LogicalPlan) -> Result<HashMap<String, JsonNativeType>> {
let mut json_types = HashMap::<String, JsonNativeType>::new();
@@ -228,11 +264,11 @@ mod tests {
#[test]
fn test_json_type_concretize_rule_rewrite() -> Result<()> {
let exprs = vec![
json_get_expr(col("k0"), path_expr("a.b"), Some(DataType::Int64))?.alias("ab"),
json_get_expr(col("k0"), path_expr("a.c"), None)?.alias("ac"),
json_get_expr(col("k0"), path_expr("d"), Some(DataType::Boolean))?.alias("d"),
json_get_expr(col("j"), path_expr("a.b"), Some(DataType::Int64))?.alias("ab"),
json_get_expr(col("j"), path_expr("a.c"), None)?.alias("ac"),
json_get_expr(col("j"), path_expr("d"), Some(DataType::Boolean))?.alias("d"),
];
let (provider, plan) = build_plan(exprs)?;
let (provider, plan) = build_json2_plan(exprs)?;
assert!(
JsonTypeConcretizeRule
@@ -253,17 +289,17 @@ mod tests {
let request = provider.scan_request();
assert_eq!(1, request.json_type_hint.len());
assert_eq!(Some(&expected), request.json_type_hint.get("k0"));
assert_eq!(Some(&expected), request.json_type_hint.get("j"));
Ok(())
}
#[test]
fn test_json_type_concretize_rule_conflict_to_variant() -> Result<()> {
let exprs = vec![
json_get_expr(col("k0"), path_expr("a"), Some(DataType::Int64))?.alias("a_num"),
json_get_expr(col("k0"), path_expr("a.b"), Some(DataType::Boolean))?.alias("a_obj"),
json_get_expr(col("j"), path_expr("a"), Some(DataType::Int64))?.alias("a_num"),
json_get_expr(col("j"), path_expr("a.b"), Some(DataType::Boolean))?.alias("a_obj"),
];
let (provider, plan) = build_plan(exprs)?;
let (provider, plan) = build_json2_plan(exprs)?;
assert!(
JsonTypeConcretizeRule
@@ -277,11 +313,26 @@ mod tests {
)]));
assert_eq!(
Some(&expected),
provider.scan_request().json_type_hint.get("k0")
provider.scan_request().json_type_hint.get("j")
);
Ok(())
}
#[test]
fn test_json_type_concretize_rule_ignores_non_json2_columns() -> Result<()> {
let exprs =
vec![json_get_expr(col("k0"), path_expr("a.b"), Some(DataType::Int64))?.alias("ab")];
let (provider, plan) = build_plan(exprs)?;
assert!(
!JsonTypeConcretizeRule
.rewrite(plan, &OptimizerContext::default())?
.transformed
);
assert!(provider.scan_request().json_type_hint.is_empty());
Ok(())
}
#[test]
fn test_json_type_concretize_rule_no_json_get() -> Result<()> {
let (provider, plan) = build_plan(vec![col("k0"), col("v0")])?;
src/table/src/table/adapter.rs
+5
View File
@@ -28,6 +28,7 @@ use datafusion_expr::TableProviderFilterPushDown as DfTableProviderFilterPushDow
use datafusion_expr::expr::Expr;
use datafusion_physical_expr::PhysicalSortExpr;
use datafusion_physical_expr::expressions::Column;
use datatypes::types::json_type::JsonNativeType;
use store_api::storage::{ScanRequest, VectorSearchRequest};
use crate::table::{TableRef, TableType};
@@ -63,6 +64,10 @@ impl DfTableProviderAdapter {
self.scan_req.lock().unwrap().vector_search = Some(hint);
}
pub fn with_json_type_hint(&self, hint: std::collections::HashMap<String, JsonNativeType>) {
self.scan_req.lock().unwrap().json_type_hint = hint;
}
pub fn get_vector_search_hint(&self) -> Option<VectorSearchRequest> {
self.scan_req.lock().unwrap().vector_search.clone()
}
tests/cases/standalone/common/types/json/json2_limit.result
+69
View File
@@ -77,6 +77,75 @@ drop table json2_whole_and_path_read;
Affected Rows: 0
create table json2_join_same_name_left (
ts timestamp time index,
k string,
j json2
)
with (
'append_mode' = 'true'
);
Affected Rows: 0
create table json2_join_same_name_right (
ts timestamp time index,
k string,
j json2
)
with (
'append_mode' = 'true'
);
Affected Rows: 0
insert into json2_join_same_name_left values
(1, 'a', '{"a": 1, "left_only": "kept"}');
Affected Rows: 1
insert into json2_join_same_name_right values
(1, 'a', '{"a": "right", "right_only": "should be kept"}');
Affected Rows: 1
admin flush_table('json2_join_same_name_left');
+------------------------------------------------+
| ADMIN flush_table('json2_join_same_name_left') |
+------------------------------------------------+
| 0 |
+------------------------------------------------+
admin flush_table('json2_join_same_name_right');
+-------------------------------------------------+
| ADMIN flush_table('json2_join_same_name_right') |
+-------------------------------------------------+
| 0 |
+-------------------------------------------------+
-- FIXME: This should return `right` and `1`. The current NULL values are caused
-- by JSON type hints losing the table qualifier in joins.
select json_get(r.j, 'a')::string, json_get(l.j, 'a')::int64
from json2_join_same_name_left l
join json2_join_same_name_right r
on l.k = r.k;
+---------------------------------------------+---------------------------------------------------+
| json_get(r.j,Utf8("a")) | arrow_cast(json_get(l.j,Utf8("a")),Utf8("Int64")) |
+---------------------------------------------+---------------------------------------------------+
| {"a":"right","right_only":"should be kept"} | |
+---------------------------------------------+---------------------------------------------------+
drop table json2_join_same_name_left;
Affected Rows: 0
drop table json2_join_same_name_right;
Affected Rows: 0
create table json2_without_append_mode (
ts timestamp time index,
j json2
tests/cases/standalone/common/types/json/json2_limit.sql
+39
View File
@@ -44,6 +44,45 @@ select j from json2_whole_and_path_read where j.a.b = 1;
drop table json2_whole_and_path_read;
create table json2_join_same_name_left (
ts timestamp time index,
k string,
j json2
)
with (
'append_mode' = 'true'
);
create table json2_join_same_name_right (
ts timestamp time index,
k string,
j json2
)
with (
'append_mode' = 'true'
);
insert into json2_join_same_name_left values
(1, 'a', '{"a": 1, "left_only": "kept"}');
insert into json2_join_same_name_right values
(1, 'a', '{"a": "right", "right_only": "should be kept"}');
admin flush_table('json2_join_same_name_left');
admin flush_table('json2_join_same_name_right');
-- FIXME: This should return `right` and `1`. The current NULL values are caused
-- by JSON type hints losing the table qualifier in joins.
select json_get(r.j, 'a')::string, json_get(l.j, 'a')::int64
from json2_join_same_name_left l
join json2_join_same_name_right r
on l.k = r.k;
drop table json2_join_same_name_left;
drop table json2_join_same_name_right;
create table json2_without_append_mode (
ts timestamp time index,
j json2