cockroachdb / cockroachdb/cockroach
sql: reduce memory usage when collecting stats on wide JSON
- Dominant language
- Go
- Stars
- 32.5k
- Forks
- 4.1k
- PR merge metrics
- PR metrics pending
Description
Consider the schema:
```sql
CREATE TABLE t (
k UUID PRIMARY KEY,
j JSON,
v STRING AS (j->>'name') VIRTUAL,
INDEX (v)
);
```
When collecting stats for this table, we'll sample rows with all 3 columns to build histograms from. For both `k` and `v` we'll build histograms with up to 200 buckets (by default). For `j` we'll only build a histogram with 2 buckets. If `j` contains wide values, retaining 10,000+ samples of the value to build only two buckets can consume a lot of memory, e.g., if the average size of `j` is 40kB, then we'll collect 400MB of samples. That's a lot of memory to consume to build a 2-bucket histogram.
Below are some concrete things we can do to reduce this memory usage. Note that they could apply to column types other than JSON, too.
#### 1. ~~Do not sample non-indexed JSON columns.~~
**UPDATE: This is done. See https://github.com/cockroachdb/cockroach/pull/139766.**
~~If a JSON column has an `INVERTED INDEX` then a histogram is calculated for the inverted index keys and used during query optimization. If a JSON column has a forward `INDEX`, then the collected histogram may be useful during optimization. But if a JSON column is not indexed, it is unlikely that the 2-bucket histogram that is collected will be useful during optimization. By default we should not generate a histogram in this case—this was actually the default behavior up until #99275 landed in 23.2 and I assume the change in behavior was unintentional. These JSON columns would not need to be sampled if no histogram is being made and memory usage during stats collection would be reduced, especially when the JSON columns are wide.~~
#### 2. Truncate sampled JSON values
For the case when JSON columns have a forward `INDEX`, we should truncate values to limit memory usage both during stats collection and during query optimization when histograms are utilized. We already truncate some types [here](https://github.com/cockroachdb/cockroach/blob/84864537545bbfa59e724cf82ce53d4b9c7da514/pkg/sql/stats/row_sampling.go#L281-L318). Truncating JSON values is non trivial. We should consider sampled value or key-encoded JSON values which would then be trivial to truncate. Note that #139380 may also be required for JSON objects that share a lot of common key/value pairs which are encoded into the prefix of a key or value.
#### 3. Do not retain underlying `jsonEncoded` values when sampling virtual computed columns
Consider a slightly different schema than above where `v` has the type `JSON`:
```sql
CREATE TABLE t (
k UUID PRIMARY KEY,
j JSON,
v JSON AS (j->'name') VIRTUAL,
INDEX (v)
);
```
Assuming we complete (1), stats collection will still use a lot of memory collecting samples of `v` when `j` is wide. As the sampler computes `j->'name'` for each row, the resulting datum will have the underlying type of `jsonEncoded` if the `j` datum is also `jsonEncoded`. Critically, the result `jsonEncoded`'s `value []byte` will be a slice of `j`'s `value []byte`. Therefore, even though the result datum will be copied [here](https://github.com/cockroachdb/cockroach/blob/84864537545bbfa59e724cf82ce53d4b9c7da514/pkg/sql/stats/row_sampling.go#L178), the copy will still reference `j`'s larger `value []byte`, preventing the underlying memory from being reclaimed by the GC. Also, our memory accounting will under-estimate the amount of memory used by samples because it only accounts for the result `jsonEncoded`'s bytes, not `j`'s.
The simplest solution is to unconditionally copy bytes withing `jsonEncoded` to a new slice when sampling. This will incur some overhead. Another option is to only copy the bytes if the original `jsonEncoded` bytes are larger than 400-bytes (the limit at which we truncate string-like types). We'd have to keep track of the total capacity of the underlying bytes of a `jsonEncoded` through operations like `->` to do this.
Jira issue: CRDB-46617
Contributor guide
Assessment
This issue has not been assessed yet.