uqa_execution/distinct/
encoding.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Canonical SQL equality encoding and borrowed-row hashing.
8
9use std::hash::{BuildHasher, Hasher};
10
11use smallvec::{Array, SmallVec};
12use uqa_core::{DecimalValue, TemporalValue, Value};
13
14use crate::{ExecError, ExecResult};
15
16pub(super) const MICROS_PER_DAY: i128 = 86_400_000_000;
17
18pub(crate) type EncodedKey = SmallVec<[u8; 64]>;
19
20/// Hash a borrowed positional SQL row in its canonical equality domain.
21///
22/// This streams encoded components straight into the caller's hasher, so it
23/// does not allocate or construct an intermediate byte key. Hash collisions
24/// remain possible; callers must verify complete [`Value`] equality before
25/// reusing an existing row or group.
26pub fn hash_canonical_row<'a, S: BuildHasher>(
27    build_hasher: &S,
28    values: impl ExactSizeIterator<Item = Option<&'a Value>>,
29) -> ExecResult<u64> {
30    let count = values.len();
31    let mut hasher = build_hasher.build_hasher();
32    {
33        let mut output = HasherOutput(&mut hasher);
34        encode_len(count, &mut output)?;
35        for value in values {
36            if let Some(value) = value {
37                encode_value(value, &mut output)?;
38            } else {
39                output.push_byte(0);
40            }
41        }
42    }
43    Ok(hasher.finish())
44}
45
46/// Pack exactly two text-or-NULL values of at most three bytes each into an injective integer key. `None` selects the general collision-safe encoder for every other row.
47pub fn try_pack_compact_text_pair<'a>(
48    values: impl ExactSizeIterator<Item = Option<&'a Value>>,
49) -> Option<u64> {
50    if values.len() != 2 {
51        return None;
52    }
53    let mut values = values;
54    let first = compact_text_component(values.next()?)?;
55    let second = compact_text_component(values.next()?)?;
56    Some(u64::from(first) << 32 | u64::from(second))
57}
58
59fn compact_text_component(value: Option<&Value>) -> Option<u32> {
60    match value {
61        None | Some(Value::Null) => Some(0),
62        Some(Value::Str(value)) if value.len() <= 3 => {
63            let mut packed = [0u8; 4];
64            packed[0] = u8::try_from(value.len()).ok()? + 1;
65            packed[1..][..value.len()].copy_from_slice(value.as_bytes());
66            Some(u32::from_be_bytes(packed))
67        }
68        Some(_) => None,
69    }
70}
71
72/// Encode positional SQL values in the exact equality domain used by DISTINCT and spill-backed row-key state. Callers that need an external exact index can persist this representation without relying on `Value`'s serialization format.
73pub fn canonical_row_key(values: &[Value]) -> ExecResult<Vec<u8>> {
74    encode_key(values)
75}
76
77/// Collision-free binary key encoding. Numeric values deliberately share one
78/// canonical domain so `1`, `1.0`, `DECIMAL '1'`, and `TRUE` retain the same
79/// equality behavior as UQA's SQL comparisons. Every structural value carries
80/// lengths/counts, preventing concatenation and nested-container collisions.
81pub(crate) fn encode_key(values: &[Value]) -> ExecResult<Vec<u8>> {
82    encode_key_borrowed(values.iter().map(Some))
83}
84
85pub(super) fn encode_key_borrowed<'a>(
86    values: impl ExactSizeIterator<Item = Option<&'a Value>>,
87) -> ExecResult<Vec<u8>> {
88    let estimated_capacity = encoded_key_capacity(values.len())?;
89    let mut output = Vec::with_capacity(estimated_capacity);
90    encode_len(values.len(), &mut output)?;
91    for value in values {
92        match value {
93            Some(value) => encode_value(value, &mut output)?,
94            None => encode_value(&Value::Null, &mut output)?,
95        }
96    }
97    Ok(output)
98}
99
100/// Encode a join probe key directly from physical slots. Single- and
101/// two-column numeric keys stay inline, and a NULL/missing component rejects
102/// the SQL equality key without allocating or cloning a `Value`.
103pub(crate) fn encode_non_null_key<'a>(
104    values: impl ExactSizeIterator<Item = Option<&'a Value>>,
105) -> ExecResult<Option<EncodedKey>> {
106    let count = values.len();
107    let mut output = EncodedKey::with_capacity(encoded_key_capacity(count)?);
108    encode_len(count, &mut output)?;
109    for value in values {
110        let Some(value) = value else {
111            return Ok(None);
112        };
113        if matches!(value, Value::Null) {
114            return Ok(None);
115        }
116        encode_value(value, &mut output)?;
117    }
118    Ok(Some(output))
119}
120
121fn encoded_key_capacity(values: usize) -> ExecResult<usize> {
122    values
123        .checked_mul(22)
124        .and_then(|bytes| bytes.checked_add(8))
125        .ok_or_else(|| encoding_error("DISTINCT key capacity overflow"))
126}
127
128trait KeyOutput {
129    fn push_byte(&mut self, value: u8);
130    fn extend_bytes(&mut self, values: &[u8]);
131}
132
133impl KeyOutput for Vec<u8> {
134    fn push_byte(&mut self, value: u8) {
135        self.push(value);
136    }
137
138    fn extend_bytes(&mut self, values: &[u8]) {
139        self.extend_from_slice(values);
140    }
141}
142
143impl<A: Array<Item = u8>> KeyOutput for SmallVec<A> {
144    fn push_byte(&mut self, value: u8) {
145        self.push(value);
146    }
147
148    fn extend_bytes(&mut self, values: &[u8]) {
149        self.extend_from_slice(values);
150    }
151}
152
153struct HasherOutput<'a, H: Hasher>(&'a mut H);
154
155impl<H: Hasher> KeyOutput for HasherOutput<'_, H> {
156    fn push_byte(&mut self, value: u8) {
157        self.0.write_u8(value);
158    }
159
160    fn extend_bytes(&mut self, values: &[u8]) {
161        self.0.write(values);
162    }
163}
164
165fn encode_value(value: &Value, output: &mut impl KeyOutput) -> ExecResult<()> {
166    match value {
167        Value::Null => output.push_byte(0),
168        Value::Void => output.push_byte(13),
169        Value::Bool(value) => {
170            encode_decimal_numeric(&DecimalValue::from_bool(*value), output)?;
171        }
172        Value::Int(value) => {
173            encode_decimal_numeric(&DecimalValue::from_i64(*value), output)?;
174        }
175        Value::Float(value) => encode_float_numeric(*value, output)?,
176        Value::Decimal(value) => encode_decimal_numeric(value, output)?,
177        Value::Str(value) => {
178            output.push_byte(2);
179            encode_bytes(value.as_bytes(), output)?;
180        }
181        Value::FixedChar(value) => {
182            output.push_byte(7);
183            encode_bytes(value.trim_end_matches(' ').as_bytes(), output)?;
184        }
185        Value::Bytes(value) => {
186            output.push_byte(3);
187            encode_bytes(value, output)?;
188        }
189        Value::Temporal(value) => encode_temporal(value, output),
190        Value::Json(value) => {
191            output.push_byte(8);
192            encode_bytes(value.as_bytes(), output)?;
193        }
194        Value::JsonB(value) => {
195            output.push_byte(9);
196            let canonical = uqa_core::jsonb_equality_key(value)
197                .ok_or_else(|| ExecError::Other("stored JSONB value is not valid JSON".into()))?;
198            encode_bytes(&canonical, output)?;
199        }
200        Value::Array(array) => {
201            output.push_byte(12);
202            encode_len(array.lower_bounds().len(), output)?;
203            for lower_bound in array.lower_bounds() {
204                output.extend_bytes(&lower_bound.to_le_bytes());
205            }
206            encode_len(array.elements().len(), output)?;
207            for value in array.elements() {
208                encode_value(value, output)?;
209            }
210        }
211        Value::List(values) => {
212            output.push_byte(5);
213            encode_len(values.len(), output)?;
214            for value in values {
215                encode_value(value, output)?;
216            }
217        }
218        Value::Row(values) => {
219            output.push_byte(10);
220            encode_len(values.len(), output)?;
221            for value in values {
222                encode_value(value, output)?;
223            }
224        }
225        Value::Record(fields) => {
226            output.push_byte(11);
227            encode_len(fields.len(), output)?;
228            for (_, value) in fields {
229                encode_value(value, output)?;
230            }
231        }
232        Value::Map(values) => {
233            output.push_byte(6);
234            encode_len(values.len(), output)?;
235            for (name, value) in values {
236                encode_bytes(name.as_bytes(), output)?;
237                encode_value(value, output)?;
238            }
239        }
240    }
241    Ok(())
242}
243
244fn encode_decimal_numeric(value: &DecimalValue, output: &mut impl KeyOutput) -> ExecResult<()> {
245    if value.is_nan() {
246        output.extend_bytes(&[1, 1]);
247    } else if value.is_negative_infinity() {
248        output.extend_bytes(&[1, 2]);
249    } else if value.is_positive_infinity() {
250        output.extend_bytes(&[1, 3]);
251    } else {
252        output.extend_bytes(&[1, 0]);
253        encode_bytes(value.to_canonical_string().as_bytes(), output)?;
254    }
255    Ok(())
256}
257
258fn encode_float_numeric(value: f64, output: &mut impl KeyOutput) -> ExecResult<()> {
259    if value.is_nan() {
260        // PostgreSQL groups all NaN values together for DISTINCT.
261        output.extend_bytes(&[1, 1]);
262    } else if value == f64::NEG_INFINITY {
263        output.extend_bytes(&[1, 2]);
264    } else if value == f64::INFINITY {
265        output.extend_bytes(&[1, 3]);
266    } else if let Some(decimal) = DecimalValue::from_f64_lossy(value) {
267        encode_decimal_numeric(&decimal, output)?;
268    } else {
269        // Preserve a finite value that cannot enter PostgreSQL's NUMERIC
270        // domain. Normalize signed zero before storing bits.
271        output.extend_bytes(&[1, 4]);
272        let normalized = if value == 0.0 { 0.0 } else { value };
273        output.extend_bytes(&normalized.to_bits().to_be_bytes());
274    }
275    Ok(())
276}
277
278fn encode_temporal(value: &TemporalValue, output: &mut impl KeyOutput) {
279    output.push_byte(4);
280    match value {
281        TemporalValue::Date { days } => {
282            output.push_byte(0);
283            output.extend_bytes(&days.to_be_bytes());
284        }
285        TemporalValue::Time { micros } => {
286            output.push_byte(1);
287            let normalized = i128::from(*micros).rem_euclid(MICROS_PER_DAY);
288            output.extend_bytes(&normalized.to_be_bytes());
289        }
290        TemporalValue::TimeTz {
291            micros,
292            offset_minutes,
293        } => {
294            output.push_byte(2);
295            let normalized = (i128::from(*micros) - i128::from(*offset_minutes) * 60_000_000)
296                .rem_euclid(MICROS_PER_DAY);
297            output.extend_bytes(&normalized.to_be_bytes());
298        }
299        TemporalValue::Timestamp { micros } => {
300            output.push_byte(3);
301            output.extend_bytes(&micros.to_be_bytes());
302        }
303        TemporalValue::TimestampTz { micros } => {
304            output.push_byte(4);
305            output.extend_bytes(&micros.to_be_bytes());
306        }
307        TemporalValue::Interval {
308            months,
309            days,
310            micros,
311        } => {
312            output.push_byte(5);
313            let normalized = (i128::from(*months) * 30 + i128::from(*days)) * MICROS_PER_DAY
314                + i128::from(*micros);
315            output.extend_bytes(&normalized.to_be_bytes());
316        }
317    }
318}
319
320fn encode_bytes(bytes: &[u8], output: &mut impl KeyOutput) -> ExecResult<()> {
321    encode_len(bytes.len(), output)?;
322    output.extend_bytes(bytes);
323    Ok(())
324}
325
326fn encode_len(length: usize, output: &mut impl KeyOutput) -> ExecResult<()> {
327    let length = u64::try_from(length)
328        .map_err(|_| encoding_error("DISTINCT key component exceeds the binary format"))?;
329    output.extend_bytes(&length.to_be_bytes());
330    Ok(())
331}
332
333fn encoding_error(message: impl Into<String>) -> ExecError {
334    ExecError::Other(message.into())
335}
uqa_execution/distinct/encoding.rs

uqa_execution/distinct/
encoding.rs
