dsq_functions/builtin/
base64_encode.rs

1use base64::{engine::general_purpose, Engine};
2use dsq_shared::value::Value;
3use dsq_shared::Result;
4use polars::prelude::*;
5use std::borrow::Cow;
6
7pub fn builtin_base64_encode(args: &[Value]) -> Result<Value> {
8    if args.len() != 1 {
9        return Err(dsq_shared::error::operation_error(
10            "base64_encode() expects 1 argument",
11        ));
12    }
13
14    match &args[0] {
15        Value::String(s) => Ok(Value::String(
16            general_purpose::STANDARD.encode(s.as_bytes()),
17        )),
18        Value::Array(arr) => {
19            let encoded: Result<Vec<Value>> = arr
20                .iter()
21                .map(|v| match v {
22                    Value::String(s) => Ok(Value::String(
23                        general_purpose::STANDARD.encode(s.as_bytes()),
24                    )),
25                    _ => Ok(Value::String(
26                        general_purpose::STANDARD.encode(v.to_string().as_bytes()),
27                    )),
28                })
29                .collect();
30            Ok(Value::Array(encoded?))
31        }
32        Value::DataFrame(df) => {
33            let mut new_series = Vec::new();
34            for col_name in df.get_column_names() {
35                if let Ok(series) = df.column(col_name) {
36                    if series.dtype() == &DataType::String {
37                        let utf8_series = series.str().map_err(|e| {
38                            dsq_shared::error::operation_error(format!(
39                                "base64_encode() failed to cast series to utf8: {}",
40                                e
41                            ))
42                        })?;
43                        let encoded_series = utf8_series
44                            .apply(|s| {
45                                s.map(|s| {
46                                    Cow::Owned(general_purpose::STANDARD.encode(s.as_bytes()))
47                                })
48                            })
49                            .into_series();
50                        let mut s = encoded_series;
51                        s.rename(col_name.clone());
52                        new_series.push(s.into());
53                    } else {
54                        let mut s = series.clone();
55                        s.rename(col_name.clone());
56                        new_series.push(s);
57                    }
58                }
59            }
60            match DataFrame::new(new_series) {
61                Ok(new_df) => Ok(Value::DataFrame(new_df)),
62                Err(e) => Err(dsq_shared::error::operation_error(format!(
63                    "base64_encode() failed on DataFrame: {}",
64                    e
65                ))),
66            }
67        }
68        Value::Series(series) => {
69            if series.dtype() == &DataType::String {
70                let utf8_series = series.str().map_err(|e| {
71                    dsq_shared::error::operation_error(format!(
72                        "base64_encode() failed to cast series to utf8: {}",
73                        e
74                    ))
75                })?;
76                let encoded_series = utf8_series
77                    .apply(|s| {
78                        s.map(|s| Cow::Owned(general_purpose::STANDARD.encode(s.as_bytes())))
79                    })
80                    .into_series();
81                Ok(Value::Series(encoded_series))
82            } else {
83                Ok(Value::Series(series.clone()))
84            }
85        }
86        _ => Err(dsq_shared::error::operation_error(
87            "base64_encode() requires string, array, DataFrame, or Series",
88        )),
89    }
90}
91
92inventory::submit! {
93    crate::FunctionRegistration {
94        name: "base64_encode",
95        func: builtin_base64_encode,
96    }
97}
dsq_functions/builtin/base64_encode.rs

dsq_functions/builtin/
base64_encode.rs
