colexechash: improve memory accounting in the hash table

[yuzefovich]

This commit improves the memory accounting in the hash table to be more
precise in the case when the distsql_workmem limit is exhausted.
Previously, we would allocate large slices first only to perform the
memory accounting after the fact, possibly running out of budget which
would result in a error being thrown. We'd end up in a situation where
the hash table is still referencing larger newly-allocated slices while
only the previous memory usage is accounted for. This commit makes it so
that we account for the needed capacity upfront, then perform the
allocation, and then reconcile the accounting if necessary. This way
we're much more likely to encounter the budget error before making the
large allocations.

Additionally, this commit accounts for some internal slices in the hash
table used only in the hash joiner case.

Also, now both the hash aggregator and the hash joiner eagerly release
references to these internal slices of the hash table when the spilling
to disk occurs (we cannot do the same for the unordered distinct because
there the hash table is actually used after the spilling too).

This required a minor change to the way the unordered distinct spills to
disk. Previously, the memory error could only occur in two spots (and
one of those would leave the hash table in an inconsistent state and we
were "smart" in how we repaired that). However, now the memory error
could occur in more spots (and we could have several different
inconsistent states), so this commit chooses a slight performance
regression of simply rebuilding the hash table from scratch, once, when
the unordered distinct spills to disk.

Addresses: #60022.
Addresses: #64906.

Release note: None

[cockroach-teamcity]

This change is 

[michae2]

Nice!

Reviewed 3 of 5 files at r1, 2 of 2 files at r2, all commit messages.
Reviewable status: complete! 1 of 0 LGTMs obtained (waiting on @DrewKimball and @yuzefovich)

---

pkg/sql/colexec/colexechash/hashtable.go line 453 at r2 (raw file):

	ht.buildFromBufferedTuplesNoAccounting()
	// Now ensure that the accounting is precise (cap's of the slices might
	// exceed len's that we've accounted for). Note that

nit: Missing the rest of this sentence.

---

pkg/sql/colexec/colexechash/hashtable.go line 1011 at r2 (raw file):

	if ht.ProbeScratch.limitedSlicesAreAccountedFor {
		ht.ProbeScratch = hashTableProbeBuffer{}
		ht.allocator.ReleaseMemory(probeBufferInternalMaxMemUsed())

nit: Would it make sense to also set ht.ProbeScratch.limitedSlicesAreAccountedFor = false here? I know we're not currently reusing the hash table after calling Release but it seems like it isn't impossible that someone would one day.

[DrewKimball]

Reviewable status: complete! 2 of 0 LGTMs obtained (waiting on @michae2 and @yuzefovich)

---

pkg/sql/colexec/colexechash/hashtable.go line 1001 at r2 (raw file):

// append-only batch.
// TODO(yuzefovich): lose the reference to the Vals as well, when possible.
func (ht *HashTable) Release() {

Should we also be releasing the slices in the DatumAlloc struct?

[nit] It might be good to nil out the remaining slices in the HashTable, or alternatively add a comment explaining that they're small enough not to matter.

---

pkg/sql/colexec/colexechash/hashtable.go line 1011 at r2 (raw file):

Previously, michae2 (Michael Erickson) wrote…

nit: Would it make sense to also set ht.ProbeScratch.limitedSlicesAreAccountedFor = false here? I know we're not currently reusing the hash table after calling Release but it seems like it isn't impossible that someone would one day.

That happens when we set ht.ProbeScratch = hashTableProbeBuffer{} but I agree it's not obvious; a comment might be a good idea here.

[yuzefovich]

Reviewable status: complete! 0 of 0 LGTMs obtained (and 2 stale) (waiting on @DrewKimball and @michae2)

---

pkg/sql/colexec/colexechash/hashtable.go line 1001 at r2 (raw file):

Previously, DrewKimball (Drew Kimball) wrote…

Should we also be releasing the slices in the DatumAlloc struct?

[nit] It might be good to nil out the remaining slices in the HashTable, or alternatively add a comment explaining that they're small enough not to matter.

Yeah, good point, it's just better to just nil everything out.

---

pkg/sql/colexec/colexechash/hashtable.go line 1011 at r2 (raw file):

Previously, DrewKimball (Drew Kimball) wrote…

That happens when we set ht.ProbeScratch = hashTableProbeBuffer{} but I agree it's not obvious; a comment might be a good idea here.

Indeed, it was already being unset implicitly. Nil-ing out the whole hash table should be clearer.

[yuzefovich]

I made minor changes to also release the memory under the append-only buffered batch, would appreciate another quick look.

Reviewable status: complete! 0 of 0 LGTMs obtained (and 2 stale) (waiting on @DrewKimball and @michae2)

[michae2]

Reviewed 3 of 3 files at r3, all commit messages.
Reviewable status: complete! 1 of 0 LGTMs obtained (and 1 stale) (waiting on @DrewKimball)

[yuzefovich]

TFTRs!

bors r+

[craig]

Build failed:

• GitHub CI (Cockroach)

[yuzefovich]

bors r+

[craig]

Build succeeded:

• GitHub CI (Cockroach)