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

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

/// Unpivot a `DataFrame` (convert columns to rows)
///
/// Equivalent to SQL's UNPIVOT operation or pandas' melt function.
///
/// # Examples
///
/// ```rust,ignore
/// use dsq_core::ops::unpivot::unpivot;
/// use dsq_core::value::Value;
///
/// let result = unpivot(
///     &dataframe_value,
///     &["id".to_string()],           // columns to keep as identifiers
///     &["col1".to_string(), "col2".to_string()], // columns to unpivot
///     "variable",                    // name for the variable column
///     "value"                        // name for the value column
/// ).unwrap();
/// ```
pub fn unpivot(
    value: &Value,
    id_columns: &[String],
    value_columns: &[String],
    variable_name: &str,
    value_name: &str,
) -> Result<Value> {
    match value {
        Value::DataFrame(df) => {
            // Use unpivot method from UnpivotDF trait
            let mut unpivoted = if id_columns.is_empty() {
                df.clone()
                    .unpivot([] as [&str; 0], value_columns)
                    .map_err(Error::from)?
            } else {
                df.clone()
                    .unpivot(id_columns, value_columns)
                    .map_err(Error::from)?
            };
            unpivoted
                .rename("variable", variable_name.into())
                .map_err(Error::from)?;
            unpivoted
                .rename("value", value_name.into())
                .map_err(Error::from)?;

            Ok(Value::DataFrame(unpivoted))
        }
        Value::LazyFrame(lf) => {
            let df = lf.clone().collect().map_err(Error::from)?;
            unpivot(
                &Value::DataFrame(df),
                id_columns,
                value_columns,
                variable_name,
                value_name,
            )
        }
        _ => Err(TypeError::UnsupportedOperation {
            operation: "unpivot".to_string(),
            typ: value.type_name().to_string(),
        }
        .into()),
    }
}