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

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

/// Pivot a `DataFrame` (convert rows to columns)
///
/// Equivalent to SQL's PIVOT operation or Excel's pivot tables.
///
/// # Examples
///
/// ```rust,ignore
/// use dsq_core::ops::pivot::pivot;
/// use dsq_core::value::Value;
///
/// let result = pivot(
///     &dataframe_value,
///     &["id".to_string()],           // index columns
///     "category",                     // column to pivot
///     "value",                       // values to aggregate
///     Some("sum")                    // aggregation function
/// ).unwrap();
/// ```
pub fn pivot(
    value: &Value,
    index_columns: &[String],
    _pivot_column: &str,
    value_column: &str,
    agg_function: Option<&str>,
) -> Result<Value> {
    match value {
        Value::DataFrame(df) => {
            let agg_expr = match agg_function {
                Some("sum") => col(value_column).sum().alias("value_sum"),
                Some("mean") => col(value_column).mean().alias("value_mean"),
                Some("count") => col(value_column).count().alias("value_count"),
                Some("min") => col(value_column).min().alias("value_min"),
                Some("max") => col(value_column).max().alias("value_max"),
                Some("first") | None => col(value_column).first().alias("value_first"),
                Some("last") => col(value_column).last().alias("value_last"),
                _ => {
                    return Err(Error::operation(format!(
                        "Unsupported aggregation function: {}",
                        agg_function.unwrap_or("")
                    )));
                }
            };

            // Pivot operation using group_by and aggregation
            // This is a simplified implementation - full pivot would require more complex logic
            let pivoted = df
                .clone()
                .lazy()
                .group_by(index_columns.iter().map(col).collect::<Vec<_>>())
                .agg([agg_expr])
                .collect()
                .map_err(Error::from)?;

            Ok(Value::DataFrame(pivoted))
        }
        Value::LazyFrame(lf) => {
            let agg_expr = match agg_function {
                Some("sum") => col(value_column).sum().alias("value_sum"),
                Some("mean") => col(value_column).mean(),
                Some("count") => col(value_column).count(),
                Some("min") => col(value_column).min(),
                Some("max") => col(value_column).max(),
                Some("first") | None => col(value_column).first(),
                Some("last") => col(value_column).last(),
                _ => {
                    return Err(Error::operation(format!(
                        "Unsupported aggregation function: {}",
                        agg_function.unwrap_or("")
                    )));
                }
            };

            // Pivot operation using group_by and aggregation
            // This is a simplified implementation - full pivot would require more complex logic
            let pivoted = lf
                .clone()
                .group_by(index_columns.iter().map(col).collect::<Vec<_>>())
                .agg([agg_expr]);

            Ok(Value::LazyFrame(Box::new(pivoted)))
        }
        _ => Err(TypeError::UnsupportedOperation {
            operation: "pivot".to_string(),
            typ: value.type_name().to_string(),
        }
        .into()),
    }
}