micromegas_datafusion_extensions/
lib.rs

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

micromegas_datafusion_extensions/
lib.rs
