Null Handling in JSON Extract (Scalar and Index Based)

[siddharthteotia]

Improvement Summary

Fixes issue - apache/pinot#18568

Both JsonExtractScalarTransformFunction and JsonExtractIndexTransformFunction had _nullHandlingEnabled plumbed through (or accessible from BaseTransformFunction) but their six typed SV methods (transformTo{Int,Long,Float,Double,BigDecimal,String}ValuesSV) ignored it and threw unconditionally when the JSON path didn't resolve and no default literal was supplied. This PR makes both transforms honor _nullHandlingEnabled: when query-level null handling is on, unresolved rows surface as SQL NULL (typed null placeholder + a bit in getNullBitmap()) instead of throwing.

All examples use a clicks table with two JSON columns and 4 rows. Both columns have a JSON index (so jsonExtractIndex is applicable):

  row 0:  flatJson   = {"country":"US", "clicks":5}
          nestedJson = {"location":{"city":"SF","country":"US"},
                        "tags":["red","blue","green"],
                        "events":[{"country":"US"},{"country":"CA"}]}

  row 1:  flatJson   = {"country":"CA", "clicks":3}
          nestedJson = {"location":{"city":"Tor"},          -- no country key
                        "tags":["green"],
                        "events":[{"country":null}]}        -- inner explicit null

  row 2:  flatJson   = {"country":null, "clicks":null}
          nestedJson = {"location":null, "tags":null, "events":null}

  row 3:  flatJson   = {}
          nestedJson = {}

SET enableNullHandling = true;

Scenario 1: Projection — flat path

SELECT jsonExtractScalar(flatJson, '$.country', 'STRING') AS c FROM clicks ORDER BY c ASC NULLS LAST;
SELECT jsonExtractIndex(flatJson, '$.country', 'STRING') AS c FROM clicks ORDER BY c ASC NULLS LAST;

- `jsonExtractScalar`
    - Before: throws `IllegalArgumentException: Cannot resolve JSON path on some records…`
    - After: `"CA"`, `"US"`, `NULL`, `NULL`

  - `jsonExtractIndex`
    - Before: throws `RuntimeException: Illegal Json Path: [$.country], for docId [N]`
    - After: `"CA"`, `"US"`, `NULL`, `NULL`

Scenario 2: Projection — nested path

SELECT jsonExtractScalar(nestedJson, '$.tags[0]', 'STRING') AS tag FROM clicks ORDER BY tag ASC NULLS LAST;
SELECT jsonExtractIndex(nestedJson, '$.tags[0]', 'STRING') AS tag FROM clicks ORDER BY tag ASC NULLS LAST;

- `jsonExtractScalar`
    - Before: throws
    - After: `"green"`, `"red"`, `NULL`, `NULL`

  - `jsonExtractIndex`
    - Before: throws
    - After: `"green"`, `"red"`, `NULL`, `NULL`

Scenario 3: DISTINCT — flat path

SELECT DISTINCT jsonExtractScalar(flatJson, '$.country', 'STRING') AS c 
              FROM clicks ORDER BY c ASC NULLS LAST;
SELECT DISTINCT jsonExtractIndex(flatJson, '$.country', 'STRING') AS c 
              FROM clicks ORDER BY c ASC NULLS LAST;

- `jsonExtractScalar`
    - Before: throws
    - After: `"CA"`, `"US"`, `NULL`

  - `jsonExtractIndex`
    - Before: **already returned `"CA"`, `"US"`, `NULL`** via `JsonIndexDistinctOperator` — operator path unchanged
    - After: `"CA"`, `"US"`, `NULL`

Scenario 4: DISTINCT - nested path

SELECT DISTINCT jsonExtractScalar(nestedJson, '$.location.country', 'STRING') AS c 
               FROM clicks ORDER BY c ASC NULLS LAST;
SELECT DISTINCT jsonExtractIndex(nestedJson, '$.location.country', 'STRING') AS c 
               FROM clicks ORDER BY c ASC NULLS LAST;

  - `jsonExtractScalar`
    - Before: throws
    - After: `"US"`, `NULL`

  - `jsonExtractIndex`
    - Before: **already returned `"US"`, `NULL`** (operator path)
    - After: `US`, `NULL`

Scenario 5: GROUP BY + COUNT — flat path

SELECT jsonExtractScalar(flatJson, '$.country', 'STRING') AS c, COUNT(*) 
              FROM clicks GROUP BY c ORDER BY c ASC NULLS LAST;
SELECT jsonExtractIndex(flatJson, '$.country', 'STRING') AS c, COUNT(*) 
              FROM clicks GROUP BY c ORDER BY c ASC NULLS LAST;

 - `jsonExtractScalar`
    - Before: throws
    - After: `("CA", 1)`, `("US", 1)`, `(NULL, 2)`

  - `jsonExtractIndex`
    - Before: throws
    - After: `("CA", 1)`, `("US", 1)`, `(NULL, 2)`

Scenario 6: GROUP BY + COUNT — nested path

SELECT jsonExtractScalar(nestedJson, '$.events[0].country', 'STRING') AS c, COUNT(*) 
              FROM clicks GROUP BY c ORDER BY c ASC NULLS LAST;
SELECT jsonExtractIndex(nestedJson, '$.events[0].country', 'STRING') AS c, COUNT(*) 
              FROM clicks GROUP BY c ORDER BY c ASC NULLS LAST;

  - `jsonExtractScalar`
    - Before: throws
    - After: `("US", 1)`, `(NULL, 3)`

  - `jsonExtractIndex`
    - Before: throws
    - After: `("US", 1)`, `(NULL, 3)`

Scenario 7: 4-arg with SQL NULL literal as default, enableNullHandling = true

• Semantically, passing NULL as the default is equivalent to passing no default at all under NH on — both should surface SQL NULL for unresolved rows.
• Before this PR, jsonExtractScalar already behaved correctly here. No change there.
• Before this PR, jsonExtractIndex did not — for STRING it silently emitted "", and for numeric types it failed at init with NumberFormatException. After the fix, init detects the SQL NULL literal and leaves _defaultValue null, so the 3-arg behavior applies.

SELECT jsonExtractScalar(nestedJson, '$.tags[0]', 'STRING', NULL) AS tag 
              FROM clicks ORDER BY tag ASC NULLS LAST;
SELECT jsonExtractIndex(nestedJson, '$.tags[0]', 'STRING', NULL) AS tag 
             FROM clicks ORDER BY tag ASC NULLS LAST;

  │     Function      │                        Before                         │                 After                  │
  ├───────────────────┼───────────────────────────────────────────────────────┼────────────────────────────────────────┤
  │ jsonExtractScalar │ "green", "red", NULL, NULL (already correct)          │ "green", "red", NULL, NULL (unchanged) │
  ├───────────────────┼───────────────────────────────────────────────────────┼────────────────────────────────────────┤
  │ jsonExtractIndex  │ "green", "red", "", "" ← empty string instead of NULL │ "green", "red", NULL, NULL ✓           │
  └───────────────────┴───────────────────────────────────────────────────────┴────────────────────────────────────────┘

Scenario 8: 4-arg with non-null default literal. Worked correctly before this PR as well

 -- NullHandling on or off — behavior is the same
 SELECT jsonExtractScalar(flatJson, '$.country', 'STRING', 'foobar') FROM clicks;
 SELECT jsonExtractIndex(flatJson, '$.country', 'STRING', 'foobar') FROM clicks;

┌─────────────────────────┬─────────────────────────────────────┐
 │       Query shape       │   Before & After (both functions)   │
 ├─────────────────────────┼─────────────────────────────────────┤
 │ Projection              │ "US", "CA", "foobar", "foobar"      │
 ├─────────────────────────┼─────────────────────────────────────┤
 │ DISTINCT (ORDER BY ASC) │ "CA", "US", "foobar"                │
 ├─────────────────────────┼─────────────────────────────────────┤
 │ GROUP BY + COUNT        │ ("CA", 1), ("US", 1), ("foobar", 2) │
 └─────────────────────────┴─────────────────────────────────────┘

Testing Done

• Added new tests to JsonExtractScalarTransformFunctionTest and JsonExtractIndexTransformFunctionTest to get 100% coverage for null handling
• Added multiple query shapes -- PROJECT, DISTINCT and GROUP BY
• Added multiple Json column shapes -- flat, arbitrarily nested

[siddharthteotia]

Out of scope — pre-existing quirks intentionally not addressed

These have been documented as deferred follow-ups; no behavior change in this PR.

NH off + SQL NULL literal as default

-- enableNullHandling stays false
SELECT jsonExtractScalar(flatJson, '$.country', 'STRING', NULL) FROM clicks;

jsonExtractScalar silently emits the typed zero (0 for INT, "" for STRING, etc.) for unresolved rows. Arguably wrong — passing SQL NULL when NH is off has no consistent semantics — but pre-existing and out of scope.

jsonExtractIndex previously had the same family of issues (STRING "", numeric init-throw); the Form B fix narrowed those to a consistent "behave like 3-arg form" under NH on but did not touch NH-off.

IMPORTANT NOTE - My ideal direction would be to get rid of nullHandlingEnabled all together and just have one function that is smart enough to support NULLs, decides when to use index or not etc. Having 1 v/s 2 functions is not that big of a deal but having an index based function and non index based function along with choice to handle NULLs or not handle them is poor SQL abstraction IMHO. Also not aligned with first principles of query planning and optimization.

For now, the PR extends the functionality so that we offer correct and complete semantics to the users in the current dialect.

Once we have had discussions on generic query planning and optimization in context of both SSE and MSE to decide upon a direction, we can fix such leaky abstractions in multiple places in the code. We will have to think of backwards compatibility as well at that point

[siddharthteotia]

The Pre-Req PR #204 has been merged. This is rebased. Will mark it ready for review once all tests have passed.

[praveenc7]

Overall looks good, left a couple of comments

[praveenc7]

Looks like we repeating this for all types can we do something like

  Object result = extractOrNull(resultExtractor, i);

  @Nullable
  private <T> T extractOrNull(IntFunction<T> extractor, int i) {
    try {
      return extractor.apply(i);
    } catch (InvalidJsonException e) {
      throw new IllegalArgumentException(
          "jsonExtractScalar received input that is not valid JSON. "
              + "This function requires a JSON-typed column or a JSON-formatted string.", e);
    } catch (Exception ignored) {
      // null/empty input or missing path → null; caller's cascade decides
      return null;
    }
  }

[praveenc7]

when a null-aware consumer needs both the values and the null bitmap for a block, the block is parsed twice (getNullBitmap and transformToIntValuesSV) I see we are caching in the JsonExtractIndexTransformFuntion class to avoid this, is there a reason of not having the same here