Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
up
Browse files Browse the repository at this point in the history
tgross35 committed May 8, 2024
1 parent bed506f commit 10b6bc7
Showing 1 changed file with 8 additions and 50 deletions.
58 changes: 8 additions & 50 deletions udf-examples/src/avg_cost.rs
Original file line number Diff line number Diff line change
@@ -11,8 +11,6 @@
#![allow(clippy::cast_precision_loss)]

use std::fmt::Display;

use udf::prelude::*;

#[derive(Debug, Default, PartialEq)]
@@ -22,37 +20,6 @@ struct AvgCost {
total_price: f64,
}

/// We just use this to show one way of handling errors a bit cleaner. Having an
/// error type and implementing `Display` for it is quite common (crates like
/// `anyhow` and `thiserror` make it more straightforward.)
///
/// This is `pub` because we reuse it for `avg2`
pub enum Errors<'a> {
WrongArgCount(usize),
FirstArgType(&'a SqlArg<'a, Init>),
SecondArgType(&'a SqlArg<'a, Init>),
}

impl<'a> Display for Errors<'a> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::WrongArgCount(n) => write!(f, "This function takes two arguments; got {n}"),
Self::FirstArgType(a) => write!(
f,
"First argument must be an integer; received {} {}",
a.value().display_name(),
a.attribute()
),
Self::SecondArgType(a) => write!(
f,
"Second argument must be an integer; received {} {}",
a.value().display_name(),
a.attribute()
),
}
}
}

#[register]
impl BasicUdf for AvgCost {
type Returns<'a> = Option<f64>
@@ -61,18 +28,14 @@ impl BasicUdf for AvgCost {

fn init(cfg: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
if args.len() != 2 {
return Err(Errors::WrongArgCount(args.len()).to_string());
return Err(format!("expected two arguments; got {}", args.len()));
}

let a0 = args.get(0).unwrap();
let a1 = args.get(1).unwrap();
let mut a0 = args.get(0).unwrap();
let mut a1 = args.get(1).unwrap();

if !a0.value().is_int() {
return Err(Errors::FirstArgType(&a0).to_string());
}
if !a1.value().is_real() {
return Err(Errors::SecondArgType(&a1).to_string());
}
a0.set_type_coercion(SqlType::Int);
a1.set_type_coercion(SqlType::Real);

cfg.set_maybe_null(true);
cfg.set_decimals(10);
@@ -91,7 +54,7 @@ impl BasicUdf for AvgCost {
if self.count == 0 || self.total_qty == 0 {
return Ok(None);
}
Ok(Some(self.total_price / self.total_qty as f64))
dbg!(Ok(Some(self.total_price / self.total_qty as f64)))
}
}

@@ -114,13 +77,8 @@ impl AggregateUdf for AvgCost {
_error: Option<NonZeroU8>,
) -> Result<(), NonZeroU8> {
// We can unwrap because we are guaranteed to have 2 args (from checks in init)
let Some(in_qty) = args.get(0).unwrap().value().as_int() else {
return Ok(());
};

let Some(mut price) = args.get(1).unwrap().value().as_real() else {
return Ok(());
};
let in_qty = args.get(0).unwrap().value().as_int().unwrap();
let mut price = args.get(1).unwrap().value().as_real().unwrap();

dbg!(&in_qty, &price, &self);

0 comments on commit 10b6bc7

Please sign in to comment.