pub mod generator;
pub mod client_sql;
use std::any;
use polars::datatypes::AnyValue;
pub mod config;
pub mod map_schema;
pub use client_sql::{ClientSql, TabGoes};
pub use config::ConfigClient;
pub type Result<T> = anyhow::Result<T>;
pub use chrono;
pub use rust_decimal;
pub use uuid;
use chrono::Days;
use polars::prelude::*;
#[derive(Debug, Clone, Copy, Default)]
pub enum IfExists {
#[default]
Append,
Replace,
Fail,
}
impl TabGoes {
pub fn join_inner(mut self, other: TabGoes, left_cols: &[&str], right_cols: &[&str]) -> Self {
let left_exprs: Vec<Expr> = left_cols.iter().map(|&c| col(c)).collect();
let right_exprs: Vec<Expr> = right_cols.iter().map(|&c| col(c)).collect();
self.lazy = self.lazy.join(
other.lazy,
left_exprs,
right_exprs,
JoinArgs::new(JoinType::Inner)
);
self._files.extend(other._files);
self
}
pub fn join_outer(mut self, other: TabGoes, left_cols: &[&str], right_cols: &[&str]) -> Self {
let left_exprs: Vec<Expr> = left_cols.iter().map(|&c| col(c)).collect();
let right_exprs: Vec<Expr> = right_cols.iter().map(|&c| col(c)).collect();
self.lazy = self.lazy.join(
other.lazy,
left_exprs,
right_exprs,
JoinArgs::new(JoinType::Left)
);
self._files.extend(other._files);
self
}
pub fn union(mut self, other: TabGoes)->Result<Self>{
self.lazy = concat(&[self.lazy, other.lazy], UnionArgs::default())?;
self._files.extend(other._files);
Ok(self)
}
pub fn append_text(mut self, column_name: &str, text: &str, in_final: bool, alias: &str) -> Self {
use polars::prelude::*;
let expr = if in_final {
concat_str([col(column_name), lit(text)], "", true)
} else {
concat_str([lit(text), col(column_name)], "", true)
};
self.lazy = self.lazy.with_columns([
expr.alias(alias)
]);
self
}
pub fn concat_columns(mut self, columns: &[&str], separator: &str, alias: &str) -> Self {
use polars::prelude::*;
let exprs: Vec<Expr> = columns.iter().map(|c| col(*c)).collect();
self.lazy = self.lazy.with_columns([
concat_str(exprs, separator, true).alias(alias)
]);
self
}
pub fn split_column_auto(mut self, column_name: &str, separator: &str, num_cols: usize) -> Self {
use polars::prelude::*;
let mut exprs = Vec::new();
for i in 0..num_cols {
let alias = if i == 0 {
column_name.to_string()
} else {
format!("{}{}", column_name, i)
};
let expr = col(column_name)
.str()
.split(lit(separator))
.list()
.get(lit(i as i64), true)
.alias(&alias);
exprs.push(expr);
}
self.lazy = self.lazy.with_columns(exprs);
self
}
pub fn explode_column(mut self, column_name: &str, separator: &str) -> Self {
use polars::prelude::*;
self.lazy = self.lazy.with_columns([
col(column_name).str().split(lit(separator)).alias(column_name)
]);
self.lazy = self.lazy.explode([col(column_name)]);
self
}
pub fn order(mut self, cols: &[&str])-> Self{
let exprs: Vec<Expr> = cols.iter().map(|&c| col(c)).collect();
self.lazy = self.lazy.select(&exprs);
self
}
pub fn pivot_column(
mut self,
index_columns: &[&str], pivot_col: &str,
value_col: &str,
) -> anyhow::Result<Self> {
use polars::prelude::*;
let df = self.lazy.collect()?;
let idx_cols: Vec<&str> = index_columns.to_vec();
let p_cols: Vec<&str> = vec![pivot_col];
let v_cols: Vec<&str> = vec![value_col];
let pivoted_df = polars_lazy::frame::pivot::pivot_stable(
&df,
idx_cols,
Some(p_cols),
Some(v_cols),
true,
Some(col(value_col).first()),
None
)?;
self.lazy = pivoted_df.lazy();
Ok(self)
}
pub fn sort_rows(mut self, column_name: &str, decrescente: bool) -> Self {
let options = SortMultipleOptions::default().with_order_descending(decrescente);
self.lazy = self.lazy.sort([column_name], options);
self
}
pub fn sort_rows_multiple(mut self, columns: &[&str], decrescente: bool) -> Self {
let options = SortMultipleOptions::default().with_order_descending(decrescente);
self.lazy = self.lazy.sort(columns.to_vec(), options);
self
}
pub async fn print(&self) -> Result<()> {
let lazy = self.lazy.clone();
let df = tokio::task::spawn_blocking(move || {
lazy.collect()
}).await??;
println!("{}", df);
Ok(())
}
pub fn join_asof(mut self, other: TabGoes, col_name: &str, margem_maxima: i64) -> Self {
let left_lazy = self.lazy.sort([col_name], SortMultipleOptions::default());
let right_lazy = other.lazy.sort([col_name], SortMultipleOptions::default());
let asof_options = AsOfOptions {
strategy: AsofStrategy::Nearest,
tolerance: Some(AnyValue::Int64(margem_maxima)),
tolerance_str: None,
left_by: None,
right_by: None,
allow_eq: true,
check_sortedness: false,
};
self.lazy = left_lazy.join(
right_lazy,
[col(col_name)],
[col(col_name)],
JoinArgs::new(JoinType::AsOf(asof_options))
);
self._files.extend(other._files);
self
}
pub fn row_number(mut self, coluna_particao: &str, alias: &str) -> Self {
self.lazy = self.lazy.with_columns([
col(coluna_particao)
.cum_count(false)
.over([col(coluna_particao)])
.alias(alias)
]);
self
}
pub fn filter<F>(mut self, predicate: F) -> Self
where F: FnOnce() -> Expr {
self.lazy = self.lazy.filter(predicate());
self
}
pub fn cast_to_string(mut self, column_name: &str) -> Self {
self.lazy = self.lazy.with_columns([
col(column_name).cast(DataType::String)
]);
self
}
pub fn cast_to_date(mut self, column_name: &str) -> Self {
self.lazy = self.lazy.with_columns([
col(column_name).cast(DataType::Date)
]);
self
}
pub fn cast_to_time(mut self, column_name: &str) -> Self {
self.lazy = self.lazy.with_columns([
col(column_name).cast(DataType::Time)
]);
self
}
pub fn add_columns(mut self, col1: &str, col2: &str, alias: &str) -> Self {
self.lazy = self.lazy.with_columns([(col(col1) + col(col2)).alias(alias)]);
self
}
pub fn sub_columns(mut self, col1: &str, col2: &str, alias: &str) -> Self {
self.lazy = self.lazy.with_columns([(col(col1) - col(col2)).alias(alias)]);
self
}
pub fn add_days(mut self, column_name: &str, days: i64, alias: &str) -> Self {
let ms = days * 86_400_000;
self.lazy = self.lazy.with_columns([
(col(column_name) + lit(ms).cast(DataType::Duration(TimeUnit::Milliseconds))).alias(alias)
]);
self
}
pub fn add_hours(mut self, column_name: &str, hours: i64, alias: &str) -> Self {
let ms = hours * 3_600_000;
self.lazy = self.lazy.with_columns([
(col(column_name) + lit(ms).cast(DataType::Duration(TimeUnit::Milliseconds))).alias(alias)
]);
self
}
pub fn diff_days(mut self, col_end: &str, col_start: &str, alias: &str) -> Self {
self.lazy = self.lazy.with_columns([
(col(col_end) - col(col_start)).dt().total_days().alias(alias)
]);
self
}
pub fn diff_hours(mut self, col_end: &str, col_start: &str, alias: &str) -> Self {
self.lazy = self.lazy.with_columns([
(col(col_end) - col(col_start)).dt().total_hours().alias(alias)
]);
self
}
pub fn drop_column(mut self, column_name: &str) -> Self {
self.lazy = self.lazy.drop([column_name]);
self
}
pub fn drop_columns(mut self, columns: &[&str]) -> Self {
self.lazy = self.lazy.drop(columns.to_vec());
self
}
pub fn diff_minutes(mut self, col_end: &str, col_start: &str, alias: &str) -> Self {
self.lazy = self.lazy.with_columns([
(col(col_end) - col(col_start)).dt().total_minutes().alias(alias)
]);
self
}
pub fn diff_seconds(mut self, col_end: &str, col_start: &str, alias: &str) -> Self {
self.lazy = self.lazy.with_columns([
(col(col_end) - col(col_start)).dt().total_seconds().alias(alias)
]);
self
}
pub fn case_custom(self, new_column: &str) -> CaseCustomBuilder {
CaseCustomBuilder {
tab: self,
new_col: new_column.to_string(),
conditions: Vec::new(),
}
}
pub fn mul_columns(mut self, col1: &str, col2: &str, alias: &str) -> Self {
self.lazy = self.lazy.with_columns([
(col(col1) * col(col2)).alias(alias)
]);
self
}
pub fn extract_hour(mut self, column_name: &str, alias: &str) -> Self {
self.lazy = self.lazy.with_columns([
col(column_name).dt().hour().alias(alias)
]);
self
}
pub fn drop_duplicates(mut self)->Self{
use polars::prelude::*;
self.lazy = self.lazy.unique(None, UniqueKeepStrategy::First);
self
}
pub fn drop_duplicates_by(mut self, columns: &[&str])->Self{
use polars::prelude::*;
let cols: Vec<String> = columns.iter().map(|&s|s.to_string()).collect();
self.lazy = self.lazy.unique(Some(cols), UniqueKeepStrategy::First);
self
}
pub fn extract_minute(mut self, column_name: &str, alias: &str) -> Self {
self.lazy = self.lazy.with_columns([
col(column_name).dt().minute().alias(alias)
]);
self
}
pub fn extract_second(mut self, column_name: &str, alias: &str) -> Self {
self.lazy = self.lazy.with_columns([
col(column_name).dt().second().alias(alias)
]);
self
}
pub fn div_columns(mut self, col1: &str, col2: &str, alias: &str) -> Self {
self.lazy = self.lazy.with_columns([
(col(col1) / col(col2)).alias(alias)
]);
self
}
pub fn substring(mut self, column_name: &str,start: i64, length: u32, alias: &str)-> Self{
self.lazy = self.lazy.with_columns([
col(column_name).str().slice(lit(start), lit(length)).alias(alias)
]);
self
}
pub fn substring_from(mut self, column_name: &str,start: i64, alias: &str)-> Self{
self.lazy = self.lazy.with_columns([
col(column_name).str().slice(lit(start), lit(u32::MAX)).alias(alias)
]);
self
}
pub fn case_int(self, new_column: &str, cond_column: &str) -> CaseIntBuilder {
CaseIntBuilder {
tab: self,
new_col: new_column.to_string(),
cond_col: cond_column.to_string(),
conditions: Vec::new(),
}
}
pub fn case_str(self, new_column: &str, cond_column: &str) -> CaseStrBuilder {
CaseStrBuilder {
tab: self,
new_col: new_column.to_string(),
cond_col: cond_column.to_string(),
conditions: Vec::new(),
}
}
pub fn add_constant_str(mut self, column_name: &str, value: &str) -> Self {
use polars::prelude::*;
self.lazy = self.lazy.with_columns([
lit(value).alias(column_name)
]);
self
}
pub fn copy_column(mut self, original_column: &str, new_column: &str) -> Self {
use polars::prelude::*;
self.lazy = self.lazy.with_columns([
col(original_column).alias(new_column)
]);
self
}
pub fn add_constant_int(mut self, column_name: &str, value: i64) -> Self {
use polars::prelude::*;
self.lazy = self.lazy.with_columns([
lit(value).alias(column_name)
]);
self
}
pub fn add_empty_column(mut self, column_name: &str) -> Self {
use polars::prelude::*;
self.lazy = self.lazy.with_columns([
lit(Null {}).cast(DataType::Null).cast(DataType::String).cast(DataType::String).alias(column_name)
]);
self
}
pub fn mul_by_number(mut self, column_name: &str, multiplier: f64, alias: &str) -> Self {
self.lazy = self.lazy.with_columns([
(col(column_name) * lit(multiplier)).alias(alias)
]);
self
}
pub fn div_by_number(mut self, column_name: &str, multiplier: f64, alias: &str) -> Self {
self.lazy = self.lazy.with_columns([
(col(column_name) / lit(multiplier)).alias(alias)
]);
self
}
pub fn rename_column(mut self, old_name: &str, new_name: &str) -> Self {
self.lazy = self.lazy.rename([old_name], [new_name], true);
self
}
pub async fn process_and_save(&self, path: String) -> Result<()> {
let lazy_clone = self.lazy.clone();
tokio::task::spawn_blocking(move || {
let mut df = lazy_clone.collect()?;
let mut file = std::fs::File::create(&path)?;
ParquetWriter::new(&mut file).finish(&mut df)?;
Ok::<(), anyhow::Error>(())
}).await??;
Ok(())
}
}
pub fn get_today_datetime(offset_days: i64) -> String {
let today = chrono::Local::now().naive_local();
let target = if offset_days >= 0 {
today.checked_add_days(Days::new(offset_days as u64)).unwrap_or(today)
} else {
today.checked_sub_days(Days::new(offset_days.unsigned_abs())).unwrap_or(today)
};
target.format("%Y-%m-%d %H:%M:%S").to_string()
}
pub fn get_today_date(offset_days: i64) -> String {
let today = chrono::Local::now().naive_local();
let target = if offset_days >= 0 {
today.checked_add_days(Days::new(offset_days as u64)).unwrap_or(today)
} else {
today.checked_sub_days(Days::new(offset_days.unsigned_abs())).unwrap_or(today)
};
target.format("%Y-%m-%d").to_string()
}
pub struct CaseIntBuilder {
tab: TabGoes,
new_col: String,
cond_col: String,
conditions: Vec<(i64, i64)>,
}
impl CaseIntBuilder {
pub fn when(mut self, if_val: i64, then_val: i64) -> Self {
self.conditions.push((if_val, then_val));
self
}
pub fn otherwise(mut self, else_val: i64) -> TabGoes {
use polars::prelude::*;
let mut expr = lit(else_val);
for (if_val, then_val) in self.conditions.iter().rev() {
expr = when(col(&self.cond_col).eq(lit(*if_val)))
.then(lit(*then_val))
.otherwise(expr);
}
self.tab.lazy = self.tab.lazy.with_columns([expr.alias(&self.new_col)]);
self.tab
}
}
pub struct CaseStrBuilder {
tab: TabGoes,
new_col: String,
cond_col: String,
conditions: Vec<(String, String)>,
}
impl CaseStrBuilder {
pub fn when(mut self, if_val: &str, then_val: &str) -> Self {
self.conditions.push((if_val.to_string(), then_val.to_string()));
self
}
pub fn otherwise(mut self, else_val: &str) -> TabGoes {
use polars::prelude::*;
let mut expr = lit(else_val);
for (if_val, then_val) in self.conditions.iter().rev() {
expr = when(col(&self.cond_col).eq(lit(if_val.as_str())))
.then(lit(then_val.as_str()))
.otherwise(expr);
}
self.tab.lazy = self.tab.lazy.with_columns([expr.alias(&self.new_col)]);
self.tab
}
}
pub trait AsExpr {
fn as_expr(self) -> polars::prelude::Expr;
}
impl AsExpr for polars::prelude::Expr {
fn as_expr(self) -> polars::prelude::Expr { self }
}
macro_rules! impl_as_expr_lit {
($($t:ty),*) => { $( impl AsExpr for $t { fn as_expr(self) -> polars::prelude::Expr { polars::prelude::lit(self) } } )* };
}
impl_as_expr_lit!(&str, String, i8, i16, i32, i64, u8, u16, u32, u64, f32, f64, bool);
pub struct CaseCustomBuilder {
tab: TabGoes,
new_col: String,
conditions: Vec<(polars::prelude::Expr, polars::prelude::Expr)>,
}
pub struct CaseCustomWhen {
builder: CaseCustomBuilder,
if_expr: polars::prelude::Expr,
}
impl CaseCustomBuilder {
pub fn when<E: AsExpr>(self, if_expr: E) -> CaseCustomWhen {
CaseCustomWhen { builder: self, if_expr: if_expr.as_expr() }
}
pub fn otherwise<E: AsExpr>(self, else_expr: E) -> TabGoes {
use polars::prelude::*;
let mut final_expr = else_expr.as_expr();
for (if_expr, then_expr) in self.conditions.into_iter().rev() {
final_expr = when(if_expr).then(then_expr).otherwise(final_expr);
}
let mut tab = self.tab;
tab.lazy = tab.lazy.with_columns([final_expr.alias(&self.new_col)]);
tab
}
}
impl CaseCustomWhen {
pub fn then<E: AsExpr>(mut self, then_expr: E) -> CaseCustomBuilder {
self.builder.conditions.push((self.if_expr, then_expr.as_expr()));
self.builder
}
}