micromegas_datafusion_extensions/
lib.rs

1/// Unified binary column accessor for Arrow arrays
2pub mod binary_column_accessor;
3/// Compute histograms from SQL
4pub mod histogram;
5/// JSONB support
6pub mod jsonb;
7/// Property UDFs
8pub mod properties;
9
10use std::sync::Arc;
11
12use datafusion::logical_expr::ScalarUDF;
13use datafusion::prelude::SessionContext;
14use histogram::{
15    accessors::{make_count_from_histogram_udf, make_sum_from_histogram_udf},
16    expand::ExpandHistogramTableFunction,
17    histogram_udaf::make_histo_udaf,
18    quantile::make_quantile_from_histogram_udf,
19    sum_histograms_udaf::sum_histograms_udaf,
20    variance::make_variance_from_histogram_udf,
21};
22use jsonb::{
23    array_elements::JsonbArrayElementsTableFunction,
24    array_length::make_jsonb_array_length_udf,
25    cast::{make_jsonb_as_f64_udf, make_jsonb_as_i64_udf, make_jsonb_as_string_udf},
26    each::JsonbEachTableFunction,
27    format_json::make_jsonb_format_json_udf,
28    get::make_jsonb_get_udf,
29    keys::make_jsonb_object_keys_udf,
30    parse::make_jsonb_parse_udf,
31    path_query::{make_jsonb_path_query_first_udf, make_jsonb_path_query_udf},
32};
33use properties::{
34    properties_udf::{PropertiesLength, PropertiesToArray},
35    property_get::PropertyGet,
36};
37
38/// Register all extension UDFs on a SessionContext.
39pub fn register_extension_udfs(ctx: &SessionContext) {
40    ctx.register_udaf(make_histo_udaf());
41    ctx.register_udaf(sum_histograms_udaf());
42    ctx.register_udf(make_quantile_from_histogram_udf());
43    ctx.register_udf(make_variance_from_histogram_udf());
44    ctx.register_udf(make_count_from_histogram_udf());
45    ctx.register_udf(make_sum_from_histogram_udf());
46    ctx.register_udtf(
47        "expand_histogram",
48        Arc::new(ExpandHistogramTableFunction::new()),
49    );
50
51    ctx.register_udf(make_jsonb_parse_udf());
52    ctx.register_udf(make_jsonb_format_json_udf());
53    ctx.register_udf(make_jsonb_get_udf());
54    ctx.register_udf(make_jsonb_as_string_udf());
55    ctx.register_udf(make_jsonb_as_f64_udf());
56    ctx.register_udf(make_jsonb_as_i64_udf());
57    ctx.register_udf(make_jsonb_array_length_udf());
58    ctx.register_udf(make_jsonb_object_keys_udf());
59    ctx.register_udf(make_jsonb_path_query_first_udf());
60    ctx.register_udf(make_jsonb_path_query_udf());
61    ctx.register_udtf(
62        "jsonb_array_elements",
63        Arc::new(JsonbArrayElementsTableFunction::new()),
64    );
65    ctx.register_udtf("jsonb_each", Arc::new(JsonbEachTableFunction::new()));
66
67    ctx.register_udf(ScalarUDF::from(PropertyGet::new()));
68    ctx.register_udf(ScalarUDF::from(PropertiesToArray::new()));
69    ctx.register_udf(ScalarUDF::from(PropertiesLength::new()));
70}
micromegas_datafusion_extensions/lib.rs

micromegas_datafusion_extensions/
lib.rs
