fix pruner

Author: b41sh

src/query/storages/common/index/src/eliminate_cast.rs

@@ -32,13 +32,31 @@ use databend_common_functions::BUILTIN_FUNCTIONS;
 
 pub(super) struct RewriteVisitor<'a> {
     pub input_domains: HashMap<String, Domain>,
+    /// Optional block-local physical types for virtual column references.
+    pub column_types: Option<&'a HashMap<String, DataType>>,
     pub func_ctx: &'a FunctionContext,
     pub fn_registry: &'a FunctionRegistry,
 }
 
 type RewriteResult = std::result::Result<Option<Expr<String>>, !>;
 
 impl ExprVisitor<String> for RewriteVisitor<'_> {
+    fn enter_column_ref(&mut self, column: &ColumnRef<String>) -> RewriteResult {
+        let Some(data_type) = self
+            .column_types
+            .and_then(|column_types| column_types.get(&column.id))
+        else {
+            return Ok(None);
+        };
+        if data_type == &column.data_type {
+            return Ok(None);
+        }
+
+        let mut column = column.clone();
+        column.data_type = data_type.clone();
+        Ok(Some(column.into()))
+    }
+
     fn enter_function_call(&mut self, call: &FunctionCall<String>) -> RewriteResult {
         if call.id.name() == "eq" {
             let result = match call.args.as_slice() {
@@ -268,6 +286,7 @@ pub fn eliminate_cast(
 ) -> Option<Expr<String>> {
     let mut visitor = RewriteVisitor {
         input_domains,
+        column_types: None,
         func_ctx: &FunctionContext::default(),
         fn_registry: &BUILTIN_FUNCTIONS,
     };

src/query/storages/common/index/src/range_index.rs

@@ -114,59 +114,58 @@ impl RangeIndex {
     where
         F: Fn(&ColumnId) -> bool,
     {
-        let mut input_domains: HashMap<String, Domain> = self
-            .expr
-            .column_refs()
-            .into_iter()
-            .map(|(name, ty)| {
-                // internal column and stream column are not actual stored columns
-                if is_internal_column(&name) || is_stream_column(&name) {
-                    return Ok((name, Domain::full(&ty)));
-                }
+        let mut input_domains = HashMap::new();
+        let mut virtual_column_types = HashMap::new();
+        for (name, ty) in self.expr.column_refs() {
+            // internal column and stream column are not actual stored columns
+            if is_internal_column(&name) || is_stream_column(&name) {
+                input_domains.insert(name, Domain::full(&ty));
+                continue;
+            }
 
-                let column_ids = self.schema.leaf_columns_of(&name);
-                if column_ids.is_empty() {
-                    // The name may refer to a virtual column (e.g. `v['a']`). Use the
-                    // block-local virtual column statistics to build the domain.
-                    // Only typed statistics are injected; everything else falls back
-                    // to a full domain to avoid wrong pruning.
-                    let virtual_domain = virtual_col_stats.and_then(|virtual_col_stats| {
-                        virtual_col_stats.get(&name).map(|stat| {
-                            let column_stat = stat.to_column_statistics();
-                            let data_type = DataType::from(&stat.data_type);
-                            statistics_to_domain(vec![&column_stat], &data_type)
-                        })
-                    });
-                    return Ok((name, virtual_domain.unwrap_or_else(|| Domain::full(&ty))));
+            let column_ids = self.schema.leaf_columns_of(&name);
+            if column_ids.is_empty() {
+                // The name may refer to a virtual column (e.g. `v['a']`). Use the
+                // block-local virtual column statistics to build the domain.
+                // Only typed statistics are injected; everything else falls back
+                // to a full domain to avoid wrong pruning.
+                if let Some(stat) = virtual_col_stats.and_then(|stats| stats.get(&name)) {
+                    let column_stat = stat.to_column_statistics();
+                    let data_type = DataType::from(&stat.data_type);
+                    let domain = statistics_to_domain(vec![&column_stat], &data_type);
+                    virtual_column_types.insert(name.clone(), data_type);
+                    input_domains.insert(name, domain);
+                } else {
+                    input_domains.insert(name, Domain::full(&ty));
                 }
+                continue;
+            }
 
-                let stats = column_ids
-                    .iter()
-                    .filter_map(|column_id| match stats.get(column_id) {
-                        None => {
-                            if column_is_default(column_id)
-                                && self.default_stats.contains_key(column_id)
-                            {
-                                Some(&self.default_stats[column_id])
-                            } else {
-                                None
-                            }
+            let stats = column_ids
+                .iter()
+                .filter_map(|column_id| match stats.get(column_id) {
+                    None => {
+                        if column_is_default(column_id)
+                            && self.default_stats.contains_key(column_id)
+                        {
+                            Some(&self.default_stats[column_id])
+                        } else {
+                            None
                         }
-                        other => other,
-                    })
-                    .collect();
-
-                let domain = statistics_to_domain(stats, &ty);
-                Ok((name, domain))
-            })
-            .collect::<Result<_>>()?;
+                    }
+                    other => other,
+                })
+                .collect();
+            input_domains.insert(name, statistics_to_domain(stats, &ty));
+        }
 
         for (name, domain) in self.spatial_predicate_domains(spatial_stats) {
             input_domains.insert(name, domain);
         }
 
         let mut visitor = RewriteVisitor {
             input_domains,
+            column_types: (!virtual_column_types.is_empty()).then_some(&virtual_column_types),
             func_ctx: &self.func_ctx,
             fn_registry: &BUILTIN_FUNCTIONS,
         };

src/query/storages/common/index/tests/it/range_pruner.rs

@@ -173,6 +173,42 @@ fn test_range_index_prunes_by_virtual_column_statistics() {
     );
 }
 
+#[test]
+fn test_range_index_prunes_try_cast_virtual_column_by_physical_type() {
+    fn n(value: u64) -> Scalar {
+        Scalar::Number(value.into())
+    }
+
+    let schema = Arc::new(TableSchema::new(vec![TableField::new(
+        "v",
+        TableDataType::Variant,
+    )]));
+    let expr = parse_expr("is_true(try_cast(v.a as uint8) = 100)", &[(
+        "v.a",
+        DataType::Variant.wrap_nullable(),
+    )]);
+    let index = RangeIndex::try_create(
+        FunctionContext::default(),
+        &expr,
+        schema,
+        Default::default(),
+    )
+    .unwrap();
+    let stats = VirtualColumnStatsOfNames::from([("v.a".to_string(), VirtualColumnStat {
+        query_column_id: 3_000_000_000,
+        min: n(1),
+        max: n(8),
+        null_count: 0,
+        data_type: TableDataType::Number(NumberDataType::UInt64),
+    })]);
+
+    assert!(
+        !index
+            .apply(&Default::default(), None, Some(&stats), |_| false)
+            .unwrap()
+    );
+}
+
 #[test]
 fn test_range_index_keeps_without_virtual_column_statistics() {
     let schema = Arc::new(TableSchema::new(vec![TableField::new(

src/query/storages/common/pruner/src/range_pruner.rs

@@ -21,7 +21,6 @@ use databend_common_expression::Expr;
 use databend_common_expression::FunctionContext;
 use databend_common_expression::Scalar;
 use databend_common_expression::TableSchemaRef;
-use databend_common_expression::infer_schema_type;
 use databend_common_expression::types::DataType;
 use databend_storages_common_index::RangeIndex;
 use databend_storages_common_index::VirtualColumnStat;
@@ -115,25 +114,28 @@ fn build_virtual_col_stats(
         let Some(column_stat) = column.column_stat.as_ref() else {
             continue;
         };
-        let Some(data_type) = column_stat
+        let Some(scalar_data_type) = column_stat
             .min
             .as_ref()
             .infer_common_type(&column_stat.max.as_ref())
         else {
             continue;
         };
-        if matches!(data_type.remove_nullable(), DataType::Variant) {
+        let physical_data_type = column.data_type();
+        let physical_expression_type = DataType::from(&physical_data_type);
+        if matches!(
+            physical_expression_type.remove_nullable(),
+            DataType::Variant
+        ) || physical_expression_type.remove_nullable() != scalar_data_type.remove_nullable()
+        {
             continue;
         }
-        let Ok(data_type) = infer_schema_type(&data_type) else {
-            continue;
-        };
         stats.insert(virtual_ref.name.clone(), VirtualColumnStat {
             query_column_id: virtual_ref.query_column_id,
             min: column_stat.min.clone(),
             max: column_stat.max.clone(),
             null_count: column_stat.null_count,
-            data_type,
+            data_type: physical_data_type,
         });
     }
     (!stats.is_empty()).then_some(stats)