dsq-core 0.2.0

Core functionality for dsq - data processing with jq syntax
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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200

//! Aggregation operations for dsq
//!
//! This module provides aggregation functions for `DataFrames` including:
//! - Group by operations
//! - Statistical aggregations (sum, mean, count, etc.)
//! - Window functions
//! - Pivot and unpivot operations
//!
//! These operations correspond to common SQL aggregations and jq's `group_by`
//! functionality, adapted for tabular data processing.

use crate::{Error, Result, Value};
use polars::prelude::*;

mod cumulative;
mod ewma;
mod group_by;
mod pivot;
mod rolling;
mod unpivot;

pub use cumulative::cumulative_agg;
pub use ewma::ewma;
pub use group_by::{group_by, group_by_agg};
pub use pivot::pivot;
pub use rolling::{rolling_agg, rolling_std, WindowFunction};
pub use unpivot::unpivot;

/// Helper function to convert `AnyValue` to Value
pub(super) fn any_value_to_value(any_val: &AnyValue) -> Result<Value> {
    use serde_json::Value as JsonValue;
    let json_val = match any_val {
        AnyValue::Null => JsonValue::Null,
        AnyValue::Boolean(b) => JsonValue::Bool(*b),
        AnyValue::Int8(i) => JsonValue::Number(serde_json::Number::from(*i)),
        AnyValue::Int16(i) => JsonValue::Number(serde_json::Number::from(*i)),
        AnyValue::Int32(i) => JsonValue::Number(serde_json::Number::from(*i)),
        AnyValue::Int64(i) => JsonValue::Number(serde_json::Number::from(*i)),
        AnyValue::UInt8(i) => JsonValue::Number(serde_json::Number::from(*i)),
        AnyValue::UInt16(i) => JsonValue::Number(serde_json::Number::from(*i)),
        AnyValue::UInt32(i) => JsonValue::Number(serde_json::Number::from(*i)),
        AnyValue::UInt64(i) => JsonValue::Number(serde_json::Number::from(*i)),
        AnyValue::Float32(f) => JsonValue::Number(
            serde_json::Number::from_f64(f64::from(*f))
                .ok_or_else(|| Error::operation("Invalid float"))?,
        ),
        AnyValue::Float64(f) => JsonValue::Number(
            serde_json::Number::from_f64(*f).ok_or_else(|| Error::operation("Invalid float"))?,
        ),
        AnyValue::String(s) => JsonValue::String((*s).to_string()),
        _ => return Err(Error::operation("Unsupported AnyValue type")),
    };
    Ok(Value::from_json(json_val))
}

/// Helper function to convert `DataFrame` to Array of Objects
pub(super) fn df_to_array(df: &DataFrame) -> Result<Vec<Value>> {
    let columns = df.get_column_names();
    let mut result = Vec::with_capacity(df.height());

    for row_idx in 0..df.height() {
        let mut obj = std::collections::HashMap::new();
        for col_name in &columns {
            let series = df.column(col_name).map_err(Error::from)?;
            let any_val = series.get(row_idx).map_err(Error::from)?;
            let value = any_value_to_value(&any_val)?;
            obj.insert(col_name.to_string(), value);
        }
        result.push(Value::Object(obj));
    }

    Ok(result)
}

/// Aggregation functions that can be applied to grouped data
#[derive(Debug, Clone)]
pub enum AggregationFunction {
    /// Count of rows in each group
    Count,
    /// Sum of values in specified column
    Sum(String),
    /// Mean/average of values in specified column
    Mean(String),
    /// Median of values in specified column
    Median(String),
    /// Minimum value in specified column
    Min(String),
    /// Maximum value in specified column
    Max(String),
    /// Standard deviation of values in specified column
    Std(String),
    /// Variance of values in specified column
    Var(String),
    /// First value in specified column (within each group)
    First(String),
    /// Last value in specified column (within each group)
    Last(String),
    /// Collect all values in specified column into a list
    List(String),
    /// Count unique values in specified column
    CountUnique(String),
    /// Concatenate string values in specified column
    StringConcat(String, Option<String>), // column, separator
}

impl AggregationFunction {
    /// Convert to Polars expression
    pub fn to_polars_expr(&self) -> Result<Expr> {
        match self {
            AggregationFunction::Count => Ok(len().alias("count")),
            AggregationFunction::Sum(col_name) => {
                Ok(col(col_name).sum().alias(format!("{col_name}_sum")))
            }
            AggregationFunction::Mean(col_name) => {
                Ok(col(col_name).mean().alias(format!("{col_name}_mean")))
            }
            AggregationFunction::Median(col_name) => {
                Ok(col(col_name).median().alias(format!("{col_name}_median")))
            }
            AggregationFunction::Min(col_name) => {
                Ok(col(col_name).min().alias(format!("{col_name}_min")))
            }
            AggregationFunction::Max(col_name) => {
                Ok(col(col_name).max().alias(format!("{col_name}_max")))
            }
            AggregationFunction::Std(col_name) => {
                Ok(col(col_name).std(1).alias(format!("{col_name}_std")))
            }
            AggregationFunction::Var(col_name) => {
                Ok(col(col_name).var(1).alias(format!("{col_name}_var")))
            }
            AggregationFunction::First(col_name) => {
                Ok(col(col_name).first().alias(format!("{col_name}_first")))
            }
            AggregationFunction::Last(col_name) => {
                Ok(col(col_name).last().alias(format!("{col_name}_last")))
            }
            AggregationFunction::List(col_name) => {
                Ok(col(col_name).alias(format!("{col_name}_list")))
            }
            AggregationFunction::CountUnique(col_name) => Ok(col(col_name)
                .n_unique()
                .alias(format!("{col_name}_nunique"))),
            AggregationFunction::StringConcat(col_name, separator) => {
                let _sep = separator.as_deref().unwrap_or(",");
                // String concatenation in groupby context requires custom aggregation
                // For now, we'll collect into a list and handle concatenation in array processing
                Ok(col(col_name).alias(format!("{col_name}_concat")))
            }
        }
    }

    /// Get the output column name for this aggregation
    #[must_use]
    pub fn output_column_name(&self) -> String {
        match self {
            AggregationFunction::Count => "count".to_string(),
            AggregationFunction::Sum(col_name) => format!("{col_name}_sum"),
            AggregationFunction::Mean(col_name) => format!("{col_name}_mean"),
            AggregationFunction::Median(col_name) => format!("{col_name}_median"),
            AggregationFunction::Min(col_name) => format!("{col_name}_min"),
            AggregationFunction::Max(col_name) => format!("{col_name}_max"),
            AggregationFunction::Std(col_name) => format!("{col_name}_std"),
            AggregationFunction::Var(col_name) => format!("{col_name}_var"),
            AggregationFunction::First(col_name) => format!("{col_name}_first"),
            AggregationFunction::Last(col_name) => format!("{col_name}_last"),
            AggregationFunction::List(col_name) => format!("{col_name}_list"),
            AggregationFunction::CountUnique(col_name) => format!("{col_name}_nunique"),
            AggregationFunction::StringConcat(col_name, _) => format!("{col_name}_concat"),
        }
    }
}

/// Compare values for ordering (used in min/max)
pub(super) fn compare_values_for_ordering(a: &Value, b: &Value) -> std::cmp::Ordering {
    use std::cmp::Ordering;

    match (a, b) {
        (Value::Null, Value::Null) => Ordering::Equal,
        (Value::Null, _) => Ordering::Less,
        (_, Value::Null) => Ordering::Greater,

        (Value::Bool(a), Value::Bool(b)) => a.cmp(b),
        (Value::Int(a), Value::Int(b)) => a.cmp(b),
        (Value::Float(a), Value::Float(b)) => a.partial_cmp(b).unwrap_or(Ordering::Equal),
        (Value::String(a), Value::String(b)) => a.cmp(b),

        // Cross-type numeric comparisons
        #[allow(clippy::cast_precision_loss)]
        (Value::Int(a), Value::Float(b)) => (*a as f64).partial_cmp(b).unwrap_or(Ordering::Equal),
        #[allow(clippy::cast_precision_loss)]
        (Value::Float(a), Value::Int(b)) => a.partial_cmp(&(*b as f64)).unwrap_or(Ordering::Equal),

        // For complex types, compare string representations
        _ => a.to_string().cmp(&b.to_string()),
    }
}

#[cfg(test)]
mod tests;