lance_jsonb/
value.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3// SPDX-FileCopyrightText: Copyright 2023 Datafuse Labs
4// Adapted from databendlabs/jsonb at commit fba895c5ebe77ce2539e187f9c652f51cbf195c3.
5
6use crate::core::JsonbItemType;
7use std::borrow::Cow;
8use std::cmp::Ordering;
9use std::collections::BTreeMap;
10use std::fmt::Debug;
11use std::fmt::Display;
12use std::fmt::Formatter;
13
14use crate::extension::ExtensionValue;
15
16use crate::Date;
17use crate::Interval;
18use crate::Number;
19use crate::Timestamp;
20use crate::TimestampTz;
21use crate::core::Encoder;
22
23pub type Object<'a> = BTreeMap<String, Value<'a>>;
24
25/// Represents a JSON or extended JSON value.
26///
27/// This enum supports both standard JSON types (Null, Bool, String, Number, Array, Object)
28/// and extended types for specialized data representation (Binary, Date, Timestamp, etc.).
29/// The extended types provide additional functionality beyond the JSON specification,
30/// making this implementation more suitable for database applications and other
31/// systems requiring richer data type support.
32#[derive(Clone, Default)]
33pub enum Value<'a> {
34    /// Represents a JSON null value
35    #[default]
36    Null,
37    /// Represents a JSON boolean value (true or false)
38    Bool(bool),
39    /// Represents a JSON string value
40    String(Cow<'a, str>),
41    /// Represents a JSON number value with various internal representations
42    Number(Number),
43    /// Extended type: Represents binary data not supported in standard JSON
44    /// Useful for storing raw bytes, images, or other binary content
45    Binary(&'a [u8]),
46    /// Extended type: Represents a calendar date (year, month, day)
47    /// Stored as days since epoch for efficient comparison and manipulation
48    Date(Date),
49    /// Extended type: Represents a timestamp without timezone information
50    /// Stored as microseconds since epoch
51    Timestamp(Timestamp),
52    /// Extended type: Represents a timestamp with timezone information
53    /// Includes both timestamp and timezone offset
54    TimestampTz(TimestampTz),
55    /// Extended type: Represents a time interval or duration
56    /// Useful for time difference calculations and scheduling
57    Interval(Interval),
58    /// Represents a JSON array of values
59    Array(Vec<Self>),
60    /// Represents a JSON object as key-value pairs
61    Object(Object<'a>),
62}
63
64impl Eq for Value<'_> {}
65
66impl PartialEq for Value<'_> {
67    fn eq(&self, other: &Self) -> bool {
68        let result = self.cmp(other);
69        result == Ordering::Equal
70    }
71}
72
73impl PartialOrd for Value<'_> {
74    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
75        Some(self.cmp(other))
76    }
77}
78
79impl Ord for Value<'_> {
80    fn cmp(&self, other: &Self) -> Ordering {
81        let self_type = self.jsonb_item_type();
82        let other_type = other.jsonb_item_type();
83
84        if let Some(ord) = self_type.partial_cmp(&other_type) {
85            return ord;
86        }
87
88        match (self, other) {
89            (Value::Null, Value::Null) => Ordering::Equal,
90            (Value::Bool(v1), Value::Bool(v2)) => v1.cmp(v2),
91            (Value::Number(v1), Value::Number(v2)) => v1.cmp(v2),
92            (Value::String(v1), Value::String(v2)) => v1.cmp(v2),
93            (Value::Array(arr1), Value::Array(arr2)) => {
94                for (v1, v2) in arr1.iter().zip(arr2.iter()) {
95                    let ord = v1.cmp(v2);
96                    if ord != Ordering::Equal {
97                        return ord;
98                    }
99                }
100                arr1.len().cmp(&arr2.len())
101            }
102            (Value::Object(obj1), Value::Object(obj2)) => {
103                for ((k1, v1), (k2, v2)) in obj1.iter().zip(obj2.iter()) {
104                    let ord = k1.cmp(k2);
105                    if ord != Ordering::Equal {
106                        return ord;
107                    }
108                    let ord = v1.cmp(v2);
109                    if ord != Ordering::Equal {
110                        return ord;
111                    }
112                }
113                obj1.len().cmp(&obj2.len())
114            }
115            (_, _) => match (self.as_extension_value(), other.as_extension_value()) {
116                (Some(self_ext), Some(other_ext)) => {
117                    if let Some(ord) = self_ext.partial_cmp(&other_ext) {
118                        return ord;
119                    }
120                    Ordering::Equal
121                }
122                (_, _) => Ordering::Equal,
123            },
124        }
125    }
126}
127
128impl Debug for Value<'_> {
129    fn fmt(&self, formatter: &mut Formatter) -> std::fmt::Result {
130        match *self {
131            Value::Null => formatter.debug_tuple("Null").finish(),
132            Value::Bool(v) => formatter.debug_tuple("Bool").field(&v).finish(),
133            Value::Number(ref v) => Debug::fmt(v, formatter),
134            Value::String(ref v) => formatter.debug_tuple("String").field(v).finish(),
135            Value::Binary(ref v) => formatter.debug_tuple("Binary").field(v).finish(),
136            Value::Date(ref v) => formatter.debug_tuple("Date").field(v).finish(),
137            Value::Timestamp(ref v) => formatter.debug_tuple("Timestamp").field(v).finish(),
138            Value::TimestampTz(ref v) => formatter.debug_tuple("TimestampTz").field(v).finish(),
139            Value::Interval(ref v) => formatter.debug_tuple("Interval").field(v).finish(),
140            Value::Array(ref v) => {
141                formatter.write_str("Array(")?;
142                Debug::fmt(v, formatter)?;
143                formatter.write_str(")")
144            }
145            Value::Object(ref v) => {
146                formatter.write_str("Object(")?;
147                Debug::fmt(v, formatter)?;
148                formatter.write_str(")")
149            }
150        }
151    }
152}
153
154/// Writes `s` as a JSON string literal (quotes included), using spec-compliant
155/// escaping.
156///
157/// Rust's `Debug` escaping for `str` is not valid JSON: it writes control
158/// characters without a short escape as `\u{N}` (braced, variable-width) and
159/// NUL as `\0`, where JSON requires `\u0000`-style escapes. `serde_json`
160/// already implements correct JSON string escaping, so we reuse it here
161/// instead of duplicating that logic.
162fn write_json_string(f: &mut Formatter<'_>, s: &str) -> std::fmt::Result {
163    let escaped = serde_json::to_string(s).map_err(|_| std::fmt::Error)?;
164    write!(f, "{escaped}")
165}
166
167impl Display for Value<'_> {
168    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
169        match self {
170            Value::Null => write!(f, "null"),
171            Value::Bool(v) => {
172                if *v {
173                    write!(f, "true")
174                } else {
175                    write!(f, "false")
176                }
177            }
178            Value::Number(v) => write!(f, "{}", v),
179            Value::String(v) => write_json_string(f, v),
180            Value::Binary(v) => {
181                write!(f, "\"")?;
182                for c in *v {
183                    write!(f, "{c:02X}")?;
184                }
185                write!(f, "\"")?;
186                Ok(())
187            }
188            Value::Date(v) => {
189                write!(f, "\"{}\"", v)
190            }
191            Value::Timestamp(v) => {
192                write!(f, "\"{}\"", v)
193            }
194            Value::TimestampTz(v) => {
195                write!(f, "\"{}\"", v)
196            }
197            Value::Interval(v) => {
198                write!(f, "\"{}\"", v)
199            }
200            Value::Array(vs) => {
201                write!(f, "[")?;
202                for (i, v) in vs.iter().enumerate() {
203                    if i > 0 {
204                        write!(f, ",")?;
205                    }
206                    write!(f, "{v}")?;
207                }
208                write!(f, "]")
209            }
210            Value::Object(vs) => {
211                write!(f, "{{")?;
212                for (i, (k, v)) in vs.iter().enumerate() {
213                    if i > 0 {
214                        write!(f, ",")?;
215                    }
216                    write_json_string(f, k)?;
217                    write!(f, ":")?;
218                    write!(f, "{v}")?;
219                }
220                write!(f, "}}")
221            }
222        }
223    }
224}
225
226impl<'a> Value<'a> {
227    /// Serializes this value into JSONB bytes, appending to `buf`.
228    pub(crate) fn write_to_vec(&self, buf: &mut Vec<u8>) {
229        let mut encoder = Encoder::new(buf);
230        encoder.encode(self);
231    }
232
233    /// Serializes this value into JSONB bytes and returns the buffer.
234    pub fn to_vec(&self) -> Vec<u8> {
235        let mut buf = Vec::new();
236        self.write_to_vec(&mut buf);
237        buf
238    }
239
240    fn jsonb_item_type(&self) -> JsonbItemType {
241        match self {
242            Value::Null => JsonbItemType::Null,
243            Value::Bool(_) => JsonbItemType::Boolean,
244            Value::Number(_) => JsonbItemType::Number,
245            Value::String(_) => JsonbItemType::String,
246            Value::Binary(_) => JsonbItemType::Extension,
247            Value::Date(_) => JsonbItemType::Extension,
248            Value::Timestamp(_) => JsonbItemType::Extension,
249            Value::TimestampTz(_) => JsonbItemType::Extension,
250            Value::Interval(_) => JsonbItemType::Extension,
251            Value::Array(arr) => JsonbItemType::Array(arr.len()),
252            Value::Object(obj) => JsonbItemType::Object(obj.len()),
253        }
254    }
255
256    fn as_extension_value(&self) -> Option<ExtensionValue<'_>> {
257        match self {
258            Value::Binary(v) => Some(ExtensionValue::Binary(v)),
259            Value::Date(v) => Some(ExtensionValue::Date(v.clone())),
260            Value::Timestamp(v) => Some(ExtensionValue::Timestamp(v.clone())),
261            Value::TimestampTz(v) => Some(ExtensionValue::TimestampTz(v.clone())),
262            Value::Interval(v) => Some(ExtensionValue::Interval(v.clone())),
263            _ => None,
264        }
265    }
266}
267
268#[cfg(test)]
269mod tests {
270    use super::*;
271
272    #[test]
273    fn test_display_escapes_object_keys() {
274        let mut obj = Object::new();
275        obj.insert("line\nbreak".to_string(), Value::Bool(false));
276        let value = Value::Object(obj);
277
278        let displayed = format!("{value}");
279        assert_eq!(displayed, "{\"line\\nbreak\":false}");
280
281        // The displayed form must itself be valid JSON that round-trips.
282        let reparsed: serde_json::Value = serde_json::from_str(&displayed).unwrap();
283        assert_eq!(reparsed, serde_json::json!({"line\nbreak": false}));
284    }
285
286    /// Rust's `Debug` escaping (used by both `Value::String` and the
287    /// `Value::Object` key display) writes control characters without a
288    /// short escape (like `\x01`) as a braced, variable-width escape, and
289    /// writes NUL as a two-character escape. Neither form is valid JSON,
290    /// which requires fixed 4-digit `\u` escapes for both.
291    #[test]
292    fn test_display_string_control_char_is_valid_json() {
293        let value = Value::String(Cow::Borrowed("\u{1}"));
294        let displayed = format!("{value}");
295        assert!(
296            serde_json::from_str::<serde_json::Value>(&displayed).is_ok(),
297            "{displayed:?} is not valid JSON"
298        );
299    }
300
301    #[test]
302    fn test_display_string_nul_is_valid_json() {
303        let value = Value::String(Cow::Borrowed("\0"));
304        let displayed = format!("{value}");
305        assert!(
306            serde_json::from_str::<serde_json::Value>(&displayed).is_ok(),
307            "{displayed:?} is not valid JSON"
308        );
309    }
310
311    #[test]
312    fn test_display_object_key_control_char_is_valid_json() {
313        let mut obj = Object::new();
314        obj.insert("\u{1}".to_string(), Value::Null);
315        let value = Value::Object(obj);
316
317        let displayed = format!("{value}");
318        let reparsed: Result<serde_json::Value, _> = serde_json::from_str(&displayed);
319        assert!(reparsed.is_ok(), "{displayed:?} is not valid JSON");
320        assert_eq!(reparsed.unwrap(), serde_json::json!({"\u{1}": null}),);
321    }
322}
lance_jsonb/value.rs

lance_jsonb/
value.rs
