1mod between;
16mod comparison;
17mod eq_list;
18mod in_list;
19mod regex_match;
20
21use std::collections::{BTreeMap, HashSet};
22
23#[cfg(test)]
24use common_function::scalars::json::json_get::JsonGetWithType;
25use common_telemetry::warn;
26use datafusion_common::ScalarValue;
27use datafusion_expr::{BinaryExpr, Expr, Operator};
28use datatypes::data_type::ConcreteDataType;
29#[cfg(test)]
30use datatypes::extension::json::Json2ExtensionType;
31use datatypes::value::Value;
32use index::inverted_index::search::predicate::Predicate;
33use index::target::IndexTarget;
34use mito_codec::index::IndexValueCodec;
35use mito_codec::row_converter::SortField;
36use object_store::ObjectStore;
37use puffin::puffin_manager::cache::PuffinMetadataCacheRef;
38use snafu::{OptionExt, ResultExt};
39use store_api::metadata::RegionMetadata;
40use store_api::region_request::PathType;
41use store_api::storage::ColumnId;
42
43use crate::cache::file_cache::FileCacheRef;
44use crate::cache::index::inverted_index::InvertedIndexCacheRef;
45use crate::error::{ColumnNotFoundSnafu, ConvertValueSnafu, EncodeSnafu, Result};
46use crate::sst::index::inverted_index::applier::{InvertedIndexApplier, index_target_data_type};
47use crate::sst::index::puffin_manager::PuffinManagerFactory;
48
49pub(crate) struct InvertedIndexApplierBuilder<'a> {
51 table_dir: String,
53
54 path_type: PathType,
56
57 object_store: ObjectStore,
59
60 file_cache: Option<FileCacheRef>,
62
63 metadata: &'a RegionMetadata,
65
66 indexed_column_ids: HashSet<ColumnId>,
68
69 output: BTreeMap<IndexTarget, Vec<Predicate>>,
71
72 puffin_manager_factory: PuffinManagerFactory,
74
75 inverted_index_cache: Option<InvertedIndexCacheRef>,
77
78 puffin_metadata_cache: Option<PuffinMetadataCacheRef>,
80}
81
82impl<'a> InvertedIndexApplierBuilder<'a> {
83 pub fn new(
85 table_dir: String,
86 path_type: PathType,
87 object_store: ObjectStore,
88 metadata: &'a RegionMetadata,
89 indexed_column_ids: HashSet<ColumnId>,
90 puffin_manager_factory: PuffinManagerFactory,
91 ) -> Self {
92 Self {
93 table_dir,
94 path_type,
95 object_store,
96 metadata,
97 indexed_column_ids,
98 output: BTreeMap::default(),
99 puffin_manager_factory,
100 file_cache: None,
101 inverted_index_cache: None,
102 puffin_metadata_cache: None,
103 }
104 }
105
106 pub fn with_file_cache(mut self, file_cache: Option<FileCacheRef>) -> Self {
108 self.file_cache = file_cache;
109 self
110 }
111
112 pub fn with_puffin_metadata_cache(
114 mut self,
115 puffin_metadata_cache: Option<PuffinMetadataCacheRef>,
116 ) -> Self {
117 self.puffin_metadata_cache = puffin_metadata_cache;
118 self
119 }
120
121 pub fn with_inverted_index_cache(
123 mut self,
124 inverted_index_cache: Option<InvertedIndexCacheRef>,
125 ) -> Self {
126 self.inverted_index_cache = inverted_index_cache;
127 self
128 }
129
130 pub fn build(mut self, exprs: &[Expr]) -> Result<Option<InvertedIndexApplier>> {
133 for expr in exprs {
134 self.traverse_and_collect(expr);
135 }
136
137 if self.output.is_empty() {
138 return Ok(None);
139 }
140
141 let expected_predicate_target_types = self.expected_predicate_target_types();
142
143 Ok(Some(
144 InvertedIndexApplier::new(
145 self.table_dir,
146 self.path_type,
147 self.object_store,
148 self.puffin_manager_factory,
149 self.output,
150 expected_predicate_target_types,
151 )?
152 .with_file_cache(self.file_cache)
153 .with_puffin_metadata_cache(self.puffin_metadata_cache)
154 .with_index_cache(self.inverted_index_cache),
155 ))
156 }
157
158 fn expected_predicate_target_types(&self) -> BTreeMap<IndexTarget, ConcreteDataType> {
164 self.output
165 .keys()
166 .filter_map(|target| {
167 let data_type = index_target_data_type(self.metadata, target)?;
168 Some((target.clone(), data_type))
169 })
170 .collect()
171 }
172
173 fn traverse_and_collect(&mut self, expr: &Expr) {
176 let res = match expr {
177 Expr::Between(between) => self.collect_between(between),
178
179 Expr::InList(in_list) => self.collect_inlist(in_list),
180 Expr::BinaryExpr(BinaryExpr { left, op, right }) => match op {
181 Operator::And => {
182 self.traverse_and_collect(left);
183 self.traverse_and_collect(right);
184 Ok(())
185 }
186 Operator::Or => self.collect_or_eq_list(left, right),
187 Operator::Eq => self.collect_eq(left, right),
188 Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq => {
189 self.collect_comparison_expr(left, op, right)
190 }
191 Operator::RegexMatch => self.collect_regex_match(left, right),
192 _ => Ok(()),
193 },
194
195 _ => Ok(()),
197 };
198
199 if let Err(err) = res {
200 warn!(err; "Failed to collect predicates, ignore it. expr: {expr}");
201 }
202 }
203
204 fn add_predicate(&mut self, target: IndexTarget, predicate: Predicate) {
206 self.output.entry(target).or_default().push(predicate);
207 }
208
209 fn column_id_and_type(
212 &self,
213 column_name: &str,
214 ) -> Result<Option<(ColumnId, ConcreteDataType)>> {
215 let column = self
216 .metadata
217 .column_by_name(column_name)
218 .context(ColumnNotFoundSnafu {
219 column: column_name,
220 })?;
221
222 if !self.indexed_column_ids.contains(&column.column_id) {
223 return Ok(None);
224 }
225
226 Ok(Some((
227 column.column_id,
228 column.column_schema.data_type.clone(),
229 )))
230 }
231
232 fn expr_to_index_target(&self, expr: &Expr) -> Result<Option<(IndexTarget, ConcreteDataType)>> {
240 if let Some(name) = Self::column_name(expr) {
241 return Ok(self
242 .column_id_and_type(name)?
243 .map(|(id, ty)| (IndexTarget::ColumnId(id), ty)));
244 }
245
246 crate::sst::index::column::json_index_target(self.metadata, expr, |hint| {
247 hint.inverted_index
248 })
249 }
250
251 fn nonnull_lit(expr: &Expr) -> Option<&ScalarValue> {
253 match expr {
254 Expr::Literal(lit, _) if !lit.is_null() => Some(lit),
255 _ => None,
256 }
257 }
258
259 fn column_name(expr: &Expr) -> Option<&str> {
261 match expr {
262 Expr::Column(column) => Some(&column.name),
263 _ => None,
264 }
265 }
266
267 fn encode_lit(lit: &ScalarValue, data_type: ConcreteDataType) -> Result<Vec<u8>> {
269 let value = Value::try_from(lit.clone()).context(ConvertValueSnafu)?;
270 let mut bytes = vec![];
271 let field = SortField::new(data_type);
272 IndexValueCodec::encode_nonnull_value(value.as_value_ref(), &field, &mut bytes)
273 .context(EncodeSnafu)?;
274 Ok(bytes)
275 }
276}
277
278#[cfg(test)]
279mod tests {
280 use std::sync::Arc;
281
282 use api::v1::SemanticType;
283 use datafusion_common::Column;
284 use datafusion_expr::{Between, Literal};
285 use datatypes::data_type::ConcreteDataType;
286 use datatypes::schema::ColumnSchema;
287 use index::inverted_index::search::predicate::{
288 Bound, Range, RangePredicate, RegexMatchPredicate,
289 };
290 use object_store::ObjectStore;
291 use object_store::services::Memory;
292 use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder};
293 use store_api::storage::RegionId;
294
295 use super::*;
296
297 #[test]
298 fn test_json_target_resolution() {
299 use datafusion_expr::col;
300 let (_dir, factory) = PuffinManagerFactory::new_for_test_block("json_target_resolution");
301 let metadata = json_metadata();
302 let builder = InvertedIndexApplierBuilder::new(
303 "test".into(),
304 PathType::Bare,
305 test_object_store(),
306 &metadata,
307 HashSet::new(),
308 factory,
309 );
310 for (path, expected) in [
311 ("$.a", vec!["a"]),
312 ("$.\"a.b\"", vec!["a.b"]),
313 ("$.nested.value", vec!["nested", "value"]),
314 ] {
315 let (target, _) = builder
316 .expr_to_index_target(&json_expr(path, None))
317 .unwrap()
318 .unwrap();
319 assert_eq!(
320 target,
321 IndexTarget::new_json_path(1, expected.into_iter().map(String::from).collect())
322 .unwrap()
323 );
324 }
325 for path in ["$.missing", "$.unindexed", "$.a[0]", "$.*", "$"] {
326 assert!(
327 builder
328 .expr_to_index_target(&json_expr(path, None))
329 .unwrap()
330 .is_none(),
331 "{path}"
332 );
333 }
334 let data_type = ConcreteDataType::int64_datatype();
335 let expected = Some((
336 IndexTarget::new_json_path(1, vec!["n".into()]).unwrap(),
337 data_type,
338 ));
339 for ty in [None, Some(ScalarValue::Int64(None))] {
340 let resolved = builder.expr_to_index_target(&json_expr("$.n", ty)).unwrap();
341 assert_eq!(resolved, expected);
342 }
343 assert!(
344 builder
345 .expr_to_index_target(&json_expr("$.n", Some(ScalarValue::Float64(None))))
346 .unwrap()
347 .is_none()
348 );
349 assert!(
350 builder
351 .expr_to_index_target(&col("attrs"))
352 .unwrap()
353 .is_none()
354 );
355 }
356
357 #[test]
358 fn test_json_plan_checks_sst_type_hints() {
359 use datafusion_expr::{col, lit};
360 use datatypes::extension::json::JsonMetadata;
361 use datatypes::json::{JsonSettings, JsonTypeHint};
362
363 let (_dir, factory) = PuffinManagerFactory::new_for_test_block("json_plan_type_hints");
364 let metadata = json_metadata();
365 for include_column in [false, true] {
366 let builder = InvertedIndexApplierBuilder::new(
367 "test".into(),
368 PathType::Bare,
369 test_object_store(),
370 &metadata,
371 HashSet::from([2]),
372 factory.clone(),
373 );
374 let mut exprs = vec![json_expr("$.n", None).eq(lit(42_i64))];
375 if include_column {
376 exprs.push(col("ts").gt(lit(ScalarValue::TimestampMillisecond(Some(0), None))));
377 }
378 let applier = builder.build(&exprs).unwrap().unwrap();
379 let plan = applier
380 .plan_for_sst(&Arc::new(metadata.clone()))
381 .unwrap()
382 .unwrap();
383 assert_eq!(plan.predicate_key, applier.default_plan.predicate_key);
384
385 for hints in [
388 vec![JsonTypeHint {
389 path: vec!["n".into()],
390 data_type: ConcreteDataType::float64_datatype(),
391 inverted_index: true,
392 skipping_index: None,
393 fulltext_index: None,
394 }],
395 vec![],
396 vec![JsonTypeHint {
397 path: vec!["n".into()],
398 data_type: ConcreteDataType::int64_datatype(),
399 inverted_index: false,
400 skipping_index: None,
401 fulltext_index: None,
402 }],
403 ] {
404 let mut sst = metadata.clone();
405 sst.column_metadatas[0].column_schema.with_extension_type(
406 &Json2ExtensionType::new(Arc::new(JsonMetadata::new(
407 JsonSettings::try_new(hints, Some(0)).unwrap(),
408 ))),
409 );
410 let plan = applier.plan_for_sst(&Arc::new(sst)).unwrap();
411 if !include_column {
412 assert!(plan.is_none());
413 continue;
414 }
415 let expected = BTreeMap::from([(
416 IndexTarget::ColumnId(2),
417 applier.predicates[&IndexTarget::ColumnId(2)].clone(),
418 )]);
419 assert_eq!(
420 plan.unwrap().predicate_key,
421 InvertedIndexApplier::build_apply_plan(&expected)
422 .unwrap()
423 .predicate_key
424 );
425 }
426 }
427 }
428
429 #[test]
430 fn test_json_predicate_collection() {
431 use datafusion_expr::lit;
432
433 let (_dir, factory) = PuffinManagerFactory::new_for_test_block("json_predicate_collection");
434 let metadata = json_metadata();
435 let mut builder = InvertedIndexApplierBuilder::new(
436 "test".into(),
437 PathType::Bare,
438 test_object_store(),
439 &metadata,
440 HashSet::new(),
441 factory,
442 );
443 let numeric_target = IndexTarget::new_json_path(1, vec!["n".into()]).unwrap();
444 let n = json_expr("$.n", Some(ScalarValue::Int64(None)));
445 for expr in [
446 n.clone().eq(lit(42_i64)),
447 lit(40_i64).lt(n.clone()),
448 n.clone().between(lit(1_i64), lit(99_i64)),
449 n.clone().in_list(vec![lit(1_i64), lit(2_i64)], false),
450 n.clone().eq(lit(1_i64)).or(n.eq(lit(2_i64))),
451 ] {
452 builder.output.clear();
453 builder.traverse_and_collect(&expr);
454 assert_eq!(builder.output.len(), 1, "{expr}");
455 assert_eq!(
456 builder.output.get(&numeric_target).unwrap().len(),
457 1,
458 "{expr}"
459 );
460 }
461 builder.output.clear();
462 let expr = json_expr("$.a", None)
463 .eq(lit("x"))
464 .or(json_expr("$.b", None).eq(lit("y")));
465 builder.traverse_and_collect(&expr);
466 assert!(builder.output.is_empty());
467 let string_target = IndexTarget::new_json_path(1, vec!["a".into()]).unwrap();
468 for pattern in [
469 ScalarValue::Utf8(Some("a.*[ip]".into())),
470 ScalarValue::Utf8View(Some("a.*[ip]".into())),
471 ScalarValue::LargeUtf8(Some("a.*[ip]".into())),
472 ] {
473 builder.output.clear();
474 builder
475 .collect_regex_match(&json_expr("$.a", None), &Expr::Literal(pattern, None))
476 .unwrap();
477 assert_eq!(builder.output.len(), 1);
478 assert_eq!(builder.output.get(&string_target).unwrap().len(), 1);
479 }
480 }
481
482 #[test]
483 fn test_collect_and_basic() {
484 let (_d, factory) = PuffinManagerFactory::new_for_test_block("test_collect_and_basic_");
485
486 let metadata = test_region_metadata();
487 let mut builder = InvertedIndexApplierBuilder::new(
488 "test".to_string(),
489 PathType::Bare,
490 test_object_store(),
491 &metadata,
492 HashSet::from_iter([1, 2, 3]),
493 factory,
494 );
495
496 let expr = Expr::BinaryExpr(BinaryExpr {
497 left: Box::new(Expr::BinaryExpr(BinaryExpr {
498 left: Box::new(tag_column()),
499 op: Operator::RegexMatch,
500 right: Box::new(string_lit("bar")),
501 })),
502 op: Operator::And,
503 right: Box::new(Expr::Between(Between {
504 expr: Box::new(tag_column2()),
505 negated: false,
506 low: Box::new(int64_lit(123)),
507 high: Box::new(int64_lit(456)),
508 })),
509 });
510
511 builder.traverse_and_collect(&expr);
512 let predicates = builder
513 .output
514 .get(&index::target::IndexTarget::ColumnId(1))
515 .unwrap();
516 assert_eq!(predicates.len(), 1);
517 assert_eq!(
518 predicates[0],
519 Predicate::RegexMatch(RegexMatchPredicate {
520 pattern: "bar".to_string()
521 })
522 );
523 let predicates = builder
524 .output
525 .get(&index::target::IndexTarget::ColumnId(2))
526 .unwrap();
527 assert_eq!(predicates.len(), 1);
528 assert_eq!(
529 predicates[0],
530 Predicate::Range(RangePredicate {
531 range: Range {
532 lower: Some(Bound {
533 inclusive: true,
534 value: encoded_int64(123),
535 }),
536 upper: Some(Bound {
537 inclusive: true,
538 value: encoded_int64(456),
539 }),
540 }
541 })
542 );
543 }
544
545 pub(crate) fn test_region_metadata() -> RegionMetadata {
546 let mut builder = RegionMetadataBuilder::new(RegionId::new(1234, 5678));
547 builder
548 .push_column_metadata(ColumnMetadata {
549 column_schema: ColumnSchema::new("a", ConcreteDataType::string_datatype(), false),
550 semantic_type: SemanticType::Tag,
551 column_id: 1,
552 })
553 .push_column_metadata(ColumnMetadata {
554 column_schema: ColumnSchema::new("b", ConcreteDataType::int64_datatype(), false),
555 semantic_type: SemanticType::Tag,
556 column_id: 2,
557 })
558 .push_column_metadata(ColumnMetadata {
559 column_schema: ColumnSchema::new("c", ConcreteDataType::string_datatype(), false),
560 semantic_type: SemanticType::Field,
561 column_id: 3,
562 })
563 .push_column_metadata(ColumnMetadata {
564 column_schema: ColumnSchema::new(
565 "d",
566 ConcreteDataType::timestamp_millisecond_datatype(),
567 false,
568 ),
569 semantic_type: SemanticType::Timestamp,
570 column_id: 4,
571 })
572 .primary_key(vec![1, 2]);
573 builder.build().unwrap()
574 }
575
576 pub(crate) fn test_object_store() -> ObjectStore {
577 ObjectStore::new(Memory::default()).unwrap()
578 }
579
580 pub(crate) fn tag_column() -> Expr {
581 Expr::Column(Column::from_name("a"))
582 }
583
584 pub(crate) fn tag_column2() -> Expr {
585 Expr::Column(Column::from_name("b"))
586 }
587
588 pub(crate) fn field_column() -> Expr {
589 Expr::Column(Column::from_name("c"))
590 }
591
592 pub(crate) fn nonexistent_column() -> Expr {
593 Expr::Column(Column::from_name("nonexistence"))
594 }
595
596 pub(crate) fn string_lit(s: impl Into<String>) -> Expr {
597 s.into().lit()
598 }
599
600 pub(crate) fn int64_lit(i: impl Into<i64>) -> Expr {
601 i.into().lit()
602 }
603
604 pub(crate) fn encoded_string(s: impl Into<String>) -> Vec<u8> {
605 let mut bytes = vec![];
606 IndexValueCodec::encode_nonnull_value(
607 Value::from(s.into()).as_value_ref(),
608 &SortField::new(ConcreteDataType::string_datatype()),
609 &mut bytes,
610 )
611 .unwrap();
612 bytes
613 }
614
615 pub(crate) fn encoded_int64(s: impl Into<i64>) -> Vec<u8> {
616 let mut bytes = vec![];
617 IndexValueCodec::encode_nonnull_value(
618 Value::from(s.into()).as_value_ref(),
619 &SortField::new(ConcreteDataType::int64_datatype()),
620 &mut bytes,
621 )
622 .unwrap();
623 bytes
624 }
625
626 fn json_expr(path: &str, ty: Option<ScalarValue>) -> Expr {
627 use common_function::scalars::udf::create_udf;
628 use datafusion_expr::expr::ScalarFunction;
629 let mut args = vec![datafusion_expr::col("attrs"), string_lit(path)];
630 if let Some(ty) = ty {
631 args.push(Expr::Literal(ty, None));
632 }
633 Expr::ScalarFunction(ScalarFunction::new_udf(
634 Arc::new(create_udf(Arc::new(JsonGetWithType::default()))),
635 args,
636 ))
637 }
638
639 fn json_metadata() -> RegionMetadata {
640 use datatypes::extension::json::JsonMetadata;
641 use datatypes::json::{JsonSettings, JsonTypeHint};
642 use datatypes::types::json_type::{JsonNativeType, JsonObjectType};
643 let mut column = ColumnSchema::new(
644 "attrs",
645 ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())),
646 true,
647 );
648 let hints = [
649 (vec!["a"], true),
650 (vec!["b"], true),
651 (vec!["a.b"], true),
652 (vec!["nested", "value"], true),
653 (vec!["unindexed"], false),
654 ]
655 .into_iter()
656 .map(|(path, inverted_index)| JsonTypeHint {
657 path: path.into_iter().map(String::from).collect(),
658 data_type: ConcreteDataType::string_datatype(),
659 inverted_index,
660 skipping_index: None,
661 fulltext_index: None,
662 })
663 .chain(std::iter::once(JsonTypeHint {
664 path: vec!["n".into()],
665 data_type: ConcreteDataType::int64_datatype(),
666 inverted_index: true,
667 skipping_index: None,
668 fulltext_index: None,
669 }))
670 .collect();
671 column.with_extension_type(&Json2ExtensionType::new(Arc::new(JsonMetadata::new(
672 JsonSettings::try_new(hints, Some(0)).unwrap(),
673 ))));
674 let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1));
675 builder
676 .push_column_metadata(ColumnMetadata {
677 column_schema: column,
678 column_id: 1,
679 semantic_type: SemanticType::Field,
680 })
681 .push_column_metadata(ColumnMetadata {
682 column_schema: ColumnSchema::new(
683 "ts",
684 ConcreteDataType::timestamp_millisecond_datatype(),
685 false,
686 ),
687 column_id: 2,
688 semantic_type: SemanticType::Timestamp,
689 });
690 builder.build().unwrap()
691 }
692}