mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-09-04 20:49:00 +00:00
feat(mito): add candidate series scanner (#8586)
* feat(mito): add candidate series scanner Signed-off-by: evenyag <realevenyag@gmail.com> * refactor(mito): exclude other ranges from candidate scan Signed-off-by: evenyag <realevenyag@gmail.com> * feat(mito): bench series candidate scan Signed-off-by: evenyag <realevenyag@gmail.com> * refactor(mito): clean up series candidate scanner Signed-off-by: evenyag <realevenyag@gmail.com> * refactor(mito): improve candidate series scan Signed-off-by: evenyag <realevenyag@gmail.com> * fix(mito2): guard primary-key-only SST reads Signed-off-by: evenyag <realevenyag@gmail.com> * docs(mito2): clarify candidate scanner invariants Signed-off-by: evenyag <realevenyag@gmail.com> * fix(mito2): compat primary-key-only reads Signed-off-by: evenyag <realevenyag@gmail.com> --------- Signed-off-by: evenyag <realevenyag@gmail.com>
This commit is contained in:
@@ -67,6 +67,13 @@ pub enum Error {
|
||||
location: Location,
|
||||
},
|
||||
|
||||
#[snafu(display("Invalid sparse primary key: {}", reason))]
|
||||
InvalidSparsePrimaryKey {
|
||||
reason: String,
|
||||
#[snafu(implicit)]
|
||||
location: Location,
|
||||
},
|
||||
|
||||
#[snafu(display("Encode null value"))]
|
||||
IndexEncodeNull {
|
||||
#[snafu(implicit)]
|
||||
@@ -94,6 +101,7 @@ impl ErrorExt for Error {
|
||||
StatusCode::InvalidArguments
|
||||
}
|
||||
NotSupportedField { .. } | UnsupportedOperation { .. } => StatusCode::Unsupported,
|
||||
InvalidSparsePrimaryKey { .. } => StatusCode::InvalidArguments,
|
||||
EvaluateFilter { source, .. } => source.status_code(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -21,13 +21,16 @@ use datatypes::prelude::ConcreteDataType;
|
||||
use datatypes::value::{Value, ValueRef};
|
||||
use memcomparable::{Deserializer, Serializer};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use snafu::ResultExt;
|
||||
use snafu::{ResultExt, ensure};
|
||||
use store_api::codec::PrimaryKeyEncoding;
|
||||
use store_api::metadata::RegionMetadataRef;
|
||||
use store_api::storage::ColumnId;
|
||||
use store_api::storage::consts::ReservedColumnId;
|
||||
|
||||
use crate::error::{DeserializeFieldSnafu, Result, SerializeFieldSnafu, UnsupportedOperationSnafu};
|
||||
use crate::error::{
|
||||
DeserializeFieldSnafu, InvalidSparsePrimaryKeySnafu, Result, SerializeFieldSnafu,
|
||||
UnsupportedOperationSnafu,
|
||||
};
|
||||
use crate::key_values::KeyValue;
|
||||
use crate::primary_key_filter::SparsePrimaryKeyFilter;
|
||||
use crate::row_converter::dense::SortField;
|
||||
@@ -346,6 +349,57 @@ impl SparsePrimaryKeyCodec {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Decodes the reserved `(table_id, tsid)` prefix from a sparse primary key.
|
||||
pub fn decode_ids(&self, bytes: &[u8]) -> Result<(u32, u64)> {
|
||||
// Two column IDs, two non-null markers, a u32 table ID, and a u64 TSID.
|
||||
const INTERNAL_PREFIX_LEN: usize = 4 + 1 + 4 + 4 + 1 + 8;
|
||||
ensure!(
|
||||
bytes.len() >= INTERNAL_PREFIX_LEN,
|
||||
InvalidSparsePrimaryKeySnafu {
|
||||
reason: format!(
|
||||
"internal prefix requires at least {INTERNAL_PREFIX_LEN} bytes, got {}",
|
||||
bytes.len()
|
||||
),
|
||||
}
|
||||
);
|
||||
|
||||
let mut deserializer = Deserializer::new(bytes);
|
||||
|
||||
let table_id_column = u32::deserialize(&mut deserializer).context(DeserializeFieldSnafu)?;
|
||||
ensure!(
|
||||
table_id_column == RESERVED_COLUMN_ID_TABLE_ID,
|
||||
InvalidSparsePrimaryKeySnafu {
|
||||
reason: format!(
|
||||
"expected table id column {}, got {}",
|
||||
RESERVED_COLUMN_ID_TABLE_ID, table_id_column
|
||||
),
|
||||
}
|
||||
);
|
||||
let table_id = self.inner.table_id_field.deserialize(&mut deserializer)?;
|
||||
|
||||
let tsid_column = u32::deserialize(&mut deserializer).context(DeserializeFieldSnafu)?;
|
||||
ensure!(
|
||||
tsid_column == RESERVED_COLUMN_ID_TSID,
|
||||
InvalidSparsePrimaryKeySnafu {
|
||||
reason: format!(
|
||||
"expected tsid column {}, got {}",
|
||||
RESERVED_COLUMN_ID_TSID, tsid_column
|
||||
),
|
||||
}
|
||||
);
|
||||
let tsid = self.inner.tsid_field.deserialize(&mut deserializer)?;
|
||||
|
||||
match (table_id, tsid) {
|
||||
(Value::UInt32(table_id), Value::UInt64(tsid)) => Ok((table_id, tsid)),
|
||||
(table_id, tsid) => InvalidSparsePrimaryKeySnafu {
|
||||
reason: format!(
|
||||
"expected UInt32 table id and UInt64 tsid, got {table_id:?} and {tsid:?}"
|
||||
),
|
||||
}
|
||||
.fail(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Decodes the given bytes into a [`SparseValues`].
|
||||
fn decode_sparse(&self, bytes: &[u8]) -> Result<SparseValues> {
|
||||
let mut deserializer = Deserializer::new(bytes);
|
||||
@@ -805,6 +859,21 @@ mod tests {
|
||||
assert_eq!(result, Value::UInt32(42));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_decode_ids() {
|
||||
let region_metadata = test_region_metadata();
|
||||
let codec = SparsePrimaryKeyCodec::new(®ion_metadata);
|
||||
let mut buffer = Vec::new();
|
||||
codec.encode_internal(42, 100, &mut buffer).unwrap();
|
||||
|
||||
assert_eq!((42, 100), codec.decode_ids(&buffer).unwrap());
|
||||
|
||||
let mut invalid = buffer.clone();
|
||||
invalid[0..4].copy_from_slice(&1_u32.to_be_bytes());
|
||||
assert!(codec.decode_ids(&invalid).is_err());
|
||||
assert!(codec.decode_ids(&buffer[..buffer.len() - 1]).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_has_column() {
|
||||
let region_metadata = test_region_metadata();
|
||||
|
||||
@@ -465,6 +465,14 @@ pub enum Error {
|
||||
error: datafusion::error::DataFusionError,
|
||||
},
|
||||
|
||||
#[snafu(display("Failed to merge candidate series"))]
|
||||
MergeCandidateSeries {
|
||||
#[snafu(implicit)]
|
||||
location: Location,
|
||||
#[snafu(source)]
|
||||
error: datafusion::error::DataFusionError,
|
||||
},
|
||||
|
||||
#[snafu(display("Failed to compute vector"))]
|
||||
ComputeVector {
|
||||
#[snafu(implicit)]
|
||||
@@ -1402,6 +1410,7 @@ impl ErrorExt for Error {
|
||||
| DecodeWal { .. }
|
||||
| ComputeArrow { .. }
|
||||
| EvalPartitionFilter { .. }
|
||||
| MergeCandidateSeries { .. }
|
||||
| BiErrors { .. }
|
||||
| StopScheduler { .. }
|
||||
| ComputeVector { .. }
|
||||
|
||||
@@ -33,6 +33,7 @@ pub mod read_columns;
|
||||
pub mod scan_region;
|
||||
pub mod scan_util;
|
||||
pub(crate) mod seq_scan;
|
||||
pub(crate) mod series_candidate;
|
||||
pub mod series_scan;
|
||||
pub mod stream;
|
||||
pub(crate) mod unordered_scan;
|
||||
|
||||
+156
-55
@@ -41,7 +41,7 @@ use store_api::storage::ColumnId;
|
||||
|
||||
use crate::error::{
|
||||
CompatReaderSnafu, ComputeArrowSnafu, ConvertValueSnafu, CreateDefaultSnafu, DecodeSnafu,
|
||||
EncodeSnafu, NewRecordBatchSnafu, Result, UnsupportedOperationSnafu,
|
||||
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};
|
||||
@@ -117,7 +117,10 @@ impl FlatCompatBatch {
|
||||
}
|
||||
|
||||
let expect_schema = mapper.batch_schema();
|
||||
if expect_schema == actual_schema {
|
||||
if expect_schema == actual_schema
|
||||
&& actual.primary_key == mapper.metadata().primary_key
|
||||
&& actual.primary_key_encoding == mapper.metadata().primary_key_encoding
|
||||
{
|
||||
// Although the SST has a different schema, but the schema after projection is the same
|
||||
// as expected schema.
|
||||
return Ok(None);
|
||||
@@ -295,11 +298,16 @@ impl FlatCompatBatch {
|
||||
)
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
|
||||
let compat_batch = RecordBatch::try_new(self.arrow_schema.clone(), columns)
|
||||
.context(NewRecordBatchSnafu)?;
|
||||
let mut columns = columns;
|
||||
let primary_key_index = primary_key_column_index(columns.len());
|
||||
columns[primary_key_index] = self.compat_primary_key(&columns[primary_key_index])?;
|
||||
|
||||
// Handles primary keys.
|
||||
self.compat_pk.compat(compat_batch)
|
||||
RecordBatch::try_new(self.arrow_schema.clone(), columns).context(NewRecordBatchSnafu)
|
||||
}
|
||||
|
||||
/// Makes an encoded primary-key array compatible with the expected metadata.
|
||||
pub(crate) fn compat_primary_key(&self, primary_key: &ArrayRef) -> Result<ArrayRef> {
|
||||
self.compat_pk.compat(primary_key)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -408,18 +416,44 @@ impl FlatRewritePrimaryKey {
|
||||
fn rewrite_key(
|
||||
&self,
|
||||
append_values: &[(ColumnId, Value)],
|
||||
batch: RecordBatch,
|
||||
) -> Result<RecordBatch> {
|
||||
let old_pk_dict_array = batch
|
||||
.column(primary_key_column_index(batch.num_columns()))
|
||||
.as_any()
|
||||
.downcast_ref::<PrimaryKeyArray>()
|
||||
.unwrap();
|
||||
let old_pk_values_array = old_pk_dict_array
|
||||
.values()
|
||||
.as_any()
|
||||
.downcast_ref::<BinaryArray>()
|
||||
.unwrap();
|
||||
primary_key: &ArrayRef,
|
||||
) -> Result<ArrayRef> {
|
||||
if let Some(old_pk_dict_array) = primary_key.as_any().downcast_ref::<PrimaryKeyArray>() {
|
||||
let old_pk_values_array = old_pk_dict_array
|
||||
.values()
|
||||
.as_any()
|
||||
.downcast_ref::<BinaryArray>()
|
||||
.context(UnexpectedSnafu {
|
||||
reason: "Primary-key dictionary values are not binary",
|
||||
})?;
|
||||
let new_pk_values_array =
|
||||
Arc::new(self.rewrite_values(append_values, old_pk_values_array)?);
|
||||
return Ok(Arc::new(PrimaryKeyArray::new(
|
||||
old_pk_dict_array.keys().clone(),
|
||||
new_pk_values_array,
|
||||
)));
|
||||
}
|
||||
|
||||
let old_pk_values_array =
|
||||
primary_key
|
||||
.as_any()
|
||||
.downcast_ref::<BinaryArray>()
|
||||
.context(UnexpectedSnafu {
|
||||
reason: format!(
|
||||
"Primary-key column is neither binary nor dictionary, got {:?}",
|
||||
primary_key.data_type()
|
||||
),
|
||||
})?;
|
||||
Ok(Arc::new(
|
||||
self.rewrite_values(append_values, old_pk_values_array)?,
|
||||
))
|
||||
}
|
||||
|
||||
fn rewrite_values(
|
||||
&self,
|
||||
append_values: &[(ColumnId, Value)],
|
||||
old_pk_values_array: &BinaryArray,
|
||||
) -> Result<BinaryArray> {
|
||||
let mut builder = BinaryBuilder::with_capacity(
|
||||
old_pk_values_array.len(),
|
||||
old_pk_values_array.value_data().len(),
|
||||
@@ -461,14 +495,7 @@ impl FlatRewritePrimaryKey {
|
||||
}
|
||||
builder.append_value(&buffer);
|
||||
}
|
||||
let new_pk_values_array = Arc::new(builder.finish());
|
||||
let new_pk_dict_array =
|
||||
PrimaryKeyArray::new(old_pk_dict_array.keys().clone(), new_pk_values_array);
|
||||
|
||||
let mut columns = batch.columns().to_vec();
|
||||
columns[primary_key_column_index(batch.num_columns())] = Arc::new(new_pk_dict_array);
|
||||
|
||||
RecordBatch::try_new(batch.schema(), columns).context(NewRecordBatchSnafu)
|
||||
Ok(builder.finish())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -538,34 +565,58 @@ impl FlatCompatPrimaryKey {
|
||||
})
|
||||
}
|
||||
|
||||
/// Makes primary key of the `batch` compatible.
|
||||
///
|
||||
/// Callers must ensure other columns except the `__primary_key` column is compatible.
|
||||
fn compat(&self, batch: RecordBatch) -> Result<RecordBatch> {
|
||||
/// Makes an encoded primary-key array compatible.
|
||||
fn compat(&self, primary_key: &ArrayRef) -> Result<ArrayRef> {
|
||||
if let Some(rewriter) = &self.rewriter {
|
||||
// If we have different encoding, rewrite the whole primary key.
|
||||
return rewriter.rewrite_key(&self.values, batch);
|
||||
return rewriter.rewrite_key(&self.values, primary_key);
|
||||
}
|
||||
|
||||
self.append_key(batch)
|
||||
self.append_key(primary_key)
|
||||
}
|
||||
|
||||
/// Appends values to the primary key of the `batch`.
|
||||
fn append_key(&self, batch: RecordBatch) -> Result<RecordBatch> {
|
||||
/// Appends values to the primary key array.
|
||||
fn append_key(&self, primary_key: &ArrayRef) -> Result<ArrayRef> {
|
||||
let Some(converter) = &self.converter else {
|
||||
return Ok(batch);
|
||||
return Ok(primary_key.clone());
|
||||
};
|
||||
|
||||
let old_pk_dict_array = batch
|
||||
.column(primary_key_column_index(batch.num_columns()))
|
||||
.as_any()
|
||||
.downcast_ref::<PrimaryKeyArray>()
|
||||
.unwrap();
|
||||
let old_pk_values_array = old_pk_dict_array
|
||||
.values()
|
||||
.as_any()
|
||||
.downcast_ref::<BinaryArray>()
|
||||
.unwrap();
|
||||
if let Some(old_pk_dict_array) = primary_key.as_any().downcast_ref::<PrimaryKeyArray>() {
|
||||
let old_pk_values_array = old_pk_dict_array
|
||||
.values()
|
||||
.as_any()
|
||||
.downcast_ref::<BinaryArray>()
|
||||
.context(UnexpectedSnafu {
|
||||
reason: "Primary-key dictionary values are not binary",
|
||||
})?;
|
||||
let new_pk_values_array =
|
||||
Arc::new(self.append_values(old_pk_values_array, converter.as_ref())?);
|
||||
return Ok(Arc::new(PrimaryKeyArray::new(
|
||||
old_pk_dict_array.keys().clone(),
|
||||
new_pk_values_array,
|
||||
)));
|
||||
}
|
||||
|
||||
let old_pk_values_array =
|
||||
primary_key
|
||||
.as_any()
|
||||
.downcast_ref::<BinaryArray>()
|
||||
.context(UnexpectedSnafu {
|
||||
reason: format!(
|
||||
"Primary-key column is neither binary nor dictionary, got {:?}",
|
||||
primary_key.data_type()
|
||||
),
|
||||
})?;
|
||||
Ok(Arc::new(
|
||||
self.append_values(old_pk_values_array, converter.as_ref())?,
|
||||
))
|
||||
}
|
||||
|
||||
fn append_values(
|
||||
&self,
|
||||
old_pk_values_array: &BinaryArray,
|
||||
converter: &dyn PrimaryKeyCodec,
|
||||
) -> Result<BinaryArray> {
|
||||
let mut builder = BinaryBuilder::with_capacity(
|
||||
old_pk_values_array.len(),
|
||||
old_pk_values_array.value_data().len()
|
||||
@@ -594,15 +645,7 @@ impl FlatCompatPrimaryKey {
|
||||
builder.append_value(&buffer);
|
||||
}
|
||||
|
||||
let new_pk_values_array = Arc::new(builder.finish());
|
||||
let new_pk_dict_array =
|
||||
PrimaryKeyArray::new(old_pk_dict_array.keys().clone(), new_pk_values_array);
|
||||
|
||||
// Overrides the primary key column.
|
||||
let mut columns = batch.columns().to_vec();
|
||||
columns[primary_key_column_index(batch.num_columns())] = Arc::new(new_pk_dict_array);
|
||||
|
||||
RecordBatch::try_new(batch.schema(), columns).context(NewRecordBatchSnafu)
|
||||
Ok(builder.finish())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -612,7 +655,7 @@ mod tests {
|
||||
|
||||
use api::v1::{OpType, SemanticType};
|
||||
use datatypes::arrow::array::{
|
||||
ArrayRef, BinaryDictionaryBuilder, Int64Array, StringDictionaryBuilder,
|
||||
ArrayRef, BinaryArray, BinaryDictionaryBuilder, Int64Array, StringDictionaryBuilder,
|
||||
TimestampMillisecondArray, UInt8Array, UInt64Array,
|
||||
};
|
||||
use datatypes::arrow::datatypes::UInt32Type;
|
||||
@@ -974,6 +1017,64 @@ mod tests {
|
||||
assert_eq!(expected_batch, result);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compat_primary_key_with_different_encoding_only() {
|
||||
let mut actual_metadata = new_metadata(
|
||||
&[
|
||||
(
|
||||
0,
|
||||
SemanticType::Timestamp,
|
||||
ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
),
|
||||
(1, SemanticType::Tag, ConcreteDataType::string_datatype()),
|
||||
(2, SemanticType::Field, ConcreteDataType::int64_datatype()),
|
||||
],
|
||||
&[1],
|
||||
);
|
||||
actual_metadata.primary_key_encoding = PrimaryKeyEncoding::Dense;
|
||||
let actual_metadata = Arc::new(actual_metadata);
|
||||
|
||||
let mut expected_metadata = (*actual_metadata).clone();
|
||||
expected_metadata.primary_key_encoding = PrimaryKeyEncoding::Sparse;
|
||||
let expected_metadata = Arc::new(expected_metadata);
|
||||
|
||||
let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap();
|
||||
let read_format = FlatReadFormat::new(
|
||||
actual_metadata,
|
||||
ReadColumns::from_deduped_column_ids([0, 1, 2]),
|
||||
None,
|
||||
"test",
|
||||
false,
|
||||
)
|
||||
.unwrap();
|
||||
let compat = FlatCompatBatch::try_new(&mapper, &read_format, false)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
let dense_key = encode_key(&[Some("tag1")]);
|
||||
let sparse_key = encode_sparse_key(&[(1, Some("tag1"))]);
|
||||
|
||||
let dictionary_key = build_flat_test_pk_array(&[&dense_key, &dense_key]);
|
||||
let result = compat.compat_primary_key(&dictionary_key).unwrap();
|
||||
let result = result.as_any().downcast_ref::<PrimaryKeyArray>().unwrap();
|
||||
let values = result
|
||||
.values()
|
||||
.as_any()
|
||||
.downcast_ref::<BinaryArray>()
|
||||
.unwrap();
|
||||
assert_eq!(values.value(result.keys().value(0) as usize), sparse_key);
|
||||
assert_eq!(values.value(result.keys().value(1) as usize), sparse_key);
|
||||
|
||||
let binary_key: ArrayRef = Arc::new(BinaryArray::from(vec![
|
||||
Some(dense_key.as_slice()),
|
||||
Some(dense_key.as_slice()),
|
||||
]));
|
||||
let result = compat.compat_primary_key(&binary_key).unwrap();
|
||||
let result = result.as_any().downcast_ref::<BinaryArray>().unwrap();
|
||||
assert_eq!(result.value(0), sparse_key);
|
||||
assert_eq!(result.value(1), sparse_key);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_flat_compat_batch_compact_sparse() {
|
||||
let mut actual_metadata = new_metadata(
|
||||
|
||||
@@ -62,11 +62,18 @@ pub(crate) struct ScanRequestFingerprint {
|
||||
append_mode: bool,
|
||||
filter_deleted: bool,
|
||||
merge_mode: MergeMode,
|
||||
stage: RangeScanStage,
|
||||
/// We keep the partition expr version to ensure we won't reuse the fingerprint after we change the partition expr.
|
||||
/// We store the version instead of the whole partition expr or partition expr filters.
|
||||
partition_expr_version: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
enum RangeScanStage {
|
||||
Data,
|
||||
CandidateSeries,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct ScanRequestFingerprintBuilder {
|
||||
pub(crate) read_columns: ReadColumns,
|
||||
@@ -105,6 +112,7 @@ impl ScanRequestFingerprintBuilder {
|
||||
append_mode,
|
||||
filter_deleted,
|
||||
merge_mode,
|
||||
stage: RangeScanStage::Data,
|
||||
partition_expr_version,
|
||||
}
|
||||
}
|
||||
@@ -154,6 +162,20 @@ impl ScanRequestFingerprint {
|
||||
append_mode: self.append_mode,
|
||||
filter_deleted: self.filter_deleted,
|
||||
merge_mode: self.merge_mode,
|
||||
stage: self.stage,
|
||||
partition_expr_version: self.partition_expr_version,
|
||||
}
|
||||
}
|
||||
|
||||
fn for_candidate_series(&self) -> Self {
|
||||
Self {
|
||||
inner: Arc::clone(&self.inner),
|
||||
time_filters: self.time_filters.clone(),
|
||||
series_row_selector: self.series_row_selector,
|
||||
append_mode: self.append_mode,
|
||||
filter_deleted: self.filter_deleted,
|
||||
merge_mode: self.merge_mode,
|
||||
stage: RangeScanStage::CandidateSeries,
|
||||
partition_expr_version: self.partition_expr_version,
|
||||
}
|
||||
}
|
||||
@@ -447,6 +469,22 @@ fn implied_from_between(
|
||||
pub(crate) fn build_range_cache_key(
|
||||
stream_ctx: &StreamContext,
|
||||
part_range: &PartitionRange,
|
||||
) -> Option<RangeScanCacheKey> {
|
||||
build_range_cache_key_inner(stream_ctx, part_range, false)
|
||||
}
|
||||
|
||||
/// Builds a cache key for a candidate-series partition-range result.
|
||||
pub(crate) fn build_candidate_range_cache_key(
|
||||
stream_ctx: &StreamContext,
|
||||
part_range: &PartitionRange,
|
||||
) -> Option<RangeScanCacheKey> {
|
||||
build_range_cache_key_inner(stream_ctx, part_range, true)
|
||||
}
|
||||
|
||||
fn build_range_cache_key_inner(
|
||||
stream_ctx: &StreamContext,
|
||||
part_range: &PartitionRange,
|
||||
candidate_series: bool,
|
||||
) -> Option<RangeScanCacheKey> {
|
||||
if !stream_ctx.input.cache_strategy.has_range_result_cache() {
|
||||
return None;
|
||||
@@ -499,6 +537,11 @@ pub(crate) fn build_range_cache_key(
|
||||
} else {
|
||||
fingerprint.clone()
|
||||
};
|
||||
let scan = if candidate_series {
|
||||
scan.for_candidate_series()
|
||||
} else {
|
||||
scan
|
||||
};
|
||||
|
||||
Some(RangeScanCacheKey {
|
||||
region_id: stream_ctx.input.region_metadata().region_id,
|
||||
@@ -1090,6 +1133,22 @@ mod tests {
|
||||
assert!(key_a.scan.time_filters().is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn candidate_series_cache_key_is_separate_from_data() {
|
||||
let partition_range = (
|
||||
Timestamp::new_millisecond(1000),
|
||||
Timestamp::new_millisecond(2000),
|
||||
);
|
||||
let (ctx, part_range) =
|
||||
new_stream_context(vec![col("k0").eq(lit("foo"))], None, partition_range).await;
|
||||
|
||||
let data_key = build_range_cache_key(&ctx, &part_range).unwrap();
|
||||
let candidate_key = build_candidate_range_cache_key(&ctx, &part_range).unwrap();
|
||||
|
||||
assert_ne!(data_key.scan, candidate_key.scan);
|
||||
assert_eq!(data_key.row_groups, candidate_key.row_groups);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disables_optimization_on_or_clause() {
|
||||
let partition_range = (
|
||||
|
||||
@@ -0,0 +1,737 @@
|
||||
// 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.
|
||||
|
||||
//! Candidate metric-series discovery for the two-stage series scan.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
use async_stream::try_stream;
|
||||
use datafusion::execution::memory_pool::{MemoryConsumer, MemoryPool};
|
||||
use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr};
|
||||
use datafusion::physical_plan::expressions::Column;
|
||||
use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet};
|
||||
use datafusion::physical_plan::sorts::streaming_merge::StreamingMergeBuilder;
|
||||
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
|
||||
use datafusion_common::DataFusionError;
|
||||
use datatypes::arrow::array::{Array, BinaryArray, BinaryBuilder};
|
||||
use datatypes::arrow::compute::SortOptions;
|
||||
use datatypes::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
|
||||
use datatypes::arrow::record_batch::RecordBatch;
|
||||
use datatypes::prelude::ConcreteDataType;
|
||||
use futures::stream::BoxStream;
|
||||
use futures::{StreamExt, TryStreamExt};
|
||||
use mito_codec::row_converter::{PrimaryKeyFilter, SparsePrimaryKeyCodec};
|
||||
use snafu::{OptionExt, ResultExt, ensure};
|
||||
use store_api::codec::PrimaryKeyEncoding;
|
||||
use store_api::region_engine::PartitionRange;
|
||||
use store_api::storage::consts::{PRIMARY_KEY_COLUMN_NAME, ReservedColumnId};
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use crate::error::{
|
||||
InvalidRequestSnafu, JoinSnafu, MergeCandidateSeriesSnafu, NewRecordBatchSnafu, Result,
|
||||
UnexpectedSnafu,
|
||||
};
|
||||
use crate::read::BoxedRecordBatchStream;
|
||||
use crate::read::pruner::{PartitionPruner, Pruner};
|
||||
use crate::read::range::RowGroupIndex;
|
||||
use crate::read::range_cache::{
|
||||
build_candidate_range_cache_key, cache_flat_range_stream, cached_flat_range_stream,
|
||||
};
|
||||
use crate::read::scan_region::StreamContext;
|
||||
use crate::read::scan_util::{PartitionMetrics, new_filter_metrics, scan_flat_mem_ranges};
|
||||
use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE;
|
||||
use crate::sst::parquet::format::PrimaryKeyArray;
|
||||
use crate::sst::parquet::prefilter::{
|
||||
CachedPrimaryKeyFilter, build_primary_key_filter, prefilter_flat_batch_by_primary_key,
|
||||
};
|
||||
use crate::sst::parquet::reader::ReaderMetrics;
|
||||
use crate::sst::parquet::row_group::ParquetFetchMetrics;
|
||||
|
||||
const CANDIDATE_SERIES_BATCH_SIZE: usize = 500;
|
||||
|
||||
/// Identifies one series in a physical metric region.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) struct MetricSeriesId {
|
||||
pub(crate) table_id: u32,
|
||||
pub(crate) tsid: u64,
|
||||
}
|
||||
|
||||
pub(crate) type MetricSeriesIdStream = BoxStream<'static, Result<Vec<MetricSeriesId>>>;
|
||||
|
||||
/// Builds candidate metric series from the ranges assigned to a [`SeriesScan`](super::series_scan::SeriesScan).
|
||||
#[allow(dead_code)]
|
||||
pub(crate) struct SeriesCandidateScanner {
|
||||
stream_ctx: Arc<StreamContext>,
|
||||
partitions: Vec<Vec<PartitionRange>>,
|
||||
pruner: Arc<Pruner>,
|
||||
range_semaphore: Arc<Semaphore>,
|
||||
memory_pool: Arc<dyn MemoryPool>,
|
||||
metrics_set: ExecutionPlanMetricsSet,
|
||||
part_metrics: PartitionMetrics,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
impl SeriesCandidateScanner {
|
||||
/// Creates a candidate-series scanner for native memtable and SST ranges.
|
||||
///
|
||||
/// Callers must fall back to the legacy series-scan path when the scan contains
|
||||
/// extension ranges. Candidate-series discovery does not support other range types.
|
||||
pub(crate) fn try_new(
|
||||
stream_ctx: Arc<StreamContext>,
|
||||
partitions: Vec<Vec<PartitionRange>>,
|
||||
pruner: Arc<Pruner>,
|
||||
range_semaphore: Arc<Semaphore>,
|
||||
memory_pool: Arc<dyn MemoryPool>,
|
||||
metrics_set: ExecutionPlanMetricsSet,
|
||||
part_metrics: PartitionMetrics,
|
||||
) -> Result<Self> {
|
||||
validate_metric_metadata(&stream_ctx)?;
|
||||
#[cfg(feature = "enterprise")]
|
||||
ensure!(
|
||||
stream_ctx.input.extension_ranges().is_empty(),
|
||||
InvalidRequestSnafu {
|
||||
region_id: stream_ctx.input.region_metadata().region_id,
|
||||
reason: "candidate-series scan does not support extension ranges; use the legacy series-scan path",
|
||||
}
|
||||
);
|
||||
Ok(Self {
|
||||
stream_ctx,
|
||||
partitions,
|
||||
pruner,
|
||||
range_semaphore,
|
||||
memory_pool,
|
||||
metrics_set,
|
||||
part_metrics,
|
||||
})
|
||||
}
|
||||
|
||||
/// Builds a globally sorted stream of candidate metric-series IDs.
|
||||
pub(crate) async fn build_stream(&self) -> Result<MetricSeriesIdStream> {
|
||||
let all_ranges = self
|
||||
.partitions
|
||||
.iter()
|
||||
.flatten()
|
||||
.copied()
|
||||
.collect::<Vec<_>>();
|
||||
self.pruner.add_partition_ranges(&all_ranges);
|
||||
let partition_pruner = Arc::new(PartitionPruner::new(self.pruner.clone(), &all_ranges));
|
||||
|
||||
let range_builder = SeriesCandidateRangeBuilder {
|
||||
stream_ctx: self.stream_ctx.clone(),
|
||||
range_semaphore: self.range_semaphore.clone(),
|
||||
memory_pool: self.memory_pool.clone(),
|
||||
metrics_set: self.metrics_set.clone(),
|
||||
part_metrics: self.part_metrics.clone(),
|
||||
};
|
||||
let mut tasks = Vec::with_capacity(all_ranges.len());
|
||||
for (range_idx, part_range) in all_ranges.into_iter().enumerate() {
|
||||
let range_builder = range_builder.clone();
|
||||
let partition_pruner = partition_pruner.clone();
|
||||
tasks.push(common_runtime::spawn_query(async move {
|
||||
let _permit = range_builder
|
||||
.range_semaphore
|
||||
.clone()
|
||||
.acquire_owned()
|
||||
.await
|
||||
.map_err(|error| {
|
||||
UnexpectedSnafu {
|
||||
reason: format!("failed to acquire candidate range permit: {error}"),
|
||||
}
|
||||
.build()
|
||||
})?;
|
||||
range_builder
|
||||
.build_range_stream(part_range, partition_pruner, range_idx)
|
||||
.await
|
||||
}));
|
||||
}
|
||||
|
||||
let mut range_streams = Vec::with_capacity(tasks.len());
|
||||
for task in tasks {
|
||||
range_streams.push(task.await.context(JoinSnafu)??);
|
||||
}
|
||||
|
||||
// Keep scanner-level merge metrics in the same synthetic partition as
|
||||
// SeriesDistributor. Output partitions occupy 0..self.partitions.len().
|
||||
let merged = merge_primary_key_streams(
|
||||
range_streams,
|
||||
self.memory_pool.clone(),
|
||||
&self.metrics_set,
|
||||
self.partitions.len(),
|
||||
"SeriesCandidateScanner::final_merge",
|
||||
)?;
|
||||
decode_metric_series(merged, self.stream_ctx.input.region_metadata().clone())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct SeriesCandidateRangeBuilder {
|
||||
stream_ctx: Arc<StreamContext>,
|
||||
range_semaphore: Arc<Semaphore>,
|
||||
memory_pool: Arc<dyn MemoryPool>,
|
||||
metrics_set: ExecutionPlanMetricsSet,
|
||||
part_metrics: PartitionMetrics,
|
||||
}
|
||||
|
||||
impl SeriesCandidateRangeBuilder {
|
||||
async fn build_range_stream(
|
||||
&self,
|
||||
part_range: PartitionRange,
|
||||
partition_pruner: Arc<PartitionPruner>,
|
||||
merge_partition: usize,
|
||||
) -> Result<BoxedRecordBatchStream> {
|
||||
let cache_key = build_candidate_range_cache_key(&self.stream_ctx, &part_range);
|
||||
if let Some(key) = cache_key.as_ref() {
|
||||
if let Some(value) = self.stream_ctx.input.cache_strategy.get_range_result(key) {
|
||||
self.part_metrics.inc_range_cache_hit();
|
||||
return Ok(cached_flat_range_stream(value));
|
||||
}
|
||||
self.part_metrics.inc_range_cache_miss();
|
||||
}
|
||||
|
||||
let range_meta = &self.stream_ctx.ranges[part_range.identifier];
|
||||
let mut sources = Vec::with_capacity(range_meta.row_group_indices.len());
|
||||
for index in &range_meta.row_group_indices {
|
||||
let source = self
|
||||
.build_source(*index, range_meta.time_range, partition_pruner.clone())
|
||||
.await?;
|
||||
if let Some(source) = source {
|
||||
sources.push(source);
|
||||
}
|
||||
}
|
||||
|
||||
let sources = self.stream_ctx.input.create_parallel_flat_sources(
|
||||
sources,
|
||||
self.range_semaphore.clone(),
|
||||
2,
|
||||
)?;
|
||||
let stream = merge_primary_key_streams(
|
||||
sources,
|
||||
self.memory_pool.clone(),
|
||||
&self.metrics_set,
|
||||
merge_partition,
|
||||
"SeriesCandidateScanner::range_merge",
|
||||
)?;
|
||||
|
||||
Ok(match cache_key {
|
||||
Some(key) => cache_flat_range_stream(
|
||||
stream,
|
||||
self.stream_ctx.input.cache_strategy.clone(),
|
||||
key,
|
||||
self.part_metrics.clone(),
|
||||
),
|
||||
None => stream,
|
||||
})
|
||||
}
|
||||
|
||||
async fn build_source(
|
||||
&self,
|
||||
index: RowGroupIndex,
|
||||
time_range: crate::sst::file::FileTimeRange,
|
||||
partition_pruner: Arc<PartitionPruner>,
|
||||
) -> Result<Option<BoxedRecordBatchStream>> {
|
||||
let metadata = self.stream_ctx.input.region_metadata().clone();
|
||||
if self.stream_ctx.is_mem_range_index(index) {
|
||||
let raw = scan_flat_mem_ranges(
|
||||
self.stream_ctx.clone(),
|
||||
self.part_metrics.clone(),
|
||||
index,
|
||||
time_range,
|
||||
);
|
||||
let filter = build_primary_key_filter(
|
||||
&metadata,
|
||||
self.stream_ctx.input.predicate_group().predicate(),
|
||||
);
|
||||
return Ok(Some(candidate_primary_key_stream(Box::pin(raw), filter)));
|
||||
}
|
||||
|
||||
if self.stream_ctx.is_file_range_index(index) {
|
||||
if partition_pruner.try_skip_manifest_pruned_file_range(index, &self.part_metrics) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let mut reader_metrics = ReaderMetrics {
|
||||
filter_metrics: new_filter_metrics(self.part_metrics.explain_verbose()),
|
||||
..Default::default()
|
||||
};
|
||||
let ranges = partition_pruner
|
||||
.build_file_ranges(index, &self.part_metrics, &mut reader_metrics)
|
||||
.await?;
|
||||
self.part_metrics.inc_num_file_ranges(ranges.len());
|
||||
self.part_metrics
|
||||
.merge_reader_metrics(&reader_metrics, None);
|
||||
|
||||
// Reuse the exact encoded-PK filter selected by the file's reader plan, but
|
||||
// execute it in `candidate_primary_key_stream` after the PK-only read.
|
||||
let filter = ranges.first().and_then(|range| range.primary_key_filter());
|
||||
let part_metrics = self.part_metrics.clone();
|
||||
let raw = Box::pin(try_stream! {
|
||||
let fetch_metrics = part_metrics
|
||||
.explain_verbose()
|
||||
.then(|| Arc::new(ParquetFetchMetrics::default()));
|
||||
let mut reader_metrics = ReaderMetrics {
|
||||
fetch_metrics: fetch_metrics.clone(),
|
||||
..Default::default()
|
||||
};
|
||||
for range in ranges {
|
||||
let build_start = Instant::now();
|
||||
let Some(mut reader) = range
|
||||
.primary_key_reader(fetch_metrics.as_deref())
|
||||
.await?
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
reader_metrics.build_cost += build_start.elapsed();
|
||||
|
||||
let scan_start = Instant::now();
|
||||
while let Some(batch) = reader.try_next().await? {
|
||||
reader_metrics.num_record_batches += 1;
|
||||
reader_metrics.num_batches += 1;
|
||||
reader_metrics.num_rows += batch.num_rows();
|
||||
yield batch;
|
||||
}
|
||||
reader_metrics.scan_cost += scan_start.elapsed();
|
||||
}
|
||||
reader_metrics.observe_rows("candidate_series");
|
||||
part_metrics.merge_reader_metrics(&reader_metrics, None);
|
||||
});
|
||||
return Ok(Some(candidate_primary_key_stream(raw, filter)));
|
||||
}
|
||||
|
||||
UnexpectedSnafu {
|
||||
reason: format!(
|
||||
"candidate-series scan received unsupported range index {}",
|
||||
index.index
|
||||
),
|
||||
}
|
||||
.fail()
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_metric_metadata(stream_ctx: &StreamContext) -> Result<()> {
|
||||
let metadata = stream_ctx.input.region_metadata();
|
||||
let valid_prefix = metadata
|
||||
.primary_key
|
||||
.starts_with(&[ReservedColumnId::table_id(), ReservedColumnId::tsid()]);
|
||||
let valid_types = metadata
|
||||
.column_by_id(ReservedColumnId::table_id())
|
||||
.zip(metadata.column_by_id(ReservedColumnId::tsid()))
|
||||
.is_some_and(|(table_id, tsid)| {
|
||||
table_id.column_schema.data_type == ConcreteDataType::uint32_datatype()
|
||||
&& tsid.column_schema.data_type == ConcreteDataType::uint64_datatype()
|
||||
});
|
||||
ensure!(
|
||||
metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse && valid_prefix && valid_types,
|
||||
InvalidRequestSnafu {
|
||||
region_id: metadata.region_id,
|
||||
reason: "candidate-series scan requires sparse (__table_id, __tsid) primary keys",
|
||||
}
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn primary_key_schema() -> SchemaRef {
|
||||
Arc::new(Schema::new(vec![Field::new(
|
||||
PRIMARY_KEY_COLUMN_NAME,
|
||||
DataType::Binary,
|
||||
false,
|
||||
)]))
|
||||
}
|
||||
|
||||
/// Filters a source by encoded-primary-key predicates and emits one binary row per local key.
|
||||
fn candidate_primary_key_stream(
|
||||
mut input: BoxedRecordBatchStream,
|
||||
mut filter: Option<CachedPrimaryKeyFilter>,
|
||||
) -> BoxedRecordBatchStream {
|
||||
Box::pin(try_stream! {
|
||||
let mut last_primary_key = Vec::new();
|
||||
let mut has_last = false;
|
||||
while let Some(batch) = input.try_next().await? {
|
||||
if let Some(batch) = normalize_candidate_batch(
|
||||
batch,
|
||||
filter.as_mut(),
|
||||
&mut last_primary_key,
|
||||
&mut has_last,
|
||||
)? {
|
||||
yield batch;
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn normalize_candidate_batch(
|
||||
mut batch: RecordBatch,
|
||||
filter: Option<&mut CachedPrimaryKeyFilter>,
|
||||
last_primary_key: &mut Vec<u8>,
|
||||
has_last: &mut bool,
|
||||
) -> Result<Option<RecordBatch>> {
|
||||
let pk_idx = batch
|
||||
.schema()
|
||||
.column_with_name(PRIMARY_KEY_COLUMN_NAME)
|
||||
.map(|(idx, _)| idx)
|
||||
.context(UnexpectedSnafu {
|
||||
reason: "candidate source does not contain __primary_key",
|
||||
})?;
|
||||
if let Some(filter) = filter {
|
||||
let Some(filtered) = prefilter_flat_batch_by_primary_key(
|
||||
batch,
|
||||
pk_idx,
|
||||
filter as &mut dyn PrimaryKeyFilter,
|
||||
)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
batch = filtered;
|
||||
}
|
||||
|
||||
let pk_column = batch.column(pk_idx);
|
||||
let mut builder = BinaryBuilder::new();
|
||||
if let Some(array) = pk_column.as_any().downcast_ref::<PrimaryKeyArray>() {
|
||||
let values = array
|
||||
.values()
|
||||
.as_any()
|
||||
.downcast_ref::<BinaryArray>()
|
||||
.context(UnexpectedSnafu {
|
||||
reason: "dictionary primary-key values are not binary",
|
||||
})?;
|
||||
for key in array.keys().values() {
|
||||
append_unique_primary_key(
|
||||
values.value(*key as usize),
|
||||
&mut builder,
|
||||
last_primary_key,
|
||||
has_last,
|
||||
);
|
||||
}
|
||||
} else if let Some(array) = pk_column.as_any().downcast_ref::<BinaryArray>() {
|
||||
for value in array.iter().flatten() {
|
||||
append_unique_primary_key(value, &mut builder, last_primary_key, has_last);
|
||||
}
|
||||
} else {
|
||||
return UnexpectedSnafu {
|
||||
reason: format!(
|
||||
"primary-key column is neither binary nor dictionary, got {:?}",
|
||||
pk_column.data_type()
|
||||
),
|
||||
}
|
||||
.fail();
|
||||
}
|
||||
|
||||
let array = builder.finish();
|
||||
if array.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
let batch = RecordBatch::try_new(primary_key_schema(), vec![Arc::new(array)])
|
||||
.context(NewRecordBatchSnafu)?;
|
||||
Ok(Some(batch))
|
||||
}
|
||||
|
||||
fn append_unique_primary_key(
|
||||
value: &[u8],
|
||||
builder: &mut BinaryBuilder,
|
||||
last_primary_key: &mut Vec<u8>,
|
||||
has_last: &mut bool,
|
||||
) {
|
||||
if !*has_last || last_primary_key != value {
|
||||
builder.append_value(value);
|
||||
last_primary_key.clear();
|
||||
last_primary_key.extend_from_slice(value);
|
||||
*has_last = true;
|
||||
}
|
||||
}
|
||||
|
||||
fn merge_primary_key_streams(
|
||||
sources: Vec<BoxedRecordBatchStream>,
|
||||
memory_pool: Arc<dyn MemoryPool>,
|
||||
metrics_set: &ExecutionPlanMetricsSet,
|
||||
partition: usize,
|
||||
consumer_name: &'static str,
|
||||
) -> Result<BoxedRecordBatchStream> {
|
||||
if sources.is_empty() {
|
||||
return Ok(Box::pin(futures::stream::empty()));
|
||||
}
|
||||
if sources.len() == 1 {
|
||||
return Ok(sources.into_iter().next().unwrap());
|
||||
}
|
||||
|
||||
let schema = primary_key_schema();
|
||||
let df_sources = sources
|
||||
.into_iter()
|
||||
.map(|source| {
|
||||
let stream = source.map_err(|error| DataFusionError::External(Box::new(error)));
|
||||
Box::pin(RecordBatchStreamAdapter::new(schema.clone(), stream)) as _
|
||||
})
|
||||
.collect();
|
||||
let ordering = LexOrdering::new([PhysicalSortExpr {
|
||||
expr: Arc::new(Column::new(PRIMARY_KEY_COLUMN_NAME, 0)),
|
||||
options: SortOptions {
|
||||
descending: false,
|
||||
nulls_first: false,
|
||||
},
|
||||
}])
|
||||
// Safe to unwrap because `LexOrdering::new` returns `None` only for empty
|
||||
// input, and this array always contains one sort expression.
|
||||
.unwrap();
|
||||
let reservation = MemoryConsumer::new(consumer_name).register(&memory_pool);
|
||||
let mut merged = StreamingMergeBuilder::new()
|
||||
.with_streams(df_sources)
|
||||
.with_schema(schema)
|
||||
.with_expressions(&ordering)
|
||||
.with_metrics(BaselineMetrics::new(metrics_set, partition))
|
||||
.with_batch_size(DEFAULT_READ_BATCH_SIZE)
|
||||
.with_reservation(reservation)
|
||||
.build()
|
||||
.context(MergeCandidateSeriesSnafu)?;
|
||||
|
||||
Ok(Box::pin(try_stream! {
|
||||
while let Some(batch) = merged.next().await {
|
||||
yield batch.context(MergeCandidateSeriesSnafu)?;
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
fn decode_metric_series(
|
||||
mut input: BoxedRecordBatchStream,
|
||||
metadata: store_api::metadata::RegionMetadataRef,
|
||||
) -> Result<MetricSeriesIdStream> {
|
||||
let codec = SparsePrimaryKeyCodec::new(&metadata);
|
||||
Ok(Box::pin(try_stream! {
|
||||
let mut last_series = None;
|
||||
let mut output = Vec::with_capacity(CANDIDATE_SERIES_BATCH_SIZE);
|
||||
while let Some(batch) = input.try_next().await? {
|
||||
let array = batch
|
||||
.column(0)
|
||||
.as_any()
|
||||
.downcast_ref::<BinaryArray>()
|
||||
.context(UnexpectedSnafu {
|
||||
reason: "merged candidate primary key is not binary",
|
||||
})?;
|
||||
for primary_key in array.iter().flatten() {
|
||||
let (table_id, tsid) = codec
|
||||
.decode_ids(primary_key)
|
||||
.context(crate::error::DecodeSnafu)?;
|
||||
let series = MetricSeriesId { table_id, tsid };
|
||||
if last_series == Some(series) {
|
||||
continue;
|
||||
}
|
||||
last_series = Some(series);
|
||||
output.push(series);
|
||||
if output.len() == CANDIDATE_SERIES_BATCH_SIZE {
|
||||
yield std::mem::replace(
|
||||
&mut output,
|
||||
Vec::with_capacity(CANDIDATE_SERIES_BATCH_SIZE),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
if !output.is_empty() {
|
||||
yield output;
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use datafusion::execution::memory_pool::UnboundedMemoryPool;
|
||||
use datafusion_expr::{col, lit};
|
||||
use datatypes::arrow::array::{ArrayRef, DictionaryArray, UInt32Array};
|
||||
use datatypes::arrow::datatypes::UInt32Type;
|
||||
use futures::TryStreamExt;
|
||||
use store_api::codec::PrimaryKeyEncoding;
|
||||
use table::predicate::Predicate;
|
||||
|
||||
use super::*;
|
||||
use crate::test_util::sst_util::sst_region_metadata_with_encoding;
|
||||
|
||||
fn binary_batch(values: &[&[u8]]) -> RecordBatch {
|
||||
RecordBatch::try_new(
|
||||
primary_key_schema(),
|
||||
vec![Arc::new(BinaryArray::from_iter_values(
|
||||
values.iter().copied(),
|
||||
))],
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn dictionary_batch(values: &[&[u8]], keys: &[u32]) -> RecordBatch {
|
||||
let dict_values: ArrayRef = Arc::new(BinaryArray::from_iter_values(values.iter().copied()));
|
||||
let dict =
|
||||
DictionaryArray::<UInt32Type>::try_new(UInt32Array::from(keys.to_vec()), dict_values)
|
||||
.unwrap();
|
||||
let schema = Arc::new(Schema::new(vec![Field::new_dictionary(
|
||||
PRIMARY_KEY_COLUMN_NAME,
|
||||
DataType::UInt32,
|
||||
DataType::Binary,
|
||||
false,
|
||||
)]));
|
||||
RecordBatch::try_new(schema, vec![Arc::new(dict)]).unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn candidate_stream_normalizes_and_deduplicates_primary_keys() {
|
||||
let input = Box::pin(futures::stream::iter(vec![
|
||||
Ok(dictionary_batch(&[b"a", b"b"], &[0, 0, 1])),
|
||||
Ok(binary_batch(&[b"b", b"c", b"c"])),
|
||||
]));
|
||||
let batches = candidate_primary_key_stream(input, None)
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
let actual = batches
|
||||
.iter()
|
||||
.flat_map(|batch| {
|
||||
batch
|
||||
.column(0)
|
||||
.as_any()
|
||||
.downcast_ref::<BinaryArray>()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.flatten()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(actual, vec![b"a".as_slice(), b"b", b"c"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn candidate_stream_filters_primary_keys_before_merge() {
|
||||
let metadata = Arc::new(sst_region_metadata_with_encoding(
|
||||
PrimaryKeyEncoding::Sparse,
|
||||
));
|
||||
let codec = SparsePrimaryKeyCodec::new(&metadata);
|
||||
let mut table_1 = Vec::new();
|
||||
let mut table_2 = Vec::new();
|
||||
codec.encode_internal(1, 10, &mut table_1).unwrap();
|
||||
codec.encode_internal(2, 20, &mut table_2).unwrap();
|
||||
|
||||
let predicate = Predicate::new(vec![
|
||||
col(store_api::metric_engine_consts::DATA_SCHEMA_TABLE_ID_COLUMN_NAME).eq(lit(1_u32)),
|
||||
]);
|
||||
let filter = build_primary_key_filter(&metadata, Some(&predicate));
|
||||
let input = Box::pin(futures::stream::iter(vec![Ok(dictionary_batch(
|
||||
&[table_1.as_slice(), table_2.as_slice()],
|
||||
&[0, 1],
|
||||
))]));
|
||||
let batches = candidate_primary_key_stream(input, filter)
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(batches.len(), 1);
|
||||
let array = batches[0]
|
||||
.column(0)
|
||||
.as_any()
|
||||
.downcast_ref::<BinaryArray>()
|
||||
.unwrap();
|
||||
assert_eq!(array.len(), 1);
|
||||
assert_eq!(array.value(0), table_1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn decode_metric_series_yields_groups_of_500() {
|
||||
let metadata = Arc::new(sst_region_metadata_with_encoding(
|
||||
PrimaryKeyEncoding::Sparse,
|
||||
));
|
||||
let codec = SparsePrimaryKeyCodec::new(&metadata);
|
||||
let primary_keys = (0..501_u64)
|
||||
.map(|tsid| {
|
||||
let mut primary_key = Vec::new();
|
||||
codec.encode_internal(1, tsid, &mut primary_key).unwrap();
|
||||
primary_key
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let batch = binary_batch(&primary_keys.iter().map(Vec::as_slice).collect::<Vec<_>>());
|
||||
let source = Box::pin(futures::stream::iter(vec![Ok(batch)]));
|
||||
let metrics = ExecutionPlanMetricsSet::new();
|
||||
let pool = Arc::new(UnboundedMemoryPool::default());
|
||||
let merged =
|
||||
merge_primary_key_streams(vec![source], pool, &metrics, 0, "candidate-test").unwrap();
|
||||
let groups = decode_metric_series(merged, metadata)
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(groups.iter().map(Vec::len).collect::<Vec<_>>(), [500, 1]);
|
||||
assert_eq!(
|
||||
groups[0][0],
|
||||
MetricSeriesId {
|
||||
table_id: 1,
|
||||
tsid: 0
|
||||
}
|
||||
);
|
||||
assert_eq!(
|
||||
groups[1][0],
|
||||
MetricSeriesId {
|
||||
table_id: 1,
|
||||
tsid: 500
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn merge_primary_keys_globally_sorts_and_deduplicates_series() {
|
||||
let metadata = Arc::new(sst_region_metadata_with_encoding(
|
||||
PrimaryKeyEncoding::Sparse,
|
||||
));
|
||||
let codec = SparsePrimaryKeyCodec::new(&metadata);
|
||||
let encode = |tsid| {
|
||||
let mut primary_key = Vec::new();
|
||||
codec.encode_internal(1, tsid, &mut primary_key).unwrap();
|
||||
primary_key
|
||||
};
|
||||
let keys_1 = [encode(1), encode(3)];
|
||||
let mut alternate_key_for_series_1 = encode(1);
|
||||
alternate_key_for_series_1.push(0);
|
||||
let keys_2 = [alternate_key_for_series_1, encode(2)];
|
||||
let sources = vec![
|
||||
Box::pin(futures::stream::iter(vec![Ok(binary_batch(
|
||||
&keys_1.iter().map(Vec::as_slice).collect::<Vec<_>>(),
|
||||
))])) as BoxedRecordBatchStream,
|
||||
Box::pin(futures::stream::iter(vec![Ok(binary_batch(
|
||||
&keys_2.iter().map(Vec::as_slice).collect::<Vec<_>>(),
|
||||
))])),
|
||||
];
|
||||
|
||||
let merged = merge_primary_key_streams(
|
||||
sources,
|
||||
Arc::new(UnboundedMemoryPool::default()),
|
||||
&ExecutionPlanMetricsSet::new(),
|
||||
0,
|
||||
"candidate-merge-test",
|
||||
)
|
||||
.unwrap();
|
||||
let groups = decode_metric_series(merged, metadata)
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
groups,
|
||||
vec![vec![
|
||||
MetricSeriesId {
|
||||
table_id: 1,
|
||||
tsid: 1,
|
||||
},
|
||||
MetricSeriesId {
|
||||
table_id: 1,
|
||||
tsid: 2,
|
||||
},
|
||||
MetricSeriesId {
|
||||
table_id: 1,
|
||||
tsid: 3,
|
||||
},
|
||||
]]
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -27,6 +27,7 @@ use datatypes::arrow::array::{Array as _, ArrayRef, BooleanArray};
|
||||
use datatypes::arrow::buffer::BooleanBuffer;
|
||||
use datatypes::arrow::record_batch::RecordBatch;
|
||||
use datatypes::schema::Schema;
|
||||
use futures::StreamExt;
|
||||
use mito_codec::row_converter::PrimaryKeyCodec;
|
||||
use parquet::arrow::arrow_reader::RowSelection;
|
||||
use parquet::file::metadata::ParquetMetaData;
|
||||
@@ -49,6 +50,8 @@ use crate::sst::file::FileHandle;
|
||||
use crate::sst::parquet::flat_format::{
|
||||
DecodedPrimaryKeys, FlatReadFormat, decode_primary_keys, time_index_column_index,
|
||||
};
|
||||
use crate::sst::parquet::json_align::ProjectedRecordBatchStream;
|
||||
use crate::sst::parquet::prefilter::CachedPrimaryKeyFilter;
|
||||
use crate::sst::parquet::reader::{
|
||||
FlatRowGroupReader, MaybeFilter, RowGroupBuildContext, RowGroupReaderBuilder,
|
||||
SimpleFilterContext,
|
||||
@@ -95,6 +98,11 @@ pub struct FileRange {
|
||||
}
|
||||
|
||||
impl FileRange {
|
||||
/// Builds the encoded-primary-key filter selected for this file.
|
||||
pub(crate) fn primary_key_filter(&self) -> Option<CachedPrimaryKeyFilter> {
|
||||
self.context.reader_builder.primary_key_filter()
|
||||
}
|
||||
|
||||
/// Creates a new [FileRange].
|
||||
pub(crate) fn new(
|
||||
context: FileRangeContextRef,
|
||||
@@ -224,6 +232,44 @@ impl FileRange {
|
||||
Ok(Some(flat_prune_reader))
|
||||
}
|
||||
|
||||
/// Creates a reader that returns only the encoded primary-key column.
|
||||
///
|
||||
/// The returned primary keys are compatible with the expected region metadata.
|
||||
pub(crate) async fn primary_key_reader(
|
||||
&self,
|
||||
fetch_metrics: Option<&ParquetFetchMetrics>,
|
||||
) -> Result<Option<ProjectedRecordBatchStream>> {
|
||||
if !self.in_dynamic_filter_range() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let stream = self
|
||||
.context
|
||||
.reader_builder
|
||||
.build_primary_key(self.context.build_context(
|
||||
self.row_group_idx,
|
||||
self.row_selection.clone(),
|
||||
fetch_metrics,
|
||||
))
|
||||
.await?;
|
||||
if self.context.compat_batch().is_none() {
|
||||
return Ok(Some(stream));
|
||||
}
|
||||
|
||||
let context = self.context.clone();
|
||||
let stream = stream
|
||||
.map(move |batch| {
|
||||
let batch = batch?;
|
||||
let compat = context.compat_batch().context(UnexpectedSnafu {
|
||||
reason: "Primary-key compatibility helper is missing",
|
||||
})?;
|
||||
let primary_key = compat.compat_primary_key(batch.column(0))?;
|
||||
RecordBatch::try_new(batch.schema(), vec![primary_key]).context(NewRecordBatchSnafu)
|
||||
})
|
||||
.boxed();
|
||||
Ok(Some(stream))
|
||||
}
|
||||
|
||||
/// Returns the helper to compat batches.
|
||||
pub(crate) fn compat_batch(&self) -> Option<&FlatCompatBatch> {
|
||||
self.context.compat_batch()
|
||||
|
||||
@@ -29,7 +29,7 @@ use datatypes::arrow::buffer::BooleanBuffer;
|
||||
use datatypes::arrow::datatypes::SchemaRef;
|
||||
use datatypes::arrow::record_batch::RecordBatch;
|
||||
use futures::StreamExt;
|
||||
use mito_codec::row_converter::{PrimaryKeyCodec, PrimaryKeyFilter};
|
||||
use mito_codec::row_converter::{PrimaryKeyCodec, PrimaryKeyFilter, build_primary_key_codec};
|
||||
use parquet::arrow::ProjectionMask;
|
||||
use parquet::arrow::arrow_reader::RowSelection;
|
||||
use parquet::schema::types::SchemaDescriptor;
|
||||
@@ -262,6 +262,34 @@ pub(crate) struct BulkFilterPlan {
|
||||
pub(crate) pk_filters: Option<Arc<Vec<SimpleFilterEvaluator>>>,
|
||||
}
|
||||
|
||||
/// Builds an encoded-primary-key filter from the supported tag predicates.
|
||||
///
|
||||
/// Predicates on fields, timestamps, or unsupported expression shapes are intentionally
|
||||
/// omitted. Callers use this as a pruning filter and must preserve the full predicate for
|
||||
/// authoritative filtering later in the scan.
|
||||
pub(crate) fn build_primary_key_filter(
|
||||
metadata: &RegionMetadataRef,
|
||||
predicate: Option<&Predicate>,
|
||||
) -> Option<CachedPrimaryKeyFilter> {
|
||||
let filters = predicate
|
||||
.into_iter()
|
||||
.flat_map(|predicate| predicate.exprs())
|
||||
.filter_map(|expr| SimpleFilterContext::new_opt(metadata, None, expr))
|
||||
.filter_map(|filter_ctx| {
|
||||
(filter_ctx.semantic_type() == SemanticType::Tag)
|
||||
.then(|| filter_ctx.filter().as_filter().cloned())
|
||||
.flatten()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
if filters.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let codec = build_primary_key_codec(metadata.as_ref());
|
||||
let filter = codec.primary_key_filter(metadata, Arc::new(filters));
|
||||
Some(CachedPrimaryKeyFilter::new(filter))
|
||||
}
|
||||
|
||||
/// How the parquet reader should apply each predicate.
|
||||
///
|
||||
/// The reader runs in two phases. Predicates routed into `prefilter_builder`
|
||||
@@ -584,12 +612,9 @@ impl PrefilterContextBuilder {
|
||||
|
||||
/// Builds a [PrefilterContext] for a specific row group.
|
||||
pub(crate) fn build(&self) -> PrefilterContext {
|
||||
let pk_filter = self.pk_filters.as_ref().map(|pk_filters| {
|
||||
let pk_filter = self
|
||||
.codec
|
||||
.primary_key_filter(&self.metadata, Arc::clone(pk_filters));
|
||||
Box::new(CachedPrimaryKeyFilter::new(pk_filter)) as Box<dyn PrimaryKeyFilter>
|
||||
});
|
||||
let pk_filter = self
|
||||
.build_primary_key_filter()
|
||||
.map(|filter| Box::new(filter) as Box<dyn PrimaryKeyFilter>);
|
||||
PrefilterContext {
|
||||
pk_filter,
|
||||
filters: self.filters.clone(),
|
||||
@@ -599,6 +624,16 @@ impl PrefilterContextBuilder {
|
||||
arrow_schema: self.arrow_schema.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Builds a fresh encoded-primary-key filter selected by this plan.
|
||||
pub(crate) fn build_primary_key_filter(&self) -> Option<CachedPrimaryKeyFilter> {
|
||||
self.pk_filters.as_ref().map(|pk_filters| {
|
||||
let filter = self
|
||||
.codec
|
||||
.primary_key_filter(&self.metadata, Arc::clone(pk_filters));
|
||||
CachedPrimaryKeyFilter::new(filter)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
const PREFILTER_COLUMN_RATIO_THRESHOLD: f64 = 0.5;
|
||||
@@ -1530,24 +1565,36 @@ mod tests {
|
||||
vec!["field_0"]
|
||||
);
|
||||
|
||||
let field_0 = metadata.column_by_name("field_0").unwrap().column_id;
|
||||
let ts = metadata.time_index_column().column_id;
|
||||
let metric_metadata: RegionMetadataRef = Arc::new(sst_region_metadata_with_encoding(
|
||||
PrimaryKeyEncoding::Sparse,
|
||||
));
|
||||
let field_0 = metric_metadata.column_by_name("field_0").unwrap().column_id;
|
||||
let ts = metric_metadata.time_index_column().column_id;
|
||||
let projected_read_format = FlatReadFormat::new(
|
||||
metadata.clone(),
|
||||
metric_metadata.clone(),
|
||||
ReadColumns::from_deduped_column_ids([field_0, ts]),
|
||||
None,
|
||||
"test",
|
||||
true,
|
||||
)
|
||||
.unwrap();
|
||||
let metric_codec = build_primary_key_codec(metric_metadata.as_ref());
|
||||
let pk_prefilter_plan = build_reader_filter_plan(
|
||||
Some(&Predicate::new(vec![col("tag_0").eq(lit("a"))])),
|
||||
None,
|
||||
PreFilterMode::All,
|
||||
&projected_read_format,
|
||||
&codec,
|
||||
&metric_codec,
|
||||
);
|
||||
assert!(pk_prefilter_plan.prefilter_builder.is_some());
|
||||
assert!(
|
||||
pk_prefilter_plan
|
||||
.prefilter_builder
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.build_primary_key_filter()
|
||||
.is_some()
|
||||
);
|
||||
assert!(pk_prefilter_plan.remaining_simple_filters.is_empty());
|
||||
}
|
||||
|
||||
|
||||
@@ -41,10 +41,11 @@ use parquet::arrow::{ProjectionMask, parquet_to_arrow_schema};
|
||||
use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData};
|
||||
use parquet::file::properties::DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT;
|
||||
use partition::expr::PartitionExpr;
|
||||
use snafu::ResultExt;
|
||||
use snafu::{OptionExt, ResultExt};
|
||||
use store_api::codec::PrimaryKeyEncoding;
|
||||
use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataRef};
|
||||
use store_api::region_request::PathType;
|
||||
use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME;
|
||||
use store_api::storage::{ColumnId, FileId};
|
||||
use table::predicate::Predicate;
|
||||
|
||||
@@ -54,6 +55,7 @@ use crate::cache::{CacheStrategy, CachedSstMeta, SstMetaPreparation, prepare_sst
|
||||
use crate::error::ApplyVectorIndexSnafu;
|
||||
use crate::error::{
|
||||
ParquetToArrowSchemaSnafu, ReadDataPartSnafu, Result, SerializePartitionExprSnafu,
|
||||
UnexpectedSnafu,
|
||||
};
|
||||
use crate::metrics::{
|
||||
PRECISE_FILTER_ROWS_TOTAL, READ_ROW_GROUPS_TOTAL, READ_ROWS_IN_ROW_GROUP_TOTAL,
|
||||
@@ -83,7 +85,7 @@ use crate::sst::parquet::format::{INTERNAL_COLUMN_NUM, need_override_sequence};
|
||||
use crate::sst::parquet::json_align::{NestedSchemaAligner, ProjectedRecordBatchStream};
|
||||
use crate::sst::parquet::metadata::MetadataLoader;
|
||||
use crate::sst::parquet::prefilter::{
|
||||
PrefilterContextBuilder, build_reader_filter_plan, execute_prefilter,
|
||||
CachedPrimaryKeyFilter, PrefilterContextBuilder, build_reader_filter_plan, execute_prefilter,
|
||||
};
|
||||
use crate::sst::parquet::push_decoder::{
|
||||
SstParquetRangeFetcher, build_sst_parquet_record_batch_stream,
|
||||
@@ -1836,6 +1838,13 @@ impl RowGroupReaderBuilder {
|
||||
self.prefilter_builder.is_some()
|
||||
}
|
||||
|
||||
/// Builds the encoded-primary-key filter selected by the reader filter plan.
|
||||
pub(crate) fn primary_key_filter(&self) -> Option<CachedPrimaryKeyFilter> {
|
||||
self.prefilter_builder
|
||||
.as_ref()
|
||||
.and_then(PrefilterContextBuilder::build_primary_key_filter)
|
||||
}
|
||||
|
||||
/// Builds a parquet record batch stream to read the row group at `row_group_idx`.
|
||||
///
|
||||
/// If prefiltering is applicable (based on `build_ctx`), this performs a two-phase read:
|
||||
@@ -1894,6 +1903,34 @@ impl RowGroupReaderBuilder {
|
||||
self.make_projected_stream(stream)
|
||||
}
|
||||
|
||||
/// Builds a stream that reads only the encoded primary-key column.
|
||||
///
|
||||
/// It preserves the normal reader's binary-or-dictionary decision. This path deliberately
|
||||
/// skips the normal prefilter pass: the caller reads `__primary_key` once and applies all
|
||||
/// encoded-primary-key filters to the returned batches.
|
||||
pub(crate) async fn build_primary_key(
|
||||
&self,
|
||||
build_ctx: RowGroupBuildContext<'_>,
|
||||
) -> Result<ProjectedRecordBatchStream> {
|
||||
let parquet_schema = self.parquet_meta.file_metadata().schema_descr();
|
||||
let primary_key_index = parquet_schema
|
||||
.columns()
|
||||
.iter()
|
||||
.position(|column| column.name() == PRIMARY_KEY_COLUMN_NAME)
|
||||
.context(UnexpectedSnafu {
|
||||
reason: "SST does not contain __primary_key",
|
||||
})?;
|
||||
let projection = ProjectionMask::leaves(parquet_schema, [primary_key_index]);
|
||||
|
||||
self.build_with_projection(
|
||||
build_ctx.row_group_idx,
|
||||
build_ctx.row_selection,
|
||||
projection,
|
||||
build_ctx.fetch_metrics,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
fn make_projected_stream(
|
||||
&self,
|
||||
stream: ProjectedRecordBatchStream,
|
||||
|
||||
Reference in New Issue
Block a user