aljaz/turso
1
0
Fork 0
mirror of https://github.com/aljazceru/turso.git synced 2026-01-25 19:04:26 +01:00
Code Issues Packages Projects Releases Wiki Activity

Merge 'Use AsValueRef in more functions' from Pedro Muniz

Browse Source 
Depends on #3932
Converting more functions to use `AsValueRef`

Reviewed-by: Jussi Saurio <jussi.saurio@gmail.com>

Closes #3934
This commit is contained in:
Jussi Saurio
2025-11-12 12:54:39 +02:00
committed by GitHub
parent 65a7dd40b3 bc06bb0415
commit 933c3112f9
11 changed files with 851 additions and 531 deletions
Download Patch File Download Diff File Expand all files Collapse all files

168
core/functions/datetime.rs
View File

@@ -1,40 +1,66 @@
use crate::types::AsValueRef;
use crate::types::Value;
use crate::LimboError::InvalidModifier;
use crate::Result;
use crate::{types::Value, vdbe::Register};
use crate::{Result, ValueRef};
use chrono::{
DateTime, Datelike, NaiveDate, NaiveDateTime, NaiveTime, TimeDelta, TimeZone, Timelike, Utc,
};
/// Execution of date/time/datetime functions
#[inline(always)]
pub fn exec_date(values: &[Register]) -> Value {
pub fn exec_date<I, E, V>(values: I) -> Value
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
exec_datetime(values, DateTimeOutput::Date)
}
#[inline(always)]
pub fn exec_time(values: &[Register]) -> Value {
pub fn exec_time<I, E, V>(values: I) -> Value
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
exec_datetime(values, DateTimeOutput::Time)
}
#[inline(always)]
pub fn exec_datetime_full(values: &[Register]) -> Value {
pub fn exec_datetime_full<I, E, V>(values: I) -> Value
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
exec_datetime(values, DateTimeOutput::DateTime)
}
#[inline(always)]
pub fn exec_strftime(values: &[Register]) -> Value {
if values.is_empty() {
pub fn exec_strftime<I, E, V>(values: I) -> Value
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut values = values.into_iter();
if values.len() == 0 {
return Value::Null;
}
let value = &values[0].get_value();
let format_str = if matches!(value, Value::Text(_) | Value::Integer(_) | Value::Float(_)) {
let value = values.next().unwrap();
let value = value.as_value_ref();
let format_str = if matches!(
value,
ValueRef::Text(_) | ValueRef::Integer(_) | ValueRef::Float(_)
) {
format!("{value}")
} else {
return Value::Null;
};
exec_datetime(&values[1..], DateTimeOutput::StrfTime(format_str))
exec_datetime(values, DateTimeOutput::StrfTime(format_str))
}
enum DateTimeOutput {
@@ -46,14 +72,22 @@ enum DateTimeOutput {
JuliaDay,
}
fn exec_datetime(values: &[Register], output_type: DateTimeOutput) -> Value {
if values.is_empty() {
let now = parse_naive_date_time(&Value::build_text("now")).unwrap();
fn exec_datetime<I, E, V>(values: I, output_type: DateTimeOutput) -> Value
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let values = values.into_iter();
if values.len() == 0 {
let now = parse_naive_date_time(Value::build_text("now")).unwrap();
return format_dt(now, output_type, false);
}
if let Some(mut dt) = parse_naive_date_time(values[0].get_value()) {
let mut values = values.peekable();
let first = values.peek().unwrap();
if let Some(mut dt) = parse_naive_date_time(first) {
// if successful, treat subsequent entries as modifiers
modify_dt(&mut dt, &values[1..], output_type)
modify_dt(&mut dt, values.skip(1), output_type)
} else {
// if the first argument is NOT a valid date/time, treat the entire set of values as modifiers.
let mut dt = chrono::Local::now().to_utc().naive_utc();
@@ -61,11 +95,17 @@ fn exec_datetime(values: &[Register], output_type: DateTimeOutput) -> Value {
}
}
fn modify_dt(dt: &mut NaiveDateTime, mods: &[Register], output_type: DateTimeOutput) -> Value {
fn modify_dt<I, E, V>(dt: &mut NaiveDateTime, mods: I, output_type: DateTimeOutput) -> Value
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mods = mods.into_iter();
let mut n_floor: i64 = 0;
let mut subsec_requested = false;
for modifier in mods {
if let Value::Text(ref text_rc) = modifier.get_value() {
if let ValueRef::Text(ref text_rc) = modifier.as_value_ref() {
// TODO: to prevent double conversion and properly support 'utc'/'localtime', we also
// need to keep track of the current timezone and apply it to the modifier.
let parsed = parse_modifier(text_rc.as_str());
@@ -339,7 +379,12 @@ fn last_day_in_month(year: i32, month: u32) -> u32 {
28
}
pub fn exec_julianday(values: &[Register]) -> Value {
pub fn exec_julianday<I, E, V>(values: I) -> Value
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
exec_datetime(values, DateTimeOutput::JuliaDay)
}
@@ -392,11 +437,12 @@ fn get_unixepoch_from_naive_datetime(value: NaiveDateTime) -> i64 {
value.and_utc().timestamp()
}
fn parse_naive_date_time(time_value: &Value) -> Option<NaiveDateTime> {
fn parse_naive_date_time(time_value: impl AsValueRef) -> Option<NaiveDateTime> {
let time_value = time_value.as_value_ref();
match time_value {
Value::Text(s) => get_date_time_from_time_value_string(s.as_str()),
Value::Integer(i) => get_date_time_from_time_value_integer(*i),
Value::Float(f) => get_date_time_from_time_value_float(*f),
ValueRef::Text(s) => get_date_time_from_time_value_string(s.as_str()),
ValueRef::Integer(i) => get_date_time_from_time_value_integer(i),
ValueRef::Float(f) => get_date_time_from_time_value_float(f),
_ => None,
}
}
@@ -767,13 +813,19 @@ fn parse_modifier(modifier: &str) -> Result<Modifier> {
}
}
pub fn exec_timediff(values: &[Register]) -> Value {
pub fn exec_timediff<I, E, V>(values: I) -> Value
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut values = values.into_iter();
if values.len() < 2 {
return Value::Null;
}
let start = parse_naive_date_time(values[0].get_value());
let end = parse_naive_date_time(values[1].get_value());
let start = parse_naive_date_time(values.next().unwrap());
let end = parse_naive_date_time(values.next().unwrap());
match (start, end) {
(Some(start), Some(end)) => {
@@ -931,7 +983,7 @@ mod tests {
];
for (input, expected) in test_cases {
let result = exec_date(&[Register::Value(input.clone())]);
let result = exec_date(&[input.clone()]);
assert_eq!(
result,
Value::build_text(expected.to_string()),
@@ -971,7 +1023,7 @@ mod tests {
];
for case in invalid_cases.iter() {
let result = exec_date(&[Register::Value(case.clone())]);
let result = exec_date([case]);
match result {
Value::Text(ref result_str) if result_str.value.is_empty() => (),
_ => panic!("Expected empty string for input: {case:?}, but got: {result:?}"),
@@ -1063,7 +1115,7 @@ mod tests {
];
for (input, expected) in test_cases {
let result = exec_time(&[Register::Value(input)]);
let result = exec_time(&[input]);
if let Value::Text(result_str) = result {
assert_eq!(result_str.as_str(), expected);
} else {
@@ -1103,7 +1155,7 @@ mod tests {
];
for case in invalid_cases {
let result = exec_time(&[Register::Value(case.clone())]);
let result = exec_time(&[case.clone()]);
match result {
Value::Text(ref result_str) if result_str.value.is_empty() => (),
_ => panic!("Expected empty string for input: {case:?}, but got: {result:?}"),
@@ -1425,8 +1477,8 @@ mod tests {
assert_eq!(dt, create_datetime(2023, 6, 15, 0, 0, 0));
}
fn text(value: &str) -> Register {
Register::Value(Value::build_text(value.to_string()))
fn text(value: &str) -> Value {
Value::build_text(value.to_string())
}
fn format(dt: NaiveDateTime) -> String {
@@ -1444,7 +1496,7 @@ mod tests {
&[text("2023-06-15 12:30:45"), text("-1 day")],
DateTimeOutput::DateTime,
);
assert_eq!(result, *text(&expected).get_value());
assert_eq!(result, text(&expected));
}
#[test]
@@ -1459,7 +1511,7 @@ mod tests {
],
DateTimeOutput::DateTime,
);
assert_eq!(result, *text(&expected).get_value());
assert_eq!(result, text(&expected));
}
#[test]
@@ -1486,7 +1538,7 @@ mod tests {
],
DateTimeOutput::DateTime,
);
assert_eq!(result, *text(&expected).get_value());
assert_eq!(result, text(&expected));
}
#[test]
@@ -1505,7 +1557,7 @@ mod tests {
],
DateTimeOutput::DateTime,
);
assert_eq!(result, *text(&expected).get_value());
assert_eq!(result, text(&expected));
}
#[test]
@@ -1525,7 +1577,7 @@ mod tests {
],
DateTimeOutput::DateTime,
);
assert_eq!(result, *text(&expected).get_value());
assert_eq!(result, text(&expected));
}
#[test]
@@ -1537,13 +1589,12 @@ mod tests {
);
assert_eq!(
result_local,
*text(
text(
&dt.and_utc()
.with_timezone(&chrono::Local)
.format("%Y-%m-%d %H:%M:%S")
.to_string()
)
.get_value()
);
// TODO: utc modifier assumes time given is not already utc
// add test when fixed in the future
@@ -1581,7 +1632,7 @@ mod tests {
.unwrap();
let expected = format(max);
let result = exec_datetime(&[text("9999-12-31 23:59:59")], DateTimeOutput::DateTime);
assert_eq!(result, *text(&expected).get_value());
assert_eq!(result, text(&expected));
}
// leap second
@@ -1593,7 +1644,7 @@ mod tests {
.unwrap();
let expected = String::new(); // SQLite ignores leap seconds
let result = exec_datetime(&[text(&leap_second.to_string())], DateTimeOutput::DateTime);
assert_eq!(result, *text(&expected).get_value());
assert_eq!(result, text(&expected));
}
#[test]
@@ -1816,61 +1867,40 @@ mod tests {
let start = Value::build_text("12:00:00");
let end = Value::build_text("14:30:45");
let expected = Value::build_text("-0000-00-00 02:30:45.000");
assert_eq!(
exec_timediff(&[Register::Value(start), Register::Value(end)]),
expected
);
assert_eq!(exec_timediff(&[start, end]), expected);
let start = Value::build_text("14:30:45");
let end = Value::build_text("12:00:00");
let expected = Value::build_text("+0000-00-00 02:30:45.000");
assert_eq!(
exec_timediff(&[Register::Value(start), Register::Value(end)]),
expected
);
assert_eq!(exec_timediff(&[start, end]), expected);
let start = Value::build_text("12:00:01.300");
let end = Value::build_text("12:00:00.500");
let expected = Value::build_text("+0000-00-00 00:00:00.800");
assert_eq!(
exec_timediff(&[Register::Value(start), Register::Value(end)]),
expected
);
assert_eq!(exec_timediff(&[start, end]), expected);
let start = Value::build_text("13:30:00");
let end = Value::build_text("16:45:30");
let expected = Value::build_text("-0000-00-00 03:15:30.000");
assert_eq!(
exec_timediff(&[Register::Value(start), Register::Value(end)]),
expected
);
assert_eq!(exec_timediff(&[start, end]), expected);
let start = Value::build_text("2023-05-10 23:30:00");
let end = Value::build_text("2023-05-11 01:15:00");
let expected = Value::build_text("-0000-00-00 01:45:00.000");
assert_eq!(
exec_timediff(&[Register::Value(start), Register::Value(end)]),
expected
);
assert_eq!(exec_timediff(&[start, end]), expected);
let start = Value::Null;
let end = Value::build_text("12:00:00");
let expected = Value::Null;
assert_eq!(
exec_timediff(&[Register::Value(start), Register::Value(end)]),
expected
);
assert_eq!(exec_timediff(&[start, end]), expected);
let start = Value::build_text("not a time");
let end = Value::build_text("12:00:00");
let expected = Value::Null;
assert_eq!(
exec_timediff(&[Register::Value(start), Register::Value(end)]),
expected
);
assert_eq!(exec_timediff(&[start, end.clone()]), expected);
let start = Value::build_text("12:00:00");
let expected = Value::Null;
assert_eq!(exec_timediff(&[Register::Value(start)]), expected);
assert_eq!(exec_timediff(&[start, end]), expected);
}
}

24
core/json/cache.rs
View File

@@ -1,6 +1,7 @@
use std::cell::{Cell, UnsafeCell};
use crate::Value;
use crate::types::AsValueRef;
use crate::{Value, ValueRef};
use super::jsonb::Jsonb;
@@ -38,25 +39,27 @@ impl JsonCache {
oldest_idx
}
pub fn insert(&mut self, key: &Value, value: &Jsonb) {
pub fn insert(&mut self, key: impl AsValueRef, value: &Jsonb) {
let key = key.as_value_ref();
if self.used < JSON_CACHE_SIZE {
self.entries[self.used] = Some((key.clone(), value.clone()));
self.entries[self.used] = Some((key.to_owned(), value.clone()));
self.age[self.used] = self.counter;
self.counter += 1;
self.used += 1
} else {
let id = self.find_oldest_entry();
self.entries[id] = Some((key.clone(), value.clone()));
self.entries[id] = Some((key.to_owned(), value.clone()));
self.age[id] = self.counter;
self.counter += 1;
}
}
pub fn lookup(&mut self, key: &Value) -> Option<Jsonb> {
pub fn lookup(&mut self, key: impl AsValueRef) -> Option<Jsonb> {
let key = key.as_value_ref();
for i in (0..self.used).rev() {
if let Some((stored_key, value)) = &self.entries[i] {
if key == stored_key {
if key == *stored_key {
self.age[i] = self.counter;
self.counter += 1;
let json = value.clone();
@@ -89,7 +92,7 @@ impl JsonCacheCell {
}
#[cfg(test)]
pub fn lookup(&self, key: &Value) -> Option<Jsonb> {
pub fn lookup(&self, key: impl AsValueRef) -> Option<Jsonb> {
assert!(!self.accessed.get());
self.accessed.set(true);
@@ -113,11 +116,12 @@ impl JsonCacheCell {
pub fn get_or_insert_with(
&self,
key: &Value,
value: impl Fn(&Value) -> crate::Result<Jsonb>,
key: impl AsValueRef,
value: impl FnOnce(ValueRef) -> crate::Result<Jsonb>,
) -> crate::Result<Jsonb> {
assert!(!self.accessed.get());
let key = key.as_value_ref();
self.accessed.set(true);
let result = unsafe {
let cache_ptr = self.inner.get();
@@ -354,7 +358,7 @@ mod tests {
// Insert the value using get_or_insert_with
let insert_result = cache_cell.get_or_insert_with(&key, |k| {
// Verify that k is the same as our key
assert_eq!(k, &key);
assert_eq!(k, key);
Ok(value.clone())
});

448
core/json/mod.rs
View File

@@ -11,8 +11,7 @@ pub use crate::json::ops::{
jsonb_replace,
};
use crate::json::path::{json_path, JsonPath, PathElement};
use crate::types::{Text, TextRef, TextSubtype, Value, ValueType};
use crate::vdbe::Register;
use crate::types::{AsValueRef, Text, TextSubtype, Value, ValueType};
use crate::{bail_constraint_error, bail_parse_error, LimboError, ValueRef};
pub use cache::JsonCacheCell;
use jsonb::{ElementType, Jsonb, JsonbHeader, PathOperationMode, SearchOperation, SetOperation};
@@ -100,17 +99,9 @@ pub fn json_from_raw_bytes_agg(data: &[u8], raw: bool) -> crate::Result<Value> {
}
}
pub fn convert_dbtype_to_jsonb(val: &Value, strict: Conv) -> crate::Result<Jsonb> {
convert_ref_dbtype_to_jsonb(
match val {
Value::Null => ValueRef::Null,
Value::Integer(x) => ValueRef::Integer(*x),
Value::Float(x) => ValueRef::Float(*x),
Value::Text(text) => ValueRef::Text(TextRef::new(text.as_str(), text.subtype)),
Value::Blob(items) => ValueRef::Blob(items.as_slice()),
},
strict,
)
pub fn convert_dbtype_to_jsonb(val: impl AsValueRef, strict: Conv) -> crate::Result<Jsonb> {
let val = val.as_value_ref();
convert_ref_dbtype_to_jsonb(val, strict)
}
fn parse_as_json_text(slice: &[u8], mode: Conv) -> crate::Result<Jsonb> {
@@ -187,18 +178,27 @@ pub fn convert_ref_dbtype_to_jsonb(val: ValueRef<'_>, strict: Conv) -> crate::Re
}
}
pub fn curry_convert_dbtype_to_jsonb(strict: Conv) -> impl Fn(&Value) -> crate::Result<Jsonb> {
pub fn curry_convert_dbtype_to_jsonb(
strict: Conv,
) -> impl FnOnce(ValueRef) -> crate::Result<Jsonb> {
move |val| convert_dbtype_to_jsonb(val, strict)
}
pub fn json_array(values: &[Register]) -> crate::Result<Value> {
pub fn json_array<I, E, V>(values: I) -> crate::Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let values = values.into_iter();
let mut json = Jsonb::make_empty_array(values.len());
for value in values.iter() {
if matches!(value.get_value(), Value::Blob(_)) {
for value in values {
let value = value.as_value_ref();
if matches!(value, ValueRef::Blob(_)) {
crate::bail_constraint_error!("JSON cannot hold BLOB values")
}
let value = convert_dbtype_to_jsonb(value.get_value(), Conv::NotStrict)?;
let value = convert_dbtype_to_jsonb(value, Conv::NotStrict)?;
json.append_jsonb_to_end(value.data());
}
json.finalize_unsafe(ElementType::ARRAY)?;
@@ -206,14 +206,21 @@ pub fn json_array(values: &[Register]) -> crate::Result<Value> {
json_string_to_db_type(json, ElementType::ARRAY, OutputVariant::ElementType)
}
pub fn jsonb_array(values: &[Register]) -> crate::Result<Value> {
pub fn jsonb_array<I, E, V>(values: I) -> crate::Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let values = values.into_iter();
let mut json = Jsonb::make_empty_array(values.len());
for value in values.iter() {
if matches!(value.get_value(), Value::Blob(_)) {
for value in values {
let value = value.as_value_ref();
if matches!(value, ValueRef::Blob(_)) {
crate::bail_constraint_error!("JSON cannot hold BLOB values")
}
let value = convert_dbtype_to_jsonb(value.get_value(), Conv::NotStrict)?;
let value = convert_dbtype_to_jsonb(value, Conv::NotStrict)?;
json.append_jsonb_to_end(value.data());
}
json.finalize_unsafe(ElementType::ARRAY)?;
@@ -246,19 +253,27 @@ pub fn json_array_length(
Ok(Value::Null)
}
pub fn json_set(args: &[Register], json_cache: &JsonCacheCell) -> crate::Result<Value> {
if args.is_empty() {
pub fn json_set<I, E, V>(args: I, json_cache: &JsonCacheCell) -> crate::Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut args = args.into_iter();
if args.len() == 0 {
return Ok(Value::Null);
}
let make_jsonb_fn = curry_convert_dbtype_to_jsonb(Conv::Strict);
let mut json = json_cache.get_or_insert_with(args[0].get_value(), make_jsonb_fn)?;
let other = args[1..].chunks_exact(2);
let mut json = json_cache.get_or_insert_with(args.next().unwrap(), make_jsonb_fn)?;
for chunk in other {
let path = json_path_from_db_value(chunk[0].get_value(), true)?;
// TODO: when `array_chunks` is stabilized we can chunk by 2 here
while args.len() > 1 {
let first = args.next().unwrap();
let value = convert_dbtype_to_jsonb(chunk[1].get_value(), Conv::NotStrict)?;
let path = json_path_from_db_value(&first, true)?;
let second = args.next().unwrap();
let value = convert_dbtype_to_jsonb(second, Conv::NotStrict)?;
let mut op = SetOperation::new(value);
if let Some(path) = path {
let _ = json.operate_on_path(&path, &mut op);
@@ -270,19 +285,27 @@ pub fn json_set(args: &[Register], json_cache: &JsonCacheCell) -> crate::Result<
json_string_to_db_type(json, el_type, OutputVariant::String)
}
pub fn jsonb_set(args: &[Register], json_cache: &JsonCacheCell) -> crate::Result<Value> {
if args.is_empty() {
pub fn jsonb_set<I, E, V>(args: I, json_cache: &JsonCacheCell) -> crate::Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut args = args.into_iter();
if args.len() == 0 {
return Ok(Value::Null);
}
let make_jsonb_fn = curry_convert_dbtype_to_jsonb(Conv::Strict);
let mut json = json_cache.get_or_insert_with(args[0].get_value(), make_jsonb_fn)?;
let other = args[1..].chunks_exact(2);
let mut json = json_cache.get_or_insert_with(args.next().unwrap(), make_jsonb_fn)?;
for chunk in other {
let path = json_path_from_db_value(chunk[0].get_value(), true)?;
// TODO: when `array_chunks` is stabilized we can chunk by 2 here
while args.len() > 1 {
let first = args.next().unwrap();
let path = json_path_from_db_value(&first, true)?;
let value = convert_dbtype_to_jsonb(chunk[1].get_value(), Conv::NotStrict)?;
let second = args.next().unwrap();
let value = convert_dbtype_to_jsonb(second, Conv::NotStrict)?;
let mut op = SetOperation::new(value);
if let Some(path) = path {
let _ = json.operate_on_path(&path, &mut op);
@@ -297,15 +320,16 @@ pub fn jsonb_set(args: &[Register], json_cache: &JsonCacheCell) -> crate::Result
/// Implements the -> operator. Always returns a proper JSON value.
/// https://sqlite.org/json1.html#the_and_operators
pub fn json_arrow_extract(
value: &Value,
path: &Value,
value: impl AsValueRef,
path: impl AsValueRef,
json_cache: &JsonCacheCell,
) -> crate::Result<Value> {
if let Value::Null = value {
let value = value.as_value_ref();
if let ValueRef::Null = value {
return Ok(Value::Null);
}
if let Some(path) = json_path_from_db_value(path, false)? {
if let Some(path) = json_path_from_db_value(&path, false)? {
let make_jsonb_fn = curry_convert_dbtype_to_jsonb(Conv::Strict);
let mut json = json_cache.get_or_insert_with(value, make_jsonb_fn)?;
let mut op = SearchOperation::new(json.len());
@@ -324,14 +348,15 @@ pub fn json_arrow_extract(
/// Implements the ->> operator. Always returns a SQL representation of the JSON subcomponent.
/// https://sqlite.org/json1.html#the_and_operators
pub fn json_arrow_shift_extract(
value: &Value,
path: &Value,
value: impl AsValueRef,
path: impl AsValueRef,
json_cache: &JsonCacheCell,
) -> crate::Result<Value> {
if let Value::Null = value {
let value = value.as_value_ref();
if let ValueRef::Null = value {
return Ok(Value::Null);
}
if let Some(path) = json_path_from_db_value(path, false)? {
if let Some(path) = json_path_from_db_value(&path, false)? {
let make_jsonb_fn = curry_convert_dbtype_to_jsonb(Conv::Strict);
let mut json = json_cache.get_or_insert_with(value, make_jsonb_fn)?;
let mut op = SearchOperation::new(json.len());
@@ -359,16 +384,23 @@ pub fn json_arrow_shift_extract(
/// Extracts a JSON value from a JSON object or array.
/// If there's only a single path, the return value might be either a TEXT or a database type.
/// https://sqlite.org/json1.html#the_json_extract_function
pub fn json_extract(
value: &Value,
paths: &[Register],
pub fn json_extract<I, E, V>(
value: impl AsValueRef,
paths: I,
json_cache: &JsonCacheCell,
) -> crate::Result<Value> {
if let Value::Null = value {
) -> crate::Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let value = value.as_value_ref();
if let ValueRef::Null = value {
return Ok(Value::Null);
}
if paths.is_empty() {
let paths = paths.into_iter();
if paths.len() == 0 {
return Ok(Value::Null);
}
let convert_to_jsonb = curry_convert_dbtype_to_jsonb(Conv::Strict);
@@ -380,16 +412,22 @@ pub fn json_extract(
Ok(result)
}
pub fn jsonb_extract(
pub fn jsonb_extract<I, E, V>(
value: &Value,
paths: &[Register],
paths: I,
json_cache: &JsonCacheCell,
) -> crate::Result<Value> {
) -> crate::Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
if let Value::Null = value {
return Ok(Value::Null);
}
if paths.is_empty() {
let paths = paths.into_iter();
if paths.len() == 0 {
return Ok(Value::Null);
}
let convert_to_jsonb = curry_convert_dbtype_to_jsonb(Conv::Strict);
@@ -401,10 +439,14 @@ pub fn jsonb_extract(
Ok(result)
}
fn jsonb_extract_internal(value: Jsonb, paths: &[Register]) -> crate::Result<(Jsonb, ElementType)> {
fn jsonb_extract_internal<E, V>(value: Jsonb, mut paths: E) -> crate::Result<(Jsonb, ElementType)>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
{
let null = Jsonb::from_raw_data(JsonbHeader::make_null().into_bytes().as_bytes());
if paths.len() == 1 {
if let Some(path) = json_path_from_db_value(paths[0].get_value(), true)? {
if let Some(path) = json_path_from_db_value(&paths.next().unwrap(), true)? {
let mut json = value;
let mut op = SearchOperation::new(json.len());
@@ -428,10 +470,8 @@ fn jsonb_extract_internal(value: Jsonb, paths: &[Register]) -> crate::Result<(Js
let mut result = Jsonb::make_empty_array(json.len());
// TODO: make an op to avoid creating new json for every path element
let paths = paths
.iter()
.map(|p| json_path_from_db_value(p.get_value(), true));
for path in paths {
let path = json_path_from_db_value(&path, true);
if let Some(path) = path? {
let mut op = SearchOperation::new(json.len());
let res = json.operate_on_path(&path, &mut op);
@@ -502,8 +542,9 @@ pub fn json_string_to_db_type(
}
}
pub fn json_type(value: &Value, path: Option<&Value>) -> crate::Result<Value> {
if let Value::Null = value {
pub fn json_type(value: impl AsValueRef, path: Option<impl AsValueRef>) -> crate::Result<Value> {
let value = value.as_value_ref();
if let ValueRef::Null = value {
return Ok(Value::Null);
}
if path.is_none() {
@@ -512,7 +553,7 @@ pub fn json_type(value: &Value, path: Option<&Value>) -> crate::Result<Value> {
return Ok(Value::Text(Text::json(element_type.into())));
}
if let Some(path) = json_path_from_db_value(path.unwrap(), true)? {
if let Some(path) = json_path_from_db_value(&path.unwrap(), true)? {
let mut json = convert_dbtype_to_jsonb(value, Conv::Strict)?;
if let Ok(mut path) = json.navigate_path(&path, PathOperationMode::ReplaceExisting) {
@@ -531,16 +572,20 @@ pub fn json_type(value: &Value, path: Option<&Value>) -> crate::Result<Value> {
}
}
fn json_path_from_db_value(path: &Value, strict: bool) -> crate::Result<Option<JsonPath<'_>>> {
fn json_path_from_db_value<'a>(
path: &'a (impl AsValueRef + 'a),
strict: bool,
) -> crate::Result<Option<JsonPath<'a>>> {
let path = path.as_value_ref();
let json_path = if strict {
match path {
Value::Text(t) => json_path(t.as_str())?,
Value::Null => return Ok(None),
ValueRef::Text(t) => json_path(t.as_str())?,
ValueRef::Null => return Ok(None),
_ => crate::bail_constraint_error!("JSON path error near: {:?}", path.to_string()),
}
} else {
match path {
Value::Text(t) => {
ValueRef::Text(t) => {
if t.as_str().starts_with("$") {
json_path(t.as_str())?
} else {
@@ -552,14 +597,14 @@ fn json_path_from_db_value(path: &Value, strict: bool) -> crate::Result<Option<J
}
}
}
Value::Null => return Ok(None),
Value::Integer(i) => JsonPath {
ValueRef::Null => return Ok(None),
ValueRef::Integer(i) => JsonPath {
elements: vec![
PathElement::Root(),
PathElement::ArrayLocator(Some(*i as i32)),
PathElement::ArrayLocator(Some(i as i32)),
],
},
Value::Float(f) => JsonPath {
ValueRef::Float(f) => JsonPath {
elements: vec![
PathElement::Root(),
PathElement::Key(Cow::Owned(f.to_string()), false),
@@ -572,9 +617,9 @@ fn json_path_from_db_value(path: &Value, strict: bool) -> crate::Result<Option<J
Ok(Some(json_path))
}
pub fn json_error_position(json: &Value) -> crate::Result<Value> {
match json {
Value::Text(t) => match Jsonb::from_str(t.as_str()) {
pub fn json_error_position(json: impl AsValueRef) -> crate::Result<Value> {
match json.as_value_ref() {
ValueRef::Text(t) => match Jsonb::from_str(t.as_str()) {
Ok(_) => Ok(Value::Integer(0)),
Err(JsonError::Message { location, .. }) => {
if let Some(loc) = location {
@@ -587,10 +632,10 @@ pub fn json_error_position(json: &Value) -> crate::Result<Value> {
}
}
},
Value::Blob(_) => {
ValueRef::Blob(_) => {
bail_parse_error!("Unsupported")
}
Value::Null => Ok(Value::Null),
ValueRef::Null => Ok(Value::Null),
_ => Ok(Value::Integer(0)),
}
}
@@ -598,19 +643,30 @@ pub fn json_error_position(json: &Value) -> crate::Result<Value> {
/// Constructs a JSON object from a list of values that represent key-value pairs.
/// The number of values must be even, and the first value of each pair (which represents the map key)
/// must be a TEXT value. The second value of each pair can be any JSON value (which represents the map value)
pub fn json_object(values: &[Register]) -> crate::Result<Value> {
pub fn json_object<I, E, V>(values: I) -> crate::Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut values = values.into_iter();
if values.len() % 2 != 0 {
bail_constraint_error!("json_object() requires an even number of arguments")
}
let mut json = Jsonb::make_empty_obj(values.len() * 50);
for chunk in values.chunks_exact(2) {
if chunk[0].get_value().value_type() != ValueType::Text {
// TODO: when `array_chunks` is stabilized we can chunk by 2 here
while values.len() > 1 {
let first = values.next().unwrap();
let first = first.as_value_ref();
if first.value_type() != ValueType::Text {
bail_constraint_error!("json_object() labels must be TEXT")
}
let key = convert_dbtype_to_jsonb(chunk[0].get_value(), Conv::ToString)?;
let key = convert_dbtype_to_jsonb(first, Conv::ToString)?;
json.append_jsonb_to_end(key.data());
let value = convert_dbtype_to_jsonb(chunk[1].get_value(), Conv::NotStrict)?;
let second = values.next().unwrap();
let value = convert_dbtype_to_jsonb(second, Conv::NotStrict)?;
json.append_jsonb_to_end(value.data());
}
@@ -619,19 +675,30 @@ pub fn json_object(values: &[Register]) -> crate::Result<Value> {
json_string_to_db_type(json, ElementType::OBJECT, OutputVariant::String)
}
pub fn jsonb_object(values: &[Register]) -> crate::Result<Value> {
pub fn jsonb_object<I, E, V>(values: I) -> crate::Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut values = values.into_iter();
if values.len() % 2 != 0 {
bail_constraint_error!("json_object() requires an even number of arguments")
}
let mut json = Jsonb::make_empty_obj(values.len() * 50);
for chunk in values.chunks_exact(2) {
if chunk[0].get_value().value_type() != ValueType::Text {
// TODO: when `array_chunks` is stabilized we can chunk by 2 here
while values.len() > 1 {
let first = values.next().unwrap();
let first = first.as_value_ref();
if first.value_type() != ValueType::Text {
bail_constraint_error!("json_object() labels must be TEXT")
}
let key = convert_dbtype_to_jsonb(chunk[0].get_value(), Conv::ToString)?;
let key = convert_dbtype_to_jsonb(first, Conv::ToString)?;
json.append_jsonb_to_end(key.data());
let value = convert_dbtype_to_jsonb(chunk[1].get_value(), Conv::NotStrict)?;
let second = values.next().unwrap();
let value = convert_dbtype_to_jsonb(second, Conv::NotStrict)?;
json.append_jsonb_to_end(value.data());
}
@@ -642,8 +709,9 @@ pub fn jsonb_object(values: &[Register]) -> crate::Result<Value> {
/// Tries to convert the value to jsonb. Returns Value::Integer(1) if it the conversion
/// succeeded, and Value::Integer(0) if it didn't.
pub fn is_json_valid(json_value: &Value) -> Value {
if matches!(json_value, Value::Null) {
pub fn is_json_valid(json_value: impl AsValueRef) -> Value {
let json_value = json_value.as_value_ref();
if matches!(json_value, ValueRef::Null) {
return Value::Null;
}
convert_dbtype_to_jsonb(json_value, Conv::Strict)
@@ -651,9 +719,10 @@ pub fn is_json_valid(json_value: &Value) -> Value {
.unwrap_or(Value::Integer(0))
}
pub fn json_quote(value: &Value) -> crate::Result<Value> {
pub fn json_quote(value: impl AsValueRef) -> crate::Result<Value> {
let value = value.as_value_ref();
match value {
Value::Text(ref t) => {
ValueRef::Text(ref t) => {
// If X is a JSON value returned by another JSON function,
// then this function is a no-op
if t.subtype == TextSubtype::Json {
@@ -678,10 +747,10 @@ pub fn json_quote(value: &Value) -> crate::Result<Value> {
Ok(Value::build_text(escaped_value))
}
// Numbers are unquoted in json
Value::Integer(ref int) => Ok(Value::Integer(int.to_owned())),
Value::Float(ref float) => Ok(Value::Float(float.to_owned())),
Value::Blob(_) => crate::bail_constraint_error!("JSON cannot hold BLOB values"),
Value::Null => Ok(Value::build_text("null")),
ValueRef::Integer(int) => Ok(Value::Integer(int)),
ValueRef::Float(float) => Ok(Value::Float(float)),
ValueRef::Blob(_) => crate::bail_constraint_error!("JSON cannot hold BLOB values"),
ValueRef::Null => Ok(Value::build_text("null")),
}
}
@@ -808,14 +877,9 @@ mod tests {
#[test]
fn test_json_array_simple() {
let text = Register::Value(Value::build_text("value1"));
let json = Register::Value(Value::Text(Text::json("\"value2\"".to_string())));
let input = vec![
text,
json,
Register::Value(Value::Integer(1)),
Register::Value(Value::Float(1.1)),
];
let text = Value::build_text("value1");
let json = Value::Text(Text::json("\"value2\"".to_string()));
let input = [text, json, Value::Integer(1), Value::Float(1.1)];
let result = json_array(&input).unwrap();
if let Value::Text(res) = result {
@@ -828,9 +892,9 @@ mod tests {
#[test]
fn test_json_array_empty() {
let input = vec![];
let input: [Value; 0] = [];
let result = json_array(&input).unwrap();
let result = json_array(input).unwrap();
if let Value::Text(res) = result {
assert_eq!(res.as_str(), "[]");
assert_eq!(res.subtype, TextSubtype::Json);
@@ -841,9 +905,9 @@ mod tests {
#[test]
fn test_json_array_blob_invalid() {
let blob = Register::Value(Value::Blob("1".as_bytes().to_vec()));
let blob = Value::Blob("1".as_bytes().to_vec());
let input = vec![blob];
let input = [blob];
let result = json_array(&input);
@@ -967,8 +1031,8 @@ mod tests {
fn test_json_extract_missing_path() {
let json_cache = JsonCacheCell::new();
let result = json_extract(
&Value::build_text("{\"a\":2}"),
&[Register::Value(Value::build_text("$.x"))],
Value::build_text("{\"a\":2}"),
&[Value::build_text("$.x")],
&json_cache,
);
@@ -980,11 +1044,7 @@ mod tests {
#[test]
fn test_json_extract_null_path() {
let json_cache = JsonCacheCell::new();
let result = json_extract(
&Value::build_text("{\"a\":2}"),
&[Register::Value(Value::Null)],
&json_cache,
);
let result = json_extract(Value::build_text("{\"a\":2}"), &[Value::Null], &json_cache);
match result {
Ok(Value::Null) => (),
@@ -996,8 +1056,8 @@ mod tests {
fn test_json_path_invalid() {
let json_cache = JsonCacheCell::new();
let result = json_extract(
&Value::build_text("{\"a\":2}"),
&[Register::Value(Value::Float(1.1))],
Value::build_text("{\"a\":2}"),
&[Value::Float(1.1)],
&json_cache,
);
@@ -1058,9 +1118,9 @@ mod tests {
#[test]
fn test_json_object_simple() {
let key = Register::Value(Value::build_text("key"));
let value = Register::Value(Value::build_text("value"));
let input = vec![key, value];
let key = Value::build_text("key");
let value = Value::build_text("value");
let input = [key, value];
let result = json_object(&input).unwrap();
let Value::Text(json_text) = result else {
@@ -1082,17 +1142,17 @@ mod tests {
let null_key = Value::build_text("null_key");
let null_value = Value::Null;
let input = vec![
Register::Value(text_key),
Register::Value(text_value),
Register::Value(json_key),
Register::Value(json_value),
Register::Value(integer_key),
Register::Value(integer_value),
Register::Value(float_key),
Register::Value(float_value),
Register::Value(null_key),
Register::Value(null_value),
let input = [
text_key,
text_value,
json_key,
json_value,
integer_key,
integer_value,
float_key,
float_value,
null_key,
null_value,
];
let result = json_object(&input).unwrap();
@@ -1107,9 +1167,9 @@ mod tests {
#[test]
fn test_json_object_json_value_is_rendered_as_json() {
let key = Register::Value(Value::build_text("key"));
let value = Register::Value(Value::Text(Text::json(r#"{"json":"value"}"#.to_string())));
let input = vec![key, value];
let key = Value::build_text("key");
let value = Value::Text(Text::json(r#"{"json":"value"}"#.to_string()));
let input = [key, value];
let result = json_object(&input).unwrap();
let Value::Text(json_text) = result else {
@@ -1120,9 +1180,9 @@ mod tests {
#[test]
fn test_json_object_json_text_value_is_rendered_as_regular_text() {
let key = Register::Value(Value::build_text("key"));
let value = Register::Value(Value::Text(Text::new(r#"{"json":"value"}"#)));
let input = vec![key, value];
let key = Value::build_text("key");
let value = Value::Text(Text::new(r#"{"json":"value"}"#));
let input = [key, value];
let result = json_object(&input).unwrap();
let Value::Text(json_text) = result else {
@@ -1133,13 +1193,13 @@ mod tests {
#[test]
fn test_json_object_nested() {
let key = Register::Value(Value::build_text("key"));
let value = Register::Value(Value::build_text("value"));
let input = vec![key, value];
let key = Value::build_text("key");
let value = Value::build_text("value");
let input = [key, value];
let parent_key = Register::Value(Value::build_text("parent_key"));
let parent_value = Register::Value(json_object(&input).unwrap());
let parent_input = vec![parent_key, parent_value];
let parent_key = Value::build_text("parent_key");
let parent_value = json_object(&input).unwrap();
let parent_input = [parent_key, parent_value];
let result = json_object(&parent_input).unwrap();
@@ -1151,9 +1211,9 @@ mod tests {
#[test]
fn test_json_object_duplicated_keys() {
let key = Register::Value(Value::build_text("key"));
let value = Register::Value(Value::build_text("value"));
let input = vec![key.clone(), value.clone(), key, value];
let key = Value::build_text("key");
let value = Value::build_text("value");
let input = [key.clone(), value.clone(), key, value];
let result = json_object(&input).unwrap();
let Value::Text(json_text) = result else {
@@ -1164,7 +1224,7 @@ mod tests {
#[test]
fn test_json_object_empty() {
let input = vec![];
let input: [Value; 0] = [];
let result = json_object(&input).unwrap();
let Value::Text(json_text) = result else {
@@ -1175,9 +1235,9 @@ mod tests {
#[test]
fn test_json_object_non_text_key() {
let key = Register::Value(Value::Integer(1));
let value = Register::Value(Value::build_text("value"));
let input = vec![key, value];
let key = Value::Integer(1);
let value = Value::build_text("value");
let input = [key, value];
match json_object(&input) {
Ok(_) => panic!("Expected error for non-TEXT key"),
@@ -1187,9 +1247,9 @@ mod tests {
#[test]
fn test_json_odd_number_of_values() {
let key = Register::Value(Value::build_text("key"));
let value = Register::Value(Value::build_text("value"));
let input = vec![key.clone(), value, key];
let key = Value::build_text("key");
let value = Value::build_text("value");
let input = [key.clone(), value, key];
assert!(json_object(&input).is_err());
}
@@ -1326,9 +1386,9 @@ mod tests {
let json_cache = JsonCacheCell::new();
let result = json_set(
&[
Register::Value(Value::build_text("{}")),
Register::Value(Value::build_text("$.field")),
Register::Value(Value::build_text("value")),
Value::build_text("{}"),
Value::build_text("$.field"),
Value::build_text("value"),
],
&json_cache,
);
@@ -1343,9 +1403,9 @@ mod tests {
let json_cache = JsonCacheCell::new();
let result = json_set(
&[
Register::Value(Value::build_text(r#"{"field":"old_value"}"#)),
Register::Value(Value::build_text("$.field")),
Register::Value(Value::build_text("new_value")),
Value::build_text(r#"{"field":"old_value"}"#),
Value::build_text("$.field"),
Value::build_text("new_value"),
],
&json_cache,
);
@@ -1363,9 +1423,9 @@ mod tests {
let json_cache = JsonCacheCell::new();
let result = json_set(
&[
Register::Value(Value::build_text("{}")),
Register::Value(Value::build_text("$.object.doesnt.exist")),
Register::Value(Value::build_text("value")),
Value::build_text("{}"),
Value::build_text("$.object.doesnt.exist"),
Value::build_text("value"),
],
&json_cache,
);
@@ -1383,9 +1443,9 @@ mod tests {
let json_cache = JsonCacheCell::new();
let result = json_set(
&[
Register::Value(Value::build_text("[]")),
Register::Value(Value::build_text("$[0]")),
Register::Value(Value::build_text("value")),
Value::build_text("[]"),
Value::build_text("$[0]"),
Value::build_text("value"),
],
&json_cache,
);
@@ -1400,9 +1460,9 @@ mod tests {
let json_cache = JsonCacheCell::new();
let result = json_set(
&[
Register::Value(Value::build_text("{}")),
Register::Value(Value::build_text("$.some_array[0]")),
Register::Value(Value::Integer(123)),
Value::build_text("{}"),
Value::build_text("$.some_array[0]"),
Value::Integer(123),
],
&json_cache,
);
@@ -1420,9 +1480,9 @@ mod tests {
let json_cache = JsonCacheCell::new();
let result = json_set(
&[
Register::Value(Value::build_text("[123]")),
Register::Value(Value::build_text("$[1]")),
Register::Value(Value::Integer(456)),
Value::build_text("[123]"),
Value::build_text("$[1]"),
Value::Integer(456),
],
&json_cache,
);
@@ -1437,9 +1497,9 @@ mod tests {
let json_cache = JsonCacheCell::new();
let result = json_set(
&[
Register::Value(Value::build_text("[123]")),
Register::Value(Value::build_text("$[200]")),
Register::Value(Value::Integer(456)),
Value::build_text("[123]"),
Value::build_text("$[200]"),
Value::Integer(456),
],
&json_cache,
);
@@ -1454,9 +1514,9 @@ mod tests {
let json_cache = JsonCacheCell::new();
let result = json_set(
&[
Register::Value(Value::build_text("[123]")),
Register::Value(Value::build_text("$[0]")),
Register::Value(Value::Integer(456)),
Value::build_text("[123]"),
Value::build_text("$[0]"),
Value::Integer(456),
],
&json_cache,
);
@@ -1470,11 +1530,7 @@ mod tests {
fn test_json_set_null_path() {
let json_cache = JsonCacheCell::new();
let result = json_set(
&[
Register::Value(Value::build_text("{}")),
Register::Value(Value::Null),
Register::Value(Value::Integer(456)),
],
&[Value::build_text("{}"), Value::Null, Value::Integer(456)],
&json_cache,
);
@@ -1488,11 +1544,11 @@ mod tests {
let json_cache = JsonCacheCell::new();
let result = json_set(
&[
Register::Value(Value::build_text("[123]")),
Register::Value(Value::build_text("$[0]")),
Register::Value(Value::Integer(456)),
Register::Value(Value::build_text("$[1]")),
Register::Value(Value::Integer(789)),
Value::build_text("[123]"),
Value::build_text("$[0]"),
Value::Integer(456),
Value::build_text("$[1]"),
Value::Integer(789),
],
&json_cache,
);
@@ -1507,9 +1563,9 @@ mod tests {
let json_cache = JsonCacheCell::new();
let result = json_set(
&[
Register::Value(Value::build_text("{}")),
Register::Value(Value::build_text("$.object[0].field")),
Register::Value(Value::Integer(123)),
Value::build_text("{}"),
Value::build_text("$.object[0].field"),
Value::Integer(123),
],
&json_cache,
);
@@ -1527,9 +1583,9 @@ mod tests {
let json_cache = JsonCacheCell::new();
let result = json_set(
&[
Register::Value(Value::build_text("{}")),
Register::Value(Value::build_text("$.object[0][0]")),
Register::Value(Value::Integer(123)),
Value::build_text("{}"),
Value::build_text("$.object[0][0]"),
Value::Integer(123),
],
&json_cache,
);
@@ -1544,11 +1600,11 @@ mod tests {
let json_cache = JsonCacheCell::new();
let result = json_set(
&[
Register::Value(Value::build_text("{}")),
Register::Value(Value::build_text("$.object[123].another")),
Register::Value(Value::build_text("value")),
Register::Value(Value::build_text("$.field")),
Register::Value(Value::build_text("value")),
Value::build_text("{}"),
Value::build_text("$.object[123].another"),
Value::build_text("value"),
Value::build_text("$.field"),
Value::build_text("value"),
],
&json_cache,
);

203
core/json/ops.rs
View File

@@ -1,4 +1,7 @@
use crate::{types::Value, vdbe::Register};
use crate::{
types::{AsValueRef, Value},
ValueRef,
};
use super::{
convert_dbtype_to_jsonb, curry_convert_dbtype_to_jsonb, json_path_from_db_value,
@@ -12,16 +15,22 @@ use super::{
/// * If the patch contains a scalar value, the target is replaced with that value
/// * If both target and patch are objects, the patch is recursively applied
/// * null values in the patch result in property removal from the target
pub fn json_patch(target: &Value, patch: &Value, cache: &JsonCacheCell) -> crate::Result<Value> {
pub fn json_patch(
target: impl AsValueRef,
patch: impl AsValueRef,
cache: &JsonCacheCell,
) -> crate::Result<Value> {
let (target, patch) = (target.as_value_ref(), patch.as_value_ref());
match (target, patch) {
(Value::Blob(_), _) | (_, Value::Blob(_)) => {
(ValueRef::Blob(_), _) | (_, ValueRef::Blob(_)) => {
crate::bail_constraint_error!("blob is not supported!");
}
_ => (),
}
let make_jsonb = curry_convert_dbtype_to_jsonb(Conv::Strict);
let mut target = cache.get_or_insert_with(target, &make_jsonb)?;
let patch = cache.get_or_insert_with(patch, &make_jsonb)?;
let mut target = cache.get_or_insert_with(target, make_jsonb)?;
let make_jsonb = curry_convert_dbtype_to_jsonb(Conv::Strict);
let patch = cache.get_or_insert_with(patch, make_jsonb)?;
target.patch(&patch)?;
@@ -30,16 +39,22 @@ pub fn json_patch(target: &Value, patch: &Value, cache: &JsonCacheCell) -> crate
json_string_to_db_type(target, element_type, OutputVariant::ElementType)
}
pub fn jsonb_patch(target: &Value, patch: &Value, cache: &JsonCacheCell) -> crate::Result<Value> {
pub fn jsonb_patch(
target: impl AsValueRef,
patch: impl AsValueRef,
cache: &JsonCacheCell,
) -> crate::Result<Value> {
let (target, patch) = (target.as_value_ref(), patch.as_value_ref());
match (target, patch) {
(Value::Blob(_), _) | (_, Value::Blob(_)) => {
(ValueRef::Blob(_), _) | (_, ValueRef::Blob(_)) => {
crate::bail_constraint_error!("blob is not supported!");
}
_ => (),
}
let make_jsonb = curry_convert_dbtype_to_jsonb(Conv::Strict);
let mut target = cache.get_or_insert_with(target, &make_jsonb)?;
let patch = cache.get_or_insert_with(patch, &make_jsonb)?;
let mut target = cache.get_or_insert_with(target, make_jsonb)?;
let make_jsonb = curry_convert_dbtype_to_jsonb(Conv::Strict);
let patch = cache.get_or_insert_with(patch, make_jsonb)?;
target.patch(&patch)?;
@@ -48,15 +63,21 @@ pub fn jsonb_patch(target: &Value, patch: &Value, cache: &JsonCacheCell) -> crat
json_string_to_db_type(target, element_type, OutputVariant::Binary)
}
pub fn json_remove(args: &[Register], json_cache: &JsonCacheCell) -> crate::Result<Value> {
if args.is_empty() {
pub fn json_remove<I, E, V>(args: I, json_cache: &JsonCacheCell) -> crate::Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut args = args.into_iter();
if args.len() == 0 {
return Ok(Value::Null);
}
let make_jsonb_fn = curry_convert_dbtype_to_jsonb(Conv::Strict);
let mut json = json_cache.get_or_insert_with(args[0].get_value(), make_jsonb_fn)?;
for arg in &args[1..] {
if let Some(path) = json_path_from_db_value(arg.get_value(), true)? {
let mut json = json_cache.get_or_insert_with(args.next().unwrap(), make_jsonb_fn)?;
for arg in args {
if let Some(path) = json_path_from_db_value(&arg, true)? {
let mut op = DeleteOperation::new();
let _ = json.operate_on_path(&path, &mut op);
}
@@ -67,15 +88,21 @@ pub fn json_remove(args: &[Register], json_cache: &JsonCacheCell) -> crate::Resu
json_string_to_db_type(json, el_type, OutputVariant::String)
}
pub fn jsonb_remove(args: &[Register], json_cache: &JsonCacheCell) -> crate::Result<Value> {
if args.is_empty() {
pub fn jsonb_remove<I, E, V>(args: I, json_cache: &JsonCacheCell) -> crate::Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut args = args.into_iter();
if args.len() == 0 {
return Ok(Value::Null);
}
let make_jsonb_fn = curry_convert_dbtype_to_jsonb(Conv::Strict);
let mut json = json_cache.get_or_insert_with(args[0].get_value(), make_jsonb_fn)?;
for arg in &args[1..] {
if let Some(path) = json_path_from_db_value(arg.get_value(), true)? {
let mut json = json_cache.get_or_insert_with(args.next().unwrap(), make_jsonb_fn)?;
for arg in args {
if let Some(path) = json_path_from_db_value(&arg, true)? {
let mut op = DeleteOperation::new();
let _ = json.operate_on_path(&path, &mut op);
}
@@ -84,18 +111,26 @@ pub fn jsonb_remove(args: &[Register], json_cache: &JsonCacheCell) -> crate::Res
Ok(Value::Blob(json.data()))
}
pub fn json_replace(args: &[Register], json_cache: &JsonCacheCell) -> crate::Result<Value> {
if args.is_empty() {
pub fn json_replace<I, E, V>(args: I, json_cache: &JsonCacheCell) -> crate::Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut args = args.into_iter();
if args.len() == 0 {
return Ok(Value::Null);
}
let make_jsonb_fn = curry_convert_dbtype_to_jsonb(Conv::Strict);
let mut json = json_cache.get_or_insert_with(args[0].get_value(), make_jsonb_fn)?;
let other = args[1..].chunks_exact(2);
for chunk in other {
let path = json_path_from_db_value(chunk[0].get_value(), true)?;
let mut json = json_cache.get_or_insert_with(args.next().unwrap(), make_jsonb_fn)?;
// TODO: when `array_chunks` is stabilized we can chunk by 2 here
while args.len() > 1 {
let first = args.next().unwrap();
let path = json_path_from_db_value(&first, true)?;
let value = convert_dbtype_to_jsonb(chunk[1].get_value(), Conv::NotStrict)?;
let second = args.next().unwrap();
let value = convert_dbtype_to_jsonb(&second, Conv::NotStrict)?;
if let Some(path) = path {
let mut op = ReplaceOperation::new(value);
@@ -108,17 +143,26 @@ pub fn json_replace(args: &[Register], json_cache: &JsonCacheCell) -> crate::Res
json_string_to_db_type(json, el_type, super::OutputVariant::String)
}
pub fn jsonb_replace(args: &[Register], json_cache: &JsonCacheCell) -> crate::Result<Value> {
if args.is_empty() {
pub fn jsonb_replace<I, E, V>(args: I, json_cache: &JsonCacheCell) -> crate::Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut args = args.into_iter();
if args.len() == 0 {
return Ok(Value::Null);
}
let make_jsonb_fn = curry_convert_dbtype_to_jsonb(Conv::Strict);
let mut json = json_cache.get_or_insert_with(args[0].get_value(), make_jsonb_fn)?;
let other = args[1..].chunks_exact(2);
for chunk in other {
let path = json_path_from_db_value(chunk[0].get_value(), true)?;
let value = convert_dbtype_to_jsonb(chunk[1].get_value(), Conv::NotStrict)?;
let mut json = json_cache.get_or_insert_with(args.next().unwrap(), make_jsonb_fn)?;
// TODO: when `array_chunks` is stabilized we can chunk by 2 here
while args.len() > 1 {
let first = args.next().unwrap();
let path = json_path_from_db_value(&first, true)?;
let second = args.next().unwrap();
let value = convert_dbtype_to_jsonb(&second, Conv::NotStrict)?;
if let Some(path) = path {
let mut op = ReplaceOperation::new(value);
@@ -131,17 +175,27 @@ pub fn jsonb_replace(args: &[Register], json_cache: &JsonCacheCell) -> crate::Re
json_string_to_db_type(json, el_type, OutputVariant::Binary)
}
pub fn json_insert(args: &[Register], json_cache: &JsonCacheCell) -> crate::Result<Value> {
if args.is_empty() {
pub fn json_insert<I, E, V>(args: I, json_cache: &JsonCacheCell) -> crate::Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut args = args.into_iter();
if args.len() == 0 {
return Ok(Value::Null);
}
let make_jsonb_fn = curry_convert_dbtype_to_jsonb(Conv::Strict);
let mut json = json_cache.get_or_insert_with(args[0].get_value(), make_jsonb_fn)?;
let other = args[1..].chunks_exact(2);
for chunk in other {
let path = json_path_from_db_value(chunk[0].get_value(), true)?;
let value = convert_dbtype_to_jsonb(chunk[1].get_value(), Conv::NotStrict)?;
let mut json = json_cache.get_or_insert_with(args.next().unwrap(), make_jsonb_fn)?;
// TODO: when `array_chunks` is stabilized we can chunk by 2 here
while args.len() > 1 {
let first = args.next().unwrap();
let path = json_path_from_db_value(&first, true)?;
let second = args.next().unwrap();
let value = convert_dbtype_to_jsonb(&second, Conv::NotStrict)?;
if let Some(path) = path {
let mut op = InsertOperation::new(value);
@@ -154,17 +208,27 @@ pub fn json_insert(args: &[Register], json_cache: &JsonCacheCell) -> crate::Resu
json_string_to_db_type(json, el_type, OutputVariant::String)
}
pub fn jsonb_insert(args: &[Register], json_cache: &JsonCacheCell) -> crate::Result<Value> {
if args.is_empty() {
pub fn jsonb_insert<I, E, V>(args: I, json_cache: &JsonCacheCell) -> crate::Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut args = args.into_iter();
if args.len() == 0 {
return Ok(Value::Null);
}
let make_jsonb_fn = curry_convert_dbtype_to_jsonb(Conv::Strict);
let mut json = json_cache.get_or_insert_with(args[0].get_value(), make_jsonb_fn)?;
let other = args[1..].chunks_exact(2);
for chunk in other {
let path = json_path_from_db_value(chunk[0].get_value(), true)?;
let value = convert_dbtype_to_jsonb(chunk[1].get_value(), Conv::NotStrict)?;
let mut json = json_cache.get_or_insert_with(args.next().unwrap(), make_jsonb_fn)?;
// TODO: when `array_chunks` is stabilized we can chunk by 2 here
while args.len() > 1 {
let first = args.next().unwrap();
let path = json_path_from_db_value(&first, true)?;
let second = args.next().unwrap();
let value = convert_dbtype_to_jsonb(&second, Conv::NotStrict)?;
if let Some(path) = path {
let mut op = InsertOperation::new(value);
@@ -282,17 +346,14 @@ mod tests {
#[test]
fn test_json_remove_empty_args() {
let args = vec![];
let args: [Value; 0] = [];
let json_cache = JsonCacheCell::new();
assert_eq!(json_remove(&args, &json_cache).unwrap(), Value::Null);
}
#[test]
fn test_json_remove_array_element() {
let args = vec![
Register::Value(create_json(r#"[1,2,3,4,5]"#)),
Register::Value(create_text("$[2]")),
];
let args = [create_json(r#"[1,2,3,4,5]"#), create_text("$[2]")];
let json_cache = JsonCacheCell::new();
let result = json_remove(&args, &json_cache).unwrap();
@@ -304,10 +365,10 @@ mod tests {
#[test]
fn test_json_remove_multiple_paths() {
let args = vec![
Register::Value(create_json(r#"{"a": 1, "b": 2, "c": 3}"#)),
Register::Value(create_text("$.a")),
Register::Value(create_text("$.c")),
let args = [
create_json(r#"{"a": 1, "b": 2, "c": 3}"#),
create_text("$.a"),
create_text("$.c"),
];
let json_cache = JsonCacheCell::new();
@@ -320,9 +381,9 @@ mod tests {
#[test]
fn test_json_remove_nested_paths() {
let args = vec![
Register::Value(create_json(r#"{"a": {"b": {"c": 1, "d": 2}}}"#)),
Register::Value(create_text("$.a.b.c")),
let args = [
create_json(r#"{"a": {"b": {"c": 1, "d": 2}}}"#),
create_text("$.a.b.c"),
];
let json_cache = JsonCacheCell::new();
@@ -335,9 +396,9 @@ mod tests {
#[test]
fn test_json_remove_duplicate_keys() {
let args = vec![
Register::Value(create_json(r#"{"a": 1, "a": 2, "a": 3}"#)),
Register::Value(create_text("$.a")),
let args = [
create_json(r#"{"a": 1, "a": 2, "a": 3}"#),
create_text("$.a"),
];
let json_cache = JsonCacheCell::new();
@@ -350,9 +411,9 @@ mod tests {
#[test]
fn test_json_remove_invalid_path() {
let args = vec![
Register::Value(create_json(r#"{"a": 1}"#)),
Register::Value(Value::Integer(42)), // Invalid path type
let args = [
create_json(r#"{"a": 1}"#),
Value::Integer(42), // Invalid path type
];
let json_cache = JsonCacheCell::new();
@@ -361,13 +422,11 @@ mod tests {
#[test]
fn test_json_remove_complex_case() {
let args = vec![
Register::Value(create_json(
r#"{"a":[1,2,3],"b":{"x":1,"x":2},"c":[{"y":1},{"y":2}]}"#,
)),
Register::Value(create_text("$.a[1]")),
Register::Value(create_text("$.b.x")),
Register::Value(create_text("$.c[0].y")),
let args = [
create_json(r#"{"a":[1,2,3],"b":{"x":1,"x":2},"c":[{"y":1},{"y":2}]}"#),
create_text("$.a[1]"),
create_text("$.b.x"),
create_text("$.c[0].y"),
];
let json_cache = JsonCacheCell::new();

6
core/mvcc/database/checkpoint_state_machine.rs
View File

@@ -450,11 +450,11 @@ impl<Clock: LogicalClock> CheckpointStateMachine<Clock> {
let record = ImmutableRecord::from_bin_record(row_version.row.data.clone());
let mut record_cursor = RecordCursor::new();
record_cursor.parse_full_header(&record).unwrap();
let values = record_cursor.get_values(&record)?;
let values = record_cursor.get_values(&record);
let mut values = values
.into_iter()
.map(|value| value.to_owned())
.collect::<Vec<_>>();
.map(|value| value.map(|v| v.to_owned()))
.collect::<Result<Vec<_>>>()?;
values[3] = Value::Integer(root_page as i64);
let record = ImmutableRecord::from_values(&values, values.len());
row_version.row.data = record.get_payload().to_owned();

10
core/mvcc/database/mod.rs
View File

@@ -1902,12 +1902,10 @@ impl<Clock: LogicalClock> MvStore<Clock> {
let row_data = row.data.clone();
let record = ImmutableRecord::from_bin_record(row_data);
let mut record_cursor = RecordCursor::new();
let record_values = record_cursor.get_values(&record).unwrap();
let ValueRef::Integer(root_page) = record_values[3] else {
panic!(
"Expected integer value for root page, got {:?}",
record_values[3]
);
let mut record_values = record_cursor.get_values(&record);
let val = record_values.nth(3).unwrap()?;
let ValueRef::Integer(root_page) = val else {
panic!("Expected integer value for root page, got {val:?}");
};
if root_page < 0 {
let table_id = self.get_table_id_from_root_page(root_page);

4
core/storage/btree.rs
View File

@@ -1618,7 +1618,7 @@ impl BTreeCursor {
.index_info
.as_ref()
.expect("indexbtree_move_to without index_info");
find_compare(&key_values, index_info)
find_compare(key_values.iter().peekable(), index_info)
};
tracing::debug!("Using record comparison strategy: {:?}", record_comparer);
let tie_breaker = get_tie_breaker_from_seek_op(cmp);
@@ -2002,7 +2002,7 @@ impl BTreeCursor {
.index_info
.as_ref()
.expect("indexbtree_seek without index_info");
find_compare(&key_values, index_info)
find_compare(key_values.iter().peekable(), index_info)
};
tracing::debug!(

21
core/translate/expr.rs
View File

@@ -4304,23 +4304,32 @@ pub fn emit_literal(
}
ast::Literal::CurrentDate => {
program.emit_insn(Insn::String8 {
value: datetime::exec_date(&[]).to_string(),
value: datetime::exec_date::<&[_; 0], std::slice::Iter<'_, Value>, &Value>(&[])
.to_string(),
dest: target_register,
});
Ok(target_register)
}
ast::Literal::CurrentTime => {
program.emit_insn(Insn::String8 {
value: datetime::exec_time(&[]).to_string(),
value: datetime::exec_time::<&[_; 0], std::slice::Iter<'_, Value>, &Value>(&[])
.to_string(),
dest: target_register,
});
Ok(target_register)
}
ast::Literal::CurrentTimestamp => {
program.emit_insn(Insn::String8 {
value: datetime::exec_datetime_full(&[]).to_string(),
dest: target_register,
});
program.emit_insn(
Insn::String8 {
value: datetime::exec_datetime_full::<
&[_; 0],
std::slice::Iter<'_, Value>,
&Value,
>(&[])
.to_string(),
dest: target_register,
},
);
Ok(target_register)
}
}

209
core/types.rs
View File

@@ -20,6 +20,7 @@ use crate::vtab::VirtualTableCursor;
use crate::{Completion, CompletionError, Result, IO};
use std::borrow::{Borrow, Cow};
use std::fmt::{Debug, Display};
use std::iter::Peekable;
use std::ops::Deref;
use std::task::Waker;
@@ -297,13 +298,6 @@ impl<'b> AsValueRef for ValueRef<'b> {
}
}
impl<'b> AsValueRef for &ValueRef<'b> {
#[inline]
fn as_value_ref<'a>(&'a self) -> ValueRef<'a> {
**self
}
}
impl AsValueRef for Value {
#[inline]
fn as_value_ref<'a>(&'a self) -> ValueRef<'a> {
@@ -311,13 +305,6 @@ impl AsValueRef for Value {
}
}
impl AsValueRef for &Value {
#[inline]
fn as_value_ref<'a>(&'a self) -> ValueRef<'a> {
self.as_ref()
}
}
impl AsValueRef for &mut Value {
#[inline]
fn as_value_ref<'a>(&'a self) -> ValueRef<'a> {
@@ -339,17 +326,9 @@ where
}
}
impl<V1, V2> AsValueRef for &Either<V1, V2>
where
V1: AsValueRef,
V2: AsValueRef,
{
#[inline]
impl<V: AsValueRef> AsValueRef for &V {
fn as_value_ref<'a>(&'a self) -> ValueRef<'a> {
match self {
Either::Left(left) => left.as_value_ref(),
Either::Right(right) => right.as_value_ref(),
}
(*self).as_value_ref()
}
}
@@ -1015,9 +994,14 @@ impl ImmutableRecord {
// TODO: inline the complete record parsing code here.
// Its probably more efficient.
// fixme(pedrocarlo): this function is very inneficient and kind of misleading because
// it always deserializes the columns
pub fn get_values<'a>(&'a self) -> Vec<ValueRef<'a>> {
let mut cursor = RecordCursor::new();
cursor.get_values(self).unwrap_or_default()
cursor
.get_values(self)
.collect::<Result<Vec<_>>>()
.unwrap_or_default()
}
pub fn from_registers<'a, I: Iterator<Item = &'a Register> + Clone>(
@@ -1032,7 +1016,7 @@ impl ImmutableRecord {
}
pub fn from_values<'a>(
values: impl IntoIterator<Item = &'a Value> + Clone,
values: impl IntoIterator<Item = impl AsValueRef + 'a> + Clone,
len: usize,
) -> Self {
let mut serials = Vec::with_capacity(len);
@@ -1042,7 +1026,7 @@ impl ImmutableRecord {
let mut serial_type_buf = [0; 9];
// write serial types
for value in values.clone() {
let serial_type = SerialType::from(value);
let serial_type = SerialType::from(value.as_value_ref());
let n = write_varint(&mut serial_type_buf[0..], serial_type.into());
serials.push((serial_type_buf, n));
@@ -1071,28 +1055,29 @@ impl ImmutableRecord {
// write content
for value in values {
let value = value.as_value_ref();
match value {
Value::Null => {}
Value::Integer(i) => {
ValueRef::Null => {}
ValueRef::Integer(i) => {
let serial_type = SerialType::from(value);
match serial_type.kind() {
SerialTypeKind::ConstInt0 | SerialTypeKind::ConstInt1 => {}
SerialTypeKind::I8 => writer.extend_from_slice(&(*i as i8).to_be_bytes()),
SerialTypeKind::I16 => writer.extend_from_slice(&(*i as i16).to_be_bytes()),
SerialTypeKind::I8 => writer.extend_from_slice(&(i as i8).to_be_bytes()),
SerialTypeKind::I16 => writer.extend_from_slice(&(i as i16).to_be_bytes()),
SerialTypeKind::I24 => {
writer.extend_from_slice(&(*i as i32).to_be_bytes()[1..])
writer.extend_from_slice(&(i as i32).to_be_bytes()[1..])
} // remove most significant byte
SerialTypeKind::I32 => writer.extend_from_slice(&(*i as i32).to_be_bytes()),
SerialTypeKind::I32 => writer.extend_from_slice(&(i as i32).to_be_bytes()),
SerialTypeKind::I48 => writer.extend_from_slice(&i.to_be_bytes()[2..]), // remove 2 most significant bytes
SerialTypeKind::I64 => writer.extend_from_slice(&i.to_be_bytes()),
other => panic!("Serial type is not an integer: {other:?}"),
}
}
Value::Float(f) => writer.extend_from_slice(&f.to_be_bytes()),
Value::Text(t) => {
ValueRef::Float(f) => writer.extend_from_slice(&f.to_be_bytes()),
ValueRef::Text(t) => {
writer.extend_from_slice(t.value.as_bytes());
}
Value::Blob(b) => {
ValueRef::Blob(b) => {
writer.extend_from_slice(b);
}
};
@@ -1462,19 +1447,48 @@ impl RecordCursor {
/// * `Ok(Vec<RefValue>)` - All values in column order
/// * `Err(LimboError)` - Parsing or deserialization failed
///
pub fn get_values<'a>(&mut self, record: &'a ImmutableRecord) -> Result<Vec<ValueRef<'a>>> {
if record.is_invalidated() {
return Ok(Vec::new());
pub fn get_values<'a, 'b>(
&'b mut self,
record: &'a ImmutableRecord,
) -> Peekable<impl ExactSizeIterator<Item = Result<ValueRef<'a>>> + use<'a, 'b>> {
struct GetValues<'a, 'b> {
cursor: &'b mut RecordCursor,
record: &'a ImmutableRecord,
idx: usize,
}
self.parse_full_header(record)?;
let mut result = Vec::with_capacity(self.serial_types.len());
impl<'a, 'b> Iterator for GetValues<'a, 'b> {
type Item = Result<ValueRef<'a>>;
for i in 0..self.serial_types.len() {
result.push(self.deserialize_column(record, i)?);
fn next(&mut self) -> Option<Self::Item> {
if self.idx == 0 {
// So that we can have the full length of serial types
if let Err(err) = self.cursor.parse_full_header(self.record) {
return Some(Err(err));
}
}
if !self.record.is_invalidated() && self.idx < self.cursor.serial_types.len() {
let res = self.cursor.deserialize_column(self.record, self.idx);
self.idx += 1;
Some(res)
} else {
None
}
}
}
Ok(result)
impl<'a, 'b> ExactSizeIterator for GetValues<'a, 'b> {
fn len(&self) -> usize {
self.cursor.serial_types.len() - self.idx
}
}
let get_values = GetValues {
cursor: self,
record,
idx: 0,
};
get_values.peekable()
}
}
@@ -1496,6 +1510,42 @@ impl<'a> ValueRef<'a> {
}
}
pub fn to_text(&self) -> Option<&'a str> {
match self {
Self::Text(t) => Some(t.as_str()),
_ => None,
}
}
pub fn as_blob(&self) -> &'a [u8] {
match self {
Self::Blob(b) => b,
_ => panic!("as_blob must be called only for Value::Blob"),
}
}
pub fn as_float(&self) -> f64 {
match self {
Self::Float(f) => *f,
Self::Integer(i) => *i as f64,
_ => panic!("as_float must be called only for Value::Float or Value::Integer"),
}
}
pub fn as_int(&self) -> Option<i64> {
match self {
Self::Integer(i) => Some(*i),
_ => None,
}
}
pub fn as_uint(&self) -> u64 {
match self {
Self::Integer(i) => (*i).cast_unsigned(),
_ => 0,
}
}
pub fn to_owned(&self) -> Value {
match self {
ValueRef::Null => Value::Null,
@@ -1508,6 +1558,16 @@ impl<'a> ValueRef<'a> {
ValueRef::Blob(b) => Value::Blob(b.to_vec()),
}
}
pub fn value_type(&self) -> ValueType {
match self {
Self::Null => ValueType::Null,
Self::Integer(_) => ValueType::Integer,
Self::Float(_) => ValueType::Float,
Self::Text(_) => ValueType::Text,
Self::Blob(_) => ValueType::Blob,
}
}
}
impl Display for ValueRef<'_> {
@@ -1542,6 +1602,13 @@ impl<'a> PartialEq<ValueRef<'a>> for ValueRef<'a> {
}
}
impl<'a> PartialEq<Value> for ValueRef<'a> {
fn eq(&self, other: &Value) -> bool {
let other = other.as_value_ref();
self.eq(&other)
}
}
impl<'a> Eq for ValueRef<'a> {}
#[expect(clippy::non_canonical_partial_ord_impl)]
@@ -1759,9 +1826,16 @@ impl RecordCompare {
}
}
pub fn find_compare(unpacked: &[ValueRef], index_info: &IndexInfo) -> RecordCompare {
if !unpacked.is_empty() && index_info.num_cols <= 13 {
match &unpacked[0] {
pub fn find_compare<I, E, V>(unpacked: I, index_info: &IndexInfo) -> RecordCompare
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = Peekable<E>, Item = V>,
{
let mut unpacked = unpacked.into_iter();
if unpacked.len() != 0 && index_info.num_cols <= 13 {
let val = unpacked.peek().unwrap();
match val.as_value_ref() {
ValueRef::Integer(_) => RecordCompare::Int,
ValueRef::Text(_) if index_info.key_info[0].collation == CollationSeq::Binary => {
RecordCompare::String
@@ -2266,23 +2340,24 @@ impl SerialType {
}
}
impl From<&Value> for SerialType {
fn from(value: &Value) -> Self {
impl<T: AsValueRef> From<T> for SerialType {
fn from(value: T) -> Self {
let value = value.as_value_ref();
match value {
Value::Null => SerialType::null(),
Value::Integer(i) => match i {
ValueRef::Null => SerialType::null(),
ValueRef::Integer(i) => match i {
0 => SerialType::const_int0(),
1 => SerialType::const_int1(),
i if *i >= I8_LOW && *i <= I8_HIGH => SerialType::i8(),
i if *i >= I16_LOW && *i <= I16_HIGH => SerialType::i16(),
i if *i >= I24_LOW && *i <= I24_HIGH => SerialType::i24(),
i if *i >= I32_LOW && *i <= I32_HIGH => SerialType::i32(),
i if *i >= I48_LOW && *i <= I48_HIGH => SerialType::i48(),
i if (I8_LOW..=I8_HIGH).contains(&i) => SerialType::i8(),
i if (I16_LOW..=I16_HIGH).contains(&i) => SerialType::i16(),
i if (I24_LOW..=I24_HIGH).contains(&i) => SerialType::i24(),
i if (I32_LOW..=I32_HIGH).contains(&i) => SerialType::i32(),
i if (I48_LOW..=I48_HIGH).contains(&i) => SerialType::i48(),
_ => SerialType::i64(),
},
Value::Float(_) => SerialType::f64(),
Value::Text(t) => SerialType::text(t.value.len() as u64),
Value::Blob(b) => SerialType::blob(b.len() as u64),
ValueRef::Float(_) => SerialType::f64(),
ValueRef::Text(t) => SerialType::text(t.value.len() as u64),
ValueRef::Blob(b) => SerialType::blob(b.len() as u64),
}
}
}
@@ -2754,7 +2829,7 @@ mod tests {
tie_breaker,
);
let comparer = find_compare(&unpacked_values, index_info);
let comparer = find_compare(unpacked_values.iter().peekable(), index_info);
let optimized_result = comparer
.compare(&serialized, &unpacked_values, index_info, 0, tie_breaker)
.unwrap();
@@ -3196,33 +3271,33 @@ mod tests {
);
let index_info_large = create_index_info(15, vec![SortOrder::Asc; 15], collations_large);
let int_values = vec![
let int_values = [
ValueRef::Integer(42),
ValueRef::Text(TextRef::new("hello", TextSubtype::Text)),
];
assert!(matches!(
find_compare(&int_values, &index_info_small),
find_compare(int_values.iter().peekable(), &index_info_small),
RecordCompare::Int
));
let string_values = vec![
let string_values = [
ValueRef::Text(TextRef::new("hello", TextSubtype::Text)),
ValueRef::Integer(42),
];
assert!(matches!(
find_compare(&string_values, &index_info_small),
find_compare(string_values.iter().peekable(), &index_info_small),
RecordCompare::String
));
let large_values: Vec<ValueRef> = (0..15).map(ValueRef::Integer).collect();
assert!(matches!(
find_compare(&large_values, &index_info_large),
find_compare(large_values.iter().peekable(), &index_info_large),
RecordCompare::Generic
));
let blob_values = vec![ValueRef::Blob(&[1, 2, 3])];
let blob_values = [ValueRef::Blob(&[1, 2, 3])];
assert!(matches!(
find_compare(&blob_values, &index_info_small),
find_compare(blob_values.iter().peekable(), &index_info_small),
RecordCompare::Generic
));
}

147
core/vdbe/execute.rs
View File

@@ -4442,7 +4442,8 @@ pub fn op_function(
| JsonFunc::JsonObject
| JsonFunc::JsonbArray
| JsonFunc::JsonbObject => {
let reg_values = &state.registers[*start_reg..*start_reg + arg_count];
let reg_values =
registers_to_ref_values(&state.registers[*start_reg..*start_reg + arg_count]);
let json_func = match json_func {
JsonFunc::JsonArray => json_array,
@@ -4463,7 +4464,9 @@ pub fn op_function(
0 => Ok(Value::Null),
_ => {
let val = &state.registers[*start_reg];
let reg_values = &state.registers[*start_reg + 1..*start_reg + arg_count];
let reg_values = registers_to_ref_values(
&state.registers[*start_reg + 1..*start_reg + arg_count],
);
json_extract(val.get_value(), reg_values, &state.json_cache)
}
@@ -4479,7 +4482,9 @@ pub fn op_function(
0 => Ok(Value::Null),
_ => {
let val = &state.registers[*start_reg];
let reg_values = &state.registers[*start_reg + 1..*start_reg + arg_count];
let reg_values = registers_to_ref_values(
&state.registers[*start_reg + 1..*start_reg + arg_count],
);
jsonb_extract(val.get_value(), reg_values, &state.json_cache)
}
@@ -4565,7 +4570,7 @@ pub fn op_function(
}
JsonFunc::JsonRemove => {
if let Ok(json) = json_remove(
&state.registers[*start_reg..*start_reg + arg_count],
registers_to_ref_values(&state.registers[*start_reg..*start_reg + arg_count]),
&state.json_cache,
) {
state.registers[*dest] = Register::Value(json);
@@ -4575,7 +4580,7 @@ pub fn op_function(
}
JsonFunc::JsonbRemove => {
if let Ok(json) = jsonb_remove(
&state.registers[*start_reg..*start_reg + arg_count],
registers_to_ref_values(&state.registers[*start_reg..*start_reg + arg_count]),
&state.json_cache,
) {
state.registers[*dest] = Register::Value(json);
@@ -4585,7 +4590,7 @@ pub fn op_function(
}
JsonFunc::JsonReplace => {
if let Ok(json) = json_replace(
&state.registers[*start_reg..*start_reg + arg_count],
registers_to_ref_values(&state.registers[*start_reg..*start_reg + arg_count]),
&state.json_cache,
) {
state.registers[*dest] = Register::Value(json);
@@ -4595,7 +4600,7 @@ pub fn op_function(
}
JsonFunc::JsonbReplace => {
if let Ok(json) = jsonb_replace(
&state.registers[*start_reg..*start_reg + arg_count],
registers_to_ref_values(&state.registers[*start_reg..*start_reg + arg_count]),
&state.json_cache,
) {
state.registers[*dest] = Register::Value(json);
@@ -4605,7 +4610,7 @@ pub fn op_function(
}
JsonFunc::JsonInsert => {
if let Ok(json) = json_insert(
&state.registers[*start_reg..*start_reg + arg_count],
registers_to_ref_values(&state.registers[*start_reg..*start_reg + arg_count]),
&state.json_cache,
) {
state.registers[*dest] = Register::Value(json);
@@ -4615,7 +4620,7 @@ pub fn op_function(
}
JsonFunc::JsonbInsert => {
if let Ok(json) = jsonb_insert(
&state.registers[*start_reg..*start_reg + arg_count],
registers_to_ref_values(&state.registers[*start_reg..*start_reg + arg_count]),
&state.json_cache,
) {
state.registers[*dest] = Register::Value(json);
@@ -4654,7 +4659,8 @@ pub fn op_function(
if arg_count % 2 == 0 {
bail_constraint_error!("json_set() needs an odd number of arguments")
}
let reg_values = &state.registers[*start_reg..*start_reg + arg_count];
let reg_values =
registers_to_ref_values(&state.registers[*start_reg..*start_reg + arg_count]);
let json_result = json_set(reg_values, &state.json_cache);
@@ -4667,7 +4673,8 @@ pub fn op_function(
if arg_count % 2 == 0 {
bail_constraint_error!("json_set() needs an odd number of arguments")
}
let reg_values = &state.registers[*start_reg..*start_reg + arg_count];
let reg_values =
registers_to_ref_values(&state.registers[*start_reg..*start_reg + arg_count]);
let json_result = jsonb_set(reg_values, &state.json_cache);
@@ -4942,11 +4949,14 @@ pub fn op_function(
state.registers[*dest] = Register::Value(result);
}
ScalarFunc::Date => {
let result = exec_date(&state.registers[*start_reg..*start_reg + arg_count]);
let values =
registers_to_ref_values(&state.registers[*start_reg..*start_reg + arg_count]);
let result = exec_date(values);
state.registers[*dest] = Register::Value(result);
}
ScalarFunc::Time => {
let values = &state.registers[*start_reg..*start_reg + arg_count];
let values =
registers_to_ref_values(&state.registers[*start_reg..*start_reg + arg_count]);
let result = exec_time(values);
state.registers[*dest] = Register::Value(result);
}
@@ -4954,13 +4964,10 @@ pub fn op_function(
if arg_count != 2 {
state.registers[*dest] = Register::Value(Value::Null);
} else {
let start = state.registers[*start_reg].get_value().clone();
let end = state.registers[*start_reg + 1].get_value().clone();
let start = state.registers[*start_reg].get_value();
let end = state.registers[*start_reg + 1].get_value();
let result = crate::functions::datetime::exec_timediff(&[
Register::Value(start),
Register::Value(end),
]);
let result = crate::functions::datetime::exec_timediff([start, end]);
state.registers[*dest] = Register::Value(result);
}
@@ -4971,12 +4978,15 @@ pub fn op_function(
state.registers[*dest] = Register::Value(Value::Integer(total_changes));
}
ScalarFunc::DateTime => {
let result =
exec_datetime_full(&state.registers[*start_reg..*start_reg + arg_count]);
let values =
registers_to_ref_values(&state.registers[*start_reg..*start_reg + arg_count]);
let result = exec_datetime_full(values);
state.registers[*dest] = Register::Value(result);
}
ScalarFunc::JulianDay => {
let result = exec_julianday(&state.registers[*start_reg..*start_reg + arg_count]);
let values =
registers_to_ref_values(&state.registers[*start_reg..*start_reg + arg_count]);
let result = exec_julianday(values);
state.registers[*dest] = Register::Value(result);
}
ScalarFunc::UnixEpoch => {
@@ -5042,7 +5052,9 @@ pub fn op_function(
program.connection.load_extension(ext)?;
}
ScalarFunc::StrfTime => {
let result = exec_strftime(&state.registers[*start_reg..*start_reg + arg_count]);
let values =
registers_to_ref_values(&state.registers[*start_reg..*start_reg + arg_count]);
let result = exec_strftime(values);
state.registers[*dest] = Register::Value(result);
}
ScalarFunc::Printf => {
@@ -5219,51 +5231,52 @@ pub fn op_function(
);
}
},
crate::function::Func::Vector(vector_func) => match vector_func {
VectorFunc::Vector => {
let result = vector32(&state.registers[*start_reg..*start_reg + arg_count])?;
state.registers[*dest] = Register::Value(result);
crate::function::Func::Vector(vector_func) => {
let values =
registers_to_ref_values(&state.registers[*start_reg..*start_reg + arg_count]);
match vector_func {
VectorFunc::Vector => {
let result = vector32(values)?;
state.registers[*dest] = Register::Value(result);
}
VectorFunc::Vector32 => {
let result = vector32(values)?;
state.registers[*dest] = Register::Value(result);
}
VectorFunc::Vector32Sparse => {
let result = vector32_sparse(values)?;
state.registers[*dest] = Register::Value(result);
}
VectorFunc::Vector64 => {
let result = vector64(values)?;
state.registers[*dest] = Register::Value(result);
}
VectorFunc::VectorExtract => {
let result = vector_extract(values)?;
state.registers[*dest] = Register::Value(result);
}
VectorFunc::VectorDistanceCos => {
let result = vector_distance_cos(values)?;
state.registers[*dest] = Register::Value(result);
}
VectorFunc::VectorDistanceL2 => {
let result = vector_distance_l2(values)?;
state.registers[*dest] = Register::Value(result);
}
VectorFunc::VectorDistanceJaccard => {
let result = vector_distance_jaccard(values)?;
state.registers[*dest] = Register::Value(result);
}
VectorFunc::VectorConcat => {
let result = vector_concat(values)?;
state.registers[*dest] = Register::Value(result);
}
VectorFunc::VectorSlice => {
let result = vector_slice(values)?;
state.registers[*dest] = Register::Value(result)
}
}
VectorFunc::Vector32 => {
let result = vector32(&state.registers[*start_reg..*start_reg + arg_count])?;
state.registers[*dest] = Register::Value(result);
}
VectorFunc::Vector32Sparse => {
let result = vector32_sparse(&state.registers[*start_reg..*start_reg + arg_count])?;
state.registers[*dest] = Register::Value(result);
}
VectorFunc::Vector64 => {
let result = vector64(&state.registers[*start_reg..*start_reg + arg_count])?;
state.registers[*dest] = Register::Value(result);
}
VectorFunc::VectorExtract => {
let result = vector_extract(&state.registers[*start_reg..*start_reg + arg_count])?;
state.registers[*dest] = Register::Value(result);
}
VectorFunc::VectorDistanceCos => {
let result =
vector_distance_cos(&state.registers[*start_reg..*start_reg + arg_count])?;
state.registers[*dest] = Register::Value(result);
}
VectorFunc::VectorDistanceL2 => {
let result =
vector_distance_l2(&state.registers[*start_reg..*start_reg + arg_count])?;
state.registers[*dest] = Register::Value(result);
}
VectorFunc::VectorDistanceJaccard => {
let result =
vector_distance_jaccard(&state.registers[*start_reg..*start_reg + arg_count])?;
state.registers[*dest] = Register::Value(result);
}
VectorFunc::VectorConcat => {
let result = vector_concat(&state.registers[*start_reg..*start_reg + arg_count])?;
state.registers[*dest] = Register::Value(result);
}
VectorFunc::VectorSlice => {
let result = vector_slice(&state.registers[*start_reg..*start_reg + arg_count])?;
state.registers[*dest] = Register::Value(result)
}
},
}
crate::function::Func::External(f) => match f.func {
ExtFunc::Scalar(f) => {
if arg_count == 0 {

142
core/vector/mod.rs
View File

@@ -1,21 +1,26 @@
use crate::types::AsValueRef;
use crate::types::Value;
use crate::types::ValueType;
use crate::vdbe::Register;
use crate::LimboError;
use crate::Result;
use crate::ValueRef;
pub mod operations;
pub mod vector_types;
use vector_types::*;
pub fn parse_vector(value: &Register, type_hint: Option<VectorType>) -> Result<Vector> {
match value.get_value().value_type() {
pub fn parse_vector<'a>(
value: &'a (impl AsValueRef + 'a),
type_hint: Option<VectorType>,
) -> Result<Vector<'a>> {
let value = value.as_value_ref();
match value.value_type() {
ValueType::Text => operations::text::vector_from_text(
type_hint.unwrap_or(VectorType::Float32Dense),
value.get_value().to_text().expect("value must be text"),
value.to_text().expect("value must be text"),
),
ValueType::Blob => {
let Some(blob) = value.get_value().to_blob() else {
let Some(blob) = value.to_blob() else {
return Err(LimboError::ConversionError(
"Invalid vector value".to_string(),
));
@@ -28,48 +33,77 @@ pub fn parse_vector(value: &Register, type_hint: Option<VectorType>) -> Result<V
}
}
pub fn vector32(args: &[Register]) -> Result<Value> {
pub fn vector32<I, E, V>(args: I) -> Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut args = args.into_iter();
if args.len() != 1 {
return Err(LimboError::ConversionError(
"vector32 requires exactly one argument".to_string(),
));
}
let vector = parse_vector(&args[0], Some(VectorType::Float32Dense))?;
let value = args.next().unwrap();
let vector = parse_vector(&value, Some(VectorType::Float32Dense))?;
let vector = operations::convert::vector_convert(vector, VectorType::Float32Dense)?;
Ok(operations::serialize::vector_serialize(vector))
}
pub fn vector32_sparse(args: &[Register]) -> Result<Value> {
pub fn vector32_sparse<I, E, V>(args: I) -> Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut args = args.into_iter();
if args.len() != 1 {
return Err(LimboError::ConversionError(
"vector32_sparse requires exactly one argument".to_string(),
));
}
let vector = parse_vector(&args[0], Some(VectorType::Float32Sparse))?;
let value = args.next().unwrap();
let vector = parse_vector(&value, Some(VectorType::Float32Sparse))?;
let vector = operations::convert::vector_convert(vector, VectorType::Float32Sparse)?;
Ok(operations::serialize::vector_serialize(vector))
}
pub fn vector64(args: &[Register]) -> Result<Value> {
pub fn vector64<I, E, V>(args: I) -> Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut args = args.into_iter();
if args.len() != 1 {
return Err(LimboError::ConversionError(
"vector64 requires exactly one argument".to_string(),
));
}
let vector = parse_vector(&args[0], Some(VectorType::Float64Dense))?;
let value = args.next().unwrap();
let vector = parse_vector(&value, Some(VectorType::Float64Dense))?;
let vector = operations::convert::vector_convert(vector, VectorType::Float64Dense)?;
Ok(operations::serialize::vector_serialize(vector))
}
pub fn vector_extract(args: &[Register]) -> Result<Value> {
pub fn vector_extract<I, E, V>(args: I) -> Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut args = args.into_iter();
if args.len() != 1 {
return Err(LimboError::ConversionError(
"vector_extract requires exactly one argument".to_string(),
));
}
let blob = match &args[0].get_value() {
Value::Blob(b) => b,
let value = args.next().unwrap();
let value = value.as_value_ref();
let blob = match value {
ValueRef::Blob(b) => b,
_ => {
return Err(LimboError::ConversionError(
"Expected blob value".to_string(),
@@ -81,78 +115,120 @@ pub fn vector_extract(args: &[Register]) -> Result<Value> {
return Ok(Value::build_text("[]"));
}
let vector = Vector::from_vec(blob.to_vec())?;
let vector = Vector::from_slice(blob)?;
Ok(Value::build_text(operations::text::vector_to_text(&vector)))
}
pub fn vector_distance_cos(args: &[Register]) -> Result<Value> {
pub fn vector_distance_cos<I, E, V>(args: I) -> Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut args = args.into_iter();
if args.len() != 2 {
return Err(LimboError::ConversionError(
"vector_distance_cos requires exactly two arguments".to_string(),
));
}
let x = parse_vector(&args[0], None)?;
let y = parse_vector(&args[1], None)?;
let value_0 = args.next().unwrap();
let value_1 = args.next().unwrap();
let x = parse_vector(&value_0, None)?;
let y = parse_vector(&value_1, None)?;
let dist = operations::distance_cos::vector_distance_cos(&x, &y)?;
Ok(Value::Float(dist))
}
pub fn vector_distance_l2(args: &[Register]) -> Result<Value> {
pub fn vector_distance_l2<I, E, V>(args: I) -> Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut args = args.into_iter();
if args.len() != 2 {
return Err(LimboError::ConversionError(
"distance_l2 requires exactly two arguments".to_string(),
));
}
let x = parse_vector(&args[0], None)?;
let y = parse_vector(&args[1], None)?;
let value_0 = args.next().unwrap();
let value_1 = args.next().unwrap();
let x = parse_vector(&value_0, None)?;
let y = parse_vector(&value_1, None)?;
let dist = operations::distance_l2::vector_distance_l2(&x, &y)?;
Ok(Value::Float(dist))
}
pub fn vector_distance_jaccard(args: &[Register]) -> Result<Value> {
pub fn vector_distance_jaccard<I, E, V>(args: I) -> Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut args = args.into_iter();
if args.len() != 2 {
return Err(LimboError::ConversionError(
"distance_jaccard requires exactly two arguments".to_string(),
));
}
let x = parse_vector(&args[0], None)?;
let y = parse_vector(&args[1], None)?;
let value_0 = args.next().unwrap();
let value_1 = args.next().unwrap();
let x = parse_vector(&value_0, None)?;
let y = parse_vector(&value_1, None)?;
let dist = operations::jaccard::vector_distance_jaccard(&x, &y)?;
Ok(Value::Float(dist))
}
pub fn vector_concat(args: &[Register]) -> Result<Value> {
pub fn vector_concat<I, E, V>(args: I) -> Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut args = args.into_iter();
if args.len() != 2 {
return Err(LimboError::InvalidArgument(
"concat requires exactly two arguments".into(),
));
}
let x = parse_vector(&args[0], None)?;
let y = parse_vector(&args[1], None)?;
let value_0 = args.next().unwrap();
let value_1 = args.next().unwrap();
let x = parse_vector(&value_0, None)?;
let y = parse_vector(&value_1, None)?;
let vector = operations::concat::vector_concat(&x, &y)?;
Ok(operations::serialize::vector_serialize(vector))
}
pub fn vector_slice(args: &[Register]) -> Result<Value> {
pub fn vector_slice<I, E, V>(args: I) -> Result<Value>
where
V: AsValueRef,
E: ExactSizeIterator<Item = V>,
I: IntoIterator<IntoIter = E, Item = V>,
{
let mut args = args.into_iter();
if args.len() != 3 {
return Err(LimboError::InvalidArgument(
"vector_slice requires exactly three arguments".into(),
));
}
let value_0 = args.next().unwrap();
let value_1 = args.next().unwrap();
let value_1 = value_1.as_value_ref();
let vector = parse_vector(&args[0], None)?;
let value_2 = args.next().unwrap();
let value_2 = value_2.as_value_ref();
let start_index = args[1]
.get_value()
let vector = parse_vector(&value_0, None)?;
let start_index = value_1
.as_int()
.ok_or_else(|| LimboError::InvalidArgument("start index must be an integer".into()))?;
let end_index = args[2]
.get_value()
let end_index = value_2
.as_int()
.ok_or_else(|| LimboError::InvalidArgument("end_index must be an integer".into()))?;