micromegas_analytics/properties/property_set_jsonb_dictionary_builder.rs
1use crate::arrow_properties::serialize_property_set_to_jsonb;
2use crate::properties::property_set::PropertySet;
3use anyhow::Result;
4use datafusion::arrow::array::{BinaryArray, DictionaryArray, Int32Array};
5use datafusion::arrow::datatypes::Int32Type;
6use datafusion::common::DataFusionError;
7use std::collections::HashMap;
8use std::sync::Arc;
9
10/// A wrapper around raw pointers that implements Send/Sync for use in HashMap keys.
11///
12/// This is safe because:
13/// 1. We only use the pointer for identity comparison (equality/hashing)
14/// 2. We never dereference the pointer
15/// 3. The parse arena keeps the underlying objects alive for the whole block,
16/// so addresses stay stable and unique while comparisons happen; after the
17/// block is parsed the pointers are never touched again.
18/// 4. The cache is scoped to single block processing (no cross-thread sharing)
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
20struct ObjectPointer(*const ());
21
22unsafe impl Send for ObjectPointer {}
23unsafe impl Sync for ObjectPointer {}
24
25/// Custom dictionary builder for PropertySet → JSONB encoding with pointer-based deduplication.
26///
27/// This builder eliminates redundant JSONB serialization and dictionary hash lookups
28/// for duplicate PropertySets by using PropertySet's `Arc<Object>` pointer addresses as keys.
29///
30/// Performance benefits over Arrow's BinaryDictionaryBuilder:
31/// - Eliminates content-based hashing: Arrow's builder hashes JSONB bytes for deduplication
32/// - Pointer-based deduplication: O(1) pointer comparison vs O(n) content hash
33/// - Serialization only when needed: Only serialize PropertySet on first encounter
34/// - Memory efficiency: Shared PropertySet references, single JSONB copy per unique set
35///
36/// # Invariant (must not outlive a single parse arena)
37///
38/// The pointer keys are only unique while every appended `PropertySet` borrows the
39/// same parse arena. A dropped arena's address can be recycled by the next block's
40/// arena, so a builder instance must be fed exactly one block / one arena: construct
41/// it fresh in each block processor and `finish()` it before the arena is dropped.
42/// Reusing one builder across arenas would let a recycled address alias a stale
43/// entry and emit the wrong JSONB. `append_property_set` debug-asserts this
44/// invariant so any future cross-arena reuse fails loudly in debug/test builds.
45pub struct PropertySetJsonbDictionaryBuilder {
46 /// Maps `Arc<Object>` pointer to dictionary index (avoids content hashing)
47 pointer_to_index: HashMap<ObjectPointer, i32>,
48 /// Pre-serialized JSONB values in dictionary
49 jsonb_values: Vec<Vec<u8>>,
50 /// Dictionary keys (indices) for each appended entry
51 keys: Vec<Option<i32>>,
52}
53
54impl PropertySetJsonbDictionaryBuilder {
55 /// Create a new builder with the specified capacity hint
56 pub fn new(capacity: usize) -> Self {
57 Self {
58 pointer_to_index: HashMap::with_capacity(capacity),
59 jsonb_values: Vec::with_capacity(capacity),
60 keys: Vec::with_capacity(capacity),
61 }
62 }
63
64 /// Append PropertySet using pointer-based deduplication
65 ///
66 /// For cache hits: reuses existing dictionary index (no serialization)
67 /// For cache misses: serializes once and stores in dictionary
68 pub fn append_property_set(&mut self, property_set: &PropertySet<'_>) -> Result<()> {
69 let ptr = ObjectPointer(property_set.object_ptr());
70
71 match self.pointer_to_index.get(&ptr) {
72 Some(&index) => {
73 // Cache hit: reuse existing dictionary index (no serialization).
74 // Invariant: an equal pointer must mean equal content. This holds only
75 // while all appended sets come from the same parse arena; if a builder
76 // is ever reused across arenas, a recycled address could alias a stale
77 // entry here. Verify in debug builds so that misuse fails loudly instead
78 // of silently emitting the wrong JSONB. Compiled out of release builds.
79 #[cfg(debug_assertions)]
80 {
81 let expected = serialize_property_set_to_jsonb(property_set)?;
82 debug_assert_eq!(
83 expected, self.jsonb_values[index as usize],
84 "pointer-dedup collision: arena address reused across blocks"
85 );
86 }
87 self.keys.push(Some(index));
88 }
89 None => {
90 // Cache miss: serialize once and store in dictionary
91 let jsonb_bytes = serialize_property_set_to_jsonb(property_set)?;
92 let new_index = self.jsonb_values.len() as i32;
93
94 self.jsonb_values.push(jsonb_bytes);
95 self.pointer_to_index.insert(ptr, new_index);
96 self.keys.push(Some(new_index));
97 }
98 }
99 Ok(())
100 }
101
102 /// Append a null value
103 pub fn append_null(&mut self) {
104 self.keys.push(None);
105 }
106
107 /// Finish building and return the DictionaryArray
108 ///
109 /// Output is identical to Arrow's BinaryDictionaryBuilder for compatibility
110 pub fn finish(self) -> Result<DictionaryArray<Int32Type>> {
111 let keys = Int32Array::from(self.keys);
112 // Convert Vec<Vec<u8>> to Vec<&[u8]> for BinaryArray::from_vec
113 let byte_slices: Vec<&[u8]> = self.jsonb_values.iter().map(|v| v.as_slice()).collect();
114 let values = Arc::new(BinaryArray::from_vec(byte_slices));
115 DictionaryArray::try_new(keys, values)
116 .map_err(|e| DataFusionError::ArrowError(Box::new(e), None).into())
117 }
118
119 /// Get the current number of appended entries
120 pub fn len(&self) -> usize {
121 self.keys.len()
122 }
123
124 /// Check if the builder is empty
125 pub fn is_empty(&self) -> bool {
126 self.keys.is_empty()
127 }
128}