Skip to content

Commit 8cf9bd1

Browse files
committed
feat(cubestore): Add XIRR aggregate function to Cube Store
1 parent d28333a commit 8cf9bd1

File tree

5 files changed

+669
-0
lines changed

5 files changed

+669
-0
lines changed

rust/cubesql/cubesql/src/compile/engine/udf/extension/xirr.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,9 @@ use datafusion::{
1414
scalar::ScalarValue,
1515
};
1616

17+
// Note: A copy/pasted and minimally(?) modified version of this is in Cubestore in udf_xirr.rs, and you might
18+
// want to update both.
19+
1720
pub const XIRR_UDAF_NAME: &str = "xirr";
1821

1922
/// Creates a XIRR Aggregate UDF.

rust/cubestore/cubestore-sql-tests/src/tests.rs

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,7 @@ pub fn sql_tests() -> Vec<(&'static str, TestFn)> {
134134
t("hyperloglog_postgres", hyperloglog_postgres),
135135
t("hyperloglog_snowflake", hyperloglog_snowflake),
136136
t("hyperloglog_databricks", hyperloglog_databricks),
137+
t("xirr", xirr),
137138
t(
138139
"aggregate_index_hll_databricks",
139140
aggregate_index_hll_databricks,
@@ -2802,6 +2803,58 @@ async fn hyperloglog_databricks(service: Box<dyn SqlClient>) {
28022803
assert_eq!(to_rows(&r), rows(&[(1, 4), (2, 4), (3, 20)]));
28032804
}
28042805

2806+
async fn xirr(service: Box<dyn SqlClient>) {
2807+
// XIRR result may differ between platforms, so we truncate the result
2808+
let r = service
2809+
.exec_query(
2810+
r#"
2811+
SELECT LEFT(XIRR(payment, date)::varchar, 10) AS xirr
2812+
FROM (
2813+
SELECT '2014-01-01'::date AS date, -10000.0 AS payment
2814+
UNION ALL
2815+
SELECT '2014-03-01'::date AS date, 2750.0 AS payment
2816+
UNION ALL
2817+
SELECT '2014-10-30'::date AS date, 4250.0 AS payment
2818+
UNION ALL
2819+
SELECT '2015-02-15'::date AS date, 3250.0 AS payment
2820+
UNION ALL
2821+
SELECT '2015-04-01'::date AS date, 2750.0 AS payment
2822+
) AS "t"
2823+
"#,
2824+
)
2825+
.await
2826+
.unwrap();
2827+
2828+
assert_eq!(to_rows(&r), rows(&["0.37485859"]));
2829+
2830+
let r = service
2831+
.exec_query(
2832+
r#"
2833+
SELECT LEFT(XIRR(payment, date)::varchar, 10) AS xirr
2834+
FROM (
2835+
SELECT '2014-01-01'::date AS date, -10000.0 AS payment
2836+
) AS "t"
2837+
WHERE 0 = 1
2838+
"#,
2839+
)
2840+
.await
2841+
.unwrap_err();
2842+
assert_eq!(r.elide_backtrace(), CubeError::internal("Arrow error: External error: Execution error: A result for XIRR couldn't be determined because the arguments are empty".to_owned()));
2843+
2844+
let r = service
2845+
.exec_query(
2846+
r#"
2847+
SELECT LEFT(XIRR(payment, date)::varchar, 10) AS xirr
2848+
FROM (
2849+
SELECT '2014-01-01'::date AS date, 10000.0 AS payment
2850+
) AS "t"
2851+
"#,
2852+
)
2853+
.await
2854+
.unwrap_err();
2855+
assert_eq!(r.elide_backtrace(), CubeError::internal("Arrow error: External error: Execution error: The XIRR function couldn't find a solution".to_owned()));
2856+
}
2857+
28052858
async fn aggregate_index_hll_databricks(service: Box<dyn SqlClient>) {
28062859
service.exec_query("CREATE SCHEMA s").await.unwrap();
28072860
service

rust/cubestore/cubestore/src/queryplanner/mod.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ pub mod now;
2323
pub mod providers;
2424
#[cfg(test)]
2525
mod test_utils;
26+
pub mod udf_xirr;
2627
pub mod udfs;
2728

2829
use crate::cachestore::CacheStore;
@@ -429,6 +430,7 @@ impl ContextProvider for MetaStoreSchemaProvider {
429430
// TODO: case-insensitive names.
430431
let kind = match name {
431432
"merge" | "MERGE" => CubeAggregateUDFKind::MergeHll,
433+
"xirr" | "XIRR" => CubeAggregateUDFKind::Xirr,
432434
_ => return None,
433435
};
434436
return Some(Arc::new(aggregate_udf_by_kind(kind).descriptor()));

0 commit comments

Comments
 (0)