feat: Support when-then-otherwise

python/test/test_polars_ternary.py

@@ -0,0 +1,79 @@
+# This is more thoroughly tested than other Exprs only because
+# Chuck is new to this, and is more comfortable testing in python than rust, for now.
+
+import pytest
+import opendp.prelude as dp
+
+
+def example_lf(margin=None, **kwargs):
+    pl = pytest.importorskip("polars")
+    domains, series = example_series()
+    lf_domain, lf = dp.lazyframe_domain(domains), pl.LazyFrame(series)
+    if margin is not None:
+        lf_domain = dp.with_margin(lf_domain, by=margin, **kwargs)
+    return lf_domain, lf
+
+
+def example_series():
+    pl = pytest.importorskip("polars")
+    return [
+        dp.series_domain("A", dp.option_domain(dp.atom_domain(T=dp.f64))),
+    ], [
+        pl.Series("A", [1.0] * 50, dtype=pl.Float64),
+    ]
+
+
+def test_when_then_otherwise():
+    pl = pytest.importorskip("polars")
+    lf_domain, lf = example_lf()
+    m_lf = dp.t.make_stable_lazyframe(
+        lf_domain,
+        dp.symmetric_distance(),
+        lf.select(
+            pl.when(pl.col("A") == 1).then(1).otherwise(0).alias('fifty'),
+            pl.when(pl.col("A") == 0).then(1).otherwise(0).alias('zero'),
+        ),
+    )
+    results = m_lf(lf).collect().sum()
+    assert results['fifty'].item() == 50
+    assert results['zero'].item() == 0
+
+
+def test_when_then_otherwise_strings():
+    pl = pytest.importorskip("polars")
+    lf_domain, lf = example_lf()
+    m_lf = dp.t.make_stable_lazyframe(
+        lf_domain,
+        dp.symmetric_distance(),
+        lf.select(
+            pl.when(pl.col("A") == 1).then(pl.lit("one")).otherwise(pl.lit("other")),
+        ),
+    )
+    assert m_lf(lf).collect()['literal'][0] == 'one'
+
+
+def test_when_then_otherwise_mismatch_types():
+    pl = pytest.importorskip("polars")
+    lf_domain, lf = example_lf()
+    with pytest.raises(dp.OpenDPException, match=r'output domains in ternary must match'):
+        dp.t.make_stable_lazyframe(
+            lf_domain,
+            dp.symmetric_distance(),
+            lf.select(
+                pl.when(pl.col("A") == 1).then(1).otherwise(pl.lit("!!!")).alias('fifty'),
+            ),
+        )
+
+
+def test_when_then_otherwise_incomplete():
+    pl = pytest.importorskip("polars")
+    lf_domain, lf = example_lf()
+    with pytest.raises(Exception, match=r'unsupported literal value: null'):
+        dp.t.make_stable_lazyframe(
+            lf_domain,
+            dp.symmetric_distance(),
+            lf.select(
+                pl.when(pl.col("A") == 1).then(1).alias('fifty'),
+            ),
+        )
+    # TODO: Should there be a better error message?

rust/src/domains/polars/expr/mod.rs

@@ -132,9 +132,16 @@ impl LazyFrameDomain {
 }
 
 #[derive(Clone)]
+/// Used internally as we wrap built-in Polars expressions
 pub struct ExprPlan {
+    /// This can just be cloned.
     pub plan: DslPlan,
+    /// Typically we will pull the fields out of an incoming expression,
+    /// process them, and then create a new equivalent expression.
     pub expr: Expr,
+    /// For expressions before aggregation in the chain, this should be None.
+    /// If this *is* an aggregation, it should be a literal.
+    /// After an aggregation, this should just duplicate `expr`.
     pub fill: Option<Expr>,
 }
 

rust/src/transformations/make_stable_expr/expr_ternary/mod.rs

@@ -0,0 +1,99 @@
+use polars::prelude::*;
+use polars_plan::dsl::Expr;
+
+use crate::core::{Function, MetricSpace, StabilityMap, Transformation};
+use crate::domains::{ExprDomain, ExprPlan, OuterMetric, WildExprDomain};
+use crate::error::*;
+use crate::transformations::DatasetMetric;
+
+use super::StableExpr;
+
+#[cfg(test)]
+mod test;
+
+/// Make a Transformation that returns a `ternary` expression
+///
+/// # Arguments
+/// * `input_domain` - Expr domain
+/// * `input_metric` - The metric under which neighboring LazyFrames are compared
+/// * `expr` - The ternary expression
+pub fn make_expr_ternary<M: OuterMetric>(
+    input_domain: WildExprDomain,
+    input_metric: M,
+    expr: Expr,
+) -> Fallible<Transformation<WildExprDomain, ExprDomain, M, M>>
+where
+    M::InnerMetric: DatasetMetric,
+    M::Distance: Clone,
+    (WildExprDomain, M): MetricSpace,
+    (ExprDomain, M): MetricSpace,
+    Expr: StableExpr<M, M>,
+{
+    let Expr::Ternary {
+        predicate,
+        truthy,
+        falsy,
+    } = expr
+    else {
+        return fallible!(MakeTransformation, "expected ternary expression");
+    };
+
+    let t_predicate = predicate
+        .as_ref()
+        .clone()
+        .make_stable(input_domain.as_row_by_row(), input_metric.clone())?;
+    let t_truthy = truthy
+        .as_ref()
+        .clone()
+        .make_stable(input_domain.as_row_by_row(), input_metric.clone())?;
+    let t_falsy = falsy
+        .as_ref()
+        .clone()
+        .make_stable(input_domain.as_row_by_row(), input_metric.clone())?;
+
+    let (truthy_domain, _truthy_metric) = t_truthy.output_space();
+    let (falsy_domain, _falsy_metric) = t_falsy.output_space();
+
+    if truthy_domain != falsy_domain {
+        return fallible!(
+            MakeTransformation,
+            "output domains in ternary must match, instead found {:?} and {:?}",
+            truthy_domain,
+            falsy_domain
+        );
+    }
+
+    if matches!(truthy_domain.column.dtype(), DataType::Categorical(_, _)) {
+        // Since literal categorical values aren't possible,
+        // not clear if this is actually reachable.
+        return fallible!(MakeTransformation, "ternary cannot be applied to categorical data, because it may trigger a data-dependent CategoricalRemappingWarning in Polars");
+    }
+
+    let mut output_domain = truthy_domain.clone();
+    output_domain.column.drop_bounds().ok();
+    output_domain.column.nullable = false;
+    output_domain.context = input_domain.context.clone();
+
+    Transformation::new(
+        input_domain,
+        output_domain,
+        Function::new_fallible(move |arg| {
+            let predicate = t_predicate.invoke(arg)?;
+            let truthy = t_truthy.invoke(arg)?;
+            let falsy = t_falsy.invoke(arg)?;
+
+            Ok(ExprPlan {
+                plan: arg.clone(),
+                expr: Expr::Ternary {
+                    predicate: Arc::new(predicate.expr),
+                    truthy: Arc::new(truthy.expr),
+                    falsy: Arc::new(falsy.expr),
+                },
+                fill: None, // Ternary is run before aggregation, so there's no empty group that needs a default filled in.
+            })
+        }),
+        input_metric.clone(),
+        input_metric,
+        StabilityMap::new(Clone::clone),
+    )
+}

rust/src/transformations/make_stable_expr/expr_ternary/test.rs

@@ -0,0 +1 @@
+

rust/src/transformations/make_stable_expr/mod.rs

@@ -59,6 +59,9 @@ mod expr_sum;
 #[cfg(feature = "contrib")]
 mod expr_to_physical;
 
+#[cfg(feature = "contrib")]
+mod expr_ternary;
+
 #[cfg(feature = "contrib")]
 mod namespace_dt;
 
@@ -174,6 +177,9 @@ where
                 ..
             } => namespace_str::make_namespace_str(input_domain, input_metric, self),
 
+            #[cfg(feature = "contrib")]
+            Expr::Ternary { .. } => expr_ternary::make_expr_ternary(input_domain, input_metric, self),
+
             expr => fallible!(
                 MakeTransformation,
                 "Expr is not recognized at this time: {:?}. {}If you would like to see this supported, please file an issue.",

rust/src/transformations/make_stable_lazyframe/source/mod.rs

@@ -30,7 +30,9 @@ where
     if &input_domain.schema() != schema.as_ref() {
         return fallible!(
             MakeTransformation,
-            "Schema mismatch. LazyFrame schema must match the schema from the input domain."
+            "Schema mismatch. LazyFrame schema must match the schema from the input domain. {:?} != {:?}",
+            &input_domain.schema(),
+            schema.as_ref()
         );
     }