Speed up JSON schema inference by ~2.8x - #9494
Conversation
0dad3c7 to
0b881ac
Compare
Rafferty97
changed the title
# Which issue does this PR close? Split out from #9494 to make review easier. It simply adds a benchmark for JSON schema inference. # Rationale for this change I have an open PR that significantly refactors the JSON schema inference code, so I want confidence that not only is the new code correct, but also has better performance than the existing code. # What changes are included in this PR? Adds a benchmark. # Are these changes tested? N/A # Are there any user-facing changes? No
…ion (#9557) # Which issue does this PR close? Another smaller PR extracted from #9494. # Rationale for this change I've moved `ValueIter` into its own module because it's already self-contained, and because that will make it easier to review the changes I have made to `arrow-json/src/reader/schema.rs`. I've also added a public `record_count` function to `ValueIter` - which can be used to simplify consuming code in Datafusion which is currently tracking it separately. # What changes are included in this PR? * Moved `ValueIter` into own module * Added `record_count` method to `ValueIter` # Are these changes tested? Yes. # Are there any user-facing changes? Addition of one new public method, `ValueIter::record_count`.
|
@Rafferty97, can you please merge up this PR to resolve the conflicts and then we can run the benchmarks again to confirm the results |
Done :) |
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
That is certainly a nice result ❤️ |
alamb
left a comment
alamb
left a comment
There was a problem hiding this comment.
Thanks again for this @Rafferty97 and for your patience
I took a look at the PR. my major comments are:
- Can you please document the design / rationale (and why are there LazyLocks being used for what seem to be very small enums)
- Can you ensure the behavior is the same as the existing code?
If we want to change the inference behavior I recommend proposing those changes in a separate PR so that we can evaluate the potential impact.
| @@ -1,4 +0,0 @@ | |||
| {"a":1, "b":[2.0, 1.3, -6.1], "c":[false, true], "d":4.1} | |||
There was a problem hiding this comment.
What is the purpose of removing this file?
There was a problem hiding this comment.
This file was used by tests related to an inference rule that coerces a mix of scalar and array values into an array type. I've removed this rule because the JSON reader can't actually do this coercion, so I figured it was better to error out instead.
I could reinstate these files and test that they cause schema inference to fail - but I'm unsure how useful that actually is?
There was a problem hiding this comment.
I mean in general, if this PR causes inference that previously used to work to start failing, I am not sure we would want to merge it
BY removing hte tests, it is harder to evaluate what effective behavior change in this PR is
There was a problem hiding this comment.
I've removed this rule because the JSON reader can't actually do this coercion, so I figured it was better to error out instead.
I don't understand -- on main test_json_infer_schema read this file and inferred its schema as
let schema = Schema::new(vec![
Field::new("a", DataType::Int64, true),
Field::new("b", list_type_of(DataType::Float64), true),
Field::new("c", list_type_of(DataType::Boolean), true),
Field::new("d", list_type_of(DataType::Utf8), true),
]);What is not supported about this?
There was a problem hiding this comment.
Given this file is so small (133 bytes), can you please unzip it to make the contents more explicit and easier to review and track changes
There was a problem hiding this comment.
The contents are identical to arrays.json. I was following the pattern set by mixed_arrays.json(.gz). I needed to create this file for tests that previously used mixed_arrays.json(.gz) which were deleted. Those files were deleted because they aren't readable by the JSON reader - they rely on coercion semantics that no longer exist.
| } | ||
|
|
||
| /// The type of a JSON value | ||
| pub enum JsonType { |
There was a problem hiding this comment.
I found it strange that the Json type and tape value are now in the infer module -- they seem more widely applicable than just for schema inference
There was a problem hiding this comment.
That's fair. I've moved them to a separate module within schema for better code organisation.
Hi @alamb, thank you for taking a look over the PR and for the detailed feedback. The The behaviour intentionally diverges from the existing code, because the existing code would perform coercions that the actual JSON reader itself doesn't do. So, when such a JSON file is encountered, the previous code would infer successfully but the actual reading into record batches would fail. This new code would return an error at inference time, which I think is more useful and less surprising to the end user. |
|
@alamb This one's ready for review again :) I've removed If you're concerned about the removal of the scalar-to-array coercion logic, I can add it back in, but I think it's better not to unless we plan to implement that coercion logic in the JSON reader itself. |
This comment has been minimized.
This comment has been minimized.
1 similar comment
This comment has been minimized.
This comment has been minimized.
This comment was marked as duplicate.
This comment was marked as duplicate.
This comment was marked as duplicate.
This comment was marked as duplicate.
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
Hi @alamb, thanks for having another look over the PR. I definitely appreciate that it's a little on the heavier side, and the aversion to changing the code's behaviour. You're right that these changes would cause some schema inferences that previously succeeded to return an error. In particular, it would now reject cases where a given key appears with both scalar and array values, which is what Here, the keys However, the JSON reader doesn't actually do this. It appears that it used to but that this functionality was removed. So, even though the schema inference succeeds, the JSON file will fail to be read at query execution time. So, my opinion is that it would actually be better for the inference itself to fail, as this is more informative to the user. With that being said, I'm open to preserving the existing behaviour for now to get the PR over the line. If that's your preference, I'll re-implement that functionality and reinstate those tests I removed. I also agree with your earlier comment about preferring not to commit gziped files in the repo itself. I'd be happy to change those tests that need them to instead do the compression themselves rather than relying on an already-compressed file existing. |
Is there some test that shows this (as in it tries and fails to read such a file)? I am sorry for ht eback and forth but I don't likely have the time to do investigation into these corner cases -- it is easy for me to review code that makes things faster or is additive behavior wise, but trying to validate / understand the potential downstream impacts of chaging behavior takes me much longer and thus is less likely to get time and also has a higher bar for inclusion I think |
|
Hey @alamb, I completely understand your time constraints and caution against introducing potentially breaking changes. To make this simpler to review and merge, I've reinstated the scalar-to-array coercion logic that the original implementation did, and thus reinstated the tests that I had changed/removed. There are still some minor behaviour changes which required some tweaks to the tests, but they're ones I wouldn't expect to have downstream impact:
If there's interest in resolving the mismatch between schema inference and file reading down the line, it can be contained to this ticket I'd opened previously: #9484 |
|
run benchmark json_reader |
This comment was marked as outdated.
This comment was marked as outdated.
|
could we merge up from main to resolve the conflict, and also update the PR body to reflect the latest changes |
bb1459d to
49bfa89
Compare
f75f077 to
4b2c31b
Compare
|
run benchmark json_reader |
This comment was marked as duplicate.
This comment was marked as duplicate.
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: Comparing json-schema (4b2c31b) to 5ec9eaf (merge-base) diff Run configurationrun benchmark json_readerCPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
run benchmark json_reader |
This comment was marked as duplicate.
This comment was marked as duplicate.
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: Comparing json-schema (4b2c31b) to 5ec9eaf (merge-base) diff Run configurationrun benchmark json_readerCPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
Jefffrey
left a comment
Jefffrey
left a comment
There was a problem hiding this comment.
i had codex do a comparison with main and found these differences, which i categorized based on being major/minor differences (these tests assert the result as of main)
major
#[test]
fn test_read_all() {
let mut data = "{}\n".repeat(1024);
data.push_str("{\"late\":true}\n");
let (schema, _) = infer_json_schema(Cursor::new(data), None).unwrap();
assert_eq!(
schema.field(0),
&Field::new("late", DataType::Boolean, true)
);
}- it seems we are only reading the first 1024 records?
#[test]
fn test_reject_invalid_json() {
let result = infer_json_schema(Cursor::new(r#"{"a":1 "b":2}"#), None);
assert!(result.is_err());
let result = infer_json_schema(Cursor::new(r#"{"a":1e}"#), None);
assert!(result.is_err());
}- these are some example json which apparently pass schema inference now
#[test]
fn test_empty() {
let values = std::iter::empty::<Result<Value, ArrowError>>();
let schema = infer_json_schema_from_iterator(values).unwrap();
assert_eq!(schema, Schema::empty());
}- empty input was producing empty schema, but not anymore
#[test]
fn test_consistent_order() {
let data = r#"{"z":1,"a":2}"#;
let (stream_schema, _) = infer_json_schema(Cursor::new(data), None).unwrap();
let value = serde_json::from_str::<Value>(data).unwrap();
let iterator_schema = infer_json_schema_from_iterator(std::iter::once(Ok(value))).unwrap();
assert_eq!(stream_schema, iterator_schema);
}- ordering of fields in inferred schema has changed, especially as now the two public functions aren't consistent with each other
#[test]
fn test_duplicate_fields() {
let data = r#"{"a":1,"a":"value"}"#;
let (schema, _) = infer_json_schema(Cursor::new(data), None).unwrap();
assert_eq!(
schema,
Schema::new(vec![Field::new("a", DataType::Utf8, true)])
);
}- duplicate fields are being inferred as two different fields; previous behaviour was to use latter value, though maybe that is also confusing
minor
max_read_records seems to be applied after reading, so it technically can overshoot how many records were read (even if theyre not used in inference). this could be surprising for consumers using infer_json_schema() as it may now move the cursor more than before. another consequence of this is invalid trailing data isnt ignored, for example:
#[test]
fn test_ignore_trailing_invalid() {
let data = b"{\"a\":1}\nthis is not JSON\n";
let (schema, record_count) = infer_json_schema(Cursor::new(data), Some(1)).unwrap();
assert_eq!(record_count, 1);
assert_eq!(schema.field(0), &Field::new("a", DataType::Int64, true));
}- we expect to read only one record, but since we read beyond that it errors
#[test]
fn test_line_delimited() {
let result = infer_json_schema(Cursor::new(r#"{"a":1} {"b":2}"#), None);
assert!(result.is_err());
}- this ties into previous points, but it seems we now allow parsing json records that arent newline delimited; it seems we didnt exactly document that we only expect newline delimited json however 🤔
#[test]
fn test_complex_null() {
let result = infer_json_schema(Cursor::new(r#"{"a":[null,{"b":1}]}"#), None);
assert!(result.is_err());
}- nulls in a complex object are now allowed - this looks to be a bug in our current code actually
#[test]
fn test_reject_only_null() {
let result = infer_json_schema(Cursor::new("null\n"), None);
assert!(result.is_err());
}- only
nullinput seems to be accepted instead of rejected
i havent taken too close a look at the code in this PR, but it does seem quite a bit of behaviour is changing (even if the benchmark numbers are pretty good 🚀). it is to be said though that our existing test suite is quite lacking to have not asserted these cases
Thanks for taking the time to review the PR. I'm a little surprised to see this many changes in behaviour, but when I get that chance I'll review them and either fix or justify these cases. |
|
Perhaps we can make a separate PR that adds the test coverage for those cases (so we are protected from regressions) and then we can update this PR to ensure the tests don't fail |
a184d8f to
96a892d
Compare
|
Hi @Jefffrey and @alamb, I've now had the time to review each of those test cases individually. Please see my questions and comments below, grouped into a few themes: Fixes
These tests exercised some edge cases I did miss, so I've fixed those issues and these tests are now green on my branch. Thanks for finding them. More tolerant parsing
My new schema inference implementation uses the With respect to accepting more than just newline-delimited JSON, I think this is a very useful feature of the JSON reader, as it permits parsing regular JSON documents as a single row without having to strip the newlines first. I think, in general, we should prefer not to reject malformed input where a sensible parse does exist and is easy to handle. So, for these two cases, I would suggest either excluding them, or rewriting them to assert a successful schema inference rather than a rejection. What do you think? Field ordering
The The latter function consumes an iterator of So, whether Duplicate fields
This is a tricky case to consider because the input in question is malformed and ambiguous, so there's an argument here that the best thing to do is just reject the input. Let me know your thoughts. Other cases
I agree that this is a bug in the original code that my PR fixes. |
96a892d to
8429123
Compare
8429123 to
15710d1
Compare
5 participants
Which issue does this PR close?
This PR sets the groundwork for implementing #9482. It also delivers an approximate 2.8x speed to JSON schema inference.
I have refactored the code that infers the schema of JSON sources, specifically:
TapeDecoder, eliminating the need to materialise rows intoserde_json::Values firstValueIterinto its own moduleRationale for this change
While working on #9482, I saw a need and opportunity to refactor the schema inference code for JSON schemas. I also discovered the bug detailed in #9484.
These changes not only make the code more readible and predictable by eliminating a lot of special case handling, but make it trivial to create a new inference function for "single field" JSON reading.
They have also provided a significant performance boost to the schema inference functions. I added a simple benchmark for
infer_json_schema, which yielded the following results on my machine, reflecting an approx. 2.8x speed up:Before changes:
infer_json_schema/1000 time: [1.4443 ms 1.4616 ms 1.4793 ms]
thrpt: [85.336 MiB/s 86.366 MiB/s 87.401 MiB/s]
After changes:
infer_json_schema/1000 time: [517.79 µs 519.10 µs 520.54 µs]
thrpt: [242.51 MiB/s 243.18 MiB/s 243.80 MiB/s]
change:
time: [−64.919% −64.485% −64.043%] (p = 0.00 < 0.05)
thrpt: [+178.11% +181.57% +185.06%]
NOTE: I haven't rerun the above benchmark since making further changes to this PR in light of feedback.
What changes are included in this PR?
At a glance:
arrow-json/src/reader/schema.rsBecause this is a somewhat sizeable PR, I've done my best to break into a logical sequence of commits to hopefully assist with the review.
Are these changes tested?
Yes, the changes pass all existing unit tests.
I have also added an additional benchmark for the schema inference performance.
Are there any user-facing changes?
There are no API changes, except for the addition of the
record_countmethod onValueIter.However, the error messages returned by infer_json_schema and its cousins will significantly change, with most of them condensed to a single "Expected {expected}, found {got}" template.
Note: Although this PR originally removed the "scalar-to-array" promotion logic that the previous implementation had, I have restored it to keep the scope of this PR narrow and easier to review. For further context on this behaviour and why it might be better to remove it, please see #9484.