Add conformance tests, support integer quantiles (0 and 1)

pola-rs · Sep 9, 2024 · a2cceac · a2cceac
1 parent c2e1e9e
commit a2cceac
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 2 deletions.
diff --git a/crates/polars-sql/src/functions.rs b/crates/polars-sql/src/functions.rs
@@ -1273,6 +1273,13 @@ impl SQLFunctionVisitor<'_> {
                                     polars_bail!(SQLSyntax: "QUANTILE_DISC value must be between 0 and 1 ({})", args[1])
                                 }
                             },
+                            Expr::Literal(LiteralValue::Int(n)) => {
+                                if (0..=1).contains(&n) {
+                                    Expr::from(n as f64)
+                                } else {
+                                    polars_bail!(SQLSyntax: "QUANTILE_DISC value must be between 0 and 1 ({})", args[1])
+                                }
+                            },
                             _ => polars_bail!(SQLSyntax: "invalid value for QUANTILE_DISC ({})", args[1])
                         };
                         Ok(e.quantile(value, QuantileInterpolOptions::Lower))
@@ -1292,6 +1299,13 @@ impl SQLFunctionVisitor<'_> {
                                     polars_bail!(SQLSyntax: "QUANTILE_CONT value must be between 0 and 1 ({})", args[1])
                                 }
                             },
+                            Expr::Literal(LiteralValue::Int(n)) => {
+                                if (0..=1).contains(&n) {
+                                    Expr::from(n as f64)
+                                } else {
+                                    polars_bail!(SQLSyntax: "QUANTILE_CONT value must be between 0 and 1 ({})", args[1])
+                                }
+                            },
                             _ => polars_bail!(SQLSyntax: "invalid value for QUANTILE_CONT ({})", args[1])
                         };
                         Ok(e.quantile(value, QuantileInterpolOptions::Linear))

diff --git a/crates/polars-sql/tests/functions_aggregate.rs b/crates/polars-sql/tests/functions_aggregate.rs
@@ -57,7 +57,38 @@ fn test_quantile_disc() {
         let sql_expr = format!("QUANTILE_DISC(Sales, {})", q);
         let (expected, actual) = create_expected(expr, &sql_expr);
 
-        assert!(expected.equals(&actual))
+        assert!(
+            expected.equals(&actual),
+            "q: {q}: expected {expected:?}, got {actual:?}"
+        )
+    }
+}
+
+#[test]
+fn test_quantile_disc_conformance() {
+    for (q, expected) in [
+        (0., 1000),
+        (0.1, 1000),
+        (0.2, 2000),
+        (0.3, 2000),
+        (0.4, 3000),
+        (0.5, 3000),
+        (0.6, 4000),
+        (0.7, 4000),
+        (0.8, 5000),
+        (0.9, 5000),
+        (1., 6000),
+    ] {
+        let expr = lit(expected);
+
+        let sql_expr = format!("QUANTILE_DISC(Sales, {q})");
+
+        let (expected, actual) = create_expected(expr, &sql_expr);
+
+        assert!(
+            expected.equals(&actual),
+            "q: {q}: expected {expected:?}, got {actual:?}"
+        );
     }
 }
 
@@ -69,6 +100,9 @@ fn test_quantile_cont() {
         let sql_expr = format!("QUANTILE_CONT(Sales, {})", q);
         let (expected, actual) = create_expected(expr, &sql_expr);
 
-        assert!(expected.equals(&actual))
+        assert!(
+            expected.equals(&actual),
+            "q: {q}: expected {expected:?}, got {actual:?}"
+        )
     }
 }