diff --git a/src/constants.rs b/src/constants.rs index 2c7fd7c..d13ba9f 100644 --- a/src/constants.rs +++ b/src/constants.rs @@ -48,3 +48,22 @@ pub(crate) const TYPE_DATE: &str = "DATE"; pub(crate) const TYPE_TIMESTAMP: &str = "TIMESTAMP"; pub(crate) const TYPE_TIMESTAMP_TZ: &str = "TIMESTAMP_TZ"; pub(crate) const TYPE_INTERVAL: &str = "INTERVAL"; + +pub(crate) const MAX_DECIMAL64_PRECISION: usize = 18; +pub(crate) const MAX_DECIMAL128_PRECISION: usize = 38; +pub(crate) const MAX_DECIMAL256_PRECISION: usize = 76; + +pub(crate) const UINT64_MIN: i128 = 0i128; +pub(crate) const UINT64_MAX: i128 = 18_446_744_073_709_551_615i128; +pub(crate) const INT64_MIN: i128 = -9_223_372_036_854_775_808i128; +pub(crate) const INT64_MAX: i128 = 9_223_372_036_854_775_807i128; +pub(crate) const DECIMAL64_MIN: i128 = -999_999_999_999_999_999i128; +pub(crate) const DECIMAL64_MAX: i128 = 999_999_999_999_999_999i128; +pub(crate) const DECIMAL128_MIN: i128 = -99_999_999_999_999_999_999_999_999_999_999_999_999i128; +pub(crate) const DECIMAL128_MAX: i128 = 99_999_999_999_999_999_999_999_999_999_999_999_999i128; + +pub(crate) const NUMBER_STRUCT_TOKEN: &str = "$jsonb::private::Number"; +pub(crate) const NUMBER_STRUCT_FIELD_SCALE: &str = "$jsonb::private::Number::Scale"; +pub(crate) const NUMBER_STRUCT_FIELD_VALUE: &str = "$jsonb::private::Number::Value"; +pub(crate) const NUMBER_STRUCT_FIELD_HIGH_VALUE: &str = "$jsonb::private::Number::High_Value"; +pub(crate) const NUMBER_STRUCT_FIELD_LOW_VALUE: &str = "$jsonb::private::Number::Low_Value"; diff --git a/src/core/databend/builder.rs b/src/core/databend/builder.rs index 0c8a482..b7370a8 100644 --- a/src/core/databend/builder.rs +++ b/src/core/databend/builder.rs @@ -20,7 +20,9 @@ use byteorder::WriteBytesExt; use super::constants::*; use super::jentry::JEntry; +use crate::core::ExtensionItem; use crate::core::JsonbItem; +use crate::core::NumberItem; use crate::error::Error; use crate::error::Result; use crate::OwnedJsonb; @@ -189,21 +191,35 @@ fn append_jsonb_item(buf: &mut Vec, jentry_index: &mut usize, item: JsonbIte }; replace_jentry(buf, jentry, jentry_index); } - JsonbItem::Number(data) => { - let jentry = JEntry::make_number_jentry(data.len()); - replace_jentry(buf, jentry, jentry_index); - buf.extend_from_slice(data); - } + JsonbItem::Number(num) => match num { + NumberItem::Raw(data) => { + let jentry = JEntry::make_number_jentry(data.len()); + replace_jentry(buf, jentry, jentry_index); + buf.extend_from_slice(data); + } + NumberItem::Number(num) => { + let len = num.compact_encode(&mut *buf)?; + let jentry = JEntry::make_number_jentry(len); + replace_jentry(buf, jentry, jentry_index); + } + }, JsonbItem::String(data) => { let jentry = JEntry::make_string_jentry(data.len()); replace_jentry(buf, jentry, jentry_index); - buf.extend_from_slice(data); - } - JsonbItem::Extension(data) => { - let jentry = JEntry::make_extension_jentry(data.len()); - replace_jentry(buf, jentry, jentry_index); - buf.extend_from_slice(data); + buf.extend_from_slice(data.as_bytes()); } + JsonbItem::Extension(ext) => match ext { + ExtensionItem::Raw(data) => { + let jentry = JEntry::make_extension_jentry(data.len()); + replace_jentry(buf, jentry, jentry_index); + buf.extend_from_slice(data); + } + ExtensionItem::Extension(ext) => { + let len = ext.compact_encode(&mut *buf)?; + let jentry = JEntry::make_extension_jentry(len); + replace_jentry(buf, jentry, jentry_index); + } + }, JsonbItem::Raw(raw_jsonb) => { append_raw_jsonb_data(buf, jentry_index, raw_jsonb)?; } diff --git a/src/core/databend/de.rs b/src/core/databend/de.rs index 1ebabb8..0d6638e 100644 --- a/src/core/databend/de.rs +++ b/src/core/databend/de.rs @@ -67,7 +67,7 @@ impl<'de> Deserializer<'de> { self.index == self.raw.len() } - fn read_header(&mut self) -> Result<(u32, u32)> { + fn read_header(&mut self) -> Result<(u32, usize)> { let header = self.raw.read_header(self.index)?; self.index += 4; Ok(header) @@ -377,12 +377,11 @@ impl<'de> de::Deserializer<'de> for &mut Deserializer<'de> { } } ARRAY_CONTAINER_TAG => { - let jentries = self.read_array_jentries(header_len as usize)?; + let jentries = self.read_array_jentries(header_len)?; visitor.visit_seq(ArrayDeserializer::new(self, jentries)) } OBJECT_CONTAINER_TAG => { - let (key_jentries, value_jentries) = - self.read_object_jentries(header_len as usize)?; + let (key_jentries, value_jentries) = self.read_object_jentries(header_len)?; let origin_index = self.index; let key_length: usize = key_jentries.iter().map(|j| j.length as usize).sum(); @@ -564,7 +563,7 @@ impl<'de> de::Deserializer<'de> for &mut Deserializer<'de> { let (header_type, header_len) = self.read_header()?; match header_type { ARRAY_CONTAINER_TAG => { - let jentries = self.read_array_jentries(header_len as usize)?; + let jentries = self.read_array_jentries(header_len)?; visitor.visit_seq(ArrayDeserializer::new(self, jentries)) } SCALAR_CONTAINER_TAG | OBJECT_CONTAINER_TAG => Err(Error::UnexpectedType), @@ -598,8 +597,7 @@ impl<'de> de::Deserializer<'de> for &mut Deserializer<'de> { let (header_type, header_len) = self.read_header()?; match header_type { OBJECT_CONTAINER_TAG => { - let (key_jentries, value_jentries) = - self.read_object_jentries(header_len as usize)?; + let (key_jentries, value_jentries) = self.read_object_jentries(header_len)?; let origin_index = self.index; let key_length: usize = key_jentries.iter().map(|j| j.length as usize).sum(); @@ -861,7 +859,7 @@ impl<'de> de::VariantAccess<'de> for EnumDeserializer<'_, 'de> { let (header_type, header_len) = self.de.read_header()?; match header_type { ARRAY_CONTAINER_TAG => { - let jentries = self.de.read_array_jentries(header_len as usize)?; + let jentries = self.de.read_array_jentries(header_len)?; visitor.visit_seq(ArrayDeserializer::new(self.de, jentries)) } SCALAR_CONTAINER_TAG | OBJECT_CONTAINER_TAG => Err(Error::UnexpectedType), @@ -883,7 +881,7 @@ impl<'de> de::VariantAccess<'de> for EnumDeserializer<'_, 'de> { match header_type { OBJECT_CONTAINER_TAG => { let (key_jentries, value_jentries) = - self.de.read_object_jentries(header_len as usize)?; + self.de.read_object_jentries(header_len)?; let origin_index = self.de.index; let key_length: usize = key_jentries.iter().map(|j| j.length as usize).sum(); diff --git a/src/core/databend/iterator.rs b/src/core/databend/iterator.rs index 3fff21a..fae6129 100644 --- a/src/core/databend/iterator.rs +++ b/src/core/databend/iterator.rs @@ -35,12 +35,12 @@ impl<'a> ArrayIterator<'a> { let (header_type, header_len) = raw_jsonb.read_header(0)?; if header_type == ARRAY_CONTAINER_TAG { let jentry_offset = 4; - let item_offset = 4 + 4 * header_len as usize; + let item_offset = 4 + 4 * header_len; Ok(Some(Self { raw_jsonb, jentry_offset, item_offset, - length: header_len as usize, + length: header_len, index: 0, })) } else { @@ -97,12 +97,12 @@ impl<'a> ObjectKeyIterator<'a> { let (header_type, header_len) = raw_jsonb.read_header(0)?; if header_type == OBJECT_CONTAINER_TAG { let jentry_offset = 4; - let item_offset = 4 + 8 * header_len as usize; + let item_offset = 4 + 8 * header_len; Ok(Some(Self { raw_jsonb, jentry_offset, item_offset, - length: header_len as usize, + length: header_len, index: 0, })) } else { @@ -159,7 +159,7 @@ impl<'a> ObjectValueIterator<'a> { let (header_type, header_len) = raw_jsonb.read_header(0)?; if header_type == OBJECT_CONTAINER_TAG { let mut jentry_offset = 4; - let mut item_offset = 4 + 8 * header_len as usize; + let mut item_offset = 4 + 8 * header_len; for _ in 0..header_len { let key_jentry = raw_jsonb.read_jentry(jentry_offset)?; jentry_offset += 4; @@ -170,7 +170,7 @@ impl<'a> ObjectValueIterator<'a> { raw_jsonb, jentry_offset, item_offset, - length: header_len as usize, + length: header_len, index: 0, })) } else { @@ -229,14 +229,14 @@ impl<'a> ObjectIterator<'a> { let (header_type, header_len) = raw_jsonb.read_header(0)?; if header_type == OBJECT_CONTAINER_TAG { let mut jentry_offset = 4; - let mut key_jentries = VecDeque::with_capacity(header_len as usize); + let mut key_jentries = VecDeque::with_capacity(header_len); for _ in 0..header_len { let key_jentry = raw_jsonb.read_jentry(jentry_offset)?; jentry_offset += 4; key_jentries.push_back(key_jentry); } let key_length: usize = key_jentries.iter().map(|j| j.length as usize).sum(); - let key_offset = 4 + 8 * header_len as usize; + let key_offset = 4 + 8 * header_len; let val_offset = key_offset + key_length; Ok(Some(Self { @@ -245,7 +245,7 @@ impl<'a> ObjectIterator<'a> { jentry_offset, key_offset, val_offset, - length: header_len as usize, + length: header_len, })) } else { Ok(None) diff --git a/src/core/databend/ser.rs b/src/core/databend/ser.rs index a2b3746..3321043 100644 --- a/src/core/databend/ser.rs +++ b/src/core/databend/ser.rs @@ -36,6 +36,26 @@ use crate::Error; use crate::OwnedJsonb; use crate::RawJsonb; +use crate::Decimal128; +use crate::Decimal256; +use crate::Decimal64; +use ethnum::i256; + +use crate::constants::NUMBER_STRUCT_FIELD_HIGH_VALUE; +use crate::constants::NUMBER_STRUCT_FIELD_LOW_VALUE; +use crate::constants::NUMBER_STRUCT_FIELD_SCALE; +use crate::constants::NUMBER_STRUCT_FIELD_VALUE; +use crate::constants::NUMBER_STRUCT_TOKEN; + +use crate::constants::DECIMAL128_MAX; +use crate::constants::DECIMAL128_MIN; +use crate::constants::DECIMAL64_MAX; +use crate::constants::DECIMAL64_MIN; + +use crate::constants::MAX_DECIMAL128_PRECISION; +use crate::constants::MAX_DECIMAL256_PRECISION; +use crate::constants::MAX_DECIMAL64_PRECISION; + /// `Serializer` is a custom serializer for JSONB data, implementing the /// `serde::ser::Serializer` trait. It allows serializing Rust data structures /// into a `Vec` representing the JSONB data. @@ -121,9 +141,9 @@ impl<'a> ser::Serializer for &'a mut Serializer { type SerializeMap = ObjectSerializer<'a>; - type SerializeStruct = ObjectSerializer<'a>; + type SerializeStruct = StructSerializer<'a>; - type SerializeStructVariant = ArraySerializer<'a>; + type SerializeStructVariant = StructSerializer<'a>; fn serialize_bool(self, v: bool) -> Result { let jentry = if v { @@ -174,6 +194,43 @@ impl<'a> ser::Serializer for &'a mut Serializer { self.write_number(Number::Float64(v)) } + fn serialize_i128(self, v: i128) -> Result { + let num = if (DECIMAL64_MIN..=DECIMAL64_MAX).contains(&v) { + Number::Decimal64(Decimal64 { + scale: 0, + value: v as i64, + }) + } else if (DECIMAL128_MIN..=DECIMAL128_MAX).contains(&v) { + Number::Decimal128(Decimal128 { scale: 0, value: v }) + } else { + Number::Decimal256(Decimal256 { + scale: 0, + value: i256::from(v), + }) + }; + self.write_number(num) + } + + fn serialize_u128(self, v: u128) -> Result { + let num = if v <= DECIMAL64_MAX as u128 { + Number::Decimal64(Decimal64 { + scale: 0, + value: v as i64, + }) + } else if v <= DECIMAL128_MAX as u128 { + Number::Decimal128(Decimal128 { + scale: 0, + value: v as i128, + }) + } else { + Number::Decimal256(Decimal256 { + scale: 0, + value: i256::from(v), + }) + }; + self.write_number(num) + } + fn serialize_char(self, v: char) -> Result { let s: String = v.to_string(); self.write_str(s.as_str()) @@ -216,23 +273,32 @@ impl<'a> ser::Serializer for &'a mut Serializer { fn serialize_newtype_struct( self, _name: &'static str, - _value: &T, + value: &T, ) -> Result { - self.serialize_unit() + T::serialize(value, self) } fn serialize_newtype_variant( self, _name: &'static str, _variant_index: u32, - _variant: &'static str, - _value: &T, + variant: &'static str, + value: &T, ) -> Result { - todo!() + let mut serializer = Serializer::new(); + value.serialize(&mut serializer)?; + let value_jsonb = OwnedJsonb::new(serializer.buffer); + + let mut builder = ObjectBuilder::new(); + builder.push_owned_jsonb(variant, value_jsonb)?; + let object_jsonb = builder.build()?; + let mut buf = object_jsonb.to_vec(); + self.buffer.append(&mut buf); + Ok(()) } fn serialize_seq(self, len: Option) -> Result { - Ok(ArraySerializer::new(&mut self.buffer, len)) + Ok(ArraySerializer::new(None, &mut self.buffer, len)) } fn serialize_tuple(self, len: usize) -> Result { @@ -242,51 +308,69 @@ impl<'a> ser::Serializer for &'a mut Serializer { fn serialize_tuple_struct( self, _name: &'static str, - _len: usize, + len: usize, ) -> Result { - todo!() + self.serialize_seq(Some(len)) } fn serialize_tuple_variant( self, _name: &'static str, _variant_index: u32, - _variant: &'static str, - _len: usize, + variant: &'static str, + len: usize, ) -> Result { - todo!() + Ok(ArraySerializer::new( + Some(variant), + &mut self.buffer, + Some(len), + )) } fn serialize_map(self, len: Option) -> Result { Ok(ObjectSerializer::new(&mut self.buffer, len)) } - fn serialize_struct(self, _name: &'static str, len: usize) -> Result { - self.serialize_map(Some(len)) + fn serialize_struct(self, name: &'static str, len: usize) -> Result { + Ok(StructSerializer::new(name, None, &mut self.buffer, len)) } fn serialize_struct_variant( self, - _name: &'static str, + name: &'static str, _variant_index: u32, - _variant: &'static str, - _len: usize, + variant: &'static str, + len: usize, ) -> Result { - todo!() + Ok(StructSerializer::new( + name, + Some(variant), + &mut self.buffer, + len, + )) + } + + fn is_human_readable(&self) -> bool { + false } } pub struct ArraySerializer<'a> { + variant: Option<&'static str>, buffer: &'a mut Vec, items: Vec, } impl<'a> ArraySerializer<'a> { - pub fn new(buffer: &'a mut Vec, len: Option) -> Self { + pub fn new(variant: Option<&'static str>, buffer: &'a mut Vec, len: Option) -> Self { let len = len.unwrap_or_default(); let items = Vec::with_capacity(len); - Self { buffer, items } + Self { + variant, + buffer, + items, + } } } @@ -297,10 +381,10 @@ impl ser::SerializeSeq for ArraySerializer<'_> { fn serialize_element(&mut self, value: &T) -> Result<()> { let mut serializer = Serializer::new(); - let res = value.serialize(&mut serializer); + value.serialize(&mut serializer)?; let item_jsonb = OwnedJsonb::new(serializer.buffer); self.items.push(item_jsonb); - res + Ok(()) } fn end(self) -> Result { @@ -329,6 +413,52 @@ impl ser::SerializeTuple for ArraySerializer<'_> { } } +impl ser::SerializeTupleStruct for ArraySerializer<'_> { + type Ok = (); + + type Error = Error; + + fn serialize_field( + &mut self, + value: &T, + ) -> std::prelude::v1::Result<(), Self::Error> { + ::serialize_element(self, value) + } + + fn end(self) -> Result { + ::end(self) + } +} + +impl ser::SerializeTupleVariant for ArraySerializer<'_> { + type Ok = (); + + type Error = Error; + + fn serialize_field(&mut self, value: &T) -> Result<()> { + ::serialize_element(self, value) + } + + fn end(self) -> Result { + let Some(variant) = self.variant else { + return Err(ser::Error::custom("Variant can not be None".to_string())); + }; + + let mut builder = ArrayBuilder::with_capacity(self.items.len()); + for item in self.items.into_iter() { + builder.push_owned_jsonb(item); + } + let array_jsonb = builder.build()?; + + let mut builder = ObjectBuilder::new(); + builder.push_owned_jsonb(variant, array_jsonb)?; + let object_jsonb = builder.build()?; + let mut buf = object_jsonb.to_vec(); + self.buffer.append(&mut buf); + Ok(()) + } +} + pub struct ObjectSerializer<'a> { buffer: &'a mut Vec, keys: Vec, @@ -356,7 +486,7 @@ impl ser::SerializeMap for ObjectSerializer<'_> { fn serialize_key(&mut self, key: &T) -> Result<()> { let mut serializer = Serializer::new(); - let res = key.serialize(&mut serializer); + key.serialize(&mut serializer)?; let key_jsonb = OwnedJsonb::new(serializer.buffer); let raw_jsonb = key_jsonb.as_raw(); let Ok(Some(key)) = raw_jsonb.as_str() else { @@ -364,16 +494,16 @@ impl ser::SerializeMap for ObjectSerializer<'_> { }; self.keys.push(key.to_string()); - res + Ok(()) } fn serialize_value(&mut self, value: &T) -> Result<()> { let mut serializer = Serializer::new(); - let res = value.serialize(&mut serializer); + value.serialize(&mut serializer)?; let value_jsonb = OwnedJsonb::new(serializer.buffer); self.values.push(value_jsonb); - res + Ok(()) } fn end(self) -> Result { @@ -393,71 +523,209 @@ impl ser::SerializeMap for ObjectSerializer<'_> { } } -impl ser::SerializeStruct for ObjectSerializer<'_> { - type Ok = (); - - type Error = Error; +pub struct StructSerializer<'a> { + name: &'static str, + variant: Option<&'static str>, + buffer: &'a mut Vec, + values: Vec<(&'static str, OwnedJsonb)>, +} - fn serialize_field( - &mut self, - key: &'static str, - value: &T, - ) -> Result<()> { - ::serialize_key(self, key)?; - ::serialize_value(self, value) - } +impl<'a> StructSerializer<'a> { + fn new( + name: &'static str, + variant: Option<&'static str>, + buffer: &'a mut Vec, + len: usize, + ) -> Self { + let values = Vec::with_capacity(len); - fn end(self) -> Result { - Ok(()) + Self { + name, + variant, + buffer, + values, + } } } -impl ser::SerializeTupleStruct for ArraySerializer<'_> { +impl ser::SerializeStruct for StructSerializer<'_> { type Ok = (); type Error = Error; fn serialize_field( &mut self, - _value: &T, - ) -> std::prelude::v1::Result<(), Self::Error> { - todo!() + key: &'static str, + value: &T, + ) -> Result<()> { + let mut serializer = Serializer::new(); + value.serialize(&mut serializer)?; + let value_jsonb = OwnedJsonb::new(serializer.buffer); + self.values.push((key, value_jsonb)); + Ok(()) } - fn end(self) -> Result { - todo!() + fn end(mut self) -> Result { + if self.name == NUMBER_STRUCT_TOKEN { + let num = if self.values.len() == 2 { + let (value_field_name, value_jsonb) = self.values.remove(1); + let (scale_field_name, scale_jsonb) = self.values.remove(0); + + if scale_field_name != NUMBER_STRUCT_FIELD_SCALE + || value_field_name != NUMBER_STRUCT_FIELD_VALUE + { + return Err(ser::Error::custom(format!( + "Invalid number struct field names: scale={}, value={}", + scale_field_name, value_field_name + ))); + } + let scale = owned_jsonb_to_u64(scale_field_name, scale_jsonb)?; + let value = owned_jsonb_to_i128(value_field_name, value_jsonb)?; + + if (DECIMAL64_MIN..=DECIMAL64_MAX).contains(&value) + && scale <= MAX_DECIMAL64_PRECISION as u64 + { + Number::Decimal64(Decimal64 { + scale: scale as u8, + value: value as i64, + }) + } else if (DECIMAL128_MIN..=DECIMAL128_MAX).contains(&value) + && scale <= MAX_DECIMAL128_PRECISION as u64 + { + Number::Decimal128(Decimal128 { + scale: scale as u8, + value, + }) + } else if scale <= MAX_DECIMAL256_PRECISION as u64 { + Number::Decimal256(Decimal256 { + scale: scale as u8, + value: i256::from(value), + }) + } else { + return Err(ser::Error::custom(format!( + "Invalid number struct scale={} value={}", + scale, value + ))); + } + } else if self.values.len() == 3 { + let (low_value_field_name, low_value_jsonb) = self.values.remove(2); + let (high_value_field_name, high_value_jsonb) = self.values.remove(1); + let (scale_field_name, scale_jsonb) = self.values.remove(0); + + if scale_field_name != NUMBER_STRUCT_FIELD_SCALE + || high_value_field_name != NUMBER_STRUCT_FIELD_HIGH_VALUE + || low_value_field_name != NUMBER_STRUCT_FIELD_LOW_VALUE + { + return Err(ser::Error::custom(format!( + "Invalid number struct field names: scale={}, high_value={} low_value={}", + scale_field_name, high_value_field_name, low_value_field_name + ))); + } + let scale = owned_jsonb_to_u64(scale_field_name, scale_jsonb)?; + let high_value = owned_jsonb_to_i128(high_value_field_name, high_value_jsonb)?; + let low_value = owned_jsonb_to_i128(low_value_field_name, low_value_jsonb)?; + + if scale <= MAX_DECIMAL256_PRECISION as u64 { + Number::Decimal256(Decimal256 { + scale: scale as u8, + value: i256::from_words(high_value, low_value), + }) + } else { + return Err(ser::Error::custom(format!( + "Invalid number struct scale={} high_value={} low_value={}", + scale, high_value, low_value + ))); + } + } else { + return Err(ser::Error::custom(format!( + "Invalid number of fields for number struct: {}", + self.values.len() + ))); + }; + + let mut serializer = Serializer::new(); + serializer.write_number(num)?; + self.buffer.append(&mut serializer.buffer); + } else { + let mut builder = ObjectBuilder::new(); + for (key_str, value) in self.values.into_iter() { + builder.push_owned_jsonb(key_str, value)?; + } + let object_jsonb = builder.build()?; + let mut buf = object_jsonb.to_vec(); + self.buffer.append(&mut buf); + } + Ok(()) } } -impl ser::SerializeTupleVariant for ArraySerializer<'_> { - type Ok = (); - - type Error = Error; - - fn serialize_field(&mut self, _value: &T) -> Result<()> { - todo!() - } +fn owned_jsonb_to_u64(field: &str, owned_jsonb: OwnedJsonb) -> Result { + let raw_jsonb = owned_jsonb.as_raw(); + let Ok(Some(num)) = raw_jsonb.as_number() else { + return Err(ser::Error::custom(format!( + "Invalid number struct field={} value={}", + field, + raw_jsonb.to_string() + ))); + }; + let Some(num) = num.as_u64() else { + return Err(ser::Error::custom(format!( + "Invalid number struct to u64 field={} value={}", + field, + raw_jsonb.to_string() + ))); + }; + Ok(num) +} - fn end(self) -> Result { - todo!() - } +fn owned_jsonb_to_i128(field: &str, owned_jsonb: OwnedJsonb) -> Result { + let raw_jsonb = owned_jsonb.as_raw(); + let Ok(Some(num)) = raw_jsonb.as_number() else { + return Err(ser::Error::custom(format!( + "Invalid number struct field={} value={}", + field, + raw_jsonb.to_string() + ))); + }; + let Some(num) = num.as_i128() else { + return Err(ser::Error::custom(format!( + "Invalid number struct to i128 field={} value={}", + field, + raw_jsonb.to_string() + ))); + }; + Ok(num) } -impl ser::SerializeStructVariant for ArraySerializer<'_> { +impl ser::SerializeStructVariant for StructSerializer<'_> { type Ok = (); type Error = Error; fn serialize_field( &mut self, - _key: &'static str, - _value: &T, + key: &'static str, + value: &T, ) -> Result<()> { - todo!() + ::serialize_field(self, key, value) } fn end(self) -> Result { - todo!() + let Some(variant) = self.variant else { + return Err(ser::Error::custom("Variant can not be None".to_string())); + }; + let mut builder = ObjectBuilder::new(); + for (key_str, value) in self.values.into_iter() { + builder.push_owned_jsonb(key_str, value)?; + } + let object_jsonb = builder.build()?; + + let mut builder = ObjectBuilder::new(); + builder.push_owned_jsonb(variant, object_jsonb)?; + let object_jsonb = builder.build()?; + let mut buf = object_jsonb.to_vec(); + self.buffer.append(&mut buf); + Ok(()) } } @@ -511,9 +779,9 @@ impl Serialize for RawJsonb<'_> { } } ARRAY_CONTAINER_TAG => { - let mut serialize_seq = serializer.serialize_seq(Some(header_len as usize))?; + let mut serialize_seq = serializer.serialize_seq(Some(header_len))?; - let mut payload_start = index + 4 * header_len as usize; + let mut payload_start = index + 4 * header_len; for _ in 0..header_len { let jentry = self .read_jentry(index) @@ -559,10 +827,10 @@ impl Serialize for RawJsonb<'_> { serialize_seq.end() } OBJECT_CONTAINER_TAG => { - let mut serialize_map = serializer.serialize_map(Some(header_len as usize))?; + let mut serialize_map = serializer.serialize_map(Some(header_len))?; - let mut keys = VecDeque::with_capacity(header_len as usize); - let mut payload_start = index + 8 * header_len as usize; + let mut keys = VecDeque::with_capacity(header_len); + let mut payload_start = index + 8 * header_len; for _ in 0..header_len { let jentry = self .read_jentry(index) @@ -962,3 +1230,432 @@ impl<'a> JsonAstEncoder<'a> { jentry } } + +#[cfg(test)] +mod tests { + #[cfg(feature = "arbitrary_precision")] + use ethnum::i256; + use serde::Serialize; + use std::collections::HashMap; + + use crate::{core::databend::ser::Serializer, Value}; + #[cfg(feature = "arbitrary_precision")] + use crate::{Decimal128, Decimal256, Decimal64, Number}; + + // Basic struct with different field types + #[derive(Serialize)] + struct BasicStruct { + int_field: i32, + float_field: f64, + string_field: String, + bool_field: bool, + option_field: Option, + } + + // Nested struct + #[derive(Serialize)] + struct NestedStruct { + name: String, + basic: BasicStruct, + count: u32, + } + + // Struct with array fields + #[derive(Serialize)] + struct ArrayStruct { + names: Vec, + values: Vec, + } + + // Struct with map fields + #[derive(Serialize)] + struct MapStruct { + properties: HashMap, + } + + // Struct with tuple fields + #[derive(Serialize)] + struct TupleStruct { + pair: (i32, String), + triple: (bool, f64, String), + } + + // Enum for testing struct variants + #[derive(Serialize)] + enum TestEnum { + Simple(String), + Struct { id: u32, name: String }, + Tuple(i32, String, bool), + } + + #[test] + fn test_basic_struct_serialization() { + let basic = BasicStruct { + int_field: 42, + float_field: std::f64::consts::PI, + string_field: "hello".to_string(), + bool_field: true, + option_field: Some("present".to_string()), + }; + + let mut serializer = Serializer::new(); + basic.serialize(&mut serializer).unwrap(); + let jsonb = serializer.to_owned_jsonb(); + let raw = jsonb.as_raw(); + + // Verify the serialized structure + let value = raw.to_value().unwrap(); + if let Value::Object(obj) = value { + assert_eq!(obj.len(), 5); + assert_eq!(obj.get("int_field").unwrap().as_i64().unwrap(), 42); + assert_eq!( + obj.get("float_field").unwrap().as_f64().unwrap(), + std::f64::consts::PI + ); + assert_eq!(obj.get("string_field").unwrap().as_str().unwrap(), "hello"); + assert!(obj.get("bool_field").unwrap().as_bool().unwrap()); + assert_eq!( + obj.get("option_field").unwrap().as_str().unwrap(), + "present" + ); + } else { + panic!("Expected object value"); + } + + // Test with None option + let basic_none = BasicStruct { + int_field: 100, + float_field: 12.34f64, + string_field: "world".to_string(), + bool_field: false, + option_field: None, + }; + + let mut serializer = Serializer::new(); + basic_none.serialize(&mut serializer).unwrap(); + let jsonb = serializer.to_owned_jsonb(); + let raw = jsonb.as_raw(); + + let value = raw.to_value().unwrap(); + if let Value::Object(obj) = value { + assert_eq!(obj.len(), 5); + assert_eq!(obj.get("int_field").unwrap().as_i64().unwrap(), 100); + assert_eq!(obj.get("float_field").unwrap().as_f64().unwrap(), 12.34f64); + assert_eq!(obj.get("string_field").unwrap().as_str().unwrap(), "world"); + assert!(!obj.get("bool_field").unwrap().as_bool().unwrap()); + assert!(obj.get("option_field").unwrap().is_null()); + } else { + panic!("Expected object value"); + } + } + + #[test] + fn test_nested_struct_serialization() { + let nested = NestedStruct { + name: "nested".to_string(), + basic: BasicStruct { + int_field: 100, + float_field: 2.71f64, + string_field: "world".to_string(), + bool_field: false, + option_field: None, + }, + count: 5, + }; + + let mut serializer = Serializer::new(); + nested.serialize(&mut serializer).unwrap(); + let jsonb = serializer.to_owned_jsonb(); + let raw = jsonb.as_raw(); + + let value = raw.to_value().unwrap(); + if let Value::Object(obj) = value { + assert_eq!(obj.len(), 3); + assert_eq!(obj.get("name").unwrap().as_str().unwrap(), "nested"); + assert_eq!(obj.get("count").unwrap().as_u64().unwrap(), 5); + + if let Value::Object(basic_obj) = obj.get("basic").unwrap() { + assert_eq!(basic_obj.len(), 5); + assert_eq!(basic_obj.get("int_field").unwrap().as_i64().unwrap(), 100); + assert_eq!( + basic_obj.get("float_field").unwrap().as_f64().unwrap(), + 2.71f64 + ); + assert_eq!( + basic_obj.get("string_field").unwrap().as_str().unwrap(), + "world" + ); + assert!(!basic_obj.get("bool_field").unwrap().as_bool().unwrap()); + assert!(basic_obj.get("option_field").unwrap().is_null()); + } else { + panic!("Expected object for basic field"); + } + } else { + panic!("Expected object value"); + } + } + + #[test] + fn test_array_struct_serialization() { + let array_struct = ArrayStruct { + names: vec![ + "Alice".to_string(), + "Bob".to_string(), + "Charlie".to_string(), + ], + values: vec![1, 2, 3, 4, 5], + }; + + let mut serializer = Serializer::new(); + array_struct.serialize(&mut serializer).unwrap(); + let jsonb = serializer.to_owned_jsonb(); + let raw = jsonb.as_raw(); + + let value = raw.to_value().unwrap(); + if let Value::Object(obj) = value { + assert_eq!(obj.len(), 2); + + if let Value::Array(names) = obj.get("names").unwrap() { + assert_eq!(names.len(), 3); + assert_eq!(names[0].as_str().unwrap(), "Alice"); + assert_eq!(names[1].as_str().unwrap(), "Bob"); + assert_eq!(names[2].as_str().unwrap(), "Charlie"); + } else { + panic!("Expected array for names field"); + } + + if let Value::Array(values) = obj.get("values").unwrap() { + assert_eq!(values.len(), 5); + for (i, val) in values.iter().enumerate() { + assert_eq!(val.as_i64().unwrap(), (i + 1) as i64); + } + } else { + panic!("Expected array for values field"); + } + } else { + panic!("Expected object value"); + } + } + + #[test] + fn test_map_struct_serialization() { + let mut properties = HashMap::new(); + properties.insert("key1".to_string(), "value1".to_string()); + properties.insert("key2".to_string(), "value2".to_string()); + properties.insert("key3".to_string(), "value3".to_string()); + + let map_struct = MapStruct { properties }; + + let mut serializer = Serializer::new(); + map_struct.serialize(&mut serializer).unwrap(); + let jsonb = serializer.to_owned_jsonb(); + let raw = jsonb.as_raw(); + + let value = raw.to_value().unwrap(); + if let Value::Object(obj) = value { + assert_eq!(obj.len(), 1); + + if let Value::Object(props) = obj.get("properties").unwrap() { + assert_eq!(props.len(), 3); + assert_eq!(props.get("key1").unwrap().as_str().unwrap(), "value1"); + assert_eq!(props.get("key2").unwrap().as_str().unwrap(), "value2"); + assert_eq!(props.get("key3").unwrap().as_str().unwrap(), "value3"); + } else { + panic!("Expected object for properties field"); + } + } else { + panic!("Expected object value"); + } + } + + #[test] + fn test_tuple_struct_serialization() { + let tuple_struct = TupleStruct { + pair: (42, "answer".to_string()), + triple: (true, std::f64::consts::PI, "pi".to_string()), + }; + + let mut serializer = Serializer::new(); + tuple_struct.serialize(&mut serializer).unwrap(); + let jsonb = serializer.to_owned_jsonb(); + let raw = jsonb.as_raw(); + + let value = raw.to_value().unwrap(); + if let Value::Object(obj) = value { + assert_eq!(obj.len(), 2); + + if let Value::Array(pair) = obj.get("pair").unwrap() { + assert_eq!(pair.len(), 2); + assert_eq!(pair[0].as_i64().unwrap(), 42); + assert_eq!(pair[1].as_str().unwrap(), "answer"); + } else { + panic!("Expected array for pair field"); + } + + if let Value::Array(triple) = obj.get("triple").unwrap() { + assert_eq!(triple.len(), 3); + assert!(triple[0].as_bool().unwrap()); + assert_eq!(triple[1].as_f64().unwrap(), std::f64::consts::PI); + assert_eq!(triple[2].as_str().unwrap(), "pi"); + } else { + panic!("Expected array for triple field"); + } + } else { + panic!("Expected object value"); + } + } + + #[test] + #[cfg(feature = "arbitrary_precision")] + fn test_number_struct_serialization() { + // Test with a value that should be serialized as Decimal64 + let decimal64_struct = Number::Decimal64(Decimal64 { + scale: 2, + value: 1234, + }); + + let mut serializer = Serializer::new(); + decimal64_struct.serialize(&mut serializer).unwrap(); + let jsonb = serializer.to_owned_jsonb(); + let raw = jsonb.as_raw(); + + let value = raw.to_value().unwrap(); + if let Value::Number(num) = value { + match num { + Number::Decimal64(dec) => { + assert_eq!(dec.scale, 2); + assert_eq!(dec.value, 1234); + } + _ => panic!("Expected Decimal64 number"), + } + } else { + panic!("Expected number value"); + } + + // Test with a value that should be serialized as Decimal128 + let decimal128_struct = Number::Decimal128(Decimal128 { + scale: 10, + value: 1000000000000000000485i128, + }); + + let mut serializer = Serializer::new(); + decimal128_struct.serialize(&mut serializer).unwrap(); + let jsonb = serializer.to_owned_jsonb(); + let raw = jsonb.as_raw(); + + let value = raw.to_value().unwrap(); + if let Value::Number(num) = value { + match num { + Number::Decimal128(dec) => { + assert_eq!(dec.scale, 10); + assert_eq!(dec.value, 1000000000000000000485i128); + } + _ => panic!("Expected Decimal128 number"), + } + } else { + panic!("Expected number value"); + } + + // Test with a value that should be serialized as Decimal256 + let decimal256_struct = Number::Decimal256(Decimal256 { + scale: 15, + value: i256::from_words( + 999999999999999999999999999999999999i128, + 999999999999999999999999999999999999i128, + ), + }); + + let mut serializer = Serializer::new(); + decimal256_struct.serialize(&mut serializer).unwrap(); + let jsonb = serializer.to_owned_jsonb(); + let raw = jsonb.as_raw(); + + let value = raw.to_value().unwrap(); + if let Value::Number(num) = value { + match num { + Number::Decimal256(dec) => { + assert_eq!(dec.scale, 15); + assert_eq!( + dec.value, + i256::from_words( + 999999999999999999999999999999999999i128, + 999999999999999999999999999999999999i128 + ) + ); + } + _ => panic!("Expected Decimal256 number"), + } + } else { + panic!("Expected number value"); + } + } + + #[test] + fn test_enum_serialization() { + // Test simple variant + let simple = TestEnum::Simple("simple value".to_string()); + let mut serializer = Serializer::new(); + simple.serialize(&mut serializer).unwrap(); + let jsonb = serializer.to_owned_jsonb(); + let raw = jsonb.as_raw(); + + let value = raw.to_value().unwrap(); + if let Value::Object(obj) = value { + assert_eq!(obj.len(), 1); + assert_eq!(obj.get("Simple").unwrap().as_str().unwrap(), "simple value"); + } else { + panic!("Expected object"); + } + + // Test struct variant + let struct_variant = TestEnum::Struct { + id: 123, + name: "test name".to_string(), + }; + let mut serializer = Serializer::new(); + struct_variant.serialize(&mut serializer).unwrap(); + let jsonb = serializer.to_owned_jsonb(); + let raw = jsonb.as_raw(); + + let value = raw.to_value().unwrap(); + if let Value::Object(obj) = value { + assert_eq!(obj.len(), 1); + + if let Value::Object(struct_obj) = obj.get("Struct").unwrap() { + assert_eq!(struct_obj.len(), 2); + assert_eq!(struct_obj.get("id").unwrap().as_u64().unwrap(), 123); + assert_eq!( + struct_obj.get("name").unwrap().as_str().unwrap(), + "test name" + ); + } else { + panic!("Expected object for Struct variant"); + } + } else { + panic!("Expected object value for Struct variant"); + } + + // Test tuple variant + let tuple_variant = TestEnum::Tuple(42, "tuple value".to_string(), true); + let mut serializer = Serializer::new(); + tuple_variant.serialize(&mut serializer).unwrap(); + let jsonb = serializer.to_owned_jsonb(); + let raw = jsonb.as_raw(); + + let value = raw.to_value().unwrap(); + if let Value::Object(obj) = value { + assert_eq!(obj.len(), 1); + + if let Value::Array(tuple_arr) = obj.get("Tuple").unwrap() { + assert_eq!(tuple_arr.len(), 3); + assert_eq!(tuple_arr[0].as_i64().unwrap(), 42); + assert_eq!(tuple_arr[1].as_str().unwrap(), "tuple value"); + assert!(tuple_arr[2].as_bool().unwrap()); + } else { + panic!("Expected array for Tuple variant"); + } + } else { + panic!("Expected object value for Tuple variant"); + } + } +} diff --git a/src/core/databend/util.rs b/src/core/databend/util.rs index 16eaf12..1a0ca65 100644 --- a/src/core/databend/util.rs +++ b/src/core/databend/util.rs @@ -22,8 +22,10 @@ use ethnum::i256; use super::constants::*; use super::jentry::JEntry; +use crate::core::ExtensionItem; use crate::core::JsonbItem; use crate::core::JsonbItemType; +use crate::core::NumberItem; use crate::error::*; use crate::extension::Date; use crate::extension::ExtensionValue; @@ -34,27 +36,28 @@ use crate::number::Decimal128; use crate::number::Decimal256; use crate::number::Decimal64; use crate::Number; +use crate::Object; use crate::OwnedJsonb; use crate::RawJsonb; +use crate::Value; +use std::collections::VecDeque; impl<'a> JsonbItem<'a> { pub(crate) fn from_raw_jsonb(raw_jsonb: RawJsonb<'a>) -> Result> { let (header_type, _) = raw_jsonb.read_header(0)?; match header_type { SCALAR_CONTAINER_TAG => { - let jentry = raw_jsonb.read_jentry(4)?; - let range = Range { - start: 8, - end: raw_jsonb.len(), - }; - let data = raw_jsonb.slice(range)?; + let (jentry, data) = raw_jsonb.read_jentry_and_data(4, 8)?; let item = match jentry.type_code { NULL_TAG => JsonbItem::Null, TRUE_TAG => JsonbItem::Boolean(true), FALSE_TAG => JsonbItem::Boolean(false), - NUMBER_TAG => JsonbItem::Number(data), - STRING_TAG => JsonbItem::String(data), - EXTENSION_TAG => JsonbItem::Extension(data), + NUMBER_TAG => JsonbItem::Number(NumberItem::Raw(data)), + STRING_TAG => { + let s = Cow::Borrowed(unsafe { std::str::from_utf8_unchecked(data) }); + JsonbItem::String(s) + } + EXTENSION_TAG => JsonbItem::Extension(ExtensionItem::Raw(data)), _ => { return Err(Error::InvalidJsonb); } @@ -68,13 +71,49 @@ impl<'a> JsonbItem<'a> { } impl<'a> RawJsonb<'a> { + /// Returns the type information of a JSONB data item + /// + /// # Returns + /// + /// Returns `Result`, where `JsonbItemType` can be one of: + /// - `JsonbItemType::Null` - represents a null value + /// - `JsonbItemType::Boolean` - represents a boolean value (true or false) + /// - `JsonbItemType::Number` - represents a numeric type + /// - `JsonbItemType::String` - represents a string type + /// - `JsonbItemType::Extension` - represents an extension type + /// - `JsonbItemType::Array(n)` - represents an array type, where n is the number of array elements + /// - `JsonbItemType::Object(n)` - represents an object type, where n is the number of key-value pairs + /// + /// # Errors + /// + /// Returns `Error::InvalidJsonb` if the JSONB data format is invalid + /// + /// # Examples + /// + /// ```rust + /// use jsonb::{OwnedJsonb, core::JsonbItemType}; + /// + /// // Create a JSONB containing an array + /// let jsonb = "[1, 2, 3]".parse::().unwrap(); + /// let raw_jsonb = jsonb.as_raw(); + /// + /// // Get type information + /// let item_type = raw_jsonb.jsonb_item_type().unwrap(); + /// assert!(matches!(item_type, JsonbItemType::Array(3))); + /// + /// // Create a JSONB containing an object + /// let jsonb = r#"{"name": "Alice", "age": 30}"#.parse::().unwrap(); + /// let raw_jsonb = jsonb.as_raw(); + /// + /// // Get type information + /// let item_type = raw_jsonb.jsonb_item_type().unwrap(); + /// assert!(matches!(item_type, JsonbItemType::Object(2))); + /// ``` pub fn jsonb_item_type(&self) -> Result { - let mut index = 0; - let (header_type, header_len) = self.read_header(index)?; - index += 4; + let (header_type, header_len) = self.read_header(0)?; match header_type { SCALAR_CONTAINER_TAG => { - let jentry = self.read_jentry(index)?; + let jentry = self.read_jentry(4)?; match jentry.type_code { NULL_TAG => Ok(JsonbItemType::Null), @@ -86,8 +125,109 @@ impl<'a> RawJsonb<'a> { _ => Err(Error::InvalidJsonb), } } - ARRAY_CONTAINER_TAG => Ok(JsonbItemType::Array(header_len as usize)), - OBJECT_CONTAINER_TAG => Ok(JsonbItemType::Object(header_len as usize)), + ARRAY_CONTAINER_TAG => Ok(JsonbItemType::Array(header_len)), + OBJECT_CONTAINER_TAG => Ok(JsonbItemType::Object(header_len)), + _ => Err(Error::InvalidJsonb), + } + } + + /// Converts JSONB binary data to a `Value` type + /// + /// This function recursively parses JSONB binary data and converts it into + /// an in-memory `Value` structure, making it convenient to access and + /// manipulate JSON data. + /// + /// # Returns + /// + /// Returns `Result>`, where `Value` can be one of: + /// - `Value::Null` - represents a null value + /// - `Value::Bool(bool)` - represents a boolean value + /// - `Value::Number(Number)` - represents a numeric value + /// - `Value::String(Cow<'a, str>)` - represents a string + /// - `Value::Binary(&'a [u8])` - represents binary data + /// - `Value::Date(Date)` - represents a date + /// - `Value::Timestamp(Timestamp)` - represents a timestamp + /// - `Value::TimestampTz(TimestampTz)` - represents a timestamp with timezone + /// - `Value::Interval(Interval)` - represents a time interval + /// - `Value::Array(Vec>)` - represents an array + /// - `Value::Object(Object)` - represents an object + /// + /// # Errors + /// + /// Returns `Error::InvalidJsonb` if the JSONB data format is invalid + /// + /// # Examples + /// + /// ```rust + /// use jsonb::{OwnedJsonb, Value}; + /// + /// // Parse a simple JSON object + /// let jsonb = r#"{"name": "Alice", "age": 30, "is_student": false}"#.parse::().unwrap(); + /// let raw_jsonb = jsonb.as_raw(); + /// + /// // Convert to Value type + /// let value = raw_jsonb.to_value().unwrap(); + /// assert!(matches!(value, Value::Object(_))); + /// + /// // Access object properties + /// if let Value::Object(obj) = value { + /// assert_eq!(obj.get("name").unwrap().as_str().unwrap(), "Alice"); + /// assert_eq!(obj.get("age").unwrap().as_i64().unwrap(), 30); + /// assert_eq!(obj.get("is_student").unwrap().as_bool().unwrap(), false); + /// } + /// ``` + pub fn to_value(&self) -> Result> { + let mut index = 0; + let (header_type, header_len) = self.read_header(index)?; + index += 4; + match header_type { + SCALAR_CONTAINER_TAG => { + let data_index = index + 4; + let (jentry, data) = self.read_jentry_and_data(index, data_index)?; + jentry_to_value(jentry, data) + } + ARRAY_CONTAINER_TAG => { + let mut arr = Vec::with_capacity(header_len); + let mut data_index = index + header_len * 4; + for _ in 0..header_len { + let (jentry, data) = self.read_jentry_and_data(index, data_index)?; + index += 4; + data_index += jentry.length as usize; + + let value = jentry_to_value(jentry, data)?; + arr.push(value); + } + Ok(Value::Array(arr)) + } + OBJECT_CONTAINER_TAG => { + let mut obj = Object::new(); + let mut keys = VecDeque::with_capacity(header_len); + let mut data_index = index + header_len * 8; + for _ in 0..header_len { + let (jentry, data) = self.read_jentry_and_data(index, data_index)?; + index += 4; + data_index += jentry.length as usize; + + let key = match jentry.type_code { + STRING_TAG => unsafe { String::from_utf8_unchecked(data.to_vec()) }, + _ => { + return Err(Error::InvalidJsonb); + } + }; + keys.push_back(key); + } + for _ in 0..header_len { + let key = keys.pop_front().unwrap(); + let (jentry, data) = self.read_jentry_and_data(index, data_index)?; + index += 4; + data_index += jentry.length as usize; + + let value = jentry_to_value(jentry, data)?; + obj.insert(key, value); + } + + Ok(Value::Object(obj)) + } _ => Err(Error::InvalidJsonb), } } @@ -97,11 +237,10 @@ impl<'a> RawJsonb<'a> { key_name: &Cow<'a, str>, eq_func: impl Fn(&[u8], &[u8]) -> bool, ) -> Result>> { - let (header_type, header_len) = self.read_header(0)?; - if header_type != OBJECT_CONTAINER_TAG || header_len == 0 { + let (header_type, length) = self.read_header(0)?; + if header_type != OBJECT_CONTAINER_TAG || length == 0 { return Ok(None); } - let length = header_len as usize; let mut index = 0; let mut jentry_offset = 4; let mut item_offset = 4 + 8 * length; @@ -156,7 +295,7 @@ impl<'a> RawJsonb<'a> { Ok(Some(value_item)) } - pub(super) fn read_header(&self, index: usize) -> Result<(u32, u32)> { + pub(super) fn read_header(&self, index: usize) -> Result<(u32, usize)> { let header = self.read_u32(index)?; let header_type = header & CONTAINER_HEADER_TYPE_MASK; match header_type { @@ -166,7 +305,7 @@ impl<'a> RawJsonb<'a> { } } let header_len = header & CONTAINER_HEADER_LEN_MASK; - Ok((header_type, header_len)) + Ok((header_type, header_len as usize)) } pub(super) fn read_jentry(&self, index: usize) -> Result { @@ -175,6 +314,20 @@ impl<'a> RawJsonb<'a> { Ok(jentry) } + pub(super) fn read_jentry_and_data( + &self, + index: usize, + data_index: usize, + ) -> Result<(JEntry, &'a [u8])> { + let jentry = self.read_jentry(index)?; + let range = Range { + start: data_index, + end: data_index + jentry.length as usize, + }; + let data = self.slice(range)?; + Ok((jentry, data)) + } + pub(super) fn read_u32(&self, idx: usize) -> Result { let bytes: [u8; 4] = self .data @@ -196,52 +349,77 @@ impl<'a> RawJsonb<'a> { impl OwnedJsonb { pub(crate) fn from_item(item: JsonbItem<'_>) -> Result { - let (jentry, data) = match item { - JsonbItem::Null => { - let jentry = JEntry::make_null_jentry(); - (jentry, None) - } - JsonbItem::Boolean(v) => { - let jentry = if v { - JEntry::make_true_jentry() - } else { - JEntry::make_false_jentry() + match item { + JsonbItem::Raw(raw_jsonb) => Ok(raw_jsonb.to_owned()), + JsonbItem::Owned(owned_jsonb) => Ok(owned_jsonb), + _ => { + let mut len = match item { + JsonbItem::Null => 0, + JsonbItem::Boolean(_) => 0, + JsonbItem::Number(NumberItem::Raw(data)) => data.len(), + JsonbItem::String(ref s) => s.len(), + JsonbItem::Extension(ExtensionItem::Raw(data)) => data.len(), + // The estimated lengths for number and extension. + _ => 10, }; - (jentry, None) - } - JsonbItem::Number(data) => { - let jentry = JEntry::make_number_jentry(data.len()); - (jentry, Some(data)) - } - JsonbItem::String(data) => { - let jentry = JEntry::make_string_jentry(data.len()); - (jentry, Some(data)) - } - JsonbItem::Extension(data) => { - let jentry = JEntry::make_extension_jentry(data.len()); - (jentry, Some(data)) - } - JsonbItem::Raw(raw_jsonb) => { - return Ok(raw_jsonb.to_owned()); - } - JsonbItem::Owned(owned_jsonb) => { - return Ok(owned_jsonb.clone()); - } - }; - let len = if let Some(data) = data { - data.len() + 8 - } else { - 8 - }; - let mut buf = Vec::with_capacity(len); - let header = SCALAR_CONTAINER_TAG; - buf.write_u32::(header)?; - buf.write_u32::(jentry.encoded())?; - if let Some(data) = data { - buf.extend_from_slice(data); + // add the length of header and jentry. + len += 8; + let mut buf = Vec::with_capacity(len); + let header = SCALAR_CONTAINER_TAG; + buf.write_u32::(header)?; + + match item { + JsonbItem::Null => { + let jentry = JEntry::make_null_jentry(); + buf.write_u32::(jentry.encoded())?; + } + JsonbItem::Boolean(v) => { + let jentry = if v { + JEntry::make_true_jentry() + } else { + JEntry::make_false_jentry() + }; + buf.write_u32::(jentry.encoded())?; + } + JsonbItem::Number(num) => match num { + NumberItem::Raw(data) => { + let jentry = JEntry::make_number_jentry(data.len()); + buf.write_u32::(jentry.encoded())?; + buf.extend_from_slice(data); + } + NumberItem::Number(num) => { + let mut data = vec![]; + let len = num.compact_encode(&mut data)?; + let jentry = JEntry::make_number_jentry(len); + buf.write_u32::(jentry.encoded())?; + buf.extend_from_slice(&data); + } + }, + JsonbItem::String(s) => { + let jentry = JEntry::make_string_jentry(s.len()); + buf.write_u32::(jentry.encoded())?; + buf.extend_from_slice(s.as_bytes()); + } + JsonbItem::Extension(ext) => match ext { + ExtensionItem::Raw(data) => { + let jentry = JEntry::make_extension_jentry(data.len()); + buf.write_u32::(jentry.encoded())?; + buf.extend_from_slice(data); + } + ExtensionItem::Extension(ext) => { + let mut data = vec![]; + let len = ext.compact_encode(&mut data)?; + let jentry = JEntry::make_extension_jentry(len); + buf.write_u32::(jentry.encoded())?; + buf.extend_from_slice(&data); + } + }, + _ => unreachable!(), + } + Ok(OwnedJsonb::new(buf)) + } } - Ok(OwnedJsonb::new(buf)) } } @@ -501,10 +679,44 @@ pub(super) fn jentry_to_jsonb_item(jentry: JEntry, data: &[u8]) -> JsonbItem<'_> NULL_TAG => JsonbItem::Null, TRUE_TAG => JsonbItem::Boolean(true), FALSE_TAG => JsonbItem::Boolean(false), - NUMBER_TAG => JsonbItem::Number(data), - STRING_TAG => JsonbItem::String(data), - EXTENSION_TAG => JsonbItem::Extension(data), + NUMBER_TAG => JsonbItem::Number(NumberItem::Raw(data)), + STRING_TAG => { + let s = Cow::Borrowed(unsafe { std::str::from_utf8_unchecked(data) }); + JsonbItem::String(s) + } + EXTENSION_TAG => JsonbItem::Extension(ExtensionItem::Raw(data)), CONTAINER_TAG => JsonbItem::Raw(RawJsonb::new(data)), _ => unreachable!(), } } + +pub(super) fn jentry_to_value(jentry: JEntry, data: &[u8]) -> Result> { + match jentry.type_code { + NULL_TAG => Ok(Value::Null), + TRUE_TAG => Ok(Value::Bool(true)), + FALSE_TAG => Ok(Value::Bool(false)), + NUMBER_TAG => { + let n = Number::decode(data)?; + Ok(Value::Number(n)) + } + STRING_TAG => { + let s = Cow::Borrowed(unsafe { std::str::from_utf8_unchecked(data) }); + Ok(Value::String(s)) + } + EXTENSION_TAG => { + let x = ExtensionValue::decode(data)?; + match x { + ExtensionValue::Binary(v) => Ok(Value::Binary(v)), + ExtensionValue::Date(v) => Ok(Value::Date(v)), + ExtensionValue::Timestamp(v) => Ok(Value::Timestamp(v)), + ExtensionValue::TimestampTz(v) => Ok(Value::TimestampTz(v)), + ExtensionValue::Interval(v) => Ok(Value::Interval(v)), + } + } + CONTAINER_TAG => { + let raw = RawJsonb::new(data); + raw.to_value() + } + _ => Err(Error::InvalidJsonb), + } +} diff --git a/src/core/item.rs b/src/core/item.rs index 18c7215..7a439e9 100644 --- a/src/core/item.rs +++ b/src/core/item.rs @@ -80,6 +80,48 @@ impl PartialOrd for JsonbItemType { } } +#[derive(Debug, Clone)] +pub(crate) enum NumberItem<'a> { + /// Represents a raw JSONB number, stored as a byte slice. + Raw(&'a [u8]), + /// Represents a JSONB number. + #[allow(dead_code)] + Number(Number), +} + +impl NumberItem<'_> { + pub(crate) fn as_number(&self) -> Result { + match self { + NumberItem::Raw(data) => { + let num = Number::decode(data)?; + Ok(num) + } + NumberItem::Number(num) => Ok(num.clone()), + } + } +} + +#[derive(Debug, Clone)] +pub(crate) enum ExtensionItem<'a> { + /// Represents a raw JSONB extension value, stored as a byte slice. + Raw(&'a [u8]), + /// Represents a raw JSONB extension value. + #[allow(dead_code)] + Extension(ExtensionValue<'a>), +} + +impl<'a> ExtensionItem<'a> { + pub(crate) fn as_extension_value(&self) -> Result> { + match self { + ExtensionItem::Raw(data) => { + let val = ExtensionValue::decode(data)?; + Ok(val) + } + ExtensionItem::Extension(val) => Ok(val.clone()), + } + } +} + /// `JsonbItem` is an internal enum used primarily within `ArrayIterator` and /// `ObjectIterator` to represent temporary values during iteration. It is also /// utilized by `ArrayBuilder` and `ObjectBuilder` to store intermediate variables @@ -96,11 +138,11 @@ pub(crate) enum JsonbItem<'a> { /// Represents a JSONB boolean value. Boolean(bool), /// Represents a JSONB number, stored as a byte slice. - Number(&'a [u8]), - /// Represents a JSONB string, stored as a byte slice. - String(&'a [u8]), + Number(NumberItem<'a>), + /// Represents a JSONB string. + String(Cow<'a, str>), /// Represents a JSONB extension values, stored as a byte slice. - Extension(&'a [u8]), + Extension(ExtensionItem<'a>), /// Represents raw JSONB data, using a borrowed slice. Raw(RawJsonb<'a>), /// Represents owned JSONB data. @@ -134,12 +176,9 @@ impl<'a> JsonbItem<'a> { } } - pub(crate) fn as_str(&self) -> Option<&'a str> { + pub(crate) fn as_str(&self) -> Option> { match self { - JsonbItem::String(data) => { - let s = unsafe { std::str::from_utf8_unchecked(data) }; - Some(s) - } + JsonbItem::String(s) => Some(s.clone()), _ => None, } } @@ -184,22 +223,22 @@ impl PartialOrd for JsonbItem<'_> { (JsonbItem::Raw(_), JsonbItem::Null) => Some(Ordering::Less), (JsonbItem::Null, JsonbItem::Raw(_)) => Some(Ordering::Greater), // compare extension - (JsonbItem::Extension(self_data), JsonbItem::Extension(other_data)) => { - let self_val = ExtensionValue::decode(self_data).ok()?; - let other_val = ExtensionValue::decode(other_data).ok()?; + (JsonbItem::Extension(self_ext), JsonbItem::Extension(other_ext)) => { + let self_val = self_ext.as_extension_value().ok()?; + let other_val = other_ext.as_extension_value().ok()?; self_val.partial_cmp(&other_val) } - (JsonbItem::Raw(self_raw), JsonbItem::Extension(other_data)) => { + (JsonbItem::Raw(self_raw), JsonbItem::Extension(other_ext)) => { let self_val = self_raw.as_extension_value(); - let other_val = ExtensionValue::decode(other_data).ok()?; + let other_val = other_ext.as_extension_value().ok()?; if let Ok(Some(self_val)) = self_val { self_val.partial_cmp(&other_val) } else { None } } - (JsonbItem::Extension(self_data), JsonbItem::Raw(other_raw)) => { - let self_val = ExtensionValue::decode(self_data).ok()?; + (JsonbItem::Extension(self_ext), JsonbItem::Raw(other_raw)) => { + let self_val = self_ext.as_extension_value().ok()?; let other_val = other_raw.as_extension_value(); if let Ok(Some(other_val)) = other_val { self_val.partial_cmp(&other_val) @@ -228,46 +267,42 @@ impl PartialOrd for JsonbItem<'_> { } } // compare number - (JsonbItem::Number(self_data), JsonbItem::Number(other_data)) => { - let self_num = Number::decode(self_data).ok()?; - let other_num = Number::decode(other_data).ok()?; - self_num.partial_cmp(&other_num) + (JsonbItem::Number(self_num), JsonbItem::Number(other_num)) => { + let self_val = self_num.as_number().ok()?; + let other_val = other_num.as_number().ok()?; + self_val.partial_cmp(&other_val) } - (JsonbItem::Raw(self_raw), JsonbItem::Number(other_data)) => { - let self_num = self_raw.as_number(); - let other_num = Number::decode(other_data).ok()?; - if let Ok(Some(self_num)) = self_num { - self_num.partial_cmp(&other_num) + (JsonbItem::Raw(self_raw), JsonbItem::Number(other_num)) => { + let self_val = self_raw.as_number(); + let other_val = other_num.as_number().ok()?; + if let Ok(Some(self_val)) = self_val { + self_val.partial_cmp(&other_val) } else { None } } - (JsonbItem::Number(self_data), JsonbItem::Raw(other_raw)) => { - let self_num = Number::decode(self_data).ok()?; - let other_num = other_raw.as_number(); - if let Ok(Some(other_num)) = other_num { - self_num.partial_cmp(&other_num) + (JsonbItem::Number(self_num), JsonbItem::Raw(other_raw)) => { + let self_val = self_num.as_number().ok()?; + let other_val = other_raw.as_number(); + if let Ok(Some(other_val)) = other_val { + self_val.partial_cmp(&other_val) } else { None } } // compare string - (JsonbItem::String(self_data), JsonbItem::String(other_data)) => { - let self_str = unsafe { std::str::from_utf8_unchecked(self_data) }; - let other_str = unsafe { std::str::from_utf8_unchecked(other_data) }; + (JsonbItem::String(self_str), JsonbItem::String(other_str)) => { self_str.partial_cmp(other_str) } - (JsonbItem::Raw(self_raw), JsonbItem::String(other_data)) => { + (JsonbItem::Raw(self_raw), JsonbItem::String(other_str)) => { let self_str = self_raw.as_str(); - let other_str = Cow::Borrowed(unsafe { std::str::from_utf8_unchecked(other_data) }); if let Ok(Some(self_str)) = self_str { - self_str.partial_cmp(&other_str) + self_str.partial_cmp(other_str) } else { None } } - (JsonbItem::String(self_data), JsonbItem::Raw(other_raw)) => { - let self_str = Cow::Borrowed(unsafe { std::str::from_utf8_unchecked(self_data) }); + (JsonbItem::String(self_str), JsonbItem::Raw(other_raw)) => { let other_str = other_raw.as_str(); if let Ok(Some(other_str)) = other_str { self_str.partial_cmp(&other_str) diff --git a/src/core/mod.rs b/src/core/mod.rs index 44c17bd..bc06af6 100644 --- a/src/core/mod.rs +++ b/src/core/mod.rs @@ -21,7 +21,7 @@ mod item; #[cfg(feature = "databend")] pub use databend::*; -pub(crate) use item::*; +pub use item::*; #[cfg(feature = "sqlite")] #[allow(unused_imports)] pub use sqlite::*; diff --git a/src/functions/operator.rs b/src/functions/operator.rs index d5cf68b..00008ee 100644 --- a/src/functions/operator.rs +++ b/src/functions/operator.rs @@ -14,6 +14,7 @@ // This file contains functions that don't neatly fit into other categories. +use std::borrow::Cow; use std::collections::BTreeMap; use std::collections::BTreeSet; use std::collections::VecDeque; @@ -21,6 +22,7 @@ use std::collections::VecDeque; use crate::constants::*; use crate::core::ArrayBuilder; use crate::core::ArrayIterator; +use crate::core::ExtensionItem; use crate::core::JsonbItem; use crate::core::JsonbItemType; use crate::core::ObjectBuilder; @@ -96,8 +98,8 @@ impl RawJsonb<'_> { JsonbItemType::Number => { let jsonb_item = JsonbItem::from_raw_jsonb(*self)?; match jsonb_item { - JsonbItem::Number(data) => { - let val = Number::decode(data)?; + JsonbItem::Number(num) => { + let val = num.as_number()?; match val { Number::UInt64(_) | Number::Int64(_) => Ok(TYPE_INTEGER), Number::Decimal64(_) @@ -113,8 +115,8 @@ impl RawJsonb<'_> { JsonbItemType::Extension => { let jsonb_item = JsonbItem::from_raw_jsonb(*self)?; match jsonb_item { - JsonbItem::Extension(data) => { - let val = ExtensionValue::decode(data)?; + JsonbItem::Extension(ext) => { + let val = ext.as_extension_value()?; match val { ExtensionValue::Binary(_v) => Ok(TYPE_BINARY), ExtensionValue::Date(_v) => Ok(TYPE_DATE), @@ -656,7 +658,7 @@ impl RawJsonb<'_> { let mut object_iter = ObjectIterator::new(*self)?.unwrap(); for result in &mut object_iter { let (key, val_item) = result?; - let key_item = JsonbItem::String(key.as_bytes()); + let key_item = JsonbItem::String(Cow::Borrowed(key)); self.jsonb_item_to_comparable_impl(depth + 1, key_item, buf)?; self.jsonb_item_to_comparable_impl(depth + 1, val_item, buf)?; } @@ -688,10 +690,10 @@ impl RawJsonb<'_> { buf.push(FALSE_LEVEL); } } - JsonbItem::Number(data) => { + JsonbItem::Number(num_item) => { buf.push(depth); buf.push(NUMBER_LEVEL); - let num = Number::decode(data)?; + let num = num_item.as_number()?; let n = num.as_f64().unwrap(); // https://github.com/rust-lang/rust/blob/9c20b2a8cc7588decb6de25ac6a7912dcef24d65/library/core/src/num/f32.rs#L1176-L1260 let s = n.to_bits() as i64; @@ -701,16 +703,23 @@ impl RawJsonb<'_> { b[0] ^= 0x80; buf.extend_from_slice(&b); } - JsonbItem::String(data) => { + JsonbItem::String(s) => { buf.push(depth); buf.push(STRING_LEVEL); - buf.extend_from_slice(data); + buf.extend_from_slice(s.as_bytes()); buf.push(0); } - JsonbItem::Extension(data) => { + JsonbItem::Extension(ext_item) => { buf.push(depth); buf.push(EXTENSION_LEVEL); - buf.extend_from_slice(data); + match ext_item { + ExtensionItem::Raw(data) => { + buf.extend_from_slice(data); + } + ExtensionItem::Extension(value) => { + value.compact_encode(&mut *buf)?; + } + } buf.push(0); } JsonbItem::Raw(raw) => { diff --git a/src/functions/path.rs b/src/functions/path.rs index 9e3be01..254b749 100644 --- a/src/functions/path.rs +++ b/src/functions/path.rs @@ -29,8 +29,11 @@ use crate::error::*; use crate::jsonpath::JsonPath; use crate::jsonpath::Selector; use crate::keypath::KeyPath; +use crate::keypath::KeyPaths; +use crate::ExtensionValue; use crate::OwnedJsonb; use crate::RawJsonb; +use crate::Value; impl RawJsonb<'_> { /// Gets the element at the specified index in a JSONB array. @@ -619,6 +622,7 @@ impl RawJsonb<'_> { /// // Deleting from an object /// let obj_jsonb = r#"{"a": 1, "b": "hello", "c": 3}"#.parse::().unwrap(); /// let raw_jsonb = obj_jsonb.as_raw(); + /// /// let deleted = raw_jsonb.delete_by_name("b").unwrap(); /// assert_eq!(deleted.to_string(), r#"{"a":1,"c":3}"#); /// @@ -1041,4 +1045,119 @@ impl RawJsonb<'_> { } Ok(false) } + + /// Extracts all scalar values from a JSONB value along with their key paths. + /// + /// This function recursively traverses the JSONB structure (both objects and arrays) + /// and collects all leaf node scalar values (null, boolean, number, string, etc.) + /// along with their corresponding key paths. The key path represents the navigation + /// path from the root to reach each scalar value. + /// + /// # Returns + /// + /// * `Result, Value<'_>)>>` - A vector of tuples, each containing: + /// - `KeyPaths`: The path to reach the scalar value + /// - `Value`: The scalar value itself + /// + /// # Examples + /// + /// ```rust + /// use jsonb::OwnedJsonb; + /// + /// let json = r#"{"user": {"name": "Alice", "scores": [85, 92, 78]}}"#; + /// let jsonb = json.parse::().unwrap(); + /// let raw_jsonb = jsonb.as_raw(); + /// let result = raw_jsonb.extract_scalar_key_values(); + /// assert!(result.is_ok()); + /// let result = result.unwrap(); + /// assert_eq!(result.len(), 4); + /// // Result contains: + /// // - path: "user", "name" -> value: "Alice" + /// // - path: "user", "scores", 0 -> value: 85 + /// // - path: "user", "scores", 1 -> value: 92 + /// // - path: "user", "scores", 2 -> value: 78 + /// ``` + pub fn extract_scalar_key_values(&self) -> Result, Value<'_>)>> { + let item = JsonbItem::from_raw_jsonb(*self)?; + let mut result = Vec::with_capacity(16); + let mut current_paths = Vec::with_capacity(3); + Self::extract_scalar_key_values_recursive(item, &mut current_paths, &mut result)?; + Ok(result) + } + + /// Helper function for `extract_scalar_key_values` that recursively traverses the JSONB structure. + /// + /// This function implements a depth-first traversal of the JSONB document, building up the + /// key path as it goes and collecting scalar values when it reaches leaf nodes. + /// + /// # Arguments + /// + /// * `current_item` - The current JSONB item being processed + /// * `current_paths` - The current path from the root to this item (modified during traversal) + /// * `result` - The collection where extracted key-value pairs are stored + /// + /// # Returns + /// + /// * `Result<()>` - Success or error during traversal + fn extract_scalar_key_values_recursive<'a>( + current_item: JsonbItem<'a>, + current_paths: &mut Vec>, + result: &mut Vec<(KeyPaths<'a>, Value<'a>)>, + ) -> Result<()> { + match current_item { + JsonbItem::Raw(raw) => { + let object_iter_opt = ObjectIterator::new(raw)?; + if let Some(mut object_iter) = object_iter_opt { + for object_result in &mut object_iter { + let (key, val_item) = object_result?; + current_paths.push(KeyPath::Name(Cow::Borrowed(key))); + // Recursively handle object values + Self::extract_scalar_key_values_recursive(val_item, current_paths, result)?; + current_paths.pop(); + } + return Ok(()); + } + let array_iter_opt = ArrayIterator::new(raw)?; + if let Some(array_iter) = array_iter_opt { + for (index, array_result) in &mut array_iter.enumerate() { + let val_item = array_result?; + current_paths.push(KeyPath::Index(index as i32)); + // Recursively handle array values + Self::extract_scalar_key_values_recursive(val_item, current_paths, result)?; + current_paths.pop(); + } + } + } + JsonbItem::Owned(_) => unreachable!(), + _ => { + // ignore scalar value + if current_paths.is_empty() { + return Ok(()); + } + let key_paths = KeyPaths { + paths: current_paths.clone(), + }; + let value = match current_item { + JsonbItem::Null => Value::Null, + JsonbItem::Boolean(val) => Value::Bool(val), + JsonbItem::String(val) => Value::String(val), + JsonbItem::Number(num) => Value::Number(num.as_number()?), + JsonbItem::Extension(ext) => { + let ext_val = ext.as_extension_value()?; + match ext_val { + ExtensionValue::Binary(val) => Value::Binary(val), + ExtensionValue::Date(val) => Value::Date(val), + ExtensionValue::Timestamp(val) => Value::Timestamp(val), + ExtensionValue::TimestampTz(val) => Value::TimestampTz(val), + ExtensionValue::Interval(val) => Value::Interval(val), + } + } + _ => unreachable!(), + }; + // Add the path and scalar value + result.push((key_paths, value)); + } + } + Ok(()) + } } diff --git a/src/functions/scalar.rs b/src/functions/scalar.rs index 71afd15..8c90734 100644 --- a/src/functions/scalar.rs +++ b/src/functions/scalar.rs @@ -309,11 +309,10 @@ impl RawJsonb<'_> { JsonbItem::Boolean(v) => { return Ok(v); } - JsonbItem::String(data) => { - let s = String::from_utf8_lossy(data); - if &s.to_lowercase() == "true" { + JsonbItem::String(s) => { + if s.eq_ignore_ascii_case("true") || s.eq_ignore_ascii_case("yes") { return Ok(true); - } else if &s.to_lowercase() == "false" { + } else if s.eq_ignore_ascii_case("false") || s.eq_ignore_ascii_case("no") { return Ok(false); } } @@ -453,9 +452,9 @@ impl RawJsonb<'_> { pub fn as_number(&self) -> Result> { let jsonb_item = JsonbItem::from_raw_jsonb(*self)?; match jsonb_item { - JsonbItem::Number(data) => { - let num = Number::decode(data)?; - Ok(Some(num)) + JsonbItem::Number(num) => { + let value = num.as_number()?; + Ok(Some(value)) } _ => Ok(None), } @@ -560,9 +559,9 @@ impl RawJsonb<'_> { pub fn as_i64(&self) -> Result> { let jsonb_item = JsonbItem::from_raw_jsonb(*self)?; match jsonb_item { - JsonbItem::Number(data) => { - let num = Number::decode(data)?; - Ok(num.as_i64()) + JsonbItem::Number(num) => { + let value = num.as_number()?; + Ok(value.as_i64()) } _ => Ok(None), } @@ -646,14 +645,13 @@ impl RawJsonb<'_> { return Ok(0_i64); } } - JsonbItem::Number(data) => { - let num = Number::decode(data)?; - if let Some(v) = num.as_i64() { + JsonbItem::Number(num) => { + let value = num.as_number()?; + if let Some(v) = value.as_i64() { return Ok(v); } } - JsonbItem::String(data) => { - let s = String::from_utf8_lossy(data); + JsonbItem::String(s) => { if let Ok(v) = s.parse::() { return Ok(v); } @@ -794,9 +792,9 @@ impl RawJsonb<'_> { pub fn as_u64(&self) -> Result> { let jsonb_item = JsonbItem::from_raw_jsonb(*self)?; match jsonb_item { - JsonbItem::Number(data) => { - let num = Number::decode(data)?; - Ok(num.as_u64()) + JsonbItem::Number(num) => { + let value = num.as_number()?; + Ok(value.as_u64()) } _ => Ok(None), } @@ -884,14 +882,13 @@ impl RawJsonb<'_> { return Ok(0_u64); } } - JsonbItem::Number(data) => { - let num = Number::decode(data)?; - if let Some(v) = num.as_u64() { + JsonbItem::Number(num) => { + let value = num.as_number()?; + if let Some(v) = value.as_u64() { return Ok(v); } } - JsonbItem::String(data) => { - let s = String::from_utf8_lossy(data); + JsonbItem::String(s) => { if let Ok(v) = s.parse::() { return Ok(v); } @@ -1024,9 +1021,9 @@ impl RawJsonb<'_> { pub fn as_f64(&self) -> Result> { let jsonb_item = JsonbItem::from_raw_jsonb(*self)?; match jsonb_item { - JsonbItem::Number(data) => { - let num = Number::decode(data)?; - Ok(num.as_f64()) + JsonbItem::Number(num) => { + let value = num.as_number()?; + Ok(value.as_f64()) } _ => Ok(None), } @@ -1107,14 +1104,13 @@ impl RawJsonb<'_> { return Ok(0_f64); } } - JsonbItem::Number(data) => { - let num = Number::decode(data)?; - if let Some(v) = num.as_f64() { + JsonbItem::Number(num) => { + let value = num.as_number()?; + if let Some(v) = value.as_f64() { return Ok(v); } } - JsonbItem::String(data) => { - let s = String::from_utf8_lossy(data); + JsonbItem::String(s) => { if let Ok(v) = s.parse::() { return Ok(v); } @@ -1243,10 +1239,7 @@ impl RawJsonb<'_> { pub fn as_str(&self) -> Result>> { let jsonb_item = JsonbItem::from_raw_jsonb(*self)?; match jsonb_item { - JsonbItem::String(data) => { - let s = unsafe { std::str::from_utf8_unchecked(data) }; - Ok(Some(Cow::Borrowed(s))) - } + JsonbItem::String(s) => Ok(Some(s)), _ => Ok(None), } } @@ -1317,14 +1310,11 @@ impl RawJsonb<'_> { Ok("false".to_string()) } } - JsonbItem::Number(data) => { - let num = Number::decode(data)?; - Ok(format!("{}", num)) - } - JsonbItem::String(data) => { - let s = unsafe { String::from_utf8_unchecked(data.to_vec()) }; - Ok(s) + JsonbItem::Number(num) => { + let value = num.as_number()?; + Ok(format!("{}", value)) } + JsonbItem::String(s) => Ok(s.to_string()), _ => Err(Error::InvalidCast), } } @@ -1502,8 +1492,8 @@ impl RawJsonb<'_> { pub fn as_extension_value(&self) -> Result>> { let jsonb_item = JsonbItem::from_raw_jsonb(*self)?; match jsonb_item { - JsonbItem::Extension(data) => { - let val = ExtensionValue::decode(data)?; + JsonbItem::Extension(ext) => { + let val = ext.as_extension_value()?; Ok(Some(val)) } _ => Ok(None), @@ -1542,8 +1532,8 @@ impl RawJsonb<'_> { JsonbItemType::Extension => { let jsonb_item = JsonbItem::from_raw_jsonb(*self)?; match jsonb_item { - JsonbItem::Extension(data) => { - let val = ExtensionValue::decode(data)?; + JsonbItem::Extension(ext) => { + let val = ext.as_extension_value()?; match val { ExtensionValue::Binary(_v) => Ok(true), _ => Ok(false), @@ -1585,8 +1575,8 @@ impl RawJsonb<'_> { pub fn as_binary(&self) -> Result>> { let jsonb_item = JsonbItem::from_raw_jsonb(*self)?; match jsonb_item { - JsonbItem::Extension(data) => { - let val = ExtensionValue::decode(data)?; + JsonbItem::Extension(ext) => { + let val = ext.as_extension_value()?; match val { ExtensionValue::Binary(v) => Ok(Some(v.to_vec())), _ => Ok(None), @@ -1629,8 +1619,8 @@ impl RawJsonb<'_> { JsonbItemType::Extension => { let jsonb_item = JsonbItem::from_raw_jsonb(*self)?; match jsonb_item { - JsonbItem::Extension(data) => { - let val = ExtensionValue::decode(data)?; + JsonbItem::Extension(ext) => { + let val = ext.as_extension_value()?; match val { ExtensionValue::Date(_v) => Ok(true), _ => Ok(false), @@ -1673,8 +1663,8 @@ impl RawJsonb<'_> { pub fn as_date(&self) -> Result> { let jsonb_item = JsonbItem::from_raw_jsonb(*self)?; match jsonb_item { - JsonbItem::Extension(data) => { - let val = ExtensionValue::decode(data)?; + JsonbItem::Extension(ext) => { + let val = ext.as_extension_value()?; match val { ExtensionValue::Date(v) => Ok(Some(v)), _ => Ok(None), @@ -1717,8 +1707,8 @@ impl RawJsonb<'_> { JsonbItemType::Extension => { let jsonb_item = JsonbItem::from_raw_jsonb(*self)?; match jsonb_item { - JsonbItem::Extension(data) => { - let val = ExtensionValue::decode(data)?; + JsonbItem::Extension(ext) => { + let val = ext.as_extension_value()?; match val { ExtensionValue::Timestamp(_v) => Ok(true), _ => Ok(false), @@ -1761,8 +1751,8 @@ impl RawJsonb<'_> { pub fn as_timestamp(&self) -> Result> { let jsonb_item = JsonbItem::from_raw_jsonb(*self)?; match jsonb_item { - JsonbItem::Extension(data) => { - let val = ExtensionValue::decode(data)?; + JsonbItem::Extension(ext) => { + let val = ext.as_extension_value()?; match val { ExtensionValue::Timestamp(v) => Ok(Some(v)), _ => Ok(None), @@ -1805,8 +1795,8 @@ impl RawJsonb<'_> { JsonbItemType::Extension => { let jsonb_item = JsonbItem::from_raw_jsonb(*self)?; match jsonb_item { - JsonbItem::Extension(data) => { - let val = ExtensionValue::decode(data)?; + JsonbItem::Extension(ext) => { + let val = ext.as_extension_value()?; match val { ExtensionValue::TimestampTz(_v) => Ok(true), _ => Ok(false), @@ -1849,8 +1839,8 @@ impl RawJsonb<'_> { pub fn as_timestamp_tz(&self) -> Result> { let jsonb_item = JsonbItem::from_raw_jsonb(*self)?; match jsonb_item { - JsonbItem::Extension(data) => { - let val = ExtensionValue::decode(data)?; + JsonbItem::Extension(ext) => { + let val = ext.as_extension_value()?; match val { ExtensionValue::TimestampTz(v) => Ok(Some(v)), _ => Ok(None), @@ -1893,8 +1883,8 @@ impl RawJsonb<'_> { JsonbItemType::Extension => { let jsonb_item = JsonbItem::from_raw_jsonb(*self)?; match jsonb_item { - JsonbItem::Extension(data) => { - let val = ExtensionValue::decode(data)?; + JsonbItem::Extension(ext) => { + let val = ext.as_extension_value()?; match val { ExtensionValue::Interval(_v) => Ok(true), _ => Ok(false), @@ -1937,8 +1927,8 @@ impl RawJsonb<'_> { pub fn as_interval(&self) -> Result> { let jsonb_item = JsonbItem::from_raw_jsonb(*self)?; match jsonb_item { - JsonbItem::Extension(data) => { - let val = ExtensionValue::decode(data)?; + JsonbItem::Extension(ext) => { + let val = ext.as_extension_value()?; match val { ExtensionValue::Interval(v) => Ok(Some(v)), _ => Ok(None), diff --git a/src/jsonpath/selector.rs b/src/jsonpath/selector.rs index e6a8eba..dd8fd37 100644 --- a/src/jsonpath/selector.rs +++ b/src/jsonpath/selector.rs @@ -490,17 +490,34 @@ impl<'a> Selector<'a> { } fn select_object_values(&mut self, parent_item: JsonbItem<'a>) -> Result<()> { - let Some(curr_raw_jsonb) = parent_item.as_raw_jsonb() else { + let jsonb_item_type = parent_item.jsonb_item_type()?; + if !matches!(jsonb_item_type, JsonbItemType::Object(_)) { return Ok(()); }; - let object_val_iter_opt = ObjectValueIterator::new(curr_raw_jsonb)?; - if let Some(mut object_val_iter) = object_val_iter_opt { - for result in &mut object_val_iter { - let val_item = result?; - self.items.push_back(val_item); + match parent_item { + JsonbItem::Raw(raw) => { + let object_val_iter_opt = ObjectValueIterator::new(raw)?; + if let Some(mut object_val_iter) = object_val_iter_opt { + for result in &mut object_val_iter { + let val_item = result?; + self.items.push_back(val_item); + } + } + } + JsonbItem::Owned(ref owned) => { + let object_val_iter_opt = ObjectValueIterator::new(owned.as_raw())?; + if let Some(mut object_val_iter) = object_val_iter_opt { + for result in &mut object_val_iter { + let val_item = result?; + let owned_item = OwnedJsonb::from_item(val_item)?; + self.items.push_back(JsonbItem::Owned(owned_item)); + } + } } + _ => {} } + Ok(()) } @@ -518,25 +535,58 @@ impl<'a> Selector<'a> { if is_match { self.items.push_back(parent_item.clone()); } - let Some(curr_raw_jsonb) = parent_item.as_raw_jsonb() else { - return Ok(()); - }; if !should_continue { return Ok(()); } - let object_val_iter_opt = ObjectValueIterator::new(curr_raw_jsonb)?; - if let Some(mut object_val_iter) = object_val_iter_opt { - for result in &mut object_val_iter { - let val_item = result?; - self.recursive_select_values(val_item, curr_level + 1, recursive_level_opt)?; + + match parent_item { + JsonbItem::Raw(raw) => { + let object_val_iter_opt = ObjectValueIterator::new(raw)?; + if let Some(mut object_val_iter) = object_val_iter_opt { + for result in &mut object_val_iter { + let val_item = result?; + self.recursive_select_values( + val_item, + curr_level + 1, + recursive_level_opt, + )?; + } + } + let array_iter_opt = ArrayIterator::new(raw)?; + if let Some(mut array_iter) = array_iter_opt { + for item_result in &mut array_iter { + let item = item_result?; + self.recursive_select_values(item, curr_level + 1, recursive_level_opt)?; + } + } } - } - let array_iter_opt = ArrayIterator::new(curr_raw_jsonb)?; - if let Some(mut array_iter) = array_iter_opt { - for item_result in &mut array_iter { - let item = item_result?; - self.recursive_select_values(item, curr_level + 1, recursive_level_opt)?; + JsonbItem::Owned(ref owned) => { + let object_val_iter_opt = ObjectValueIterator::new(owned.as_raw())?; + if let Some(mut object_val_iter) = object_val_iter_opt { + for result in &mut object_val_iter { + let val_item = result?; + let owned_item = OwnedJsonb::from_item(val_item)?; + self.recursive_select_values( + JsonbItem::Owned(owned_item), + curr_level + 1, + recursive_level_opt, + )?; + } + } + let array_iter_opt = ArrayIterator::new(owned.as_raw())?; + if let Some(mut array_iter) = array_iter_opt { + for item_result in &mut array_iter { + let item = item_result?; + let owned_item = OwnedJsonb::from_item(item)?; + self.recursive_select_values( + JsonbItem::Owned(owned_item), + curr_level + 1, + recursive_level_opt, + )?; + } + } } + _ => {} } Ok(()) } @@ -546,37 +596,63 @@ impl<'a> Selector<'a> { parent_item: JsonbItem<'a>, name: &'a str, ) -> Result<()> { - let Some(curr_raw_jsonb) = parent_item.as_raw_jsonb() else { + let jsonb_item_type = parent_item.jsonb_item_type()?; + if !matches!(jsonb_item_type, JsonbItemType::Object(_)) { return Ok(()); }; let key_name = Cow::Borrowed(name); - if let Some(val_item) = - curr_raw_jsonb.get_object_value_by_key_name(&key_name, |name, key| key.eq(name))? - { - self.items.push_back(val_item); + match parent_item { + JsonbItem::Raw(raw) => { + if let Some(val_item) = + raw.get_object_value_by_key_name(&key_name, |name, key| key.eq(name))? + { + self.items.push_back(val_item); + } + } + JsonbItem::Owned(ref owned) => { + let raw = owned.as_raw(); + if let Some(val_item) = + raw.get_object_value_by_key_name(&key_name, |name, key| key.eq(name))? + { + let owned_item = OwnedJsonb::from_item(val_item)?; + self.items.push_back(JsonbItem::Owned(owned_item)); + } + } + _ => {} } Ok(()) } fn select_array_values(&mut self, parent_item: JsonbItem<'a>) -> Result<()> { - let Some(curr_raw_jsonb) = parent_item.as_raw_jsonb() else { + let jsonb_item_type = parent_item.jsonb_item_type()?; + if !matches!(jsonb_item_type, JsonbItemType::Array(_)) { // In lax mode, bracket wildcard allow Scalar and Object value. self.items.push_back(parent_item); return Ok(()); }; - let array_iter_opt = ArrayIterator::new(curr_raw_jsonb)?; - if let Some(mut array_iter) = array_iter_opt { - for item_result in &mut array_iter { - let item = item_result?; - self.items.push_back(item); + match parent_item { + JsonbItem::Raw(raw) => { + let array_iter_opt = ArrayIterator::new(raw)?; + if let Some(mut array_iter) = array_iter_opt { + for item_result in &mut array_iter { + let item = item_result?; + self.items.push_back(item); + } + } } - } else { - // In lax mode, bracket wildcard allow Scalar and Object value. - // convert to Jsonb item to compare with other path values. - let item = JsonbItem::from_raw_jsonb(curr_raw_jsonb)?; - self.items.push_back(item); + JsonbItem::Owned(ref owned) => { + let array_iter_opt = ArrayIterator::new(owned.as_raw())?; + if let Some(mut array_iter) = array_iter_opt { + for item_result in &mut array_iter { + let item = item_result?; + let owned_item = OwnedJsonb::from_item(item)?; + self.items.push_back(JsonbItem::Owned(owned_item)); + } + } + } + _ => {} } Ok(()) } @@ -586,11 +662,7 @@ impl<'a> Selector<'a> { parent_item: JsonbItem<'a>, array_indices: &Vec, ) -> Result<()> { - let Some(curr_raw_jsonb) = parent_item.as_raw_jsonb() else { - return Ok(()); - }; - - let jsonb_item_type = curr_raw_jsonb.jsonb_item_type()?; + let jsonb_item_type = parent_item.jsonb_item_type()?; let JsonbItemType::Array(arr_len) = jsonb_item_type else { return Ok(()); }; @@ -599,14 +671,31 @@ impl<'a> Selector<'a> { if indices.is_empty() { continue; } - let array_iter_opt = ArrayIterator::new(curr_raw_jsonb)?; - if let Some(array_iter) = array_iter_opt { - for (i, item_result) in &mut array_iter.enumerate() { - let item = item_result?; - if indices.contains(&i) { - self.items.push_back(item); + match parent_item { + JsonbItem::Raw(raw) => { + let array_iter_opt = ArrayIterator::new(raw)?; + if let Some(array_iter) = array_iter_opt { + for (i, item_result) in &mut array_iter.enumerate() { + let item = item_result?; + if indices.contains(&i) { + self.items.push_back(item); + } + } } } + JsonbItem::Owned(ref owned) => { + let array_iter_opt = ArrayIterator::new(owned.as_raw())?; + if let Some(array_iter) = array_iter_opt { + for (i, item_result) in &mut array_iter.enumerate() { + let item = item_result?; + if indices.contains(&i) { + let owned_item = OwnedJsonb::from_item(item)?; + self.items.push_back(JsonbItem::Owned(owned_item)); + } + } + } + } + _ => {} } } Ok(()) @@ -836,13 +925,11 @@ impl<'a> Selector<'a> { let value = match item { JsonbItem::Null => PathValue::Null, JsonbItem::Boolean(v) => PathValue::Boolean(v), - JsonbItem::Number(data) => { - let n = Number::decode(data)?; + JsonbItem::Number(num) => { + let n = num.as_number()?; PathValue::Number(n) } - JsonbItem::String(data) => PathValue::String(Cow::Borrowed(unsafe { - std::str::from_utf8_unchecked(data) - })), + JsonbItem::String(s) => PathValue::String(s), JsonbItem::Raw(raw) => { // collect values in the array. let array_iter_opt = ArrayIterator::new(raw)?; @@ -852,15 +939,11 @@ impl<'a> Selector<'a> { let value = match item { JsonbItem::Null => PathValue::Null, JsonbItem::Boolean(v) => PathValue::Boolean(v), - JsonbItem::Number(data) => { - let n = Number::decode(data)?; + JsonbItem::Number(num) => { + let n = num.as_number()?; PathValue::Number(n) } - JsonbItem::String(data) => { - PathValue::String(Cow::Borrowed(unsafe { - std::str::from_utf8_unchecked(data) - })) - } + JsonbItem::String(s) => PathValue::String(s), JsonbItem::Raw(raw) => PathValue::Raw(raw), _ => { continue; @@ -869,7 +952,21 @@ impl<'a> Selector<'a> { values.push(value); } } else { - values.push(PathValue::Raw(raw)); + let jsonb_item = JsonbItem::from_raw_jsonb(raw)?; + let value = match jsonb_item { + JsonbItem::Null => PathValue::Null, + JsonbItem::Boolean(v) => PathValue::Boolean(v), + JsonbItem::Number(num) => { + let n = num.as_number()?; + PathValue::Number(n) + } + JsonbItem::String(s) => PathValue::String(s), + JsonbItem::Raw(raw) => PathValue::Raw(raw), + _ => { + continue; + } + }; + values.push(value); } continue; } @@ -949,6 +1046,8 @@ impl<'a> Selector<'a> { } }; Some(res) + } else if matches!(op, BinaryOperator::NotEq) { + Some(true) } else { None } diff --git a/src/keypath.rs b/src/keypath.rs index dfb19ea..6f056a1 100644 --- a/src/keypath.rs +++ b/src/keypath.rs @@ -34,13 +34,13 @@ use crate::Error; /// Represents a set of key path chains. /// Compatible with PostgreSQL extracts JSON sub-object paths syntax. -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone, Eq, PartialEq, Hash)] pub struct KeyPaths<'a> { pub paths: Vec>, } /// Represents a valid key path. -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone, Eq, PartialEq, Hash)] pub enum KeyPath<'a> { /// represents the index of an Array, allow negative indexing. Index(i32), diff --git a/src/number.rs b/src/number.rs index e62e833..470fefe 100644 --- a/src/number.rs +++ b/src/number.rs @@ -288,52 +288,98 @@ impl Serialize for Number { use std::io::Write; const NUMBER_TOKEN: &str = "$serde_json::private::Number"; - struct WriteAdapter<'a>(&'a mut std::io::Cursor<&'a mut [u8]>); + if serializer.is_human_readable() { + struct WriteAdapter<'a>(&'a mut std::io::Cursor<&'a mut [u8]>); - impl std::fmt::Write for WriteAdapter<'_> { - fn write_str(&mut self, s: &str) -> std::fmt::Result { - self.0.write_all(s.as_bytes()).map_err(|_| std::fmt::Error) + impl std::fmt::Write for WriteAdapter<'_> { + fn write_str(&mut self, s: &str) -> std::fmt::Result { + self.0.write_all(s.as_bytes()).map_err(|_| std::fmt::Error) + } } - } - impl WriteAdapter<'_> { - fn position(&self) -> usize { - self.0.position() as usize + impl WriteAdapter<'_> { + fn position(&self) -> usize { + self.0.position() as usize + } } - } - - let mut buffer = [0u8; 128]; - let mut cursor = std::io::Cursor::new(&mut buffer[..]); - let mut adapter = WriteAdapter(&mut cursor); - match self { - Number::Decimal64(v) => { - format_decimal_i128(&mut adapter, v.value as i128, v.scale as usize) - .map_err(|e| { - serde::ser::Error::custom(format!("Format decimal64 error: {e}")) - })?; - } - Number::Decimal128(v) => { - format_decimal_i128(&mut adapter, v.value, v.scale as usize).map_err( - |e| serde::ser::Error::custom(format!("Format decimal128 error: {e}")), - )?; - } - Number::Decimal256(v) => { - format_decimal_i256(&mut adapter, v.value, v.scale as usize).map_err( - |e| serde::ser::Error::custom(format!("Format decimal256 error: {e}")), - )?; + let mut buffer = [0u8; 128]; + let mut cursor = std::io::Cursor::new(&mut buffer[..]); + let mut adapter = WriteAdapter(&mut cursor); + + match self { + Number::Decimal64(v) => { + format_decimal_i128(&mut adapter, v.value as i128, v.scale as usize) + .map_err(|e| { + serde::ser::Error::custom(format!( + "Format decimal64 error: {e}" + )) + })?; + } + Number::Decimal128(v) => { + format_decimal_i128(&mut adapter, v.value, v.scale as usize).map_err( + |e| { + serde::ser::Error::custom(format!( + "Format decimal128 error: {e}" + )) + }, + )?; + } + Number::Decimal256(v) => { + format_decimal_i256(&mut adapter, v.value, v.scale as usize).map_err( + |e| { + serde::ser::Error::custom(format!( + "Format decimal256 error: {e}" + )) + }, + )?; + } + _ => unreachable!(), } - _ => unreachable!(), - } - let pos = adapter.position(); - let num_str = std::str::from_utf8(&buffer[..pos]).map_err(|e| { - serde::ser::Error::custom(format!("Invalid decimal number: {e}")) - })?; + let pos = adapter.position(); + let num_str = std::str::from_utf8(&buffer[..pos]).map_err(|e| { + serde::ser::Error::custom(format!("Invalid decimal number: {e}")) + })?; - let mut serialize_struct = serializer.serialize_struct(NUMBER_TOKEN, 1)?; - serialize_struct.serialize_field(NUMBER_TOKEN, num_str)?; - serialize_struct.end() + let mut serialize_struct = serializer.serialize_struct(NUMBER_TOKEN, 1)?; + serialize_struct.serialize_field(NUMBER_TOKEN, num_str)?; + serialize_struct.end() + } else { + use crate::constants::NUMBER_STRUCT_FIELD_HIGH_VALUE; + use crate::constants::NUMBER_STRUCT_FIELD_LOW_VALUE; + use crate::constants::NUMBER_STRUCT_FIELD_SCALE; + use crate::constants::NUMBER_STRUCT_FIELD_VALUE; + use crate::constants::NUMBER_STRUCT_TOKEN; + + let mut serialize_struct = + serializer.serialize_struct(NUMBER_STRUCT_TOKEN, 2)?; + match self { + Number::Decimal64(v) => { + serialize_struct + .serialize_field(NUMBER_STRUCT_FIELD_SCALE, &v.scale)?; + serialize_struct + .serialize_field(NUMBER_STRUCT_FIELD_VALUE, &v.value)?; + } + Number::Decimal128(v) => { + serialize_struct + .serialize_field(NUMBER_STRUCT_FIELD_SCALE, &v.scale)?; + serialize_struct + .serialize_field(NUMBER_STRUCT_FIELD_VALUE, &v.value)?; + } + Number::Decimal256(v) => { + serialize_struct + .serialize_field(NUMBER_STRUCT_FIELD_SCALE, &v.scale)?; + let (high_value, low_value) = v.value.into_words(); + serialize_struct + .serialize_field(NUMBER_STRUCT_FIELD_HIGH_VALUE, &high_value)?; + serialize_struct + .serialize_field(NUMBER_STRUCT_FIELD_LOW_VALUE, &low_value)?; + } + _ => unreachable!(), + } + serialize_struct.end() + } } #[cfg(not(feature = "arbitrary_precision"))] Number::Decimal64(_) | Number::Decimal128(_) | Number::Decimal256(_) => { @@ -352,6 +398,41 @@ impl Serialize for Number { } impl Number { + /// Returns the i128 representation of the number, if possible. + /// + /// This method returns None if the number cannot be represented as an i64. + pub fn as_i128(&self) -> Option { + match self { + Number::Int64(v) => Some(*v as i128), + Number::UInt64(v) => Some(*v as i128), + Number::Float64(_) => None, + Number::Decimal64(v) => { + if v.scale == 0 { + Some(v.value as i128) + } else { + None + } + } + Number::Decimal128(v) => { + if v.scale == 0 { + Some(v.value) + } else { + None + } + } + Number::Decimal256(v) => { + if v.scale == 0 + && v.value >= i256::from(i128::MIN) + && v.value <= i256::from(i128::MAX) + { + Some(v.value.as_i128()) + } else { + None + } + } + } + } + /// Returns the i64 representation of the number, if possible. /// /// This method returns None if the number cannot be represented as an i64. diff --git a/src/owned.rs b/src/owned.rs index dd3f718..2b3bddb 100644 --- a/src/owned.rs +++ b/src/owned.rs @@ -334,7 +334,7 @@ impl Ord for OwnedJsonb { /// }; /// /// let owned_jsonb: OwnedJsonb = to_owned_jsonb(&person).unwrap(); -/// +/// assert_eq!(format!("{}", owned_jsonb), "{\"age\":42,\"name\":\"Bob\"}"); /// println!("JSONB data: {}", owned_jsonb); /// ``` pub fn to_owned_jsonb(value: &T) -> Result diff --git a/src/parser.rs b/src/parser.rs index b3c76fc..eb9afc7 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -35,23 +35,24 @@ use crate::OwnedJsonb; #[cfg(feature = "arbitrary_precision")] use ethnum::i256; +use crate::constants::MAX_DECIMAL128_PRECISION; +use crate::constants::MAX_DECIMAL256_PRECISION; #[cfg(feature = "arbitrary_precision")] -const MAX_DECIMAL64_PRECISION: usize = 18; -const MAX_DECIMAL128_PRECISION: usize = 38; -const MAX_DECIMAL256_PRECISION: usize = 76; - -const UINT64_MIN: i128 = 0i128; -const UINT64_MAX: i128 = 18_446_744_073_709_551_615i128; -const INT64_MIN: i128 = -9_223_372_036_854_775_808i128; -const INT64_MAX: i128 = 9_223_372_036_854_775_807i128; +use crate::constants::MAX_DECIMAL64_PRECISION; + #[cfg(feature = "arbitrary_precision")] -const DECIMAL64_MIN: i128 = -999999999999999999i128; +use crate::constants::DECIMAL128_MAX; #[cfg(feature = "arbitrary_precision")] -const DECIMAL64_MAX: i128 = 999999999999999999i128; +use crate::constants::DECIMAL128_MIN; #[cfg(feature = "arbitrary_precision")] -const DECIMAL128_MIN: i128 = -99999999999999999999999999999999999999i128; +use crate::constants::DECIMAL64_MAX; #[cfg(feature = "arbitrary_precision")] -const DECIMAL128_MAX: i128 = 99999999999999999999999999999999999999i128; +use crate::constants::DECIMAL64_MIN; + +use crate::constants::INT64_MAX; +use crate::constants::INT64_MIN; +use crate::constants::UINT64_MAX; +use crate::constants::UINT64_MIN; // JSON literal constants const NULL_LOWERCASE: [u8; 4] = [b'n', b'u', b'l', b'l']; diff --git a/src/raw.rs b/src/raw.rs index 3800015..06df2cb 100644 --- a/src/raw.rs +++ b/src/raw.rs @@ -346,7 +346,7 @@ impl Ord for RawJsonb<'_> { /// use jsonb::RawJsonb; /// use serde::Deserialize; /// -/// #[derive(Deserialize, Debug)] +/// #[derive(Deserialize, Debug, PartialEq, Eq)] /// struct Person { /// name: String, /// age: u32, @@ -356,6 +356,7 @@ impl Ord for RawJsonb<'_> { /// let raw_jsonb = owned_jsonb.as_raw(); /// /// let person: Person = from_raw_jsonb(&raw_jsonb).unwrap(); +/// assert_eq!(person, Person { name: "Alice".to_string(), age: 20 }); /// println!("{:?}", person); // Output: Person { name: "Alice", age: 20 } /// ``` pub fn from_raw_jsonb<'de, T>(raw_jsonb: &'de RawJsonb) -> Result diff --git a/src/value.rs b/src/value.rs index 0825fc8..a39fa72 100644 --- a/src/value.rs +++ b/src/value.rs @@ -12,13 +12,16 @@ // See the License for the specific language governing permissions and // limitations under the License. +use crate::core::JsonbItemType; use std::borrow::Cow; +use std::cmp::Ordering; use std::collections::BTreeMap; use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; use std::mem::discriminant; +use crate::ExtensionValue; use rand::distr::Alphanumeric; use rand::distr::SampleString; use rand::rng; @@ -43,7 +46,7 @@ pub type Object<'a> = BTreeMap>; /// The extended types provide additional functionality beyond the JSON specification, /// making this implementation more suitable for database applications and other /// systems requiring richer data type support. -#[derive(Clone, PartialEq, Default, Eq)] +#[derive(Clone, Default)] pub enum Value<'a> { /// Represents a JSON null value #[default] @@ -75,6 +78,70 @@ pub enum Value<'a> { Object(Object<'a>), } +impl Eq for Value<'_> {} + +impl PartialEq for Value<'_> { + fn eq(&self, other: &Self) -> bool { + let result = self.cmp(other); + result == Ordering::Equal + } +} + +impl PartialOrd for Value<'_> { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for Value<'_> { + fn cmp(&self, other: &Self) -> Ordering { + let self_type = self.jsonb_item_type(); + let other_type = other.jsonb_item_type(); + + if let Some(ord) = self_type.partial_cmp(&other_type) { + return ord; + } + + match (self, other) { + (Value::Null, Value::Null) => Ordering::Equal, + (Value::Bool(v1), Value::Bool(v2)) => v1.cmp(v2), + (Value::Number(v1), Value::Number(v2)) => v1.cmp(v2), + (Value::String(v1), Value::String(v2)) => v1.cmp(v2), + (Value::Array(arr1), Value::Array(arr2)) => { + for (v1, v2) in arr1.iter().zip(arr2.iter()) { + let ord = v1.cmp(v2); + if ord != Ordering::Equal { + return ord; + } + } + arr1.len().cmp(&arr2.len()) + } + (Value::Object(obj1), Value::Object(obj2)) => { + for ((k1, v1), (k2, v2)) in obj1.iter().zip(obj2.iter()) { + let ord = k1.cmp(k2); + if ord != Ordering::Equal { + return ord; + } + let ord = v1.cmp(v2); + if ord != Ordering::Equal { + return ord; + } + } + obj1.len().cmp(&obj2.len()) + } + (_, _) => match (self.as_extension_value(), other.as_extension_value()) { + (Some(self_ext), Some(other_ext)) => { + if let Some(ord) = self_ext.partial_cmp(&other_ext) { + return ord; + } + Ordering::Equal + } + (_, _) => Ordering::Equal, + }, + } + } +} + impl Debug for Value<'_> { fn fmt(&self, formatter: &mut Formatter) -> std::fmt::Result { match *self { @@ -457,4 +524,31 @@ impl<'a> Value<'a> { }; val } + + fn jsonb_item_type(&self) -> JsonbItemType { + match self { + Value::Null => JsonbItemType::Null, + Value::Bool(_) => JsonbItemType::Boolean, + Value::Number(_) => JsonbItemType::Number, + Value::String(_) => JsonbItemType::String, + Value::Binary(_) => JsonbItemType::Extension, + Value::Date(_) => JsonbItemType::Extension, + Value::Timestamp(_) => JsonbItemType::Extension, + Value::TimestampTz(_) => JsonbItemType::Extension, + Value::Interval(_) => JsonbItemType::Extension, + Value::Array(arr) => JsonbItemType::Array(arr.len()), + Value::Object(obj) => JsonbItemType::Object(obj.len()), + } + } + + fn as_extension_value(&self) -> Option> { + match self { + Value::Binary(v) => Some(ExtensionValue::Binary(v)), + Value::Date(v) => Some(ExtensionValue::Date(v.clone())), + Value::Timestamp(v) => Some(ExtensionValue::Timestamp(v.clone())), + Value::TimestampTz(v) => Some(ExtensionValue::TimestampTz(v.clone())), + Value::Interval(v) => Some(ExtensionValue::Interval(v.clone())), + _ => None, + } + } } diff --git a/tests/it/functions.rs b/tests/it/functions.rs index d2d4b7c..d3117f9 100644 --- a/tests/it/functions.rs +++ b/tests/it/functions.rs @@ -18,10 +18,13 @@ use std::collections::BTreeMap; use std::collections::BTreeSet; use ethnum::I256; +use jsonb::core::JsonbItemType; use jsonb::from_raw_jsonb; use jsonb::from_slice; use jsonb::jsonpath::parse_json_path; use jsonb::keypath::parse_key_paths; +use jsonb::keypath::KeyPath; +use jsonb::keypath::KeyPaths; use jsonb::parse_value; use jsonb::Date; use jsonb::Decimal128; @@ -227,7 +230,7 @@ fn test_path_exists_expr() { #[test] fn test_select_by_path() { - let source = r#"{"name":"Fred","phones":[{"type":"home","number":3720453},{"type":"work","number":5062051}],"car_no":123,"测试\"\uD83D\uDC8E":"ab","numbers":[2,3,4]}"#; + let source = r#"{"name":"Fred","phones":[{"type":"home","number":3720453},{"type":"work","number":5062051}],"car_no":123,"测试\"\uD83D\uDC8E":"ab","numbers":[2,3,4],"key":null}"#; let paths = vec![ (r#"$.name"#, vec![r#""Fred""#]), @@ -299,6 +302,10 @@ fn test_select_by_path() { r#"$.phones[0 to last]?(@.number == 3720453 && @.type == "work")"#, vec![], ), + ( + r#"$.car_no?($.name == "Fred" && $.car_no != null)"#, + vec![r#"123"#], + ), (r#"$.car_no"#, vec![r#"123"#]), (r#"$.测试\"\uD83D\uDC8E"#, vec![r#""ab""#]), // predicates return the result of the filter expression. @@ -1887,6 +1894,299 @@ fn test_to_serde_json() { } } +#[test] +fn test_extract_scalar_key_values() { + // Test case 1: Simple object with scalar values + let json = r#"{"name": "John", "age": 30, "active": true}"#; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + + let result = raw_jsonb.extract_scalar_key_values().unwrap(); + assert_eq!(result.len(), 3); + + let expected = vec![ + ( + KeyPaths { + paths: vec![KeyPath::Name(Cow::Borrowed("active"))], + }, + Value::Bool(true), + ), + ( + KeyPaths { + paths: vec![KeyPath::Name(Cow::Borrowed("age"))], + }, + Value::Number(Number::UInt64(30)), + ), + ( + KeyPaths { + paths: vec![KeyPath::Name(Cow::Borrowed("name"))], + }, + Value::String(Cow::Borrowed("John")), + ), + ]; + for ((key_paths, value), (expected_key_paths, expected_value)) in + result.into_iter().zip(expected.into_iter()) + { + assert_eq!(key_paths, expected_key_paths); + assert_eq!(value, expected_value); + } + + // Test case 2: Nested object with array + let json = r#"{"user": {"name": "Alice", "scores": [85, 92, 78]}}"#; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + + let result = raw_jsonb.extract_scalar_key_values().unwrap(); + assert_eq!(result.len(), 4); + + let expected = vec![ + ( + KeyPaths { + paths: vec![ + KeyPath::Name(Cow::Borrowed("user")), + KeyPath::Name(Cow::Borrowed("name")), + ], + }, + Value::String(Cow::Borrowed("Alice")), + ), + ( + KeyPaths { + paths: vec![ + KeyPath::Name(Cow::Borrowed("user")), + KeyPath::Name(Cow::Borrowed("scores")), + KeyPath::Index(0), + ], + }, + Value::Number(Number::UInt64(85)), + ), + ( + KeyPaths { + paths: vec![ + KeyPath::Name(Cow::Borrowed("user")), + KeyPath::Name(Cow::Borrowed("scores")), + KeyPath::Index(1), + ], + }, + Value::Number(Number::UInt64(92)), + ), + ( + KeyPaths { + paths: vec![ + KeyPath::Name(Cow::Borrowed("user")), + KeyPath::Name(Cow::Borrowed("scores")), + KeyPath::Index(2), + ], + }, + Value::Number(Number::UInt64(78)), + ), + ]; + for ((key_paths, value), (expected_key_paths, expected_value)) in + result.into_iter().zip(expected.into_iter()) + { + assert_eq!(key_paths, expected_key_paths); + assert_eq!(value, expected_value); + } + + // Test case 3: Complex nested structure + let json = r#"{"k1": [{"k2": "v2"}, {"k3": "v3"}]}"#; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + + let result = raw_jsonb.extract_scalar_key_values().unwrap(); + assert_eq!(result.len(), 2); + + let expected = vec![ + ( + KeyPaths { + paths: vec![ + KeyPath::Name(Cow::Borrowed("k1")), + KeyPath::Index(0), + KeyPath::Name(Cow::Borrowed("k2")), + ], + }, + Value::String(Cow::Borrowed("v2")), + ), + ( + KeyPaths { + paths: vec![ + KeyPath::Name(Cow::Borrowed("k1")), + KeyPath::Index(1), + KeyPath::Name(Cow::Borrowed("k3")), + ], + }, + Value::String(Cow::Borrowed("v3")), + ), + ]; + for ((key_paths, value), (expected_key_paths, expected_value)) in + result.into_iter().zip(expected.into_iter()) + { + assert_eq!(key_paths, expected_key_paths); + assert_eq!(value, expected_value); + } +} + +#[test] +fn test_jsonb_item_type() { + // Test null value + let json = "null"; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + let item_type = raw_jsonb.jsonb_item_type().unwrap(); + assert!(matches!(item_type, JsonbItemType::Null)); + + // Test boolean values + let json = "true"; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + let item_type = raw_jsonb.jsonb_item_type().unwrap(); + assert!(matches!(item_type, JsonbItemType::Boolean)); + + let json = "false"; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + let item_type = raw_jsonb.jsonb_item_type().unwrap(); + assert!(matches!(item_type, JsonbItemType::Boolean)); + + // Test number value + let json = "123.45"; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + let item_type = raw_jsonb.jsonb_item_type().unwrap(); + assert!(matches!(item_type, JsonbItemType::Number)); + + // Test string value + let json = r#""hello world""#; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + let item_type = raw_jsonb.jsonb_item_type().unwrap(); + assert!(matches!(item_type, JsonbItemType::String)); + + // Test array value + let json = "[1, 2, 3, 4, 5]"; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + let item_type = raw_jsonb.jsonb_item_type().unwrap(); + assert!(matches!(item_type, JsonbItemType::Array(5))); + + // Test empty array + let json = "[]"; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + let item_type = raw_jsonb.jsonb_item_type().unwrap(); + assert!(matches!(item_type, JsonbItemType::Array(0))); + + // Test object value + let json = r#"{"name": "Alice", "age": 30, "active": true}"#; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + let item_type = raw_jsonb.jsonb_item_type().unwrap(); + assert!(matches!(item_type, JsonbItemType::Object(3))); + + // Test empty object + let json = "{}"; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + let item_type = raw_jsonb.jsonb_item_type().unwrap(); + assert!(matches!(item_type, JsonbItemType::Object(0))); +} + +#[test] +fn test_to_value() { + // Test null value + let json = "null"; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + let value = raw_jsonb.to_value().unwrap(); + assert!(value.is_null()); + + // Test boolean values + let json = "true"; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + let value = raw_jsonb.to_value().unwrap(); + assert!(value.as_bool().unwrap()); + + let json = "false"; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + let value = raw_jsonb.to_value().unwrap(); + assert!(!value.as_bool().unwrap()); + + // Test number values + let json = "123"; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + let value = raw_jsonb.to_value().unwrap(); + assert_eq!(value.as_i64().unwrap(), 123); + + let json = "123.45"; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + let value = raw_jsonb.to_value().unwrap(); + assert_eq!(value.as_f64().unwrap(), 123.45); + + // Test string value + let json = r#""hello world""#; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + let value = raw_jsonb.to_value().unwrap(); + assert_eq!(value.as_str().unwrap(), "hello world"); + + // Test array value + let json = "[1, 2, 3, 4, 5]"; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + let value = raw_jsonb.to_value().unwrap(); + if let Value::Array(arr) = value { + assert_eq!(arr.len(), 5); + for (i, val) in arr.iter().enumerate() { + assert_eq!(val.as_i64().unwrap(), (i + 1) as i64); + } + } else { + panic!("Expected array value"); + } + + // Test simple object + let json = r#"{"name": "Alice", "age": 30, "active": true}"#; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + let value = raw_jsonb.to_value().unwrap(); + if let Value::Object(obj) = value { + assert_eq!(obj.len(), 3); + assert_eq!(obj.get("name").unwrap().as_str().unwrap(), "Alice"); + assert_eq!(obj.get("age").unwrap().as_i64().unwrap(), 30); + assert!(obj.get("active").unwrap().as_bool().unwrap()); + } else { + panic!("Expected object value"); + } + + // Test nested object with array + let json = r#"{"user": {"name": "Bob", "scores": [85, 90, 95]}}"#; + let jsonb = json.parse::().unwrap(); + let raw_jsonb = jsonb.as_raw(); + let value = raw_jsonb.to_value().unwrap(); + if let Value::Object(obj) = value { + assert_eq!(obj.len(), 1); + if let Value::Object(user) = obj.get("user").unwrap() { + assert_eq!(user.len(), 2); + assert_eq!(user.get("name").unwrap().as_str().unwrap(), "Bob"); + + if let Value::Array(scores) = user.get("scores").unwrap() { + assert_eq!(scores.len(), 3); + assert_eq!(scores[0].as_i64().unwrap(), 85); + assert_eq!(scores[1].as_i64().unwrap(), 90); + assert_eq!(scores[2].as_i64().unwrap(), 95); + } else { + panic!("Expected array for scores"); + } + } else { + panic!("Expected object for user"); + } + } else { + panic!("Expected object value"); + } +} + fn init_object<'a>(entries: Vec<(&str, Value<'a>)>) -> Value<'a> { let mut map = BTreeMap::new(); for (key, val) in entries {