u70b3 commented on code in PR #4971:
URL: https://github.com/apache/datafusion-comet/pull/4971#discussion_r4133228690
##########
native/spark-expr/src/string_funcs/get_json_object.rs:
##########
@@ -246,68 +306,269 @@ fn parse_json_path(path: &str) -> Option<ParsedPath> {
}
}
- Some(ParsedPath {
- segments,
- has_wildcard,
- })
+ Some(ParsedPath { segments })
+}
+
+/// Find the end of a string body starting at `i` (just past the opening
+/// quote). Short bodies are scanned inline; long bodies use memchr2, the same
+/// approach as serde_json's `ignore_str`. Returns the index just past the
+/// closing quote, or None for an unterminated string (the parser rejects the
+/// document anyway).
+#[inline]
+fn skip_string_body(bytes: &[u8], mut i: usize) -> Option<usize> {
+ const SHORT_STRING: usize = 32;
+ if bytes.len().checked_sub(i)? <= SHORT_STRING {
+ while i < bytes.len() {
+ match bytes[i] {
+ b'"' => return Some(i + 1),
+ b'\\' => i = i.checked_add(2)?,
+ _ => i += 1,
+ }
+ }
+ return None;
+ }
+ loop {
+ match memchr::memchr2(b'"', b'\\', bytes.get(i..)?) {
+ Some(off) if bytes[i + off] == b'"' => return Some(i + off + 1),
+ Some(off) => i = i.checked_add(off)?.checked_add(2)?, // escaped
byte
+ None => return None,
+ }
+ }
+}
+
+/// Spark 3.5+ bundles Jackson versions that reject numbers beyond the default
+/// 1000-digit limit, including values this evaluation skips. serde_json's
+/// `IgnoredAny` enforces no such limit, so inspect number tokens before
parsing.
+/// Spark 3.4's Jackson has no default limit and bypasses this scan.
+fn has_oversized_number(json: &str) -> bool {
+ const MAX_NUMBER_DIGITS: usize = 1000;
+ const JACKSON_READER_BUFFER_UNITS: usize = 4000;
+ let bytes = json.as_bytes();
+ let mut i = 0;
+ // Only near-limit floats need the Reader's UTF-16 position. Keep the
+ // prefix count between candidates so many long numbers stay linear.
+ let mut counted_through = 0;
+ let mut utf16_units = 0;
+ while i < bytes.len() {
+ match bytes[i] {
+ // Skip string bodies: Jackson applies no numeric constraint to
+ // string content.
+ b'"' => match skip_string_body(bytes, i + 1) {
+ Some(end) => i = end,
+ None => return false,
+ },
+ b'-' | b'0'..=b'9' => {
+ let mut j = i + usize::from(bytes[i] == b'-');
+ let int_start = j;
+ while j < bytes.len() && bytes[j].is_ascii_digit() {
+ j += 1;
+ }
+ let int_len = j - int_start;
+ let mut fract_len = 0;
+ if j < bytes.len() && bytes[j] == b'.' {
+ j += 1;
+ let start = j;
+ while j < bytes.len() && bytes[j].is_ascii_digit() {
+ j += 1;
+ }
+ fract_len = j - start;
+ }
+ let mut exp_len = 0;
+ if j < bytes.len() && (bytes[j] | 0x20) == b'e' {
+ j += 1;
+ if j < bytes.len() && (bytes[j] == b'+' || bytes[j] ==
b'-') {
+ j += 1;
+ }
+ let start = j;
+ while j < bytes.len() && bytes[j].is_ascii_digit() {
+ j += 1;
+ }
+ exp_len = j - start;
+ }
+ let is_float = fract_len > 0 || exp_len > 0;
+ let digit_count = if is_float {
+ // Spark parses UTF8String via InputStreamReader, then
+ // ReaderBasedJsonParser with a 4000-UTF-16-unit buffer.
+ // Its slow path uses -1 for an absent fraction or
exponent,
+ // forgiving one digit when a number reaches a buffer edge
+ // or EOF. Numbers starting with zero always take that
path.
+ let starts_with_zero = int_len == 1 && bytes[int_start] ==
b'0';
+ let reaches_buffer_edge = if int_len + fract_len + exp_len
> MAX_NUMBER_DIGITS {
+ utf16_units +=
json[counted_through..i].encode_utf16().count();
+ counted_through = i;
+ utf16_units % JACKSON_READER_BUFFER_UNITS + (j - i)
+ >= JACKSON_READER_BUFFER_UNITS
+ } else {
+ false
+ };
+ let slow_path = starts_with_zero || reaches_buffer_edge ||
j == bytes.len();
+ let absent_component = fract_len == 0 || exp_len == 0;
+ int_len + fract_len + exp_len - usize::from(slow_path &&
absent_component)
+ } else {
+ int_len
+ };
+ if digit_count > MAX_NUMBER_DIGITS {
+ return true;
+ }
+ i = j;
+ }
+ _ => i += 1,
+ }
+ }
+ false
}
/// Evaluate a parsed JSONPath against a JSON string.
/// Returns the result as a string, or None if no match.
+#[cfg(test)]
fn evaluate_path(json_str: &str, path: &ParsedPath) -> Option<String> {
- if !path.has_wildcard {
- return value_into_string(extract_no_wildcard(json_str,
&path.segments)?);
- }
-
- let value: Value = serde_json::from_str(json_str).ok()?;
+ evaluate_path_with_number_limit(json_str, path, true)
+}
- // Wildcard path: may return multiple results
- let results = evaluate_with_wildcard(&value, &path.segments);
+fn evaluate_path_with_number_limit(
+ json_str: &str,
+ path: &ParsedPath,
+ check_number_length: bool,
+) -> Option<String> {
+ if check_number_length && has_oversized_number(json_str) {
Review Comment:
done, Skipped the numeric-validation pass for JSON inputs of at most 1,000
bytes; the short-record benchmark is back near baseline.
##########
native/spark-expr/src/string_funcs/get_json_object.rs:
##########
@@ -246,68 +306,269 @@ fn parse_json_path(path: &str) -> Option<ParsedPath> {
}
}
- Some(ParsedPath {
- segments,
- has_wildcard,
- })
+ Some(ParsedPath { segments })
+}
+
+/// Find the end of a string body starting at `i` (just past the opening
+/// quote). Short bodies are scanned inline; long bodies use memchr2, the same
+/// approach as serde_json's `ignore_str`. Returns the index just past the
+/// closing quote, or None for an unterminated string (the parser rejects the
+/// document anyway).
+#[inline]
+fn skip_string_body(bytes: &[u8], mut i: usize) -> Option<usize> {
+ const SHORT_STRING: usize = 32;
+ if bytes.len().checked_sub(i)? <= SHORT_STRING {
+ while i < bytes.len() {
+ match bytes[i] {
+ b'"' => return Some(i + 1),
+ b'\\' => i = i.checked_add(2)?,
+ _ => i += 1,
+ }
+ }
+ return None;
+ }
+ loop {
+ match memchr::memchr2(b'"', b'\\', bytes.get(i..)?) {
+ Some(off) if bytes[i + off] == b'"' => return Some(i + off + 1),
+ Some(off) => i = i.checked_add(off)?.checked_add(2)?, // escaped
byte
+ None => return None,
+ }
+ }
+}
+
+/// Spark 3.5+ bundles Jackson versions that reject numbers beyond the default
+/// 1000-digit limit, including values this evaluation skips. serde_json's
+/// `IgnoredAny` enforces no such limit, so inspect number tokens before
parsing.
+/// Spark 3.4's Jackson has no default limit and bypasses this scan.
+fn has_oversized_number(json: &str) -> bool {
+ const MAX_NUMBER_DIGITS: usize = 1000;
+ const JACKSON_READER_BUFFER_UNITS: usize = 4000;
+ let bytes = json.as_bytes();
+ let mut i = 0;
+ // Only near-limit floats need the Reader's UTF-16 position. Keep the
+ // prefix count between candidates so many long numbers stay linear.
+ let mut counted_through = 0;
+ let mut utf16_units = 0;
+ while i < bytes.len() {
+ match bytes[i] {
+ // Skip string bodies: Jackson applies no numeric constraint to
+ // string content.
+ b'"' => match skip_string_body(bytes, i + 1) {
+ Some(end) => i = end,
+ None => return false,
+ },
+ b'-' | b'0'..=b'9' => {
+ let mut j = i + usize::from(bytes[i] == b'-');
+ let int_start = j;
+ while j < bytes.len() && bytes[j].is_ascii_digit() {
+ j += 1;
+ }
+ let int_len = j - int_start;
+ let mut fract_len = 0;
+ if j < bytes.len() && bytes[j] == b'.' {
+ j += 1;
+ let start = j;
+ while j < bytes.len() && bytes[j].is_ascii_digit() {
+ j += 1;
+ }
+ fract_len = j - start;
+ }
+ let mut exp_len = 0;
+ if j < bytes.len() && (bytes[j] | 0x20) == b'e' {
+ j += 1;
+ if j < bytes.len() && (bytes[j] == b'+' || bytes[j] ==
b'-') {
+ j += 1;
+ }
+ let start = j;
+ while j < bytes.len() && bytes[j].is_ascii_digit() {
+ j += 1;
+ }
+ exp_len = j - start;
+ }
+ let is_float = fract_len > 0 || exp_len > 0;
+ let digit_count = if is_float {
+ // Spark parses UTF8String via InputStreamReader, then
+ // ReaderBasedJsonParser with a 4000-UTF-16-unit buffer.
+ // Its slow path uses -1 for an absent fraction or
exponent,
+ // forgiving one digit when a number reaches a buffer edge
+ // or EOF. Numbers starting with zero always take that
path.
+ let starts_with_zero = int_len == 1 && bytes[int_start] ==
b'0';
+ let reaches_buffer_edge = if int_len + fract_len + exp_len
> MAX_NUMBER_DIGITS {
+ utf16_units +=
json[counted_through..i].encode_utf16().count();
+ counted_through = i;
+ utf16_units % JACKSON_READER_BUFFER_UNITS + (j - i)
+ >= JACKSON_READER_BUFFER_UNITS
+ } else {
+ false
+ };
+ let slow_path = starts_with_zero || reaches_buffer_edge ||
j == bytes.len();
+ let absent_component = fract_len == 0 || exp_len == 0;
+ int_len + fract_len + exp_len - usize::from(slow_path &&
absent_component)
+ } else {
+ int_len
+ };
+ if digit_count > MAX_NUMBER_DIGITS {
+ return true;
+ }
+ i = j;
+ }
+ _ => i += 1,
+ }
+ }
+ false
}
/// Evaluate a parsed JSONPath against a JSON string.
/// Returns the result as a string, or None if no match.
+#[cfg(test)]
fn evaluate_path(json_str: &str, path: &ParsedPath) -> Option<String> {
- if !path.has_wildcard {
- return value_into_string(extract_no_wildcard(json_str,
&path.segments)?);
- }
-
- let value: Value = serde_json::from_str(json_str).ok()?;
+ evaluate_path_with_number_limit(json_str, path, true)
+}
- // Wildcard path: may return multiple results
- let results = evaluate_with_wildcard(&value, &path.segments);
+fn evaluate_path_with_number_limit(
+ json_str: &str,
+ path: &ParsedPath,
+ check_number_length: bool,
+) -> Option<String> {
+ if check_number_length && has_oversized_number(json_str) {
+ return None;
+ }
- match results.len() {
- 0 => None,
- 1 => {
- // Single wildcard match: Spark preserves JSON serialization format
- // (strings keep their quotes, numbers don't)
- if results[0].is_null() {
- None
- } else {
- serde_json::to_string(results[0]).ok()
- }
- }
- // Multiple results: wrap in JSON array. A slice of `&Value` serializes
- // as a JSON array, so no clone into an owned `Value::Array` is needed.
- _ => serde_json::to_string(&results).ok(),
+ let result = extract_path(json_str, &path.segments)?;
+ if !result.matched {
+ return None;
}
+ // The top level is not an array context. Jackson's generator separates
+ // consecutive root-level writes with a single space, so join with one.
+ Some(PathResult::join(result.writes, " "))
}
-/// Evaluation for paths without wildcards.
-///
/// Descends into the document while it is being parsed, so only the matched
-/// subtree is materialized as a `Value`; everything else is skipped by the
-/// parser without allocating. The whole document is still consumed, so
-/// malformed JSON anywhere in the input yields no match, as a full parse
would.
-fn extract_no_wildcard(json_str: &str, segments: &[PathSegment]) ->
Option<Value> {
+/// subtrees are materialized as `Value`s; everything else is skipped by the
+/// parser without allocating. The whole document is still consumed, so
malformed
+/// JSON anywhere in the input yields no match, as a full parse would.
+fn extract_path(json_str: &str, segments: &[PathSegment]) ->
Option<PathResult> {
let mut de = serde_json::Deserializer::from_str(json_str);
- let found = PathSeed { segments }.deserialize(&mut de).ok()?;
+ let found = PathSeed {
+ segments,
+ style: Style::Raw,
+ reject_direct_null: false,
+ }
+ .deserialize(&mut de)
+ .ok()?;
de.end().ok()?;
- found
+ Some(found)
}
/// Deserializes the value at `segments`, discarding everything else.
struct PathSeed<'a> {
segments: &'a [PathSegment],
+ /// The output style in effect, mirroring the `style` parameter Spark
+ /// threads through `evaluatePath`.
+ style: Style,
+ /// A JSON null directly below a named field is not a match in Spark. Nulls
+ /// reached through array traversal are matches and serialize as `null`.
+ reject_direct_null: bool,
+}
+
+/// The outcome of applying (part of) a path, modeled on Spark's generator
+/// protocol: `writes` holds one rendered fragment per generator write and
+/// `matched` is Spark's dirty flag.
+///
+/// The two can diverge: the wildcard arms that write directly to the generator
+/// emit their array wrapper even when nothing inside matched, so an unmatched
+/// result can still carry writes. Spark's generator keeps those bytes — a
later
+/// occurrence of a duplicated field can build on them — so they are preserved
+/// here rather than discarded.
+#[derive(Default)]
+struct PathResult {
+ // A simple field or index lookup produces one write; keep it inline while
+ // descending through nested objects and arrays.
+ writes: SmallVec<[String; 1]>,
+ matched: bool,
+}
+
+impl PathResult {
+ fn join(mut writes: SmallVec<[String; 1]>, separator: &str) -> String {
+ if writes.len() == 1 {
+ writes.pop().unwrap()
+ } else {
+ writes.join(separator)
+ }
+ }
+
+ /// A single verbatim write of a matched value, honoring the output style:
+ /// a string in Raw style is written unquoted (Spark's scalar-unwrap arm),
+ /// everything else keeps JSON serialization.
+ fn write(value: Value, style: Style) -> Self {
+ match value {
+ Value::String(s) if style == Style::Raw => Self {
+ writes: smallvec![s],
+ matched: true,
+ },
+ value => Self {
+ writes: smallvec![value.to_string()],
Review Comment:
done, Terminal wildcards now serialize into one shared output buffer while
preserving Spark styles and match behavior; added regressions and benchmarks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
