ndatafusion 0.1.1

Extensions and support for linear algebra in DataFusion
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61

//! Public aggregate-UDF constructors and catalog assembly helpers.
//!
//! Use this module when you need direct access to individual aggregate UDF definitions or when you
//! want to register only the grouped aggregate surface instead of calling
//! [`crate::register_all`].

use std::sync::Arc;

use datafusion::logical_expr::AggregateUDF;

pub use crate::udaf::{
    linear_regression_fit_udaf, vector_correlation_agg_udaf, vector_covariance_agg_udaf,
    vector_pca_fit_udaf,
};

/// Return the aggregate catalog currently registered by [`crate::register_all`].
#[must_use]
pub fn all_default_aggregates() -> Vec<Arc<AggregateUDF>> {
    vec![
        vector_covariance_agg_udaf(),
        vector_correlation_agg_udaf(),
        vector_pca_fit_udaf(),
        linear_regression_fit_udaf(),
    ]
}

#[cfg(test)]
mod tests {
    use super::{
        all_default_aggregates, linear_regression_fit_udaf, vector_correlation_agg_udaf,
        vector_covariance_agg_udaf, vector_pca_fit_udaf,
    };

    #[test]
    fn default_udaf_catalog_matches_current_surface() {
        assert_eq!(all_default_aggregates().len(), 4);
    }

    #[test]
    fn representative_udafs_expose_expected_parameter_names() {
        assert_eq!(
            linear_regression_fit_udaf().signature().parameter_names.as_deref(),
            Some(
                ["design".to_string(), "response".to_string(), "add_intercept".to_string(),]
                    .as_slice()
            )
        );
    }

    #[test]
    fn documented_udafs_expose_documentation() {
        for udaf in [
            vector_covariance_agg_udaf(),
            vector_correlation_agg_udaf(),
            vector_pca_fit_udaf(),
            linear_regression_fit_udaf(),
        ] {
            assert!(udaf.documentation().is_some(), "missing docs for {}", udaf.name());
        }
    }
}