1use std::collections::HashSet;
23use std::sync::Arc;
24
25use arrow::array::AsArray;
26use arrow::{
27 array::{ArrayRef, BooleanArray, new_null_array},
28 datatypes::{DataType, Field, Schema, SchemaRef},
29 record_batch::{RecordBatch, RecordBatchOptions},
30};
31pub use datafusion_common::pruning::PruningStatistics;
33use datafusion_physical_expr::simplifier::PhysicalExprSimplifier;
34use datafusion_physical_plan::metrics::Count;
35use log::{debug, trace};
36
37use datafusion_common::error::Result;
38use datafusion_common::tree_node::{TransformedResult, TreeNodeRecursion};
39use datafusion_common::{
40 _internal_datafusion_err, Column, DFSchema, assert_eq_or_internal_err,
41};
42use datafusion_common::{
43 ScalarValue, internal_datafusion_err, plan_datafusion_err, plan_err,
44 tree_node::{Transformed, TreeNode},
45};
46use datafusion_expr_common::casts::try_cast_literal_to_type;
47use datafusion_expr_common::operator::Operator;
48use datafusion_physical_expr::utils::{Guarantee, LiteralGuarantee};
49use datafusion_physical_expr::{PhysicalExprRef, expressions as phys_expr};
50use datafusion_physical_expr_common::physical_expr::snapshot_physical_expr_opt;
51use datafusion_physical_plan::{ColumnarValue, PhysicalExpr};
52
53#[derive(Debug, Clone)]
365pub struct PruningPredicate {
366 schema: SchemaRef,
368 predicate_expr: Arc<dyn PhysicalExpr>,
371 required_columns: RequiredColumns,
373 orig_expr: Arc<dyn PhysicalExpr>,
376 literal_guarantees: Vec<LiteralGuarantee>,
381}
382
383pub fn build_pruning_predicate(
389 predicate: Arc<dyn PhysicalExpr>,
390 file_schema: &SchemaRef,
391 predicate_creation_errors: &Count,
392) -> Option<Arc<PruningPredicate>> {
393 PruningPredicateBuilder::new()
394 .with_file_schema(Arc::clone(file_schema))
395 .with_error_counter(predicate_creation_errors)
396 .build(predicate)
397}
398
399#[derive(Default)]
413pub struct PruningPredicateBuilder<'a> {
414 file_schema: Option<SchemaRef>,
415 error_counter: Option<&'a Count>,
416 max_in_list_size: usize,
417}
418
419impl<'a> PruningPredicateBuilder<'a> {
420 pub fn new() -> Self {
422 Self {
423 file_schema: None,
424 error_counter: None,
425 max_in_list_size: MAX_IN_LIST_SIZE,
426 }
427 }
428
429 pub fn with_file_schema(mut self, file_schema: SchemaRef) -> Self {
432 self.file_schema = Some(file_schema);
433 self
434 }
435
436 pub fn with_error_counter(mut self, error_counter: &'a Count) -> Self {
440 self.error_counter = Some(error_counter);
441 self
442 }
443
444 pub fn with_max_in_list_size(mut self, max_in_list_size: usize) -> Self {
451 self.max_in_list_size = max_in_list_size;
452 self
453 }
454
455 pub fn build(
460 self,
461 predicate: Arc<dyn PhysicalExpr>,
462 ) -> Option<Arc<PruningPredicate>> {
463 let error_counter = self.error_counter;
464 match self.try_build(predicate) {
465 Ok(pruning_predicate) => {
466 if !pruning_predicate.always_true() {
467 return Some(Arc::new(pruning_predicate));
468 }
469 }
470 Err(e) => {
471 debug!("Could not create pruning predicate for: {e}");
472 if let Some(counter) = error_counter {
473 counter.add(1);
474 }
475 }
476 }
477 None
478 }
479
480 pub fn try_build(
484 self,
485 mut predicate: Arc<dyn PhysicalExpr>,
486 ) -> Result<PruningPredicate> {
487 let file_schema = self.file_schema.ok_or_else(|| {
488 _internal_datafusion_err!(
489 "PruningPredicateBuilder requires a file schema (call `with_file_schema`)"
490 )
491 })?;
492
493 let tf = snapshot_physical_expr_opt(predicate)?;
497 if tf.transformed {
498 let simplifier = PhysicalExprSimplifier::new(&file_schema);
505 predicate = simplifier.simplify(tf.data)?;
506 } else {
507 predicate = tf.data;
508 }
509 let unhandled_hook = Arc::new(ConstantUnhandledPredicateHook::default()) as _;
510
511 let mut required_columns = RequiredColumns::new();
513 let predicate_expr = build_predicate_expression(
514 &predicate,
515 &file_schema,
516 &mut required_columns,
517 &unhandled_hook,
518 self.max_in_list_size,
519 );
520 let predicate_schema = required_columns.schema();
521 let predicate_expr =
523 PhysicalExprSimplifier::new(&predicate_schema).simplify(predicate_expr)?;
524 let literal_guarantees = LiteralGuarantee::analyze(&predicate);
525
526 Ok(PruningPredicate {
527 schema: file_schema,
528 predicate_expr,
529 required_columns,
530 orig_expr: predicate,
531 literal_guarantees,
532 })
533 }
534}
535
536pub trait UnhandledPredicateHook {
540 fn handle(&self, expr: &Arc<dyn PhysicalExpr>) -> Arc<dyn PhysicalExpr>;
543}
544
545#[derive(Debug, Clone)]
548struct ConstantUnhandledPredicateHook {
549 default: Arc<dyn PhysicalExpr>,
550}
551
552impl Default for ConstantUnhandledPredicateHook {
553 fn default() -> Self {
554 Self {
555 default: Arc::new(phys_expr::Literal::new(ScalarValue::from(true))),
556 }
557 }
558}
559
560impl UnhandledPredicateHook for ConstantUnhandledPredicateHook {
561 fn handle(&self, _expr: &Arc<dyn PhysicalExpr>) -> Arc<dyn PhysicalExpr> {
562 Arc::clone(&self.default)
563 }
564}
565
566impl PruningPredicate {
567 #[deprecated(since = "55.0.0", note = "Use PruningPredicateBuilder instead")]
598 pub fn try_new(expr: Arc<dyn PhysicalExpr>, schema: SchemaRef) -> Result<Self> {
599 PruningPredicateBuilder::new()
600 .with_file_schema(schema)
601 .try_build(expr)
602 }
603
604 pub fn prune<S: PruningStatistics + ?Sized>(
619 &self,
620 statistics: &S,
621 ) -> Result<Vec<bool>> {
622 let mut builder = BoolVecBuilder::new(statistics.num_containers());
623
624 for literal_guarantee in &self.literal_guarantees {
627 let LiteralGuarantee {
628 column,
629 guarantee,
630 literals,
631 } = literal_guarantee;
632 if let Some(results) = statistics.contained(column, literals) {
633 match guarantee {
634 Guarantee::In => builder.combine_array(&results),
639 Guarantee::NotIn => {
645 builder.combine_array(&arrow::compute::not(&results)?)
646 }
647 }
648 if builder.check_all_pruned() {
651 return Ok(builder.build());
652 }
653 }
654 }
655
656 let statistics_batch =
662 build_statistics_record_batch(statistics, &self.required_columns)?;
663
664 builder.combine_value(self.predicate_expr.evaluate(&statistics_batch)?);
666
667 Ok(builder.build())
668 }
669
670 pub fn schema(&self) -> &SchemaRef {
672 &self.schema
673 }
674
675 pub fn orig_expr(&self) -> &Arc<dyn PhysicalExpr> {
677 &self.orig_expr
678 }
679
680 pub fn predicate_expr(&self) -> &Arc<dyn PhysicalExpr> {
682 &self.predicate_expr
683 }
684
685 pub fn literal_guarantees(&self) -> &[LiteralGuarantee] {
691 &self.literal_guarantees
692 }
693
694 pub fn always_true(&self) -> bool {
701 is_always_true(&self.predicate_expr) && self.literal_guarantees.is_empty()
702 }
703
704 pub fn required_columns(&self) -> &RequiredColumns {
705 &self.required_columns
706 }
707
708 pub fn literal_columns(&self) -> Vec<String> {
716 let mut seen = HashSet::new();
717 self.literal_guarantees
718 .iter()
719 .map(|e| &e.column.name)
720 .filter(|name| seen.insert(*name))
722 .map(|s| s.to_string())
723 .collect()
724 }
725}
726
727#[derive(Debug)]
729struct BoolVecBuilder {
730 inner: Vec<bool>,
734}
735
736impl BoolVecBuilder {
737 fn new(num_containers: usize) -> Self {
739 Self {
740 inner: vec![true; num_containers],
742 }
743 }
744
745 fn combine_array(&mut self, array: &BooleanArray) {
753 assert_eq!(array.len(), self.inner.len());
754 for (cur, new) in self.inner.iter_mut().zip(array.iter()) {
755 if let Some(false) = new {
759 *cur = false;
760 }
761 }
762 }
763
764 fn combine_value(&mut self, value: ColumnarValue) {
770 match value {
771 ColumnarValue::Array(array) => {
772 self.combine_array(array.as_boolean());
773 }
774 ColumnarValue::Scalar(ScalarValue::Boolean(Some(false))) => {
775 self.inner = vec![false; self.inner.len()];
777 }
778 _ => {
779 }
782 }
783 }
784
785 fn build(self) -> Vec<bool> {
787 self.inner
788 }
789
790 fn check_all_pruned(&self) -> bool {
792 self.inner.iter().all(|&x| !x)
793 }
794}
795
796fn is_always_true(expr: &Arc<dyn PhysicalExpr>) -> bool {
797 expr.downcast_ref::<phys_expr::Literal>()
798 .map(|l| matches!(l.value(), ScalarValue::Boolean(Some(true))))
799 .unwrap_or_default()
800}
801
802fn is_always_false(expr: &Arc<dyn PhysicalExpr>) -> bool {
803 expr.downcast_ref::<phys_expr::Literal>()
804 .map(|l| matches!(l.value(), ScalarValue::Boolean(Some(false))))
805 .unwrap_or_default()
806}
807
808#[derive(Debug, Default, Clone)]
818pub struct RequiredColumns {
819 columns: Vec<(phys_expr::Column, StatisticsType, Field)>,
825}
826
827impl RequiredColumns {
828 fn new() -> Self {
829 Self::default()
830 }
831
832 pub fn single_column(&self) -> Option<&phys_expr::Column> {
841 if self.columns.windows(2).all(|w| {
842 let c1 = &w[0].0;
844 let c2 = &w[1].0;
845 c1 == c2
846 }) {
847 self.columns.first().map(|r| &r.0)
848 } else {
849 None
850 }
851 }
852
853 fn schema(&self) -> Schema {
860 let fields = self
861 .columns
862 .iter()
863 .map(|(_c, _t, f)| f.clone())
864 .collect::<Vec<_>>();
865 Schema::new(fields)
866 }
867
868 pub(crate) fn iter(
871 &self,
872 ) -> impl Iterator<Item = &(phys_expr::Column, StatisticsType, Field)> {
873 self.columns.iter()
874 }
875
876 fn find_stat_column(
877 &self,
878 column: &phys_expr::Column,
879 statistics_type: StatisticsType,
880 ) -> Option<usize> {
881 match statistics_type {
882 StatisticsType::RowCount => {
883 self.columns
885 .iter()
886 .enumerate()
887 .find(|(_i, (_c, t, _f))| t == &statistics_type)
888 .map(|(i, (_c, _t, _f))| i)
889 }
890 _ => self
891 .columns
892 .iter()
893 .enumerate()
894 .find(|(_i, (c, t, _f))| c == column && t == &statistics_type)
895 .map(|(i, (_c, _t, _f))| i),
896 }
897 }
898
899 fn stat_column_expr(
908 &mut self,
909 column: &phys_expr::Column,
910 column_expr: &Arc<dyn PhysicalExpr>,
911 field: &Field,
912 stat_type: StatisticsType,
913 ) -> Result<Arc<dyn PhysicalExpr>> {
914 let (idx, need_to_insert) = match self.find_stat_column(column, stat_type) {
915 Some(idx) => (idx, false),
916 None => (self.columns.len(), true),
917 };
918
919 let column_name = column.name();
920 let stat_column_name = match stat_type {
921 StatisticsType::Min => format!("{column_name}_min"),
922 StatisticsType::Max => format!("{column_name}_max"),
923 StatisticsType::NullCount => format!("{column_name}_null_count"),
924 StatisticsType::RowCount => "row_count".to_string(),
925 };
926
927 let stat_column = phys_expr::Column::new(&stat_column_name, idx);
928
929 if need_to_insert {
931 let nullable = true;
933 let stat_field =
934 Field::new(stat_column.name(), field.data_type().clone(), nullable);
935 self.columns.push((column.clone(), stat_type, stat_field));
936 }
937 rewrite_column_expr(Arc::clone(column_expr), column, &stat_column)
938 }
939
940 fn min_column_expr(
942 &mut self,
943 column: &phys_expr::Column,
944 column_expr: &Arc<dyn PhysicalExpr>,
945 field: &Field,
946 ) -> Result<Arc<dyn PhysicalExpr>> {
947 self.stat_column_expr(column, column_expr, field, StatisticsType::Min)
948 }
949
950 fn max_column_expr(
952 &mut self,
953 column: &phys_expr::Column,
954 column_expr: &Arc<dyn PhysicalExpr>,
955 field: &Field,
956 ) -> Result<Arc<dyn PhysicalExpr>> {
957 self.stat_column_expr(column, column_expr, field, StatisticsType::Max)
958 }
959
960 fn null_count_column_expr(
962 &mut self,
963 column: &phys_expr::Column,
964 column_expr: &Arc<dyn PhysicalExpr>,
965 field: &Field,
966 ) -> Result<Arc<dyn PhysicalExpr>> {
967 self.stat_column_expr(column, column_expr, field, StatisticsType::NullCount)
968 }
969
970 fn row_count_column_expr(
972 &mut self,
973 column: &phys_expr::Column,
974 column_expr: &Arc<dyn PhysicalExpr>,
975 field: &Field,
976 ) -> Result<Arc<dyn PhysicalExpr>> {
977 self.stat_column_expr(column, column_expr, field, StatisticsType::RowCount)
978 }
979}
980
981impl From<Vec<(phys_expr::Column, StatisticsType, Field)>> for RequiredColumns {
982 fn from(columns: Vec<(phys_expr::Column, StatisticsType, Field)>) -> Self {
983 Self { columns }
984 }
985}
986
987fn build_statistics_record_batch<S: PruningStatistics + ?Sized>(
1013 statistics: &S,
1014 required_columns: &RequiredColumns,
1015) -> Result<RecordBatch> {
1016 let mut arrays = Vec::<ArrayRef>::new();
1017 for (column, statistics_type, stat_field) in required_columns.iter() {
1019 let column = Column::from_name(column.name());
1020 let data_type = stat_field.data_type();
1021
1022 let num_containers = statistics.num_containers();
1023
1024 let array = match statistics_type {
1025 StatisticsType::Min => statistics.min_values(&column),
1026 StatisticsType::Max => statistics.max_values(&column),
1027 StatisticsType::NullCount => statistics.null_counts(&column),
1028 StatisticsType::RowCount => statistics.row_counts(),
1029 };
1030 let array = array.unwrap_or_else(|| new_null_array(data_type, num_containers));
1031
1032 assert_eq_or_internal_err!(
1033 num_containers,
1034 array.len(),
1035 "mismatched statistics length. Expected {}, got {}",
1036 num_containers,
1037 array.len()
1038 );
1039
1040 let array = arrow::compute::cast(&array, data_type)?;
1043
1044 arrays.push(array);
1045 }
1046
1047 let schema = Arc::new(required_columns.schema());
1048 let mut options = RecordBatchOptions::default();
1050 options.row_count = Some(statistics.num_containers());
1051
1052 trace!("Creating statistics batch for {required_columns:#?} with {arrays:#?}");
1053
1054 RecordBatch::try_new_with_options(schema, arrays, &options).map_err(|err| {
1055 plan_datafusion_err!("Can not create statistics record batch: {err}")
1056 })
1057}
1058
1059struct PruningExpressionBuilder<'a> {
1060 column: phys_expr::Column,
1061 column_expr: Arc<dyn PhysicalExpr>,
1062 op: Operator,
1063 scalar_expr: Arc<dyn PhysicalExpr>,
1064 field: &'a Field,
1065 required_columns: &'a mut RequiredColumns,
1066}
1067
1068impl<'a> PruningExpressionBuilder<'a> {
1069 fn try_new(
1070 left: &'a Arc<dyn PhysicalExpr>,
1071 right: &'a Arc<dyn PhysicalExpr>,
1072 left_columns: ColumnReferenceCount,
1073 right_columns: ColumnReferenceCount,
1074 op: Operator,
1075 schema: &'a SchemaRef,
1076 required_columns: &'a mut RequiredColumns,
1077 ) -> Result<Self> {
1078 let (column_expr, scalar_expr, column, correct_operator) = match (
1080 left_columns,
1081 right_columns,
1082 ) {
1083 (ColumnReferenceCount::One(column), ColumnReferenceCount::Zero) => {
1084 (left, right, column, op)
1085 }
1086 (ColumnReferenceCount::Zero, ColumnReferenceCount::One(column)) => {
1087 (right, left, column, reverse_operator(op)?)
1088 }
1089 (ColumnReferenceCount::One(_), ColumnReferenceCount::One(_)) => {
1090 return plan_err!(
1092 "Expression not supported for pruning: left has 1 column, right has 1 column"
1093 );
1094 }
1095 (ColumnReferenceCount::Zero, ColumnReferenceCount::Zero) => {
1096 return plan_err!(
1098 "Pruning literal expressions is not supported, please call PhysicalExprSimplifier first"
1099 );
1100 }
1101 (ColumnReferenceCount::Many, _) | (_, ColumnReferenceCount::Many) => {
1102 return plan_err!(
1103 "Expression not supported for pruning: left or right has multiple columns"
1104 );
1105 }
1106 };
1107
1108 let df_schema = DFSchema::try_from(Arc::clone(schema))?;
1109 let (column_expr, correct_operator, scalar_expr) = rewrite_expr_to_prunable(
1110 column_expr,
1111 correct_operator,
1112 scalar_expr,
1113 df_schema,
1114 )?;
1115 let field = match schema.column_with_name(column.name()) {
1116 Some((_, f)) => f,
1117 _ => {
1118 return plan_err!("Field not found in schema");
1119 }
1120 };
1121
1122 Ok(Self {
1123 column,
1124 column_expr,
1125 op: correct_operator,
1126 scalar_expr,
1127 field,
1128 required_columns,
1129 })
1130 }
1131
1132 fn op(&self) -> Operator {
1133 self.op
1134 }
1135
1136 fn scalar_expr(&self) -> &Arc<dyn PhysicalExpr> {
1137 &self.scalar_expr
1138 }
1139
1140 fn min_column_expr(&mut self) -> Result<Arc<dyn PhysicalExpr>> {
1141 self.required_columns
1142 .min_column_expr(&self.column, &self.column_expr, self.field)
1143 }
1144
1145 fn max_column_expr(&mut self) -> Result<Arc<dyn PhysicalExpr>> {
1146 self.required_columns
1147 .max_column_expr(&self.column, &self.column_expr, self.field)
1148 }
1149
1150 fn null_count_column_expr(&mut self) -> Result<Arc<dyn PhysicalExpr>> {
1157 let column_expr = Arc::new(self.column.clone()) as _;
1159
1160 let null_count_field = &Field::new(self.field.name(), DataType::UInt64, true);
1162
1163 self.required_columns.null_count_column_expr(
1164 &self.column,
1165 &column_expr,
1166 null_count_field,
1167 )
1168 }
1169
1170 fn row_count_column_expr(&mut self) -> Result<Arc<dyn PhysicalExpr>> {
1177 let column_expr = Arc::new(self.column.clone()) as _;
1179
1180 let row_count_field = &Field::new(self.field.name(), DataType::UInt64, true);
1182
1183 self.required_columns.row_count_column_expr(
1184 &self.column,
1185 &column_expr,
1186 row_count_field,
1187 )
1188 }
1189}
1190
1191fn rewrite_expr_to_prunable(
1204 column_expr: &PhysicalExprRef,
1205 op: Operator,
1206 scalar_expr: &PhysicalExprRef,
1207 schema: DFSchema,
1208) -> Result<(PhysicalExprRef, Operator, PhysicalExprRef)> {
1209 if !is_compare_op(op) {
1210 return plan_err!("rewrite_expr_to_prunable only support compare expression");
1211 }
1212
1213 if column_expr.downcast_ref::<phys_expr::Column>().is_some() {
1214 Ok((Arc::clone(column_expr), op, Arc::clone(scalar_expr)))
1216 } else if let Some(cast) = column_expr.downcast_ref::<phys_expr::CastExpr>() {
1217 let (left, op, right) = rewrite_cast_child_to_prunable(
1219 cast.expr(),
1220 cast.cast_type(),
1221 op,
1222 scalar_expr,
1223 schema,
1224 )?;
1225 let left = Arc::new(phys_expr::CastExpr::new_with_target_field(
1226 left,
1227 Arc::clone(cast.target_field()),
1228 None,
1229 ));
1230 Ok((left, op, right))
1235 } else if let Some(try_cast) = column_expr.downcast_ref::<phys_expr::TryCastExpr>() {
1236 let (left, op, right) = rewrite_cast_child_to_prunable(
1238 try_cast.expr(),
1239 try_cast.cast_type(),
1240 op,
1241 scalar_expr,
1242 schema,
1243 )?;
1244 let left = Arc::new(phys_expr::TryCastExpr::new(
1245 left,
1246 try_cast.cast_type().clone(),
1247 ));
1248 Ok((left, op, right))
1249 } else if let Some(neg) = column_expr.downcast_ref::<phys_expr::NegativeExpr>() {
1250 let (left, op, right) =
1252 rewrite_expr_to_prunable(neg.arg(), op, scalar_expr, schema)?;
1253 let right = Arc::new(phys_expr::NegativeExpr::new(right));
1254 Ok((left, reverse_operator(op)?, right))
1255 } else if let Some(not) = column_expr.downcast_ref::<phys_expr::NotExpr>() {
1256 if !matches!(
1258 op,
1259 Operator::Eq
1260 | Operator::NotEq
1261 | Operator::IsDistinctFrom
1262 | Operator::IsNotDistinctFrom
1263 ) {
1264 return plan_err!(
1265 "Not with operator other than Eq / NotEq / IsDistinctFrom / IsNotDistinctFrom is not supported"
1266 );
1267 }
1268 if not.arg().downcast_ref::<phys_expr::Column>().is_some() {
1269 let left = Arc::clone(not.arg());
1270 let right = Arc::new(phys_expr::NotExpr::new(Arc::clone(scalar_expr)));
1271 Ok((left, reverse_operator(op)?, right))
1272 } else {
1273 plan_err!("Not with complex expression {column_expr:?} is not supported")
1274 }
1275 } else {
1276 plan_err!("column expression {column_expr:?} is not supported")
1277 }
1278}
1279
1280fn rewrite_cast_child_to_prunable(
1281 cast_child_expr: &PhysicalExprRef,
1282 cast_type: &DataType,
1283 op: Operator,
1284 scalar_expr: &PhysicalExprRef,
1285 schema: DFSchema,
1286) -> Result<(PhysicalExprRef, Operator, PhysicalExprRef)> {
1287 verify_support_type_for_prune(
1288 &cast_child_expr.data_type(schema.as_arrow())?,
1289 cast_type,
1290 )?;
1291 rewrite_expr_to_prunable(cast_child_expr, op, scalar_expr, schema)
1292}
1293
1294fn is_compare_op(op: Operator) -> bool {
1295 matches!(
1296 op,
1297 Operator::Eq
1298 | Operator::NotEq
1299 | Operator::Lt
1300 | Operator::LtEq
1301 | Operator::Gt
1302 | Operator::GtEq
1303 | Operator::IsDistinctFrom
1304 | Operator::IsNotDistinctFrom
1305 | Operator::LikeMatch
1306 | Operator::NotLikeMatch
1307 )
1308}
1309
1310fn verify_support_type_for_prune(from_type: &DataType, to_type: &DataType) -> Result<()> {
1315 let from_type = match from_type {
1317 DataType::Dictionary(_, t) => {
1318 return verify_support_type_for_prune(t.as_ref(), to_type);
1319 }
1320 _ => from_type,
1321 };
1322 let to_type = match to_type {
1323 DataType::Dictionary(_, t) => {
1324 return verify_support_type_for_prune(from_type, t.as_ref());
1325 }
1326 _ => to_type,
1327 };
1328 if from_type.is_string() == to_type.is_string() {
1332 Ok(())
1333 } else {
1334 plan_err!(
1335 "Try Cast/Cast with from type {from_type} to type {to_type} is not supported"
1336 )
1337 }
1338}
1339
1340fn rewrite_column_expr(
1342 e: Arc<dyn PhysicalExpr>,
1343 column_old: &phys_expr::Column,
1344 column_new: &phys_expr::Column,
1345) -> Result<Arc<dyn PhysicalExpr>> {
1346 e.transform(|expr| {
1347 if let Some(column) = expr.downcast_ref::<phys_expr::Column>()
1348 && column == column_old
1349 {
1350 return Ok(Transformed::yes(Arc::new(column_new.clone())));
1351 }
1352
1353 Ok(Transformed::no(expr))
1354 })
1355 .data()
1356}
1357
1358fn reverse_operator(op: Operator) -> Result<Operator> {
1359 op.swap().ok_or_else(|| {
1360 internal_datafusion_err!(
1361 "Could not reverse operator {op} while building pruning predicate"
1362 )
1363 })
1364}
1365
1366fn build_single_column_expr(
1371 column: &phys_expr::Column,
1372 schema: &Schema,
1373 required_columns: &mut RequiredColumns,
1374 is_not: bool, ) -> Option<Arc<dyn PhysicalExpr>> {
1376 let field = schema.field_with_name(column.name()).ok()?;
1377
1378 if *field.data_type() == DataType::Boolean {
1379 let col_ref = Arc::new(column.clone()) as _;
1380
1381 let min = required_columns
1382 .min_column_expr(column, &col_ref, field)
1383 .ok()?;
1384 let max = required_columns
1385 .max_column_expr(column, &col_ref, field)
1386 .ok()?;
1387
1388 if is_not {
1392 Some(Arc::new(phys_expr::NotExpr::new(Arc::new(
1395 phys_expr::BinaryExpr::new(min, Operator::And, max),
1396 ))))
1397 } else {
1398 Some(Arc::new(phys_expr::BinaryExpr::new(min, Operator::Or, max)))
1401 }
1402 } else {
1403 None
1404 }
1405}
1406
1407fn build_is_null_column_expr(
1416 expr: &Arc<dyn PhysicalExpr>,
1417 schema: &Schema,
1418 required_columns: &mut RequiredColumns,
1419 with_not: bool,
1420) -> Option<Arc<dyn PhysicalExpr>> {
1421 if let Some(col) = expr.downcast_ref::<phys_expr::Column>() {
1422 let field = schema.field_with_name(col.name()).ok()?;
1423
1424 let null_count_field = &Field::new(field.name(), DataType::UInt64, true);
1425 if with_not {
1426 if let Ok(row_count_expr) =
1427 required_columns.row_count_column_expr(col, expr, null_count_field)
1428 {
1429 required_columns
1430 .null_count_column_expr(col, expr, null_count_field)
1431 .map(|null_count_column_expr| {
1432 Arc::new(phys_expr::BinaryExpr::new(
1434 null_count_column_expr,
1435 Operator::NotEq,
1436 row_count_expr,
1437 )) as _
1438 })
1439 .ok()
1440 } else {
1441 None
1442 }
1443 } else {
1444 required_columns
1445 .null_count_column_expr(col, expr, null_count_field)
1446 .map(|null_count_column_expr| {
1447 Arc::new(phys_expr::BinaryExpr::new(
1449 null_count_column_expr,
1450 Operator::Gt,
1451 Arc::new(phys_expr::Literal::new(ScalarValue::UInt64(Some(0)))),
1452 )) as _
1453 })
1454 .ok()
1455 }
1456 } else {
1457 None
1458 }
1459}
1460
1461pub const MAX_IN_LIST_SIZE: usize = 20;
1468
1469pub struct PredicateRewriter {
1472 unhandled_hook: Arc<dyn UnhandledPredicateHook>,
1473 max_in_list_size: usize,
1474}
1475
1476impl Default for PredicateRewriter {
1477 fn default() -> Self {
1478 Self {
1479 unhandled_hook: Arc::new(ConstantUnhandledPredicateHook::default()),
1480 max_in_list_size: MAX_IN_LIST_SIZE,
1481 }
1482 }
1483}
1484
1485impl PredicateRewriter {
1486 pub fn new() -> Self {
1488 Self::default()
1489 }
1490
1491 pub fn with_unhandled_hook(
1493 mut self,
1494 unhandled_hook: Arc<dyn UnhandledPredicateHook>,
1495 ) -> Self {
1496 self.unhandled_hook = unhandled_hook;
1497 self
1498 }
1499
1500 pub fn with_max_in_list_size(mut self, max_in_list_size: usize) -> Self {
1509 self.max_in_list_size = max_in_list_size;
1510 self
1511 }
1512
1513 pub fn rewrite_predicate_to_statistics_predicate(
1524 &self,
1525 expr: &Arc<dyn PhysicalExpr>,
1526 schema: &Schema,
1527 ) -> Arc<dyn PhysicalExpr> {
1528 let mut required_columns = RequiredColumns::new();
1529 build_predicate_expression(
1530 expr,
1531 &Arc::new(schema.clone()),
1532 &mut required_columns,
1533 &self.unhandled_hook,
1534 self.max_in_list_size,
1535 )
1536 }
1537}
1538
1539fn build_predicate_expression(
1551 expr: &Arc<dyn PhysicalExpr>,
1552 schema: &SchemaRef,
1553 required_columns: &mut RequiredColumns,
1554 unhandled_hook: &Arc<dyn UnhandledPredicateHook>,
1555 max_in_list_size: usize,
1556) -> Arc<dyn PhysicalExpr> {
1557 if is_always_false(expr) {
1558 return Arc::clone(expr);
1561 }
1562 if let Some(is_null) = expr.downcast_ref::<phys_expr::IsNullExpr>() {
1564 return build_is_null_column_expr(is_null.arg(), schema, required_columns, false)
1565 .unwrap_or_else(|| unhandled_hook.handle(expr));
1566 }
1567 if let Some(is_not_null) = expr.downcast_ref::<phys_expr::IsNotNullExpr>() {
1568 return build_is_null_column_expr(
1569 is_not_null.arg(),
1570 schema,
1571 required_columns,
1572 true,
1573 )
1574 .unwrap_or_else(|| unhandled_hook.handle(expr));
1575 }
1576 if let Some(col) = expr.downcast_ref::<phys_expr::Column>() {
1577 return build_single_column_expr(col, schema, required_columns, false)
1578 .unwrap_or_else(|| unhandled_hook.handle(expr));
1579 }
1580 if let Some(not) = expr.downcast_ref::<phys_expr::NotExpr>() {
1581 if let Some(col) = not.arg().downcast_ref::<phys_expr::Column>() {
1583 return build_single_column_expr(col, schema, required_columns, true)
1584 .unwrap_or_else(|| unhandled_hook.handle(expr));
1585 } else {
1586 return unhandled_hook.handle(expr);
1587 }
1588 }
1589 if let Some(in_list) = expr.downcast_ref::<phys_expr::InListExpr>() {
1590 if !in_list.list().is_empty() && in_list.list().len() <= max_in_list_size {
1591 let eq_op = if in_list.negated() {
1592 Operator::NotEq
1593 } else {
1594 Operator::Eq
1595 };
1596 let re_op = if in_list.negated() {
1597 Operator::And
1598 } else {
1599 Operator::Or
1600 };
1601 let change_expr = in_list
1602 .list()
1603 .iter()
1604 .map(|e| {
1605 Arc::new(phys_expr::BinaryExpr::new(
1606 Arc::clone(in_list.expr()),
1607 eq_op,
1608 Arc::clone(e),
1609 )) as _
1610 })
1611 .reduce(|a, b| Arc::new(phys_expr::BinaryExpr::new(a, re_op, b)) as _)
1612 .unwrap();
1613 return build_predicate_expression(
1614 &change_expr,
1615 schema,
1616 required_columns,
1617 unhandled_hook,
1618 max_in_list_size,
1619 );
1620 } else {
1621 return unhandled_hook.handle(expr);
1622 }
1623 }
1624
1625 let (left, op, right) = {
1626 if let Some(bin_expr) = expr.downcast_ref::<phys_expr::BinaryExpr>() {
1627 (
1628 Arc::clone(bin_expr.left()),
1629 *bin_expr.op(),
1630 Arc::clone(bin_expr.right()),
1631 )
1632 } else if let Some(like_expr) = expr.downcast_ref::<phys_expr::LikeExpr>() {
1633 if like_expr.case_insensitive() {
1634 return unhandled_hook.handle(expr);
1635 }
1636 let op = match (like_expr.negated(), like_expr.case_insensitive()) {
1637 (false, false) => Operator::LikeMatch,
1638 (true, false) => Operator::NotLikeMatch,
1639 (false, true) => Operator::ILikeMatch,
1640 (true, true) => Operator::NotILikeMatch,
1641 };
1642 (
1643 Arc::clone(like_expr.expr()),
1644 op,
1645 Arc::clone(like_expr.pattern()),
1646 )
1647 } else {
1648 return unhandled_hook.handle(expr);
1649 }
1650 };
1651
1652 if op == Operator::And || op == Operator::Or {
1653 let left_expr = build_predicate_expression(
1654 &left,
1655 schema,
1656 required_columns,
1657 unhandled_hook,
1658 max_in_list_size,
1659 );
1660 let right_expr = build_predicate_expression(
1661 &right,
1662 schema,
1663 required_columns,
1664 unhandled_hook,
1665 max_in_list_size,
1666 );
1667 let expr = match (&left_expr, op, &right_expr) {
1669 (left, Operator::And, right)
1670 if is_always_false(left) || is_always_false(right) =>
1671 {
1672 Arc::new(phys_expr::Literal::new(ScalarValue::Boolean(Some(false))))
1673 }
1674 (left, Operator::And, _) if is_always_true(left) => right_expr,
1675 (_, Operator::And, right) if is_always_true(right) => left_expr,
1676 (left, Operator::Or, right)
1677 if is_always_true(left) || is_always_true(right) =>
1678 {
1679 Arc::new(phys_expr::Literal::new(ScalarValue::Boolean(Some(true))))
1680 }
1681 (left, Operator::Or, _) if is_always_false(left) => right_expr,
1682 (_, Operator::Or, right) if is_always_false(right) => left_expr,
1683
1684 _ => Arc::new(phys_expr::BinaryExpr::new(left_expr, op, right_expr)),
1685 };
1686 return expr;
1687 }
1688
1689 let left_columns = ColumnReferenceCount::from_expression(&left);
1690 let right_columns = ColumnReferenceCount::from_expression(&right);
1691 let expr_builder = PruningExpressionBuilder::try_new(
1692 &left,
1693 &right,
1694 left_columns,
1695 right_columns,
1696 op,
1697 schema,
1698 required_columns,
1699 );
1700 let mut expr_builder = match expr_builder {
1701 Ok(builder) => builder,
1702 Err(e) => {
1705 debug!("Error building pruning expression: {e}");
1706 return unhandled_hook.handle(expr);
1707 }
1708 };
1709
1710 build_statistics_expr(&mut expr_builder)
1711 .unwrap_or_else(|_| unhandled_hook.handle(expr))
1712}
1713
1714#[derive(Debug, PartialEq, Eq)]
1724enum ColumnReferenceCount {
1725 Zero,
1727 One(phys_expr::Column),
1729 Many,
1731}
1732
1733impl ColumnReferenceCount {
1734 fn from_expression(expr: &Arc<dyn PhysicalExpr>) -> Self {
1736 let mut seen = HashSet::<phys_expr::Column>::new();
1737 expr.apply(|expr| {
1738 if let Some(column) = expr.downcast_ref::<phys_expr::Column>() {
1739 seen.insert(column.clone());
1740 if seen.len() > 1 {
1741 return Ok(TreeNodeRecursion::Stop);
1742 }
1743 }
1744 Ok(TreeNodeRecursion::Continue)
1745 })
1746 .expect("no way to return error during recursion");
1748 match seen.len() {
1749 0 => ColumnReferenceCount::Zero,
1750 1 => ColumnReferenceCount::One(
1751 seen.into_iter().next().expect("just checked len==1"),
1752 ),
1753 _ => ColumnReferenceCount::Many,
1754 }
1755 }
1756}
1757
1758fn build_statistics_expr(
1759 expr_builder: &mut PruningExpressionBuilder,
1760) -> Result<Arc<dyn PhysicalExpr>> {
1761 let statistics_expr: Arc<dyn PhysicalExpr> = match expr_builder.op() {
1762 Operator::NotEq => build_ne_statistics_expr(expr_builder)?,
1763 Operator::Eq => {
1764 build_eq_statistics_expr(expr_builder)?
1767 }
1768 Operator::IsDistinctFrom => return build_is_distinct_from(expr_builder),
1769 Operator::IsNotDistinctFrom => return build_is_not_distinct_from(expr_builder),
1770 Operator::NotLikeMatch => build_not_like_match(expr_builder)?,
1771 Operator::LikeMatch => build_like_match(expr_builder).ok_or_else(|| {
1772 plan_datafusion_err!(
1773 "LIKE expression with wildcard at the beginning is not supported"
1774 )
1775 })?,
1776 Operator::Gt => {
1777 Arc::new(phys_expr::BinaryExpr::new(
1779 expr_builder.max_column_expr()?,
1780 Operator::Gt,
1781 Arc::clone(expr_builder.scalar_expr()),
1782 ))
1783 }
1784 Operator::GtEq => {
1785 Arc::new(phys_expr::BinaryExpr::new(
1787 expr_builder.max_column_expr()?,
1788 Operator::GtEq,
1789 Arc::clone(expr_builder.scalar_expr()),
1790 ))
1791 }
1792 Operator::Lt => {
1793 Arc::new(phys_expr::BinaryExpr::new(
1795 expr_builder.min_column_expr()?,
1796 Operator::Lt,
1797 Arc::clone(expr_builder.scalar_expr()),
1798 ))
1799 }
1800 Operator::LtEq => {
1801 Arc::new(phys_expr::BinaryExpr::new(
1803 expr_builder.min_column_expr()?,
1804 Operator::LtEq,
1805 Arc::clone(expr_builder.scalar_expr()),
1806 ))
1807 }
1808 _ => {
1810 return plan_err!(
1811 "expressions other than (neq, eq, gt, gteq, lt, lteq) are not supported"
1812 );
1813 }
1814 };
1815 let statistics_expr = wrap_null_count_check_expr(statistics_expr, expr_builder)?;
1816 Ok(statistics_expr)
1817}
1818
1819fn binary_expr(
1820 left: Arc<dyn PhysicalExpr>,
1821 op: Operator,
1822 right: Arc<dyn PhysicalExpr>,
1823) -> Arc<dyn PhysicalExpr> {
1824 Arc::new(phys_expr::BinaryExpr::new(left, op, right))
1825}
1826
1827fn and_expr(
1828 left: Arc<dyn PhysicalExpr>,
1829 right: Arc<dyn PhysicalExpr>,
1830) -> Arc<dyn PhysicalExpr> {
1831 binary_expr(left, Operator::And, right)
1832}
1833
1834fn or_expr(
1835 left: Arc<dyn PhysicalExpr>,
1836 right: Arc<dyn PhysicalExpr>,
1837) -> Arc<dyn PhysicalExpr> {
1838 binary_expr(left, Operator::Or, right)
1839}
1840
1841fn build_eq_statistics_expr(
1842 expr_builder: &mut PruningExpressionBuilder,
1843) -> Result<Arc<dyn PhysicalExpr>> {
1844 let min_column_expr = expr_builder.min_column_expr()?;
1845 let max_column_expr = expr_builder.max_column_expr()?;
1846 Ok(and_expr(
1847 binary_expr(
1848 min_column_expr,
1849 Operator::LtEq,
1850 Arc::clone(expr_builder.scalar_expr()),
1851 ),
1852 binary_expr(
1853 Arc::clone(expr_builder.scalar_expr()),
1854 Operator::LtEq,
1855 max_column_expr,
1856 ),
1857 ))
1858}
1859
1860fn build_ne_statistics_expr(
1861 expr_builder: &mut PruningExpressionBuilder,
1862) -> Result<Arc<dyn PhysicalExpr>> {
1863 let min_column_expr = expr_builder.min_column_expr()?;
1864 let max_column_expr = expr_builder.max_column_expr()?;
1865 Ok(or_expr(
1866 binary_expr(
1867 min_column_expr,
1868 Operator::NotEq,
1869 Arc::clone(expr_builder.scalar_expr()),
1870 ),
1871 binary_expr(
1872 Arc::clone(expr_builder.scalar_expr()),
1873 Operator::NotEq,
1874 max_column_expr,
1875 ),
1876 ))
1877}
1878
1879fn column_has_nulls_expr(
1880 expr_builder: &mut PruningExpressionBuilder,
1881) -> Result<Arc<dyn PhysicalExpr>> {
1882 Ok(binary_expr(
1883 expr_builder.null_count_column_expr()?,
1884 Operator::Gt,
1885 Arc::new(phys_expr::Literal::new(ScalarValue::UInt64(Some(0)))),
1886 ))
1887}
1888
1889fn column_has_non_nulls_expr(
1890 expr_builder: &mut PruningExpressionBuilder,
1891) -> Result<Arc<dyn PhysicalExpr>> {
1892 Ok(binary_expr(
1893 expr_builder.null_count_column_expr()?,
1894 Operator::NotEq,
1895 expr_builder.row_count_column_expr()?,
1896 ))
1897}
1898
1899fn build_is_distinct_from(
1900 expr_builder: &mut PruningExpressionBuilder,
1901) -> Result<Arc<dyn PhysicalExpr>> {
1902 let scalar_expr = Arc::clone(expr_builder.scalar_expr());
1903
1904 Ok(or_expr(
1905 and_expr(
1906 Arc::new(phys_expr::IsNullExpr::new(Arc::clone(&scalar_expr))),
1907 column_has_non_nulls_expr(expr_builder)?,
1908 ),
1909 and_expr(
1910 Arc::new(phys_expr::IsNotNullExpr::new(scalar_expr)),
1911 or_expr(
1912 column_has_nulls_expr(expr_builder)?,
1913 build_ne_statistics_expr(expr_builder)?,
1914 ),
1915 ),
1916 ))
1917}
1918
1919fn build_is_not_distinct_from(
1920 expr_builder: &mut PruningExpressionBuilder,
1921) -> Result<Arc<dyn PhysicalExpr>> {
1922 let scalar_expr = Arc::clone(expr_builder.scalar_expr());
1923
1924 Ok(or_expr(
1925 and_expr(
1926 Arc::new(phys_expr::IsNullExpr::new(Arc::clone(&scalar_expr))),
1927 column_has_nulls_expr(expr_builder)?,
1928 ),
1929 and_expr(
1930 Arc::new(phys_expr::IsNotNullExpr::new(scalar_expr)),
1931 and_expr(
1932 column_has_non_nulls_expr(expr_builder)?,
1933 build_eq_statistics_expr(expr_builder)?,
1934 ),
1935 ),
1936 ))
1937}
1938
1939fn unpack_string(s: &ScalarValue) -> Option<&str> {
1941 s.try_as_str().flatten()
1942}
1943
1944fn extract_string_literal(expr: &Arc<dyn PhysicalExpr>) -> Option<&str> {
1945 if let Some(lit) = expr.downcast_ref::<phys_expr::Literal>() {
1946 let s = unpack_string(lit.value())?;
1947 return Some(s);
1948 }
1949 None
1950}
1951
1952fn string_literal_as(value: String, target_type: &DataType) -> Arc<dyn PhysicalExpr> {
1954 let utf8 = ScalarValue::Utf8(Some(value));
1955 let scalar = try_cast_literal_to_type(&utf8, target_type).unwrap_or(utf8);
1956 Arc::new(phys_expr::Literal::new(scalar))
1957}
1958
1959fn build_like_match(
1963 expr_builder: &mut PruningExpressionBuilder,
1964) -> Option<Arc<dyn PhysicalExpr>> {
1965 let min_column_expr = expr_builder.min_column_expr().ok()?;
1976 let max_column_expr = expr_builder.max_column_expr().ok()?;
1977 let scalar_expr = expr_builder.scalar_expr();
1978 let target_type = expr_builder.field.data_type();
1980 let s = extract_string_literal(scalar_expr)?;
1982 let (decoded_prefix, rest) = split_constant_prefix(s);
1984 let has_wildcard = !rest.is_empty();
1985 if has_wildcard && decoded_prefix.is_empty() {
1986 return None;
1988 }
1989 let (lower_bound, upper_bound) = if has_wildcard {
1990 let incremented_prefix = increment_utf8(&decoded_prefix)?;
1991 let lower_bound_lit = string_literal_as(decoded_prefix, target_type);
1992 let upper_bound_lit = string_literal_as(incremented_prefix, target_type);
1993 (lower_bound_lit, upper_bound_lit)
1994 } else {
1995 let bound = string_literal_as(decoded_prefix, target_type);
1997 (Arc::clone(&bound), bound)
1998 };
1999 let lower_bound_expr = Arc::new(phys_expr::BinaryExpr::new(
2000 lower_bound,
2001 Operator::LtEq,
2002 Arc::clone(&max_column_expr),
2003 ));
2004 let upper_bound_expr = Arc::new(phys_expr::BinaryExpr::new(
2005 Arc::clone(&min_column_expr),
2006 Operator::LtEq,
2007 upper_bound,
2008 ));
2009 let combined = Arc::new(phys_expr::BinaryExpr::new(
2010 upper_bound_expr,
2011 Operator::And,
2012 lower_bound_expr,
2013 ));
2014 Some(combined)
2015}
2016
2017fn build_not_like_match(
2023 expr_builder: &mut PruningExpressionBuilder<'_>,
2024) -> Result<Arc<dyn PhysicalExpr>> {
2025 let min_column_expr = expr_builder.min_column_expr()?;
2028 let max_column_expr = expr_builder.max_column_expr()?;
2029
2030 let scalar_expr = expr_builder.scalar_expr();
2031
2032 let pattern = extract_string_literal(scalar_expr).ok_or_else(|| {
2033 plan_datafusion_err!("cannot extract literal from NOT LIKE expression")
2034 })?;
2035
2036 let (const_prefix, remaining) = split_constant_prefix(pattern);
2037 if const_prefix.is_empty() || remaining != "%" {
2038 return Err(plan_datafusion_err!(
2050 "NOT LIKE expressions only support constant_prefix+wildcard`%`"
2051 ));
2052 }
2053
2054 let min_col_not_like_epxr = Arc::new(phys_expr::LikeExpr::new(
2055 true,
2056 false,
2057 Arc::clone(&min_column_expr),
2058 Arc::clone(scalar_expr),
2059 ));
2060
2061 let max_col_not_like_expr = Arc::new(phys_expr::LikeExpr::new(
2062 true,
2063 false,
2064 Arc::clone(&max_column_expr),
2065 Arc::clone(scalar_expr),
2066 ));
2067
2068 Ok(Arc::new(phys_expr::BinaryExpr::new(
2069 min_col_not_like_epxr,
2070 Operator::Or,
2071 max_col_not_like_expr,
2072 )))
2073}
2074
2075fn split_constant_prefix(pattern: &str) -> (String, &str) {
2077 let mut prefix = String::with_capacity(pattern.len());
2078 let mut iter = pattern.char_indices();
2079 while let Some((idx, c)) = iter.next() {
2080 match c {
2081 '%' | '_' => return (prefix, &pattern[idx..]),
2082 '\\' => match iter.next() {
2083 Some((_, escaped)) => prefix.push(escaped),
2084 None => prefix.push('\\'),
2085 },
2086 _ => prefix.push(c),
2087 }
2088 }
2089 (prefix, "")
2090}
2091
2092fn increment_utf8(data: &str) -> Option<String> {
2100 fn is_valid_unicode(c: char) -> bool {
2102 let cp = c as u32;
2103
2104 if [0xFFFE, 0xFFFF].contains(&cp) || (0xFDD0..=0xFDEF).contains(&cp) {
2106 return false;
2107 }
2108
2109 if cp >= 0x110000 {
2111 return false;
2112 }
2113
2114 true
2115 }
2116
2117 let mut code_points: Vec<char> = data.chars().collect();
2119
2120 for idx in (0..code_points.len()).rev() {
2122 let original = code_points[idx] as u32;
2123
2124 if let Some(next_char) = char::from_u32(original + 1)
2126 && is_valid_unicode(next_char)
2127 {
2128 code_points[idx] = next_char;
2129 code_points.truncate(idx + 1);
2131 return Some(code_points.into_iter().collect());
2132 }
2133 }
2134
2135 None
2136}
2137
2138fn wrap_null_count_check_expr(
2159 statistics_expr: Arc<dyn PhysicalExpr>,
2160 expr_builder: &mut PruningExpressionBuilder,
2161) -> Result<Arc<dyn PhysicalExpr>> {
2162 Ok(and_expr(
2164 column_has_non_nulls_expr(expr_builder)?,
2165 statistics_expr,
2166 ))
2167}
2168
2169#[derive(Debug, Copy, Clone, PartialEq, Eq)]
2170pub(crate) enum StatisticsType {
2171 Min,
2172 Max,
2173 NullCount,
2174 RowCount,
2175}
2176
2177#[cfg(test)]
2178mod tests {
2179 use std::collections::HashMap;
2180 use std::ops::{Not, Rem};
2181
2182 use super::*;
2183 use datafusion_common::test_util::batches_to_string;
2184 use datafusion_expr::{and, col, lit, or};
2185 use datafusion_physical_expr::utils::collect_columns;
2186 use insta::assert_snapshot;
2187
2188 use arrow::array::Decimal128Array;
2189 use arrow::{
2190 array::{BinaryArray, Int32Array, Int64Array, StringArray, UInt64Array},
2191 datatypes::TimeUnit,
2192 };
2193 use datafusion_expr::expr::InList;
2194 use datafusion_expr::{BinaryExpr, Expr, cast, is_null, try_cast};
2195 use datafusion_functions_nested::expr_fn::{array_has, make_array};
2196 use datafusion_physical_expr::expressions::{
2197 self as phys_expr, DynamicFilterPhysicalExpr,
2198 };
2199 use datafusion_physical_expr::planner::logical2physical;
2200 use itertools::Itertools;
2201
2202 #[derive(Debug, Default)]
2203 struct ContainerStats {
2211 min: Option<ArrayRef>,
2212 max: Option<ArrayRef>,
2213 null_counts: Option<ArrayRef>,
2215 row_counts: Option<ArrayRef>,
2216 contained: Vec<(HashSet<ScalarValue>, BooleanArray)>,
2220 }
2221
2222 impl ContainerStats {
2223 fn new() -> Self {
2224 Default::default()
2225 }
2226 fn new_decimal128(
2227 min: impl IntoIterator<Item = Option<i128>>,
2228 max: impl IntoIterator<Item = Option<i128>>,
2229 precision: u8,
2230 scale: i8,
2231 ) -> Self {
2232 Self::new()
2233 .with_min(Arc::new(
2234 min.into_iter()
2235 .collect::<Decimal128Array>()
2236 .with_precision_and_scale(precision, scale)
2237 .unwrap(),
2238 ))
2239 .with_max(Arc::new(
2240 max.into_iter()
2241 .collect::<Decimal128Array>()
2242 .with_precision_and_scale(precision, scale)
2243 .unwrap(),
2244 ))
2245 }
2246
2247 fn new_i64(
2248 min: impl IntoIterator<Item = Option<i64>>,
2249 max: impl IntoIterator<Item = Option<i64>>,
2250 ) -> Self {
2251 Self::new()
2252 .with_min(Arc::new(min.into_iter().collect::<Int64Array>()))
2253 .with_max(Arc::new(max.into_iter().collect::<Int64Array>()))
2254 }
2255
2256 fn new_i32(
2257 min: impl IntoIterator<Item = Option<i32>>,
2258 max: impl IntoIterator<Item = Option<i32>>,
2259 ) -> Self {
2260 Self::new()
2261 .with_min(Arc::new(min.into_iter().collect::<Int32Array>()))
2262 .with_max(Arc::new(max.into_iter().collect::<Int32Array>()))
2263 }
2264
2265 fn new_utf8<'a>(
2266 min: impl IntoIterator<Item = Option<&'a str>>,
2267 max: impl IntoIterator<Item = Option<&'a str>>,
2268 ) -> Self {
2269 Self::new()
2270 .with_min(Arc::new(min.into_iter().collect::<StringArray>()))
2271 .with_max(Arc::new(max.into_iter().collect::<StringArray>()))
2272 }
2273
2274 fn new_bool(
2275 min: impl IntoIterator<Item = Option<bool>>,
2276 max: impl IntoIterator<Item = Option<bool>>,
2277 ) -> Self {
2278 Self::new()
2279 .with_min(Arc::new(min.into_iter().collect::<BooleanArray>()))
2280 .with_max(Arc::new(max.into_iter().collect::<BooleanArray>()))
2281 }
2282
2283 fn min(&self) -> Option<ArrayRef> {
2284 self.min.clone()
2285 }
2286
2287 fn max(&self) -> Option<ArrayRef> {
2288 self.max.clone()
2289 }
2290
2291 fn null_counts(&self) -> Option<ArrayRef> {
2292 self.null_counts.clone()
2293 }
2294
2295 fn row_counts(&self) -> Option<ArrayRef> {
2296 self.row_counts.clone()
2297 }
2298
2299 fn arrays(&self) -> Vec<ArrayRef> {
2301 let contained_arrays = self
2302 .contained
2303 .iter()
2304 .map(|(_values, contained)| Arc::new(contained.clone()) as ArrayRef);
2305
2306 [
2307 self.min.as_ref().cloned(),
2308 self.max.as_ref().cloned(),
2309 self.null_counts.as_ref().cloned(),
2310 self.row_counts.as_ref().cloned(),
2311 ]
2312 .into_iter()
2313 .flatten()
2314 .chain(contained_arrays)
2315 .collect()
2316 }
2317
2318 fn len(&self) -> usize {
2322 self.arrays().iter().map(|a| a.len()).next().unwrap_or(0)
2324 }
2325
2326 fn assert_invariants(&self) {
2328 let mut prev_len = None;
2329
2330 for len in self.arrays().iter().map(|a| a.len()) {
2331 match prev_len {
2333 None => {
2334 prev_len = Some(len);
2335 }
2336 Some(prev_len) => {
2337 assert_eq!(prev_len, len);
2338 }
2339 }
2340 }
2341 }
2342
2343 fn with_min(mut self, min: ArrayRef) -> Self {
2345 self.min = Some(min);
2346 self
2347 }
2348
2349 fn with_max(mut self, max: ArrayRef) -> Self {
2351 self.max = Some(max);
2352 self
2353 }
2354
2355 fn with_null_counts(
2358 mut self,
2359 counts: impl IntoIterator<Item = Option<u64>>,
2360 ) -> Self {
2361 let null_counts: ArrayRef =
2362 Arc::new(counts.into_iter().collect::<UInt64Array>());
2363
2364 self.assert_invariants();
2365 self.null_counts = Some(null_counts);
2366 self
2367 }
2368
2369 fn with_row_counts(
2372 mut self,
2373 counts: impl IntoIterator<Item = Option<u64>>,
2374 ) -> Self {
2375 let row_counts: ArrayRef =
2376 Arc::new(counts.into_iter().collect::<UInt64Array>());
2377
2378 self.assert_invariants();
2379 self.row_counts = Some(row_counts);
2380 self
2381 }
2382
2383 #[allow(clippy::allow_attributes, clippy::mutable_key_type)] pub fn with_contained(
2386 mut self,
2387 values: impl IntoIterator<Item = ScalarValue>,
2388 contained: impl IntoIterator<Item = Option<bool>>,
2389 ) -> Self {
2390 let contained: BooleanArray = contained.into_iter().collect();
2391 let values: HashSet<_> = values.into_iter().collect();
2392
2393 self.contained.push((values, contained));
2394 self.assert_invariants();
2395 self
2396 }
2397
2398 #[allow(clippy::allow_attributes, clippy::mutable_key_type)] fn contained(&self, find_values: &HashSet<ScalarValue>) -> Option<BooleanArray> {
2401 self.contained
2403 .iter()
2404 .find(|(values, _contained)| values == find_values)
2405 .map(|(_values, contained)| contained.clone())
2406 }
2407 }
2408
2409 #[derive(Debug, Default)]
2410 struct TestStatistics {
2411 stats: HashMap<Column, ContainerStats>,
2413 }
2414
2415 impl TestStatistics {
2416 fn new() -> Self {
2417 Self::default()
2418 }
2419
2420 fn with(
2421 mut self,
2422 name: impl Into<String>,
2423 container_stats: ContainerStats,
2424 ) -> Self {
2425 let col = Column::from_name(name.into());
2426 self.stats.insert(col, container_stats);
2427 self
2428 }
2429
2430 fn with_null_counts(
2434 mut self,
2435 name: impl Into<String>,
2436 counts: impl IntoIterator<Item = Option<u64>>,
2437 ) -> Self {
2438 let col = Column::from_name(name.into());
2439
2440 let container_stats = self
2442 .stats
2443 .remove(&col)
2444 .unwrap_or_default()
2445 .with_null_counts(counts);
2446
2447 self.stats.insert(col, container_stats);
2449 self
2450 }
2451
2452 fn with_row_counts(
2456 mut self,
2457 name: impl Into<String>,
2458 counts: impl IntoIterator<Item = Option<u64>>,
2459 ) -> Self {
2460 let col = Column::from_name(name.into());
2461
2462 let container_stats = self
2464 .stats
2465 .remove(&col)
2466 .unwrap_or_default()
2467 .with_row_counts(counts);
2468
2469 self.stats.insert(col, container_stats);
2471 self
2472 }
2473
2474 fn with_contained(
2476 mut self,
2477 name: impl Into<String>,
2478 values: impl IntoIterator<Item = ScalarValue>,
2479 contained: impl IntoIterator<Item = Option<bool>>,
2480 ) -> Self {
2481 let col = Column::from_name(name.into());
2482
2483 let container_stats = self
2485 .stats
2486 .remove(&col)
2487 .unwrap_or_default()
2488 .with_contained(values, contained);
2489
2490 self.stats.insert(col, container_stats);
2492 self
2493 }
2494 }
2495
2496 impl PruningStatistics for TestStatistics {
2497 fn min_values(&self, column: &Column) -> Option<ArrayRef> {
2498 self.stats
2499 .get(column)
2500 .map(|container_stats| container_stats.min())
2501 .unwrap_or(None)
2502 }
2503
2504 fn max_values(&self, column: &Column) -> Option<ArrayRef> {
2505 self.stats
2506 .get(column)
2507 .map(|container_stats| container_stats.max())
2508 .unwrap_or(None)
2509 }
2510
2511 fn num_containers(&self) -> usize {
2512 self.stats
2513 .values()
2514 .next()
2515 .map(|container_stats| container_stats.len())
2516 .unwrap_or(0)
2517 }
2518
2519 fn null_counts(&self, column: &Column) -> Option<ArrayRef> {
2520 self.stats
2521 .get(column)
2522 .map(|container_stats| container_stats.null_counts())
2523 .unwrap_or(None)
2524 }
2525
2526 fn row_counts(&self) -> Option<ArrayRef> {
2527 self.stats
2528 .values()
2529 .find_map(|container_stats| container_stats.row_counts())
2530 }
2531
2532 fn contained(
2533 &self,
2534 column: &Column,
2535 values: &HashSet<ScalarValue>,
2536 ) -> Option<BooleanArray> {
2537 self.stats
2538 .get(column)
2539 .and_then(|container_stats| container_stats.contained(values))
2540 }
2541 }
2542
2543 struct OneContainerStats {
2545 min_values: Option<ArrayRef>,
2546 max_values: Option<ArrayRef>,
2547 num_containers: usize,
2548 }
2549
2550 impl PruningStatistics for OneContainerStats {
2551 fn min_values(&self, _column: &Column) -> Option<ArrayRef> {
2552 self.min_values.clone()
2553 }
2554
2555 fn max_values(&self, _column: &Column) -> Option<ArrayRef> {
2556 self.max_values.clone()
2557 }
2558
2559 fn num_containers(&self) -> usize {
2560 self.num_containers
2561 }
2562
2563 fn null_counts(&self, _column: &Column) -> Option<ArrayRef> {
2564 None
2565 }
2566
2567 fn row_counts(&self) -> Option<ArrayRef> {
2568 None
2569 }
2570
2571 fn contained(
2572 &self,
2573 _column: &Column,
2574 _values: &HashSet<ScalarValue>,
2575 ) -> Option<BooleanArray> {
2576 None
2577 }
2578 }
2579
2580 #[test]
2583 fn test_unique_row_count_field_and_column() {
2584 let schema: SchemaRef = Arc::new(Schema::new(vec![
2586 Field::new("c1", DataType::Int32, true),
2587 Field::new("c2", DataType::Int32, true),
2588 ]));
2589 let expr = col("c1").eq(lit(100)).and(col("c2").eq(lit(200)));
2590 let expr = logical2physical(&expr, &schema);
2591 let p = PruningPredicateBuilder::new()
2592 .with_file_schema(Arc::clone(&schema))
2593 .try_build(expr)
2594 .unwrap();
2595 assert_eq!(
2597 "c1_null_count@2 != row_count@3 AND c1_min@0 <= 100 AND 100 <= c1_max@1 AND c2_null_count@6 != row_count@3 AND c2_min@4 <= 200 AND 200 <= c2_max@5",
2598 p.predicate_expr.to_string()
2599 );
2600
2601 let mut fields = HashSet::new();
2604 for (_col, _ty, field) in p.required_columns().iter() {
2605 let was_new = fields.insert(field);
2606 if !was_new {
2607 panic!(
2608 "Duplicate field in required schema: {field:?}. Previous fields:\n{fields:#?}"
2609 );
2610 }
2611 }
2612 }
2613
2614 #[test]
2615 fn prune_all_rows_null_counts() {
2616 let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
2619 let statistics = TestStatistics::new().with(
2620 "i",
2621 ContainerStats::new_i32(
2622 vec![Some(0)], vec![Some(0)], )
2625 .with_null_counts(vec![Some(1)])
2626 .with_row_counts(vec![Some(1)]),
2627 );
2628 let expected_ret = &[false];
2629 prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret);
2630
2631 let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
2633 let container_stats = ContainerStats {
2634 min: Some(Arc::new(Int32Array::from(vec![None]))),
2635 max: Some(Arc::new(Int32Array::from(vec![None]))),
2636 null_counts: Some(Arc::new(UInt64Array::from(vec![Some(1)]))),
2637 row_counts: Some(Arc::new(UInt64Array::from(vec![Some(1)]))),
2638 ..ContainerStats::default()
2639 };
2640 let statistics = TestStatistics::new().with("i", container_stats);
2641 let expected_ret = &[false];
2642 prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret);
2643
2644 let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
2646 let container_stats = ContainerStats {
2647 min: Some(Arc::new(Int32Array::from(vec![Some(0)]))),
2648 max: Some(Arc::new(Int32Array::from(vec![Some(0)]))),
2649 null_counts: Some(Arc::new(UInt64Array::from(vec![None]))),
2650 row_counts: Some(Arc::new(UInt64Array::from(vec![Some(1)]))),
2651 ..ContainerStats::default()
2652 };
2653 let statistics = TestStatistics::new().with("i", container_stats);
2654 let expected_ret = &[true];
2655 prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret);
2656 let expected_ret = &[false];
2657 prune_with_expr(col("i").gt(lit(0)), &schema, &statistics, expected_ret);
2658
2659 let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
2661 let container_stats = ContainerStats {
2662 min: Some(Arc::new(Int32Array::from(vec![Some(0)]))),
2663 max: Some(Arc::new(Int32Array::from(vec![Some(0)]))),
2664 null_counts: Some(Arc::new(UInt64Array::from(vec![Some(1)]))),
2665 row_counts: Some(Arc::new(UInt64Array::from(vec![None]))),
2666 ..ContainerStats::default()
2667 };
2668 let statistics = TestStatistics::new().with("i", container_stats);
2669 let expected_ret = &[true];
2670 prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret);
2671 let expected_ret = &[false];
2672 prune_with_expr(col("i").gt(lit(0)), &schema, &statistics, expected_ret);
2673 }
2674
2675 #[test]
2676 fn prune_missing_statistics() {
2677 let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
2680 let container_stats = ContainerStats {
2681 min: Some(Arc::new(Int32Array::from(vec![None, Some(0)]))),
2682 max: Some(Arc::new(Int32Array::from(vec![Some(0), None]))),
2683 null_counts: Some(Arc::new(UInt64Array::from(vec![Some(0), Some(0)]))),
2684 row_counts: Some(Arc::new(UInt64Array::from(vec![Some(1), Some(1)]))),
2685 ..ContainerStats::default()
2686 };
2687 let statistics = TestStatistics::new().with("i", container_stats);
2688 let expected_ret = &[true, true];
2689 prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret);
2690 let expected_ret = &[false, true];
2691 prune_with_expr(col("i").gt(lit(0)), &schema, &statistics, expected_ret);
2692 let expected_ret = &[true, false];
2693 prune_with_expr(col("i").lt(lit(0)), &schema, &statistics, expected_ret);
2694 }
2695
2696 #[test]
2697 fn prune_null_stats() {
2698 let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
2701
2702 let statistics = TestStatistics::new().with(
2703 "i",
2704 ContainerStats::new_i32(
2705 vec![Some(0)], vec![Some(0)], )
2708 .with_null_counts(vec![Some(1)])
2709 .with_row_counts(vec![Some(1)]),
2710 );
2711
2712 let expected_ret = &[false];
2713
2714 prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret);
2716 }
2717
2718 #[test]
2719 fn test_build_statistics_record_batch() {
2720 let required_columns = RequiredColumns::from(vec![
2722 (
2724 phys_expr::Column::new("s1", 1),
2725 StatisticsType::Min,
2726 Field::new("s1_min", DataType::Int32, true),
2727 ),
2728 (
2730 phys_expr::Column::new("s2", 2),
2731 StatisticsType::Max,
2732 Field::new("s2_max", DataType::Int32, true),
2733 ),
2734 (
2736 phys_expr::Column::new("s3", 3),
2737 StatisticsType::Max,
2738 Field::new("s3_max", DataType::Utf8, true),
2739 ),
2740 (
2742 phys_expr::Column::new("s3", 3),
2743 StatisticsType::Min,
2744 Field::new("s3_min", DataType::Utf8, true),
2745 ),
2746 ]);
2747
2748 let statistics = TestStatistics::new()
2749 .with(
2750 "s1",
2751 ContainerStats::new_i32(
2752 vec![None, None, Some(9), None], vec![Some(10), None, None, None], ),
2755 )
2756 .with(
2757 "s2",
2758 ContainerStats::new_i32(
2759 vec![Some(2), None, None, None], vec![Some(20), None, None, None], ),
2762 )
2763 .with(
2764 "s3",
2765 ContainerStats::new_utf8(
2766 vec![Some("a"), None, None, None], vec![Some("q"), None, Some("r"), None], ),
2769 );
2770
2771 let batch =
2772 build_statistics_record_batch(&statistics, &required_columns).unwrap();
2773 assert_snapshot!(batches_to_string(&[batch]), @r"
2774 +--------+--------+--------+--------+
2775 | s1_min | s2_max | s3_max | s3_min |
2776 +--------+--------+--------+--------+
2777 | | 20 | q | a |
2778 | | | | |
2779 | 9 | | r | |
2780 | | | | |
2781 +--------+--------+--------+--------+
2782 ");
2783 }
2784
2785 #[test]
2786 fn test_build_statistics_casting() {
2787 let required_columns = RequiredColumns::from(vec![(
2792 phys_expr::Column::new("s3", 3),
2793 StatisticsType::Min,
2794 Field::new(
2795 "s1_min",
2796 DataType::Timestamp(TimeUnit::Nanosecond, None),
2797 true,
2798 ),
2799 )]);
2800
2801 let statistics = OneContainerStats {
2803 min_values: Some(Arc::new(Int64Array::from(vec![Some(10)]))),
2804 max_values: Some(Arc::new(Int64Array::from(vec![Some(20)]))),
2805 num_containers: 1,
2806 };
2807
2808 let batch =
2809 build_statistics_record_batch(&statistics, &required_columns).unwrap();
2810
2811 assert_snapshot!(batches_to_string(&[batch]), @r"
2812 +-------------------------------+
2813 | s1_min |
2814 +-------------------------------+
2815 | 1970-01-01T00:00:00.000000010 |
2816 +-------------------------------+
2817 ");
2818 }
2819
2820 #[test]
2821 fn test_build_statistics_no_required_stats() {
2822 let required_columns = RequiredColumns::new();
2823
2824 let statistics = OneContainerStats {
2825 min_values: Some(Arc::new(Int64Array::from(vec![Some(10)]))),
2826 max_values: Some(Arc::new(Int64Array::from(vec![Some(20)]))),
2827 num_containers: 1,
2828 };
2829
2830 let batch =
2831 build_statistics_record_batch(&statistics, &required_columns).unwrap();
2832 assert_eq!(batch.num_rows(), 1); }
2834
2835 #[test]
2836 fn test_build_statistics_inconsistent_types() {
2837 let required_columns = RequiredColumns::from(vec![(
2841 phys_expr::Column::new("s3", 3),
2842 StatisticsType::Min,
2843 Field::new("s1_min", DataType::Utf8, true),
2844 )]);
2845
2846 let statistics = OneContainerStats {
2848 min_values: Some(Arc::new(BinaryArray::from(vec![&[255u8] as &[u8]]))),
2849 max_values: None,
2850 num_containers: 1,
2851 };
2852
2853 let batch =
2854 build_statistics_record_batch(&statistics, &required_columns).unwrap();
2855 assert_snapshot!(batches_to_string(&[batch]), @r"
2856 +--------+
2857 | s1_min |
2858 +--------+
2859 | |
2860 +--------+
2861 ");
2862 }
2863
2864 #[test]
2865 fn test_build_statistics_inconsistent_length() {
2866 let required_columns = RequiredColumns::from(vec![(
2868 phys_expr::Column::new("s1", 3),
2869 StatisticsType::Min,
2870 Field::new("s1_min", DataType::Int64, true),
2871 )]);
2872
2873 let statistics = OneContainerStats {
2875 min_values: Some(Arc::new(Int64Array::from(vec![Some(10)]))),
2876 max_values: Some(Arc::new(Int64Array::from(vec![Some(20)]))),
2877 num_containers: 3,
2878 };
2879
2880 let result =
2881 build_statistics_record_batch(&statistics, &required_columns).unwrap_err();
2882 assert!(
2883 result
2884 .to_string()
2885 .contains("mismatched statistics length. Expected 3, got 1"),
2886 "{}",
2887 result
2888 );
2889 }
2890
2891 #[test]
2892 fn row_group_predicate_eq() -> Result<()> {
2893 let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
2894 let expected_expr =
2895 "c1_null_count@2 != row_count@3 AND c1_min@0 <= 1 AND 1 <= c1_max@1";
2896
2897 let expr = col("c1").eq(lit(1));
2899 let predicate_expr =
2900 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2901 assert_eq!(predicate_expr.to_string(), expected_expr);
2902
2903 let expr = lit(1).eq(col("c1"));
2905 let predicate_expr =
2906 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2907 assert_eq!(predicate_expr.to_string(), expected_expr);
2908
2909 Ok(())
2910 }
2911
2912 #[test]
2913 fn row_group_predicate_not_eq() -> Result<()> {
2914 let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
2915 let expected_expr =
2916 "c1_null_count@2 != row_count@3 AND (c1_min@0 != 1 OR 1 != c1_max@1)";
2917
2918 let expr = col("c1").not_eq(lit(1));
2920 let predicate_expr =
2921 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2922 assert_eq!(predicate_expr.to_string(), expected_expr);
2923
2924 let expr = lit(1).not_eq(col("c1"));
2926 let predicate_expr =
2927 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2928 assert_eq!(predicate_expr.to_string(), expected_expr);
2929
2930 Ok(())
2931 }
2932
2933 #[test]
2934 fn row_group_predicate_gt() -> Result<()> {
2935 let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
2936 let expected_expr = "c1_null_count@1 != row_count@2 AND c1_max@0 > 1";
2937
2938 let expr = col("c1").gt(lit(1));
2940 let predicate_expr =
2941 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2942 assert_eq!(predicate_expr.to_string(), expected_expr);
2943
2944 let expr = lit(1).lt(col("c1"));
2946 let predicate_expr =
2947 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2948 assert_eq!(predicate_expr.to_string(), expected_expr);
2949
2950 Ok(())
2951 }
2952
2953 #[test]
2954 fn row_group_predicate_gt_eq() -> Result<()> {
2955 let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
2956 let expected_expr = "c1_null_count@1 != row_count@2 AND c1_max@0 >= 1";
2957
2958 let expr = col("c1").gt_eq(lit(1));
2960 let predicate_expr =
2961 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2962 assert_eq!(predicate_expr.to_string(), expected_expr);
2963 let expr = lit(1).lt_eq(col("c1"));
2965 let predicate_expr =
2966 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2967 assert_eq!(predicate_expr.to_string(), expected_expr);
2968
2969 Ok(())
2970 }
2971
2972 #[test]
2973 fn row_group_predicate_lt() -> Result<()> {
2974 let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
2975 let expected_expr = "c1_null_count@1 != row_count@2 AND c1_min@0 < 1";
2976
2977 let expr = col("c1").lt(lit(1));
2979 let predicate_expr =
2980 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2981 assert_eq!(predicate_expr.to_string(), expected_expr);
2982
2983 let expr = lit(1).gt(col("c1"));
2985 let predicate_expr =
2986 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
2987 assert_eq!(predicate_expr.to_string(), expected_expr);
2988
2989 Ok(())
2990 }
2991
2992 #[test]
2993 fn row_group_predicate_lt_eq() -> Result<()> {
2994 let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
2995 let expected_expr = "c1_null_count@1 != row_count@2 AND c1_min@0 <= 1";
2996
2997 let expr = col("c1").lt_eq(lit(1));
2999 let predicate_expr =
3000 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3001 assert_eq!(predicate_expr.to_string(), expected_expr);
3002 let expr = lit(1).gt_eq(col("c1"));
3004 let predicate_expr =
3005 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3006 assert_eq!(predicate_expr.to_string(), expected_expr);
3007
3008 Ok(())
3009 }
3010
3011 #[test]
3012 fn row_group_predicate_and() -> Result<()> {
3013 let schema = Schema::new(vec![
3014 Field::new("c1", DataType::Int32, false),
3015 Field::new("c2", DataType::Int32, false),
3016 Field::new("c3", DataType::Int32, false),
3017 ]);
3018 let expr = col("c1").lt(lit(1)).and(col("c2").lt(col("c3")));
3020 let expected_expr = "c1_null_count@1 != row_count@2 AND c1_min@0 < 1";
3021 let predicate_expr =
3022 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3023 assert_eq!(predicate_expr.to_string(), expected_expr);
3024
3025 Ok(())
3026 }
3027
3028 #[test]
3029 fn row_group_predicate_or() -> Result<()> {
3030 let schema = Schema::new(vec![
3031 Field::new("c1", DataType::Int32, false),
3032 Field::new("c2", DataType::Int32, false),
3033 ]);
3034 let expr = col("c1").lt(lit(1)).or(col("c2").rem(lit(2)).eq(lit(0)));
3036 let expected_expr = "true";
3037 let predicate_expr =
3038 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3039 assert_eq!(predicate_expr.to_string(), expected_expr);
3040
3041 Ok(())
3042 }
3043
3044 #[test]
3045 fn row_group_predicate_not() -> Result<()> {
3046 let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
3047 let expected_expr = "true";
3048
3049 let expr = col("c1").not();
3050 let predicate_expr =
3051 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3052 assert_eq!(predicate_expr.to_string(), expected_expr);
3053
3054 Ok(())
3055 }
3056
3057 #[test]
3058 fn row_group_predicate_not_bool() -> Result<()> {
3059 let schema = Schema::new(vec![Field::new("c1", DataType::Boolean, false)]);
3060 let expected_expr = "NOT c1_min@0 AND c1_max@1";
3061
3062 let expr = col("c1").not();
3063 let predicate_expr =
3064 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3065 assert_eq!(predicate_expr.to_string(), expected_expr);
3066
3067 Ok(())
3068 }
3069
3070 #[test]
3071 fn row_group_predicate_bool() -> Result<()> {
3072 let schema = Schema::new(vec![Field::new("c1", DataType::Boolean, false)]);
3073 let expected_expr = "c1_min@0 OR c1_max@1";
3074
3075 let expr = col("c1");
3076 let predicate_expr =
3077 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3078 assert_eq!(predicate_expr.to_string(), expected_expr);
3079
3080 Ok(())
3081 }
3082
3083 #[test]
3085 fn row_group_predicate_non_boolean() {
3086 let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)]));
3087 let statistics = TestStatistics::new()
3088 .with("c1", ContainerStats::new_i32(vec![Some(0)], vec![Some(10)]));
3089 let expected_ret = &[true];
3090 prune_with_expr(lit(1), &schema, &statistics, expected_ret);
3091 }
3092
3093 #[test]
3097 fn row_group_predicate_literal_false() {
3098 let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)]));
3100 let statistics = TestStatistics::new()
3101 .with("c1", ContainerStats::new_i32(vec![Some(0)], vec![Some(10)]));
3102 let expected_ret = &[false];
3103 prune_with_simplified_expr(lit(1).eq(lit(2)), &schema, &statistics, expected_ret);
3104 }
3105
3106 #[test]
3109 fn row_group_predicate_literal_true() {
3110 let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)]));
3112 let statistics = TestStatistics::new()
3113 .with("c1", ContainerStats::new_i32(vec![Some(0)], vec![Some(10)]));
3114 let expected_ret = &[true];
3115 prune_with_simplified_expr(lit(1).eq(lit(1)), &schema, &statistics, expected_ret);
3116 }
3117
3118 #[test]
3121 fn row_group_predicate_literal_null() {
3122 let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)]));
3124 let statistics = TestStatistics::new()
3125 .with("c1", ContainerStats::new_i32(vec![Some(0)], vec![Some(10)]));
3126 let expected_ret = &[true];
3127 prune_with_simplified_expr(
3128 lit(1).eq(lit(ScalarValue::Null)),
3129 &schema,
3130 &statistics,
3131 expected_ret,
3132 );
3133 }
3134
3135 #[test]
3138 fn row_group_predicate_complex_literals() {
3139 let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)]));
3140 let statistics = TestStatistics::new()
3141 .with("c1", ContainerStats::new_i32(vec![Some(0)], vec![Some(10)]));
3142
3143 prune_with_simplified_expr(
3145 (lit(1) + lit(2)).gt(lit(0)),
3146 &schema,
3147 &statistics,
3148 &[true],
3149 );
3150
3151 prune_with_simplified_expr(
3153 (lit(1) + lit(2)).lt(lit(0)),
3154 &schema,
3155 &statistics,
3156 &[false],
3157 );
3158
3159 prune_with_simplified_expr(
3161 lit(true).and(lit(false)),
3162 &schema,
3163 &statistics,
3164 &[false],
3165 );
3166
3167 prune_with_simplified_expr(
3169 lit(true).or(lit(false)),
3170 &schema,
3171 &statistics,
3172 &[true],
3173 );
3174
3175 prune_with_simplified_expr(
3177 lit(1).lt(lit(2)).and(lit(3).gt(lit(1))),
3178 &schema,
3179 &statistics,
3180 &[true],
3181 );
3182
3183 prune_with_simplified_expr(
3185 lit(1).gt(lit(2)).or(lit(3).lt(lit(1))),
3186 &schema,
3187 &statistics,
3188 &[false],
3189 );
3190 }
3191
3192 #[test]
3194 fn row_group_predicate_dynamic_filter_with_literals() {
3195 let schema = Arc::new(Schema::new(vec![
3196 Field::new("c1", DataType::Int32, true),
3197 Field::new("part", DataType::Utf8, true),
3198 ]));
3199 let statistics = TestStatistics::new()
3200 .with_row_counts("c1", vec![Some(10)]);
3202 let dynamic_filter_expr = col("c1").gt(lit(5)).and(col("part").eq(lit("B")));
3203 let phys_expr = logical2physical(&dynamic_filter_expr, &schema);
3204 let children = collect_columns(&phys_expr)
3205 .iter()
3206 .map(|c| Arc::new(c.clone()) as Arc<dyn PhysicalExpr>)
3207 .collect_vec();
3208 let dynamic_phys_expr =
3209 Arc::new(DynamicFilterPhysicalExpr::new(children, phys_expr))
3210 as Arc<dyn PhysicalExpr>;
3211 let remapped_expr = dynamic_phys_expr
3213 .children()
3214 .into_iter()
3215 .map(|child_expr| {
3216 let Some(col_expr) = child_expr.downcast_ref::<phys_expr::Column>()
3217 else {
3218 return Arc::clone(child_expr);
3219 };
3220 if col_expr.name() == "part" {
3221 Arc::new(phys_expr::Literal::new(ScalarValue::Utf8(Some(
3223 "A".to_string(),
3224 )))) as Arc<dyn PhysicalExpr>
3225 } else {
3226 Arc::clone(child_expr)
3227 }
3228 })
3229 .collect_vec();
3230 let dynamic_filter_expr =
3231 dynamic_phys_expr.with_new_children(remapped_expr).unwrap();
3232 let expected = &[false];
3234 let p = PruningPredicateBuilder::new()
3235 .with_file_schema(Arc::clone(&schema))
3236 .try_build(dynamic_filter_expr)
3237 .unwrap();
3238 let result = p.prune(&statistics).unwrap();
3239 assert_eq!(result, expected);
3240 }
3241
3242 #[test]
3243 fn row_group_predicate_lt_bool() -> Result<()> {
3244 let schema = Schema::new(vec![Field::new("c1", DataType::Boolean, false)]);
3245 let expected_expr = "c1_null_count@1 != row_count@2 AND c1_min@0 < true";
3246
3247 let expr = col("c1").lt(lit(true));
3250 let predicate_expr =
3251 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3252 assert_eq!(predicate_expr.to_string(), expected_expr);
3253
3254 Ok(())
3255 }
3256
3257 #[test]
3258 fn row_group_predicate_required_columns() -> Result<()> {
3259 let schema = Schema::new(vec![
3260 Field::new("c1", DataType::Int32, false),
3261 Field::new("c2", DataType::Int32, false),
3262 ]);
3263 let mut required_columns = RequiredColumns::new();
3264 let expr = col("c1")
3266 .lt(lit(1))
3267 .and(col("c2").eq(lit(2)).or(col("c2").eq(lit(3))));
3268 let expected_expr = "c1_null_count@1 != row_count@2 AND c1_min@0 < 1 AND (c2_null_count@5 != row_count@2 AND c2_min@3 <= 2 AND 2 <= c2_max@4 OR c2_null_count@5 != row_count@2 AND c2_min@3 <= 3 AND 3 <= c2_max@4)";
3269 let predicate_expr =
3270 test_build_predicate_expression(&expr, &schema, &mut required_columns);
3271 assert_eq!(predicate_expr.to_string(), expected_expr);
3272 println!("required_columns: {required_columns:#?}"); let c1_min_field = Field::new("c1_min", DataType::Int32, false);
3275 assert_eq!(
3276 required_columns.columns[0],
3277 (
3278 phys_expr::Column::new("c1", 0),
3279 StatisticsType::Min,
3280 c1_min_field.with_nullable(true) )
3282 );
3283 let c1_null_count_field = Field::new("c1_null_count", DataType::UInt64, false);
3285 assert_eq!(
3286 required_columns.columns[1],
3287 (
3288 phys_expr::Column::new("c1", 0),
3289 StatisticsType::NullCount,
3290 c1_null_count_field.with_nullable(true) )
3292 );
3293 let row_count_field = Field::new("row_count", DataType::UInt64, false);
3295 assert_eq!(
3296 required_columns.columns[2],
3297 (
3298 phys_expr::Column::new("c1", 0),
3299 StatisticsType::RowCount,
3300 row_count_field.with_nullable(true) )
3302 );
3303 let c2_min_field = Field::new("c2_min", DataType::Int32, false);
3305 assert_eq!(
3306 required_columns.columns[3],
3307 (
3308 phys_expr::Column::new("c2", 1),
3309 StatisticsType::Min,
3310 c2_min_field.with_nullable(true) )
3312 );
3313 let c2_max_field = Field::new("c2_max", DataType::Int32, false);
3314 assert_eq!(
3315 required_columns.columns[4],
3316 (
3317 phys_expr::Column::new("c2", 1),
3318 StatisticsType::Max,
3319 c2_max_field.with_nullable(true) )
3321 );
3322 let c2_null_count_field = Field::new("c2_null_count", DataType::UInt64, false);
3324 assert_eq!(
3325 required_columns.columns[5],
3326 (
3327 phys_expr::Column::new("c2", 1),
3328 StatisticsType::NullCount,
3329 c2_null_count_field.with_nullable(true) )
3331 );
3332 let row_count_field = Field::new("row_count", DataType::UInt64, false);
3334 assert_eq!(
3335 required_columns.columns[2],
3336 (
3337 phys_expr::Column::new("c1", 0),
3338 StatisticsType::RowCount,
3339 row_count_field.with_nullable(true) )
3341 );
3342 assert_eq!(required_columns.columns.len(), 6);
3344
3345 Ok(())
3346 }
3347
3348 #[test]
3349 fn row_group_predicate_in_list() -> Result<()> {
3350 let schema = Schema::new(vec![
3351 Field::new("c1", DataType::Int32, false),
3352 Field::new("c2", DataType::Int32, false),
3353 ]);
3354 let expr = Expr::InList(InList::new(
3356 Box::new(col("c1")),
3357 vec![lit(1), lit(2), lit(3)],
3358 false,
3359 ));
3360 let expected_expr = "c1_null_count@2 != row_count@3 AND c1_min@0 <= 1 AND 1 <= c1_max@1 OR c1_null_count@2 != row_count@3 AND c1_min@0 <= 2 AND 2 <= c1_max@1 OR c1_null_count@2 != row_count@3 AND c1_min@0 <= 3 AND 3 <= c1_max@1";
3361 let predicate_expr =
3362 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3363 assert_eq!(predicate_expr.to_string(), expected_expr);
3364
3365 Ok(())
3366 }
3367
3368 #[test]
3369 fn row_group_predicate_in_list_empty() -> Result<()> {
3370 let schema = Schema::new(vec![
3371 Field::new("c1", DataType::Int32, false),
3372 Field::new("c2", DataType::Int32, false),
3373 ]);
3374 let expr = Expr::InList(InList::new(Box::new(col("c1")), vec![], false));
3376 let expected_expr = "true";
3377 let predicate_expr =
3378 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3379 assert_eq!(predicate_expr.to_string(), expected_expr);
3380
3381 Ok(())
3382 }
3383
3384 #[test]
3385 fn row_group_predicate_in_list_negated() -> Result<()> {
3386 let schema = Schema::new(vec![
3387 Field::new("c1", DataType::Int32, false),
3388 Field::new("c2", DataType::Int32, false),
3389 ]);
3390 let expr = Expr::InList(InList::new(
3392 Box::new(col("c1")),
3393 vec![lit(1), lit(2), lit(3)],
3394 true,
3395 ));
3396 let expected_expr = "c1_null_count@2 != row_count@3 AND (c1_min@0 != 1 OR 1 != c1_max@1) AND c1_null_count@2 != row_count@3 AND (c1_min@0 != 2 OR 2 != c1_max@1) AND c1_null_count@2 != row_count@3 AND (c1_min@0 != 3 OR 3 != c1_max@1)";
3397 let predicate_expr =
3398 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3399 assert_eq!(predicate_expr.to_string(), expected_expr);
3400
3401 Ok(())
3402 }
3403
3404 #[test]
3405 fn row_group_predicate_between() -> Result<()> {
3406 let schema = Schema::new(vec![
3407 Field::new("c1", DataType::Int32, false),
3408 Field::new("c2", DataType::Int32, false),
3409 ]);
3410
3411 let expr1 = col("c1").between(lit(1), lit(5));
3413
3414 let expr2 = col("c1").gt_eq(lit(1)).and(col("c1").lt_eq(lit(5)));
3416
3417 let predicate_expr1 =
3418 test_build_predicate_expression(&expr1, &schema, &mut RequiredColumns::new());
3419
3420 let predicate_expr2 =
3421 test_build_predicate_expression(&expr2, &schema, &mut RequiredColumns::new());
3422 assert_eq!(predicate_expr1.to_string(), predicate_expr2.to_string());
3423
3424 Ok(())
3425 }
3426
3427 #[test]
3428 fn row_group_predicate_between_with_in_list() -> Result<()> {
3429 let schema = Schema::new(vec![
3430 Field::new("c1", DataType::Int32, false),
3431 Field::new("c2", DataType::Int32, false),
3432 ]);
3433 let expr1 = col("c1").in_list(vec![lit(1), lit(2)], false);
3435
3436 let expr2 = col("c2").between(lit(4), lit(5));
3438
3439 let expr3 = expr1.and(expr2);
3441
3442 let expected_expr = "(c1_null_count@2 != row_count@3 AND c1_min@0 <= 1 AND 1 <= c1_max@1 OR c1_null_count@2 != row_count@3 AND c1_min@0 <= 2 AND 2 <= c1_max@1) AND c2_null_count@5 != row_count@3 AND c2_max@4 >= 4 AND c2_null_count@5 != row_count@3 AND c2_min@6 <= 5";
3443 let predicate_expr =
3444 test_build_predicate_expression(&expr3, &schema, &mut RequiredColumns::new());
3445 assert_eq!(predicate_expr.to_string(), expected_expr);
3446
3447 Ok(())
3448 }
3449
3450 #[test]
3451 fn row_group_predicate_in_list_to_many_values() -> Result<()> {
3452 let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
3453 let expr = col("c1").in_list((1..=21).map(lit).collect(), false);
3457
3458 let expected_expr = "true";
3459 let predicate_expr =
3460 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3461 assert_eq!(predicate_expr.to_string(), expected_expr);
3462
3463 Ok(())
3464 }
3465
3466 #[test]
3472 fn row_group_predicate_in_list_rewritten_at_raised_cap() -> Result<()> {
3473 let schema =
3474 Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)]));
3475 let expr = col("c1").in_list((1..=25).map(lit).collect(), false);
3477 let physical = logical2physical(&expr, &schema);
3478 let rewriter = PredicateRewriter::new().with_max_in_list_size(32);
3479 let predicate_expr =
3480 rewriter.rewrite_predicate_to_statistics_predicate(&physical, &schema);
3481 assert_ne!(
3485 predicate_expr.to_string(),
3486 "true",
3487 "IN(25) with raised cap must rewrite into a statistics-based predicate, not fall through to `true`"
3488 );
3489 assert!(
3491 predicate_expr.to_string().contains(" <= 1 ")
3492 && predicate_expr.to_string().contains(" <= 25 "),
3493 "rewritten predicate should include per-value bounds for each IN entry, got: {predicate_expr}"
3494 );
3495 Ok(())
3496 }
3497
3498 #[test]
3502 fn row_group_predicate_in_list_disabled_at_zero_cap() -> Result<()> {
3503 let schema =
3504 Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)]));
3505 let expr = col("c1").in_list(vec![lit(1), lit(2), lit(3)], false);
3506 let physical = logical2physical(&expr, &schema);
3507 let rewriter = PredicateRewriter::new().with_max_in_list_size(0);
3508 let predicate_expr =
3509 rewriter.rewrite_predicate_to_statistics_predicate(&physical, &schema);
3510 assert_eq!(
3511 predicate_expr.to_string(),
3512 "true",
3513 "cap=0 must skip IN rewrite even for small lists"
3514 );
3515 Ok(())
3516 }
3517
3518 #[test]
3523 fn pruning_predicate_builder_threads_max_in_list_size() -> Result<()> {
3524 let schema =
3525 Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)]));
3526 let expr = col("c1").in_list((1..=25).map(lit).collect(), false);
3527 let physical = logical2physical(&expr, &schema);
3528
3529 let default_pp = PruningPredicateBuilder::new()
3533 .with_file_schema(Arc::clone(&schema))
3534 .try_build(Arc::clone(&physical))?;
3535 assert_eq!(
3536 default_pp.predicate_expr().to_string(),
3537 "true",
3538 "default cap must fall through to `true` for 25-item IN"
3539 );
3540
3541 let raised_pp = PruningPredicateBuilder::new()
3544 .with_file_schema(Arc::clone(&schema))
3545 .with_max_in_list_size(32)
3546 .try_build(physical)?;
3547 let raised_expr = raised_pp.predicate_expr().to_string();
3548 assert_ne!(
3549 raised_expr, "true",
3550 "raised cap must produce a real statistics predicate for 25-item IN"
3551 );
3552 assert!(
3553 raised_expr.contains(" <= 1 ") && raised_expr.contains(" <= 25 "),
3554 "raised-cap predicate should include per-value bounds, got: {raised_expr}"
3555 );
3556 Ok(())
3557 }
3558
3559 #[test]
3560 #[expect(deprecated)]
3561 fn deprecated_try_new_delegates_to_builder() -> Result<()> {
3562 let schema =
3563 Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)]));
3564 let expr = logical2physical(&col("c1").eq(lit(1)), &schema);
3565
3566 let deprecated =
3567 PruningPredicate::try_new(Arc::clone(&expr), Arc::clone(&schema))?;
3568 let builder = PruningPredicateBuilder::new()
3569 .with_file_schema(schema)
3570 .try_build(expr)?;
3571
3572 assert_eq!(
3573 deprecated.predicate_expr().to_string(),
3574 builder.predicate_expr().to_string()
3575 );
3576 assert_eq!(
3577 deprecated.required_columns().schema(),
3578 builder.required_columns().schema()
3579 );
3580 Ok(())
3581 }
3582
3583 #[test]
3584 fn row_group_predicate_cast_int_int() -> Result<()> {
3585 let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
3586 let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Int64) <= 1 AND 1 <= CAST(c1_max@1 AS Int64)";
3587
3588 let expr = cast(col("c1"), DataType::Int64).eq(lit(ScalarValue::Int64(Some(1))));
3591 let predicate_expr =
3592 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3593 assert_eq!(predicate_expr.to_string(), expected_expr);
3594
3595 let expr = lit(ScalarValue::Int64(Some(1))).eq(cast(col("c1"), DataType::Int64));
3597 let predicate_expr =
3598 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3599 assert_eq!(predicate_expr.to_string(), expected_expr);
3600
3601 let expected_expr =
3602 "c1_null_count@1 != row_count@2 AND TRY_CAST(c1_max@0 AS Int64) > 1";
3603
3604 let expr =
3606 try_cast(col("c1"), DataType::Int64).gt(lit(ScalarValue::Int64(Some(1))));
3607 let predicate_expr =
3608 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3609 assert_eq!(predicate_expr.to_string(), expected_expr);
3610
3611 let expr =
3613 lit(ScalarValue::Int64(Some(1))).lt(try_cast(col("c1"), DataType::Int64));
3614 let predicate_expr =
3615 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3616 assert_eq!(predicate_expr.to_string(), expected_expr);
3617
3618 Ok(())
3619 }
3620
3621 #[test]
3622 fn row_group_predicate_cast_string_string() -> Result<()> {
3623 let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]);
3624 let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Utf8) <= 1 AND 1 <= CAST(c1_max@1 AS Utf8)";
3625
3626 let expr = cast(col("c1"), DataType::Utf8)
3628 .eq(lit(ScalarValue::Utf8(Some("1".to_string()))));
3629 let predicate_expr =
3630 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3631 assert_eq!(predicate_expr.to_string(), expected_expr);
3632
3633 let expr = lit(ScalarValue::Utf8(Some("1".to_string())))
3635 .eq(cast(col("c1"), DataType::Utf8));
3636 let predicate_expr =
3637 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3638 assert_eq!(predicate_expr.to_string(), expected_expr);
3639
3640 Ok(())
3641 }
3642
3643 #[test]
3644 fn row_group_predicate_cast_string_int() -> Result<()> {
3645 let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]);
3646 let expected_expr = "true";
3647
3648 let expr = cast(col("c1"), DataType::Int32).eq(lit(ScalarValue::Int32(Some(1))));
3650 let predicate_expr =
3651 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3652 assert_eq!(predicate_expr.to_string(), expected_expr);
3653
3654 let expr = lit(ScalarValue::Int32(Some(1))).eq(cast(col("c1"), DataType::Int32));
3656 let predicate_expr =
3657 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3658 assert_eq!(predicate_expr.to_string(), expected_expr);
3659
3660 Ok(())
3661 }
3662
3663 #[test]
3664 fn row_group_predicate_cast_int_string() -> Result<()> {
3665 let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
3666 let expected_expr = "true";
3667
3668 let expr = cast(col("c1"), DataType::Utf8)
3670 .eq(lit(ScalarValue::Utf8(Some("1".to_string()))));
3671 let predicate_expr =
3672 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3673 assert_eq!(predicate_expr.to_string(), expected_expr);
3674
3675 let expr = lit(ScalarValue::Utf8(Some("1".to_string())))
3677 .eq(cast(col("c1"), DataType::Utf8));
3678 let predicate_expr =
3679 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3680 assert_eq!(predicate_expr.to_string(), expected_expr);
3681
3682 Ok(())
3683 }
3684
3685 #[test]
3686 fn row_group_predicate_date_date() -> Result<()> {
3687 let schema = Schema::new(vec![Field::new("c1", DataType::Date32, false)]);
3688 let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Date64) <= 1970-01-01 AND 1970-01-01 <= CAST(c1_max@1 AS Date64)";
3689
3690 let expr =
3692 cast(col("c1"), DataType::Date64).eq(lit(ScalarValue::Date64(Some(123))));
3693 let predicate_expr =
3694 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3695 assert_eq!(predicate_expr.to_string(), expected_expr);
3696
3697 let expr =
3699 lit(ScalarValue::Date64(Some(123))).eq(cast(col("c1"), DataType::Date64));
3700 let predicate_expr =
3701 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3702 assert_eq!(predicate_expr.to_string(), expected_expr);
3703
3704 Ok(())
3705 }
3706
3707 #[test]
3708 fn row_group_predicate_dict_string_date() -> Result<()> {
3709 let schema = Schema::new(vec![Field::new("c1", DataType::Date32, false)]);
3711 let expected_expr = "true";
3712
3713 let expr = cast(
3715 col("c1"),
3716 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
3717 )
3718 .eq(lit(ScalarValue::Utf8(Some("2024-01-01".to_string()))));
3719 let predicate_expr =
3720 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3721 assert_eq!(predicate_expr.to_string(), expected_expr);
3722
3723 let expr = lit(ScalarValue::Utf8(Some("2024-01-01".to_string()))).eq(cast(
3725 col("c1"),
3726 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
3727 ));
3728 let predicate_expr =
3729 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3730 assert_eq!(predicate_expr.to_string(), expected_expr);
3731
3732 Ok(())
3733 }
3734
3735 #[test]
3736 fn row_group_predicate_date_dict_string() -> Result<()> {
3737 let schema = Schema::new(vec![Field::new(
3739 "c1",
3740 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
3741 false,
3742 )]);
3743 let expected_expr = "true";
3744
3745 let expr =
3747 cast(col("c1"), DataType::Date32).eq(lit(ScalarValue::Date32(Some(123))));
3748 let predicate_expr =
3749 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3750 assert_eq!(predicate_expr.to_string(), expected_expr);
3751
3752 let expr =
3754 lit(ScalarValue::Date32(Some(123))).eq(cast(col("c1"), DataType::Date32));
3755 let predicate_expr =
3756 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3757 assert_eq!(predicate_expr.to_string(), expected_expr);
3758
3759 Ok(())
3760 }
3761
3762 #[test]
3763 fn row_group_predicate_dict_dict_same_value_type() -> Result<()> {
3764 let schema = Schema::new(vec![Field::new(
3766 "c1",
3767 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
3768 false,
3769 )]);
3770
3771 let expr = col("c1").eq(lit(ScalarValue::Utf8(Some("test".to_string()))));
3773 let predicate_expr =
3774 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3775 let expected_expr =
3776 "c1_null_count@2 != row_count@3 AND c1_min@0 <= test AND test <= c1_max@1";
3777 assert_eq!(predicate_expr.to_string(), expected_expr);
3778
3779 let expr = cast(
3781 col("c1"),
3782 DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
3783 )
3784 .eq(lit(ScalarValue::Utf8(Some("test".to_string()))));
3785 let predicate_expr =
3786 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3787 let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Dictionary(UInt16, Utf8)) <= test AND test <= CAST(c1_max@1 AS Dictionary(UInt16, Utf8))";
3788 assert_eq!(predicate_expr.to_string(), expected_expr);
3789
3790 Ok(())
3791 }
3792
3793 #[test]
3794 fn row_group_predicate_dict_dict_different_value_type() -> Result<()> {
3795 let schema = Schema::new(vec![Field::new(
3797 "c1",
3798 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Int32)),
3799 false,
3800 )]);
3801 let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Int64) <= 123 AND 123 <= CAST(c1_max@1 AS Int64)";
3802
3803 let expr =
3805 cast(col("c1"), DataType::Int64).eq(lit(ScalarValue::Int64(Some(123))));
3806 let predicate_expr =
3807 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3808 assert_eq!(predicate_expr.to_string(), expected_expr);
3809
3810 Ok(())
3811 }
3812
3813 #[test]
3814 fn row_group_predicate_nested_dict() -> Result<()> {
3815 let schema = Schema::new(vec![Field::new(
3817 "c1",
3818 DataType::Dictionary(
3819 Box::new(DataType::UInt8),
3820 Box::new(DataType::Dictionary(
3821 Box::new(DataType::UInt16),
3822 Box::new(DataType::Utf8),
3823 )),
3824 ),
3825 false,
3826 )]);
3827 let expected_expr =
3828 "c1_null_count@2 != row_count@3 AND c1_min@0 <= test AND test <= c1_max@1";
3829
3830 let expr = col("c1").eq(lit(ScalarValue::Utf8(Some("test".to_string()))));
3832 let predicate_expr =
3833 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3834 assert_eq!(predicate_expr.to_string(), expected_expr);
3835
3836 Ok(())
3837 }
3838
3839 #[test]
3840 fn row_group_predicate_dict_date_dict_date() -> Result<()> {
3841 let schema = Schema::new(vec![Field::new(
3843 "c1",
3844 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Date32)),
3845 false,
3846 )]);
3847 let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Dictionary(UInt16, Date64)) <= 1970-01-01 AND 1970-01-01 <= CAST(c1_max@1 AS Dictionary(UInt16, Date64))";
3848
3849 let expr = cast(
3851 col("c1"),
3852 DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Date64)),
3853 )
3854 .eq(lit(ScalarValue::Date64(Some(123))));
3855 let predicate_expr =
3856 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3857 assert_eq!(predicate_expr.to_string(), expected_expr);
3858
3859 Ok(())
3860 }
3861
3862 #[test]
3863 fn row_group_predicate_date_string() -> Result<()> {
3864 let schema = Schema::new(vec![Field::new("c1", DataType::Utf8, false)]);
3865 let expected_expr = "true";
3866
3867 let expr =
3869 cast(col("c1"), DataType::Date32).eq(lit(ScalarValue::Date32(Some(123))));
3870 let predicate_expr =
3871 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3872 assert_eq!(predicate_expr.to_string(), expected_expr);
3873
3874 let expr =
3876 lit(ScalarValue::Date32(Some(123))).eq(cast(col("c1"), DataType::Date32));
3877 let predicate_expr =
3878 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3879 assert_eq!(predicate_expr.to_string(), expected_expr);
3880
3881 Ok(())
3882 }
3883
3884 #[test]
3885 fn row_group_predicate_string_date() -> Result<()> {
3886 let schema = Schema::new(vec![Field::new("c1", DataType::Date32, false)]);
3887 let expected_expr = "true";
3888
3889 let expr = cast(col("c1"), DataType::Utf8)
3891 .eq(lit(ScalarValue::Utf8(Some("2024-01-01".to_string()))));
3892 let predicate_expr =
3893 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3894 assert_eq!(predicate_expr.to_string(), expected_expr);
3895
3896 let expr = lit(ScalarValue::Utf8(Some("2024-01-01".to_string())))
3898 .eq(cast(col("c1"), DataType::Utf8));
3899 let predicate_expr =
3900 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3901 assert_eq!(predicate_expr.to_string(), expected_expr);
3902
3903 Ok(())
3904 }
3905
3906 #[test]
3907 fn row_group_predicate_cast_list() -> Result<()> {
3908 let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
3909 let expr = Expr::InList(InList::new(
3911 Box::new(cast(col("c1"), DataType::Int64)),
3912 vec![
3913 lit(ScalarValue::Int64(Some(1))),
3914 lit(ScalarValue::Int64(Some(2))),
3915 lit(ScalarValue::Int64(Some(3))),
3916 ],
3917 false,
3918 ));
3919 let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Int64) <= 1 AND 1 <= CAST(c1_max@1 AS Int64) OR c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Int64) <= 2 AND 2 <= CAST(c1_max@1 AS Int64) OR c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Int64) <= 3 AND 3 <= CAST(c1_max@1 AS Int64)";
3920 let predicate_expr =
3921 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3922 assert_eq!(predicate_expr.to_string(), expected_expr);
3923
3924 let expr = Expr::InList(InList::new(
3925 Box::new(cast(col("c1"), DataType::Int64)),
3926 vec![
3927 lit(ScalarValue::Int64(Some(1))),
3928 lit(ScalarValue::Int64(Some(2))),
3929 lit(ScalarValue::Int64(Some(3))),
3930 ],
3931 true,
3932 ));
3933 let expected_expr = "c1_null_count@2 != row_count@3 AND (CAST(c1_min@0 AS Int64) != 1 OR 1 != CAST(c1_max@1 AS Int64)) AND c1_null_count@2 != row_count@3 AND (CAST(c1_min@0 AS Int64) != 2 OR 2 != CAST(c1_max@1 AS Int64)) AND c1_null_count@2 != row_count@3 AND (CAST(c1_min@0 AS Int64) != 3 OR 3 != CAST(c1_max@1 AS Int64))";
3934 let predicate_expr =
3935 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
3936 assert_eq!(predicate_expr.to_string(), expected_expr);
3937
3938 Ok(())
3939 }
3940
3941 #[test]
3942 fn prune_decimal_data() {
3943 let schema = Arc::new(Schema::new(vec![Field::new(
3945 "s1",
3946 DataType::Decimal128(9, 2),
3947 true,
3948 )]));
3949
3950 prune_with_expr(
3951 col("s1").gt(lit(ScalarValue::Decimal128(Some(500), 9, 2))),
3953 &schema,
3954 &TestStatistics::new().with(
3957 "s1",
3958 ContainerStats::new_i32(
3959 vec![Some(0), Some(4), None, Some(3)], vec![Some(5), Some(6), Some(4), None], ),
3962 ),
3963 &[false, true, false, true],
3964 );
3965
3966 prune_with_expr(
3967 cast(col("s1"), DataType::Decimal128(14, 3))
3969 .gt(lit(ScalarValue::Decimal128(Some(5000), 14, 3))),
3970 &schema,
3971 &TestStatistics::new().with(
3972 "s1",
3973 ContainerStats::new_i32(
3974 vec![Some(0), Some(4), None, Some(3)], vec![Some(5), Some(6), Some(4), None], ),
3977 ),
3978 &[false, true, false, true],
3979 );
3980
3981 prune_with_expr(
3982 try_cast(col("s1"), DataType::Decimal128(14, 3))
3984 .gt(lit(ScalarValue::Decimal128(Some(5000), 14, 3))),
3985 &schema,
3986 &TestStatistics::new().with(
3987 "s1",
3988 ContainerStats::new_i32(
3989 vec![Some(0), Some(4), None, Some(3)], vec![Some(5), Some(6), Some(4), None], ),
3992 ),
3993 &[false, true, false, true],
3994 );
3995
3996 let schema = Arc::new(Schema::new(vec![Field::new(
3998 "s1",
3999 DataType::Decimal128(18, 2),
4000 true,
4001 )]));
4002 prune_with_expr(
4003 col("s1").gt(lit(ScalarValue::Decimal128(Some(500), 18, 2))),
4005 &schema,
4006 &TestStatistics::new().with(
4009 "s1",
4010 ContainerStats::new_i64(
4011 vec![Some(0), Some(4), None, Some(3)], vec![Some(5), Some(6), Some(4), None], ),
4014 ),
4015 &[false, true, false, true],
4016 );
4017
4018 let schema = Arc::new(Schema::new(vec![Field::new(
4020 "s1",
4021 DataType::Decimal128(23, 2),
4022 true,
4023 )]));
4024
4025 prune_with_expr(
4026 col("s1").gt(lit(ScalarValue::Decimal128(Some(500), 23, 2))),
4028 &schema,
4029 &TestStatistics::new().with(
4030 "s1",
4031 ContainerStats::new_decimal128(
4032 vec![Some(0), Some(400), None, Some(300)], vec![Some(500), Some(600), Some(400), None], 23,
4035 2,
4036 ),
4037 ),
4038 &[false, true, false, true],
4039 );
4040 }
4041
4042 #[test]
4043 fn prune_api() {
4044 let schema = Arc::new(Schema::new(vec![
4045 Field::new("s1", DataType::Utf8, true),
4046 Field::new("s2", DataType::Int32, true),
4047 ]));
4048
4049 let statistics = TestStatistics::new().with(
4050 "s2",
4051 ContainerStats::new_i32(
4052 vec![Some(0), Some(4), None, Some(3)], vec![Some(5), Some(6), None, None], ),
4055 );
4056 prune_with_expr(
4057 col("s2").gt(lit(5)),
4059 &schema,
4060 &statistics,
4061 &[false, true, true, true],
4066 );
4067
4068 prune_with_expr(
4069 cast(col("s2"), DataType::Int64).gt(lit(ScalarValue::Int64(Some(5)))),
4071 &schema,
4072 &statistics,
4073 &[false, true, true, true],
4074 );
4075 }
4076
4077 #[test]
4078 fn prune_not_eq_data() {
4079 let schema = Arc::new(Schema::new(vec![Field::new("s1", DataType::Utf8, true)]));
4080
4081 prune_with_expr(
4082 col("s1").not_eq(lit("M")),
4084 &schema,
4085 &TestStatistics::new().with(
4086 "s1",
4087 ContainerStats::new_utf8(
4088 vec![Some("A"), Some("A"), Some("N"), Some("M"), None, Some("A")], vec![Some("Z"), Some("L"), Some("Z"), Some("M"), None, None], ),
4091 ),
4092 &[true, true, true, false, true, true],
4099 );
4100 }
4101
4102 fn bool_setup() -> (SchemaRef, TestStatistics, Vec<bool>, Vec<bool>) {
4118 let schema =
4119 Arc::new(Schema::new(vec![Field::new("b1", DataType::Boolean, true)]));
4120
4121 let statistics = TestStatistics::new().with(
4122 "b1",
4123 ContainerStats::new_bool(
4124 vec![Some(false), Some(false), Some(true), None, Some(false)], vec![Some(false), Some(true), Some(true), None, None], ),
4127 );
4128 let expected_true = vec![false, true, true, true, true];
4129 let expected_false = vec![true, true, false, true, true];
4130
4131 (schema, statistics, expected_true, expected_false)
4132 }
4133
4134 #[test]
4135 fn prune_bool_const_expr() {
4136 let (schema, statistics, _, _) = bool_setup();
4137
4138 prune_with_expr(
4139 lit(true),
4141 &schema,
4142 &statistics,
4143 &[true, true, true, true, true],
4144 );
4145
4146 prune_with_expr(
4147 lit(false),
4149 &schema,
4150 &statistics,
4151 &[false, false, false, false, false],
4152 );
4153 }
4154
4155 #[test]
4156 fn prune_bool_column() {
4157 let (schema, statistics, expected_true, _) = bool_setup();
4158
4159 prune_with_expr(
4160 col("b1"),
4162 &schema,
4163 &statistics,
4164 &expected_true,
4165 );
4166 }
4167
4168 #[test]
4169 fn prune_bool_not_column() {
4170 let (schema, statistics, _, expected_false) = bool_setup();
4171
4172 prune_with_expr(
4173 col("b1").not(),
4175 &schema,
4176 &statistics,
4177 &expected_false,
4178 );
4179 }
4180
4181 #[test]
4182 fn prune_bool_column_eq_true() {
4183 let (schema, statistics, expected_true, _) = bool_setup();
4184
4185 prune_with_expr(
4186 col("b1").eq(lit(true)),
4188 &schema,
4189 &statistics,
4190 &expected_true,
4191 );
4192 }
4193
4194 #[test]
4195 fn prune_bool_not_column_eq_true() {
4196 let (schema, statistics, _, expected_false) = bool_setup();
4197
4198 prune_with_expr(
4199 col("b1").not().eq(lit(true)),
4201 &schema,
4202 &statistics,
4203 &expected_false,
4204 );
4205 }
4206
4207 fn int32_setup() -> (SchemaRef, TestStatistics) {
4217 let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
4218
4219 let statistics = TestStatistics::new().with(
4220 "i",
4221 ContainerStats::new_i32(
4222 vec![Some(-5), Some(1), Some(-11), None, Some(1)], vec![Some(5), Some(11), Some(-1), None, None], ),
4225 );
4226 (schema, statistics)
4227 }
4228
4229 #[test]
4230 fn prune_int32_col_gt_zero() {
4231 let (schema, statistics) = int32_setup();
4232
4233 let expected_ret = &[true, true, false, true, true];
4240
4241 prune_with_expr(col("i").gt(lit(0)), &schema, &statistics, expected_ret);
4243
4244 prune_with_expr(
4246 Expr::Negative(Box::new(col("i"))).lt(lit(0)),
4247 &schema,
4248 &statistics,
4249 expected_ret,
4250 );
4251 }
4252
4253 #[test]
4254 fn prune_int32_col_lte_zero() {
4255 let (schema, statistics) = int32_setup();
4256
4257 let expected_ret = &[true, false, true, true, false];
4264
4265 prune_with_expr(
4266 col("i").lt_eq(lit(0)),
4268 &schema,
4269 &statistics,
4270 expected_ret,
4271 );
4272
4273 prune_with_expr(
4274 Expr::Negative(Box::new(col("i"))).gt_eq(lit(0)),
4276 &schema,
4277 &statistics,
4278 expected_ret,
4279 );
4280 }
4281
4282 #[test]
4283 fn prune_int32_col_lte_zero_cast() {
4284 let (schema, statistics) = int32_setup();
4285
4286 let expected_ret = &[true, true, true, true, true];
4293
4294 prune_with_expr(
4295 cast(col("i"), DataType::Utf8).lt_eq(lit("0")),
4297 &schema,
4298 &statistics,
4299 expected_ret,
4300 );
4301
4302 prune_with_expr(
4303 try_cast(col("i"), DataType::Utf8).lt_eq(lit("0")),
4305 &schema,
4306 &statistics,
4307 expected_ret,
4308 );
4309
4310 prune_with_expr(
4311 cast(Expr::Negative(Box::new(col("i"))), DataType::Utf8).gt_eq(lit("0")),
4313 &schema,
4314 &statistics,
4315 expected_ret,
4316 );
4317
4318 prune_with_expr(
4319 try_cast(Expr::Negative(Box::new(col("i"))), DataType::Utf8).gt_eq(lit("0")),
4321 &schema,
4322 &statistics,
4323 expected_ret,
4324 );
4325 }
4326
4327 #[test]
4328 fn prune_int32_col_eq_zero() {
4329 let (schema, statistics) = int32_setup();
4330
4331 let expected_ret = &[true, false, false, true, false];
4338
4339 prune_with_expr(
4340 col("i").eq(lit(0)),
4342 &schema,
4343 &statistics,
4344 expected_ret,
4345 );
4346 }
4347
4348 #[test]
4349 fn prune_int32_col_is_not_distinct_from() {
4350 let (schema, statistics) = int32_setup();
4351
4352 let expected_ret = &[true, false, false, true, false];
4355
4356 prune_with_expr(
4357 is_not_distinct_from(col("i"), lit(0)),
4358 &schema,
4359 &statistics,
4360 expected_ret,
4361 );
4362
4363 prune_with_expr(
4366 is_not_distinct_from(lit(0), col("i")),
4367 &schema,
4368 &statistics,
4369 expected_ret,
4370 );
4371
4372 let statistics = statistics
4373 .with_row_counts("i", vec![Some(10), Some(9), None, Some(4), Some(10)])
4374 .with_null_counts("i", vec![Some(0), Some(1), None, Some(4), Some(0)]);
4375
4376 let expected_ret = &[true, false, false, false, false];
4377 prune_with_expr(
4378 is_not_distinct_from(col("i"), lit(0)),
4379 &schema,
4380 &statistics,
4381 expected_ret,
4382 );
4383
4384 let expected_ret = &[false, true, true, true, false];
4385 prune_with_expr(
4386 is_not_distinct_from(col("i"), lit(ScalarValue::Int32(None))),
4387 &schema,
4388 &statistics,
4389 expected_ret,
4390 );
4391 }
4392
4393 #[test]
4394 fn prune_int32_col_is_distinct_from() {
4395 let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
4396 let statistics = TestStatistics::new().with(
4397 "i",
4398 ContainerStats::new_i32(
4399 vec![Some(0), Some(0), Some(5), None],
4400 vec![Some(0), Some(2), Some(5), None],
4401 )
4402 .with_row_counts(vec![Some(2), Some(2), Some(2), Some(2)])
4403 .with_null_counts(vec![Some(0), Some(0), Some(0), Some(2)]),
4404 );
4405
4406 let expected_ret = &[false, true, true, true];
4407 prune_with_expr(
4408 is_distinct_from(col("i"), lit(0)),
4409 &schema,
4410 &statistics,
4411 expected_ret,
4412 );
4413
4414 prune_with_expr(
4417 is_distinct_from(lit(0), col("i")),
4418 &schema,
4419 &statistics,
4420 expected_ret,
4421 );
4422
4423 let expected_ret = &[true, true, true, false];
4424 prune_with_expr(
4425 is_distinct_from(col("i"), lit(ScalarValue::Int32(None))),
4426 &schema,
4427 &statistics,
4428 expected_ret,
4429 );
4430 }
4431
4432 #[test]
4433 fn prune_int32_col_eq_zero_cast() {
4434 let (schema, statistics) = int32_setup();
4435
4436 let expected_ret = &[true, false, false, true, false];
4443
4444 prune_with_expr(
4445 cast(col("i"), DataType::Int64).eq(lit(0i64)),
4446 &schema,
4447 &statistics,
4448 expected_ret,
4449 );
4450
4451 prune_with_expr(
4452 try_cast(col("i"), DataType::Int64).eq(lit(0i64)),
4453 &schema,
4454 &statistics,
4455 expected_ret,
4456 );
4457 }
4458
4459 #[test]
4460 fn prune_int32_col_eq_zero_cast_as_str() {
4461 let (schema, statistics) = int32_setup();
4462
4463 let expected_ret = &[true, true, true, true, true];
4473
4474 prune_with_expr(
4475 cast(col("i"), DataType::Utf8).eq(lit("0")),
4476 &schema,
4477 &statistics,
4478 expected_ret,
4479 );
4480 }
4481
4482 #[test]
4483 fn prune_int32_col_lt_neg_one() {
4484 let (schema, statistics) = int32_setup();
4485
4486 let expected_ret = &[true, true, false, true, true];
4493
4494 prune_with_expr(
4495 col("i").gt(lit(-1)),
4497 &schema,
4498 &statistics,
4499 expected_ret,
4500 );
4501
4502 prune_with_expr(
4503 Expr::Negative(Box::new(col("i"))).lt(lit(1)),
4505 &schema,
4506 &statistics,
4507 expected_ret,
4508 );
4509 }
4510
4511 #[test]
4512 fn prune_int32_is_null() {
4513 let (schema, statistics) = int32_setup();
4514
4515 let expected_ret = &[true, true, true, true, true];
4518
4519 prune_with_expr(
4520 col("i").is_null(),
4522 &schema,
4523 &statistics,
4524 expected_ret,
4525 );
4526
4527 let statistics = statistics.with_null_counts(
4529 "i",
4530 vec![
4531 Some(0), Some(1), None, None, Some(0), ],
4537 );
4538
4539 let expected_ret = &[false, true, true, true, false];
4540
4541 prune_with_expr(
4542 col("i").is_null(),
4544 &schema,
4545 &statistics,
4546 expected_ret,
4547 );
4548 }
4549
4550 #[test]
4551 fn prune_int32_column_is_known_all_null() {
4552 let (schema, statistics) = int32_setup();
4553
4554 let expected_ret = &[true, false, true, true, false];
4561
4562 prune_with_expr(
4563 col("i").lt(lit(0)),
4565 &schema,
4566 &statistics,
4567 expected_ret,
4568 );
4569
4570 let statistics = statistics.with_row_counts(
4572 "i",
4573 vec![
4574 Some(10), Some(9), None, Some(4),
4578 Some(10),
4579 ],
4580 );
4581
4582 prune_with_expr(
4584 col("i").lt(lit(0)),
4586 &schema,
4587 &statistics,
4588 expected_ret,
4589 );
4590
4591 let statistics = statistics.with_null_counts(
4593 "i",
4594 vec![
4595 Some(0), Some(1), None, Some(4), Some(0), ],
4601 );
4602
4603 let expected_ret = &[true, false, true, false, false];
4612
4613 prune_with_expr(
4614 col("i").lt(lit(0)),
4616 &schema,
4617 &statistics,
4618 expected_ret,
4619 );
4620 }
4621
4622 #[test]
4623 fn prune_cast_scalar() {
4624 let (schema, statistics) = int32_setup();
4626 let expected_ret = &[true, true, false, true, true];
4627
4628 prune_with_expr(
4629 col("i").gt(cast(lit(ScalarValue::Int64(Some(0))), DataType::Int32)),
4631 &schema,
4632 &statistics,
4633 expected_ret,
4634 );
4635
4636 prune_with_expr(
4637 cast(col("i"), DataType::Int64).gt(lit(ScalarValue::Int64(Some(0)))),
4639 &schema,
4640 &statistics,
4641 expected_ret,
4642 );
4643
4644 prune_with_expr(
4645 try_cast(col("i"), DataType::Int64).gt(lit(ScalarValue::Int64(Some(0)))),
4647 &schema,
4648 &statistics,
4649 expected_ret,
4650 );
4651
4652 prune_with_expr(
4653 Expr::Negative(Box::new(cast(col("i"), DataType::Int64)))
4655 .lt(lit(ScalarValue::Int64(Some(0)))),
4656 &schema,
4657 &statistics,
4658 expected_ret,
4659 );
4660 }
4661
4662 #[test]
4663 fn test_increment_utf8() {
4664 assert_eq!(increment_utf8("abc").unwrap(), "abd");
4666 assert_eq!(increment_utf8("abz").unwrap(), "ab{");
4667
4668 assert_eq!(increment_utf8("~").unwrap(), "\u{7f}"); assert_eq!(increment_utf8("\u{7f}").unwrap(), "\u{80}"); assert_eq!(increment_utf8("ß").unwrap(), "à"); assert_eq!(increment_utf8("℣").unwrap(), "ℤ"); assert_eq!(increment_utf8("\u{7FF}").unwrap(), "\u{800}"); assert_eq!(increment_utf8("\u{FFFF}").unwrap(), "\u{10000}"); assert!(increment_utf8("").is_none());
4684 assert!(increment_utf8("\u{10FFFF}").is_none()); assert_eq!(increment_utf8("a\u{10FFFF}").unwrap(), "b");
4688
4689 assert_eq!(increment_utf8("a\u{D7FF}").unwrap(), "b");
4691 assert!(increment_utf8("\u{D7FF}").is_none());
4692
4693 assert_eq!(increment_utf8("a\u{FDCF}").unwrap(), "b");
4695 assert!(increment_utf8("\u{FDCF}").is_none());
4696
4697 assert_eq!(increment_utf8("a\u{10FFFF}").unwrap(), "b");
4699 assert!(increment_utf8("\u{10FFFF}").is_none()); }
4701
4702 fn utf8_setup() -> (SchemaRef, TestStatistics) {
4715 let schema = Arc::new(Schema::new(vec![Field::new("s1", DataType::Utf8, true)]));
4716
4717 let statistics = TestStatistics::new().with(
4718 "s1",
4719 ContainerStats::new_utf8(
4720 vec![
4721 Some("A"),
4722 Some("A"),
4723 Some("N"),
4724 Some("M"),
4725 None,
4726 Some("A"),
4727 Some(""),
4728 Some(""),
4729 Some("AB"),
4730 Some("A\u{10ffff}\u{10ffff}"),
4731 ], vec![
4733 Some("Z"),
4734 Some("L"),
4735 Some("Z"),
4736 Some("M"),
4737 None,
4738 None,
4739 Some("A"),
4740 Some(""),
4741 Some("A\u{10ffff}\u{10ffff}\u{10ffff}"),
4742 Some("A\u{10ffff}\u{10ffff}"),
4743 ], ),
4745 );
4746 (schema, statistics)
4747 }
4748
4749 #[test]
4750 fn prune_utf8_eq() {
4751 let (schema, statistics) = utf8_setup();
4752
4753 let expr = col("s1").eq(lit("A"));
4754 #[rustfmt::skip]
4755 let expected_ret = &[
4756 true,
4758 true,
4760 false,
4762 false,
4764 true,
4766 true,
4768 true,
4770 false,
4772 false,
4774 false,
4776 ];
4777 prune_with_expr(expr, &schema, &statistics, expected_ret);
4778
4779 let expr = col("s1").eq(lit(""));
4780 #[rustfmt::skip]
4781 let expected_ret = &[
4782 false,
4784 false,
4786 false,
4788 false,
4790 true,
4792 false,
4794 true,
4796 true,
4798 false,
4800 false,
4802 ];
4803 prune_with_expr(expr, &schema, &statistics, expected_ret);
4804 }
4805
4806 #[test]
4807 fn prune_utf8_not_eq() {
4808 let (schema, statistics) = utf8_setup();
4809
4810 let expr = col("s1").not_eq(lit("A"));
4811 #[rustfmt::skip]
4812 let expected_ret = &[
4813 true,
4815 true,
4817 true,
4819 true,
4821 true,
4823 true,
4825 true,
4827 true,
4829 true,
4831 true,
4833 ];
4834 prune_with_expr(expr, &schema, &statistics, expected_ret);
4835
4836 let expr = col("s1").not_eq(lit(""));
4837 #[rustfmt::skip]
4838 let expected_ret = &[
4839 true,
4841 true,
4843 true,
4845 true,
4847 true,
4849 true,
4851 true,
4853 false,
4855 true,
4857 true,
4859 ];
4860 prune_with_expr(expr, &schema, &statistics, expected_ret);
4861 }
4862
4863 #[test]
4864 fn prune_utf8_like_one() {
4865 let (schema, statistics) = utf8_setup();
4866
4867 let expr = col("s1").like(lit("A_"));
4868 #[rustfmt::skip]
4869 let expected_ret = &[
4870 true,
4872 true,
4874 false,
4876 false,
4878 true,
4880 true,
4882 true,
4884 false,
4886 true,
4888 true,
4890 ];
4891 prune_with_expr(expr, &schema, &statistics, expected_ret);
4892
4893 let expr = col("s1").like(lit("_A_"));
4894 #[rustfmt::skip]
4895 let expected_ret = &[
4896 true,
4898 true,
4900 true,
4902 true,
4904 true,
4906 true,
4908 true,
4910 true,
4912 true,
4914 true,
4916 ];
4917 prune_with_expr(expr, &schema, &statistics, expected_ret);
4918
4919 let expr = col("s1").like(lit("_"));
4920 #[rustfmt::skip]
4921 let expected_ret = &[
4922 true,
4924 true,
4926 true,
4928 true,
4930 true,
4932 true,
4934 true,
4936 true,
4938 true,
4940 true,
4942 ];
4943 prune_with_expr(expr, &schema, &statistics, expected_ret);
4944
4945 let expr = col("s1").like(lit(""));
4946 #[rustfmt::skip]
4947 let expected_ret = &[
4948 false,
4950 false,
4952 false,
4954 false,
4956 true,
4958 false,
4960 true,
4962 true,
4964 false,
4966 false,
4968 ];
4969 prune_with_expr(expr, &schema, &statistics, expected_ret);
4970 }
4971
4972 #[test]
4973 fn prune_utf8_like_many() {
4974 let (schema, statistics) = utf8_setup();
4975
4976 let expr = col("s1").like(lit("A%"));
4977 #[rustfmt::skip]
4978 let expected_ret = &[
4979 true,
4981 true,
4983 false,
4985 false,
4987 true,
4989 true,
4991 true,
4993 false,
4995 true,
4997 true,
4999 ];
5000 prune_with_expr(expr, &schema, &statistics, expected_ret);
5001
5002 let expr = col("s1").like(lit("%A%"));
5003 #[rustfmt::skip]
5004 let expected_ret = &[
5005 true,
5007 true,
5009 true,
5011 true,
5013 true,
5015 true,
5017 true,
5019 true,
5021 true,
5023 true,
5025 ];
5026 prune_with_expr(expr, &schema, &statistics, expected_ret);
5027
5028 let expr = col("s1").like(lit("%"));
5029 #[rustfmt::skip]
5030 let expected_ret = &[
5031 true,
5033 true,
5035 true,
5037 true,
5039 true,
5041 true,
5043 true,
5045 true,
5047 true,
5049 true,
5051 ];
5052 prune_with_expr(expr, &schema, &statistics, expected_ret);
5053
5054 let expr = col("s1").like(lit(""));
5055 #[rustfmt::skip]
5056 let expected_ret = &[
5057 false,
5059 false,
5061 false,
5063 false,
5065 true,
5067 false,
5069 true,
5071 true,
5073 false,
5075 false,
5077 ];
5078 prune_with_expr(expr, &schema, &statistics, expected_ret);
5079 }
5080
5081 #[test]
5084 fn prune_utf8_like_escaped_chars() {
5085 let schema = Arc::new(Schema::new(vec![Field::new("s1", DataType::Utf8, true)]));
5086 let statistics = TestStatistics::new().with(
5087 "s1",
5088 ContainerStats::new_utf8(
5089 vec![
5090 Some("foo_aaa"),
5091 Some(r#"foo\aaa"#),
5092 Some("foo"),
5093 Some("bar"),
5094 Some("foo%aaa"),
5095 Some("%foo_aaa"),
5096 ], vec![
5098 Some("foo_zzz"),
5099 Some(r#"foo\zzz"#),
5100 Some("foozzz"),
5101 Some("baz"),
5102 Some("foo%zzz"),
5103 Some("%foo_zzz"),
5104 ], ),
5106 );
5107
5108 let expr = col("s1").like(lit(r#"foo\_%"#));
5109 #[rustfmt::skip]
5110 let expected_ret = &[
5111 true,
5114 false,
5116 true,
5119 false,
5121 false,
5123 false,
5125 ];
5126 prune_with_expr(expr, &schema, &statistics, expected_ret);
5127
5128 let expr = col("s1").like(lit(r#"foo\\%"#));
5129 #[rustfmt::skip]
5130 let expected_ret = &[
5131 false,
5133 true,
5136 true,
5139 false,
5141 false,
5143 false,
5145 ];
5146 prune_with_expr(expr, &schema, &statistics, expected_ret);
5147
5148 let expr = col("s1").like(lit(r#"foo\%%"#));
5149 #[rustfmt::skip]
5150 let expected_ret = &[
5151 false,
5153 false,
5155 true,
5157 false,
5159 true,
5162 false,
5164 ];
5165 prune_with_expr(expr, &schema, &statistics, expected_ret);
5166
5167 let expr = col("s1").like(lit(r#"foo\_"#));
5170 #[rustfmt::skip]
5171 let expected_ret = &[
5172 false,
5174 false,
5176 true,
5178 false,
5180 false,
5182 false,
5184 ];
5185 prune_with_expr(expr, &schema, &statistics, expected_ret);
5186
5187 let expr = col("s1").like(lit(r#"\%foo%"#));
5190 #[rustfmt::skip]
5191 let expected_ret = &[
5192 false,
5194 false,
5196 false,
5198 false,
5200 false,
5202 true,
5205 ];
5206 prune_with_expr(expr, &schema, &statistics, expected_ret);
5207
5208 let expr = col("s1").like(lit(r#"foo\%\_"#));
5210 #[rustfmt::skip]
5211 let expected_ret = &[
5212 false,
5214 false,
5216 true,
5218 false,
5220 false,
5222 false,
5224 ];
5225 prune_with_expr(expr, &schema, &statistics, expected_ret);
5226
5227 let expr = col("s1").like(lit(r#"foo\\bar%"#));
5230 #[rustfmt::skip]
5231 let expected_ret = &[
5232 false,
5234 true,
5237 true,
5239 false,
5241 false,
5243 false,
5245 ];
5246 prune_with_expr(expr, &schema, &statistics, expected_ret);
5247 }
5248
5249 #[test]
5250 fn prune_utf8_not_like_one() {
5251 let (schema, statistics) = utf8_setup();
5252
5253 let expr = col("s1").not_like(lit("A\u{10ffff}_"));
5254 #[rustfmt::skip]
5255 let expected_ret = &[
5256 true,
5258 true,
5260 true,
5262 true,
5264 true,
5266 true,
5268 true,
5270 true,
5272 true,
5274 true,
5277 ];
5278 prune_with_expr(expr, &schema, &statistics, expected_ret);
5279 }
5280
5281 #[test]
5282 fn prune_utf8_not_like_many() {
5283 let (schema, statistics) = utf8_setup();
5284
5285 let expr = col("s1").not_like(lit("A\u{10ffff}%"));
5286 #[rustfmt::skip]
5287 let expected_ret = &[
5288 true,
5290 true,
5292 true,
5294 true,
5296 true,
5298 true,
5300 true,
5302 true,
5304 true,
5306 false,
5308 ];
5309 prune_with_expr(expr, &schema, &statistics, expected_ret);
5310
5311 let expr = col("s1").not_like(lit("A\u{10ffff}%\u{10ffff}"));
5312 #[rustfmt::skip]
5313 let expected_ret = &[
5314 true,
5316 true,
5318 true,
5320 true,
5322 true,
5324 true,
5326 true,
5328 true,
5330 true,
5332 true,
5334 ];
5335 prune_with_expr(expr, &schema, &statistics, expected_ret);
5336
5337 let expr = col("s1").not_like(lit("A\u{10ffff}%\u{10ffff}_"));
5338 #[rustfmt::skip]
5339 let expected_ret = &[
5340 true,
5342 true,
5344 true,
5346 true,
5348 true,
5350 true,
5352 true,
5354 true,
5356 true,
5358 true,
5360 ];
5361 prune_with_expr(expr, &schema, &statistics, expected_ret);
5362
5363 let expr = col("s1").not_like(lit("A\\%%"));
5364 let statistics = TestStatistics::new().with(
5365 "s1",
5366 ContainerStats::new_utf8(
5367 vec![Some("A%a"), Some("A")],
5368 vec![Some("A%c"), Some("A")],
5369 ),
5370 );
5371 let expected_ret = &[false, true];
5372 prune_with_expr(expr, &schema, &statistics, expected_ret);
5373 }
5374
5375 #[test]
5376 fn test_rewrite_expr_to_prunable() {
5377 let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
5378 let df_schema = DFSchema::try_from(schema.clone()).unwrap();
5379
5380 let left_input = col("a");
5382 let left_input = logical2physical(&left_input, &schema);
5383 let right_input = lit(ScalarValue::Int32(Some(12)));
5384 let right_input = logical2physical(&right_input, &schema);
5385 let (result_left, _, result_right) = rewrite_expr_to_prunable(
5386 &left_input,
5387 Operator::Eq,
5388 &right_input,
5389 df_schema.clone(),
5390 )
5391 .unwrap();
5392 assert_eq!(result_left.to_string(), left_input.to_string());
5393 assert_eq!(result_right.to_string(), right_input.to_string());
5394
5395 let left_input = cast(col("a"), DataType::Decimal128(20, 3));
5397 let left_input = logical2physical(&left_input, &schema);
5398 let right_input = lit(ScalarValue::Decimal128(Some(12), 20, 3));
5399 let right_input = logical2physical(&right_input, &schema);
5400 let (result_left, _, result_right) = rewrite_expr_to_prunable(
5401 &left_input,
5402 Operator::Gt,
5403 &right_input,
5404 df_schema.clone(),
5405 )
5406 .unwrap();
5407 assert_eq!(result_left.to_string(), left_input.to_string());
5408 assert_eq!(result_right.to_string(), right_input.to_string());
5409
5410 let left_input = try_cast(col("a"), DataType::Int64);
5412 let left_input = logical2physical(&left_input, &schema);
5413 let right_input = lit(ScalarValue::Int64(Some(12)));
5414 let right_input = logical2physical(&right_input, &schema);
5415 let (result_left, _, result_right) =
5416 rewrite_expr_to_prunable(&left_input, Operator::Gt, &right_input, df_schema)
5417 .unwrap();
5418 assert_eq!(result_left.to_string(), left_input.to_string());
5419 assert_eq!(result_right.to_string(), right_input.to_string());
5420
5421 }
5423
5424 #[test]
5425 fn test_rewrite_expr_to_prunable_custom_unhandled_hook() {
5426 struct CustomUnhandledHook;
5427
5428 impl UnhandledPredicateHook for CustomUnhandledHook {
5429 fn handle(&self, _expr: &Arc<dyn PhysicalExpr>) -> Arc<dyn PhysicalExpr> {
5433 Arc::new(phys_expr::Literal::new(ScalarValue::Int32(Some(42))))
5434 }
5435 }
5436
5437 let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
5438 let schema_with_b = Schema::new(vec![
5439 Field::new("a", DataType::Int32, true),
5440 Field::new("b", DataType::Int32, true),
5441 ]);
5442
5443 let rewriter = PredicateRewriter::new()
5444 .with_unhandled_hook(Arc::new(CustomUnhandledHook {}));
5445
5446 let transform_expr = |expr| {
5447 let expr = logical2physical(&expr, &schema_with_b);
5448 rewriter.rewrite_predicate_to_statistics_predicate(&expr, &schema)
5449 };
5450
5451 let known_expression = col("a").eq(lit(12));
5453 let known_expression_transformed = PredicateRewriter::new()
5454 .rewrite_predicate_to_statistics_predicate(
5455 &logical2physical(&known_expression, &schema),
5456 &schema,
5457 );
5458
5459 let input = col("b").eq(lit(12));
5461 let expected = logical2physical(&lit(42), &schema);
5462 let transformed = transform_expr(input.clone());
5463 assert_eq!(transformed.to_string(), expected.to_string());
5464
5465 let input = known_expression.clone().and(input.clone());
5467 let expected = phys_expr::BinaryExpr::new(
5468 Arc::<dyn PhysicalExpr>::clone(&known_expression_transformed),
5469 Operator::And,
5470 logical2physical(&lit(42), &schema),
5471 );
5472 let transformed = transform_expr(input.clone());
5473 assert_eq!(transformed.to_string(), expected.to_string());
5474
5475 let input = array_has(make_array(vec![lit(1)]), col("a"));
5477 let expected = logical2physical(&lit(42), &schema);
5478 let transformed = transform_expr(input.clone());
5479 assert_eq!(transformed.to_string(), expected.to_string());
5480
5481 let input = known_expression.and(input);
5483 let expected = phys_expr::BinaryExpr::new(
5484 Arc::<dyn PhysicalExpr>::clone(&known_expression_transformed),
5485 Operator::And,
5486 logical2physical(&lit(42), &schema),
5487 );
5488 let transformed = transform_expr(input.clone());
5489 assert_eq!(transformed.to_string(), expected.to_string());
5490 }
5491
5492 #[test]
5493 fn test_rewrite_expr_to_prunable_error() {
5494 let schema = Schema::new(vec![Field::new("a", DataType::Utf8, true)]);
5497 let df_schema = DFSchema::try_from(schema.clone()).unwrap();
5498 let left_input = cast(col("a"), DataType::Int64);
5499 let left_input = logical2physical(&left_input, &schema);
5500 let right_input = lit(ScalarValue::Int64(Some(12)));
5501 let right_input = logical2physical(&right_input, &schema);
5502 let result = rewrite_expr_to_prunable(
5503 &left_input,
5504 Operator::Gt,
5505 &right_input,
5506 df_schema.clone(),
5507 );
5508 assert!(result.is_err());
5509
5510 let left_input = is_null(col("a"));
5512 let left_input = logical2physical(&left_input, &schema);
5513 let right_input = lit(ScalarValue::Int64(Some(12)));
5514 let right_input = logical2physical(&right_input, &schema);
5515 let result =
5516 rewrite_expr_to_prunable(&left_input, Operator::Gt, &right_input, df_schema);
5517 assert!(result.is_err());
5518 }
5520
5521 #[test]
5522 fn prune_with_contained_one_column() {
5523 let schema = Arc::new(Schema::new(vec![Field::new("s1", DataType::Utf8, true)]));
5524
5525 let statistics = TestStatistics::new()
5527 .with_contained(
5528 "s1",
5529 [ScalarValue::from("foo")],
5530 [
5531 Some(true),
5533 Some(false),
5535 None,
5537 Some(true),
5539 Some(false),
5541 None,
5543 Some(true),
5545 Some(false),
5547 None,
5549 ],
5550 )
5551 .with_contained(
5552 "s1",
5553 [ScalarValue::from("bar")],
5554 [
5555 Some(true),
5557 Some(true),
5558 Some(true),
5559 Some(false),
5561 Some(false),
5562 Some(false),
5563 None,
5565 None,
5566 None,
5567 ],
5568 )
5569 .with_contained(
5570 "s1",
5573 [ScalarValue::from("foo"), ScalarValue::from("bar")],
5574 [
5575 None,
5577 None,
5578 None,
5579 Some(true),
5581 Some(true),
5582 Some(true),
5583 Some(false),
5585 Some(false),
5586 Some(false),
5587 ],
5588 );
5589
5590 prune_with_expr(
5592 col("s1").eq(lit("foo")),
5593 &schema,
5594 &statistics,
5595 &[true, false, true, true, false, true, true, false, true],
5597 );
5598
5599 prune_with_expr(
5601 col("s1").eq(lit("bar")),
5602 &schema,
5603 &statistics,
5604 &[true, true, true, false, false, false, true, true, true],
5606 );
5607
5608 prune_with_expr(
5610 col("s1").eq(lit("baz")),
5611 &schema,
5612 &statistics,
5613 &[true, true, true, true, true, true, true, true, true],
5615 );
5616
5617 prune_with_expr(
5619 col("s1").eq(lit("foo")).and(col("s1").eq(lit("bar"))),
5620 &schema,
5621 &statistics,
5622 &[true, true, true, true, true, true, true, true, true],
5626 );
5627
5628 prune_with_expr(
5630 col("s1").eq(lit("foo")).or(col("s1").eq(lit("bar"))),
5631 &schema,
5632 &statistics,
5633 &[true, true, true, true, true, true, false, false, false],
5635 );
5636
5637 prune_with_expr(
5639 col("s1").eq(lit("foo")).or(col("s1").eq(lit("baz"))),
5640 &schema,
5641 &statistics,
5642 &[true, true, true, true, true, true, true, true, true],
5644 );
5645
5646 prune_with_expr(
5648 col("s1")
5649 .eq(lit("foo"))
5650 .or(col("s1").eq(lit("bar")))
5651 .or(col("s1").eq(lit("baz"))),
5652 &schema,
5653 &statistics,
5654 &[true, true, true, true, true, true, true, true, true],
5657 );
5658
5659 prune_with_expr(
5661 col("s1").not_eq(lit("foo")),
5662 &schema,
5663 &statistics,
5664 &[false, true, true, false, true, true, false, true, true],
5666 );
5667
5668 prune_with_expr(
5670 col("s1").not_eq(lit("bar")),
5671 &schema,
5672 &statistics,
5673 &[false, false, false, true, true, true, true, true, true],
5675 );
5676
5677 prune_with_expr(
5679 col("s1")
5680 .not_eq(lit("foo"))
5681 .and(col("s1").not_eq(lit("bar"))),
5682 &schema,
5683 &statistics,
5684 &[true, true, true, false, false, false, true, true, true],
5686 );
5687
5688 prune_with_expr(
5690 col("s1")
5691 .not_eq(lit("foo"))
5692 .and(col("s1").not_eq(lit("bar")))
5693 .and(col("s1").not_eq(lit("baz"))),
5694 &schema,
5695 &statistics,
5696 &[true, true, true, true, true, true, true, true, true],
5698 );
5699
5700 prune_with_expr(
5702 col("s1")
5703 .not_eq(lit("foo"))
5704 .or(col("s1").not_eq(lit("bar"))),
5705 &schema,
5706 &statistics,
5707 &[true, true, true, true, true, true, true, true, true],
5709 );
5710
5711 prune_with_expr(
5713 col("s1")
5714 .not_eq(lit("foo"))
5715 .or(col("s1").not_eq(lit("bar")))
5716 .or(col("s1").not_eq(lit("baz"))),
5717 &schema,
5718 &statistics,
5719 &[true, true, true, true, true, true, true, true, true],
5721 );
5722 }
5723
5724 #[test]
5725 fn prune_with_contained_two_columns() {
5726 let schema = Arc::new(Schema::new(vec![
5727 Field::new("s1", DataType::Utf8, true),
5728 Field::new("s2", DataType::Utf8, true),
5729 ]));
5730
5731 let statistics = TestStatistics::new()
5733 .with_contained(
5734 "s1",
5735 [ScalarValue::from("foo")],
5736 [
5737 Some(true),
5739 Some(false),
5741 None,
5743 Some(true),
5745 Some(false),
5747 None,
5749 Some(true),
5751 Some(false),
5753 None,
5755 ],
5756 )
5757 .with_contained(
5758 "s2", [ScalarValue::from("bar")],
5760 [
5761 Some(true),
5763 Some(true),
5764 Some(true),
5765 Some(false),
5767 Some(false),
5768 Some(false),
5769 None,
5771 None,
5772 None,
5773 ],
5774 );
5775
5776 prune_with_expr(
5778 col("s1").eq(lit("foo")),
5779 &schema,
5780 &statistics,
5781 &[true, false, true, true, false, true, true, false, true],
5783 );
5784
5785 let expr = col("s1").eq(lit("foo")).or(col("s2").eq(lit("bar")));
5787 prune_with_expr(
5788 expr,
5789 &schema,
5790 &statistics,
5791 &[true, true, true, true, true, true, true, true, true],
5793 );
5794
5795 prune_with_expr(
5797 col("s1").eq(lit("foo")).and(col("s2").not_eq(lit("bar"))),
5798 &schema,
5799 &statistics,
5800 &[false, false, false, true, false, true, true, false, true],
5804 );
5805
5806 prune_with_expr(
5808 col("s1")
5809 .not_eq(lit("foo"))
5810 .and(col("s2").not_eq(lit("bar"))),
5811 &schema,
5812 &statistics,
5813 &[false, false, false, false, true, true, false, true, true],
5817 );
5818
5819 prune_with_expr(
5821 col("s1")
5822 .not_eq(lit("foo"))
5823 .and(col("s2").eq(lit("bar")).or(col("s2").eq(lit("baz")))),
5824 &schema,
5825 &statistics,
5826 &[false, true, true, false, true, true, false, true, true],
5829 );
5830
5831 prune_with_expr(
5833 col("s1").like(lit("foo%bar%")),
5834 &schema,
5835 &statistics,
5836 &[true, true, true, true, true, true, true, true, true],
5838 );
5839
5840 prune_with_expr(
5842 col("s1")
5843 .like(lit("foo%bar%"))
5844 .and(col("s2").eq(lit("bar"))),
5845 &schema,
5846 &statistics,
5847 &[true, true, true, false, false, false, true, true, true],
5849 );
5850
5851 prune_with_expr(
5853 col("s1").like(lit("foo%bar%")).or(col("s2").eq(lit("bar"))),
5854 &schema,
5855 &statistics,
5856 &[true, true, true, true, true, true, true, true, true],
5859 );
5860 }
5861
5862 #[test]
5863 fn prune_with_range_and_contained() {
5864 let schema = Arc::new(Schema::new(vec![
5866 Field::new("i", DataType::Int32, true),
5867 Field::new("s", DataType::Utf8, true),
5868 ]));
5869
5870 let statistics = TestStatistics::new()
5871 .with(
5872 "i",
5873 ContainerStats::new_i32(
5874 vec![
5878 Some(-5),
5879 Some(10),
5880 None,
5881 Some(-5),
5882 Some(10),
5883 None,
5884 Some(-5),
5885 Some(10),
5886 None,
5887 ], vec![
5889 Some(5),
5890 Some(20),
5891 None,
5892 Some(5),
5893 Some(20),
5894 None,
5895 Some(5),
5896 Some(20),
5897 None,
5898 ], ),
5900 )
5901 .with_contained(
5903 "s",
5904 [ScalarValue::from("foo")],
5905 [
5906 Some(true),
5908 Some(true),
5909 Some(true),
5910 Some(false),
5912 Some(false),
5913 Some(false),
5914 None,
5916 None,
5917 None,
5918 ],
5919 );
5920
5921 prune_with_expr(
5923 col("i").eq(lit(0)).and(col("s").eq(lit("foo"))),
5924 &schema,
5925 &statistics,
5926 &[true, false, true, false, false, false, true, false, true],
5931 );
5932
5933 prune_with_expr(
5935 col("i").eq(lit(0)).and(col("s").not_eq(lit("foo"))),
5936 &schema,
5937 &statistics,
5938 &[false, false, false, true, false, true, true, false, true],
5942 );
5943
5944 prune_with_expr(
5946 col("i").eq(lit(0)).or(col("s").eq(lit("foo"))),
5947 &schema,
5948 &statistics,
5949 &[true, true, true, true, true, true, true, true, true],
5952 );
5953 }
5954
5955 fn prune_with_expr(
5962 expr: Expr,
5963 schema: &SchemaRef,
5964 statistics: &TestStatistics,
5965 expected: &[bool],
5966 ) {
5967 println!("Pruning with expr: {expr}");
5968 let expr = logical2physical(&expr, schema);
5969 let p = PruningPredicateBuilder::new()
5970 .with_file_schema(Arc::<Schema>::clone(schema))
5971 .try_build(expr)
5972 .unwrap();
5973 let result = p.prune(statistics).unwrap();
5974 assert_eq!(result, expected);
5975 }
5976
5977 fn prune_with_simplified_expr(
5978 expr: Expr,
5979 schema: &SchemaRef,
5980 statistics: &TestStatistics,
5981 expected: &[bool],
5982 ) {
5983 println!("Pruning with expr: {expr}");
5984 let expr = logical2physical(&expr, schema);
5985 let simplifier = PhysicalExprSimplifier::new(schema);
5986 let expr = simplifier.simplify(expr).unwrap();
5987 let p = PruningPredicateBuilder::new()
5988 .with_file_schema(Arc::<Schema>::clone(schema))
5989 .try_build(expr)
5990 .unwrap();
5991 let result = p.prune(statistics).unwrap();
5992 assert_eq!(result, expected);
5993 }
5994
5995 fn is_not_distinct_from(left: Expr, right: Expr) -> Expr {
5996 Expr::BinaryExpr(BinaryExpr::new(
5997 Box::new(left),
5998 Operator::IsNotDistinctFrom,
5999 Box::new(right),
6000 ))
6001 }
6002
6003 fn is_distinct_from(left: Expr, right: Expr) -> Expr {
6004 Expr::BinaryExpr(BinaryExpr::new(
6005 Box::new(left),
6006 Operator::IsDistinctFrom,
6007 Box::new(right),
6008 ))
6009 }
6010
6011 fn test_build_predicate_expression(
6012 expr: &Expr,
6013 schema: &Schema,
6014 required_columns: &mut RequiredColumns,
6015 ) -> Arc<dyn PhysicalExpr> {
6016 let expr = logical2physical(expr, schema);
6017 let unhandled_hook = Arc::new(ConstantUnhandledPredicateHook::default()) as _;
6018 build_predicate_expression(
6019 &expr,
6020 &Arc::new(schema.clone()),
6021 required_columns,
6022 &unhandled_hook,
6023 MAX_IN_LIST_SIZE,
6024 )
6025 }
6026
6027 #[test]
6028 fn test_build_predicate_expression_with_false() {
6029 let expr = lit(ScalarValue::Boolean(Some(false)));
6030 let schema = Schema::empty();
6031 let res =
6032 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
6033 let expected = logical2physical(&expr, &schema);
6034 assert_eq!(&res, &expected);
6035 }
6036
6037 #[test]
6038 fn test_build_predicate_expression_with_and_false() {
6039 let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]);
6040 let expr = and(
6041 col("c1").eq(lit("a")),
6042 lit(ScalarValue::Boolean(Some(false))),
6043 );
6044 let res =
6045 test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new());
6046 let expected = logical2physical(&lit(ScalarValue::Boolean(Some(false))), &schema);
6047 assert_eq!(&res, &expected);
6048 }
6049
6050 #[test]
6051 fn test_build_predicate_expression_with_or_false() {
6052 let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]);
6053 let left_expr = col("c1").eq(lit("a"));
6054 let right_expr = lit(ScalarValue::Boolean(Some(false)));
6055 let res = test_build_predicate_expression(
6056 &or(left_expr.clone(), right_expr.clone()),
6057 &schema,
6058 &mut RequiredColumns::new(),
6059 );
6060 let expected =
6061 "c1_null_count@2 != row_count@3 AND c1_min@0 <= a AND a <= c1_max@1";
6062 assert_eq!(res.to_string(), expected);
6063 }
6064}