Skip to content

Commit ef546e4

Browse files
committed
refactor: validate quantile value
Ensures the quantile values is between 0 and 1, emitting a plan error if not.
1 parent faa8094 commit ef546e4

File tree

2 files changed

+33
-1
lines changed

2 files changed

+33
-1
lines changed

datafusion/src/physical_plan/aggregates.rs

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -547,7 +547,7 @@ mod tests {
547547
Arc::new(
548548
expressions::Column::new_with_schema("c1", &input_schema).unwrap(),
549549
),
550-
Arc::new(expressions::Literal::new(ScalarValue::Float64(Some(4.2)))),
550+
Arc::new(expressions::Literal::new(ScalarValue::Float64(Some(0.2)))),
551551
];
552552
let result_agg_phy_exprs = create_aggregate_expr(
553553
&AggregateFunction::ApproxQuantile,
@@ -567,6 +567,30 @@ mod tests {
567567
}
568568
}
569569

570+
#[test]
571+
fn test_agg_approx_quantile_invalid_phy_expr() {
572+
for data_type in NUMERICS {
573+
let input_schema =
574+
Schema::new(vec![Field::new("c1", data_type.clone(), true)]);
575+
let input_phy_exprs: Vec<Arc<dyn PhysicalExpr>> = vec![
576+
Arc::new(
577+
expressions::Column::new_with_schema("c1", &input_schema).unwrap(),
578+
),
579+
Arc::new(expressions::Literal::new(ScalarValue::Float64(Some(4.2)))),
580+
];
581+
let err = create_aggregate_expr(
582+
&AggregateFunction::ApproxQuantile,
583+
false,
584+
&input_phy_exprs[..],
585+
&input_schema,
586+
"c1",
587+
)
588+
.expect_err("should fail due to invalid quantile");
589+
590+
assert!(matches!(err, DataFusionError::Plan(_)));
591+
}
592+
}
593+
570594
#[test]
571595
fn test_min_max_expr() -> Result<()> {
572596
let funcs = vec![AggregateFunction::Min, AggregateFunction::Max];

datafusion/src/physical_plan/expressions/approx_quantile.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,14 @@ impl ApproxQuantile {
8989
)))
9090
};
9191

92+
// Ensure the quantile is between 0 and 1.
93+
if !(0.0..1.0).contains(&quantile) {
94+
return Err(DataFusionError::Plan(format!(
95+
"Quantile value must be between 0.0 and 1.0, {} is invalid",
96+
quantile
97+
)));
98+
}
99+
92100
Ok(Self {
93101
name: name.into(),
94102
input_data_type,

0 commit comments

Comments
 (0)