micromegas_datafusion_extensions/
lib.rs

1#![doc = include_str!("../README.md")]
2/// Unified binary column accessor for Arrow arrays
3pub mod binary_column_accessor;
4/// Binning UDFs (bin_center)
5pub mod binning;
6/// Color UDFs (rgba, lerp_color, color_scale)
7pub mod color;
8/// Compute histograms from SQL
9pub mod histogram;
10/// JSONB support
11pub mod jsonb;
12/// Scalar math UDFs (lerp, unlerp)
13pub mod math;
14/// Property UDFs
15pub mod properties;
16
17use std::sync::Arc;
18
19use binning::bin_center::make_bin_center_udf;
20use color::{
21    color_scale::make_color_scale_udf, lerp_color::make_lerp_color_udf, rgba::make_rgba_udf,
22};
23use datafusion::logical_expr::ScalarUDF;
24use datafusion::prelude::SessionContext;
25use histogram::{
26    accessors::{make_count_from_histogram_udf, make_sum_from_histogram_udf},
27    expand::ExpandHistogramTableFunction,
28    histogram_udaf::make_histo_udaf,
29    quantile::make_quantile_from_histogram_udf,
30    sum_histograms_udaf::sum_histograms_udaf,
31    variance::make_variance_from_histogram_udf,
32};
33use jsonb::{
34    array_elements::JsonbArrayElementsTableFunction,
35    array_length::make_jsonb_array_length_udf,
36    cast::{make_jsonb_as_f64_udf, make_jsonb_as_i64_udf, make_jsonb_as_string_udf},
37    each::JsonbEachTableFunction,
38    format_json::make_jsonb_format_json_udf,
39    get::make_jsonb_get_udf,
40    keys::make_jsonb_object_keys_udf,
41    list_udfs::{make_jsonb_elements_udf, make_jsonb_entries_udf, make_jsonb_path_elements_udf},
42    parse::make_jsonb_parse_udf,
43    path_query::{make_jsonb_path_query_first_udf, make_jsonb_path_query_udf},
44};
45use math::{lerp::make_lerp_udf, unlerp::make_unlerp_udf};
46use properties::{
47    properties_udf::{PropertiesLength, PropertiesToArray},
48    property_get::PropertyGet,
49};
50
51/// Register all extension UDFs on a SessionContext.
52pub fn register_extension_udfs(ctx: &SessionContext) {
53    ctx.register_udaf(make_histo_udaf());
54    ctx.register_udaf(sum_histograms_udaf());
55    ctx.register_udf(make_quantile_from_histogram_udf());
56    ctx.register_udf(make_variance_from_histogram_udf());
57    ctx.register_udf(make_count_from_histogram_udf());
58    ctx.register_udf(make_sum_from_histogram_udf());
59    ctx.register_udtf(
60        "expand_histogram",
61        Arc::new(ExpandHistogramTableFunction::new()),
62    );
63
64    ctx.register_udf(make_jsonb_parse_udf());
65    ctx.register_udf(make_jsonb_format_json_udf());
66    ctx.register_udf(make_jsonb_get_udf());
67    ctx.register_udf(make_jsonb_as_string_udf());
68    ctx.register_udf(make_jsonb_as_f64_udf());
69    ctx.register_udf(make_jsonb_as_i64_udf());
70    ctx.register_udf(make_jsonb_array_length_udf());
71    ctx.register_udf(make_jsonb_object_keys_udf());
72    ctx.register_udf(make_jsonb_path_query_first_udf());
73    ctx.register_udf(make_jsonb_path_query_udf());
74    ctx.register_udf(make_jsonb_entries_udf());
75    ctx.register_udf(make_jsonb_elements_udf());
76    ctx.register_udf(make_jsonb_path_elements_udf());
77    ctx.register_udtf(
78        "jsonb_array_elements",
79        Arc::new(JsonbArrayElementsTableFunction::new()),
80    );
81    ctx.register_udtf("jsonb_each", Arc::new(JsonbEachTableFunction::new()));
82
83    ctx.register_udf(ScalarUDF::from(PropertyGet::new()));
84    ctx.register_udf(ScalarUDF::from(PropertiesToArray::new()));
85    ctx.register_udf(ScalarUDF::from(PropertiesLength::new()));
86
87    ctx.register_udf(make_rgba_udf());
88    ctx.register_udf(make_lerp_color_udf());
89    ctx.register_udf(make_color_scale_udf());
90
91    ctx.register_udf(make_lerp_udf());
92    ctx.register_udf(make_unlerp_udf());
93
94    ctx.register_udf(make_bin_center_udf());
95}
micromegas_datafusion_extensions/lib.rs

micromegas_datafusion_extensions/
lib.rs
