Use registry NDV for equality/inequality selectivity on expressions

Author: asolimando

datafusion/physical-expr/src/expression_analyzer/default.rs

@@ -52,6 +52,27 @@ impl DefaultExpressionAnalyzer {
             .and_then(|idx| input_stats.column_statistics.get(idx))
     }
 
+    /// Resolve NDV for a binary expression: try direct column stats first,
+    /// then fall back to the registry for arbitrary expressions
+    fn resolve_ndv(
+        left: &Arc<dyn PhysicalExpr>,
+        right: &Arc<dyn PhysicalExpr>,
+        input_stats: &Statistics,
+        registry: &ExpressionAnalyzerRegistry,
+    ) -> Option<usize> {
+        Self::get_column_stats(left, input_stats)
+            .or_else(|| Self::get_column_stats(right, input_stats))
+            .and_then(|s| s.distinct_count.get_value())
+            .filter(|&&ndv| ndv > 0)
+            .copied()
+            .or_else(|| {
+                let l = registry.get_distinct_count(left, input_stats);
+                let r = registry.get_distinct_count(right, input_stats);
+                l.max(r)
+            })
+            .filter(|&n| n > 0)
+    }
+
     /// Recursive selectivity estimation through the registry chain
     fn estimate_selectivity_recursive(
         &self,
@@ -103,24 +124,26 @@ impl ExpressionAnalyzer for DefaultExpressionAnalyzer {
 
                 // Equality: selectivity = 1/NDV
                 Operator::Eq => {
-                    let ndv = Self::get_column_stats(binary.left(), input_stats)
-                        .or_else(|| Self::get_column_stats(binary.right(), input_stats))
-                        .and_then(|s| s.distinct_count.get_value())
-                        .filter(|&&ndv| ndv > 0);
-                    if let Some(ndv) = ndv {
-                        return AnalysisResult::Computed(1.0 / (*ndv as f64));
+                    if let Some(ndv) = Self::resolve_ndv(
+                        binary.left(),
+                        binary.right(),
+                        input_stats,
+                        registry,
+                    ) {
+                        return AnalysisResult::Computed(1.0 / (ndv as f64));
                     }
                     0.1 // Default equality selectivity
                 }
 
                 // Inequality: selectivity = 1 - 1/NDV
                 Operator::NotEq => {
-                    let ndv = Self::get_column_stats(binary.left(), input_stats)
-                        .or_else(|| Self::get_column_stats(binary.right(), input_stats))
-                        .and_then(|s| s.distinct_count.get_value())
-                        .filter(|&&ndv| ndv > 0);
-                    if let Some(ndv) = ndv {
-                        return AnalysisResult::Computed(1.0 - (1.0 / (*ndv as f64)));
+                    if let Some(ndv) = Self::resolve_ndv(
+                        binary.left(),
+                        binary.right(),
+                        input_stats,
+                        registry,
+                    ) {
+                        return AnalysisResult::Computed(1.0 - (1.0 / (ndv as f64)));
                     }
                     0.9
                 }

datafusion/physical-expr/src/expression_analyzer/tests.rs

@@ -162,6 +162,44 @@ fn test_not_selectivity() {
     assert!((sel - 0.99).abs() < 0.001); // 1 - 0.01
 }
 
+#[test]
+fn test_equality_selectivity_expression_eq_literal() {
+    let stats = make_stats_with_ndv(1000, 100);
+    let col = Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>;
+    let one =
+        Arc::new(Literal::new(ScalarValue::Int32(Some(1)))) as Arc<dyn PhysicalExpr>;
+    let forty_two =
+        Arc::new(Literal::new(ScalarValue::Int32(Some(42)))) as Arc<dyn PhysicalExpr>;
+    let a_plus_1 =
+        Arc::new(BinaryExpr::new(col, Operator::Plus, one)) as Arc<dyn PhysicalExpr>;
+    let eq = Arc::new(BinaryExpr::new(a_plus_1, Operator::Eq, forty_two))
+        as Arc<dyn PhysicalExpr>;
+
+    let registry = ExpressionAnalyzerRegistry::new();
+    let sel = registry.get_selectivity(&eq, &stats).unwrap();
+    // NDV(a + 1) = NDV(a) = 100, so selectivity = 1/100 = 0.01
+    assert!((sel - 0.01).abs() < 0.001);
+}
+
+#[test]
+fn test_inequality_selectivity_expression_neq_literal() {
+    let stats = make_stats_with_ndv(1000, 100);
+    let col = Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>;
+    let one =
+        Arc::new(Literal::new(ScalarValue::Int32(Some(1)))) as Arc<dyn PhysicalExpr>;
+    let forty_two =
+        Arc::new(Literal::new(ScalarValue::Int32(Some(42)))) as Arc<dyn PhysicalExpr>;
+    let a_plus_1 =
+        Arc::new(BinaryExpr::new(col, Operator::Plus, one)) as Arc<dyn PhysicalExpr>;
+    let neq = Arc::new(BinaryExpr::new(a_plus_1, Operator::NotEq, forty_two))
+        as Arc<dyn PhysicalExpr>;
+
+    let registry = ExpressionAnalyzerRegistry::new();
+    let sel = registry.get_selectivity(&neq, &stats).unwrap();
+    // NDV(a + 1) = 100, selectivity = 1 - 1/100 = 0.99
+    assert!((sel - 0.99).abs() < 0.001);
+}
+
 // Min/max tests
 
 #[test]