Skip to content

Commit 4f0b63a

Browse files
michaelmiormichaelmior
committed
refactor(query): replace serde_json::Value with jsonb_schema::Value
1 parent 2d469c2 commit 4f0b63a

File tree

11 files changed

+496
-241
lines changed

11 files changed

+496
-241
lines changed

‎benches/db.rs‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ fn generate_query_plan(collection_name: &str, selectivity: f64, total_docs: usiz
2828
let predicate = Expression::Binary {
2929
left: Box::new(Expression::FieldReference("value".to_string())),
3030
op: BinaryOperator::Lt,
31-
right: Box::new(Expression::Literal(json!(filter_value))),
31+
right: Box::new(Expression::Literal(json!(filter_value).into())),
3232
};
3333

3434
LogicalPlan::Filter {
@@ -58,7 +58,7 @@ fn insertion_benchmark(c: &mut Criterion) {
5858
// Benchmarked routine: Insert documents
5959
for i in 0..(max_docs / num_keys) {
6060
let doc = generate_doc_with_keys(*num_keys, i);
61-
db.insert("test", hint::black_box(doc)).unwrap();
61+
db.insert("test", hint::black_box(doc.into())).unwrap();
6262
}
6363

6464
total_duration += start.elapsed();
@@ -95,7 +95,8 @@ fn query_benchmark(c: &mut Criterion) {
9595
let mut db = DB::new(dir.path().to_str().unwrap(), num_docs + 1, 10, 1024, None); // Don't flush
9696
db.create_collection(collection_name).unwrap();
9797
for i in 0..num_docs {
98-
db.insert(collection_name, json!({"value": i})).unwrap();
98+
db.insert(collection_name, json!({"value": i}).into())
99+
.unwrap();
99100
}
100101
let db_arc = std::sync::Arc::new(std::sync::Mutex::new(db));
101102

‎benches/logging.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ fn logging_benchmark(c: &mut Criterion) {
1515
for _ in 0..iters {
1616
let op = Operation::Insert {
1717
id: "test_doc_id".to_string(),
18-
doc: serde_json::json!({"key": "value"}),
18+
doc: serde_json::json!({"key": "value"}).into(),
1919
};
2020

2121
logger.log(op).unwrap();

‎src/bin/argusdb.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ use tokio::net::TcpListener;
1616
use tokio::sync::Mutex;
1717
use tracing::{Level, info, span};
1818

19+
use argusdb::Value;
1920
use argusdb::db::DB;
2021
use argusdb::parser as argus_parser;
2122
use argusdb::query::{Statement, execute_plan};
@@ -151,7 +152,7 @@ impl SimpleQueryHandler for ArgusHandler {
151152
let obj = doc.as_object().unwrap();
152153
for field in fields.iter() {
153154
let key = field.name();
154-
let val = obj.get(key).unwrap_or(&serde_json::Value::Null);
155+
let val = obj.get(key).unwrap_or(&Value::Null);
155156
encoder
156157
.encode_field(&val.to_string())
157158
.map_err(|e| PgWireError::ApiError(Box::new(e)))?;

‎src/db.rs‎

Lines changed: 76 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1+
use crate::Value;
12
use crate::jstable;
23
use crate::log::{Log, LogEntry, Logger, NullLogger, Operation};
34
use crate::storage::MemTable;
4-
use serde_json::Value;
55
use std::collections::HashMap;
66
use std::fmt::Debug;
77
use std::fs;
@@ -59,10 +59,12 @@ impl<'a> Iterator for MergedIterator<'a> {
5959
}
6060
}
6161

62-
if let Some(doc) = result
63-
&& !doc.is_null()
64-
{
65-
return Some((min_id, doc));
62+
// Check for tombstone (Null)
63+
if let Some(doc) = result {
64+
use jsonb_schema::Value as JsonbValue;
65+
if !matches!(doc, JsonbValue::Null) {
66+
return Some((min_id, doc));
67+
}
6668
}
6769
// If null (tombstone), loop again
6870
}
@@ -266,7 +268,8 @@ impl Collection {
266268
fn get(&self, id: &str) -> Option<Value> {
267269
// 1. Check MemTable
268270
if let Some(doc) = self.memtable.documents.get(id) {
269-
if doc.is_null() {
271+
use jsonb_schema::Value as JsonbValue;
272+
if matches!(doc, JsonbValue::Null) {
270273
return None; // Tombstone
271274
}
272275
return Some(doc.clone());
@@ -281,29 +284,30 @@ impl Collection {
281284
};
282285

283286
for i in (0..self.jstable_count).rev() {
284-
if let Some(table) = self.tables.get(i as usize) {
285-
if table.filter.contains(&hash) {
286-
// Possible match, find offset using index
287-
let index = &table.index;
288-
// Find first key > id. We want the one before that.
289-
let idx = index.partition_point(|(k, _)| k.as_str() <= id);
290-
let start_offset = if idx > 0 { index[idx - 1].1 } else { 0 };
291-
292-
let path = self.dir.join(format!("jstable-{}", i));
293-
if let Ok(mut iter) = jstable::JSTableIterator::new(path.to_str().unwrap()) {
294-
if iter.seek(start_offset).is_ok() {
295-
for (rid, doc) in iter.flatten() {
296-
if rid == id {
297-
if doc.is_null() {
298-
return None; // Tombstone
299-
}
300-
return Some(doc);
301-
}
302-
if rid.as_str() > id {
303-
// Not found in this table (sorted)
304-
break;
305-
}
287+
if let Some(table) = self.tables.get(i as usize)
288+
&& table.filter.contains(&hash)
289+
{
290+
// Possible match, find offset using index
291+
let index = &table.index;
292+
// Find first key > id. We want the one before that.
293+
let idx = index.partition_point(|(k, _)| k.as_str() <= id);
294+
let start_offset = if idx > 0 { index[idx - 1].1 } else { 0 };
295+
296+
let path = self.dir.join(format!("jstable-{}", i));
297+
if let Ok(mut iter) = jstable::JSTableIterator::new(path.to_str().unwrap())
298+
&& iter.seek(start_offset).is_ok()
299+
{
300+
for (rid, doc) in iter.flatten() {
301+
if rid == id {
302+
use jsonb_schema::Value as JsonbValue;
303+
if matches!(doc, JsonbValue::Null) {
304+
return None; // Tombstone
306305
}
306+
return Some(doc);
307+
}
308+
if rid.as_str() > id {
309+
// Not found in this table (sorted)
310+
break;
307311
}
308312
}
309313
}
@@ -485,6 +489,7 @@ impl DB {
485489
#[cfg(test)]
486490
mod tests {
487491
use super::*;
492+
use crate::serde_to_jsonb;
488493
use serde_json::json;
489494
use tempfile::tempdir;
490495

@@ -505,13 +510,15 @@ mod tests {
505510
db.create_collection("test").unwrap();
506511

507512
for i in 0..MEMTABLE_THRESHOLD {
508-
db.insert("test", json!({ "a": i })).unwrap();
513+
db.insert("test", serde_to_jsonb(json!({ "a": i })))
514+
.unwrap();
509515
}
510516
let col = db.collections.get("test").unwrap();
511517
assert_eq!(col.memtable.len(), MEMTABLE_THRESHOLD);
512518
assert_eq!(col.jstable_count, 0);
513519

514-
db.insert("test", json!({"a": MEMTABLE_THRESHOLD})).unwrap();
520+
db.insert("test", serde_to_jsonb(json!({"a": MEMTABLE_THRESHOLD})))
521+
.unwrap();
515522
let col = db.collections.get("test").unwrap();
516523
assert_eq!(col.memtable.len(), 1);
517524
assert_eq!(col.jstable_count, 1);
@@ -533,10 +540,10 @@ mod tests {
533540
Some(1024 * 1024),
534541
);
535542
db.create_collection("test").unwrap();
536-
let doc1 = json!({"a": 1});
543+
let doc1 = serde_to_jsonb(json!({"a": 1}));
537544
let id1 = db.insert("test", doc1.clone()).unwrap();
538545

539-
let doc2 = json!({"b": "hello"});
546+
let doc2 = serde_to_jsonb(json!({"b": "hello"}));
540547
db.update("test", &id1, doc2.clone()).unwrap();
541548

542549
db.delete("test", &id1).unwrap();
@@ -582,10 +589,10 @@ mod tests {
582589
Some(1024 * 1024),
583590
);
584591
db.create_collection("test").unwrap();
585-
let doc1 = json!({"a": 1});
592+
let doc1 = serde_to_jsonb(json!({"a": 1}));
586593
let id1 = db.insert("test", doc1.clone()).unwrap();
587594

588-
let doc2 = json!({"b": "hello"});
595+
let doc2 = serde_to_jsonb(json!({"b": "hello"}));
589596
let id2 = db.insert("test", doc2.clone()).unwrap();
590597

591598
db.delete("test", &id1).unwrap();
@@ -602,7 +609,11 @@ mod tests {
602609

603610
assert_eq!(col.memtable.len(), 2);
604611
assert_eq!(*col.memtable.documents.get(&id2).unwrap(), doc2);
605-
assert!(col.memtable.documents.get(&id1).unwrap().is_null());
612+
use jsonb_schema::Value as JsonbValue;
613+
assert!(matches!(
614+
col.memtable.documents.get(&id1).unwrap(),
615+
JsonbValue::Null
616+
));
606617
}
607618

608619
#[test]
@@ -618,12 +629,14 @@ mod tests {
618629
db.create_collection("test").unwrap();
619630

620631
for i in 0..(MEMTABLE_THRESHOLD * JSTABLE_THRESHOLD as usize) {
621-
db.insert("test", json!({ "a": i })).unwrap();
632+
db.insert("test", serde_to_jsonb(json!({ "a": i })))
633+
.unwrap();
622634
}
623635

624636
let col = db.collections.get("test").unwrap();
625637
assert_eq!(col.jstable_count, JSTABLE_THRESHOLD - 1);
626-
db.insert("test", json!({ "a": 999 })).unwrap();
638+
db.insert("test", serde_to_jsonb(json!({ "a": 999 })))
639+
.unwrap();
627640

628641
let col = db.collections.get("test").unwrap();
629642
assert_eq!(col.jstable_count, 1);
@@ -640,32 +653,39 @@ mod tests {
640653
Some(1024 * 1024),
641654
);
642655
db.create_collection("test").unwrap();
643-
let id_to_delete = db.insert("test", json!({ "a": 100 })).unwrap();
656+
let id_to_delete = db
657+
.insert("test", serde_to_jsonb(json!({ "a": 100 })))
658+
.unwrap();
644659

645660
for i in 0..9 {
646-
db.insert("test", json!({ "fill": i })).unwrap();
661+
db.insert("test", serde_to_jsonb(json!({ "fill": i })))
662+
.unwrap();
647663
}
648-
db.insert("test", json!({ "trigger_1": 1 })).unwrap();
664+
db.insert("test", serde_to_jsonb(json!({ "trigger_1": 1 })))
665+
.unwrap();
649666

650667
let col = db.collections.get("test").unwrap();
651668
assert_eq!(col.jstable_count, 1);
652669

653670
db.delete("test", &id_to_delete).unwrap();
654671

655672
for i in 0..8 {
656-
db.insert("test", json!({ "fill_2": i })).unwrap();
673+
db.insert("test", serde_to_jsonb(json!({ "fill_2": i })))
674+
.unwrap();
657675
}
658-
db.insert("test", json!({ "trigger_2": 1 })).unwrap();
676+
db.insert("test", serde_to_jsonb(json!({ "trigger_2": 1 })))
677+
.unwrap();
659678

660679
let col = db.collections.get("test").unwrap();
661680
assert_eq!(col.jstable_count, 2);
662681

663682
for t in 0..3 {
664683
for i in 0..9 {
665-
db.insert("test", json!({ "fill_more": t, "i": i }))
684+
db.insert("test", serde_to_jsonb(json!({ "fill_more": t, "i": i })))
666685
.unwrap();
667686
}
668-
db.insert("test", json!({ "trigger_more": t })).unwrap();
687+
db.insert("test", serde_to_jsonb(json!({ "trigger_more": t })))
688+
.unwrap();
669689
}
670690

671691
let col = db.collections.get("test").unwrap();
@@ -690,21 +710,16 @@ mod tests {
690710
db.create_collection("test").unwrap();
691711

692712
for i in 0..MEMTABLE_THRESHOLD {
693-
db.insert("test", json!({"val": i})).unwrap();
713+
db.insert("test", serde_to_jsonb(json!({"val": i})))
714+
.unwrap();
694715
}
695-
db.insert("test", json!({"val": 10})).unwrap();
716+
db.insert("test", serde_to_jsonb(json!({"val": 10})))
717+
.unwrap();
696718

697719
let results: HashMap<String, Value> = db.scan("test").unwrap().collect();
698720
assert_eq!(results.len(), 11);
699721
}
700722

701-
#[test]
702-
fn test_sanitize() {
703-
assert_eq!(sanitize_filename("valid"), "valid");
704-
assert_eq!(sanitize_filename("foo/bar"), "foo_2fbar");
705-
assert_eq!(sanitize_filename("test.1"), "test_2e1");
706-
}
707-
708723
#[test]
709724
fn test_create_collection() {
710725
let dir = tempdir().unwrap();
@@ -792,7 +807,7 @@ mod tests {
792807
INDEX_THRESHOLD,
793808
Some(1024 * 1024),
794809
);
795-
let res = db.insert("test", json!({ "a": 1 }));
810+
let res = db.insert("test", serde_to_jsonb(json!({ "a": 1 })));
796811
assert!(res.is_err());
797812
}
798813

@@ -829,18 +844,21 @@ mod tests {
829844
Some(1024 * 1024),
830845
);
831846
db.create_collection("test").unwrap();
832-
let id = db.insert("test", json!({ "a": 1 })).unwrap();
847+
let id = db
848+
.insert("test", serde_to_jsonb(json!({ "a": 1 })))
849+
.unwrap();
833850

834851
let doc = db.get("test", &id).unwrap().unwrap();
835-
assert_eq!(doc, json!({ "a": 1 }));
852+
assert_eq!(doc, serde_to_jsonb(json!({ "a": 1 })));
836853

837854
// Flush to force creation of JSTable
838855
for i in 0..MEMTABLE_THRESHOLD {
839-
db.insert("test", json!({ "fill": i })).unwrap();
856+
db.insert("test", serde_to_jsonb(json!({ "fill": i })))
857+
.unwrap();
840858
}
841859

842860
let doc = db.get("test", &id).unwrap().unwrap();
843-
assert_eq!(doc, json!({ "a": 1 }));
861+
assert_eq!(doc, serde_to_jsonb(json!({ "a": 1 })));
844862

845863
assert!(db.get("test", "non-existent").unwrap().is_none());
846864
}

0 commit comments

Comments
 (0)