datafusion_roaring/
lib.rs

1//! Roaring bitmap scalar and aggregate UDFs for Apache DataFusion.
2//!
3//! The SQL API uses Arrow `Binary` for 32-bit [`roaring::RoaringBitmap`] values.
4//! Serialized values carry a versioned header and can be stored in Arrow IPC or Parquet.
5
6#![deny(missing_docs)]
7
8mod aggregate;
9mod codec;
10mod scalar;
11
12use aggregate::{AggregateKind, RoaringAggregate};
13use datafusion_expr::{AggregateUDF, ScalarUDF};
14use scalar::{RoaringScalar, ScalarKind};
15
16pub use codec::{decode_bitmap, encode_bitmap};
17
18/// Builds `roaring_agg(UInt32) -> Binary`.
19pub fn roaring_agg_udaf() -> AggregateUDF {
20    AggregateUDF::from(RoaringAggregate::new("roaring_agg", AggregateKind::Values))
21}
22
23/// Builds `roaring_or_agg(Binary) -> Binary`.
24pub fn roaring_or_agg_udaf() -> AggregateUDF {
25    AggregateUDF::from(RoaringAggregate::new(
26        "roaring_or_agg",
27        AggregateKind::Union,
28    ))
29}
30
31/// Builds every aggregate UDF exported by this crate, including documented aliases.
32pub fn all_udafs() -> Vec<AggregateUDF> {
33    vec![
34        roaring_agg_udaf(),
35        roaring_or_agg_udaf(),
36        AggregateUDF::from(RoaringAggregate::new(
37            "roaring_union_agg",
38            AggregateKind::Union,
39        )),
40    ]
41}
42
43fn scalar(name: &'static str, kind: ScalarKind) -> ScalarUDF {
44    ScalarUDF::from(RoaringScalar::new(name, kind))
45}
46
47/// Builds `roaring_cardinality(Binary) -> UInt64`.
48pub fn roaring_cardinality_udf() -> ScalarUDF {
49    scalar("roaring_cardinality", ScalarKind::Cardinality)
50}
51
52/// Builds `roaring_contains(Binary, UInt32) -> Boolean`.
53pub fn roaring_contains_udf() -> ScalarUDF {
54    scalar("roaring_contains", ScalarKind::Contains)
55}
56
57/// Builds `roaring_union(Binary, Binary) -> Binary`.
58pub fn roaring_union_udf() -> ScalarUDF {
59    scalar("roaring_union", ScalarKind::Union)
60}
61
62/// Builds `roaring_intersection(Binary, Binary) -> Binary`.
63pub fn roaring_intersection_udf() -> ScalarUDF {
64    scalar("roaring_intersection", ScalarKind::Intersection)
65}
66
67/// Builds `roaring_difference(Binary, Binary) -> Binary`.
68pub fn roaring_difference_udf() -> ScalarUDF {
69    scalar("roaring_difference", ScalarKind::Difference)
70}
71
72/// Builds `roaring_xor(Binary, Binary) -> Binary`.
73pub fn roaring_xor_udf() -> ScalarUDF {
74    scalar("roaring_xor", ScalarKind::Xor)
75}
76
77/// Builds every scalar UDF exported by this crate, including documented aliases.
78pub fn all_udfs() -> Vec<ScalarUDF> {
79    use ScalarKind::*;
80    [
81        ("roaring_cardinality", Cardinality),
82        ("roaring_contains", Contains),
83        ("roaring_min", Min),
84        ("roaring_max", Max),
85        ("roaring_rank", Rank),
86        ("roaring_select", Select),
87        ("roaring_is_empty", IsEmpty),
88        ("roaring_serialized_size", SerializedSize),
89        ("roaring_contains_range", ContainsRange),
90        ("roaring_range_cardinality", RangeCardinality),
91        ("roaring_empty", Empty),
92        ("roaring_from_range", FromRange),
93        ("roaring_from_uint_list", FromList),
94        ("roaring_or", Union),
95        ("roaring_union", Union),
96        ("roaring_and", Intersection),
97        ("roaring_intersection", Intersection),
98        ("roaring_xor", Xor),
99        ("roaring_symmetric_difference", Xor),
100        ("roaring_andnot", Difference),
101        ("roaring_difference", Difference),
102        ("roaring_union_cardinality", UnionCardinality),
103        ("roaring_or_cardinality", UnionCardinality),
104        ("roaring_intersection_cardinality", IntersectionCardinality),
105        ("roaring_and_cardinality", IntersectionCardinality),
106        ("roaring_symmetric_difference_cardinality", XorCardinality),
107        ("roaring_xor_cardinality", XorCardinality),
108        ("roaring_difference_cardinality", DifferenceCardinality),
109        ("roaring_andnot_cardinality", DifferenceCardinality),
110        ("roaring_intersects", Intersects),
111        ("roaring_is_disjoint", IsDisjoint),
112        ("roaring_is_subset", IsSubset),
113        ("roaring_is_superset", IsSuperset),
114        ("roaring_insert", Insert),
115        ("roaring_remove", Remove),
116        ("roaring_insert_range", InsertRange),
117        ("roaring_remove_range", RemoveRange),
118        ("roaring_is_valid", IsValid),
119        ("roaring_validate", Validate),
120        ("roaring_to_uint_list", ToList),
121    ]
122    .into_iter()
123    .map(|(name, kind)| scalar(name, kind))
124    .collect()
125}
datafusion_roaring/lib.rs

datafusion_roaring/
lib.rs
