unable to encode table key: *tree.DJSON error during CREATE STATISTICS

[gaoxk]

Describe the problem
Running CREATE STATISTICS ... results in the error ERROR: unable to encode table key: *tree.DJSON when partial indexes get involved.

To Reproduce

1. Set up CockroachDB
2. Create the following table

CREATE TABLE public.policies (
    id UUID NOT NULL,
    principal_user_id UUID NOT NULL,
    serialized JSONB NOT NULL,
    created_at TIMESTAMPTZ NOT NULL DEFAULT now():::TIMESTAMPTZ,
    expires_at TIMESTAMPTZ NULL,
    revoked_at TIMESTAMPTZ NULL,
    CONSTRAINT "primary" PRIMARY KEY (id ASC),
    CONSTRAINT fk_principal_user_id_ref_users FOREIGN KEY (principal_user_id) REFERENCES public.users(id) ON DELETE CASCADE,
    INDEX policies_principal_user_id_expires_at_idx (principal_user_id ASC, expires_at ASC),
    INDEX policies_superuser (principal_user_id ASC, expires_at ASC, revoked_at ASC) WHERE (serialized->>'type':::STRING) LIKE 'SUPERUSER_%':::STRING,
    FAMILY "primary" (id, principal_user_id, serialized, created_at, expires_at, revoked_at)
)

3. Run CREATE STATISTICS policy_stats FROM policies;
4. See error ERROR: unable to encode table key: *tree.DJSON

Expected behavior
I expect no error to occur.

Additional data / screenshots
Output of SHOW STATISTICS FOR TABLE policies;:

 statistics_name |          column_names          |             created              | row_count | distinct_count | null_count |    histogram_id
------------------+--------------------------------+----------------------------------+-----------+----------------+------------+---------------------
  __auto__        | {id}                           | 2021-03-09 16:57:22.255134+00:00 |         0 |              0 |          0 |               NULL
  __auto__        | {principal_user_id}            | 2021-03-09 16:57:22.255134+00:00 |         0 |              0 |          0 |               NULL
  __auto__        | {expires_at}                   | 2021-03-09 16:57:22.255134+00:00 |         0 |              0 |          0 |               NULL
  __auto__        | {principal_user_id,expires_at} | 2021-03-09 16:57:22.255134+00:00 |         0 |              0 |          0 |               NULL
  __auto__        | {serialized}                   | 2021-03-09 16:57:22.255134+00:00 |         0 |              0 |          0 |               NULL
  __auto__        | {created_at}                   | 2021-03-09 16:57:22.255134+00:00 |         0 |              0 |          0 |               NULL
  __auto__        | {id}                           | 2021-04-05 13:50:50.501794+00:00 |         0 |              0 |          0 |               NULL
  __auto__        | {principal_user_id}            | 2021-04-05 13:50:50.501794+00:00 |         0 |              0 |          0 |               NULL
  __auto__        | {expires_at}                   | 2021-04-05 13:50:50.501794+00:00 |         0 |              0 |          0 |               NULL
  __auto__        | {principal_user_id,expires_at} | 2021-04-05 13:50:50.501794+00:00 |         0 |              0 |          0 |               NULL
  __auto__        | {serialized}                   | 2021-04-05 13:50:50.501794+00:00 |         0 |              0 |          0 |               NULL
  __auto__        | {created_at}                   | 2021-04-05 13:50:50.501794+00:00 |         0 |              0 |          0 |               NULL
  __auto__        | {revoked_at}                   | 2021-04-05 18:32:02.452387+00:00 |         0 |              0 |          0 |               NULL
  __auto__        | {id}                           | 2021-04-05 18:32:02.452387+00:00 |         0 |              0 |          0 |               NULL
  __auto__        | {expires_at}                   | 2021-04-05 18:32:02.452387+00:00 |         0 |              0 |          0 |               NULL
  __auto__        | {principal_user_id,expires_at} | 2021-04-05 18:32:02.452387+00:00 |         0 |              0 |          0 |               NULL
  __auto__        | {serialized}                   | 2021-04-05 18:32:02.452387+00:00 |         0 |              0 |          0 |               NULL
  __auto__        | {created_at}                   | 2021-04-05 18:32:02.452387+00:00 |         0 |              0 |          0 |               NULL
  __auto__        | {principal_user_id}            | 2021-04-05 18:32:02.452387+00:00 |         0 |              0 |          0 |               NULL
  __auto__        | {id}                           | 2021-04-07 15:10:15.273846+00:00 |         1 |              1 |          0 | 647947273111896065
  __auto__        | {principal_user_id}            | 2021-04-07 15:10:15.273846+00:00 |         1 |              1 |          0 | 647947273276719105
  __auto__        | {expires_at}                   | 2021-04-07 15:10:15.273846+00:00 |         1 |              1 |          1 | 647947273439674369
  __auto__        | {principal_user_id,expires_at} | 2021-04-07 15:10:15.273846+00:00 |         1 |              1 |          0 |               NULL
  __auto__        | {serialized}                   | 2021-04-07 15:10:15.273846+00:00 |         1 |              1 |          0 |               NULL
  __auto__        | {created_at}                   | 2021-04-07 15:10:15.273846+00:00 |         1 |              1 |          0 | 647947273714663425
  __auto__        | {revoked_at}                   | 2021-04-07 15:10:15.273846+00:00 |         1 |              1 |          1 | 647947273802940417
(26 rows)

Environment:

• CockroachDB CCL v20.2.3 (x86_64-apple-darwin14, built 2020/12/14 18:37:58, go1.13.14) (same version as client)
• Server OS: Mac Big Sur
  -cockroach sql

[mgartner]

Thanks for the report! It looks like this has been fixed on master, so I'll bisect to figure out which commit fixed it and try to backport that commit to 20.2.

[mgartner]

Fixed by #64226.