refactor: validate quantile value

Ensures the quantile values is between 0 and 1, emitting a plan error if not.
apache · Jan 27, 2022 · ef546e4 · ef546e4
1 parent faa8094
commit ef546e4
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 1 deletion.
diff --git a/datafusion/src/physical_plan/aggregates.rs b/datafusion/src/physical_plan/aggregates.rs
@@ -547,7 +547,7 @@ mod tests {
                 Arc::new(
                     expressions::Column::new_with_schema("c1", &input_schema).unwrap(),
                 ),
-                Arc::new(expressions::Literal::new(ScalarValue::Float64(Some(4.2)))),
+                Arc::new(expressions::Literal::new(ScalarValue::Float64(Some(0.2)))),
             ];
             let result_agg_phy_exprs = create_aggregate_expr(
                 &AggregateFunction::ApproxQuantile,
@@ -567,6 +567,30 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_agg_approx_quantile_invalid_phy_expr() {
+        for data_type in NUMERICS {
+            let input_schema =
+                Schema::new(vec![Field::new("c1", data_type.clone(), true)]);
+            let input_phy_exprs: Vec<Arc<dyn PhysicalExpr>> = vec![
+                Arc::new(
+                    expressions::Column::new_with_schema("c1", &input_schema).unwrap(),
+                ),
+                Arc::new(expressions::Literal::new(ScalarValue::Float64(Some(4.2)))),
+            ];
+            let err = create_aggregate_expr(
+                &AggregateFunction::ApproxQuantile,
+                false,
+                &input_phy_exprs[..],
+                &input_schema,
+                "c1",
+            )
+            .expect_err("should fail due to invalid quantile");
+
+            assert!(matches!(err, DataFusionError::Plan(_)));
+        }
+    }
+
     #[test]
     fn test_min_max_expr() -> Result<()> {
         let funcs = vec![AggregateFunction::Min, AggregateFunction::Max];

diff --git a/datafusion/src/physical_plan/expressions/approx_quantile.rs b/datafusion/src/physical_plan/expressions/approx_quantile.rs
@@ -89,6 +89,14 @@ impl ApproxQuantile {
             )))
         };
 
+        // Ensure the quantile is between 0 and 1.
+        if !(0.0..1.0).contains(&quantile) {
+            return Err(DataFusionError::Plan(format!(
+                "Quantile value must be between 0.0 and 1.0, {} is invalid",
+                quantile
+            )));
+        }
+
         Ok(Self {
             name: name.into(),
             input_data_type,