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.
| assert_eq!(small_field.data_type(), &DataType::Float64); | ||
| } | ||
|
|
||
| #[test] |
There was a problem hiding this comment.
These tests pertain to coercion logic that I've removed, due to them being inconsistent with the JSON reader, which is incapable of doing these coercions.
| } | ||
|
|
||
| /// 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.
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing json-schema (469b437) to 322f9ce (merge-base) diff File an issue against this benchmark runner |
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing json-schema (469b437) to 322f9ce (merge-base) diff 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 |
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
Hey @alamb, thank you again for reviewing my PR. I've done my best to address your concerns either by fixing the code or provide explanations for why things have changed, but please let me know if there's anything else I could do to help move this along. Otherwise, do you think this might be ready to merge? |
|
Hi @Rafferty97 -- thanks for the ping. i will try and find time to review this |
|
Hi @alamb - hope I'm not pestering, but I'm still keen to get this work merged if possible :) |
alamb
left a comment
alamb
left a comment
There was a problem hiding this comment.
Thanks @Rafferty97 -- I am sorry for the delay
I find it very hard to understand what this PR is doing -- the title says "speed up JSON inference" but then it also seems to change the JSON readers behavior in certain cases (illustrated by removing / changing the existing tests)
In general I think we would be very welcoming to improvements in performance (accompanied by benchmark results, as this PR has), but probably not very interested in changes to behavior without some way to restore the old behavior -- we have many production systems now built on this reader, so changing its behavior is not something I think we can do lightly
| @@ -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.
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
| @@ -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.
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?
|
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 |
4 participants
Which issue does this PR close?
This PR fixes #9484, and also 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%]
What changes are included in this PR?
At a glance:
Because 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 - except for one intentionally removed due to the change in behaviour related to #9484 (removing scalar-to-array promotion).
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.
Finally, some files that used to generate a valid schema will now return errors. However, this is desirable because those files would have failed to be read by the actual JSON reader anyway - due to the lack of support for scalar-to-array promotion in the JSON reader. (See #9484)