dsq_functions/builtin/
base58_encode.rs

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