Keep the cast when folding json_get, plus property tests for the rewriter

[adriangb]

Stacked on #127. The diff shown is this change alone, on top of #127's head.
#127 is approved and mergeable; merge it first and this retargets to main.

The bug

JsonFunctionRewriter replaced CAST(json_get(x,'k') AS T) with the bare accessor for T. The accessors return one Arrow type per JSON type — json_get_int is always Int64, json_get_float always Float64, json_get_str always Utf8 — so whenever T was not that type, the expression silently came out as the accessor's type and the narrowing the cast promised never happened:

CAST(json_get(x,'k') AS INT)           -> Int64,   not Int32
CAST(json_get(x,'k') AS REAL)          -> Float64, not Float32
CAST(json_get(x,'k') AS DECIMAL(10,2)) -> Float64, not Decimal128(10,2)
CAST(json_get(x,'k') AS VARCHAR)       -> Utf8,    not Utf8View

Values escaped the declared type too. On {"a": 3000000000}:

	before	after
CAST(json_get(x,'a') AS INT)	3000000000 (Int64)	overflow error, as a real cast does
TRY_CAST(json_get(x,'a') AS INT)	3000000000 (Int64)	NULL (Int32), as TRY_CAST promises
CAST(json_get(x,'a') AS DECIMAL(10,2))	3000000000.0 (Float64)	overflow error

The CAST half predates #127; #127 routes TRY_CAST through the same type table, which is how the TRY_CAST row above arises.

The fix

Replace only the cast's input, and drop the cast only where the accessor already returns the type that was asked for:

CAST(json_get(x,'k') AS BIGINT)  ->  json_get_int(x,'k')                      (unchanged)
CAST(json_get(x,'k') AS INT)     ->  CAST(json_get_int(x,'k') AS Int32)       (new)

::bigint, ::double and ::boolean plan exactly as before. Where the cast stays it costs one primitive-to-primitive cast, and the JSON union is still never materialized — which was the whole point of the rewrite.

That removes the reason arrow_cast folded only the four types an accessor returns exactly: arrow_cast(x,'Int32') now folds to CAST(json_get_int(x) AS Int32) and keeps its Int32, so it stops materializing the union for no reason. Both paths share one type table now, with the arrow_* type argument parsed the way DataFusion parses it itself in ArrowCastFunc::return_field_from_args. ArrowCastFunc::simplify already drops the cast when the accessor's type matches, so the exact arrow_cast cases keep their old plans as well.

The set of Arrow types that fold is unchanged — smallint and friends are still left alone.

Pushdown is unaffected

JsonGet::placement moves a json_get over a column with literal path arguments towards the leaf nodes, and a cast sitting on top of the accessor is exactly the shape that could have stranded it higher up the plan. It does not: the accessor still lands in the leaf projection for every folding target, and unwrap_cast_in_comparison removes the added cast again for an integer comparison, so where (json_data->'foo')::int = 5 plans identically to before. narrowing_cast_still_reaches_the_leaf_nodes pins this.

Tests

tests/main.rs expectations move from the accessor's type to the type the query asked for (::int → Int32, ::float → Float32, ::numeric → Decimal128(38,10), ::string → Utf8View), and the plan tests show the retained cast. test_plan_arrow_cast_fn_inexact_type_not_folded is now test_plan_arrow_cast_fn_narrowing_type_keeps_cast — its premise is inverted by this change.

tests/rewrite_differential.rs (new)

Property tests over generated JSON documents × 5 cast spellings (CAST, ::, TRY_CAST, arrow_cast, arrow_try_cast) × 8 target types.

prop_fold_is_invisible — folding must equal applying the same cast to the typed accessor: CAST(json_get(x,'k') AS T) == CAST(json_get_<T>(x,'k') AS T). The right-hand side is not itself folded (the rewriter only folds casts over json_get), so it works as a reference implementation. This is the property that found the bug above; it now covers every folding pair rather than only the type-exact ones.

prop_unnest_is_invisible — flattening json_get(json_get(j,'a'),'b') must not change the answer, checked against the genuinely un-rewritten plan.

Getting that un-rewritten plan needs no second SessionContext. Function rewrites run in the analyzer, before the optimizer merges projections, so a cast over a subquery column is never folded — and merge_projection then inlines the json_get back underneath. barrier_really_blocks_folding asserts the trick still works, so the differential cannot quietly become a tautology.

For cast folding that plain "same with and without the rewrite" equality is false by design, and unfolded_union_cast_is_weaker records why: casting the JSON union is far weaker than the accessors (CAST(union AS Float64) is NULL even for 42), so the rewrite is deliberately better than what it replaces. Hence the typed-accessor reference instead.

Also every_cast_spelling_folds_as_documented (plan-shape table over the whole matrix), narrowing_cast_preserves_type_and_narrows, narrowing_cast_still_avoids_the_union and narrowing_cast_still_reaches_the_leaf_nodes.

tests/json_robustness.rs (new)

Three properties across all 11 path-taking UDFs, so a change to any one of them or to the shared scanning code in common.rs has to keep them true:

• prop_no_panic_on_arbitrary_json — valid JSON, truncated JSON, JSON with a byte removed, and arbitrary text must come back as NULL, never a panic or an error. This is the shape of the bug fixed in Fix panic on JSON integers outside the i64 fast path #124 (a JSON integer too wide for jiter's fast path reached a todo!()); the generator emits those integers, so it would have caught it.
• prop_scalar_and_array_paths_agree — a literal document is const-evaluated through each UDF's ScalarValue path, a column goes through the array path.
• prop_input_encoding_does_not_change_result — Utf8, LargeUtf8, Utf8View and dictionary-encoded input must agree.

That last one found a separate, unrelated bug — json_get_array fails on every dictionary-encoded column — which is pinned here by json_get_array_errors_on_dictionary_input and fixed in its own PR. Delete that test and the json_get_array exclusion when the fix lands.

Case counts default low so the suite stays fast (~14s); PROPTEST_CASES overrides, which is how these get run as a fuzzer:

PROPTEST_CASES=50000 cargo test --release --test rewrite_differential

Verification

• cargo test and cargo test --release green (189 tests). The release run is not optional here — it is how the profile-dependent error path in the json_get_array characterization test was caught.
• cargo clippy --all-targets --all-features clean (the crate denies clippy::pedantic).
• 50k proptest cases on the rewriter properties and 25k on the UDF properties, clean.
• Mutation testing (cargo-mutants, src/rewrite.rs): 14 missed / 45 viable on Fold json_get casts written as TRY_CAST, arrow_cast or arrow_try_cast #127's head, 10 after adding these tests. The four newly killed are exactly the fold's type arms — three of the four type arms Fold json_get casts written as TRY_CAST, arrow_cast or arrow_try_cast #127 adds to optimise_json_get_arrow_cast could be deleted with Fold json_get casts written as TRY_CAST, arrow_cast or arrow_try_cast #127's own tests still green. The 10 that remain are equivalent mutants: expr_to_sql_repr's eight integer arms are byte-identical to its _ => scalar.to_string() fallback, and name() only feeds diagnostics.
• Coverage (cargo-llvm-cov): 87.9% → 93.5% lines overall, functions-missed 23 → 8.

🤖 Generated with Claude Code

[codecov-commenter]

Codecov Report

❌ Patch coverage is 97.67442% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 91.17%. Comparing base (5ff1fd6) to head (e9b5823).

Files with missing lines	Patch %	Lines
src/rewrite.rs	97.67%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@                    Coverage Diff                    @@
##           fix-try-cast-json-get     #129      +/-   ##
=========================================================
+ Coverage                  84.70%   91.17%   +6.47%     
=========================================================
  Files                         18       18              
  Lines                       1608     1620      +12     
  Branches                    1608     1620      +12     
=========================================================
+ Hits                        1362     1477     +115     
+ Misses                       173       97      -76     
+ Partials                      73       46      -27     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.