opendp 0.14.2-dev.20260401.2

A library of differential privacy algorithms for the statistical analysis of sensitive private data.
use polars::{
    df,
    lazy::frame::{IntoLazy, LazyFrame},
    prelude::col,
};

use crate::{
    domains::{AtomDomain, LazyFrameDomain, SeriesDomain},
    metrics::{FrameDistance, SymmetricDistance},
    transformations::make_stable_lazyframe,
};

use super::*;

fn get_alias_data() -> Fallible<(LazyFrameDomain, LazyFrame)> {
    let lf_domain = LazyFrameDomain::new(vec![
        SeriesDomain::new("A", AtomDomain::<i32>::default()),
        SeriesDomain::new("B", AtomDomain::<bool>::default()),
    ])?;
    let ints = [1, 2, 3, 4, 5];
    let bools = [true, false, true, true, false];
    let lf = df!("A" => ints, "B" => bools)?.lazy();
    Ok((lf_domain, lf))
}

macro_rules! test_query {
    ($expr:expr) => {{
        let (lf_domain, lf) = get_alias_data()?;

        let query = lf.clone().with_column($expr);
        let t_lf = make_stable_lazyframe(lf_domain, FrameDistance(SymmetricDistance), query)?;

        assert_eq!(
            &t_lf.output_domain.schema(),
            &**t_lf.invoke(&lf)?.collect()?.schema()
        );
        Ok(())
    }};
}

#[test]
fn test_alias_shadow_forward() -> Fallible<()> {
    test_query!(col("A").alias("B"))
}

#[test]
fn test_alias_shadow_backward() -> Fallible<()> {
    test_query!(col("B").alias("A"))
}

#[test]
fn test_alias_new_column() -> Fallible<()> {
    test_query!(col("A").alias("C"))
}

#[test]
fn test_alias_same_column() -> Fallible<()> {
    test_query!(col("A").alias("A"))
}

#[test]
fn test_alias() -> Fallible<()> {
    let ints = [1, 2, 3, 4, 5];
    let bools = [true, false, true, true, false];
    let lf = df!("A" => ints, "B" => bools)?.lazy();

    let expr_domain = LazyFrameDomain::new(vec![
        SeriesDomain::new("A", AtomDomain::<i32>::default()),
        SeriesDomain::new("B", AtomDomain::<bool>::default()),
    ])?
    .row_by_row();

    let t_ab = col("A")
        .alias("B")
        .make_stable(expr_domain.clone(), FrameDistance(SymmetricDistance))?;
    let t_bc = col("B")
        .alias("C")
        .make_stable(expr_domain, FrameDistance(SymmetricDistance))?;

    let expr_ab = t_ab.invoke(&lf.logical_plan)?.expr;
    let expr_bc = t_bc.invoke(&lf.logical_plan)?.expr;

    let actual = lf.with_columns([expr_ab, expr_bc]).collect()?;
    let expect = df!("A" => ints, "B" => ints, "C" => bools)?;
    assert_eq!(actual, expect);

    Ok(())
}