lance_jsonb/functions/path.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
6// This file contains functions that dealing with path-based access to JSONB data.
7
8use std::borrow::Cow;
9
10use crate::OwnedJsonb;
11use crate::RawJsonb;
12use crate::core::ArrayIterator;
13use crate::error::*;
14
15impl RawJsonb<'_> {
16 /// Gets the element at the specified index in a JSONB array.
17 ///
18 /// If the JSONB value is an array, this function returns the element at the given `index` as an `OwnedJsonb`.
19 /// If the `index` is out of bounds, it returns `Ok(None)`.
20 /// If the JSONB value is not an array (e.g., it's an object or a scalar), this function also returns `Ok(None)`.
21 ///
22 /// # Arguments
23 ///
24 /// * `self` - The JSONB value.
25 /// * `index` - The index of the desired element.
26 ///
27 /// # Returns
28 ///
29 /// * `Ok(Some(OwnedJsonb))` - The element at the specified index as an `OwnedJsonb` if the input is an array and the index is valid.
30 /// * `Ok(None)` - If the input is not an array, or if the index is out of bounds.
31 /// * `Err(Error)` - If an error occurred during decoding (e.g., invalid JSONB data).
32 ///
33 /// # Examples
34 ///
35 /// ```rust
36 /// use lance_jsonb::OwnedJsonb;
37 ///
38 /// let arr_jsonb = r#"[1, "hello", {"a": 1}]"#.parse::<OwnedJsonb>().unwrap();
39 /// let raw_jsonb = arr_jsonb.as_raw();
40 ///
41 /// let element0 = raw_jsonb.get_by_index(0).unwrap();
42 /// assert_eq!(element0.unwrap().to_string(), "1");
43 ///
44 /// let element1 = raw_jsonb.get_by_index(1).unwrap();
45 /// assert_eq!(element1.unwrap().to_string(), r#""hello""#);
46 ///
47 /// let element2 = raw_jsonb.get_by_index(2).unwrap();
48 /// assert_eq!(element2.unwrap().to_string(), r#"{"a":1}"#);
49 ///
50 /// let element3 = raw_jsonb.get_by_index(3).unwrap();
51 /// assert!(element3.is_none()); // Index out of bounds
52 ///
53 /// let obj_jsonb = r#"{"a": 1}"#.parse::<OwnedJsonb>().unwrap();
54 /// let raw_jsonb = obj_jsonb.as_raw();
55 /// let element = raw_jsonb.get_by_index(0).unwrap();
56 /// assert!(element.is_none()); // Not an array
57 /// ```
58 pub fn get_by_index(&self, index: usize) -> Result<Option<OwnedJsonb>> {
59 let array_iter_opt = ArrayIterator::new(*self)?;
60 if let Some(mut array_iter) = array_iter_opt
61 && let Some(item_result) = array_iter.nth(index)
62 {
63 let item = item_result?;
64 let value = OwnedJsonb::from_item(item)?;
65 return Ok(Some(value));
66 }
67 Ok(None)
68 }
69
70 /// Gets the value associated with a given key in a JSONB object.
71 ///
72 /// If the JSONB value is an object, this function searches for a key matching the provided `name`
73 /// and returns the associated value as an `OwnedJsonb`.
74 /// The `ignore_case` parameter controls whether the key search is case-sensitive.
75 /// If the key is not found, it returns `Ok(None)`.
76 /// If the JSONB value is not an object (e.g., it's an array or a scalar), this function also returns `Ok(None)`.
77 ///
78 /// # Arguments
79 ///
80 /// * `self` - The JSONB value.
81 /// * `name` - The key to search for.
82 /// * `ignore_case` - Whether the key search should be case-insensitive.
83 ///
84 /// # Returns
85 ///
86 /// * `Ok(Some(OwnedJsonb))` - The value associated with the key as an `OwnedJsonb`, if the input is an object and the key is found.
87 /// * `Ok(None)` - If the input is not an object, or if the key is not found.
88 /// * `Err(Error)` - If an error occurred during decoding (e.g., invalid JSONB data).
89 ///
90 /// # Examples
91 ///
92 /// ```rust
93 /// use lance_jsonb::OwnedJsonb;
94 ///
95 /// let obj_jsonb = r#"{"a": 1, "b": "hello", "c": [1, 2]}"#.parse::<OwnedJsonb>().unwrap();
96 /// let raw_jsonb = obj_jsonb.as_raw();
97 ///
98 /// let value_a = raw_jsonb.get_by_name("a", false).unwrap();
99 /// assert_eq!(value_a.unwrap().to_string(), "1");
100 ///
101 /// let value_b = raw_jsonb.get_by_name("b", false).unwrap();
102 /// assert_eq!(value_b.unwrap().to_string(), r#""hello""#);
103 ///
104 /// let value_c = raw_jsonb.get_by_name("c", false).unwrap();
105 /// assert_eq!(value_c.unwrap().to_string(), "[1,2]");
106 ///
107 /// let value_d = raw_jsonb.get_by_name("d", false).unwrap();
108 /// assert!(value_d.is_none()); // Key not found
109 ///
110 /// // Case-insensitive search
111 /// let value_a_case_insensitive = raw_jsonb.get_by_name("A", true).unwrap();
112 /// assert_eq!(value_a_case_insensitive.unwrap().to_string(), "1");
113 ///
114 /// let arr_jsonb = "[1, 2, 3]".parse::<OwnedJsonb>().unwrap();
115 /// let raw_jsonb = arr_jsonb.as_raw();
116 /// let value = raw_jsonb.get_by_name("a", false).unwrap();
117 /// assert!(value.is_none()); // Not an object
118 /// ```
119 pub fn get_by_name(&self, name: &str, ignore_case: bool) -> Result<Option<OwnedJsonb>> {
120 let key_name = Cow::Borrowed(name);
121 if let Some(val_item) =
122 self.get_object_value_by_key_name(&key_name, |name, key| key.eq(name))?
123 {
124 let value = OwnedJsonb::from_item(val_item)?;
125 return Ok(Some(value));
126 }
127 if ignore_case
128 && let Some(val_item) = self.get_object_value_by_key_name(&key_name, |name, key| {
129 key.eq_ignore_ascii_case(name)
130 })?
131 {
132 let value = OwnedJsonb::from_item(val_item)?;
133 return Ok(Some(value));
134 }
135 Ok(None)
136 }
137}