opt: inverted-index accelerate filters of the form j->'a' = '{"b": "c"}

[mgartner]

The optimizer currently plans JSON inverted index scans for query filters with a -> JSON fetch value operator on the right side of a = equality operator. However, inverted index scans are only planned with the right side of the = is a JSON boolean, string, number, or null, but not when it is an object or array.

For example:

EXPLAIN (OPT, VERBOSE) SELECT k FROM t WHERE j->'a' = '1';
                                  info
------------------------------------------------------------------------
  project
   ├── columns: k:1
   ├── immutable
   ├── stats: [rows=333.333333]
   ├── cost: 125.131111
   ├── key: (1)
   ├── prune: (1)
   └── scan t@t_j_idx
        ├── columns: k:1
        ├── inverted constraint: /5/1
        │    └── spans: ["7a\x00\x01*\x02\x00", "7a\x00\x01*\x02\x00"]
        ├── stats: [rows=111.111111, distinct(5)=100, null(5)=0]
        ├── cost: 121.787778
        └── key: (1)

EXPLAIN (OPT, VERBOSE) SELECT k FROM t WHERE j->'a' = '{"a": "b"}';
                               info
-------------------------------------------------------------------
  project
   ├── columns: k:1
   ├── immutable
   ├── stats: [rows=333.333333]
   ├── cost: 1097.38333
   ├── key: (1)
   ├── prune: (1)
   └── select
        ├── columns: k:1 j:2
        ├── immutable
        ├── stats: [rows=333.333333]
        ├── cost: 1094.04
        ├── key: (1)
        ├── fd: (1)-->(2)
        ├── scan t
        │    ├── columns: k:1 j:2
        │    ├── stats: [rows=1000, distinct(1)=1000, null(1)=0]
        │    ├── cost: 1084.02
        │    ├── key: (1)
        │    ├── fd: (1)-->(2)
        │    └── prune: (1,2)
        └── filters
             └── (j:2->'a') = '{"a": "b"}' [outer=(2), immutable]

It should be possible to modify jsonOrArrayFilterPlanner. extractJSONFetchValEqCondition so that inverted index scans are generated for the second query above.

Some things to think about:

• Is it possible to support objects with more than 1 key/value pair, e.g., j->'a' = '{"b": "c", "d": "e"}?
• Is it possible to support nested objects, e.g., ``j->'a' = '{"b": {"c": "d"}}`?
• What about arrays on the right side?
• What about an array nested within the object?
• What about an object nested within an array?

[jordanlewis]

Excited about this issue!

[angelazxu]

I spent some time thinking about the questions you brought up, and have some ideas of how to approach this, but also there are some things I'm confused about. Would love to hear your thoughts @mgartner :)

• Is it possible to support objects with more than 1 key/value pair, e.g., j->'a' = '{"b": "c", "d": "e"}'?

  • Maybe? I was looking at the Inverted Indexes RFC (here) which explained that each key/value pair is its own row in the inverted index, so would we have to check for each key/value pair individually? Not sure if my understanding is correct here, but I'm assuming we would need to convert the expression into something like  j->'a' = '{"b": "c"}' AND j->'a' = '{"d": "e"}' to have it work.
  • I did try experimenting with this and running some SQL queries, and it seems like they might not be equivalent expressions, but maybe that's because we're not using the inverted index?

• Is it possible to support nested objects, e.g., j->'a' = '{"b": '{"c": "d"}'}'?

  • I think so, this would just be one row in the inverted index, so supporting this should be pretty similar to supporting a single object. I'm trying to think of the best way to do this - would it be equivalent to convert it to j-> 'a' -> 'b' -> 'c' = "d"? Then it should just follow the existing logic in jsonOrArrayFilterPlanner.extractJSONFetchValEqCondition since the a, b, c would be part of the left side and the right side would just be "d" which is just a string.
  • If those two expressions are equivalent, could that be the general approach for all cases where the right side is a JSON object (without arrays) - just shift all the keys to the left and keep the value on the right? (and would this just be another normalization rule then?)

• What about arrays on the right side? / What about an array nested within the object? / What about an object nested within an array?

  • For arrays, I'm not sure if inverted indexes could be used since order of elements is preserved. From what I understand, every element is stored in its own row in the inverted index, and we're only able to tell that it's part of AN array / the level of nesting in the array, but not the position of it. So if we want to filter something like j->'a' = '["b", '{"c": "d"}']', and there was an object containing {"a" : ['{"c": "d"}', "b"]}, that would also be returned if we used the inverted index.
  • However, also in the RFC, it mentions that for queries with multiple constraints, each row in the intersection must be rechecked against the index condition. Could we possibly implement something like this, where we find all rows containing all elements in the expression using the inverted index, and then recheck the output to filter out the ones that don't match the correct order?

[mgartner]

Not sure if my understanding is correct here, but I'm assuming we would need to convert the expression into something like j->'a' = '{"b": "c"}' AND j->'a' = '{"d": "e"}' to have it work.

That sounds right to me. In invertedidx. extractJSONOrArrayContainsCondition we currently handle expressions like j->'a' = '1' by turning it into a JSON object, {"a": 1} and pass the object to getInvertedExprForJSONOrArrayIndex to generate an inverted expression, which ends up determining the inverted spans to scan. If we turn j->'a' = '{"b": "c", "d": "e"} into some object, getInvertedExprForJSONOrArrayIndex it might handle the rest?

If those two expressions are equivalent, could that be the general approach for all cases where the right side is a JSON object (without arrays) - just shift all the keys to the left and keep the value on the right? (and would this just be another normalization rule then?)

It's a good idea, though i'm not sure it's necessary if we can turn the expression into a JSON object that leads to an equivalent inverted index scan after passing to getInvertedExprForJSONOrArrayIndex, like I mentioned above. One thing to be cautious of is that -> results in NULL if the JSON on the left-hand-side does not contain the key on the right-hand-side (e.g., '{"a": "b"}'::JSON->'c' is NULL) while @> always returns true or false, never NULL. We previously had a normalization rule that turned expressions with -> into expressions with @>, but it was buggy because of this NULL difference. This is the PR where the normalization rule was removed to fix the bug: #55316. It is mind-bending for me to wrap my head around every time it comes up, so don't worry if it bends your mind too—just something to be aware of.

However, also in the RFC, it mentions that for queries with multiple constraints, each row in the intersection must be rechecked against the index condition. Could we possibly implement something like this, where we find all rows containing all elements in the expression using the inverted index, and then recheck the output to filter out the ones that don't match the correct order?

Good find! You can see this in action in the opt test below. Notice that we scan two spans (one for {"a": "b"} and one for {"c: "d"}), which will produce all rows that have either a/b or c/d. So, there is an inverted-filterer after the scan to produce only the rows that have a/b and c/d. It's possible we could follow a similar pattern for something like j->3 = '"a" to first retrieve all rows that have "a" as some element of the j array, and then fetch the entire JSON column to determine if "a" is the element at index 3. It might be more complicated with nested arrays/objects though...

exec-ddl
CREATE TABLE t (
  k INT PRIMARY KEY,
  j JSON,
  INVERTED INDEX (j)
)
----

opt disable=GenerateInvertedIndexZigzagJoins
SELECT k FROM t WHERE j @> '{"a": "b"}' AND j @> '{"c": "d"}'
----
project
 ├── columns: k:1!null
 ├── immutable
 ├── key: (1)
 └── inverted-filter
      ├── columns: k:1!null
      ├── inverted expression: /4
      │    ├── tight: true, unique: true
      │    ├── union spans: empty
      │    └── INTERSECTION
      │         ├── span expression
      │         │    ├── tight: true, unique: true
      │         │    └── union spans: ["7a\x00\x01\x12b\x00\x01", "7a\x00\x01\x12b\x00\x01"]
      │         └── span expression
      │              ├── tight: true, unique: true
      │              └── union spans: ["7c\x00\x01\x12d\x00\x01", "7c\x00\x01\x12d\x00\x01"]
      ├── key: (1)
      └── scan t@secondary
           ├── columns: k:1!null j_inverted_key:4!null
           ├── inverted constraint: /4/1
           │    └── spans
           │         ├── ["7a\x00\x01\x12b\x00\x01", "7a\x00\x01\x12b\x00\x01"]
           │         └── ["7c\x00\x01\x12d\x00\x01", "7c\x00\x01\x12d\x00\x01"]
           ├── key: (1)
           └── fd: (1)-->(4)

You've gotten a great understanding of how inverted indexes work from the RFC and the code. Another tool that has been helpful for me in understanding how JSON inverted indexes are encoded is a kvtrace logic test. I have a __test file (setup to be ignored by git) in pkg/sql/opt/exec/execbuilder/testdata/ for running tests like below with make test PKG='./pkg/sql/opt/exec/execbuilder' TESTS='TestExecBuild/local/__test' TESTFLAGS='--rewrite=true.

statement ok
CREATE TABLE t (
  k INT PRIMARY KEY,
  j JSON,
  INVERTED INDEX (j),
  FAMILY (k, j)
)

query T kvtrace
INSERT INTO t VALUES (5, '{"a": "b", "c": "d"}')
----
CPut /Table/53/1/5/0 -> /TUPLE/
InitPut /Table/53/2/"a"/"b"/5/0 -> /BYTES/
InitPut /Table/53/2/"c"/"d"/5/0 -> /BYTES/

query T kvtrace
SELECT k FROM t WHERE j @> '{"a": "b"}'
----
Scan /Table/53/2/"a"/"b"{-/PrefixEnd}

You can play around with different INSERT statements to see how a row is encoded and written to an inverted index. Above, you can see two InitPuts for the inverted index. In /Table/53/2/"a"/"b"/1/0 -> /BYTES/, 53 is a unique identifier for the table t, 2 indicates the second index of the table (the inverted index), "a"/"b" is one of the encodings for the inserted value for j, and 5 is the primary key of the inserted row.

Lastly, don't feel like you have to add index acceleration for all these complex filters in a single PR—you can start with the most simple expression and iteratively add support for more complex expressions.