u70b3 commented on code in PR #4971:
URL: https://github.com/apache/datafusion-comet/pull/4971#discussion_r4133214916
##########
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)
Review Comment:
done, Removed the fixed-buffer approximation. The opt-in path now uses a
deterministic 1,000-digit policy and documents the recycler-dependent
difference; exact JVM-history parity is not possible.
##########
native/spark-expr/src/string_funcs/get_json_object.rs:
##########
@@ -372,25 +644,107 @@ impl<'de> Visitor<'de> for SegmentVisitor<'_> {
}
fn visit_seq<A: SeqAccess<'de>>(self, mut seq: A) -> Result<Self::Value,
A::Error> {
- let PathSegment::Index(idx) = &self.segments[0] else {
- IgnoredAny.visit_seq(seq)?;
- return Ok(None);
- };
-
- for _ in 0..*idx {
- if seq.next_element::<IgnoredAny>()?.is_none() {
- return Ok(None);
+ match &self.segments[0] {
+ PathSegment::Index(idx) => {
+ // Spark switches to Quoted style for the "one or more results"
+ // case: an index immediately followed by a subscript wildcard.
+ let child_style = match self.segments.get(1) {
+ Some(PathSegment::SubscriptWildcard) |
Some(PathSegment::DoubleWildcard) => {
+ Style::Quoted
+ }
+ _ => self.style,
+ };
+ for _ in 0..*idx {
+ if seq.next_element::<IgnoredAny>()?.is_none() {
+ return Ok(PathResult::default());
+ }
+ }
+ let found = seq
+ .next_element_seed(PathSeed {
+ segments: &self.segments[1..],
+ style: child_style,
+ reject_direct_null: false,
+ })?
+ .unwrap_or_default();
+ // The remaining elements are still visited, so that a
malformed element
+ // after the match yields no match, as a full parse would.
+ IgnoredAny.visit_seq(seq)?;
+ Ok(found)
+ }
+ PathSegment::DoubleWildcard => {
+ // Spark consumes both wildcards of `[*][*]` at once: the
+ // remaining path applies to the outer elements in flatten
+ // style, and the collected writes always form a single array,
+ // even when there is only one element or none matched.
+ let mut writes = SmallVec::new();
+ let mut matched = false;
+ while let Some(mut result) = seq.next_element_seed(PathSeed {
+ segments: &self.segments[1..],
+ style: Style::Flatten,
+ reject_direct_null: false,
+ })? {
+ matched |= result.matched;
+ writes.append(&mut result.writes);
+ }
+ Ok(PathResult::wrap(writes, matched))
+ }
+ PathSegment::SubscriptWildcard => match self.style {
+ // Quoted style: the array wrapper is always kept, even for a
+ // single match.
+ Style::Quoted => {
+ let mut writes = SmallVec::new();
+ let mut matched = false;
+ while let Some(mut result) =
seq.next_element_seed(PathSeed {
+ segments: &self.segments[1..],
+ style: Style::Quoted,
+ reject_direct_null: false,
+ })? {
+ matched |= result.matched;
+ writes.append(&mut result.writes);
+ }
+ Ok(PathResult::wrap(writes, matched))
+ }
+ // Raw or Flatten style: Spark buffers the element writes into
+ // a temporary array and only emits it when more than one
+ // element wrote; a lone writer's brackets are stripped, and
+ // nothing at all is written when no element matched.
+ Style::Raw | Style::Flatten => {
+ let child_style = if self.style == Style::Raw {
+ Style::Quoted
+ } else {
+ Style::Flatten
+ };
+ let mut writers = 0;
+ let mut writes = SmallVec::new();
+ while let Some(mut result) =
seq.next_element_seed(PathSeed {
Review Comment:
done, Enforced the Spark 3.5+ nesting limit even in skipped wildcard
subtrees; added both boundary cases while preserving Spark 3.4 behavior.
--
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]
