alopex-dataframe 0.7.2

Polars-compatible DataFrame API for Alopex DB (v0.1)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

use std::collections::HashMap;
use std::sync::Arc;

use alopex_dataframe::expr::col;
use alopex_dataframe::{DataFrame, DataFrameError, Series};
use arrow::array::{Array, ArrayRef, Float64Array, Int64Array, StringArray};
use arrow::datatypes::DataType;

fn df() -> DataFrame {
    let g: ArrayRef = Arc::new(StringArray::from(vec![
        Some("x"),
        Some("x"),
        Some("y"),
        Some("y"),
        Some("y"),
    ]));
    let v: ArrayRef = Arc::new(Int64Array::from(vec![
        Some(10_i64),
        None,
        Some(3),
        Some(7),
        None,
    ]));

    DataFrame::new(vec![
        Series::from_arrow("g", vec![g]).unwrap(),
        Series::from_arrow("v", vec![v]).unwrap(),
    ])
    .unwrap()
}

#[cfg_attr(not(feature = "lane_ci"), ignore)]
#[test]
fn group_by_agg_semantics_and_column_order() {
    let out = df()
        .lazy()
        .group_by(vec![col("g")])
        .agg(vec![
            col("v").sum().alias("sum_v"),
            col("v").mean().alias("mean_v"),
            col("v").min().alias("min_v"),
            col("v").max().alias("max_v"),
            col("v").count().alias("cnt_v"),
        ])
        .collect()
        .unwrap();

    let names: Vec<_> = out
        .schema()
        .fields()
        .iter()
        .map(|f| f.name().clone())
        .collect();
    assert_eq!(
        names,
        vec!["g", "sum_v", "mean_v", "min_v", "max_v", "cnt_v"]
    );

    assert_eq!(out.column("mean_v").unwrap().dtype(), DataType::Float64);
    assert_eq!(out.column("cnt_v").unwrap().dtype(), DataType::Int64);

    let sums = group_i64_map(&out, "sum_v");
    let mins = group_i64_map(&out, "min_v");
    let maxs = group_i64_map(&out, "max_v");
    let cnts = group_i64_map(&out, "cnt_v");
    let means = group_f64_map(&out, "mean_v");

    assert_eq!(sums.get("x").copied(), Some(10));
    assert_eq!(mins.get("x").copied(), Some(10));
    assert_eq!(maxs.get("x").copied(), Some(10));
    assert_eq!(cnts.get("x").copied(), Some(1));
    assert_eq!(means.get("x").copied(), Some(10.0));

    assert_eq!(sums.get("y").copied(), Some(10));
    assert_eq!(mins.get("y").copied(), Some(3));
    assert_eq!(maxs.get("y").copied(), Some(7));
    assert_eq!(cnts.get("y").copied(), Some(2));
    assert_eq!(means.get("y").copied(), Some(5.0));
}

#[cfg_attr(not(feature = "lane_ci"), ignore)]
#[test]
fn non_numeric_aggregation_is_type_mismatch() {
    let g: ArrayRef = Arc::new(StringArray::from(vec!["x", "x"]));
    let s: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
    let df = DataFrame::new(vec![
        Series::from_arrow("g", vec![g]).unwrap(),
        Series::from_arrow("s", vec![s]).unwrap(),
    ])
    .unwrap();

    let err = df
        .lazy()
        .group_by(vec![col("g")])
        .agg(vec![col("s").sum().alias("sum_s")])
        .collect()
        .unwrap_err();
    assert!(matches!(err, DataFrameError::TypeMismatch { .. }));
    assert!(err.to_string().contains("expected"));
}

fn group_i64_map(df: &DataFrame, value_col: &str) -> HashMap<String, i64> {
    let keys = df.column("g").unwrap().to_arrow();
    let vals = df.column(value_col).unwrap().to_arrow();
    assert_eq!(keys.len(), 1);
    assert_eq!(vals.len(), 1);

    let keys = keys[0].as_any().downcast_ref::<StringArray>().unwrap();
    let vals = vals[0].as_any().downcast_ref::<Int64Array>().unwrap();

    let mut out = HashMap::new();
    for i in 0..keys.len() {
        out.insert(keys.value(i).to_string(), vals.value(i));
    }
    out
}

fn group_f64_map(df: &DataFrame, value_col: &str) -> HashMap<String, f64> {
    let keys = df.column("g").unwrap().to_arrow();
    let vals = df.column(value_col).unwrap().to_arrow();
    assert_eq!(keys.len(), 1);
    assert_eq!(vals.len(), 1);

    let keys = keys[0].as_any().downcast_ref::<StringArray>().unwrap();
    let vals = vals[0].as_any().downcast_ref::<Float64Array>().unwrap();

    let mut out = HashMap::new();
    for i in 0..keys.len() {
        out.insert(keys.value(i).to_string(), vals.value(i));
    }
    out
}