Skip to content

Commit

Permalink
try using simplified jiter value
Browse files Browse the repository at this point in the history
  • Loading branch information
davidhewitt committed Feb 13, 2025
1 parent 741961c commit 983dcd2
Show file tree
Hide file tree
Showing 4 changed files with 48 additions and 24 deletions.
15 changes: 7 additions & 8 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ idna = "1.0.3"
base64 = "0.22.1"
num-bigint = "0.4.6"
uuid = "1.12.1"
jiter = { version = "0.8.2", features = ["python"] }
jiter = { git = "https://github.com/pydantic/jiter", branch = "dh/simpler-value", features = ["python"] }
hex = "0.4.3"

[lib]
Expand Down
28 changes: 18 additions & 10 deletions src/input/input_json.rs
Original file line number Diff line number Diff line change
@@ -1,9 +1,8 @@
use std::borrow::Cow;

use jiter::{JsonArray, JsonObject, JsonValue, LazyIndexMap};
use jiter::{JsonArray, JsonObject, JsonValue};
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList, PyString};
use smallvec::SmallVec;
use speedate::MicrosecondsPrecisionOverflowBehavior;
use strum::EnumMessage;

Expand Down Expand Up @@ -62,7 +61,9 @@ impl<'py, 'data> Input<'py> for JsonValue<'data> {
match self {
JsonValue::Object(object) => {
let dict = PyDict::new(py);
for (k, v) in LazyIndexMap::iter(object) {
for (k, v) in object.as_slice() {
// TODO: jiter doesn't deduplicate keys, so we should probably do that here to
// avoid potential wasted work creating Python objects.
dict.set_item(k, v).unwrap();
}
Some(dict)
Expand Down Expand Up @@ -253,7 +254,14 @@ impl<'py, 'data> Input<'py> for JsonValue<'data> {
JsonValue::Str(s) => Ok(string_to_vec(s).into()),
JsonValue::Object(object) => {
// return keys iterator to match python's behavior
let keys: JsonArray = JsonArray::new(object.keys().map(|k| JsonValue::Str(k.clone())).collect());
// FIXME jiter doesn't deduplicate keys, should probably do that here before iteration.
let keys: JsonArray = JsonArray::new(
object
.as_slice()
.iter()
.map(|(k, _)| JsonValue::Str(k.clone()))
.collect(),
);
Ok(GenericIterator::from(keys).into_static())
}
_ => Err(ValError::new(ErrorTypeDefaults::IterableType, self)),
Expand Down Expand Up @@ -543,19 +551,19 @@ impl<'data> ValidatedDict<'_> for &'_ JsonObject<'data> {
&'a self,
consumer: impl ConsumeIterator<ValResult<(Self::Key<'a>, Self::Item<'a>)>, Output = R>,
) -> ValResult<R> {
Ok(consumer.consume_iterator(LazyIndexMap::iter(self).map(|(k, v)| Ok((k.as_ref(), v)))))
Ok(consumer.consume_iterator(self.as_slice().iter().map(|(k, v)| Ok((k.as_ref(), v)))))
}

fn last_key(&self) -> Option<Self::Key<'_>> {
self.keys().last().map(AsRef::as_ref)
self.last().map(|(k, _)| k.as_ref())
}
}

impl<'a, 'py, 'data> ValidatedList<'py> for &'a JsonArray<'data> {
type Item = &'a JsonValue<'data>;

fn len(&self) -> Option<usize> {
Some(SmallVec::len(self))
Some(Vec::len(self))
}
fn iterate<R>(self, consumer: impl ConsumeIterator<PyResult<Self::Item>, Output = R>) -> ValResult<R> {
Ok(consumer.consume_iterator(self.iter().map(Ok)))
Expand All @@ -569,7 +577,7 @@ impl<'a, 'data> ValidatedTuple<'_> for &'a JsonArray<'data> {
type Item = &'a JsonValue<'data>;

fn len(&self) -> Option<usize> {
Some(SmallVec::len(self))
Some(Vec::len(self))
}
fn iterate<R>(self, consumer: impl ConsumeIterator<PyResult<Self::Item>, Output = R>) -> ValResult<R> {
Ok(consumer.consume_iterator(self.iter().map(Ok)))
Expand Down Expand Up @@ -637,12 +645,12 @@ impl<'data> KeywordArgs<'_> for JsonObject<'data> {
Self: 'a;

fn len(&self) -> usize {
LazyIndexMap::len(self)
Vec::len(self)
}
fn get_item<'k>(&self, key: &'k LookupKey) -> ValResult<Option<(&'k LookupPath, Self::Item<'_>)>> {
key.json_get(self)
}
fn iter(&self) -> impl Iterator<Item = ValResult<(Self::Key<'_>, Self::Item<'_>)>> {
LazyIndexMap::iter(self).map(|(k, v)| Ok((k.as_ref(), v)))
self.as_slice().iter().map(|(k, v)| Ok((k.as_ref(), v)))
}
}
27 changes: 22 additions & 5 deletions src/lookup_key.rs
Original file line number Diff line number Diff line change
Expand Up @@ -262,20 +262,33 @@ impl LookupKey {
&'s self,
dict: &'a JsonObject<'data>,
) -> ValResult<Option<(&'s LookupPath, &'a JsonValue<'data>)>> {
// FIXME: use of find_map in here probably leads to quadratic complexity
match self {
Self::Simple(path) => match dict.get(path.first_key()) {
Self::Simple(path) => match dict
.iter()
.rev()
.find_map(|(k, v)| (k == path.first_key()).then_some(v))
{
Some(value) => {
debug_assert!(path.rest.is_empty());
Ok(Some((path, value)))
}
None => Ok(None),
},
Self::Choice { path1, path2 } => match dict.get(path1.first_key()) {
Self::Choice { path1, path2 } => match dict
.iter()
.rev()
.find_map(|(k, v)| (k == path1.first_key()).then_some(v))
{
Some(value) => {
debug_assert!(path1.rest.is_empty());
Ok(Some((path1, value)))
}
None => match dict.get(path2.first_key()) {
None => match dict
.iter()
.rev()
.find_map(|(k, v)| (k == path2.first_key()).then_some(v))
{
Some(value) => {
debug_assert!(path2.rest.is_empty());
Ok(Some((path2, value)))
Expand All @@ -287,7 +300,11 @@ impl LookupKey {
for path in path_choices {
// first step is different from the rest as we already know dict is JsonObject
// because of above checks, we know that path should have at least one element, hence unwrap
let v: &JsonValue = match dict.get(path.first_item.key.as_str()) {
let v: &JsonValue = match dict
.iter()
.rev()
.find_map(|(k, v)| (k == path.first_key()).then_some(v))
{
Some(v) => v,
None => continue,
};
Expand Down Expand Up @@ -527,7 +544,7 @@ impl PathItem {

pub fn json_obj_get<'a, 'data>(&self, json_obj: &'a JsonObject<'data>) -> Option<&'a JsonValue<'data>> {
match self {
Self::S(PathItemString { key, .. }) => json_obj.get(key.as_str()),
Self::S(PathItemString { key, .. }) => json_obj.iter().rev().find_map(|(k, v)| (k == key).then_some(v)),
_ => None,
}
}
Expand Down

0 comments on commit 983dcd2

Please sign in to comment.