From 1179bce50315c3e7294599307f2fcc735f9c5c94 Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 12:17:32 +0800 Subject: [PATCH 01/61] feat: add allocation-accounted mpool ownership --- pkg/common/mpool/allocation_account.go | 480 ++++++++++++ .../mpool/allocation_account_mpool_test.go | 741 ++++++++++++++++++ pkg/common/mpool/allocation_account_test.go | 304 +++++++ pkg/common/mpool/mpool.go | 559 +++++++++++-- pkg/common/mpool/mpool_test.go | 36 +- 5 files changed, 2050 insertions(+), 70 deletions(-) create mode 100644 pkg/common/mpool/allocation_account.go create mode 100644 pkg/common/mpool/allocation_account_mpool_test.go create mode 100644 pkg/common/mpool/allocation_account_test.go diff --git a/pkg/common/mpool/allocation_account.go b/pkg/common/mpool/allocation_account.go new file mode 100644 index 0000000000000..603728d49a64b --- /dev/null +++ b/pkg/common/mpool/allocation_account.go @@ -0,0 +1,480 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mpool + +import ( + "errors" + "fmt" + "math" + "runtime" + "sync" + "sync/atomic" +) + +// AllocationOwner and AllocationSite are bounded diagnostic dimensions. +// Callers assign stable values in their own allocation-site ledger. Zero is +// reserved so an accounted allocation can never be published without an +// explicit owner and site. +type AllocationOwner uint8 +type AllocationSite uint8 + +const ( + AllocationOwnerMin AllocationOwner = 1 + AllocationOwnerMax AllocationOwner = 63 + AllocationSiteMin AllocationSite = 1 + AllocationSiteMax AllocationSite = math.MaxUint8 +) + +var ( + ErrAllocationAccountCapacity = errors.New("allocation account capacity exceeded") + ErrAllocationAccountSealed = errors.New("allocation account is sealed") + ErrAllocationAccountInvalid = errors.New("invalid allocation account") + ErrAllocationAccountStale = errors.New("stale allocation account handle") + ErrAllocationMetadataSlots = errors.New("allocation metadata slots exhausted") + ErrAllocationGenerationSlots = errors.New("allocation account generation slots exhausted") + ErrAllocationAccountLive = errors.New("allocation account still owns memory") +) + +const ( + allocationAccountSealedBit = uint64(1) << 63 + allocationAccountUsedMask = allocationAccountSealedBit - 1 +) + +// AllocationAccountHandle identifies one use of a reusable registry slot. +// The upper 32 bits are the slot generation and the lower 32 bits are the slot. +type AllocationAccountHandle uint64 + +func newAllocationAccountHandle(slot, generation uint32) AllocationAccountHandle { + return AllocationAccountHandle(uint64(generation)<<32 | uint64(slot)) +} + +func (h AllocationAccountHandle) slot() uint32 { + return uint32(h) +} + +func (h AllocationAccountHandle) generation() uint32 { + return uint32(uint64(h) >> 32) +} + +// AllocationAccountSnapshot is immutable observation state. A terminal owner +// may publish it after Seal and exact zero; taking a snapshot does not mutate +// account lifecycle. +type AllocationAccountSnapshot struct { + Handle AllocationAccountHandle + Limit uint64 + Used uint64 + Peak uint64 + Sealed bool +} + +// AllocationCapacityController lets an account share a higher-level aggregate +// cap during migration. The controller owns cap policy only; physical MPool +// metadata remains the sole release owner. +type AllocationCapacityController interface { + AcquireAllocationCapacity(uint64) error + ReleaseAllocationCapacity(uint64) +} + +// AllocationAccount owns physical allocation capacity for one execution +// generation. state packs the sealed bit and used bytes into one atomic word, +// so Acquire and Seal have one unambiguous linearization point. +type AllocationAccount struct { + registry *AllocationAccountRegistry + handle AllocationAccountHandle + limit uint64 + control AllocationCapacityController + + state atomic.Uint64 + peak atomic.Uint64 + inflight atomic.Int64 +} + +func (a *AllocationAccount) Handle() AllocationAccountHandle { + if a == nil { + return 0 + } + return a.handle +} + +func (a *AllocationAccount) Snapshot() AllocationAccountSnapshot { + if a == nil { + return AllocationAccountSnapshot{} + } + state := a.state.Load() + return AllocationAccountSnapshot{ + Handle: a.handle, + Limit: a.limit, + Used: state & allocationAccountUsedMask, + Peak: a.peak.Load(), + Sealed: state&allocationAccountSealedBit != 0, + } +} + +func (a *AllocationAccount) acquire(capacity uint64) error { + if a == nil || a.registry == nil || a.handle == 0 { + return ErrAllocationAccountInvalid + } + if capacity == 0 { + return nil + } + state := a.state.Load() + if state&allocationAccountSealedBit != 0 { + return ErrAllocationAccountSealed + } + used := state & allocationAccountUsedMask + if used > a.limit || capacity > a.limit-used { + return newAllocationAccountCapacityError(used, capacity, a.limit) + } + + // Register before consulting the shared controller. Once Seal publishes the + // sealed bit, it either observes this transaction or this transaction + // observes sealed before acquiring controller capacity. + a.inflight.Add(1) + defer a.inflight.Add(-1) + state = a.state.Load() + if state&allocationAccountSealedBit != 0 { + return ErrAllocationAccountSealed + } + used = state & allocationAccountUsedMask + if used > a.limit || capacity > a.limit-used { + return newAllocationAccountCapacityError(used, capacity, a.limit) + } + if a.control != nil { + if err := a.control.AcquireAllocationCapacity(capacity); err != nil { + return err + } + } + acquired := false + defer func() { + if !acquired && a.control != nil { + a.control.ReleaseAllocationCapacity(capacity) + } + }() + + for { + state = a.state.Load() + if state&allocationAccountSealedBit != 0 { + return ErrAllocationAccountSealed + } + used = state & allocationAccountUsedMask + if used > a.limit || capacity > a.limit-used { + return newAllocationAccountCapacityError(used, capacity, a.limit) + } + next := used + capacity + if a.state.CompareAndSwap(state, next) { + for { + peak := a.peak.Load() + if next <= peak || a.peak.CompareAndSwap(peak, next) { + acquired = true + return nil + } + } + } + } +} + +func newAllocationAccountCapacityError( + used uint64, + requested uint64, + limit uint64, +) error { + return fmt.Errorf( + "%w: used=%d requested=%d limit=%d", + ErrAllocationAccountCapacity, + used, + requested, + limit, + ) +} + +func (a *AllocationAccount) release(capacity uint64) { + if capacity == 0 { + return + } + // Keep the local charge until the higher-level policy charge is gone. With + // metadata released by allocationLease first, exact local zero is therefore + // also a complete-release boundary. + if a.control != nil { + state := a.state.Load() + if capacity > state&allocationAccountUsedMask { + panic("allocation account release underflow") + } + a.control.ReleaseAllocationCapacity(capacity) + } + for { + state := a.state.Load() + used := state & allocationAccountUsedMask + if capacity > used { + panic("allocation account release underflow") + } + next := state - capacity + if a.state.CompareAndSwap(state, next) { + return + } + } +} + +// Seal prevents every later acquisition. It waits only for acquisitions that +// linearized before the sealed bit was published to finish updating peak. +func (a *AllocationAccount) Seal() AllocationAccountSnapshot { + if a == nil { + return AllocationAccountSnapshot{} + } + for { + state := a.state.Load() + if state&allocationAccountSealedBit != 0 || + a.state.CompareAndSwap(state, state|allocationAccountSealedBit) { + break + } + } + for a.inflight.Load() != 0 { + runtime.Gosched() + } + return a.Snapshot() +} + +type allocationAccountRegistrySlot struct { + account atomic.Pointer[AllocationAccount] +} + +// AllocationAccountRegistry bounds live generations and accounted-allocation +// metadata for one CN. Registry slots are reused only after Seal and exact +// zero. Their generation counters never wrap. +type AllocationAccountRegistry struct { + mu sync.Mutex + + slots []allocationAccountRegistrySlot + generations []uint32 + free []uint32 + + maxAllocations uint64 + liveAllocations atomic.Uint64 + peakAllocations atomic.Uint64 +} + +func NewAllocationAccountRegistry( + generationSlots uint32, + allocationSlots uint64, +) (*AllocationAccountRegistry, error) { + if generationSlots == 0 || uint64(generationSlots) >= uint64(math.MaxInt) { + return nil, ErrAllocationAccountInvalid + } + registry := &AllocationAccountRegistry{ + slots: make([]allocationAccountRegistrySlot, uint64(generationSlots)+1), + generations: make([]uint32, uint64(generationSlots)+1), + free: make([]uint32, generationSlots), + maxAllocations: allocationSlots, + } + for i := uint32(0); i < generationSlots; i++ { + registry.free[i] = generationSlots - i + } + return registry, nil +} + +func (r *AllocationAccountRegistry) Open( + limit uint64, +) (*AllocationAccount, error) { + return r.OpenWithController(limit, nil) +} + +func (r *AllocationAccountRegistry) OpenWithController( + limit uint64, + control AllocationCapacityController, +) (*AllocationAccount, error) { + if r == nil || limit > allocationAccountUsedMask { + return nil, ErrAllocationAccountInvalid + } + + r.mu.Lock() + defer r.mu.Unlock() + for len(r.free) > 0 { + index := len(r.free) - 1 + slot := r.free[index] + r.free = r.free[:index] + generation := r.generations[slot] + if generation == math.MaxUint32 { + continue + } + generation++ + r.generations[slot] = generation + account := &AllocationAccount{ + registry: r, + handle: newAllocationAccountHandle(slot, generation), + limit: limit, + control: control, + } + r.slots[slot].account.Store(account) + return account, nil + } + return nil, ErrAllocationGenerationSlots +} + +func (r *AllocationAccountRegistry) Resolve( + handle AllocationAccountHandle, +) (*AllocationAccount, bool) { + if r == nil { + return nil, false + } + slot := handle.slot() + if slot == 0 || uint64(slot) >= uint64(len(r.slots)) { + return nil, false + } + account := r.slots[slot].account.Load() + return account, account != nil && account.handle == handle +} + +// Finalize removes a sealed, empty account and makes its slot reusable. A live +// account remains resolvable so physical Free can still release its charge. +func (r *AllocationAccountRegistry) Finalize( + account *AllocationAccount, +) (AllocationAccountSnapshot, error) { + if r == nil || account == nil || account.registry != r { + return AllocationAccountSnapshot{}, ErrAllocationAccountInvalid + } + snapshot := account.Snapshot() + if !snapshot.Sealed || snapshot.Used != 0 || + account.inflight.Load() != 0 { + return snapshot, ErrAllocationAccountLive + } + + r.mu.Lock() + defer r.mu.Unlock() + slot := account.handle.slot() + if slot == 0 || uint64(slot) >= uint64(len(r.slots)) || + r.slots[slot].account.Load() != account { + return snapshot, ErrAllocationAccountStale + } + if current := account.Snapshot(); !current.Sealed || current.Used != 0 || + account.inflight.Load() != 0 { + return current, ErrAllocationAccountLive + } + r.slots[slot].account.Store(nil) + if r.generations[slot] != math.MaxUint32 { + r.free = append(r.free, slot) + } + return account.Snapshot(), nil +} + +func (r *AllocationAccountRegistry) reserveMetadata() error { + if r == nil { + return ErrAllocationAccountInvalid + } + for { + live := r.liveAllocations.Load() + if live >= r.maxAllocations { + return ErrAllocationMetadataSlots + } + if r.liveAllocations.CompareAndSwap(live, live+1) { + next := live + 1 + for { + peak := r.peakAllocations.Load() + if next <= peak || + r.peakAllocations.CompareAndSwap(peak, next) { + break + } + } + return nil + } + } +} + +func (r *AllocationAccountRegistry) releaseMetadata() { + for { + live := r.liveAllocations.Load() + if live == 0 { + panic("allocation metadata slot release underflow") + } + if r.liveAllocations.CompareAndSwap(live, live-1) { + return + } + } +} + +// LiveAllocationMetadata returns published and currently in-flight metadata +// slots. A failed unpublished transaction returns its slot before returning. +func (r *AllocationAccountRegistry) LiveAllocationMetadata() uint64 { + if r == nil { + return 0 + } + return r.liveAllocations.Load() +} + +// PeakAllocationMetadata returns the exact high-water slot count. +func (r *AllocationAccountRegistry) PeakAllocationMetadata() uint64 { + if r == nil { + return 0 + } + return r.peakAllocations.Load() +} + +type allocationAccountRequest struct { + account *AllocationAccount + owner AllocationOwner + site AllocationSite + // checkpoint is nil for every public caller. Same-package fault tests use + // it to prove rollback at each unpublished transaction boundary. + checkpoint func(allocationCheckpoint) error +} + +type allocationCheckpoint uint8 + +const ( + allocationAfterAccount allocationCheckpoint = iota + 1 + allocationAfterMetadata + allocationAfterGlobalStats + allocationAfterPoolStats + allocationAfterPhysical + allocationAfterHeader +) + +func (r allocationAccountRequest) validate() error { + if r.account == nil || r.account.registry == nil || + r.owner < AllocationOwnerMin || r.owner > AllocationOwnerMax || + r.site < AllocationSiteMin { + return ErrAllocationAccountInvalid + } + resolved, ok := r.account.registry.Resolve(r.account.handle) + if !ok || resolved != r.account { + return ErrAllocationAccountStale + } + return nil +} + +func (r allocationAccountRequest) reach( + checkpoint allocationCheckpoint, +) error { + if r.checkpoint == nil { + return nil + } + return r.checkpoint(checkpoint) +} + +type allocationLease struct { + account *AllocationAccount + owner AllocationOwner + site AllocationSite + _ [6]byte +} + +func (l allocationLease) release(capacity uint64) { + if l.account == nil || l.account.registry == nil { + panic("invalid allocation account lease") + } + // Return finite metadata first. account.release retains the local charge + // until controller cleanup completes, so exact zero is a complete-release + // boundary. + l.account.registry.releaseMetadata() + l.account.release(capacity) +} diff --git a/pkg/common/mpool/allocation_account_mpool_test.go b/pkg/common/mpool/allocation_account_mpool_test.go new file mode 100644 index 0000000000000..cd2366a19405e --- /dev/null +++ b/pkg/common/mpool/allocation_account_mpool_test.go @@ -0,0 +1,741 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mpool + +import ( + "bytes" + "errors" + "fmt" + "sync" + "testing" + "unsafe" + + "github.com/stretchr/testify/require" +) + +const ( + testAllocationOwner AllocationOwner = 1 + testAllocationSite AllocationSite = 1 +) + +func newTestAllocationAccount( + t testing.TB, + limit uint64, + metadataSlots uint64, +) (*AllocationAccountRegistry, *AllocationAccount) { + t.Helper() + registry, err := NewAllocationAccountRegistry(4, metadataSlots) + require.NoError(t, err) + account, err := registry.Open(limit) + require.NoError(t, err) + return registry, account +} + +func finalizeTestAllocationAccount( + t testing.TB, + registry *AllocationAccountRegistry, + account *AllocationAccount, +) { + t.Helper() + account.Seal() + _, err := registry.Finalize(account) + require.NoError(t, err) +} + +func TestMPoolAccountedAllocGrowFree(t *testing.T) { + require.Equal(t, uintptr(kMemHdrSz), unsafe.Sizeof(memHdr{})) + require.Equal(t, uintptr(16), unsafe.Sizeof(allocationLease{})) + require.Equal(t, uintptr(64), unsafe.Sizeof(AllocationAccount{})) + require.Equal(t, uintptr(8), unsafe.Sizeof(allocationAccountRegistrySlot{})) + + registry, account := newTestAllocationAccount(t, 1024, 8) + mp := MustNew("accounted-alloc-grow") + defer DeleteMPool(mp) + + empty, err := mp.AllocAccounted( + 0, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + require.Nil(t, empty) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + + buffer, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + require.Equal(t, uint64(64), account.Snapshot().Used) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + + var lease allocationLease + hdr, ok := mp.getPtrMetadata( + unsafe.Pointer(unsafe.SliceData(buffer)), + &lease, + ) + require.True(t, ok) + require.True(t, hdr.isOffHeap()) + require.True(t, hdr.isAccounted()) + require.Same(t, account, lease.account) + require.Equal(t, testAllocationOwner, lease.owner) + require.Equal(t, testAllocationSite, lease.site) + + same, err := mp.Grow(buffer, 32, true) + require.NoError(t, err) + require.Equal( + t, + unsafe.Pointer(unsafe.SliceData(buffer)), + unsafe.Pointer(unsafe.SliceData(same)), + ) + require.Equal(t, uint64(64), account.Snapshot().Used) + + grown, err := mp.Grow(same, 128, true) + require.NoError(t, err) + require.Equal(t, uint64(cap(grown)), account.Snapshot().Used) + require.Equal( + t, + uint64(64+cap(grown)), + account.Snapshot().Peak, + ) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + require.Equal(t, uint64(2), registry.PeakAllocationMetadata()) + + mp.Free(grown) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) +} + +func TestMPoolAccountedRollback(t *testing.T) { + t.Run("account-capacity", func(t *testing.T) { + registry, account := newTestAllocationAccount(t, 63, 1) + mp := MustNew("accounted-capacity") + defer DeleteMPool(mp) + + _, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.ErrorIs(t, err, ErrAllocationAccountCapacity) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) + }) + + t.Run("metadata-capacity", func(t *testing.T) { + registry, account := newTestAllocationAccount(t, 64, 0) + mp := MustNew("accounted-metadata-capacity") + defer DeleteMPool(mp) + + _, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.ErrorIs(t, err, ErrAllocationMetadataSlots) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) + }) + + t.Run("pool-capacity", func(t *testing.T) { + const allocationSize = 768 << 10 + registry, account := newTestAllocationAccount( + t, + 2<<20, + 2, + ) + mp, err := NewMPool("accounted-pool-capacity", 1<<20, NoFixed) + require.NoError(t, err) + defer DeleteMPool(mp) + + first, err := mp.AllocAccounted( + allocationSize, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + _, err = mp.AllocAccounted( + allocationSize, + account, + testAllocationOwner, + testAllocationSite, + ) + require.Error(t, err) + require.Equal(t, uint64(allocationSize), account.Snapshot().Used) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + require.Equal(t, uint64(2), registry.PeakAllocationMetadata()) + + mp.Free(first) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) + }) + + t.Run("global-capacity", func(t *testing.T) { + registry, account := newTestAllocationAccount(t, 64, 1) + mp := MustNew("accounted-global-capacity") + defer DeleteMPool(mp) + + oldGlobalCap := globalCap.Load() + globalBefore := GlobalStats().NumCurrBytes.Load() + globalCap.Store(globalBefore + 63) + t.Cleanup(func() { + globalCap.Store(oldGlobalCap) + }) + + _, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.Error(t, err) + require.Equal(t, globalBefore, GlobalStats().NumCurrBytes.Load()) + require.Zero(t, mp.CurrNB()) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) + }) + + t.Run("grow-metadata-capacity", func(t *testing.T) { + registry, account := newTestAllocationAccount(t, 512, 1) + mp := MustNew("accounted-grow-metadata-capacity") + defer DeleteMPool(mp) + + buffer, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + copy(buffer, bytes.Repeat([]byte{0x5a}, len(buffer))) + + _, err = mp.Grow(buffer, 128, true) + require.ErrorIs(t, err, ErrAllocationMetadataSlots) + require.Equal(t, bytes.Repeat([]byte{0x5a}, len(buffer)), buffer) + require.Equal(t, uint64(64), account.Snapshot().Used) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + + mp.Free(buffer) + finalizeTestAllocationAccount(t, registry, account) + }) + + t.Run("realloc-zero-metadata-capacity", func(t *testing.T) { + registry, account := newTestAllocationAccount(t, 512, 1) + mp := MustNew("accounted-realloc-zero-metadata-capacity") + defer DeleteMPool(mp) + + buffer, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + copy(buffer, bytes.Repeat([]byte{0x5a}, len(buffer))) + + _, err = mp.ReallocZero(buffer, 128, true) + require.ErrorIs(t, err, ErrAllocationMetadataSlots) + require.Equal(t, bytes.Repeat([]byte{0x5a}, len(buffer)), buffer) + require.Equal(t, uint64(64), account.Snapshot().Used) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + + mp.Free(buffer) + finalizeTestAllocationAccount(t, registry, account) + }) + + t.Run("allocator-size-limit-precedes-admission", func(t *testing.T) { + registry, account := newTestAllocationAccount(t, 0, 1) + mp := MustNew("accounted-allocator-size-limit") + defer DeleteMPool(mp) + + _, err := mp.AllocAccounted( + CapLimit-kMemHdrSz, + account, + testAllocationOwner, + testAllocationSite, + ) + require.ErrorIs(t, err, ErrAllocationAccountCapacity) + _, err = mp.AllocAccounted( + CapLimit-kMemHdrSz+1, + account, + testAllocationOwner, + testAllocationSite, + ) + require.Error(t, err) + require.NotErrorIs(t, err, ErrAllocationAccountCapacity) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) + }) +} + +func TestMPoolAccountedGrowCapacityBoundary(t *testing.T) { + const ( + oldCapacity = 64 + required = 65 + ) + newCapacity, ok := GrowCapacity(oldCapacity, required) + require.True(t, ok) + require.Greater(t, newCapacity, int64(required)) + + for _, testCase := range []struct { + name string + limit uint64 + wantError bool + }{ + { + name: "exact-old-plus-rounded-new", + limit: oldCapacity + uint64(newCapacity), + }, + { + name: "one-byte-short", + limit: oldCapacity + uint64(newCapacity) - 1, + wantError: true, + }, + } { + t.Run(testCase.name, func(t *testing.T) { + registry, account := newTestAllocationAccount( + t, + testCase.limit, + 2, + ) + mp := MustNew("accounted-grow-capacity-boundary") + defer DeleteMPool(mp) + + buffer, err := mp.AllocAccounted( + oldCapacity, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + grown, err := mp.Grow(buffer, required, true) + if testCase.wantError { + require.ErrorIs(t, err, ErrAllocationAccountCapacity) + require.Equal(t, uint64(oldCapacity), account.Snapshot().Used) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + mp.Free(buffer) + } else { + require.NoError(t, err) + require.Equal(t, newCapacity, int64(cap(grown))) + require.Equal(t, uint64(newCapacity), account.Snapshot().Used) + require.Equal(t, testCase.limit, account.Snapshot().Peak) + mp.Free(grown) + } + finalizeTestAllocationAccount(t, registry, account) + }) + } +} + +func TestMPoolAccountedReallocZero(t *testing.T) { + registry, account := newTestAllocationAccount(t, 512, 2) + mp := MustNew("accounted-realloc-zero") + defer DeleteMPool(mp) + + buffer, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + copy(buffer, bytes.Repeat([]byte{0x5a}, 64)) + + _, err = mp.ReallocZero(buffer, 128, false) + require.ErrorIs(t, err, ErrAllocationAccountInvalid) + require.Equal(t, uint64(64), account.Snapshot().Used) + + replacement, err := mp.ReallocZero(buffer, 128, true) + require.NoError(t, err) + require.Equal(t, bytes.Repeat([]byte{0x5a}, 64), replacement[:64]) + require.Equal(t, make([]byte, 64), replacement[64:]) + require.Equal(t, uint64(128), account.Snapshot().Used) + require.Equal(t, uint64(192), account.Snapshot().Peak) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + require.Equal(t, uint64(2), registry.PeakAllocationMetadata()) + + mp.Free(replacement) + finalizeTestAllocationAccount(t, registry, account) +} + +func TestMPoolAccountedSealRejectsGrowth(t *testing.T) { + registry, account := newTestAllocationAccount(t, 512, 2) + mp := MustNew("accounted-sealed-growth") + defer DeleteMPool(mp) + + buffer, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + account.Seal() + _, err = mp.Grow(buffer, 128, true) + require.ErrorIs(t, err, ErrAllocationAccountSealed) + require.Equal(t, uint64(64), account.Snapshot().Used) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + + mp.Free(buffer) + _, err = registry.Finalize(account) + require.NoError(t, err) +} + +func TestMPoolAccountedCrossPoolAndTeardown(t *testing.T) { + t.Run("cross-pool", func(t *testing.T) { + registry, account := newTestAllocationAccount(t, 64, 1) + owner := MustNew("accounted-cross-owner") + other := MustNew("accounted-cross-other") + defer DeleteMPool(owner) + defer DeleteMPool(other) + + buffer, err := owner.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + other.Free(buffer) + require.Panics(t, func() { + other.Free(buffer) + }) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) + }) + + t.Run("deleted-owner", func(t *testing.T) { + registry, account := newTestAllocationAccount(t, 64, 1) + owner := MustNew("accounted-deleted-owner") + other := MustNew("accounted-deleted-other") + defer DeleteMPool(other) + + globalBefore := GlobalStats().NumCurrBytes.Load() + buffer, err := owner.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + DeleteMPool(owner) + require.Equal(t, globalBefore+64, GlobalStats().NumCurrBytes.Load()) + require.Equal(t, uint64(64), account.Snapshot().Used) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + + other.Free(buffer) + require.Equal(t, globalBefore, GlobalStats().NumCurrBytes.Load()) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) + }) + + t.Run("no-lock-teardown", func(t *testing.T) { + registry, account := newTestAllocationAccount(t, 64, 1) + mp := MustNewNoLock("accounted-no-lock-teardown") + globalBefore := GlobalStats().NumCurrBytes.Load() + _, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + DeleteMPool(mp) + require.Equal(t, globalBefore, GlobalStats().NumCurrBytes.Load()) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) + }) +} + +func TestMPoolAccountedConcurrentAllocFree(t *testing.T) { + const ( + workers = 32 + rounds = 200 + size = 64 + ) + registry, account := newTestAllocationAccount( + t, + workers*size, + workers, + ) + mp := MustNew("accounted-concurrent") + defer DeleteMPool(mp) + + var wait sync.WaitGroup + wait.Add(workers) + for range workers { + go func() { + defer wait.Done() + for range rounds { + buffer, err := mp.AllocAccounted( + size, + account, + testAllocationOwner, + testAllocationSite, + ) + if err != nil { + t.Errorf("accounted alloc: %v", err) + return + } + mp.Free(buffer) + } + }() + } + wait.Wait() + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) +} + +func TestMPoolAccountedTransactionRollback(t *testing.T) { + countGlobalMetadata := func() (headers int, leases int) { + for i := range globalPtrShards { + shard := &globalPtrShards[i] + shard.mu.Lock() + headers += len(shard.m) + leases += len(shard.leases) + shard.mu.Unlock() + } + return headers, leases + } + + injectedError := errors.New("injected allocation error") + stages := []struct { + name string + checkpoint allocationCheckpoint + }{ + {name: "account", checkpoint: allocationAfterAccount}, + {name: "metadata", checkpoint: allocationAfterMetadata}, + {name: "global-stats", checkpoint: allocationAfterGlobalStats}, + {name: "pool-stats", checkpoint: allocationAfterPoolStats}, + {name: "physical", checkpoint: allocationAfterPhysical}, + {name: "header-publication", checkpoint: allocationAfterHeader}, + } + for _, fault := range []struct { + name string + trigger func() error + wantPanic bool + }{ + { + name: "error", + trigger: func() error { + return injectedError + }, + }, + { + name: "panic", + trigger: func() error { + panic("injected allocation panic") + }, + wantPanic: true, + }, + } { + for _, stage := range stages { + for _, noLock := range []bool{false, true} { + poolKind := "sharded" + if noLock { + poolKind = "no-lock" + } + t.Run(fault.name+"/"+stage.name+"/"+poolKind, func(t *testing.T) { + registry, account := newTestAllocationAccount(t, 64, 1) + var mp *MPool + if noLock { + mp = MustNewNoLock("accounted-rollback-no-lock") + } else { + mp = MustNew("accounted-rollback") + } + defer DeleteMPool(mp) + + globalBytesBefore := GlobalStats().NumCurrBytes.Load() + headersBefore, leasesBefore := countGlobalMetadata() + request := allocationAccountRequest{ + account: account, + owner: testAllocationOwner, + site: testAllocationSite, + checkpoint: func( + reached allocationCheckpoint, + ) error { + if reached != stage.checkpoint { + return nil + } + return fault.trigger() + }, + } + if fault.wantPanic { + require.PanicsWithValue(t, "injected allocation panic", func() { + _, _ = mp.allocAccountedWithDetailK("", 64, request) + }) + } else { + _, err := mp.allocAccountedWithDetailK("", 64, request) + require.ErrorIs(t, err, injectedError) + } + + require.Zero(t, mp.CurrNB()) + require.Equal( + t, + globalBytesBefore, + GlobalStats().NumCurrBytes.Load(), + ) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + if noLock { + require.Empty(t, mp.ptrs) + require.Empty(t, mp.leases) + } else { + headersAfter, leasesAfter := countGlobalMetadata() + require.Equal(t, headersBefore, headersAfter) + require.Equal(t, leasesBefore, leasesAfter) + } + finalizeTestAllocationAccount(t, registry, account) + }) + } + } + } +} + +func BenchmarkMPoolAccountedAllocation(b *testing.B) { + for _, accounted := range []bool{false, true} { + mode := "unaccounted" + if accounted { + mode = "accounted" + } + for _, size := range []int{64, 4 << 10, 16 << 10, 64 << 10} { + b.Run(fmt.Sprintf("%s/alloc-free/%d", mode, size), func(b *testing.B) { + mp := MustNew("benchmark-allocation-account") + defer DeleteMPool(mp) + var registry *AllocationAccountRegistry + var account *AllocationAccount + if accounted { + var err error + registry, err = NewAllocationAccountRegistry(1, 1) + require.NoError(b, err) + account, err = registry.Open(1 << 60) + require.NoError(b, err) + } + b.ReportAllocs() + b.SetBytes(int64(size)) + b.ResetTimer() + for range b.N { + var buffer []byte + var allocErr error + if accounted { + buffer, allocErr = mp.AllocAccounted( + size, + account, + testAllocationOwner, + testAllocationSite, + ) + } else { + buffer, allocErr = mp.Alloc(size, true) + } + if allocErr != nil { + b.Fatal(allocErr) + } + mp.Free(buffer) + } + b.StopTimer() + if accounted { + finalizeTestAllocationAccount(b, registry, account) + } + }) + } + + b.Run(mode+"/grow-replacement", func(b *testing.B) { + mp := MustNew("benchmark-allocation-account-grow") + defer DeleteMPool(mp) + var registry *AllocationAccountRegistry + var account *AllocationAccount + if accounted { + var err error + registry, err = NewAllocationAccountRegistry(1, 2) + require.NoError(b, err) + account, err = registry.Open(1 << 60) + require.NoError(b, err) + } + b.ReportAllocs() + b.SetBytes(64 << 10) + b.ResetTimer() + for range b.N { + var buffer []byte + var allocErr error + if accounted { + buffer, allocErr = mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + } else { + buffer, allocErr = mp.Alloc(64, true) + } + if allocErr != nil { + b.Fatal(allocErr) + } + buffer, allocErr = mp.Grow(buffer, 64<<10, true) + if allocErr != nil { + b.Fatal(allocErr) + } + mp.Free(buffer) + } + b.StopTimer() + if accounted { + finalizeTestAllocationAccount(b, registry, account) + } + }) + } + + b.Run("accounted/parallel-alloc-free/65536", func(b *testing.B) { + mp := MustNew("benchmark-allocation-account-parallel") + defer DeleteMPool(mp) + registry, err := NewAllocationAccountRegistry(1, 1024) + require.NoError(b, err) + account, err := registry.Open(1 << 60) + require.NoError(b, err) + b.ReportAllocs() + b.SetBytes(64 << 10) + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + buffer, allocErr := mp.AllocAccounted( + 64<<10, + account, + testAllocationOwner, + testAllocationSite, + ) + if allocErr != nil { + b.Fatal(allocErr) + } + mp.Free(buffer) + } + }) + b.StopTimer() + finalizeTestAllocationAccount(b, registry, account) + }) +} diff --git a/pkg/common/mpool/allocation_account_test.go b/pkg/common/mpool/allocation_account_test.go new file mode 100644 index 0000000000000..a92877d0815a0 --- /dev/null +++ b/pkg/common/mpool/allocation_account_test.go @@ -0,0 +1,304 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mpool + +import ( + "errors" + "math" + "runtime" + "sync" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" +) + +type testAllocationCapacityController struct { + used atomic.Uint64 + reject atomic.Bool + acquireStarted chan struct{} + acquireProceed chan struct{} + releaseStarted chan struct{} + releaseProceed chan struct{} +} + +func (c *testAllocationCapacityController) AcquireAllocationCapacity( + capacity uint64, +) error { + if c.reject.Load() { + return ErrAllocationAccountCapacity + } + c.used.Add(capacity) + if c.acquireStarted != nil { + close(c.acquireStarted) + <-c.acquireProceed + } + return nil +} + +func (c *testAllocationCapacityController) ReleaseAllocationCapacity( + capacity uint64, +) { + c.used.Add(^uint64(capacity - 1)) + if c.releaseStarted != nil { + close(c.releaseStarted) + <-c.releaseProceed + } +} + +func TestAllocationAccountRegistryLifecycle(t *testing.T) { + registry, err := NewAllocationAccountRegistry(1, 2) + require.NoError(t, err) + + account, err := registry.Open(128) + require.NoError(t, err) + firstHandle := account.Handle() + resolved, ok := registry.Resolve(firstHandle) + require.True(t, ok) + require.Same(t, account, resolved) + + _, err = registry.Open(128) + require.ErrorIs(t, err, ErrAllocationGenerationSlots) + _, err = registry.Finalize(account) + require.ErrorIs(t, err, ErrAllocationAccountLive) + + require.NoError(t, account.acquire(64)) + require.NoError(t, account.acquire(64)) + err = account.acquire(1) + require.ErrorIs(t, err, ErrAllocationAccountCapacity) + require.Equal(t, AllocationAccountSnapshot{ + Handle: firstHandle, + Limit: 128, + Used: 128, + Peak: 128, + }, account.Snapshot()) + + account.release(64) + sealed := account.Seal() + require.True(t, sealed.Sealed) + require.Equal(t, uint64(64), sealed.Used) + require.ErrorIs(t, account.acquire(1), ErrAllocationAccountSealed) + // A zero-byte request is always a no-op, including after Seal. + require.NoError(t, account.acquire(0)) + + account.release(64) + final, err := registry.Finalize(account) + require.NoError(t, err) + require.True(t, final.Sealed) + require.Zero(t, final.Used) + _, ok = registry.Resolve(firstHandle) + require.False(t, ok) + _, err = registry.Finalize(account) + require.ErrorIs(t, err, ErrAllocationAccountStale) + + reused, err := registry.Open(128) + require.NoError(t, err) + require.NotEqual(t, firstHandle, reused.Handle()) + require.Equal(t, firstHandle.slot(), reused.Handle().slot()) + require.Equal(t, firstHandle.generation()+1, reused.Handle().generation()) + reused.Seal() + _, err = registry.Finalize(reused) + require.NoError(t, err) +} + +func TestAllocationAccountControllerRollback(t *testing.T) { + registry, err := NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + controller := &testAllocationCapacityController{} + account, err := registry.OpenWithController(8, controller) + require.NoError(t, err) + + controller.reject.Store(true) + require.ErrorIs(t, account.acquire(1), ErrAllocationAccountCapacity) + require.Zero(t, controller.used.Load()) + require.Zero(t, account.Snapshot().Used) + + controller.reject.Store(false) + controller.acquireStarted = make(chan struct{}) + controller.acquireProceed = make(chan struct{}) + acquireResult := make(chan error, 1) + go func() { + acquireResult <- account.acquire(1) + }() + <-controller.acquireStarted + sealed := make(chan struct{}) + go func() { + account.Seal() + close(sealed) + }() + for !account.Snapshot().Sealed { + runtime.Gosched() + } + close(controller.acquireProceed) + require.ErrorIs(t, <-acquireResult, ErrAllocationAccountSealed) + <-sealed + require.Zero(t, controller.used.Load()) + require.Zero(t, account.Snapshot().Used) + _, err = registry.Finalize(account) + require.NoError(t, err) +} + +func TestAllocationAccountFinalizeWaitsForRelease(t *testing.T) { + registry, err := NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + controller := &testAllocationCapacityController{ + releaseStarted: make(chan struct{}), + releaseProceed: make(chan struct{}), + } + t.Cleanup(func() { + select { + case <-controller.releaseProceed: + default: + close(controller.releaseProceed) + } + }) + account, err := registry.OpenWithController(1, controller) + require.NoError(t, err) + require.NoError(t, account.acquire(1)) + require.NoError(t, registry.reserveMetadata()) + account.Seal() + + released := make(chan struct{}) + go func() { + defer close(released) + allocationLease{account: account, owner: 1, site: 1}.release(1) + }() + <-controller.releaseStarted + require.Equal(t, uint64(1), account.Snapshot().Used) + _, err = registry.Finalize(account) + require.ErrorIs(t, err, ErrAllocationAccountLive) + + close(controller.releaseProceed) + <-released + _, err = registry.Finalize(account) + require.NoError(t, err) +} + +func TestAllocationAccountRegistryBounds(t *testing.T) { + _, err := NewAllocationAccountRegistry(0, 1) + require.ErrorIs(t, err, ErrAllocationAccountInvalid) + + registry, err := NewAllocationAccountRegistry(1, 0) + require.NoError(t, err) + _, err = registry.Open(allocationAccountUsedMask + 1) + require.ErrorIs(t, err, ErrAllocationAccountInvalid) + + account, err := registry.Open(1) + require.NoError(t, err) + require.NoError(t, account.acquire(1)) + require.ErrorIs(t, registry.reserveMetadata(), ErrAllocationMetadataSlots) + account.release(1) + account.Seal() + _, err = registry.Finalize(account) + require.NoError(t, err) + + registry.mu.Lock() + registry.generations[1] = math.MaxUint32 + registry.mu.Unlock() + _, err = registry.Open(1) + require.ErrorIs(t, err, ErrAllocationGenerationSlots) +} + +func TestAllocationAccountRequestValidation(t *testing.T) { + registry, account := newTestAllocationAccount(t, 64, 1) + mp := MustNew("accounted-request-validation") + defer DeleteMPool(mp) + + for _, testCase := range []struct { + name string + account *AllocationAccount + owner AllocationOwner + site AllocationSite + }{ + {name: "nil-account", owner: 1, site: 1}, + {name: "zero-owner", account: account, site: 1}, + {name: "owner-out-of-range", account: account, owner: 64, site: 1}, + {name: "zero-site", account: account, owner: 1}, + } { + t.Run(testCase.name, func(t *testing.T) { + _, err := mp.AllocAccounted( + 64, + testCase.account, + testCase.owner, + testCase.site, + ) + require.ErrorIs(t, err, ErrAllocationAccountInvalid) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + }) + } + finalizeTestAllocationAccount(t, registry, account) + _, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.ErrorIs(t, err, ErrAllocationAccountStale) +} + +func TestAllocationAccountSealLinearization(t *testing.T) { + const contenders = 128 + + registry, err := NewAllocationAccountRegistry(1, contenders) + require.NoError(t, err) + account, err := registry.Open(contenders) + require.NoError(t, err) + + start := make(chan struct{}) + sealed := make(chan struct{}) + var wait sync.WaitGroup + var acquired atomic.Uint64 + wait.Add(contenders) + for range contenders { + go func() { + defer wait.Done() + <-start + err := account.acquire(1) + if err == nil { + acquired.Add(1) + <-sealed + account.release(1) + return + } + if !errors.Is(err, ErrAllocationAccountSealed) { + t.Errorf("unexpected acquire error: %v", err) + } + }() + } + + close(start) + snapshot := account.Seal() + close(sealed) + wait.Wait() + require.Equal(t, acquired.Load(), snapshot.Used) + require.Equal(t, acquired.Load(), snapshot.Peak) + require.Zero(t, account.Snapshot().Used) + _, err = registry.Finalize(account) + require.NoError(t, err) +} + +func TestAllocationAccountReleaseUnderflow(t *testing.T) { + registry, err := NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + account, err := registry.Open(1) + require.NoError(t, err) + require.Panics(t, func() { + account.release(1) + }) + account.Seal() + _, err = registry.Finalize(account) + require.NoError(t, err) +} diff --git a/pkg/common/mpool/mpool.go b/pkg/common/mpool/mpool.go index ce33dd8c2cf30..adf1e20f59ae0 100644 --- a/pkg/common/mpool/mpool.go +++ b/pkg/common/mpool/mpool.go @@ -228,15 +228,29 @@ type memHdr struct { poolId int64 allocSz int32 guard [3]uint8 - offHeap bool + kind uint8 } +const ( + memKindOnHeap uint8 = iota + memKindOffHeap + memKindAccountedOffHeap +) + func init() { if unsafe.Sizeof(memHdr{}) != kMemHdrSz { panic("memory header size assertion failed") } } +func (pHdr memHdr) isOffHeap() bool { + return pHdr.kind != memKindOnHeap +} + +func (pHdr memHdr) isAccounted() bool { + return pHdr.kind == memKindAccountedOffHeap +} + func (pHdr *memHdr) SetGuard() { pHdr.guard[0] = 0xDE pHdr.guard[1] = 0xAD @@ -324,6 +338,7 @@ type MPool struct { noLock bool ptrs map[unsafe.Pointer]memHdr + leases map[unsafe.Pointer]allocationLease } const ( @@ -334,14 +349,51 @@ const ( func (mp *MPool) recordPtrHdr(ptr unsafe.Pointer, pHdr memHdr) error { if !mp.noLock { return gRecordPtr(ptr, pHdr) - } else { - _, ok := mp.ptrs[ptr] - if ok { - return moerr.NewInternalErrorNoCtx("ptr already recorded") + } + if _, ok := mp.ptrs[ptr]; ok { + return moerr.NewInternalErrorNoCtx("ptr already recorded") + } + mp.ptrs[ptr] = pHdr + return nil +} + +func (mp *MPool) recordAccountedPtrMetadata( + ptr unsafe.Pointer, + pHdr memHdr, + lease allocationLease, + request allocationAccountRequest, +) error { + if !mp.noLock { + return gRecordAccountedPtrMetadata( + ptr, + pHdr, + lease, + request, + ) + } + if _, ok := mp.ptrs[ptr]; ok { + return moerr.NewInternalErrorNoCtx("ptr already recorded") + } + if _, ok := mp.leases[ptr]; ok { + return moerr.NewInternalErrorNoCtx("account lease already recorded") + } + committed := false + defer func() { + if !committed { + delete(mp.ptrs, ptr) + delete(mp.leases, ptr) } - mp.ptrs[ptr] = pHdr - return nil + }() + mp.ptrs[ptr] = pHdr + if err := request.reach(allocationAfterHeader); err != nil { + return err + } + if mp.leases == nil { + mp.leases = make(map[unsafe.Pointer]allocationLease) } + mp.leases[ptr] = lease + committed = true + return nil } func (mp *MPool) getPtrHdr(ptr unsafe.Pointer) (memHdr, bool) { @@ -352,25 +404,81 @@ func (mp *MPool) getPtrHdr(ptr unsafe.Pointer) (memHdr, bool) { return hdr, ok } -func (mp *MPool) removePtrHdr(ptr unsafe.Pointer) (memHdr, bool) { +func (mp *MPool) getPtrMetadata( + ptr unsafe.Pointer, + lease *allocationLease, +) (memHdr, bool) { if !mp.noLock { - return gRemovePtr(ptr) - } else { - hdr, ok := mp.ptrs[ptr] + return gGetPtrMetadata(ptr, lease) + } + hdr, ok := mp.ptrs[ptr] + if !ok { + return memHdr{}, false + } + if !hdr.isAccounted() { + return hdr, true + } + accountedLease, hasLease := mp.leases[ptr] + if !hasLease { + panic(moerr.NewInternalErrorNoCtx( + "accounted allocation has no account lease", + )) + } + *lease = accountedLease + return hdr, true +} + +func (mp *MPool) removePtrMetadata( + ptr unsafe.Pointer, + lease *allocationLease, +) (memHdr, bool) { + if !mp.noLock { + return gRemovePtrMetadata(ptr, lease) + } + hdr, ok := mp.ptrs[ptr] + if !ok { + if _, hasLease := mp.leases[ptr]; hasLease { + panic(moerr.NewInternalErrorNoCtx( + "account lease exists without allocation header", + )) + } + return memHdr{}, false + } + if !hdr.isAccounted() { delete(mp.ptrs, ptr) - return hdr, ok + return hdr, true + } + accountedLease, hasLease := mp.leases[ptr] + if !hasLease { + panic(moerr.NewInternalErrorNoCtx( + "accounted allocation has no account lease", + )) } + delete(mp.ptrs, ptr) + delete(mp.leases, ptr) + *lease = accountedLease + return hdr, true } func (mp *MPool) deallocateAllPtrs() { for ptr, hdr := range mp.ptrs { - if hdr.offHeap { + lease, hasLease := mp.leases[ptr] + if hdr.isAccounted() != hasLease { + panic(moerr.NewInternalErrorNoCtx( + "allocation header and account lease disagree during teardown", + )) + } + if hdr.isOffHeap() { sz := int(hdr.allocSz) profileRecordFree(uintptr(ptr), int64(sz)) simpleCAllocator().Deallocate(unsafe.Slice((*byte)(ptr), sz), uint64(sz)) + if hasLease { + lease.release(uint64(sz)) + } } } mp.ptrs = nil + mp.leases = nil } func (mp *MPool) EnableDetailRecording() { @@ -403,33 +511,23 @@ func (mp *MPool) Cap() int64 { return mp.cap } -const ( - xxxIWouldRatherUseAfterFreeCrashLaterThanLeak = true -) - func (mp *MPool) destroy() { - if mp.stats.NumAlloc.Load() < mp.stats.NumFree.Load() { - // this is a memory leak, + liveBytes := mp.stats.NumCurrBytes.Load() + if liveBytes != 0 { logutil.Errorf("mp error: %s", mp.stats.Report("")) - - // here we MUST free all the memories allocated by this mpool. - // otherwise it is a memory leak. Whoever still holds - // a pointer of this mpool is a bug (the cross pool case). - // - // We are so messed up because the cross pool free. - // If a pointer is handed out to someone else and we free here - // it will be a use after free. We risk a crash or a leak. - // Eitherway we are screwed. - if xxxIWouldRatherUseAfterFreeCrashLaterThanLeak { + if mp.noLock { + // A noLock pool exclusively owns its local pointer metadata, so its + // teardown is the physical deallocation event. + nfree := mp.stats.NumAlloc.Load() - mp.stats.NumFree.Load() mp.deallocateAllPtrs() + mp.stats.RecordManyFrees(mp.tag, nfree, liveBytes) + globalStats.RecordManyFrees(mp.tag, nfree, liveBytes) + mp.resource.recordFree(liveBytes) } + // A normal pool may have handed allocations to another owner. Its + // global pointer metadata and account lease remain authoritative until + // a later physical Free, including after this pool is unregistered. } - - // Here we just compensate whatever left over in mp.stats - // into globalStats. - globalStats.RecordManyFrees(mp.tag, - mp.stats.NumAlloc.Load()-mp.stats.NumFree.Load(), - mp.stats.NumCurrBytes.Load()) } // New a MPool. Tag is user supplied, used for debugging/diagnostics. @@ -628,8 +726,9 @@ var globalPools sync.Map const numPtrShards = 128 type ptrShard struct { - mu sync.Mutex - m map[unsafe.Pointer]memHdr + mu sync.Mutex + m map[unsafe.Pointer]memHdr + leases map[unsafe.Pointer]allocationLease } var globalPtrShards [numPtrShards]ptrShard @@ -667,14 +766,36 @@ func GlobalCap() int64 { var CapLimit = math.MaxInt32 // 2GB - 1 +func maxAllocationSize() int64 { + return int64(CapLimit) - kMemHdrSz +} + func (mp *MPool) Alloc(sz int, offHeap bool) ([]byte, error) { detailk := mp.getDetailK() return mp.allocWithDetailK(detailk, int64(sz), offHeap) } +// AllocAccounted allocates off-heap memory owned by account. Owner and site +// are bounded diagnostics; zero is invalid. Existing unaccounted callers keep +// using Alloc. +func (mp *MPool) AllocAccounted( + sz int, + account *AllocationAccount, + owner AllocationOwner, + site AllocationSite, +) ([]byte, error) { + detailk := mp.getDetailK() + request := allocationAccountRequest{ + account: account, + owner: owner, + site: site, + } + return mp.allocAccountedWithDetailK(detailk, int64(sz), request) +} + func (mp *MPool) allocWithDetailK(detailk string, sz int64, offHeap bool) ([]byte, error) { // reject unexpected alloc size. - if sz < 0 || sz > int64(CapLimit)-kMemHdrSz { + if sz < 0 || sz > maxAllocationSize() { logutil.Errorf("mpool memory allocation exceed limit with requested size %d: %s", sz, string(debug.Stack())) return nil, moerr.NewInternalErrorNoCtxf("mpool memory allocation exceed limit with requested size %d", sz) } @@ -684,48 +805,188 @@ func (mp *MPool) allocWithDetailK(detailk string, sz int64, offHeap bool) ([]byt return mp.alloc(detailk, sz, offHeap) } -func (mp *MPool) alloc(detailk string, sz int64, offHeap bool) ([]byte, error) { +func (mp *MPool) allocAccountedWithDetailK( + detailk string, + sz int64, + request allocationAccountRequest, +) ([]byte, error) { + if err := request.validate(); err != nil { + return nil, err + } + // reject unexpected alloc size. + if sz < 0 || sz > maxAllocationSize() { + logutil.Errorf("mpool memory allocation exceed limit with requested size %d: %s", sz, string(debug.Stack())) + return nil, moerr.NewInternalErrorNoCtxf("mpool memory allocation exceed limit with requested size %d", sz) + } + if sz == 0 { + return nil, nil + } + return mp.allocAccounted(detailk, sz, request) +} + +func (mp *MPool) alloc( + detailk string, + sz int64, + offHeap bool, +) ([]byte, error) { var bs []byte var err error hdr := memHdr{ poolId: mp.id, allocSz: int32(sz), - offHeap: offHeap, + } + if offHeap { + hdr.kind = memKindOffHeap } hdr.SetGuard() if offHeap { gcurr := globalStats.RecordAlloc("global", sz) if gcurr > GlobalCap() { - // compensate global globalStats.RecordFree("global", sz) return nil, moerr.NewOOMNoCtx() } mycurr := mp.stats.RecordAlloc(mp.tag, sz) if mycurr > mp.Cap() { - // compensate both global and my mp.stats.RecordFree(mp.tag, sz) globalStats.RecordFree("global", sz) return nil, moerr.NewInternalErrorNoCtxf("mpool out of space, alloc %d bytes, cap %d", sz, mp.cap) } bs, err = simpleCAllocator().Allocate(uint64(sz)) if err != nil { - panic(err) + mp.stats.RecordFree(mp.tag, sz) + globalStats.RecordFree("global", sz) + return nil, err + } + } else { + bs = make([]byte, sz) + } + + ptr := unsafe.Pointer(&bs[0]) + if err = mp.recordPtrHdr(ptr, hdr); err != nil { + if offHeap { + simpleCAllocator().Deallocate(bs, uint64(sz)) + mp.stats.RecordFree(mp.tag, sz) + globalStats.RecordFree("global", sz) } + return nil, err + } + if offHeap { mp.recordResourcePeak(mp.resource.recordAlloc(sz)) if mp.details != nil { mp.details.recordAlloc(detailk, sz) } - } else { - bs = make([]byte, sz) + profileRecordAlloc(3, uintptr(ptr), sz) } + return bs, nil +} - // always record the ptr, offHeap or not. - mp.recordPtrHdr(unsafe.Pointer(&bs[0]), hdr) - if offHeap { - profileRecordAlloc(3, uintptr(unsafe.Pointer(&bs[0])), sz) +func (mp *MPool) allocAccounted( + detailk string, + sz int64, + request allocationAccountRequest, +) ([]byte, error) { + var bs []byte + var err error + accountHeld := false + metadataHeld := false + globalHeld := false + poolHeld := false + physicalHeld := false + published := false + defer func() { + if published { + return + } + if physicalHeld { + simpleCAllocator().Deallocate(bs, uint64(sz)) + } + if poolHeld { + mp.stats.RecordFree(mp.tag, sz) + } + if globalHeld { + globalStats.RecordFree("global", sz) + } + if metadataHeld { + request.account.registry.releaseMetadata() + } + if accountHeld { + request.account.release(uint64(sz)) + } + }() + + if err = request.account.acquire(uint64(sz)); err != nil { + return nil, err + } + accountHeld = true + if err = request.reach(allocationAfterAccount); err != nil { + return nil, err + } + if err = request.account.registry.reserveMetadata(); err != nil { + return nil, err + } + metadataHeld = true + if err = request.reach(allocationAfterMetadata); err != nil { + return nil, err + } + + hdr := memHdr{ + poolId: mp.id, + allocSz: int32(sz), + kind: memKindAccountedOffHeap, + } + hdr.SetGuard() + lease := allocationLease{ + account: request.account, + owner: request.owner, + site: request.site, + } + + gcurr := globalStats.RecordAlloc("global", sz) + globalHeld = true + if err = request.reach(allocationAfterGlobalStats); err != nil { + return nil, err + } + if gcurr > GlobalCap() { + return nil, moerr.NewOOMNoCtx() + } + mycurr := mp.stats.RecordAlloc(mp.tag, sz) + poolHeld = true + if err = request.reach(allocationAfterPoolStats); err != nil { + return nil, err + } + if mycurr > mp.Cap() { + return nil, moerr.NewInternalErrorNoCtxf( + "mpool out of space, alloc %d bytes, cap %d", + sz, + mp.cap, + ) } + bs, err = simpleCAllocator().Allocate(uint64(sz)) + if err != nil { + return nil, err + } + physicalHeld = true + if err = request.reach(allocationAfterPhysical); err != nil { + return nil, err + } + + ptr := unsafe.Pointer(&bs[0]) + if err = mp.recordAccountedPtrMetadata( + ptr, + hdr, + lease, + request, + ); err != nil { + return nil, err + } + published = true + mp.recordResourcePeak(mp.resource.recordAlloc(sz)) + if mp.details != nil { + mp.details.recordAlloc(detailk, sz) + } + profileRecordAlloc(3, uintptr(ptr), sz) return bs, nil } @@ -744,7 +1005,8 @@ func (mp *MPool) freeWithDetailK(detailk string, bs []byte) { } func (mp *MPool) freePtr(detailk string, ptr unsafe.Pointer) { - hdr, ok := mp.removePtrHdr(ptr) + var lease allocationLease + hdr, ok := mp.removePtrMetadata(ptr, &lease) if !ok { // this is a double free. panic(moerr.NewInternalErrorNoCtx("invalid ptr, double free")) @@ -761,25 +1023,38 @@ func (mp *MPool) freePtr(detailk string, ptr unsafe.Pointer) { // Call profileRecordFree and the full globalStats.RecordFree // (not just NumCurrBytes.Add) so NumFree/NumFreeBytes stay // consistent with freePtrInternal. - if hdr.offHeap { + if hdr.isOffHeap() { sz := int64(hdr.allocSz) profileRecordFree(uintptr(ptr), sz) globalStats.RecordFree("global", sz) simpleCAllocator().Deallocate(unsafe.Slice((*byte)(ptr), sz), uint64(sz)) + if hdr.isAccounted() { + lease.release(uint64(sz)) + } } } else { owner := otherPool.(*MPool) owner.resource.crossPoolFree.Add(1) - owner.freePtrInternal(detailk, ptr, hdr) + owner.freePtrInternal(detailk, ptr, hdr, lease) } return } - mp.freePtrInternal(detailk, ptr, hdr) + mp.freePtrInternal(detailk, ptr, hdr, lease) } -func (mp *MPool) freePtrInternal(detailk string, ptr unsafe.Pointer, hdr memHdr) { - if !hdr.offHeap { +func (mp *MPool) freePtrInternal( + detailk string, + ptr unsafe.Pointer, + hdr memHdr, + lease allocationLease, +) { + if !hdr.isOffHeap() { + if hdr.isAccounted() { + panic(moerr.NewInternalErrorNoCtx( + "accounted allocation is not off-heap", + )) + } return } sz := int64(hdr.allocSz) @@ -798,6 +1073,9 @@ func (mp *MPool) freePtrInternal(detailk string, ptr unsafe.Pointer, hdr memHdr) } simpleCAllocator().Deallocate(unsafe.Slice((*byte)(ptr), sz), uint64(sz)) + if hdr.isAccounted() { + lease.release(uint64(sz)) + } } func (mp *MPool) reAllocWithDetailK(detailk string, old []byte, sz int64, offHeap bool, bufferMore bool) ([]byte, error) { @@ -814,7 +1092,35 @@ func (mp *MPool) reAllocWithDetailK(detailk string, old []byte, sz int64, offHea } } - ret, err := mp.allocWithDetailK(detailk, int64(newSz), offHeap) + var request *allocationAccountRequest + if ptr := unsafe.Pointer(unsafe.SliceData(old)); ptr != nil { + var lease allocationLease + hdr, ok := mp.getPtrMetadata(ptr, &lease) + if !ok { + return nil, moerr.NewInternalErrorNoCtx( + "invalid grow pointer: allocation metadata not found", + ) + } + if hdr.isAccounted() { + if !offHeap { + return nil, ErrAllocationAccountInvalid + } + accounted := allocationAccountRequest{ + account: lease.account, + owner: lease.owner, + site: lease.site, + } + request = &accounted + } + } + + var ret []byte + var err error + if request != nil { + ret, err = mp.allocAccountedWithDetailK(detailk, newSz, *request) + } else { + ret, err = mp.allocWithDetailK(detailk, newSz, offHeap) + } if err != nil { return nil, err } @@ -847,7 +1153,7 @@ func (mp *MPool) Grow2(old []byte, old2 []byte, sz int, offHeap bool) ([]byte, e // ReallocZero is like Realloc, but it clears the memory. func (mp *MPool) ReallocZero(old []byte, sz int, offHeap bool) ([]byte, error) { detailk := mp.getDetailK() - if sz < 0 || sz > CapLimit-kMemHdrSz { + if int64(sz) < 0 || int64(sz) > maxAllocationSize() { return nil, moerr.NewInternalErrorNoCtxf( "mpool memory allocation exceed limit with requested size %d", sz, @@ -870,9 +1176,10 @@ func (mp *MPool) ReallocZero(old []byte, sz int, offHeap bool) ([]byte, error) { oldptr := unsafe.Pointer(unsafe.SliceData(old)) var hdr memHdr + var lease allocationLease var ok bool if oldptr != nil { - hdr, ok = mp.getPtrHdr(oldptr) + hdr, ok = mp.getPtrMetadata(oldptr, &lease) } if !ok { if len(old) != 0 || cap(old) != 0 { @@ -913,10 +1220,33 @@ func (mp *MPool) ReallocZero(old []byte, sz int, offHeap bool) ([]byte, error) { return resized, nil } + if hdr.isAccounted() { + if !offHeap { + return nil, ErrAllocationAccountInvalid + } + request := allocationAccountRequest{ + account: lease.account, + owner: lease.owner, + site: lease.site, + } + replacement, err := mp.allocAccountedWithDetailK( + detailk, + int64(sz), + request, + ) + if err != nil { + return nil, err + } + copy(replacement, fullAllocation[:oldLength]) + clear(replacement[oldLength:]) + mp.freeWithDetailK(detailk, fullAllocation) + return replacement, nil + } + // Only resize in place when the source and destination are off-heap and // owned by this pool. Other provenance/ownership transitions use the normal // allocate-copy-free path so accounting and cross-pool cleanup stay correct. - if !hdr.offHeap || !offHeap || hdr.poolId != mp.id { + if !hdr.isOffHeap() || !offHeap || hdr.poolId != mp.id { return mp.reAllocWithDetailK( detailk, fullAllocation[:oldLength], @@ -950,16 +1280,22 @@ func (mp *MPool) ReallocZero(old []byte, sz int, offHeap bool) ([]byte, error) { return nil, err } newptr := unsafe.Pointer(&newbs[0]) - removedHdr, removed := mp.removePtrHdr(oldptr) + var removedLease allocationLease + removedHdr, removed := mp.removePtrMetadata(oldptr, &removedLease) if !removed || removedHdr != hdr { panic(moerr.NewInternalErrorNoCtx( "allocation metadata changed during realloc", )) } + if removedHdr.isAccounted() || removedLease.account != nil { + panic(moerr.NewInternalErrorNoCtx( + "unaccounted realloc removed an account lease", + )) + } newHdr := memHdr{ poolId: hdr.poolId, allocSz: int32(sz), - offHeap: true, + kind: memKindOffHeap, } newHdr.SetGuard() if err := mp.recordPtrHdr(newptr, newHdr); err != nil { @@ -1083,15 +1419,51 @@ func init() { } } -func gRecordPtr(ptr unsafe.Pointer, hdr memHdr) error { +func gRecordPtr( + ptr unsafe.Pointer, + hdr memHdr, +) error { + shard := getPtrShard(ptr) + shard.mu.Lock() + defer shard.mu.Unlock() + if _, ok := shard.m[ptr]; ok { + return moerr.NewInternalErrorNoCtx("ptr already recorded") + } + shard.m[ptr] = hdr + return nil +} + +func gRecordAccountedPtrMetadata( + ptr unsafe.Pointer, + hdr memHdr, + lease allocationLease, + request allocationAccountRequest, +) error { shard := getPtrShard(ptr) shard.mu.Lock() defer shard.mu.Unlock() - _, ok := shard.m[ptr] - if ok { + if _, ok := shard.m[ptr]; ok { return moerr.NewInternalErrorNoCtx("ptr already recorded") } + if _, ok := shard.leases[ptr]; ok { + return moerr.NewInternalErrorNoCtx("account lease already recorded") + } + committed := false + defer func() { + if !committed { + delete(shard.m, ptr) + delete(shard.leases, ptr) + } + }() shard.m[ptr] = hdr + if err := request.reach(allocationAfterHeader); err != nil { + return err + } + if shard.leases == nil { + shard.leases = make(map[unsafe.Pointer]allocationLease) + } + shard.leases[ptr] = lease + committed = true return nil } @@ -1103,13 +1475,60 @@ func gGetPtr(ptr unsafe.Pointer) (memHdr, bool) { return hdr, ok } -func gRemovePtr(ptr unsafe.Pointer) (memHdr, bool) { +func gGetPtrMetadata( + ptr unsafe.Pointer, + lease *allocationLease, +) (memHdr, bool) { shard := getPtrShard(ptr) shard.mu.Lock() defer shard.mu.Unlock() hdr, ok := shard.m[ptr] + if !ok { + return memHdr{}, false + } + if !hdr.isAccounted() { + return hdr, true + } + accountedLease, hasLease := shard.leases[ptr] + if !hasLease { + panic(moerr.NewInternalErrorNoCtx( + "accounted allocation has no account lease", + )) + } + *lease = accountedLease + return hdr, true +} + +func gRemovePtrMetadata( + ptr unsafe.Pointer, + lease *allocationLease, +) (memHdr, bool) { + shard := getPtrShard(ptr) + shard.mu.Lock() + defer shard.mu.Unlock() + hdr, ok := shard.m[ptr] + if !ok { + if _, hasLease := shard.leases[ptr]; hasLease { + panic(moerr.NewInternalErrorNoCtx( + "account lease exists without allocation header", + )) + } + return memHdr{}, false + } + if !hdr.isAccounted() { + delete(shard.m, ptr) + return hdr, true + } + accountedLease, hasLease := shard.leases[ptr] + if !hasLease { + panic(moerr.NewInternalErrorNoCtx( + "accounted allocation has no account lease", + )) + } delete(shard.m, ptr) - return hdr, ok + delete(shard.leases, ptr) + *lease = accountedLease + return hdr, true } // alignUp rounds n up to a multiple of a. a must be a power of 2. @@ -1144,7 +1563,9 @@ func roundupsize(size int64) int64 { // request. Callers which must reserve memory before growing a slice use this // helper so admission and allocation share the same growth calculation. func GrowCapacity(oldCap int64, requiredSize int64) (int64, bool) { - if oldCap < 0 || requiredSize < 0 { + maxCapacity := maxAllocationSize() + if oldCap < 0 || requiredSize < 0 || + oldCap > maxCapacity || requiredSize > maxCapacity { return 0, false } if requiredSize <= oldCap { @@ -1175,8 +1596,8 @@ func GrowCapacity(oldCap int64, requiredSize int64) (int64, bool) { if newcap < requiredSize { return 0, false } - if newcap > int64(CapLimit) && requiredSize <= int64(CapLimit) { - newcap = int64(CapLimit) + if newcap > maxCapacity { + newcap = maxCapacity } return newcap, true } diff --git a/pkg/common/mpool/mpool_test.go b/pkg/common/mpool/mpool_test.go index 133eacece5833..565a2af9967cc 100644 --- a/pkg/common/mpool/mpool_test.go +++ b/pkg/common/mpool/mpool_test.go @@ -381,6 +381,13 @@ func TestGrowCapacityValidation(t *testing.T) { capacity, ok := GrowCapacity(128, 64) require.True(t, ok) require.Equal(t, int64(128), capacity) + + maxCapacity := maxAllocationSize() + capacity, ok = GrowCapacity(maxCapacity, maxCapacity) + require.True(t, ok) + require.Equal(t, maxCapacity, capacity) + _, ok = GrowCapacity(maxCapacity, maxCapacity+1) + require.False(t, ok) } func TestUseMalloc(t *testing.T) { @@ -534,6 +541,33 @@ func TestCrossPoolFreeOnHeap(t *testing.T) { DeleteMPool(mp2) } +func TestMPoolTeardownTracksPhysicalLifetime(t *testing.T) { + t.Run("normal-pool-late-free", func(t *testing.T) { + owner := MustNew("teardown-normal-owner") + other := MustNew("teardown-normal-other") + defer DeleteMPool(other) + + globalBefore := GlobalStats().NumCurrBytes.Load() + buffer, err := owner.Alloc(64, true) + require.NoError(t, err) + DeleteMPool(owner) + require.Equal(t, globalBefore+64, GlobalStats().NumCurrBytes.Load()) + + other.Free(buffer) + require.Equal(t, globalBefore, GlobalStats().NumCurrBytes.Load()) + }) + + t.Run("no-lock-pool-owns-teardown", func(t *testing.T) { + mp := MustNewNoLock("teardown-no-lock-owner") + globalBefore := GlobalStats().NumCurrBytes.Load() + _, err := mp.Alloc(64, true) + require.NoError(t, err) + + DeleteMPool(mp) + require.Equal(t, globalBefore, GlobalStats().NumCurrBytes.Load()) + }) +} + // TestDoubleFree tests that double free is detected and panics. func TestDoubleFree(t *testing.T) { mp := MustNew("double-free-test") @@ -812,7 +846,7 @@ func TestMPoolReallocZeroUsesRecordedSourceProvenance(t *testing.T) { hdr, ok := mp.getPtrHdr(unsafe.Pointer(unsafe.SliceData(resized))) require.True(t, ok) - require.Equal(t, testCase.targetOffHeap, hdr.offHeap) + require.Equal(t, testCase.targetOffHeap, hdr.isOffHeap()) require.Equal(t, int32(newSize), hdr.allocSz) mp.Free(resized) From f18cabeedbf80aa889786ae58856bdd3ab817cb4 Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 12:17:39 +0800 Subject: [PATCH 02/61] feat: bridge allocation accounts to hash build budget --- pkg/vm/process/hashbuild_budget.go | 87 +++++++++++++-- pkg/vm/process/hashbuild_budget_test.go | 138 +++++++++++++++++++++++- 2 files changed, 216 insertions(+), 9 deletions(-) diff --git a/pkg/vm/process/hashbuild_budget.go b/pkg/vm/process/hashbuild_budget.go index dd8dba0548975..963435b7fc012 100644 --- a/pkg/vm/process/hashbuild_budget.go +++ b/pkg/vm/process/hashbuild_budget.go @@ -828,6 +828,7 @@ type HashBuildBudgetGeneration struct { id uint64 cap uint64 used uint64 + allocationUsed uint64 closed bool spillDiskCap, spillDiskUsed uint64 spillFDConfiguredCap, spillFDCap, spillFDUsed uint64 @@ -835,9 +836,12 @@ type HashBuildBudgetGeneration struct { peakUsed uint64 } +var _ mpool.AllocationCapacityController = (*HashBuildBudgetGeneration)(nil) + // HashBuildBudgetGenerationSnapshot is an immutable fixed-cardinality view. type HashBuildBudgetGenerationSnapshot struct { ID, Cap, Used, PeakUsed uint64 + AllocationUsed uint64 ReserveCount, RejectCount, ReconcileCount, ReleaseCount uint64 SpillDiskCap, SpillDiskUsed, SpillFDCap uint64 SpillFDUsed uint64 @@ -1092,7 +1096,8 @@ func (g *HashBuildBudgetGeneration) Snapshot() HashBuildBudgetGenerationSnapshot defer g.budget.mu.Unlock() return HashBuildBudgetGenerationSnapshot{ ID: g.id, Cap: g.cap, Used: g.used, PeakUsed: g.peakUsed, - ReserveCount: g.reserveCount, RejectCount: g.rejectCount, ReconcileCount: g.reconcileCount, ReleaseCount: g.releaseCount, + AllocationUsed: g.allocationUsed, + ReserveCount: g.reserveCount, RejectCount: g.rejectCount, ReconcileCount: g.reconcileCount, ReleaseCount: g.releaseCount, SpillDiskCap: g.spillDiskCap, SpillDiskUsed: g.spillDiskUsed, SpillFDCap: g.spillFDCap, SpillFDUsed: g.spillFDUsed, Closed: g.closed || g.budget.closed, } @@ -1130,10 +1135,53 @@ func (g *HashBuildBudgetGeneration) Close() { g.budget.mu.Unlock() } +// AcquireAllocationCapacity adapts allocation-accounted MPool ownership into +// the existing HashBuild query/CN policy during migration. It creates no +// independently releasable reservation token: the physical allocation lease +// is the sole release owner. +func (g *HashBuildBudgetGeneration) AcquireAllocationCapacity(size uint64) error { + if size == 0 { + return nil + } + _, err := g.reserve(size, true) + return err +} + +// ReleaseAllocationCapacity is called only by physical MPool Free through the +// allocation account. A mismatch is an ownership invariant failure. +func (g *HashBuildBudgetGeneration) ReleaseAllocationCapacity(size uint64) { + if size == 0 { + return + } + if g == nil || g.budget == nil { + panic("nil hash build allocation capacity controller") + } + b := g.budget + b.mu.Lock() + if g.allocationUsed < size || g.used < size || b.aggregateUsed < size { + b.mu.Unlock() + panic("hash build allocation capacity release underflow") + } + g.allocationUsed -= size + g.used -= size + b.aggregateUsed -= size + g.releaseCount++ + b.mu.Unlock() + observeHashBuildBudget("memory", "release", "query", size) + observeHashBuildBudget("memory", "release", "cn", size) +} + // Reserve performs the required two-level sequence: charge CN aggregate, // then charge query-CN. If query-CN rejects, aggregate is rolled back before // returning, so callers never observe a partial reservation. func (g *HashBuildBudgetGeneration) Reserve(size uint64) (*HashBuildReservation, error) { + return g.reserve(size, false) +} + +func (g *HashBuildBudgetGeneration) reserve( + size uint64, + allocationOwned bool, +) (*HashBuildReservation, error) { if g == nil || g.budget == nil { return nil, &HashBuildBudgetError{Kind: HashBuildBudgetErrorInvalid, Message: "nil hash build generation"} } @@ -1158,9 +1206,13 @@ func (g *HashBuildBudgetGeneration) Reserve(size uint64) (*HashBuildReservation, b.mu.Unlock() return nil, err } - token, firstErr, aggregateRejected := g.reserveLocked(size, false) + token, firstErr, aggregateRejected := g.reserveLocked( + size, + false, + allocationOwned, + ) b.mu.Unlock() - if token != nil { + if firstErr == nil && !aggregateRejected { observeHashBuildBudget("memory", "reserve", "query", size) observeHashBuildBudget("memory", "reserve", "cn", size) } @@ -1178,9 +1230,13 @@ func (g *HashBuildBudgetGeneration) Reserve(size uint64) (*HashBuildReservation, return nil, err } b.mu.Lock() - token, firstErr, aggregateRejected := g.reserveLocked(size, false) + token, firstErr, aggregateRejected := g.reserveLocked( + size, + false, + allocationOwned, + ) b.mu.Unlock() - if token != nil { + if firstErr == nil && !aggregateRejected { observeHashBuildBudget("memory", "reserve", "query", size) observeHashBuildBudget("memory", "reserve", "cn", size) } @@ -1202,9 +1258,13 @@ func (g *HashBuildBudgetGeneration) Reserve(size uint64) (*HashBuildReservation, } b.mu.Lock() - token, err, aggregateRejected := g.reserveLocked(size, true) + token, err, aggregateRejected := g.reserveLocked( + size, + true, + allocationOwned, + ) b.mu.Unlock() - if token != nil { + if err == nil && !aggregateRejected { observeHashBuildBudget("memory", "reserve", "query", size) observeHashBuildBudget("memory", "reserve", "cn", size) } @@ -1217,7 +1277,11 @@ func (g *HashBuildBudgetGeneration) Reserve(size uint64) (*HashBuildReservation, // reserveLocked attempts one memory reservation. b.mu must be held. The bool // result identifies an aggregate-cap failure so Reserve can trigger a forced // live-ceiling refresh without counting a transient failure as a rejection. -func (g *HashBuildBudgetGeneration) reserveLocked(size uint64, recordAggregateReject bool) (*HashBuildReservation, error, bool) { +func (g *HashBuildBudgetGeneration) reserveLocked( + size uint64, + recordAggregateReject bool, + allocationOwned bool, +) (*HashBuildReservation, error, bool) { b := g.budget if b.closed || g.closed { g.rejectCount++ @@ -1246,6 +1310,13 @@ func (g *HashBuildBudgetGeneration) reserveLocked(size uint64, recordAggregateRe if g.used > g.peakUsed { g.peakUsed = g.used } + if allocationOwned { + if size > math.MaxUint64-g.allocationUsed { + panic("hash build allocation capacity overflow") + } + g.allocationUsed += size + return nil, nil, false + } return &HashBuildReservation{budget: b, generation: g, core: &hashBuildReservationCore{size: size}}, nil, false } diff --git a/pkg/vm/process/hashbuild_budget_test.go b/pkg/vm/process/hashbuild_budget_test.go index 118c3fdace694..2c3cc29266bba 100644 --- a/pkg/vm/process/hashbuild_budget_test.go +++ b/pkg/vm/process/hashbuild_budget_test.go @@ -115,6 +115,107 @@ func TestHashBuildBudgetAdmissionIdentifiesResource(t *testing.T) { } } +func TestHashBuildBudgetAllocationAccountAdapter(t *testing.T) { + budget := MustNewHashBuildBudget(10, 10) + generation, err := budget.OpenGeneration(1) + if err != nil { + t.Fatal(err) + } + legacy, err := generation.Reserve(4) + if err != nil { + t.Fatal(err) + } + + registry, err := commonmpool.NewAllocationAccountRegistry(1, 2) + if err != nil { + t.Fatal(err) + } + account, err := registry.OpenWithController(10, generation) + if err != nil { + t.Fatal(err) + } + mp := commonmpool.MustNew("hash-build-allocation-account-adapter") + defer commonmpool.DeleteMPool(mp) + + noMetadataRegistry, err := commonmpool.NewAllocationAccountRegistry(1, 0) + if err != nil { + t.Fatal(err) + } + noMetadataAccount, err := noMetadataRegistry.OpenWithController( + 10, + generation, + ) + if err != nil { + t.Fatal(err) + } + if _, err = mp.AllocAccounted( + 1, + noMetadataAccount, + 1, + 1, + ); !errors.Is(err, commonmpool.ErrAllocationMetadataSlots) { + t.Fatalf("metadata admission error = %v", err) + } + if generation.Used() != 4 || + generation.Snapshot().AllocationUsed != 0 { + t.Fatalf("metadata failure leaked controller capacity: %+v", + generation.Snapshot()) + } + noMetadataAccount.Seal() + if _, err = noMetadataRegistry.Finalize(noMetadataAccount); err != nil { + t.Fatal(err) + } + + buffer, err := mp.AllocAccounted(6, account, 1, 1) + if err != nil { + t.Fatal(err) + } + snapshot := generation.Snapshot() + if snapshot.Used != 10 || snapshot.AllocationUsed != 6 { + t.Fatalf("unexpected generation snapshot: %+v", snapshot) + } + if account.Snapshot().Used != 6 { + t.Fatalf("account used = %d, want 6", account.Snapshot().Used) + } + + if _, err = mp.AllocAccounted(1, account, 1, 1); !errors.Is( + err, + ErrHashBuildBudgetAdmission, + ) { + t.Fatalf("combined legacy/exact admission error = %v", err) + } + if account.Snapshot().Used != 6 || + registry.LiveAllocationMetadata() != 1 { + t.Fatal("failed adapter admission did not roll back") + } + + generation.Close() + if _, err = mp.AllocAccounted(1, account, 1, 1); !errors.Is( + err, + ErrHashBuildBudgetClosed, + ) { + t.Fatalf("closed generation admission error = %v", err) + } + if account.Snapshot().Used != 6 || + registry.LiveAllocationMetadata() != 1 { + t.Fatal("closed adapter admission did not roll back") + } + + // Close rejects new capacity but cannot invalidate a live physical lease. + mp.Free(buffer) + if generation.Used() != 4 || generation.Snapshot().AllocationUsed != 0 { + t.Fatalf("allocation release did not retain only legacy charge: %+v", + generation.Snapshot()) + } + if !legacy.Release() || generation.Used() != 0 { + t.Fatal("legacy reservation did not release") + } + account.Seal() + if _, err = registry.Finalize(account); err != nil { + t.Fatal(err) + } +} + func TestHashBuildBudgetQueryRejectRollsBackCN(t *testing.T) { b := MustNewHashBuildBudget(10, 4) g1, _ := b.OpenGeneration(1) @@ -1702,7 +1803,7 @@ func TestHashBuildBudgetDefensiveAndProviderFailurePaths(t *testing.T) { } closedGeneration.closed = true closedBudget.mu.Lock() - _, err, rejected := closedGeneration.reserveLocked(1, true) + _, err, rejected := closedGeneration.reserveLocked(1, true, false) closedBudget.mu.Unlock() if rejected || !errors.Is(err, ErrHashBuildBudgetClosed) { t.Fatalf("closed reserveLocked: rejected=%v err=%v", rejected, err) @@ -1859,6 +1960,41 @@ func BenchmarkHashBuildBudgetReserveCachedProvider(b *testing.B) { b.ReportMetric(float64(calls.Load()), "provider-calls") } +func BenchmarkHashBuildBudgetAllocationAccount(b *testing.B) { + const capacity = uint64(1 << 60) + budget := MustNewHashBuildBudget(capacity, capacity) + generation, err := budget.OpenGeneration(1) + if err != nil { + b.Fatal(err) + } + registry, err := commonmpool.NewAllocationAccountRegistry(1, 1) + if err != nil { + b.Fatal(err) + } + account, err := registry.OpenWithController(capacity, generation) + if err != nil { + b.Fatal(err) + } + mp := commonmpool.MustNew("hash-build-allocation-account-benchmark") + defer commonmpool.DeleteMPool(mp) + + b.ReportAllocs() + b.SetBytes(64 << 10) + b.ResetTimer() + for range b.N { + buffer, allocErr := mp.AllocAccounted(64<<10, account, 1, 1) + if allocErr != nil { + b.Fatal(allocErr) + } + mp.Free(buffer) + } + b.StopTimer() + account.Seal() + if _, err = registry.Finalize(account); err != nil { + b.Fatal(err) + } +} + func TestResolveHashBuildCeiling(t *testing.T) { const gib = uint64(1 << 30) got, err := ResolveHashBuildCeiling(HashBuildCeilingInputs{ From 5676b9b46ae5bcb33a0d951653d2f580a84370cb Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 13:40:00 +0800 Subject: [PATCH 03/61] feat: propagate allocation accounts through vectors --- .../batch/allocation_account_test.go | 256 +++++++++ pkg/container/batch/batch.go | 94 ++- pkg/container/batch/types.go | 4 + pkg/container/vector/allocation_account.go | 279 +++++++++ .../vector/allocation_account_test.go | 534 ++++++++++++++++++ pkg/container/vector/tools.go | 2 +- pkg/container/vector/vector.go | 232 ++++++-- pkg/container/vector/versions.go | 9 + 8 files changed, 1352 insertions(+), 58 deletions(-) create mode 100644 pkg/container/batch/allocation_account_test.go create mode 100644 pkg/container/vector/allocation_account.go create mode 100644 pkg/container/vector/allocation_account_test.go diff --git a/pkg/container/batch/allocation_account_test.go b/pkg/container/batch/allocation_account_test.go new file mode 100644 index 0000000000000..ddc74d21be964 --- /dev/null +++ b/pkg/container/batch/allocation_account_test.go @@ -0,0 +1,256 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package batch + +import ( + "bytes" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/stretchr/testify/require" +) + +type testBatchAllocationAccount struct { + registry *mpool.AllocationAccountRegistry + account *mpool.AllocationAccount + selection *vector.AllocationAccountSelection +} + +func newTestBatchAllocationAccount( + t *testing.T, + allocationSlots uint64, +) testBatchAllocationAccount { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, allocationSlots) + require.NoError(t, err) + account, err := registry.Open(16 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2) + require.NoError(t, err) + return testBatchAllocationAccount{ + registry: registry, + account: account, + selection: selection, + } +} + +func finalizeTestBatchAllocationAccount( + t *testing.T, + state testBatchAllocationAccount, +) { + t.Helper() + snapshot := state.account.Seal() + require.Zero(t, snapshot.Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + _, err := state.registry.Finalize(state.account) + require.NoError(t, err) +} + +func newBatchAllocationTestSource( + t *testing.T, + mp *mpool.MPool, + selection *vector.AllocationAccountSelection, +) *Batch { + t.Helper() + bat := NewWithSchema( + true, + []string{"id", "value"}, + []types.Type{types.T_int64.ToType(), types.T_varchar.ToType()}, + ) + if selection != nil { + require.NoError(t, bat.SetAllocationAccount(selection)) + } + for i := 0; i < 32; i++ { + require.NoError(t, vector.AppendFixed(bat.Vecs[0], int64(i), false, mp)) + require.NoError( + t, + vector.AppendBytes( + bat.Vecs[1], + []byte("batch allocation payload that is not inline"), + false, + mp, + ), + ) + } + bat.SetRowCount(32) + return bat +} + +func TestBatchAllocationAccountCloneDupAndWindow(t *testing.T) { + state := newTestBatchAllocationAccount(t, 64) + mp := mpool.MustNewZero() + source := newBatchAllocationTestSource(t, mp, state.selection) + sourceUsed := state.account.Snapshot().Used + require.NotZero(t, sourceUsed) + _, err := source.Clone(mp, false) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + + cloned, err := source.Clone(mp, true) + require.NoError(t, err) + require.Same(t, state.selection, cloned.AllocationAccountSelection()) + require.Greater(t, state.account.Snapshot().Used, sourceUsed) + cloned.Clean(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + dup, err := source.Dup(mp) + require.NoError(t, err) + require.Same(t, state.selection, dup.AllocationAccountSelection()) + dup.Clean(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + selectedColumns, err := source.CloneSelectedColumns( + []int{1}, + []string{"value"}, + mp, + ) + require.NoError(t, err) + require.Same(t, state.selection, selectedColumns.AllocationAccountSelection()) + selectedColumns.Clean(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + window, err := source.Window(4, 12) + require.NoError(t, err) + require.Same(t, state.selection, window.AllocationAccountSelection()) + for _, vec := range window.Vecs { + require.Nil(t, vec.AllocationAccountSelection()) + } + window.Clean(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + var encoded bytes.Buffer + data, err := source.MarshalBinaryWithBuffer(&encoded, true) + require.NoError(t, err) + decoded := NewOffHeapEmpty() + require.NoError(t, decoded.SetAllocationAccount(state.selection)) + require.NoError(t, decoded.UnmarshalFromReader(bytes.NewReader(data), mp)) + require.Same(t, state.selection, decoded.AllocationAccountSelection()) + for _, vec := range decoded.Vecs { + require.Same(t, state.selection, vec.AllocationAccountSelection()) + } + decoded.Clean(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + aliasDecoded := NewWithSchema( + true, + source.Attrs, + []types.Type{types.T_int64.ToType(), types.T_varchar.ToType()}, + ) + require.NoError(t, aliasDecoded.SetAllocationAccount(state.selection)) + require.NoError(t, aliasDecoded.UnmarshalBinaryWithAnyMp(data, mp)) + require.Same(t, state.selection, aliasDecoded.AllocationAccountSelection()) + for _, vec := range aliasDecoded.Vecs { + require.Nil(t, vec.AllocationAccountSelection()) + } + aliasDecoded.Clean(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchAllocationAccountDestinationCloneUnionAndReuse(t *testing.T) { + state := newTestBatchAllocationAccount(t, 64) + mp := mpool.MustNewZero() + source := newBatchAllocationTestSource(t, mp, nil) + + destination := NewWithSchema( + true, + source.Attrs, + []types.Type{types.T_int64.ToType(), types.T_varchar.ToType()}, + ) + require.NoError(t, destination.SetAllocationAccount(state.selection)) + require.NoError(t, source.CloneTo(destination, mp)) + require.NotZero(t, state.account.Snapshot().Used) + require.Equal(t, source.RowCount(), destination.RowCount()) + + destination.FreeColumns(mp) + require.Zero(t, state.account.Snapshot().Used) + for _, vec := range destination.Vecs { + require.Same(t, state.selection, vec.AllocationAccountSelection()) + } + + require.NoError(t, destination.Union(source, []int64{1, 3, 5, 7}, mp)) + require.Equal(t, 4, destination.RowCount()) + require.NotZero(t, state.account.Snapshot().Used) + + destination.Clean(mp) + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchAllocationAccountCloneRollback(t *testing.T) { + state := newTestBatchAllocationAccount(t, 1) + mp := mpool.MustNewZero() + source := NewWithSchema( + true, + []string{"left", "right"}, + []types.Type{types.T_int64.ToType(), types.T_int64.ToType()}, + ) + for i := range source.Vecs { + require.NoError(t, vector.AppendFixed(source.Vecs[i], int64(i), false, mp)) + } + source.SetRowCount(1) + + destination := NewWithSchema( + true, + source.Attrs, + []types.Type{types.T_int64.ToType(), types.T_int64.ToType()}, + ) + require.NoError(t, destination.SetAllocationAccount(state.selection)) + err := source.CloneTo(destination, mp) + require.ErrorIs(t, err, mpool.ErrAllocationMetadataSlots) + require.Zero(t, state.account.Snapshot().Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + require.Nil(t, destination.AllocationAccountSelection()) + + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchAllocationAccountConfigurationIsAtomic(t *testing.T) { + state := newTestBatchAllocationAccount(t, 8) + mp := mpool.MustNewZero() + + onHeap := NewWithSchema( + false, + nil, + []types.Type{types.T_int64.ToType()}, + ) + require.ErrorIs( + t, + onHeap.SetAllocationAccount(state.selection), + mpool.ErrAllocationAccountInvalid, + ) + + offHeap := NewWithSchema( + true, + nil, + []types.Type{types.T_int64.ToType(), types.T_int64.ToType()}, + ) + require.NoError(t, vector.AppendFixed(offHeap.Vecs[1], int64(1), false, mp)) + require.ErrorIs( + t, + offHeap.SetAllocationAccount(state.selection), + mpool.ErrAllocationAccountInvalid, + ) + require.Nil(t, offHeap.Vecs[0].AllocationAccountSelection()) + require.Nil(t, offHeap.AllocationAccountSelection()) + + offHeap.Clean(mp) + onHeap.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} diff --git a/pkg/container/batch/batch.go b/pkg/container/batch/batch.go index ece832ab7d7de..589fe21f66600 100644 --- a/pkg/container/batch/batch.go +++ b/pkg/container/batch/batch.go @@ -200,6 +200,7 @@ func (c *batchUnmarshalCursor) readUint32() (uint32, error) { } func (bat *Batch) UnmarshalBinaryWithAnyMp(data []byte, mp *mpool.MPool) (err error) { + allocationAccount := bat.allocationAccount cursor := batchUnmarshalCursor{data: data} rowCount, err := cursor.readInt64() if err != nil { @@ -234,6 +235,7 @@ func (bat *Batch) UnmarshalBinaryWithAnyMp(data []byte, mp *mpool.MPool) (err er } } bat.Clean(mp) + bat.allocationAccount = allocationAccount } bat.Vecs = make([]*vector.Vector, vecsLen) } @@ -290,6 +292,14 @@ func (bat *Batch) UnmarshalBinaryWithAnyMp(data []byte, mp *mpool.MPool) (err er } vecs[i].Free(mp) } + // UnmarshalBinary installs aliases into vecData. An empty accounted + // receiver must explicitly drop its future-allocation selection first; + // the Batch retains the destination context for a later owned copy. + if vecs[i].AllocationAccountSelection() != nil { + if err := vecs[i].SetAllocationAccount(nil); err != nil { + return err + } + } if err := vecs[i].UnmarshalBinary(vecData); err != nil { return err } @@ -400,6 +410,7 @@ func (bat *Batch) UnmarshalBinaryWithAnyMp(data []byte, mp *mpool.MPool) (err er } func (bat *Batch) UnmarshalFromReader(r io.Reader, mp *mpool.MPool) (err error) { + allocationAccount := bat.allocationAccount i64, err := types.ReadInt64(r) if err != nil { return err @@ -413,11 +424,17 @@ func (bat *Batch) UnmarshalFromReader(r io.Reader, mp *mpool.MPool) (err error) if l != len(bat.Vecs) { if len(bat.Vecs) > 0 { bat.Clean(mp) + bat.allocationAccount = allocationAccount } bat.Vecs = make([]*vector.Vector, l) for i := range bat.Vecs { if bat.offHeap { bat.Vecs[i] = vector.NewOffHeapVec() + if allocationAccount != nil { + if err := bat.Vecs[i].SetAllocationAccount(allocationAccount); err != nil { + return err + } + } } else { bat.Vecs[i] = vector.NewVecFromReuse() } @@ -544,27 +561,70 @@ func (bat *Batch) SetAttributes(attrs []string) { bat.Attrs = attrs } +// AllocationAccountSelection returns the immutable destination selection used +// by this batch's owned off-heap vectors. +func (bat *Batch) AllocationAccountSelection() *vector.AllocationAccountSelection { + if bat == nil { + return nil + } + return bat.allocationAccount +} + +// SetAllocationAccount configures every existing empty destination vector as +// one transaction. Existing physical allocations are never relabeled. +func (bat *Batch) SetAllocationAccount( + selection *vector.AllocationAccountSelection, +) error { + if bat == nil || (selection != nil && !bat.offHeap) { + return mpool.ErrAllocationAccountInvalid + } + for _, vec := range bat.Vecs { + if vec != nil { + if err := vec.CanSetAllocationAccount(selection); err != nil { + return err + } + } + } + for _, vec := range bat.Vecs { + if vec != nil { + if err := vec.SetAllocationAccount(selection); err != nil { + panic(err) + } + } + } + bat.allocationAccount = selection + return nil +} + +func (bat *Batch) configureOwnedVector(vec *vector.Vector) { + if vec == nil { + return + } + vec.SetOffHeap(bat.offHeap) + if bat.allocationAccount != nil { + if err := vec.SetAllocationAccount(bat.allocationAccount); err != nil { + panic(err) + } + } +} + func (bat *Batch) InsertVector( pos int32, attr string, vec *vector.Vector, ) { + bat.configureOwnedVector(vec) bat.Vecs = append(bat.Vecs, nil) copy(bat.Vecs[pos+1:], bat.Vecs[pos:]) bat.Vecs[pos] = vec - if vec != nil { - vec.SetOffHeap(bat.offHeap) - } bat.Attrs = append(bat.Attrs, "") copy(bat.Attrs[pos+1:], bat.Attrs[pos:]) bat.Attrs[pos] = attr } func (bat *Batch) SetVector(pos int32, vec *vector.Vector) { + bat.configureOwnedVector(vec) bat.Vecs[pos] = vec - if vec != nil { - vec.SetOffHeap(bat.offHeap) - } } func (bat *Batch) GetVector(pos int32) *vector.Vector { @@ -587,6 +647,11 @@ func (bat *Batch) CloneSelectedColumns( cloned.Vecs[idx] = vector.NewVec(typ) } } + if bat.allocationAccount != nil { + if err = cloned.SetAllocationAccount(bat.allocationAccount); err != nil { + return nil, err + } + } if err = bat.CloneSelectedColumnsTo(selectCols, cloned, mp); err != nil { cloned.Clean(mp) cloned = nil @@ -626,6 +691,7 @@ func (bat *Batch) SelectColumns(cols []int, attrs []string) *Batch { rbat := NewWithSize(len(cols)) rbat.Attrs = attrs rbat.offHeap = bat.offHeap + rbat.allocationAccount = bat.allocationAccount for i, col := range cols { rbat.Vecs[i] = bat.Vecs[col] } @@ -650,6 +716,7 @@ func (bat *Batch) Clean(m *mpool.MPool) { bat.Attrs = nil bat.ExtraBuf = nil bat.SetRowCount(0) + bat.allocationAccount = nil } func (bat *Batch) Last() bool { @@ -681,6 +748,11 @@ func (bat *Batch) FreeColumns(m *mpool.MPool) { for _, vec := range bat.Vecs { if vec != nil { vec.Free(m) + if bat.allocationAccount != nil { + if err := vec.SetAllocationAccount(bat.allocationAccount); err != nil { + panic(err) + } + } } } } @@ -705,11 +777,19 @@ func (bat *Batch) GetSchema() (attrs []string, attrTypes []types.Type) { } func (bat *Batch) Clone(mp *mpool.MPool, offHeap bool) (*Batch, error) { + if bat.allocationAccount != nil && !offHeap { + return nil, mpool.ErrAllocationAccountInvalid + } var ( cloned *Batch attrs, attrTypes = bat.GetSchema() ) cloned = NewWithSchema(offHeap, attrs, attrTypes) + if offHeap && bat.allocationAccount != nil { + if err := cloned.SetAllocationAccount(bat.allocationAccount); err != nil { + return nil, err + } + } cloned.Recursive = bat.Recursive err := bat.CloneTo(cloned, mp) if err != nil { @@ -870,6 +950,8 @@ func (bat *Batch) Window(start, end int) (*Batch, error) { b := NewWithSize(len(bat.Vecs)) var err error b.Attrs = bat.Attrs + b.offHeap = bat.offHeap + b.allocationAccount = bat.allocationAccount for i, vec := range bat.Vecs { b.Vecs[i], err = vec.Window(start, end) if err != nil { diff --git a/pkg/container/batch/types.go b/pkg/container/batch/types.go index 6229e3f98b22a..f3d1a708579b4 100644 --- a/pkg/container/batch/types.go +++ b/pkg/container/batch/types.go @@ -56,4 +56,8 @@ type Batch struct { // row count of batch, to instead of old len(Zs). rowCount int offHeap bool + + // allocationAccount is the destination selection for owned off-heap + // vectors created or reused by this batch. Alias vectors do not copy it. + allocationAccount *vector.AllocationAccountSelection } diff --git a/pkg/container/vector/allocation_account.go b/pkg/container/vector/allocation_account.go new file mode 100644 index 0000000000000..cdbf963c3c5bf --- /dev/null +++ b/pkg/container/vector/allocation_account.go @@ -0,0 +1,279 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package vector + +import ( + "fmt" + "io" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" +) + +// AllocationAccountSelection is an immutable choice for the first owned +// off-heap data and area allocations of a Vector. The physical MPool +// allocation metadata remains the sole owner of the resulting charge. +// +// A selection may be shared by all vectors owned by one Batch. Views do not +// copy it: they share storage and therefore must not create a second charge. +type AllocationAccountSelection struct { + account *mpool.AllocationAccount + owner mpool.AllocationOwner + dataSite mpool.AllocationSite + areaSite mpool.AllocationSite +} + +func NewAllocationAccountSelection( + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, + dataSite mpool.AllocationSite, + areaSite mpool.AllocationSite, +) (*AllocationAccountSelection, error) { + selection := &AllocationAccountSelection{ + account: account, + owner: owner, + dataSite: dataSite, + areaSite: areaSite, + } + if err := selection.validate(); err != nil { + return nil, err + } + return selection, nil +} + +func (s *AllocationAccountSelection) validate() error { + if s == nil || s.account == nil || s.account.Handle() == 0 || + s.owner < mpool.AllocationOwnerMin || + s.owner > mpool.AllocationOwnerMax || + s.dataSite < mpool.AllocationSiteMin || + s.areaSite < mpool.AllocationSiteMin { + return mpool.ErrAllocationAccountInvalid + } + return nil +} + +// AllocationAccountSelection returns the immutable selection used by this +// vector's future owned allocations. It is nil for legacy vectors and views. +func (v *Vector) AllocationAccountSelection() *AllocationAccountSelection { + if v == nil { + return nil + } + return v.allocationAccount +} + +// CanSetAllocationAccount reports whether selection can be installed without +// converting or relabeling an existing physical allocation. +func (v *Vector) CanSetAllocationAccount( + selection *AllocationAccountSelection, +) error { + if v == nil { + return mpool.ErrAllocationAccountInvalid + } + if selection != nil { + if err := selection.validate(); err != nil { + return err + } + if !v.offHeap { + return fmt.Errorf( + "%w: allocation-accounted vector must be off-heap", + mpool.ErrAllocationAccountInvalid, + ) + } + } + if v.allocationAccount == selection { + return nil + } + if v.hasBackingStorage() { + return fmt.Errorf( + "%w: vector already has backing storage", + mpool.ErrAllocationAccountInvalid, + ) + } + return nil +} + +func (v *Vector) hasBackingStorage() bool { + return cap(v.data) != 0 || cap(v.area) != 0 +} + +// SetAllocationAccount selects the account used by future owned data and area +// allocations. It is intentionally explicit and is legal only before the +// first backing allocation. Reset retains the selection; Free clears it. +func (v *Vector) SetAllocationAccount( + selection *AllocationAccountSelection, +) error { + if err := v.CanSetAllocationAccount(selection); err != nil { + return err + } + v.allocationAccount = selection + return nil +} + +func (v *Vector) allocData(mp *mpool.MPool, size int) ([]byte, error) { + return v.allocOwned(mp, size, v.offHeap, true) +} + +func (v *Vector) allocArea(mp *mpool.MPool, size int) ([]byte, error) { + return v.allocOwned(mp, size, v.offHeap, false) +} + +func (v *Vector) allocOwned( + mp *mpool.MPool, + size int, + offHeap bool, + data bool, +) ([]byte, error) { + if mp == nil { + return nil, moerr.NewInternalErrorNoCtx( + "vector allocation does not have a mpool", + ) + } + if v.allocationAccount == nil { + return mp.Alloc(size, offHeap) + } + if !offHeap { + return nil, fmt.Errorf( + "%w: accounted allocation must be off-heap", + mpool.ErrAllocationAccountInvalid, + ) + } + site := v.allocationAccount.areaSite + if data { + site = v.allocationAccount.dataSite + } + return mp.AllocAccounted( + size, + v.allocationAccount.account, + v.allocationAccount.owner, + site, + ) +} + +func (v *Vector) growData(mp *mpool.MPool, size int) ([]byte, error) { + return v.growOwned(mp, v.data, size, true) +} + +func (v *Vector) growArea(mp *mpool.MPool, size int) ([]byte, error) { + return v.growOwned(mp, v.area, size, false) +} + +func (v *Vector) growOwned( + mp *mpool.MPool, + old []byte, + size int, + data bool, +) ([]byte, error) { + if size <= cap(old) { + return old[:size], nil + } + if mp == nil { + return nil, moerr.NewInternalErrorNoCtx( + "vector growth does not have a mpool", + ) + } + if cap(old) != 0 || v.allocationAccount == nil { + return mp.Grow(old, size, v.offHeap) + } + + capacity, ok := mpool.GrowCapacity(0, int64(size)) + if !ok { + return nil, moerr.NewInternalErrorNoCtxf( + "invalid mpool grow capacity, old %d, required %d", + cap(old), + size, + ) + } + buf, err := v.allocOwned(mp, int(capacity), true, data) + if err != nil { + return nil, err + } + return buf[:size], nil +} + +func (v *Vector) growArea2( + mp *mpool.MPool, + src []byte, + size int, +) ([]byte, error) { + oldLen := len(v.area) + if size < oldLen+len(src) { + return nil, moerr.NewInternalErrorNoCtxf( + "mpool grow2 actually shrinks, %d+%d, %d", + oldLen, + len(src), + size, + ) + } + grown, err := v.growArea(mp, size) + if err != nil { + return nil, err + } + copy(grown[oldLen:oldLen+len(src)], src) + return grown, nil +} + +func (v *Vector) readSizeBytes( + r io.Reader, + mp *mpool.MPool, + data bool, +) (int32, []byte, error) { + size, err := types.ReadInt32(r) + if err != nil { + return 0, nil, err + } + var old []byte + if data { + old = v.data + } else { + old = v.area + } + if size == 0 { + if old != nil { + old = old[:0] + } + if data { + v.data = old + } else { + v.area = old + } + return 0, old, nil + } + if size < 0 { + return size, nil, moerr.NewInvalidInputNoCtx( + "negative vector buffer size", + ) + } + var buf []byte + if data { + buf, err = v.growData(mp, int(size)) + } else { + buf, err = v.growArea(mp, int(size)) + } + if err != nil { + return 0, nil, err + } + // Grow may already have freed the old allocation. Publish its replacement + // before reading so a short reader still leaves one reachable cleanup owner. + if data { + v.data = buf + } else { + v.area = buf + } + if _, err = io.ReadFull(r, buf); err != nil { + return size, buf, err + } + return size, buf, nil +} diff --git a/pkg/container/vector/allocation_account_test.go b/pkg/container/vector/allocation_account_test.go new file mode 100644 index 0000000000000..7accd35d508d3 --- /dev/null +++ b/pkg/container/vector/allocation_account_test.go @@ -0,0 +1,534 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package vector + +import ( + "bytes" + "errors" + "math/rand" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/stretchr/testify/require" +) + +const ( + testVectorAllocationOwner mpool.AllocationOwner = 1 + testVectorDataAllocationSite mpool.AllocationSite = 1 + testVectorAreaAllocationSite mpool.AllocationSite = 2 +) + +type testVectorAllocationAccount struct { + registry *mpool.AllocationAccountRegistry + account *mpool.AllocationAccount + selection *AllocationAccountSelection +} + +func newTestVectorAllocationAccount( + t testing.TB, + limit uint64, + allocationSlots uint64, +) testVectorAllocationAccount { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, allocationSlots) + require.NoError(t, err) + account, err := registry.Open(limit) + require.NoError(t, err) + selection, err := NewAllocationAccountSelection( + account, + testVectorAllocationOwner, + testVectorDataAllocationSite, + testVectorAreaAllocationSite, + ) + require.NoError(t, err) + return testVectorAllocationAccount{ + registry: registry, + account: account, + selection: selection, + } +} + +func finalizeTestVectorAllocationAccount( + t testing.TB, + state testVectorAllocationAccount, +) { + t.Helper() + snapshot := state.account.Seal() + require.Zero(t, snapshot.Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + _, err := state.registry.Finalize(state.account) + require.NoError(t, err) +} + +func newAccountedTestVector( + t testing.TB, + typ types.Type, + selection *AllocationAccountSelection, +) *Vector { + t.Helper() + vec := NewOffHeapVecWithType(typ) + require.NoError(t, vec.SetAllocationAccount(selection)) + return vec +} + +func TestVectorAllocationAccountConfiguration(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 8) + mp := mpool.MustNewZero() + + _, err := NewAllocationAccountSelection( + nil, + testVectorAllocationOwner, + testVectorDataAllocationSite, + testVectorAreaAllocationSite, + ) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + + onHeap := NewVec(types.T_int64.ToType()) + require.ErrorIs( + t, + onHeap.SetAllocationAccount(state.selection), + mpool.ErrAllocationAccountInvalid, + ) + + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.Same(t, state.selection, vec.AllocationAccountSelection()) + require.NoError(t, AppendFixed(vec, int64(1), false, mp)) + require.ErrorIs( + t, + vec.SetAllocationAccount(nil), + mpool.ErrAllocationAccountInvalid, + ) + require.Panics(t, func() { + vec.SetOffHeap(false) + }) + _, err = vec.Dup(mp) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + _, err = vec.CloneToFlatCompact(mp) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + + vec.Free(mp) + require.Nil(t, vec.AllocationAccountSelection()) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountFixedResetReuseAndFree(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 16) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + + require.NoError(t, vec.PreExtend(128, mp)) + initial := state.account.Snapshot() + require.Equal(t, uint64(cap(vec.data)), initial.Used) + require.Equal(t, uint64(1), state.registry.LiveAllocationMetadata()) + + for i := 0; i < 64; i++ { + require.NoError(t, AppendFixed(vec, int64(i), false, mp)) + } + require.Equal(t, initial.Used, state.account.Snapshot().Used) + + vec.ResetWithSameType() + require.Equal(t, initial.Used, state.account.Snapshot().Used) + for i := 0; i < 64; i++ { + require.NoError(t, AppendFixed(vec, int64(i*2), false, mp)) + } + require.Equal(t, initial.Used, state.account.Snapshot().Used) + + growAt := vec.Capacity() + 1 + for vec.Length() < growAt { + require.NoError(t, AppendFixed(vec, int64(vec.Length()), false, mp)) + } + grown := state.account.Snapshot() + require.Equal(t, uint64(cap(vec.data)), grown.Used) + require.Greater(t, grown.Peak, grown.Used) + require.Equal(t, uint64(1), state.registry.LiveAllocationMetadata()) + + require.NoError(t, vec.Shuffle([]int64{0, 2, 4, 6}, mp)) + shuffled := state.account.Snapshot() + require.Equal(t, uint64(cap(vec.data)), shuffled.Used) + require.Equal(t, uint64(1), state.registry.LiveAllocationMetadata()) + + vec.Free(mp) + require.Zero(t, state.account.Snapshot().Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountVarlenaDataAndArea(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 16) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_varchar.ToType(), state.selection) + + require.NoError(t, vec.PreExtendWithArea(64, 4096, mp)) + initial := state.account.Snapshot() + require.Equal(t, uint64(cap(vec.data)+cap(vec.area)), initial.Used) + require.Equal(t, uint64(2), state.registry.LiveAllocationMetadata()) + + for i := 0; i < 32; i++ { + require.NoError(t, AppendBytes(vec, bytes.Repeat([]byte{byte(i)}, 64), false, mp)) + } + require.Equal(t, initial.Used, state.account.Snapshot().Used) + + vec.ResetWithSameType() + require.Equal(t, initial.Used, state.account.Snapshot().Used) + require.NoError(t, AppendBytes(vec, bytes.Repeat([]byte("r"), 128), false, mp)) + require.Equal(t, initial.Used, state.account.Snapshot().Used) + + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountLeavesGoBitmapsUnaccounted(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 8) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, vec.PreExtend(128, mp)) + before := state.account.Snapshot().Used + + // Null/group bitmaps still use Go []uint64. They are an explicit activation + // blocker and must not be mislabeled as part of the off-heap vector charge. + vec.SetAllNulls(32 * 1024) + vec.GetGrouping().AddRange(0, 32*1024) + require.Equal(t, before, state.account.Snapshot().Used) + + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountViewAndDeepCopy(t *testing.T) { + stateA := newTestVectorAllocationAccount(t, 1<<20, 32) + stateB := newTestVectorAllocationAccount(t, 1<<20, 32) + mp := mpool.MustNewZero() + + source := newAccountedTestVector(t, types.T_varchar.ToType(), stateA.selection) + for _, value := range [][]byte{ + []byte("first value that is not inline"), + []byte("second value that is not inline"), + []byte("third value that is not inline"), + } { + require.NoError(t, AppendBytes(source, value, false, mp)) + } + sourceUsed := stateA.account.Snapshot().Used + + view, err := source.Window(1, 3) + require.NoError(t, err) + require.Nil(t, view.AllocationAccountSelection()) + require.Equal(t, sourceUsed, stateA.account.Snapshot().Used) + view.Free(mp) + require.Equal(t, sourceUsed, stateA.account.Snapshot().Used) + + dup, err := source.DupOffHeap(mp) + require.NoError(t, err) + require.Same(t, stateA.selection, dup.AllocationAccountSelection()) + require.Greater(t, stateA.account.Snapshot().Used, sourceUsed) + dup.Free(mp) + require.Equal(t, sourceUsed, stateA.account.Snapshot().Used) + + crossOwner, err := source.DupOffHeapWithAllocation(mp, stateB.selection) + require.NoError(t, err) + require.Equal(t, sourceUsed, stateA.account.Snapshot().Used) + require.NotZero(t, stateB.account.Snapshot().Used) + crossOwner.Free(mp) + require.Zero(t, stateB.account.Snapshot().Used) + + window, err := source.CloneWindowWithAllocation(1, 3, mp, stateB.selection) + require.NoError(t, err) + require.Equal(t, source.GetBytesAt(1), window.GetBytesAt(0)) + require.Equal(t, source.GetBytesAt(2), window.GetBytesAt(1)) + window.Free(mp) + + compact, err := source.CloneToFlatCompactWithAllocation(mp, stateB.selection) + require.NoError(t, err) + require.Equal(t, source.Length(), compact.Length()) + compact.Free(mp) + + source.Free(mp) + finalizeTestVectorAllocationAccount(t, stateA) + finalizeTestVectorAllocationAccount(t, stateB) +} + +func TestVectorAllocationAccountCopyRollback(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 1) + mp := mpool.MustNewZero() + source := NewVec(types.T_varchar.ToType()) + require.NoError( + t, + AppendBytes( + source, + bytes.Repeat([]byte("payload"), 32), + false, + mp, + ), + ) + + _, err := source.DupOffHeapWithAllocation(mp, state.selection) + require.ErrorIs(t, err, mpool.ErrAllocationMetadataSlots) + require.Zero(t, state.account.Snapshot().Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + + source.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountCrossPoolFreeAndSeal(t *testing.T) { + t.Run("cross pool free", func(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 8) + ownerPool := mpool.MustNewZero() + freeingPool := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, vec.PreExtend(128, ownerPool)) + require.NotZero(t, state.account.Snapshot().Used) + + vec.Free(freeingPool) + require.Zero(t, state.account.Snapshot().Used) + finalizeTestVectorAllocationAccount(t, state) + }) + + t.Run("sealed before allocation", func(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 8) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + state.account.Seal() + + err := AppendFixed(vec, int64(1), false, mp) + require.ErrorIs(t, err, mpool.ErrAllocationAccountSealed) + require.Zero(t, state.account.Snapshot().Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + vec.Free(mp) + _, err = state.registry.Finalize(state.account) + require.NoError(t, err) + }) +} + +func TestVectorAllocationAccountRandomizedAppendAndSelection(t *testing.T) { + state := newTestVectorAllocationAccount(t, 8<<20, 64) + mp := mpool.MustNewZero() + rng := rand.New(rand.NewSource(26459)) + + fixed := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + var fixedExpected []int64 + for i := 0; i < 2_000; i++ { + value := rng.Int63() + fixedExpected = append(fixedExpected, value) + require.NoError(t, AppendFixed(fixed, value, false, mp)) + } + require.Equal(t, fixedExpected, MustFixedColNoTypeCheck[int64](fixed)) + + varlen := newAccountedTestVector(t, types.T_varchar.ToType(), state.selection) + var expected [][]byte + for i := 0; i < 1_000; i++ { + size := rng.Intn(96) + value := make([]byte, size) + _, err := rng.Read(value) + require.NoError(t, err) + expected = append(expected, append([]byte(nil), value...)) + require.NoError(t, AppendBytes(varlen, value, false, mp)) + } + for i := range expected { + require.True(t, bytes.Equal(expected[i], varlen.GetBytesAt(i))) + } + + selected := newAccountedTestVector(t, types.T_varchar.ToType(), state.selection) + sels := []int64{1, 3, 7, 11, 23, 101, 509, 999} + require.NoError(t, selected.Union(varlen, sels, mp)) + for i, sel := range sels { + require.True(t, bytes.Equal(expected[sel], selected.GetBytesAt(i))) + } + + fixed.Free(mp) + varlen.Free(mp) + selected.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountDecodeCopyAndReader(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 16) + mp := mpool.MustNewZero() + source := NewVec(types.T_varchar.ToType()) + require.NoError( + t, + AppendBytes(source, []byte("decoded payload that is not inline"), false, mp), + ) + encoded, err := source.MarshalBinary() + require.NoError(t, err) + + copied := newAccountedTestVector(t, types.T_varchar.ToType(), state.selection) + require.NoError(t, copied.UnmarshalBinaryWithCopy(encoded, mp)) + require.Equal(t, source.GetBytesAt(0), copied.GetBytesAt(0)) + require.Equal(t, uint64(2), state.registry.LiveAllocationMetadata()) + require.ErrorIs( + t, + copied.UnmarshalBinaryWithCopy(encoded, mp), + mpool.ErrAllocationAccountInvalid, + ) + require.ErrorIs( + t, + copied.UnmarshalBinary(encoded), + mpool.ErrAllocationAccountInvalid, + ) + copied.Free(mp) + + aliased := newAccountedTestVector(t, types.T_varchar.ToType(), state.selection) + require.ErrorIs( + t, + aliased.UnmarshalBinary(encoded), + mpool.ErrAllocationAccountInvalid, + ) + aliased.Free(mp) + + readerDecoded := newAccountedTestVector(t, types.T_varchar.ToType(), state.selection) + require.NoError(t, readerDecoded.UnmarshalWithReader(bytes.NewReader(encoded), mp)) + require.Equal(t, source.GetBytesAt(0), readerDecoded.GetBytesAt(0)) + readerDecoded.Free(mp) + + short := newAccountedTestVector(t, types.T_varchar.ToType(), state.selection) + dataHeader := 1 + types.TSize + 4 + 4 + require.Error( + t, + short.UnmarshalWithReader( + bytes.NewReader(encoded[:dataHeader+1]), + mp, + ), + ) + require.NotZero(t, state.account.Snapshot().Used) + short.Free(mp) + + source.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func BenchmarkVectorAllocationAccount(b *testing.B) { + const rows = 8192 + mp := mpool.MustNewZero() + state := newTestVectorAllocationAccount(b, 1<<40, 64) + + b.Run("legacy-fixed-preextend-free", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + vec := NewOffHeapVecWithType(types.T_int64.ToType()) + if err := vec.PreExtend(rows, mp); err != nil { + b.Fatal(err) + } + vec.Free(mp) + } + }) + b.Run("accounted-fixed-preextend-free", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + vec := NewOffHeapVecWithType(types.T_int64.ToType()) + if err := vec.SetAllocationAccount(state.selection); err != nil { + b.Fatal(err) + } + if err := vec.PreExtend(rows, mp); err != nil { + b.Fatal(err) + } + vec.Free(mp) + } + }) + b.Run("legacy-varlen-preextend-free", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + vec := NewOffHeapVecWithType(types.T_varchar.ToType()) + if err := vec.PreExtendWithArea(rows, 1<<20, mp); err != nil { + b.Fatal(err) + } + vec.Free(mp) + } + }) + b.Run("accounted-varlen-preextend-free", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + vec := NewOffHeapVecWithType(types.T_varchar.ToType()) + if err := vec.SetAllocationAccount(state.selection); err != nil { + b.Fatal(err) + } + if err := vec.PreExtendWithArea(rows, 1<<20, mp); err != nil { + b.Fatal(err) + } + vec.Free(mp) + } + }) + b.Run("accounted-fixed-reset-reuse", func(b *testing.B) { + vec := newAccountedTestVector(b, types.T_int64.ToType(), state.selection) + if err := vec.PreExtend(rows, mp); err != nil { + b.Fatal(err) + } + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + vec.ResetWithSameType() + } + b.StopTimer() + vec.Free(mp) + }) + + finalizeTestVectorAllocationAccount(b, state) +} + +func TestVectorAllocationAccountErrorsAreTyped(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1, 1) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + err := vec.PreExtend(1, mp) + require.True(t, errors.Is(err, mpool.ErrAllocationAccountCapacity)) + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountHelperBoundaries(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 8) + mp := mpool.MustNewZero() + var nilVec *Vector + require.Nil(t, nilVec.AllocationAccountSelection()) + require.ErrorIs( + t, + nilVec.CanSetAllocationAccount(state.selection), + mpool.ErrAllocationAccountInvalid, + ) + + vec := newAccountedTestVector(t, types.T_varchar.ToType(), state.selection) + _, err := vec.allocOwned(mp, 1, false, true) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + _, err = vec.growData(nil, 1) + require.Error(t, err) + _, err = vec.growData(mp, mpool.CapLimit) + require.Error(t, err) + _, err = vec.growArea2(mp, []byte{1}, 0) + require.Error(t, err) + + require.NoError(t, vec.PreExtendWithArea(1, 128, mp)) + used := state.account.Snapshot().Used + zero := int32(0) + size, area, err := vec.readSizeBytes( + bytes.NewReader(types.EncodeInt32(&zero)), + mp, + false, + ) + require.NoError(t, err) + require.Zero(t, size) + require.Empty(t, area) + require.Equal(t, used, state.account.Snapshot().Used) + + negative := int32(-1) + _, _, err = vec.readSizeBytes( + bytes.NewReader(types.EncodeInt32(&negative)), + mp, + false, + ) + require.Error(t, err) + + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} diff --git a/pkg/container/vector/tools.go b/pkg/container/vector/tools.go index 8920131cf729e..3c731326a091d 100644 --- a/pkg/container/vector/tools.go +++ b/pkg/container/vector/tools.go @@ -210,7 +210,7 @@ func extend(v *Vector, rows int, m *mpool.MPool) error { tgtLen := v.length + rows tgtDataCap := tgtLen * v.typ.TypeSize() if tgtDataCap > cap(v.data) { - ndata, err := m.Grow(v.data, tgtDataCap, v.offHeap) + ndata, err := v.growData(m, tgtDataCap) if err != nil { return err } diff --git a/pkg/container/vector/vector.go b/pkg/container/vector/vector.go index f05b9024ed47b..d5074f4eda726 100644 --- a/pkg/container/vector/vector.go +++ b/pkg/container/vector/vector.go @@ -67,6 +67,10 @@ type Vector struct { isBin bool offHeap bool + + // allocationAccount selects the account for this vector's first owned + // off-heap data and area allocations. Physical MPool metadata owns release. + allocationAccount *AllocationAccountSelection } func toSliceOfLengthNoTypeCheck[T any](vec *Vector, length int) []T { @@ -230,6 +234,9 @@ func (v *Vector) SetTypeAndFixData(typ types.Type, mp *mpool.MPool) { } func (v *Vector) SetOffHeap(offHeap bool) { + if !offHeap && v.allocationAccount != nil { + panic("allocation-accounted vector must remain off-heap") + } v.offHeap = offHeap } @@ -520,7 +527,7 @@ func NewVecWithDataCopy( vec.length = length var err error if len(data) > 0 { - vec.data, err = mp.Alloc(len(data), false) + vec.data, err = vec.allocData(mp, len(data)) if err != nil { vec.Free(mp) return nil, err @@ -528,7 +535,7 @@ func NewVecWithDataCopy( copy(vec.data, data) } if len(area) > 0 { - vec.area, err = mp.Alloc(len(area), false) + vec.area, err = vec.allocArea(mp, len(area)) if err != nil { vec.Free(mp) return nil, err @@ -727,6 +734,7 @@ func (v *Vector) Free(mp *mpool.MPool) { v.gsp.Reset() v.sorted = false v.isBin = false + v.allocationAccount = nil // if !v.OnUsed || v.OnPut { // panic("free vector which unalloc or in put list") @@ -820,6 +828,12 @@ func (v *Vector) UnmarshalBinaryTrusted(data []byte) error { } func (v *Vector) unmarshalBinary(data []byte, validateValues bool) error { + if v.allocationAccount != nil { + return fmt.Errorf( + "%w: cannot install aliases in an accounted vector", + mpool.ErrAllocationAccountInvalid, + ) + } read := func(size int) ([]byte, error) { if size < 0 || size > len(data) { return nil, io.ErrUnexpectedEOF @@ -900,6 +914,9 @@ func (v *Vector) unmarshalBinary(data []byte, validateValues bool) error { v.cantFreeData = true v.cantFreeArea = true + // The decoded buffers alias the input byte slice. They have no physical + // MPool ownership and therefore cannot retain an allocation selection. + v.allocationAccount = nil return nil } @@ -1023,6 +1040,12 @@ func canonicalVectorTypeSize(typ types.Type) (int, error) { } func (v *Vector) UnmarshalBinaryWithCopy(data []byte, mp *mpool.MPool) error { + if v.allocationAccount != nil && v.hasBackingStorage() { + return fmt.Errorf( + "%w: cannot replace accounted vector storage without Free", + mpool.ErrAllocationAccountInvalid, + ) + } var err error // read class @@ -1041,7 +1064,7 @@ func (v *Vector) UnmarshalBinaryWithCopy(data []byte, mp *mpool.MPool) error { dataLen := int(types.DecodeUint32(data[:4])) data = data[4:] if dataLen > 0 { - v.data, err = mp.Alloc(dataLen, v.offHeap) + v.data, err = v.allocData(mp, dataLen) if err != nil { return err } @@ -1053,7 +1076,7 @@ func (v *Vector) UnmarshalBinaryWithCopy(data []byte, mp *mpool.MPool) error { areaLen := int(types.DecodeUint32(data[:4])) data = data[4:] if areaLen > 0 { - v.area, err = mp.Alloc(areaLen, v.offHeap) + v.area, err = v.allocArea(mp, areaLen) if err != nil { return err } @@ -1095,7 +1118,7 @@ func (v *Vector) UnmarshalWithReader(r io.Reader, mp *mpool.MPool) error { } // read data - dataLen, dataBuf, err := types.ReadSizeBytesMp(r, v.data, mp, v.offHeap) + dataLen, dataBuf, err := v.readSizeBytes(r, mp, true) if err != nil { return err } @@ -1104,7 +1127,7 @@ func (v *Vector) UnmarshalWithReader(r io.Reader, mp *mpool.MPool) error { } // read area - areaLen, areaBuf, err := types.ReadSizeBytesMp(r, v.area, mp, v.offHeap) + areaLen, areaBuf, err := v.readSizeBytes(r, mp, false) if err != nil { return err } @@ -1170,7 +1193,7 @@ func (v *Vector) PreExtendWithArea(rows int, extraAreaSize int, mp *mpool.MPool) // grow area var err error oldSz := len(area1) - area1, err = mp.Grow(area1, voff+extraAreaSize, v.offHeap) + area1, err = v.growArea(mp, voff+extraAreaSize) if err != nil { return err } @@ -1184,29 +1207,53 @@ func (v *Vector) PreExtendWithArea(rows int, extraAreaSize int, mp *mpool.MPool) // Dup use to copy an identical vector func (v *Vector) Dup(mp *mpool.MPool) (*Vector, error) { - return v.dup(mp, false, v.offHeap) + if v.allocationAccount != nil { + return nil, fmt.Errorf( + "%w: accounted vector duplication requires an off-heap destination", + mpool.ErrAllocationAccountInvalid, + ) + } + return v.dup(mp, false, v.offHeap, nil) } // DupOffHeap copies a vector with all owned backing data allocated off-heap. func (v *Vector) DupOffHeap(mp *mpool.MPool) (*Vector, error) { - return v.dup(mp, true, true) + return v.dup(mp, true, true, v.allocationAccount) } -func (v *Vector) dup(mp *mpool.MPool, offHeap, areaOffHeap bool) (*Vector, error) { - if v.IsConstNull() { - return NewConstNull(v.typ, v.Length(), mp), nil - } - - var err error +// DupOffHeapWithAllocation copies a vector into an explicitly selected +// destination account. Passing nil creates a legacy unaccounted destination. +func (v *Vector) DupOffHeapWithAllocation( + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (*Vector, error) { + return v.dup(mp, true, true, selection) +} +func (v *Vector) dup( + mp *mpool.MPool, + offHeap bool, + areaOffHeap bool, + selection *AllocationAccountSelection, +) (*Vector, error) { w := NewVecFromReuse() w.offHeap = offHeap + if selection != nil { + if err := w.SetAllocationAccount(selection); err != nil { + return nil, err + } + } w.class = v.class w.typ = v.typ w.length = v.length w.sorted = v.sorted w.GetNulls().InitWith(v.GetNulls()) + if v.IsConstNull() { + return w, nil + } + + var err error dataLen := v.typ.TypeSize() if v.IsConst() { if err := extend(w, 1, mp); err != nil { @@ -1223,7 +1270,7 @@ func (v *Vector) dup(mp *mpool.MPool, offHeap, areaOffHeap bool) (*Vector, error copy(w.data, v.data[:dataLen]) if len(v.area) > 0 { - if w.area, err = mp.Alloc(len(v.area), areaOffHeap); err != nil { + if w.area, err = w.allocOwned(mp, len(v.area), areaOffHeap, false); err != nil { w.Free(mp) return nil, err } @@ -1236,7 +1283,37 @@ func (v *Vector) dup(mp *mpool.MPool, offHeap, areaOffHeap bool) (*Vector, error // retains varlen payload referenced by the vector's logical rows, so stale or // unreferenced bytes in area are not propagated into batch memory accounting. func (v *Vector) CloneToFlatCompact(mp *mpool.MPool) (*Vector, error) { - w := NewVec(v.typ) + if v.allocationAccount != nil { + return nil, fmt.Errorf( + "%w: accounted compact clone requires a destination selection", + mpool.ErrAllocationAccountInvalid, + ) + } + return v.cloneToFlatCompact(mp, nil) +} + +// CloneToFlatCompactWithAllocation creates an off-heap compact copy under the +// explicit destination selection. +func (v *Vector) CloneToFlatCompactWithAllocation( + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (*Vector, error) { + return v.cloneToFlatCompact(mp, selection) +} + +func (v *Vector) cloneToFlatCompact( + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (*Vector, error) { + var w *Vector + if selection == nil { + w = NewVec(v.typ) + } else { + w = NewOffHeapVecWithType(v.typ) + if err := w.SetAllocationAccount(selection); err != nil { + return nil, err + } + } if v.class != FLAT || (!v.typ.IsFixedLen() && !v.typ.IsVarlen()) { if err := GetUnionAllFunction(v.typ, mp)(w, v); err != nil { w.Free(mp) @@ -1276,7 +1353,7 @@ func (v *Vector) CloneToFlatCompact(mp *mpool.MPool) (*Vector, error) { } if totalArea > 0 { var err error - w.area, err = mp.Alloc(totalArea, w.offHeap) + w.area, err = w.allocArea(mp, totalArea) if err != nil { w.Free(mp) return nil, err @@ -2391,7 +2468,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err return err } if sz := len(v.area) + len(w.area); sz > cap(v.area) { - area, err := mp.Grow(v.area, sz, v.offHeap) + area, err := v.growArea(mp, sz) if err != nil { return err } @@ -2796,7 +2873,7 @@ func pregrowVarlenaArea(vec *Vector, totalBytes int, mp *mpool.MPool) error { return nil } origLen := len(vec.area) - grown, err := mp.Grow(vec.area, need, vec.offHeap) + grown, err := vec.growArea(mp, need) if err != nil { return err } @@ -3120,10 +3197,20 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp if len(w.area) > 0 { // preserve mpool semantics: append within cap, else mpool Grow2 (so // v.area stays mpool-tracked rather than escaping to the Go heap). - if baseOff+len(w.area) <= cap(v.area) || mp == nil { + if baseOff+len(w.area) <= cap(v.area) { v.area = append(v.area, w.area...) - } else if v.area, err = mp.Grow2(v.area, w.area, baseOff+len(w.area), v.offHeap); err != nil { - return err + } else if mp == nil { + if v.allocationAccount != nil { + return moerr.NewInternalErrorNoCtx( + "accounted vector area growth does not have a mpool", + ) + } + v.area = append(v.area, w.area...) + } else { + v.area, err = v.growArea2(mp, w.area, baseOff+len(w.area)) + if err != nil { + return err + } } } // one memmove of the header array; inline varlenas carry their bytes here. @@ -4244,7 +4331,7 @@ func shuffleFixedNoTypeCheck[T types.FixedSizeT](v *Vector, sels []int64, mp *mp ns := len(sels) var vs []T ToFixedColNoTypeCheck(v, &vs) - data, err := mp.Alloc(ns*v.GetType().TypeSize(), v.offHeap) + data, err := v.allocData(mp, ns*v.GetType().TypeSize()) if err != nil { return err } @@ -4341,31 +4428,44 @@ func (v *Vector) Window(start, end int) (*Vector, error) { // CloneWindow Deep copies the content from start to end into another vector. Afterwise it's safe to destroy the original one. func (v *Vector) CloneWindow(start, end int, mp *mpool.MPool) (*Vector, error) { + return v.CloneWindowWithAllocation( + start, + end, + mp, + v.allocationAccount, + ) +} + +// CloneWindowWithAllocation deep-copies a window into an explicitly selected +// off-heap destination account. +func (v *Vector) CloneWindowWithAllocation( + start int, + end int, + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (*Vector, error) { if start == end { - return NewOffHeapVecWithType(v.typ), nil + w := NewOffHeapVecWithType(v.typ) + if selection != nil { + if err := w.SetAllocationAccount(selection); err != nil { + return nil, err + } + } + return w, nil } if end > v.Length() { panic(fmt.Sprintf("CloneWindow end %d >= length %d", end, v.Length())) } - if v.IsConstNull() { - return NewConstNull(v.typ, end-start, mp), nil - } else if v.IsConst() { - if v.typ.IsVarlen() { - return NewConstBytes(v.typ, v.GetBytesAt(0), end-start, mp) - } else { - vec := NewOffHeapVecWithType(v.typ) - vec.class = v.class - vec.data = make([]byte, len(v.data)) - copy(vec.data, v.data) - vec.length = end - start - vec.cantFreeArea = true - vec.cantFreeData = true - vec.sorted = v.sorted - return vec, nil + w := NewOffHeapVecWithType(v.typ) + if selection != nil { + if err := w.SetAllocationAccount(selection); err != nil { + return nil, err } } - w := NewOffHeapVecWithType(v.typ) if err := v.CloneWindowTo(w, start, end, mp); err != nil { + if mp != nil { + w.Free(mp) + } return nil, err } return w, nil @@ -4383,15 +4483,24 @@ func (v *Vector) CloneWindowTo(w *Vector, start, end int, mp *mpool.MPool) error } else if v.IsConst() { if v.typ.IsVarlen() { w.class = CONSTANT - SetConstBytes(v, v.GetBytesAt(0), end-start, mp) - return nil + return SetConstBytes(w, v.GetBytesAt(0), end-start, mp) } else { + if mp == nil { + if w.allocationAccount != nil { + return moerr.NewInternalErrorNoCtx( + "accounted vector clone does not have a mpool", + ) + } + w.data = make([]byte, len(v.data)) + w.cantFreeData = true + } else { + if err := w.PreExtend(1, mp); err != nil { + return err + } + copy(w.data, v.data) + } w.class = v.class - w.data = make([]byte, len(v.data)) - copy(w.data, v.data) w.length = end - start - w.cantFreeArea = true - w.cantFreeData = true w.sorted = v.sorted return nil } @@ -4399,6 +4508,11 @@ func (v *Vector) CloneWindowTo(w *Vector, start, end int, mp *mpool.MPool) error nulls.Range(&v.nsp, uint64(start), uint64(end), uint64(start), &w.nsp) length := (end - start) * v.typ.TypeSize() if mp == nil { + if w.allocationAccount != nil { + return moerr.NewInternalErrorNoCtx( + "accounted vector clone does not have a mpool", + ) + } w.data = make([]byte, length) copy(w.data, v.data[start*v.typ.TypeSize():end*v.typ.TypeSize()]) w.length = end - start @@ -5482,14 +5596,25 @@ func BuildVarlenaNoInline(vec *Vector, v1 *types.Varlena, bs *[]byte, m *mpool.M vlen := len(*bs) area1 := vec.GetArea() voff := len(area1) - if voff+vlen <= cap(area1) || m == nil { + if voff+vlen <= cap(area1) { + area1 = append(area1, *bs...) + v1.SetOffsetLen(uint32(voff), uint32(vlen)) + vec.area = area1 + return nil + } + if m == nil { + if vec.allocationAccount != nil { + return moerr.NewInternalErrorNoCtx( + "accounted vector area growth does not have a mpool", + ) + } area1 = append(area1, *bs...) v1.SetOffsetLen(uint32(voff), uint32(vlen)) vec.area = area1 return nil } var err error - area1, err = m.Grow2(area1, *bs, voff+vlen, vec.offHeap) + area1, err = vec.growArea2(m, *bs, voff+vlen) if err != nil { return err } @@ -5507,13 +5632,18 @@ func BuildVarlenaNoInlineFromByteJson(vec *Vector, v1 *types.Varlena, bj bytejso if voff+vlen > cap(area1) && m != nil { // Pass nil to Grow2, we can grow area1 to voff+vlen without // copy bytejson data. - area1, err = m.Grow2(area1, nil, voff+vlen, vec.offHeap) + area1, err = vec.growArea2(m, nil, voff+vlen) if err != nil { return err } area1[voff] = byte(bj.Type) copy(area1[voff+1:voff+vlen], bj.Data) } else { + if voff+vlen > cap(area1) && vec.allocationAccount != nil { + return moerr.NewInternalErrorNoCtx( + "accounted vector area growth does not have a mpool", + ) + } area1 = append(area1, byte(bj.Type)) area1 = append(area1, bj.Data...) } @@ -5606,7 +5736,7 @@ func BuildVarlenaFromByteJsonEncoded( } if int(newAreaLen) > cap(vec.area) { - newArea, err := m.Grow2(vec.area, nil, int(newAreaLen), vec.offHeap) + newArea, err := vec.growArea2(m, nil, int(newAreaLen)) if err != nil { return err } diff --git a/pkg/container/vector/versions.go b/pkg/container/vector/versions.go index 7abc1d4b3a7fb..b9fdcd7abf317 100644 --- a/pkg/container/vector/versions.go +++ b/pkg/container/vector/versions.go @@ -16,7 +16,9 @@ package vector import ( "bytes" + "fmt" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/types" ) @@ -70,6 +72,12 @@ func (v *Vector) MarshalBinaryWithBufferV1(buf *bytes.Buffer) error { } func (v *Vector) UnmarshalBinaryV1(data []byte) error { + if v.allocationAccount != nil { + return fmt.Errorf( + "%w: cannot install aliases in an accounted vector", + mpool.ErrAllocationAccountInvalid, + ) + } // read class v.class = int(data[0]) data = data[1:] @@ -116,6 +124,7 @@ func (v *Vector) UnmarshalBinaryV1(data []byte) error { v.cantFreeData = true v.cantFreeArea = true + v.allocationAccount = nil return nil } From d2464556c00a354deb88c9c8b7c2785ac1536c41 Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 14:39:33 +0800 Subject: [PATCH 04/61] feat: propagate allocation accounts through expressions --- .../mpool/allocation_account_mpool_test.go | 64 ++ pkg/common/mpool/mpool.go | 33 + pkg/container/vector/allocation_account.go | 99 +++ pkg/container/vector/functionTools.go | 105 +++- .../vector/function_result_allocation_test.go | 110 ++++ pkg/sql/colexec/evalExpression.go | 585 +++++++++++++++--- pkg/sql/colexec/evalExpressionReset.go | 64 +- pkg/sql/colexec/eval_expression_allocation.go | 236 +++++++ .../eval_expression_allocation_test.go | 510 +++++++++++++++ pkg/sql/util/eval_expr_util.go | 34 +- pkg/sql/util/eval_expr_util_test.go | 49 ++ 11 files changed, 1746 insertions(+), 143 deletions(-) create mode 100644 pkg/container/vector/function_result_allocation_test.go create mode 100644 pkg/sql/colexec/eval_expression_allocation.go create mode 100644 pkg/sql/colexec/eval_expression_allocation_test.go diff --git a/pkg/common/mpool/allocation_account_mpool_test.go b/pkg/common/mpool/allocation_account_mpool_test.go index cd2366a19405e..83c23cf9c8067 100644 --- a/pkg/common/mpool/allocation_account_mpool_test.go +++ b/pkg/common/mpool/allocation_account_mpool_test.go @@ -123,6 +123,70 @@ func TestMPoolAccountedAllocGrowFree(t *testing.T) { finalizeTestAllocationAccount(t, registry, account) } +func TestMPoolMakeSliceAccounted(t *testing.T) { + registry, account := newTestAllocationAccount(t, 64, 2) + mp := MustNew("accounted-typed-slice") + defer DeleteMPool(mp) + + values, err := MakeSliceAccounted[int64]( + 4, + mp, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + require.Len(t, values, 4) + require.Equal(t, uint64(32), account.Snapshot().Used) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + for i := range values { + values[i] = int64(i + 1) + } + require.Equal(t, []int64{1, 2, 3, 4}, values) + + _, err = MakeSliceAccounted[int64]( + 5, + mp, + account, + testAllocationOwner, + testAllocationSite, + ) + require.ErrorIs(t, err, ErrAllocationAccountCapacity) + require.Equal(t, uint64(32), account.Snapshot().Used) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + + empty, err := MakeSliceAccounted[int64]( + 0, + mp, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + require.Nil(t, empty) + _, err = MakeSliceAccounted[int64]( + -1, + mp, + account, + testAllocationOwner, + testAllocationSite, + ) + require.ErrorIs(t, err, ErrAllocationAccountInvalid) + _, err = MakeSliceAccounted[struct{}]( + 1, + mp, + account, + testAllocationOwner, + testAllocationSite, + ) + require.ErrorIs(t, err, ErrAllocationAccountInvalid) + + FreeSlice(mp, values[:0]) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) +} + func TestMPoolAccountedRollback(t *testing.T) { t.Run("account-capacity", func(t *testing.T) { registry, account := newTestAllocationAccount(t, 63, 1) diff --git a/pkg/common/mpool/mpool.go b/pkg/common/mpool/mpool.go index adf1e20f59ae0..d281385ff60f9 100644 --- a/pkg/common/mpool/mpool.go +++ b/pkg/common/mpool/mpool.go @@ -1327,6 +1327,38 @@ func MakeSlice[T any](n int, mp *MPool, offHeap bool) ([]T, error) { return makeSliceWithCapWithDetailK[T](detailk, n, n, mp, offHeap) } +// MakeSliceAccounted allocates an off-heap typed slice whose physical +// allocation is owned by account. FreeSlice releases the resulting charge. +func MakeSliceAccounted[T any]( + n int, + mp *MPool, + account *AllocationAccount, + owner AllocationOwner, + site AllocationSite, +) ([]T, error) { + if n < 0 { + return nil, ErrAllocationAccountInvalid + } + if n == 0 { + return nil, nil + } + var value T + elementSize := unsafe.Sizeof(value) + maxSize := maxAllocationSize() + if elementSize == 0 || + maxSize <= 0 || + uint64(n) > uint64(maxSize)/uint64(elementSize) { + return nil, ErrAllocationAccountInvalid + } + size := int(uint64(n) * uint64(elementSize)) + bs, err := mp.AllocAccounted(size, account, owner, site) + if err != nil { + return nil, err + } + values := unsafe.Slice((*T)(unsafe.Pointer(&bs[0])), n) + return values[:n:n], nil +} + func MakeSliceArgs[T any](mp *MPool, offHeap bool, args ...T) ([]T, error) { detailk := mp.getDetailK() ret, err := makeSliceWithCapWithDetailK[T](detailk, len(args), len(args), mp, offHeap) @@ -1341,6 +1373,7 @@ func FreeSlice[T any](mp *MPool, bs []T) { if cap(bs) == 0 { return } + bs = bs[:1] detailk := mp.getDetailK() mp.freePtr(detailk, unsafe.Pointer(&bs[0])) } diff --git a/pkg/container/vector/allocation_account.go b/pkg/container/vector/allocation_account.go index cdbf963c3c5bf..96ed56ad09194 100644 --- a/pkg/container/vector/allocation_account.go +++ b/pkg/container/vector/allocation_account.go @@ -54,6 +54,105 @@ func NewAllocationAccountSelection( return selection, nil } +// NewOffHeapVecWithTypeAndAllocation constructs an empty owning Vector whose +// future data and area allocations use selection. +func NewOffHeapVecWithTypeAndAllocation( + typ types.Type, + selection *AllocationAccountSelection, +) (*Vector, error) { + vec := NewOffHeapVecWithType(typ) + if err := vec.SetAllocationAccount(selection); err != nil { + vec.Free(nil) + return nil, err + } + return vec, nil +} + +// NewConstNullWithAllocation constructs a constant NULL Vector with dormant +// allocation provenance for any future owned backing. +func NewConstNullWithAllocation( + typ types.Type, + length int, + selection *AllocationAccountSelection, +) (*Vector, error) { + vec, err := NewOffHeapVecWithTypeAndAllocation(typ, selection) + if err != nil { + return nil, err + } + vec.class = CONSTANT + vec.length = length + return vec, nil +} + +// NewConstFixedWithAllocation constructs an off-heap constant fixed-width +// Vector and charges its physical backing to selection. +func NewConstFixedWithAllocation[T any]( + typ types.Type, + value T, + length int, + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (*Vector, error) { + vec, err := NewOffHeapVecWithTypeAndAllocation(typ, selection) + if err != nil { + return nil, err + } + vec.class = CONSTANT + if length > 0 { + if err = SetConstFixed(vec, value, length, mp); err != nil { + vec.Free(mp) + return nil, err + } + } + return vec, nil +} + +// NewConstBytesWithAllocation constructs an off-heap constant varlen Vector +// and charges data and area independently through selection. +func NewConstBytesWithAllocation( + typ types.Type, + value []byte, + length int, + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (*Vector, error) { + vec, err := NewOffHeapVecWithTypeAndAllocation(typ, selection) + if err != nil { + return nil, err + } + vec.class = CONSTANT + if length > 0 { + if err = SetConstBytes(vec, value, length, mp); err != nil { + vec.Free(mp) + return nil, err + } + } + return vec, nil +} + +// NewConstArrayWithAllocation constructs an off-heap constant array Vector +// and charges data and area independently through selection. +func NewConstArrayWithAllocation[T types.ArrayElement]( + typ types.Type, + value []T, + length int, + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (*Vector, error) { + vec, err := NewOffHeapVecWithTypeAndAllocation(typ, selection) + if err != nil { + return nil, err + } + vec.class = CONSTANT + if length > 0 { + if err = SetConstArray(vec, value, length, mp); err != nil { + vec.Free(mp) + return nil, err + } + } + return vec, nil +} + func (s *AllocationAccountSelection) validate() error { if s == nil || s.account == nil || s.account.Handle() == 0 || s.owner < mpool.AllocationOwnerMin || diff --git a/pkg/container/vector/functionTools.go b/pkg/container/vector/functionTools.go index 9605a068a0ad4..2c34ce57fb336 100644 --- a/pkg/container/vector/functionTools.go +++ b/pkg/container/vector/functionTools.go @@ -575,9 +575,10 @@ type FunctionResult[T types.FixedSizeT] struct { vec *Vector mp *mpool.MPool - isVarlena bool - cols []T - length uint64 + allocationAccount *AllocationAccountSelection + isVarlena bool + cols []T + length uint64 // convenientParam save parameter wrappers for easy getting row values. // @@ -594,11 +595,15 @@ func MustFunctionResult[T types.FixedSizeT](wrapper FunctionResultWrapper) *Func } func newResultFunc[T types.FixedSizeT]( - resultType types.Type, mp *mpool.MPool) *FunctionResult[T] { + resultType types.Type, + mp *mpool.MPool, + allocationAccount *AllocationAccountSelection, +) *FunctionResult[T] { f := &FunctionResult[T]{ - typ: resultType, - mp: mp, + typ: resultType, + mp: mp, + allocationAccount: allocationAccount, } var tempT T @@ -621,7 +626,18 @@ func (fr *FunctionResult[T]) getConvenientParamList() []reusableParameterWrapper func (fr *FunctionResult[T]) PreExtendAndReset(targetSize int) error { if fr.vec == nil { - fr.vec = NewOffHeapVecWithType(fr.typ) + var err error + if fr.allocationAccount == nil { + fr.vec = NewOffHeapVecWithType(fr.typ) + } else { + fr.vec, err = NewOffHeapVecWithTypeAndAllocation( + fr.typ, + fr.allocationAccount, + ) + if err != nil { + return err + } + } } oldLength := fr.vec.Length() @@ -757,65 +773,88 @@ func (fr *FunctionResult[T]) Free() { fr.vec.Free(fr.mp) fr.vec = nil } + fr.allocationAccount = nil fr.convenientParam = nil } func NewFunctionResultWrapper(typ types.Type, mp *mpool.MPool) FunctionResultWrapper { + return newFunctionResultWrapper(typ, mp, nil) +} + +// NewFunctionResultWrapperWithAllocation constructs a result owner whose +// lazily allocated Vector data and area use selection. Existing callers remain +// on the legacy path through NewFunctionResultWrapper. +func NewFunctionResultWrapperWithAllocation( + typ types.Type, + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (FunctionResultWrapper, error) { + if err := selection.validate(); err != nil { + return nil, err + } + return newFunctionResultWrapper(typ, mp, selection), nil +} + +func newFunctionResultWrapper( + typ types.Type, + mp *mpool.MPool, + selection *AllocationAccountSelection, +) FunctionResultWrapper { if typ.IsVarlen() { - return newResultFunc[types.Varlena](typ, mp) + return newResultFunc[types.Varlena](typ, mp, selection) } switch typ.Oid { case types.T_bool: - return newResultFunc[bool](typ, mp) + return newResultFunc[bool](typ, mp, selection) case types.T_bit: - return newResultFunc[uint64](typ, mp) + return newResultFunc[uint64](typ, mp, selection) case types.T_int8: - return newResultFunc[int8](typ, mp) + return newResultFunc[int8](typ, mp, selection) case types.T_int16: - return newResultFunc[int16](typ, mp) + return newResultFunc[int16](typ, mp, selection) case types.T_int32: - return newResultFunc[int32](typ, mp) + return newResultFunc[int32](typ, mp, selection) case types.T_int64: - return newResultFunc[int64](typ, mp) + return newResultFunc[int64](typ, mp, selection) case types.T_uint8: - return newResultFunc[uint8](typ, mp) + return newResultFunc[uint8](typ, mp, selection) case types.T_uint16: - return newResultFunc[uint16](typ, mp) + return newResultFunc[uint16](typ, mp, selection) case types.T_uint32: - return newResultFunc[uint32](typ, mp) + return newResultFunc[uint32](typ, mp, selection) case types.T_uint64: - return newResultFunc[uint64](typ, mp) + return newResultFunc[uint64](typ, mp, selection) case types.T_float32: - return newResultFunc[float32](typ, mp) + return newResultFunc[float32](typ, mp, selection) case types.T_float64: - return newResultFunc[float64](typ, mp) + return newResultFunc[float64](typ, mp, selection) case types.T_date: - return newResultFunc[types.Date](typ, mp) + return newResultFunc[types.Date](typ, mp, selection) case types.T_year: - return newResultFunc[types.MoYear](typ, mp) + return newResultFunc[types.MoYear](typ, mp, selection) case types.T_datetime: - return newResultFunc[types.Datetime](typ, mp) + return newResultFunc[types.Datetime](typ, mp, selection) case types.T_time: - return newResultFunc[types.Time](typ, mp) + return newResultFunc[types.Time](typ, mp, selection) case types.T_timestamp: - return newResultFunc[types.Timestamp](typ, mp) + return newResultFunc[types.Timestamp](typ, mp, selection) case types.T_decimal64: - return newResultFunc[types.Decimal64](typ, mp) + return newResultFunc[types.Decimal64](typ, mp, selection) case types.T_decimal128: - return newResultFunc[types.Decimal128](typ, mp) + return newResultFunc[types.Decimal128](typ, mp, selection) case types.T_decimal256: - return newResultFunc[types.Decimal256](typ, mp) + return newResultFunc[types.Decimal256](typ, mp, selection) case types.T_TS: - return newResultFunc[types.TS](typ, mp) + return newResultFunc[types.TS](typ, mp, selection) case types.T_Rowid: - return newResultFunc[types.Rowid](typ, mp) + return newResultFunc[types.Rowid](typ, mp, selection) case types.T_Blockid: - return newResultFunc[types.Blockid](typ, mp) + return newResultFunc[types.Blockid](typ, mp, selection) case types.T_uuid: - return newResultFunc[types.Uuid](typ, mp) + return newResultFunc[types.Uuid](typ, mp, selection) case types.T_enum: - return newResultFunc[types.Enum](typ, mp) + return newResultFunc[types.Enum](typ, mp, selection) } panic(fmt.Sprintf("unexpected type %s for function result", typ)) } diff --git a/pkg/container/vector/function_result_allocation_test.go b/pkg/container/vector/function_result_allocation_test.go new file mode 100644 index 0000000000000..c33ec0186816a --- /dev/null +++ b/pkg/container/vector/function_result_allocation_test.go @@ -0,0 +1,110 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package vector + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/stretchr/testify/require" +) + +func TestFunctionResultAllocationAccountLifecycle(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 8) + mp := mpool.MustNew("function-result-allocation") + defer mpool.DeleteMPool(mp) + + fixed, err := NewFunctionResultWrapperWithAllocation( + types.T_int64.ToType(), + mp, + state.selection, + ) + require.NoError(t, err) + require.NoError(t, fixed.PreExtendAndReset(8)) + fixedVector := fixed.GetResultVector() + require.Same( + t, + state.selection, + fixedVector.AllocationAccountSelection(), + ) + firstUsed := state.account.Snapshot().Used + require.Positive(t, firstUsed) + + require.NoError(t, fixed.PreExtendAndReset(2)) + require.Equal(t, firstUsed, state.account.Snapshot().Used) + + fixed.SetResultVector(nil) + require.NoError(t, fixed.PreExtendAndReset(16)) + transferredUsed := state.account.Snapshot().Used + require.Greater(t, transferredUsed, firstUsed) + fixed.Free() + require.Equal(t, firstUsed, state.account.Snapshot().Used) + fixedVector.Free(mp) + require.Zero(t, state.account.Snapshot().Used) + + varlen, err := NewFunctionResultWrapperWithAllocation( + types.T_varchar.ToType(), + mp, + state.selection, + ) + require.NoError(t, err) + require.NoError(t, varlen.PreExtendAndReset(2)) + result := MustFunctionResult[types.Varlena](varlen) + require.NoError(t, result.AppendBytes(make([]byte, 256), false)) + require.NoError(t, result.AppendBytes([]byte("small"), false)) + varlenUsed := state.account.Snapshot().Used + require.Positive(t, varlenUsed) + + require.NoError(t, varlen.PreExtendAndReset(1)) + require.Equal(t, varlenUsed, state.account.Snapshot().Used) + varlen.Free() + require.Zero(t, state.account.Snapshot().Used) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestFunctionResultAllocationAccountFailure(t *testing.T) { + zeroMP := mpool.MustNewZero() + defer mpool.DeleteMPool(zeroMP) + require.ErrorIs( + t, + func() error { + _, err := NewFunctionResultWrapperWithAllocation( + types.T_int64.ToType(), + zeroMP, + nil, + ) + return err + }(), + mpool.ErrAllocationAccountInvalid, + ) + + state := newTestVectorAllocationAccount(t, 7, 1) + mp := mpool.MustNew("function-result-allocation-failure") + defer mpool.DeleteMPool(mp) + result, err := NewFunctionResultWrapperWithAllocation( + types.T_int64.ToType(), + mp, + state.selection, + ) + require.NoError(t, err) + + err = result.PreExtendAndReset(1) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Zero(t, state.account.Snapshot().Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + result.Free() + finalizeTestVectorAllocationAccount(t, state) +} diff --git a/pkg/sql/colexec/evalExpression.go b/pkg/sql/colexec/evalExpression.go index 01d9329d00f73..fe79419df630b 100644 --- a/pkg/sql/colexec/evalExpression.go +++ b/pkg/sql/colexec/evalExpression.go @@ -92,9 +92,39 @@ type ExpressionExecutor interface { } func NewExpressionExecutorsFromPlanExpressions(proc *process.Process, planExprs []*plan.Expr) (executors []ExpressionExecutor, err error) { + return newExpressionExecutorsFromPlanExpressions(proc, planExprs, nil) +} + +// NewExpressionExecutorsFromPlanExpressionsWithAllocation constructs dormant +// allocation-accounted expression trees. No production caller uses this path +// until the complete expression allocation-site ledger is closed. +func NewExpressionExecutorsFromPlanExpressionsWithAllocation( + proc *process.Process, + planExprs []*plan.Expr, + allocation *ExpressionAllocationAccount, +) (executors []ExpressionExecutor, err error) { + if err = allocation.validate(); err != nil { + return nil, err + } + return newExpressionExecutorsFromPlanExpressions( + proc, + planExprs, + allocation, + ) +} + +func newExpressionExecutorsFromPlanExpressions( + proc *process.Process, + planExprs []*plan.Expr, + allocation *ExpressionAllocationAccount, +) (executors []ExpressionExecutor, err error) { executors = make([]ExpressionExecutor, len(planExprs)) for i := range executors { - executors[i], err = NewExpressionExecutor(proc, planExprs[i]) + executors[i], err = newExpressionExecutor( + proc, + planExprs[i], + allocation, + ) if err != nil { for j := 0; j < i; j++ { executors[j].Free() @@ -106,10 +136,39 @@ func NewExpressionExecutorsFromPlanExpressions(proc *process.Process, planExprs } func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (ExpressionExecutor, error) { + return newExpressionExecutor(proc, planExpr, nil) +} + +// NewExpressionExecutorWithAllocation is the single-root counterpart of +// NewExpressionExecutorsFromPlanExpressionsWithAllocation. +func NewExpressionExecutorWithAllocation( + proc *process.Process, + planExpr *plan.Expr, + allocation *ExpressionAllocationAccount, +) (ExpressionExecutor, error) { + if err := allocation.validate(); err != nil { + return nil, err + } + return newExpressionExecutor(proc, planExpr, allocation) +} + +func newExpressionExecutor( + proc *process.Process, + planExpr *plan.Expr, + allocation *ExpressionAllocationAccount, +) (ExpressionExecutor, error) { + if planExpr == nil { + return nil, moerr.NewInvalidInput(proc.Ctx, "nil expression") + } switch t := planExpr.Expr.(type) { case *plan.Expr_Lit: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) - vec, err := generateConstExpressionExecutor(proc, typ, t.Lit) + vec, err := generateConstExpressionExecutor( + proc, + typ, + t.Lit, + allocation, + ) if err != nil { return nil, err } @@ -117,17 +176,25 @@ func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (Expressi case *plan.Expr_T: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) - vec := vector.NewConstNull(typ, 1, proc.Mp()) + var selection *vector.AllocationAccountSelection + if allocation != nil { + selection = allocation.constant + } + vec, err := newExpressionConstNull(typ, 1, selection, proc.Mp()) + if err != nil { + return nil, err + } return NewFixedVectorExpressionExecutor(proc.Mp(), false, vec), nil case *plan.Expr_Col: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) ce := NewColumnExpressionExecutor() *ce = ColumnExpressionExecutor{ - mp: proc.Mp(), - relIndex: int(t.Col.RelPos), - colIndex: int(t.Col.ColPos), - typ: typ, + mp: proc.Mp(), + relIndex: int(t.Col.RelPos), + colIndex: int(t.Col.ColPos), + typ: typ, + allocation: allocation, } // [issue#19574] // if < 0, it's special for agg or others. @@ -138,24 +205,42 @@ func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (Expressi case *plan.Expr_P: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) - return NewParamExpressionExecutor(proc.Mp(), int(t.P.Pos), typ), nil + executor := NewParamExpressionExecutor(proc.Mp(), int(t.P.Pos), typ) + executor.allocation = allocation + return executor, nil case *plan.Expr_V: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) ve := NewVarExpressionExecutor() *ve = VarExpressionExecutor{ - mp: proc.Mp(), - name: t.V.Name, - system: t.V.System, - global: t.V.Global, - typ: typ, + mp: proc.Mp(), + name: t.V.Name, + system: t.V.System, + global: t.V.Global, + typ: typ, + allocation: allocation, } return ve, nil case *plan.Expr_Vec: - vec := vector.NewVec(types.T_any.ToType()) - err := vec.UnmarshalBinary(t.Vec.Data) + var vec *vector.Vector + var err error + if allocation == nil { + vec = vector.NewVec(types.T_any.ToType()) + err = vec.UnmarshalBinary(t.Vec.Data) + } else { + vec, err = newExpressionVector( + types.T_any.ToType(), + allocation.constant, + ) + if err == nil { + err = vec.UnmarshalBinaryWithCopy(t.Vec.Data, proc.Mp()) + } + } if err != nil { + if vec != nil { + vec.Free(proc.Mp()) + } return nil, err } return NewFixedVectorExpressionExecutor(proc.Mp(), true, vec), nil @@ -164,9 +249,16 @@ func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (Expressi executor := NewListExpressionExecutor() resultVecTyp := t.List.List[0].GetTyp() typ := types.New(types.T(resultVecTyp.Id), resultVecTyp.Width, resultVecTyp.Scale) - executor.Init(proc, typ, len(t.List.List)) + if err := executor.init(proc, typ, len(t.List.List), allocation); err != nil { + executor.Free() + return nil, err + } for i := range executor.parameterExecutor { - subExecutor, paramErr := NewExpressionExecutor(proc, t.List.List[i]) + subExecutor, paramErr := newExpressionExecutor( + proc, + t.List.List[i], + allocation, + ) if paramErr != nil { executor.Free() return nil, paramErr @@ -196,13 +288,17 @@ func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (Expressi } typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) - if err = executor.Init(proc, len(t.F.Args), typ); err != nil { + if err = executor.init(proc, len(t.F.Args), typ, allocation); err != nil { executor.Free() return nil, err } for i := range executor.parameterExecutor { - subExecutor, paramErr := NewExpressionExecutor(proc, t.F.Args[i]) + subExecutor, paramErr := newExpressionExecutor( + proc, + t.F.Args[i], + allocation, + ) if paramErr != nil { executor.Free() return nil, paramErr @@ -230,7 +326,8 @@ type FixedVectorExpressionExecutor struct { } type FunctionExpressionExecutor struct { - m *mpool.MPool + m *mpool.MPool + allocation *ExpressionAllocationAccount // resultType is the declared function return type. Some built-ins refine // result metadata (for example temporal scale or decimal width/scale) at // runtime, so reusable result vectors must start each evaluation from this @@ -262,9 +359,10 @@ type FunctionExpressionExecutor struct { } type ColumnExpressionExecutor struct { - mp *mpool.MPool - relIndex int - colIndex int + mp *mpool.MPool + relIndex int + colIndex int + allocation *ExpressionAllocationAccount // result type. typ types.Type @@ -283,8 +381,9 @@ func (expr *ColumnExpressionExecutor) GetColIndex() int { } type ParamExpressionExecutor struct { - mp *mpool.MPool - null *vector.Vector + mp *mpool.MPool + allocation *ExpressionAllocationAccount + null *vector.Vector // maskedNull is separate from null/vec because it is not a resolved // parameter value and must never participate in the folded-value cache. maskedNull *vector.Vector @@ -298,7 +397,20 @@ type ParamExpressionExecutor struct { func (expr *ParamExpressionExecutor) Eval(proc *process.Process, batches []*batch.Batch, selectList []bool) (*vector.Vector, error) { if noRowsSelected(selectList, expressionRowCount(batches)) { if expr.maskedNull == nil { - expr.maskedNull = vector.NewConstNull(expr.typ, 1, proc.GetMPool()) + var selection *vector.AllocationAccountSelection + if expr.allocation != nil { + selection = expr.allocation.result + } + var err error + expr.maskedNull, err = newExpressionConstNull( + expr.typ, + 1, + selection, + proc.GetMPool(), + ) + if err != nil { + return nil, err + } } return expr.maskedNull, nil } @@ -319,13 +431,35 @@ func (expr *ParamExpressionExecutor) Eval(proc *process.Process, batches []*batc if val == nil { if expr.null == nil { - expr.null = vector.NewConstNull(expr.typ, 1, proc.GetMPool()) + var selection *vector.AllocationAccountSelection + if expr.allocation != nil { + selection = expr.allocation.result + } + expr.null, err = newExpressionConstNull( + expr.typ, + 1, + selection, + proc.GetMPool(), + ) + if err != nil { + return nil, err + } } return expr.null, nil } if expr.vec == nil { - expr.vec, err = vector.NewConstBytes(expr.typ, val, 1, proc.Mp()) + var selection *vector.AllocationAccountSelection + if expr.allocation != nil { + selection = expr.allocation.result + } + expr.vec, err = newExpressionConstBytes( + expr.typ, + val, + 1, + proc.Mp(), + selection, + ) } else { err = vector.SetConstBytes(expr.vec, val, 1, proc.GetMPool()) } @@ -373,8 +507,9 @@ func (expr *ParamExpressionExecutor) IsColumnExpr() bool { } type VarExpressionExecutor struct { - mp *mpool.MPool - null *vector.Vector + mp *mpool.MPool + allocation *ExpressionAllocationAccount + null *vector.Vector // maskedNull lets a skipped variable avoid the resolver without changing // the value cache used by a later selected evaluation. maskedNull *vector.Vector @@ -389,7 +524,20 @@ type VarExpressionExecutor struct { func (expr *VarExpressionExecutor) Eval(proc *process.Process, batches []*batch.Batch, selectList []bool) (*vector.Vector, error) { if noRowsSelected(selectList, expressionRowCount(batches)) { if expr.maskedNull == nil { - expr.maskedNull = vector.NewConstNull(expr.typ, 1, proc.GetMPool()) + var selection *vector.AllocationAccountSelection + if expr.allocation != nil { + selection = expr.allocation.result + } + var err error + expr.maskedNull, err = newExpressionConstNull( + expr.typ, + 1, + selection, + proc.GetMPool(), + ) + if err != nil { + return nil, err + } } return expr.maskedNull, nil } @@ -411,7 +559,16 @@ func (expr *VarExpressionExecutor) Eval(proc *process.Process, batches []*batch. if val == nil { if expr.null == nil { - expr.null, err = util.GenVectorByVarValue(proc, expr.typ, nil) + var selection *vector.AllocationAccountSelection + if expr.allocation != nil { + selection = expr.allocation.result + } + expr.null, err = util.GenVectorByVarValueWithAllocation( + proc, + expr.typ, + nil, + selection, + ) } if err == nil { expr.null.SetIsBin(isBin) @@ -420,7 +577,16 @@ func (expr *VarExpressionExecutor) Eval(proc *process.Process, batches []*batch. } if expr.vec == nil { - expr.vec, err = util.GenVectorByVarValue(proc, expr.typ, val) + var selection *vector.AllocationAccountSelection + if expr.allocation != nil { + selection = expr.allocation.result + } + expr.vec, err = util.GenVectorByVarValueWithAllocation( + proc, + expr.typ, + val, + selection, + ) } else { switch v := val.(type) { case []byte: @@ -474,7 +640,8 @@ func (expr *VarExpressionExecutor) IsColumnExpr() bool { } type ListExpressionExecutor struct { - mp *mpool.MPool + mp *mpool.MPool + allocation *ExpressionAllocationAccount typ types.Type resultVector *vector.Vector @@ -484,11 +651,24 @@ type ListExpressionExecutor struct { func (expr *ListExpressionExecutor) Eval(proc *process.Process, batches []*batch.Batch, selectList []bool) (*vector.Vector, error) { if expr.resultVector == nil { - expr.resultVector = vector.NewOffHeapVecWithType(expr.typ) + var selection *vector.AllocationAccountSelection + if expr.allocation != nil { + selection = expr.allocation.result + } + var err error + expr.resultVector, err = newExpressionVector(expr.typ, selection) + if err != nil { + return nil, err + } } else { expr.resultVector.CleanOnlyData() } - expr.resultVector.PreExtend(len(expr.parameterExecutor), proc.Mp()) + if err := expr.resultVector.PreExtend( + len(expr.parameterExecutor), + proc.Mp(), + ); err != nil { + return nil, err + } for i := range expr.parameterExecutor { vec, err := expr.parameterExecutor[i].Eval(proc, batches, selectList) if err != nil { @@ -531,12 +711,30 @@ func (expr *ListExpressionExecutor) IsColumnExpr() bool { } func (expr *ListExpressionExecutor) Init(proc *process.Process, typ types.Type, parameterNum int) { + if err := expr.init(proc, typ, parameterNum, nil); err != nil { + panic(err) + } +} + +func (expr *ListExpressionExecutor) init( + proc *process.Process, + typ types.Type, + parameterNum int, + allocation *ExpressionAllocationAccount, +) error { m := proc.Mp() expr.typ = typ expr.mp = m + expr.allocation = allocation expr.parameterExecutor = make([]ExpressionExecutor, parameterNum) - expr.resultVector = vector.NewOffHeapVecWithType(typ) + var selection *vector.AllocationAccountSelection + if allocation != nil { + selection = allocation.result + } + var err error + expr.resultVector, err = newExpressionVector(typ, selection) + return err } func (expr *ListExpressionExecutor) SetParameter(index int, executor ExpressionExecutor) { @@ -553,14 +751,32 @@ func (expr *FunctionExpressionExecutor) Init( proc *process.Process, parameterNum int, retType types.Type) (err error) { + return expr.init(proc, parameterNum, retType, nil) +} + +func (expr *FunctionExpressionExecutor) init( + proc *process.Process, + parameterNum int, + retType types.Type, + allocation *ExpressionAllocationAccount, +) (err error) { m := proc.Mp() expr.m = m + expr.allocation = allocation expr.resultType = retType expr.parameterResults = make([]*vector.Vector, parameterNum) expr.parameterExecutor = make([]ExpressionExecutor, parameterNum) - expr.resultVector = vector.NewFunctionResultWrapper(retType, m) + if allocation == nil { + expr.resultVector = vector.NewFunctionResultWrapper(retType, m) + return nil + } + expr.resultVector, err = vector.NewFunctionResultWrapperWithAllocation( + retType, + m, + allocation.result, + ) return err } @@ -588,8 +804,26 @@ func (expr *FunctionExpressionExecutor) EvalIff(proc *process.Process, batches [ } rowCount := expressionRowCount(batches) if len(expr.selectList1) < rowCount { - expr.selectList1 = make([]bool, rowCount) - expr.selectList2 = make([]bool, rowCount) + expr.selectList1, err = ensureExpressionSlice( + expr.selectList1, + rowCount, + expr.m, + expr.allocation, + ExpressionAllocationSiteSelection, + ) + if err != nil { + return err + } + expr.selectList2, err = ensureExpressionSlice( + expr.selectList2, + rowCount, + expr.m, + expr.allocation, + ExpressionAllocationSiteSelection, + ) + if err != nil { + return err + } } trueBranch := expr.selectList1[:rowCount] @@ -624,14 +858,17 @@ func (expr *FunctionExpressionExecutor) EvalIff(proc *process.Process, batches [ return err } } else { - expr.parameterResults[1] = expr.iffNullResult(0, rowCount) + expr.parameterResults[1], err = expr.iffNullResult(0, rowCount) + if err != nil { + return err + } } if hasSelectedRows(falseBranch) { expr.parameterResults[2], err = expr.parameterExecutor[2].Eval(proc, batches, falseBranch) return err } - expr.parameterResults[2] = expr.iffNullResult(1, rowCount) - return nil + expr.parameterResults[2], err = expr.iffNullResult(1, rowCount) + return err } func hasSelectedRows(selectList []bool) bool { @@ -643,26 +880,60 @@ func hasSelectedRows(selectList []bool) bool { return false } -func (expr *FunctionExpressionExecutor) iffNullResult(index, length int) *vector.Vector { +func (expr *FunctionExpressionExecutor) iffNullResult( + index int, + length int, +) (*vector.Vector, error) { typ := expr.resultType result := expr.iffNullResults[index] if result == nil || *result.GetType() != typ { if result != nil { result.Free(expr.m) } - result = vector.NewConstNull(typ, length, expr.m) + var selection *vector.AllocationAccountSelection + if expr.allocation != nil { + selection = expr.allocation.result + } + var err error + result, err = newExpressionConstNull( + typ, + length, + selection, + expr.m, + ) + if err != nil { + return nil, err + } expr.iffNullResults[index] = result } else { result.SetLength(length) } - return result + return result, nil } func (expr *FunctionExpressionExecutor) EvalCase(proc *process.Process, batches []*batch.Batch, selectList []bool) (err error) { rowCount := expressionRowCount(batches) if len(expr.selectList1) < rowCount { - expr.selectList1 = make([]bool, rowCount) - expr.selectList2 = make([]bool, rowCount) + expr.selectList1, err = ensureExpressionSlice( + expr.selectList1, + rowCount, + expr.m, + expr.allocation, + ExpressionAllocationSiteSelection, + ) + if err != nil { + return err + } + expr.selectList2, err = ensureExpressionSlice( + expr.selectList2, + rowCount, + expr.m, + expr.allocation, + ExpressionAllocationSiteSelection, + ) + if err != nil { + return err + } } remaining := expr.selectList1[:rowCount] selectedBranch := expr.selectList2[:rowCount] @@ -702,7 +973,16 @@ func (expr *FunctionExpressionExecutor) EvalCase(proc *process.Process, batches func (expr *FunctionExpressionExecutor) EvalCoalesce(proc *process.Process, batches []*batch.Batch, selectList []bool) (err error) { rowCount := expressionRowCount(batches) if len(expr.selectList1) < rowCount { - expr.selectList1 = make([]bool, rowCount) + expr.selectList1, err = ensureExpressionSlice( + expr.selectList1, + rowCount, + expr.m, + expr.allocation, + ExpressionAllocationSiteSelection, + ) + if err != nil { + return err + } } remaining := expr.selectList1[:rowCount] if selectList != nil { @@ -760,6 +1040,19 @@ func (expr *FunctionExpressionExecutor) evalSelectedRows( rowCount int, selectList []bool, ) (*vector.Vector, error) { + var err error + if expr.allocation != nil { + expr.selectedRows, err = ensureExpressionSlice( + expr.selectedRows, + rowCount, + expr.m, + expr.allocation, + ExpressionAllocationSiteSelectedRows, + ) + if err != nil { + return nil, err + } + } expr.selectedRows = expr.selectedRows[:0] for row := 0; row < rowCount; row++ { if selectList[row] { @@ -786,7 +1079,17 @@ func (expr *FunctionExpressionExecutor) evalSelectedRows( if rowAligned && !parameter.IsConst() { selected := expr.selectedParameterVectors[i] if selected == nil { - selected = vector.NewOffHeapVecWithType(*parameter.GetType()) + var selection *vector.AllocationAccountSelection + if expr.allocation != nil { + selection = expr.allocation.scratch + } + selected, err = newExpressionVector( + *parameter.GetType(), + selection, + ) + if err != nil { + return nil, err + } expr.selectedParameterVectors[i] = selected } else { selected.Reset(*parameter.GetType()) @@ -806,7 +1109,22 @@ func (expr *FunctionExpressionExecutor) evalSelectedRows( return nil, err } if expr.selectedResult == nil { - expr.selectedResult = vector.NewFunctionResultWrapper(expr.resultType, expr.m) + if expr.allocation == nil { + expr.selectedResult = vector.NewFunctionResultWrapper( + expr.resultType, + expr.m, + ) + } else { + expr.selectedResult, err = + vector.NewFunctionResultWrapperWithAllocation( + expr.resultType, + expr.m, + expr.allocation.scratch, + ) + if err != nil { + return nil, err + } + } } expr.resetResultType(expr.selectedResult) if err := expr.selectedResult.PreExtendAndReset(selectedCount); err != nil { @@ -826,7 +1144,19 @@ func (expr *FunctionExpressionExecutor) evalSelectedRows( result.SetIsBin(runtimeIsBin) result.ResetWithSameType() if expr.selectedNullResult == nil { - expr.selectedNullResult = vector.NewConstNull(runtimeType, 1, expr.m) + var selection *vector.AllocationAccountSelection + if expr.allocation != nil { + selection = expr.allocation.scratch + } + expr.selectedNullResult, err = newExpressionConstNull( + runtimeType, + 1, + selection, + expr.m, + ) + if err != nil { + return nil, err + } } else { expr.selectedNullResult.SetType(runtimeType) expr.selectedNullResult.SetLength(1) @@ -908,7 +1238,16 @@ func (expr *FunctionExpressionExecutor) Eval(proc *process.Process, batches []*b return nil, err } if selectList != nil && len(expr.selectList.SelectList) < rowCount { - expr.selectList.SelectList = make([]bool, rowCount) + expr.selectList.SelectList, err = ensureExpressionSlice( + expr.selectList.SelectList, + rowCount, + expr.m, + expr.allocation, + ExpressionAllocationSiteSelection, + ) + if err != nil { + return nil, err + } } if selectList == nil { expr.selectList.AnyNull = false @@ -943,6 +1282,9 @@ func (expr *FunctionExpressionExecutor) EvalWithoutResultReusing(proc *process.P return nil, err } if expr.folded.canFold { + if vec.AllocationAccountSelection() != nil { + return vec.DupOffHeap(proc.Mp()) + } return vec.Dup(proc.Mp()) } expr.resultVector.SetResultVector(nil) @@ -971,6 +1313,18 @@ func (expr *FunctionExpressionExecutor) Free() { parameter.Free(expr.m) } } + freeExpressionSlice(expr.selectList1, expr.m, expr.allocation) + freeExpressionSlice(expr.selectList2, expr.m, expr.allocation) + freeExpressionSlice( + expr.selectList.SelectList, + expr.m, + expr.allocation, + ) + freeExpressionSlice(expr.selectedRows, expr.m, expr.allocation) + expr.selectList1 = nil + expr.selectList2 = nil + expr.selectList.SelectList = nil + expr.selectedRows = nil for _, p := range expr.parameterExecutor { if p != nil { @@ -1008,19 +1362,39 @@ func (expr *ColumnExpressionExecutor) Eval(_ *process.Process, batches []*batch. vec := batches[relIndex].Vecs[expr.colIndex] if vec.IsConstNull() { - vec = expr.getConstNullVec(expr.typ, vec.Length()) + var err error + vec, err = expr.getConstNullVec(expr.typ, vec.Length()) + if err != nil { + return nil, err + } } return vec, nil } -func (expr *ColumnExpressionExecutor) getConstNullVec(typ types.Type, length int) *vector.Vector { +func (expr *ColumnExpressionExecutor) getConstNullVec( + typ types.Type, + length int, +) (*vector.Vector, error) { if expr.nullVecCache != nil { expr.nullVecCache.SetType(typ) expr.nullVecCache.SetLength(length) } else { - expr.nullVecCache = vector.NewConstNull(typ, length, expr.mp) + var selection *vector.AllocationAccountSelection + if expr.allocation != nil { + selection = expr.allocation.result + } + var err error + expr.nullVecCache, err = newExpressionConstNull( + typ, + length, + selection, + expr.mp, + ) + if err != nil { + return nil, err + } } - return expr.nullVecCache + return expr.nullVecCache, nil } func (expr *ColumnExpressionExecutor) EvalWithoutResultReusing(proc *process.Process, batches []*batch.Batch, _ []bool) (*vector.Vector, error) { @@ -1058,6 +1432,9 @@ func (expr *FixedVectorExpressionExecutor) EvalWithoutResultReusing(proc *proces if err != nil { return nil, err } + if vec.AllocationAccountSelection() != nil { + return vec.DupOffHeap(proc.Mp()) + } return vec.Dup(proc.Mp()) } @@ -1077,111 +1454,135 @@ func (expr *FixedVectorExpressionExecutor) IsColumnExpr() bool { return false } -func generateConstExpressionExecutor(proc *process.Process, typ types.Type, con *plan.Literal) (vec *vector.Vector, err error) { +func generateConstExpressionExecutor( + proc *process.Process, + typ types.Type, + con *plan.Literal, + allocation *ExpressionAllocationAccount, +) (vec *vector.Vector, err error) { + var selection *vector.AllocationAccountSelection + if allocation != nil { + selection = allocation.constant + } if con.GetIsnull() { - vec = vector.NewConstNull(typ, 1, proc.Mp()) + vec, err = newExpressionConstNull(typ, 1, selection, proc.Mp()) } else { switch val := con.GetValue().(type) { case *plan.Literal_Bval: - vec, err = vector.NewConstFixed(constBType, val.Bval, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constBType, val.Bval, 1, proc.Mp(), selection) case *plan.Literal_I8Val: - vec, err = vector.NewConstFixed(constI8Type, int8(val.I8Val), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constI8Type, int8(val.I8Val), 1, proc.Mp(), selection) case *plan.Literal_I16Val: - vec, err = vector.NewConstFixed(constI16Type, int16(val.I16Val), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constI16Type, int16(val.I16Val), 1, proc.Mp(), selection) case *plan.Literal_I32Val: - vec, err = vector.NewConstFixed(constI32Type, val.I32Val, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constI32Type, val.I32Val, 1, proc.Mp(), selection) case *plan.Literal_I64Val: - vec, err = vector.NewConstFixed(constI64Type, val.I64Val, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constI64Type, val.I64Val, 1, proc.Mp(), selection) case *plan.Literal_U8Val: - vec, err = vector.NewConstFixed(constU8Type, uint8(val.U8Val), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constU8Type, uint8(val.U8Val), 1, proc.Mp(), selection) case *plan.Literal_U16Val: - vec, err = vector.NewConstFixed(constU16Type, uint16(val.U16Val), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constU16Type, uint16(val.U16Val), 1, proc.Mp(), selection) case *plan.Literal_U32Val: - vec, err = vector.NewConstFixed(constU32Type, val.U32Val, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constU32Type, val.U32Val, 1, proc.Mp(), selection) case *plan.Literal_U64Val: if typ.Oid == types.T_bit { - vec, err = vector.NewConstFixed(typ, val.U64Val, 1, proc.Mp()) + vec, err = newExpressionConstFixed(typ, val.U64Val, 1, proc.Mp(), selection) } else { - vec, err = vector.NewConstFixed(constU64Type, val.U64Val, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constU64Type, val.U64Val, 1, proc.Mp(), selection) } case *plan.Literal_Fval: - vec, err = vector.NewConstFixed(constFType, val.Fval, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constFType, val.Fval, 1, proc.Mp(), selection) case *plan.Literal_Dval: - vec, err = vector.NewConstFixed(constDType, val.Dval, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constDType, val.Dval, 1, proc.Mp(), selection) case *plan.Literal_Dateval: - vec, err = vector.NewConstFixed(constDateType, types.Date(val.Dateval), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constDateType, types.Date(val.Dateval), 1, proc.Mp(), selection) case *plan.Literal_Timeval: - vec, err = vector.NewConstFixed(typ, types.Time(val.Timeval), 1, proc.Mp()) + vec, err = newExpressionConstFixed(typ, types.Time(val.Timeval), 1, proc.Mp(), selection) case *plan.Literal_Datetimeval: - vec, err = vector.NewConstFixed(typ, types.Datetime(val.Datetimeval), 1, proc.Mp()) + vec, err = newExpressionConstFixed(typ, types.Datetime(val.Datetimeval), 1, proc.Mp(), selection) case *plan.Literal_Decimal64Val: cd64 := val.Decimal64Val d64 := types.Decimal64(cd64.A) - vec, err = vector.NewConstFixed(typ, d64, 1, proc.Mp()) + vec, err = newExpressionConstFixed(typ, d64, 1, proc.Mp(), selection) case *plan.Literal_Decimal128Val: cd128 := val.Decimal128Val d128 := types.Decimal128{B0_63: uint64(cd128.A), B64_127: uint64(cd128.B)} - vec, err = vector.NewConstFixed(typ, d128, 1, proc.Mp()) + vec, err = newExpressionConstFixed(typ, d128, 1, proc.Mp(), selection) case *plan.Literal_Timestampval: scale := typ.Scale if scale < 0 || scale > 6 { return nil, moerr.NewErrTooBigPrecision(proc.Ctx, int64(scale), "TIMESTAMP", 6) } - vec, err = vector.NewConstFixed(constTimestampTypes[scale], types.Timestamp(val.Timestampval), 1, proc.Mp()) + vec, err = newExpressionConstFixed( + constTimestampTypes[scale], + types.Timestamp(val.Timestampval), + 1, + proc.Mp(), + selection, + ) case *plan.Literal_Sval: sval := val.Sval // Distinguish binary with non-binary string. if typ.Oid == types.T_binary || typ.Oid == types.T_varbinary || typ.Oid == types.T_blob { - vec, err = vector.NewConstBytes(constBinType, []byte(sval), 1, proc.Mp()) + vec, err = newExpressionConstBytes(constBinType, []byte(sval), 1, proc.Mp(), selection) } else if typ.Oid == types.T_geometry { - vec, err = vector.NewConstBytes(typ, []byte(sval), 1, proc.Mp()) + vec, err = newExpressionConstBytes(typ, []byte(sval), 1, proc.Mp(), selection) } else if typ.Oid == types.T_array_float32 { array, err1 := types.StringToArray[float32](sval) if err1 != nil { return nil, err1 } - vec, err = vector.NewConstArray(typ, array, 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, array, 1, proc.Mp(), selection) } else if typ.Oid == types.T_array_float64 { array, err1 := types.StringToArray[float64](sval) if err1 != nil { return nil, err1 } - vec, err = vector.NewConstArray(typ, array, 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, array, 1, proc.Mp(), selection) } else if typ.Oid == types.T_datalink { _, _, err1 := datalink.ParseDatalink(sval, proc) if err1 != nil { return nil, err1 } - vec, err = vector.NewConstBytes(constBinType, []byte(sval), 1, proc.Mp()) + vec, err = newExpressionConstBytes(constBinType, []byte(sval), 1, proc.Mp(), selection) } else { - vec, err = vector.NewConstBytes(constSType, []byte(sval), 1, proc.Mp()) + vec, err = newExpressionConstBytes(constSType, []byte(sval), 1, proc.Mp(), selection) } case *plan.Literal_Defaultval: defaultVal := val.Defaultval - vec, err = vector.NewConstFixed(constBType, defaultVal, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constBType, defaultVal, 1, proc.Mp(), selection) case *plan.Literal_EnumVal: - vec, err = vector.NewConstFixed(constEnumType, types.Enum(val.EnumVal), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constEnumType, types.Enum(val.EnumVal), 1, proc.Mp(), selection) case *plan.Literal_VecVal: switch typ.Oid { case types.T_array_float32: - vec, err = vector.NewConstArray(typ, types.BytesToArray[float32]([]byte(val.VecVal)), 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, types.BytesToArray[float32]([]byte(val.VecVal)), 1, proc.Mp(), selection) case types.T_array_float64: - vec, err = vector.NewConstArray(typ, types.BytesToArray[float64]([]byte(val.VecVal)), 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, types.BytesToArray[float64]([]byte(val.VecVal)), 1, proc.Mp(), selection) case types.T_array_bf16: - vec, err = vector.NewConstArray(typ, types.BytesToArray[types.BF16]([]byte(val.VecVal)), 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, types.BytesToArray[types.BF16]([]byte(val.VecVal)), 1, proc.Mp(), selection) case types.T_array_float16: - vec, err = vector.NewConstArray(typ, types.BytesToArray[types.Float16]([]byte(val.VecVal)), 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, types.BytesToArray[types.Float16]([]byte(val.VecVal)), 1, proc.Mp(), selection) case types.T_array_int8: - vec, err = vector.NewConstArray(typ, types.BytesToArray[int8]([]byte(val.VecVal)), 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, types.BytesToArray[int8]([]byte(val.VecVal)), 1, proc.Mp(), selection) case types.T_array_uint8: - vec, err = vector.NewConstArray(typ, types.BytesToArray[uint8]([]byte(val.VecVal)), 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, types.BytesToArray[uint8]([]byte(val.VecVal)), 1, proc.Mp(), selection) } default: return nil, moerr.NewNYI(proc.Ctx, fmt.Sprintf("const expression %v", con.GetValue())) } + if err != nil { + return nil, err + } + if vec == nil { + return nil, moerr.NewNYI( + proc.Ctx, + fmt.Sprintf("const expression %v", con.GetValue()), + ) + } vec.SetIsBin(con.IsBin) } - return vec, err + return vec, nil } func GenerateConstListExpressionExecutor(proc *process.Process, exprs []*plan.Expr) (*vector.Vector, error) { diff --git a/pkg/sql/colexec/evalExpressionReset.go b/pkg/sql/colexec/evalExpressionReset.go index da236e888e8c8..e17f8c29c940a 100644 --- a/pkg/sql/colexec/evalExpressionReset.go +++ b/pkg/sql/colexec/evalExpressionReset.go @@ -209,13 +209,27 @@ func (expr *FunctionExpressionExecutor) tryFoldFlowControl( return false, nil } -func (expr *FunctionExpressionExecutor) fillSkippedFlowControlParameters() func() { +func (expr *FunctionExpressionExecutor) fillSkippedFlowControlParameters() ( + func(), + error, +) { // The registered kernels still receive their complete argument list. Supply // typed NULLs for branches that lazy folding deliberately did not evaluate; // the selected conditions make those placeholders unobservable. var boolNull *vector.Vector var resultNull *vector.Vector temporaryIndexes := make([]int, 0, len(expr.parameterResults)) + cleanup := func() { + for _, i := range temporaryIndexes { + expr.parameterResults[i] = nil + } + if boolNull != nil { + boolNull.Free(expr.m) + } + if resultNull != nil { + resultNull.Free(expr.m) + } + } parameterCount := len(expr.parameterResults) for i := range expr.parameterResults { if expr.parameterResults[i] != nil { @@ -227,28 +241,45 @@ func (expr *FunctionExpressionExecutor) fillSkippedFlowControlParameters() func( } if isCondition { if boolNull == nil { - boolNull = vector.NewConstNull(types.T_bool.ToType(), 1, expr.m) + var selection *vector.AllocationAccountSelection + if expr.allocation != nil { + selection = expr.allocation.result + } + var err error + boolNull, err = newExpressionConstNull( + types.T_bool.ToType(), + 1, + selection, + expr.m, + ) + if err != nil { + return nil, err + } } expr.parameterResults[i] = boolNull } else { if resultNull == nil { - resultNull = vector.NewConstNull(expr.resultType, 1, expr.m) + var selection *vector.AllocationAccountSelection + if expr.allocation != nil { + selection = expr.allocation.result + } + var err error + resultNull, err = newExpressionConstNull( + expr.resultType, + 1, + selection, + expr.m, + ) + if err != nil { + cleanup() + return nil, err + } } expr.parameterResults[i] = resultNull } temporaryIndexes = append(temporaryIndexes, i) } - return func() { - for _, i := range temporaryIndexes { - expr.parameterResults[i] = nil - } - if boolNull != nil { - boolNull.Free(expr.m) - } - if resultNull != nil { - resultNull.Free(expr.m) - } - } + return cleanup, nil } func (expr *FunctionExpressionExecutor) finishFolding(proc *process.Process, execLen int) error { @@ -281,7 +312,10 @@ func (expr *FunctionExpressionExecutor) doFold(proc *process.Process, atRuntime if err != nil || !folded { return err } - cleanup := expr.fillSkippedFlowControlParameters() + cleanup, err := expr.fillSkippedFlowControlParameters() + if err != nil { + return err + } defer cleanup() return expr.finishFolding(proc, 1) } diff --git a/pkg/sql/colexec/eval_expression_allocation.go b/pkg/sql/colexec/eval_expression_allocation.go new file mode 100644 index 0000000000000..42b958b3776f2 --- /dev/null +++ b/pkg/sql/colexec/eval_expression_allocation.go @@ -0,0 +1,236 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package colexec + +import ( + "math" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" +) + +// Expression allocation sites are stable diagnostics within the owner chosen +// by the caller. The API is dormant: legacy expression constructors do not +// create or select an account. +const ( + ExpressionAllocationSiteConstantData mpool.AllocationSite = iota + 1 + ExpressionAllocationSiteConstantArea + ExpressionAllocationSiteResultData + ExpressionAllocationSiteResultArea + ExpressionAllocationSiteScratchData + ExpressionAllocationSiteScratchArea + ExpressionAllocationSiteSelection + ExpressionAllocationSiteSelectedRows +) + +// ExpressionAllocationAccount is the immutable allocation provenance shared +// by one expression tree. Vector selections remain separate so diagnostics +// distinguish constants, results, and selected-row scratch. +type ExpressionAllocationAccount struct { + account *mpool.AllocationAccount + owner mpool.AllocationOwner + + constant *vector.AllocationAccountSelection + result *vector.AllocationAccountSelection + scratch *vector.AllocationAccountSelection +} + +func NewExpressionAllocationAccount( + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, +) (*ExpressionAllocationAccount, error) { + constant, err := vector.NewAllocationAccountSelection( + account, + owner, + ExpressionAllocationSiteConstantData, + ExpressionAllocationSiteConstantArea, + ) + if err != nil { + return nil, err + } + result, err := vector.NewAllocationAccountSelection( + account, + owner, + ExpressionAllocationSiteResultData, + ExpressionAllocationSiteResultArea, + ) + if err != nil { + return nil, err + } + scratch, err := vector.NewAllocationAccountSelection( + account, + owner, + ExpressionAllocationSiteScratchData, + ExpressionAllocationSiteScratchArea, + ) + if err != nil { + return nil, err + } + return &ExpressionAllocationAccount{ + account: account, + owner: owner, + constant: constant, + result: result, + scratch: scratch, + }, nil +} + +func (a *ExpressionAllocationAccount) validate() error { + if a == nil || a.account == nil || a.account.Handle() == 0 || + a.owner < mpool.AllocationOwnerMin || + a.owner > mpool.AllocationOwnerMax || + a.constant == nil || a.result == nil || a.scratch == nil { + return mpool.ErrAllocationAccountInvalid + } + return nil +} + +func newExpressionVector( + typ types.Type, + selection *vector.AllocationAccountSelection, +) (*vector.Vector, error) { + if selection == nil { + return vector.NewOffHeapVecWithType(typ), nil + } + return vector.NewOffHeapVecWithTypeAndAllocation(typ, selection) +} + +func newExpressionConstNull( + typ types.Type, + length int, + selection *vector.AllocationAccountSelection, + mp *mpool.MPool, +) (*vector.Vector, error) { + if selection == nil { + return vector.NewConstNull(typ, length, mp), nil + } + return vector.NewConstNullWithAllocation(typ, length, selection) +} + +func newExpressionConstFixed[T any]( + typ types.Type, + value T, + length int, + mp *mpool.MPool, + selection *vector.AllocationAccountSelection, +) (*vector.Vector, error) { + if selection == nil { + return vector.NewConstFixed(typ, value, length, mp) + } + return vector.NewConstFixedWithAllocation( + typ, + value, + length, + mp, + selection, + ) +} + +func newExpressionConstBytes( + typ types.Type, + value []byte, + length int, + mp *mpool.MPool, + selection *vector.AllocationAccountSelection, +) (*vector.Vector, error) { + if selection == nil { + return vector.NewConstBytes(typ, value, length, mp) + } + return vector.NewConstBytesWithAllocation( + typ, + value, + length, + mp, + selection, + ) +} + +func newExpressionConstArray[T types.ArrayElement]( + typ types.Type, + value []T, + length int, + mp *mpool.MPool, + selection *vector.AllocationAccountSelection, +) (*vector.Vector, error) { + if selection == nil { + return vector.NewConstArray(typ, value, length, mp) + } + return vector.NewConstArrayWithAllocation( + typ, + value, + length, + mp, + selection, + ) +} + +func ensureExpressionSlice[T any]( + values []T, + length int, + mp *mpool.MPool, + allocation *ExpressionAllocationAccount, + site mpool.AllocationSite, +) ([]T, error) { + if length < 0 { + return nil, mpool.ErrAllocationAccountInvalid + } + if length <= cap(values) { + return values[:length], nil + } + if allocation == nil { + return make([]T, length), nil + } + if err := allocation.validate(); err != nil { + return nil, err + } + + newCapacity := cap(values) + if newCapacity == 0 { + newCapacity = 1 + } + for newCapacity < length { + if newCapacity > math.MaxInt/2 { + newCapacity = length + break + } + newCapacity *= 2 + } + next, err := mpool.MakeSliceAccounted[T]( + newCapacity, + mp, + allocation.account, + allocation.owner, + site, + ) + if err != nil { + return nil, err + } + copy(next, values) + if cap(values) > 0 { + mpool.FreeSlice(mp, values) + } + return next[:length], nil +} + +func freeExpressionSlice[T any]( + values []T, + mp *mpool.MPool, + allocation *ExpressionAllocationAccount, +) { + if allocation != nil && cap(values) > 0 { + mpool.FreeSlice(mp, values) + } +} diff --git a/pkg/sql/colexec/eval_expression_allocation_test.go b/pkg/sql/colexec/eval_expression_allocation_test.go new file mode 100644 index 0000000000000..ff63d06c8fb51 --- /dev/null +++ b/pkg/sql/colexec/eval_expression_allocation_test.go @@ -0,0 +1,510 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package colexec + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/pb/plan" + "github.com/matrixorigin/matrixone/pkg/sql/plan/function" + "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/matrixorigin/matrixone/pkg/vm/process" + "github.com/stretchr/testify/require" +) + +type testExpressionAllocationAccount struct { + registry *mpool.AllocationAccountRegistry + account *mpool.AllocationAccount + allocation *ExpressionAllocationAccount +} + +func newTestExpressionAllocationAccount( + t testing.TB, + limit uint64, + metadataSlots uint64, +) testExpressionAllocationAccount { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, metadataSlots) + require.NoError(t, err) + account, err := registry.Open(limit) + require.NoError(t, err) + allocation, err := NewExpressionAllocationAccount(account, 1) + require.NoError(t, err) + return testExpressionAllocationAccount{ + registry: registry, + account: account, + allocation: allocation, + } +} + +func finalizeTestExpressionAllocationAccount( + t testing.TB, + state testExpressionAllocationAccount, +) { + t.Helper() + snapshot := state.account.Seal() + require.Zero(t, snapshot.Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + _, err := state.registry.Finalize(state.account) + require.NoError(t, err) +} + +func expressionAllocationColumn(pos int32, typ types.Type) *plan.Expr { + return &plan.Expr{ + Typ: plan.Type{ + Id: int32(typ.Oid), + Width: typ.Width, + Scale: typ.Scale, + }, + Expr: &plan.Expr_Col{ + Col: &plan.ColRef{RelPos: 0, ColPos: pos}, + }, + } +} + +func expressionAllocationString(value string) *plan.Expr { + return &plan.Expr{ + Typ: plan.Type{ + Id: int32(types.T_varchar), + NotNullable: true, + }, + Expr: &plan.Expr_Lit{Lit: &plan.Literal{ + Value: &plan.Literal_Sval{Sval: value}, + }}, + } +} + +func expressionAllocationFunction( + t testing.TB, + proc *process.Process, + name string, + args ...*plan.Expr, +) *plan.Expr { + t.Helper() + argTypes := make([]types.Type, len(args)) + for i := range args { + argTypes[i] = types.New( + types.T(args[i].Typ.Id), + args[i].Typ.Width, + args[i].Typ.Scale, + ) + } + fn, err := function.GetFunctionByName(proc.Ctx, name, argTypes) + require.NoError(t, err) + retType := fn.GetReturnType() + return &plan.Expr{ + Typ: plan.Type{ + Id: int32(retType.Oid), + Width: retType.Width, + Scale: retType.Scale, + }, + Expr: &plan.Expr_F{F: &plan.Function{ + Func: &plan.ObjectRef{ + Obj: fn.GetEncodedOverloadID(), + ObjName: name, + }, + Args: args, + }}, + } +} + +func expressionAllocationCast( + t testing.TB, + proc *process.Process, + source *plan.Expr, + targetType types.Type, +) *plan.Expr { + t.Helper() + target := &plan.Expr{ + Typ: plan.Type{ + Id: int32(targetType.Oid), + Width: targetType.Width, + Scale: targetType.Scale, + NotNullable: true, + }, + Expr: &plan.Expr_T{T: &plan.TargetType{}}, + } + return expressionAllocationFunction(t, proc, "cast", source, target) +} + +func TestExpressionAllocationAccountNestedSelectedLifecycle(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("expression-allocation-lifecycle"), + ) + defer proc.Free() + state := newTestExpressionAllocationAccount(t, 64<<20, 128) + + input := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.NewVector( + 4, + types.T_bool.ToType(), + proc.Mp(), + false, + []bool{true, false, true, false}, + ), + testutil.NewVector( + 4, + types.T_varchar.ToType(), + proc.Mp(), + false, + []string{"a", "b", "c", "d"}, + ), + testutil.NewVector( + 4, + types.T_int64.ToType(), + proc.Mp(), + false, + []int64{10, 20, 30, 40}, + ), + }, nil) + defer input.Clean(proc.Mp()) + + thenExpr := expressionAllocationFunction( + t, + proc, + "concat", + expressionAllocationCast( + t, + proc, + expressionAllocationColumn(2, types.T_int64.ToType()), + types.T_varchar.ToType(), + ), + expressionAllocationString("-then-payload-longer-than-inline"), + ) + elseExpr := expressionAllocationFunction( + t, + proc, + "concat", + expressionAllocationColumn(1, types.T_varchar.ToType()), + expressionAllocationString("-else-payload-longer-than-inline"), + ) + caseExpr := expressionAllocationFunction( + t, + proc, + "case", + expressionAllocationColumn(0, types.T_bool.ToType()), + thenExpr, + elseExpr, + ) + executor, err := NewExpressionExecutorWithAllocation( + proc, + caseExpr, + state.allocation, + ) + require.NoError(t, err) + require.Positive(t, state.account.Snapshot().Used) + + result, err := executor.Eval( + proc, + []*batch.Batch{input}, + []bool{true, false, true, false}, + ) + require.NoError(t, err) + require.NotNil(t, result.AllocationAccountSelection()) + require.Equal(t, "10-then-payload-longer-than-inline", result.GetStringAt(0)) + require.True(t, result.IsNull(1)) + require.Equal(t, "30-then-payload-longer-than-inline", result.GetStringAt(2)) + require.True(t, result.IsNull(3)) + + root := executor.(*FunctionExpressionExecutor) + require.GreaterOrEqual(t, cap(root.selectList1), input.RowCount()) + require.GreaterOrEqual(t, cap(root.selectList2), input.RowCount()) + require.GreaterOrEqual(t, cap(root.selectedRows), input.RowCount()) + require.NotNil(t, root.selectedResult) + require.NotNil(t, root.selectedResult.GetResultVector()) + require.NotNil( + t, + root.selectedResult.GetResultVector().AllocationAccountSelection(), + ) + usedAfterPartial := state.account.Snapshot().Used + + executor.ResetForNextQuery() + result, err = executor.Eval( + proc, + []*batch.Batch{input}, + []bool{true, false, true, false}, + ) + require.NoError(t, err) + require.Equal(t, usedAfterPartial, state.account.Snapshot().Used) + + transferred, err := executor.EvalWithoutResultReusing( + proc, + []*batch.Batch{input}, + nil, + ) + require.NoError(t, err) + require.NotNil(t, transferred.AllocationAccountSelection()) + executor.Free() + require.Positive(t, state.account.Snapshot().Used) + transferred.Free(proc.Mp()) + require.Zero(t, state.account.Snapshot().Used) + finalizeTestExpressionAllocationAccount(t, state) +} + +func TestExpressionAllocationAccountConstantKinds(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("expression-allocation-constants"), + ) + defer proc.Free() + state := newTestExpressionAllocationAccount(t, 1<<20, 8) + + fixed := &plan.Expr{ + Typ: plan.Type{ + Id: int32(types.T_int64), + NotNullable: true, + }, + Expr: &plan.Expr_Lit{Lit: &plan.Literal{ + Value: &plan.Literal_I64Val{I64Val: 42}, + }}, + } + null := &plan.Expr{ + Typ: plan.Type{Id: int32(types.T_int64)}, + Expr: &plan.Expr_Lit{Lit: &plan.Literal{Isnull: true}}, + } + executors, err := + NewExpressionExecutorsFromPlanExpressionsWithAllocation( + proc, + []*plan.Expr{fixed, null}, + state.allocation, + ) + require.NoError(t, err) + require.Positive(t, state.account.Snapshot().Used) + for _, executor := range executors { + fixedExecutor := executor.(*FixedVectorExpressionExecutor) + require.NotNil( + t, + fixedExecutor.resultVector.AllocationAccountSelection(), + ) + _, err = executor.Eval( + proc, + []*batch.Batch{batch.EmptyForConstFoldBatch}, + nil, + ) + require.NoError(t, err) + } + for _, executor := range executors { + executor.Free() + } + require.Zero(t, state.account.Snapshot().Used) + finalizeTestExpressionAllocationAccount(t, state) +} + +func TestExpressionAllocationAccountDecodedVectorTransfer(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("expression-allocation-decoded-vector"), + ) + defer proc.Free() + state := newTestExpressionAllocationAccount(t, 1<<20, 8) + + source := testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()) + data, err := source.MarshalBinary() + require.NoError(t, err) + source.Free(proc.Mp()) + + executor, err := NewExpressionExecutorWithAllocation( + proc, + &plan.Expr{ + Typ: plan.Type{Id: int32(types.T_int32)}, + Expr: &plan.Expr_Vec{Vec: &plan.LiteralVec{ + Len: 3, + Data: data, + }}, + }, + state.allocation, + ) + require.NoError(t, err) + fixed := executor.(*FixedVectorExpressionExecutor) + require.NotNil( + t, + fixed.resultVector.AllocationAccountSelection(), + ) + require.False(t, fixed.resultVector.NeedDup()) + require.Positive(t, state.account.Snapshot().Used) + + transferred, err := executor.EvalWithoutResultReusing( + proc, + []*batch.Batch{batch.EmptyForConstFoldBatch}, + nil, + ) + require.NoError(t, err) + executor.Free() + require.Positive(t, state.account.Snapshot().Used) + transferred.Free(proc.Mp()) + require.Zero(t, state.account.Snapshot().Used) + finalizeTestExpressionAllocationAccount(t, state) +} + +func TestExpressionAllocationAccountFoldedTransfer(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("expression-allocation-folded-transfer"), + ) + defer proc.Free() + state := newTestExpressionAllocationAccount(t, 1<<20, 16) + + expr := expressionAllocationFunction( + t, + proc, + "concat", + expressionAllocationString("folded-payload-longer-than-inline"), + expressionAllocationString("-suffix"), + ) + executor, err := NewExpressionExecutorWithAllocation( + proc, + expr, + state.allocation, + ) + require.NoError(t, err) + + transferred, err := executor.EvalWithoutResultReusing( + proc, + []*batch.Batch{batch.EmptyForConstFoldBatch}, + nil, + ) + require.NoError(t, err) + require.Equal( + t, + "folded-payload-longer-than-inline-suffix", + transferred.GetStringAt(0), + ) + require.NotNil(t, transferred.AllocationAccountSelection()) + executor.Free() + require.Positive(t, state.account.Snapshot().Used) + transferred.Free(proc.Mp()) + require.Zero(t, state.account.Snapshot().Used) + finalizeTestExpressionAllocationAccount(t, state) +} + +func TestExpressionAllocationAccountConstructionRollback(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("expression-allocation-construction"), + ) + defer proc.Free() + state := newTestExpressionAllocationAccount(t, 1<<20, 1) + + expr := expressionAllocationFunction( + t, + proc, + "concat", + expressionAllocationString("left"), + expressionAllocationString("right"), + ) + _, err := NewExpressionExecutorWithAllocation( + proc, + expr, + state.allocation, + ) + require.ErrorIs(t, err, mpool.ErrAllocationMetadataSlots) + require.Zero(t, state.account.Snapshot().Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + finalizeTestExpressionAllocationAccount(t, state) +} + +func TestExpressionAllocationAccountScratchFailureCleanup(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("expression-allocation-scratch-failure"), + ) + defer proc.Free() + state := newTestExpressionAllocationAccount(t, 8, 8) + + input := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.NewVector( + 8, + types.T_bool.ToType(), + proc.Mp(), + false, + []bool{true, false, true, false, true, false, true, false}, + ), + testutil.NewVector( + 8, + types.T_int64.ToType(), + proc.Mp(), + false, + []int64{1, 2, 3, 4, 5, 6, 7, 8}, + ), + }, nil) + defer input.Clean(proc.Mp()) + + expr := expressionAllocationFunction( + t, + proc, + "case", + expressionAllocationColumn(0, types.T_bool.ToType()), + expressionAllocationColumn(1, types.T_int64.ToType()), + expressionAllocationColumn(1, types.T_int64.ToType()), + ) + executor, err := NewExpressionExecutorWithAllocation( + proc, + expr, + state.allocation, + ) + require.NoError(t, err) + _, err = executor.Eval(proc, []*batch.Batch{input}, nil) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Equal(t, uint64(8), state.account.Snapshot().Used) + + executor.Free() + require.Zero(t, state.account.Snapshot().Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + finalizeTestExpressionAllocationAccount(t, state) +} + +func TestExpressionAllocationAccountZeroLengthScratchGrowth(t *testing.T) { + mp := mpool.MustNew("expression-allocation-zero-length-scratch") + defer mpool.DeleteMPool(mp) + state := newTestExpressionAllocationAccount(t, 1<<20, 4) + + values, err := ensureExpressionSlice( + []int64(nil), + 4, + mp, + state.allocation, + ExpressionAllocationSiteSelectedRows, + ) + require.NoError(t, err) + require.Equal(t, uint64(32), state.account.Snapshot().Used) + + values = values[:0] + values, err = ensureExpressionSlice( + values, + 8, + mp, + state.allocation, + ExpressionAllocationSiteSelectedRows, + ) + require.NoError(t, err) + require.Len(t, values, 8) + require.Equal(t, uint64(64), state.account.Snapshot().Used) + + values = values[:0] + freeExpressionSlice(values, mp, state.allocation) + require.Zero(t, state.account.Snapshot().Used) + finalizeTestExpressionAllocationAccount(t, state) +} diff --git a/pkg/sql/util/eval_expr_util.go b/pkg/sql/util/eval_expr_util.go index da58ac7b71d38..2de311ce9baf5 100644 --- a/pkg/sql/util/eval_expr_util.go +++ b/pkg/sql/util/eval_expr_util.go @@ -111,12 +111,40 @@ func DecodeBinaryString(s string) ([]byte, error) { } func GenVectorByVarValue(proc *process.Process, typ types.Type, val any) (*vector.Vector, error) { + return GenVectorByVarValueWithAllocation(proc, typ, val, nil) +} + +// GenVectorByVarValueWithAllocation is the dormant allocation-accounted +// variant used by expression executors. A nil selection preserves the legacy +// allocation mode. +func GenVectorByVarValueWithAllocation( + proc *process.Process, + typ types.Type, + val any, + selection *vector.AllocationAccountSelection, +) (*vector.Vector, error) { if val == nil { - vec := vector.NewConstNull(typ, 1, proc.Mp()) - return vec, nil + if selection == nil { + return vector.NewConstNull(typ, 1, proc.Mp()), nil + } + return vector.NewConstNullWithAllocation(typ, 1, selection) } else { strVal := getVal(val) - return vector.NewConstBytes(typ, []byte(strVal), 1, proc.Mp()) + if selection == nil { + return vector.NewConstBytes( + typ, + []byte(strVal), + 1, + proc.Mp(), + ) + } + return vector.NewConstBytesWithAllocation( + typ, + []byte(strVal), + 1, + proc.Mp(), + selection, + ) } } diff --git a/pkg/sql/util/eval_expr_util_test.go b/pkg/sql/util/eval_expr_util_test.go index 679c5c1d738e9..46f87f38c685d 100644 --- a/pkg/sql/util/eval_expr_util_test.go +++ b/pkg/sql/util/eval_expr_util_test.go @@ -19,7 +19,9 @@ import ( "testing" "time" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/sql/parsers/tree" "github.com/matrixorigin/matrixone/pkg/testutil" "github.com/matrixorigin/matrixone/pkg/vm/process" @@ -42,6 +44,53 @@ func TestHexToInt(t *testing.T) { require.Error(t, err) } +func TestGenVectorByVarValueWithAllocation(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 4) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection( + account, + 1, + 1, + 2, + ) + require.NoError(t, err) + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("variable-value-allocation"), + ) + defer proc.Free() + + nullVec, err := GenVectorByVarValueWithAllocation( + proc, + types.T_varchar.ToType(), + nil, + selection, + ) + require.NoError(t, err) + require.Same(t, selection, nullVec.AllocationAccountSelection()) + + valueVec, err := GenVectorByVarValueWithAllocation( + proc, + types.T_varchar.ToType(), + "variable payload longer than the inline varlena limit", + selection, + ) + require.NoError(t, err) + require.Same(t, selection, valueVec.AllocationAccountSelection()) + require.Positive(t, account.Snapshot().Used) + + nullVec.Free(proc.Mp()) + valueVec.Free(proc.Mp()) + snapshot := account.Seal() + require.Zero(t, snapshot.Used) + require.Zero(t, registry.LiveAllocationMetadata()) + _, err = registry.Finalize(account) + require.NoError(t, err) +} + func TestSetInsertValueStringBinaryHexPadding(t *testing.T) { proc := testutil.NewProcess(t) From 3c9b70143d49f81ff41213fda5937c68cc20cda4 Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 14:39:42 +0800 Subject: [PATCH 05/61] feat: propagate allocation accounts through spill scratch --- .../colexec/spillutil/allocation_account.go | 173 +++++++++++ .../spillutil/allocation_account_test.go | 276 ++++++++++++++++++ pkg/sql/colexec/spillutil/join_spill.go | 205 +++++++++++-- 3 files changed, 634 insertions(+), 20 deletions(-) create mode 100644 pkg/sql/colexec/spillutil/allocation_account.go create mode 100644 pkg/sql/colexec/spillutil/allocation_account_test.go diff --git a/pkg/sql/colexec/spillutil/allocation_account.go b/pkg/sql/colexec/spillutil/allocation_account.go new file mode 100644 index 0000000000000..8cdfb2ad6613f --- /dev/null +++ b/pkg/sql/colexec/spillutil/allocation_account.go @@ -0,0 +1,173 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package spillutil + +import ( + "math" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/sql/colexec" +) + +// Spill allocation sites use a range disjoint from colexec expression sites +// when both subsystems share one logical owner. +const ( + SpillAllocationSiteDecodedData mpool.AllocationSite = iota + 32 + SpillAllocationSiteDecodedArea + SpillAllocationSiteSelectedData + SpillAllocationSiteSelectedArea + SpillAllocationSiteHashValues + SpillAllocationSiteRowIDs +) + +// SpillAllocationAccount is the dormant allocation provenance for one spill +// engine. Serialization buffers remain a named activation blocker. +type SpillAllocationAccount struct { + account *mpool.AllocationAccount + owner mpool.AllocationOwner + + decoded *vector.AllocationAccountSelection + selected *vector.AllocationAccountSelection + expression *colexec.ExpressionAllocationAccount +} + +func NewSpillAllocationAccount( + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, +) (*SpillAllocationAccount, error) { + decoded, err := vector.NewAllocationAccountSelection( + account, + owner, + SpillAllocationSiteDecodedData, + SpillAllocationSiteDecodedArea, + ) + if err != nil { + return nil, err + } + selected, err := vector.NewAllocationAccountSelection( + account, + owner, + SpillAllocationSiteSelectedData, + SpillAllocationSiteSelectedArea, + ) + if err != nil { + return nil, err + } + expression, err := colexec.NewExpressionAllocationAccount(account, owner) + if err != nil { + return nil, err + } + return &SpillAllocationAccount{ + account: account, + owner: owner, + decoded: decoded, + selected: selected, + expression: expression, + }, nil +} + +func (a *SpillAllocationAccount) validate() error { + if a == nil || a.account == nil || a.account.Handle() == 0 || + a.owner < mpool.AllocationOwnerMin || + a.owner > mpool.AllocationOwnerMax || + a.decoded == nil || a.selected == nil || a.expression == nil { + return mpool.ErrAllocationAccountInvalid + } + return nil +} + +func newSpillBatch( + size int, + selection *vector.AllocationAccountSelection, +) (*batch.Batch, error) { + bat := batch.NewOffHeapWithSize(size) + if selection != nil { + if err := bat.SetAllocationAccount(selection); err != nil { + bat.Clean(nil) + return nil, err + } + } + return bat, nil +} + +func newSpillVector( + typ types.Type, + selection *vector.AllocationAccountSelection, +) (*vector.Vector, error) { + if selection == nil { + return vector.NewOffHeapVecWithType(typ), nil + } + return vector.NewOffHeapVecWithTypeAndAllocation(typ, selection) +} + +func growSpillSlice[T any]( + values []T, + length int, + mp *mpool.MPool, + allocation *SpillAllocationAccount, + site mpool.AllocationSite, +) ([]T, error) { + if length < 0 { + return nil, mpool.ErrAllocationAccountInvalid + } + if length <= cap(values) { + return values[:length], nil + } + if allocation == nil { + return make([]T, length), nil + } + if err := allocation.validate(); err != nil { + return nil, err + } + newCapacity := cap(values) + if newCapacity == 0 { + newCapacity = 1 + } + for newCapacity < length { + if newCapacity > math.MaxInt/2 { + newCapacity = length + break + } + newCapacity *= 2 + } + next, err := mpool.MakeSliceAccounted[T]( + newCapacity, + mp, + allocation.account, + allocation.owner, + site, + ) + if err != nil { + return nil, err + } + copy(next, values) + if cap(values) > 0 { + mpool.FreeSlice(mp, values) + } + return next[:length], nil +} + +func freeSpillSlice[T any]( + values []T, + mp *mpool.MPool, + allocation *SpillAllocationAccount, +) { + if allocation != nil && cap(values) > 0 { + mpool.FreeSlice(mp, values) + } +} diff --git a/pkg/sql/colexec/spillutil/allocation_account_test.go b/pkg/sql/colexec/spillutil/allocation_account_test.go new file mode 100644 index 0000000000000..6d74f701f3d5b --- /dev/null +++ b/pkg/sql/colexec/spillutil/allocation_account_test.go @@ -0,0 +1,276 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package spillutil + +import ( + "bytes" + "os" + "path/filepath" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/matrixorigin/matrixone/pkg/vm/process" + "github.com/stretchr/testify/require" +) + +type testSpillAllocationAccount struct { + registry *mpool.AllocationAccountRegistry + account *mpool.AllocationAccount + allocation *SpillAllocationAccount +} + +func newTestSpillAllocationAccount( + t testing.TB, + limit uint64, + metadataSlots uint64, +) testSpillAllocationAccount { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, metadataSlots) + require.NoError(t, err) + account, err := registry.Open(limit) + require.NoError(t, err) + allocation, err := NewSpillAllocationAccount(account, 2) + require.NoError(t, err) + return testSpillAllocationAccount{ + registry: registry, + account: account, + allocation: allocation, + } +} + +func finalizeTestSpillAllocationAccount( + t testing.TB, + state testSpillAllocationAccount, +) { + t.Helper() + snapshot := state.account.Seal() + require.Zero(t, snapshot.Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + _, err := state.registry.Finalize(state.account) + require.NoError(t, err) +} + +func writeSpillAllocationTestFile( + t testing.TB, + bat *batch.Batch, + truncate int, +) *os.File { + t.Helper() + var encoded bytes.Buffer + require.NoError(t, marshalSpillRecord(bat, &encoded)) + payload := encoded.Bytes() + if truncate > 0 { + payload = payload[:len(payload)-truncate] + } + path := filepath.Join(t.TempDir(), "spill.bin") + require.NoError(t, os.WriteFile(path, payload, 0o600)) + file, err := os.Open(path) + require.NoError(t, err) + return file +} + +func writeSpillAllocationTestRecords( + t testing.TB, + batches ...*batch.Batch, +) *os.File { + t.Helper() + var payload bytes.Buffer + for _, bat := range batches { + var encoded bytes.Buffer + require.NoError(t, marshalSpillRecord(bat, &encoded)) + _, err := payload.Write(encoded.Bytes()) + require.NoError(t, err) + } + path := filepath.Join(t.TempDir(), "spill-records.bin") + require.NoError(t, os.WriteFile(path, payload.Bytes(), 0o600)) + file, err := os.Open(path) + require.NoError(t, err) + return file +} + +func TestSpillAllocationAccountDecodedBatchLifecycle(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-decoded"), + ) + defer proc.Free() + state := newTestSpillAllocationAccount(t, 8<<20, 16) + source := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.NewVector( + 4, + types.T_int64.ToType(), + proc.Mp(), + false, + []int64{1, 2, 3, 4}, + ), + testutil.NewVector( + 4, + types.T_varchar.ToType(), + proc.Mp(), + false, + []string{"a", "payload-longer-than-inline", "c", "d"}, + ), + }, nil) + defer source.Clean(proc.Mp()) + + reader := BucketReader{ + fd: writeSpillAllocationTestFile(t, source, 0), + allocation: state.allocation, + } + reuse := batch.NewOffHeapWithSize(0) + decoded, err := reader.ReadBatch(proc, reuse) + require.NoError(t, err) + require.Same(t, reuse, decoded) + require.NotNil(t, decoded.AllocationAccountSelection()) + for _, vec := range decoded.Vecs { + require.NotNil(t, vec.AllocationAccountSelection()) + } + require.Positive(t, state.account.Snapshot().Used) + + reader.Close() + reuse.Clean(proc.Mp()) + require.Zero(t, state.account.Snapshot().Used) + + reader = BucketReader{ + fd: writeSpillAllocationTestRecords(t, source, source), + mergeRecords: true, + allocation: state.allocation, + } + reuse = batch.NewOffHeapWithSize(0) + decoded, err = reader.ReadBatch(proc, reuse) + require.NoError(t, err) + require.Equal(t, 2*source.RowCount(), decoded.RowCount()) + require.Positive(t, state.account.Snapshot().Used) + reader.Close() + reuse.Clean(proc.Mp()) + require.Zero(t, state.account.Snapshot().Used) + + reader = BucketReader{ + fd: writeSpillAllocationTestFile(t, source, 1), + allocation: state.allocation, + } + reuse = batch.NewOffHeapWithSize(0) + _, err = reader.ReadBatch(proc, reuse) + require.Error(t, err) + require.Zero(t, state.account.Snapshot().Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + reader.Close() + finalizeTestSpillAllocationAccount(t, state) +} + +func TestSpillAllocationAccountScatterScratchLifecycle(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-scatter"), + ) + defer proc.Free() + state := newTestSpillAllocationAccount(t, 1<<20, 8) + engine, err := NewSpillEngineWithAllocation( + SpillEngineConfig{}, + state.allocation, + ) + require.NoError(t, err) + + source := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.NewVector( + 8, + types.T_int64.ToType(), + proc.Mp(), + false, + []int64{1, 2, 3, 4, 5, 6, 7, 8}, + ), + }, nil) + defer source.Clean(proc.Mp()) + writers := MakeBucketWriters("spill_allocation_scatter") + defer func() { + for i := range writers { + writers[i].Close() + } + }() + require.NoError(t, engine.scatterBatchBounded( + proc, + source, + source.Vecs, + writers, + 0, + false, + process.NewAnalyzer(0, false, false, "test"), + )) + require.Len(t, engine.scatterHashValues, source.RowCount()) + require.Len(t, engine.scatterBucketRowIds, source.RowCount()) + snapshot := state.account.Snapshot() + require.Equal(t, uint64(source.RowCount()*(8+4)), snapshot.Used) + require.Greater(t, snapshot.Peak, snapshot.Used) + + engine.releaseScatterScratch() + require.Zero(t, state.account.Snapshot().Used) + engine.Cleanup(proc) + finalizeTestSpillAllocationAccount(t, state) +} + +func TestSpillAllocationAccountScatterFailureCleanup(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-scatter-failure"), + ) + defer proc.Free() + const rows = 8 + state := newTestSpillAllocationAccount(t, rows*(8+4), 8) + engine, err := NewSpillEngineWithAllocation( + SpillEngineConfig{}, + state.allocation, + ) + require.NoError(t, err) + + source := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.NewVector( + rows, + types.T_int64.ToType(), + proc.Mp(), + false, + []int64{1, 2, 3, 4, 5, 6, 7, 8}, + ), + }, nil) + defer source.Clean(proc.Mp()) + writers := MakeBucketWriters("spill_allocation_scatter_failure") + defer func() { + for i := range writers { + writers[i].Close() + } + }() + err = engine.scatterBatchBounded( + proc, + source, + source.Vecs, + writers, + 0, + false, + process.NewAnalyzer(0, false, false, "test"), + ) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Equal(t, uint64(rows*(8+4)), state.account.Snapshot().Used) + + engine.releaseScatterScratch() + require.Zero(t, state.account.Snapshot().Used) + engine.Cleanup(proc) + finalizeTestSpillAllocationAccount(t, state) +} diff --git a/pkg/sql/colexec/spillutil/join_spill.go b/pkg/sql/colexec/spillutil/join_spill.go index 70f0aa7f1e468..263a7608cbc18 100644 --- a/pkg/sql/colexec/spillutil/join_spill.go +++ b/pkg/sql/colexec/spillutil/join_spill.go @@ -88,6 +88,7 @@ type BucketReader struct { batchCharge uint64 spillFile *message.SpillFile mergeRecords bool + allocation *SpillAllocationAccount } func (r *BucketReader) ReadBatch(proc *process.Process, reuseBat *batch.Batch) (*batch.Batch, error) { @@ -97,6 +98,19 @@ func (r *BucketReader) ReadBatch(proc *process.Process, reuseBat *batch.Batch) ( if r.fd == nil { return nil, io.EOF } + if reuseBat == nil { + return nil, moerr.NewInvalidInput( + proc.Ctx, + "spill batch reader requires a reuse batch", + ) + } + if r.allocation != nil { + if err := reuseBat.SetAllocationAccount( + r.allocation.decoded, + ); err != nil { + return nil, err + } + } if r.reader == nil { r.reader = bufio.NewReaderSize(r.fd, 4*1024*1024) } @@ -155,7 +169,20 @@ func (r *BucketReader) ReadBatch(proc *process.Process, reuseBat *batch.Batch) ( if nextRows > int64(colexec.DefaultBatchSize-reuseBat.RowCount()) { break } - next := batch.NewOffHeapWithSize(0) + var selection *vector.AllocationAccountSelection + if r.allocation != nil { + selection = r.allocation.decoded + } + next, err := newSpillBatch(0, selection) + if err != nil { + return nil, r.mergeReadError( + proc, + reuseBat, + nil, + nil, + err, + ) + } _, nextToken, _, err := r.readBatchRecord(proc, next, nil, 0, false) if err != nil { return nil, r.mergeReadError(proc, reuseBat, next, nextToken, err) @@ -556,6 +583,13 @@ func (r *BucketReader) readBatchRecord( // A caller-provided reuse batch has no budget ownership on the first // read. Drop it before admitting the decoded payload. reuseBat.Clean(proc.Mp()) + if r.allocation != nil { + if err := reuseBat.SetAllocationAccount( + r.allocation.decoded, + ); err != nil { + return nil, nil, 0, err + } + } var err error token, err = r.budget.Reserve(projected) if err != nil { @@ -580,6 +614,14 @@ func (r *BucketReader) readBatchRecord( } if !retainedOK || !peakOK || growErr != nil { reuseBat.Clean(proc.Mp()) + if r.allocation != nil { + if err := reuseBat.SetAllocationAccount( + r.allocation.decoded, + ); err != nil { + token.Release() + return nil, nil, 0, err + } + } token.Release() token = nil var err error @@ -1421,8 +1463,22 @@ func (e *SpillEngine) scatterBatchBounded( } } - if cap(e.scatterHashValues) < rows { - e.scatterHashValues = make([]uint64, rows) + if e.allocation != nil { + if e.allocationMP != nil && e.allocationMP != proc.Mp() { + return mpool.ErrAllocationAccountInvalid + } + e.allocationMP = proc.Mp() + } + var err error + e.scatterHashValues, err = growSpillSlice( + e.scatterHashValues, + rows, + proc.Mp(), + e.allocation, + SpillAllocationSiteHashValues, + ) + if err != nil { + return err } hashValues := e.scatterHashValues[:rows] // Re-spill must consume fresh hash bits. Merely changing the initial seed @@ -1437,8 +1493,15 @@ func (e *SpillEngine) scatterBatchBounded( if shift >= 64 { return process.ErrHashBuildBudgetInvalid } - if cap(e.scatterBucketRowIds) < rows { - e.scatterBucketRowIds = make([]int32, rows) + e.scatterBucketRowIds, err = growSpillSlice( + e.scatterBucketRowIds, + rows, + proc.Mp(), + e.allocation, + SpillAllocationSiteRowIDs, + ) + if err != nil { + return err } if cap(e.keyVecs) < len(keyVecs) { e.keyVecs = make([]*vector.Vector, len(keyVecs)) @@ -1456,9 +1519,22 @@ func (e *SpillEngine) scatterBatchBounded( } sels := e.scatterBucketRowIds[start:end] if selected == nil { - selected = batch.NewOffHeapWithSize(len(bat.Vecs)) + var selection *vector.AllocationAccountSelection + if e.allocation != nil { + selection = e.allocation.selected + } + selected, err = newSpillBatch(len(bat.Vecs), selection) + if err != nil { + return err + } for j, vec := range bat.Vecs { - selected.Vecs[j] = vector.NewOffHeapVecWithType(*vec.GetType()) + selected.Vecs[j], err = newSpillVector( + *vec.GetType(), + selection, + ) + if err != nil { + return err + } } } selected.CleanOnlyData() @@ -1571,8 +1647,19 @@ func (e *SpillEngine) discardScatterBuffers() { // charged while the next child hashmap is rebuilt. Cleanup calls this method // as an idempotent fallback for cancellation paths. func (e *SpillEngine) releaseScatterScratch() { + freeSpillSlice( + e.scatterHashValues, + e.allocationMP, + e.allocation, + ) + freeSpillSlice( + e.scatterBucketRowIds, + e.allocationMP, + e.allocation, + ) e.scatterHashValues = nil e.scatterBucketRowIds = nil + e.allocationMP = nil e.keyVecs = nil e.scatterWriteBuf = bytes.Buffer{} for i := range e.scatterBucketCounts { @@ -1697,9 +1784,11 @@ const ( // SpillEngine owns the spill bucket queue and drives the probe-batch loop. type SpillEngine struct { - cfg SpillEngineConfig - buckets []SpillBucket - spillFS spillFileServiceCache + cfg SpillEngineConfig + buckets []SpillBucket + spillFS spillFileServiceCache + allocation *SpillAllocationAccount + allocationMP *mpool.MPool // Current bucket state buildReader BucketReader @@ -1741,10 +1830,35 @@ type SpillEngine struct { // NewSpillEngine creates an engine from configuration. Call InitFromSpilledMap next. func NewSpillEngine(cfg SpillEngineConfig) *SpillEngine { + return newSpillEngine(cfg, nil) +} + +// NewSpillEngineWithAllocation constructs the dormant allocation-accounted +// spill path. Legacy production callers continue to use NewSpillEngine. +func NewSpillEngineWithAllocation( + cfg SpillEngineConfig, + allocation *SpillAllocationAccount, +) (*SpillEngine, error) { + if err := allocation.validate(); err != nil { + return nil, err + } + return newSpillEngine(cfg, allocation), nil +} + +func newSpillEngine( + cfg SpillEngineConfig, + allocation *SpillAllocationAccount, +) *SpillEngine { if cfg.MaxQueue <= 0 { cfg.MaxQueue = SpillNumBuckets * SpillNumBuckets } - return &SpillEngine{cfg: cfg} + engine := &SpillEngine{ + cfg: cfg, + allocation: allocation, + } + engine.buildReader.allocation = allocation + engine.probeReader.allocation = allocation + return engine } func (e *SpillEngine) makeBucketWriters(prefix string) []BucketWriter { @@ -1935,7 +2049,15 @@ func (e *SpillEngine) NextProbeBatch(proc *process.Process) (*batch.Batch, error return nil, nil } if e.probeReadBatch == nil { - e.probeReadBatch = batch.NewOffHeapWithSize(0) + var selection *vector.AllocationAccountSelection + if e.allocation != nil { + selection = e.allocation.decoded + } + var err error + e.probeReadBatch, err = newSpillBatch(0, selection) + if err != nil { + return nil, err + } } e.probeReader.mergeRecords = e.cfg.MergeProbeBatches || e.cfg.IsDedup bat, err := e.probeReader.ReadBatch(proc, e.probeReadBatch) @@ -2062,7 +2184,16 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana e.buckets[0].BuildFd = nil // prevent Cleanup double-close on error defer e.buildReader.closeCurrentFile() if e.buildReadBatch == nil { - e.buildReadBatch = batch.NewOffHeapWithSize(0) + var selection *vector.AllocationAccountSelection + if e.allocation != nil { + selection = e.allocation.decoded + } + readBatch, err := newSpillBatch(0, selection) + if err != nil { + builder.Free(proc) + return nil, BucketSkip, err + } + e.buildReadBatch = readBatch } // A rebuild may pre-admit one scatter workspace so a retained-copy reject // can still repartition the batches already owned by the builder. Release it @@ -2287,12 +2418,38 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal // Cache key executors. if len(e.keyExecs) != len(e.cfg.BuildKeyExprs) { - execs, lease, err := hashbuild.NewBudgetedExpressionExecutors( - proc, - e.cfg.Budget, - e.cfg.BuildKeyExprs, - false, - ) + var execs []colexec.ExpressionExecutor + var lease *hashbuild.ExpressionMemoryLease + var err error + if e.allocation == nil { + execs, lease, err = + hashbuild.NewBudgetedExpressionExecutors( + proc, + e.cfg.Budget, + e.cfg.BuildKeyExprs, + false, + ) + } else { + execs, err = + colexec.NewExpressionExecutorsFromPlanExpressionsWithAllocation( + proc, + e.cfg.BuildKeyExprs, + e.allocation.expression, + ) + if err == nil { + lease, err = hashbuild.NewExpressionMemoryLease( + nil, + e.cfg.BuildKeyExprs, + execs, + false, + ) + } + if err != nil { + for _, exec := range execs { + exec.Free() + } + } + } if err != nil { return nil, err } @@ -2385,7 +2542,15 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal } if e.probeReadBatch == nil { - e.probeReadBatch = batch.NewOffHeapWithSize(0) + var selection *vector.AllocationAccountSelection + if e.allocation != nil { + selection = e.allocation.decoded + } + readBatch, err := newSpillBatch(0, selection) + if err != nil { + return nil, err + } + e.probeReadBatch = readBatch } // Scatter probe file. Reuse reader's 4 MiB buffer from the build pass. From 70c231947057f8b013ad00cc6fd83426300f313e Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 15:03:59 +0800 Subject: [PATCH 06/61] feat: add allocation-accounted streaming buffers --- pkg/common/bitmap/bitmap.go | 41 +++++- pkg/common/bitmap/bitmap_test.go | 16 ++- pkg/common/mpool/accounted_buffer.go | 166 ++++++++++++++++++++++ pkg/common/mpool/accounted_buffer_test.go | 111 +++++++++++++++ pkg/container/batch/batch.go | 141 +++++++++++++++--- pkg/container/batch/batch_test.go | 28 ++++ pkg/container/nulls/nulls.go | 15 ++ pkg/container/nulls/nulls_test.go | 20 +++ pkg/container/vector/vector.go | 116 ++++++++++++--- pkg/container/vector/vector_test.go | 32 +++++ 10 files changed, 636 insertions(+), 50 deletions(-) create mode 100644 pkg/common/mpool/accounted_buffer.go create mode 100644 pkg/common/mpool/accounted_buffer_test.go diff --git a/pkg/common/bitmap/bitmap.go b/pkg/common/bitmap/bitmap.go index 80ed391dff5a0..7f8a3150ae09e 100644 --- a/pkg/common/bitmap/bitmap.go +++ b/pkg/common/bitmap/bitmap.go @@ -18,6 +18,7 @@ import ( "bytes" "encoding" "fmt" + "io" "math/bits" "unsafe" @@ -394,15 +395,43 @@ func (n *Bitmap) ToI64Array(out *[]int64) []int64 { func (n *Bitmap) Marshal() []byte { var buf bytes.Buffer - u1 := uint64(n.len) - u2 := uint64(len(n.data) * 8) - buf.Write(types.EncodeInt64(&n.count)) - buf.Write(types.EncodeUint64(&u1)) - buf.Write(types.EncodeUint64(&u2)) - buf.Write(types.EncodeSlice(n.data)) + _ = n.MarshalTo(&buf) return buf.Bytes() } +func (n *Bitmap) MarshalSize() int { + if n == nil { + return 0 + } + return 24 + len(n.data)*8 +} + +func (n *Bitmap) MarshalTo(w io.Writer) error { + if n == nil { + return nil + } + if w == nil { + return io.ErrClosedPipe + } + bitLength := uint64(n.len) + dataLength := uint64(len(n.data) * 8) + for _, value := range [][]byte{ + types.EncodeInt64(&n.count), + types.EncodeUint64(&bitLength), + types.EncodeUint64(&dataLength), + types.EncodeSlice(n.data), + } { + written, err := w.Write(value) + if err != nil { + return err + } + if written != len(value) { + return io.ErrShortWrite + } + } + return nil +} + // MarshalV1 in version 1, Bitmap.emptyFlag is type int32, now we use Bitmap.count replace it func (n *Bitmap) MarshalV1() []byte { var buf bytes.Buffer diff --git a/pkg/common/bitmap/bitmap_test.go b/pkg/common/bitmap/bitmap_test.go index 35bd38cb43b0d..c187be6bba6c2 100644 --- a/pkg/common/bitmap/bitmap_test.go +++ b/pkg/common/bitmap/bitmap_test.go @@ -15,7 +15,9 @@ package bitmap import ( + "bytes" "fmt" + "io" "testing" "github.com/stretchr/testify/require" @@ -32,6 +34,12 @@ func newBm(n int) *Bitmap { return &bm } +type shortMarshalWriter struct{} + +func (shortMarshalWriter) Write(value []byte) (int, error) { + return len(value) - 1, nil +} + func TestNulls(t *testing.T) { np := newBm(Rows) np.AddRange(0, 0) @@ -68,7 +76,13 @@ func TestNulls(t *testing.T) { fmt.Printf("numbers: %v\n", np.Count()) nq := newBm(Rows) - nq.Unmarshal(np.Marshal()) + encoded := np.Marshal() + var streamed bytes.Buffer + require.NoError(t, np.MarshalTo(&streamed)) + require.Equal(t, encoded, streamed.Bytes()) + require.Equal(t, len(encoded), np.MarshalSize()) + require.ErrorIs(t, np.MarshalTo(shortMarshalWriter{}), io.ErrShortWrite) + nq.Unmarshal(encoded) require.Equal(t, np.ToArray(), nq.ToArray()) diff --git a/pkg/common/mpool/accounted_buffer.go b/pkg/common/mpool/accounted_buffer.go new file mode 100644 index 0000000000000..af49b6504bedf --- /dev/null +++ b/pkg/common/mpool/accounted_buffer.go @@ -0,0 +1,166 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mpool + +import ( + "math" +) + +// AccountedBuffer is a non-copyable, allocation-accounted off-heap byte +// buffer. Reset retains physical capacity; Free is its terminal release. +type AccountedBuffer struct { + data []byte + mp *MPool + + account *AllocationAccount + owner AllocationOwner + site AllocationSite +} + +func NewAccountedBuffer( + mp *MPool, + account *AllocationAccount, + owner AllocationOwner, + site AllocationSite, +) (*AccountedBuffer, error) { + if mp == nil { + return nil, ErrAllocationAccountInvalid + } + request := allocationAccountRequest{ + account: account, + owner: owner, + site: site, + } + if err := request.validate(); err != nil { + return nil, err + } + return &AccountedBuffer{ + mp: mp, + account: account, + owner: owner, + site: site, + }, nil +} + +func (b *AccountedBuffer) Bytes() []byte { + if b == nil { + return nil + } + return b.data +} + +func (b *AccountedBuffer) Len() int { + if b == nil { + return 0 + } + return len(b.data) +} + +func (b *AccountedBuffer) Cap() int { + if b == nil { + return 0 + } + return cap(b.data) +} + +// EnsureCapacity admits and allocates an absolute retained capacity. +func (b *AccountedBuffer) EnsureCapacity(required int) error { + if b == nil || b.mp == nil || b.account == nil || required < 0 { + return ErrAllocationAccountInvalid + } + if required <= cap(b.data) { + return nil + } + if int64(required) > maxAllocationSize() { + return ErrAllocationAccountInvalid + } + + oldLength := len(b.data) + capacity, ok := GrowCapacity(int64(cap(b.data)), int64(required)) + if !ok || capacity > int64(math.MaxInt) { + return ErrAllocationAccountInvalid + } + if cap(b.data) == 0 { + data, err := b.mp.AllocAccounted( + int(capacity), + b.account, + b.owner, + b.site, + ) + if err != nil { + return err + } + b.data = data[:oldLength] + return nil + } + + data, err := b.mp.Grow(b.data, int(capacity), true) + if err != nil { + return err + } + b.data = data[:oldLength] + return nil +} + +func (b *AccountedBuffer) Write(value []byte) (int, error) { + if b == nil { + return 0, ErrAllocationAccountInvalid + } + if len(value) > math.MaxInt-len(b.data) { + return 0, ErrAllocationAccountInvalid + } + oldLength := len(b.data) + required := oldLength + len(value) + if err := b.EnsureCapacity(required); err != nil { + return 0, err + } + b.data = b.data[:required] + copy(b.data[oldLength:], value) + return len(value), nil +} + +func (b *AccountedBuffer) WriteString(value string) (int, error) { + if b == nil || len(value) > math.MaxInt-len(b.data) { + return 0, ErrAllocationAccountInvalid + } + oldLength := len(b.data) + required := oldLength + len(value) + if err := b.EnsureCapacity(required); err != nil { + return 0, err + } + b.data = b.data[:required] + copy(b.data[oldLength:], value) + return len(value), nil +} + +func (b *AccountedBuffer) Reset() { + if b != nil { + b.data = b.data[:0] + } +} + +func (b *AccountedBuffer) Free() { + if b == nil { + return + } + if cap(b.data) > 0 { + b.mp.Free(b.data) + } + b.data = nil + b.mp = nil + b.account = nil + b.owner = 0 + b.site = 0 +} diff --git a/pkg/common/mpool/accounted_buffer_test.go b/pkg/common/mpool/accounted_buffer_test.go new file mode 100644 index 0000000000000..5fa9d896fdb6a --- /dev/null +++ b/pkg/common/mpool/accounted_buffer_test.go @@ -0,0 +1,111 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mpool + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestAccountedBufferLifecycle(t *testing.T) { + registry, account := newTestAllocationAccount(t, 1<<20, 4) + mp := MustNew("accounted-buffer") + defer DeleteMPool(mp) + buffer, err := NewAccountedBuffer( + mp, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + + require.NoError(t, buffer.EnsureCapacity(32)) + firstCapacity := buffer.Cap() + require.GreaterOrEqual(t, firstCapacity, 32) + require.Equal(t, uint64(firstCapacity), account.Snapshot().Used) + _, err = buffer.WriteString("accounted") + require.NoError(t, err) + require.Equal(t, "accounted", string(buffer.Bytes())) + + buffer.Reset() + require.Zero(t, buffer.Len()) + require.Equal(t, firstCapacity, buffer.Cap()) + _, err = buffer.Write([]byte("reuse")) + require.NoError(t, err) + require.Equal(t, uint64(firstCapacity), account.Snapshot().Used) + + require.NoError(t, buffer.EnsureCapacity(firstCapacity+1)) + secondCapacity := buffer.Cap() + require.Greater(t, secondCapacity, firstCapacity) + snapshot := account.Snapshot() + require.Equal(t, uint64(secondCapacity), snapshot.Used) + require.GreaterOrEqual( + t, + snapshot.Peak, + uint64(firstCapacity+secondCapacity), + ) + require.Equal(t, "reuse", string(buffer.Bytes())) + + buffer.Free() + buffer.Free() + require.Zero(t, account.Snapshot().Used) + finalizeTestAllocationAccount(t, registry, account) +} + +func TestAccountedBufferFailureRetainsPublishedData(t *testing.T) { + registry, account := newTestAllocationAccount(t, 64, 2) + mp := MustNew("accounted-buffer-failure") + defer DeleteMPool(mp) + buffer, err := NewAccountedBuffer( + mp, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + _, err = buffer.Write([]byte("published")) + require.NoError(t, err) + before := append([]byte(nil), buffer.Bytes()...) + snapshot := account.Snapshot() + used := snapshot.Used + + err = buffer.EnsureCapacity(usedSizeToInt(t, snapshot.Limit)) + require.ErrorIs(t, err, ErrAllocationAccountCapacity) + require.Equal(t, before, buffer.Bytes()) + require.Equal(t, used, account.Snapshot().Used) + + buffer.Free() + finalizeTestAllocationAccount(t, registry, account) +} + +func TestAccountedBufferConfiguration(t *testing.T) { + _, err := NewAccountedBuffer(nil, nil, 0, 0) + require.ErrorIs(t, err, ErrAllocationAccountInvalid) + var buffer *AccountedBuffer + require.Nil(t, buffer.Bytes()) + require.Zero(t, buffer.Len()) + require.Zero(t, buffer.Cap()) + _, err = buffer.Write([]byte("x")) + require.ErrorIs(t, err, ErrAllocationAccountInvalid) + buffer.Reset() + buffer.Free() +} + +func usedSizeToInt(t testing.TB, value uint64) int { + t.Helper() + require.LessOrEqual(t, value, uint64(^uint(0)>>1)) + return int(value) +} diff --git a/pkg/container/batch/batch.go b/pkg/container/batch/batch.go index 589fe21f66600..a1c8cf0724d53 100644 --- a/pkg/container/batch/batch.go +++ b/pkg/container/batch/batch.go @@ -112,50 +112,145 @@ func (bat *Batch) MarshalBinary() ([]byte, error) { } func (bat *Batch) MarshalBinaryWithBuffer(w *bytes.Buffer, reset bool) ([]byte, error) { - // reset the buffer if caller wants to. if reset { w.Reset() } + if err := bat.MarshalBinaryTo(w); err != nil { + return nil, err + } + return w.Bytes(), nil +} + +func (bat *Batch) MarshalBinarySize() (int, error) { + if bat == nil { + return 0, moerr.NewInvalidInputNoCtx("invalid batch for marshal") + } + const fixedSize = uint64(8 + 4 + 4 + 4 + 4 + 4) + total := fixedSize + add := func(value uint64) bool { + if value > uint64(^uint(0)>>1)-total { + return false + } + total += value + return true + } + if uint64(len(bat.Vecs)) > uint64(^uint32(0)>>1) || + uint64(len(bat.Attrs)) > uint64(^uint32(0)>>1) || + uint64(len(bat.ExtraBuf)) > uint64(^uint32(0)>>1) { + return 0, moerr.NewInvalidInputNoCtx( + "batch field exceeds marshal format", + ) + } + for _, vec := range bat.Vecs { + if vec == nil { + return 0, moerr.NewInvalidInputNoCtx( + "cannot marshal a nil batch vector", + ) + } + size, err := vec.MarshalBinarySize() + if err != nil { + return 0, err + } + if uint64(size) > uint64(^uint32(0)) || + !add(4+uint64(size)) { + return 0, moerr.NewInvalidInputNoCtx( + "batch vector exceeds marshal format", + ) + } + } + for _, attr := range bat.Attrs { + if uint64(len(attr)) > uint64(^uint32(0)>>1) || + !add(4+uint64(len(attr))) { + return 0, moerr.NewInvalidInputNoCtx( + "batch attribute exceeds marshal format", + ) + } + } + if !add(uint64(len(bat.ExtraBuf))) { + return 0, moerr.NewInvalidInputNoCtx( + "batch marshal size exceeds platform limit", + ) + } + return int(total), nil +} - // row count. +func (bat *Batch) MarshalBinaryTo(w io.Writer) error { + if bat == nil || w == nil { + return io.ErrClosedPipe + } + if _, err := bat.MarshalBinarySize(); err != nil { + return err + } rl := int64(bat.rowCount) - w.Write(types.EncodeInt64(&rl)) + if err := writeBatchMarshalBytes(w, types.EncodeInt64(&rl)); err != nil { + return err + } - // Vecs l := int32(len(bat.Vecs)) - w.Write(types.EncodeInt32(&l)) + if err := writeBatchMarshalBytes(w, types.EncodeInt32(&l)); err != nil { + return err + } for i := 0; i < int(l); i++ { - var size uint32 - offset := w.Len() - w.Write(types.EncodeUint32(&size)) - err := bat.Vecs[i].MarshalBinaryWithBuffer(w) + size, err := bat.Vecs[i].MarshalBinarySize() if err != nil { - return nil, err + return err + } + wireSize := uint32(size) + if err := writeBatchMarshalBytes( + w, + types.EncodeUint32(&wireSize), + ); err != nil { + return err + } + if err := bat.Vecs[i].MarshalBinaryTo(w); err != nil { + return err } - size = uint32(w.Len() - offset - 4) - buf := w.Bytes() - copy(buf[offset:], types.EncodeUint32(&size)) } - // Attrs l = int32(len(bat.Attrs)) - w.Write(types.EncodeInt32(&l)) + if err := writeBatchMarshalBytes(w, types.EncodeInt32(&l)); err != nil { + return err + } for i := 0; i < int(l); i++ { size := int32(len(bat.Attrs[i])) - w.Write(types.EncodeInt32(&size)) - n, _ := w.WriteString(bat.Attrs[i]) + if err := writeBatchMarshalBytes(w, types.EncodeInt32(&size)); err != nil { + return err + } + n, err := io.WriteString(w, bat.Attrs[i]) + if err != nil { + return err + } if int32(n) != size { - panic("unexpected length for string") + return io.ErrShortWrite } } - // ExtraBuf - types.WriteSizeBytes(bat.ExtraBuf, w) + extraSize := int32(len(bat.ExtraBuf)) + if err := writeBatchMarshalBytes(w, types.EncodeInt32(&extraSize)); err != nil { + return err + } + if err := writeBatchMarshalBytes(w, bat.ExtraBuf); err != nil { + return err + } - w.Write(types.EncodeInt32(&bat.Recursive)) - w.Write(types.EncodeInt32(&bat.ShuffleIDX)) + if err := writeBatchMarshalBytes( + w, + types.EncodeInt32(&bat.Recursive), + ); err != nil { + return err + } + return writeBatchMarshalBytes(w, types.EncodeInt32(&bat.ShuffleIDX)) +} - return w.Bytes(), nil +func writeBatchMarshalBytes(w io.Writer, value []byte) error { + written, err := w.Write(value) + if err != nil { + return err + } + if written != len(value) { + return io.ErrShortWrite + } + return nil } func (bat *Batch) UnmarshalBinary(data []byte) (err error) { diff --git a/pkg/container/batch/batch_test.go b/pkg/container/batch/batch_test.go index 0044033f8c1e6..673f03796960c 100644 --- a/pkg/container/batch/batch_test.go +++ b/pkg/container/batch/batch_test.go @@ -17,6 +17,7 @@ package batch import ( "bytes" "fmt" + "io" "testing" "github.com/matrixorigin/matrixone/pkg/common/mpool" @@ -51,6 +52,17 @@ func TestBatchMarshalAndUnmarshal(t *testing.T) { for _, tc := range tcs { data, err := tc.bat.MarshalBinary() require.NoError(t, err) + size, err := tc.bat.MarshalBinarySize() + require.NoError(t, err) + require.Equal(t, len(data), size) + var streamed bytes.Buffer + require.NoError(t, tc.bat.MarshalBinaryTo(&streamed)) + require.Equal(t, data, streamed.Bytes()) + require.ErrorIs( + t, + tc.bat.MarshalBinaryTo(shortBatchMarshalWriter{}), + io.ErrShortWrite, + ) rbat := new(Batch) err = rbat.UnmarshalBinary(data) @@ -85,6 +97,22 @@ func TestBatchMarshalAndUnmarshal(t *testing.T) { } } +type shortBatchMarshalWriter struct{} + +func (shortBatchMarshalWriter) Write(value []byte) (int, error) { + return len(value) - 1, nil +} + +func TestMarshalBinarySizeRejectsInvalidBatch(t *testing.T) { + var nilBatch *Batch + _, err := nilBatch.MarshalBinarySize() + require.Error(t, err) + + invalid := NewWithSize(1) + _, err = invalid.MarshalBinarySize() + require.Error(t, err) +} + func TestBatch(t *testing.T) { for _, tc := range tcs { data, err := types.Encode(tc.bat) diff --git a/pkg/container/nulls/nulls.go b/pkg/container/nulls/nulls.go index f6b11432c6d62..fc16f895b6265 100644 --- a/pkg/container/nulls/nulls.go +++ b/pkg/container/nulls/nulls.go @@ -19,6 +19,7 @@ package nulls import ( "fmt" + "io" "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/util" @@ -372,6 +373,20 @@ func (nsp *Nulls) Show() ([]byte, error) { return nsp.np.Marshal(), nil } +func (nsp *Nulls) MarshalSize() int { + if nsp == nil || nsp.np.EmptyByFlag() { + return 0 + } + return nsp.np.MarshalSize() +} + +func (nsp *Nulls) MarshalTo(w io.Writer) error { + if nsp == nil || nsp.np.EmptyByFlag() { + return nil + } + return nsp.np.MarshalTo(w) +} + // ShowV1 in version 1, bitmap is v1 func (nsp *Nulls) ShowV1() ([]byte, error) { if nsp.np.EmptyByFlag() { diff --git a/pkg/container/nulls/nulls_test.go b/pkg/container/nulls/nulls_test.go index 6c9d24cc916d6..f4f4f895434e9 100644 --- a/pkg/container/nulls/nulls_test.go +++ b/pkg/container/nulls/nulls_test.go @@ -15,9 +15,11 @@ package nulls import ( + "bytes" "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestOr(t *testing.T) { @@ -68,6 +70,24 @@ func TestAny(t *testing.T) { }) } +func TestMarshalTo(t *testing.T) { + var n Nulls + var empty bytes.Buffer + require.NoError(t, n.MarshalTo(&empty)) + require.Zero(t, n.MarshalSize()) + require.Zero(t, empty.Len()) + + n.InitWithSize(128) + n.Add(3) + n.Add(65) + encoded, err := n.Show() + require.NoError(t, err) + var streamed bytes.Buffer + require.NoError(t, n.MarshalTo(&streamed)) + require.Equal(t, encoded, streamed.Bytes()) + require.Equal(t, len(encoded), n.MarshalSize()) +} + func TestSize(t *testing.T) { t.Run("Size test", func(t *testing.T) { var n Nulls diff --git a/pkg/container/vector/vector.go b/pkg/container/vector/vector.go index d5074f4eda726..66e5775238141 100644 --- a/pkg/container/vector/vector.go +++ b/pkg/container/vector/vector.go @@ -759,50 +759,126 @@ func (v *Vector) MarshalBinary() ([]byte, error) { } func (v *Vector) MarshalBinaryWithBuffer(buf *bytes.Buffer) error { + return v.MarshalBinaryTo(buf) +} - // write class - buf.WriteByte(uint8(v.class)) +func (v *Vector) MarshalBinarySize() (int, error) { + if v == nil || v.length < 0 { + return 0, moerr.NewInvalidInputNoCtx("invalid vector for marshal") + } + const maxWireBuffer = uint64(^uint32(0)) + if uint64(v.length) > maxWireBuffer { + return 0, moerr.NewInvalidInputNoCtx( + "vector length exceeds marshal format", + ) + } + typeSize := v.typ.TypeSize() + if typeSize < 0 { + return 0, moerr.NewInvalidInputNoCtx( + "vector type has invalid marshal size", + ) + } + dataLength := uint64(typeSize) + if !v.IsConst() { + if v.length != 0 && + dataLength > ^uint64(0)/uint64(v.length) { + return 0, moerr.NewInvalidInputNoCtx( + "vector data exceeds marshal format", + ) + } + dataLength *= uint64(v.length) + } else if v.IsConstNull() { + dataLength = 0 + } + areaLength := uint64(len(v.area)) + nullLength := uint64(v.nsp.MarshalSize()) + if dataLength > maxWireBuffer || + areaLength > maxWireBuffer || + nullLength > maxWireBuffer { + return 0, moerr.NewInvalidInputNoCtx( + "vector buffer exceeds marshal format", + ) + } + if dataLength > uint64(len(v.data)) { + return 0, moerr.NewInvalidInputNoCtx( + "vector data is shorter than its marshal length", + ) + } + total := uint64(1+types.TSize+4+4+4+4+1) + + dataLength + areaLength + nullLength + if total > uint64(^uint(0)>>1) { + return 0, moerr.NewInvalidInputNoCtx( + "vector marshal size exceeds platform limit", + ) + } + return int(total), nil +} - // write type - data := types.EncodeType(&v.typ) - buf.Write(data) +func (v *Vector) MarshalBinaryTo(w io.Writer) error { + if w == nil { + return io.ErrClosedPipe + } + if _, err := v.MarshalBinarySize(); err != nil { + return err + } + if err := writeVectorMarshalBytes(w, []byte{uint8(v.class)}); err != nil { + return err + } + if err := writeVectorMarshalBytes(w, types.EncodeType(&v.typ)); err != nil { + return err + } - // write length length := uint32(v.length) - buf.Write(types.EncodeUint32(&length)) + if err := writeVectorMarshalBytes(w, types.EncodeUint32(&length)); err != nil { + return err + } - // write dataLen, data dataLen := uint32(v.typ.TypeSize()) if !v.IsConst() { dataLen *= uint32(v.length) } else if v.IsConstNull() { dataLen = 0 } - buf.Write(types.EncodeUint32(&dataLen)) + if err := writeVectorMarshalBytes(w, types.EncodeUint32(&dataLen)); err != nil { + return err + } if dataLen > 0 { - buf.Write(v.data[:dataLen]) + if err := writeVectorMarshalBytes(w, v.data[:dataLen]); err != nil { + return err + } } - // write areaLen, area areaLen := uint32(len(v.area)) - buf.Write(types.EncodeUint32(&areaLen)) + if err := writeVectorMarshalBytes(w, types.EncodeUint32(&areaLen)); err != nil { + return err + } if areaLen > 0 { - buf.Write(v.area) + if err := writeVectorMarshalBytes(w, v.area); err != nil { + return err + } } - // write nspLen, nsp - nspData, err := v.nsp.Show() - if err != nil { + nspLen := uint32(v.nsp.MarshalSize()) + if err := writeVectorMarshalBytes(w, types.EncodeUint32(&nspLen)); err != nil { return err } - nspLen := uint32(len(nspData)) - buf.Write(types.EncodeUint32(&nspLen)) if nspLen > 0 { - buf.Write(nspData) + if err := v.nsp.MarshalTo(w); err != nil { + return err + } } - buf.Write(types.EncodeBool(&v.sorted)) + return writeVectorMarshalBytes(w, types.EncodeBool(&v.sorted)) +} +func writeVectorMarshalBytes(w io.Writer, value []byte) error { + written, err := w.Write(value) + if err != nil { + return err + } + if written != len(value) { + return io.ErrShortWrite + } return nil } diff --git a/pkg/container/vector/vector_test.go b/pkg/container/vector/vector_test.go index 32ecd739b2930..cd034561dc54d 100644 --- a/pkg/container/vector/vector_test.go +++ b/pkg/container/vector/vector_test.go @@ -15,7 +15,9 @@ package vector import ( + "bytes" "fmt" + "io" "slices" "strings" "testing" @@ -1687,6 +1689,13 @@ func TestMarshalAndUnMarshal(t *testing.T) { require.NoError(t, err) data, err := v.MarshalBinary() require.NoError(t, err) + size, err := v.MarshalBinarySize() + require.NoError(t, err) + require.Equal(t, len(data), size) + var streamed bytes.Buffer + require.NoError(t, v.MarshalBinaryTo(&streamed)) + require.Equal(t, data, streamed.Bytes()) + require.ErrorIs(t, v.MarshalBinaryTo(shortVectorMarshalWriter{}), io.ErrShortWrite) w := NewVecFromReuse() err = w.UnmarshalBinary(data) require.NoError(t, err) @@ -1701,6 +1710,29 @@ func TestMarshalAndUnMarshal(t *testing.T) { require.Equal(t, int64(0), mp.CurrNB()) } +type shortVectorMarshalWriter struct{} + +func (shortVectorMarshalWriter) Write(value []byte) (int, error) { + return len(value) - 1, nil +} + +func TestMarshalBinarySizeRejectsInvalidVector(t *testing.T) { + var nilVector *Vector + _, err := nilVector.MarshalBinarySize() + require.Error(t, err) + + typ := types.T_int64.ToType() + typ.Size = -1 + invalidType := NewVec(typ) + _, err = invalidType.MarshalBinarySize() + require.Error(t, err) + + shortData := NewVec(types.T_int64.ToType()) + shortData.SetLength(1) + _, err = shortData.MarshalBinarySize() + require.Error(t, err) +} + func TestUnmarshalBinaryAcceptsNullBitmapCoveragePastLength(t *testing.T) { mp := mpool.MustNewZero() source := NewVec(types.T_int64.ToType()) From 4cb90d6bed3456949db6d117a81f241732bcec5d Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 15:04:04 +0800 Subject: [PATCH 07/61] feat: account spill serialization buffers --- .../colexec/spillutil/allocation_account.go | 19 +- .../spillutil/allocation_account_test.go | 107 +++++++- pkg/sql/colexec/spillutil/join_spill.go | 255 +++++++++++++++--- pkg/sql/colexec/spillutil/join_spill_test.go | 16 +- 4 files changed, 346 insertions(+), 51 deletions(-) diff --git a/pkg/sql/colexec/spillutil/allocation_account.go b/pkg/sql/colexec/spillutil/allocation_account.go index 8cdfb2ad6613f..afa4cf9924cca 100644 --- a/pkg/sql/colexec/spillutil/allocation_account.go +++ b/pkg/sql/colexec/spillutil/allocation_account.go @@ -33,10 +33,12 @@ const ( SpillAllocationSiteSelectedArea SpillAllocationSiteHashValues SpillAllocationSiteRowIDs + SpillAllocationSiteMarshalBuffer + SpillAllocationSiteCoalesceBuffer ) // SpillAllocationAccount is the dormant allocation provenance for one spill -// engine. Serialization buffers remain a named activation blocker. +// engine. type SpillAllocationAccount struct { account *mpool.AllocationAccount owner mpool.AllocationOwner @@ -171,3 +173,18 @@ func freeSpillSlice[T any]( mpool.FreeSlice(mp, values) } } + +func (a *SpillAllocationAccount) newBuffer( + mp *mpool.MPool, + site mpool.AllocationSite, +) (*mpool.AccountedBuffer, error) { + if err := a.validate(); err != nil { + return nil, err + } + return mpool.NewAccountedBuffer( + mp, + a.account, + a.owner, + site, + ) +} diff --git a/pkg/sql/colexec/spillutil/allocation_account_test.go b/pkg/sql/colexec/spillutil/allocation_account_test.go index 6d74f701f3d5b..7e4f5fa1432f5 100644 --- a/pkg/sql/colexec/spillutil/allocation_account_test.go +++ b/pkg/sql/colexec/spillutil/allocation_account_test.go @@ -205,6 +205,7 @@ func TestSpillAllocationAccountScatterScratchLifecycle(t *testing.T) { writers[i].Close() } }() + analyzer := process.NewAnalyzer(0, false, false, "test") require.NoError(t, engine.scatterBatchBounded( proc, source, @@ -212,13 +213,24 @@ func TestSpillAllocationAccountScatterScratchLifecycle(t *testing.T) { writers, 0, false, - process.NewAnalyzer(0, false, false, "test"), + analyzer, )) require.Len(t, engine.scatterHashValues, source.RowCount()) require.Len(t, engine.scatterBucketRowIds, source.RowCount()) snapshot := state.account.Snapshot() - require.Equal(t, uint64(source.RowCount()*(8+4)), snapshot.Used) + require.Greater( + t, + snapshot.Used, + uint64(source.RowCount()*(8+4)), + ) require.Greater(t, snapshot.Peak, snapshot.Used) + require.NoError(t, engine.flushScatterBuffers(proc, writers, analyzer)) + var writtenRows int64 + for i := range writers { + writtenRows += writers[i].Rows + } + require.Equal(t, int64(source.RowCount()), writtenRows) + require.Equal(t, snapshot.Used, state.account.Snapshot().Used) engine.releaseScatterScratch() require.Zero(t, state.account.Snapshot().Used) @@ -226,6 +238,97 @@ func TestSpillAllocationAccountScatterScratchLifecycle(t *testing.T) { finalizeTestSpillAllocationAccount(t, state) } +func TestSpillAllocationAccountMarshalBufferLifecycle(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-marshal"), + ) + defer proc.Free() + state := newTestSpillAllocationAccount(t, 1<<20, 8) + source := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.NewVector( + 4, + types.T_int64.ToType(), + proc.Mp(), + false, + []int64{1, 2, 3, 4}, + ), + }, nil) + defer source.Clean(proc.Mp()) + + accounted, err := state.allocation.newBuffer( + proc.Mp(), + SpillAllocationSiteMarshalBuffer, + ) + require.NoError(t, err) + require.NoError(t, marshalSpillRecordTo(source, accounted)) + var legacy bytes.Buffer + require.NoError(t, marshalSpillRecord(source, &legacy)) + require.Equal(t, legacy.Bytes(), accounted.Bytes()) + used := state.account.Snapshot().Used + require.Positive(t, used) + + accounted.Reset() + require.Zero(t, accounted.Len()) + require.Equal(t, used, state.account.Snapshot().Used) + require.NoError(t, marshalSpillRecordTo(source, accounted)) + require.Equal(t, used, state.account.Snapshot().Used) + + accounted.Free() + require.Zero(t, state.account.Snapshot().Used) + finalizeTestSpillAllocationAccount(t, state) +} + +func TestSpillAllocationAccountCoalesceAdmissionFallback(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-coalesce-fallback"), + ) + defer proc.Free() + // The record buffer consumes the only metadata slot. Coalescing is an + // optional optimization, so its admission failure must fall back to one + // direct write instead of failing the scatter. + state := newTestSpillAllocationAccount(t, 1<<20, 1) + engine, err := NewSpillEngineWithAllocation( + SpillEngineConfig{}, + state.allocation, + ) + require.NoError(t, err) + source := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.NewVector( + 2, + types.T_int64.ToType(), + proc.Mp(), + false, + []int64{1, 2}, + ), + }, nil) + defer source.Clean(proc.Mp()) + writers := MakeBucketWriters("spill_allocation_coalesce_fallback") + defer func() { + for i := range writers { + writers[i].Close() + } + }() + + require.NoError(t, engine.appendScatterRecord( + proc, + source, + &writers[0], + 0, + process.NewAnalyzer(0, false, false, "test"), + )) + require.Equal(t, int64(source.RowCount()), writers[0].Rows) + require.Nil(t, engine.scatterAccountedWriteBuffers[0].Bytes()) + require.Equal(t, uint64(1), state.registry.PeakAllocationMetadata()) + + engine.releaseScatterScratch() + engine.Cleanup(proc) + finalizeTestSpillAllocationAccount(t, state) +} + func TestSpillAllocationAccountScatterFailureCleanup(t *testing.T) { proc := testutil.NewProcessWithMPool( t, diff --git a/pkg/sql/colexec/spillutil/join_spill.go b/pkg/sql/colexec/spillutil/join_spill.go index 263a7608cbc18..b7d0d210c4179 100644 --- a/pkg/sql/colexec/spillutil/join_spill.go +++ b/pkg/sql/colexec/spillutil/join_spill.go @@ -503,18 +503,6 @@ func maxIntValue() int { return int(^uint(0) >> 1) } -func marshalSpillRecordGrowBytes(bat *batch.Batch) (uint64, bool) { - base := uint64(bat.Allocated()) - if size := uint64(bat.Size()); size > base { - base = size - } - columns := uint64(len(bat.Vecs)) - if columns > (math.MaxUint64-24)/128 { - return 0, false - } - return addUint64(base, columns*128+24) -} - func (r *BucketReader) releaseReadBatch(proc *process.Process, bat *batch.Batch, token *process.HashBuildReservation) { if bat != nil { bat.Clean(proc.Mp()) @@ -890,34 +878,99 @@ func FlushBucketBatch(proc *process.Process, bat *batch.Batch, w *BucketWriter, return writeBucketPayload(proc, bucketBuf.Bytes(), cnt, w, analyzer) } +type spillRecordBuffer interface { + io.Writer + Bytes() []byte + EnsureCapacity(int) error + Len() int + Reset() +} + +type legacySpillRecordBuffer struct { + buffer *bytes.Buffer +} + +func (b legacySpillRecordBuffer) Write(value []byte) (int, error) { + return b.buffer.Write(value) +} + +func (b legacySpillRecordBuffer) Bytes() []byte { + return b.buffer.Bytes() +} + +func (b legacySpillRecordBuffer) EnsureCapacity(required int) error { + if b.buffer.Cap() < required { + *b.buffer = *bytes.NewBuffer(make([]byte, 0, required)) + } + return nil +} + +func (b legacySpillRecordBuffer) Len() int { + return b.buffer.Len() +} + +func (b legacySpillRecordBuffer) Reset() { + b.buffer.Reset() +} + func marshalSpillRecord(bat *batch.Batch, buf *bytes.Buffer) error { + if buf == nil { + return process.ErrHashBuildBudgetInvalid + } + return marshalSpillRecordTo( + bat, + legacySpillRecordBuffer{buffer: buf}, + ) +} + +func marshalSpillRecordTo( + bat *batch.Batch, + buf spillRecordBuffer, +) error { if bat == nil || bat.RowCount() == 0 { return nil } cnt := int64(bat.RowCount()) buf.Reset() - grow, ok := marshalSpillRecordGrowBytes(bat) - if !ok || grow > uint64(maxIntValue()) { + batchSize, err := bat.MarshalBinarySize() + if err != nil || batchSize > maxIntValue()-24 { + if err != nil { + return err + } return process.ErrHashBuildBudgetInvalid } - if uint64(buf.Cap()) < grow { - // Allocate the final serialization capacity in one step. Retaining a - // smaller bytes.Buffer while it grows geometrically would invalidate the - // single-payload admission estimate. - *buf = *bytes.NewBuffer(make([]byte, 0, int(grow))) + if err := buf.EnsureCapacity(batchSize + 24); err != nil { + return err + } + if err := writeSpillRecordBytes(buf, types.EncodeInt64(&cnt)); err != nil { + return err } - buf.Write(types.EncodeInt64(&cnt)) batchSizePos := buf.Len() var zero int64 - buf.Write(types.EncodeInt64(&zero)) + if err := writeSpillRecordBytes(buf, types.EncodeInt64(&zero)); err != nil { + return err + } batchStart := buf.Len() - if _, err := bat.MarshalBinaryWithBuffer(buf, false); err != nil { + if err := bat.MarshalBinaryTo(buf); err != nil { return err } - batchSize := int64(buf.Len() - batchStart) - copy(buf.Bytes()[batchSizePos:batchSizePos+8], types.EncodeInt64(&batchSize)) + serializedSize := int64(buf.Len() - batchStart) + copy( + buf.Bytes()[batchSizePos:batchSizePos+8], + types.EncodeInt64(&serializedSize), + ) magic := uint64(SpillMagic) - buf.Write(types.EncodeUint64(&magic)) + return writeSpillRecordBytes(buf, types.EncodeUint64(&magic)) +} + +func writeSpillRecordBytes(w io.Writer, value []byte) error { + written, err := w.Write(value) + if err != nil { + return err + } + if written != len(value) { + return io.ErrShortWrite + } return nil } @@ -1558,6 +1611,16 @@ func (e *SpillEngine) appendScatterRecord(proc *process.Process, bat *batch.Batc return process.ErrHashBuildBudgetInvalid } cnt := int64(bat.RowCount()) + if e.allocation != nil { + return e.appendAccountedScatterRecord( + proc, + bat, + writer, + bucket, + cnt, + analyzer, + ) + } if err := marshalSpillRecord(bat, &e.scatterWriteBuf); err != nil { return err } @@ -1587,6 +1650,86 @@ func (e *SpillEngine) appendScatterRecord(proc *process.Process, bat *batch.Batc return nil } +func (e *SpillEngine) appendAccountedScatterRecord( + proc *process.Process, + bat *batch.Batch, + writer *BucketWriter, + bucket int, + rows int64, + analyzer process.Analyzer, +) error { + if e.allocationMP != nil && e.allocationMP != proc.Mp() { + return mpool.ErrAllocationAccountInvalid + } + e.allocationMP = proc.Mp() + if e.scatterAccountedWriteBuf == nil { + var err error + e.scatterAccountedWriteBuf, err = e.allocation.newBuffer( + proc.Mp(), + SpillAllocationSiteMarshalBuffer, + ) + if err != nil { + return err + } + } + if err := marshalSpillRecordTo( + bat, + e.scatterAccountedWriteBuf, + ); err != nil { + return err + } + payload := e.scatterAccountedWriteBuf.Bytes() + buf := e.scatterAccountedWriteBuffers[bucket] + if buf != nil && buf.Len() > 0 && + buf.Len()+len(payload) > spillWriteCoalesceSize { + if err := e.flushPendingScatterBucket( + proc, + writer, + bucket, + analyzer, + ); err != nil { + return err + } + } + if len(payload) > spillWriteCoalesceSize { + return writeBucketPayload(proc, payload, rows, writer, analyzer) + } + if buf == nil { + var err error + buf, err = e.allocation.newBuffer( + proc.Mp(), + SpillAllocationSiteCoalesceBuffer, + ) + if err != nil { + return err + } + e.scatterAccountedWriteBuffers[bucket] = buf + } + if buf.Len() == 0 && buf.Cap() < spillWriteCoalesceSize { + if err := buf.EnsureCapacity(spillWriteCoalesceSize); err != nil { + if errors.Is(err, mpool.ErrAllocationAccountCapacity) || + errors.Is(err, mpool.ErrAllocationMetadataSlots) { + return writeBucketPayload( + proc, + payload, + rows, + writer, + analyzer, + ) + } + return err + } + } + if _, err := buf.Write(payload); err != nil { + return err + } + e.scatterWriteRows[bucket] += rows + if buf.Len() >= spillWriteCoalesceSize { + return e.flushPendingScatterBucket(proc, writer, bucket, analyzer) + } + return nil +} + func (e *SpillEngine) ensureScatterCoalesceCapacity(buf *bytes.Buffer) bool { if buf == nil || buf.Cap() >= spillWriteCoalesceSize { return true @@ -1605,6 +1748,22 @@ func (e *SpillEngine) flushPendingScatterBucket(proc *process.Process, writer *B if bucket < 0 || bucket >= SpillNumBuckets || writer == nil { return process.ErrHashBuildBudgetInvalid } + if e.allocation != nil { + buf := e.scatterAccountedWriteBuffers[bucket] + if buf == nil || buf.Len() == 0 { + return nil + } + err := writeBucketPayload( + proc, + buf.Bytes(), + e.scatterWriteRows[bucket], + writer, + analyzer, + ) + buf.Reset() + e.scatterWriteRows[bucket] = 0 + return err + } buf := &e.scatterWriteBuffers[bucket] if buf.Len() == 0 { return nil @@ -1621,7 +1780,12 @@ func (e *SpillEngine) flushPendingScatterBucket(proc *process.Process, writer *B func (e *SpillEngine) flushScatterBuffers(proc *process.Process, writers []BucketWriter, analyzer process.Analyzer) error { var firstErr error for bucket := 0; bucket < SpillNumBuckets; bucket++ { - if e.scatterWriteBuffers[bucket].Len() == 0 { + pending := e.scatterWriteBuffers[bucket].Len() + if e.allocation != nil && + e.scatterAccountedWriteBuffers[bucket] != nil { + pending = e.scatterAccountedWriteBuffers[bucket].Len() + } + if pending == 0 { continue } var writer *BucketWriter @@ -1638,6 +1802,9 @@ func (e *SpillEngine) flushScatterBuffers(proc *process.Process, writers []Bucke func (e *SpillEngine) discardScatterBuffers() { for bucket := range e.scatterWriteBuffers { e.scatterWriteBuffers[bucket].Reset() + if e.scatterAccountedWriteBuffers[bucket] != nil { + e.scatterAccountedWriteBuffers[bucket].Reset() + } e.scatterWriteRows[bucket] = 0 } } @@ -1659,9 +1826,12 @@ func (e *SpillEngine) releaseScatterScratch() { ) e.scatterHashValues = nil e.scatterBucketRowIds = nil - e.allocationMP = nil e.keyVecs = nil e.scatterWriteBuf = bytes.Buffer{} + if e.scatterAccountedWriteBuf != nil { + e.scatterAccountedWriteBuf.Free() + e.scatterAccountedWriteBuf = nil + } for i := range e.scatterBucketCounts { e.scatterBucketCounts[i] = 0 } @@ -1670,8 +1840,13 @@ func (e *SpillEngine) releaseScatterScratch() { } for i := range e.scatterWriteBuffers { e.scatterWriteBuffers[i] = bytes.Buffer{} + if e.scatterAccountedWriteBuffers[i] != nil { + e.scatterAccountedWriteBuffers[i].Free() + e.scatterAccountedWriteBuffers[i] = nil + } e.scatterWriteRows[i] = 0 } + e.allocationMP = nil if e.scatterScratchReservation != nil { e.scatterScratchReservation.Release() e.scatterScratchReservation = nil @@ -1680,8 +1855,8 @@ func (e *SpillEngine) releaseScatterScratch() { } // reconcileScatterScratch leaves only the capacities retained by the engine -// charged after a batch completes. The source batch, selected vectors, and -// marshal buffer are transient and must not pin budget across the queue. +// charged after a batch completes. Source and selected vectors are transient; +// reusable marshal and coalesce buffers remain charged only for the phase. func (e *SpillEngine) reconcileScatterScratch() error { if e.scatterScratchReservation == nil { return nil @@ -1806,16 +1981,18 @@ type SpillEngine struct { buildExprLease *hashbuild.ExpressionMemoryLease // Reusable scatter buffers to avoid per-batch allocations. - scatterHashValues []uint64 - scatterBucketRowIds []int32 - scatterBucketCounts [SpillNumBuckets]int32 - scatterBucketOffsets [SpillNumBuckets + 1]int32 - scatterWriteBuf bytes.Buffer - scatterWriteBuffers [SpillNumBuckets]bytes.Buffer - scatterWriteRows [SpillNumBuckets]int64 - // The lease follows reusable scratch capacities within one rebuild/scatter - // phase. releaseScatterScratch drops both the backing arrays and this token; - // Cleanup is the idempotent terminal fallback. + scatterHashValues []uint64 + scatterBucketRowIds []int32 + scatterBucketCounts [SpillNumBuckets]int32 + scatterBucketOffsets [SpillNumBuckets + 1]int32 + scatterWriteBuf bytes.Buffer + scatterWriteBuffers [SpillNumBuckets]bytes.Buffer + scatterAccountedWriteBuf *mpool.AccountedBuffer + scatterAccountedWriteBuffers [SpillNumBuckets]*mpool.AccountedBuffer + scatterWriteRows [SpillNumBuckets]int64 + // The lease follows the reusable scratch capacities for the engine + // lifetime. It is released only by Cleanup, after all backing arrays have + // been dropped. scatterScratchReservation *process.HashBuildReservation // scatterScratchFloor is pre-admitted only while rebuilding an already // spilled bucket. It keeps one bounded repartition workspace available if the diff --git a/pkg/sql/colexec/spillutil/join_spill_test.go b/pkg/sql/colexec/spillutil/join_spill_test.go index 3ca3f369d8d18..5473ed6437e18 100644 --- a/pkg/sql/colexec/spillutil/join_spill_test.go +++ b/pkg/sql/colexec/spillutil/join_spill_test.go @@ -1463,7 +1463,7 @@ func TestReaderBatchLeaseUsesSinglePayloadEstimate(t *testing.T) { require.Zero(t, generation.Used()) } -func TestMarshalSpillRecordPreallocatesSinglePayload(t *testing.T) { +func TestMarshalSpillRecordPreallocatesExactPayload(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() bat := batch.NewWithSize(1) @@ -1477,11 +1477,9 @@ func TestMarshalSpillRecordPreallocatesSinglePayload(t *testing.T) { buf := bytes.NewBuffer(make([]byte, 0, 1<<20)) require.NoError(t, marshalSpillRecord(bat, buf)) - base := uint64(bat.Allocated()) - if size := uint64(bat.Size()); size > base { - base = size - } - require.Equal(t, base+128+24, uint64(buf.Cap())) + size, err := bat.MarshalBinarySize() + require.NoError(t, err) + require.Equal(t, size+24, buf.Cap()) small := batch.NewWithSize(1) small.Vecs[0], err = vector.NewConstBytes( @@ -4109,9 +4107,9 @@ func TestScatterPeakDoesNotDoubleChargeReservedSource(t *testing.T) { require.True(t, ok) growth, ok := emptyEngine.scatterCapacityGrowthBytes(bat.RowCount(), 1) require.True(t, ok) - marshalOverlap, ok := marshalSpillRecordGrowBytes(bat) - require.True(t, ok) - capacity := source + retained + growth + charged + marshalOverlap + marshalSize, err := bat.MarshalBinarySize() + require.NoError(t, err) + capacity := source + retained + growth + charged + uint64(marshalSize+24) budget, err := process.NewHashBuildBudget(capacity, capacity) require.NoError(t, err) generation, err := budget.OpenGeneration(capacity) From 54c973b83f348fa2ae904cd037f776f8c7364439 Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 15:21:51 +0800 Subject: [PATCH 08/61] feat: close retained vector allocation ownership gaps --- pkg/container/pSpool/buffer.go | 36 ++- pkg/container/pSpool/copy.go | 208 +++++++++++++++--- pkg/container/pSpool/copy_benchmark_test.go | 58 +++++ pkg/container/pSpool/sender_test.go | 158 +++++++++++++ .../vector/allocation_account_test.go | 113 ++++++++++ pkg/container/vector/pSpoolTools.go | 153 +++++++++++-- pkg/container/vector/vector.go | 26 ++- pkg/sql/plan/function/func_binary.go | 28 ++- 8 files changed, 719 insertions(+), 61 deletions(-) create mode 100644 pkg/container/pSpool/copy_benchmark_test.go diff --git a/pkg/container/pSpool/buffer.go b/pkg/container/pSpool/buffer.go index 686eba17ca736..541c9ef9acee4 100644 --- a/pkg/container/pSpool/buffer.go +++ b/pkg/container/pSpool/buffer.go @@ -63,14 +63,30 @@ func (b *spoolBuffer) putCacheID(mp *mpool.MPool, id uint32, bat *batch.Batch) { vec.Free(mp) } - data := vector.GetAndClearVecData(vec) - area := vector.GetAndClearVecArea(vec) - - if data != nil { - b.bytesCache[id].bs = append(b.bytesCache[id].bs, data) - } - if area != nil { - b.bytesCache[id].bs = append(b.bytesCache[id].bs, area) + if vec.AllocationAccountSelection() == nil { + data := vector.DetachLegacyVectorData(vec) + area := vector.DetachLegacyVectorArea(vec) + if data != nil { + b.bytesCache[id].bs = append(b.bytesCache[id].bs, data) + } + if area != nil { + b.bytesCache[id].bs = append(b.bytesCache[id].bs, area) + } + } else { + data := vector.DetachVectorData(vec) + area := vector.DetachVectorArea(vec) + if data.Capacity() != 0 { + b.bytesCache[id].buffers = append( + b.bytesCache[id].buffers, + data, + ) + } + if area.Capacity() != 0 { + b.bytesCache[id].buffers = append( + b.bytesCache[id].buffers, + area, + ) + } } bat.ReplaceVector(vec, nil, i) @@ -103,5 +119,9 @@ func (b *spoolBuffer) clean(mp *mpool.MPool) { mp.Free(b.bytesCache[i].bs[j]) } b.bytesCache[i].bs = nil + for j := range b.bytesCache[i].buffers { + b.bytesCache[i].buffers[j].Free(mp) + } + b.bytesCache[i].buffers = nil } } diff --git a/pkg/container/pSpool/copy.go b/pkg/container/pSpool/copy.go index e3797eeabe75c..38b327d7a3a9f 100644 --- a/pkg/container/pSpool/copy.go +++ b/pkg/container/pSpool/copy.go @@ -36,8 +36,10 @@ type cachedBatch struct { } type oneBatchMemoryCache struct { - // bytes to copy vector's data and area to. + // bs keeps the allocation-unaccounted production fast path unchanged. bs [][]byte + // buffers copy vector data and area while preserving allocation provenance. + buffers []vector.DetachedBuffer } func initCachedBatch(mp *mpool.MPool, capacity uint32) *cachedBatch { @@ -73,6 +75,12 @@ func (cb *cachedBatch) GetCopiedBatch( cacheID, dst = cb.buffer.getCacheID() dst.Recursive = src.Recursive dst.ShuffleIDX = src.ShuffleIDX + if sourceSelection := src.AllocationAccountSelection(); sourceSelection != dst.AllocationAccountSelection() { + if err = dst.SetAllocationAccount(sourceSelection); err != nil { + cb.CacheBatch(true, cacheID, dst) + return nil, false, 0, err + } + } if cap(dst.Vecs) >= len(src.Vecs) { dst.Vecs = dst.Vecs[:len(src.Vecs)] @@ -101,7 +109,19 @@ func (cb *cachedBatch) GetCopiedBatch( } typ := *vec.GetType() - dst.Vecs[i] = vector.NewOffHeapVecWithType(typ) + selection := vec.AllocationAccountSelection() + if selection == nil { + dst.Vecs[i] = vector.NewOffHeapVecWithType(typ) + } else { + dst.Vecs[i], err = vector.NewOffHeapVecWithTypeAndAllocation( + typ, + selection, + ) + if err != nil { + cb.CacheBatch(true, cacheID, dst) + return nil, false, 0, err + } + } if vec.IsConst() { if err = vector.GetConstSetFunction(typ, cb.mp)(dst.Vecs[i], vec, 0, vec.Length()); err != nil { @@ -110,8 +130,15 @@ func (cb *cachedBatch) GetCopiedBatch( } } else { - cb.buffer.bytesCache[cacheID].setSuitableDataAreaToVector( - len(vec.GetData()), len(vec.GetArea()), dst.Vecs[i]) + if err = cb.buffer.bytesCache[cacheID]. + setSuitableDataAreaToVector( + len(vec.GetData()), + len(vec.GetArea()), + dst.Vecs[i], + ); err != nil { + cb.CacheBatch(true, cacheID, dst) + return nil, false, 0, err + } dst.Vecs[i].Reset(typ) if err = vector.GetUnionAllFunction(typ, cb.mp)( dst.Vecs[i], @@ -144,25 +171,48 @@ func (cb *cachedBatch) GetCopiedBatch( // setSuitableDataAreaToVector get two long-enough bytes slices from the cache, and set them to the vector. // if not found, set the last one to the vector. func (mc *oneBatchMemoryCache) setSuitableDataAreaToVector( - dataSize, areaSize int, vec *vector.Vector) { + dataSize, areaSize int, + vec *vector.Vector, +) error { + if vec.AllocationAccountSelection() == nil { + mc.setSuitableLegacyDataAreaToVector(dataSize, areaSize, vec) + return nil + } + return mc.setSuitableAccountedDataAreaToVector( + dataSize, + areaSize, + vec, + ) +} + +func (mc *oneBatchMemoryCache) setSuitableAccountedDataAreaToVector( + dataSize, areaSize int, + vec *vector.Vector, +) error { // return directly once cache was empty. - if len(mc.bs) == 0 { - return + if len(mc.buffers) == 0 { + return nil } setDataFirst := dataSize >= areaSize first, second := dataSize, areaSize + firstKind := vector.DetachedDataBuffer + secondKind := vector.DetachedAreaBuffer if !setDataFirst { first, second = areaSize, dataSize + firstKind, secondKind = secondKind, firstKind } if first > 0 { suitIdx := -1 suitDifference := math.MaxInt - for i, bs := range mc.bs { - if difference := cap(bs) - first; difference > 0 { + for i := range mc.buffers { + if !mc.buffers[i].CanAttachTo(vec, firstKind) { + continue + } + if difference := mc.buffers[i].Capacity() - first; difference > 0 { if difference < suitDifference { suitIdx = i suitDifference = difference @@ -172,10 +222,9 @@ func (mc *oneBatchMemoryCache) setSuitableDataAreaToVector( if suitIdx != -1 { mem := mc.removeItemAndArrange(suitIdx) - if setDataFirst { - vector.SetVecData(vec, mem) - } else { - vector.SetVecArea(vec, mem) + if err := mem.AttachTo(vec, firstKind); err != nil { + mc.buffers = append(mc.buffers, mem) + return err } } } @@ -184,8 +233,11 @@ func (mc *oneBatchMemoryCache) setSuitableDataAreaToVector( suitIdx := -1 suitDifference := math.MaxInt - for i, bs := range mc.bs { - if difference := cap(bs) - second; difference > 0 { + for i := range mc.buffers { + if !mc.buffers[i].CanAttachTo(vec, secondKind) { + continue + } + if difference := mc.buffers[i].Capacity() - second; difference > 0 { if difference < suitDifference { suitIdx = i suitDifference = difference @@ -195,34 +247,136 @@ func (mc *oneBatchMemoryCache) setSuitableDataAreaToVector( if suitIdx != -1 { mem := mc.removeItemAndArrange(suitIdx) + if err := mem.AttachTo(vec, secondKind); err != nil { + mc.buffers = append(mc.buffers, mem) + return err + } + } + } + + if cap(vec.GetData()) == 0 && dataSize > 0 { + if idx := mc.lastAttachable( + vec, + vector.DetachedDataBuffer, + ); idx >= 0 { + mem := mc.removeItemAndArrange(idx) + if err := mem.AttachTo( + vec, + vector.DetachedDataBuffer, + ); err != nil { + mc.buffers = append(mc.buffers, mem) + return err + } + } + } + if cap(vec.GetArea()) == 0 && areaSize > 0 { + if idx := mc.lastAttachable( + vec, + vector.DetachedAreaBuffer, + ); idx >= 0 { + mem := mc.removeItemAndArrange(idx) + if err := mem.AttachTo( + vec, + vector.DetachedAreaBuffer, + ); err != nil { + mc.buffers = append(mc.buffers, mem) + return err + } + } + } + return nil +} + +func (mc *oneBatchMemoryCache) setSuitableLegacyDataAreaToVector( + dataSize, areaSize int, + vec *vector.Vector, +) { + if len(mc.bs) == 0 { + return + } + + setDataFirst := dataSize >= areaSize + first, second := dataSize, areaSize + if !setDataFirst { + first, second = areaSize, dataSize + } + + if first > 0 { + if idx := mc.bestLegacyBuffer(first); idx >= 0 { + mem := mc.removeLegacyBuffer(idx) if setDataFirst { - vector.SetVecArea(vec, mem) + vector.AttachLegacyVectorData(vec, mem) } else { - vector.SetVecData(vec, mem) + vector.AttachLegacyVectorArea(vec, mem) + } + } + } + if second > 0 { + if idx := mc.bestLegacyBuffer(second); idx >= 0 { + mem := mc.removeLegacyBuffer(idx) + if setDataFirst { + vector.AttachLegacyVectorArea(vec, mem) + } else { + vector.AttachLegacyVectorData(vec, mem) } } } - if len(mc.bs) > 0 && cap(vec.GetData()) == 0 && dataSize > 0 { - vector.SetVecData(vec, mc.bs[len(mc.bs)-1]) - mc.bs = mc.bs[:len(mc.bs)-1] + vector.AttachLegacyVectorData(vec, mc.removeLegacyBuffer(len(mc.bs)-1)) } if len(mc.bs) > 0 && cap(vec.GetArea()) == 0 && areaSize > 0 { - vector.SetVecArea(vec, mc.bs[len(mc.bs)-1]) - mc.bs = mc.bs[:len(mc.bs)-1] + vector.AttachLegacyVectorArea(vec, mc.removeLegacyBuffer(len(mc.bs)-1)) } } -// removeItemAndArrange return and remove the idx item of cache. -func (mc *oneBatchMemoryCache) removeItemAndArrange(idx int) []byte { - last := len(mc.bs) - 1 - dst := mc.bs[idx] +func (mc *oneBatchMemoryCache) bestLegacyBuffer(size int) int { + best := -1 + difference := math.MaxInt + for i, buffer := range mc.bs { + if current := cap(buffer) - size; current > 0 && + current < difference { + best = i + difference = current + } + } + return best +} +func (mc *oneBatchMemoryCache) removeLegacyBuffer(idx int) []byte { + last := len(mc.bs) - 1 + buffer := mc.bs[idx] if idx != last { mc.bs[idx] = mc.bs[last] - mc.bs = mc.bs[:last] } + mc.bs[last] = nil mc.bs = mc.bs[:last] + return buffer +} + +func (mc *oneBatchMemoryCache) lastAttachable( + vec *vector.Vector, + kind vector.DetachedBufferKind, +) int { + for i := len(mc.buffers) - 1; i >= 0; i-- { + if mc.buffers[i].CanAttachTo(vec, kind) { + return i + } + } + return -1 +} + +// removeItemAndArrange return and remove the idx item of cache. +func (mc *oneBatchMemoryCache) removeItemAndArrange( + idx int, +) vector.DetachedBuffer { + last := len(mc.buffers) - 1 + dst := mc.buffers[idx] + + if idx != last { + mc.buffers[idx] = mc.buffers[last] + } + mc.buffers[last] = vector.DetachedBuffer{} + mc.buffers = mc.buffers[:last] return dst } diff --git a/pkg/container/pSpool/copy_benchmark_test.go b/pkg/container/pSpool/copy_benchmark_test.go new file mode 100644 index 0000000000000..e7b4e82009b19 --- /dev/null +++ b/pkg/container/pSpool/copy_benchmark_test.go @@ -0,0 +1,58 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pSpool + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" +) + +func BenchmarkCachedBatchReuse(b *testing.B) { + mp := mpool.MustNewZero() + source := batch.NewWithSize(1) + source.Vecs[0] = vector.NewVec(types.T_int64.ToType()) + values := make([]int64, 8192) + for i := range values { + values[i] = int64(i) + } + if err := vector.AppendFixedList( + source.Vecs[0], + values, + nil, + mp, + ); err != nil { + b.Fatal(err) + } + source.SetRowCount(len(values)) + cache := initCachedBatch(mp, 1) + b.Cleanup(func() { + source.Clean(mp) + cache.free() + }) + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + copied, useCache, cacheID, err := cache.GetCopiedBatch(source) + if err != nil { + b.Fatal(err) + } + cache.CacheBatch(useCache, cacheID, copied) + } +} diff --git a/pkg/container/pSpool/sender_test.go b/pkg/container/pSpool/sender_test.go index cbb6ee8aa9d8a..99c766c376bee 100644 --- a/pkg/container/pSpool/sender_test.go +++ b/pkg/container/pSpool/sender_test.go @@ -99,6 +99,164 @@ func TestPipelineSpoolForceCleanupRetainsUntilReceiversDrained(t *testing.T) { require.Equal(t, int64(0), mp.CurrNB()) } +func TestCachedBatchPreservesAllocationProvenance(t *testing.T) { + mp := mpool.MustNewZero() + registry, err := mpool.NewAllocationAccountRegistry(2, 32) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection( + account, + 1, + 1, + 2, + ) + require.NoError(t, err) + otherAccount, err := registry.Open(1 << 20) + require.NoError(t, err) + otherSelection, err := vector.NewAllocationAccountSelection( + otherAccount, + 1, + 1, + 2, + ) + require.NoError(t, err) + newSource := func( + value string, + target *vector.AllocationAccountSelection, + ) *batch.Batch { + source := batch.NewOffHeapWithSize(1) + require.NoError(t, source.SetAllocationAccount(target)) + vec := vector.NewOffHeapVecWithType(types.T_varchar.ToType()) + source.SetVector(0, vec) + require.NoError(t, vector.AppendBytes( + vec, + []byte(value), + false, + mp, + )) + source.SetRowCount(1) + return source + } + firstSource := newSource("first cached allocation payload", selection) + secondSource := newSource("second", selection) + cache := initCachedBatch(mp, 1) + + first, useCache, cacheID, err := cache.GetCopiedBatch(firstSource) + require.NoError(t, err) + require.True(t, useCache) + require.Same(t, selection, first.AllocationAccountSelection()) + require.Same( + t, + selection, + first.Vecs[0].AllocationAccountSelection(), + ) + cache.CacheBatch(useCache, cacheID, first) + beforeReuse := account.Snapshot().Used + + second, useCache, cacheID, err := cache.GetCopiedBatch(secondSource) + require.NoError(t, err) + require.True(t, useCache) + require.Equal(t, beforeReuse, account.Snapshot().Used) + require.Same(t, selection, second.AllocationAccountSelection()) + require.Same( + t, + selection, + second.Vecs[0].AllocationAccountSelection(), + ) + cache.CacheBatch(useCache, cacheID, second) + + otherSource := newSource("other allocation account", otherSelection) + firstAccountBeforeOther := account.Snapshot().Used + otherBeforeCopy := otherAccount.Snapshot().Used + other, useCache, cacheID, err := cache.GetCopiedBatch(otherSource) + require.NoError(t, err) + require.Equal(t, firstAccountBeforeOther, account.Snapshot().Used) + require.Greater(t, otherAccount.Snapshot().Used, otherBeforeCopy) + require.Same( + t, + otherSelection, + other.Vecs[0].AllocationAccountSelection(), + ) + cache.CacheBatch(useCache, cacheID, other) + + firstSource.Clean(mp) + secondSource.Clean(mp) + otherSource.Clean(mp) + cache.free() + require.Zero(t, account.Snapshot().Used) + require.Zero(t, otherAccount.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + account.Seal() + otherAccount.Seal() + _, err = registry.Finalize(account) + require.NoError(t, err) + _, err = registry.Finalize(otherAccount) + require.NoError(t, err) +} + +func TestCachedBatchAllocationFailureReturnsCacheOwnership(t *testing.T) { + mp := mpool.MustNewZero() + registry, err := mpool.NewAllocationAccountRegistry(1, 4) + require.NoError(t, err) + account, err := registry.Open(64) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection( + account, + 1, + 1, + 2, + ) + require.NoError(t, err) + source := batch.NewOffHeapWithSize(1) + require.NoError(t, source.SetAllocationAccount(selection)) + source.SetVector(0, vector.NewOffHeapVecWithType(types.T_int64.ToType())) + require.NoError(t, vector.AppendFixedList( + source.Vecs[0], + []int64{1, 2, 3, 4, 5, 6, 7, 8}, + nil, + mp, + )) + source.SetRowCount(8) + require.Equal(t, uint64(64), account.Snapshot().Used) + cache := initCachedBatch(mp, 1) + + _, _, _, err = cache.GetCopiedBatch(source) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Len(t, cache.buffer.readyToUse, 1) + require.Equal(t, uint64(64), account.Snapshot().Used) + + source.Clean(mp) + cache.free() + require.Zero(t, account.Snapshot().Used) + account.Seal() + _, err = registry.Finalize(account) + require.NoError(t, err) +} + +func TestLegacyCacheNonLastSelectionRetainsOwnership(t *testing.T) { + mp := mpool.MustNewZero() + cache := oneBatchMemoryCache{} + for _, size := range []int{64, 128, 256} { + buffer, err := mp.Alloc(size, true) + require.NoError(t, err) + cache.bs = append(cache.bs, buffer) + } + vec := vector.NewOffHeapVecWithType(types.T_int8.ToType()) + require.NoError(t, cache.setSuitableDataAreaToVector( + 100, + 0, + vec, + )) + require.Len(t, cache.bs, 2) + + vec.Free(mp) + for i := range cache.bs { + mp.Free(cache.bs[i]) + } + require.Zero(t, mp.CurrNB()) +} + func TestPipelineSpoolForceCleanupAfterTerminalSignalDoesNotNeedNilEndMessage(t *testing.T) { mp := mpool.MustNewZeroNoFixed() t.Cleanup(func() { diff --git a/pkg/container/vector/allocation_account_test.go b/pkg/container/vector/allocation_account_test.go index 7accd35d508d3..7fb139493e2a6 100644 --- a/pkg/container/vector/allocation_account_test.go +++ b/pkg/container/vector/allocation_account_test.go @@ -487,6 +487,119 @@ func TestVectorAllocationAccountErrorsAreTyped(t *testing.T) { finalizeTestVectorAllocationAccount(t, state) } +func TestDetachedBufferPreservesAllocationProvenance(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 8) + mp := mpool.MustNewZero() + source := newAccountedTestVector( + t, + types.T_varchar.ToType(), + state.selection, + ) + require.NoError(t, AppendBytes( + source, + []byte("detached allocation payload that uses the vector area"), + false, + mp, + )) + used := state.account.Snapshot().Used + require.Positive(t, used) + require.Panics(t, func() { + DetachLegacyVectorData(source) + }) + require.Panics(t, func() { + DetachLegacyVectorArea(source) + }) + + data := DetachVectorData(source) + area := DetachVectorArea(source) + require.Positive(t, data.Capacity()) + require.Positive(t, area.Capacity()) + source.Free(mp) + require.Equal(t, used, state.account.Snapshot().Used) + + destination := newAccountedTestVector( + t, + types.T_varchar.ToType(), + state.selection, + ) + require.True(t, data.CanAttachTo(destination, DetachedDataBuffer)) + require.False(t, data.CanAttachTo(destination, DetachedAreaBuffer)) + require.NoError(t, data.AttachTo(destination, DetachedDataBuffer)) + require.NoError(t, area.AttachTo(destination, DetachedAreaBuffer)) + require.Zero(t, data.Capacity()) + require.Zero(t, area.Capacity()) + require.Equal(t, used, state.account.Snapshot().Used) + + destination.Free(mp) + require.Zero(t, state.account.Snapshot().Used) + data.Free(mp) + area.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestSetTypeAndFixDataAllocationFailureIsAtomic(t *testing.T) { + state := newTestVectorAllocationAccount(t, 512, 4) + mp := mpool.MustNewZero() + vec := newAccountedTestVector( + t, + types.T_date.ToType(), + state.selection, + ) + require.NoError(t, vec.PreExtend(128, mp)) + vec.SetLength(128) + used := state.account.Snapshot().Used + require.Equal(t, uint64(512), used) + + err := vec.SetTypeAndFixData(types.T_datetime.ToType(), mp) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Equal(t, types.T_date, vec.GetType().Oid) + require.Equal(t, 128, vec.Length()) + require.Equal(t, used, state.account.Snapshot().Used) + + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestDetachedLegacyBufferAndTypeChange(t *testing.T) { + mp := mpool.MustNewZero() + source := NewOffHeapVecWithType(types.T_varchar.ToType()) + require.NoError(t, AppendBytes( + source, + []byte("legacy detached allocation payload"), + false, + mp, + )) + data := DetachLegacyVectorData(source) + area := DetachLegacyVectorArea(source) + source.Free(mp) + + destination := NewOffHeapVecWithType(types.T_varchar.ToType()) + AttachLegacyVectorData(destination, data) + AttachLegacyVectorArea(destination, area) + destination.Free(mp) + require.Zero(t, mp.CurrNB()) + + fixed := NewOffHeapVecWithType(types.T_date.ToType()) + require.NoError(t, AppendFixed( + fixed, + types.Date(1), + false, + mp, + )) + require.NoError(t, fixed.SetTypeAndFixData( + types.T_datetime.ToType(), + mp, + )) + require.Equal(t, types.T_datetime, fixed.GetType().Oid) + require.Equal(t, 1, fixed.Length()) + require.Error(t, fixed.SetTypeAndFixData( + types.T_varchar.ToType(), + mp, + )) + fixed.Free(mp) + require.Zero(t, mp.CurrNB()) +} + func TestVectorAllocationAccountHelperBoundaries(t *testing.T) { state := newTestVectorAllocationAccount(t, 1<<20, 8) mp := mpool.MustNewZero() diff --git a/pkg/container/vector/pSpoolTools.go b/pkg/container/vector/pSpoolTools.go index 8548c4e7d8916..c62c482f337a9 100644 --- a/pkg/container/vector/pSpoolTools.go +++ b/pkg/container/vector/pSpoolTools.go @@ -14,27 +14,152 @@ package vector -// SetVecData is dangerous and should be used with caution. -func SetVecData(v *Vector, data []byte) { - data = data[:cap(data)] - v.data = data +import ( + "fmt" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" +) + +// DetachedBuffer transfers one owned Vector backing allocation through the +// pipeline spool without losing its immutable allocation provenance. +// A non-empty value must be attached or freed exactly once. +type DetachedBuffer struct { + data []byte + selection *AllocationAccountSelection + kind DetachedBufferKind +} + +type DetachedBufferKind uint8 + +const ( + DetachedDataBuffer DetachedBufferKind = iota + DetachedAreaBuffer +) + +// DetachLegacyVectorData is the allocation-unaccounted spool fast path. The +// explicit guard prevents raw ownership transfer from dropping provenance. +func DetachLegacyVectorData(v *Vector) []byte { + if v.allocationAccount != nil { + panic("cannot detach accounted vector data without provenance") + } + data := v.data + v.data = nil + return data +} + +func DetachLegacyVectorArea(v *Vector) []byte { + if v.allocationAccount != nil { + panic("cannot detach accounted vector area without provenance") + } + area := v.area + v.area = nil + return area +} + +func AttachLegacyVectorData(v *Vector, data []byte) { + if v.allocationAccount != nil || cap(v.data) != 0 { + panic("cannot attach legacy vector data") + } + v.data = data[:cap(data)] } -// SetVecArea is dangerous and should be used with caution. -func SetVecArea(v *Vector, area []byte) { +func AttachLegacyVectorArea(v *Vector, area []byte) { + if v.allocationAccount != nil || cap(v.area) != 0 { + panic("cannot attach legacy vector area") + } v.area = area } -// GetAndClearVecData is a dangerous function that may cause data leakage. -func GetAndClearVecData(v *Vector) []byte { - s := v.data +func DetachVectorData(v *Vector) DetachedBuffer { + if v == nil { + return DetachedBuffer{} + } + buffer := DetachedBuffer{ + data: v.data, + selection: v.allocationAccount, + } v.data = nil - return s + return buffer } -// GetAndClearVecArea is a dangerous function that may cause data leakage. -func GetAndClearVecArea(v *Vector) []byte { - s := v.area +func DetachVectorArea(v *Vector) DetachedBuffer { + if v == nil { + return DetachedBuffer{} + } + buffer := DetachedBuffer{ + data: v.area, + selection: v.allocationAccount, + kind: DetachedAreaBuffer, + } v.area = nil - return s + return buffer +} + +func (b *DetachedBuffer) Capacity() int { + if b == nil { + return 0 + } + return cap(b.data) +} + +// CanAttachTo preserves data/area site provenance for accounted allocations. +// Legacy buffers have no site identity and retain the historical ability to +// serve either backing. +func (b *DetachedBuffer) CanAttachTo( + v *Vector, + kind DetachedBufferKind, +) bool { + if b == nil || v == nil || cap(b.data) == 0 || + b.selection != v.allocationAccount || + kind > DetachedAreaBuffer { + return false + } + return b.selection == nil || b.kind == kind +} + +func (b *DetachedBuffer) AttachTo( + v *Vector, + kind DetachedBufferKind, +) error { + if !b.CanAttachTo(v, kind) { + return fmt.Errorf( + "%w: detached vector buffer provenance mismatch", + mpool.ErrAllocationAccountInvalid, + ) + } + if kind == DetachedAreaBuffer { + if cap(v.area) != 0 { + return fmt.Errorf( + "%w: vector area already has backing storage", + mpool.ErrAllocationAccountInvalid, + ) + } + v.area = b.data + } else { + if cap(v.data) != 0 { + return fmt.Errorf( + "%w: vector data already has backing storage", + mpool.ErrAllocationAccountInvalid, + ) + } + v.data = b.data[:cap(b.data)] + } + b.clear() + return nil +} + +func (b *DetachedBuffer) Free(mp *mpool.MPool) { + if b == nil { + return + } + if cap(b.data) != 0 { + mp.Free(b.data) + } + b.clear() +} + +func (b *DetachedBuffer) clear() { + b.data = nil + b.selection = nil + b.kind = DetachedDataBuffer } diff --git a/pkg/container/vector/vector.go b/pkg/container/vector/vector.go index 66e5775238141..eedf1c04314ea 100644 --- a/pkg/container/vector/vector.go +++ b/pkg/container/vector/vector.go @@ -212,25 +212,35 @@ func (v *Vector) SetType(typ types.Type) { v.typ = typ } -// Bug #23240 -// Neither this function, nor the SetType function are good -// Maybe we should just disallow. -func (v *Vector) SetTypeAndFixData(typ types.Type, mp *mpool.MPool) { +// SetTypeAndFixData changes a fixed-width result type and grows its owned data +// before publishing the new type. A failed growth leaves the original vector +// type, length, and backing allocation intact. +func (v *Vector) SetTypeAndFixData( + typ types.Type, + mp *mpool.MPool, +) error { if v.typ.IsVarlen() && typ.IsVarlen() { v.typ = typ - return + return nil } if v.typ.IsVarlen() || typ.IsVarlen() { - // this is a weird thing to do, we should not allow it. - panic("SetTypeAndFixData is not allowed to change from/to varlen type") + return moerr.NewInvalidInputNoCtx( + "SetTypeAndFixData cannot change from or to a varlen type", + ) } + oldType := v.typ v.typ = typ oldLength := v.length v.length = 0 - extend(v, oldLength, mp) + if err := extend(v, oldLength, mp); err != nil { + v.typ = oldType + v.length = oldLength + return err + } v.length = oldLength + return nil } func (v *Vector) SetOffHeap(offHeap bool) { diff --git a/pkg/sql/plan/function/func_binary.go b/pkg/sql/plan/function/func_binary.go index 4fbd6e60ae8f0..eba30b6207385 100644 --- a/pkg/sql/plan/function/func_binary.go +++ b/pkg/sql/plan/function/func_binary.go @@ -2079,7 +2079,12 @@ func TimestampAddDate(ivecs []*vector.Vector, result vector.FunctionResultWrappe if resultType == types.T_date { // Result wrapper is DATE, but we need to return DATETIME // Convert to DATETIME type - vec.SetTypeAndFixData(types.New(types.T_datetime, 0, scale), proc.GetMPool()) + if err := vec.SetTypeAndFixData( + types.New(types.T_datetime, 0, scale), + proc.GetMPool(), + ); err != nil { + return err + } rss := vector.MustFixedColNoTypeCheck[types.Datetime](vec) rsNull := vec.GetNulls() @@ -2163,7 +2168,12 @@ func TimestampAddDate(ivecs []*vector.Vector, result vector.FunctionResultWrappe } else { // Result wrapper is DATETIME (backward compatibility) // Use SetType to change vector type to DATE - vec.SetTypeAndFixData(types.New(types.T_date, 0, 0), proc.GetMPool()) + if err := vec.SetTypeAndFixData( + types.New(types.T_date, 0, 0), + proc.GetMPool(), + ); err != nil { + return err + } rss := vector.MustFixedColNoTypeCheck[types.Date](vec) rsNull := vec.GetNulls() @@ -2221,7 +2231,12 @@ func TimestampAddDate(ivecs []*vector.Vector, result vector.FunctionResultWrappe scale := maxScale if resultType == types.T_date { // Result wrapper is DATE, but we need to return DATETIME - vec.SetTypeAndFixData(types.New(types.T_datetime, 0, scale), proc.GetMPool()) + if err := vec.SetTypeAndFixData( + types.New(types.T_datetime, 0, scale), + proc.GetMPool(), + ); err != nil { + return err + } rss := vector.MustFixedColNoTypeCheck[types.Datetime](vec) rsNull := vec.GetNulls() @@ -2318,7 +2333,12 @@ func TimestampAddDate(ivecs []*vector.Vector, result vector.FunctionResultWrappe } } else { // Result wrapper is DATETIME, but all units are date units, so return DATE - vec.SetTypeAndFixData(types.New(types.T_date, 0, 0), proc.GetMPool()) + if err := vec.SetTypeAndFixData( + types.New(types.T_date, 0, 0), + proc.GetMPool(), + ); err != nil { + return err + } rss := vector.MustFixedColNoTypeCheck[types.Date](vec) rsNull := vec.GetNulls() From 6c1dcb8bc9bbfef159928830058d464c841e982b Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 16:25:23 +0800 Subject: [PATCH 09/61] feat: account expression bitmap and conversion scratch --- pkg/common/bitmap/bitmap.go | 256 ++++++++++-- pkg/common/bitmap/bitmap_test.go | 72 +++- pkg/common/bitmap/cbitmap.go | 4 +- pkg/common/bitmap/types.go | 7 +- pkg/common/mpool/accounted_buffer.go | 13 + pkg/common/mpool/accounted_buffer_test.go | 22 + pkg/container/pSpool/buffer.go | 4 + pkg/container/pSpool/copy.go | 15 + pkg/container/pSpool/sender_test.go | 15 +- pkg/container/vector/allocation_account.go | 200 ++++++++- .../vector/allocation_account_test.go | 259 +++++++++++- pkg/container/vector/functionTools.go | 390 ++++++++++++++---- .../vector/function_result_allocation_test.go | 178 +++++++- pkg/container/vector/tools.go | 3 + pkg/container/vector/vector.go | 118 +++++- pkg/sql/colexec/aggexec/maxby.go | 3 +- pkg/sql/colexec/evalExpression.go | 6 +- pkg/sql/colexec/eval_expression_allocation.go | 48 ++- .../colexec/spillutil/allocation_account.go | 12 +- .../spillutil/allocation_account_test.go | 2 +- pkg/sql/plan/function/baseTemplate.go | 128 ++++-- pkg/sql/plan/function/func_binary.go | 60 ++- pkg/sql/plan/function/func_compare.go | 10 +- 23 files changed, 1610 insertions(+), 215 deletions(-) diff --git a/pkg/common/bitmap/bitmap.go b/pkg/common/bitmap/bitmap.go index 7f8a3150ae09e..06603cdac077b 100644 --- a/pkg/common/bitmap/bitmap.go +++ b/pkg/common/bitmap/bitmap.go @@ -19,6 +19,7 @@ import ( "encoding" "fmt" "io" + "math" "math/bits" "unsafe" @@ -33,6 +34,8 @@ import ( type bitmask = uint64 +const MarshalHeaderSize = 24 + /* * Array giving the position of the right-most set bit for each possible * byte value. count the right-most position as the 0th bit, and the @@ -58,16 +61,62 @@ var rightmost_one_pos_8 = [256]uint8{ 4, 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, } +func encodeTaggedLen(length int64, external bool) int64 { + if length < 0 { + panic("negative bitmap length") + } + if external { + return ^length + } + return length +} + +func (n *Bitmap) logicalLen() int64 { + if n.taggedLen < 0 { + return ^n.taggedLen + } + return n.taggedLen +} + +func (n *Bitmap) setLogicalLen(length int64) { + n.taggedLen = encodeTaggedLen(length, n.HasExternalStorage()) +} + func (n *Bitmap) InitWith(m *Bitmap) { - n.len = m.len + if n == m { + return + } + n.setLogicalLen(m.logicalLen()) n.count = m.count + if n.HasExternalStorage() { + if len(m.data) > cap(n.data) { + panic("bitmap external storage capacity exceeded") + } + previousLength := len(n.data) + storage := n.data[:cap(n.data)] + clear(storage[len(m.data):max(previousLength, len(m.data))]) + n.data = storage[:len(m.data)] + copy(n.data, m.data) + return + } n.data = append([]uint64(nil), m.data...) } -func (n *Bitmap) InitWithSize(len int64) { - n.len = len +func (n *Bitmap) InitWithSize(length int64) { + n.setLogicalLen(length) n.count = 0 - n.data = make([]uint64, (len+63)/64) + words := int((length + 63) / 64) + if n.HasExternalStorage() { + if words > cap(n.data) { + panic("bitmap external storage capacity exceeded") + } + previousLength := len(n.data) + storage := n.data[:cap(n.data)] + clear(storage[:max(previousLength, words)]) + n.data = storage[:words] + return + } + n.data = make([]uint64, words) } func (n *Bitmap) Clone() *Bitmap { @@ -114,7 +163,7 @@ func (itr *BitmapIterator) hasNext(i uint64) (uint64, bool) { // if the uint64 is not 0, then calculate the rightest_one position in a word, add up prev result and return. // when there is 1 in Bitmap, return true, otherwise Bitmap is empty and return false. // either case loop over words not bits - nwords := (itr.bm.len + 63) / 64 + nwords := (itr.bm.logicalLen() + 63) / 64 current_word := i >> 6 mask := (^(bitmask)(0)) << (i & 0x3F) // ignore bits check before var result uint64 @@ -159,14 +208,69 @@ func (itr *BitmapIterator) Next() uint64 { // Reset set n.data to nil func (n *Bitmap) Reset() { - n.len = 0 + n.setLogicalLen(0) n.count = 0 + if n.HasExternalStorage() { + clear(n.data) + storage := n.data[:cap(n.data)] + n.data = storage[:0] + return + } n.data = nil } +// InstallExternalStorage replaces the bitmap backing with caller-owned +// storage while preserving the logical bitmap. The caller remains responsible +// for releasing the returned previous external storage, if any. +func (n *Bitmap) InstallExternalStorage(storage []uint64) []uint64 { + required := len(n.data) + if required > cap(storage) { + panic("bitmap external storage capacity exceeded") + } + var previous []uint64 + if n.HasExternalStorage() && cap(n.data) > 0 { + previous = n.data[:cap(n.data)] + } + target := storage[:cap(storage)] + if len(target) > required { + clear(target[required:]) + } + copy(target[:required], n.data) + n.data = target[:required] + n.taggedLen = encodeTaggedLen(n.logicalLen(), true) + return previous +} + +// ReleaseExternalStorage detaches caller-owned storage and clears the bitmap. +// It returns nil for a legacy Go-owned bitmap. +func (n *Bitmap) ReleaseExternalStorage() []uint64 { + if !n.HasExternalStorage() { + return nil + } + var storage []uint64 + if cap(n.data) > 0 { + storage = n.data[:cap(n.data)] + } + n.count = 0 + n.taggedLen = 0 + n.data = nil + return storage +} + +func (n *Bitmap) ExternalStorageCapacity() int { + if n == nil || !n.HasExternalStorage() { + return 0 + } + return cap(n.data) +} + +func (n *Bitmap) HasExternalStorage() bool { + return n != nil && n.taggedLen < 0 +} + // Len returns the number of bits in the Bitmap. func (n *Bitmap) Len() int64 { - return n.len + return n.logicalLen() } // Size return number of bytes in n.data @@ -211,7 +315,7 @@ func (n *Bitmap) AddMany(rows []uint64) { } func (n *Bitmap) Remove(row uint64) { - if row >= uint64(n.len) { + if row >= uint64(n.logicalLen()) { return } if n.data[row>>6]&(1<<(row&0x3F)) != 0 { @@ -222,7 +326,7 @@ func (n *Bitmap) Remove(row uint64) { // Contains returns true if the row is contained in the Bitmap func (n *Bitmap) Contains(row uint64) bool { - if row >= uint64(n.len) { + if row >= uint64(n.logicalLen()) { return false } idx := row >> 6 @@ -256,8 +360,8 @@ func (n *Bitmap) AddRange(start, end uint64) { } func (n *Bitmap) RemoveRange(start, end uint64) { - if end > uint64(n.len) { - end = uint64(n.len) + if end > uint64(n.logicalLen()) { + end = uint64(n.logicalLen()) } if start >= end { return @@ -298,7 +402,7 @@ func (n *Bitmap) IsSame(b *Bitmap) bool { func (n *Bitmap) Or(b *Bitmap) { n.TryExpand(b) - size := (int(b.len) + 63) / 64 + size := (int(b.logicalLen()) + 63) / 64 for i := range size { cnt := bits.OnesCount64(n.data[i]) n.data[i] |= b.data[i] @@ -309,7 +413,7 @@ func (n *Bitmap) Or(b *Bitmap) { func (n *Bitmap) And(b *Bitmap) { n.TryExpand(b) n.count = 0 - size := (int(b.len) + 63) / 64 + size := (int(b.logicalLen()) + 63) / 64 for i := range size { n.data[i] &= b.data[i] n.count += int64(bits.OnesCount64(n.data[i])) @@ -320,7 +424,7 @@ func (n *Bitmap) And(b *Bitmap) { } func (n *Bitmap) Negate() { - nBlock, nTail := int(n.len)/64, int(n.len)%64 + nBlock, nTail := int(n.logicalLen())/64, int(n.logicalLen())%64 n.count = 0 for i := range nBlock { n.data[i] = ^n.data[i] @@ -334,16 +438,19 @@ func (n *Bitmap) Negate() { } func (n *Bitmap) TryExpand(m *Bitmap) { - n.TryExpandWithSize(int(m.len)) + n.TryExpandWithSize(int(m.logicalLen())) } func (n *Bitmap) TryExpandWithSize(size int) { - if int(n.len) >= size { + if int(n.logicalLen()) >= size { return } newCap := (size + 63) / 64 - n.len = int64(size) + n.setLogicalLen(int64(size)) if newCap > cap(n.data) { + if n.HasExternalStorage() { + panic("bitmap external storage capacity exceeded") + } data := make([]uint64, newCap) copy(data, n.data) n.data = data @@ -356,7 +463,7 @@ func (n *Bitmap) TryExpandWithSize(size int) { func (n *Bitmap) Filter(sels []int64) *Bitmap { var b Bitmap - b.InitWithSize(n.len) + b.InitWithSize(n.logicalLen()) for i, sel := range sels { if n.Contains(uint64(sel)) { b.Add(uint64(i)) @@ -403,7 +510,60 @@ func (n *Bitmap) MarshalSize() int { if n == nil { return 0 } - return 24 + len(n.data)*8 + return MarshalHeaderSize + len(n.data)*8 +} + +// DecodeMarshalHeader validates the fixed bitmap wire header. +func DecodeMarshalHeader(data []byte) ( + count int64, + bitLength int64, + dataSize int, + err error, +) { + if len(data) < MarshalHeaderSize { + return 0, 0, 0, io.ErrUnexpectedEOF + } + count = types.DecodeInt64(data[:8]) + rawBitLength := types.DecodeUint64(data[8:16]) + rawDataSize := types.DecodeUint64(data[16:24]) + if count < 0 || + rawBitLength > math.MaxInt64 || + rawDataSize > math.MaxInt || + rawDataSize%8 != 0 { + return 0, 0, 0, fmt.Errorf("invalid bitmap wire header") + } + bitLength = int64(rawBitLength) + dataSize = int(rawDataSize) + if count > bitLength || + uint64(dataSize/8) != (rawBitLength+63)/64 { + return 0, 0, 0, fmt.Errorf("invalid bitmap wire header") + } + return count, bitLength, dataSize, nil +} + +// PrepareExternalUnmarshal publishes a validated bitmap header into existing +// caller-owned storage and returns the payload bytes to fill. +func (n *Bitmap) PrepareExternalUnmarshal( + header []byte, + totalSize int, +) ([]byte, error) { + if !n.HasExternalStorage() { + return nil, fmt.Errorf("bitmap does not use external storage") + } + count, bitLength, dataSize, err := DecodeMarshalHeader(header) + if err != nil { + return nil, err + } + if totalSize != MarshalHeaderSize+dataSize || + dataSize/8 > cap(n.data) { + return nil, fmt.Errorf("invalid bitmap external storage capacity") + } + storage := n.data[:cap(n.data)] + clear(storage) + n.data = storage[:dataSize/8] + n.count = count + n.setLogicalLen(bitLength) + return types.EncodeSlice(n.data), nil } func (n *Bitmap) MarshalTo(w io.Writer) error { @@ -413,7 +573,7 @@ func (n *Bitmap) MarshalTo(w io.Writer) error { if w == nil { return io.ErrClosedPipe } - bitLength := uint64(n.len) + bitLength := uint64(n.logicalLen()) dataLength := uint64(len(n.data) * 8) for _, value := range [][]byte{ types.EncodeInt64(&n.count), @@ -436,7 +596,7 @@ func (n *Bitmap) MarshalTo(w io.Writer) error { func (n *Bitmap) MarshalV1() []byte { var buf bytes.Buffer empty := int32(0) - u1 := uint64(n.len) + u1 := uint64(n.logicalLen()) u2 := uint64(len(n.data) * 8) buf.Write(types.EncodeInt32(&empty)) buf.Write(types.EncodeUint64(&u1)) @@ -448,21 +608,41 @@ func (n *Bitmap) MarshalV1() []byte { func (n *Bitmap) Unmarshal(data []byte) { n.count = types.DecodeInt64(data[:8]) data = data[8:] - n.len = int64(types.DecodeUint64(data[:8])) + n.setLogicalLen(int64(types.DecodeUint64(data[:8]))) data = data[8:] size := int(types.DecodeUint64(data[:8])) data = data[8:] if size == 0 { - n.data = nil + if n.HasExternalStorage() { + storage := n.data[:cap(n.data)] + clear(storage) + n.data = storage[:0] + } else { + n.data = nil + } } else { + if n.HasExternalStorage() { + words := size / 8 + if size%8 != 0 || words > cap(n.data) { + panic("bitmap external storage capacity exceeded") + } + storage := n.data[:cap(n.data)] + clear(storage) + n.data = storage[:words] + copy(n.data, types.DecodeSlice[uint64](data[:size])) + return + } n.data = types.DecodeSlice[uint64](data[:size]) } } func (n *Bitmap) UnmarshalNoCopy(data []byte) { + if n.HasExternalStorage() { + panic("cannot install alias into bitmap external storage") + } n.count = types.DecodeInt64(data[:8]) data = data[8:] - n.len = int64(types.DecodeUint64(data[:8])) + n.setLogicalLen(int64(types.DecodeUint64(data[:8]))) data = data[8:] size := int(types.DecodeUint64(data[:8])) data = data[8:] @@ -476,14 +656,31 @@ func (n *Bitmap) UnmarshalNoCopy(data []byte) { // UnmarshalV1 in version 1, Bitmap.emptyFlag is type int32, now we use Bitmap.count replace it func (n *Bitmap) UnmarshalV1(data []byte) { data = data[4:] - n.len = int64(types.DecodeUint64(data[:8])) + n.setLogicalLen(int64(types.DecodeUint64(data[:8]))) data = data[8:] size := int(types.DecodeUint64(data[:8])) data = data[8:] if size == 0 { - n.data = nil + if n.HasExternalStorage() { + storage := n.data[:cap(n.data)] + clear(storage) + n.data = storage[:0] + } else { + n.data = nil + } } else { - n.data = types.DecodeSlice[uint64](data[:size]) + if n.HasExternalStorage() { + words := size / 8 + if size%8 != 0 || words > cap(n.data) { + panic("bitmap external storage capacity exceeded") + } + storage := n.data[:cap(n.data)] + clear(storage) + n.data = storage[:words] + copy(n.data, types.DecodeSlice[uint64](data[:size])) + } else { + n.data = types.DecodeSlice[uint64](data[:size]) + } } n.count = 0 for i := 0; i < len(n.data); i++ { @@ -492,8 +689,11 @@ func (n *Bitmap) UnmarshalV1(data []byte) { } func (n *Bitmap) UnmarshalNoCopyV1(data []byte) { + if n.HasExternalStorage() { + panic("cannot install alias into bitmap external storage") + } data = data[4:] - n.len = int64(types.DecodeUint64(data[:8])) + n.setLogicalLen(int64(types.DecodeUint64(data[:8]))) data = data[8:] size := int(types.DecodeUint64(data[:8])) data = data[8:] diff --git a/pkg/common/bitmap/bitmap_test.go b/pkg/common/bitmap/bitmap_test.go index c187be6bba6c2..39a55ac82366d 100644 --- a/pkg/common/bitmap/bitmap_test.go +++ b/pkg/common/bitmap/bitmap_test.go @@ -19,6 +19,7 @@ import ( "fmt" "io" "testing" + "unsafe" "github.com/stretchr/testify/require" ) @@ -173,8 +174,8 @@ func TestBitmap_Compatibility(t *testing.T) { np.AddMany(rows) npV1 := &Bitmap{ - len: np.len, - data: np.data, + taggedLen: np.taggedLen, + data: np.data, } data := npV1.MarshalV1() @@ -232,3 +233,70 @@ func TestBitmap_And2(t *testing.T) { np.And(np2) require.Equal(t, 100, np.Count()) } + +func TestBitmapExternalStorageLifecycle(t *testing.T) { + require.Equal( + t, + 2*unsafe.Sizeof(int64(0))+unsafe.Sizeof([]uint64(nil)), + unsafe.Sizeof(Bitmap{}), + ) + storage := []uint64{^uint64(0), ^uint64(0)} + var value Bitmap + require.Nil(t, value.InstallExternalStorage(storage)) + require.True(t, value.HasExternalStorage()) + require.Equal(t, 2, value.ExternalStorageCapacity()) + + value.InitWithSize(65) + require.Equal(t, []uint64{0, 0}, storage) + value.Add(64) + require.True(t, value.Contains(64)) + + value.Reset() + require.True(t, value.HasExternalStorage()) + require.Equal(t, 2, value.ExternalStorageCapacity()) + require.Equal(t, []uint64{0, 0}, storage) + value.TryExpandWithSize(128) + value.Add(127) + require.True(t, value.Contains(127)) + require.Panics(t, func() { + value.TryExpandWithSize(129) + }) + + released := value.ReleaseExternalStorage() + require.Equal(t, storage, released) + require.False(t, value.HasExternalStorage()) + value.TryExpandWithSize(129) + value.Add(128) + require.True(t, value.Contains(128)) +} + +func TestBitmapExternalStorageUnmarshal(t *testing.T) { + source := newBm(128) + source.Add(1) + source.Add(127) + encoded := source.Marshal() + + storage := make([]uint64, 2) + var copied Bitmap + copied.InstallExternalStorage(storage) + copied.Unmarshal(encoded) + require.True(t, copied.IsSame(source)) + require.True(t, copied.HasExternalStorage()) + + var streamed Bitmap + streamStorage := make([]uint64, 2) + streamed.InstallExternalStorage(streamStorage) + payload, err := streamed.PrepareExternalUnmarshal( + encoded[:MarshalHeaderSize], + len(encoded), + ) + require.NoError(t, err) + copy(payload, encoded[MarshalHeaderSize:]) + require.True(t, streamed.IsSame(source)) + + _, err = streamed.PrepareExternalUnmarshal( + encoded[:MarshalHeaderSize], + len(encoded)-1, + ) + require.Error(t, err) +} diff --git a/pkg/common/bitmap/cbitmap.go b/pkg/common/bitmap/cbitmap.go index 1edc7ffb8d595..2576718f2a2fb 100644 --- a/pkg/common/bitmap/cbitmap.go +++ b/pkg/common/bitmap/cbitmap.go @@ -27,7 +27,7 @@ func (n *Bitmap) cPtr() *C.uint64_t { return (*C.uint64_t)(unsafe.Pointer(&n.data[0])) } func (n *Bitmap) cLen() C.uint64_t { - return C.uint64_t(n.len) + return C.uint64_t(n.logicalLen()) } func (n *Bitmap) C_IsEmpty() bool { @@ -53,5 +53,5 @@ func (n *Bitmap) RawPtrLen() (uintptr, uintptr) { if n == nil || len(n.data) == 0 { return 0, 0 } - return uintptr(unsafe.Pointer(&n.data[0])), uintptr(n.len) + return uintptr(unsafe.Pointer(&n.data[0])), uintptr(n.logicalLen()) } diff --git a/pkg/common/bitmap/types.go b/pkg/common/bitmap/types.go index bf1141d9fc19d..883625d144f86 100644 --- a/pkg/common/bitmap/types.go +++ b/pkg/common/bitmap/types.go @@ -22,8 +22,11 @@ type Iterator interface { type Bitmap struct { count int64 //in version 1, we use emptyFlag with type int32 to indicate whether it is empty - len int64 - data []uint64 + // taggedLen stores the logical length directly for legacy backing and its + // bitwise complement for caller-owned backing. This keeps Bitmap's legacy + // footprint unchanged while making backing ownership explicit. + taggedLen int64 + data []uint64 } type BitmapIterator struct { diff --git a/pkg/common/mpool/accounted_buffer.go b/pkg/common/mpool/accounted_buffer.go index af49b6504bedf..4c688117ed9cd 100644 --- a/pkg/common/mpool/accounted_buffer.go +++ b/pkg/common/mpool/accounted_buffer.go @@ -114,6 +114,19 @@ func (b *AccountedBuffer) EnsureCapacity(required int) error { return nil } +// Resize changes the logical length after admitting any required retained +// capacity. Existing bytes are preserved. +func (b *AccountedBuffer) Resize(length int) error { + if b == nil || length < 0 { + return ErrAllocationAccountInvalid + } + if err := b.EnsureCapacity(length); err != nil { + return err + } + b.data = b.data[:length] + return nil +} + func (b *AccountedBuffer) Write(value []byte) (int, error) { if b == nil { return 0, ErrAllocationAccountInvalid diff --git a/pkg/common/mpool/accounted_buffer_test.go b/pkg/common/mpool/accounted_buffer_test.go index 5fa9d896fdb6a..855cf9e8409a2 100644 --- a/pkg/common/mpool/accounted_buffer_test.go +++ b/pkg/common/mpool/accounted_buffer_test.go @@ -36,6 +36,13 @@ func TestAccountedBufferLifecycle(t *testing.T) { firstCapacity := buffer.Cap() require.GreaterOrEqual(t, firstCapacity, 32) require.Equal(t, uint64(firstCapacity), account.Snapshot().Used) + require.NoError(t, buffer.Resize(16)) + require.Equal(t, 16, buffer.Len()) + require.Equal(t, firstCapacity, buffer.Cap()) + copy(buffer.Bytes(), "retained") + require.NoError(t, buffer.Resize(8)) + require.Equal(t, "retained", string(buffer.Bytes())) + buffer.Reset() _, err = buffer.WriteString("accounted") require.NoError(t, err) require.Equal(t, "accounted", string(buffer.Bytes())) @@ -98,10 +105,25 @@ func TestAccountedBufferConfiguration(t *testing.T) { require.Nil(t, buffer.Bytes()) require.Zero(t, buffer.Len()) require.Zero(t, buffer.Cap()) + require.ErrorIs(t, buffer.Resize(0), ErrAllocationAccountInvalid) _, err = buffer.Write([]byte("x")) require.ErrorIs(t, err, ErrAllocationAccountInvalid) buffer.Reset() buffer.Free() + + registry, account := newTestAllocationAccount(t, 1<<20, 1) + mp := MustNew("accounted-buffer-invalid-resize") + defer DeleteMPool(mp) + buffer, err = NewAccountedBuffer( + mp, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + require.ErrorIs(t, buffer.Resize(-1), ErrAllocationAccountInvalid) + buffer.Free() + finalizeTestAllocationAccount(t, registry, account) } func usedSizeToInt(t testing.TB, value uint64) int { diff --git a/pkg/container/pSpool/buffer.go b/pkg/container/pSpool/buffer.go index 541c9ef9acee4..c152fc0d07cf1 100644 --- a/pkg/container/pSpool/buffer.go +++ b/pkg/container/pSpool/buffer.go @@ -88,6 +88,10 @@ func (b *spoolBuffer) putCacheID(mp *mpool.MPool, id uint32, bat *batch.Batch) { ) } } + // data/area ownership has moved to the cache. Release the remaining + // Vector-owned state, including allocation-accounted bitmap backing, + // before dropping the Vector pointer. + vec.Free(mp) bat.ReplaceVector(vec, nil, i) } diff --git a/pkg/container/pSpool/copy.go b/pkg/container/pSpool/copy.go index 38b327d7a3a9f..81c7116aa30f7 100644 --- a/pkg/container/pSpool/copy.go +++ b/pkg/container/pSpool/copy.go @@ -149,6 +149,21 @@ func (cb *cachedBatch) GetCopiedBatch( dst.Vecs[i].SetSorted(vec.GetSorted()) } + if vec.HasGrouping() { + groupingRows := vec.GetGrouping().GetBitmap().Len() + if groupingRows < 0 || groupingRows > int64(math.MaxInt) { + cb.CacheBatch(true, cacheID, dst) + return nil, false, 0, mpool.ErrAllocationAccountInvalid + } + if err = dst.Vecs[i].PreExtendBitmap( + int(groupingRows), + cb.mp, + ); err != nil { + cb.CacheBatch(true, cacheID, dst) + return nil, false, 0, err + } + dst.Vecs[i].SetGrouping(vec.GetGrouping()) + } dst.Vecs[i].SetIsBin(vec.GetIsBin()) // range src and found the same vector. diff --git a/pkg/container/pSpool/sender_test.go b/pkg/container/pSpool/sender_test.go index 99c766c376bee..516e06c64eb2b 100644 --- a/pkg/container/pSpool/sender_test.go +++ b/pkg/container/pSpool/sender_test.go @@ -105,20 +105,24 @@ func TestCachedBatchPreservesAllocationProvenance(t *testing.T) { require.NoError(t, err) account, err := registry.Open(1 << 20) require.NoError(t, err) - selection, err := vector.NewAllocationAccountSelection( + selection, err := vector.NewAllocationAccountSelectionWithBitmaps( account, 1, 1, 2, + 3, + 4, ) require.NoError(t, err) otherAccount, err := registry.Open(1 << 20) require.NoError(t, err) - otherSelection, err := vector.NewAllocationAccountSelection( + otherSelection, err := vector.NewAllocationAccountSelectionWithBitmaps( otherAccount, 1, 1, 2, + 3, + 4, ) require.NoError(t, err) newSource := func( @@ -135,11 +139,13 @@ func TestCachedBatchPreservesAllocationProvenance(t *testing.T) { false, mp, )) + vec.GetGrouping().Add(0) source.SetRowCount(1) return source } firstSource := newSource("first cached allocation payload", selection) secondSource := newSource("second", selection) + secondSource.Vecs[0].ToConst() cache := initCachedBatch(mp, 1) first, useCache, cacheID, err := cache.GetCopiedBatch(firstSource) @@ -151,20 +157,23 @@ func TestCachedBatchPreservesAllocationProvenance(t *testing.T) { selection, first.Vecs[0].AllocationAccountSelection(), ) + require.True(t, first.Vecs[0].GetGrouping().Contains(0)) cache.CacheBatch(useCache, cacheID, first) beforeReuse := account.Snapshot().Used second, useCache, cacheID, err := cache.GetCopiedBatch(secondSource) require.NoError(t, err) require.True(t, useCache) - require.Equal(t, beforeReuse, account.Snapshot().Used) + require.Greater(t, account.Snapshot().Used, beforeReuse) require.Same(t, selection, second.AllocationAccountSelection()) require.Same( t, selection, second.Vecs[0].AllocationAccountSelection(), ) + require.True(t, second.Vecs[0].GetGrouping().Contains(0)) cache.CacheBatch(useCache, cacheID, second) + require.Equal(t, beforeReuse, account.Snapshot().Used) otherSource := newSource("other allocation account", otherSelection) firstAccountBeforeOther := account.Snapshot().Used diff --git a/pkg/container/vector/allocation_account.go b/pkg/container/vector/allocation_account.go index 96ed56ad09194..32849404bd986 100644 --- a/pkg/container/vector/allocation_account.go +++ b/pkg/container/vector/allocation_account.go @@ -17,23 +17,64 @@ package vector import ( "fmt" "io" + "math" + "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/types" ) // AllocationAccountSelection is an immutable choice for the first owned -// off-heap data and area allocations of a Vector. The physical MPool -// allocation metadata remains the sole owner of the resulting charge. +// off-heap allocations of a Vector. The physical MPool allocation metadata +// remains the sole owner of the resulting charge. // // A selection may be shared by all vectors owned by one Batch. Views do not // copy it: they share storage and therefore must not create a second charge. type AllocationAccountSelection struct { - account *mpool.AllocationAccount - owner mpool.AllocationOwner - dataSite mpool.AllocationSite - areaSite mpool.AllocationSite + account *mpool.AllocationAccount + owner mpool.AllocationOwner + dataSite mpool.AllocationSite + areaSite mpool.AllocationSite + nullsSite mpool.AllocationSite + groupingSite mpool.AllocationSite + accountBitmaps bool +} + +// FunctionParameterAllocation is the immutable allocation provenance for +// row-scaled function-parameter conversion scratch. +type FunctionParameterAllocation struct { + account *mpool.AllocationAccount + owner mpool.AllocationOwner + site mpool.AllocationSite +} + +func NewFunctionParameterAllocation( + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, + site mpool.AllocationSite, +) (*FunctionParameterAllocation, error) { + allocation := &FunctionParameterAllocation{ + account: account, + owner: owner, + site: site, + } + if err := allocation.validate(); err != nil { + return nil, err + } + return allocation, nil +} + +func (a *FunctionParameterAllocation) validate() error { + if a == nil || + a.account == nil || + a.account.Handle() == 0 || + a.owner < mpool.AllocationOwnerMin || + a.owner > mpool.AllocationOwnerMax || + a.site < mpool.AllocationSiteMin { + return mpool.ErrAllocationAccountInvalid + } + return nil } func NewAllocationAccountSelection( @@ -54,8 +95,33 @@ func NewAllocationAccountSelection( return selection, nil } +// NewAllocationAccountSelectionWithBitmaps additionally selects physical +// allocation sites for the Vector null and grouping bitmap backing. +func NewAllocationAccountSelectionWithBitmaps( + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, + dataSite mpool.AllocationSite, + areaSite mpool.AllocationSite, + nullsSite mpool.AllocationSite, + groupingSite mpool.AllocationSite, +) (*AllocationAccountSelection, error) { + selection := &AllocationAccountSelection{ + account: account, + owner: owner, + dataSite: dataSite, + areaSite: areaSite, + nullsSite: nullsSite, + groupingSite: groupingSite, + accountBitmaps: true, + } + if err := selection.validate(); err != nil { + return nil, err + } + return selection, nil +} + // NewOffHeapVecWithTypeAndAllocation constructs an empty owning Vector whose -// future data and area allocations use selection. +// future allocations use selection. func NewOffHeapVecWithTypeAndAllocation( typ types.Type, selection *AllocationAccountSelection, @@ -108,7 +174,7 @@ func NewConstFixedWithAllocation[T any]( } // NewConstBytesWithAllocation constructs an off-heap constant varlen Vector -// and charges data and area independently through selection. +// and charges its backing independently through selection. func NewConstBytesWithAllocation( typ types.Type, value []byte, @@ -131,7 +197,7 @@ func NewConstBytesWithAllocation( } // NewConstArrayWithAllocation constructs an off-heap constant array Vector -// and charges data and area independently through selection. +// and charges its backing independently through selection. func NewConstArrayWithAllocation[T types.ArrayElement]( typ types.Type, value []T, @@ -158,7 +224,10 @@ func (s *AllocationAccountSelection) validate() error { s.owner < mpool.AllocationOwnerMin || s.owner > mpool.AllocationOwnerMax || s.dataSite < mpool.AllocationSiteMin || - s.areaSite < mpool.AllocationSiteMin { + s.areaSite < mpool.AllocationSiteMin || + (s.accountBitmaps && + (s.nullsSite < mpool.AllocationSiteMin || + s.groupingSite < mpool.AllocationSiteMin)) { return mpool.ErrAllocationAccountInvalid } return nil @@ -205,22 +274,125 @@ func (v *Vector) CanSetAllocationAccount( } func (v *Vector) hasBackingStorage() bool { - return cap(v.data) != 0 || cap(v.area) != 0 + return cap(v.data) != 0 || + cap(v.area) != 0 || + v.nsp.GetBitmap().Size() != 0 || + v.gsp.GetBitmap().Size() != 0 || + v.nsp.GetBitmap().ExternalStorageCapacity() != 0 || + v.gsp.GetBitmap().ExternalStorageCapacity() != 0 } -// SetAllocationAccount selects the account used by future owned data and area -// allocations. It is intentionally explicit and is legal only before the -// first backing allocation. Reset retains the selection; Free clears it. +// SetAllocationAccount selects the account used by future owned allocations. +// It is intentionally explicit and is legal only before the first backing +// allocation. Reset retains the selection; Free clears it. func (v *Vector) SetAllocationAccount( selection *AllocationAccountSelection, ) error { if err := v.CanSetAllocationAccount(selection); err != nil { return err } + if v.allocationAccount == selection { + return nil + } + if v.allocationAccount != nil && + v.allocationAccount.accountBitmaps && + (selection == nil || !selection.accountBitmaps) { + v.nsp.GetBitmap().ReleaseExternalStorage() + v.gsp.GetBitmap().ReleaseExternalStorage() + } v.allocationAccount = selection + if selection != nil && selection.accountBitmaps { + v.nsp.GetBitmap().InstallExternalStorage(nil) + v.gsp.GetBitmap().InstallExternalStorage(nil) + } return nil } +func (v *Vector) ensureBitmapCapacity(rows int, mp *mpool.MPool) error { + if v.allocationAccount == nil || !v.allocationAccount.accountBitmaps { + return nil + } + if rows < 0 || rows > math.MaxInt-64 || mp == nil { + return mpool.ErrAllocationAccountInvalid + } + // Nulls.AddRange currently expands through end+1 even though end is + // exclusive. Keep one admitted sentinel bit so raw bitmap mutation cannot + // escape to a Go allocation at the vector's logical row boundary. + if rows > 0 { + rows++ + } + nulls, err := v.allocateBitmapGrowth( + v.nsp.GetBitmap(), + rows, + mp, + v.allocationAccount.nullsSite, + ) + if err != nil { + return err + } + grouping, err := v.allocateBitmapGrowth( + v.gsp.GetBitmap(), + rows, + mp, + v.allocationAccount.groupingSite, + ) + if err != nil { + mpool.FreeSlice(mp, nulls) + return err + } + if cap(nulls) > 0 { + previous := v.nsp.GetBitmap().InstallExternalStorage(nulls) + mpool.FreeSlice(mp, previous) + } + if cap(grouping) > 0 { + previous := v.gsp.GetBitmap().InstallExternalStorage(grouping) + mpool.FreeSlice(mp, previous) + } + return nil +} + +func (v *Vector) allocateBitmapGrowth( + value *bitmap.Bitmap, + rows int, + mp *mpool.MPool, + site mpool.AllocationSite, +) ([]uint64, error) { + requiredWords := (rows + 63) / 64 + if requiredWords <= value.ExternalStorageCapacity() { + return nil, nil + } + requiredBytes := int64(requiredWords) * 8 + oldBytes := int64(value.ExternalStorageCapacity()) * 8 + newBytes, ok := mpool.GrowCapacity(oldBytes, requiredBytes) + if !ok || newBytes > int64(math.MaxInt) || newBytes%8 != 0 { + return nil, mpool.ErrAllocationAccountInvalid + } + next, err := mpool.MakeSliceAccounted[uint64]( + int(newBytes/8), + mp, + v.allocationAccount.account, + v.allocationAccount.owner, + site, + ) + if err != nil { + return nil, err + } + clear(next) + return next, nil +} + +func (v *Vector) freeBitmapStorage(mp *mpool.MPool) { + for _, value := range []*bitmap.Bitmap{ + v.nsp.GetBitmap(), + v.gsp.GetBitmap(), + } { + storage := value.ReleaseExternalStorage() + if cap(storage) > 0 { + mpool.FreeSlice(mp, storage) + } + } +} + func (v *Vector) allocData(mp *mpool.MPool, size int) ([]byte, error) { return v.allocOwned(mp, size, v.offHeap, true) } diff --git a/pkg/container/vector/allocation_account_test.go b/pkg/container/vector/allocation_account_test.go index 7fb139493e2a6..253343a2314a7 100644 --- a/pkg/container/vector/allocation_account_test.go +++ b/pkg/container/vector/allocation_account_test.go @@ -26,15 +26,78 @@ import ( ) const ( - testVectorAllocationOwner mpool.AllocationOwner = 1 - testVectorDataAllocationSite mpool.AllocationSite = 1 - testVectorAreaAllocationSite mpool.AllocationSite = 2 + testVectorAllocationOwner mpool.AllocationOwner = 1 + testVectorDataAllocationSite mpool.AllocationSite = 1 + testVectorAreaAllocationSite mpool.AllocationSite = 2 + testVectorNullAllocationSite mpool.AllocationSite = 3 + testVectorGroupAllocationSite mpool.AllocationSite = 4 + testVectorParamAllocationSite mpool.AllocationSite = 5 ) type testVectorAllocationAccount struct { registry *mpool.AllocationAccountRegistry account *mpool.AllocationAccount selection *AllocationAccountSelection + parameter *FunctionParameterAllocation +} + +func newTestVectorParameterAllocationAccount( + t testing.TB, + limit uint64, + allocationSlots uint64, +) testVectorAllocationAccount { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, allocationSlots) + require.NoError(t, err) + account, err := registry.Open(limit) + require.NoError(t, err) + selection, err := NewAllocationAccountSelectionWithBitmaps( + account, + testVectorAllocationOwner, + testVectorDataAllocationSite, + testVectorAreaAllocationSite, + testVectorNullAllocationSite, + testVectorGroupAllocationSite, + ) + require.NoError(t, err) + parameter, err := NewFunctionParameterAllocation( + account, + testVectorAllocationOwner, + testVectorParamAllocationSite, + ) + require.NoError(t, err) + return testVectorAllocationAccount{ + registry: registry, + account: account, + selection: selection, + parameter: parameter, + } +} + +func newTestVectorBitmapAllocationAccount( + t testing.TB, + limit uint64, + allocationSlots uint64, +) testVectorAllocationAccount { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, allocationSlots) + require.NoError(t, err) + account, err := registry.Open(limit) + require.NoError(t, err) + selection, err := NewAllocationAccountSelectionWithBitmaps( + account, + testVectorAllocationOwner, + testVectorDataAllocationSite, + testVectorAreaAllocationSite, + testVectorNullAllocationSite, + testVectorGroupAllocationSite, + ) + require.NoError(t, err) + return testVectorAllocationAccount{ + registry: registry, + account: account, + selection: selection, + } } func newTestVectorAllocationAccount( @@ -207,6 +270,164 @@ func TestVectorAllocationAccountLeavesGoBitmapsUnaccounted(t *testing.T) { finalizeTestVectorAllocationAccount(t, state) } +func TestVectorAllocationAccountBitmapResetReuseAndFree(t *testing.T) { + state := newTestVectorBitmapAllocationAccount(t, 8<<20, 16) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + + require.NoError(t, vec.PreExtend(32*1024, mp)) + vec.SetLength(32 * 1024) + vec.SetAllNulls(32 * 1024) + vec.GetGrouping().AddRange(0, 32*1024) + require.True(t, vec.nsp.GetBitmap().HasExternalStorage()) + require.True(t, vec.gsp.GetBitmap().HasExternalStorage()) + require.Equal(t, 32*1024, vec.GetNulls().Count()) + require.Equal(t, 32*1024, vec.GetGrouping().Count()) + + initial := state.account.Snapshot() + expected := cap(vec.data) + + 8*vec.nsp.GetBitmap().ExternalStorageCapacity() + + 8*vec.gsp.GetBitmap().ExternalStorageCapacity() + require.Equal(t, uint64(expected), initial.Used) + require.Equal(t, uint64(3), state.registry.LiveAllocationMetadata()) + + vec.ResetWithSameType() + require.True(t, vec.GetNulls().IsEmpty()) + require.True(t, vec.GetGrouping().IsEmpty()) + require.Equal(t, initial.Used, state.account.Snapshot().Used) + require.NoError(t, vec.PreExtend(32*1024, mp)) + vec.SetNull(32*1024 - 1) + require.True(t, vec.IsNull(32*1024-1)) + require.Equal(t, initial.Used, state.account.Snapshot().Used) + + vec.ResetWithSameType() + require.NoError(t, vec.PreExtend(64*1024, mp)) + grown := state.account.Snapshot() + require.Greater(t, grown.Used, initial.Used) + require.Greater(t, grown.Peak, grown.Used) + require.Equal(t, uint64(3), state.registry.LiveAllocationMetadata()) + + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountBitmapGrowthFailurePreservesOwner(t *testing.T) { + state := newTestVectorBitmapAllocationAccount(t, 1000, 8) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, vec.PreExtend(64, mp)) + vec.SetLength(64) + vec.SetNull(7) + vec.GetGrouping().Add(9) + + used := state.account.Snapshot().Used + dataCapacity := cap(vec.data) + nullCapacity := vec.nsp.GetBitmap().ExternalStorageCapacity() + groupCapacity := vec.gsp.GetBitmap().ExternalStorageCapacity() + // The null replacement fits by itself, but admitting the grouping + // replacement would exceed the account. Neither replacement is published. + err := vec.PreExtend(2*1024, mp) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Equal(t, used, state.account.Snapshot().Used) + require.Equal(t, dataCapacity, cap(vec.data)) + require.Equal(t, nullCapacity, vec.nsp.GetBitmap().ExternalStorageCapacity()) + require.Equal(t, groupCapacity, vec.gsp.GetBitmap().ExternalStorageCapacity()) + require.True(t, vec.IsNull(7)) + require.True(t, vec.GetGrouping().Contains(9)) + + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountBitmapRejectsUnadmittedRawGrowth(t *testing.T) { + state := newTestVectorBitmapAllocationAccount(t, 1<<20, 8) + mp := mpool.MustNewZero() + legacy := NewOffHeapVecWithType(types.T_int64.ToType()) + legacy.GetNulls().Add(0) + require.ErrorIs( + t, + legacy.SetAllocationAccount(state.selection), + mpool.ErrAllocationAccountInvalid, + ) + legacy.Free(mp) + + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.Panics(t, func() { + vec.GetNulls().Add(0) + }) + require.Zero(t, state.account.Snapshot().Used) + + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountBitmapCopyDecode(t *testing.T) { + state := newTestVectorBitmapAllocationAccount(t, 1<<20, 32) + mp := mpool.MustNewZero() + source := NewOffHeapVecWithType(types.T_int64.ToType()) + for i := 0; i < 128; i++ { + require.NoError(t, AppendFixed(source, int64(i), i%7 == 0, mp)) + } + source.GetGrouping().Add(2, 9, 64) + encoded, err := source.MarshalBinary() + require.NoError(t, err) + + copied := newAccountedTestVector( + t, + types.T_int64.ToType(), + state.selection, + ) + require.NoError(t, copied.UnmarshalBinaryWithCopy(encoded, mp)) + require.Equal(t, source.Length(), copied.Length()) + require.True(t, copied.GetNulls().IsSame(source.GetNulls())) + require.True(t, copied.nsp.GetBitmap().HasExternalStorage()) + require.True(t, copied.gsp.GetBitmap().HasExternalStorage()) + copied.Free(mp) + + fromReader := newAccountedTestVector( + t, + types.T_int64.ToType(), + state.selection, + ) + require.NoError(t, fromReader.UnmarshalWithReader(bytes.NewReader(encoded), mp)) + require.Equal(t, source.Length(), fromReader.Length()) + require.True(t, fromReader.GetNulls().IsSame(source.GetNulls())) + require.True(t, fromReader.nsp.GetBitmap().HasExternalStorage()) + fromReader.Free(mp) + + duplicate, err := source.DupOffHeapWithAllocation(mp, state.selection) + require.NoError(t, err) + require.True(t, duplicate.GetNulls().IsSame(source.GetNulls())) + require.True(t, duplicate.GetGrouping().IsSame(source.GetGrouping())) + duplicate.Free(mp) + + window, err := source.CloneWindowWithAllocation( + 1, + 65, + mp, + state.selection, + ) + require.NoError(t, err) + require.Equal(t, 64, window.Length()) + require.True(t, window.IsNull(6)) + require.True(t, window.IsNull(13)) + window.Free(mp) + + rollup := NewOffHeapVecWithType(types.T_int64.ToType()) + rollup.SetLength(128) + rollup.GetGrouping().AddRange(0, 128) + rollup.ToConst() + rollupCopy, err := rollup.DupOffHeapWithAllocation(mp, state.selection) + require.NoError(t, err) + require.True(t, rollupCopy.IsConstNull()) + require.True(t, rollupCopy.GetGrouping().IsSame(rollup.GetGrouping())) + rollupCopy.Free(mp) + rollup.Free(mp) + + source.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + func TestVectorAllocationAccountViewAndDeepCopy(t *testing.T) { stateA := newTestVectorAllocationAccount(t, 1<<20, 32) stateB := newTestVectorAllocationAccount(t, 1<<20, 32) @@ -413,6 +634,7 @@ func BenchmarkVectorAllocationAccount(b *testing.B) { const rows = 8192 mp := mpool.MustNewZero() state := newTestVectorAllocationAccount(b, 1<<40, 64) + bitmapState := newTestVectorBitmapAllocationAccount(b, 1<<40, 64) b.Run("legacy-fixed-preextend-free", func(b *testing.B) { b.ReportAllocs() @@ -437,6 +659,19 @@ func BenchmarkVectorAllocationAccount(b *testing.B) { vec.Free(mp) } }) + b.Run("accounted-bitmap-fixed-preextend-free", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + vec := NewOffHeapVecWithType(types.T_int64.ToType()) + if err := vec.SetAllocationAccount(bitmapState.selection); err != nil { + b.Fatal(err) + } + if err := vec.PreExtend(rows, mp); err != nil { + b.Fatal(err) + } + vec.Free(mp) + } + }) b.Run("legacy-varlen-preextend-free", func(b *testing.B) { b.ReportAllocs() for i := 0; i < b.N; i++ { @@ -473,8 +708,26 @@ func BenchmarkVectorAllocationAccount(b *testing.B) { b.StopTimer() vec.Free(mp) }) + b.Run("accounted-bitmap-fixed-reset-reuse", func(b *testing.B) { + vec := newAccountedTestVector( + b, + types.T_int64.ToType(), + bitmapState.selection, + ) + if err := vec.PreExtend(rows, mp); err != nil { + b.Fatal(err) + } + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + vec.ResetWithSameType() + } + b.StopTimer() + vec.Free(mp) + }) finalizeTestVectorAllocationAccount(b, state) + finalizeTestVectorAllocationAccount(b, bitmapState) } func TestVectorAllocationAccountErrorsAreTyped(t *testing.T) { diff --git a/pkg/container/vector/functionTools.go b/pkg/container/vector/functionTools.go index 2c34ce57fb336..dc5c48bb80337 100644 --- a/pkg/container/vector/functionTools.go +++ b/pkg/container/vector/functionTools.go @@ -16,9 +16,12 @@ package vector import ( "fmt" + "math" + "unsafe" "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/common/util" "github.com/matrixorigin/matrixone/pkg/container/bytejson" "github.com/matrixorigin/matrixone/pkg/container/nulls" "github.com/matrixorigin/matrixone/pkg/container/types" @@ -66,106 +69,45 @@ func GenerateFunctionFixedTypeParameter[T types.FixedSizeTExceptStrType](v *Vect } } - // Special handling for type conversions to decimal128 var cols []T - var convertedType types.Type + parameterType := *t var anyT T switch (any)(anyT).(type) { case types.Decimal128: - convertedType = types.T_decimal128.ToType() - convertedType.Width = 38 - if t.Oid == types.T_decimal64 { - convertedType.Scale = t.Scale - d64Cols := MustFixedColWithTypeCheck[types.Decimal64](v) - // Optimize: for const vector, only convert one element - if v.IsConst() { - d128 := functionUtil.ConvertD64ToD128(d64Cols[0]) - return &FunctionParameterScalar[T]{ - typ: convertedType, - sourceVector: v, - scalarValue: any(d128).(T), - } - } - cols = make([]T, len(d64Cols)) - for i, d64 := range d64Cols { - cols[i] = any(functionUtil.ConvertD64ToD128(d64)).(T) - } - } else if t.Oid == types.T_float64 { - convertedType.Scale = 16 - f64Cols := MustFixedColWithTypeCheck[float64](v) - // Optimize: for const vector, only convert one element - if v.IsConst() { - d128, err := types.Decimal128FromFloat64(f64Cols[0], 38, 16) - if err != nil { - // Conversion failed, use zero value (similar to MySQL behavior for invalid conversions) - d128 = types.Decimal128{B0_63: 0, B64_127: 0} - } - return &FunctionParameterScalar[T]{ - typ: convertedType, - sourceVector: v, - scalarValue: any(d128).(T), - } - } - cols = make([]T, len(f64Cols)) - for i, f64 := range f64Cols { - d128, err := types.Decimal128FromFloat64(f64, 38, 16) - if err != nil { - // Conversion failed, use zero value - d128 = types.Decimal128{B0_63: 0, B64_127: 0} - } - cols[i] = any(d128).(T) - } - } else if t.Oid == types.T_float32 { - convertedType.Scale = 7 - f32Cols := MustFixedColWithTypeCheck[float32](v) - // Optimize: for const vector, only convert one element - if v.IsConst() { - d128, err := types.Decimal128FromFloat64(float64(f32Cols[0]), 38, 7) - if err != nil { - // Conversion failed, use zero value - d128 = types.Decimal128{B0_63: 0, B64_127: 0} - } - return &FunctionParameterScalar[T]{ - typ: convertedType, - sourceVector: v, - scalarValue: any(d128).(T), - } - } - cols = make([]T, len(f32Cols)) - for i, f32 := range f32Cols { - d128, err := types.Decimal128FromFloat64(float64(f32), 38, 7) - if err != nil { - // Conversion failed, use zero value - d128 = types.Decimal128{B0_63: 0, B64_127: 0} - } - cols[i] = any(d128).(T) + if needsDecimal128ParameterConversion[T](v) { + parameter, err := generateDecimal128ParameterWithScratch[T]( + nil, + 0, + v, + nil, + ) + if err != nil { + panic(err) } - } else { - convertedType = *t - cols = MustFixedColWithTypeCheck[T](v) + return parameter } + cols = MustFixedColWithTypeCheck[T](v) default: - convertedType = *t cols = MustFixedColWithTypeCheck[T](v) } if v.IsConst() { return &FunctionParameterScalar[T]{ - typ: convertedType, + typ: parameterType, sourceVector: v, scalarValue: cols[0], } } if !v.nsp.IsEmpty() { return &FunctionParameterNormal[T]{ - typ: convertedType, + typ: parameterType, sourceVector: v, values: cols, nullMap: v.GetNulls().GetBitmap(), } } return &FunctionParameterWithoutNull[T]{ - typ: convertedType, + typ: parameterType, sourceVector: v, values: cols, } @@ -541,19 +483,225 @@ type reusableParameterWrapper interface{} type OptFunctionResultWrapper interface { UseOptFunctionParamFrame(paramCount int) getConvenientParamList() []reusableParameterWrapper + hasParameterScratch() bool + resizeParameterScratch(idx int, size int) ([]byte, error) + setParameterAllocation(allocation *FunctionParameterAllocation) } func OptGetParamFromWrapper[ParamType types.FixedSizeTExceptStrType]( - wrapper FunctionResultWrapper, idx int, src *Vector) FunctionParameterWrapper[ParamType] { + wrapper FunctionResultWrapper, + idx int, + src *Vector, +) (FunctionParameterWrapper[ParamType], error) { ws := wrapper.getConvenientParamList() + if needsDecimal128ParameterConversion[ParamType](src) { + if !wrapper.hasParameterScratch() { + fr := GenerateFunctionFixedTypeParameter[ParamType](src) + ws[idx] = fr + return fr, nil + } + fr, err := generateDecimal128ParameterWithScratch[ParamType]( + wrapper, + idx, + src, + ws[idx], + ) + if err != nil { + return nil, err + } + ws[idx] = fr + return fr, nil + } if fr, ok := ws[idx].(FunctionParameterWrapper[ParamType]); ok && ReuseFunctionFixedTypeParameter(src, fr) { - return fr + return fr, nil } fr := GenerateFunctionFixedTypeParameter[ParamType](src) ws[idx] = fr - return fr + return fr, nil +} + +func needsDecimal128ParameterConversion[T types.FixedSizeTExceptStrType]( + src *Vector, +) bool { + var value T + if _, ok := any(value).(types.Decimal128); !ok { + return false + } + switch src.GetType().Oid { + case types.T_decimal64, types.T_float32, types.T_float64: + return true + default: + return false + } +} + +func generateDecimal128ParameterWithScratch[ + T types.FixedSizeTExceptStrType, +]( + wrapper FunctionResultWrapper, + idx int, + src *Vector, + reuse reusableParameterWrapper, +) (FunctionParameterWrapper[T], error) { + if src.IsConstNull() { + if parameter, ok := reuse.(*FunctionParameterScalarNull[T]); ok { + parameter.typ = *src.GetType() + parameter.sourceVector = src + return parameter, nil + } + return &FunctionParameterScalarNull[T]{ + typ: *src.GetType(), + sourceVector: src, + }, nil + } + convertedType := types.T_decimal128.ToType() + convertedType.Width = 38 + switch src.GetType().Oid { + case types.T_decimal64: + convertedType.Scale = src.GetType().Scale + case types.T_float32: + convertedType.Scale = 7 + case types.T_float64: + convertedType.Scale = 16 + default: + return nil, mpool.ErrAllocationAccountInvalid + } + + convert := func(row int) types.Decimal128 { + switch src.GetType().Oid { + case types.T_decimal64: + values := MustFixedColWithTypeCheck[types.Decimal64](src) + return functionUtil.ConvertD64ToD128(values[row]) + case types.T_float32: + values := MustFixedColWithTypeCheck[float32](src) + value, err := types.Decimal128FromFloat64( + float64(values[row]), + 38, + 7, + ) + if err == nil { + return value + } + case types.T_float64: + values := MustFixedColWithTypeCheck[float64](src) + value, err := types.Decimal128FromFloat64(values[row], 38, 16) + if err == nil { + return value + } + } + return types.Decimal128{} + } + + if src.IsConst() { + value := any(convert(0)).(T) + if parameter, ok := reuse.(*FunctionParameterScalar[T]); ok { + parameter.typ = convertedType + parameter.sourceVector = src + parameter.scalarValue = value + return parameter, nil + } + return &FunctionParameterScalar[T]{ + typ: convertedType, + sourceVector: src, + scalarValue: value, + }, nil + } + + var values []T + if wrapper == nil { + values = make([]T, src.Length()) + } else { + var err error + values, err = parameterScratchSlice[T](wrapper, idx, src.Length()) + if err != nil { + return nil, err + } + } + switch src.GetType().Oid { + case types.T_decimal64: + source := MustFixedColWithTypeCheck[types.Decimal64](src) + for row := range values { + values[row] = any( + functionUtil.ConvertD64ToD128(source[row]), + ).(T) + } + case types.T_float32: + source := MustFixedColWithTypeCheck[float32](src) + for row := range values { + value, conversionErr := types.Decimal128FromFloat64( + float64(source[row]), + 38, + 7, + ) + if conversionErr != nil { + value = types.Decimal128{} + } + values[row] = any(value).(T) + } + case types.T_float64: + source := MustFixedColWithTypeCheck[float64](src) + for row := range values { + value, conversionErr := types.Decimal128FromFloat64( + source[row], + 38, + 16, + ) + if conversionErr != nil { + value = types.Decimal128{} + } + values[row] = any(value).(T) + } + } + if !src.nsp.IsEmpty() { + if parameter, ok := reuse.(*FunctionParameterNormal[T]); ok { + parameter.typ = convertedType + parameter.sourceVector = src + parameter.values = values + parameter.nullMap = src.GetNulls().GetBitmap() + return parameter, nil + } + return &FunctionParameterNormal[T]{ + typ: convertedType, + sourceVector: src, + values: values, + nullMap: src.GetNulls().GetBitmap(), + }, nil + } + if parameter, ok := reuse.(*FunctionParameterWithoutNull[T]); ok { + parameter.typ = convertedType + parameter.sourceVector = src + parameter.values = values + return parameter, nil + } + return &FunctionParameterWithoutNull[T]{ + typ: convertedType, + sourceVector: src, + values: values, + }, nil +} + +func parameterScratchSlice[T types.FixedSizeTExceptStrType]( + wrapper FunctionResultWrapper, + idx int, + length int, +) ([]T, error) { + var value T + elementSize := unsafe.Sizeof(value) + if length < 0 || + elementSize == 0 || + uint64(length) > uint64(math.MaxInt)/uint64(elementSize) { + return nil, mpool.ErrAllocationAccountInvalid + } + data, err := wrapper.resizeParameterScratch( + idx, + int(uint64(length)*uint64(elementSize)), + ) + if err != nil { + return nil, err + } + return util.UnsafeSliceCastToLength[T](data, length), nil } func OptGetBytesParamFromWrapper(wrapper FunctionResultWrapper, idx int, src *Vector) FunctionParameterWrapper[types.Varlena] { @@ -575,16 +723,18 @@ type FunctionResult[T types.FixedSizeT] struct { vec *Vector mp *mpool.MPool - allocationAccount *AllocationAccountSelection - isVarlena bool - cols []T - length uint64 + allocationAccount *AllocationAccountSelection + parameterAllocation *FunctionParameterAllocation + isVarlena bool + cols []T + length uint64 // convenientParam save parameter wrappers for easy getting row values. // // this field is for optimisation to reduce the allocation of FunctionParameterWrapper pointer. // there are still many built-in functions don't use it now, and will be fixed in the future. - convenientParam []reusableParameterWrapper + convenientParam []reusableParameterWrapper + parameterScratch []*mpool.AccountedBuffer } func MustFunctionResult[T types.FixedSizeT](wrapper FunctionResultWrapper) *FunctionResult[T] { @@ -618,12 +768,55 @@ func (fr *FunctionResult[T]) UseOptFunctionParamFrame(paramCount int) { if fr.convenientParam == nil { fr.convenientParam = make([]reusableParameterWrapper, paramCount) } + if fr.allocationAccount != nil && + fr.parameterAllocation != nil && + fr.parameterScratch == nil { + fr.parameterScratch = make([]*mpool.AccountedBuffer, paramCount) + } } func (fr *FunctionResult[T]) getConvenientParamList() []reusableParameterWrapper { return fr.convenientParam } +func (fr *FunctionResult[T]) hasParameterScratch() bool { + return fr.parameterAllocation != nil +} + +func (fr *FunctionResult[T]) setParameterAllocation( + allocation *FunctionParameterAllocation, +) { + fr.parameterAllocation = allocation +} + +func (fr *FunctionResult[T]) resizeParameterScratch( + idx int, + size int, +) ([]byte, error) { + if !fr.hasParameterScratch() { + return nil, mpool.ErrAllocationAccountInvalid + } + if idx < 0 || idx >= len(fr.parameterScratch) { + return nil, mpool.ErrAllocationAccountInvalid + } + if fr.parameterScratch[idx] == nil { + buffer, err := mpool.NewAccountedBuffer( + fr.mp, + fr.parameterAllocation.account, + fr.parameterAllocation.owner, + fr.parameterAllocation.site, + ) + if err != nil { + return nil, err + } + fr.parameterScratch[idx] = buffer + } + if err := fr.parameterScratch[idx].Resize(size); err != nil { + return nil, err + } + return fr.parameterScratch[idx].Bytes(), nil +} + func (fr *FunctionResult[T]) PreExtendAndReset(targetSize int) error { if fr.vec == nil { var err error @@ -773,8 +966,16 @@ func (fr *FunctionResult[T]) Free() { fr.vec.Free(fr.mp) fr.vec = nil } + for i := range fr.parameterScratch { + if fr.parameterScratch[i] != nil { + fr.parameterScratch[i].Free() + fr.parameterScratch[i] = nil + } + } fr.allocationAccount = nil + fr.parameterAllocation = nil fr.convenientParam = nil + fr.parameterScratch = nil } func NewFunctionResultWrapper(typ types.Type, mp *mpool.MPool) FunctionResultWrapper { @@ -795,6 +996,27 @@ func NewFunctionResultWrapperWithAllocation( return newFunctionResultWrapper(typ, mp, selection), nil } +func NewFunctionResultWrapperWithParameterAllocation( + typ types.Type, + mp *mpool.MPool, + selection *AllocationAccountSelection, + parameterAllocation *FunctionParameterAllocation, +) (FunctionResultWrapper, error) { + if err := selection.validate(); err != nil { + return nil, err + } + if err := parameterAllocation.validate(); err != nil { + return nil, err + } + if selection.account != parameterAllocation.account || + selection.owner != parameterAllocation.owner { + return nil, mpool.ErrAllocationAccountInvalid + } + result := newFunctionResultWrapper(typ, mp, selection) + result.setParameterAllocation(parameterAllocation) + return result, nil +} + func newFunctionResultWrapper( typ types.Type, mp *mpool.MPool, diff --git a/pkg/container/vector/function_result_allocation_test.go b/pkg/container/vector/function_result_allocation_test.go index c33ec0186816a..4a32d6c478f2b 100644 --- a/pkg/container/vector/function_result_allocation_test.go +++ b/pkg/container/vector/function_result_allocation_test.go @@ -91,7 +91,36 @@ func TestFunctionResultAllocationAccountFailure(t *testing.T) { mpool.ErrAllocationAccountInvalid, ) - state := newTestVectorAllocationAccount(t, 7, 1) + state := newTestVectorAllocationAccount(t, 1<<20, 4) + _, err := NewFunctionResultWrapperWithParameterAllocation( + types.T_int64.ToType(), + zeroMP, + state.selection, + nil, + ) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + otherOwner, err := NewFunctionParameterAllocation( + state.account, + testVectorAllocationOwner+1, + testVectorParamAllocationSite, + ) + require.NoError(t, err) + _, err = NewFunctionResultWrapperWithParameterAllocation( + types.T_int64.ToType(), + zeroMP, + state.selection, + otherOwner, + ) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + _, err = NewFunctionParameterAllocation( + nil, + testVectorAllocationOwner, + testVectorParamAllocationSite, + ) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + finalizeTestVectorAllocationAccount(t, state) + + state = newTestVectorAllocationAccount(t, 7, 1) mp := mpool.MustNew("function-result-allocation-failure") defer mpool.DeleteMPool(mp) result, err := NewFunctionResultWrapperWithAllocation( @@ -108,3 +137,150 @@ func TestFunctionResultAllocationAccountFailure(t *testing.T) { result.Free() finalizeTestVectorAllocationAccount(t, state) } + +func TestFunctionResultAllocationAccountDecimalParameterScratch(t *testing.T) { + state := newTestVectorParameterAllocationAccount(t, 1<<20, 16) + mp := mpool.MustNew("function-parameter-allocation") + defer mpool.DeleteMPool(mp) + result, err := NewFunctionResultWrapperWithParameterAllocation( + types.T_bool.ToType(), + mp, + state.selection, + state.parameter, + ) + require.NoError(t, err) + result.UseOptFunctionParamFrame(1) + + source := NewOffHeapVecWithType(types.T_decimal64.ToType()) + for i := int64(1); i <= 32; i++ { + require.NoError(t, AppendFixed( + source, + types.Decimal64(i), + i%7 == 0, + mp, + )) + } + parameter, err := OptGetParamFromWrapper[types.Decimal128]( + result, + 0, + source, + ) + require.NoError(t, err) + values := parameter.UnSafeGetAllValue() + require.Len(t, values, source.Length()) + require.Equal(t, types.Decimal128{B0_63: 1}, values[0]) + _, isNull := parameter.GetValue(6) + require.True(t, isNull) + first := state.account.Snapshot() + require.Positive(t, first.Used) + require.Equal(t, uint64(1), state.registry.LiveAllocationMetadata()) + + reused, err := OptGetParamFromWrapper[types.Decimal128]( + result, + 0, + source, + ) + require.NoError(t, err) + require.Same(t, parameter, reused) + require.Equal(t, first.Used, state.account.Snapshot().Used) + + float32Source := NewOffHeapVecWithType(types.T_float32.ToType()) + require.NoError(t, AppendFixed(float32Source, float32(1.25), false, mp)) + float32Parameter, err := OptGetParamFromWrapper[types.Decimal128]( + result, + 0, + float32Source, + ) + require.NoError(t, err) + expectedFloat32, err := types.Decimal128FromFloat64(1.25, 38, 7) + require.NoError(t, err) + require.Equal(t, expectedFloat32, float32Parameter.UnSafeGetAllValue()[0]) + require.Equal(t, int32(7), float32Parameter.GetType().Scale) + + float64Source := NewOffHeapVecWithType(types.T_float64.ToType()) + require.NoError(t, AppendFixed(float64Source, 2.5, false, mp)) + float64Parameter, err := OptGetParamFromWrapper[types.Decimal128]( + result, + 0, + float64Source, + ) + require.NoError(t, err) + expectedFloat64, err := types.Decimal128FromFloat64(2.5, 38, 16) + require.NoError(t, err) + require.Equal(t, expectedFloat64, float64Parameter.UnSafeGetAllValue()[0]) + require.Equal(t, int32(16), float64Parameter.GetType().Scale) + + constSource, err := NewConstFixed( + types.T_float64.ToType(), + 3.5, + 8, + mp, + ) + require.NoError(t, err) + beforeConst := state.account.Snapshot().Used + constParameter, err := OptGetParamFromWrapper[types.Decimal128]( + result, + 0, + constSource, + ) + require.NoError(t, err) + expectedConst, err := types.Decimal128FromFloat64(3.5, 38, 16) + require.NoError(t, err) + value, isNull := constParameter.GetValue(7) + require.False(t, isNull) + require.Equal(t, expectedConst, value) + require.Equal(t, beforeConst, state.account.Snapshot().Used) + + result.Free() + require.Zero(t, state.account.Snapshot().Used) + source.Free(mp) + float32Source.Free(mp) + float64Source.Free(mp) + constSource.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestFunctionResultAllocationAccountDecimalParameterFailure(t *testing.T) { + state := newTestVectorParameterAllocationAccount(t, 127, 4) + mp := mpool.MustNew("function-parameter-allocation-failure") + defer mpool.DeleteMPool(mp) + result, err := NewFunctionResultWrapperWithParameterAllocation( + types.T_bool.ToType(), + mp, + state.selection, + state.parameter, + ) + require.NoError(t, err) + result.UseOptFunctionParamFrame(1) + + source := NewOffHeapVecWithType(types.T_decimal64.ToType()) + for i := 0; i < 8; i++ { + require.NoError(t, AppendFixed( + source, + types.Decimal64(i), + false, + mp, + )) + } + _, err = OptGetParamFromWrapper[types.Decimal128]( + result, + 0, + source, + ) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Zero(t, state.account.Snapshot().Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + + source.SetLength(7) + _, err = OptGetParamFromWrapper[types.Decimal128]( + result, + 0, + source, + ) + require.NoError(t, err) + require.Positive(t, state.account.Snapshot().Used) + + result.Free() + source.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} diff --git a/pkg/container/vector/tools.go b/pkg/container/vector/tools.go index 3c731326a091d..ccce3160db174 100644 --- a/pkg/container/vector/tools.go +++ b/pkg/container/vector/tools.go @@ -208,6 +208,9 @@ func extend(v *Vector, rows int, m *mpool.MPool) error { } tgtLen := v.length + rows + if err := v.ensureBitmapCapacity(tgtLen, m); err != nil { + return err + } tgtDataCap := tgtLen * v.typ.TypeSize() if tgtDataCap > cap(v.data) { ndata, err := v.growData(m, tgtDataCap) diff --git a/pkg/container/vector/vector.go b/pkg/container/vector/vector.go index eedf1c04314ea..20691e3a43239 100644 --- a/pkg/container/vector/vector.go +++ b/pkg/container/vector/vector.go @@ -18,6 +18,7 @@ import ( "bytes" "fmt" "io" + "math" "math/bits" "slices" "sort" @@ -187,7 +188,10 @@ func (v *Vector) Capacity() int { // Allocated returns the total allocated memory size of the vector. // it can be used to estimate the memory usage of the vector. func (v *Vector) Allocated() int { - return cap(v.data) + cap(v.area) + return cap(v.data) + + cap(v.area) + + 8*v.nsp.GetBitmap().ExternalStorageCapacity() + + 8*v.gsp.GetBitmap().ExternalStorageCapacity() } func (v *Vector) SetLength(n int) { @@ -733,6 +737,7 @@ func (v *Vector) Free(mp *mpool.MPool) { if !v.cantFreeArea { mp.Free(v.area) } + v.freeBitmapStorage(mp) v.class = FLAT v.data = nil v.area = nil @@ -1068,14 +1073,15 @@ func validateVectorNullBitmap(data []byte, validateValues bool) error { if len(data) == 0 { return nil } - if len(data) < 24 { + if len(data) < bitmap.MarshalHeaderSize { return io.ErrUnexpectedEOF } count := types.DecodeInt64(data[:8]) bitmapLen := types.DecodeUint64(data[8:16]) bitmapDataLen := types.DecodeUint64(data[16:24]) if count < 0 || bitmapLen > uint64(1<<63-1) || uint64(count) > bitmapLen || - bitmapDataLen%8 != 0 || bitmapDataLen != uint64(len(data)-24) { + bitmapDataLen%8 != 0 || + bitmapDataLen != uint64(len(data)-bitmap.MarshalHeaderSize) { return moerr.NewInvalidInputNoCtx("invalid vector null bitmap") } if bitmapDataLen != ((bitmapLen+63)/64)*8 { @@ -1084,7 +1090,7 @@ func validateVectorNullBitmap(data []byte, validateValues bool) error { if !validateValues { return nil } - words := types.DecodeSlice[uint64](data[24:]) + words := types.DecodeSlice[uint64](data[bitmap.MarshalHeaderSize:]) actualCount := int64(0) for i, word := range words { if i == len(words)-1 && bitmapLen%64 != 0 && word>>uint(bitmapLen%64) != 0 { @@ -1145,6 +1151,9 @@ func (v *Vector) UnmarshalBinaryWithCopy(data []byte, mp *mpool.MPool) error { // read length v.length = int(types.DecodeUint32(data[:4])) data = data[4:] + if err = v.ensureBitmapCapacity(v.length, mp); err != nil { + return err + } // read data dataLen := int(types.DecodeUint32(data[:4])) @@ -1202,6 +1211,9 @@ func (v *Vector) UnmarshalWithReader(r io.Reader, mp *mpool.MPool) error { if v.length, err = types.ReadInt32AsInt(r); err != nil { return err } + if err = v.ensureBitmapCapacity(v.length, mp); err != nil { + return err + } // read data dataLen, dataBuf, err := v.readSizeBytes(r, mp, true) @@ -1221,18 +1233,9 @@ func (v *Vector) UnmarshalWithReader(r io.Reader, mp *mpool.MPool) error { v.area = areaBuf } - // read nsp, do not use mpool. nspBuf is different because - // it is not managed by vector. In the following, it will - // be unmarshalled into v.nsp - nspLen, nspBuf, err := types.ReadSizeBytes(r) - if err != nil { + if err = v.readNullsWithReader(r); err != nil { return err } - if nspLen > 0 { - v.nsp.Read(nspBuf) - } else { - v.nsp.Reset() - } v.sorted, err = types.ReadBool(r) if err != nil { @@ -1242,6 +1245,51 @@ func (v *Vector) UnmarshalWithReader(r io.Reader, mp *mpool.MPool) error { return nil } +func (v *Vector) readNullsWithReader(r io.Reader) error { + if v.allocationAccount == nil || !v.allocationAccount.accountBitmaps { + nspLen, nspBuf, err := types.ReadSizeBytes(r) + if err != nil { + return err + } + if nspLen > 0 { + return v.nsp.Read(nspBuf) + } + v.nsp.Reset() + return nil + } + + size, err := types.ReadInt32(r) + if err != nil { + return err + } + if size == 0 { + v.nsp.Reset() + return nil + } + if size < bitmap.MarshalHeaderSize { + return moerr.NewInvalidInputNoCtx("invalid bitmap wire size") + } + var header [bitmap.MarshalHeaderSize]byte + if _, err = io.ReadFull(r, header[:]); err != nil { + return err + } + _, bitLength, _, err := bitmap.DecodeMarshalHeader(header[:]) + if err != nil || bitLength > int64(v.length)+1 { + return moerr.NewInvalidInputNoCtx("invalid vector null bitmap") + } + payload, err := v.nsp.GetBitmap().PrepareExternalUnmarshal( + header[:], + int(size), + ) + if err != nil { + return err + } + if _, err = io.ReadFull(r, payload); err != nil { + v.nsp.Reset() + } + return err +} + func (v *Vector) ToConst() { v.class = CONSTANT } @@ -1256,6 +1304,12 @@ func (v *Vector) PreExtend(rows int, mp *mpool.MPool) error { return extend(v, rows, mp) } +// PreExtendBitmap ensures allocation-accounted null and grouping storage can +// represent rows without allocating vector data. Legacy vectors are unchanged. +func (v *Vector) PreExtendBitmap(rows int, mp *mpool.MPool) error { + return v.ensureBitmapCapacity(rows, mp) +} + // PreExtendArea use to expand the mpool and area of vector // extraAreaSize: the size of area to be extended // mp: mpool @@ -1331,11 +1385,25 @@ func (v *Vector) dup( } w.class = v.class w.typ = v.typ - w.length = v.length w.sorted = v.sorted - w.GetNulls().InitWith(v.GetNulls()) if v.IsConstNull() { + w.length = v.length + if v.HasGrouping() { + groupingRows := v.GetGrouping().GetBitmap().Len() + if groupingRows < 0 || groupingRows > int64(math.MaxInt) { + w.Free(mp) + return nil, mpool.ErrAllocationAccountInvalid + } + if err := w.ensureBitmapCapacity( + int(groupingRows), + mp, + ); err != nil { + w.Free(mp) + return nil, err + } + w.GetGrouping().InitWith(v.GetGrouping()) + } return w, nil } @@ -1353,6 +1421,21 @@ func (v *Vector) dup( } dataLen *= v.length } + bitmapRows := max( + v.GetNulls().GetBitmap().Len(), + v.GetGrouping().GetBitmap().Len(), + ) + if bitmapRows < 0 || bitmapRows > int64(math.MaxInt) { + w.Free(mp) + return nil, mpool.ErrAllocationAccountInvalid + } + if err := w.ensureBitmapCapacity(int(bitmapRows), mp); err != nil { + w.Free(mp) + return nil, err + } + w.length = v.length + w.GetNulls().InitWith(v.GetNulls()) + w.GetGrouping().InitWith(v.GetGrouping()) copy(w.data, v.data[:dataLen]) if len(v.area) > 0 { @@ -4591,6 +4674,9 @@ func (v *Vector) CloneWindowTo(w *Vector, start, end int, mp *mpool.MPool) error return nil } } + if err := w.PreExtendBitmap(end-start, mp); err != nil { + return err + } nulls.Range(&v.nsp, uint64(start), uint64(end), uint64(start), &w.nsp) length := (end - start) * v.typ.TypeSize() if mp == nil { diff --git a/pkg/sql/colexec/aggexec/maxby.go b/pkg/sql/colexec/aggexec/maxby.go index 0044e57cc8420..d39e17bd9cfb9 100644 --- a/pkg/sql/colexec/aggexec/maxby.go +++ b/pkg/sql/colexec/aggexec/maxby.go @@ -269,8 +269,7 @@ func compactMaxByStateVector(vec *vector.Vector, mp *mpool.MPool) error { if vec == nil || !vec.GetType().IsVarlen() { return nil } - fixedCapacity := vec.Capacity() * vec.GetType().TypeSize() - areaCapacity := vec.Allocated() - fixedCapacity + areaCapacity := cap(vec.GetArea()) if areaCapacity <= maxByVarlenaCompactionSlack { return nil } diff --git a/pkg/sql/colexec/evalExpression.go b/pkg/sql/colexec/evalExpression.go index fe79419df630b..2863b10e8c053 100644 --- a/pkg/sql/colexec/evalExpression.go +++ b/pkg/sql/colexec/evalExpression.go @@ -772,10 +772,11 @@ func (expr *FunctionExpressionExecutor) init( expr.resultVector = vector.NewFunctionResultWrapper(retType, m) return nil } - expr.resultVector, err = vector.NewFunctionResultWrapperWithAllocation( + expr.resultVector, err = vector.NewFunctionResultWrapperWithParameterAllocation( retType, m, allocation.result, + allocation.parameter, ) return err } @@ -1116,10 +1117,11 @@ func (expr *FunctionExpressionExecutor) evalSelectedRows( ) } else { expr.selectedResult, err = - vector.NewFunctionResultWrapperWithAllocation( + vector.NewFunctionResultWrapperWithParameterAllocation( expr.resultType, expr.m, expr.allocation.scratch, + expr.allocation.parameter, ) if err != nil { return nil, err diff --git a/pkg/sql/colexec/eval_expression_allocation.go b/pkg/sql/colexec/eval_expression_allocation.go index 42b958b3776f2..37f7687d83b8d 100644 --- a/pkg/sql/colexec/eval_expression_allocation.go +++ b/pkg/sql/colexec/eval_expression_allocation.go @@ -34,6 +34,13 @@ const ( ExpressionAllocationSiteScratchArea ExpressionAllocationSiteSelection ExpressionAllocationSiteSelectedRows + ExpressionAllocationSiteConstantNulls + ExpressionAllocationSiteConstantGrouping + ExpressionAllocationSiteResultNulls + ExpressionAllocationSiteResultGrouping + ExpressionAllocationSiteScratchNulls + ExpressionAllocationSiteScratchGrouping + ExpressionAllocationSiteParameterConversion ) // ExpressionAllocationAccount is the immutable allocation provenance shared @@ -43,48 +50,64 @@ type ExpressionAllocationAccount struct { account *mpool.AllocationAccount owner mpool.AllocationOwner - constant *vector.AllocationAccountSelection - result *vector.AllocationAccountSelection - scratch *vector.AllocationAccountSelection + constant *vector.AllocationAccountSelection + result *vector.AllocationAccountSelection + scratch *vector.AllocationAccountSelection + parameter *vector.FunctionParameterAllocation } func NewExpressionAllocationAccount( account *mpool.AllocationAccount, owner mpool.AllocationOwner, ) (*ExpressionAllocationAccount, error) { - constant, err := vector.NewAllocationAccountSelection( + constant, err := vector.NewAllocationAccountSelectionWithBitmaps( account, owner, ExpressionAllocationSiteConstantData, ExpressionAllocationSiteConstantArea, + ExpressionAllocationSiteConstantNulls, + ExpressionAllocationSiteConstantGrouping, ) if err != nil { return nil, err } - result, err := vector.NewAllocationAccountSelection( + result, err := vector.NewAllocationAccountSelectionWithBitmaps( account, owner, ExpressionAllocationSiteResultData, ExpressionAllocationSiteResultArea, + ExpressionAllocationSiteResultNulls, + ExpressionAllocationSiteResultGrouping, ) if err != nil { return nil, err } - scratch, err := vector.NewAllocationAccountSelection( + scratch, err := vector.NewAllocationAccountSelectionWithBitmaps( account, owner, ExpressionAllocationSiteScratchData, ExpressionAllocationSiteScratchArea, + ExpressionAllocationSiteScratchNulls, + ExpressionAllocationSiteScratchGrouping, + ) + if err != nil { + return nil, err + } + parameter, err := vector.NewFunctionParameterAllocation( + account, + owner, + ExpressionAllocationSiteParameterConversion, ) if err != nil { return nil, err } return &ExpressionAllocationAccount{ - account: account, - owner: owner, - constant: constant, - result: result, - scratch: scratch, + account: account, + owner: owner, + constant: constant, + result: result, + scratch: scratch, + parameter: parameter, }, nil } @@ -92,7 +115,8 @@ func (a *ExpressionAllocationAccount) validate() error { if a == nil || a.account == nil || a.account.Handle() == 0 || a.owner < mpool.AllocationOwnerMin || a.owner > mpool.AllocationOwnerMax || - a.constant == nil || a.result == nil || a.scratch == nil { + a.constant == nil || a.result == nil || a.scratch == nil || + a.parameter == nil { return mpool.ErrAllocationAccountInvalid } return nil diff --git a/pkg/sql/colexec/spillutil/allocation_account.go b/pkg/sql/colexec/spillutil/allocation_account.go index afa4cf9924cca..1be4e02d81776 100644 --- a/pkg/sql/colexec/spillutil/allocation_account.go +++ b/pkg/sql/colexec/spillutil/allocation_account.go @@ -35,6 +35,10 @@ const ( SpillAllocationSiteRowIDs SpillAllocationSiteMarshalBuffer SpillAllocationSiteCoalesceBuffer + SpillAllocationSiteDecodedNulls + SpillAllocationSiteDecodedGrouping + SpillAllocationSiteSelectedNulls + SpillAllocationSiteSelectedGrouping ) // SpillAllocationAccount is the dormant allocation provenance for one spill @@ -52,20 +56,24 @@ func NewSpillAllocationAccount( account *mpool.AllocationAccount, owner mpool.AllocationOwner, ) (*SpillAllocationAccount, error) { - decoded, err := vector.NewAllocationAccountSelection( + decoded, err := vector.NewAllocationAccountSelectionWithBitmaps( account, owner, SpillAllocationSiteDecodedData, SpillAllocationSiteDecodedArea, + SpillAllocationSiteDecodedNulls, + SpillAllocationSiteDecodedGrouping, ) if err != nil { return nil, err } - selected, err := vector.NewAllocationAccountSelection( + selected, err := vector.NewAllocationAccountSelectionWithBitmaps( account, owner, SpillAllocationSiteSelectedData, SpillAllocationSiteSelectedArea, + SpillAllocationSiteSelectedNulls, + SpillAllocationSiteSelectedGrouping, ) if err != nil { return nil, err diff --git a/pkg/sql/colexec/spillutil/allocation_account_test.go b/pkg/sql/colexec/spillutil/allocation_account_test.go index 7e4f5fa1432f5..cbc59c88bd7bf 100644 --- a/pkg/sql/colexec/spillutil/allocation_account_test.go +++ b/pkg/sql/colexec/spillutil/allocation_account_test.go @@ -182,7 +182,7 @@ func TestSpillAllocationAccountScatterScratchLifecycle(t *testing.T) { mpool.MustNew("spill-allocation-scatter"), ) defer proc.Free() - state := newTestSpillAllocationAccount(t, 1<<20, 8) + state := newTestSpillAllocationAccount(t, 1<<20, 64) engine, err := NewSpillEngineWithAllocation( SpillEngineConfig{}, state.allocation, diff --git a/pkg/sql/plan/function/baseTemplate.go b/pkg/sql/plan/function/baseTemplate.go index ad6f163f554cf..c56c4ece2fd38 100644 --- a/pkg/sql/plan/function/baseTemplate.go +++ b/pkg/sql/plan/function/baseTemplate.go @@ -58,8 +58,14 @@ func generalFunctionTemplateFactor[T1 templateTp1, T2 templateTr1]( return func(parameters []*vector.Vector, result vector.FunctionResultWrapper, _ *process.Process, length int, selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[T2](result) - p1 := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) - p2 := vector.OptGetParamFromWrapper[T1](rs, 1, parameters[1]) + p1, err := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) + if err != nil { + return err + } + p2, err := vector.OptGetParamFromWrapper[T1](rs, 1, parameters[1]) + if err != nil { + return err + } rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[T2](rsVec) @@ -155,8 +161,14 @@ func generalFunctionTemplateFactor[T1 templateTp1, T2 templateTr1]( return func(parameters []*vector.Vector, result vector.FunctionResultWrapper, _ *process.Process, length int, selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[T2](result) - p1 := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) - p2 := vector.OptGetParamFromWrapper[T1](rs, 1, parameters[1]) + p1, err := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) + if err != nil { + return err + } + p2, err := vector.OptGetParamFromWrapper[T1](rs, 1, parameters[1]) + if err != nil { + return err + } rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[T2](rsVec) @@ -383,8 +395,14 @@ func decimalBatchArith[TIn templateDec, TOut templateDecOut](parameters []*vecto arithFn func(v1, v2 []TIn, rs []TOut, scale1, scale2 int32, rsnull *nulls.Nulls) error, selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[TOut](result) - p1 := vector.OptGetParamFromWrapper[TIn](rs, 0, parameters[0]) - p2 := vector.OptGetParamFromWrapper[TIn](rs, 1, parameters[1]) + p1, err := vector.OptGetParamFromWrapper[TIn](rs, 0, parameters[0]) + if err != nil { + return err + } + p2, err := vector.OptGetParamFromWrapper[TIn](rs, 1, parameters[1]) + if err != nil { + return err + } rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[TOut](rsVec) scale1 := p1.GetType().Scale @@ -441,7 +459,7 @@ func decimalBatchArith[TIn templateDec, TOut templateDecOut](parameters []*vecto } else { v2 = p2.UnSafeGetAllValue() } - err := arithFn(v1, v2, rss, scale1, scale2, rsNull) + err = arithFn(v1, v2, rss, scale1, scale2, rsNull) if err != nil { if moerr.IsMoErrCode(err, moerr.ErrInvalidInput) { return moerr.NewOutOfRange(proc.Ctx, "DECIMAL", err.Error()) @@ -467,8 +485,14 @@ func opBinaryFixedFixedToFixed[ resultFn func(v1 T1, v2 T2) Tr, selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[Tr](result) - p1 := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) - p2 := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) + p1, err := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) + if err != nil { + return err + } + p2, err := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) + if err != nil { + return err + } rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[Tr](rsVec) @@ -590,8 +614,14 @@ func opBinaryFixedFixedToFixedWithErrorCheck[ resultFn func(v1 T1, v2 T2) (Tr, error), selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[Tr](result) - p1 := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) - p2 := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) + p1, err := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) + if err != nil { + return err + } + p2, err := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) + if err != nil { + return err + } rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[Tr](rsVec) @@ -740,8 +770,14 @@ func opBinaryFixedFixedToFixedWithNullOnError[ resultFn func(v1 T1, v2 T2) (Tr, error), selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[Tr](result) - p1 := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) - p2 := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) + p1, err := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) + if err != nil { + return err + } + p2, err := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) + if err != nil { + return err + } rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[Tr](rsVec) @@ -897,7 +933,10 @@ func opBinaryStrFixedToFixedWithErrorCheck[ result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[Tr](result) p1 := vector.OptGetBytesParamFromWrapper(rs, 0, parameters[0]) - p2 := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) + p2, err := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) + if err != nil { + return err + } rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[Tr](rsVec) @@ -1047,7 +1086,10 @@ func opBinaryStrFixedToStrWithErrorCheck[ result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[types.Varlena](result) p1 := vector.OptGetBytesParamFromWrapper(rs, 0, parameters[0]) - p2 := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) + p2, err := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) + if err != nil { + return err + } rsVec := rs.GetResultVector() c1, c2 := parameters[0].IsConst(), parameters[1].IsConst() @@ -1210,7 +1252,10 @@ func opBinaryFixedStrToFixedWithErrorCheck[ resultFn func(v1 T1, v2 string) (Tr, error), selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[Tr](result) - p1 := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) + p1, err := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) + if err != nil { + return err + } p2 := vector.OptGetBytesParamFromWrapper(rs, 1, parameters[1]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[Tr](rsVec) @@ -1359,8 +1404,14 @@ func specialTemplateForModFunction[ modFn func(v1, v2 T) T, selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[T](result) - p1 := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) - p2 := vector.OptGetParamFromWrapper[T](rs, 1, parameters[1]) + p1, err := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) + if err != nil { + return err + } + p2, err := vector.OptGetParamFromWrapper[T](rs, 1, parameters[1]) + if err != nil { + return err + } rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[T](rsVec) @@ -1631,8 +1682,14 @@ func specialTemplateForDivFunction[ divFn func(v1, v2 T) (T2, error), selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[T2](result) - p1 := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) - p2 := vector.OptGetParamFromWrapper[T](rs, 1, parameters[1]) + p1, err := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) + if err != nil { + return err + } + p2, err := vector.OptGetParamFromWrapper[T](rs, 1, parameters[1]) + if err != nil { + return err + } rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[T2](rsVec) @@ -2586,7 +2643,10 @@ func opUnaryFixedToFixed[ resultFn func(v T) Tr, selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(1) rs := vector.MustFunctionResult[Tr](result) - p1 := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) + p1, err := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) + if err != nil { + return err + } rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[Tr](rsVec) @@ -2997,7 +3057,10 @@ func opUnaryFixedToStr[ resultFn func(v T) string, selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(1) rs := vector.MustFunctionResult[types.Varlena](result) - p1 := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) + p1, err := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) + if err != nil { + return err + } rsVec := rs.GetResultVector() c1 := parameters[0].IsConst() @@ -3074,7 +3137,10 @@ func opUnaryFixedToStrWithNullOnError[ resultFn func(v T) (string, error), selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(1) rs := vector.MustFunctionResult[types.Varlena](result) - p1 := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) + p1, err := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) + if err != nil { + return err + } var constValue []byte constNull := false @@ -3137,7 +3203,10 @@ func opUnaryFixedToStrWithErrorCheck[ resultFn func(v T) (string, error), selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(1) rs := vector.MustFunctionResult[types.Varlena](result) - p1 := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) + p1, err := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) + if err != nil { + return err + } rsVec := rs.GetResultVector() c1 := parameters[0].IsConst() @@ -3564,7 +3633,10 @@ func opUnaryFixedToFixedWithErrorCheck[ resultFn func(v T) (Tr, error), selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(1) rs := vector.MustFunctionResult[Tr](result) - p1 := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) + p1, err := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) + if err != nil { + return err + } rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[Tr](rsVec) @@ -3604,7 +3676,6 @@ func opUnaryFixedToFixedWithErrorCheck[ } // basic case. - var err error if p1.WithAnyNullValue() || rsAnyNull { nulls.Or(rsNull, parameters[0].GetNulls(), rsNull) rowCount := uint64(length) @@ -3638,7 +3709,10 @@ func opUnaryFixedToFixedWithNullOnError[ resultFn func(v T) (Tr, error), selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(1) rs := vector.MustFunctionResult[Tr](result) - p1 := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) + p1, err := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) + if err != nil { + return err + } rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[Tr](rsVec) diff --git a/pkg/sql/plan/function/func_binary.go b/pkg/sql/plan/function/func_binary.go index eba30b6207385..6532c8146be0c 100644 --- a/pkg/sql/plan/function/func_binary.go +++ b/pkg/sql/plan/function/func_binary.go @@ -1818,8 +1818,14 @@ func DateAdd(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc * // Use custom implementation to handle maximum overflow (return NULL) result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[types.Date](result) - p1 := vector.OptGetParamFromWrapper[types.Date](rs, 0, ivecs[0]) - p2 := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) + p1, err := vector.OptGetParamFromWrapper[types.Date](rs, 0, ivecs[0]) + if err != nil { + return err + } + p2, err := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) + if err != nil { + return err + } rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[types.Date](rsVec) rsNull := rsVec.GetNulls() @@ -1863,8 +1869,14 @@ func DatetimeAdd(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pr // Use custom implementation to handle maximum overflow (return NULL) result.UseOptFunctionParamFrame(2) - p1 := vector.OptGetParamFromWrapper[types.Datetime](rs, 0, ivecs[0]) - p2 := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) + p1, err := vector.OptGetParamFromWrapper[types.Datetime](rs, 0, ivecs[0]) + if err != nil { + return err + } + p2, err := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) + if err != nil { + return err + } rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[types.Datetime](rsVec) rsNull := rsVec.GetNulls() @@ -1985,8 +1997,14 @@ func TimestampAdd(ivecs []*vector.Vector, result vector.FunctionResultWrapper, p rs.TempSetType(types.New(types.T_timestamp, 0, scale)) result.UseOptFunctionParamFrame(2) - p1 := vector.OptGetParamFromWrapper[types.Timestamp](rs, 0, ivecs[0]) - p2 := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) + p1, err := vector.OptGetParamFromWrapper[types.Timestamp](rs, 0, ivecs[0]) + if err != nil { + return err + } + p2, err := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) + if err != nil { + return err + } rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[types.Timestamp](rsVec) rsNull := rsVec.GetNulls() @@ -4446,8 +4464,14 @@ func DateSub(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc * result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[types.Date](result) - p1 := vector.OptGetParamFromWrapper[types.Date](rs, 0, ivecs[0]) - p2 := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) + p1, err := vector.OptGetParamFromWrapper[types.Date](rs, 0, ivecs[0]) + if err != nil { + return err + } + p2, err := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) + if err != nil { + return err + } rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[types.Date](rsVec) rsNull := rsVec.GetNulls() @@ -4595,8 +4619,14 @@ func DatetimeSub(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pr // Use custom implementation to handle maximum overflow (return NULL) result.UseOptFunctionParamFrame(2) - p1 := vector.OptGetParamFromWrapper[types.Datetime](rs, 0, ivecs[0]) - p2 := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) + p1, err := vector.OptGetParamFromWrapper[types.Datetime](rs, 0, ivecs[0]) + if err != nil { + return err + } + p2, err := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) + if err != nil { + return err + } rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[types.Datetime](rsVec) rsNull := rsVec.GetNulls() @@ -4718,8 +4748,14 @@ func TimestampSub(ivecs []*vector.Vector, result vector.FunctionResultWrapper, p // Use custom implementation to handle maximum overflow (return NULL) result.UseOptFunctionParamFrame(2) - p1 := vector.OptGetParamFromWrapper[types.Timestamp](rs, 0, ivecs[0]) - p2 := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) + p1, err := vector.OptGetParamFromWrapper[types.Timestamp](rs, 0, ivecs[0]) + if err != nil { + return err + } + p2, err := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) + if err != nil { + return err + } rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[types.Timestamp](rsVec) rsNull := rsVec.GetNulls() diff --git a/pkg/sql/plan/function/func_compare.go b/pkg/sql/plan/function/func_compare.go index b22d59c60d6bd..01b1b55de2c7f 100644 --- a/pkg/sql/plan/function/func_compare.go +++ b/pkg/sql/plan/function/func_compare.go @@ -104,8 +104,14 @@ func opBinaryFixedFixedToFixedNullSafe[T types.FixedSizeTExceptStrType]( ) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[bool](result) - p1 := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) - p2 := vector.OptGetParamFromWrapper[T](rs, 1, parameters[1]) + p1, err := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) + if err != nil { + return err + } + p2, err := vector.OptGetParamFromWrapper[T](rs, 1, parameters[1]) + if err != nil { + return err + } rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[bool](rsVec) From d153bd6f227c8b311c4a7b370b946d4b800013e6 Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 16:31:28 +0800 Subject: [PATCH 10/61] perf: remove row scratch from field and values --- pkg/sql/plan/function/func_binary.go | 49 +++++++++++----------------- pkg/sql/plan/function/func_unary.go | 14 ++++---- 2 files changed, 26 insertions(+), 37 deletions(-) diff --git a/pkg/sql/plan/function/func_binary.go b/pkg/sql/plan/function/func_binary.go index 6532c8146be0c..ddc06f6bb9584 100644 --- a/pkg/sql/plan/function/func_binary.go +++ b/pkg/sql/plan/function/func_binary.go @@ -4861,21 +4861,26 @@ func fieldCheck(overloads []overload, inputs []types.Type) checkResult { } func FieldNumber[T number](ivecs []*vector.Vector, result vector.FunctionResultWrapper, _ *process.Process, length int, selectList *FunctionSelectList) (err error) { + result.UseOptFunctionParamFrame(len(ivecs)) rs := vector.MustFunctionResult[uint64](result) - - fs := make([]vector.FunctionParameterWrapper[T], len(ivecs)) - for i := range ivecs { - fs[i] = vector.GenerateFunctionFixedTypeParameter[T](ivecs[i]) + first, err := vector.OptGetParamFromWrapper[T](rs, 0, ivecs[0]) + if err != nil { + return err } - nums := make([]uint64, length) + nums := vector.MustFixedColNoTypeCheck[uint64](rs.GetResultVector()) + clear(nums[:length]) for j := 1; j < len(ivecs); j++ { + candidate, err := vector.OptGetParamFromWrapper[T](rs, j, ivecs[j]) + if err != nil { + return err + } for i := uint64(0); i < uint64(length); i++ { - v1, null1 := fs[0].GetValue(i) - v2, null2 := fs[j].GetValue(i) + v1, null1 := first.GetValue(i) + v2, null2 := candidate.GetValue(i) if (nums[i] != 0) || (null1 || null2) { continue @@ -4887,32 +4892,23 @@ func FieldNumber[T number](ivecs []*vector.Vector, result vector.FunctionResultW } } - - for i := uint64(0); i < uint64(length); i++ { - if err := rs.Append(nums[i], false); err != nil { - return err - } - } - return nil } func FieldString(ivecs []*vector.Vector, result vector.FunctionResultWrapper, _ *process.Process, length int, selectList *FunctionSelectList) (err error) { + result.UseOptFunctionParamFrame(len(ivecs)) rs := vector.MustFunctionResult[uint64](result) - - fs := make([]vector.FunctionParameterWrapper[types.Varlena], len(ivecs)) - for i := range ivecs { - fs[i] = vector.GenerateFunctionStrParameter(ivecs[i]) - } - - nums := make([]uint64, length) + first := vector.OptGetBytesParamFromWrapper(rs, 0, ivecs[0]) + nums := vector.MustFixedColNoTypeCheck[uint64](rs.GetResultVector()) + clear(nums[:length]) for j := 1; j < len(ivecs); j++ { + candidate := vector.OptGetBytesParamFromWrapper(rs, j, ivecs[j]) for i := uint64(0); i < uint64(length); i++ { - v1, null1 := fs[0].GetStrValue(i) - v2, null2 := fs[j].GetStrValue(i) + v1, null1 := first.GetStrValue(i) + v2, null2 := candidate.GetStrValue(i) if (nums[i] != 0) || (null1 || null2) { continue @@ -4924,13 +4920,6 @@ func FieldString(ivecs []*vector.Vector, result vector.FunctionResultWrapper, _ } } - - for i := uint64(0); i < uint64(length); i++ { - if err := rs.Append(nums[i], false); err != nil { - return err - } - } - return nil } diff --git a/pkg/sql/plan/function/func_unary.go b/pkg/sql/plan/function/func_unary.go index f9c59f806cdbf..ff33c71066fbe 100644 --- a/pkg/sql/plan/function/func_unary.go +++ b/pkg/sql/plan/function/func_unary.go @@ -4672,13 +4672,13 @@ func Values(parameters []*vector.Vector, result vector.FunctionResultWrapper, pr toVec := result.GetResultVector() toVec.Reset(*toVec.GetType()) - sels := make([]int64, fromVec.Length()) - for j := 0; j < len(sels); j++ { - sels[j] = int64(j) - } - - err := toVec.Union(fromVec, sels, proc.GetMPool()) - return err + return toVec.UnionBatch( + fromVec, + 0, + fromVec.Length(), + nil, + proc.GetMPool(), + ) } func builtInNameConst(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { From f4ccabfd8719ea30ac10be38f92b68650d566584 Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 17:39:46 +0800 Subject: [PATCH 11/61] perf: remove row-scaled function scratch --- pkg/container/vector/functionTools.go | 54 +++ .../vector/function_result_allocation_test.go | 48 ++ pkg/sql/plan/function/func_binary.go | 387 ++++++++++------ pkg/sql/plan/function/func_binary_aes_test.go | 120 ++++- .../func_binary_array_distance_gpu_test.go | 16 +- .../func_binary_array_distance_test.go | 133 +++++- pkg/sql/plan/function/func_builtin_jq.go | 436 +++++++++--------- .../function/func_builtin_json_row_test.go | 179 +++++++ pkg/sql/plan/function/func_unary.go | 201 ++++---- pkg/vectorindex/metric/cpu.go | 18 + pkg/vectorindex/metric/gpu.go | 141 +++++- pkg/vectorindex/metric/pairwise.go | 59 ++- pkg/vectorindex/metric/pairwise_test.go | 32 ++ 13 files changed, 1324 insertions(+), 500 deletions(-) create mode 100644 pkg/sql/plan/function/func_builtin_json_row_test.go diff --git a/pkg/container/vector/functionTools.go b/pkg/container/vector/functionTools.go index dc5c48bb80337..3ae33ddfb9712 100644 --- a/pkg/container/vector/functionTools.go +++ b/pkg/container/vector/functionTools.go @@ -880,6 +880,60 @@ func (fr *FunctionResult[T]) AppendBytes(val []byte, isnull bool) error { return nil } +// AppendBytesWithFill appends one non-null varlena value and lets fill write +// directly into the result Vector's admitted backing storage. The provided +// slice is valid only during fill. A panic rolls the unpublished row and area +// length back before propagating. +func (fr *FunctionResult[T]) AppendBytesWithFill( + size int, + fill func([]byte), +) error { + if !fr.isVarlena || + fr.vec == nil || + fr.vec.IsConst() || + size < 0 || + fill == nil { + return mpool.ErrAllocationAccountInvalid + } + oldAreaLen := len(fr.vec.area) + if uint64(oldAreaLen)+uint64(size) > uint64(math.MaxUint32) { + return mpool.ErrAllocationAccountInvalid + } + areaSize := size + if size <= types.VarlenaInlineSize { + areaSize = 0 + } + if err := fr.vec.PreExtendWithArea(1, areaSize, fr.mp); err != nil { + return err + } + + index := fr.vec.length + values := toSliceOfLengthNoTypeCheck[types.Varlena](fr.vec, index+1) + oldValue := values[index] + values[index] = types.Varlena{} + var target []byte + if size <= types.VarlenaInlineSize { + values[index][0] = byte(size) + target = values[index][1 : 1+size] + } else { + fr.vec.area = fr.vec.area[:oldAreaLen+size] + values[index].SetOffsetLen(uint32(oldAreaLen), uint32(size)) + target = fr.vec.area[oldAreaLen:] + } + + published := false + defer func() { + if !published { + fr.vec.area = fr.vec.area[:oldAreaLen] + values[index] = oldValue + } + }() + fill(target) + fr.vec.length++ + published = true + return nil +} + func (fr *FunctionResult[T]) AppendByteJson(bj bytejson.ByteJson, isnull bool) error { if !fr.vec.IsConst() { return AppendByteJson(fr.vec, bj, isnull, fr.mp) diff --git a/pkg/container/vector/function_result_allocation_test.go b/pkg/container/vector/function_result_allocation_test.go index 4a32d6c478f2b..7b7aed3ebc568 100644 --- a/pkg/container/vector/function_result_allocation_test.go +++ b/pkg/container/vector/function_result_allocation_test.go @@ -138,6 +138,54 @@ func TestFunctionResultAllocationAccountFailure(t *testing.T) { finalizeTestVectorAllocationAccount(t, state) } +func TestFunctionResultAppendBytesWithFillLifecycle(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 8) + mp := mpool.MustNew("function-result-fill") + defer mpool.DeleteMPool(mp) + wrapper, err := NewFunctionResultWrapperWithAllocation( + types.T_varchar.ToType(), + mp, + state.selection, + ) + require.NoError(t, err) + require.NoError(t, wrapper.PreExtendAndReset(3)) + result := MustFunctionResult[types.Varlena](wrapper) + + require.NoError(t, result.AppendBytesWithFill(5, func(dst []byte) { + copy(dst, "small") + })) + large := make([]byte, 256) + for idx := range large { + large[idx] = byte(idx) + } + require.NoError(t, result.AppendBytesWithFill(len(large), func(dst []byte) { + copy(dst, large) + })) + require.Equal(t, []byte("small"), wrapper.GetResultVector().GetBytesAt(0)) + require.Equal(t, large, wrapper.GetResultVector().GetBytesAt(1)) + + beforeLength := wrapper.GetResultVector().Length() + beforeAreaLength := len(wrapper.GetResultVector().GetArea()) + require.Panics(t, func() { + _ = result.AppendBytesWithFill(512, func(dst []byte) { + dst[0] = 1 + panic("injected fill failure") + }) + }) + require.Equal(t, beforeLength, wrapper.GetResultVector().Length()) + require.Equal(t, beforeAreaLength, len(wrapper.GetResultVector().GetArea())) + + require.NoError(t, result.AppendBytesWithFill(4, func(dst []byte) { + copy(dst, "last") + })) + require.Equal(t, []byte("last"), wrapper.GetResultVector().GetBytesAt(2)) + require.Positive(t, state.account.Snapshot().Used) + + wrapper.Free() + require.Zero(t, state.account.Snapshot().Used) + finalizeTestVectorAllocationAccount(t, state) +} + func TestFunctionResultAllocationAccountDecimalParameterScratch(t *testing.T) { state := newTestVectorParameterAllocationAccount(t, 1<<20, 16) mp := mpool.MustNew("function-parameter-allocation") diff --git a/pkg/sql/plan/function/func_binary.go b/pkg/sql/plan/function/func_binary.go index ddc06f6bb9584..e9efab50fe848 100644 --- a/pkg/sql/plan/function/func_binary.go +++ b/pkg/sql/plan/function/func_binary.go @@ -39,6 +39,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/clusterservice" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/util" "github.com/matrixorigin/matrixone/pkg/container/nulls" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" @@ -8603,6 +8604,7 @@ func batchArrayDistanceSync[T types.RealNumbers]( length int, m metric.MetricType, proc *process.Process, + dist []float32, ) ([]float32, bool, error) { c0, c1 := ivecs[0].IsConst(), ivecs[1].IsConst() if c0 == c1 { @@ -8626,15 +8628,9 @@ func batchArrayDistanceSync[T types.RealNumbers]( if len(queryBytes) == 0 { return nil, false, nil } - x := [][]T{types.BytesToArray[T](queryBytes)} + query := types.BytesToArray[T](queryBytes) col := ivecs[colIdx] - y := make([][]T, length) - for i := range y { - y[i] = types.BytesToArray[T](col.GetBytesAt(i)) - } - - dist := make([]float32, length) // proc is non-nil under SQL execution; the nil branch keeps unit // tests (which don't synthesize a process) compiling and lets // EffectiveGpuMode fall back to the build-tag default. @@ -8643,7 +8639,17 @@ func batchArrayDistanceSync[T types.RealNumbers]( resolver = proc.GetResolveVariableFunc() } gpuMode := gpumode.EffectiveGpuMode(resolver) - handle, err := metric.PairwiseDistanceLaunch(x, y, m, dist, metric.GPUThresholdSQL, gpuMode) + handle, err := metric.PairwiseDistanceLaunchOneToMany( + query, + length, + func(row int) []T { + return types.BytesToArray[T](col.GetBytesAt(row)) + }, + m, + dist, + metric.GPUThresholdSQL, + gpuMode, + ) if err != nil { return nil, false, err } @@ -8654,15 +8660,59 @@ func batchArrayDistanceSync[T types.RealNumbers]( return dist, true, nil } +func tryBatchArrayDistance[T types.RealNumbers]( + ivecs []*vector.Vector, + result vector.FunctionResultWrapper, + proc *process.Process, + length int, + m metric.MetricType, + cosineSimilarity bool, +) (bool, error) { + rs := vector.MustFunctionResult[float64](result) + output := vector.MustFixedColNoTypeCheck[float64](rs.GetResultVector()) + if len(output) < length { + return false, moerr.NewInternalErrorNoCtx( + "array distance result is smaller than the input batch", + ) + } + + // The float64 result already owns 8*length admitted bytes. Pairwise distance + // needs 4*length temporary bytes, so use the upper half of that same backing + // store and convert forward after Wait. Each float64 write can only overwrite + // float32 values that were already read. + outputBytes := util.UnsafeSliceCast[float32](output[:length]) + distScratch := outputBytes[length:] + dist, ok, err := batchArrayDistanceSync[T]( + ivecs, + length, + m, + proc, + distScratch, + ) + if err != nil || !ok { + return ok, err + } + for idx, value := range dist { + if cosineSimilarity { + output[idx] = 1 - float64(value) + } else { + output[idx] = float64(value) + } + } + return true, nil +} + func InnerProductArray[T types.RealNumbers](ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - if dist, ok, err := batchArrayDistanceSync[T](ivecs, length, metric.Metric_InnerProduct, proc); err != nil { + if ok, err := tryBatchArrayDistance[T]( + ivecs, + result, + proc, + length, + metric.Metric_InnerProduct, + false, + ); err != nil { return err } else if ok { - rs := vector.MustFunctionResult[float64](result) - rss := vector.MustFixedColNoTypeCheck[float64](rs.GetResultVector()) - for i, d := range dist { - rss[i] = float64(d) - } return nil } return opBinaryBytesBytesToFixedWithErrorCheck[float64](ivecs, result, proc, length, func(v1, v2 []byte) (out float64, err error) { @@ -8674,14 +8724,16 @@ func InnerProductArray[T types.RealNumbers](ivecs []*vector.Vector, result vecto func CosineSimilarityArray[T types.RealNumbers](ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { // Use Metric_CosineDistance and convert: similarity = 1 - distance. - if dist, ok, err := batchArrayDistanceSync[T](ivecs, length, metric.Metric_CosineDistance, proc); err != nil { + if ok, err := tryBatchArrayDistance[T]( + ivecs, + result, + proc, + length, + metric.Metric_CosineDistance, + true, + ); err != nil { return err } else if ok { - rs := vector.MustFunctionResult[float64](result) - rss := vector.MustFixedColNoTypeCheck[float64](rs.GetResultVector()) - for i, d := range dist { - rss[i] = 1.0 - float64(d) - } return nil } return opBinaryBytesBytesToFixedWithErrorCheck[float64](ivecs, result, proc, length, func(v1, v2 []byte) (out float64, err error) { @@ -8692,14 +8744,16 @@ func CosineSimilarityArray[T types.RealNumbers](ivecs []*vector.Vector, result v } func L2DistanceArray[T types.RealNumbers](ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - if dist, ok, err := batchArrayDistanceSync[T](ivecs, length, metric.Metric_L2Distance, proc); err != nil { + if ok, err := tryBatchArrayDistance[T]( + ivecs, + result, + proc, + length, + metric.Metric_L2Distance, + false, + ); err != nil { return err } else if ok { - rs := vector.MustFunctionResult[float64](result) - rss := vector.MustFixedColNoTypeCheck[float64](rs.GetResultVector()) - for i, d := range dist { - rss[i] = float64(d) - } return nil } return opBinaryBytesBytesToFixedWithErrorCheck[float64](ivecs, result, proc, length, func(v1, v2 []byte) (out float64, err error) { @@ -12171,14 +12225,16 @@ func sameGeometryPoint(a, b geometryPoint2D) bool { } func L2DistanceSqArray[T types.RealNumbers](ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - if dist, ok, err := batchArrayDistanceSync[T](ivecs, length, metric.Metric_L2sqDistance, proc); err != nil { + if ok, err := tryBatchArrayDistance[T]( + ivecs, + result, + proc, + length, + metric.Metric_L2sqDistance, + false, + ); err != nil { return err } else if ok { - rs := vector.MustFunctionResult[float64](result) - rss := vector.MustFixedColNoTypeCheck[float64](rs.GetResultVector()) - for i, d := range dist { - rss[i] = float64(d) - } return nil } return opBinaryBytesBytesToFixedWithErrorCheck[float64](ivecs, result, proc, length, func(v1, v2 []byte) (out float64, err error) { @@ -12189,14 +12245,16 @@ func L2DistanceSqArray[T types.RealNumbers](ivecs []*vector.Vector, result vecto } func CosineDistanceArray[T types.RealNumbers](ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - if dist, ok, err := batchArrayDistanceSync[T](ivecs, length, metric.Metric_CosineDistance, proc); err != nil { + if ok, err := tryBatchArrayDistance[T]( + ivecs, + result, + proc, + length, + metric.Metric_CosineDistance, + false, + ); err != nil { return err } else if ok { - rs := vector.MustFunctionResult[float64](result) - rss := vector.MustFixedColNoTypeCheck[float64](rs.GetResultVector()) - for i, d := range dist { - rss[i] = float64(d) - } return nil } return opBinaryBytesBytesToFixedWithErrorCheck[float64](ivecs, result, proc, length, func(v1, v2 []byte) (out float64, err error) { @@ -12293,107 +12351,110 @@ func generateAESKey(key []byte, keyLen int) ([]byte, error) { return out, nil } -// pkcs7Padding adds PKCS7 padding to the data -func pkcs7Padding(data []byte, blockSize int) []byte { - padding := blockSize - len(data)%blockSize - padtext := make([]byte, padding) - for i := range padtext { - padtext[i] = byte(padding) - } - return append(data, padtext...) -} - -// pkcs7Unpadding removes PKCS7 padding from the data -func pkcs7Unpadding(data []byte) ([]byte, error) { - if len(data) == 0 { - return nil, moerr.NewInvalidInputNoCtx("invalid padding") - } - padding := int(data[len(data)-1]) - if padding > len(data) || padding == 0 { - return nil, moerr.NewInvalidInputNoCtx("invalid padding") - } - // Verify padding - for i := len(data) - padding; i < len(data); i++ { - if data[i] != byte(padding) { - return nil, moerr.NewInvalidInputNoCtx("invalid padding") - } - } - return data[:len(data)-padding], nil -} - -// encryptECB encrypts data using AES-128-ECB mode -func encryptECB(plaintext, key []byte) ([]byte, error) { - block, err := aes.NewCipher(key) - if err != nil { - return nil, err - } - - // Add PKCS7 padding - padded := pkcs7Padding(plaintext, aes.BlockSize) - - // Encrypt each block independently (ECB mode) - ciphertext := make([]byte, len(padded)) - for i := 0; i < len(padded); i += aes.BlockSize { - block.Encrypt(ciphertext[i:i+aes.BlockSize], padded[i:i+aes.BlockSize]) - } - - return ciphertext, nil -} - -// decryptECB decrypts data using AES-128-ECB mode -func decryptECB(ciphertext, key []byte) ([]byte, error) { - block, err := aes.NewCipher(key) - if err != nil { - return nil, err - } - - // Check that ciphertext length is a multiple of block size - if len(ciphertext)%aes.BlockSize != 0 { - return nil, moerr.NewInvalidInputNoCtx("invalid ciphertext length") +func aesPaddedSize(plaintextSize int) (int, error) { + padding := aes.BlockSize - plaintextSize%aes.BlockSize + if plaintextSize > int(^uint(0)>>1)-padding { + return 0, moerr.NewInvalidInputNoCtx("plaintext is too large") } - - // Decrypt each block independently (ECB mode) - plaintext := make([]byte, len(ciphertext)) - for i := 0; i < len(ciphertext); i += aes.BlockSize { - block.Decrypt(plaintext[i:i+aes.BlockSize], ciphertext[i:i+aes.BlockSize]) - } - - // Remove PKCS7 padding - return pkcs7Unpadding(plaintext) + return plaintextSize + padding, nil } -// encryptCBC encrypts data using AES-CBC mode -func encryptCBC(plaintext, key, iv []byte) ([]byte, error) { - block, err := aes.NewCipher(key) - if err != nil { - return nil, err +func encryptAESPKCS7Into( + block cipher.Block, + plaintext []byte, + iv []byte, + useCBC bool, + ciphertext []byte, +) { + fullBytes := len(plaintext) - len(plaintext)%aes.BlockSize + var finalBlock [aes.BlockSize]byte + copy(finalBlock[:], plaintext[fullBytes:]) + padding := byte(aes.BlockSize - len(plaintext)%aes.BlockSize) + for idx := len(plaintext) % aes.BlockSize; idx < aes.BlockSize; idx++ { + finalBlock[idx] = padding } - if len(iv) < aes.BlockSize { - return nil, moerr.NewInvalidInputNoCtx("invalid iv length") - } - padded := pkcs7Padding(plaintext, aes.BlockSize) - ciphertext := make([]byte, len(padded)) - mode := cipher.NewCBCEncrypter(block, iv[:aes.BlockSize]) - mode.CryptBlocks(ciphertext, padded) - return ciphertext, nil -} -// decryptCBC decrypts data using AES-CBC mode -func decryptCBC(ciphertext, key, iv []byte) ([]byte, error) { - block, err := aes.NewCipher(key) - if err != nil { - return nil, err - } - if len(iv) < aes.BlockSize { - return nil, moerr.NewInvalidInputNoCtx("invalid iv length") + if useCBC { + mode := cipher.NewCBCEncrypter(block, iv[:aes.BlockSize]) + if fullBytes > 0 { + mode.CryptBlocks(ciphertext[:fullBytes], plaintext[:fullBytes]) + } + mode.CryptBlocks(ciphertext[fullBytes:], finalBlock[:]) + return } - if len(ciphertext)%aes.BlockSize != 0 { - return nil, moerr.NewInvalidInputNoCtx("invalid ciphertext length") + for offset := 0; offset < fullBytes; offset += aes.BlockSize { + block.Encrypt( + ciphertext[offset:offset+aes.BlockSize], + plaintext[offset:offset+aes.BlockSize], + ) + } + block.Encrypt(ciphertext[fullBytes:], finalBlock[:]) +} + +func decryptAESPKCS7LastBlock( + block cipher.Block, + ciphertext []byte, + iv []byte, + useCBC bool, +) ([aes.BlockSize]byte, int, error) { + var finalBlock [aes.BlockSize]byte + if len(ciphertext) == 0 || len(ciphertext)%aes.BlockSize != 0 { + return finalBlock, 0, moerr.NewInvalidInputNoCtx( + "invalid ciphertext length", + ) + } + lastOffset := len(ciphertext) - aes.BlockSize + block.Decrypt(finalBlock[:], ciphertext[lastOffset:]) + if useCBC { + previous := iv[:aes.BlockSize] + if lastOffset > 0 { + previous = ciphertext[lastOffset-aes.BlockSize : lastOffset] + } + for idx := range finalBlock { + finalBlock[idx] ^= previous[idx] + } + } + + padding := int(finalBlock[aes.BlockSize-1]) + if padding == 0 || padding > aes.BlockSize { + return finalBlock, 0, moerr.NewInvalidInputNoCtx("invalid padding") + } + for idx := aes.BlockSize - padding; idx < aes.BlockSize; idx++ { + if finalBlock[idx] != byte(padding) { + return finalBlock, 0, moerr.NewInvalidInputNoCtx( + "invalid padding", + ) + } + } + return finalBlock, padding, nil +} + +func decryptAESPKCS7Into( + block cipher.Block, + ciphertext []byte, + iv []byte, + useCBC bool, + finalBlock [aes.BlockSize]byte, + padding int, + plaintext []byte, +) { + fullBytes := len(ciphertext) - aes.BlockSize + if useCBC { + if fullBytes > 0 { + cipher.NewCBCDecrypter( + block, + iv[:aes.BlockSize], + ).CryptBlocks(plaintext[:fullBytes], ciphertext[:fullBytes]) + } + } else { + for offset := 0; offset < fullBytes; offset += aes.BlockSize { + block.Decrypt( + plaintext[offset:offset+aes.BlockSize], + ciphertext[offset:offset+aes.BlockSize], + ) + } } - plaintext := make([]byte, len(ciphertext)) - mode := cipher.NewCBCDecrypter(block, iv[:aes.BlockSize]) - mode.CryptBlocks(plaintext, ciphertext) - return pkcs7Unpadding(plaintext) + copy(plaintext[fullBytes:], finalBlock[:aes.BlockSize-padding]) } type aesModeInfo struct { @@ -12472,14 +12533,9 @@ func AESEncrypt(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro continue } - var ciphertext []byte - var encErr error - if modeInfo.useCBC { - ciphertext, encErr = encryptCBC(str, aesKey, iv) - } else { - ciphertext, encErr = encryptECB(str, aesKey) - } - if encErr != nil { + block, blockErr := aes.NewCipher(aesKey) + ciphertextSize, sizeErr := aesPaddedSize(len(str)) + if blockErr != nil || sizeErr != nil { // On error, return NULL (MySQL behavior) if err := rs.AppendBytes(nil, true); err != nil { return err @@ -12487,7 +12543,18 @@ func AESEncrypt(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro continue } - if err := rs.AppendBytes(ciphertext, false); err != nil { + if err := rs.AppendBytesWithFill( + ciphertextSize, + func(ciphertext []byte) { + encryptAESPKCS7Into( + block, + str, + iv, + modeInfo.useCBC, + ciphertext, + ) + }, + ); err != nil { return err } } @@ -12545,14 +12612,8 @@ func AESDecrypt(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro continue } - var plaintext []byte - var decErr error - if modeInfo.useCBC { - plaintext, decErr = decryptCBC(crypt, aesKey, iv) - } else { - plaintext, decErr = decryptECB(crypt, aesKey) - } - if decErr != nil { + block, blockErr := aes.NewCipher(aesKey) + if blockErr != nil { // On error, return NULL (MySQL behavior) if err := rs.AppendBytes(nil, true); err != nil { return err @@ -12560,7 +12621,33 @@ func AESDecrypt(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro continue } - if err := rs.AppendBytes(plaintext, false); err != nil { + finalBlock, padding, decErr := decryptAESPKCS7LastBlock( + block, + crypt, + iv, + modeInfo.useCBC, + ) + if decErr != nil { + if err := rs.AppendBytes(nil, true); err != nil { + return err + } + continue + } + plaintextSize := len(crypt) - padding + if err := rs.AppendBytesWithFill( + plaintextSize, + func(plaintext []byte) { + decryptAESPKCS7Into( + block, + crypt, + iv, + modeInfo.useCBC, + finalBlock, + padding, + plaintext, + ) + }, + ); err != nil { return err } } diff --git a/pkg/sql/plan/function/func_binary_aes_test.go b/pkg/sql/plan/function/func_binary_aes_test.go index c30f901349bde..4f25c712a6e0d 100644 --- a/pkg/sql/plan/function/func_binary_aes_test.go +++ b/pkg/sql/plan/function/func_binary_aes_test.go @@ -15,10 +15,13 @@ package function import ( + "crypto/aes" "fmt" + "strings" "testing" "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/testutil" "github.com/matrixorigin/matrixone/pkg/vm/process" "github.com/stretchr/testify/require" @@ -42,8 +45,12 @@ func TestAESEncryptDecryptECB(t *testing.T) { aesKey, err := generateAESKey([]byte(key), 16) require.NoError(t, err) - ciphertext, err := encryptECB([]byte(plain), aesKey) + block, err := aes.NewCipher(aesKey) require.NoError(t, err) + ciphertextSize, err := aesPaddedSize(len(plain)) + require.NoError(t, err) + ciphertext := make([]byte, ciphertextSize) + encryptAESPKCS7Into(block, []byte(plain), nil, false, ciphertext) encryptCase := NewFunctionTestCase(proc, []FunctionTestInput{ @@ -76,8 +83,18 @@ func TestAESEncryptDecryptCBC(t *testing.T) { aesKey, err := generateAESKey([]byte(key), 32) require.NoError(t, err) - ciphertext, err := encryptCBC([]byte(plain), aesKey, []byte(iv)) + block, err := aes.NewCipher(aesKey) + require.NoError(t, err) + ciphertextSize, err := aesPaddedSize(len(plain)) require.NoError(t, err) + ciphertext := make([]byte, ciphertextSize) + encryptAESPKCS7Into( + block, + []byte(plain), + []byte(iv), + true, + ciphertext, + ) encryptCase := NewFunctionTestCase(proc, []FunctionTestInput{ @@ -104,6 +121,105 @@ func TestAESEncryptDecryptCBC(t *testing.T) { require.True(t, ok, fmt.Sprintf("decrypt cbc failed: %s", info)) } +func TestAESEncryptDecryptBlockBoundaries(t *testing.T) { + plaintexts := []string{ + "", + strings.Repeat("a", aes.BlockSize-1), + strings.Repeat("b", aes.BlockSize), + strings.Repeat("c", aes.BlockSize+1), + strings.Repeat("d", 2*aes.BlockSize), + } + for _, tc := range []struct { + name string + mode string + key string + iv string + }{ + {name: "ecb", mode: "aes-128-ecb", key: "boundary-key"}, + {name: "cbc", mode: "aes-256-cbc", key: "boundary-key", iv: "0123456789abcdef"}, + } { + t.Run(tc.name, func(t *testing.T) { + proc := newAESProcess(t, tc.mode) + mp := proc.Mp() + plainVec := newVectorByType( + mp, + types.T_varchar.ToType(), + plaintexts, + nil, + ) + defer plainVec.Free(mp) + keys := make([]string, len(plaintexts)) + for idx := range keys { + keys[idx] = tc.key + } + keyVec := newVectorByType( + mp, + types.T_varchar.ToType(), + keys, + nil, + ) + defer keyVec.Free(mp) + params := []*vector.Vector{plainVec, keyVec} + if tc.iv != "" { + ivs := make([]string, len(plaintexts)) + for idx := range ivs { + ivs[idx] = tc.iv + } + ivVec := newVectorByType( + mp, + types.T_varchar.ToType(), + ivs, + nil, + ) + defer ivVec.Free(mp) + params = append(params, ivVec) + } + + encrypted := vector.NewFunctionResultWrapper( + types.T_blob.ToType(), + mp, + ) + defer encrypted.Free() + require.NoError(t, encrypted.PreExtendAndReset(len(plaintexts))) + require.NoError(t, AESEncrypt( + params, + encrypted, + proc, + len(plaintexts), + nil, + )) + + decryptParams := []*vector.Vector{ + encrypted.GetResultVector(), + keyVec, + } + if len(params) == 3 { + decryptParams = append(decryptParams, params[2]) + } + decrypted := vector.NewFunctionResultWrapper( + types.T_varchar.ToType(), + mp, + ) + defer decrypted.Free() + require.NoError(t, decrypted.PreExtendAndReset(len(plaintexts))) + require.NoError(t, AESDecrypt( + decryptParams, + decrypted, + proc, + len(plaintexts), + nil, + )) + for idx, plaintext := range plaintexts { + require.Equal( + t, + []byte(plaintext), + decrypted.GetResultVector().GetBytesAt(idx), + ) + } + }) + } +} + func TestAESEncryptCBCMissingIV(t *testing.T) { proc := newAESProcess(t, "aes-256-cbc") plain := "missing iv" diff --git a/pkg/sql/plan/function/func_binary_array_distance_gpu_test.go b/pkg/sql/plan/function/func_binary_array_distance_gpu_test.go index 0b9aee421d619..c019e8619f124 100644 --- a/pkg/sql/plan/function/func_binary_array_distance_gpu_test.go +++ b/pkg/sql/plan/function/func_binary_array_distance_gpu_test.go @@ -60,8 +60,8 @@ func TestBatchArrayDistanceSync_GPU_L2sq(t *testing.T) { constVec := makeConstArrayVec[float32](t, mp, query, N) colVec := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - gpuDist, ok, err := batchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, N, metric.Metric_L2sqDistance, nil) + gpuDist, ok, err := testBatchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, N, metric.Metric_L2sqDistance) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(gpuDist)) @@ -102,8 +102,8 @@ func TestBatchArrayDistanceSync_GPU_InnerProduct(t *testing.T) { constVec := makeConstArrayVec[float32](t, mp, query, N) colVec := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - gpuDist, ok, err := batchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, N, metric.Metric_InnerProduct, nil) + gpuDist, ok, err := testBatchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, N, metric.Metric_InnerProduct) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(gpuDist)) @@ -143,8 +143,8 @@ func TestBatchArrayDistanceSync_GPU_CosineDistance(t *testing.T) { constVec := makeConstArrayVec[float32](t, mp, query, N) colVec := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - gpuDist, ok, err := batchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, N, metric.Metric_CosineDistance, nil) + gpuDist, ok, err := testBatchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, N, metric.Metric_CosineDistance) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(gpuDist)) @@ -184,8 +184,8 @@ func TestBatchArrayDistanceSync_GPU_L2Distance(t *testing.T) { constVec := makeConstArrayVec[float32](t, mp, query, N) colVec := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - gpuDist, ok, err := batchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, N, metric.Metric_L2Distance, nil) + gpuDist, ok, err := testBatchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, N, metric.Metric_L2Distance) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(gpuDist)) diff --git a/pkg/sql/plan/function/func_binary_array_distance_test.go b/pkg/sql/plan/function/func_binary_array_distance_test.go index 8d8b1f6283e7a..ea0fc73128ef1 100644 --- a/pkg/sql/plan/function/func_binary_array_distance_test.go +++ b/pkg/sql/plan/function/func_binary_array_distance_test.go @@ -26,7 +26,7 @@ import ( ) // makeConstArrayVec creates a constant vector holding a single array value repeated length times. -func makeConstArrayVec[T types.RealNumbers](t *testing.T, mp *mpool.MPool, arr []T, length int) *vector.Vector { +func makeConstArrayVec[T types.RealNumbers](t testing.TB, mp *mpool.MPool, arr []T, length int) *vector.Vector { t.Helper() b := types.ArrayToBytes[T](arr) v, err := vector.NewConstBytes(types.T_array_float32.ToType(), b, length, mp) @@ -35,7 +35,7 @@ func makeConstArrayVec[T types.RealNumbers](t *testing.T, mp *mpool.MPool, arr [ } // makeConstArrayVec64 is the float64 variant. -func makeConstArrayVec64(t *testing.T, mp *mpool.MPool, arr []float64, length int) *vector.Vector { +func makeConstArrayVec64(t testing.TB, mp *mpool.MPool, arr []float64, length int) *vector.Vector { t.Helper() b := types.ArrayToBytes[float64](arr) v, err := vector.NewConstBytes(types.T_array_float64.ToType(), b, length, mp) @@ -44,7 +44,7 @@ func makeConstArrayVec64(t *testing.T, mp *mpool.MPool, arr []float64, length in } // makeColArrayVec creates a column vector holding one array per row. -func makeColArrayVec[T types.RealNumbers](t *testing.T, mp *mpool.MPool, typ types.Type, rows [][]T) *vector.Vector { +func makeColArrayVec[T types.RealNumbers](t testing.TB, mp *mpool.MPool, typ types.Type, rows [][]T) *vector.Vector { t.Helper() v := vector.NewVec(typ) for _, row := range rows { @@ -66,6 +66,56 @@ func approxEqF32(a, b float32) bool { return diff/avg < 1e-4 } +func testBatchArrayDistanceSync[T types.RealNumbers]( + ivecs []*vector.Vector, + length int, + m metric.MetricType, +) ([]float32, bool, error) { + return batchArrayDistanceSync[T]( + ivecs, + length, + m, + nil, + make([]float32, length), + ) +} + +func BenchmarkBatchArrayDistanceSync8192(b *testing.B) { + mp := mpool.MustNewZero() + defer mpool.DeleteMPool(mp) + + const dimension = 128 + query := make([]float32, dimension) + rows := make([][]float32, 8192) + for row := range rows { + rows[row] = make([]float32, dimension) + for col := range rows[row] { + rows[row][col] = float32((row + col) % 17) + } + } + constVec := makeConstArrayVec[float32](b, mp, query, len(rows)) + defer constVec.Free(mp) + colVec := makeColArrayVec[float32](b, mp, types.T_array_float32.ToType(), rows) + defer colVec.Free(mp) + inputs := []*vector.Vector{constVec, colVec} + distScratch := make([]float32, len(rows)) + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + dist, ok, err := batchArrayDistanceSync[float32]( + inputs, + len(rows), + metric.Metric_L2sqDistance, + nil, + distScratch, + ) + require.NoError(b, err) + require.True(b, ok) + require.Len(b, dist, len(rows)) + } +} + // TestBatchArrayDistanceSync_L2Sq verifies batchArrayDistanceSync with Metric_L2sqDistance // on a small const-vs-column input (always CPU path). func TestBatchArrayDistanceSync_L2Sq(t *testing.T) { @@ -81,8 +131,8 @@ func TestBatchArrayDistanceSync_L2Sq(t *testing.T) { constVec := makeConstArrayVec[float32](t, mp, query, N) colVec := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - dist, ok, err := batchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, N, metric.Metric_L2sqDistance, nil) + dist, ok, err := testBatchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, N, metric.Metric_L2sqDistance) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(dist)) @@ -91,6 +141,43 @@ func TestBatchArrayDistanceSync_L2Sq(t *testing.T) { } } +func TestBatchArrayDistanceResultScratchAlias(t *testing.T) { + mp := mpool.MustNewZero() + defer mpool.DeleteMPool(mp) + + query := []float32{1, 0, 0} + rows := [][]float32{ + {1, 0, 0}, + {0, 1, 0}, + {0, 0, 1}, + } + constVec := makeConstArrayVec[float32](t, mp, query, len(rows)) + defer constVec.Free(mp) + colVec := makeColArrayVec[float32]( + t, + mp, + types.T_array_float32.ToType(), + rows, + ) + defer colVec.Free(mp) + + result := vector.NewFunctionResultWrapper(types.T_float64.ToType(), mp) + defer result.Free() + require.NoError(t, result.PreExtendAndReset(len(rows))) + require.NoError(t, L2DistanceSqArray[float32]( + []*vector.Vector{constVec, colVec}, + result, + nil, + len(rows), + nil, + )) + require.Equal( + t, + []float64{0, 2, 2}, + vector.MustFixedColNoTypeCheck[float64](result.GetResultVector()), + ) +} + // TestBatchArrayDistanceSync_L2 verifies batchArrayDistanceSync with Metric_L2Distance. func TestBatchArrayDistanceSync_L2(t *testing.T) { mp := mpool.MustNewZero() @@ -105,8 +192,8 @@ func TestBatchArrayDistanceSync_L2(t *testing.T) { constVec := makeConstArrayVec[float32](t, mp, query, N) colVec := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - dist, ok, err := batchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, N, metric.Metric_L2Distance, nil) + dist, ok, err := testBatchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, N, metric.Metric_L2Distance) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(dist)) @@ -130,8 +217,8 @@ func TestBatchArrayDistanceSync_InnerProduct(t *testing.T) { constVec := makeConstArrayVec[float32](t, mp, query, N) colVec := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - dist, ok, err := batchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, N, metric.Metric_InnerProduct, nil) + dist, ok, err := testBatchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, N, metric.Metric_InnerProduct) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(dist)) @@ -155,8 +242,8 @@ func TestBatchArrayDistanceSync_CosineDistance(t *testing.T) { constVec := makeConstArrayVec[float32](t, mp, query, N) colVec := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - dist, ok, err := batchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, N, metric.Metric_CosineDistance, nil) + dist, ok, err := testBatchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, N, metric.Metric_CosineDistance) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(dist)) @@ -180,8 +267,8 @@ func TestBatchArrayDistanceSync_QueryAsSecondArg(t *testing.T) { constVec := makeConstArrayVec[float32](t, mp, query, N) // Note: const is ivecs[1], column is ivecs[0] - dist, ok, err := batchArrayDistanceSync[float32]( - []*vector.Vector{colVec, constVec}, N, metric.Metric_L2sqDistance, nil) + dist, ok, err := testBatchArrayDistanceSync[float32]( + []*vector.Vector{colVec, constVec}, N, metric.Metric_L2sqDistance) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(dist)) @@ -203,8 +290,8 @@ func TestBatchArrayDistanceSync_Float64(t *testing.T) { constVec := makeConstArrayVec64(t, mp, query, N) colVec := makeColArrayVec[float64](t, mp, types.T_array_float64.ToType(), rows) - dist, ok, err := batchArrayDistanceSync[float64]( - []*vector.Vector{constVec, colVec}, N, metric.Metric_L2sqDistance, nil) + dist, ok, err := testBatchArrayDistanceSync[float64]( + []*vector.Vector{constVec, colVec}, N, metric.Metric_L2sqDistance) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(dist)) @@ -224,8 +311,8 @@ func TestBatchArrayDistanceSync_BothConst(t *testing.T) { v1, err := vector.NewConstBytes(types.T_array_float32.ToType(), b, 4, mp) require.NoError(t, err) - _, ok, err := batchArrayDistanceSync[float32]( - []*vector.Vector{v0, v1}, 4, metric.Metric_L2sqDistance, nil) + _, ok, err := testBatchArrayDistanceSync[float32]( + []*vector.Vector{v0, v1}, 4, metric.Metric_L2sqDistance) require.NoError(t, err) require.False(t, ok, "both-const should return ok=false") } @@ -239,8 +326,8 @@ func TestBatchArrayDistanceSync_BothCol(t *testing.T) { v0 := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) v1 := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - _, ok, err := batchArrayDistanceSync[float32]( - []*vector.Vector{v0, v1}, 2, metric.Metric_L2sqDistance, nil) + _, ok, err := testBatchArrayDistanceSync[float32]( + []*vector.Vector{v0, v1}, 2, metric.Metric_L2sqDistance) require.NoError(t, err) require.False(t, ok, "col-vs-col should return ok=false") } @@ -254,8 +341,8 @@ func TestBatchArrayDistanceSync_NullConst(t *testing.T) { rows := [][]float32{{1, 0, 0}, {0, 1, 0}, {0, 0, 1}, {1, 1, 0}} colVec := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - _, ok, err := batchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, 4, metric.Metric_L2sqDistance, nil) + _, ok, err := testBatchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, 4, metric.Metric_L2sqDistance) require.NoError(t, err) require.False(t, ok, "null const should return ok=false") } @@ -274,8 +361,8 @@ func TestBatchArrayDistanceSync_NullInColumn(t *testing.T) { require.NoError(t, vector.AppendBytes(colVec, nil, true, mp)) // null row require.NoError(t, vector.AppendBytes(colVec, types.ArrayToBytes[float32]([]float32{0, 1, 0}), false, mp)) - _, ok, err := batchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, 3, metric.Metric_L2sqDistance, nil) + _, ok, err := testBatchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, 3, metric.Metric_L2sqDistance) require.NoError(t, err) require.False(t, ok, "column with nulls should return ok=false") } diff --git a/pkg/sql/plan/function/func_builtin_jq.go b/pkg/sql/plan/function/func_builtin_jq.go index ec33dbe8f0442..5d1b5a2481e03 100644 --- a/pkg/sql/plan/function/func_builtin_jq.go +++ b/pkg/sql/plan/function/func_builtin_jq.go @@ -267,7 +267,7 @@ func (op *opBuiltInJq) getJqCode(jq string) (*gojq.Code, error) { // We removed all the terminal color related code and we write to buffer w // and do not flush until the encoding is done. type JqEncoder struct { - w *bytes.Buffer + w bytes.Buffer tab bool indent int depth int @@ -275,7 +275,7 @@ type JqEncoder struct { } func (e *JqEncoder) intialize(tab bool, indent int) { - e.w = new(bytes.Buffer) + e.w.Reset() e.tab = tab e.indent = indent } @@ -490,281 +490,277 @@ func (e *JqEncoder) writeIndentInternal(n int, spaces string) { } type opBuiltInJsonRow struct { - enc []JqEncoder + enc JqEncoder + columns []jsonRowColumnEncoder } func newOpBuiltInJsonRow() *opBuiltInJsonRow { var op opBuiltInJsonRow + op.enc.intialize(false, 0) return &op } -func (op *opBuiltInJsonRow) grow(length int) { - if len(op.enc) == 0 { - op.enc = make([]JqEncoder, length) - for i := 0; i < length; i++ { - op.enc[i].intialize(false, 0) - } - } else if length > len(op.enc) { - for i := len(op.enc); i < length; i++ { - op.enc = append(op.enc, JqEncoder{}) - op.enc[i].intialize(false, 0) - } - } -} - func (op *opBuiltInJsonRow) jsonRow(params []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - op.grow(length) rs := vector.MustFunctionResult[types.Varlena](result) - ulen := uint64(length) - - for j := 0; j < length; j++ { - op.enc[j].w.WriteByte('[') + if cap(op.columns) < len(params) { + op.columns = make([]jsonRowColumnEncoder, len(params)) + } else { + op.columns = op.columns[:len(params)] + } + for idx, param := range params { + column, err := prepareJSONRowColumn(param, proc) + if err != nil { + clear(op.columns) + return err + } + op.columns[idx] = column } + defer clear(op.columns) - for i := 0; i < len(params); i++ { - // write separator first - if i > 0 { - for j := 0; j < length; j++ { - op.enc[j].w.WriteByte(',') + op.enc.done() + defer op.enc.done() + for row := uint64(0); row < uint64(length); row++ { + if selectList.Contains(row) { + if err := rs.AppendBytes(nil, true); err != nil { + return err } + continue } - - // oh the dreaded type switch - fromType := params[i].GetType() - switch fromType.Oid { - case types.T_any: // scalar null - op.encodeScalarNull(ulen) - case types.T_bool: - op.encodeBool(params[i], ulen) - case types.T_int8: - encodeInt[int8](op, params[i], ulen) - case types.T_int16: - encodeInt[int16](op, params[i], ulen) - case types.T_int32: - encodeInt[int32](op, params[i], ulen) - case types.T_int64: - encodeInt[int64](op, params[i], ulen) - case types.T_uint8: - encodeInt[uint8](op, params[i], ulen) - case types.T_uint16: - encodeInt[uint16](op, params[i], ulen) - case types.T_uint32: - encodeInt[uint32](op, params[i], ulen) - case types.T_uint64: - encodeInt[uint64](op, params[i], ulen) - case types.T_float32: - encodeFloat[float32](op, params[i], ulen) - case types.T_float64: - encodeFloat[float64](op, params[i], ulen) - case types.T_decimal64: - encodeDecimal[types.Decimal64](op, params[i], ulen) - case types.T_decimal128: - encodeDecimal[types.Decimal128](op, params[i], ulen) - case types.T_date: - encodeFixedStringer[types.Date](op, params[i], ulen) - case types.T_time: - encodeFixedStringer[types.Time](op, params[i], ulen) - case types.T_datetime: - encodeFixedStringer[types.Datetime](op, params[i], ulen) - case types.T_timestamp: - encodeFixedStringer[types.Timestamp](op, params[i], ulen) - case types.T_char, types.T_varchar, types.T_text: - encodeString(op, params[i], ulen) - case types.T_binary, types.T_varbinary, types.T_blob: - // well, in cast, we handle binary as if they are string. - // However it id deemed too dangerous to do so in json_row. - return moerr.NewInvalidInputf(proc.Ctx, "binary data not supported json_row: %v", fromType.String()) - case types.T_array_float32: - // vector of float, we will encode them as json array - encodeFloatArray[float32](op, params[i], ulen) - case types.T_array_float64: - // vector of float, we will encode them as json array - encodeFloatArray[float64](op, params[i], ulen) - case types.T_array_bf16: - encodeNarrowArray[types.BF16](op, params[i], ulen, - func(x types.BF16) float64 { return float64(x.ToFloat32()) }) - case types.T_array_float16: - encodeNarrowArray[types.Float16](op, params[i], ulen, - func(x types.Float16) float64 { return float64(x.ToFloat32()) }) - case types.T_array_int8: - encodeNarrowArray[int8](op, params[i], ulen, - func(x int8) float64 { return float64(x) }) - case types.T_array_uint8: - encodeNarrowArray[uint8](op, params[i], ulen, - func(x uint8) float64 { return float64(x) }) - case types.T_uuid: - encodeFixedStringer[types.Uuid](op, params[i], ulen) - case types.T_json: - if err := encodeJson(op, params[i], ulen); err != nil { + op.enc.w.WriteByte('[') + for paramIdx := range op.columns { + if paramIdx > 0 { + op.enc.w.WriteByte(',') + } + if err := op.columns[paramIdx](&op.enc, row); err != nil { return err } - default: - return moerr.NewInvalidInputf(proc.Ctx, "unsupported type for json_row: %v", fromType.String()) } - } - - for j := 0; j < length; j++ { - op.enc[j].w.WriteByte(']') - if selectList.Contains(uint64(j)) { - rs.AppendBytes(nil, true) - } else { - rs.AppendBytes(op.enc[j].bytes(), false) + op.enc.w.WriteByte(']') + if err := rs.AppendBytes(op.enc.bytes(), false); err != nil { + return err } - op.enc[j].done() + op.enc.done() } return nil } -func (op *opBuiltInJsonRow) encodeScalarNull(length uint64) { - for i := uint64(0); i < length; i++ { - op.enc[i].w.WriteString("null") - } +type jsonRowColumnEncoder func(*JqEncoder, uint64) error + +func encodeJSONRowNull(e *JqEncoder, _ uint64) error { + e.w.WriteString("null") + return nil } -func (op *opBuiltInJsonRow) encodeBool(v *vector.Vector, length uint64) { - p := vector.GenerateFunctionFixedTypeParameter[bool](v) - for i := uint64(0); i < length; i++ { - v, null := p.GetValue(i) - if null { - op.enc[i].w.WriteString("null") - } else { - if v { - op.enc[i].w.WriteString("true") +func prepareJSONRowColumn( + v *vector.Vector, + proc *process.Process, +) (jsonRowColumnEncoder, error) { + switch fromType := v.GetType(); fromType.Oid { + case types.T_any: + return encodeJSONRowNull, nil + case types.T_bool: + param := vector.GenerateFunctionFixedTypeParameter[bool](v) + return func(e *JqEncoder, row uint64) error { + value, isNull := param.GetValue(row) + if isNull { + return encodeJSONRowNull(e, row) + } + if value { + e.w.WriteString("true") } else { - op.enc[i].w.WriteString("false") + e.w.WriteString("false") } - } + return nil + }, nil + case types.T_int8: + return prepareJSONRowSignedColumn[int8](v), nil + case types.T_int16: + return prepareJSONRowSignedColumn[int16](v), nil + case types.T_int32: + return prepareJSONRowSignedColumn[int32](v), nil + case types.T_int64: + return prepareJSONRowSignedColumn[int64](v), nil + case types.T_uint8: + return prepareJSONRowUnsignedColumn[uint8](v), nil + case types.T_uint16: + return prepareJSONRowUnsignedColumn[uint16](v), nil + case types.T_uint32: + return prepareJSONRowUnsignedColumn[uint32](v), nil + case types.T_uint64: + return prepareJSONRowUnsignedColumn[uint64](v), nil + case types.T_float32: + return prepareJSONRowFloatColumn[float32](v), nil + case types.T_float64: + return prepareJSONRowFloatColumn[float64](v), nil + case types.T_decimal64: + return prepareJSONRowDecimalColumn[types.Decimal64](v), nil + case types.T_decimal128: + return prepareJSONRowDecimalColumn[types.Decimal128](v), nil + case types.T_date: + return prepareJSONRowStringerColumn[types.Date](v), nil + case types.T_time: + return prepareJSONRowStringerColumn[types.Time](v), nil + case types.T_datetime: + return prepareJSONRowStringerColumn[types.Datetime](v), nil + case types.T_timestamp: + return prepareJSONRowStringerColumn[types.Timestamp](v), nil + case types.T_char, types.T_varchar, types.T_text: + return prepareJSONRowStringColumn(v), nil + case types.T_array_float32: + return prepareJSONRowArrayColumn(v, + func(value float32) float64 { return float64(value) }), nil + case types.T_array_float64: + return prepareJSONRowArrayColumn(v, + func(value float64) float64 { return value }), nil + case types.T_array_bf16: + return prepareJSONRowArrayColumn(v, + func(value types.BF16) float64 { return float64(value.ToFloat32()) }), nil + case types.T_array_float16: + return prepareJSONRowArrayColumn(v, + func(value types.Float16) float64 { return float64(value.ToFloat32()) }), nil + case types.T_array_int8: + return prepareJSONRowArrayColumn(v, + func(value int8) float64 { return float64(value) }), nil + case types.T_array_uint8: + return prepareJSONRowArrayColumn(v, + func(value uint8) float64 { return float64(value) }), nil + case types.T_uuid: + return prepareJSONRowStringerColumn[types.Uuid](v), nil + case types.T_json: + param := vector.GenerateFunctionStrParameter(v) + return func(e *JqEncoder, row uint64) error { + value, isNull := param.GetStrValue(row) + if isNull { + return encodeJSONRowNull(e, row) + } + jsonValue, err := types.DecodeJson(value).MarshalJSON() + if err != nil { + return err + } + e.w.Write(jsonValue) + return nil + }, nil + case types.T_binary, types.T_varbinary, types.T_blob: + return nil, moerr.NewInvalidInputf(proc.Ctx, + "binary data not supported json_row: %v", + fromType.String()) + default: + return nil, moerr.NewInvalidInputf(proc.Ctx, + "unsupported type for json_row: %v", + fromType.String()) } } -func encodeInt[T constraints.Integer](op *opBuiltInJsonRow, v *vector.Vector, length uint64) { - p := vector.GenerateFunctionFixedTypeParameter[T](v) - for i := uint64(0); i < length; i++ { - v, null := p.GetValue(i) - if null { - op.enc[i].w.WriteString("null") - } else { - op.enc[i].w.Write(strconv.AppendInt(op.enc[i].buf[:0], int64(v), 10)) +func prepareJSONRowSignedColumn[T constraints.Signed]( + v *vector.Vector, +) jsonRowColumnEncoder { + param := vector.GenerateFunctionFixedTypeParameter[T](v) + return func(e *JqEncoder, row uint64) error { + value, isNull := param.GetValue(row) + if isNull { + return encodeJSONRowNull(e, row) } + e.w.Write(strconv.AppendInt(e.buf[:0], int64(value), 10)) + return nil } } -func encodeFloat[T constraints.Float](op *opBuiltInJsonRow, v *vector.Vector, length uint64) { - p := vector.GenerateFunctionFixedTypeParameter[T](v) - for i := uint64(0); i < length; i++ { - v, null := p.GetValue(i) - if null { - op.enc[i].w.WriteString("null") - } else { - op.enc[i].encodeFloat64(float64(v)) +func prepareJSONRowUnsignedColumn[T constraints.Unsigned]( + v *vector.Vector, +) jsonRowColumnEncoder { + param := vector.GenerateFunctionFixedTypeParameter[T](v) + return func(e *JqEncoder, row uint64) error { + value, isNull := param.GetValue(row) + if isNull { + return encodeJSONRowNull(e, row) } + e.w.Write(strconv.AppendUint(e.buf[:0], uint64(value), 10)) + return nil } } -func encodeDecimal[T types.DecimalWithFormat](op *opBuiltInJsonRow, v *vector.Vector, length uint64) { - p := vector.GenerateFunctionFixedTypeParameter[T](v) - fromTyp := v.GetType() - for i := uint64(0); i < length; i++ { - v, null := p.GetValue(i) - if null { - op.enc[i].w.WriteString("null") - } else { - bs := []byte(v.Format(fromTyp.Scale)) - op.enc[i].w.Write(bs) +func prepareJSONRowFloatColumn[T constraints.Float]( + v *vector.Vector, +) jsonRowColumnEncoder { + param := vector.GenerateFunctionFixedTypeParameter[T](v) + return func(e *JqEncoder, row uint64) error { + value, isNull := param.GetValue(row) + if isNull { + return encodeJSONRowNull(e, row) } + e.encodeFloat64(float64(value)) + return nil } } -func encodeFixedStringer[T types.FixedWithStringer](op *opBuiltInJsonRow, v *vector.Vector, length uint64) { - p := vector.GenerateFunctionFixedTypeParameter[T](v) - for i := uint64(0); i < length; i++ { - v, null := p.GetValue(i) - if null { - op.enc[i].w.WriteString("null") - } else { - op.enc[i].encodeString(v.String()) +func prepareJSONRowDecimalColumn[T types.DecimalWithFormat]( + v *vector.Vector, +) jsonRowColumnEncoder { + param := vector.GenerateFunctionFixedTypeParameter[T](v) + scale := v.GetType().Scale + return func(e *JqEncoder, row uint64) error { + value, isNull := param.GetValue(row) + if isNull { + return encodeJSONRowNull(e, row) } + e.w.WriteString(value.Format(scale)) + return nil } } -func encodeString(op *opBuiltInJsonRow, v *vector.Vector, length uint64) { - p := vector.GenerateFunctionStrParameter(v) - for i := uint64(0); i < length; i++ { - v, null := p.GetStrValue(i) - if null { - op.enc[i].w.WriteString("null") - } else { - op.enc[i].encodeString(string(v)) +func prepareJSONRowStringerColumn[T types.FixedWithStringer]( + v *vector.Vector, +) jsonRowColumnEncoder { + param := vector.GenerateFunctionFixedTypeParameter[T](v) + return func(e *JqEncoder, row uint64) error { + value, isNull := param.GetValue(row) + if isNull { + return encodeJSONRowNull(e, row) } + e.encodeString(value.String()) + return nil } } -func encodeFloatArray[T constraints.Float](op *opBuiltInJsonRow, v *vector.Vector, length uint64) { - // GenStrParam: array is varlena also. - p := vector.GenerateFunctionStrParameter(v) - for i := uint64(0); i < length; i++ { - v, null := p.GetStrValue(i) - if null { - op.enc[i].w.WriteString("null") - } else { - vv := types.BytesToArray[T](v) - op.enc[i].w.WriteByte('[') - for j, val := range vv { - if j > 0 { - op.enc[i].w.WriteByte(',') - } - ff := float64(val) - op.enc[i].encodeFloat64(ff) - } - op.enc[i].w.WriteByte(']') +func prepareJSONRowStringColumn(v *vector.Vector) jsonRowColumnEncoder { + param := vector.GenerateFunctionStrParameter(v) + return func(e *JqEncoder, row uint64) error { + value, isNull := param.GetStrValue(row) + if isNull { + return encodeJSONRowNull(e, row) } + e.encodeString(string(value)) + return nil } } -// encodeNarrowArray is encodeFloatArray for element types that are not -// constraints.Float: BF16/Float16 need ToFloat32(), int8/uint8 are plain -// integers. The JSON shape emitted is identical to the f32/f64 arrays. -func encodeNarrowArray[T types.ArrayElement](op *opBuiltInJsonRow, v *vector.Vector, length uint64, toF64 func(T) float64) { - p := vector.GenerateFunctionStrParameter(v) - for i := uint64(0); i < length; i++ { - v, null := p.GetStrValue(i) - if null { - op.enc[i].w.WriteString("null") - } else { - vv := types.BytesToArray[T](v) - op.enc[i].w.WriteByte('[') - for j, val := range vv { - if j > 0 { - op.enc[i].w.WriteByte(',') - } - op.enc[i].encodeFloat64(toF64(val)) - } - op.enc[i].w.WriteByte(']') +func prepareJSONRowArrayColumn[T types.ArrayElement]( + v *vector.Vector, + toFloat64 func(T) float64, +) jsonRowColumnEncoder { + param := vector.GenerateFunctionStrParameter(v) + return func(e *JqEncoder, row uint64) error { + value, isNull := param.GetStrValue(row) + if isNull { + return encodeJSONRowNull(e, row) } + encodeJSONRowArray(e, types.BytesToArray[T](value), toFloat64) + return nil } } -func encodeJson(op *opBuiltInJsonRow, v *vector.Vector, length uint64) error { - // GenStrParam: json is varlena also. - p := vector.GenerateFunctionStrParameter(v) - for i := uint64(0); i < length; i++ { - v, null := p.GetStrValue(i) - if null { - op.enc[i].w.WriteString("null") - } else { - bj := types.DecodeJson(v) - val, err := bj.MarshalJSON() - // this should a valid json and we should never - // error here. Check it anyway. - if err != nil { - return err - } - // note here we already have a valid json string - // do NOT use encodeString, which will escape - // the string again. - op.enc[i].w.Write(val) +func encodeJSONRowArray[T types.ArrayElement]( + e *JqEncoder, + values []T, + toFloat64 func(T) float64, +) { + e.w.WriteByte('[') + for idx, value := range values { + if idx > 0 { + e.w.WriteByte(',') } + e.encodeFloat64(toFloat64(value)) } - return nil + e.w.WriteByte(']') } diff --git a/pkg/sql/plan/function/func_builtin_json_row_test.go b/pkg/sql/plan/function/func_builtin_json_row_test.go new file mode 100644 index 0000000000000..0c1f53397095e --- /dev/null +++ b/pkg/sql/plan/function/func_builtin_json_row_test.go @@ -0,0 +1,179 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package function + +import ( + "strconv" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/stretchr/testify/require" +) + +func TestJSONRowStreamsRowsAndResetsAfterError(t *testing.T) { + mp := mpool.MustNewZero() + defer mpool.DeleteMPool(mp) + proc := testutil.NewProcessWithMPool(t, "", mp) + + ints := vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixedList( + ints, + []int64{1, 2, 3}, + []bool{false, true, false}, + mp, + )) + defer ints.Free(mp) + strings := vector.NewVec(types.T_varchar.ToType()) + require.NoError(t, vector.AppendBytesList( + strings, + [][]byte{[]byte("a"), []byte("b"), []byte("c")}, + nil, + mp, + )) + defer strings.Free(mp) + bools := vector.NewVec(types.T_bool.ToType()) + require.NoError(t, vector.AppendFixedList( + bools, + []bool{true, false, true}, + nil, + mp, + )) + defer bools.Free(mp) + uints := vector.NewVec(types.T_uint64.ToType()) + require.NoError(t, vector.AppendFixedList( + uints, + []uint64{^uint64(0), 2, 3}, + nil, + mp, + )) + defer uints.Free(mp) + + result := vector.NewFunctionResultWrapper(types.T_varchar.ToType(), mp) + defer result.Free() + op := newOpBuiltInJsonRow() + require.NoError(t, result.PreExtendAndReset(3)) + require.NoError(t, op.jsonRow( + []*vector.Vector{ints, strings, bools, uints}, + result, + proc, + 3, + &FunctionSelectList{AnyNull: true, SelectList: []bool{true, false, true}}, + )) + out := result.GetResultVector() + require.Equal( + t, + []byte(`[1,"a",true,18446744073709551615]`), + out.GetBytesAt(0), + ) + require.True(t, out.IsNull(1)) + require.Equal(t, []byte(`[3,"c",true,3]`), out.GetBytesAt(2)) + + binary := vector.NewVec(types.T_binary.ToType()) + require.NoError(t, vector.AppendBytesList(binary, [][]byte{[]byte("x")}, nil, mp)) + defer binary.Free(mp) + require.NoError(t, result.PreExtendAndReset(1)) + require.Error(t, op.jsonRow( + []*vector.Vector{binary}, + result, + proc, + 1, + nil, + )) + for _, column := range op.columns { + require.Nil(t, column) + } + require.Zero(t, op.enc.w.Len()) + + require.NoError(t, result.PreExtendAndReset(3)) + require.NoError(t, op.jsonRow( + []*vector.Vector{ints, strings}, + result, + proc, + 3, + nil, + )) + out = result.GetResultVector() + require.Equal(t, []byte(`[1,"a"]`), out.GetBytesAt(0)) + require.Equal(t, []byte(`[null,"b"]`), out.GetBytesAt(1)) + require.Equal(t, []byte(`[3,"c"]`), out.GetBytesAt(2)) +} + +func newJSONRowBenchmarkParameters( + b *testing.B, + mp *mpool.MPool, + rows int, +) []*vector.Vector { + b.Helper() + ints := vector.NewVec(types.T_int64.ToType()) + intValues := make([]int64, rows) + for i := range intValues { + intValues[i] = int64(i) + } + require.NoError(b, vector.AppendFixedList(ints, intValues, nil, mp)) + + strings := vector.NewVec(types.T_varchar.ToType()) + stringValues := make([][]byte, rows) + for i := range stringValues { + stringValues[i] = []byte("value-" + strconv.Itoa(i%100)) + } + require.NoError(b, vector.AppendBytesList(strings, stringValues, nil, mp)) + return []*vector.Vector{ints, strings} +} + +func BenchmarkJSONRowFreshOperator8192(b *testing.B) { + mp := mpool.MustNewZero() + defer mpool.DeleteMPool(mp) + proc := testutil.NewProcessWithMPool(b, "", mp) + params := newJSONRowBenchmarkParameters(b, mp, 8192) + defer params[0].Free(mp) + defer params[1].Free(mp) + result := vector.NewFunctionResultWrapper(types.T_varchar.ToType(), mp) + defer result.Free() + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + require.NoError(b, result.PreExtendAndReset(8192)) + require.NoError(b, newOpBuiltInJsonRow().jsonRow( + params, + result, + proc, + 8192, + nil, + )) + } +} + +func BenchmarkJSONRowReusedOperator8192(b *testing.B) { + mp := mpool.MustNewZero() + defer mpool.DeleteMPool(mp) + proc := testutil.NewProcessWithMPool(b, "", mp) + params := newJSONRowBenchmarkParameters(b, mp, 8192) + defer params[0].Free(mp) + defer params[1].Free(mp) + result := vector.NewFunctionResultWrapper(types.T_varchar.ToType(), mp) + defer result.Free() + op := newOpBuiltInJsonRow() + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + require.NoError(b, result.PreExtendAndReset(8192)) + require.NoError(b, op.jsonRow(params, result, proc, 8192, nil)) + } +} diff --git a/pkg/sql/plan/function/func_unary.go b/pkg/sql/plan/function/func_unary.go index ff33c71066fbe..7231ee4b92fa7 100644 --- a/pkg/sql/plan/function/func_unary.go +++ b/pkg/sql/plan/function/func_unary.go @@ -216,120 +216,146 @@ func AbsArray[T types.RealNumbers](ivecs []*vector.Vector, result vector.Functio }, selectList) } -var ( - arrayF32Pool = sync.Pool{ - New: func() interface{} { - s := make([]float32, 128) - return &s - }, - } - - arrayF64Pool = sync.Pool{ - New: func() interface{} { - s := make([]float64, 128) - return &s - }, - } -) - func NormalizeL2Array[T types.ArrayElement](parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { source := vector.GenerateFunctionStrParameter(parameters[0]) rs := vector.MustFunctionResult[types.Varlena](result) rowCount := uint64(length) - var inArrayF32 []float32 - var outArrayF32Ptr *[]float32 - var outArrayF32 []float32 - - var inArrayF64 []float64 - var outArrayF64Ptr *[]float64 - var outArrayF64 []float64 - - var data []byte - var null bool - for i := uint64(0); i < rowCount; i++ { - data, null = source.GetStrValue(i) + data, null := source.GetStrValue(i) if null { - _ = rs.AppendMustNullForBytesResult() + if err := rs.AppendMustNullForBytesResult(); err != nil { + return err + } continue } switch t := parameters[0].GetType().Oid; t { case types.T_array_float32: - inArrayF32 = types.BytesToArray[float32](data) - - outArrayF32Ptr = arrayF32Pool.Get().(*[]float32) - outArrayF32 = *outArrayF32Ptr - - if cap(outArrayF32) < len(inArrayF32) { - outArrayF32 = make([]float32, len(inArrayF32)) - } else { - outArrayF32 = outArrayF32[:len(inArrayF32)] + if err := appendNormalizedRealArray[float32](rs, data); err != nil { + return err } - _ = moarray.NormalizeL2(inArrayF32, outArrayF32) - _ = rs.AppendBytes(types.ArrayToBytes[float32](outArrayF32), false) - - *outArrayF32Ptr = outArrayF32 - arrayF32Pool.Put(outArrayF32Ptr) case types.T_array_float64: - inArrayF64 = types.BytesToArray[float64](data) - - outArrayF64Ptr = arrayF64Pool.Get().(*[]float64) - outArrayF64 = *outArrayF64Ptr - - if cap(outArrayF64) < len(inArrayF64) { - outArrayF64 = make([]float64, len(inArrayF64)) - } else { - outArrayF64 = outArrayF64[:len(inArrayF64)] + if err := appendNormalizedRealArray[float64](rs, data); err != nil { + return err } - _ = moarray.NormalizeL2(inArrayF64, outArrayF64) - _ = rs.AppendBytes(types.ArrayToBytes[float64](outArrayF64), false) - - *outArrayF64Ptr = outArrayF64 - arrayF64Pool.Put(outArrayF64Ptr) case types.T_array_bf16: - _ = appendNormalizedNarrowArray[types.BF16](rs, data) + if err := appendNormalizedNarrowArray( + rs, + data, + types.BF16.ToFloat32, + types.BF16FromFloat32, + ); err != nil { + return err + } case types.T_array_float16: - _ = appendNormalizedNarrowArray[types.Float16](rs, data) + if err := appendNormalizedNarrowArray( + rs, + data, + types.Float16.ToFloat32, + types.Float16FromFloat32, + ); err != nil { + return err + } case types.T_array_int8: // A normalized vector is a unit vector, which cannot be represented in // an integer element type (components round to 0/±1 and the norm is no // longer 1), so int8/uint8 normalize_l2 widens the result to vecf32. // The overload's retType is T_array_float32 to match (see list_builtIn). - _ = appendNormalizedIntArrayAsFloat32[int8](rs, data) + if err := appendNormalizedArrayAsFloat32( + rs, + data, + func(value int8) float32 { return float32(value) }, + ); err != nil { + return err + } case types.T_array_uint8: - _ = appendNormalizedIntArrayAsFloat32[uint8](rs, data) + if err := appendNormalizedArrayAsFloat32( + rs, + data, + func(value uint8) float32 { return float32(value) }, + ); err != nil { + return err + } } - } return nil } -// appendNormalizedNarrowArray normalizes a bf16/f16 vector by upcasting to -// float32, normalizing in float32, then narrowing back to T. bf16/f16 are -// floating-point so they can hold a (near-)unit vector; int8/uint8 cannot and -// use appendNormalizedIntArrayAsFloat32 instead. -func appendNormalizedNarrowArray[T types.ArrayElement](rs *vector.FunctionResult[types.Varlena], data []byte) error { - in := types.ToFloat32Array[T](types.BytesToArray[T](data)) - out := make([]float32, len(in)) - _ = moarray.NormalizeL2(in, out) - return rs.AppendBytes(types.ArrayToBytes[T](types.FromFloat32Array[T](out)), false) +func appendNormalizedRealArray[T types.RealNumbers]( + rs *vector.FunctionResult[types.Varlena], + data []byte, +) error { + input := types.BytesToArray[T](data) + return rs.AppendBytesWithFill(len(data), func(dst []byte) { + _ = moarray.NormalizeL2(input, types.BytesToArray[T](dst)) + }) +} + +// appendNormalizedNarrowArray normalizes a bf16/f16 vector through float32 +// arithmetic while writing the narrowed values directly into the result. +func appendNormalizedNarrowArray[T types.ArrayElement]( + rs *vector.FunctionResult[types.Varlena], + data []byte, + toFloat32 func(T) float32, + fromFloat32 func(float32) T, +) error { + input := types.BytesToArray[T](data) + return rs.AppendBytesWithFill(len(data), func(dst []byte) { + output := types.BytesToArray[T](dst) + normalizeArrayInto(input, output, toFloat32, fromFloat32) + }) } -// appendNormalizedIntArrayAsFloat32 normalizes an integer-typed (int8/uint8) +// appendNormalizedArrayAsFloat32 normalizes an integer-typed (int8/uint8) // vector and writes the result as float32. A unit vector cannot be represented // in an integer element type — narrowing back would round components to 0/±1 so // the norm is no longer 1 (e.g. normalize_l2([0,1,2,3]::vecuint8) would become // [0,0,1,1], whose norm is √2). Widening the result to vecf32 keeps the unit-norm // contract; the int8/uint8 overloads declare retType T_array_float32 to match. -func appendNormalizedIntArrayAsFloat32[T types.ArrayElement](rs *vector.FunctionResult[types.Varlena], data []byte) error { - in := types.ToFloat32Array[T](types.BytesToArray[T](data)) - out := make([]float32, len(in)) - _ = moarray.NormalizeL2(in, out) - return rs.AppendBytes(types.ArrayToBytes[float32](out), false) +func appendNormalizedArrayAsFloat32[T types.ArrayElement]( + rs *vector.FunctionResult[types.Varlena], + data []byte, + toFloat32 func(T) float32, +) error { + input := types.BytesToArray[T](data) + if len(input) > int(^uint(0)>>1)/4 { + return moerr.NewInternalErrorNoCtx("normalized array result is too large") + } + return rs.AppendBytesWithFill(len(input)*4, func(dst []byte) { + output := types.BytesToArray[float32](dst) + normalizeArrayInto( + input, + output, + toFloat32, + func(value float32) float32 { return value }, + ) + }) +} + +func normalizeArrayInto[TIn, TOut types.ArrayElement]( + input []TIn, + output []TOut, + toFloat32 func(TIn) float32, + fromFloat32 func(float32) TOut, +) { + var sumSquares float64 + for _, value := range input { + converted := float64(toFloat32(value)) + sumSquares += converted * converted + } + norm := math.Sqrt(sumSquares) + if norm == 0 { + for idx, value := range input { + output[idx] = fromFloat32(toFloat32(value)) + } + return + } + for idx, value := range input { + output[idx] = fromFloat32(float32(float64(toFloat32(value)) / norm)) + } } func L1NormArray[T types.RealNumbers](ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { @@ -6205,9 +6231,12 @@ func generateSHAKey(key []byte) []byte { } func generateInitializationVector(key []byte, length int) []byte { - data := append(key, byte(length)) - hash := sha256.Sum256(data) - return hash[:aes.BlockSize] + hasher := sha256.New() + _, _ = hasher.Write(key) + var lengthByte [1]byte + lengthByte[0] = byte(length) + _, _ = hasher.Write(lengthByte[:]) + return hasher.Sum(nil)[:aes.BlockSize] } // encode function encrypts a string, returns a binary string of the same length of the original string. @@ -6222,10 +6251,10 @@ func encodeByAES(plaintext []byte, key []byte, null bool, rs *vector.FunctionRes return err } initializationVector := generateInitializationVector(key, len(plaintext)) - ciphertext := make([]byte, len(plaintext)) stream := cipher.NewCTR(block, initializationVector) - stream.XORKeyStream(ciphertext, plaintext) - return rs.AppendMustBytesValue(ciphertext) + return rs.AppendBytesWithFill(len(plaintext), func(ciphertext []byte) { + stream.XORKeyStream(ciphertext, plaintext) + }) } func Encode(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { @@ -6257,10 +6286,10 @@ func decodeByAES(ciphertext []byte, key []byte, null bool, rs *vector.FunctionRe return err } iv := generateInitializationVector(key, len(ciphertext)) - plaintext := make([]byte, len(ciphertext)) stream := cipher.NewCTR(block, iv) - stream.XORKeyStream(plaintext, ciphertext) - return rs.AppendMustBytesValue(plaintext) + return rs.AppendBytesWithFill(len(ciphertext), func(plaintext []byte) { + stream.XORKeyStream(plaintext, ciphertext) + }) } func Decode(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { diff --git a/pkg/vectorindex/metric/cpu.go b/pkg/vectorindex/metric/cpu.go index 3ca9222647712..70134166ad6d1 100644 --- a/pkg/vectorindex/metric/cpu.go +++ b/pkg/vectorindex/metric/cpu.go @@ -49,6 +49,24 @@ func PairwiseDistanceLaunch[T types.ArrayElement]( return PairwiseDistanceLaunchCPU(x, y, metric, dist) } +func PairwiseDistanceLaunchOneToMany[T types.RealNumbers]( + query []T, + rowCount int, + rowAt func(int) []T, + metric MetricType, + dist []float32, + _ uint64, + _ bool, +) (PairwiseJobHandle, error) { + return PairwiseDistanceLaunchOneToManyCPU( + query, + rowCount, + rowAt, + metric, + dist, + ) +} + func PairwiseDistanceWait(handle PairwiseJobHandle, metric MetricType) ([]float32, error) { return PairwiseDistanceWaitCPU(handle, metric) } diff --git a/pkg/vectorindex/metric/gpu.go b/pkg/vectorindex/metric/gpu.go index 5c474bd0ab76e..1bb56ac0a5c62 100644 --- a/pkg/vectorindex/metric/gpu.go +++ b/pkg/vectorindex/metric/gpu.go @@ -21,6 +21,7 @@ import ( "sync" "github.com/matrixorigin/matrixone/pkg/common/malloc" + "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/common/util" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/cuvs" @@ -202,6 +203,75 @@ func PairwiseDistanceLaunch[T types.ArrayElement]( return PairwiseDistanceLaunchCPU(x, y, metric, dist) } +func PairwiseDistanceLaunchOneToMany[T types.RealNumbers]( + query []T, + rowCount int, + rowAt func(int) []T, + metric MetricType, + dist []float32, + minWorkSize uint64, + gpuMode bool, +) (PairwiseJobHandle, error) { + if !gpuMode { + return PairwiseDistanceLaunchOneToManyCPU( + query, + rowCount, + rowAt, + metric, + dist, + ) + } + if rowCount < 0 || len(dist) < rowCount { + return 0, moerr.NewInternalErrorNoCtx( + "pairwise distance output is smaller than the row count", + ) + } + if rowCount == 0 { + return PairwiseDistanceLaunchOneToManyCPU( + query, + rowCount, + rowAt, + metric, + dist, + ) + } + + dim := len(query) + work := uint64(rowCount) + if dim != 0 && work > ^uint64(0)/uint64(dim) { + work = ^uint64(0) + } else { + work *= uint64(dim) + } + cuvsMetric, supportedMetric := MetricTypeToCuvsMetric[metric] + if supportedMetric && + work >= minWorkSize { + if typedQuery, ok := any(query).([]float32); ok { + return gpuPairwiseLaunchRows( + 1, + rowCount, + dim, + func(_ int) []float32 { + return typedQuery + }, + func(row int) []float32 { + return any(rowAt(row)).([]float32) + }, + cuvsMetric, + dist[:rowCount], + 4, + ) + } + } + return PairwiseDistanceLaunchOneToManyCPU( + query, + rowCount, + rowAt, + metric, + dist, + ) +} + // gpuPairwiseLaunch flattens [][]C into a C-allocator buffer (elemSize bytes per // element) and launches the async cuVS pairwise distance. C is float32 (4B) or // cuvs.Float16 (2B). Mirrors the old f32-only path, generalized over the element. @@ -212,27 +282,88 @@ func gpuPairwiseLaunch[C cuvs.VectorType]( dist []float32, elemSize int, ) (PairwiseJobHandle, error) { - nX, nY := len(x), len(y) + return gpuPairwiseLaunchRows( + len(x), + len(y), + dim, + func(row int) []C { + return x[row] + }, + func(row int) []C { + return y[row] + }, + cuvsMetric, + dist, + elemSize, + ) +} + +func gpuPairwiseLaunchRows[C cuvs.VectorType]( + nX, nY, dim int, + xAt, yAt func(int) []C, + cuvsMetric cuvs.DistanceType, + dist []float32, + elemSize int, +) (PairwiseJobHandle, error) { + if nX < 0 || + nY < 0 || + dim < 0 || + elemSize <= 0 || + uint64(dim) > uint64(^uint32(0)) || + uint64(dim) > ^uint64(0)/uint64(elemSize) { + return 0, moerr.NewInternalErrorNoCtx( + "pairwise distance input is too large", + ) + } + rowBytes := uint64(dim) * uint64(elemSize) + if rowBytes != 0 && + (uint64(nX) > ^uint64(0)/rowBytes || + uint64(nY) > ^uint64(0)/rowBytes) { + return 0, moerr.NewInternalErrorNoCtx( + "pairwise distance input is too large", + ) + } allocator := malloc.NewCAllocator() // 1. Flatten Y - yBuf, yDeallocator, err := allocator.Allocate(uint64(nY*dim*elemSize), malloc.NoClear) + yBuf, yDeallocator, err := allocator.Allocate( + uint64(nY)*rowBytes, + malloc.NoClear, + ) if err != nil { return 0, err } yf := util.UnsafeSliceCast[C](yBuf) - for i, v := range y { + for i := 0; i < nY; i++ { + v := yAt(i) + if len(v) != dim { + yDeallocator.Deallocate() + return 0, moerr.NewInternalErrorNoCtx( + "vector dimension not matched", + ) + } copy(yf[i*dim:(i+1)*dim], v) } // 2. Flatten X - xBuf, xDeallocator, err := allocator.Allocate(uint64(nX*dim*elemSize), malloc.NoClear) + xBuf, xDeallocator, err := allocator.Allocate( + uint64(nX)*rowBytes, + malloc.NoClear, + ) if err != nil { yDeallocator.Deallocate() return 0, err } xf := util.UnsafeSliceCast[C](xBuf) - for i, v := range x { + for i := 0; i < nX; i++ { + v := xAt(i) + if len(v) != dim { + xDeallocator.Deallocate() + yDeallocator.Deallocate() + return 0, moerr.NewInternalErrorNoCtx( + "vector dimension not matched", + ) + } copy(xf[i*dim:(i+1)*dim], v) } diff --git a/pkg/vectorindex/metric/pairwise.go b/pkg/vectorindex/metric/pairwise.go index a012c5f5ade2c..e0c866aeda284 100644 --- a/pkg/vectorindex/metric/pairwise.go +++ b/pkg/vectorindex/metric/pairwise.go @@ -72,18 +72,15 @@ func PairwiseDistanceLaunchCPU[T types.ArrayElement]( dist = make([]float32, nX*nY) } - job := &pairWiseJob{ - dist: dist, - } - // One unified loop over any ArrayElement type — the resolver handles f32/f64 // and the narrow kernels (bf16/f16/int8/uint8) uniformly. + var jobErr error for r := 0; r < nX; r++ { xr := x[r] for c := 0; c < nY; c++ { d, err := distFn(xr, y[c]) if err != nil { - job.err = err + jobErr = err goto DONE } dist[r*nY+c] = d @@ -97,6 +94,56 @@ func PairwiseDistanceLaunchCPU[T types.ArrayElement]( } DONE: + return registerPairwiseCPUJob(dist, jobErr), nil +} + +// PairwiseDistanceLaunchOneToManyCPU computes the distances between one query +// and rowCount caller-owned rows. The caller supplies the output storage so the +// SQL expression path does not need a row-scaled [][]T descriptor slice or an +// additional result allocation. +func PairwiseDistanceLaunchOneToManyCPU[T types.RealNumbers]( + query []T, + rowCount int, + rowAt func(int) []T, + metric MetricType, + dist []float32, +) (PairwiseJobHandle, error) { + if rowCount < 0 || len(dist) < rowCount { + return 0, moerr.NewInternalErrorNoCtx( + "pairwise distance output is smaller than the row count", + ) + } + dist = dist[:rowCount] + distFn, err := ResolveDistanceFn[T, float32](metric) + if err != nil { + return 0, err + } + + var jobErr error + for row := 0; row < rowCount; row++ { + value, err := distFn(query, rowAt(row)) + if err != nil { + jobErr = err + break + } + dist[row] = value + } + if jobErr == nil && metric == Metric_L2Distance { + for idx := range dist { + dist[idx] = float32(math.Sqrt(float64(dist[idx]))) + } + } + return registerPairwiseCPUJob(dist, jobErr), nil +} + +func registerPairwiseCPUJob( + dist []float32, + err error, +) PairwiseJobHandle { + job := &pairWiseJob{ + dist: dist, + err: err, + } jobMu.Lock() id := nextID nextID++ @@ -107,7 +154,7 @@ DONE: jobMap[uint64(handle)] = job jobMu.Unlock() - return handle, nil + return handle } // PairwiseDistanceWaitCPU returns the results of the pairwise distance calculation diff --git a/pkg/vectorindex/metric/pairwise_test.go b/pkg/vectorindex/metric/pairwise_test.go index 8aab26132f5b1..aba18191b4417 100644 --- a/pkg/vectorindex/metric/pairwise_test.go +++ b/pkg/vectorindex/metric/pairwise_test.go @@ -126,6 +126,38 @@ func TestPairwiseDistanceLaunchWaitCPU_Float64_L2(t *testing.T) { require.InDelta(t, 1.0, float64(out[3]), 1e-5) } +func TestPairwiseDistanceLaunchOneToManyCPUUsesCallerOutput(t *testing.T) { + query := []float32{1, 0} + rows := [][]float32{{1, 0}, {1, 1}, {0, 1}} + dist := make([]float32, len(rows)) + + handle, err := PairwiseDistanceLaunchOneToManyCPU( + query, + len(rows), + func(row int) []float32 { + return rows[row] + }, + Metric_L2sqDistance, + dist, + ) + require.NoError(t, err) + out, err := PairwiseDistanceWaitCPU(handle, Metric_L2sqDistance) + require.NoError(t, err) + require.Equal(t, []float32{0, 1, 2}, out) + require.Equal(t, &dist[0], &out[0]) + + _, err = PairwiseDistanceLaunchOneToManyCPU( + query, + len(rows), + func(row int) []float32 { + return rows[row] + }, + Metric_L2sqDistance, + dist[:len(rows)-1], + ) + require.Error(t, err) +} + func TestPairwiseDistanceWaitCPU_InvalidHandle(t *testing.T) { _, err := PairwiseDistanceWaitCPU(PairwiseJobHandle(0), Metric_L2sqDistance) require.Error(t, err) From f2b0645a543fe847afb6dad8e7855c2bdbd9b9b2 Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 19:24:02 +0800 Subject: [PATCH 12/61] feat: add allocation-accounted function scratch --- pkg/container/vector/allocation_account.go | 35 +++-- .../vector/allocation_account_test.go | 22 +-- pkg/container/vector/functionTools.go | 134 +++++++++++++----- .../vector/function_result_allocation_test.go | 97 +++++++++++-- pkg/sql/colexec/evalExpression.go | 28 ++-- pkg/sql/colexec/eval_expression_allocation.go | 26 ++-- .../eval_expression_allocation_test.go | 35 +++++ 7 files changed, 282 insertions(+), 95 deletions(-) diff --git a/pkg/container/vector/allocation_account.go b/pkg/container/vector/allocation_account.go index 32849404bd986..0daaf8151d543 100644 --- a/pkg/container/vector/allocation_account.go +++ b/pkg/container/vector/allocation_account.go @@ -41,23 +41,27 @@ type AllocationAccountSelection struct { accountBitmaps bool } -// FunctionParameterAllocation is the immutable allocation provenance for -// row-scaled function-parameter conversion scratch. -type FunctionParameterAllocation struct { - account *mpool.AllocationAccount - owner mpool.AllocationOwner - site mpool.AllocationSite +// FunctionAllocation is the immutable allocation provenance for row-scaled +// function-owned scratch. Parameter conversion and general function scratch +// use distinct sites so the physical charges remain diagnosable. +type FunctionAllocation struct { + account *mpool.AllocationAccount + owner mpool.AllocationOwner + parameterSite mpool.AllocationSite + scratchSite mpool.AllocationSite } -func NewFunctionParameterAllocation( +func NewFunctionAllocation( account *mpool.AllocationAccount, owner mpool.AllocationOwner, - site mpool.AllocationSite, -) (*FunctionParameterAllocation, error) { - allocation := &FunctionParameterAllocation{ - account: account, - owner: owner, - site: site, + parameterSite mpool.AllocationSite, + scratchSite mpool.AllocationSite, +) (*FunctionAllocation, error) { + allocation := &FunctionAllocation{ + account: account, + owner: owner, + parameterSite: parameterSite, + scratchSite: scratchSite, } if err := allocation.validate(); err != nil { return nil, err @@ -65,13 +69,14 @@ func NewFunctionParameterAllocation( return allocation, nil } -func (a *FunctionParameterAllocation) validate() error { +func (a *FunctionAllocation) validate() error { if a == nil || a.account == nil || a.account.Handle() == 0 || a.owner < mpool.AllocationOwnerMin || a.owner > mpool.AllocationOwnerMax || - a.site < mpool.AllocationSiteMin { + a.parameterSite < mpool.AllocationSiteMin || + a.scratchSite < mpool.AllocationSiteMin { return mpool.ErrAllocationAccountInvalid } return nil diff --git a/pkg/container/vector/allocation_account_test.go b/pkg/container/vector/allocation_account_test.go index 253343a2314a7..9594112db793a 100644 --- a/pkg/container/vector/allocation_account_test.go +++ b/pkg/container/vector/allocation_account_test.go @@ -26,22 +26,23 @@ import ( ) const ( - testVectorAllocationOwner mpool.AllocationOwner = 1 - testVectorDataAllocationSite mpool.AllocationSite = 1 - testVectorAreaAllocationSite mpool.AllocationSite = 2 - testVectorNullAllocationSite mpool.AllocationSite = 3 - testVectorGroupAllocationSite mpool.AllocationSite = 4 - testVectorParamAllocationSite mpool.AllocationSite = 5 + testVectorAllocationOwner mpool.AllocationOwner = 1 + testVectorDataAllocationSite mpool.AllocationSite = 1 + testVectorAreaAllocationSite mpool.AllocationSite = 2 + testVectorNullAllocationSite mpool.AllocationSite = 3 + testVectorGroupAllocationSite mpool.AllocationSite = 4 + testVectorParamAllocationSite mpool.AllocationSite = 5 + testVectorScratchAllocationSite mpool.AllocationSite = 6 ) type testVectorAllocationAccount struct { registry *mpool.AllocationAccountRegistry account *mpool.AllocationAccount selection *AllocationAccountSelection - parameter *FunctionParameterAllocation + function *FunctionAllocation } -func newTestVectorParameterAllocationAccount( +func newTestVectorFunctionAllocationAccount( t testing.TB, limit uint64, allocationSlots uint64, @@ -60,17 +61,18 @@ func newTestVectorParameterAllocationAccount( testVectorGroupAllocationSite, ) require.NoError(t, err) - parameter, err := NewFunctionParameterAllocation( + function, err := NewFunctionAllocation( account, testVectorAllocationOwner, testVectorParamAllocationSite, + testVectorScratchAllocationSite, ) require.NoError(t, err) return testVectorAllocationAccount{ registry: registry, account: account, selection: selection, - parameter: parameter, + function: function, } } diff --git a/pkg/container/vector/functionTools.go b/pkg/container/vector/functionTools.go index 3ae33ddfb9712..a0ff0af783531 100644 --- a/pkg/container/vector/functionTools.go +++ b/pkg/container/vector/functionTools.go @@ -485,7 +485,9 @@ type OptFunctionResultWrapper interface { getConvenientParamList() []reusableParameterWrapper hasParameterScratch() bool resizeParameterScratch(idx int, size int) ([]byte, error) - setParameterAllocation(allocation *FunctionParameterAllocation) + setFunctionAllocation(allocation *FunctionAllocation) + HasFunctionScratch() bool + ResizeFunctionScratch(size int) ([]byte, bool, error) } func OptGetParamFromWrapper[ParamType types.FixedSizeTExceptStrType]( @@ -723,11 +725,11 @@ type FunctionResult[T types.FixedSizeT] struct { vec *Vector mp *mpool.MPool - allocationAccount *AllocationAccountSelection - parameterAllocation *FunctionParameterAllocation - isVarlena bool - cols []T - length uint64 + allocationAccount *AllocationAccountSelection + functionAllocation *FunctionAllocation + isVarlena bool + cols []T + length uint64 // convenientParam save parameter wrappers for easy getting row values. // @@ -735,6 +737,7 @@ type FunctionResult[T types.FixedSizeT] struct { // there are still many built-in functions don't use it now, and will be fixed in the future. convenientParam []reusableParameterWrapper parameterScratch []*mpool.AccountedBuffer + functionScratch *mpool.AccountedBuffer } func MustFunctionResult[T types.FixedSizeT](wrapper FunctionResultWrapper) *FunctionResult[T] { @@ -769,7 +772,7 @@ func (fr *FunctionResult[T]) UseOptFunctionParamFrame(paramCount int) { fr.convenientParam = make([]reusableParameterWrapper, paramCount) } if fr.allocationAccount != nil && - fr.parameterAllocation != nil && + fr.functionAllocation != nil && fr.parameterScratch == nil { fr.parameterScratch = make([]*mpool.AccountedBuffer, paramCount) } @@ -780,13 +783,17 @@ func (fr *FunctionResult[T]) getConvenientParamList() []reusableParameterWrapper } func (fr *FunctionResult[T]) hasParameterScratch() bool { - return fr.parameterAllocation != nil + return fr.functionAllocation != nil } -func (fr *FunctionResult[T]) setParameterAllocation( - allocation *FunctionParameterAllocation, +func (fr *FunctionResult[T]) setFunctionAllocation( + allocation *FunctionAllocation, ) { - fr.parameterAllocation = allocation + fr.functionAllocation = allocation +} + +func (fr *FunctionResult[T]) HasFunctionScratch() bool { + return fr.functionAllocation != nil } func (fr *FunctionResult[T]) resizeParameterScratch( @@ -802,9 +809,9 @@ func (fr *FunctionResult[T]) resizeParameterScratch( if fr.parameterScratch[idx] == nil { buffer, err := mpool.NewAccountedBuffer( fr.mp, - fr.parameterAllocation.account, - fr.parameterAllocation.owner, - fr.parameterAllocation.site, + fr.functionAllocation.account, + fr.functionAllocation.owner, + fr.functionAllocation.parameterSite, ) if err != nil { return nil, err @@ -817,6 +824,33 @@ func (fr *FunctionResult[T]) resizeParameterScratch( return fr.parameterScratch[idx].Bytes(), nil } +// ResizeFunctionScratch returns retained allocation-accounted off-heap +// scratch for data-scaled function internals. Legacy results report selected +// false so callers preserve their existing allocator path. +func (fr *FunctionResult[T]) ResizeFunctionScratch( + size int, +) ([]byte, bool, error) { + if fr.functionAllocation == nil { + return nil, false, nil + } + if fr.functionScratch == nil { + buffer, err := mpool.NewAccountedBuffer( + fr.mp, + fr.functionAllocation.account, + fr.functionAllocation.owner, + fr.functionAllocation.scratchSite, + ) + if err != nil { + return nil, true, err + } + fr.functionScratch = buffer + } + if err := fr.functionScratch.Resize(size); err != nil { + return nil, true, err + } + return fr.functionScratch.Bytes(), true, nil +} + func (fr *FunctionResult[T]) PreExtendAndReset(targetSize int) error { if fr.vec == nil { var err error @@ -883,24 +917,40 @@ func (fr *FunctionResult[T]) AppendBytes(val []byte, isnull bool) error { // AppendBytesWithFill appends one non-null varlena value and lets fill write // directly into the result Vector's admitted backing storage. The provided // slice is valid only during fill. A panic rolls the unpublished row and area -// length back before propagating. +// length back before propagating. Use AppendBytesWithBuilder when construction +// can return an error or a shorter value. func (fr *FunctionResult[T]) AppendBytesWithFill( size int, fill func([]byte), +) error { + return fr.AppendBytesWithBuilder(size, func(dst []byte) (int, error) { + fill(dst) + return size, nil + }) +} + +// AppendBytesWithBuilder reserves capacity for one non-null varlena value and +// lets build return the number of bytes it initialized. This supports codecs +// whose exact output is known only after encoding without allocating a second +// payload buffer. Reserved capacity remains owned by the result across reuse; +// only the published area length is reduced to the actual value size. +func (fr *FunctionResult[T]) AppendBytesWithBuilder( + capacity int, + build func([]byte) (int, error), ) error { if !fr.isVarlena || fr.vec == nil || fr.vec.IsConst() || - size < 0 || - fill == nil { + capacity < 0 || + build == nil { return mpool.ErrAllocationAccountInvalid } oldAreaLen := len(fr.vec.area) - if uint64(oldAreaLen)+uint64(size) > uint64(math.MaxUint32) { + if uint64(oldAreaLen)+uint64(capacity) > uint64(math.MaxUint32) { return mpool.ErrAllocationAccountInvalid } - areaSize := size - if size <= types.VarlenaInlineSize { + areaSize := capacity + if capacity <= types.VarlenaInlineSize { areaSize = 0 } if err := fr.vec.PreExtendWithArea(1, areaSize, fr.mp); err != nil { @@ -912,12 +962,10 @@ func (fr *FunctionResult[T]) AppendBytesWithFill( oldValue := values[index] values[index] = types.Varlena{} var target []byte - if size <= types.VarlenaInlineSize { - values[index][0] = byte(size) - target = values[index][1 : 1+size] + if capacity <= types.VarlenaInlineSize { + target = values[index][1 : 1+capacity] } else { - fr.vec.area = fr.vec.area[:oldAreaLen+size] - values[index].SetOffsetLen(uint32(oldAreaLen), uint32(size)) + fr.vec.area = fr.vec.area[:oldAreaLen+capacity] target = fr.vec.area[oldAreaLen:] } @@ -928,7 +976,23 @@ func (fr *FunctionResult[T]) AppendBytesWithFill( values[index] = oldValue } }() - fill(target) + written, err := build(target) + if err != nil { + return err + } + if written < 0 || written > capacity { + return mpool.ErrAllocationAccountInvalid + } + if written <= types.VarlenaInlineSize { + if capacity > types.VarlenaInlineSize { + copy(values[index][1:1+written], target[:written]) + fr.vec.area = fr.vec.area[:oldAreaLen] + } + values[index][0] = byte(written) + } else { + fr.vec.area = fr.vec.area[:oldAreaLen+written] + values[index].SetOffsetLen(uint32(oldAreaLen), uint32(written)) + } fr.vec.length++ published = true return nil @@ -1026,8 +1090,12 @@ func (fr *FunctionResult[T]) Free() { fr.parameterScratch[i] = nil } } + if fr.functionScratch != nil { + fr.functionScratch.Free() + fr.functionScratch = nil + } fr.allocationAccount = nil - fr.parameterAllocation = nil + fr.functionAllocation = nil fr.convenientParam = nil fr.parameterScratch = nil } @@ -1050,24 +1118,24 @@ func NewFunctionResultWrapperWithAllocation( return newFunctionResultWrapper(typ, mp, selection), nil } -func NewFunctionResultWrapperWithParameterAllocation( +func NewFunctionResultWrapperWithFunctionAllocation( typ types.Type, mp *mpool.MPool, selection *AllocationAccountSelection, - parameterAllocation *FunctionParameterAllocation, + functionAllocation *FunctionAllocation, ) (FunctionResultWrapper, error) { if err := selection.validate(); err != nil { return nil, err } - if err := parameterAllocation.validate(); err != nil { + if err := functionAllocation.validate(); err != nil { return nil, err } - if selection.account != parameterAllocation.account || - selection.owner != parameterAllocation.owner { + if selection.account != functionAllocation.account || + selection.owner != functionAllocation.owner { return nil, mpool.ErrAllocationAccountInvalid } result := newFunctionResultWrapper(typ, mp, selection) - result.setParameterAllocation(parameterAllocation) + result.setFunctionAllocation(functionAllocation) return result, nil } diff --git a/pkg/container/vector/function_result_allocation_test.go b/pkg/container/vector/function_result_allocation_test.go index 7b7aed3ebc568..50a987a93f72c 100644 --- a/pkg/container/vector/function_result_allocation_test.go +++ b/pkg/container/vector/function_result_allocation_test.go @@ -15,6 +15,7 @@ package vector import ( + "errors" "testing" "github.com/matrixorigin/matrixone/pkg/common/mpool" @@ -92,30 +93,32 @@ func TestFunctionResultAllocationAccountFailure(t *testing.T) { ) state := newTestVectorAllocationAccount(t, 1<<20, 4) - _, err := NewFunctionResultWrapperWithParameterAllocation( + _, err := NewFunctionResultWrapperWithFunctionAllocation( types.T_int64.ToType(), zeroMP, state.selection, nil, ) require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) - otherOwner, err := NewFunctionParameterAllocation( + otherOwner, err := NewFunctionAllocation( state.account, testVectorAllocationOwner+1, testVectorParamAllocationSite, + testVectorScratchAllocationSite, ) require.NoError(t, err) - _, err = NewFunctionResultWrapperWithParameterAllocation( + _, err = NewFunctionResultWrapperWithFunctionAllocation( types.T_int64.ToType(), zeroMP, state.selection, otherOwner, ) require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) - _, err = NewFunctionParameterAllocation( + _, err = NewFunctionAllocation( nil, testVectorAllocationOwner, testVectorParamAllocationSite, + testVectorScratchAllocationSite, ) require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) finalizeTestVectorAllocationAccount(t, state) @@ -175,10 +178,25 @@ func TestFunctionResultAppendBytesWithFillLifecycle(t *testing.T) { require.Equal(t, beforeLength, wrapper.GetResultVector().Length()) require.Equal(t, beforeAreaLength, len(wrapper.GetResultVector().GetArea())) - require.NoError(t, result.AppendBytesWithFill(4, func(dst []byte) { + fillErr := errors.New("injected fill error") + require.ErrorIs(t, result.AppendBytesWithBuilder(128, func(dst []byte) (int, error) { + dst[0] = 2 + return 0, fillErr + }), fillErr) + require.Equal(t, beforeLength, wrapper.GetResultVector().Length()) + require.Equal(t, beforeAreaLength, len(wrapper.GetResultVector().GetArea())) + require.ErrorIs(t, result.AppendBytesWithBuilder(128, func([]byte) (int, error) { + return 129, nil + }), mpool.ErrAllocationAccountInvalid) + require.Equal(t, beforeLength, wrapper.GetResultVector().Length()) + require.Equal(t, beforeAreaLength, len(wrapper.GetResultVector().GetArea())) + + require.NoError(t, result.AppendBytesWithBuilder(512, func(dst []byte) (int, error) { copy(dst, "last") + return 4, nil })) require.Equal(t, []byte("last"), wrapper.GetResultVector().GetBytesAt(2)) + require.Equal(t, beforeAreaLength, len(wrapper.GetResultVector().GetArea())) require.Positive(t, state.account.Snapshot().Used) wrapper.Free() @@ -187,14 +205,14 @@ func TestFunctionResultAppendBytesWithFillLifecycle(t *testing.T) { } func TestFunctionResultAllocationAccountDecimalParameterScratch(t *testing.T) { - state := newTestVectorParameterAllocationAccount(t, 1<<20, 16) + state := newTestVectorFunctionAllocationAccount(t, 1<<20, 16) mp := mpool.MustNew("function-parameter-allocation") defer mpool.DeleteMPool(mp) - result, err := NewFunctionResultWrapperWithParameterAllocation( + result, err := NewFunctionResultWrapperWithFunctionAllocation( types.T_bool.ToType(), mp, state.selection, - state.parameter, + state.function, ) require.NoError(t, err) result.UseOptFunctionParamFrame(1) @@ -289,14 +307,14 @@ func TestFunctionResultAllocationAccountDecimalParameterScratch(t *testing.T) { } func TestFunctionResultAllocationAccountDecimalParameterFailure(t *testing.T) { - state := newTestVectorParameterAllocationAccount(t, 127, 4) + state := newTestVectorFunctionAllocationAccount(t, 127, 4) mp := mpool.MustNew("function-parameter-allocation-failure") defer mpool.DeleteMPool(mp) - result, err := NewFunctionResultWrapperWithParameterAllocation( + result, err := NewFunctionResultWrapperWithFunctionAllocation( types.T_bool.ToType(), mp, state.selection, - state.parameter, + state.function, ) require.NoError(t, err) result.UseOptFunctionParamFrame(1) @@ -332,3 +350,60 @@ func TestFunctionResultAllocationAccountDecimalParameterFailure(t *testing.T) { source.Free(mp) finalizeTestVectorAllocationAccount(t, state) } + +func TestFunctionResultAllocationAccountFunctionScratch(t *testing.T) { + state := newTestVectorFunctionAllocationAccount(t, 1<<20, 8) + mp := mpool.MustNew("function-scratch-allocation") + defer mpool.DeleteMPool(mp) + result, err := NewFunctionResultWrapperWithFunctionAllocation( + types.T_bool.ToType(), + mp, + state.selection, + state.function, + ) + require.NoError(t, err) + + scratch, selected, err := result.ResizeFunctionScratch(128) + require.NoError(t, err) + require.True(t, selected) + require.Len(t, scratch, 128) + used := state.account.Snapshot().Used + require.Positive(t, used) + + scratch, selected, err = result.ResizeFunctionScratch(64) + require.NoError(t, err) + require.True(t, selected) + require.Len(t, scratch, 64) + require.Equal(t, used, state.account.Snapshot().Used) + + result.Free() + require.Zero(t, state.account.Snapshot().Used) + finalizeTestVectorAllocationAccount(t, state) + + legacy := NewFunctionResultWrapper(types.T_bool.ToType(), mp) + scratch, selected, err = legacy.ResizeFunctionScratch(128) + require.NoError(t, err) + require.False(t, selected) + require.Nil(t, scratch) + legacy.Free() +} + +func TestFunctionResultAllocationAccountFunctionScratchFailure(t *testing.T) { + state := newTestVectorFunctionAllocationAccount(t, 127, 2) + mp := mpool.MustNew("function-scratch-allocation-failure") + defer mpool.DeleteMPool(mp) + result, err := NewFunctionResultWrapperWithFunctionAllocation( + types.T_bool.ToType(), + mp, + state.selection, + state.function, + ) + require.NoError(t, err) + + _, selected, err := result.ResizeFunctionScratch(128) + require.True(t, selected) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Zero(t, state.account.Snapshot().Used) + result.Free() + finalizeTestVectorAllocationAccount(t, state) +} diff --git a/pkg/sql/colexec/evalExpression.go b/pkg/sql/colexec/evalExpression.go index 2863b10e8c053..1ae0acb4f2ff6 100644 --- a/pkg/sql/colexec/evalExpression.go +++ b/pkg/sql/colexec/evalExpression.go @@ -772,11 +772,11 @@ func (expr *FunctionExpressionExecutor) init( expr.resultVector = vector.NewFunctionResultWrapper(retType, m) return nil } - expr.resultVector, err = vector.NewFunctionResultWrapperWithParameterAllocation( + expr.resultVector, err = vector.NewFunctionResultWrapperWithFunctionAllocation( retType, m, allocation.result, - allocation.parameter, + allocation.function, ) return err } @@ -1117,11 +1117,11 @@ func (expr *FunctionExpressionExecutor) evalSelectedRows( ) } else { expr.selectedResult, err = - vector.NewFunctionResultWrapperWithParameterAllocation( + vector.NewFunctionResultWrapperWithFunctionAllocation( expr.resultType, expr.m, expr.allocation.scratch, - expr.allocation.parameter, + expr.allocation.function, ) if err != nil { return nil, err @@ -1573,17 +1573,17 @@ func generateConstExpressionExecutor( default: return nil, moerr.NewNYI(proc.Ctx, fmt.Sprintf("const expression %v", con.GetValue())) } - if err != nil { - return nil, err - } - if vec == nil { - return nil, moerr.NewNYI( - proc.Ctx, - fmt.Sprintf("const expression %v", con.GetValue()), - ) - } - vec.SetIsBin(con.IsBin) } + if err != nil { + return nil, err + } + if vec == nil { + return nil, moerr.NewNYI( + proc.Ctx, + fmt.Sprintf("const expression %v", con.GetValue()), + ) + } + vec.SetIsBin(con.IsBin) return vec, nil } diff --git a/pkg/sql/colexec/eval_expression_allocation.go b/pkg/sql/colexec/eval_expression_allocation.go index 37f7687d83b8d..aa65a8fbabd2e 100644 --- a/pkg/sql/colexec/eval_expression_allocation.go +++ b/pkg/sql/colexec/eval_expression_allocation.go @@ -41,6 +41,7 @@ const ( ExpressionAllocationSiteScratchNulls ExpressionAllocationSiteScratchGrouping ExpressionAllocationSiteParameterConversion + ExpressionAllocationSiteFunctionScratch ) // ExpressionAllocationAccount is the immutable allocation provenance shared @@ -50,10 +51,10 @@ type ExpressionAllocationAccount struct { account *mpool.AllocationAccount owner mpool.AllocationOwner - constant *vector.AllocationAccountSelection - result *vector.AllocationAccountSelection - scratch *vector.AllocationAccountSelection - parameter *vector.FunctionParameterAllocation + constant *vector.AllocationAccountSelection + result *vector.AllocationAccountSelection + scratch *vector.AllocationAccountSelection + function *vector.FunctionAllocation } func NewExpressionAllocationAccount( @@ -93,21 +94,22 @@ func NewExpressionAllocationAccount( if err != nil { return nil, err } - parameter, err := vector.NewFunctionParameterAllocation( + function, err := vector.NewFunctionAllocation( account, owner, ExpressionAllocationSiteParameterConversion, + ExpressionAllocationSiteFunctionScratch, ) if err != nil { return nil, err } return &ExpressionAllocationAccount{ - account: account, - owner: owner, - constant: constant, - result: result, - scratch: scratch, - parameter: parameter, + account: account, + owner: owner, + constant: constant, + result: result, + scratch: scratch, + function: function, }, nil } @@ -116,7 +118,7 @@ func (a *ExpressionAllocationAccount) validate() error { a.owner < mpool.AllocationOwnerMin || a.owner > mpool.AllocationOwnerMax || a.constant == nil || a.result == nil || a.scratch == nil || - a.parameter == nil { + a.function == nil { return mpool.ErrAllocationAccountInvalid } return nil diff --git a/pkg/sql/colexec/eval_expression_allocation_test.go b/pkg/sql/colexec/eval_expression_allocation_test.go index ff63d06c8fb51..af0c5bf983595 100644 --- a/pkg/sql/colexec/eval_expression_allocation_test.go +++ b/pkg/sql/colexec/eval_expression_allocation_test.go @@ -425,6 +425,41 @@ func TestExpressionAllocationAccountConstructionRollback(t *testing.T) { finalizeTestExpressionAllocationAccount(t, state) } +func TestExpressionAllocationAccountConstNullPreservesBinaryFlag(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("expression-allocation-const-null"), + ) + defer proc.Free() + state := newTestExpressionAllocationAccount(t, 1<<20, 4) + + expr := &plan.Expr{ + Typ: plan.Type{Id: int32(types.T_varbinary)}, + Expr: &plan.Expr_Lit{Lit: &plan.Literal{ + Isnull: true, + IsBin: true, + }}, + } + executor, err := NewExpressionExecutorWithAllocation( + proc, + expr, + state.allocation, + ) + require.NoError(t, err) + result, err := executor.Eval( + proc, + []*batch.Batch{batch.EmptyForConstFoldBatch}, + nil, + ) + require.NoError(t, err) + require.True(t, result.IsConstNull()) + require.True(t, result.GetIsBin()) + + executor.Free() + finalizeTestExpressionAllocationAccount(t, state) +} + func TestExpressionAllocationAccountScratchFailureCleanup(t *testing.T) { proc := testutil.NewProcessWithMPool( t, From 9fb49b3b09ac20dbd2772918cb11f2f2ef0f1836 Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 19:24:06 +0800 Subject: [PATCH 13/61] perf: close data-scaled function scratch paths --- .../bytejson/bytejson_composite_plan.go | 358 +++++++ .../bytejson/bytejson_composite_plan_test.go | 125 +++ pkg/container/bytejson/bytejson_keys_plan.go | 85 ++ .../bytejson/bytejson_keys_plan_test.go | 68 ++ .../bytejson/bytejson_scalar_plan.go | 66 ++ .../bytejson/bytejson_text_writer.go | 295 ++++++ .../bytejson/bytejson_text_writer_test.go | 60 ++ pkg/sql/plan/function/func_binary.go | 357 +++++-- .../func_binary_array_distance_test.go | 2 + pkg/sql/plan/function/func_builtin.go | 139 ++- pkg/sql/plan/function/func_builtin_jq.go | 273 +++++- pkg/sql/plan/function/func_builtin_json.go | 912 +++++++++++------- pkg/sql/plan/function/func_cast.go | 58 +- pkg/sql/plan/function/func_compare.go | 46 +- pkg/sql/plan/function/func_prefix.go | 143 ++- .../plan/function/func_string_complex_test.go | 4 +- pkg/sql/plan/function/func_unary.go | 406 ++++++-- .../function/func_unary_codec_scratch_test.go | 161 ++++ .../function_allocation_scratch_test.go | 763 +++++++++++++++ pkg/sql/plan/function/operator_in.go | 278 +++++- pkg/vectorindex/metric/cpu.go | 32 + pkg/vectorindex/metric/gpu.go | 178 +++- 22 files changed, 4139 insertions(+), 670 deletions(-) create mode 100644 pkg/container/bytejson/bytejson_composite_plan.go create mode 100644 pkg/container/bytejson/bytejson_composite_plan_test.go create mode 100644 pkg/container/bytejson/bytejson_keys_plan.go create mode 100644 pkg/container/bytejson/bytejson_keys_plan_test.go create mode 100644 pkg/container/bytejson/bytejson_scalar_plan.go create mode 100644 pkg/container/bytejson/bytejson_text_writer.go create mode 100644 pkg/container/bytejson/bytejson_text_writer_test.go create mode 100644 pkg/sql/plan/function/func_unary_codec_scratch_test.go create mode 100644 pkg/sql/plan/function/function_allocation_scratch_test.go diff --git a/pkg/container/bytejson/bytejson_composite_plan.go b/pkg/container/bytejson/bytejson_composite_plan.go new file mode 100644 index 0000000000000..f37df516f7c62 --- /dev/null +++ b/pkg/container/bytejson/bytejson_composite_plan.go @@ -0,0 +1,358 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package bytejson + +import ( + "bytes" + "encoding/base64" + "encoding/binary" + "math" + "slices" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" +) + +type fixedDataEncoder struct { + typeCode TpCode + data [numberSize]byte + size uint32 +} + +func NewLiteralDataEncoder(literal byte) ByteJsonDataEncoder { + return &fixedDataEncoder{typeCode: TpCodeLiteral, data: [numberSize]byte{literal}, size: 1} +} + +func NewInt64DataEncoder(value int64) ByteJsonDataEncoder { + encoder := &fixedDataEncoder{typeCode: TpCodeInt64, size: numberSize} + endian.PutUint64(encoder.data[:], uint64(value)) + return encoder +} + +func NewUint64DataEncoder(value uint64) ByteJsonDataEncoder { + encoder := &fixedDataEncoder{typeCode: TpCodeUint64, size: numberSize} + endian.PutUint64(encoder.data[:], value) + return encoder +} + +func NewFloat64DataEncoder(value float64) ByteJsonDataEncoder { + encoder := &fixedDataEncoder{typeCode: TpCodeFloat64, size: numberSize} + endian.PutUint64(encoder.data[:], math.Float64bits(value)) + return encoder +} + +func (e *fixedDataEncoder) TypeCode() TpCode { return e.typeCode } +func (e *fixedDataEncoder) DataSize() uint32 { return e.size } +func (e *fixedDataEncoder) EncodeDataInto(dst []byte) (int, error) { + if e == nil || len(dst) != int(e.size) { + return 0, moerr.NewInvalidArgNoCtx("JSON scalar", "result size mismatch") + } + return copy(dst, e.data[:e.size]), nil +} + +type rawDataEncoder struct { + value ByteJson +} + +// NewRawDataEncoder references an already storage-compatible value. +func NewRawDataEncoder(value ByteJson) (ByteJsonDataEncoder, error) { + if value.requiresLegacyBinaryEncoding() { + return nil, moerr.NewInvalidArgNoCtx( + "JSON value", + "value is not storage compatible", + ) + } + if uint64(len(value.Data)) > math.MaxUint32 { + return nil, moerr.NewInvalidArgNoCtx("JSON value", "value is too large") + } + return &rawDataEncoder{value: value}, nil +} + +func (e *rawDataEncoder) TypeCode() TpCode { return e.value.Type } +func (e *rawDataEncoder) DataSize() uint32 { return uint32(len(e.value.Data)) } +func (e *rawDataEncoder) EncodeDataInto(dst []byte) (int, error) { + if e == nil || len(dst) != len(e.value.Data) { + return 0, moerr.NewInvalidArgNoCtx("JSON value", "result size mismatch") + } + return copy(dst, e.value.Data), nil +} + +type typedStringDataEncoder struct { + typeCode TpCode + value []byte + dataSize uint32 +} + +func NewTypedStringDataEncoder(tp TpCode, value []byte) (ByteJsonDataEncoder, error) { + switch tp { + case TpCodeString, TpCodeDecimal, TpCodeDate, TpCodeTime, TpCodeDatetime, TpCodeBlob: + default: + return nil, moerr.NewInvalidArgNoCtx("JSON string", "invalid type code") + } + dataSize, err := binaryStringDataSize(len(value)) + if err != nil { + return nil, err + } + return &typedStringDataEncoder{typeCode: tp, value: value, dataSize: dataSize}, nil +} + +func (e *typedStringDataEncoder) TypeCode() TpCode { return e.typeCode } +func (e *typedStringDataEncoder) DataSize() uint32 { return e.dataSize } +func (e *typedStringDataEncoder) EncodeDataInto(dst []byte) (int, error) { + if e == nil || len(dst) != int(e.dataSize) { + return 0, moerr.NewInvalidArgNoCtx("JSON string", "result size mismatch") + } + written := binary.PutUvarint(dst, uint64(len(e.value))) + written += copy(dst[written:], e.value) + return written, nil +} + +type binaryDataEncoder struct { + value []byte + prefix string + dataSize uint32 +} + +func NewOpaqueDataEncoder(value []byte) (ByteJsonDataEncoder, error) { + return newBinaryDataEncoder(value, "") +} + +func NewBitDataEncoder(value []byte) (ByteJsonDataEncoder, error) { + return newBinaryDataEncoder(value, persistedBitPrefix) +} + +func newBinaryDataEncoder(value []byte, prefix string) (ByteJsonDataEncoder, error) { + encodedLength := uint64(base64.StdEncoding.EncodedLen(len(value))) + uint64(len(prefix)) + if encodedLength > math.MaxInt { + return nil, moerr.NewInvalidArgNoCtx("JSON binary", "value is too large") + } + dataSize, err := binaryStringDataSize(int(encodedLength)) + if err != nil { + return nil, err + } + return &binaryDataEncoder{value: value, prefix: prefix, dataSize: dataSize}, nil +} + +func (e *binaryDataEncoder) TypeCode() TpCode { return TpCodeBlob } +func (e *binaryDataEncoder) DataSize() uint32 { return e.dataSize } +func (e *binaryDataEncoder) EncodeDataInto(dst []byte) (int, error) { + if e == nil || len(dst) != int(e.dataSize) { + return 0, moerr.NewInvalidArgNoCtx("JSON binary", "result size mismatch") + } + encodedLength := base64.StdEncoding.EncodedLen(len(e.value)) + len(e.prefix) + written := binary.PutUvarint(dst, uint64(encodedLength)) + written += copy(dst[written:], e.prefix) + base64.StdEncoding.Encode(dst[written:], e.value) + return written + base64.StdEncoding.EncodedLen(len(e.value)), nil +} + +func binaryStringDataSize(length int) (uint32, error) { + if length < 0 { + return 0, moerr.NewInvalidArgNoCtx("JSON string", "invalid length") + } + var lengthBuffer [binary.MaxVarintLen64]byte + lengthSize := binary.PutUvarint(lengthBuffer[:], uint64(length)) + total := uint64(lengthSize) + uint64(length) + if total > math.MaxUint32 { + return 0, moerr.NewInvalidArgNoCtx("JSON string", "value is too large") + } + return uint32(total), nil +} + +type ArrayDataEncoder struct { + values []ByteJsonDataEncoder + dataSize uint32 +} + +func NewArrayDataEncoder(values []ByteJsonDataEncoder) (*ArrayDataEncoder, error) { + total := uint64(headerSize) + uint64(len(values))*valEntrySize + for _, value := range values { + if value == nil { + return nil, moerr.NewInvalidArgNoCtx("JSON array", "nil value encoder") + } + if value.TypeCode() != TpCodeLiteral { + total += uint64(value.DataSize()) + } + if total > math.MaxUint32 { + return nil, moerr.NewInvalidArgNoCtx("JSON array", "result is too large") + } + } + return &ArrayDataEncoder{values: values, dataSize: uint32(total)}, nil +} + +func (e *ArrayDataEncoder) TypeCode() TpCode { return TpCodeArray } +func (e *ArrayDataEncoder) DataSize() uint32 { return e.dataSize } +func (e *ArrayDataEncoder) EncodeDataInto(dst []byte) (int, error) { + if e == nil || len(dst) != int(e.dataSize) { + return 0, moerr.NewInvalidArgNoCtx("JSON array", "result size mismatch") + } + headerEnd := headerSize + len(e.values)*valEntrySize + clear(dst[:headerEnd]) + endian.PutUint32(dst, uint32(len(e.values))) + endian.PutUint32(dst[docSizeOff:], e.dataSize) + payloadOffset := headerEnd + for idx, value := range e.values { + entryOffset := headerSize + idx*valEntrySize + dst[entryOffset] = byte(value.TypeCode()) + if value.TypeCode() == TpCodeLiteral { + if _, err := value.EncodeDataInto(dst[entryOffset+valTypeSize : entryOffset+valTypeSize+1]); err != nil { + return 0, err + } + continue + } + endian.PutUint32(dst[entryOffset+valTypeSize:], uint32(payloadOffset)) + size := int(value.DataSize()) + written, err := value.EncodeDataInto(dst[payloadOffset : payloadOffset+size]) + if err != nil { + return 0, err + } + if written != size { + return 0, moerr.NewInvalidArgNoCtx("JSON array", "value size mismatch") + } + payloadOffset += size + } + return payloadOffset, nil +} + +type IndexedFloatArrayDataEncoder struct { + count int + valueAt func(int) float64 + dataSize uint32 +} + +func NewIndexedFloatArrayDataEncoder( + count int, + valueAt func(int) float64, +) (*IndexedFloatArrayDataEncoder, error) { + if count < 0 || valueAt == nil { + return nil, moerr.NewInvalidArgNoCtx("JSON array", "invalid value accessor") + } + total := uint64(headerSize) + uint64(count)*(valEntrySize+numberSize) + if total > math.MaxUint32 { + return nil, moerr.NewInvalidArgNoCtx("JSON array", "result is too large") + } + return &IndexedFloatArrayDataEncoder{ + count: count, valueAt: valueAt, dataSize: uint32(total), + }, nil +} + +func (e *IndexedFloatArrayDataEncoder) TypeCode() TpCode { return TpCodeArray } +func (e *IndexedFloatArrayDataEncoder) DataSize() uint32 { return e.dataSize } +func (e *IndexedFloatArrayDataEncoder) EncodeDataInto(dst []byte) (int, error) { + if e == nil || len(dst) != int(e.dataSize) { + return 0, moerr.NewInvalidArgNoCtx("JSON array", "result size mismatch") + } + headerEnd := headerSize + e.count*valEntrySize + clear(dst[:headerEnd]) + endian.PutUint32(dst, uint32(e.count)) + endian.PutUint32(dst[docSizeOff:], e.dataSize) + payloadOffset := headerEnd + for idx := 0; idx < e.count; idx++ { + entryOffset := headerSize + idx*valEntrySize + dst[entryOffset] = byte(TpCodeFloat64) + endian.PutUint32(dst[entryOffset+valTypeSize:], uint32(payloadOffset)) + endian.PutUint64(dst[payloadOffset:], math.Float64bits(e.valueAt(idx))) + payloadOffset += numberSize + } + return payloadOffset, nil +} + +type ObjectDataEncoderEntry struct { + Key []byte + Value ByteJsonDataEncoder + order int +} + +type ObjectDataEncoder struct { + entries []ObjectDataEncoderEntry + dataSize uint32 +} + +func NewObjectDataEncoder(entries []ObjectDataEncoderEntry) (*ObjectDataEncoder, error) { + for idx := range entries { + entries[idx].order = idx + if entries[idx].Value == nil || len(entries[idx].Key) > math.MaxUint16 { + return nil, moerr.NewInvalidArgNoCtx("JSON object", "invalid entry") + } + } + slices.SortFunc(entries, func(left, right ObjectDataEncoderEntry) int { + if order := bytes.Compare(left.Key, right.Key); order != 0 { + return order + } + return left.order - right.order + }) + unique := entries[:0] + for _, entry := range entries { + if len(unique) > 0 && bytes.Equal(unique[len(unique)-1].Key, entry.Key) { + unique[len(unique)-1] = entry + } else { + unique = append(unique, entry) + } + } + entries = unique + total := uint64(headerSize) + uint64(len(entries))*(keyEntrySize+valEntrySize) + for _, entry := range entries { + total += uint64(len(entry.Key)) + if entry.Value.TypeCode() != TpCodeLiteral { + total += uint64(entry.Value.DataSize()) + } + if total > math.MaxUint32 { + return nil, moerr.NewInvalidArgNoCtx("JSON object", "result is too large") + } + } + return &ObjectDataEncoder{entries: entries, dataSize: uint32(total)}, nil +} + +func (e *ObjectDataEncoder) TypeCode() TpCode { return TpCodeObject } +func (e *ObjectDataEncoder) DataSize() uint32 { return e.dataSize } +func (e *ObjectDataEncoder) EncodeDataInto(dst []byte) (int, error) { + if e == nil || len(dst) != int(e.dataSize) { + return 0, moerr.NewInvalidArgNoCtx("JSON object", "result size mismatch") + } + count := len(e.entries) + keyEntryBegin := headerSize + valueEntryBegin := keyEntryBegin + count*keyEntrySize + payloadOffset := valueEntryBegin + count*valEntrySize + clear(dst[:payloadOffset]) + endian.PutUint32(dst, uint32(count)) + endian.PutUint32(dst[docSizeOff:], e.dataSize) + for idx, entry := range e.entries { + entryOffset := keyEntryBegin + idx*keyEntrySize + endian.PutUint32(dst[entryOffset:], uint32(payloadOffset)) + endian.PutUint16(dst[entryOffset+keyOriginOff:], uint16(len(entry.Key))) + payloadOffset += copy(dst[payloadOffset:], entry.Key) + } + for idx, entry := range e.entries { + value := entry.Value + entryOffset := valueEntryBegin + idx*valEntrySize + dst[entryOffset] = byte(value.TypeCode()) + if value.TypeCode() == TpCodeLiteral { + if _, err := value.EncodeDataInto(dst[entryOffset+valTypeSize : entryOffset+valTypeSize+1]); err != nil { + return 0, err + } + continue + } + endian.PutUint32(dst[entryOffset+valTypeSize:], uint32(payloadOffset)) + size := int(value.DataSize()) + written, err := value.EncodeDataInto(dst[payloadOffset : payloadOffset+size]) + if err != nil { + return 0, err + } + if written != size { + return 0, moerr.NewInvalidArgNoCtx("JSON object", "value size mismatch") + } + payloadOffset += size + } + return payloadOffset, nil +} diff --git a/pkg/container/bytejson/bytejson_composite_plan_test.go b/pkg/container/bytejson/bytejson_composite_plan_test.go new file mode 100644 index 0000000000000..ee3f3cf8376b1 --- /dev/null +++ b/pkg/container/bytejson/bytejson_composite_plan_test.go @@ -0,0 +1,125 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package bytejson + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func encodedByteJSON(t *testing.T, encoder ByteJsonDataEncoder) ByteJson { + t.Helper() + data := make([]byte, encoder.DataSize()) + written, err := encoder.EncodeDataInto(data) + require.NoError(t, err) + require.Equal(t, len(data), written) + return ByteJson{Type: encoder.TypeCode(), Data: data} +} + +func TestArrayDataEncoderMatchesCreateByteJSON(t *testing.T) { + stringEncoder, err := NewStringDataEncoder([]byte("value")) + require.NoError(t, err) + values := []ByteJsonDataEncoder{ + NewLiteralDataEncoder(LiteralNull), + NewLiteralDataEncoder(LiteralTrue), + NewInt64DataEncoder(-7), + NewUint64DataEncoder(9), + NewFloat64DataEncoder(1.25), + stringEncoder, + } + encoder, err := NewArrayDataEncoder(values) + require.NoError(t, err) + want, err := CreateByteJSON([]any{nil, true, int64(-7), uint64(9), 1.25, "value"}) + require.NoError(t, err) + require.Equal(t, want, encodedByteJSON(t, encoder)) +} + +func TestIndexedFloatArrayDataEncoder(t *testing.T) { + values := []float64{1.5, -2, 3.25} + encoder, err := NewIndexedFloatArrayDataEncoder( + len(values), + func(idx int) float64 { return values[idx] }, + ) + require.NoError(t, err) + want, err := CreateByteJSON([]any{1.5, -2.0, 3.25}) + require.NoError(t, err) + require.Equal(t, want, encodedByteJSON(t, encoder)) +} + +func TestObjectDataEncoderSortsAndKeepsLastDuplicate(t *testing.T) { + first, err := NewStringDataEncoder([]byte("first")) + require.NoError(t, err) + last, err := NewStringDataEncoder([]byte("last")) + require.NoError(t, err) + encoder, err := NewObjectDataEncoder([]ObjectDataEncoderEntry{ + {Key: []byte("z"), Value: NewInt64DataEncoder(1)}, + {Key: []byte("a"), Value: first}, + {Key: []byte("a"), Value: last}, + }) + require.NoError(t, err) + want, err := CreateByteJSON(map[string]any{"a": "last", "z": int64(1)}) + require.NoError(t, err) + require.Equal(t, want, encodedByteJSON(t, encoder)) +} + +func TestBinaryDataEncodersAreStorageCompatible(t *testing.T) { + raw := []byte{0, 1, 2, 250, 251} + for _, constructor := range []func([]byte) (ByteJsonDataEncoder, error){ + NewOpaqueDataEncoder, + NewBitDataEncoder, + } { + encoder, err := constructor(raw) + require.NoError(t, err) + value := encodedByteJSON(t, encoder) + require.Equal(t, TpCodeBlob, value.Type) + require.False(t, value.requiresLegacyBinaryEncoding()) + } +} + +func TestCompositeDataEncodersPreserveNestedAndTypedValues(t *testing.T) { + nested, err := CreateByteJSON(map[string]any{ + "key": []any{int64(1), "value"}, + }) + require.NoError(t, err) + raw, err := NewRawDataEncoder(nested) + require.NoError(t, err) + date, err := NewTypedStringDataEncoder(TpCodeDate, []byte("2026-07-31")) + require.NoError(t, err) + decimal, err := NewTypedStringDataEncoder(TpCodeDecimal, []byte("123.450")) + require.NoError(t, err) + array, err := NewArrayDataEncoder([]ByteJsonDataEncoder{raw, date, decimal}) + require.NoError(t, err) + encoded := encodedByteJSON(t, array) + require.Equal(t, nested, encoded.GetArrayElem(0)) + require.Equal(t, TpCodeDate, encoded.GetArrayElem(1).Type) + require.Equal(t, []byte("2026-07-31"), encoded.GetArrayElem(1).GetString()) + require.Equal(t, TpCodeDecimal, encoded.GetArrayElem(2).Type) + require.Equal(t, []byte("123.450"), encoded.GetArrayElem(2).GetString()) +} + +func TestCompositeDataEncodersRejectInvalidPlansAndDestinations(t *testing.T) { + _, err := NewArrayDataEncoder([]ByteJsonDataEncoder{nil}) + require.Error(t, err) + _, err = NewObjectDataEncoder([]ObjectDataEncoderEntry{{Key: []byte("key")}}) + require.Error(t, err) + _, err = NewTypedStringDataEncoder(TpCodeObject, []byte("invalid")) + require.Error(t, err) + + encoder, err := NewArrayDataEncoder([]ByteJsonDataEncoder{NewInt64DataEncoder(1)}) + require.NoError(t, err) + _, err = encoder.EncodeDataInto(make([]byte, encoder.DataSize()-1)) + require.Error(t, err) +} diff --git a/pkg/container/bytejson/bytejson_keys_plan.go b/pkg/container/bytejson/bytejson_keys_plan.go new file mode 100644 index 0000000000000..95a313e4dc915 --- /dev/null +++ b/pkg/container/bytejson/bytejson_keys_plan.go @@ -0,0 +1,85 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package bytejson + +import ( + "encoding/binary" + "math" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" +) + +// ObjectKeysArrayEncoder exposes an object's sorted keys as a JSON array +// without materializing []any or a second binary-JSON payload. +type ObjectKeysArrayEncoder struct { + object ByteJson + dataSize uint32 +} + +func NewObjectKeysArrayEncoder(object ByteJson) (*ObjectKeysArrayEncoder, error) { + if object.Type != TpCodeObject { + return nil, moerr.NewInvalidArgNoCtx("json_keys", "JSON value is not an object") + } + count := object.GetElemCnt() + total := uint64(headerSize) + uint64(count)*uint64(valEntrySize) + var lengthBuffer [binary.MaxVarintLen64]byte + for idx := 0; idx < count; idx++ { + key := object.GetObjectKey(idx) + lengthSize := binary.PutUvarint(lengthBuffer[:], uint64(len(key))) + total += uint64(lengthSize) + uint64(len(key)) + if total > math.MaxUint32 { + return nil, moerr.NewInvalidArgNoCtx("json_keys", "JSON result is too large") + } + } + return &ObjectKeysArrayEncoder{ + object: object, + dataSize: uint32(total), + }, nil +} + +func (e *ObjectKeysArrayEncoder) TypeCode() TpCode { + return TpCodeArray +} + +func (e *ObjectKeysArrayEncoder) DataSize() uint32 { + if e == nil { + return 0 + } + return e.dataSize +} + +func (e *ObjectKeysArrayEncoder) EncodeDataInto(dst []byte) (int, error) { + if e == nil || uint64(len(dst)) != uint64(e.dataSize) { + return 0, moerr.NewInvalidArgNoCtx("json_keys", "JSON result size mismatch") + } + count := e.object.GetElemCnt() + headerEnd := headerSize + count*valEntrySize + clear(dst[:headerEnd]) + endian.PutUint32(dst, uint32(count)) + endian.PutUint32(dst[docSizeOff:], e.dataSize) + payloadOffset := headerEnd + for idx := 0; idx < count; idx++ { + entryOffset := headerSize + idx*valEntrySize + dst[entryOffset] = byte(TpCodeString) + endian.PutUint32(dst[entryOffset+valTypeSize:], uint32(payloadOffset)) + key := e.object.GetObjectKey(idx) + payloadOffset += binary.PutUvarint(dst[payloadOffset:], uint64(len(key))) + payloadOffset += copy(dst[payloadOffset:], key) + } + if payloadOffset != len(dst) { + return 0, moerr.NewInvalidArgNoCtx("json_keys", "JSON result size mismatch") + } + return payloadOffset, nil +} diff --git a/pkg/container/bytejson/bytejson_keys_plan_test.go b/pkg/container/bytejson/bytejson_keys_plan_test.go new file mode 100644 index 0000000000000..608925d2f1cdc --- /dev/null +++ b/pkg/container/bytejson/bytejson_keys_plan_test.go @@ -0,0 +1,68 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package bytejson + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestObjectKeysArrayEncoder(t *testing.T) { + object, err := CreateByteJSON(map[string]any{ + "z": int64(1), + "a": true, + "longer": nil, + }) + require.NoError(t, err) + encoder, err := NewObjectKeysArrayEncoder(object) + require.NoError(t, err) + encoded := make([]byte, encoder.DataSize()) + written, err := encoder.EncodeDataInto(encoded) + require.NoError(t, err) + require.Equal(t, len(encoded), written) + result := ByteJson{Type: encoder.TypeCode(), Data: encoded} + visible, err := result.MarshalJSON() + require.NoError(t, err) + require.JSONEq(t, `["a", "longer", "z"]`, string(visible)) +} + +func TestObjectKeysArrayEncoderRejectsInvalidInputAndSize(t *testing.T) { + array, err := CreateByteJSON([]any{int64(1)}) + require.NoError(t, err) + _, err = NewObjectKeysArrayEncoder(array) + require.Error(t, err) + + object, err := CreateByteJSON(map[string]any{"key": int64(1)}) + require.NoError(t, err) + encoder, err := NewObjectKeysArrayEncoder(object) + require.NoError(t, err) + _, err = encoder.EncodeDataInto(make([]byte, encoder.DataSize()-1)) + require.Error(t, err) +} + +func TestStringDataEncoder(t *testing.T) { + encoder, err := NewStringDataEncoder([]byte("a value \" with unicode 世界")) + require.NoError(t, err) + encoded := make([]byte, encoder.DataSize()) + written, err := encoder.EncodeDataInto(encoded) + require.NoError(t, err) + require.Equal(t, len(encoded), written) + result := ByteJson{Type: encoder.TypeCode(), Data: encoded} + require.Equal(t, []byte("a value \" with unicode 世界"), result.GetString()) + + _, err = encoder.EncodeDataInto(encoded[:len(encoded)-1]) + require.Error(t, err) +} diff --git a/pkg/container/bytejson/bytejson_scalar_plan.go b/pkg/container/bytejson/bytejson_scalar_plan.go new file mode 100644 index 0000000000000..164edac97c0b6 --- /dev/null +++ b/pkg/container/bytejson/bytejson_scalar_plan.go @@ -0,0 +1,66 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package bytejson + +import ( + "encoding/binary" + "math" + "unicode/utf8" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" +) + +// StringDataEncoder writes a binary-JSON string from caller-owned bytes. +// The source must remain valid until EncodeDataInto returns. +type StringDataEncoder struct { + value []byte + dataSize uint32 +} + +func NewStringDataEncoder(value []byte) (*StringDataEncoder, error) { + if !utf8.Valid(value) { + return nil, moerr.NewInvalidArgNoCtx("JSON string", "invalid UTF-8") + } + var lengthBuffer [binary.MaxVarintLen64]byte + lengthSize := binary.PutUvarint(lengthBuffer[:], uint64(len(value))) + total := uint64(lengthSize) + uint64(len(value)) + if total > math.MaxUint32 { + return nil, moerr.NewInvalidArgNoCtx("JSON string", "value is too large") + } + return &StringDataEncoder{ + value: value, + dataSize: uint32(total), + }, nil +} + +func (e *StringDataEncoder) TypeCode() TpCode { + return TpCodeString +} + +func (e *StringDataEncoder) DataSize() uint32 { + if e == nil { + return 0 + } + return e.dataSize +} + +func (e *StringDataEncoder) EncodeDataInto(dst []byte) (int, error) { + if e == nil || uint64(len(dst)) != uint64(e.dataSize) { + return 0, moerr.NewInvalidArgNoCtx("JSON string", "result size mismatch") + } + written := binary.PutUvarint(dst, uint64(len(e.value))) + written += copy(dst[written:], e.value) + return written, nil +} diff --git a/pkg/container/bytejson/bytejson_text_writer.go b/pkg/container/bytejson/bytejson_text_writer.go new file mode 100644 index 0000000000000..d91674766654f --- /dev/null +++ b/pkg/container/bytejson/bytejson_text_writer.go @@ -0,0 +1,295 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package bytejson + +import ( + "encoding/base64" + "fmt" + "io" + "math" + "strconv" + "unicode/utf8" +) + +// WriteJSONText writes the visible JSON representation without allocating a +// payload-sized intermediate slice. It is byte-for-byte equivalent to +// ByteJson.MarshalJSON. +func WriteJSONText(w io.Writer, value ByteJson) error { + switch value.Type { + case TpCodeArray: + if err := writeByte(w, '['); err != nil { + return err + } + for idx := 0; idx < value.GetElemCnt(); idx++ { + if idx > 0 { + if err := writeString(w, ", "); err != nil { + return err + } + } + if err := WriteJSONText(w, value.GetArrayElem(idx)); err != nil { + return err + } + } + return writeByte(w, ']') + case TpCodeObject: + if err := writeByte(w, '{'); err != nil { + return err + } + for idx := 0; idx < value.GetElemCnt(); idx++ { + if idx > 0 { + if err := writeString(w, ", "); err != nil { + return err + } + } + if err := WriteJSONString(w, value.GetObjectKey(idx)); err != nil { + return err + } + if err := writeString(w, ": "); err != nil { + return err + } + if err := WriteJSONText(w, value.GetObjectVal(idx)); err != nil { + return err + } + } + return writeByte(w, '}') + case TpCodeInt64: + var buf [32]byte + return writeBytes(w, strconv.AppendInt(buf[:0], value.GetInt64(), 10)) + case TpCodeUint64: + var buf [32]byte + return writeBytes(w, strconv.AppendUint(buf[:0], value.GetUint64(), 10)) + case TpCodeLiteral: + if len(value.Data) == 0 { + return fmt.Errorf("invalid JSON literal") + } + switch value.Data[0] { + case LiteralNull: + return writeString(w, "null") + case LiteralTrue: + return writeString(w, "true") + case LiteralFalse: + return writeString(w, "false") + default: + return fmt.Errorf("invalid JSON literal %d", value.Data[0]) + } + case TpCodeFloat64: + f := value.GetFloat64() + if math.IsInf(f, 0) || math.IsNaN(f) { + return fmt.Errorf("invalid JSON float64 %f", f) + } + format := byte('e') + abs := math.Abs(f) + if abs == 0 || 1e-6 <= abs && abs < 1e21 { + format = 'f' + } + var buf [32]byte + return writeBytes(w, strconv.AppendFloat(buf[:0], f, format, -1, 64)) + case TpCodeString: + return WriteJSONString(w, value.GetString()) + case TpCodeDecimal: + return writeBytes(w, value.GetString()) + case TpCodeDate, TpCodeTime, TpCodeDatetime: + if err := writeByte(w, '"'); err != nil { + return err + } + if err := writeBytes(w, value.GetString()); err != nil { + return err + } + return writeByte(w, '"') + case TpCodeBlob: + if err := writeByte(w, '"'); err != nil { + return err + } + if err := writeBinaryJSONText(w, value); err != nil { + return err + } + return writeByte(w, '"') + case TpCodeOpaque, TpCodeBit: + if err := writeByte(w, '"'); err != nil { + return err + } + if err := writeBinaryJSONText(w, value); err != nil { + return err + } + return writeByte(w, '"') + default: + return fmt.Errorf("invalid JSON type %d", value.Type) + } +} + +// WriteJSONObjectKeyText applies JSON_OBJECT's key coercion to an existing +// binary-JSON value without allocating its visible representation. +func WriteJSONObjectKeyText(w io.Writer, value ByteJson) error { + switch value.Type { + case TpCodeString, TpCodeDate, TpCodeTime, TpCodeDatetime: + return writeBytes(w, value.GetString()) + case TpCodeBlob, TpCodeOpaque, TpCodeBit: + return writeBinaryJSONText(w, value) + default: + return WriteJSONText(w, value) + } +} + +// WriteJSONBase64Text writes the visible text of a binary JSON scalar without +// surrounding quotes. +func WriteJSONBase64Text(w io.Writer, value []byte) error { + return writeRawBase64(w, value) +} + +func writeBinaryJSONText(w io.Writer, value ByteJson) error { + if value.Type == TpCodeOpaque || value.Type == TpCodeBit { + return writeRawBase64(w, value.GetString()) + } + data := value.GetString() + if len(data) >= len(persistedBitPrefix) && + string(data[:len(persistedBitPrefix)]) == persistedBitPrefix { + encoded := data[len(persistedBitPrefix):] + if _, ok := base64DecodedLen(encoded); ok { + return writeNormalizedBase64(w, encoded) + } + } + return writeBytes(w, data) +} + +// WriteJSONString writes one JSON string without allocating an escaped copy. +func WriteJSONString(w io.Writer, value []byte) error { + if err := writeByte(w, '"'); err != nil { + return err + } + start := 0 + for offset := 0; offset < len(value); { + b := value[offset] + if b < utf8.RuneSelf { + if b >= ' ' && b != '"' && b != '\\' { + offset++ + continue + } + if err := writeBytes(w, value[start:offset]); err != nil { + return err + } + var escaped string + switch b { + case '"': + escaped = `\"` + case '\\': + escaped = `\\` + case '\b': + escaped = `\b` + case '\f': + escaped = `\f` + case '\n': + escaped = `\n` + case '\r': + escaped = `\r` + case '\t': + escaped = `\t` + default: + const hex = "0123456789abcdef" + var escapedControl = [6]byte{'\\', 'u', '0', '0', hex[b>>4], hex[b&0xf]} + if err := writeBytes(w, escapedControl[:]); err != nil { + return err + } + offset++ + start = offset + continue + } + if err := writeString(w, escaped); err != nil { + return err + } + offset++ + start = offset + continue + } + _, size := utf8.DecodeRune(value[offset:]) + if size == 1 { + return fmt.Errorf("invalid UTF-8") + } + offset += size + } + if err := writeBytes(w, value[start:]); err != nil { + return err + } + return writeByte(w, '"') +} + +func writeNormalizedBase64(w io.Writer, encoded []byte) error { + var decoded [binaryJSONCompareDecodedChunkSize]byte + for offset := 0; offset < len(encoded); { + n, next, ok := decodeBase64Chunk(encoded, offset, decoded[:]) + if !ok { + return fmt.Errorf("invalid base64 JSON value") + } + if err := writeRawBase64(w, decoded[:n]); err != nil { + return err + } + offset = next + } + return nil +} + +func writeRawBase64(w io.Writer, raw []byte) error { + const decodedChunk = 3 * 256 + const encodedChunk = 4 * 256 + var encoded [encodedChunk]byte + for len(raw) > 0 { + length := min(len(raw), decodedChunk) + if length < len(raw) { + length -= length % 3 + } + written := base64.StdEncoding.EncodedLen(length) + base64.StdEncoding.Encode(encoded[:written], raw[:length]) + if err := writeBytes(w, encoded[:written]); err != nil { + return err + } + raw = raw[length:] + } + return nil +} + +func writeBytes(w io.Writer, value []byte) error { + for len(value) > 0 { + written, err := w.Write(value) + if err != nil { + return err + } + if written <= 0 || written > len(value) { + return io.ErrShortWrite + } + value = value[written:] + } + return nil +} + +func writeString(w io.Writer, value string) error { + if stringWriter, ok := w.(io.StringWriter); ok { + written, err := stringWriter.WriteString(value) + if err != nil { + return err + } + if written != len(value) { + return io.ErrShortWrite + } + return nil + } + return writeBytes(w, []byte(value)) +} + +func writeByte(w io.Writer, value byte) error { + if byteWriter, ok := w.(io.ByteWriter); ok { + return byteWriter.WriteByte(value) + } + buffer := [1]byte{value} + return writeBytes(w, buffer[:]) +} diff --git a/pkg/container/bytejson/bytejson_text_writer_test.go b/pkg/container/bytejson/bytejson_text_writer_test.go new file mode 100644 index 0000000000000..5519cc477c1b5 --- /dev/null +++ b/pkg/container/bytejson/bytejson_text_writer_test.go @@ -0,0 +1,60 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package bytejson + +import ( + "bytes" + "encoding/base64" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestWriteJSONTextMatchesMarshalJSON(t *testing.T) { + nested, err := CreateByteJSON(map[string]any{ + "array": []any{nil, true, int64(-7), uint64(9), 1.25, "a\n\"中"}, + "empty": map[string]any{}, + }) + require.NoError(t, err) + raw := []byte{0, 1, 2, 3, 250, 251, 252} + values := []ByteJson{ + nested, + {Type: TpCodeOpaque, Data: appendBinaryString(nil, string(raw))}, + {Type: TpCodeBit, Data: appendBinaryString(nil, string(raw))}, + { + Type: TpCodeBlob, + Data: appendBinaryString( + nil, + persistedBitPrefix+base64.StdEncoding.EncodeToString(raw), + ), + }, + { + Type: TpCodeBlob, + Data: appendBinaryString(nil, persistedBitPrefix+"not-base64!"), + }, + } + for _, value := range values { + want, err := value.MarshalJSON() + require.NoError(t, err) + var got bytes.Buffer + require.NoError(t, WriteJSONText(&got, value)) + require.Equal(t, want, got.Bytes()) + } +} + +func TestWriteJSONStringRejectsInvalidUTF8(t *testing.T) { + var output bytes.Buffer + require.Error(t, WriteJSONString(&output, []byte{'a', 0xff})) +} diff --git a/pkg/sql/plan/function/func_binary.go b/pkg/sql/plan/function/func_binary.go index e9efab50fe848..f22eaf04c2e22 100644 --- a/pkg/sql/plan/function/func_binary.go +++ b/pkg/sql/plan/function/func_binary.go @@ -25,6 +25,7 @@ import ( "encoding/hex" "encoding/json" "fmt" + "io" "math" "math/big" "math/bits" @@ -33,12 +34,14 @@ import ( "strconv" "strings" "time" + "unicode/utf8" "github.com/matrixorigin/matrixone/pkg/util/fault" "go.uber.org/zap" "github.com/matrixorigin/matrixone/pkg/clusterservice" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/util" "github.com/matrixorigin/matrixone/pkg/container/nulls" "github.com/matrixorigin/matrixone/pkg/container/types" @@ -3684,7 +3687,7 @@ func DateFormat(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro } //format := "%b %D %M" -> []func{func1,func2, func3} - var buf bytes.Buffer + var legacy bytes.Buffer for i := uint64(0); i < uint64(length); i++ { d, null1 := dates.GetValue(i) if null1 || null2 { @@ -3692,15 +3695,19 @@ func DateFormat(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro return err } } else { - buf.Reset() var isNull bool - if isNull, err = dateFmtOperator(proc.Ctx, d, string(fmt), &buf); err != nil { + if isNull, err = appendFormattedBytesForResult( + result, + rs, + &legacy, + func(buf formatBuffer) (bool, error) { + return dateFmtOperator(proc.Ctx, d, string(fmt), buf) + }, + ); err != nil { return err } if isNull { err = rs.AppendBytes(nil, true) - } else { - err = rs.AppendBytes(buf.Bytes(), false) } if err != nil { return err @@ -3710,11 +3717,117 @@ func DateFormat(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro return nil } -type DateFormatFunc func(ctx context.Context, datetime types.Datetime, format string, buf *bytes.Buffer) (isNull bool, err error) +type formatBuffer interface { + io.Writer + io.StringWriter + io.ByteWriter + WriteRune(rune) (int, error) + Grow(int) +} + +type countingFormatBuffer struct { + written int + err error +} + +func (w *countingFormatBuffer) add(size int) (int, error) { + if w.err != nil { + return 0, w.err + } + if size < 0 || size > math.MaxInt-w.written { + w.err = io.ErrShortBuffer + return 0, w.err + } + w.written += size + return size, nil +} + +func (w *countingFormatBuffer) Write(value []byte) (int, error) { + return w.add(len(value)) +} + +func (w *countingFormatBuffer) WriteString(value string) (int, error) { + return w.add(len(value)) +} + +func (w *countingFormatBuffer) WriteByte(byte) error { + _, err := w.add(1) + return err +} + +func (w *countingFormatBuffer) WriteRune(value rune) (int, error) { + size := utf8.RuneLen(value) + if size < 0 { + size = utf8.RuneLen(utf8.RuneError) + } + return w.add(size) +} + +func (w *countingFormatBuffer) Grow(int) {} + +func appendFormattedBytes( + rs *vector.FunctionResult[types.Varlena], + build func(formatBuffer) (bool, error), +) (bool, error) { + var counter countingFormatBuffer + isNull, err := build(&counter) + if err != nil || isNull { + return isNull, err + } + if counter.err != nil { + return false, counter.err + } + var output fixedSliceWriter + err = rs.AppendBytesWithBuilder(counter.written, func(dst []byte) (int, error) { + output.Reset(dst) + secondNull, buildErr := build(&output) + if buildErr != nil { + return 0, buildErr + } + if output.Err() != nil { + return 0, output.Err() + } + if secondNull { + return 0, moerr.NewInternalErrorNoCtx( + "format result changed between sizing and encoding", + ) + } + if output.Written() != counter.written { + return 0, moerr.NewInternalErrorNoCtx( + "format result size changed between sizing and encoding", + ) + } + return output.Written(), nil + }) + return false, err +} + +// appendFormattedBytesForResult preserves the legacy one-pass buffer path and +// uses exact two-pass publication only when dormant allocation accounting is +// selected. This avoids imposing duplicate formatting work on production +// callers before the expression owner is activated. +func appendFormattedBytesForResult( + result vector.FunctionResultWrapper, + rs *vector.FunctionResult[types.Varlena], + legacy *bytes.Buffer, + build func(formatBuffer) (bool, error), +) (bool, error) { + if result.HasFunctionScratch() { + return appendFormattedBytes(rs, build) + } + legacy.Reset() + isNull, err := build(legacy) + if err != nil || isNull { + return isNull, err + } + return false, rs.AppendBytes(legacy.Bytes(), false) +} + +type DateFormatFunc func(ctx context.Context, datetime types.Datetime, format string, buf formatBuffer) (isNull bool, err error) // DATE_FORMAT datetime // handle '%d/%m/%Y' -> 22/04/2021 -func date_format_combine_pattern1(_ context.Context, t types.Datetime, format string, buf *bytes.Buffer) (bool, error) { +func date_format_combine_pattern1(_ context.Context, t types.Datetime, format string, buf formatBuffer) (bool, error) { month := int(t.Month()) day := int(t.Day()) year := int(t.Year()) @@ -3739,7 +3852,7 @@ func date_format_combine_pattern1(_ context.Context, t types.Datetime, format st } // handle '%Y%m%d' -> 20210422 -func date_format_combine_pattern2(_ context.Context, t types.Datetime, format string, buf *bytes.Buffer) (bool, error) { +func date_format_combine_pattern2(_ context.Context, t types.Datetime, format string, buf formatBuffer) (bool, error) { year := t.Year() month := int(t.Month()) day := int(t.Day()) @@ -3764,7 +3877,7 @@ func date_format_combine_pattern2(_ context.Context, t types.Datetime, format st } // handle '%Y' -> 2021 -func date_format_combine_pattern3(_ context.Context, t types.Datetime, format string, buf *bytes.Buffer) (bool, error) { +func date_format_combine_pattern3(_ context.Context, t types.Datetime, format string, buf formatBuffer) (bool, error) { year := t.Year() // Year conversion buf.WriteByte(byte('0' + (year / 1000 % 10))) @@ -3775,7 +3888,7 @@ func date_format_combine_pattern3(_ context.Context, t types.Datetime, format st } // %Y-%m-%d 2021-04-22 -func date_format_combine_pattern4(_ context.Context, t types.Datetime, format string, buf *bytes.Buffer) (bool, error) { +func date_format_combine_pattern4(_ context.Context, t types.Datetime, format string, buf formatBuffer) (bool, error) { year := t.Year() month := int(t.Month()) day := int(t.Day()) @@ -3803,7 +3916,7 @@ func date_format_combine_pattern4(_ context.Context, t types.Datetime, format st // handle '%Y-%m-%d %H:%i:%s' -> 2004-04-03 13:11:10 // handle ' %Y-%m-%d %T' -> 2004-04-03 13:11:10 -func date_format_combine_pattern5(_ context.Context, t types.Datetime, format string, buf *bytes.Buffer) (bool, error) { +func date_format_combine_pattern5(_ context.Context, t types.Datetime, format string, buf formatBuffer) (bool, error) { year := int(t.Year()) month := int(t.Month()) day := int(t.Day()) @@ -3853,7 +3966,7 @@ func date_format_combine_pattern5(_ context.Context, t types.Datetime, format st } // handle '%Y/%m/%d' -> 2010/01/07 -func date_format_combine_pattern6(_ context.Context, t types.Datetime, format string, buf *bytes.Buffer) (bool, error) { +func date_format_combine_pattern6(_ context.Context, t types.Datetime, format string, buf formatBuffer) (bool, error) { year := t.Year() month := int(t.Month()) day := int(t.Day()) @@ -3881,7 +3994,7 @@ func date_format_combine_pattern6(_ context.Context, t types.Datetime, format st // handle '%Y/%m/%d %H:%i:%s' -> 2010/01/07 23:12:34 // handle '%Y/%m/%d %T' -> 2010/01/07 23:12:34 -func date_format_combine_pattern7(_ context.Context, t types.Datetime, format string, buf *bytes.Buffer) (bool, error) { +func date_format_combine_pattern7(_ context.Context, t types.Datetime, format string, buf formatBuffer) (bool, error) { year := int(t.Year()) month := int(t.Month()) day := int(t.Day()) @@ -3931,7 +4044,7 @@ func date_format_combine_pattern7(_ context.Context, t types.Datetime, format st } // datetimeFormat: format the datetime value according to the format string. -func datetimeFormat(ctx context.Context, datetime types.Datetime, format string, buf *bytes.Buffer) (bool, error) { +func datetimeFormat(ctx context.Context, datetime types.Datetime, format string, buf formatBuffer) (bool, error) { inPatternMatch := false for _, b := range format { if inPatternMatch { @@ -4003,7 +4116,7 @@ var ( ) // makeDateFormat: Get the format string corresponding to the date according to a single format character -func makeDateFormat(_ context.Context, t types.Datetime, b rune, buf *bytes.Buffer) (bool, error) { +func makeDateFormat(_ context.Context, t types.Datetime, b rune, buf formatBuffer) (bool, error) { switch b { case 'b': m := t.Month() @@ -4169,7 +4282,7 @@ func TimeFormat(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro fmt, null2 := formats.GetStrValue(0) emptyFormat := len(fmt) == 0 - var buf bytes.Buffer + var legacy bytes.Buffer for i := uint64(0); i < uint64(length); i++ { t, null1 := times.GetValue(i) if null1 || null2 || emptyFormat { @@ -4177,11 +4290,15 @@ func TimeFormat(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro return err } } else { - buf.Reset() - if err = timeFormat(proc.Ctx, t, string(fmt), &buf); err != nil { - return err - } - if err = rs.AppendBytes(buf.Bytes(), false); err != nil { + _, err = appendFormattedBytesForResult( + result, + rs, + &legacy, + func(buf formatBuffer) (bool, error) { + return false, timeFormat(proc.Ctx, t, string(fmt), buf) + }, + ) + if err != nil { return err } } @@ -4191,7 +4308,7 @@ func TimeFormat(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro // timeFormat: Get the format string corresponding to the time according to format specifiers // Only supports time-related format specifiers: %H, %h, %I, %i, %k, %l, %S, %s, %f, %p, %r, %T -func timeFormat(ctx context.Context, t types.Time, format string, buf *bytes.Buffer) error { +func timeFormat(ctx context.Context, t types.Time, format string, buf formatBuffer) error { hour, minute, sec, msec, isNeg := t.ClockFormat() if isNeg && len(format) > 0 { buf.WriteByte('-') @@ -4218,7 +4335,7 @@ func timeFormat(ctx context.Context, t types.Time, format string, buf *bytes.Buf // makeTimeFormat: Get the format string corresponding to the time according to a single format character // Only supports time-related format specifiers -func makeTimeFormat(ctx context.Context, hour uint64, minute, sec uint8, msec uint64, b rune, buf *bytes.Buffer) error { +func makeTimeFormat(ctx context.Context, hour uint64, minute, sec uint8, msec uint64, b rune, buf formatBuffer) error { switch b { case 'f': fmt.Fprintf(buf, "%06d", msec) @@ -4288,7 +4405,7 @@ func FormatIntByWidth(num, n int) string { return builder.String() } -func FormatInt2BufByWidth(num, n int, buf *bytes.Buffer) { +func FormatInt2BufByWidth(num, n int, buf formatBuffer) { numStr := strconv.Itoa(num) if len(numStr) >= n { buf.WriteString(numStr) @@ -5231,21 +5348,46 @@ func MakeSet(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc * continue } - // Build the result string by checking each bit position - var parts []string + resultSize := 0 + partCount := 0 for j := 0; j < len(strParams); j++ { // Check if bit j is set (0-based, so bit 0 corresponds to str1, bit 1 to str2, etc.) if (bitsUint>>uint(j))&1 == 1 { str, null := strParams[j].GetStrValue(i) if !null { - parts = append(parts, functionUtil.QuickBytesToStr(str)) + if len(str) > math.MaxInt-resultSize { + return moerr.NewInvalidInputNoCtx("MAKE_SET result is too large") + } + resultSize += len(str) + partCount++ } } } - - // Join with comma separator - resultStr := strings.Join(parts, ",") - if err := rs.AppendBytes([]byte(resultStr), false); err != nil { + if partCount > 1 { + if partCount-1 > math.MaxInt-resultSize { + return moerr.NewInvalidInputNoCtx("MAKE_SET result is too large") + } + resultSize += partCount - 1 + } + if err := rs.AppendBytesWithFill(resultSize, func(dst []byte) { + written := 0 + parts := 0 + for j := 0; j < len(strParams); j++ { + if (bitsUint>>uint(j))&1 == 0 { + continue + } + str, null := strParams[j].GetStrValue(i) + if null { + continue + } + if parts > 0 { + dst[written] = ',' + written++ + } + written += copy(dst[written:], str) + parts++ + } + }); err != nil { return err } } @@ -5525,18 +5667,38 @@ func ExportSet(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc } } - // Build the result string - var parts []string + resultSize := 0 for j := int64(0); j < numberOfBits; j++ { + partSize := len(off) if (bitsUint>>uint(j))&1 == 1 { - parts = append(parts, functionUtil.QuickBytesToStr(on)) - } else { - parts = append(parts, functionUtil.QuickBytesToStr(off)) + partSize = len(on) + } + if partSize > math.MaxInt-resultSize { + return moerr.NewInvalidInputNoCtx("EXPORT_SET result is too large") } + resultSize += partSize } - - resultStr := strings.Join(parts, separator) - if err := rs.AppendBytes([]byte(resultStr), false); err != nil { + separatorBytes := functionUtil.QuickStrToBytes(separator) + if numberOfBits > 1 { + separatorSize := uint64(numberOfBits-1) * uint64(len(separatorBytes)) + if separatorSize > uint64(math.MaxInt-resultSize) { + return moerr.NewInvalidInputNoCtx("EXPORT_SET result is too large") + } + resultSize += int(separatorSize) + } + if err := rs.AppendBytesWithFill(resultSize, func(dst []byte) { + written := 0 + for j := int64(0); j < numberOfBits; j++ { + if j > 0 { + written += copy(dst[written:], separatorBytes) + } + part := off + if (bitsUint>>uint(j))&1 == 1 { + part = on + } + written += copy(dst[written:], part) + } + }); err != nil { return err } } @@ -5894,7 +6056,7 @@ func FromUnixTimeInt64Format(ivecs []*vector.Vector, result vector.FunctionResul formatMask, null1 := vector.GenerateFunctionStrParameter(ivecs[1]).GetStrValue(0) f := string(formatMask) - var buf bytes.Buffer + var legacy bytes.Buffer for i := uint64(0); i < uint64(length); i++ { v, null := vs.GetValue(i) @@ -5903,12 +6065,10 @@ func FromUnixTimeInt64Format(ivecs []*vector.Vector, result vector.FunctionResul return err } } else { - buf.Reset() r := types.DatetimeFromUnix(proc.GetSessionInfo().TimeZone, v) - if _, err = datetimeFormat(proc.Ctx, r, f, &buf); err != nil { - return err - } - if err = rs.AppendBytes(buf.Bytes(), false); err != nil { + if _, err = appendFormattedBytesForResult(result, rs, &legacy, func(buf formatBuffer) (bool, error) { + return datetimeFormat(proc.Ctx, r, f, buf) + }); err != nil { return err } } @@ -5926,7 +6086,7 @@ func FromUnixTimeUint64Format(ivecs []*vector.Vector, result vector.FunctionResu formatMask, null1 := vector.GenerateFunctionStrParameter(ivecs[1]).GetStrValue(0) f := string(formatMask) - var buf bytes.Buffer + var legacy bytes.Buffer for i := uint64(0); i < uint64(length); i++ { v, null := vs.GetValue(i) @@ -5935,12 +6095,10 @@ func FromUnixTimeUint64Format(ivecs []*vector.Vector, result vector.FunctionResu return err } } else { - buf.Reset() r := types.DatetimeFromUnix(proc.GetSessionInfo().TimeZone, int64(v)) - if _, err = datetimeFormat(proc.Ctx, r, f, &buf); err != nil { - return err - } - if err = rs.AppendBytes(buf.Bytes(), false); err != nil { + if _, err = appendFormattedBytesForResult(result, rs, &legacy, func(buf formatBuffer) (bool, error) { + return datetimeFormat(proc.Ctx, r, f, buf) + }); err != nil { return err } } @@ -5958,7 +6116,7 @@ func FromUnixTimeFloat64Format(ivecs []*vector.Vector, result vector.FunctionRes formatMask, null1 := vector.GenerateFunctionStrParameter(ivecs[1]).GetStrValue(0) f := string(formatMask) - var buf bytes.Buffer + var legacy bytes.Buffer for i := uint64(0); i < uint64(length); i++ { v, null := vs.GetValue(i) @@ -5967,13 +6125,11 @@ func FromUnixTimeFloat64Format(ivecs []*vector.Vector, result vector.FunctionRes return err } } else { - buf.Reset() x, y := splitDecimalToIntAndFrac(v) r := types.DatetimeFromUnixWithNsec(proc.GetSessionInfo().TimeZone, x, y) - if _, err = datetimeFormat(proc.Ctx, r, f, &buf); err != nil { - return err - } - if err = rs.AppendBytes(buf.Bytes(), false); err != nil { + if _, err = appendFormattedBytesForResult(result, rs, &legacy, func(buf formatBuffer) (bool, error) { + return datetimeFormat(proc.Ctx, r, f, buf) + }); err != nil { return err } } @@ -5992,7 +6148,7 @@ func FromUnixTimeDecimal256Format(ivecs []*vector.Vector, result vector.Function formatMask, null1 := vector.GenerateFunctionStrParameter(ivecs[1]).GetStrValue(0) f := string(formatMask) - var buf bytes.Buffer + var legacy bytes.Buffer for i := uint64(0); i < uint64(length); i++ { v, null := vs.GetValue(i) sec, nsec, ok, convErr := decimal256UnixTimeParts(v, scale) @@ -6005,12 +6161,10 @@ func FromUnixTimeDecimal256Format(ivecs []*vector.Vector, result vector.Function return err } } else { - buf.Reset() r := types.DatetimeFromUnixWithNsec(proc.GetSessionInfo().TimeZone, sec, nsec) - if _, err = datetimeFormat(proc.Ctx, r, f, &buf); err != nil { - return err - } - if err = rs.AppendBytes(buf.Bytes(), false); err != nil { + if _, err = appendFormattedBytesForResult(result, rs, &legacy, func(buf formatBuffer) (bool, error) { + return datetimeFormat(proc.Ctx, r, f, buf) + }); err != nil { return err } } @@ -8605,6 +8759,7 @@ func batchArrayDistanceSync[T types.RealNumbers]( m metric.MetricType, proc *process.Process, dist []float32, + result vector.FunctionResultWrapper, ) ([]float32, bool, error) { c0, c1 := ivecs[0].IsConst(), ivecs[1].IsConst() if c0 == c1 { @@ -8639,7 +8794,28 @@ func batchArrayDistanceSync[T types.RealNumbers]( resolver = proc.GetResolveVariableFunc() } gpuMode := gpumode.EffectiveGpuMode(resolver) - handle, err := metric.PairwiseDistanceLaunchOneToMany( + scratchSize, usesGPU, err := metric.PairwiseDistanceOneToManyScratchSize( + query, + length, + m, + metric.GPUThresholdSQL, + gpuMode, + ) + if err != nil { + return nil, false, err + } + var scratch []byte + if usesGPU && result != nil { + var selected bool + scratch, selected, err = result.ResizeFunctionScratch(scratchSize) + if err != nil { + return nil, false, err + } + if !selected { + scratch = nil + } + } + handle, err := metric.PairwiseDistanceLaunchOneToManyWithScratch( query, length, func(row int) []T { @@ -8649,6 +8825,7 @@ func batchArrayDistanceSync[T types.RealNumbers]( dist, metric.GPUThresholdSQL, gpuMode, + scratch, ) if err != nil { return nil, false, err @@ -8688,6 +8865,7 @@ func tryBatchArrayDistance[T types.RealNumbers]( m, proc, distScratch, + result, ) if err != nil || !ok { return ok, err @@ -12304,6 +12482,28 @@ func arrayDistanceNarrow[T types.ArrayElement]( func arrayDistanceViaF32[T types.ArrayElement]( ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList, kernel func(v1, v2 []float32) (float64, error)) error { + if result.HasFunctionScratch() { + return opBinaryBytesBytesToFixedWithErrorCheck[float64](ivecs, result, proc, length, func(v1, v2 []byte) (float64, error) { + left := types.BytesToArray[T](v1) + right := types.BytesToArray[T](v2) + if len(right) > math.MaxInt-len(left) || len(left)+len(right) > math.MaxInt/4 { + return 0, mpool.ErrAllocationAccountInvalid + } + scratch, selected, err := result.ResizeFunctionScratch((len(left) + len(right)) * 4) + if err != nil { + return 0, err + } + if !selected { + return 0, mpool.ErrAllocationAccountInvalid + } + values := util.UnsafeSliceCast[float32](scratch) + leftValues := values[:len(left)] + rightValues := values[len(left):] + arrayToFloat32Into(leftValues, left) + arrayToFloat32Into(rightValues, right) + return kernel(leftValues, rightValues) + }, selectList) + } return opBinaryBytesBytesToFixedWithErrorCheck[float64](ivecs, result, proc, length, func(v1, v2 []byte) (float64, error) { f1 := types.ToFloat32Array[T](types.BytesToArray[T](v1)) f2 := types.ToFloat32Array[T](types.BytesToArray[T](v2)) @@ -12311,6 +12511,35 @@ func arrayDistanceViaF32[T types.ArrayElement]( }, selectList) } +func arrayToFloat32Into[T types.ArrayElement](dst []float32, src []T) { + switch values := any(src).(type) { + case []float32: + copy(dst, values) + case []float64: + for idx, value := range values { + dst[idx] = float32(value) + } + case []types.BF16: + for idx, value := range values { + dst[idx] = value.ToFloat32() + } + case []types.Float16: + for idx, value := range values { + dst[idx] = value.ToFloat32() + } + case []int8: + for idx, value := range values { + dst[idx] = float32(value) + } + case []uint8: + for idx, value := range values { + dst[idx] = float32(value) + } + default: + panic(moerr.NewInternalErrorNoCtx("unsupported array element type")) + } +} + func L2DistanceArrayViaF32[T types.ArrayElement](ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { return arrayDistanceNarrow[T](ivecs, result, proc, length, selectList, metric.Metric_L2Distance, true) } diff --git a/pkg/sql/plan/function/func_binary_array_distance_test.go b/pkg/sql/plan/function/func_binary_array_distance_test.go index ea0fc73128ef1..0810d0bec2b54 100644 --- a/pkg/sql/plan/function/func_binary_array_distance_test.go +++ b/pkg/sql/plan/function/func_binary_array_distance_test.go @@ -77,6 +77,7 @@ func testBatchArrayDistanceSync[T types.RealNumbers]( m, nil, make([]float32, length), + nil, ) } @@ -109,6 +110,7 @@ func BenchmarkBatchArrayDistanceSync8192(b *testing.B) { metric.Metric_L2sqDistance, nil, distScratch, + nil, ) require.NoError(b, err) require.True(b, ok) diff --git a/pkg/sql/plan/function/func_builtin.go b/pkg/sql/plan/function/func_builtin.go index 5bb92efb50ed5..da66af88d3aa7 100644 --- a/pkg/sql/plan/function/func_builtin.go +++ b/pkg/sql/plan/function/func_builtin.go @@ -29,6 +29,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/hashmap" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/runtime" "github.com/matrixorigin/matrixone/pkg/common/util" "github.com/matrixorigin/matrixone/pkg/config" @@ -232,8 +233,8 @@ func parseLeadingInteger(s string) (int64, bool) { return v, true } -// encodeCharBytes converts an int64 argument for MySQL CHAR() into big-endian -// bytes. MySQL treats CHAR(N) values as unsigned 32-bit integers and expands +// appendCharBytes appends one MySQL CHAR() integer as big-endian bytes. MySQL +// treats CHAR(N) values as unsigned 32-bit integers and expands // values > 255 into multiple big-endian bytes: // // CHAR(256) → 0x0100 (two bytes) @@ -241,10 +242,10 @@ func parseLeadingInteger(s string) (int64, bool) { // CHAR(-1) → 0xFFFFFFFF (four bytes, via two's complement uint32) // // See MySQL docs: https://dev.mysql.com/doc/refman/8.4/en/string-functions.html#function_char -func encodeCharBytes(v int64) []byte { +func appendCharBytes(dst []byte, v int64) []byte { uv := uint32(v) if uv == 0 { - return []byte{0} + return append(dst, 0) } // Encode as big-endian 32-bit, then strip leading zero bytes. var buf [4]byte @@ -257,7 +258,7 @@ func encodeCharBytes(v int64) []byte { for start < 3 && buf[start] == 0 { start++ } - return buf[start:] + return append(dst, buf[start:]...) } const ( @@ -1239,25 +1240,21 @@ func builtInChar(parameters []*vector.Vector, result vector.FunctionResultWrappe continue } - var resultBytes []byte - - // Process all arguments - for _, getter := range getters { - v, null := getter(i) - if null { - // MySQL skips NULL arguments instead of returning NULL - continue - } - // Convert argument to big-endian multi-byte sequence. - // MySQL treats CHAR(N) as unsigned 32-bit ints and expands - // values > 255 into multiple big-endian bytes (CHAR(256) → 0x0100). - // Negative values use two's complement uint32 (CHAR(-1) → 0xFFFFFFFF). - resultBytes = append(resultBytes, encodeCharBytes(v)...) + if len(getters) > math.MaxInt/4 { + return moerr.NewInvalidInputNoCtx("CHAR has too many arguments") } - - // resultBytes is empty only when every argument was NULL. MySQL returns - // an empty (non-NULL) string in that case, e.g. CHAR(NULL, NULL) -> ''. - if err := rs.AppendBytes(resultBytes, false); err != nil { + // MySQL skips NULL arguments and expands each remaining value to at + // most four bytes. Build directly in the admitted result backing. + if err := rs.AppendBytesWithBuilder(len(getters)*4, func(dst []byte) (int, error) { + output := dst[:0] + for _, getter := range getters { + value, null := getter(i) + if !null { + output = appendCharBytes(output, value) + } + } + return len(output), nil + }); err != nil { return err } } @@ -2085,8 +2082,95 @@ func builtInUnixTimestampVarcharToDecimal128(parameters []*vector.Vector, result return nil } +func builtInHashAccounted( + parameters []*vector.Vector, + result vector.FunctionResultWrapper, + length int, + appendState func(uint64), +) error { + var keys [hashmap.UnitLimit][]byte + var states [hashmap.UnitLimit][3]uint64 + var keySizes [hashmap.UnitLimit]int + for start := 0; start < length; start += hashmap.UnitLimit { + count := min(length-start, hashmap.UnitLimit) + total := 0 + for localRow := 0; localRow < count; localRow++ { + row := start + localRow + size := 0 + for _, parameter := range parameters { + if size == math.MaxInt { + return moerr.NewInvalidInputNoCtx("HASH input is too large") + } + size++ // one NULL marker per parameter + if !parameter.IsNull(uint64(row)) { + valueSize := len(parameter.GetRawBytesAt(row)) + if valueSize > math.MaxInt-size { + return moerr.NewInvalidInputNoCtx("HASH input is too large") + } + size += valueSize + } + } + if size < len(hashtable.StrKeyPadding) { + size = len(hashtable.StrKeyPadding) + } + if size > math.MaxInt-total { + return moerr.NewInvalidInputNoCtx("HASH input is too large") + } + keySizes[localRow] = size + total += size + } + + scratch, selected, err := result.ResizeFunctionScratch(total) + if err != nil { + return err + } + if !selected { + return mpool.ErrAllocationAccountInvalid + } + offset := 0 + for localRow := 0; localRow < count; localRow++ { + size := keySizes[localRow] + keys[localRow] = scratch[offset : offset+size] + offset += size + } + for localRow := 0; localRow < count; localRow++ { + row := start + localRow + key := keys[localRow] + written := 0 + for _, parameter := range parameters { + isNull := parameter.IsNull(uint64(row)) + if isNull { + key[written] = 1 + written++ + continue + } + key[written] = 0 + written++ + written += copy(key[written:], parameter.GetRawBytesAt(row)) + } + if written < len(hashtable.StrKeyPadding) { + copy(key[written:], hashtable.StrKeyPadding[written:]) + } + } + hashtable.BytesBatchGenHashStates(&keys[0], &states[0], count) + for localRow := 0; localRow < count; localRow++ { + appendState(states[localRow][0]) + } + } + return nil +} + // XXX I just copy this function. func builtInHash(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { + if result.HasFunctionScratch() { + rs := vector.MustFunctionResult[int64](result) + return builtInHashAccounted( + parameters, + result, + length, + func(state uint64) { rs.AppendMustValue(int64(state)) }, + ) + } fillStringGroupStr := func(keys [][]byte, vec *vector.Vector, n int, start int) { if vec.IsConst() { area := vec.GetArea() @@ -2200,6 +2284,15 @@ func builtInHash(parameters []*vector.Vector, result vector.FunctionResultWrappe // builtInHashPartition mirrors builtInHash but returns uint64 so downstream modulo results are non-negative. func builtInHashPartition(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { + if result.HasFunctionScratch() { + rs := vector.MustFunctionResult[uint64](result) + return builtInHashAccounted( + parameters, + result, + length, + rs.AppendMustValue, + ) + } fillStringGroupStr := func(keys [][]byte, vec *vector.Vector, n int, start int) { if vec.IsConst() { area := vec.GetArea() diff --git a/pkg/sql/plan/function/func_builtin_jq.go b/pkg/sql/plan/function/func_builtin_jq.go index 5d1b5a2481e03..9c76af9ae3a8f 100644 --- a/pkg/sql/plan/function/func_builtin_jq.go +++ b/pkg/sql/plan/function/func_builtin_jq.go @@ -18,7 +18,9 @@ import ( "bytes" "cmp" "encoding/json" + "errors" "fmt" + "io" "math" "math/big" "slices" @@ -27,8 +29,11 @@ import ( "github.com/itchyny/gojq" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/bytejson" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/sql/plan/function/functionUtil" "github.com/matrixorigin/matrixone/pkg/vm/process" "golang.org/x/exp/constraints" ) @@ -72,6 +77,13 @@ func (op *opBuiltInJq) tryJq(params []*vector.Vector, result vector.FunctionResu func (op *opBuiltInJq) tryJqImpl(params []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList, isTry bool) error { + var scratchOutput functionScratchOutput + if result.HasFunctionScratch() { + scratchOutput.result = result + op.enc.useWriter(&scratchOutput) + defer op.enc.restoreWriter() + } + p1 := vector.GenerateFunctionStrParameter(params[0]) p2 := vector.GenerateFunctionStrParameter(params[1]) rs := vector.MustFunctionResult[types.Varlena](result) @@ -95,14 +107,16 @@ func (op *opBuiltInJq) tryJqImpl(params []*vector.Vector, result vector.Function err = op.jqImpl(v1, code) } if err != nil { - if isTry { + if isTry && !isJqOutputError(err) { rs.AddNullRange(0, uint64(length)) return nil } else { return err } } - rs.AppendBytes(op.enc.bytes(), false) + if err := rs.AppendBytes(op.enc.bytes(), false); err != nil { + return err + } op.enc.done() } return nil @@ -117,20 +131,26 @@ func (op *opBuiltInJq) tryJqImpl(params []*vector.Vector, result vector.Function for i := uint64(0); i < uint64(length); i++ { v2, null2 := p2.GetStrValue(i) if null2 || selectList.Contains(i) { - rs.AppendBytes(nil, true) + if err := rs.AppendBytes(nil, true); err != nil { + return err + } } else { code, err := op.getJqCode(string(v2)) if err == nil { err = op.jqImpl(v1, code) } if err != nil { - if isTry { - rs.AppendBytes(nil, true) + if isTry && !isJqOutputError(err) { + if err := rs.AppendBytes(nil, true); err != nil { + return err + } } else { return err } } else { - rs.AppendBytes(op.enc.bytes(), false) + if err := rs.AppendBytes(op.enc.bytes(), false); err != nil { + return err + } op.enc.done() } } @@ -157,17 +177,23 @@ func (op *opBuiltInJq) tryJqImpl(params []*vector.Vector, result vector.Function for i := uint64(0); i < uint64(length); i++ { v1, null1 := p1.GetStrValue(i) if null1 || selectList.Contains(i) { - rs.AppendBytes(nil, true) + if err := rs.AppendBytes(nil, true); err != nil { + return err + } } else { err = op.jqImpl(v1, code) if err != nil { - if isTry { - rs.AppendBytes(nil, true) + if isTry && !isJqOutputError(err) { + if err := rs.AppendBytes(nil, true); err != nil { + return err + } } else { return err } } else { - rs.AppendBytes(op.enc.bytes(), false) + if err := rs.AppendBytes(op.enc.bytes(), false); err != nil { + return err + } op.enc.done() } } @@ -178,7 +204,9 @@ func (op *opBuiltInJq) tryJqImpl(params []*vector.Vector, result vector.Function v1, null1 := p1.GetStrValue(i) v2, null2 := p2.GetStrValue(i) if null1 || null2 || selectList.Contains(i) { - rs.AppendBytes(nil, true) + if err := rs.AppendBytes(nil, true); err != nil { + return err + } } else { code, err := op.getJqCode(string(v2)) if err == nil { @@ -186,14 +214,18 @@ func (op *opBuiltInJq) tryJqImpl(params []*vector.Vector, result vector.Function } if err != nil { - if isTry { - rs.AppendBytes(nil, true) + if isTry && !isJqOutputError(err) { + if err := rs.AppendBytes(nil, true); err != nil { + return err + } // continue } else { return err } } else { - rs.AppendBytes(op.enc.bytes(), false) + if err := rs.AppendBytes(op.enc.bytes(), false); err != nil { + return err + } op.enc.done() } } @@ -267,19 +299,162 @@ func (op *opBuiltInJq) getJqCode(jq string) (*gojq.Code, error) { // We removed all the terminal color related code and we write to buffer w // and do not flush until the encoding is done. type JqEncoder struct { - w bytes.Buffer + legacy jqLegacyOutput + w jqOutput tab bool indent int depth int buf [64]byte } +type jqOutputError struct { + err error +} + +func (e *jqOutputError) Error() string { return e.err.Error() } +func (e *jqOutputError) Unwrap() error { return e.err } + +func isJqOutputError(err error) bool { + var outputErr *jqOutputError + return errors.As(err, &outputErr) +} + +type jqOutput interface { + formatBuffer + Bytes() []byte + Len() int + Reset() + Err() error +} + +type jqLegacyOutput struct { + bytes.Buffer +} + +func (*jqLegacyOutput) Err() error { return nil } + +type functionScratchOutput struct { + result vector.FunctionResultWrapper + data []byte + written int + err error +} + +func (w *functionScratchOutput) ensure(required int) error { + if w.err != nil { + return w.err + } + if required < 0 { + w.err = mpool.ErrAllocationAccountInvalid + return w.err + } + if required <= cap(w.data) { + w.data = w.data[:required] + return nil + } + capacity, ok := mpool.GrowCapacity(int64(cap(w.data)), int64(required)) + if !ok || capacity > int64(math.MaxInt) { + w.err = mpool.ErrAllocationAccountInvalid + return w.err + } + data, selected, err := w.result.ResizeFunctionScratch(int(capacity)) + if err != nil { + w.err = err + return err + } + if !selected { + w.err = mpool.ErrAllocationAccountInvalid + return w.err + } + w.data = data[:required] + return nil +} + +func (w *functionScratchOutput) Write(value []byte) (int, error) { + if len(value) > math.MaxInt-w.written { + w.err = io.ErrShortBuffer + return 0, w.err + } + if err := w.ensure(w.written + len(value)); err != nil { + return 0, err + } + copy(w.data[w.written:], value) + w.written += len(value) + return len(value), nil +} + +func (w *functionScratchOutput) WriteString(value string) (int, error) { + if len(value) > math.MaxInt-w.written { + w.err = io.ErrShortBuffer + return 0, w.err + } + if err := w.ensure(w.written + len(value)); err != nil { + return 0, err + } + copy(w.data[w.written:], value) + w.written += len(value) + return len(value), nil +} + +func (w *functionScratchOutput) WriteByte(value byte) error { + if w.written == math.MaxInt { + w.err = io.ErrShortBuffer + return w.err + } + if err := w.ensure(w.written + 1); err != nil { + return err + } + w.data[w.written] = value + w.written++ + return nil +} + +func (w *functionScratchOutput) WriteRune(value rune) (int, error) { + var encoded [utf8.UTFMax]byte + size := utf8.EncodeRune(encoded[:], value) + return w.Write(encoded[:size]) +} + +func (w *functionScratchOutput) Grow(size int) { + if size < 0 || size > math.MaxInt-w.written { + w.err = io.ErrShortBuffer + return + } + _ = w.ensure(w.written + size) +} + +func (w *functionScratchOutput) Bytes() []byte { + return w.data[:w.written] +} + +func (w *functionScratchOutput) Len() int { return w.written } + +func (w *functionScratchOutput) Reset() { + w.data = w.data[:0] + w.written = 0 + w.err = nil +} + +func (w *functionScratchOutput) Err() error { return w.err } + func (e *JqEncoder) intialize(tab bool, indent int) { - e.w.Reset() + e.legacy.Reset() + e.w = &e.legacy e.tab = tab e.indent = indent } +func (e *JqEncoder) useWriter(w jqOutput) { + e.w = w + e.done() +} + +func (e *JqEncoder) restoreWriter() { + e.done() + e.w = &e.legacy + e.legacy.Reset() +} + func (e *JqEncoder) bytes() []byte { return e.w.Bytes() } @@ -288,6 +463,13 @@ func (e *JqEncoder) done() { e.depth = 0 } +func (e *JqEncoder) err() error { + if err := e.w.Err(); err != nil { + return &jqOutputError{err: err} + } + return nil +} + func (e *JqEncoder) encode(v any) error { switch v := v.(type) { case nil: @@ -317,7 +499,7 @@ func (e *JqEncoder) encode(v any) error { default: panic(fmt.Sprintf("invalid type: %[1]T (%[1]v)", v)) } - return nil + return e.err() } // ref: floatEncoder in encoding/json @@ -348,16 +530,22 @@ func (e *JqEncoder) encodeFloat64(f float64) { // ref: encodeState#string in encoding/json func (e *JqEncoder) encodeString(s string) { + e.encodeBytes(functionUtil.QuickStrToBytes(s)) +} + +// encodeBytes preserves JSON_ROW's legacy replacement behavior for invalid +// UTF-8 while avoiding a per-row []byte-to-string allocation. +func (e *JqEncoder) encodeBytes(value []byte) { e.w.WriteByte('"') start := 0 - for i := 0; i < len(s); { - if b := s[i]; b < utf8.RuneSelf { + for i := 0; i < len(value); { + if b := value[i]; b < utf8.RuneSelf { if ' ' <= b && b <= '~' && b != '"' && b != '\\' { i++ continue } if start < i { - e.w.WriteString(s[start:i]) + e.w.Write(value[start:i]) } switch b { case '"': @@ -384,20 +572,20 @@ func (e *JqEncoder) encodeString(s string) { start = i continue } - c, size := utf8.DecodeRuneInString(s[i:]) + c, size := utf8.DecodeRune(value[i:]) if c == utf8.RuneError && size == 1 { if start < i { - e.w.WriteString(s[start:i]) + e.w.Write(value[start:i]) } e.w.WriteString(`\ufffd`) - i += size + i++ start = i continue } i += size } - if start < len(s) { - e.w.WriteString(s[start:]) + if start < len(value) { + e.w.Write(value[start:]) } e.w.WriteByte('"') } @@ -476,16 +664,10 @@ func (e *JqEncoder) writeIndent() { } func (e *JqEncoder) writeIndentInternal(n int, spaces string) { - if l := len(spaces); n <= l { - e.w.WriteString(spaces[:n]) - } else { - e.w.WriteString(spaces) - for n -= l; n > 0; n, l = n-l, l*2 { - if n < l { - l = n - } - e.w.Write(e.w.Bytes()[e.w.Len()-l:]) - } + for n > 0 { + length := min(n, len(spaces)) + e.w.WriteString(spaces[:length]) + n -= length } } @@ -502,6 +684,13 @@ func newOpBuiltInJsonRow() *opBuiltInJsonRow { func (op *opBuiltInJsonRow) jsonRow(params []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { + var scratchOutput functionScratchOutput + if result.HasFunctionScratch() { + scratchOutput.result = result + op.enc.useWriter(&scratchOutput) + defer op.enc.restoreWriter() + } + rs := vector.MustFunctionResult[types.Varlena](result) if cap(op.columns) < len(params) { op.columns = make([]jsonRowColumnEncoder, len(params)) @@ -537,6 +726,9 @@ func (op *opBuiltInJsonRow) jsonRow(params []*vector.Vector, result vector.Funct } } op.enc.w.WriteByte(']') + if err := op.enc.err(); err != nil { + return err + } if err := rs.AppendBytes(op.enc.bytes(), false); err != nil { return err } @@ -634,12 +826,7 @@ func prepareJSONRowColumn( if isNull { return encodeJSONRowNull(e, row) } - jsonValue, err := types.DecodeJson(value).MarshalJSON() - if err != nil { - return err - } - e.w.Write(jsonValue) - return nil + return bytejson.WriteJSONText(e.w, types.DecodeJson(value)) }, nil case types.T_binary, types.T_varbinary, types.T_blob: return nil, moerr.NewInvalidInputf(proc.Ctx, @@ -730,8 +917,8 @@ func prepareJSONRowStringColumn(v *vector.Vector) jsonRowColumnEncoder { if isNull { return encodeJSONRowNull(e, row) } - e.encodeString(string(value)) - return nil + e.encodeBytes(value) + return e.err() } } diff --git a/pkg/sql/plan/function/func_builtin_json.go b/pkg/sql/plan/function/func_builtin_json.go index 8579a6e4c4f31..7889b3556a404 100644 --- a/pkg/sql/plan/function/func_builtin_json.go +++ b/pkg/sql/plan/function/func_builtin_json.go @@ -19,7 +19,7 @@ import ( "context" "encoding/binary" "encoding/json" - "fmt" + "math" "strconv" "strings" "time" @@ -86,11 +86,12 @@ func encodeJsonOrderingParam(value []byte) ([]byte, error) { } type opBuiltInJsonExtract struct { - allConst bool - npath int - pathStrs []string - paths []*bytejson.Path - simple bool + allConst bool + npath int + pathStrs []string + paths []*bytejson.Path + pathWrappers []vector.FunctionParameterWrapper[types.Varlena] + simple bool } type opBuiltInJsonContains struct{} @@ -979,7 +980,7 @@ func computeStringJsonRemove(json []byte, paths []*bytejson.Path) (bytejson.Byte return bj.Remove(paths) } -func (op *opBuiltInJsonExtract) buildPath(params []*vector.Vector, length int, selectList *FunctionSelectList) error { +func (op *opBuiltInJsonExtract) buildPath(params []*vector.Vector, _ int, selectList *FunctionSelectList) error { op.npath = len(params) - 1 if op.npath == 0 { return nil @@ -1025,19 +1026,20 @@ func (op *opBuiltInJsonExtract) buildPath(params []*vector.Vector, length int, s return nil } } - } else { - op.pathStrs = make([]string, op.npath*length) - op.paths = make([]*bytejson.Path, op.npath*length) + } else if len(op.pathStrs) != op.npath { + op.pathStrs = make([]string, op.npath) + op.paths = make([]*bytejson.Path, op.npath) } - // Do it! - pathWrapers := make([]vector.FunctionParameterWrapper[types.Varlena], op.npath) + if len(op.pathWrappers) != op.npath { + op.pathWrappers = make([]vector.FunctionParameterWrapper[types.Varlena], op.npath) + } for i := 0; i < op.npath; i++ { - pathWrapers[i] = vector.GenerateFunctionStrParameter(params[i+1]) + op.pathWrappers[i] = vector.GenerateFunctionStrParameter(params[i+1]) } if op.allConst { - if err := op.buildOnePath(pathWrapers, 0, op.pathStrs, op.paths); err != nil { + if _, err := op.buildOnePath(0); err != nil { return err } op.simple = true @@ -1049,63 +1051,45 @@ func (op *opBuiltInJsonExtract) buildPath(params []*vector.Vector, length int, s } return nil } else { - op.simple = true - for i := 0; i < length; i++ { - strs := op.pathStrs[i*op.npath : (i+1)*op.npath] - paths := op.paths[i*op.npath : (i+1)*op.npath] - if selectList.Contains(uint64(i)) { - for j := 0; j < op.npath; j++ { - strs[j] = "" - paths[j] = nil - } - continue - } - if err := op.buildOnePath(pathWrapers, i, strs, paths); err != nil { - return err - } - for _, p := range paths { - if p == nil { - continue - } - op.simple = op.simple && p.IsSimple() - } - } + clear(op.pathStrs) + clear(op.paths) + op.simple = false } return nil } func (op *opBuiltInJsonExtract) getPaths(i uint64) []*bytejson.Path { - if op.allConst { - return op.paths + if !op.allConst && len(op.paths) > op.npath { + return op.paths[i*uint64(op.npath) : (i+1)*uint64(op.npath)] } - return op.paths[i*uint64(op.npath) : (i+1)*uint64(op.npath)] + return op.paths } -func (op *opBuiltInJsonExtract) buildOnePath(paramWrappers []vector.FunctionParameterWrapper[types.Varlena], i int, strs []string, paths []*bytejson.Path) error { +func (op *opBuiltInJsonExtract) buildOnePath(i int) (bool, error) { skip := false - for j := 0; j < len(paramWrappers); j++ { - pathBytes, pIsNull := paramWrappers[j].GetStrValue(uint64(i)) + simple := true + for j := 0; j < len(op.pathWrappers); j++ { + pathBytes, pIsNull := op.pathWrappers[j].GetStrValue(uint64(i)) if pIsNull { skip = true break } - strs[j] = string(pathBytes) - p, err := types.ParseStringToPath(strs[j]) + op.pathStrs[j] = string(pathBytes) + p, err := types.ParseStringToPath(op.pathStrs[j]) if err != nil { - return err + return false, err } - paths[j] = &p + op.paths[j] = &p + simple = simple && p.IsSimple() } if skip { - for j := 0; j < len(paramWrappers); j++ { - strs[j] = "" - paths[j] = nil - } + clear(op.pathStrs) + clear(op.paths) } - return nil + return simple, nil } func (op *opBuiltInJsonExtract) jsonExtract(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { @@ -1151,6 +1135,13 @@ func (op *opBuiltInJsonExtract) jsonExtract(parameters []*vector.Vector, result } continue } + rowSimple := op.simple + if !op.allConst { + rowSimple, err = op.buildOnePath(int(i)) + if err != nil { + return err + } + } jsonBytes, jIsNull := jsonWrapper.GetStrValue(i) if jIsNull { if err = rs.AppendBytes(nil, true); err != nil { @@ -1166,7 +1157,15 @@ func (op *opBuiltInJsonExtract) jsonExtract(parameters []*vector.Vector, result } continue } else { - out, exists, err := fn(jsonBytes, paths) + rowFn := fn + if rowSimple { + if jsonVec.GetType().Oid == types.T_json { + rowFn = computeJsonSimpleWithExists + } else { + rowFn = computeStringSimpleWithExists + } + } + out, exists, err := rowFn(jsonBytes, paths) if err != nil { return err } @@ -1207,7 +1206,7 @@ func (op *opBuiltInJsonExtract) jsonExtractString(parameters []*vector.Vector, r return err } - if !op.simple || op.npath > 1 { + if op.allConst && (!op.simple || op.npath > 1) { return moerr.NewInvalidInput(proc.Ctx, "json_extract_string should use a path that retrives a single value") } if jsonVec.GetType().Oid == types.T_json { @@ -1223,6 +1222,16 @@ func (op *opBuiltInJsonExtract) jsonExtractString(parameters []*vector.Vector, r } continue } + rowSimple := op.simple + if !op.allConst { + rowSimple, err = op.buildOnePath(int(i)) + if err != nil { + return err + } + } + if !rowSimple || op.npath > 1 { + return moerr.NewInvalidInput(proc.Ctx, "json_extract_string should use a path that retrives a single value") + } jsonBytes, jIsNull := jsonWrapper.GetStrValue(i) if jIsNull { if err = rs.AppendBytes(nil, true); err != nil { @@ -1290,7 +1299,7 @@ func (op *opBuiltInJsonExtract) jsonExtractFloat64(parameters []*vector.Vector, if err = op.buildPath(parameters, length, selectList); err != nil { return err } - if !op.simple || op.npath > 1 { + if op.allConst && (!op.simple || op.npath > 1) { return moerr.NewInvalidInput(proc.Ctx, "json_extract_float64 should use a path that retrives a single value") } @@ -1307,6 +1316,16 @@ func (op *opBuiltInJsonExtract) jsonExtractFloat64(parameters []*vector.Vector, } continue } + rowSimple := op.simple + if !op.allConst { + rowSimple, err = op.buildOnePath(int(i)) + if err != nil { + return err + } + } + if !rowSimple || op.npath > 1 { + return moerr.NewInvalidInput(proc.Ctx, "json_extract_float64 should use a path that retrives a single value") + } jsonBytes, jIsNull := jsonWrapper.GetStrValue(i) if jIsNull { if err = rs.Append(0, true); err != nil { @@ -1738,6 +1757,13 @@ func (op *opBuiltInJsonSet) buildJsonFunction(parameters []*vector.Vector, resul jsonVec := parameters[0] jsonWrapper := vector.GenerateFunctionStrParameter(jsonVec) rs := vector.MustFunctionResult[types.Varlena](result) + valueCount := (len(parameters) - 1) / 2 + pathExprs := make([]*bytejson.Path, valueCount) + valExprs := make([]bytejson.ByteJson, valueCount) + valueEncoders := make([]bytejson.ByteJsonDataEncoder, valueCount) + defer clear(pathExprs) + defer clear(valueEncoders) + accountedScratch := result.HasFunctionScratch() if selectList.IgnoreAllRow() { for i := 0; i < length; i++ { @@ -1788,7 +1814,6 @@ rowLoop: } // build all paths - pathExprs := make([]*bytejson.Path, 0, (len(parameters)-1)/2+1) for j := 1; j < len(parameters); j += 2 { pathBytes, pIsNull := vector.GenerateFunctionStrParameter(parameters[j]).GetStrValue(uint64(i)) if pIsNull { @@ -1807,17 +1832,53 @@ rowLoop: return moerr.NewInvalidArg(proc.Ctx, jsonModifyFunctionName(jsonFuncType), "invalid path expression") } - pathExprs = append(pathExprs, &p) + pathExprs[j/2] = &p } - // build all values - valExprs := make([]bytejson.ByteJson, 0, (len(parameters)-1)/2+1) + // Build one storage-compatible representation for all values. Exact + // execution owns this backing through FunctionResult scratch; legacy + // execution preserves a row-local Go buffer. + valueBytes := 0 for j := 2; j < len(parameters); j += 2 { - val, err := op.buildJsonModifyValue(proc, parameters[j], int(i)) + encoder, err := (&opBuiltInJsonArray{}).buildValueEncoder( + proc, + parameters[j], + int(i), + ) + if err != nil { + return err + } + valueEncoders[j/2-1] = encoder + size := uint64(encoder.DataSize()) + 1 + if size > uint64(math.MaxInt-valueBytes) { + return moerr.NewInvalidArg(proc.Ctx, jsonModifyFunctionName(jsonFuncType), "JSON value is too large") + } + valueBytes += int(size) + } + var valueStorage []byte + if accountedScratch { + valueStorage, _, err = result.ResizeFunctionScratch(valueBytes) if err != nil { return err } - valExprs = append(valExprs, val) + } else { + valueStorage = make([]byte, valueBytes) + } + offset := 0 + for idx, encoder := range valueEncoders { + valueStorage[offset] = byte(encoder.TypeCode()) + size := int(encoder.DataSize()) + written, encodeErr := encoder.EncodeDataInto( + valueStorage[offset+1 : offset+1+size], + ) + if encodeErr != nil { + return encodeErr + } + if written != size { + return moerr.NewInternalErrorNoCtx("JSON value encoder size mismatch") + } + valExprs[idx] = types.DecodeJson(valueStorage[offset : offset+1+size]) + offset += size + 1 } out, err := fn(jsonBytes, pathExprs, valExprs) @@ -1850,14 +1911,6 @@ func jsonModifyFunctionName(jsonFuncType bytejson.JsonModifyType) string { } } -func (op *opBuiltInJsonSet) buildJsonModifyValue(proc *process.Process, v *vector.Vector, row int) (bytejson.ByteJson, error) { - elem, err := (&opBuiltInJsonArray{}).convertToAny(proc, v, row) - if err != nil { - return bytejson.Null, err - } - return bytejson.CreateByteJSON(elem) -} - type opBuiltInJsonArray struct{} func newOpBuiltInJsonArray() *opBuiltInJsonArray { @@ -1867,6 +1920,8 @@ func newOpBuiltInJsonArray() *opBuiltInJsonArray { func (op *opBuiltInJsonArray) jsonArray(params []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { rs := vector.MustFunctionResult[types.Varlena](result) + encoders := make([]bytejson.ByteJsonDataEncoder, len(params)) + defer clear(encoders) if selectList != nil && selectList.IgnoreAllRow() { for j := 0; j < length; j++ { @@ -1884,250 +1939,225 @@ func (op *opBuiltInJsonArray) jsonArray(params []*vector.Vector, result vector.F } continue } - elems := make([]any, 0, len(params)) - for i := 0; i < len(params); i++ { - elem, err := op.convertToAny(proc, params[i], j) + for i := range params { + encoder, err := op.buildValueEncoder(proc, params[i], j) if err != nil { return err } - elems = append(elems, elem) - } - - bj, err := bytejson.CreateByteJSON(elems) - if err != nil { - return err + encoders[i] = encoder } - dt, err := bj.Marshal() + encoder, err := bytejson.NewArrayDataEncoder(encoders) if err != nil { return err } - if err := rs.AppendBytes(dt, false); err != nil { + if err := rs.AppendByteJsonEncoded(encoder); err != nil { return err } } return nil } -func (op *opBuiltInJsonArray) convertToAny(proc *process.Process, v *vector.Vector, row int) (any, error) { - ctx := context.Background() - if proc != nil { - ctx = proc.Ctx +func (op *opBuiltInJsonArray) buildValueEncoder( + proc *process.Process, + v *vector.Vector, + row int, +) (bytejson.ByteJsonDataEncoder, error) { + if v.IsNull(uint64(row)) { + return bytejson.NewLiteralDataEncoder(bytejson.LiteralNull), nil } fromType := v.GetType() switch fromType.Oid { case types.T_bool: - if v.IsNull(uint64(row)) { - return nil, nil + literal := bytejson.LiteralFalse + if vector.GetFixedAtNoTypeCheck[bool](v, row) { + literal = bytejson.LiteralTrue } - return vector.GetFixedAtNoTypeCheck[bool](v, row), nil + return bytejson.NewLiteralDataEncoder(literal), nil case types.T_int8: - if v.IsNull(uint64(row)) { - return nil, nil - } - return int64(vector.GetFixedAtNoTypeCheck[int8](v, row)), nil + return bytejson.NewInt64DataEncoder( + int64(vector.GetFixedAtNoTypeCheck[int8](v, row)), + ), nil case types.T_int16: - if v.IsNull(uint64(row)) { - return nil, nil - } - return int64(vector.GetFixedAtNoTypeCheck[int16](v, row)), nil + return bytejson.NewInt64DataEncoder( + int64(vector.GetFixedAtNoTypeCheck[int16](v, row)), + ), nil case types.T_int32: - if v.IsNull(uint64(row)) { - return nil, nil - } - return int64(vector.GetFixedAtNoTypeCheck[int32](v, row)), nil + return bytejson.NewInt64DataEncoder( + int64(vector.GetFixedAtNoTypeCheck[int32](v, row)), + ), nil case types.T_int64: - if v.IsNull(uint64(row)) { - return nil, nil - } - return vector.GetFixedAtNoTypeCheck[int64](v, row), nil + return bytejson.NewInt64DataEncoder( + vector.GetFixedAtNoTypeCheck[int64](v, row), + ), nil case types.T_uint8: - if v.IsNull(uint64(row)) { - return nil, nil - } - return uint64(vector.GetFixedAtNoTypeCheck[uint8](v, row)), nil + return bytejson.NewUint64DataEncoder( + uint64(vector.GetFixedAtNoTypeCheck[uint8](v, row)), + ), nil case types.T_uint16: - if v.IsNull(uint64(row)) { - return nil, nil - } - return uint64(vector.GetFixedAtNoTypeCheck[uint16](v, row)), nil + return bytejson.NewUint64DataEncoder( + uint64(vector.GetFixedAtNoTypeCheck[uint16](v, row)), + ), nil case types.T_uint32: - if v.IsNull(uint64(row)) { - return nil, nil - } - return uint64(vector.GetFixedAtNoTypeCheck[uint32](v, row)), nil + return bytejson.NewUint64DataEncoder( + uint64(vector.GetFixedAtNoTypeCheck[uint32](v, row)), + ), nil case types.T_uint64: - if v.IsNull(uint64(row)) { - return nil, nil - } - return vector.GetFixedAtNoTypeCheck[uint64](v, row), nil + return bytejson.NewUint64DataEncoder( + vector.GetFixedAtNoTypeCheck[uint64](v, row), + ), nil case types.T_float32: - if v.IsNull(uint64(row)) { - return nil, nil - } - return float64(vector.GetFixedAtNoTypeCheck[float32](v, row)), nil + return bytejson.NewFloat64DataEncoder( + float64(vector.GetFixedAtNoTypeCheck[float32](v, row)), + ), nil case types.T_float64: - if v.IsNull(uint64(row)) { - return nil, nil - } - return vector.GetFixedAtNoTypeCheck[float64](v, row), nil - case types.T_char, types.T_varchar, types.T_text: - if v.IsNull(uint64(row)) { - return nil, nil - } - return string(v.GetBytesAt(row)), nil + return bytejson.NewFloat64DataEncoder( + vector.GetFixedAtNoTypeCheck[float64](v, row), + ), nil + case types.T_char, types.T_varchar, types.T_text, types.T_geometry: + return bytejson.NewTypedStringDataEncoder( + bytejson.TpCodeString, + v.GetBytesAt(row), + ) case types.T_json: - if v.IsNull(uint64(row)) { - return nil, nil - } data := v.GetBytesAt(row) if len(data) == 0 { - return nil, nil + return bytejson.NewLiteralDataEncoder(bytejson.LiteralNull), nil } - bj := types.DecodeJson(data) - return bj, nil + return bytejson.NewRawDataEncoder(types.DecodeJson(data)) case types.T_date: - if v.IsNull(uint64(row)) { - return nil, nil - } - return newTypedByteJson(bytejson.TpCodeDate, vector.GetFixedAtNoTypeCheck[types.Date](v, row).String()), nil + return newJSONTypedStringEncoder( + bytejson.TpCodeDate, + vector.GetFixedAtNoTypeCheck[types.Date](v, row).String(), + ) case types.T_time: - if v.IsNull(uint64(row)) { - return nil, nil - } - return newTypedByteJson(bytejson.TpCodeTime, vector.GetFixedAtNoTypeCheck[types.Time](v, row).String2(fromType.Scale)), nil + return newJSONTypedStringEncoder( + bytejson.TpCodeTime, + vector.GetFixedAtNoTypeCheck[types.Time](v, row).String2(fromType.Scale), + ) case types.T_datetime: - if v.IsNull(uint64(row)) { - return nil, nil - } - return newTypedByteJson(bytejson.TpCodeDatetime, vector.GetFixedAtNoTypeCheck[types.Datetime](v, row).String2(fromType.Scale)), nil + return newJSONTypedStringEncoder( + bytejson.TpCodeDatetime, + vector.GetFixedAtNoTypeCheck[types.Datetime](v, row).String2(fromType.Scale), + ) case types.T_timestamp: - if v.IsNull(uint64(row)) { - return nil, nil - } - return newTypedByteJson(bytejson.TpCodeDatetime, vector.GetFixedAtNoTypeCheck[types.Timestamp](v, row).String2(jsonSessionTimeZone(proc), fromType.Scale)), nil + return newJSONTypedStringEncoder( + bytejson.TpCodeDatetime, + vector.GetFixedAtNoTypeCheck[types.Timestamp](v, row).String2( + jsonSessionTimeZone(proc), + fromType.Scale, + ), + ) case types.T_decimal64: - if v.IsNull(uint64(row)) { - return nil, nil - } - val := vector.GetFixedAtNoTypeCheck[types.Decimal64](v, row) - return newTypedByteJson(bytejson.TpCodeDecimal, string(val.Format(fromType.Scale))), nil + return newJSONTypedStringEncoder( + bytejson.TpCodeDecimal, + string(vector.GetFixedAtNoTypeCheck[types.Decimal64](v, row).Format(fromType.Scale)), + ) case types.T_decimal128: - if v.IsNull(uint64(row)) { - return nil, nil - } - val := vector.GetFixedAtNoTypeCheck[types.Decimal128](v, row) - return newTypedByteJson(bytejson.TpCodeDecimal, string(val.Format(fromType.Scale))), nil - case types.T_binary, types.T_varbinary, types.T_blob: - if v.IsNull(uint64(row)) { - return nil, nil - } - return newTypedByteJson(bytejson.TpCodeOpaque, string(v.GetBytesAt(row))), nil + return newJSONTypedStringEncoder( + bytejson.TpCodeDecimal, + string(vector.GetFixedAtNoTypeCheck[types.Decimal128](v, row).Format(fromType.Scale)), + ) case types.T_decimal256: - if v.IsNull(uint64(row)) { - return nil, nil - } - val := vector.GetFixedAtNoTypeCheck[types.Decimal256](v, row) - return newTypedByteJson(bytejson.TpCodeDecimal, string(val.Format(fromType.Scale))), nil + return newJSONTypedStringEncoder( + bytejson.TpCodeDecimal, + string(vector.GetFixedAtNoTypeCheck[types.Decimal256](v, row).Format(fromType.Scale)), + ) + case types.T_binary, types.T_varbinary, types.T_blob: + return bytejson.NewOpaqueDataEncoder(v.GetBytesAt(row)) case types.T_year: - if v.IsNull(uint64(row)) { - return nil, nil - } - val := vector.GetFixedAtNoTypeCheck[int16](v, row) - return strconv.FormatInt(int64(val), 10), nil + return newJSONTypedStringEncoder( + bytejson.TpCodeString, + strconv.FormatInt( + int64(vector.GetFixedAtNoTypeCheck[int16](v, row)), + 10, + ), + ) case types.T_bit: - if v.IsNull(uint64(row)) { - return nil, nil + width := fromType.Width + if width <= 0 { + width = 1 } - ctx := context.Background() - if proc != nil && proc.Ctx != nil { - ctx = proc.Ctx - } - return bitToJSON(vector.GetFixedAtNoTypeCheck[uint64](v, row), fromType.Width, ctx) - case types.T_enum: - if v.IsNull(uint64(row)) { - return nil, nil + if width > 64 { + ctx := context.Background() + if proc != nil && proc.Ctx != nil { + ctx = proc.Ctx + } + return nil, moerr.NewInvalidInputf(ctx, "cannot cast BIT(%d) to json", width) } - val := vector.GetFixedAtNoTypeCheck[types.Enum](v, row) - return val.String(), nil - case types.T_geometry: - if v.IsNull(uint64(row)) { - return nil, nil + value := vector.GetFixedAtNoTypeCheck[uint64](v, row) + if width < 64 { + value &= uint64(1)<= 2 && bj[0] == '"' { - key = string(bj[1 : len(bj)-1]) - } else { - key = fmt.Sprint(v) - } - default: - if bj, err := v.MarshalJSON(); err == nil { - key = string(bj) - } else { - key = fmt.Sprint(v) - } - } - case nil: - return moerr.NewInvalidInputf(proc.Ctx, "JSON documents may not contain NULL member names") - default: - key = fmt.Sprint(v) + if err := keyOutput.Err(); err != nil { + return err } + keyOffsets[entryIdx] = [2]int{start, keyOutput.Len()} - elem, err := arrayOp.convertToAny(proc, params[i+1], j) + value, err := arrayOp.buildValueEncoder(proc, params[i+1], j) if err != nil { return err } - obj[key] = elem + entries[entryIdx].Value = value } - - bj, err := bytejson.CreateByteJSON(obj) - if err != nil { - return err + keyBytes := keyOutput.Bytes() + for idx, offsets := range keyOffsets { + entries[idx].Key = keyBytes[offsets[0]:offsets[1]] } - dt, err := bj.Marshal() + encoder, err := bytejson.NewObjectDataEncoder(entries) if err != nil { return err } - if err := rs.AppendBytes(dt, false); err != nil { + if err := rs.AppendByteJsonEncoded(encoder); err != nil { return err } } return nil } +func writeJSONObjectKey( + w formatBuffer, + proc *process.Process, + v *vector.Vector, + row int, +) error { + fromType := v.GetType() + var numeric [64]byte + write := func(value []byte) error { + _, err := w.Write(value) + return err + } + switch fromType.Oid { + case types.T_bool: + return write(strconv.AppendBool( + numeric[:0], + vector.GetFixedAtNoTypeCheck[bool](v, row), + )) + case types.T_int8: + return write(strconv.AppendInt(numeric[:0], int64(vector.GetFixedAtNoTypeCheck[int8](v, row)), 10)) + case types.T_int16: + return write(strconv.AppendInt(numeric[:0], int64(vector.GetFixedAtNoTypeCheck[int16](v, row)), 10)) + case types.T_int32: + return write(strconv.AppendInt(numeric[:0], int64(vector.GetFixedAtNoTypeCheck[int32](v, row)), 10)) + case types.T_int64: + return write(strconv.AppendInt(numeric[:0], vector.GetFixedAtNoTypeCheck[int64](v, row), 10)) + case types.T_uint8: + return write(strconv.AppendUint(numeric[:0], uint64(vector.GetFixedAtNoTypeCheck[uint8](v, row)), 10)) + case types.T_uint16: + return write(strconv.AppendUint(numeric[:0], uint64(vector.GetFixedAtNoTypeCheck[uint16](v, row)), 10)) + case types.T_uint32: + return write(strconv.AppendUint(numeric[:0], uint64(vector.GetFixedAtNoTypeCheck[uint32](v, row)), 10)) + case types.T_uint64: + return write(strconv.AppendUint(numeric[:0], vector.GetFixedAtNoTypeCheck[uint64](v, row), 10)) + case types.T_float32: + return write(strconv.AppendFloat(numeric[:0], float64(vector.GetFixedAtNoTypeCheck[float32](v, row)), 'g', -1, 64)) + case types.T_float64: + return write(strconv.AppendFloat(numeric[:0], vector.GetFixedAtNoTypeCheck[float64](v, row), 'g', -1, 64)) + case types.T_char, types.T_varchar, types.T_text, types.T_geometry: + return write(v.GetBytesAt(row)) + case types.T_json: + return bytejson.WriteJSONObjectKeyText(w, types.DecodeJson(v.GetBytesAt(row))) + case types.T_date: + _, err := w.WriteString(vector.GetFixedAtNoTypeCheck[types.Date](v, row).String()) + return err + case types.T_time: + _, err := w.WriteString(vector.GetFixedAtNoTypeCheck[types.Time](v, row).String2(fromType.Scale)) + return err + case types.T_datetime: + _, err := w.WriteString(vector.GetFixedAtNoTypeCheck[types.Datetime](v, row).String2(fromType.Scale)) + return err + case types.T_timestamp: + _, err := w.WriteString(vector.GetFixedAtNoTypeCheck[types.Timestamp](v, row).String2(jsonSessionTimeZone(proc), fromType.Scale)) + return err + case types.T_decimal64: + _, err := w.WriteString(vector.GetFixedAtNoTypeCheck[types.Decimal64](v, row).Format(fromType.Scale)) + return err + case types.T_decimal128: + _, err := w.WriteString(vector.GetFixedAtNoTypeCheck[types.Decimal128](v, row).Format(fromType.Scale)) + return err + case types.T_decimal256: + _, err := w.WriteString(vector.GetFixedAtNoTypeCheck[types.Decimal256](v, row).Format(fromType.Scale)) + return err + case types.T_binary, types.T_varbinary, types.T_blob: + return bytejson.WriteJSONBase64Text(w, v.GetBytesAt(row)) + case types.T_year: + return write(strconv.AppendInt(numeric[:0], int64(vector.GetFixedAtNoTypeCheck[int16](v, row)), 10)) + case types.T_bit: + width := fromType.Width + if width <= 0 { + width = 1 + } + if width > 64 { + ctx := context.Background() + if proc != nil && proc.Ctx != nil { + ctx = proc.Ctx + } + return moerr.NewInvalidInputf(ctx, "cannot cast BIT(%d) to json", width) + } + value := vector.GetFixedAtNoTypeCheck[uint64](v, row) + if width < 64 { + value &= uint64(1)< 0 { + if err := w.WriteByte(' '); err != nil { + return err + } + } + value := strconv.AppendFloat(numeric[:0], valueAt(idx), 'g', -1, 64) + if _, err := w.Write(value); err != nil { + return err + } + } + return w.WriteByte(']') +} + type opBuiltInJsonType struct{} func newOpBuiltInJsonType() *opBuiltInJsonType { @@ -2419,19 +2572,25 @@ func jsonKeysRoot(ivecs []*vector.Vector, result vector.FunctionResultWrapper, p if selectList != nil && selectList.IgnoreAllRow() { for i := 0; i < length; i++ { - rs.AppendMustNullForBytesResult() + if err := rs.AppendMustNullForBytesResult(); err != nil { + return err + } } return nil } for i := uint64(0); i < uint64(length); i++ { if selectList.Contains(i) { - rs.AppendMustNullForBytesResult() + if err := rs.AppendMustNullForBytesResult(); err != nil { + return err + } continue } v, null := p1.GetStrValue(i) if null { - rs.AppendMustNullForBytesResult() + if err := rs.AppendMustNullForBytesResult(); err != nil { + return err + } continue } var bj bytejson.ByteJson @@ -2444,14 +2603,14 @@ func jsonKeysRoot(ivecs []*vector.Vector, result vector.FunctionResultWrapper, p if err != nil { return moerr.NewInvalidArg(proc.Ctx, "json_keys", "invalid JSON document") } - keysArray, err := buildJsonKeysArray(bj) - if err != nil || keysArray.IsNull() { + appended, err := appendJsonKeysArray(rs, bj) + if err != nil { + return err + } + if !appended { rs.AppendMustNullForBytesResult() continue } - if err := rs.AppendByteJson(keysArray, false); err != nil { - return err - } } return nil } @@ -2501,28 +2660,33 @@ func jsonKeysWithPath(ivecs []*vector.Vector, result vector.FunctionResultWrappe rs.AppendMustNullForBytesResult() continue } - keysArray, err := buildJsonKeysArray(val) - if err != nil || keysArray.IsNull() { + appended, err := appendJsonKeysArray(rs, val) + if err != nil { + return err + } + if !appended { rs.AppendMustNullForBytesResult() continue } - if err := rs.AppendByteJson(keysArray, false); err != nil { - return err - } } return nil } -func buildJsonKeysArray(bj bytejson.ByteJson) (bytejson.ByteJson, error) { +func appendJsonKeysArray( + rs *vector.FunctionResult[types.Varlena], + bj bytejson.ByteJson, +) (bool, error) { if bj.Type != bytejson.TpCodeObject { - return bytejson.Null, nil + return false, nil } - cnt := bj.GetElemCnt() - keys := make([]any, cnt) - for i := 0; i < cnt; i++ { - keys[i] = string(bj.GetObjectKey(i)) + encoder, err := bytejson.NewObjectKeysArrayEncoder(bj) + if err != nil { + return false, err + } + if err := rs.AppendByteJsonEncoded(encoder); err != nil { + return false, err } - return bytejson.CreateByteJSON(keys) + return true, nil } // JSON_PRETTY @@ -2531,6 +2695,7 @@ func JsonPretty(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro result.UseOptFunctionParamFrame(1) rs := vector.MustFunctionResult[types.Varlena](result) p1 := vector.OptGetBytesParamFromWrapper(rs, 0, ivecs[0]) + var legacy bytes.Buffer if selectList != nil && selectList.IgnoreAllRow() { for i := 0; i < length; i++ { @@ -2559,25 +2724,17 @@ func JsonPretty(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro if err != nil { return moerr.NewInvalidArg(proc.Ctx, "json_pretty", "invalid JSON document") } - out, err := jsonPrettyPrint(bj, 0) + _, err = appendFormattedBytesForResult(result, rs, &legacy, func(w formatBuffer) (bool, error) { + return false, jsonPrettyPrintTo(w, bj, 0) + }) if err != nil { return err } - rs.AppendMustBytesValue(out) } return nil } -func jsonPrettyPrint(bj bytejson.ByteJson, depth int) ([]byte, error) { - var buf bytes.Buffer - err := jsonPrettyPrintTo(&buf, bj, depth) - if err != nil { - return nil, err - } - return buf.Bytes(), nil -} - -func jsonPrettyPrintTo(w *bytes.Buffer, bj bytejson.ByteJson, depth int) error { +func jsonPrettyPrintTo(w formatBuffer, bj bytejson.ByteJson, depth int) error { switch bj.Type { case bytejson.TpCodeObject: return prettyPrintObject(w, bj, depth) @@ -2588,69 +2745,94 @@ func jsonPrettyPrintTo(w *bytes.Buffer, bj bytejson.ByteJson, depth int) error { } } -func prettyPrintObject(w *bytes.Buffer, bj bytejson.ByteJson, depth int) error { +func prettyPrintObject(w formatBuffer, bj bytejson.ByteJson, depth int) error { cnt := bj.GetElemCnt() if cnt == 0 { - w.WriteString("{}") - return nil + _, err := w.WriteString("{}") + return err + } + if _, err := w.WriteString("{\n"); err != nil { + return err } - indent := strings.Repeat(" ", depth+1) - w.WriteString("{\n") for i := 0; i < cnt; i++ { key := bj.GetObjectKey(i) - // Escape key the same way JSON_QUOTE would. - keyJSON, _ := json.Marshal(string(key)) - w.WriteString(indent) - w.Write(keyJSON) - w.WriteString(": ") + if err := writePrettyIndent(w, depth+1); err != nil { + return err + } + if err := bytejson.WriteJSONString(w, key); err != nil { + return err + } + if _, err := w.WriteString(": "); err != nil { + return err + } val := bj.GetObjectVal(i) if err := jsonPrettyPrintTo(w, val, depth+1); err != nil { return err } if i < cnt-1 { - w.WriteString(",") + if err := w.WriteByte(','); err != nil { + return err + } + } + if err := w.WriteByte('\n'); err != nil { + return err } - w.WriteString("\n") } - w.WriteString(strings.Repeat(" ", depth)) - w.WriteString("}") - return nil + if err := writePrettyIndent(w, depth); err != nil { + return err + } + return w.WriteByte('}') } -func prettyPrintArray(w *bytes.Buffer, bj bytejson.ByteJson, depth int) error { +func prettyPrintArray(w formatBuffer, bj bytejson.ByteJson, depth int) error { cnt := bj.GetElemCnt() if cnt == 0 { - w.WriteString("[]") - return nil + _, err := w.WriteString("[]") + return err + } + if _, err := w.WriteString("[\n"); err != nil { + return err } - indent := strings.Repeat(" ", depth+1) - w.WriteString("[\n") for i := 0; i < cnt; i++ { - w.WriteString(indent) + if err := writePrettyIndent(w, depth+1); err != nil { + return err + } elem := bj.GetArrayElem(i) if err := jsonPrettyPrintTo(w, elem, depth+1); err != nil { return err } if i < cnt-1 { - w.WriteString(",") + if err := w.WriteByte(','); err != nil { + return err + } + } + if err := w.WriteByte('\n'); err != nil { + return err } - w.WriteString("\n") } - w.WriteString(strings.Repeat(" ", depth)) - w.WriteString("]") - return nil + if err := writePrettyIndent(w, depth); err != nil { + return err + } + return w.WriteByte(']') } -func prettyPrintScalar(w *bytes.Buffer, bj bytejson.ByteJson) error { - // Use MarshalJSON to get properly formatted/escaped scalar value. - text, err := bj.MarshalJSON() - if err != nil { - return err +func writePrettyIndent(w formatBuffer, depth int) error { + const spaces = " " + remaining := depth * 2 + for remaining > 0 { + length := min(remaining, len(spaces)) + if _, err := w.WriteString(spaces[:length]); err != nil { + return err + } + remaining -= length } - w.Write(text) return nil } +func prettyPrintScalar(w formatBuffer, bj bytejson.ByteJson) error { + return bytejson.WriteJSONText(w, bj) +} + // JSON_SCHEMA_VALID func JsonSchemaValid(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) diff --git a/pkg/sql/plan/function/func_cast.go b/pkg/sql/plan/function/func_cast.go index 29616b6810cc4..30fc6d0a9d8eb 100644 --- a/pkg/sql/plan/function/func_cast.go +++ b/pkg/sql/plan/function/func_cast.go @@ -7475,11 +7475,18 @@ func arrayToArray[I types.ArrayElement, O types.ArrayElement]( // upcast the source element type to []float32, then narrow to the // target element type (int8 rounds+clamps; bf16/f16 round-to-even). // This replaces moarray.Cast[I,O], which only handled float pairs. - _v := types.BytesToArray[I](v) - f32 := types.ToFloat32Array[I](_v) - out := types.FromFloat32Array[O](f32) - bytes := types.ArrayToBytes[O](out) - if err := to.AppendBytes(bytes, false); err != nil { + values := types.BytesToArray[I](v) + var outputElement O + elementSize := int(unsafe.Sizeof(outputElement)) + if len(values) > math.MaxInt/elementSize { + return moerr.NewInvalidInputNoCtx("array cast result is too large") + } + if err := to.AppendBytesWithFill(len(values)*elementSize, func(dst []byte) { + output := util.UnsafeSliceCast[O](dst) + for idx, value := range values { + output[idx] = float32ToArrayElement[O](arrayElementToFloat32(value)) + } + }); err != nil { return err } } @@ -7488,6 +7495,47 @@ func arrayToArray[I types.ArrayElement, O types.ArrayElement]( return nil } +func arrayElementToFloat32[T types.ArrayElement](value T) float32 { + switch typed := any(value).(type) { + case float32: + return typed + case float64: + return float32(typed) + case types.BF16: + return typed.ToFloat32() + case types.Float16: + return typed.ToFloat32() + case int8: + return float32(typed) + case uint8: + return float32(typed) + default: + panic(moerr.NewInternalErrorNoCtx("unsupported array element type")) + } +} + +func float32ToArrayElement[T types.ArrayElement](value float32) T { + var output any + var zero T + switch any(zero).(type) { + case float32: + output = value + case float64: + output = float64(value) + case types.BF16: + output = types.BF16FromFloat32(value) + case types.Float16: + output = types.Float16FromFloat32(value) + case int8: + output = types.Float32ToInt8(value) + case uint8: + output = types.Float32ToUint8(value) + default: + panic(moerr.NewInternalErrorNoCtx("unsupported array element type")) + } + return output.(T) +} + func uuidToStr( ctx context.Context, from vector.FunctionParameterWrapper[types.Uuid], diff --git a/pkg/sql/plan/function/func_compare.go b/pkg/sql/plan/function/func_compare.go index 01b1b55de2c7f..bc4cd2d05f3dd 100644 --- a/pkg/sql/plan/function/func_compare.go +++ b/pkg/sql/plan/function/func_compare.go @@ -1581,7 +1581,7 @@ func operatorOpInt64Uint64Fn( func operatorOpStrFn( parameters []*vector.Vector, result vector.FunctionResultWrapper, _ *process.Process, length int, - fn func([]byte, []byte) ([]byte, error)) error { + fn func([]byte, []byte, []byte)) error { p1 := vector.GenerateFunctionStrParameter(parameters[0]) p2 := vector.GenerateFunctionStrParameter(parameters[1]) rs := vector.MustFunctionResult[types.Varlena](result) @@ -1593,11 +1593,14 @@ func operatorOpStrFn( return err } } else { - rv, err := fn(v1, v2) - if err != nil { - return err + if len(v1) != len(v2) { + return moerr.NewInternalErrorNoCtx( + "Binary operands of bitwise operators must be of equal length", + ) } - if err = rs.AppendBytes(rv, false); err != nil { + if err := rs.AppendBytesWithFill(len(v1), func(dst []byte) { + fn(dst, v1, v2) + }); err != nil { return err } } @@ -1622,15 +1625,10 @@ func operatorOpBitAndInt64Uint64Fn(parameters []*vector.Vector, result vector.Fu } func operatorOpBitAndStrFn(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - return operatorOpStrFn(parameters, result, proc, length, func(i []byte, i2 []byte) ([]byte, error) { - if len(i) != len(i2) { - return nil, moerr.NewInternalErrorNoCtx("Binary operands of bitwise operators must be of equal length") - } - rv := make([]byte, len(i)) - for j := range rv { - rv[j] = i[j] & i2[j] + return operatorOpStrFn(parameters, result, proc, length, func(dst, left, right []byte) { + for idx := range dst { + dst[idx] = left[idx] & right[idx] } - return rv, nil }) } @@ -1651,15 +1649,10 @@ func operatorOpBitXorInt64Uint64Fn(parameters []*vector.Vector, result vector.Fu } func operatorOpBitXorStrFn(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - return operatorOpStrFn(parameters, result, proc, length, func(i []byte, i2 []byte) ([]byte, error) { - if len(i) != len(i2) { - return nil, moerr.NewInternalErrorNoCtx("Binary operands of bitwise operators must be of equal length") + return operatorOpStrFn(parameters, result, proc, length, func(dst, left, right []byte) { + for idx := range dst { + dst[idx] = left[idx] ^ right[idx] } - rv := make([]byte, len(i)) - for j := range rv { - rv[j] = i[j] ^ i2[j] - } - return rv, nil }) } @@ -1680,15 +1673,10 @@ func operatorOpBitOrInt64Uint64Fn(parameters []*vector.Vector, result vector.Fun } func operatorOpBitOrStrFn(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - return operatorOpStrFn(parameters, result, proc, length, func(i []byte, i2 []byte) ([]byte, error) { - if len(i) != len(i2) { - return nil, moerr.NewInternalErrorNoCtx("Binary operands of bitwise operators must be of equal length") - } - rv := make([]byte, len(i)) - for j := range rv { - rv[j] = i[j] | i2[j] + return operatorOpStrFn(parameters, result, proc, length, func(dst, left, right []byte) { + for idx := range dst { + dst[idx] = left[idx] | right[idx] } - return rv, nil }) } diff --git a/pkg/sql/plan/function/func_prefix.go b/pkg/sql/plan/function/func_prefix.go index 7c7429e75f906..761415482b083 100644 --- a/pkg/sql/plan/function/func_prefix.go +++ b/pkg/sql/plan/function/func_prefix.go @@ -16,6 +16,8 @@ package function import ( "bytes" + "encoding/binary" + "math" "sort" "github.com/matrixorigin/matrixone/pkg/common/mpool" @@ -84,8 +86,10 @@ func PrefixInRange(parameters []*vector.Vector, result vector.FunctionResultWrap } type implPrefixIn struct { - ready bool - vals [][]byte + ready bool + vals [][]byte + scratch []byte + scratchCount int } func newImplPrefixIn() *implPrefixIn { @@ -93,7 +97,6 @@ func newImplPrefixIn() *implPrefixIn { } func (op *implPrefixIn) init(rvec *vector.Vector, mp *mpool.MPool) error { - op.ready = true op.vals = make([][]byte, rvec.Length()) vlen := 0 @@ -123,12 +126,124 @@ func (op *implPrefixIn) init(rvec *vector.Vector, mp *mpool.MPool) error { } } op.vals = op.vals[:vlen] + op.ready = true return nil } +const prefixScratchEntrySize = 8 + +type prefixScratchEntries struct { + data []byte + count int +} + +func (e prefixScratchEntries) Len() int { + return e.count +} + +func (e prefixScratchEntries) Less(left, right int) bool { + return bytes.Compare(e.value(left), e.value(right)) < 0 +} + +func (e prefixScratchEntries) Swap(left, right int) { + leftEntry := e.data[left*prefixScratchEntrySize : (left+1)*prefixScratchEntrySize] + rightEntry := e.data[right*prefixScratchEntrySize : (right+1)*prefixScratchEntrySize] + var saved [prefixScratchEntrySize]byte + copy(saved[:], leftEntry) + copy(leftEntry, rightEntry) + copy(rightEntry, saved[:]) +} + +func (e prefixScratchEntries) value(index int) []byte { + entry := e.data[index*prefixScratchEntrySize:] + offset := binary.LittleEndian.Uint32(entry) + length := binary.LittleEndian.Uint32(entry[4:]) + return e.data[int(offset):int(offset+length)] +} + +func (op *implPrefixIn) initAccounted( + rvec *vector.Vector, + result vector.FunctionResultWrapper, +) error { + rowCount := rvec.Length() + if rowCount < 0 || rowCount > math.MaxInt/prefixScratchEntrySize { + return mpool.ErrAllocationAccountInvalid + } + total := rowCount * prefixScratchEntrySize + for row := 0; row < rowCount; row++ { + valueSize := len(rvec.GetBytesAt(row)) + if valueSize > math.MaxInt-total { + return mpool.ErrAllocationAccountInvalid + } + total += valueSize + } + if uint64(total) > math.MaxUint32 { + return mpool.ErrAllocationAccountInvalid + } + scratch, selected, err := result.ResizeFunctionScratch(total) + if err != nil { + return err + } + if !selected { + return mpool.ErrAllocationAccountInvalid + } + entries := prefixScratchEntries{data: scratch, count: rowCount} + payloadOffset := rowCount * prefixScratchEntrySize + for row := 0; row < rowCount; row++ { + value := rvec.GetBytesAt(row) + entry := scratch[row*prefixScratchEntrySize:] + binary.LittleEndian.PutUint32(entry, uint32(payloadOffset)) + binary.LittleEndian.PutUint32(entry[4:], uint32(len(value))) + payloadOffset += copy(scratch[payloadOffset:], value) + } + if !rvec.GetSorted() { + sort.Sort(entries) + } + compactCount := 0 + for row := 0; row < rowCount; row++ { + value := entries.value(row) + if compactCount != 0 && bytes.HasPrefix(value, entries.value(compactCount-1)) { + continue + } + if compactCount != row { + copy( + scratch[compactCount*prefixScratchEntrySize:], + scratch[row*prefixScratchEntrySize:(row+1)*prefixScratchEntrySize], + ) + } + compactCount++ + } + op.scratch = scratch + op.scratchCount = compactCount + op.ready = true + return nil +} + +func (op *implPrefixIn) valueCount() int { + if op.scratch != nil { + return op.scratchCount + } + return len(op.vals) +} + +func (op *implPrefixIn) valueAt(index int) []byte { + if op.scratch != nil { + return (prefixScratchEntries{ + data: op.scratch, + count: op.scratchCount, + }).value(index) + } + return op.vals[index] +} + func (op *implPrefixIn) doPrefixIn(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { if !op.ready { - err := op.init(parameters[1], proc.Mp()) + var err error + if result.HasFunctionScratch() { + err = op.initAccounted(parameters[1], result) + } else { + err = op.init(parameters[1], proc.Mp()) + } if err != nil { return err } @@ -136,7 +251,7 @@ func (op *implPrefixIn) doPrefixIn(parameters []*vector.Vector, result vector.Fu lvec := parameters[0] res := vector.MustFixedColWithTypeCheck[bool](result.GetResultVector()) - if len(op.vals) == 0 { + if op.valueCount() == 0 { for i := range length { res[i] = false } @@ -147,9 +262,9 @@ func (op *implPrefixIn) doPrefixIn(parameters []*vector.Vector, result vector.Fu lvecHasNull := lvec.HasNull() if lvec.GetSorted() && !lvecHasNull { - rval := op.vals[0] + rval := op.valueAt(0) rpos := 0 - rlen := len(op.vals) + rlen := op.valueCount() for i := range length { lval := lcol[i].GetByteSlice(larea) @@ -162,7 +277,7 @@ func (op *implPrefixIn) doPrefixIn(parameters []*vector.Vector, result vector.Fu return nil } - rval = op.vals[rpos] + rval = op.valueAt(rpos) } res[i] = bytes.HasPrefix(lval, rval) @@ -177,21 +292,21 @@ func (op *implPrefixIn) doPrefixIn(parameters []*vector.Vector, result vector.Fu rNulls.Add(i) } else { lval := lcol[i].GetByteSlice(larea) - rpos, _ := sort.Find(len(op.vals), func(j int) int { - return types.PrefixCompare(lval, op.vals[j]) + rpos, _ := sort.Find(op.valueCount(), func(j int) int { + return types.PrefixCompare(lval, op.valueAt(j)) }) - res[i] = rpos < len(op.vals) && bytes.HasPrefix(lval, op.vals[rpos]) + res[i] = rpos < op.valueCount() && bytes.HasPrefix(lval, op.valueAt(rpos)) } } } else { for i := range length { lval := lcol[i].GetByteSlice(larea) - rpos, _ := sort.Find(len(op.vals), func(j int) int { - return types.PrefixCompare(lval, op.vals[j]) + rpos, _ := sort.Find(op.valueCount(), func(j int) int { + return types.PrefixCompare(lval, op.valueAt(j)) }) - res[i] = rpos < len(op.vals) && bytes.HasPrefix(lval, op.vals[rpos]) + res[i] = rpos < op.valueCount() && bytes.HasPrefix(lval, op.valueAt(rpos)) } } } diff --git a/pkg/sql/plan/function/func_string_complex_test.go b/pkg/sql/plan/function/func_string_complex_test.go index 88877aecd6f00..3712a3f96a825 100644 --- a/pkg/sql/plan/function/func_string_complex_test.go +++ b/pkg/sql/plan/function/func_string_complex_test.go @@ -838,8 +838,8 @@ func Test_EncodeCharBytes(t *testing.T) { {math.MinInt64, []byte{0x00}}, } for _, c := range cases { - got := encodeCharBytes(c.input) - require.Equal(t, c.want, got, "encodeCharBytes(%d)", c.input) + got := appendCharBytes(nil, c.input) + require.Equal(t, c.want, got, "appendCharBytes(%d)", c.input) } } diff --git a/pkg/sql/plan/function/func_unary.go b/pkg/sql/plan/function/func_unary.go index 7231ee4b92fa7..1bafbdda271b6 100644 --- a/pkg/sql/plan/function/func_unary.go +++ b/pkg/sql/plan/function/func_unary.go @@ -50,6 +50,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/system" "github.com/matrixorigin/matrixone/pkg/common/util" + "github.com/matrixorigin/matrixone/pkg/container/bytejson" "github.com/matrixorigin/matrixone/pkg/container/nulls" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" @@ -1010,15 +1011,25 @@ func Empty(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *pr } func JsonQuote(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - single := func(str string) ([]byte, error) { - bj, err := types.ParseStringToByteJson(strconv.Quote(str)) + source := vector.GenerateFunctionStrParameter(ivecs[0]) + rs := vector.MustFunctionResult[types.Varlena](result) + for row := uint64(0); row < uint64(length); row++ { + value, isNull := source.GetStrValue(row) + if isNull || selectList.Contains(row) { + if err := rs.AppendBytes(nil, true); err != nil { + return err + } + continue + } + encoder, err := bytejson.NewStringDataEncoder(value) if err != nil { - return nil, err + return err + } + if err := rs.AppendByteJsonEncoded(encoder); err != nil { + return err } - return bj.Marshal() } - - return opUnaryStrToBytesWithErrorCheck(ivecs, result, proc, length, single, selectList) + return nil } func JsonUnquote(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { @@ -1047,6 +1058,11 @@ func JsonUnquote(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pr // Escapes single quotes by doubling them, backslashes, and control characters func QuoteString(str string) string { var result strings.Builder + writeQuotedString(&result, str) + return result.String() +} + +func writeQuotedString(result formatBuffer, str string) { result.WriteByte('\'') for _, r := range str { @@ -1079,15 +1095,28 @@ func QuoteString(str string) string { } result.WriteByte('\'') - return result.String() } func Quote(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - return opUnaryBytesToBytes(ivecs, result, proc, length, func(v []byte) []byte { - str := functionUtil.QuickBytesToStr(v) - quoted := QuoteString(str) - return functionUtil.QuickStrToBytes(quoted) - }, selectList) + source := vector.GenerateFunctionStrParameter(ivecs[0]) + rs := vector.MustFunctionResult[types.Varlena](result) + var legacy bytes.Buffer + for row := uint64(0); row < uint64(length); row++ { + value, isNull := source.GetStrValue(row) + if isNull || selectList.Contains(row) { + if err := rs.AppendBytes(nil, true); err != nil { + return err + } + continue + } + if _, err := appendFormattedBytesForResult(result, rs, &legacy, func(output formatBuffer) (bool, error) { + writeQuotedString(output, functionUtil.QuickBytesToStr(value)) + return false, nil + }); err != nil { + return err + } + } + return nil } func StAsText(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { @@ -4955,11 +4984,24 @@ func HexFloat64(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro } func HexArray(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - return opUnaryBytesToBytesWithErrorCheck(ivecs, result, proc, length, func(data []byte) ([]byte, error) { - buf := make([]byte, hex.EncodedLen(len(functionUtil.QuickBytesToStr(data)))) - hex.Encode(buf, data) - return buf, nil - }, selectList) + source := vector.GenerateFunctionStrParameter(ivecs[0]) + rs := vector.MustFunctionResult[types.Varlena](result) + for row := uint64(0); row < uint64(length); row++ { + data, isNull := source.GetStrValue(row) + if isNull || selectList.Contains(row) { + if err := rs.AppendBytes(nil, true); err != nil { + return err + } + continue + } + encodedSize := hex.EncodedLen(len(data)) + if err := rs.AppendBytesWithFill(encodedSize, func(dst []byte) { + hex.Encode(dst, data) + }); err != nil { + return err + } + } + return nil } func hexEncodeString(xs []byte) string { @@ -5783,17 +5825,45 @@ func unhexToBytes(data []byte, null bool, rs *vector.FunctionResult[types.Varlen return rs.AppendMustNullForBytesResult() } - // Add a '0' to the front, if the length is not the multiple of 2 - str := functionUtil.QuickBytesToStr(data) - if len(str)%2 != 0 { - str = "0" + str + decodedSize := (len(data) + 1) / 2 + var decodeErr error + err := rs.AppendBytesWithBuilder(decodedSize, func(dst []byte) (int, error) { + written, err := decodeHexInto(dst, data) + decodeErr = err + return written, err + }) + if decodeErr != nil { + return rs.AppendMustNullForBytesResult() } + return err +} - bs, err := hex.DecodeString(str) - if err != nil { - return rs.AppendMustNullForBytesResult() +func decodeHexInto(dst, src []byte) (int, error) { + written := 0 + if len(src)%2 != 0 { + value, ok := decodeHexNibble(src[0]) + if !ok { + return 0, hex.InvalidByteError(src[0]) + } + dst[0] = value + written = 1 + src = src[1:] + } + n, err := hex.Decode(dst[written:], src) + return written + n, err +} + +func decodeHexNibble(value byte) (byte, bool) { + switch { + case value >= '0' && value <= '9': + return value - '0', true + case value >= 'a' && value <= 'f': + return value - 'a' + 10, true + case value >= 'A' && value <= 'F': + return value - 'A' + 10, true + default: + return 0, false } - return rs.AppendMustBytesValue(bs) } func Unhex(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { @@ -5812,10 +5882,24 @@ func Unhex(parameters []*vector.Vector, result vector.FunctionResultWrapper, pro } func Md5(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - return opUnaryBytesToBytes(parameters, result, proc, length, func(data []byte) []byte { + source := vector.GenerateFunctionStrParameter(parameters[0]) + rs := vector.MustFunctionResult[types.Varlena](result) + for row := uint64(0); row < uint64(length); row++ { + data, isNull := source.GetStrValue(row) + if isNull || selectList.Contains(row) { + if err := rs.AppendBytes(nil, true); err != nil { + return err + } + continue + } sum := md5.Sum(data) - return []byte(hex.EncodeToString(sum[:])) - }, selectList) + if err := rs.AppendBytesWithFill(hex.EncodedLen(len(sum)), func(dst []byte) { + hex.Encode(dst, sum[:]) + }); err != nil { + return err + } + } + return nil } @@ -5843,11 +5927,24 @@ func (content *crc32ExecContext) builtInCrc32(parameters []*vector.Vector, resul } func ToBase64(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) (err error) { - return opUnaryBytesToBytesWithErrorCheck(ivecs, result, proc, length, func(data []byte) ([]byte, error) { - buf := make([]byte, base64.StdEncoding.EncodedLen(len(functionUtil.QuickBytesToStr(data)))) - base64.StdEncoding.Encode(buf, data) - return buf, nil - }, selectList) + source := vector.GenerateFunctionStrParameter(ivecs[0]) + rs := vector.MustFunctionResult[types.Varlena](result) + for row := uint64(0); row < uint64(length); row++ { + data, isNull := source.GetStrValue(row) + if isNull || selectList.Contains(row) { + if err = rs.AppendBytes(nil, true); err != nil { + return err + } + continue + } + encodedSize := base64.StdEncoding.EncodedLen(len(data)) + if err = rs.AppendBytesWithFill(encodedSize, func(dst []byte) { + base64.StdEncoding.Encode(dst, data) + }); err != nil { + return err + } + } + return nil } func FromBase64(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { @@ -5857,16 +5954,31 @@ func FromBase64(parameters []*vector.Vector, result vector.FunctionResultWrapper rowCount := uint64(length) for i := uint64(0); i < rowCount; i++ { data, null := source.GetStrValue(i) - if null { - return rs.AppendMustNullForBytesResult() + if null || selectList.Contains(i) { + if err := rs.AppendMustNullForBytesResult(); err != nil { + return err + } + continue } - buf := make([]byte, base64.StdEncoding.DecodedLen(len(functionUtil.QuickBytesToStr(data)))) - _, err := base64.StdEncoding.Decode(buf, data) + var decodeErr error + err := rs.AppendBytesWithBuilder( + base64.StdEncoding.DecodedLen(len(data)), + func(dst []byte) (int, error) { + var written int + written, decodeErr = base64.StdEncoding.Decode(dst, data) + return written, decodeErr + }, + ) + if decodeErr != nil { + if err = rs.AppendMustNullForBytesResult(); err != nil { + return err + } + continue + } if err != nil { - return rs.AppendMustNullForBytesResult() + return err } - _ = rs.AppendMustBytesValue(buf) } return nil @@ -5907,11 +6019,10 @@ func VecFromBase64[T types.ArrayElement](parameters []*vector.Vector, result vec } } - var buf []byte rowCount := uint64(length) for i := uint64(0); i < rowCount; i++ { data, null := source.GetStrValue(i) - if null { + if null || selectList.Contains(i) { if err := rs.AppendBytes(nil, true); err != nil { return err } @@ -5919,21 +6030,22 @@ func VecFromBase64[T types.ArrayElement](parameters []*vector.Vector, result vec } need := base64.StdEncoding.DecodedLen(len(data)) - if cap(buf) < need { - buf = make([]byte, need) - } else { - buf = buf[:need] - } - n, err := base64.StdEncoding.Decode(buf, data) - if err != nil { - return moerr.NewInternalErrorNoCtx("vec_from_base64: invalid base64 input") - } - - if n%elemSize != 0 { - return moerr.NewInternalErrorNoCtxf("vec_from_base64: decoded length %d is not a multiple of %d bytes", n, elemSize) - } - - if err = rs.AppendBytes(buf[:n], false); err != nil { + if err := rs.AppendBytesWithBuilder(need, func(dst []byte) (int, error) { + written, err := base64.StdEncoding.Decode(dst, data) + if err != nil { + return 0, moerr.NewInternalErrorNoCtx( + "vec_from_base64: invalid base64 input", + ) + } + if written%elemSize != 0 { + return 0, moerr.NewInternalErrorNoCtxf( + "vec_from_base64: decoded length %d is not a multiple of %d bytes", + written, + elemSize, + ) + } + return written, nil + }); err != nil { return err } } @@ -5946,6 +6058,8 @@ func VecFromBase64[T types.ArrayElement](parameters []*vector.Vector, result vec func Compress(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { source := vector.GenerateFunctionStrParameter(parameters[0]) rs := vector.MustFunctionResult[types.Varlena](result) + var writer *flate.Writer + var output fixedSliceWriter rowCount := uint64(length) for i := uint64(0); i < rowCount; i++ { @@ -5964,49 +6078,126 @@ func Compress(parameters []*vector.Vector, result vector.FunctionResultWrapper, continue } - // Compress using zlib (flate) - var buf bytes.Buffer - writer, err := flate.NewWriter(&buf, flate.DefaultCompression) + capacity, err := flateResultCapacity(len(data)) if err != nil { - if err := rs.AppendBytes(nil, true); err != nil { + return err + } + if writer == nil { + writer, err = flate.NewWriter(io.Discard, flate.DefaultCompression) + if err != nil { return err } - continue } - - _, err = writer.Write(data) - if err != nil { - writer.Close() - if err := rs.AppendBytes(nil, true); err != nil { - return err + var compressionErr error + err = rs.AppendBytesWithBuilder(capacity, func(dst []byte) (int, error) { + output.Reset(dst[4:]) + writer.Reset(&output) + if _, compressionErr = writer.Write(data); compressionErr == nil { + compressionErr = writer.Close() + } else { + _ = writer.Close() + } + if compressionErr != nil { + return 0, compressionErr + } + binary.LittleEndian.PutUint32(dst[:4], uint32(len(data))) + return 4 + output.Written(), nil + }) + if compressionErr != nil { + if nullErr := rs.AppendBytes(nil, true); nullErr != nil { + return nullErr } continue } - - err = writer.Close() if err != nil { - if err := rs.AppendBytes(nil, true); err != nil { - return err - } - continue + return err } + } - compressed := buf.Bytes() + return nil +} - // MySQL format: 4-byte length (little-endian) + compressed data - originalLen := uint32(len(data)) - result := make([]byte, 4+len(compressed)) - binary.LittleEndian.PutUint32(result[0:4], originalLen) - copy(result[4:], compressed) +type fixedSliceWriter struct { + dst []byte + written int + err error +} - if err := rs.AppendBytes(result, false); err != nil { - return err - } +func (w *fixedSliceWriter) Reset(dst []byte) { + w.dst = dst + w.written = 0 + w.err = nil +} + +func (w *fixedSliceWriter) Write(value []byte) (int, error) { + if w.err != nil { + return 0, w.err + } + if len(value) > len(w.dst)-w.written { + w.err = io.ErrShortBuffer + return 0, w.err } + copy(w.dst[w.written:], value) + w.written += len(value) + return len(value), nil +} +func (w *fixedSliceWriter) WriteString(value string) (int, error) { + if w.err != nil { + return 0, w.err + } + if len(value) > len(w.dst)-w.written { + w.err = io.ErrShortBuffer + return 0, w.err + } + copy(w.dst[w.written:], value) + w.written += len(value) + return len(value), nil +} + +func (w *fixedSliceWriter) WriteByte(value byte) error { + if w.err != nil { + return w.err + } + if w.written == len(w.dst) { + w.err = io.ErrShortBuffer + return w.err + } + w.dst[w.written] = value + w.written++ return nil } +func (w *fixedSliceWriter) WriteRune(value rune) (int, error) { + var encoded [utf8.UTFMax]byte + size := utf8.EncodeRune(encoded[:], value) + return w.Write(encoded[:size]) +} + +func (w *fixedSliceWriter) Grow(int) {} + +func (w *fixedSliceWriter) Written() int { + return w.written +} + +func (w *fixedSliceWriter) Err() error { + return w.err +} + +func flateResultCapacity(inputSize int) (int, error) { + if inputSize < 0 || uint64(inputSize) > math.MaxUint32 { + return 0, moerr.NewInvalidInputNoCtx("compress input is too large") + } + // zlib's compressBound is also an upper bound for the contained raw + // DEFLATE stream; add four bytes for MySQL's original-length prefix. + size := uint64(inputSize) + bound := size + (size >> 12) + (size >> 14) + (size >> 25) + 13 + if bound > uint64(^uint(0)>>1)-4 { + return 0, moerr.NewInvalidInputNoCtx("compress result is too large") + } + return int(bound) + 4, nil +} + // Uncompress: UNCOMPRESS(string) - Uncompresses a string compressed by COMPRESS() // Reads 4-byte length, then decompresses the rest func Uncompress(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { @@ -6043,30 +6234,31 @@ func Uncompress(parameters []*vector.Vector, result vector.FunctionResultWrapper originalLen := binary.LittleEndian.Uint32(data[0:4]) compressed := data[4:] - // Decompress using zlib (flate) - reader := flate.NewReader(bytes.NewReader(compressed)) - decompressed := make([]byte, originalLen) - n, err := reader.Read(decompressed) - reader.Close() - - if err != nil && err != io.EOF { - // Decompression failed, return NULL - if err := rs.AppendBytes(nil, true); err != nil { - return err + var decodeErr error + err := rs.AppendBytesWithBuilder(int(originalLen), func(dst []byte) (int, error) { + reader := flate.NewReader(bytes.NewReader(compressed)) + defer reader.Close() + if _, decodeErr = io.ReadFull(reader, dst); decodeErr != nil { + return 0, decodeErr + } + var extra [1]byte + n, err := reader.Read(extra[:]) + if err != io.EOF || n != 0 { + if err == nil { + err = errors.New("decompressed length exceeds header") + } + decodeErr = err + return 0, err } - continue - } - - // Check if we got the expected length - if uint32(n) != originalLen { - // Length mismatch, return NULL + return len(dst), nil + }) + if decodeErr != nil { if err := rs.AppendBytes(nil, true); err != nil { return err } continue } - - if err := rs.AppendBytes(decompressed, false); err != nil { + if err != nil { return err } } @@ -6419,18 +6611,20 @@ func RandomBytes(parameters []*vector.Vector, result vector.FunctionResultWrappe continue } - // Generate random bytes using crypto/rand - randomBytes := make([]byte, lenVal) - _, err := rand.Read(randomBytes) - if err != nil { + var randomErr error + err := rs.AppendBytesWithBuilder(int(lenVal), func(dst []byte) (int, error) { + var written int + written, randomErr = rand.Read(dst) + return written, randomErr + }) + if randomErr != nil { // On error, return NULL if err := rs.AppendBytes(nil, true); err != nil { return err } continue } - - if err := rs.AppendBytes(randomBytes, false); err != nil { + if err != nil { return err } } diff --git a/pkg/sql/plan/function/func_unary_codec_scratch_test.go b/pkg/sql/plan/function/func_unary_codec_scratch_test.go new file mode 100644 index 0000000000000..c9fca4abd8734 --- /dev/null +++ b/pkg/sql/plan/function/func_unary_codec_scratch_test.go @@ -0,0 +1,161 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package function + +import ( + "strings" + "testing" + + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/stretchr/testify/require" +) + +func TestCompressUncompressDirectOutput(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + values := []string{ + "", + "a", + strings.Repeat("compressible-value-", 4096), + } + source := newVectorByType(mp, types.T_blob.ToType(), values, nil) + defer source.Free(mp) + + compressed := vector.NewFunctionResultWrapper(types.T_blob.ToType(), mp) + defer compressed.Free() + require.NoError(t, compressed.PreExtendAndReset(len(values))) + require.NoError(t, Compress( + []*vector.Vector{source}, + compressed, + proc, + len(values), + nil, + )) + + decoded := vector.NewFunctionResultWrapper(types.T_blob.ToType(), mp) + defer decoded.Free() + require.NoError(t, decoded.PreExtendAndReset(len(values))) + require.NoError(t, Uncompress( + []*vector.Vector{compressed.GetResultVector()}, + decoded, + proc, + len(values), + nil, + )) + for row, value := range values { + require.Equal(t, []byte(value), decoded.GetResultVector().GetBytesAt(row)) + } +} + +func TestCompressDirectOutputBound(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + for _, size := range []int{ + 0, 1, 15, 16, 17, 127, 128, 255, 256, + 4095, 4096, 4097, 16383, 16384, 16385, + 65534, 65535, 65536, 1 << 20, + } { + value := make([]byte, size) + state := uint64(size) + 1 + for idx := range value { + state = state*6364136223846793005 + 1442695040888963407 + value[idx] = byte(state >> 56) + } + source := newVectorByType( + mp, + types.T_blob.ToType(), + []string{string(value)}, + nil, + ) + compressed := vector.NewFunctionResultWrapper(types.T_blob.ToType(), mp) + require.NoError(t, compressed.PreExtendAndReset(1)) + require.NoErrorf(t, Compress( + []*vector.Vector{source}, + compressed, + proc, + 1, + nil, + ), "size %d", size) + + decoded := vector.NewFunctionResultWrapper(types.T_blob.ToType(), mp) + require.NoError(t, decoded.PreExtendAndReset(1)) + require.NoErrorf(t, Uncompress( + []*vector.Vector{compressed.GetResultVector()}, + decoded, + proc, + 1, + nil, + ), "size %d", size) + require.Equalf(t, value, decoded.GetResultVector().GetBytesAt(0), "size %d", size) + decoded.Free() + compressed.Free() + source.Free(mp) + } +} + +func TestFromBase64DirectOutputNullAndSelection(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + source := vector.NewVec(types.T_varchar.ToType()) + defer source.Free(mp) + require.NoError(t, vector.AppendBytes(source, []byte("YWJj"), false, mp)) + require.NoError(t, vector.AppendBytes(source, nil, true, mp)) + require.NoError(t, vector.AppendBytes(source, []byte("ZGVm"), false, mp)) + + result := vector.NewFunctionResultWrapper(types.T_blob.ToType(), mp) + defer result.Free() + require.NoError(t, result.PreExtendAndReset(3)) + require.NoError(t, FromBase64( + []*vector.Vector{source}, + result, + proc, + 3, + &FunctionSelectList{ + AnyNull: true, + SelectList: []bool{true, true, false}, + }, + )) + require.Equal(t, []byte("abc"), result.GetResultVector().GetBytesAt(0)) + require.True(t, result.GetResultVector().IsNull(1)) + require.True(t, result.GetResultVector().IsNull(2)) +} + +func TestRandomBytesDirectOutput(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + source := newVectorByType( + mp, + types.T_int64.ToType(), + []int64{1, 1024, 0}, + nil, + ) + defer source.Free(mp) + + result := vector.NewFunctionResultWrapper(types.T_blob.ToType(), mp) + defer result.Free() + require.NoError(t, result.PreExtendAndReset(3)) + require.NoError(t, RandomBytes( + []*vector.Vector{source}, + result, + proc, + 3, + nil, + )) + require.Len(t, result.GetResultVector().GetBytesAt(0), 1) + require.Len(t, result.GetResultVector().GetBytesAt(1), 1024) + require.True(t, result.GetResultVector().IsNull(2)) +} diff --git a/pkg/sql/plan/function/function_allocation_scratch_test.go b/pkg/sql/plan/function/function_allocation_scratch_test.go new file mode 100644 index 0000000000000..1a0678f67e1df --- /dev/null +++ b/pkg/sql/plan/function/function_allocation_scratch_test.go @@ -0,0 +1,763 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package function + +import ( + "math" + "strings" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/nulls" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/stretchr/testify/require" +) + +const ( + testFunctionOwner mpool.AllocationOwner = 1 + testFunctionResultData mpool.AllocationSite = 1 + testFunctionResultArea mpool.AllocationSite = 2 + testFunctionResultNulls mpool.AllocationSite = 3 + testFunctionResultGroup mpool.AllocationSite = 4 + testFunctionParam mpool.AllocationSite = 5 + testFunctionScratch mpool.AllocationSite = 6 +) + +func newAccountedFunctionResult( + t *testing.T, + typ types.Type, + mp *mpool.MPool, + limit uint64, +) (vector.FunctionResultWrapper, *mpool.AllocationAccountRegistry, *mpool.AllocationAccount) { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, 32) + require.NoError(t, err) + account, err := registry.Open(limit) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelectionWithBitmaps( + account, + testFunctionOwner, + testFunctionResultData, + testFunctionResultArea, + testFunctionResultNulls, + testFunctionResultGroup, + ) + require.NoError(t, err) + allocation, err := vector.NewFunctionAllocation( + account, + testFunctionOwner, + testFunctionParam, + testFunctionScratch, + ) + require.NoError(t, err) + result, err := vector.NewFunctionResultWrapperWithFunctionAllocation( + typ, + mp, + selection, + allocation, + ) + require.NoError(t, err) + return result, registry, account +} + +func finalizeAccountedFunctionResult( + t *testing.T, + result vector.FunctionResultWrapper, + registry *mpool.AllocationAccountRegistry, + account *mpool.AllocationAccount, +) { + t.Helper() + result.Free() + require.Zero(t, account.Snapshot().Used) + account.Seal() + _, err := registry.Finalize(account) + require.NoError(t, err) + require.Zero(t, registry.LiveAllocationMetadata()) +} + +func TestAppendFormattedBytesRollsBackChangedSecondPass(t *testing.T) { + proc := testutil.NewProcess(t) + result, registry, account := newAccountedFunctionResult( + t, + types.T_varchar.ToType(), + proc.Mp(), + 1<<20, + ) + require.NoError(t, result.PreExtendAndReset(1)) + calls := 0 + _, err := appendFormattedBytes( + vector.MustFunctionResult[types.Varlena](result), + func(w formatBuffer) (bool, error) { + calls++ + if calls == 1 { + _, err := w.WriteString("two") + return false, err + } + _, err := w.WriteString("x") + return false, err + }, + ) + require.Error(t, err) + require.Zero(t, result.GetResultVector().Length()) + finalizeAccountedFunctionResult(t, result, registry, account) +} + +func TestBuiltInHashAccountedScratchMatchesLegacy(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + rows := 3 + nsp := nulls.NewWithSize(rows) + nsp.Set(1) + stringsInput := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{"alpha", "unused", strings.Repeat("wide-value", 256)}, + nsp, + ) + defer stringsInput.Free(mp) + integersInput := newVectorByType( + mp, + types.T_int64.ToType(), + []int64{7, 8, 9}, + nil, + ) + defer integersInput.Free(mp) + inputs := []*vector.Vector{stringsInput, integersInput} + + legacy := vector.NewFunctionResultWrapper(types.T_int64.ToType(), mp) + require.NoError(t, legacy.PreExtendAndReset(rows)) + require.NoError(t, builtInHash(inputs, legacy, proc, rows, nil)) + want := append( + []int64(nil), + vector.MustFixedColWithTypeCheck[int64](legacy.GetResultVector())..., + ) + legacy.Free() + + accounted, registry, account := newAccountedFunctionResult( + t, + types.T_int64.ToType(), + mp, + 1<<20, + ) + require.NoError(t, accounted.PreExtendAndReset(rows)) + require.NoError(t, builtInHash(inputs, accounted, proc, rows, nil)) + require.Equal( + t, + want, + vector.MustFixedColWithTypeCheck[int64](accounted.GetResultVector()), + ) + require.Positive(t, account.Snapshot().Used) + finalizeAccountedFunctionResult(t, accounted, registry, account) +} + +func TestBuiltInHashAccountedScratchRejectsCapacity(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + input := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{strings.Repeat("x", 4096)}, + nil, + ) + defer input.Free(mp) + result, registry, account := newAccountedFunctionResult( + t, + types.T_int64.ToType(), + mp, + 1024, + ) + require.NoError(t, result.PreExtendAndReset(1)) + err := builtInHash([]*vector.Vector{input}, result, proc, 1, nil) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + finalizeAccountedFunctionResult(t, result, registry, account) +} + +func TestPrefixInAccountedScratchMatchesLegacy(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + left := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{"abc", "zzz", "other"}, + nil, + ) + defer left.Free(mp) + right := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{"z", "ab", "a"}, + nil, + ) + defer right.Free(mp) + inputs := []*vector.Vector{left, right} + + legacy := vector.NewFunctionResultWrapper(types.T_bool.ToType(), mp) + require.NoError(t, legacy.PreExtendAndReset(3)) + require.NoError(t, newImplPrefixIn().doPrefixIn(inputs, legacy, proc, 3, nil)) + want := append( + []bool(nil), + vector.MustFixedColWithTypeCheck[bool](legacy.GetResultVector())..., + ) + legacy.Free() + + accounted, registry, account := newAccountedFunctionResult( + t, + types.T_bool.ToType(), + mp, + 1<<20, + ) + require.NoError(t, accounted.PreExtendAndReset(3)) + op := newImplPrefixIn() + require.NoError(t, op.doPrefixIn(inputs, accounted, proc, 3, nil)) + require.Equal( + t, + want, + vector.MustFixedColWithTypeCheck[bool](accounted.GetResultVector()), + ) + require.NotEmpty(t, op.scratch) + require.Empty(t, op.vals) + require.Positive(t, account.Snapshot().Used) + + require.NoError(t, accounted.PreExtendAndReset(3)) + require.NoError(t, op.doPrefixIn(inputs, accounted, proc, 3, nil)) + require.Equal( + t, + want, + vector.MustFixedColWithTypeCheck[bool](accounted.GetResultVector()), + ) + finalizeAccountedFunctionResult(t, accounted, registry, account) +} + +func TestPrefixInAccountedScratchRejectsCapacity(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + left := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{"value"}, + nil, + ) + defer left.Free(mp) + right := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{strings.Repeat("x", 4096)}, + nil, + ) + defer right.Free(mp) + result, registry, account := newAccountedFunctionResult( + t, + types.T_bool.ToType(), + mp, + 1024, + ) + require.NoError(t, result.PreExtendAndReset(1)) + err := newImplPrefixIn().doPrefixIn( + []*vector.Vector{left, right}, + result, + proc, + 1, + nil, + ) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + finalizeAccountedFunctionResult(t, result, registry, account) +} + +func TestJqAccountedOutputMatchesLegacy(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + jsonInput := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{`{"values":[1,2,3]}`}, + nil, + ) + defer jsonInput.Free(mp) + queryInput := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{`.values | map(. * 2)`}, + nil, + ) + defer queryInput.Free(mp) + params := []*vector.Vector{jsonInput, queryInput} + + legacy := vector.NewFunctionResultWrapper(types.T_varchar.ToType(), mp) + require.NoError(t, legacy.PreExtendAndReset(1)) + require.NoError(t, newOpBuiltInJq().jq(params, legacy, proc, 1, nil)) + want := append([]byte(nil), legacy.GetResultVector().GetBytesAt(0)...) + legacy.Free() + require.Equal(t, []byte(`[2,4,6]`), want) + + accounted, registry, account := newAccountedFunctionResult( + t, + types.T_varchar.ToType(), + mp, + 1<<20, + ) + require.NoError(t, accounted.PreExtendAndReset(1)) + op := newOpBuiltInJq() + require.NoError(t, op.jq(params, accounted, proc, 1, nil)) + require.Equal(t, want, accounted.GetResultVector().GetBytesAt(0)) + require.Positive(t, account.Snapshot().Used) + finalizeAccountedFunctionResult(t, accounted, registry, account) +} + +func TestJqAccountedOutputRejectsCapacity(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + jsonInput := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{`"` + strings.Repeat("x", 4096) + `"`}, + nil, + ) + defer jsonInput.Free(mp) + queryInput := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{"."}, + nil, + ) + defer queryInput.Free(mp) + params := []*vector.Vector{jsonInput, queryInput} + for _, test := range []struct { + name string + run func(*opBuiltInJq, vector.FunctionResultWrapper) error + }{ + { + name: "jq", + run: func(op *opBuiltInJq, result vector.FunctionResultWrapper) error { + return op.jq(params, result, proc, 1, nil) + }, + }, + { + name: "try_jq", + run: func(op *opBuiltInJq, result vector.FunctionResultWrapper) error { + return op.tryJq(params, result, proc, 1, nil) + }, + }, + } { + t.Run(test.name, func(t *testing.T) { + result, registry, account := newAccountedFunctionResult( + t, + types.T_varchar.ToType(), + mp, + 1024, + ) + require.NoError(t, result.PreExtendAndReset(1)) + err := test.run(newOpBuiltInJq(), result) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + finalizeAccountedFunctionResult(t, result, registry, account) + }) + } +} + +func TestJSONRowAccountedOutputMatchesLegacy(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + integers := newVectorByType( + mp, + types.T_int64.ToType(), + []int64{7, 8}, + nil, + ) + defer integers.Free(mp) + stringsInput := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{"alpha", "beta"}, + nil, + ) + defer stringsInput.Free(mp) + params := []*vector.Vector{integers, stringsInput} + + legacy := vector.NewFunctionResultWrapper(types.T_varchar.ToType(), mp) + require.NoError(t, legacy.PreExtendAndReset(2)) + require.NoError(t, newOpBuiltInJsonRow().jsonRow(params, legacy, proc, 2, nil)) + want0 := append([]byte(nil), legacy.GetResultVector().GetBytesAt(0)...) + want1 := append([]byte(nil), legacy.GetResultVector().GetBytesAt(1)...) + legacy.Free() + + accounted, registry, account := newAccountedFunctionResult( + t, + types.T_varchar.ToType(), + mp, + 1<<20, + ) + require.NoError(t, accounted.PreExtendAndReset(2)) + require.NoError(t, newOpBuiltInJsonRow().jsonRow( + params, + accounted, + proc, + 2, + nil, + )) + require.Equal(t, want0, accounted.GetResultVector().GetBytesAt(0)) + require.Equal(t, want1, accounted.GetResultVector().GetBytesAt(1)) + require.Positive(t, account.Snapshot().Used) + finalizeAccountedFunctionResult(t, accounted, registry, account) +} + +func TestJSONObjectAccountedKeyScratchMatchesLegacy(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + keys := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{strings.Repeat("key", 256), "second"}, + nil, + ) + defer keys.Free(mp) + values := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{"first", "value"}, + nil, + ) + defer values.Free(mp) + params := []*vector.Vector{keys, values} + + legacy := vector.NewFunctionResultWrapper(types.T_json.ToType(), mp) + require.NoError(t, legacy.PreExtendAndReset(2)) + require.NoError(t, newOpBuiltInJsonObject().jsonObject(params, legacy, proc, 2, nil)) + want0 := append([]byte(nil), legacy.GetResultVector().GetBytesAt(0)...) + want1 := append([]byte(nil), legacy.GetResultVector().GetBytesAt(1)...) + legacy.Free() + + accounted, registry, account := newAccountedFunctionResult( + t, + types.T_json.ToType(), + mp, + 1<<20, + ) + require.NoError(t, accounted.PreExtendAndReset(2)) + require.NoError(t, newOpBuiltInJsonObject().jsonObject( + params, + accounted, + proc, + 2, + nil, + )) + require.Equal(t, want0, accounted.GetResultVector().GetBytesAt(0)) + require.Equal(t, want1, accounted.GetResultVector().GetBytesAt(1)) + require.Positive(t, account.Snapshot().Used) + finalizeAccountedFunctionResult(t, accounted, registry, account) +} + +func TestJSONObjectAccountedKeyScratchRejectsCapacity(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + keys := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{strings.Repeat("key", 2048)}, + nil, + ) + defer keys.Free(mp) + values := newVectorByType( + mp, + types.T_int64.ToType(), + []int64{1}, + nil, + ) + defer values.Free(mp) + result, registry, account := newAccountedFunctionResult( + t, + types.T_json.ToType(), + mp, + 1024, + ) + require.NoError(t, result.PreExtendAndReset(1)) + err := newOpBuiltInJsonObject().jsonObject( + []*vector.Vector{keys, values}, + result, + proc, + 1, + nil, + ) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + finalizeAccountedFunctionResult(t, result, registry, account) +} + +func TestJSONModifyAccountedValueScratchRejectsCapacity(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + document := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{`{"value":0}`}, + nil, + ) + defer document.Free(mp) + path := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{"$.value"}, + nil, + ) + defer path.Free(mp) + arrayType := types.T_array_float32.ToType() + array := vector.NewVec(arrayType) + values := make([]float32, 1024) + for idx := range values { + values[idx] = float32(idx) + } + require.NoError(t, vector.AppendBytes( + array, + types.ArrayToBytes(values), + false, + mp, + )) + defer array.Free(mp) + result, registry, account := newAccountedFunctionResult( + t, + types.T_json.ToType(), + mp, + 1024, + ) + require.NoError(t, result.PreExtendAndReset(1)) + err := newOpBuiltInJsonSet().buildJsonSet( + []*vector.Vector{document, path, array}, + result, + proc, + 1, + nil, + ) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + finalizeAccountedFunctionResult(t, result, registry, account) +} + +func TestFixedInAccountedScratchMatchesMapSemantics(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + left := newVectorByType( + mp, + types.T_float64.ToType(), + []float64{math.NaN(), math.Copysign(0, -1), 1, 2}, + nil, + ) + defer left.Free(mp) + tupleNulls := nulls.NewWithSize(3) + tupleNulls.Set(2) + tuple := newVectorByType( + mp, + types.T_float64.ToType(), + []float64{math.NaN(), 0, 0}, + tupleNulls, + ) + defer tuple.Free(mp) + params := []*vector.Vector{left, tuple} + + legacy := vector.NewFunctionResultWrapper(types.T_bool.ToType(), mp) + require.NoError(t, legacy.PreExtendAndReset(4)) + require.NoError(t, newOpOperatorFixedIn[float64]().operatorIn( + params, + legacy, + proc, + 4, + nil, + )) + wantValues := append( + []bool(nil), + vector.MustFixedColWithTypeCheck[bool](legacy.GetResultVector())..., + ) + wantNulls := make([]bool, 4) + for row := range wantNulls { + wantNulls[row] = legacy.GetResultVector().IsNull(uint64(row)) + } + legacy.Free() + + accounted, registry, account := newAccountedFunctionResult( + t, + types.T_bool.ToType(), + mp, + 1<<20, + ) + require.NoError(t, accounted.PreExtendAndReset(4)) + op := newOpOperatorFixedIn[float64]() + require.NoError(t, op.operatorIn(params, accounted, proc, 4, nil)) + require.True(t, op.accounted) + require.Nil(t, op.mp) + require.Equal( + t, + wantValues, + vector.MustFixedColWithTypeCheck[bool](accounted.GetResultVector()), + ) + for row, wantNull := range wantNulls { + require.Equal(t, wantNull, accounted.GetResultVector().IsNull(uint64(row))) + } + finalizeAccountedFunctionResult(t, accounted, registry, account) +} + +func TestFixedInAccountedScratchRejectsCapacity(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + left := newVectorByType(mp, types.T_int64.ToType(), []int64{1}, nil) + defer left.Free(mp) + tupleValues := make([]int64, 1024) + for idx := range tupleValues { + tupleValues[idx] = int64(idx) + } + tuple := newVectorByType(mp, types.T_int64.ToType(), tupleValues, nil) + defer tuple.Free(mp) + result, registry, account := newAccountedFunctionResult( + t, + types.T_bool.ToType(), + mp, + 1024, + ) + require.NoError(t, result.PreExtendAndReset(1)) + err := newOpOperatorFixedIn[int64]().operatorIn( + []*vector.Vector{left, tuple}, + result, + proc, + 1, + nil, + ) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + finalizeAccountedFunctionResult(t, result, registry, account) +} + +func TestStringInAccountedScratchMatchesMapSemantics(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + left := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{"alpha", "missing", "", "wide"}, + nil, + ) + defer left.Free(mp) + tupleNulls := nulls.NewWithSize(5) + tupleNulls.Set(4) + tuple := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{"wide", "alpha", "alpha", "", "unused"}, + tupleNulls, + ) + defer tuple.Free(mp) + params := []*vector.Vector{left, tuple} + + legacy := vector.NewFunctionResultWrapper(types.T_bool.ToType(), mp) + require.NoError(t, legacy.PreExtendAndReset(4)) + require.NoError(t, newOpOperatorStrIn().operatorIn( + params, + legacy, + proc, + 4, + nil, + )) + wantValues := append( + []bool(nil), + vector.MustFixedColWithTypeCheck[bool](legacy.GetResultVector())..., + ) + wantNulls := make([]bool, 4) + for row := range wantNulls { + wantNulls[row] = legacy.GetResultVector().IsNull(uint64(row)) + } + legacy.Free() + + accounted, registry, account := newAccountedFunctionResult( + t, + types.T_bool.ToType(), + mp, + 1<<20, + ) + require.NoError(t, accounted.PreExtendAndReset(4)) + op := newOpOperatorStrIn() + require.NoError(t, op.operatorIn(params, accounted, proc, 4, nil)) + require.True(t, op.accounted) + require.Nil(t, op.mp) + require.Equal( + t, + wantValues, + vector.MustFixedColWithTypeCheck[bool](accounted.GetResultVector()), + ) + for row, wantNull := range wantNulls { + require.Equal(t, wantNull, accounted.GetResultVector().IsNull(uint64(row))) + } + finalizeAccountedFunctionResult(t, accounted, registry, account) +} + +func TestStringInAccountedScratchRejectsCapacity(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + left := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{"value"}, + nil, + ) + defer left.Free(mp) + tuple := newVectorByType( + mp, + types.T_varchar.ToType(), + []string{strings.Repeat("x", 4096)}, + nil, + ) + defer tuple.Free(mp) + result, registry, account := newAccountedFunctionResult( + t, + types.T_bool.ToType(), + mp, + 1024, + ) + require.NoError(t, result.PreExtendAndReset(1)) + err := newOpOperatorStrIn().operatorIn( + []*vector.Vector{left, tuple}, + result, + proc, + 1, + nil, + ) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + finalizeAccountedFunctionResult(t, result, registry, account) +} + +func TestNarrowCosineSimilarityAccountedScratchRejectsCapacity(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + arrayType := types.T_array_uint8.ToType() + left := vector.NewVec(arrayType) + right := vector.NewVec(arrayType) + values := make([]uint8, 4096) + for idx := range values { + values[idx] = uint8(idx) + } + require.NoError(t, vector.AppendBytes(left, types.ArrayToBytes(values), false, mp)) + require.NoError(t, vector.AppendBytes(right, types.ArrayToBytes(values), false, mp)) + defer left.Free(mp) + defer right.Free(mp) + result, registry, account := newAccountedFunctionResult( + t, + types.T_float64.ToType(), + mp, + 1024, + ) + require.NoError(t, result.PreExtendAndReset(1)) + err := CosineSimilarityArrayViaF32[uint8]( + []*vector.Vector{left, right}, + result, + proc, + 1, + nil, + ) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + finalizeAccountedFunctionResult(t, result, registry, account) +} diff --git a/pkg/sql/plan/function/operator_in.go b/pkg/sql/plan/function/operator_in.go index 7e57b1158940b..6e5b99f69fcab 100644 --- a/pkg/sql/plan/function/operator_in.go +++ b/pkg/sql/plan/function/operator_in.go @@ -15,6 +15,16 @@ package function import ( + "bytes" + "cmp" + "encoding/binary" + "math" + "slices" + "sort" + "unsafe" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/common/util" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/vm/process" @@ -27,15 +37,205 @@ type TGenericOfIn interface { } type opOperatorFixedIn[T TGenericOfIn] struct { - ready bool - hasNull bool - mp map[T]bool + ready bool + hasNull bool + mp map[T]bool + accounted bool + scratch []byte + count int } type opOperatorStrIn struct { - ready bool - hasNull bool - mp map[string]bool + ready bool + hasNull bool + mp map[string]bool + accounted bool + scratch []byte + count int +} + +func compareInValues[T TGenericOfIn](left, right T) int { + switch value := any(left).(type) { + case uint8: + return cmp.Compare(value, any(right).(uint8)) + case uint16: + return cmp.Compare(value, any(right).(uint16)) + case uint32: + return cmp.Compare(value, any(right).(uint32)) + case uint64: + return cmp.Compare(value, any(right).(uint64)) + case int8: + return cmp.Compare(value, any(right).(int8)) + case int16: + return cmp.Compare(value, any(right).(int16)) + case int32: + return cmp.Compare(value, any(right).(int32)) + case int64: + return cmp.Compare(value, any(right).(int64)) + case float32: + return compareInFloat64(float64(value), float64(any(right).(float32))) + case float64: + return compareInFloat64(value, any(right).(float64)) + case bool: + other := any(right).(bool) + if value == other { + return 0 + } + if !value { + return -1 + } + return 1 + case types.Uuid: + return types.CompareUuid(value, any(right).(types.Uuid)) + case types.Time: + return cmp.Compare(value, any(right).(types.Time)) + case types.Timestamp: + return cmp.Compare(value, any(right).(types.Timestamp)) + case types.Date: + return cmp.Compare(value, any(right).(types.Date)) + case types.Datetime: + return cmp.Compare(value, any(right).(types.Datetime)) + case types.Decimal64: + return value.Compare(any(right).(types.Decimal64)) + case types.Decimal128: + return value.Compare(any(right).(types.Decimal128)) + case types.Decimal256: + return value.Compare(any(right).(types.Decimal256)) + case types.MoYear: + return cmp.Compare(value, any(right).(types.MoYear)) + default: + panic("unsupported IN value type") + } +} + +func compareInFloat64(left, right float64) int { + leftNaN := math.IsNaN(left) + rightNaN := math.IsNaN(right) + switch { + case leftNaN && rightNaN: + return cmp.Compare(math.Float64bits(left), math.Float64bits(right)) + case leftNaN: + return 1 + case rightNaN: + return -1 + default: + return cmp.Compare(left, right) + } +} + +func (op *opOperatorFixedIn[T]) initAccounted( + tuple *vector.Vector, + result vector.FunctionResultWrapper, +) error { + op.hasNull = false + count := 0 + parameter := vector.GenerateFunctionFixedTypeParameter[T](tuple) + for row := uint64(0); row < uint64(tuple.Length()); row++ { + _, isNull := parameter.GetValue(row) + if isNull { + op.hasNull = true + } else { + count++ + } + } + var zero T + elementSize := int(unsafe.Sizeof(zero)) + if count > math.MaxInt/elementSize { + return mpool.ErrAllocationAccountInvalid + } + scratch, selected, err := result.ResizeFunctionScratch(count * elementSize) + if err != nil { + return err + } + if !selected { + return mpool.ErrAllocationAccountInvalid + } + values := util.UnsafeSliceCast[T](scratch)[:count] + write := 0 + for row := uint64(0); row < uint64(tuple.Length()); row++ { + value, isNull := parameter.GetValue(row) + if !isNull { + values[write] = value + write++ + } + } + slices.SortFunc(values, compareInValues[T]) + op.accounted = true + op.scratch = scratch + op.count = count + op.ready = true + return nil +} + +func (op *opOperatorFixedIn[T]) containsAccounted(value T) bool { + values := util.UnsafeSliceCast[T](op.scratch)[:op.count] + idx, found := slices.BinarySearchFunc(values, value, compareInValues[T]) + return found && values[idx] == value +} + +func (op *opOperatorStrIn) initAccounted( + tuple *vector.Vector, + result vector.FunctionResultWrapper, +) error { + op.hasNull = false + count := 0 + payloadSize := 0 + parameter := vector.GenerateFunctionStrParameter(tuple) + for row := uint64(0); row < uint64(tuple.Length()); row++ { + value, isNull := parameter.GetStrValue(row) + if isNull { + op.hasNull = true + continue + } + if len(value) > math.MaxInt-payloadSize { + return mpool.ErrAllocationAccountInvalid + } + payloadSize += len(value) + count++ + } + if count > math.MaxInt/prefixScratchEntrySize || + payloadSize > math.MaxInt-count*prefixScratchEntrySize { + return mpool.ErrAllocationAccountInvalid + } + total := count*prefixScratchEntrySize + payloadSize + if uint64(total) > math.MaxUint32 { + return mpool.ErrAllocationAccountInvalid + } + scratch, selected, err := result.ResizeFunctionScratch(total) + if err != nil { + return err + } + if !selected { + return mpool.ErrAllocationAccountInvalid + } + entries := prefixScratchEntries{data: scratch, count: count} + payloadOffset := count * prefixScratchEntrySize + write := 0 + for row := uint64(0); row < uint64(tuple.Length()); row++ { + value, isNull := parameter.GetStrValue(row) + if isNull { + continue + } + entry := scratch[write*prefixScratchEntrySize:] + binary.LittleEndian.PutUint32(entry, uint32(payloadOffset)) + binary.LittleEndian.PutUint32(entry[4:], uint32(len(value))) + payloadOffset += copy(scratch[payloadOffset:], value) + write++ + } + sort.Sort(entries) + op.accounted = true + op.scratch = scratch + op.count = count + op.ready = true + return nil +} + +func (op *opOperatorStrIn) containsAccounted(value []byte) bool { + entries := prefixScratchEntries{data: op.scratch, count: op.count} + idx := sort.Search(op.count, func(idx int) bool { + return bytes.Compare(entries.value(idx), value) >= 0 + }) + return idx < op.count && bytes.Equal(entries.value(idx), value) } func newOpOperatorFixedIn[T TGenericOfIn]() *opOperatorFixedIn[T] { @@ -118,9 +318,51 @@ func (op *opOperatorStrIn) init(tuple *vector.Vector) { } } +func (op *opOperatorFixedIn[T]) ensureInitialized( + tuple *vector.Vector, + result vector.FunctionResultWrapper, +) error { + if op.ready { + return nil + } + if result.HasFunctionScratch() { + return op.initAccounted(tuple, result) + } + op.init(tuple) + return nil +} + +func (op *opOperatorStrIn) ensureInitialized( + tuple *vector.Vector, + result vector.FunctionResultWrapper, +) error { + if op.ready { + return nil + } + if result.HasFunctionScratch() { + return op.initAccounted(tuple, result) + } + op.init(tuple) + return nil +} + +func (op *opOperatorFixedIn[T]) contains(value T) bool { + if op.accounted { + return op.containsAccounted(value) + } + return op.mp[value] +} + +func (op *opOperatorStrIn) contains(value []byte) bool { + if op.accounted { + return op.containsAccounted(value) + } + return op.mp[string(value)] +} + func (op *opOperatorFixedIn[T]) operatorIn(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - if !op.ready { - op.init(parameters[1]) + if err := op.ensureInitialized(parameters[1], result); err != nil { + return err } p := vector.GenerateFunctionFixedTypeParameter[T](parameters[0]) @@ -132,7 +374,7 @@ func (op *opOperatorFixedIn[T]) operatorIn(parameters []*vector.Vector, result v return err } } else { - _, ok := op.mp[v] + ok := op.contains(v) if !ok && op.hasNull { if err := rs.Append(false, true); err != nil { return err @@ -148,8 +390,8 @@ func (op *opOperatorFixedIn[T]) operatorIn(parameters []*vector.Vector, result v } func (op *opOperatorFixedIn[T]) operatorNotIn(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - if !op.ready { - op.init(parameters[1]) + if err := op.ensureInitialized(parameters[1], result); err != nil { + return err } p := vector.GenerateFunctionFixedTypeParameter[T](parameters[0]) @@ -161,7 +403,7 @@ func (op *opOperatorFixedIn[T]) operatorNotIn(parameters []*vector.Vector, resul return err } } else { - _, ok := op.mp[v] + ok := op.contains(v) if !ok && op.hasNull { if err := rs.Append(false, true); err != nil { return err @@ -177,8 +419,8 @@ func (op *opOperatorFixedIn[T]) operatorNotIn(parameters []*vector.Vector, resul } func (op *opOperatorStrIn) operatorIn(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - if !op.ready { - op.init(parameters[1]) + if err := op.ensureInitialized(parameters[1], result); err != nil { + return err } p := vector.GenerateFunctionStrParameter(parameters[0]) @@ -190,7 +432,7 @@ func (op *opOperatorStrIn) operatorIn(parameters []*vector.Vector, result vector return err } } else { - _, ok := op.mp[string(v)] + ok := op.contains(v) if !ok && op.hasNull { if err := rs.Append(false, true); err != nil { return err @@ -206,8 +448,8 @@ func (op *opOperatorStrIn) operatorIn(parameters []*vector.Vector, result vector } func (op *opOperatorStrIn) operatorNotIn(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - if !op.ready { - op.init(parameters[1]) + if err := op.ensureInitialized(parameters[1], result); err != nil { + return err } p := vector.GenerateFunctionStrParameter(parameters[0]) @@ -219,7 +461,7 @@ func (op *opOperatorStrIn) operatorNotIn(parameters []*vector.Vector, result vec return err } } else { - _, ok := op.mp[string(v)] + ok := op.contains(v) if !ok && op.hasNull { if err := rs.Append(false, true); err != nil { return err diff --git a/pkg/vectorindex/metric/cpu.go b/pkg/vectorindex/metric/cpu.go index 70134166ad6d1..19acc1e02077e 100644 --- a/pkg/vectorindex/metric/cpu.go +++ b/pkg/vectorindex/metric/cpu.go @@ -57,6 +57,38 @@ func PairwiseDistanceLaunchOneToMany[T types.RealNumbers]( dist []float32, _ uint64, _ bool, +) (PairwiseJobHandle, error) { + return PairwiseDistanceLaunchOneToManyWithScratch( + query, + rowCount, + rowAt, + metric, + dist, + 0, + false, + nil, + ) +} + +func PairwiseDistanceOneToManyScratchSize[T types.RealNumbers]( + _ []T, + _ int, + _ MetricType, + _ uint64, + _ bool, +) (int, bool, error) { + return 0, false, nil +} + +func PairwiseDistanceLaunchOneToManyWithScratch[T types.RealNumbers]( + query []T, + rowCount int, + rowAt func(int) []T, + metric MetricType, + dist []float32, + _ uint64, + _ bool, + _ []byte, ) (PairwiseJobHandle, error) { return PairwiseDistanceLaunchOneToManyCPU( query, diff --git a/pkg/vectorindex/metric/gpu.go b/pkg/vectorindex/metric/gpu.go index 1bb56ac0a5c62..03a4ace196f97 100644 --- a/pkg/vectorindex/metric/gpu.go +++ b/pkg/vectorindex/metric/gpu.go @@ -106,6 +106,7 @@ type gpuJob struct { cuvsJobID uint64 deallocators []malloc.Deallocator dist []float32 + scratch []byte } type gpuJobManager struct { @@ -142,6 +143,20 @@ func (m *gpuJobManager) update(jobID uint64, cuvsID uint64, d ...malloc.Dealloca } } +func (m *gpuJobManager) updateScratch( + jobID uint64, + cuvsID uint64, + scratch []byte, +) { + m.mu.Lock() + defer m.mu.Unlock() + job := m.jobs[jobID] + if job != nil { + job.cuvsJobID = cuvsID + job.scratch = scratch + } +} + func (m *gpuJobManager) pop(jobID uint64) *gpuJob { m.mu.Lock() defer m.mu.Unlock() @@ -211,6 +226,71 @@ func PairwiseDistanceLaunchOneToMany[T types.RealNumbers]( dist []float32, minWorkSize uint64, gpuMode bool, +) (PairwiseJobHandle, error) { + return PairwiseDistanceLaunchOneToManyWithScratch( + query, + rowCount, + rowAt, + metric, + dist, + minWorkSize, + gpuMode, + nil, + ) +} + +func PairwiseDistanceOneToManyScratchSize[T types.RealNumbers]( + query []T, + rowCount int, + metric MetricType, + minWorkSize uint64, + gpuMode bool, +) (int, bool, error) { + if !gpuMode || rowCount <= 0 { + return 0, false, nil + } + dim := len(query) + work := uint64(rowCount) + if dim != 0 && work > ^uint64(0)/uint64(dim) { + work = ^uint64(0) + } else { + work *= uint64(dim) + } + _, supportedMetric := MetricTypeToCuvsMetric[metric] + if !supportedMetric || work < minWorkSize { + return 0, false, nil + } + if _, ok := any(query).([]float32); !ok { + return 0, false, nil + } + rows := uint64(rowCount) + 1 + if rows == 0 || dim != 0 && rows > ^uint64(0)/uint64(dim) { + return 0, false, moerr.NewInternalErrorNoCtx( + "pairwise distance input is too large", + ) + } + elements := rows * uint64(dim) + if elements > uint64(^uint(0)>>1)/4 { + return 0, false, moerr.NewInternalErrorNoCtx( + "pairwise distance input is too large", + ) + } + return int(elements * 4), true, nil +} + +// PairwiseDistanceLaunchOneToManyWithScratch uses caller-owned scratch for +// GPU input flattening when provided. The caller must retain the buffer until +// PairwiseDistanceWait returns. A nil buffer preserves the legacy C-allocator +// path. +func PairwiseDistanceLaunchOneToManyWithScratch[T types.RealNumbers]( + query []T, + rowCount int, + rowAt func(int) []T, + metric MetricType, + dist []float32, + minWorkSize uint64, + gpuMode bool, + scratch []byte, ) (PairwiseJobHandle, error) { if !gpuMode { return PairwiseDistanceLaunchOneToManyCPU( @@ -247,7 +327,7 @@ func PairwiseDistanceLaunchOneToMany[T types.RealNumbers]( if supportedMetric && work >= minWorkSize { if typedQuery, ok := any(query).([]float32); ok { - return gpuPairwiseLaunchRows( + return gpuPairwiseLaunchRowsWithScratch( 1, rowCount, dim, @@ -260,6 +340,7 @@ func PairwiseDistanceLaunchOneToMany[T types.RealNumbers]( cuvsMetric, dist[:rowCount], 4, + scratch, ) } } @@ -304,6 +385,27 @@ func gpuPairwiseLaunchRows[C cuvs.VectorType]( cuvsMetric cuvs.DistanceType, dist []float32, elemSize int, +) (PairwiseJobHandle, error) { + return gpuPairwiseLaunchRowsWithScratch( + nX, + nY, + dim, + xAt, + yAt, + cuvsMetric, + dist, + elemSize, + nil, + ) +} + +func gpuPairwiseLaunchRowsWithScratch[C cuvs.VectorType]( + nX, nY, dim int, + xAt, yAt func(int) []C, + cuvsMetric cuvs.DistanceType, + dist []float32, + elemSize int, + scratch []byte, ) (PairwiseJobHandle, error) { if nX < 0 || nY < 0 || @@ -323,6 +425,20 @@ func gpuPairwiseLaunchRows[C cuvs.VectorType]( "pairwise distance input is too large", ) } + if scratch != nil { + return gpuPairwiseLaunchRowsFromScratch( + nX, + nY, + dim, + xAt, + yAt, + cuvsMetric, + dist, + elemSize, + rowBytes, + scratch, + ) + } allocator := malloc.NewCAllocator() // 1. Flatten Y @@ -392,6 +508,66 @@ func gpuPairwiseLaunchRows[C cuvs.VectorType]( return PairwiseJobHandle(gpuID), nil } +func gpuPairwiseLaunchRowsFromScratch[C cuvs.VectorType]( + nX, nY, dim int, + xAt, yAt func(int) []C, + cuvsMetric cuvs.DistanceType, + dist []float32, + elemSize int, + rowBytes uint64, + scratch []byte, +) (PairwiseJobHandle, error) { + yBytes := uint64(nY) * rowBytes + xBytes := uint64(nX) * rowBytes + if yBytes > uint64(len(scratch)) || xBytes > uint64(len(scratch))-yBytes { + return 0, moerr.NewInternalErrorNoCtx( + "pairwise distance scratch is smaller than the flattened input", + ) + } + yf := util.UnsafeSliceCast[C](scratch[:yBytes]) + for row := 0; row < nY; row++ { + value := yAt(row) + if len(value) != dim { + return 0, moerr.NewInternalErrorNoCtx( + "vector dimension not matched", + ) + } + copy(yf[row*dim:(row+1)*dim], value) + } + xScratch := scratch[yBytes : yBytes+xBytes] + xf := util.UnsafeSliceCast[C](xScratch) + for row := 0; row < nX; row++ { + value := xAt(row) + if len(value) != dim { + return 0, moerr.NewInternalErrorNoCtx( + "vector dimension not matched", + ) + } + copy(xf[row*dim:(row+1)*dim], value) + } + + gpuID := globalGpuJobManager.add(dist) + cuvsID, err := cuvs.PairwiseDistanceLaunch( + xf, + uint64(nX), + yf, + uint64(nY), + uint32(dim), + cuvsMetric, + dist, + ) + if err != nil { + globalGpuJobManager.pop(gpuID) + return 0, err + } + globalGpuJobManager.updateScratch( + gpuID, + cuvsID, + scratch[:yBytes+xBytes], + ) + return PairwiseJobHandle(gpuID), nil +} + // PairwiseDistanceWait waits for the completion of the asynchronous GPU distance // calculation initiated by Launch. func PairwiseDistanceWait(handle PairwiseJobHandle, metric MetricType) ([]float32, error) { From 2f5870e9e3e1fe0ffa74f7fd417d3cc6544ab376 Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 20:49:54 +0800 Subject: [PATCH 14/61] executor: finalize allocation accounts per attempt --- pkg/common/mpool/accounted_buffer.go | 4 +- pkg/common/mpool/allocation_account.go | 327 +++++++++++++++++- .../mpool/allocation_account_mpool_test.go | 8 +- pkg/common/mpool/allocation_account_test.go | 184 ++++++++++ pkg/common/mpool/mpool.go | 19 +- .../compile/allocation_account_lifecycle.go | 103 ++++++ .../allocation_account_lifecycle_test.go | 321 +++++++++++++++++ pkg/sql/compile/compile.go | 20 ++ pkg/sql/compile/compile2.go | 40 ++- pkg/sql/compile/compile_test.go | 22 ++ pkg/sql/compile/remoterunServer.go | 12 +- pkg/sql/compile/types.go | 5 + pkg/vm/message/message.go | 106 +++++- pkg/vm/message/message_test.go | 95 +++++ 14 files changed, 1233 insertions(+), 33 deletions(-) create mode 100644 pkg/sql/compile/allocation_account_lifecycle.go create mode 100644 pkg/sql/compile/allocation_account_lifecycle_test.go diff --git a/pkg/common/mpool/accounted_buffer.go b/pkg/common/mpool/accounted_buffer.go index 4c688117ed9cd..ee59f29eb7314 100644 --- a/pkg/common/mpool/accounted_buffer.go +++ b/pkg/common/mpool/accounted_buffer.go @@ -84,13 +84,13 @@ func (b *AccountedBuffer) EnsureCapacity(required int) error { return nil } if int64(required) > maxAllocationSize() { - return ErrAllocationAccountInvalid + return ErrAllocationAllocatorLimit } oldLength := len(b.data) capacity, ok := GrowCapacity(int64(cap(b.data)), int64(required)) if !ok || capacity > int64(math.MaxInt) { - return ErrAllocationAccountInvalid + return ErrAllocationAllocatorLimit } if cap(b.data) == 0 { data, err := b.mp.AllocAccounted( diff --git a/pkg/common/mpool/allocation_account.go b/pkg/common/mpool/allocation_account.go index 603728d49a64b..ec82418fc0eee 100644 --- a/pkg/common/mpool/allocation_account.go +++ b/pkg/common/mpool/allocation_account.go @@ -38,10 +38,18 @@ const ( ) var ( - ErrAllocationAccountCapacity = errors.New("allocation account capacity exceeded") - ErrAllocationAccountSealed = errors.New("allocation account is sealed") - ErrAllocationAccountInvalid = errors.New("invalid allocation account") - ErrAllocationAccountStale = errors.New("stale allocation account handle") + ErrAllocationAccountCapacity = errors.New("allocation account capacity exceeded") + ErrAllocationAccountSealed = errors.New("allocation account is sealed") + ErrAllocationAccountInvalid = errors.New("invalid allocation account") + ErrAllocationAccountStale = errors.New("stale allocation account handle") + ErrAllocationAccountMismatch = errors.New("allocation account ownership mismatch") + ErrAllocationAllocatorLimit = errors.New("allocation exceeds allocator size limit") + ErrAllocationAccountInvariant = errors.New( + "allocation account invariant failure", + ) + ErrAllocationAdmissionSuspended = errors.New( + "allocation account admission is suspended", + ) ErrAllocationMetadataSlots = errors.New("allocation metadata slots exhausted") ErrAllocationGenerationSlots = errors.New("allocation account generation slots exhausted") ErrAllocationAccountLive = errors.New("allocation account still owns memory") @@ -79,6 +87,73 @@ type AllocationAccountSnapshot struct { Sealed bool } +// AllocationAccountTerminalState classifies the one immutable snapshot +// exported by an execution generation. A nonzero account at terminal cleanup +// is an ownership invariant failure, not recoverable capacity pressure. +type AllocationAccountTerminalState uint8 + +const ( + AllocationAccountTerminalValid AllocationAccountTerminalState = iota + 1 + AllocationAccountTerminalInvariantFailure +) + +// AllocationFailureReason is the non-overlapping control-flow reason exposed +// to later pressure handling. Only Capacity is eligible for reclaim/spill or +// a smaller operation retry; every other reason is terminal for that logical +// operation or generation. +type AllocationFailureReason uint8 + +const ( + AllocationFailureNone AllocationFailureReason = iota + AllocationFailureCapacity + AllocationFailureSealed + AllocationFailureMismatch + AllocationFailureAllocatorLimit + AllocationFailureInvariant + AllocationFailureSuspended +) + +func AllocationFailureReasonOf(err error) AllocationFailureReason { + switch { + case errors.Is(err, ErrAllocationAccountInvariant): + return AllocationFailureInvariant + case errors.Is(err, ErrAllocationAccountMismatch): + return AllocationFailureMismatch + case errors.Is(err, ErrAllocationAccountSealed): + return AllocationFailureSealed + case errors.Is(err, ErrAllocationAllocatorLimit): + return AllocationFailureAllocatorLimit + case errors.Is(err, ErrAllocationAdmissionSuspended): + return AllocationFailureSuspended + case errors.Is(err, ErrAllocationAccountCapacity), + errors.Is(err, ErrAllocationMetadataSlots): + return AllocationFailureCapacity + default: + return AllocationFailureNone + } +} + +func IsRetryableAllocationCapacity(err error) bool { + return AllocationFailureReasonOf(err) == AllocationFailureCapacity +} + +// AllocationAccountTerminalSnapshot is the immutable terminal observation of +// one generation. Failure snapshots retain the live-byte value observed at +// the terminal boundary even if a later physical Free drains the tombstone. +type AllocationAccountTerminalSnapshot struct { + AllocationAccountSnapshot + State AllocationAccountTerminalState +} + +// AllocationAccountCheckpoint records the physical live-byte boundary before +// a retryable logical operation. The owner performs its own private-allocation +// rollback, then ValidateRollback proves that the same generation returned to +// this exact boundary before a retry can begin. +type AllocationAccountCheckpoint struct { + Handle AllocationAccountHandle + Used uint64 +} + // AllocationCapacityController lets an account share a higher-level aggregate // cap during migration. The controller owns cap policy only; physical MPool // metadata remains the sole release owner. @@ -122,6 +197,78 @@ func (a *AllocationAccount) Snapshot() AllocationAccountSnapshot { } } +func (a *AllocationAccount) Checkpoint() (AllocationAccountCheckpoint, error) { + if a == nil || a.registry == nil || a.handle == 0 { + return AllocationAccountCheckpoint{}, ErrAllocationAccountInvalid + } + resolved, ok := a.registry.Resolve(a.handle) + if !ok || resolved != a { + return AllocationAccountCheckpoint{}, ErrAllocationAccountStale + } + snapshot := a.Snapshot() + if snapshot.Sealed { + return AllocationAccountCheckpoint{}, ErrAllocationAccountSealed + } + return AllocationAccountCheckpoint{ + Handle: snapshot.Handle, + Used: snapshot.Used, + }, nil +} + +// ValidateRollback proves that an owner restored its complete physical +// allocation boundary. It never mutates accounting: only physical Free owns a +// release, so a helper cannot hide a leaked allocation by decrementing usage. +func (a *AllocationAccount) ValidateRollback( + checkpoint AllocationAccountCheckpoint, +) error { + if a == nil || checkpoint.Handle == 0 { + return ErrAllocationAccountInvalid + } + if checkpoint.Handle != a.handle { + return fmt.Errorf( + "%w: checkpoint=%d account=%d", + ErrAllocationAccountMismatch, + checkpoint.Handle, + a.handle, + ) + } + snapshot := a.Snapshot() + if snapshot.Sealed { + return ErrAllocationAccountSealed + } + if snapshot.Used != checkpoint.Used { + return fmt.Errorf( + "%w: checkpoint-used=%d current-used=%d", + ErrAllocationAccountInvariant, + checkpoint.Used, + snapshot.Used, + ) + } + return nil +} + +// RollbackToCheckpoint runs the owner's physical cleanup and then proves the +// exact generation boundary. It deliberately does not own or synthesize any +// release: MPool allocation metadata remains the sole release authority. +func (a *AllocationAccount) RollbackToCheckpoint( + checkpoint AllocationAccountCheckpoint, + rollback func() error, +) error { + if rollback == nil { + return ErrAllocationAccountInvalid + } + if a == nil || checkpoint.Handle != a.handle { + return ErrAllocationAccountMismatch + } + if a.Snapshot().Sealed { + return ErrAllocationAccountSealed + } + if err := rollback(); err != nil { + return err + } + return a.ValidateRollback(checkpoint) +} + func (a *AllocationAccount) acquire(capacity uint64) error { if a == nil || a.registry == nil || a.handle == 0 { return ErrAllocationAccountInvalid @@ -221,6 +368,9 @@ func (a *AllocationAccount) release(capacity uint64) { } next := state - capacity if a.state.CompareAndSwap(state, next) { + if next&allocationAccountUsedMask == 0 { + a.registry.tryDrainTombstone(a) + } return } } @@ -247,6 +397,9 @@ func (a *AllocationAccount) Seal() AllocationAccountSnapshot { type allocationAccountRegistrySlot struct { account atomic.Pointer[AllocationAccount] + // terminal and tombstone are protected by AllocationAccountRegistry.mu. + terminal *AllocationAccountTerminalSnapshot + tombstone bool } // AllocationAccountRegistry bounds live generations and accounted-allocation @@ -258,6 +411,8 @@ type AllocationAccountRegistry struct { slots []allocationAccountRegistrySlot generations []uint32 free []uint32 + suspended bool + tombstones uint32 maxAllocations uint64 liveAllocations atomic.Uint64 @@ -299,6 +454,9 @@ func (r *AllocationAccountRegistry) OpenWithController( r.mu.Lock() defer r.mu.Unlock() + if r.suspended { + return nil, ErrAllocationAdmissionSuspended + } for len(r.free) > 0 { index := len(r.free) - 1 slot := r.free[index] @@ -321,6 +479,156 @@ func (r *AllocationAccountRegistry) OpenWithController( return nil, ErrAllocationGenerationSlots } +// CompleteTerminal seals one generation and publishes its immutable terminal +// state exactly once. A nonzero terminal state remains resolvable as a +// release-capable tombstone and suspends new generations on this registry. +// The tombstone is removed automatically after the last physical Free. +// +// first is true only for the call that created the immutable snapshot. A +// repeated call while a tombstone is live returns the same snapshot with +// first=false. +func (r *AllocationAccountRegistry) CompleteTerminal( + account *AllocationAccount, +) (snapshot AllocationAccountTerminalSnapshot, first bool, err error) { + if r == nil || account == nil || account.registry != r { + return snapshot, false, ErrAllocationAccountInvalid + } + account.Seal() + + r.mu.Lock() + defer r.mu.Unlock() + slot := account.handle.slot() + if slot == 0 || uint64(slot) >= uint64(len(r.slots)) || + r.slots[slot].account.Load() != account { + return snapshot, false, ErrAllocationAccountStale + } + entry := &r.slots[slot] + if entry.terminal != nil { + snapshot = *entry.terminal + if snapshot.State == AllocationAccountTerminalInvariantFailure { + return snapshot, false, newAllocationTerminalInvariantError(snapshot) + } + return snapshot, false, nil + } + + current := account.Snapshot() + if !current.Sealed || account.inflight.Load() != 0 { + return AllocationAccountTerminalSnapshot{ + AllocationAccountSnapshot: current, + State: AllocationAccountTerminalInvariantFailure, + }, false, fmt.Errorf( + "%w: terminal account is not quiescent", + ErrAllocationAccountInvariant, + ) + } + snapshot = AllocationAccountTerminalSnapshot{ + AllocationAccountSnapshot: current, + State: AllocationAccountTerminalValid, + } + if current.Used == 0 { + entry.terminal = &snapshot + r.removeSlotLocked(slot, account) + return snapshot, true, nil + } + + snapshot.State = AllocationAccountTerminalInvariantFailure + entry.terminal = &snapshot + entry.tombstone = true + r.tombstones++ + r.suspended = true + // A physical Free may have raced the terminal observation. The immutable + // failure snapshot remains truthful at its linearization point, while a + // now-empty tombstone can be removed immediately. + if account.Snapshot().Used == 0 && account.inflight.Load() == 0 { + r.removeTombstoneLocked(slot, account) + } + return snapshot, true, newAllocationTerminalInvariantError(snapshot) +} + +func newAllocationTerminalInvariantError( + snapshot AllocationAccountTerminalSnapshot, +) error { + return fmt.Errorf( + "%w: handle=%d used=%d peak=%d limit=%d", + ErrAllocationAccountInvariant, + snapshot.Handle, + snapshot.Used, + snapshot.Peak, + snapshot.Limit, + ) +} + +func (r *AllocationAccountRegistry) removeSlotLocked( + slot uint32, + account *AllocationAccount, +) { + entry := &r.slots[slot] + if entry.account.Load() != account { + return + } + entry.account.Store(nil) + entry.terminal = nil + entry.tombstone = false + if r.generations[slot] != math.MaxUint32 { + r.free = append(r.free, slot) + } +} + +func (r *AllocationAccountRegistry) removeTombstoneLocked( + slot uint32, + account *AllocationAccount, +) { + entry := &r.slots[slot] + if entry.account.Load() != account || !entry.tombstone { + return + } + if r.tombstones == 0 { + panic("allocation account tombstone underflow") + } + r.tombstones-- + r.removeSlotLocked(slot, account) + r.suspended = r.tombstones != 0 +} + +func (r *AllocationAccountRegistry) tryDrainTombstone( + account *AllocationAccount, +) { + if r == nil || account == nil || account.Snapshot().Used != 0 || + account.inflight.Load() != 0 { + return + } + r.mu.Lock() + defer r.mu.Unlock() + slot := account.handle.slot() + if slot == 0 || uint64(slot) >= uint64(len(r.slots)) { + return + } + if r.slots[slot].account.Load() == account && + r.slots[slot].tombstone && + account.Snapshot().Used == 0 && + account.inflight.Load() == 0 { + r.removeTombstoneLocked(slot, account) + } +} + +func (r *AllocationAccountRegistry) AdmissionSuspended() bool { + if r == nil { + return false + } + r.mu.Lock() + defer r.mu.Unlock() + return r.suspended +} + +func (r *AllocationAccountRegistry) LiveTombstones() uint32 { + if r == nil { + return 0 + } + r.mu.Lock() + defer r.mu.Unlock() + return r.tombstones +} + func (r *AllocationAccountRegistry) Resolve( handle AllocationAccountHandle, ) (*AllocationAccount, bool) { @@ -360,10 +668,7 @@ func (r *AllocationAccountRegistry) Finalize( account.inflight.Load() != 0 { return current, ErrAllocationAccountLive } - r.slots[slot].account.Store(nil) - if r.generations[slot] != math.MaxUint32 { - r.free = append(r.free, slot) - } + r.removeSlotLocked(slot, account) return account.Snapshot(), nil } @@ -420,9 +725,9 @@ func (r *AllocationAccountRegistry) PeakAllocationMetadata() uint64 { } type allocationAccountRequest struct { - account *AllocationAccount - owner AllocationOwner - site AllocationSite + account *AllocationAccount + owner AllocationOwner + site AllocationSite // checkpoint is nil for every public caller. Same-package fault tests use // it to prove rollback at each unpublished transaction boundary. checkpoint func(allocationCheckpoint) error diff --git a/pkg/common/mpool/allocation_account_mpool_test.go b/pkg/common/mpool/allocation_account_mpool_test.go index 83c23cf9c8067..cea0b62194d8b 100644 --- a/pkg/common/mpool/allocation_account_mpool_test.go +++ b/pkg/common/mpool/allocation_account_mpool_test.go @@ -58,7 +58,7 @@ func TestMPoolAccountedAllocGrowFree(t *testing.T) { require.Equal(t, uintptr(kMemHdrSz), unsafe.Sizeof(memHdr{})) require.Equal(t, uintptr(16), unsafe.Sizeof(allocationLease{})) require.Equal(t, uintptr(64), unsafe.Sizeof(AllocationAccount{})) - require.Equal(t, uintptr(8), unsafe.Sizeof(allocationAccountRegistrySlot{})) + require.LessOrEqual(t, unsafe.Sizeof(allocationAccountRegistrySlot{}), uintptr(32)) registry, account := newTestAllocationAccount(t, 1024, 8) mp := MustNew("accounted-alloc-grow") @@ -179,7 +179,7 @@ func TestMPoolMakeSliceAccounted(t *testing.T) { testAllocationOwner, testAllocationSite, ) - require.ErrorIs(t, err, ErrAllocationAccountInvalid) + require.ErrorIs(t, err, ErrAllocationAllocatorLimit) FreeSlice(mp, values[:0]) require.Zero(t, account.Snapshot().Used) @@ -349,7 +349,7 @@ func TestMPoolAccountedRollback(t *testing.T) { testAllocationOwner, testAllocationSite, ) - require.Error(t, err) + require.ErrorIs(t, err, ErrAllocationAllocatorLimit) require.NotErrorIs(t, err, ErrAllocationAccountCapacity) require.Zero(t, account.Snapshot().Used) require.Zero(t, registry.LiveAllocationMetadata()) @@ -430,7 +430,7 @@ func TestMPoolAccountedReallocZero(t *testing.T) { copy(buffer, bytes.Repeat([]byte{0x5a}, 64)) _, err = mp.ReallocZero(buffer, 128, false) - require.ErrorIs(t, err, ErrAllocationAccountInvalid) + require.ErrorIs(t, err, ErrAllocationAccountMismatch) require.Equal(t, uint64(64), account.Snapshot().Used) replacement, err := mp.ReallocZero(buffer, 128, true) diff --git a/pkg/common/mpool/allocation_account_test.go b/pkg/common/mpool/allocation_account_test.go index a92877d0815a0..7729b2a4d64ea 100644 --- a/pkg/common/mpool/allocation_account_test.go +++ b/pkg/common/mpool/allocation_account_test.go @@ -302,3 +302,187 @@ func TestAllocationAccountReleaseUnderflow(t *testing.T) { _, err = registry.Finalize(account) require.NoError(t, err) } + +func TestAllocationAccountTerminalTombstoneSuspendsAdmission(t *testing.T) { + registry, err := NewAllocationAccountRegistry(3, 2) + require.NoError(t, err) + account, err := registry.Open(64) + require.NoError(t, err) + require.NoError(t, account.acquire(64)) + + snapshot, first, err := registry.CompleteTerminal(account) + require.ErrorIs(t, err, ErrAllocationAccountInvariant) + require.True(t, first) + require.Equal(t, AllocationAccountTerminalInvariantFailure, snapshot.State) + require.Equal(t, uint64(64), snapshot.Used) + require.True(t, snapshot.Sealed) + require.True(t, registry.AdmissionSuspended()) + require.Equal(t, uint32(1), registry.LiveTombstones()) + require.ErrorIs(t, account.acquire(1), ErrAllocationAccountSealed) + _, err = registry.Open(64) + require.ErrorIs(t, err, ErrAllocationAdmissionSuspended) + + repeated, first, err := registry.CompleteTerminal(account) + require.ErrorIs(t, err, ErrAllocationAccountInvariant) + require.False(t, first) + require.Equal(t, snapshot, repeated) + + account.release(64) + require.False(t, registry.AdmissionSuspended()) + require.Zero(t, registry.LiveTombstones()) + _, ok := registry.Resolve(snapshot.Handle) + require.False(t, ok) + + next, err := registry.Open(64) + require.NoError(t, err) + require.NotEqual(t, snapshot.Handle, next.Handle()) + valid, first, err := registry.CompleteTerminal(next) + require.NoError(t, err) + require.True(t, first) + require.Equal(t, AllocationAccountTerminalValid, valid.State) + require.Zero(t, valid.Used) + _, ok = registry.Resolve(valid.Handle) + require.False(t, ok) +} + +func TestAllocationAccountMultipleTombstonesDrainBeforeResume(t *testing.T) { + registry, err := NewAllocationAccountRegistry(3, 2) + require.NoError(t, err) + first, err := registry.Open(64) + require.NoError(t, err) + second, err := registry.Open(64) + require.NoError(t, err) + require.NoError(t, first.acquire(1)) + require.NoError(t, second.acquire(1)) + + _, created, err := registry.CompleteTerminal(first) + require.ErrorIs(t, err, ErrAllocationAccountInvariant) + require.True(t, created) + _, created, err = registry.CompleteTerminal(second) + require.ErrorIs(t, err, ErrAllocationAccountInvariant) + require.True(t, created) + require.Equal(t, uint32(2), registry.LiveTombstones()) + + first.release(1) + require.True(t, registry.AdmissionSuspended()) + require.Equal(t, uint32(1), registry.LiveTombstones()) + second.release(1) + require.False(t, registry.AdmissionSuspended()) + require.Zero(t, registry.LiveTombstones()) +} + +func TestAllocationAccountOpenSuspendLinearization(t *testing.T) { + const contenders = 128 + + registry, err := NewAllocationAccountRegistry(contenders+1, 1) + require.NoError(t, err) + leaked, err := registry.Open(1) + require.NoError(t, err) + require.NoError(t, leaked.acquire(1)) + + start := make(chan struct{}) + opened := make(chan *AllocationAccount, contenders) + var wait sync.WaitGroup + wait.Add(contenders) + for range contenders { + go func() { + defer wait.Done() + <-start + account, openErr := registry.Open(1) + if openErr == nil { + opened <- account + return + } + if !errors.Is(openErr, ErrAllocationAdmissionSuspended) && + !errors.Is(openErr, ErrAllocationGenerationSlots) { + t.Errorf("unexpected open error: %v", openErr) + } + }() + } + close(start) + _, first, err := registry.CompleteTerminal(leaked) + require.ErrorIs(t, err, ErrAllocationAccountInvariant) + require.True(t, first) + wait.Wait() + close(opened) + + _, err = registry.Open(1) + require.ErrorIs(t, err, ErrAllocationAdmissionSuspended) + for account := range opened { + _, _, finishErr := registry.CompleteTerminal(account) + require.NoError(t, finishErr) + } + leaked.release(1) + require.False(t, registry.AdmissionSuspended()) +} + +func TestAllocationAccountCheckpointValidation(t *testing.T) { + registry, err := NewAllocationAccountRegistry(2, 2) + require.NoError(t, err) + account, err := registry.Open(8) + require.NoError(t, err) + other, err := registry.Open(8) + require.NoError(t, err) + + checkpoint, err := account.Checkpoint() + require.NoError(t, err) + require.NoError(t, account.acquire(1)) + require.ErrorIs(t, account.ValidateRollback(checkpoint), ErrAllocationAccountInvariant) + require.NoError(t, account.RollbackToCheckpoint(checkpoint, func() error { + account.release(1) + return nil + })) + require.NoError(t, account.ValidateRollback(checkpoint)) + require.ErrorIs( + t, + account.RollbackToCheckpoint(checkpoint, nil), + ErrAllocationAccountInvalid, + ) + + otherCheckpoint, err := other.Checkpoint() + require.NoError(t, err) + require.ErrorIs(t, account.ValidateRollback(otherCheckpoint), ErrAllocationAccountMismatch) + called := false + require.ErrorIs(t, account.RollbackToCheckpoint(otherCheckpoint, func() error { + called = true + return nil + }), ErrAllocationAccountMismatch) + require.False(t, called) + account.Seal() + require.ErrorIs(t, account.ValidateRollback(checkpoint), ErrAllocationAccountSealed) + _, err = registry.Finalize(account) + require.NoError(t, err) + other.Seal() + _, err = registry.Finalize(other) + require.NoError(t, err) +} + +func TestAllocationFailureReasonsAreNonOverlapping(t *testing.T) { + testCases := []struct { + err error + reason AllocationFailureReason + retryable bool + }{ + {ErrAllocationAccountCapacity, AllocationFailureCapacity, true}, + {ErrAllocationMetadataSlots, AllocationFailureCapacity, true}, + {ErrAllocationAccountSealed, AllocationFailureSealed, false}, + {ErrAllocationAccountMismatch, AllocationFailureMismatch, false}, + {ErrAllocationAllocatorLimit, AllocationFailureAllocatorLimit, false}, + {ErrAllocationAccountInvariant, AllocationFailureInvariant, false}, + {ErrAllocationAdmissionSuspended, AllocationFailureSuspended, false}, + {errors.New("unrelated"), AllocationFailureNone, false}, + } + for _, testCase := range testCases { + require.Equal(t, testCase.reason, AllocationFailureReasonOf(testCase.err)) + require.Equal(t, testCase.retryable, IsRetryableAllocationCapacity(testCase.err)) + } + + // A terminal invariant dominates a joined underlying capacity error, so a + // failed terminal cleanup can never enter the pressure retry loop. + joined := errors.Join( + ErrAllocationAccountCapacity, + ErrAllocationAccountInvariant, + ) + require.Equal(t, AllocationFailureInvariant, AllocationFailureReasonOf(joined)) + require.False(t, IsRetryableAllocationCapacity(joined)) +} diff --git a/pkg/common/mpool/mpool.go b/pkg/common/mpool/mpool.go index d281385ff60f9..4e88b54b1f070 100644 --- a/pkg/common/mpool/mpool.go +++ b/pkg/common/mpool/mpool.go @@ -816,7 +816,12 @@ func (mp *MPool) allocAccountedWithDetailK( // reject unexpected alloc size. if sz < 0 || sz > maxAllocationSize() { logutil.Errorf("mpool memory allocation exceed limit with requested size %d: %s", sz, string(debug.Stack())) - return nil, moerr.NewInternalErrorNoCtxf("mpool memory allocation exceed limit with requested size %d", sz) + return nil, fmt.Errorf( + "%w: requested=%d maximum=%d", + ErrAllocationAllocatorLimit, + sz, + maxAllocationSize(), + ) } if sz == 0 { return nil, nil @@ -1103,7 +1108,7 @@ func (mp *MPool) reAllocWithDetailK(detailk string, old []byte, sz int64, offHea } if hdr.isAccounted() { if !offHeap { - return nil, ErrAllocationAccountInvalid + return nil, ErrAllocationAccountMismatch } accounted := allocationAccountRequest{ account: lease.account, @@ -1222,7 +1227,7 @@ func (mp *MPool) ReallocZero(old []byte, sz int, offHeap bool) ([]byte, error) { if hdr.isAccounted() { if !offHeap { - return nil, ErrAllocationAccountInvalid + return nil, ErrAllocationAccountMismatch } request := allocationAccountRequest{ account: lease.account, @@ -1348,7 +1353,13 @@ func MakeSliceAccounted[T any]( if elementSize == 0 || maxSize <= 0 || uint64(n) > uint64(maxSize)/uint64(elementSize) { - return nil, ErrAllocationAccountInvalid + return nil, fmt.Errorf( + "%w: elements=%d element-size=%d maximum=%d", + ErrAllocationAllocatorLimit, + n, + elementSize, + maxSize, + ) } size := int(uint64(n) * uint64(elementSize)) bs, err := mp.AllocAccounted(size, account, owner, site) diff --git a/pkg/sql/compile/allocation_account_lifecycle.go b/pkg/sql/compile/allocation_account_lifecycle.go new file mode 100644 index 0000000000000..5feb71e218535 --- /dev/null +++ b/pkg/sql/compile/allocation_account_lifecycle.go @@ -0,0 +1,103 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package compile + +import ( + "sync" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/vm/message" +) + +// statementAllocationAttempt owns one local execution generation. The +// MessageBoard pointer is captured at open so prepared/retry Reset cannot make +// terminal cleanup drain a newer board. +type statementAllocationAttempt struct { + registry *mpool.AllocationAccountRegistry + account *mpool.AllocationAccount + board *message.MessageBoard + exporter func(mpool.AllocationAccountTerminalSnapshot) + + once sync.Once + snapshot mpool.AllocationAccountTerminalSnapshot + err error +} + +func (c *Compile) beginAllocationAccountAttempt() ( + *statementAllocationAttempt, + error, +) { + if c == nil || c.allocationAccountRegistry == nil { + return nil, nil + } + if c.proc == nil || c.MessageBoard == nil || c.allocationAttempt != nil || + c.allocationTerminalExporter == nil { + return nil, mpool.ErrAllocationAccountInvariant + } + account, err := c.allocationAccountRegistry.Open(c.allocationAccountLimit) + if err != nil { + return nil, err + } + attempt := &statementAllocationAttempt{ + registry: c.allocationAccountRegistry, + account: account, + board: c.MessageBoard, + exporter: c.allocationTerminalExporter, + } + c.allocationAttempt = attempt + return attempt, nil +} + +func (a *statementAllocationAttempt) finish() ( + mpool.AllocationAccountTerminalSnapshot, + error, +) { + if a == nil { + return mpool.AllocationAccountTerminalSnapshot{}, nil + } + a.once.Do(func() { + // Scope.Run/MergeRun and remote notifier barriers must have returned + // before this point. Draining the board first releases queued JoinMap + // and spill payload ownership through their normal Destroy methods. + a.board.CloseAndDrain() + var first bool + a.snapshot, first, a.err = a.registry.CompleteTerminal(a.account) + if first && a.exporter != nil { + a.exporter(a.snapshot) + } + }) + return a.snapshot, a.err +} + +func (c *Compile) finishAllocationAccountAttempt() error { + if c == nil || c.allocationAttempt == nil { + return nil + } + attempt := c.allocationAttempt + c.allocationAttempt = nil + _, err := attempt.finish() + return err +} + +func (c *Compile) copyAllocationAccountLifecycleTo(dst *Compile) { + if c == nil || dst == nil { + return + } + dst.ConfigureAllocationAccountLifecycle( + c.allocationAccountRegistry, + c.allocationAccountLimit, + c.allocationTerminalExporter, + ) +} diff --git a/pkg/sql/compile/allocation_account_lifecycle_test.go b/pkg/sql/compile/allocation_account_lifecycle_test.go new file mode 100644 index 0000000000000..12add81745975 --- /dev/null +++ b/pkg/sql/compile/allocation_account_lifecycle_test.go @@ -0,0 +1,321 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package compile + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/golang/mock/gomock" + "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/pb/plan" + "github.com/matrixorigin/matrixone/pkg/pb/txn" + "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/matrixorigin/matrixone/pkg/vm" + "github.com/matrixorigin/matrixone/pkg/vm/message" + "github.com/matrixorigin/matrixone/pkg/vm/process" + "github.com/stretchr/testify/require" +) + +type allocationLifecycleErrorOperator struct { + *colexec.MockOperator + err error +} + +func (op *allocationLifecycleErrorOperator) Call( + *process.Process, +) (vm.CallResult, error) { + return vm.CancelResult, op.err +} + +func newRunLifecycleCompile( + t *testing.T, + exporter func(mpool.AllocationAccountTerminalSnapshot), +) (*Compile, *mpool.AllocationAccountRegistry) { + t.Helper() + proc := testutil.NewProcess(t) + ctrl := gomock.NewController(t) + txnClient, txnOperator := newTestTxnClientAndOpWithIsolation( + ctrl, + txn.TxnIsolation_RC, + ) + proc.Base.TxnClient = txnClient + proc.Base.TxnOperator = txnOperator + proc.ReplaceTopCtx(context.Background()) + c := NewCompile( + "local", + "", + "select 1", + "", + "", + nil, + proc, + nil, + false, + nil, + time.Now(), + ) + c.pn = &plan.Plan{Plan: &plan.Plan_Query{Query: &plan.Query{}}} + c.anal = newAnalyzeModule() + registry, err := mpool.NewAllocationAccountRegistry(2, 2) + require.NoError(t, err) + c.ConfigureAllocationAccountLifecycle(registry, 1<<20, exporter) + return c, registry +} + +func newTestAllocationLifecycleCompile( + t *testing.T, + registry *mpool.AllocationAccountRegistry, + exporter func(mpool.AllocationAccountTerminalSnapshot), +) *Compile { + t.Helper() + return &Compile{ + proc: testutil.NewProcess(t), + MessageBoard: message.NewMessageBoard(), + allocationAccountRegistry: registry, + allocationAccountLimit: 1 << 20, + allocationTerminalExporter: exporter, + } +} + +func TestStatementAllocationAttemptZeroTerminalExportsOnce(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(2, 2) + require.NoError(t, err) + var exported []mpool.AllocationAccountTerminalSnapshot + c := newTestAllocationLifecycleCompile(t, registry, func( + snapshot mpool.AllocationAccountTerminalSnapshot, + ) { + exported = append(exported, snapshot) + }) + + attempt, err := c.beginAllocationAccountAttempt() + require.NoError(t, err) + require.NotNil(t, attempt) + buffer, err := c.proc.Mp().AllocAccounted( + 64, + attempt.account, + 1, + 1, + ) + require.NoError(t, err) + c.proc.Mp().Free(buffer) + + require.NoError(t, c.finishAllocationAccountAttempt()) + require.Len(t, exported, 1) + require.Equal(t, mpool.AllocationAccountTerminalValid, exported[0].State) + require.Zero(t, exported[0].Used) + require.Equal(t, uint64(64), exported[0].Peak) + _, ok := registry.Resolve(exported[0].Handle) + require.False(t, ok) + + repeated, err := attempt.finish() + require.NoError(t, err) + require.Equal(t, exported[0], repeated) + require.Len(t, exported, 1) + require.NotSame(t, c.MessageBoard, c.MessageBoard.Reset()) +} + +func TestStatementAllocationAttemptLateFreeDrainsTombstone(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(2, 2) + require.NoError(t, err) + var exported []mpool.AllocationAccountTerminalSnapshot + c := newTestAllocationLifecycleCompile(t, registry, func( + snapshot mpool.AllocationAccountTerminalSnapshot, + ) { + exported = append(exported, snapshot) + }) + attempt, err := c.beginAllocationAccountAttempt() + require.NoError(t, err) + buffer, err := c.proc.Mp().AllocAccounted( + 64, + attempt.account, + 1, + 1, + ) + require.NoError(t, err) + + err = c.finishAllocationAccountAttempt() + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + require.Len(t, exported, 1) + require.Equal( + t, + mpool.AllocationAccountTerminalInvariantFailure, + exported[0].State, + ) + require.Equal(t, uint64(cap(buffer)), exported[0].Used) + require.True(t, registry.AdmissionSuspended()) + _, err = registry.Open(1) + require.ErrorIs(t, err, mpool.ErrAllocationAdmissionSuspended) + + // The physical allocation retains its original account after the producer + // process has detached the generation. Its normal Free drains the + // tombstone; no synthetic release is needed. + c.proc.Mp().Free(buffer) + require.False(t, registry.AdmissionSuspended()) + _, ok := registry.Resolve(exported[0].Handle) + require.False(t, ok) + + c.MessageBoard = c.MessageBoard.Reset() + next, err := c.beginAllocationAccountAttempt() + require.NoError(t, err) + require.NotEqual(t, attempt.account.Handle(), next.account.Handle()) + require.NoError(t, c.finishAllocationAccountAttempt()) + require.Len(t, exported, 2) +} + +func TestStatementAllocationAttemptConcurrentTerminalIsOneShot(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + var exports atomic.Int32 + c := newTestAllocationLifecycleCompile(t, registry, func( + mpool.AllocationAccountTerminalSnapshot, + ) { + exports.Add(1) + }) + attempt, err := c.beginAllocationAccountAttempt() + require.NoError(t, err) + + const contenders = 128 + start := make(chan struct{}) + errs := make(chan error, contenders) + var wait sync.WaitGroup + wait.Add(contenders) + for range contenders { + go func() { + defer wait.Done() + <-start + _, finishErr := attempt.finish() + errs <- finishErr + }() + } + close(start) + wait.Wait() + close(errs) + for finishErr := range errs { + require.NoError(t, finishErr) + } + require.Equal(t, int32(1), exports.Load()) + c.allocationAttempt = nil +} + +func TestStatementAllocationAttemptRejectsOverlappingOpen(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + c := newTestAllocationLifecycleCompile(t, registry, func( + mpool.AllocationAccountTerminalSnapshot, + ) { + }) + attempt, err := c.beginAllocationAccountAttempt() + require.NoError(t, err) + + _, err = c.beginAllocationAccountAttempt() + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + require.NoError(t, c.finishAllocationAccountAttempt()) + _, finishErr := attempt.finish() + require.NoError(t, finishErr) + + next, err := registry.Open(1) + require.NoError(t, err, "rejected overlap must not consume a slot") + _, _, err = registry.CompleteTerminal(next) + require.NoError(t, err) +} + +func TestStatementAllocationAttemptRequiresTerminalExporter(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + c := newTestAllocationLifecycleCompile(t, registry, nil) + + _, err = c.beginAllocationAccountAttempt() + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + account, err := registry.Open(1) + require.NoError(t, err, "failed begin must not consume a generation slot") + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestCompileRunFinalizesAllocationAttemptOnCancellation(t *testing.T) { + var snapshots []mpool.AllocationAccountTerminalSnapshot + c, registry := newRunLifecycleCompile(t, func( + snapshot mpool.AllocationAccountTerminalSnapshot, + ) { + snapshots = append(snapshots, snapshot) + }) + // The canceled outer context is observed after runOnce, after the + // allocation generation has opened. + canceled, cancel := context.WithCancel(context.Background()) + cancel() + c.proc.ReplaceTopCtx(canceled) + c.scopes = []*Scope{newScope(magicType(255))} + + _, err := c.Run(0) + require.ErrorIs(t, err, context.Canceled) + require.Len(t, snapshots, 1) + require.Equal(t, mpool.AllocationAccountTerminalValid, snapshots[0].State) + require.Zero(t, snapshots[0].Used) + _, ok := registry.Resolve(snapshots[0].Handle) + require.False(t, ok) + c.Release() +} + +func TestCompileRunFinalizesAllocationAttemptOnExecutionError(t *testing.T) { + var snapshots []mpool.AllocationAccountTerminalSnapshot + c, registry := newRunLifecycleCompile(t, func( + snapshot mpool.AllocationAccountTerminalSnapshot, + ) { + snapshots = append(snapshots, snapshot) + }) + executionErr := moerr.NewInternalErrorNoCtx("allocation lifecycle test") + scope := newScope(Normal) + scope.Proc = c.proc.NewNoContextChildProc(0) + scope.RootOp = &allocationLifecycleErrorOperator{ + MockOperator: colexec.NewMockOperator(), + err: executionErr, + } + c.scopes = []*Scope{scope} + + _, err := c.Run(0) + require.ErrorIs(t, err, executionErr) + require.Len(t, snapshots, 1) + require.Equal(t, mpool.AllocationAccountTerminalValid, snapshots[0].State) + _, ok := registry.Resolve(snapshots[0].Handle) + require.False(t, ok) + c.Release() +} + +func TestCompileRunFinalizesAllocationAttemptOnPanic(t *testing.T) { + var snapshots []mpool.AllocationAccountTerminalSnapshot + c, registry := newRunLifecycleCompile(t, func( + snapshot mpool.AllocationAccountTerminalSnapshot, + ) { + snapshots = append(snapshots, snapshot) + }) + c.scopes = []*Scope{newScope(magicType(255))} + // Force the panic after beginAllocationAccountAttempt and before runOnce. + c.lockMeta = nil + + require.Panics(t, func() { + _, _ = c.Run(0) + }) + require.Len(t, snapshots, 1) + require.Equal(t, mpool.AllocationAccountTerminalValid, snapshots[0].State) + _, ok := registry.Resolve(snapshots[0].Handle) + require.False(t, ok) + c.Release() +} diff --git a/pkg/sql/compile/compile.go b/pkg/sql/compile/compile.go index 07a3692630ecd..66d96eb05c66a 100644 --- a/pkg/sql/compile/compile.go +++ b/pkg/sql/compile/compile.go @@ -284,6 +284,9 @@ func (c *Compile) clear() { c.fuzzys[i].release() } + if err := c.finishAllocationAccountAttempt(); err != nil { + logutil.Errorf("allocation account terminal cleanup failed: %v", err) + } c.MessageBoard = c.MessageBoard.Reset() c.fuzzys = c.fuzzys[:0] c.scopes = c.scopes[:0] @@ -317,6 +320,10 @@ func (c *Compile) clear() { c.needLockMeta = false c.isInternal = false c.resourceAttemptOwnerEligible = false + c.allocationAccountRegistry = nil + c.allocationAccountLimit = 0 + c.allocationTerminalExporter = nil + c.allocationAttempt = nil c.isPrepare = false c.hasMergeOp = false c.needBlock = false @@ -6988,6 +6995,19 @@ func (c *Compile) SetResourceAttemptOwnerEligible() { c.resourceAttemptOwnerEligible = true } +// ConfigureAllocationAccountLifecycle installs the dormant generation +// provider used by owner-activation PRs. A nil registry keeps production on +// the legacy path and opens no generation. +func (c *Compile) ConfigureAllocationAccountLifecycle( + registry *mpool.AllocationAccountRegistry, + limit uint64, + exporter func(mpool.AllocationAccountTerminalSnapshot), +) { + c.allocationAccountRegistry = registry + c.allocationAccountLimit = limit + c.allocationTerminalExporter = exporter +} + func (c *Compile) SetBuildPlanFunc(buildPlanFunc func(ctx context.Context) (*plan2.Plan, error)) { c.buildPlanFunc = buildPlanFunc } diff --git a/pkg/sql/compile/compile2.go b/pkg/sql/compile/compile2.go index f8b44fe215a9e..4c87325b0858b 100644 --- a/pkg/sql/compile/compile2.go +++ b/pkg/sql/compile/compile2.go @@ -17,6 +17,7 @@ package compile import ( "context" "encoding/hex" + "errors" "math" gotrace "runtime/trace" "strings" @@ -258,8 +259,22 @@ func (c *Compile) Run(_ uint64) (queryResult *util2.RunResult, err error) { attemptAnal := runC.anal var coordinatorPhaseStart time.Time var coordinatorPhaseBase time.Duration + var allocationAttempt *statementAllocationAttempt + finishAllocationAttempt := func() error { + if allocationAttempt == nil { + return nil + } + attempt := allocationAttempt + allocationAttempt = nil + if runC != nil && runC.allocationAttempt == attempt { + runC.allocationAttempt = nil + } + _, finishErr := attempt.finish() + return finishErr + } defer func() { if recovered := recover(); recovered != nil { + _ = finishAllocationAttempt() if attemptOpen { if !coordinatorPhaseStart.IsZero() { attemptPreRunWall = coordinatorPhaseBase + time.Since(coordinatorPhaseStart) @@ -288,7 +303,11 @@ func (c *Compile) Run(_ uint64) (queryResult *util2.RunResult, err error) { // Before compile.runOnce, Reset the 'StatsInfo' execution related resources in context // running. - if err = runC.prePipelineInitializer(); err == nil { + allocationAttempt, err = runC.beginAllocationAccountAttempt() + if err == nil { + err = runC.prePipelineInitializer() + } + if err == nil { preRunWall = carriedPreRunWall + time.Since(preRunOnceStart) attemptPreRunWall = preRunWall runC.MessageBoard.BeforeRunonce() @@ -315,6 +334,15 @@ func (c *Compile) Run(_ uint64) (queryResult *util2.RunResult, err error) { attemptPreRunWall = preRunWall coordinatorPhaseStart = time.Time{} coordinatorPhaseBase = 0 + if terminalErr := finishAllocationAttempt(); terminalErr != nil { + err = errors.Join(err, terminalErr) + resourceRecorder.finishAttempt( + uint64(retryTimes), attemptStart, preRunWall, attemptRemoteWait, stats, + attemptScopes, attemptAnal, c.addr, false, + ) + attemptOpen = false + return nil, err + } c.fatalLog(retryTimes, err) if !c.canRetry(err) { @@ -430,6 +458,15 @@ func (c *Compile) Run(_ uint64) (queryResult *util2.RunResult, err error) { coordinatorPhaseStart = time.Time{} coordinatorPhaseBase = 0 } + if terminalErr := finishAllocationAttempt(); terminalErr != nil { + err = errors.Join(err, terminalErr) + resourceRecorder.finishAttempt( + uint64(retryTimes), attemptStart, attemptPreRunWall, attemptRemoteWait, stats, + attemptScopes, attemptAnal, c.addr, false, + ) + attemptOpen = false + return nil, err + } queryResult.AffectRows = runC.getAffectedRows() if c.uid != "mo_logger" && strings.Contains(strings.ToLower(c.sql), "insert") && @@ -732,6 +769,7 @@ func (c *Compile) buildRetryCompile(defChanged bool) (*Compile, error) { var e error runC := NewCompile(c.addr, c.db, c.sql, c.tenant, c.uid, c.e, c.proc, c.stmt, c.isInternal, c.cnLabel, c.startAt) + c.copyAllocationAccountLifecycleTo(runC) runC.SetQuerySchedulingIntent(c.querySchedulingIntent) runC.SetSchedulingTraceRecorder(c.schedulingTrace) runC.SetOriginSQL(c.originSQL) diff --git a/pkg/sql/compile/compile_test.go b/pkg/sql/compile/compile_test.go index 4e93fd2eeec8b..7cc7877d0b5b6 100644 --- a/pkg/sql/compile/compile_test.go +++ b/pkg/sql/compile/compile_test.go @@ -128,6 +128,16 @@ func TestCompileRunPreservesBinaryPrepareParamAcrossRetries(t *testing.T) { } c := NewCompile("test", "test", "select ?", "", "", newStubEngine(), proc, stmts[0], false, nil, time.Now()) + registry, err := mpool.NewAllocationAccountRegistry(4, 4) + require.NoError(t, err) + var terminalSnapshots []mpool.AllocationAccountTerminalSnapshot + c.ConfigureAllocationAccountLifecycle( + registry, + 1<<20, + func(snapshot mpool.AllocationAccountTerminalSnapshot) { + terminalSnapshots = append(terminalSnapshots, snapshot) + }, + ) require.NoError(t, c.Compile(ctx, pn, fill)) _, err = c.Run(0) require.NoError(t, err) @@ -136,6 +146,18 @@ func TestCompileRunPreservesBinaryPrepareParamAcrossRetries(t *testing.T) { require.Zero(t, params.Length()) require.Nil(t, params.GetData()) require.Nil(t, params.GetArea()) + require.Len(t, terminalSnapshots, 3) + seenHandles := make(map[mpool.AllocationAccountHandle]struct{}, 3) + for _, snapshot := range terminalSnapshots { + require.Equal(t, mpool.AllocationAccountTerminalValid, snapshot.State) + require.Zero(t, snapshot.Used) + require.True(t, snapshot.Sealed) + _, duplicate := seenHandles[snapshot.Handle] + require.False(t, duplicate) + seenHandles[snapshot.Handle] = struct{}{} + _, ok := registry.Resolve(snapshot.Handle) + require.False(t, ok) + } c.Release() proc.Free() diff --git a/pkg/sql/compile/remoterunServer.go b/pkg/sql/compile/remoterunServer.go index 297a4657d40e1..3debcf7fa45c1 100644 --- a/pkg/sql/compile/remoterunServer.go +++ b/pkg/sql/compile/remoterunServer.go @@ -188,7 +188,7 @@ func (receiver *messageReceiverOnServer) waitUntilDisconnectedOrCancelled() { } } -func handlePipelineMessage(receiver *messageReceiverOnServer) error { +func handlePipelineMessage(receiver *messageReceiverOnServer) (err error) { switch receiver.messageTyp { case pipeline.Method_PrepareDoneNotifyMessage: @@ -251,6 +251,7 @@ func handlePipelineMessage(receiver *messageReceiverOnServer) error { if errBuildCompile != nil { return errBuildCompile } + var allocationAttempt *statementAllocationAttempt var runErr error defer func() { // Capture operator and descendant facts before cleanup. The MPool @@ -262,6 +263,11 @@ func handlePipelineMessage(receiver *messageReceiverOnServer) error { descendant := runCompile.anal.remoteResourceSummary() expectedDirect := countExpectedRemoteScopes(runCompile.scopes, receiver.cnInformation.cnAddr) memoryPool := runCompile.proc.Mp() + if allocationAttempt != nil { + runCompile.allocationAttempt = nil + _, terminalErr := allocationAttempt.finish() + err = errors.Join(err, terminalErr) + } runCompile.clear() localMemory, localMemoryQuality := memoryPool.ResourceSnapshot() aggregate := composeRemoteResourceAggregate( @@ -301,6 +307,10 @@ func handlePipelineMessage(receiver *messageReceiverOnServer) error { } runCompile.scopes = []*Scope{s} + allocationAttempt, runErr = runCompile.beginAllocationAccountAttempt() + if runErr != nil { + return runErr + } runCompile.InitPipelineContextToExecuteQuery() normalizeRemoteDispatchReceiverAddresses(s, runCompile.addr) diff --git a/pkg/sql/compile/types.go b/pkg/sql/compile/types.go index 4b7af939ae88b..a596dfa427d4e 100644 --- a/pkg/sql/compile/types.go +++ b/pkg/sql/compile/types.go @@ -21,6 +21,7 @@ import ( "github.com/google/uuid" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" icebergapi "github.com/matrixorigin/matrixone/pkg/iceberg/api" @@ -343,6 +344,10 @@ type Compile struct { // resourceAttemptOwnerEligible is set only for the top-level statement // Compile. The statement root still arbitrates the single actual owner. resourceAttemptOwnerEligible bool + allocationAccountRegistry *mpool.AllocationAccountRegistry + allocationAccountLimit uint64 + allocationTerminalExporter func(mpool.AllocationAccountTerminalSnapshot) + allocationAttempt *statementAllocationAttempt hasMergeOp bool // ncpu set as system.GoRoutines() while NewCompile, instead of global static value. diff --git a/pkg/vm/message/message.go b/pkg/vm/message/message.go index f8590dd81f33b..22c2a500c3b40 100644 --- a/pkg/vm/message/message.go +++ b/pkg/vm/message/message.go @@ -75,6 +75,7 @@ type MessageCenter struct { type MessageBoard struct { reset bool // for debug purpose + closed bool multiCN bool stmtId uuid.UUID messageCenter *MessageCenter @@ -98,10 +99,18 @@ func (m *MessageBoard) finalize() { } func (m *MessageBoard) DebugString() string { + if m == nil || m.rwMutex == nil { + return "messageBoard is nil\n" + } + m.rwMutex.RLock() + defer m.rwMutex.RUnlock() buf := bytes.NewBuffer(make([]byte, 0, 400)) if m.reset { buf.WriteString("messageBoard has been reseted!\n") } + if m.closed { + buf.WriteString("messageBoard has been closed!\n") + } if m.multiCN { buf.WriteString("messageBoard on MultiCN\n") } else { @@ -135,26 +144,79 @@ func (m *MessageBoard) BeforeRunonce() { // call this before runonce m.rwMutex.Lock() defer m.rwMutex.Unlock() - m.reset = false + if !m.closed { + m.reset = false + } } func (m *MessageBoard) Reset() *MessageBoard { - if m.multiCN { - m.messageCenter.RwMutex.Lock() - delete(m.messageCenter.StmtIDToBoard, m.stmtId) - m.messageCenter.RwMutex.Unlock() + m.rwMutex.RLock() + multiCN := m.multiCN + center := m.messageCenter + stmtID := m.stmtId + m.rwMutex.RUnlock() + if multiCN { + center.RwMutex.Lock() + delete(center.StmtIDToBoard, stmtID) + center.RwMutex.Unlock() // other pipeline could still access thie messageBoard // so reset current message board to a new one return NewMessageBoard() } m.rwMutex.Lock() defer m.rwMutex.Unlock() + if m.closed { + return NewMessageBoard() + } m.cleanupQueuedMessagesLocked() m.multiCN = false m.reset = true return m } +// CloseAndDrain is the terminal MessageBoard boundary for one execution +// attempt. Callers invoke it only after all scope and remote-notifier producers +// are quiescent. It removes a multi-CN registration, destroys every queued +// ownership-bearing message, and prevents a late producer from republishing +// into the closed generation. The operation is idempotent; true identifies +// the call that performed the close. +func (m *MessageBoard) CloseAndDrain() bool { + if m == nil || m.rwMutex == nil { + return false + } + m.rwMutex.RLock() + multiCN := m.multiCN + center := m.messageCenter + stmtID := m.stmtId + m.rwMutex.RUnlock() + if multiCN && center != nil { + center.RwMutex.Lock() + if center.StmtIDToBoard[stmtID] == m { + delete(center.StmtIDToBoard, stmtID) + } + center.RwMutex.Unlock() + } + + m.rwMutex.Lock() + defer m.rwMutex.Unlock() + if m.closed { + return false + } + m.closed = true + m.reset = true + for _, waiter := range m.waiters { + if waiter == nil { + continue + } + select { + case waiter <- true: + default: + } + } + m.cleanupQueuedMessagesLocked() + return true +} + func (m *MessageBoard) cleanupQueuedMessages() { if m == nil || m.rwMutex == nil { return @@ -198,6 +260,11 @@ func NewMessageReceiver(tags []int32, addr MessageAddress, mb *MessageBoard) *Me func SendMessage(m Message, mb *MessageBoard) { if m.GetReceiverAddr().CnAddr == CURRENTCN { // message for current CN mb.rwMutex.Lock() + if mb.closed { + mb.rwMutex.Unlock() + m.Destroy() + return + } mb.messages = append(mb.messages, &m) if m.NeedBlock() { // broadcast for block message @@ -214,7 +281,7 @@ func SendMessage(m Message, mb *MessageBoard) { } } -func (mr *MessageReceiver) receiveMessageNonBlock() []Message { +func (mr *MessageReceiver) receiveMessageNonBlock() ([]Message, bool) { mr.mb.rwMutex.RLock() defer mr.mb.rwMutex.RUnlock() var result []Message @@ -235,25 +302,44 @@ func (mr *MessageReceiver) receiveMessageNonBlock() []Message { } } } - return result + return result, mr.mb.closed } func (mr *MessageReceiver) ReceiveMessage(needBlock bool, ctx context.Context) ([]Message, bool, error) { - var result = mr.receiveMessageNonBlock() + result, closed := mr.receiveMessageNonBlock() if !needBlock || len(result) > 0 { return result, false, nil } + if closed { + return result, false, moerr.NewInternalErrorNoCtx( + "message board is closed", + ) + } if mr.waiter == nil { mr.waiter = make(chan bool, 1) mr.mb.rwMutex.Lock() - mr.mb.waiters = append(mr.mb.waiters, mr.waiter) + if mr.mb.closed { + closed = true + } else { + mr.mb.waiters = append(mr.mb.waiters, mr.waiter) + } mr.mb.rwMutex.Unlock() + if closed { + return result, false, moerr.NewInternalErrorNoCtx( + "message board is closed", + ) + } } for { - result = mr.receiveMessageNonBlock() + result, closed = mr.receiveMessageNonBlock() if len(result) > 0 { break } + if closed { + return result, false, moerr.NewInternalErrorNoCtx( + "message board is closed", + ) + } timeout := messageTimeout if mr.debug { timeout = 1 * time.Second diff --git a/pkg/vm/message/message_test.go b/pkg/vm/message/message_test.go index 2a5c2adfcd450..011392d217b18 100644 --- a/pkg/vm/message/message_test.go +++ b/pkg/vm/message/message_test.go @@ -15,6 +15,7 @@ package message import ( + "context" "os" "runtime" "runtime/debug" @@ -35,6 +36,32 @@ type testMessage struct { destroyed *atomic.Int32 } +type accountedTestMessage struct { + mp *mpool.MPool + buffer []byte +} + +func (m *accountedTestMessage) Serialize() []byte { return nil } + +func (m *accountedTestMessage) Deserialize([]byte) Message { return m } + +func (m *accountedTestMessage) NeedBlock() bool { return true } + +func (m *accountedTestMessage) GetMsgTag() int32 { return 1 } + +func (m *accountedTestMessage) GetReceiverAddr() MessageAddress { + return AddrBroadCastOnCurrentCN() +} + +func (m *accountedTestMessage) DebugString() string { return "accounted test message" } + +func (m *accountedTestMessage) Destroy() { + if m.buffer != nil { + m.mp.Free(m.buffer) + m.buffer = nil + } +} + func (m testMessage) Serialize() []byte { return nil } @@ -97,6 +124,74 @@ func TestMessageBoardResetDestroysQueuedMessages(t *testing.T) { require.Empty(t, mb.waiters) } +func TestMessageBoardCloseAndDrainRejectsLateMessages(t *testing.T) { + var destroyed atomic.Int32 + mb := NewMessageBoard() + SendMessage(testMessage{tag: 1, destroyed: &destroyed}, mb) + + receiver := NewMessageReceiver( + []int32{2}, + AddrBroadCastOnCurrentCN(), + mb, + ) + waiting := make(chan error, 1) + go func() { + _, _, err := receiver.ReceiveMessage(true, context.Background()) + waiting <- err + }() + + require.True(t, mb.CloseAndDrain()) + require.False(t, mb.CloseAndDrain()) + require.ErrorContains(t, <-waiting, "message board is closed") + require.Equal(t, int32(1), destroyed.Load()) + require.Empty(t, mb.messages) + require.Empty(t, mb.waiters) + + SendMessage(testMessage{tag: 2, destroyed: &destroyed}, mb) + require.Equal(t, int32(2), destroyed.Load()) + require.NotSame(t, mb, mb.Reset()) +} + +func TestMessageBoardCloseAndDrainRemovesMultiCNRegistration(t *testing.T) { + center := &MessageCenter{ + StmtIDToBoard: make(map[uuid.UUID]*MessageBoard), + RwMutex: &sync.Mutex{}, + } + stmtID := uuid.New() + mb := NewMessageBoard().SetMultiCN(center, stmtID) + require.Same(t, mb, center.StmtIDToBoard[stmtID]) + + require.True(t, mb.CloseAndDrain()) + _, ok := center.StmtIDToBoard[stmtID] + require.False(t, ok) +} + +func TestClosedMessageBoardLatePayloadDrainsOriginalGeneration(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + mp := mpool.MustNewZero() + buffer, err := mp.AllocAccounted(64, account, 1, 1) + require.NoError(t, err) + + mb := NewMessageBoard() + require.True(t, mb.CloseAndDrain()) + terminal, first, err := registry.CompleteTerminal(account) + require.True(t, first) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + require.Equal(t, uint64(cap(buffer)), terminal.Used) + require.True(t, registry.AdmissionSuspended()) + + // A producer that already owns a payload cannot republish it after the + // attempt boundary. Destroy performs the physical Free against the account + // captured by the allocation, even though that generation is now sealed. + SendMessage(&accountedTestMessage{mp: mp, buffer: buffer}, mb) + require.False(t, registry.AdmissionSuspended()) + _, ok := registry.Resolve(account.Handle()) + require.False(t, ok) +} + func TestMessageBoardFinalizerDestroysQueuedMessages(t *testing.T) { var destroyed atomic.Int32 From 0938eea790c7355175cbadccd2665c2c0424a72c Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 21:20:24 +0800 Subject: [PATCH 15/61] executor: activate allocation-accounted hash tables --- pkg/common/hashmap/inthashmap.go | 10 +- pkg/common/hashmap/strhashmap.go | 10 +- pkg/common/mpool/allocation_account.go | 39 +- pkg/common/mpool/allocation_account_test.go | 18 + pkg/container/hashtable/allocation_account.go | 114 ++++++ .../hashtable/allocation_account_test.go | 365 ++++++++++++++++++ pkg/container/hashtable/int64_hash_map.go | 108 +++++- pkg/container/hashtable/string_hash_map.go | 102 ++++- pkg/sql/colexec/hashbuild/budget.go | 46 +++ pkg/sql/colexec/hashbuild/hashmap.go | 61 ++- pkg/sql/colexec/hashbuild/hashmap_test.go | 130 +++++++ pkg/sql/colexec/hashbuild/types.go | 58 +++ pkg/sql/colexec/rightdedupjoin/join.go | 7 + pkg/sql/colexec/rightdedupjoin/types.go | 38 ++ .../compile/allocation_account_lifecycle.go | 202 +++++++++- .../allocation_account_lifecycle_test.go | 139 +++++++ pkg/sql/compile/analyze_module.go | 9 + pkg/sql/compile/compile.go | 18 + pkg/sql/compile/compile2.go | 11 +- pkg/sql/compile/remoterunClient.go | 1 + pkg/sql/compile/remoterunServer.go | 23 +- pkg/sql/compile/resource_accounting.go | 55 ++- pkg/sql/compile/resource_accounting_test.go | 68 +++- pkg/sql/compile/types.go | 2 + pkg/util/resource/summary.go | 118 +++++- pkg/util/resource/usage_test.go | 24 ++ pkg/vm/process/hashbuild_budget.go | 56 ++- pkg/vm/process/hashbuild_budget_test.go | 38 ++ 28 files changed, 1789 insertions(+), 81 deletions(-) create mode 100644 pkg/container/hashtable/allocation_account.go create mode 100644 pkg/container/hashtable/allocation_account_test.go diff --git a/pkg/common/hashmap/inthashmap.go b/pkg/common/hashmap/inthashmap.go index 22de287951b43..10e7d26fec42f 100644 --- a/pkg/common/hashmap/inthashmap.go +++ b/pkg/common/hashmap/inthashmap.go @@ -32,8 +32,16 @@ func init() { } func NewIntHashMap(hasNull bool, memPool *mpool.MPool) (*IntHashMap, error) { + return NewIntHashMapWithAllocation(hasNull, memPool, nil) +} + +func NewIntHashMapWithAllocation( + hasNull bool, + memPool *mpool.MPool, + allocation *hashtable.AllocationAccountSelection, +) (*IntHashMap, error) { mp := &hashtable.Int64HashMap{} - if err := mp.Init(memPool); err != nil { + if err := mp.InitWithAllocation(memPool, allocation); err != nil { return nil, err } return &IntHashMap{ diff --git a/pkg/common/hashmap/strhashmap.go b/pkg/common/hashmap/strhashmap.go index 98addc3841f0f..6fa09dbf1ba56 100644 --- a/pkg/common/hashmap/strhashmap.go +++ b/pkg/common/hashmap/strhashmap.go @@ -38,8 +38,16 @@ func init() { } func NewStrHashMap(hasNull bool, memPool *mpool.MPool) (*StrHashMap, error) { + return NewStrHashMapWithAllocation(hasNull, memPool, nil) +} + +func NewStrHashMapWithAllocation( + hasNull bool, + memPool *mpool.MPool, + allocation *hashtable.AllocationAccountSelection, +) (*StrHashMap, error) { mp := &hashtable.StringHashMap{} - if err := mp.Init(memPool); err != nil { + if err := mp.InitWithAllocation(memPool, allocation); err != nil { return nil, err } return &StrHashMap{ diff --git a/pkg/common/mpool/allocation_account.go b/pkg/common/mpool/allocation_account.go index ec82418fc0eee..25584c257b437 100644 --- a/pkg/common/mpool/allocation_account.go +++ b/pkg/common/mpool/allocation_account.go @@ -489,6 +489,17 @@ func (r *AllocationAccountRegistry) OpenWithController( // first=false. func (r *AllocationAccountRegistry) CompleteTerminal( account *AllocationAccount, +) (snapshot AllocationAccountTerminalSnapshot, first bool, err error) { + return r.CompleteTerminalWithError(account, nil) +} + +// CompleteTerminalWithError additionally records an owner-lifecycle invariant +// discovered after physical producers quiesced. A zero-live failure is removed +// immediately (there is no provenance to retain); a nonzero failure follows +// the same tombstone/suspension path as CompleteTerminal. +func (r *AllocationAccountRegistry) CompleteTerminalWithError( + account *AllocationAccount, + terminalCause error, ) (snapshot AllocationAccountTerminalSnapshot, first bool, err error) { if r == nil || account == nil || account.registry != r { return snapshot, false, ErrAllocationAccountInvalid @@ -525,9 +536,18 @@ func (r *AllocationAccountRegistry) CompleteTerminal( AllocationAccountSnapshot: current, State: AllocationAccountTerminalValid, } + if terminalCause != nil { + snapshot.State = AllocationAccountTerminalInvariantFailure + } if current.Used == 0 { entry.terminal = &snapshot r.removeSlotLocked(slot, account) + if snapshot.State == AllocationAccountTerminalInvariantFailure { + return snapshot, true, errors.Join( + terminalCause, + newAllocationTerminalInvariantError(snapshot), + ) + } return snapshot, true, nil } @@ -542,7 +562,10 @@ func (r *AllocationAccountRegistry) CompleteTerminal( if account.Snapshot().Used == 0 && account.inflight.Load() == 0 { r.removeTombstoneLocked(slot, account) } - return snapshot, true, newAllocationTerminalInvariantError(snapshot) + return snapshot, true, errors.Join( + terminalCause, + newAllocationTerminalInvariantError(snapshot), + ) } func newAllocationTerminalInvariantError( @@ -724,6 +747,20 @@ func (r *AllocationAccountRegistry) PeakAllocationMetadata() uint64 { return r.peakAllocations.Load() } +func (r *AllocationAccountRegistry) MaxAllocationMetadata() uint64 { + if r == nil { + return 0 + } + return r.maxAllocations +} + +func (r *AllocationAccountRegistry) GenerationCapacity() uint32 { + if r == nil || len(r.slots) == 0 { + return 0 + } + return uint32(len(r.slots) - 1) +} + type allocationAccountRequest struct { account *AllocationAccount owner AllocationOwner diff --git a/pkg/common/mpool/allocation_account_test.go b/pkg/common/mpool/allocation_account_test.go index 7729b2a4d64ea..c0d07192538d6 100644 --- a/pkg/common/mpool/allocation_account_test.go +++ b/pkg/common/mpool/allocation_account_test.go @@ -345,6 +345,24 @@ func TestAllocationAccountTerminalTombstoneSuspendsAdmission(t *testing.T) { require.False(t, ok) } +func TestAllocationAccountZeroLiveOwnerInvariantExportsFailure(t *testing.T) { + registry, err := NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + account, err := registry.Open(1) + require.NoError(t, err) + ownerErr := errors.New("owner teardown failed") + + snapshot, first, err := registry.CompleteTerminalWithError(account, ownerErr) + require.True(t, first) + require.ErrorIs(t, err, ownerErr) + require.ErrorIs(t, err, ErrAllocationAccountInvariant) + require.Equal(t, AllocationAccountTerminalInvariantFailure, snapshot.State) + require.Zero(t, snapshot.Used) + require.False(t, registry.AdmissionSuspended()) + _, ok := registry.Resolve(account.Handle()) + require.False(t, ok) +} + func TestAllocationAccountMultipleTombstonesDrainBeforeResume(t *testing.T) { registry, err := NewAllocationAccountRegistry(3, 2) require.NoError(t, err) diff --git a/pkg/container/hashtable/allocation_account.go b/pkg/container/hashtable/allocation_account.go new file mode 100644 index 0000000000000..df40ba29f6cb5 --- /dev/null +++ b/pkg/container/hashtable/allocation_account.go @@ -0,0 +1,114 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package hashtable + +import ( + "unsafe" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" +) + +func HashMapBlockDescriptorBytes() uint64 { + return uint64(unsafe.Sizeof([]Int64HashMapCell(nil))) +} + +// AllocationAccountSelection is the immutable provenance for one hash table's +// cell blocks and outer descriptor storage. The outer []slice-header backing +// is itself data-scaled ownership and therefore uses a distinct physical site. +type AllocationAccountSelection struct { + account *mpool.AllocationAccount + owner mpool.AllocationOwner + cellSite mpool.AllocationSite + descriptorSite mpool.AllocationSite +} + +func NewAllocationAccountSelection( + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, + cellSite mpool.AllocationSite, + descriptorSite mpool.AllocationSite, +) (*AllocationAccountSelection, error) { + selection := &AllocationAccountSelection{ + account: account, + owner: owner, + cellSite: cellSite, + descriptorSite: descriptorSite, + } + if err := selection.validate(); err != nil { + return nil, err + } + return selection, nil +} + +func (s *AllocationAccountSelection) validate() error { + if s == nil || s.account == nil || s.account.Handle() == 0 || + s.owner < mpool.AllocationOwnerMin || + s.owner > mpool.AllocationOwnerMax || + s.cellSite < mpool.AllocationSiteMin || + s.descriptorSite < mpool.AllocationSiteMin || + s.cellSite == s.descriptorSite { + return mpool.ErrAllocationAccountInvalid + } + return nil +} + +func makeHashTableCellSlice[T any]( + length int, + mp *mpool.MPool, + selection *AllocationAccountSelection, + site mpool.AllocationSite, +) ([]T, error) { + if selection == nil { + return mpool.MakeSlice[T](length, mp, true) + } + if err := selection.validate(); err != nil { + return nil, err + } + return mpool.MakeSliceAccounted[T]( + length, + mp, + selection.account, + selection.owner, + site, + ) +} + +func makeHashTableDescriptorSlice[T any]( + length int, + mp *mpool.MPool, + selection *AllocationAccountSelection, + site mpool.AllocationSite, +) ([]T, error) { + if selection == nil { + return make([]T, length), nil + } + return makeHashTableCellSlice[T](length, mp, selection, site) +} + +func freeHashTableCellSlice[T any](mp *mpool.MPool, values []T) { + if cap(values) > 0 { + mpool.FreeSlice(mp, values) + } +} + +func freeHashTableDescriptorSlice[T any]( + mp *mpool.MPool, + values []T, + selection *AllocationAccountSelection, +) { + if selection != nil { + freeHashTableCellSlice(mp, values) + } +} diff --git a/pkg/container/hashtable/allocation_account_test.go b/pkg/container/hashtable/allocation_account_test.go new file mode 100644 index 0000000000000..178261c973ef4 --- /dev/null +++ b/pkg/container/hashtable/allocation_account_test.go @@ -0,0 +1,365 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package hashtable + +import ( + "testing" + "unsafe" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/stretchr/testify/require" +) + +const ( + testHashTableOwner mpool.AllocationOwner = 1 + testHashTableCellSite mpool.AllocationSite = 24 + testHashTableDescriptorSite mpool.AllocationSite = 25 +) + +func newHashTableAllocation( + t testing.TB, + limit uint64, + metadataSlots uint64, +) (*mpool.AllocationAccountRegistry, *mpool.AllocationAccount, *AllocationAccountSelection) { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, metadataSlots) + require.NoError(t, err) + account, err := registry.Open(limit) + require.NoError(t, err) + selection, err := NewAllocationAccountSelection( + account, + testHashTableOwner, + testHashTableCellSite, + testHashTableDescriptorSite, + ) + require.NoError(t, err) + return registry, account, selection +} + +func completeHashTableAllocation( + t testing.TB, + registry *mpool.AllocationAccountRegistry, + account *mpool.AllocationAccount, +) { + t.Helper() + snapshot, first, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.True(t, first) + require.Equal(t, mpool.AllocationAccountTerminalValid, snapshot.State) + require.Zero(t, snapshot.Used) +} + +func TestHashTableAllocationAccountsCellAndDescriptorStorage(t *testing.T) { + t.Run("int", func(t *testing.T) { + registry, account, selection := newHashTableAllocation(t, 64<<20, 64) + mp := mpool.MustNewZero() + var table Int64HashMap + require.NoError(t, table.InitWithAllocation(mp, selection)) + descriptorBytes := uint64(unsafe.Sizeof([]Int64HashMapCell(nil))) + require.Equal( + t, + Int64HashMapInitialAllocationBytes()+descriptorBytes, + account.Snapshot().Used, + ) + table.Free() + completeHashTableAllocation(t, registry, account) + }) + + t.Run("string", func(t *testing.T) { + registry, account, selection := newHashTableAllocation(t, 64<<20, 64) + mp := mpool.MustNewZero() + var table StringHashMap + require.NoError(t, table.InitWithAllocation(mp, selection)) + descriptorBytes := uint64(unsafe.Sizeof([]StringHashMapCell(nil))) + require.Equal( + t, + StringHashMapInitialAllocationBytes()+descriptorBytes, + account.Snapshot().Used, + ) + table.Free() + completeHashTableAllocation(t, registry, account) + }) +} + +func TestIntHashTableAccountedReplacementPeakAndRollback(t *testing.T) { + const requestedRows = uint64(20_000) + probeMP := mpool.MustNewZero() + var probe Int64HashMap + require.NoError(t, probe.Init(probeMP)) + plan := probe.PlanResize(requestedRows) + require.False(t, plan.Noop) + require.False(t, plan.ReuseCurrentBlocks) + probe.Free() + mpool.DeleteMPool(probeMP) + + descriptorSize := uint64(unsafe.Sizeof([]Int64HashMapCell(nil))) + initialUsed := Int64HashMapInitialAllocationBytes() + descriptorSize + targetDescriptors := plan.TargetBlockCount * descriptorSize + expectedPeak := initialUsed + plan.AdditionalBytes + targetDescriptors + + t.Run("commit", func(t *testing.T) { + registry, account, selection := newHashTableAllocation(t, expectedPeak, 64) + mp := mpool.MustNewZero() + var table Int64HashMap + require.NoError(t, table.InitWithAllocation(mp, selection)) + require.NoError(t, table.ResizeWithPlan(table.PlanResize(requestedRows))) + require.Equal(t, expectedPeak, account.Snapshot().Peak) + require.Equal( + t, + plan.NewBytes+targetDescriptors, + account.Snapshot().Used, + ) + table.Free() + completeHashTableAllocation(t, registry, account) + }) + + t.Run("one byte short", func(t *testing.T) { + registry, account, selection := newHashTableAllocation(t, expectedPeak-1, 64) + mp := mpool.MustNewZero() + var table Int64HashMap + require.NoError(t, table.InitWithAllocation(mp, selection)) + beforeCells := table.cells + before := account.Snapshot() + err := table.ResizeWithPlan(table.PlanResize(requestedRows)) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Same(t, &beforeCells[0][0], &table.cells[0][0]) + require.Equal(t, before.Used, account.Snapshot().Used) + require.Equal(t, uint64(2), registry.LiveAllocationMetadata()) + table.Free() + completeHashTableAllocation(t, registry, account) + }) + + t.Run("metadata rollback", func(t *testing.T) { + // Initial descriptor+cell consume two slots. The replacement descriptor + // consumes the last one; the first replacement cell must reject and the + // complete private replacement rolls back before returning. + registry, account, selection := newHashTableAllocation(t, expectedPeak, 3) + mp := mpool.MustNewZero() + var table Int64HashMap + require.NoError(t, table.InitWithAllocation(mp, selection)) + beforeCells := table.cells + before := account.Snapshot() + err := table.ResizeWithPlan(table.PlanResize(requestedRows)) + require.ErrorIs(t, err, mpool.ErrAllocationMetadataSlots) + require.Same(t, &beforeCells[0][0], &table.cells[0][0]) + require.Equal(t, before.Used, account.Snapshot().Used) + require.Equal(t, uint64(2), registry.LiveAllocationMetadata()) + table.Free() + completeHashTableAllocation(t, registry, account) + }) +} + +func TestIntHashTableAccountedSegmentedGrowth(t *testing.T) { + registry, account, selection := newHashTableAllocation(t, 96<<20, 128) + mp := mpool.MustNewZero() + var table Int64HashMap + require.NoError(t, table.InitWithAllocation(mp, selection)) + + firstTarget := maxElemCnt(maxIntCellCntPerBlock, intCellSize) + require.NoError(t, table.ResizeOnDemand(int(firstTarget))) + plan := table.PlanResize(firstTarget + 1) + require.True(t, plan.ReuseCurrentBlocks) + before := account.Snapshot().Used + descriptorSize := uint64(unsafe.Sizeof([]Int64HashMapCell(nil))) + expectedPeak := before + plan.AdditionalBytes + plan.TargetBlockCount*descriptorSize + require.NoError(t, table.ResizeWithPlan(plan)) + require.Equal(t, expectedPeak, account.Snapshot().Peak) + require.Equal( + t, + plan.NewBytes+plan.TargetBlockCount*descriptorSize, + account.Snapshot().Used, + ) + + table.Free() + completeHashTableAllocation(t, registry, account) +} + +func TestStringHashTableAccountedReplacementPeakAndRollback(t *testing.T) { + const requestedRows = uint64(20_000) + probeMP := mpool.MustNewZero() + var probe StringHashMap + require.NoError(t, probe.Init(probeMP)) + plan := probe.PlanResize(requestedRows) + require.False(t, plan.Noop) + require.False(t, plan.ReuseCurrentBlocks) + probe.Free() + mpool.DeleteMPool(probeMP) + + descriptorSize := uint64(unsafe.Sizeof([]StringHashMapCell(nil))) + initialUsed := StringHashMapInitialAllocationBytes() + descriptorSize + targetDescriptors := plan.TargetBlockCount * descriptorSize + expectedPeak := initialUsed + plan.AdditionalBytes + targetDescriptors + + for _, tc := range []struct { + name string + limit uint64 + ok bool + }{ + {name: "commit", limit: expectedPeak, ok: true}, + {name: "one byte short", limit: expectedPeak - 1}, + } { + t.Run(tc.name, func(t *testing.T) { + registry, account, selection := newHashTableAllocation(t, tc.limit, 64) + mp := mpool.MustNewZero() + var table StringHashMap + require.NoError(t, table.InitWithAllocation(mp, selection)) + beforeCells := table.cells + before := account.Snapshot() + err := table.ResizeWithPlan(table.PlanResize(requestedRows)) + if tc.ok { + require.NoError(t, err) + require.Equal(t, expectedPeak, account.Snapshot().Peak) + require.Equal( + t, + plan.NewBytes+targetDescriptors, + account.Snapshot().Used, + ) + } else { + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Same(t, &beforeCells[0][0], &table.cells[0][0]) + require.Equal(t, before.Used, account.Snapshot().Used) + require.Equal(t, uint64(2), registry.LiveAllocationMetadata()) + } + table.Free() + completeHashTableAllocation(t, registry, account) + }) + } +} + +func TestStringHashTableAccountedSegmentedGrowth(t *testing.T) { + registry, account, selection := newHashTableAllocation(t, 192<<20, 128) + mp := mpool.MustNewZero() + var table StringHashMap + require.NoError(t, table.InitWithAllocation(mp, selection)) + + firstTarget := maxElemCnt(maxStrCellCntPerBlock, strCellSize) + require.NoError(t, table.ResizeOnDemand(firstTarget)) + plan := table.PlanResize(firstTarget + 1) + require.True(t, plan.ReuseCurrentBlocks) + before := account.Snapshot().Used + descriptorSize := uint64(unsafe.Sizeof([]StringHashMapCell(nil))) + expectedPeak := before + plan.AdditionalBytes + plan.TargetBlockCount*descriptorSize + require.NoError(t, table.ResizeWithPlan(plan)) + require.Equal(t, expectedPeak, account.Snapshot().Peak) + require.Equal( + t, + plan.NewBytes+plan.TargetBlockCount*descriptorSize, + account.Snapshot().Used, + ) + + table.Free() + completeHashTableAllocation(t, registry, account) +} + +func TestAccountedHashTableNoopAndStalePlanDoNotChangeCharge(t *testing.T) { + registry, account, selection := newHashTableAllocation(t, 64<<20, 64) + mp := mpool.MustNewZero() + var table Int64HashMap + require.NoError(t, table.InitWithAllocation(mp, selection)) + + before := account.Snapshot() + noop := table.PlanResize(1) + require.True(t, noop.Noop) + require.NoError(t, table.ResizeWithPlan(noop)) + require.Equal(t, before, account.Snapshot()) + + stale := table.PlanResize(20_000) + require.NoError(t, table.ResizeOnDemand(2_000)) + before = account.Snapshot() + err := table.ResizeWithPlan(stale) + require.ErrorIs(t, err, ErrStaleResizePlan) + require.Equal(t, before, account.Snapshot()) + + table.Free() + completeHashTableAllocation(t, registry, account) +} + +func TestHashTableAccountedHighCardinalityResizeReturnsToZero(t *testing.T) { + const rows = 1_000_000 + for _, tc := range []struct { + name string + run func(*mpool.MPool, *AllocationAccountSelection) error + }{ + { + name: "int", + run: func(mp *mpool.MPool, selection *AllocationAccountSelection) error { + var table Int64HashMap + if err := table.InitWithAllocation(mp, selection); err != nil { + return err + } + defer table.Free() + return table.ResizeOnDemand(rows) + }, + }, + { + name: "string", + run: func(mp *mpool.MPool, selection *AllocationAccountSelection) error { + var table StringHashMap + if err := table.InitWithAllocation(mp, selection); err != nil { + return err + } + defer table.Free() + return table.ResizeOnDemand(rows) + }, + }, + } { + t.Run(tc.name, func(t *testing.T) { + registry, account, selection := newHashTableAllocation(t, 256<<20, 256) + mp := mpool.MustNewZero() + require.NoError(t, tc.run(mp, selection)) + require.Positive(t, account.Snapshot().Peak) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + completeHashTableAllocation(t, registry, account) + }) + } +} + +func BenchmarkHashTableResizeAccounting(b *testing.B) { + const rows = 100_000 + b.Run("legacy", func(b *testing.B) { + mp := mpool.MustNewZero() + b.ReportAllocs() + b.ResetTimer() + for range b.N { + var table Int64HashMap + if err := table.Init(mp); err != nil { + b.Fatal(err) + } + if err := table.ResizeOnDemand(rows); err != nil { + b.Fatal(err) + } + table.Free() + } + }) + b.Run("accounted", func(b *testing.B) { + registry, account, selection := newHashTableAllocation(b, 256<<20, 256) + mp := mpool.MustNewZero() + b.ReportAllocs() + b.ResetTimer() + for range b.N { + var table Int64HashMap + if err := table.InitWithAllocation(mp, selection); err != nil { + b.Fatal(err) + } + if err := table.ResizeOnDemand(rows); err != nil { + b.Fatal(err) + } + table.Free() + } + b.StopTimer() + completeHashTableAllocation(b, registry, account) + }) +} diff --git a/pkg/container/hashtable/int64_hash_map.go b/pkg/container/hashtable/int64_hash_map.go index 29b1b28bc319b..d9b791f52a326 100644 --- a/pkg/container/hashtable/int64_hash_map.go +++ b/pkg/container/hashtable/int64_hash_map.go @@ -38,6 +38,7 @@ type Int64HashMap struct { cellCnt uint64 elemCnt uint64 cells [][]Int64HashMapCell + account *AllocationAccountSelection version uint64 admit ResizeAdmission @@ -68,19 +69,34 @@ func (ht *Int64HashMap) cellAt(index uint64) *Int64HashMapCell { func (ht *Int64HashMap) Free() { ht.freeCells(ht.cells) ht.cells = nil + ht.account = nil } func (ht *Int64HashMap) freeCells(cells [][]Int64HashMapCell) { for i, block := range cells { - mpool.FreeSlice(ht.mp, block) + freeHashTableCellSlice(ht.mp, block) cells[i] = nil } + freeHashTableDescriptorSlice(ht.mp, cells, ht.account) } func (ht *Int64HashMap) allocateCells(blockCount int, blockCellCnt uint64) ([][]Int64HashMapCell, error) { - cells := make([][]Int64HashMapCell, blockCount) + cells, err := makeHashTableDescriptorSlice[[]Int64HashMapCell]( + blockCount, + ht.mp, + ht.account, + ht.descriptorSite(), + ) + if err != nil { + return nil, err + } for i := range cells { - block, err := mpool.MakeSlice[Int64HashMapCell](int(blockCellCnt), ht.mp, true) + block, err := makeHashTableCellSlice[Int64HashMapCell]( + int(blockCellCnt), + ht.mp, + ht.account, + ht.cellSite(), + ) if err != nil { ht.freeCells(cells) return nil, err @@ -90,36 +106,83 @@ func (ht *Int64HashMap) allocateCells(blockCount int, blockCellCnt uint64) ([][] return cells, nil } -func (ht *Int64HashMap) allocate(index int, ncells int) error { - if ht.cells[index] != nil { - panic("overwriting") - } - - cell, err := mpool.MakeSlice[Int64HashMapCell](ncells, ht.mp, true) +func (ht *Int64HashMap) appendCells( + blockCount int, + blockCellCnt uint64, +) ([][]Int64HashMapCell, error) { + cells, err := makeHashTableDescriptorSlice[[]Int64HashMapCell]( + blockCount, + ht.mp, + ht.account, + ht.descriptorSite(), + ) if err != nil { - return err + return nil, err } - ht.cells[index] = cell - return nil + copy(cells, ht.cells) + for i := len(ht.cells); i < len(cells); i++ { + block, allocErr := makeHashTableCellSlice[Int64HashMapCell]( + int(blockCellCnt), + ht.mp, + ht.account, + ht.cellSite(), + ) + if allocErr != nil { + for j := len(ht.cells); j < i; j++ { + freeHashTableCellSlice(ht.mp, cells[j]) + cells[j] = nil + } + freeHashTableDescriptorSlice(ht.mp, cells, ht.account) + return nil, allocErr + } + cells[i] = block + } + return cells, nil } func (ht *Int64HashMap) Init(mp *mpool.MPool) (err error) { + return ht.InitWithAllocation(mp, nil) +} + +func (ht *Int64HashMap) InitWithAllocation( + mp *mpool.MPool, + account *AllocationAccountSelection, +) (err error) { + if account != nil { + if err = account.validate(); err != nil { + return err + } + } ht.mp = mp + ht.account = account ht.blockCellCntBits = kInitialCellCntBits ht.cellCntMask = kInitialCellCnt - 1 ht.elemCnt = 0 ht.cellCnt = kInitialCellCnt ht.version = 0 - ht.cells = make([][]Int64HashMapCell, 1) - - if err = ht.allocate(0, int(ht.blockCellCnt())); err != nil { + if ht.cells, err = ht.allocateCells(1, ht.blockCellCnt()); err != nil { + ht.account = nil return err } return } +func (ht *Int64HashMap) cellSite() mpool.AllocationSite { + if ht.account == nil { + return 0 + } + return ht.account.cellSite +} + +func (ht *Int64HashMap) descriptorSite() mpool.AllocationSite { + if ht.account == nil { + return 0 + } + return ht.account.descriptorSite +} + func (ht *Int64HashMap) InsertBatch(n int, hashes []uint64, keysPtr unsafe.Pointer, values []uint64) error { if n <= 0 { return nil @@ -279,15 +342,17 @@ func (ht *Int64HashMap) ResizeWithPlan(plan ResizePlan) error { }() if plan.ReuseCurrentBlocks { - newBlocks, err := ht.allocateCells( - int(plan.TargetBlockCount-plan.CurrentBlockCount), + newCells, err := ht.appendCells( + int(plan.TargetBlockCount), plan.TargetBlockCellCount, ) if err != nil { return err } oldCellCnt := ht.cellCnt - ht.cells = append(ht.cells, newBlocks...) + oldDescriptors := ht.cells + ht.cells = newCells + freeHashTableDescriptorSlice(ht.mp, oldDescriptors, ht.account) ht.cellCnt = plan.TargetCellCount ht.cellCntMask = ht.cellCnt - 1 ht.version++ @@ -344,8 +409,11 @@ func (ht *Int64HashMap) Size() int64 { ret := int64(41) for i := range ht.cells { ret += int64(len(ht.cells[i]) * int(intCellSize)) - // 16 is the len of ht.cells[i] - ret += 16 + } + if ht.account != nil { + ret += int64(len(ht.cells)) * int64(unsafe.Sizeof([]Int64HashMapCell(nil))) + } else { + ret += int64(len(ht.cells)) * 16 } return ret } diff --git a/pkg/container/hashtable/string_hash_map.go b/pkg/container/hashtable/string_hash_map.go index db32e3b582a43..464b2d05f3134 100644 --- a/pkg/container/hashtable/string_hash_map.go +++ b/pkg/container/hashtable/string_hash_map.go @@ -45,6 +45,7 @@ type StringHashMap struct { cellCnt uint64 elemCnt uint64 cells [][]StringHashMapCell + account *AllocationAccountSelection version uint64 admit ResizeAdmission @@ -75,19 +76,34 @@ func (ht *StringHashMap) cellAt(index uint64) *StringHashMapCell { func (ht *StringHashMap) Free() { ht.freeCells(ht.cells) ht.cells = nil + ht.account = nil } func (ht *StringHashMap) freeCells(cells [][]StringHashMapCell) { for i, block := range cells { - mpool.FreeSlice(ht.mp, block) + freeHashTableCellSlice(ht.mp, block) cells[i] = nil } + freeHashTableDescriptorSlice(ht.mp, cells, ht.account) } func (ht *StringHashMap) allocateCells(blockCount int, blockCellCnt uint64) ([][]StringHashMapCell, error) { - cells := make([][]StringHashMapCell, blockCount) + cells, err := makeHashTableDescriptorSlice[[]StringHashMapCell]( + blockCount, + ht.mp, + ht.account, + ht.descriptorSite(), + ) + if err != nil { + return nil, err + } for i := range cells { - block, err := mpool.MakeSlice[StringHashMapCell](int(blockCellCnt), ht.mp, true) + block, err := makeHashTableCellSlice[StringHashMapCell]( + int(blockCellCnt), + ht.mp, + ht.account, + ht.cellSite(), + ) if err != nil { ht.freeCells(cells) return nil, err @@ -97,34 +113,83 @@ func (ht *StringHashMap) allocateCells(blockCount int, blockCellCnt uint64) ([][ return cells, nil } -func (ht *StringHashMap) allocate(index int, ncells int) error { - if ht.cells[index] != nil { - panic("overwriting") - } - c, err := mpool.MakeSlice[StringHashMapCell](ncells, ht.mp, true) +func (ht *StringHashMap) appendCells( + blockCount int, + blockCellCnt uint64, +) ([][]StringHashMapCell, error) { + cells, err := makeHashTableDescriptorSlice[[]StringHashMapCell]( + blockCount, + ht.mp, + ht.account, + ht.descriptorSite(), + ) if err != nil { - return err + return nil, err } - ht.cells[index] = c - return nil + copy(cells, ht.cells) + for i := len(ht.cells); i < len(cells); i++ { + block, allocErr := makeHashTableCellSlice[StringHashMapCell]( + int(blockCellCnt), + ht.mp, + ht.account, + ht.cellSite(), + ) + if allocErr != nil { + for j := len(ht.cells); j < i; j++ { + freeHashTableCellSlice(ht.mp, cells[j]) + cells[j] = nil + } + freeHashTableDescriptorSlice(ht.mp, cells, ht.account) + return nil, allocErr + } + cells[i] = block + } + return cells, nil } func (ht *StringHashMap) Init(mp *mpool.MPool) (err error) { + return ht.InitWithAllocation(mp, nil) +} + +func (ht *StringHashMap) InitWithAllocation( + mp *mpool.MPool, + account *AllocationAccountSelection, +) (err error) { + if account != nil { + if err = account.validate(); err != nil { + return err + } + } ht.mp = mp + ht.account = account ht.blockCellCntBits = kInitialCellCntBits ht.elemCnt = 0 ht.cellCnt = kInitialCellCnt ht.version = 0 ht.cellCntMask = kInitialCellCnt - 1 - ht.cells = make([][]StringHashMapCell, 1) - if err := ht.allocate(0, int(ht.blockCellCnt())); err != nil { + if ht.cells, err = ht.allocateCells(1, ht.blockCellCnt()); err != nil { + ht.account = nil return err } return } +func (ht *StringHashMap) cellSite() mpool.AllocationSite { + if ht.account == nil { + return 0 + } + return ht.account.cellSite +} + +func (ht *StringHashMap) descriptorSite() mpool.AllocationSite { + if ht.account == nil { + return 0 + } + return ht.account.descriptorSite +} + func (ht *StringHashMap) InsertStringBatch(states [][3]uint64, keys [][]byte, values []uint64) error { if len(keys) == 0 { return nil @@ -278,15 +343,17 @@ func (ht *StringHashMap) ResizeWithPlan(plan ResizePlan) error { }() if plan.ReuseCurrentBlocks { - newBlocks, err := ht.allocateCells( - int(plan.TargetBlockCount-plan.CurrentBlockCount), + newCells, err := ht.appendCells( + int(plan.TargetBlockCount), plan.TargetBlockCellCount, ) if err != nil { return err } oldCellCnt := ht.cellCnt - ht.cells = append(ht.cells, newBlocks...) + oldDescriptors := ht.cells + ht.cells = newCells + freeHashTableDescriptorSlice(ht.mp, oldDescriptors, ht.account) ht.cellCnt = plan.TargetCellCount ht.cellCntMask = ht.cellCnt - 1 ht.version++ @@ -340,6 +407,9 @@ func (ht *StringHashMap) Size() int64 { for i := range ht.cells { ret += int64(int(strCellSize) * len(ht.cells[i])) } + if ht.account != nil { + ret += int64(len(ht.cells)) * int64(unsafe.Sizeof([]StringHashMapCell(nil))) + } return ret } diff --git a/pkg/sql/colexec/hashbuild/budget.go b/pkg/sql/colexec/hashbuild/budget.go index 8cdf08174cae3..481a50fea627f 100644 --- a/pkg/sql/colexec/hashbuild/budget.go +++ b/pkg/sql/colexec/hashbuild/budget.go @@ -174,6 +174,52 @@ func NewBudgetedEmptyJoinMap( return jm, nil } +// NewAccountedEmptyJoinMap creates the consumer-grown empty-map variant under +// the statement allocation generation. Physical cell/descriptor Free is its +// only memory release owner; the account's controller already charges the +// shared query/CN policy, so no legacy reservation is stacked on it. +func NewAccountedEmptyJoinMap( + keyWidth int, + account *mpool.AllocationAccount, + mp *mpool.MPool, +) (*message.JoinMap, error) { + if account == nil || mp == nil { + return nil, mpool.ErrAllocationAccountInvalid + } + selection, err := hashtable.NewAllocationAccountSelection( + account, + HashBuildAllocationOwner, + HashBuildAllocationSiteHashCell, + HashBuildAllocationSiteHashDescriptor, + ) + if err != nil { + return nil, err + } + var ( + intHashMap *hashmap.IntHashMap + strHashMap *hashmap.StrHashMap + ) + if keyWidth <= 8 { + intHashMap, err = hashmap.NewIntHashMapWithAllocation(false, mp, selection) + } else { + strHashMap, err = hashmap.NewStrHashMapWithAllocation(false, mp, selection) + } + if err != nil { + return nil, err + } + + jm := message.NewJoinMap( + message.GroupSels{}, + intHashMap, + strHashMap, + nil, + nil, + mp, + ) + jm.IncRef(1) + return jm, nil +} + func (hb *HashmapBuilder) attachIntHashMapAdmission(m *hashmap.IntHashMap) error { owner := hb.mapReservation budget := hb.budget diff --git a/pkg/sql/colexec/hashbuild/hashmap.go b/pkg/sql/colexec/hashbuild/hashmap.go index c356136a04a04..b9248a014f3d6 100644 --- a/pkg/sql/colexec/hashbuild/hashmap.go +++ b/pkg/sql/colexec/hashbuild/hashmap.go @@ -79,7 +79,6 @@ type HashmapBuilder struct { auxReservation *process.HashBuildReservation keyExprs []*plan.Expr expressionLease *ExpressionMemoryLease - // Exact runtime-filter keys are an optional owner inside the mandatory // JoinMap build. The fallback bit is observed by HashBuild for diagnostics. // @@ -90,6 +89,8 @@ type HashmapBuilder struct { // must not be retried or re-spilled. runtimeFilterCollectionFallback bool retainedBatchRecoverySafe bool + mapAllocationAccount *mpool.AllocationAccount + mapAllocation *hashtable.AllocationAccountSelection } func (hb *HashmapBuilder) GetSize() int64 { @@ -242,6 +243,8 @@ func (hb *HashmapBuilder) Reset(proc *process.Process, hashTableHasNotSent bool) // Free them before releasing expression reservations; Prepare recreates the // executor set for the next generation. hb.FreeExecutors() + hb.mapAllocationAccount = nil + hb.mapAllocation = nil } func (hb *HashmapBuilder) Free(proc *process.Process) { @@ -263,6 +266,8 @@ func (hb *HashmapBuilder) Free(proc *process.Process) { } hb.UniqueJoinKeys = nil hb.uniqueKeySlots = nil + hb.mapAllocationAccount = nil + hb.mapAllocation = nil } func (hb *HashmapBuilder) FreeExecutors() { @@ -649,18 +654,29 @@ func (hb *HashmapBuilder) buildHashmap( var err error var itr hashmap.Iterator if hb.keyWidth <= 8 { - if err = hb.reserveInitialMap(int64(hashtable.Int64HashMapInitialAllocationBytes())); err != nil { - return err + if hb.mapAllocation == nil { + if err = hb.reserveInitialMap(int64(hashtable.Int64HashMapInitialAllocationBytes())); err != nil { + return err + } + hb.IntHashMap, err = hashmap.NewIntHashMap(false, proc.Mp()) + if err == nil { + err = hb.attachIntHashMapAdmission(hb.IntHashMap) + } + } else { + hb.IntHashMap, err = hashmap.NewIntHashMapWithAllocation( + false, + proc.Mp(), + hb.mapAllocation, + ) } - if hb.IntHashMap, err = hashmap.NewIntHashMap(false, proc.Mp()); err != nil { + if err != nil { + if hb.IntHashMap != nil { + hb.IntHashMap.Free() + hb.IntHashMap = nil + } hb.releaseMapReservation() return err } - if err = hb.attachIntHashMapAdmission(hb.IntHashMap); err != nil { - hb.IntHashMap.Free() - hb.IntHashMap = nil - return err - } if hb.cachedIntIterator != nil { hashmap.IteratorChangeOwner(hb.cachedIntIterator, hb.IntHashMap) itr = hb.cachedIntIterator @@ -669,18 +685,29 @@ func (hb *HashmapBuilder) buildHashmap( hb.cachedIntIterator = itr } } else { - if err = hb.reserveInitialMap(int64(hashtable.StringHashMapInitialAllocationBytes())); err != nil { - return err + if hb.mapAllocation == nil { + if err = hb.reserveInitialMap(int64(hashtable.StringHashMapInitialAllocationBytes())); err != nil { + return err + } + hb.StrHashMap, err = hashmap.NewStrHashMap(false, proc.Mp()) + if err == nil { + err = hb.attachStrHashMapAdmission(hb.StrHashMap) + } + } else { + hb.StrHashMap, err = hashmap.NewStrHashMapWithAllocation( + false, + proc.Mp(), + hb.mapAllocation, + ) } - if hb.StrHashMap, err = hashmap.NewStrHashMap(false, proc.Mp()); err != nil { + if err != nil { + if hb.StrHashMap != nil { + hb.StrHashMap.Free() + hb.StrHashMap = nil + } hb.releaseMapReservation() return err } - if err = hb.attachStrHashMapAdmission(hb.StrHashMap); err != nil { - hb.StrHashMap.Free() - hb.StrHashMap = nil - return err - } if hb.cachedStrIterator != nil { hashmap.IteratorChangeOwner(hb.cachedStrIterator, hb.StrHashMap) itr = hb.cachedStrIterator diff --git a/pkg/sql/colexec/hashbuild/hashmap_test.go b/pkg/sql/colexec/hashbuild/hashmap_test.go index 4b0a3ccfa3f42..716182b73bc13 100644 --- a/pkg/sql/colexec/hashbuild/hashmap_test.go +++ b/pkg/sql/colexec/hashbuild/hashmap_test.go @@ -343,6 +343,67 @@ func TestPublishedJoinMapResizeKeepsReservationWithConsumer(t *testing.T) { hb.Reset(proc, false) } +func TestHashmapBuilderAccountedCellsDoNotStackLegacyMapReservation(t *testing.T) { + const budgetCap = uint64(16 << 20) + budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) + require.NoError(t, err) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.OpenWithController(budgetCap, generation) + require.NoError(t, err) + + var op HashBuild + op.NeedHashMap = true + require.NoError(t, op.SetAllocationAccount(account)) + hb := &op.ctr.hashmapBuilder + hb.setBudget(generation) + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + require.NoError(t, hb.Prepare( + []*plan.Expr{newExpr(0, types.T_int32.ToType())}, + -1, + -1, + nil, + proc, + )) + input := testutil.NewBatch( + []types.Type{types.T_int32.ToType()}, + true, + 10_000, + proc.Mp(), + ) + require.NoError(t, hb.copyBuildBatch(input, proc)) + hb.InputBatchRowCount = input.RowCount() + input.Clean(proc.Mp()) + + require.NoError(t, hb.BuildHashmap(false, false, false, proc)) + require.Nil(t, hb.mapReservation) + require.Positive(t, account.Snapshot().Used) + require.Equal( + t, + account.Snapshot().Used, + generation.Snapshot().AllocationUsed, + ) + + jm := hb.GetJoinMap(proc.Mp()) + require.NotNil(t, jm) + jm.IncRef(1) + beforeResize := account.Snapshot().Used + require.NoError(t, jm.PreAlloc(100_000)) + require.Greater(t, account.Snapshot().Used, beforeResize) + jm.Free() + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + + hb.Reset(proc, false) + terminal, first, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.True(t, first) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) +} + func TestHashMapReservationOwnerRetainsSegmentedGrowthTokens(t *testing.T) { budget, err := process.NewHashBuildBudget(1<<20, 1<<20) require.NoError(t, err) @@ -412,6 +473,75 @@ func TestBudgetedEmptyJoinMapRejectsUnadmittedAllocationAndResize(t *testing.T) } } +func TestAccountedEmptyJoinMapUsesPhysicalAllocationAsSoleCharge(t *testing.T) { + for _, tc := range []struct { + name string + keyWidth int + initialBytes uint64 + }{ + {name: "int", keyWidth: 4, initialBytes: hashtable.Int64HashMapInitialAllocationBytes()}, + {name: "string", keyWidth: 128, initialBytes: hashtable.StringHashMapInitialAllocationBytes()}, + } { + t.Run(tc.name, func(t *testing.T) { + const capBytes = uint64(64 << 20) + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.OpenWithController(capBytes, generation) + require.NoError(t, err) + mp := mpool.MustNewZero() + + jm, err := NewAccountedEmptyJoinMap(tc.keyWidth, account, mp) + require.NoError(t, err) + descriptorBytes := hashtable.HashMapBlockDescriptorBytes() + expectedInitial := tc.initialBytes + descriptorBytes + require.Equal(t, expectedInitial, account.Snapshot().Used) + require.Equal(t, expectedInitial, generation.Used()) + require.Equal(t, expectedInitial, generation.Snapshot().AllocationUsed) + + require.NoError(t, jm.PreAlloc(10_000)) + require.Equal(t, account.Snapshot().Used, generation.Used()) + require.Equal( + t, + account.Snapshot().Used, + generation.Snapshot().AllocationUsed, + ) + jm.Free() + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + terminal, first, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.True(t, first) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) + }) + } +} + +func TestAccountedEmptyJoinMapInitialFailureRollsBackController(t *testing.T) { + initial := hashtable.Int64HashMapInitialAllocationBytes() + + hashtable.HashMapBlockDescriptorBytes() + budget := process.MustNewHashBuildBudget(initial, initial) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 2) + require.NoError(t, err) + account, err := registry.OpenWithController(initial-1, generation) + require.NoError(t, err) + mp := mpool.MustNewZero() + + jm, err := NewAccountedEmptyJoinMap(4, account, mp) + require.Nil(t, jm) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.Zero(t, registry.LiveAllocationMetadata()) + require.Zero(t, mp.CurrNB()) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + func TestCopyBuildBatchBudgetsSmallIngressAfterFullBatches(t *testing.T) { const budgetCap = uint64(32 << 20) budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) diff --git a/pkg/sql/colexec/hashbuild/types.go b/pkg/sql/colexec/hashbuild/types.go index 60f2523e9ddfe..8485750b70191 100644 --- a/pkg/sql/colexec/hashbuild/types.go +++ b/pkg/sql/colexec/hashbuild/types.go @@ -21,8 +21,10 @@ import ( "sync/atomic" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/reuse" "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/hashtable" "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/fileservice" "github.com/matrixorigin/matrixone/pkg/logutil" @@ -46,6 +48,13 @@ const ( SendSucceed ) +const HashBuildAllocationOwner mpool.AllocationOwner = 1 + +const ( + HashBuildAllocationSiteHashCell mpool.AllocationSite = iota + 24 + HashBuildAllocationSiteHashDescriptor +) + type container struct { state int runtimeFilterIn bool @@ -267,6 +276,55 @@ func (hashBuild *HashBuild) GetOperatorBase() *vm.OperatorBase { return &hashBuild.OperatorBase } +func (hashBuild *HashBuild) AllocationAccountEnabled() bool { + return hashBuild != nil && hashBuild.NeedHashMap +} + +// SetAllocationAccount selects immutable provenance for the hash-table owner +// before Prepare. Compile invokes it once for each execution attempt; Reset +// clears the selection only after producer or JoinMap ownership has moved on. +func (hashBuild *HashBuild) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + builder := &hashBuild.ctr.hashmapBuilder + if builder.mapAllocationAccount != nil { + if builder.mapAllocationAccount == account { + return nil + } + return mpool.ErrAllocationAccountMismatch + } + selection, err := hashtable.NewAllocationAccountSelection( + account, + HashBuildAllocationOwner, + HashBuildAllocationSiteHashCell, + HashBuildAllocationSiteHashDescriptor, + ) + if err != nil { + return err + } + builder.mapAllocationAccount = account + builder.mapAllocation = selection + return nil +} + +func (hashBuild *HashBuild) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + builder := &hashBuild.ctr.hashmapBuilder + if builder.mapAllocationAccount == nil { + return nil + } + if builder.mapAllocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if builder.IntHashMap != nil || builder.StrHashMap != nil { + return mpool.ErrAllocationAccountInvariant + } + builder.mapAllocationAccount = nil + builder.mapAllocation = nil + return nil +} + func init() { reuse.CreatePool[HashBuild]( func() *HashBuild { diff --git a/pkg/sql/colexec/rightdedupjoin/join.go b/pkg/sql/colexec/rightdedupjoin/join.go index 1289d9c6dee4b..52df5361a3527 100644 --- a/pkg/sql/colexec/rightdedupjoin/join.go +++ b/pkg/sql/colexec/rightdedupjoin/join.go @@ -290,6 +290,13 @@ func (rightDedupJoin *RightDedupJoin) newEmptyJoinMap(proc *process.Process) (*m if err != nil { return nil, err } + if rightDedupJoin.allocationAccount != nil { + return hashbuild.NewAccountedEmptyJoinMap( + keyWidth, + rightDedupJoin.allocationAccount, + proc.Mp(), + ) + } return hashbuild.NewBudgetedEmptyJoinMap(keyWidth, budget, proc.Mp()) } diff --git a/pkg/sql/colexec/rightdedupjoin/types.go b/pkg/sql/colexec/rightdedupjoin/types.go index e5f7d820657a7..5489c1e98e602 100644 --- a/pkg/sql/colexec/rightdedupjoin/types.go +++ b/pkg/sql/colexec/rightdedupjoin/types.go @@ -17,6 +17,7 @@ package rightdedupjoin import ( "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/hashmap" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/reuse" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" @@ -90,10 +91,45 @@ type RightDedupJoin struct { DelColIdx int32 UpdateColIdxList []int32 UpdateColExprList []*plan.Expr + allocationAccount *mpool.AllocationAccount vm.OperatorBase } +func (rightDedupJoin *RightDedupJoin) AllocationAccountEnabled() bool { + return rightDedupJoin != nil +} + +func (rightDedupJoin *RightDedupJoin) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if account == nil || account.Handle() == 0 { + return mpool.ErrAllocationAccountInvalid + } + if rightDedupJoin.allocationAccount != nil && + rightDedupJoin.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + rightDedupJoin.allocationAccount = account + return nil +} + +func (rightDedupJoin *RightDedupJoin) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if rightDedupJoin.allocationAccount == nil { + return nil + } + if rightDedupJoin.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if rightDedupJoin.ctr.mp != nil { + return mpool.ErrAllocationAccountInvariant + } + rightDedupJoin.allocationAccount = nil + return nil +} + func (rightDedupJoin *RightDedupJoin) GetOperatorBase() *vm.OperatorBase { return &rightDedupJoin.OperatorBase } @@ -150,6 +186,7 @@ func (rightDedupJoin *RightDedupJoin) Reset(proc *process.Process, pipelineFaile ctr.resetEvalVectors() } ctr.state = Build + rightDedupJoin.allocationAccount = nil } func (rightDedupJoin *RightDedupJoin) Free(proc *process.Process, pipelineFailed bool, err error) { @@ -164,6 +201,7 @@ func (rightDedupJoin *RightDedupJoin) Free(proc *process.Process, pipelineFailed } ctr.cleanEvalVectors() ctr.releaseProbeExpressionLease() + rightDedupJoin.allocationAccount = nil } func (rightDedupJoin *RightDedupJoin) ExecProjection(proc *process.Process, input *batch.Batch) (*batch.Batch, error) { diff --git a/pkg/sql/compile/allocation_account_lifecycle.go b/pkg/sql/compile/allocation_account_lifecycle.go index 5feb71e218535..0d797ca709a5c 100644 --- a/pkg/sql/compile/allocation_account_lifecycle.go +++ b/pkg/sql/compile/allocation_account_lifecycle.go @@ -15,12 +15,21 @@ package compile import ( + "errors" "sync" "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/vm" "github.com/matrixorigin/matrixone/pkg/vm/message" + "github.com/matrixorigin/matrixone/pkg/vm/process" ) +type executionAllocationAccountOwner interface { + AllocationAccountEnabled() bool + SetAllocationAccount(*mpool.AllocationAccount) error + ClearAllocationAccount(*mpool.AllocationAccount) error +} + // statementAllocationAttempt owns one local execution generation. The // MessageBoard pointer is captured at open so prepared/retry Reset cannot make // terminal cleanup drain a newer board. @@ -29,6 +38,7 @@ type statementAllocationAttempt struct { account *mpool.AllocationAccount board *message.MessageBoard exporter func(mpool.AllocationAccountTerminalSnapshot) + owners []executionAllocationAccountOwner once sync.Once snapshot mpool.AllocationAccountTerminalSnapshot @@ -46,20 +56,194 @@ func (c *Compile) beginAllocationAccountAttempt() ( c.allocationTerminalExporter == nil { return nil, mpool.ErrAllocationAccountInvariant } - account, err := c.allocationAccountRegistry.Open(c.allocationAccountLimit) + var controller mpool.AllocationCapacityController + var err error + if c.allocationControllerProvider != nil { + controller, err = c.allocationControllerProvider() + if err != nil { + return nil, err + } + if controller == nil { + return nil, mpool.ErrAllocationAccountInvariant + } + } + account, err := c.allocationAccountRegistry.OpenWithController( + c.allocationAccountLimit, + controller, + ) if err != nil { return nil, err } + owners, err := configureAllocationAccountOwners(c.scopes, account) + if err != nil { + snapshot, first, finalizeErr := c.allocationAccountRegistry. + CompleteTerminalWithError(account, err) + if first { + c.allocationTerminalExporter(snapshot) + } + if finalizeErr != nil { + return nil, finalizeErr + } + return nil, err + } attempt := &statementAllocationAttempt{ registry: c.allocationAccountRegistry, account: account, board: c.MessageBoard, exporter: c.allocationTerminalExporter, + owners: owners, } c.allocationAttempt = attempt return attempt, nil } +func configureAllocationAccountOwners( + scopes []*Scope, + account *mpool.AllocationAccount, +) ([]executionAllocationAccountOwner, error) { + var configured []executionAllocationAccountOwner + isConfigured := func(candidate executionAllocationAccountOwner) bool { + for _, owner := range configured { + if owner == candidate { + return true + } + } + return false + } + rollback := func(cause error) error { + for i := len(configured) - 1; i >= 0; i-- { + cause = errors.Join( + cause, + configured[i].ClearAllocationAccount(account), + ) + } + return cause + } + var configure func(*Scope) error + configure = func(scope *Scope) error { + if scope == nil { + return nil + } + if err := vm.HandleAllOp( + scope.RootOp, + func(_ vm.Operator, op vm.Operator) error { + if owner, ok := op.(executionAllocationAccountOwner); ok && + owner.AllocationAccountEnabled() { + if isConfigured(owner) { + return nil + } + if err := owner.SetAllocationAccount(account); err != nil { + return err + } + configured = append(configured, owner) + } + return nil + }, + ); err != nil { + return err + } + for _, preScope := range scope.PreScopes { + if err := configure(preScope); err != nil { + return err + } + } + return nil + } + for _, scope := range scopes { + if err := configure(scope); err != nil { + return nil, rollback(err) + } + } + return configured, nil +} + +func hasAllocationAccountOwner(scopes []*Scope) bool { + var inspect func(*Scope) bool + inspect = func(scope *Scope) bool { + if scope == nil { + return false + } + found := false + _ = vm.HandleAllOp(scope.RootOp, func(_ vm.Operator, op vm.Operator) error { + if owner, ok := op.(executionAllocationAccountOwner); ok && + owner.AllocationAccountEnabled() { + found = true + } + return nil + }) + if found { + return true + } + for _, preScope := range scope.PreScopes { + if inspect(preScope) { + return true + } + } + return false + } + for _, scope := range scopes { + if inspect(scope) { + return true + } + } + return false +} + +// ensureAllocationAccountLifecycle activates accounting only when the physical +// plan contains a complete migrated owner. Legacy plans never open a registry +// slot or initialize the HashBuild budget. +func (c *Compile) ensureAllocationAccountLifecycle( + exporter func(mpool.AllocationAccountTerminalSnapshot), +) error { + if c == nil { + return nil + } + if !hasAllocationAccountOwner(c.scopes) { + if c.allocationLifecycleAutomatic { + c.allocationAccountRegistry = nil + c.allocationAccountLimit = 0 + c.allocationControllerProvider = nil + c.allocationTerminalExporter = nil + c.allocationLifecycleAutomatic = false + } + return nil + } + if c.allocationAccountRegistry != nil && !c.allocationLifecycleAutomatic { + return nil + } + if exporter == nil { + return mpool.ErrAllocationAccountInvariant + } + if c.proc == nil { + return mpool.ErrAllocationAccountInvariant + } + budget, err := c.proc.GetHashBuildBudget() + if err != nil { + return err + } + registry, err := budget.AllocationAccountRegistry() + if err != nil { + return err + } + limit := budget.Snapshot().Cap + if limit == 0 { + return mpool.ErrAllocationAccountInvariant + } + c.ConfigureAllocationAccountLifecycleWithController( + registry, + limit, + func() (mpool.AllocationCapacityController, error) { + if budget.Closed() { + return nil, process.ErrHashBuildBudgetClosed + } + return budget, nil + }, + exporter, + ) + c.allocationLifecycleAutomatic = true + return nil +} + func (a *statementAllocationAttempt) finish() ( mpool.AllocationAccountTerminalSnapshot, error, @@ -72,8 +256,20 @@ func (a *statementAllocationAttempt) finish() ( // before this point. Draining the board first releases queued JoinMap // and spill payload ownership through their normal Destroy methods. a.board.CloseAndDrain() + for i := len(a.owners) - 1; i >= 0; i-- { + a.err = errors.Join( + a.err, + a.owners[i].ClearAllocationAccount(a.account), + ) + } + a.owners = nil var first bool - a.snapshot, first, a.err = a.registry.CompleteTerminal(a.account) + var terminalErr error + a.snapshot, first, terminalErr = a.registry.CompleteTerminalWithError( + a.account, + a.err, + ) + a.err = terminalErr if first && a.exporter != nil { a.exporter(a.snapshot) } @@ -100,4 +296,6 @@ func (c *Compile) copyAllocationAccountLifecycleTo(dst *Compile) { c.allocationAccountLimit, c.allocationTerminalExporter, ) + dst.allocationControllerProvider = c.allocationControllerProvider + dst.allocationLifecycleAutomatic = c.allocationLifecycleAutomatic } diff --git a/pkg/sql/compile/allocation_account_lifecycle_test.go b/pkg/sql/compile/allocation_account_lifecycle_test.go index 12add81745975..883d9737bddf1 100644 --- a/pkg/sql/compile/allocation_account_lifecycle_test.go +++ b/pkg/sql/compile/allocation_account_lifecycle_test.go @@ -27,6 +27,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/pb/txn" "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" "github.com/matrixorigin/matrixone/pkg/testutil" "github.com/matrixorigin/matrixone/pkg/vm" "github.com/matrixorigin/matrixone/pkg/vm/message" @@ -39,6 +40,48 @@ type allocationLifecycleErrorOperator struct { err error } +type allocationLifecycleOwnerOperator struct { + *colexec.MockOperator + account *mpool.AllocationAccount + failSet bool + failClear bool + clears int +} + +func (op *allocationLifecycleOwnerOperator) AllocationAccountEnabled() bool { + return true +} + +func (op *allocationLifecycleOwnerOperator) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if op.failSet { + return mpool.ErrAllocationAccountMismatch + } + if op.account != nil && op.account != account { + return mpool.ErrAllocationAccountMismatch + } + op.account = account + return nil +} + +func (op *allocationLifecycleOwnerOperator) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if op.account == nil { + return nil + } + if op.failClear { + return mpool.ErrAllocationAccountInvariant + } + if op.account != account { + return mpool.ErrAllocationAccountMismatch + } + op.account = nil + op.clears++ + return nil +} + func (op *allocationLifecycleErrorOperator) Call( *process.Process, ) (vm.CallResult, error) { @@ -250,6 +293,102 @@ func TestStatementAllocationAttemptRequiresTerminalExporter(t *testing.T) { require.NoError(t, err) } +func TestStatementAllocationAttemptOwnerConfigurationRollsBack(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + var exported []mpool.AllocationAccountTerminalSnapshot + c := newTestAllocationLifecycleCompile(t, registry, func( + snapshot mpool.AllocationAccountTerminalSnapshot, + ) { + exported = append(exported, snapshot) + }) + configured := &allocationLifecycleOwnerOperator{ + MockOperator: colexec.NewMockOperator(), + } + rejected := &allocationLifecycleOwnerOperator{ + MockOperator: colexec.NewMockOperator(), + failSet: true, + } + c.scopes = []*Scope{ + {RootOp: configured}, + {RootOp: rejected}, + } + + _, err = c.beginAllocationAccountAttempt() + require.ErrorIs(t, err, mpool.ErrAllocationAccountMismatch) + require.Nil(t, configured.account) + require.Equal(t, 1, configured.clears) + require.Nil(t, c.allocationAttempt) + require.Len(t, exported, 1) + require.Equal( + t, + mpool.AllocationAccountTerminalInvariantFailure, + exported[0].State, + ) + + account, err := registry.Open(1) + require.NoError(t, err, "failed owner configuration leaked its registry slot") + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestStatementAllocationAttemptOwnerTeardownFailureExportsFailure(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + var exported []mpool.AllocationAccountTerminalSnapshot + c := newTestAllocationLifecycleCompile(t, registry, func( + snapshot mpool.AllocationAccountTerminalSnapshot, + ) { + exported = append(exported, snapshot) + }) + owner := &allocationLifecycleOwnerOperator{ + MockOperator: colexec.NewMockOperator(), + failClear: true, + } + c.scopes = []*Scope{{RootOp: owner}} + + attempt, err := c.beginAllocationAccountAttempt() + require.NoError(t, err) + err = c.finishAllocationAccountAttempt() + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + require.Len(t, exported, 1) + require.Equal( + t, + mpool.AllocationAccountTerminalInvariantFailure, + exported[0].State, + ) + require.Zero(t, exported[0].Used) + require.False(t, registry.AdmissionSuspended()) + _, ok := registry.Resolve(attempt.account.Handle()) + require.False(t, ok) +} + +func TestCompileAutomaticallyActivatesCompleteHashTableOwner(t *testing.T) { + proc := testutil.NewProcess(t) + c := &Compile{ + proc: proc, + MessageBoard: message.NewMessageBoard(), + } + owner := hashbuild.NewArgument() + owner.NeedHashMap = true + c.scopes = []*Scope{{RootOp: owner}} + + require.NoError(t, c.ensureAllocationAccountLifecycle(func( + mpool.AllocationAccountTerminalSnapshot, + ) { + })) + require.True(t, c.allocationLifecycleAutomatic) + require.NotNil(t, c.allocationAccountRegistry) + attempt, err := c.beginAllocationAccountAttempt() + require.NoError(t, err) + require.NotNil(t, attempt) + require.NoError(t, owner.ClearAllocationAccount(attempt.account)) + require.NoError(t, c.finishAllocationAccountAttempt()) + _, ok := c.allocationAccountRegistry.Resolve(attempt.account.Handle()) + require.False(t, ok) + owner.Release() +} + func TestCompileRunFinalizesAllocationAttemptOnCancellation(t *testing.T) { var snapshots []mpool.AllocationAccountTerminalSnapshot c, registry := newRunLifecycleCompile(t, func( diff --git a/pkg/sql/compile/analyze_module.go b/pkg/sql/compile/analyze_module.go index 4009c26681899..f7f39436eea30 100644 --- a/pkg/sql/compile/analyze_module.go +++ b/pkg/sql/compile/analyze_module.go @@ -49,6 +49,7 @@ type AnalyzeModule struct { explainPhyBuffer *bytes.Buffer remoteUsage resource.Usage remoteMemory resource.MemoryTotals + remoteAllocation resource.AllocationAccountTotals remoteQuality resource.QualityFlags remoteMissingFragments uint64 remoteMissingMemoryDomains uint64 @@ -62,6 +63,7 @@ type AnalyzeModule struct { type remoteResourceSnapshot struct { Usage resource.Usage Memory resource.MemoryTotals + Allocation resource.AllocationAccountTotals Quality resource.QualityFlags MissingFragmentCount uint64 MissingMemoryDomainCount uint64 @@ -82,6 +84,7 @@ func (anal *AnalyzeModule) Reset(isPrepare bool, isTpQuery bool) { anal.retryTimes = 0 anal.remoteUsage = resource.Usage{} anal.remoteMemory = resource.MemoryTotals{} + anal.remoteAllocation = resource.AllocationAccountTotals{} anal.remoteQuality = 0 anal.remoteMissingFragments = 0 anal.remoteMissingMemoryDomains = 0 @@ -102,6 +105,7 @@ func (anal *AnalyzeModule) Reset(isPrepare bool, isTpQuery bool) { func (anal *AnalyzeModule) appendRemoteResource( delta resource.Delta, memory resource.MemoryTotals, + allocation resource.AllocationAccountTotals, missingFragments uint64, missingMemoryDomains uint64, ) { @@ -113,6 +117,10 @@ func (anal *AnalyzeModule) appendRemoteResource( quality := delta.Quality quality |= resource.MergeUsage(&anal.remoteUsage, delta.Usage) quality |= resource.MergeMemoryTotals(&anal.remoteMemory, memory) + quality |= resource.MergeAllocationAccountTotals( + &anal.remoteAllocation, + allocation, + ) anal.remoteMissingFragments, quality = addCheckedRemoteCounter( anal.remoteMissingFragments, missingFragments, quality) anal.remoteMissingMemoryDomains, quality = addCheckedRemoteCounter( @@ -130,6 +138,7 @@ func (anal *AnalyzeModule) remoteResourceSummary() remoteResourceSnapshot { return remoteResourceSnapshot{ Usage: anal.remoteUsage, Memory: anal.remoteMemory, + Allocation: anal.remoteAllocation, Quality: anal.remoteQuality, MissingFragmentCount: anal.remoteMissingFragments, MissingMemoryDomainCount: anal.remoteMissingMemoryDomains, diff --git a/pkg/sql/compile/compile.go b/pkg/sql/compile/compile.go index 66d96eb05c66a..8d7cc309fb5c0 100644 --- a/pkg/sql/compile/compile.go +++ b/pkg/sql/compile/compile.go @@ -322,7 +322,9 @@ func (c *Compile) clear() { c.resourceAttemptOwnerEligible = false c.allocationAccountRegistry = nil c.allocationAccountLimit = 0 + c.allocationControllerProvider = nil c.allocationTerminalExporter = nil + c.allocationLifecycleAutomatic = false c.allocationAttempt = nil c.isPrepare = false c.hasMergeOp = false @@ -7002,10 +7004,26 @@ func (c *Compile) ConfigureAllocationAccountLifecycle( registry *mpool.AllocationAccountRegistry, limit uint64, exporter func(mpool.AllocationAccountTerminalSnapshot), +) { + c.ConfigureAllocationAccountLifecycleWithController( + registry, + limit, + nil, + exporter, + ) +} + +func (c *Compile) ConfigureAllocationAccountLifecycleWithController( + registry *mpool.AllocationAccountRegistry, + limit uint64, + controllerProvider func() (mpool.AllocationCapacityController, error), + exporter func(mpool.AllocationAccountTerminalSnapshot), ) { c.allocationAccountRegistry = registry c.allocationAccountLimit = limit + c.allocationControllerProvider = controllerProvider c.allocationTerminalExporter = exporter + c.allocationLifecycleAutomatic = false } func (c *Compile) SetBuildPlanFunc(buildPlanFunc func(ctx context.Context) (*plan2.Plan, error)) { diff --git a/pkg/sql/compile/compile2.go b/pkg/sql/compile/compile2.go index 4c87325b0858b..c12e920192924 100644 --- a/pkg/sql/compile/compile2.go +++ b/pkg/sql/compile/compile2.go @@ -24,6 +24,7 @@ import ( "time" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" commonutil "github.com/matrixorigin/matrixone/pkg/common/util" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/defines" @@ -303,7 +304,15 @@ func (c *Compile) Run(_ uint64) (queryResult *util2.RunResult, err error) { // Before compile.runOnce, Reset the 'StatsInfo' execution related resources in context // running. - allocationAttempt, err = runC.beginAllocationAccountAttempt() + exporter := func(snapshot mpool.AllocationAccountTerminalSnapshot) { + if resourceRecorder != nil { + resourceRecorder.recordAllocationAccountTerminal(snapshot) + } + } + err = runC.ensureAllocationAccountLifecycle(exporter) + if err == nil { + allocationAttempt, err = runC.beginAllocationAccountAttempt() + } if err == nil { err = runC.prePipelineInitializer() } diff --git a/pkg/sql/compile/remoterunClient.go b/pkg/sql/compile/remoterunClient.go index 1d99d825b7b53..16b84fe9aab1c 100644 --- a/pkg/sql/compile/remoterunClient.go +++ b/pkg/sql/compile/remoterunClient.go @@ -805,6 +805,7 @@ func (sender *messageSenderOnClient) dealRemoteTerminal(data []byte) error { sender.anal.appendRemoteResource( envelope.Delta, envelope.Memory, + envelope.Allocation, envelope.MissingFragmentCount, envelope.MissingMemoryDomainCount, ) diff --git a/pkg/sql/compile/remoterunServer.go b/pkg/sql/compile/remoterunServer.go index 3debcf7fa45c1..d7bfa5b87b05e 100644 --- a/pkg/sql/compile/remoterunServer.go +++ b/pkg/sql/compile/remoterunServer.go @@ -252,6 +252,8 @@ func handlePipelineMessage(receiver *messageReceiverOnServer) (err error) { return errBuildCompile } var allocationAttempt *statementAllocationAttempt + var localAllocation resource.AllocationAccountTotals + var localAllocationQuality resource.QualityFlags var runErr error defer func() { // Capture operator and descendant facts before cleanup. The MPool @@ -277,8 +279,14 @@ func handlePipelineMessage(receiver *messageReceiverOnServer) (err error) { descendant, expectedDirect, ) + aggregate.Delta.Quality |= localAllocationQuality | + resource.MergeAllocationAccountTotals( + &aggregate.Allocation, + localAllocation, + ) receiver.resourceDelta = aggregate.Delta receiver.resourceMemory = aggregate.Memory + receiver.resourceAllocation = aggregate.Allocation receiver.resourceMissingFragments = aggregate.MissingFragmentCount receiver.resourceMissingMemoryDomains = aggregate.MissingMemoryDomainCount @@ -307,7 +315,18 @@ func handlePipelineMessage(receiver *messageReceiverOnServer) (err error) { } runCompile.scopes = []*Scope{s} - allocationAttempt, runErr = runCompile.beginAllocationAccountAttempt() + runErr = runCompile.ensureAllocationAccountLifecycle(func( + snapshot mpool.AllocationAccountTerminalSnapshot, + ) { + localAllocationQuality |= localAllocation.AddGeneration( + snapshot.Peak, + snapshot.Used, + snapshot.State == mpool.AllocationAccountTerminalValid, + ) + }) + if runErr == nil { + allocationAttempt, runErr = runCompile.beginAllocationAccountAttempt() + } if runErr != nil { return runErr } @@ -569,6 +588,7 @@ type messageReceiverOnServer struct { phyPlan *models.PhyPlan resourceDelta resource.Delta resourceMemory resource.MemoryTotals + resourceAllocation resource.AllocationAccountTotals resourceMissingFragments uint64 resourceMissingMemoryDomains uint64 } @@ -871,6 +891,7 @@ func (receiver *messageReceiverOnServer) setTerminalAnalysis(message *pipeline.M TerminalResourceVersion: remoteTerminalResourceVersion, Delta: receiver.resourceDelta, Memory: receiver.resourceMemory, + Allocation: receiver.resourceAllocation, MissingFragmentCount: receiver.resourceMissingFragments, MissingMemoryDomainCount: receiver.resourceMissingMemoryDomains, } diff --git a/pkg/sql/compile/resource_accounting.go b/pkg/sql/compile/resource_accounting.go index e0cf9d1fa61dc..1d695199ae88e 100644 --- a/pkg/sql/compile/resource_accounting.go +++ b/pkg/sql/compile/resource_accounting.go @@ -20,6 +20,7 @@ import ( "sync/atomic" "time" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/sql/models" "github.com/matrixorigin/matrixone/pkg/util/resource" "github.com/matrixorigin/matrixone/pkg/util/trace/impl/motrace/statistic" @@ -27,14 +28,16 @@ import ( ) type executionResourceRecorder struct { - root *resource.Root - stats *statistic.StatsInfo - execution resource.ExecutionSummary - published bool - ownsAttempts bool + root *resource.Root + stats *statistic.StatsInfo + execution resource.ExecutionSummary + published bool + ownsAttempts bool + pendingAllocation resource.AllocationAccountTotals + pendingAllocationQuality resource.QualityFlags } -const remoteTerminalResourceVersion = 1 +const remoteTerminalResourceVersion = 2 // remoteTerminalEnvelope keeps PhyPlan fields at the top level so clients from // before resource accounting can still decode the terminal plan during a @@ -42,11 +45,12 @@ const remoteTerminalResourceVersion = 1 // appended resource facts from a legacy bare PhyPlan payload. type remoteTerminalEnvelope struct { models.PhyPlan - TerminalResourceVersion uint32 `json:"terminal_resource_version,omitempty"` - Delta resource.Delta `json:"resource_delta"` - Memory resource.MemoryTotals `json:"memory"` - MissingFragmentCount uint64 `json:"missing_fragment_count,omitempty"` - MissingMemoryDomainCount uint64 `json:"missing_memory_domain_count,omitempty"` + TerminalResourceVersion uint32 `json:"terminal_resource_version,omitempty"` + Delta resource.Delta `json:"resource_delta"` + Memory resource.MemoryTotals `json:"memory"` + Allocation resource.AllocationAccountTotals `json:"allocation_account"` + MissingFragmentCount uint64 `json:"missing_fragment_count,omitempty"` + MissingMemoryDomainCount uint64 `json:"missing_memory_domain_count,omitempty"` } // remoteResourceAggregate is the already-reduced terminal output sent by one @@ -54,10 +58,24 @@ type remoteTerminalEnvelope struct { type remoteResourceAggregate struct { Delta resource.Delta Memory resource.MemoryTotals + Allocation resource.AllocationAccountTotals MissingFragmentCount uint64 MissingMemoryDomainCount uint64 } +func (r *executionResourceRecorder) recordAllocationAccountTerminal( + snapshot mpool.AllocationAccountTerminalSnapshot, +) { + if r == nil { + return + } + r.pendingAllocationQuality |= r.pendingAllocation.AddGeneration( + snapshot.Peak, + snapshot.Used, + snapshot.State == mpool.AllocationAccountTerminalValid, + ) +} + func newExecutionResourceRecorder( ctx context.Context, attemptOwnerEligible bool, @@ -170,6 +188,17 @@ func (r *executionResourceRecorder) finishAttempt( summary := resource.AttemptSummary{WallNS: wallNS} summary.Quality |= delta.Quality | resource.MergeUsage(&summary.Usage, delta.Usage) summary.Quality |= resource.MergeMemoryTotals(&summary.Memory, remoteAggregate.Memory) + summary.Quality |= resource.MergeAllocationAccountTotals( + &summary.Allocation, + remoteAggregate.Allocation, + ) + summary.Quality |= r.pendingAllocationQuality | + resource.MergeAllocationAccountTotals( + &summary.Allocation, + r.pendingAllocation, + ) + r.pendingAllocation = resource.AllocationAccountTotals{} + r.pendingAllocationQuality = 0 var quality resource.QualityFlags summary.MissingFragmentCount, quality = addCheckedRemoteCounter( summary.MissingFragmentCount, remoteAggregate.MissingFragmentCount, summary.Quality) @@ -236,6 +265,10 @@ func composeRemoteResourceAggregate( result.Delta.Quality |= resource.MergeUsage(&result.Delta.Usage, descendant.Usage) result.Delta.Quality |= resource.MergeMemoryDomain(&result.Memory, localMemory) result.Delta.Quality |= resource.MergeMemoryTotals(&result.Memory, descendant.Memory) + result.Delta.Quality |= resource.MergeAllocationAccountTotals( + &result.Allocation, + descendant.Allocation, + ) result.MissingFragmentCount = descendant.MissingFragmentCount result.MissingMemoryDomainCount = descendant.MissingMemoryDomainCount if descendant.MissingFragmentCount > 0 { diff --git a/pkg/sql/compile/resource_accounting_test.go b/pkg/sql/compile/resource_accounting_test.go index 216afa990782d..c313049e8a99d 100644 --- a/pkg/sql/compile/resource_accounting_test.go +++ b/pkg/sql/compile/resource_accounting_test.go @@ -22,6 +22,7 @@ import ( "testing" "time" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/sql/colexec/value_scan" "github.com/matrixorigin/matrixone/pkg/sql/models" "github.com/matrixorigin/matrixone/pkg/testutil" @@ -51,6 +52,7 @@ func TestExecutionResourceRecorder(t *testing.T) { MaxDomainPeakLiveBytes: 80, SumDomainPeakLiveBytesBound: 80, }, + resource.AllocationAccountTotals{}, 0, 0, ) @@ -75,6 +77,42 @@ func TestExecutionResourceRecorder(t *testing.T) { require.Zero(t, summary.Quality&resource.QualityMissingFragment) } +func TestExecutionResourceRecorderPublishesAllocationTerminal(t *testing.T) { + root := resource.NewRoot(resource.ConnExternal) + recorder := newExecutionResourceRecorder( + resource.ContextWithRoot(context.Background(), root), + true, + ) + require.NotNil(t, recorder) + recorder.recordAllocationAccountTerminal( + mpool.AllocationAccountTerminalSnapshot{ + AllocationAccountSnapshot: mpool.AllocationAccountSnapshot{ + Peak: 64, + }, + State: mpool.AllocationAccountTerminalValid, + }, + ) + recorder.finishAttempt( + 0, + time.Now(), + 0, + 0, + nil, + nil, + nil, + "", + false, + ) + recorder.publish() + + summary := root.PreResponseSummary() + require.Equal(t, uint64(1), summary.Allocation.GenerationCount) + require.Equal(t, uint64(1), summary.Allocation.ValidGenerationCount) + require.Equal(t, uint64(64), summary.Allocation.MaxGenerationPeak) + require.Zero(t, summary.Allocation.LiveBytesAtTerminal) + require.Zero(t, summary.Quality&resource.QualityInvariantFailure) +} + func TestExplainPhyBufferUsesPublishedCurrentAttempt(t *testing.T) { root := resource.NewRoot(resource.ConnExternal) require.True(t, root.MergeExecution(resource.ExecutionSummary{ @@ -258,6 +296,12 @@ func TestRemoteTerminalEnvelope(t *testing.T) { MaxDomainPeakLiveBytes: 15, SumDomainPeakLiveBytesBound: 15, }, + Allocation: resource.AllocationAccountTotals{ + GenerationCount: 1, + ValidGenerationCount: 1, + MaxGenerationPeak: 17, + SumGenerationPeak: 17, + }, } data, err := json.Marshal(envelope) require.NoError(t, err) @@ -269,6 +313,8 @@ func TestRemoteTerminalEnvelope(t *testing.T) { require.Equal(t, uint64(11), summary.Usage.ExclusiveActiveNS) require.Equal(t, uint64(12), summary.Usage.S3ReadBytes) require.Equal(t, uint64(15), summary.Memory.MaxDomainPeakLiveBytes) + require.Equal(t, uint64(1), summary.Allocation.GenerationCount) + require.Equal(t, uint64(17), summary.Allocation.MaxGenerationPeak) require.NotZero(t, summary.Quality&resource.QualityPartial) require.Len(t, anal.remotePhyPlans, 1) require.Equal(t, "Merge", anal.remotePhyPlans[0].LocalScope[0].Magic) @@ -444,7 +490,13 @@ func TestRemoteResourceCounterSaturates(t *testing.T) { remoteMissingMemoryDomains: math.MaxUint64, remoteReports: math.MaxUint64, } - anal.appendRemoteResource(resource.Delta{}, resource.MemoryTotals{}, 1, 1) + anal.appendRemoteResource( + resource.Delta{}, + resource.MemoryTotals{}, + resource.AllocationAccountTotals{}, + 1, + 1, + ) snapshot := anal.remoteResourceSummary() require.Equal(t, uint64(math.MaxUint64), snapshot.MissingFragmentCount) require.Equal(t, uint64(math.MaxUint64), snapshot.MissingMemoryDomainCount) @@ -466,7 +518,11 @@ func TestAnalyzeModuleResetClearsRemoteResourceAggregate(t *testing.T) { anal := &AnalyzeModule{} anal.appendRemoteResource( resource.Delta{Usage: resource.Usage{S3ReadBytes: 11}, Quality: resource.QualityPartial}, - resource.MemoryTotals{MaxDomainPeakLiveBytes: 8}, 2, 3) + resource.MemoryTotals{MaxDomainPeakLiveBytes: 8}, + resource.AllocationAccountTotals{}, + 2, + 3, + ) anal.Reset(false, false) snapshot := anal.remoteResourceSummary() require.Equal(t, remoteResourceSnapshot{}, snapshot) @@ -480,7 +536,13 @@ func TestAnalyzeModuleRemoteResourceConcurrentAccess(t *testing.T) { go func() { defer wg.Done() for j := 0; j < 100; j++ { - anal.appendRemoteResource(resource.Delta{Usage: resource.Usage{SpillBytes: 1}}, resource.MemoryTotals{}, 1, 1) + anal.appendRemoteResource( + resource.Delta{Usage: resource.Usage{SpillBytes: 1}}, + resource.MemoryTotals{}, + resource.AllocationAccountTotals{}, + 1, + 1, + ) _ = anal.remoteResourceSummary() } }() diff --git a/pkg/sql/compile/types.go b/pkg/sql/compile/types.go index a596dfa427d4e..cd7a13db799ec 100644 --- a/pkg/sql/compile/types.go +++ b/pkg/sql/compile/types.go @@ -346,7 +346,9 @@ type Compile struct { resourceAttemptOwnerEligible bool allocationAccountRegistry *mpool.AllocationAccountRegistry allocationAccountLimit uint64 + allocationControllerProvider func() (mpool.AllocationCapacityController, error) allocationTerminalExporter func(mpool.AllocationAccountTerminalSnapshot) + allocationLifecycleAutomatic bool allocationAttempt *statementAllocationAttempt hasMergeOp bool diff --git a/pkg/util/resource/summary.go b/pkg/util/resource/summary.go index 3eef2ffd9613a..05c69fd8541b0 100644 --- a/pkg/util/resource/summary.go +++ b/pkg/util/resource/summary.go @@ -57,6 +57,106 @@ type MemoryTotals struct { CrossPoolFreeCount uint64 } +// AllocationAccountTotals is the fixed-size terminal observation of activated +// allocation generations. It is diagnostic only: these bytes are a subset of +// allocator memory and are never added to MemoryTotals or fed back into +// admission. +type AllocationAccountTotals struct { + GenerationCount uint64 + ValidGenerationCount uint64 + FailedGenerationCount uint64 + MaxGenerationPeak uint64 + SumGenerationPeak uint64 + LiveBytesAtTerminal uint64 +} + +func (t *AllocationAccountTotals) AddGeneration( + peak uint64, + liveAtTerminal uint64, + valid bool, +) QualityFlags { + var quality QualityFlags + t.GenerationCount, quality = addChecked(t.GenerationCount, 1, quality) + if valid { + t.ValidGenerationCount, quality = addChecked( + t.ValidGenerationCount, + 1, + quality, + ) + } else { + t.FailedGenerationCount, quality = addChecked( + t.FailedGenerationCount, + 1, + quality|QualityInvariantFailure, + ) + } + if peak > t.MaxGenerationPeak { + t.MaxGenerationPeak = peak + } + t.SumGenerationPeak, quality = addChecked( + t.SumGenerationPeak, + peak, + quality, + ) + t.LiveBytesAtTerminal, quality = addChecked( + t.LiveBytesAtTerminal, + liveAtTerminal, + quality, + ) + if liveAtTerminal != 0 { + quality |= QualityNonZeroLiveAtSeal | QualityInvariantFailure + } + return quality +} + +func MergeAllocationAccountTotals( + dst *AllocationAccountTotals, + delta AllocationAccountTotals, +) QualityFlags { + var quality QualityFlags + dst.GenerationCount, quality = addChecked( + dst.GenerationCount, + delta.GenerationCount, + quality, + ) + dst.ValidGenerationCount, quality = addChecked( + dst.ValidGenerationCount, + delta.ValidGenerationCount, + quality, + ) + dst.FailedGenerationCount, quality = addChecked( + dst.FailedGenerationCount, + delta.FailedGenerationCount, + quality, + ) + dst.SumGenerationPeak, quality = addChecked( + dst.SumGenerationPeak, + delta.SumGenerationPeak, + quality, + ) + dst.LiveBytesAtTerminal, quality = addChecked( + dst.LiveBytesAtTerminal, + delta.LiveBytesAtTerminal, + quality, + ) + if delta.MaxGenerationPeak > dst.MaxGenerationPeak { + dst.MaxGenerationPeak = delta.MaxGenerationPeak + } + if delta.FailedGenerationCount != 0 || delta.LiveBytesAtTerminal != 0 { + quality |= QualityInvariantFailure + } + if delta.ValidGenerationCount > delta.GenerationCount || + delta.FailedGenerationCount > + delta.GenerationCount-delta.ValidGenerationCount || + delta.MaxGenerationPeak > delta.SumGenerationPeak { + quality |= QualityInvariantFailure + } + if delta.LiveBytesAtTerminal != 0 { + quality |= QualityNonZeroLiveAtSeal + } + return quality +} + // MergeMemoryDomain merges one physical domain exactly once. func MergeMemoryDomain(dst *MemoryTotals, domain MemoryDomainSummary) QualityFlags { flags := domain.Validate() @@ -74,8 +174,9 @@ func MergeMemoryDomain(dst *MemoryTotals, domain MemoryDomainSummary) QualityFla // AttemptSummary is the immutable result of one compile/run generation. type AttemptSummary struct { - Usage Usage - Memory MemoryTotals + Usage Usage + Memory MemoryTotals + Allocation AllocationAccountTotals WallNS uint64 MissingFragmentCount uint64 @@ -85,8 +186,9 @@ type AttemptSummary struct { // ExecutionSummary is fixed-size in retry count. type ExecutionSummary struct { - Usage Usage - Memory MemoryTotals + Usage Usage + Memory MemoryTotals + Allocation AllocationAccountTotals AttemptCount uint64 RetryWallNS uint64 @@ -100,6 +202,7 @@ type ExecutionSummary struct { func (s *ExecutionSummary) AddAttempt(attempt AttemptSummary, retried bool) { s.Quality |= attempt.Quality | MergeUsage(&s.Usage, attempt.Usage) s.Quality |= MergeMemoryTotals(&s.Memory, attempt.Memory) + s.Quality |= MergeAllocationAccountTotals(&s.Allocation, attempt.Allocation) s.AttemptCount, s.Quality = addChecked(s.AttemptCount, 1, s.Quality) if retried { s.RetryWallNS, s.Quality = addChecked(s.RetryWallNS, attempt.WallNS, s.Quality) @@ -123,8 +226,9 @@ const ( // algebra. Serialization and plan diagnostics consume this value but never add // resources independently. type StatementResourceSummary struct { - Usage Usage - Memory MemoryTotals + Usage Usage + Memory MemoryTotals + Allocation AllocationAccountTotals StatementWallNS uint64 AttemptCount uint64 @@ -140,6 +244,7 @@ type StatementResourceSummary struct { func (s *StatementResourceSummary) MergeExecution(execution ExecutionSummary) { s.Quality |= execution.Quality | MergeUsage(&s.Usage, execution.Usage) s.Quality |= MergeMemoryTotals(&s.Memory, execution.Memory) + s.Quality |= MergeAllocationAccountTotals(&s.Allocation, execution.Allocation) s.AttemptCount, s.Quality = addChecked(s.AttemptCount, execution.AttemptCount, s.Quality) s.RetryWallNS, s.Quality = addChecked(s.RetryWallNS, execution.RetryWallNS, s.Quality) s.MissingFragmentCount, s.Quality = addChecked( @@ -154,6 +259,7 @@ func (s *StatementResourceSummary) MergeExecution(execution ExecutionSummary) { func (s *StatementResourceSummary) Merge(other StatementResourceSummary) { s.Quality |= other.Quality | QualityAggregated | MergeUsage(&s.Usage, other.Usage) s.Quality |= MergeMemoryTotals(&s.Memory, other.Memory) + s.Quality |= MergeAllocationAccountTotals(&s.Allocation, other.Allocation) s.StatementWallNS, s.Quality = addChecked(s.StatementWallNS, other.StatementWallNS, s.Quality) s.AttemptCount, s.Quality = addChecked(s.AttemptCount, other.AttemptCount, s.Quality) s.RetryWallNS, s.Quality = addChecked(s.RetryWallNS, other.RetryWallNS, s.Quality) diff --git a/pkg/util/resource/usage_test.go b/pkg/util/resource/usage_test.go index 3968db4711d0d..1bf6d14ecdbe3 100644 --- a/pkg/util/resource/usage_test.go +++ b/pkg/util/resource/usage_test.go @@ -85,6 +85,30 @@ func TestMergeUsageOverflowSaturates(t *testing.T) { } } +func TestAllocationAccountTotalsMergeAndFailureQuality(t *testing.T) { + var valid AllocationAccountTotals + if quality := valid.AddGeneration(10, 0, true); quality != 0 { + t.Fatalf("valid generation quality = %v", quality) + } + var failed AllocationAccountTotals + quality := failed.AddGeneration(20, 5, false) + if quality&QualityInvariantFailure == 0 || + quality&QualityNonZeroLiveAtSeal == 0 { + t.Fatalf("failure quality = %v", quality) + } + + quality = MergeAllocationAccountTotals(&valid, failed) + if valid.GenerationCount != 2 || valid.ValidGenerationCount != 1 || + valid.FailedGenerationCount != 1 || valid.MaxGenerationPeak != 20 || + valid.SumGenerationPeak != 30 || valid.LiveBytesAtTerminal != 5 { + t.Fatalf("merged allocation totals = %+v", valid) + } + if quality&QualityInvariantFailure == 0 || + quality&QualityNonZeroLiveAtSeal == 0 { + t.Fatalf("merge quality = %v", quality) + } +} + func TestLocalRecorder(t *testing.T) { var recorder LocalRecorder recorder.AddActiveInterval(100, 10, 20) diff --git a/pkg/vm/process/hashbuild_budget.go b/pkg/vm/process/hashbuild_budget.go index 963435b7fc012..dc0db22802451 100644 --- a/pkg/vm/process/hashbuild_budget.go +++ b/pkg/vm/process/hashbuild_budget.go @@ -15,6 +15,7 @@ package process import ( + "errors" "fmt" "math" "sync" @@ -30,6 +31,11 @@ import ( const hashBuildMinimumReserve = uint64(4 << 30) +const ( + hashBuildAllocationGenerationSlots = uint32(131_072) + hashBuildMinimumCellBlockBytes = uint64(16 << 10) +) + const ( // Keep a process-wide reserve for listeners, RPC connections, object // storage, logs, and other descriptors that are not represented by the @@ -257,6 +263,10 @@ type HashBuildBudget struct { spillFDConfiguredCap uint64 spillFDCap uint64 spillFDUsed uint64 + + allocationRegistryOnce sync.Once + allocationRegistry *mpool.AllocationAccountRegistry + allocationRegistryErr error } // NewHashBuildBudget creates a local-CN budget. Both caps are finite and @@ -1124,6 +1134,40 @@ func (g *HashBuildBudgetGeneration) Closed() bool { return g.closed || g.budget.closed } +// AllocationAccountRegistry returns the bounded CN-local registry shared by +// every activated HashBuild generation under this aggregate budget. The slot +// formula covers every minimum-size cell block, its published descriptor, and +// one private replacement descriptor/block transaction. +func (g *HashBuildBudgetGeneration) AllocationAccountRegistry() ( + *mpool.AllocationAccountRegistry, + error, +) { + if g == nil || g.budget == nil { + return nil, ErrHashBuildBudgetInvalid + } + b := g.budget + b.allocationRegistryOnce.Do(func() { + capBytes := b.AggregateCap() + blocks := capBytes / hashBuildMinimumCellBlockBytes + if capBytes%hashBuildMinimumCellBlockBytes != 0 { + blocks++ + } + if blocks == 0 { + blocks = 1 + } + if blocks > math.MaxUint64/3 { + b.allocationRegistryErr = ErrHashBuildBudgetInvalid + return + } + b.allocationRegistry, b.allocationRegistryErr = + mpool.NewAllocationAccountRegistry( + hashBuildAllocationGenerationSlots, + blocks*3, + ) + }) + return b.allocationRegistry, b.allocationRegistryErr +} + // Close rejects future reservations for this generation while allowing all // currently live tokens to release. It is idempotent. func (g *HashBuildBudgetGeneration) Close() { @@ -1144,7 +1188,17 @@ func (g *HashBuildBudgetGeneration) AcquireAllocationCapacity(size uint64) error return nil } _, err := g.reserve(size, true) - return err + if err == nil { + return nil + } + switch { + case errors.Is(err, ErrHashBuildBudgetClosed): + return errors.Join(mpool.ErrAllocationAccountSealed, err) + case errors.Is(err, ErrHashBuildBudgetAdmission): + return errors.Join(mpool.ErrAllocationAccountCapacity, err) + default: + return errors.Join(mpool.ErrAllocationAccountInvariant, err) + } } // ReleaseAllocationCapacity is called only by physical MPool Free through the diff --git a/pkg/vm/process/hashbuild_budget_test.go b/pkg/vm/process/hashbuild_budget_test.go index 2c3cc29266bba..a3e9e8f7ee027 100644 --- a/pkg/vm/process/hashbuild_budget_test.go +++ b/pkg/vm/process/hashbuild_budget_test.go @@ -184,6 +184,11 @@ func TestHashBuildBudgetAllocationAccountAdapter(t *testing.T) { ) { t.Fatalf("combined legacy/exact admission error = %v", err) } + if !errors.Is(err, commonmpool.ErrAllocationAccountCapacity) || + commonmpool.AllocationFailureReasonOf(err) != + commonmpool.AllocationFailureCapacity { + t.Fatalf("adapter did not type policy pressure as capacity: %v", err) + } if account.Snapshot().Used != 6 || registry.LiveAllocationMetadata() != 1 { t.Fatal("failed adapter admission did not roll back") @@ -196,6 +201,10 @@ func TestHashBuildBudgetAllocationAccountAdapter(t *testing.T) { ) { t.Fatalf("closed generation admission error = %v", err) } + if !errors.Is(err, commonmpool.ErrAllocationAccountSealed) || + commonmpool.IsRetryableAllocationCapacity(err) { + t.Fatalf("closed adapter error entered capacity retry: %v", err) + } if account.Snapshot().Used != 6 || registry.LiveAllocationMetadata() != 1 { t.Fatal("closed adapter admission did not roll back") @@ -216,6 +225,35 @@ func TestHashBuildBudgetAllocationAccountAdapter(t *testing.T) { } } +func TestHashBuildAllocationAccountRegistryUsesBoundedFormula(t *testing.T) { + budget := MustNewHashBuildBudget(16<<10, 16<<10) + first, err := budget.OpenGeneration(1) + if err != nil { + t.Fatal(err) + } + registry, err := first.AllocationAccountRegistry() + if err != nil { + t.Fatal(err) + } + if registry.GenerationCapacity() != hashBuildAllocationGenerationSlots { + t.Fatalf("generation slots = %d", registry.GenerationCapacity()) + } + if registry.MaxAllocationMetadata() != 3 { + t.Fatalf("allocation slots = %d, want 3", registry.MaxAllocationMetadata()) + } + second, err := budget.OpenGeneration(2) + if err != nil { + t.Fatal(err) + } + secondRegistry, err := second.AllocationAccountRegistry() + if err != nil { + t.Fatal(err) + } + if secondRegistry != registry { + t.Fatal("one CN budget created multiple allocation registries") + } +} + func TestHashBuildBudgetQueryRejectRollsBackCN(t *testing.T) { b := MustNewHashBuildBudget(10, 4) g1, _ := b.OpenGeneration(1) From c74d0857ac6e8c822fd3eaaca29b89b336d6934d Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 21:34:36 +0800 Subject: [PATCH 16/61] executor: account retained join-map batches --- pkg/sql/colexec/hashbuild/budget.go | 30 ++- pkg/sql/colexec/hashbuild/hashmap.go | 21 +- pkg/sql/colexec/hashbuild/hashmap_test.go | 257 +++++++++++++++++++++- pkg/sql/colexec/hashbuild/types.go | 21 +- pkg/sql/colexec/join_util.go | 34 ++- pkg/vm/message/group_sels_test.go | 63 ++++++ pkg/vm/message/joinMapMsg.go | 81 ++++++- pkg/vm/process/hashbuild_budget.go | 14 +- pkg/vm/process/hashbuild_budget_test.go | 8 +- pkg/vm/process/process.go | 32 ++- 10 files changed, 528 insertions(+), 33 deletions(-) diff --git a/pkg/sql/colexec/hashbuild/budget.go b/pkg/sql/colexec/hashbuild/budget.go index 481a50fea627f..59dfc9fd59b9e 100644 --- a/pkg/sql/colexec/hashbuild/budget.go +++ b/pkg/sql/colexec/hashbuild/budget.go @@ -322,6 +322,13 @@ func batchCopyAllocatedDelta( } func (hb *HashmapBuilder) copyBuildBatch(src *batch.Batch, proc *process.Process) error { + if hb.batchAllocation != nil { + return hb.Batches.CopyIntoBatchesWithAllocation( + src, + proc, + hb.batchAllocation, + ) + } if hb.budget == nil { return hb.Batches.CopyIntoBatches(src, proc) } @@ -588,13 +595,14 @@ func (hb *HashmapBuilder) cleanBatches(proc *process.Process) { func (hb *HashmapBuilder) buildAuxBytes( needUniqueVec bool, + needAllocateSels ...bool, ) (uint64, error) { uniqueBytes, err := hb.uniqueJoinKeyBytes() if err != nil { return 0, err } return hb.buildAuxBytesWithUniqueProjection( - needUniqueVec, uniqueBytes) + needUniqueVec, uniqueBytes, needAllocateSels...) } func (hb *HashmapBuilder) uniqueJoinKeyBytes() (uint64, error) { @@ -616,6 +624,7 @@ func (hb *HashmapBuilder) uniqueJoinKeyBytes() (uint64, error) { func (hb *HashmapBuilder) buildAuxBytesWithUniqueProjection( needUniqueVec bool, uniqueBytes uint64, + needAllocateSels ...bool, ) (uint64, error) { // Covers mandatory hashmap/sels scratch plus the selected runtime-filter // key vectors' actual persistent capacities. Before their first append, a @@ -644,15 +653,26 @@ func (hb *HashmapBuilder) buildAuxBytesWithUniqueProjection( rowCount = hb.hashMapRowCount } rows := uint64(rowCount) + perRowBytes := uint64(64) + if len(needAllocateSels) > 0 && needAllocateSels[0] && + hb.batchAllocation != nil { + // GroupSels' physical slices are charged by batchAllocation. + perRowBytes -= 16 + } const iteratorScratch = uint64(640 << 10) - if rows > math.MaxUint64/64 || bytes > math.MaxUint64-rows*64 || bytes+rows*64 > math.MaxUint64-iteratorScratch { + if rows > math.MaxUint64/perRowBytes || + bytes > math.MaxUint64-rows*perRowBytes || + bytes+rows*perRowBytes > math.MaxUint64-iteratorScratch { return 0, process.ErrHashBuildBudgetInvalid } - bytes += rows*64 + iteratorScratch + bytes += rows*perRowBytes + iteratorScratch return bytes, nil } -func (hb *HashmapBuilder) reserveBuildAux(needUniqueVec bool) error { +func (hb *HashmapBuilder) reserveBuildAux( + needUniqueVec bool, + needAllocateSels ...bool, +) error { if hb.budget == nil { return nil } @@ -663,7 +683,7 @@ func (hb *HashmapBuilder) reserveBuildAux(needUniqueVec bool) error { // worse, collecting optional keys under a mandatory-only charge). return hb.resizeBuildAuxReservation(needUniqueVec) } - bytes, err := hb.buildAuxBytes(needUniqueVec) + bytes, err := hb.buildAuxBytes(needUniqueVec, needAllocateSels...) if err != nil { return err } diff --git a/pkg/sql/colexec/hashbuild/hashmap.go b/pkg/sql/colexec/hashbuild/hashmap.go index b9248a014f3d6..7f72692e241cb 100644 --- a/pkg/sql/colexec/hashbuild/hashmap.go +++ b/pkg/sql/colexec/hashbuild/hashmap.go @@ -91,6 +91,7 @@ type HashmapBuilder struct { retainedBatchRecoverySafe bool mapAllocationAccount *mpool.AllocationAccount mapAllocation *hashtable.AllocationAccountSelection + batchAllocation *vector.AllocationAccountSelection } func (hb *HashmapBuilder) GetSize() int64 { @@ -245,6 +246,7 @@ func (hb *HashmapBuilder) Reset(proc *process.Process, hashTableHasNotSent bool) hb.FreeExecutors() hb.mapAllocationAccount = nil hb.mapAllocation = nil + hb.batchAllocation = nil } func (hb *HashmapBuilder) Free(proc *process.Process) { @@ -268,6 +270,7 @@ func (hb *HashmapBuilder) Free(proc *process.Process) { hb.uniqueKeySlots = nil hb.mapAllocationAccount = nil hb.mapAllocation = nil + hb.batchAllocation = nil } func (hb *HashmapBuilder) FreeExecutors() { @@ -616,7 +619,7 @@ func (hb *HashmapBuilder) buildHashmap( if hb.InputBatchRowCount == 0 { return nil } - if err := hb.reserveBuildAux(needUniqueVec); err != nil { + if err := hb.reserveBuildAux(needUniqueVec, needAllocateSels); err != nil { if !needUniqueVec { return err } @@ -628,7 +631,7 @@ func (hb *HashmapBuilder) buildHashmap( // retention. Retry the admission in place without that owner before // allocating or mutating the mandatory map. needUniqueVec = false - if err = hb.reserveBuildAux(false); err != nil { + if err = hb.reserveBuildAux(false, needAllocateSels); err != nil { return err } // Linearize the fallback only after mandatory admission succeeds. A @@ -733,7 +736,19 @@ func (hb *HashmapBuilder) buildHashmap( } if needAllocateSels { - if err := hb.Sels.Init(hb.InputBatchRowCount, proc.Mp()); err != nil { + var err error + if hb.batchAllocation == nil { + err = hb.Sels.Init(hb.InputBatchRowCount, proc.Mp()) + } else { + err = hb.Sels.InitWithAllocation( + hb.InputBatchRowCount, + proc.Mp(), + hb.mapAllocationAccount, + HashBuildAllocationOwner, + HashBuildAllocationSiteGroupSels, + ) + } + if err != nil { return err } } diff --git a/pkg/sql/colexec/hashbuild/hashmap_test.go b/pkg/sql/colexec/hashbuild/hashmap_test.go index 716182b73bc13..3366ce427e66c 100644 --- a/pkg/sql/colexec/hashbuild/hashmap_test.go +++ b/pkg/sql/colexec/hashbuild/hashmap_test.go @@ -343,7 +343,7 @@ func TestPublishedJoinMapResizeKeepsReservationWithConsumer(t *testing.T) { hb.Reset(proc, false) } -func TestHashmapBuilderAccountedCellsDoNotStackLegacyMapReservation(t *testing.T) { +func TestHashmapBuilderAccountedJoinMapDoesNotStackLegacyReservations(t *testing.T) { const budgetCap = uint64(16 << 20) budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) require.NoError(t, err) @@ -375,6 +375,11 @@ func TestHashmapBuilderAccountedCellsDoNotStackLegacyMapReservation(t *testing.T proc.Mp(), ) require.NoError(t, hb.copyBuildBatch(input, proc)) + require.Empty(t, hb.batchReservations) + require.NotEmpty(t, hb.Batches.Buf) + for _, copied := range hb.Batches.Buf { + require.Same(t, hb.batchAllocation, copied.AllocationAccountSelection()) + } hb.InputBatchRowCount = input.RowCount() input.Clean(proc.Mp()) @@ -389,15 +394,132 @@ func TestHashmapBuilderAccountedCellsDoNotStackLegacyMapReservation(t *testing.T jm := hb.GetJoinMap(proc.Mp()) require.NotNil(t, jm) - jm.IncRef(1) + jm.IncRef(2) + hb.Reset(proc, false) beforeResize := account.Snapshot().Used require.NoError(t, jm.PreAlloc(100_000)) require.Greater(t, account.Snapshot().Used, beforeResize) + beforeFirstConsumer := account.Snapshot().Used + jm.Free() + require.Equal(t, beforeFirstConsumer, account.Snapshot().Used) jm.Free() require.Zero(t, account.Snapshot().Used) require.Zero(t, generation.Used()) + terminal, first, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.True(t, first) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) +} + +func TestHashmapBuilderAccountedBatchCopyOneByteShortRollsBack(t *testing.T) { + const budgetCap = uint64(64 << 20) + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + input := testutil.NewBatch( + []types.Type{types.T_int32.ToType(), types.T_varchar.ToType()}, + true, + 10_000, + proc.Mp(), + ) + defer input.Clean(proc.Mp()) + + measure := func(limit uint64, metadataSlots uint64) ( + mpool.AllocationAccountSnapshot, + uint64, + error, + ) { + budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, metadataSlots) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + var op HashBuild + op.NeedHashMap = true + require.NoError(t, op.SetAllocationAccount(account)) + hb := &op.ctr.hashmapBuilder + hb.setBudget(generation) + + copyErr := hb.copyBuildBatch(input, proc) + snapshot := account.Snapshot() + metadataPeak := registry.PeakAllocationMetadata() + if copyErr == nil { + hb.cleanBatches(proc) + } + require.Empty(t, hb.Batches.Buf) + require.Empty(t, hb.batchReservations) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.NoError(t, op.ClearAllocationAccount(account)) + terminal, first, terminalErr := registry.CompleteTerminal(account) + require.NoError(t, terminalErr) + require.True(t, first) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) + return snapshot, metadataPeak, copyErr + } + + probe, metadataPeak, err := measure(budgetCap, 128) + require.NoError(t, err) + require.Positive(t, probe.Peak) + require.Positive(t, metadataPeak) + rejected, _, err := measure(probe.Peak-1, 128) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Zero(t, rejected.Used) + rejected, _, err = measure(budgetCap, metadataPeak-1) + require.ErrorIs(t, err, mpool.ErrAllocationMetadataSlots) + require.Zero(t, rejected.Used) +} + +func TestAccountedJoinMapTransfersBatchesAndGroupSelsToLastConsumer(t *testing.T) { + const budgetCap = uint64(16 << 20) + budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.OpenWithController(budgetCap, generation) + require.NoError(t, err) + var op HashBuild + op.NeedHashMap = true + require.NoError(t, op.SetAllocationAccount(account)) + hb := &op.ctr.hashmapBuilder + hb.setBudget(generation) + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + require.NoError(t, hb.Prepare( + []*plan.Expr{newExpr(0, types.T_int32.ToType())}, + -1, + -1, + nil, + proc, + )) + input := makeIntKeyValueBatch( + proc, + []int32{1, 1, 2, 2}, + []int32{10, 20, 30, 40}, + ) + require.NoError(t, hb.copyBuildBatch(input, proc)) + hb.InputBatchRowCount = input.RowCount() + input.Clean(proc.Mp()) + require.NoError(t, hb.BuildHashmap(false, true, false, proc)) + require.Positive(t, hb.Sels.Size()) + + jm := hb.GetJoinMap(proc.Mp()) + require.NotNil(t, jm) + jm.IncRef(2) hb.Reset(proc, false) + require.Equal(t, []int32{0, 1}, jm.GetSels(0)) + require.Equal(t, []int32{2, 3}, jm.GetSels(1)) + live := account.Snapshot().Used + require.Positive(t, live) + jm.Free() + require.Equal(t, live, account.Snapshot().Used) + jm.Free() + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + terminal, first, err := registry.CompleteTerminal(account) require.NoError(t, err) require.True(t, first) @@ -542,6 +664,37 @@ func TestAccountedEmptyJoinMapInitialFailureRollsBackController(t *testing.T) { require.NoError(t, err) } +func TestAccountedJoinMapLateFreeKeepsOriginalGeneration(t *testing.T) { + const capBytes = uint64(64 << 20) + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + firstGeneration, err := budget.OpenGeneration(1) + require.NoError(t, err) + secondGeneration, err := budget.OpenGeneration(2) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(2, 16) + require.NoError(t, err) + firstAccount, err := registry.OpenWithController(capBytes, firstGeneration) + require.NoError(t, err) + mp := mpool.MustNewZero() + jm, err := NewAccountedEmptyJoinMap(4, firstAccount, mp) + require.NoError(t, err) + firstUsed := firstGeneration.Used() + require.Positive(t, firstUsed) + + secondAccount, err := registry.OpenWithController(capBytes, secondGeneration) + require.NoError(t, err) + require.Zero(t, secondGeneration.Used()) + jm.Free() + require.Zero(t, firstGeneration.Used()) + require.Zero(t, firstAccount.Snapshot().Used) + require.Zero(t, secondGeneration.Used()) + + _, _, err = registry.CompleteTerminal(firstAccount) + require.NoError(t, err) + _, _, err = registry.CompleteTerminal(secondAccount) + require.NoError(t, err) +} + func TestCopyBuildBatchBudgetsSmallIngressAfterFullBatches(t *testing.T) { const budgetCap = uint64(32 << 20) budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) @@ -986,7 +1139,7 @@ func TestReserveBuildAuxChargesOneRetainedCopy(t *testing.T) { generation, err := budget.OpenGeneration(1) require.NoError(t, err) hb.setBudget(generation) - require.NoError(t, hb.reserveBuildAux(true)) + require.NoError(t, hb.reserveBuildAux(true, false)) require.Equal(t, want, generation.Used()) hb.releaseReservations() require.Zero(t, generation.Used()) @@ -994,6 +1147,34 @@ func TestReserveBuildAuxChargesOneRetainedCopy(t *testing.T) { hb.Batches.Buf = nil } +func TestReserveBuildAuxDoesNotStackAccountedGroupSels(t *testing.T) { + const rows = 10_000 + const iteratorScratch = uint64(640 << 10) + want := uint64(rows)*48 + iteratorScratch + budget := process.MustNewHashBuildBudget(want, want) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.OpenWithController(want, generation) + require.NoError(t, err) + var op HashBuild + op.NeedHashMap = true + require.NoError(t, op.SetAllocationAccount(account)) + hb := &op.ctr.hashmapBuilder + hb.InputBatchRowCount = rows + hb.setBudget(generation) + + require.NoError(t, hb.reserveBuildAux(false, true)) + require.Equal(t, want, generation.Used()) + require.Zero(t, generation.Snapshot().AllocationUsed) + hb.releaseReservations() + require.Zero(t, generation.Used()) + require.NoError(t, op.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + func TestReserveUniqueAppendOverlapChargesReplacedCapacity(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() @@ -2259,6 +2440,76 @@ func BenchmarkBuildHashmapColdStr(b *testing.B) { } } +func BenchmarkCopyBuildBatchAccounting(b *testing.B) { + const capBytes = uint64(256 << 20) + for _, accounted := range []bool{false, true} { + name := "legacy" + if accounted { + name = "accounted" + } + b.Run(name, func(b *testing.B) { + proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) + defer proc.Free() + input := testutil.NewBatch( + []types.Type{types.T_int32.ToType(), types.T_varchar.ToType()}, + true, + colexec.DefaultBatchSize, + proc.Mp(), + ) + defer input.Clean(proc.Mp()) + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err := budget.OpenGeneration(1) + if err != nil { + b.Fatal(err) + } + var hb HashmapBuilder + hb.setBudget(generation) + var ( + registry *mpool.AllocationAccountRegistry + account *mpool.AllocationAccount + op HashBuild + ) + if accounted { + registry, err = mpool.NewAllocationAccountRegistry(1, 64) + if err != nil { + b.Fatal(err) + } + account, err = registry.OpenWithController(capBytes, generation) + if err != nil { + b.Fatal(err) + } + op.NeedHashMap = true + if err = op.SetAllocationAccount(account); err != nil { + b.Fatal(err) + } + hb.batchAllocation = op.ctr.hashmapBuilder.batchAllocation + } + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if err = hb.copyBuildBatch(input, proc); err != nil { + b.Fatal(err) + } + hb.cleanBatches(proc) + } + b.StopTimer() + if generation.Used() != 0 { + b.Fatalf("generation used = %d", generation.Used()) + } + if accounted { + op.ctr.hashmapBuilder.batchAllocation = nil + if err = op.ClearAllocationAccount(account); err != nil { + b.Fatal(err) + } + if _, _, err = registry.CompleteTerminal(account); err != nil { + b.Fatal(err) + } + } + }) + } +} + func TestExtractRestoreCachedIterators(t *testing.T) { var hb HashmapBuilder mp := mpool.MustNewZero() diff --git a/pkg/sql/colexec/hashbuild/types.go b/pkg/sql/colexec/hashbuild/types.go index 8485750b70191..83139dca12c30 100644 --- a/pkg/sql/colexec/hashbuild/types.go +++ b/pkg/sql/colexec/hashbuild/types.go @@ -53,6 +53,11 @@ const HashBuildAllocationOwner mpool.AllocationOwner = 1 const ( HashBuildAllocationSiteHashCell mpool.AllocationSite = iota + 24 HashBuildAllocationSiteHashDescriptor + HashBuildAllocationSiteBatchData + HashBuildAllocationSiteBatchArea + HashBuildAllocationSiteBatchNulls + HashBuildAllocationSiteBatchGrouping + HashBuildAllocationSiteGroupSels ) type container struct { @@ -302,8 +307,20 @@ func (hashBuild *HashBuild) SetAllocationAccount( if err != nil { return err } + batchSelection, err := vector.NewAllocationAccountSelectionWithBitmaps( + account, + HashBuildAllocationOwner, + HashBuildAllocationSiteBatchData, + HashBuildAllocationSiteBatchArea, + HashBuildAllocationSiteBatchNulls, + HashBuildAllocationSiteBatchGrouping, + ) + if err != nil { + return err + } builder.mapAllocationAccount = account builder.mapAllocation = selection + builder.batchAllocation = batchSelection return nil } @@ -317,11 +334,13 @@ func (hashBuild *HashBuild) ClearAllocationAccount( if builder.mapAllocationAccount != account { return mpool.ErrAllocationAccountMismatch } - if builder.IntHashMap != nil || builder.StrHashMap != nil { + if builder.IntHashMap != nil || builder.StrHashMap != nil || + len(builder.Batches.Buf) != 0 || builder.Sels.Size() != 0 { return mpool.ErrAllocationAccountInvariant } builder.mapAllocationAccount = nil builder.mapAllocation = nil + builder.batchAllocation = nil return nil } diff --git a/pkg/sql/colexec/join_util.go b/pkg/sql/colexec/join_util.go index fcd63a9099f90..73c973e5ba69f 100644 --- a/pkg/sql/colexec/join_util.go +++ b/pkg/sql/colexec/join_util.go @@ -62,6 +62,18 @@ func (bs *Batches) Reset() { // the batches structure hold data in fix size 8192 rows, and continue to append from next batch // if error return , the batches will clean itself func (bs *Batches) CopyIntoBatches(src *batch.Batch, proc *process.Process) (err error) { + return bs.CopyIntoBatchesWithAllocation(src, proc, nil) +} + +// CopyIntoBatchesWithAllocation selects provenance for every retained vector +// destination. The Go descriptors are bounded by one Batch per 8,192 rows and +// one Vector pointer per input column; physical data, area, null, and grouping +// buffers are allocation-accounted and remain owned by the copied batches. +func (bs *Batches) CopyIntoBatchesWithAllocation( + src *batch.Batch, + proc *process.Process, + selection *vector.AllocationAccountSelection, +) (err error) { defer func() { if err != nil { bs.Clean(proc.Mp()) @@ -71,10 +83,21 @@ func (bs *Batches) CopyIntoBatches(src *batch.Batch, proc *process.Process) (err if bs.Buf == nil { bs.Buf = make([]*batch.Batch, 0, 16) } + if len(bs.Buf) > 0 && + bs.Buf[len(bs.Buf)-1].AllocationAccountSelection() != selection { + return mpool.ErrAllocationAccountMismatch + } var tmp *batch.Batch if src.RowCount() == DefaultBatchSize { - tmp, err = src.Dup(proc.Mp()) + if selection == nil { + tmp, err = src.Dup(proc.Mp()) + } else { + tmp, err = proc.NewBatchFromSrcWithAllocation(src, 0, selection) + if err == nil { + err = src.CloneTo(tmp, proc.Mp()) + } + } if err != nil { return err } @@ -96,12 +119,19 @@ func (bs *Batches) CopyIntoBatches(src *batch.Batch, proc *process.Process) (err lenBuf := len(bs.Buf) if lenBuf > 0 && bs.Buf[lenBuf-1].RowCount() != DefaultBatchSize { tmp = bs.Buf[lenBuf-1] + if tmp.AllocationAccountSelection() != selection { + return mpool.ErrAllocationAccountMismatch + } } else { preAllocSize := length - offset if preAllocSize > DefaultBatchSize { preAllocSize = DefaultBatchSize } - tmp, err = proc.NewBatchFromSrc(src, preAllocSize) + tmp, err = proc.NewBatchFromSrcWithAllocation( + src, + preAllocSize, + selection, + ) if err != nil { return err } diff --git a/pkg/vm/message/group_sels_test.go b/pkg/vm/message/group_sels_test.go index d6b3aa5132711..91940adb45ec2 100644 --- a/pkg/vm/message/group_sels_test.go +++ b/pkg/vm/message/group_sels_test.go @@ -112,3 +112,66 @@ func TestGroupSels_NullsSkipped(t *testing.T) { require.ElementsMatch(t, []int32{3}, js.Get(2)) js.Free(mp) } + +func TestGroupSelsAllocationAccountLifecycleAndRollback(t *testing.T) { + const ( + owner mpool.AllocationOwner = 1 + site mpool.AllocationSite = 30 + ) + for _, tc := range []struct { + name string + limit uint64 + metadataSlots uint64 + wantErr error + }{ + {name: "exact", limit: 64, metadataSlots: 3}, + { + name: "one byte short", + limit: 63, + metadataSlots: 3, + wantErr: mpool.ErrAllocationAccountCapacity, + }, + { + name: "one metadata slot short", + limit: 64, + metadataSlots: 2, + wantErr: mpool.ErrAllocationMetadataSlots, + }, + } { + t.Run(tc.name, func(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, tc.metadataSlots) + require.NoError(t, err) + account, err := registry.Open(tc.limit) + require.NoError(t, err) + mp := testMp() + var sels GroupSels + require.NoError(t, sels.InitWithAllocation(4, mp, account, owner, site)) + sels.Insert(0, 0) + sels.Insert(0, 1) + sels.Insert(1, 2) + sels.Insert(1, 3) + + err = sels.Finalize(2, 4, mp) + if tc.wantErr == nil { + require.NoError(t, err) + require.Equal(t, uint64(64), account.Snapshot().Peak) + require.Equal(t, uint64(32), account.Snapshot().Used) + require.ElementsMatch(t, []int32{0, 1}, sels.Get(0)) + require.ElementsMatch(t, []int32{2, 3}, sels.Get(1)) + } else { + require.ErrorIs(t, err, tc.wantErr) + require.Nil(t, sels.offsets) + require.Nil(t, sels.vals) + require.NotNil(t, sels.tmp) + require.Equal(t, uint64(32), account.Snapshot().Used) + } + sels.Free(mp) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + snapshot, first, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.True(t, first) + require.Equal(t, mpool.AllocationAccountTerminalValid, snapshot.State) + }) + } +} diff --git a/pkg/vm/message/joinMapMsg.go b/pkg/vm/message/joinMapMsg.go index 3e0e1ed1fb2cb..ed3864b2d06e8 100644 --- a/pkg/vm/message/joinMapMsg.go +++ b/pkg/vm/message/joinMapMsg.go @@ -17,6 +17,7 @@ package message import ( "bytes" "context" + "math" "os" "strconv" "sync" @@ -39,6 +40,10 @@ type GroupSels struct { // tmp holds (groupID, rowID) pairs during the build phase, before Finalize. tmp []int32 + + account *mpool.AllocationAccount + owner mpool.AllocationOwner + site mpool.AllocationSite } func freeSlice[T any](mp *mpool.MPool, s []T) { @@ -46,15 +51,64 @@ func freeSlice[T any](mp *mpool.MPool, s []T) { } func (sels *GroupSels) Init(n int, mp *mpool.MPool) error { + return sels.InitWithAllocation(n, mp, nil, 0, 0) +} + +// InitWithAllocation makes the complete temporary/final row-index owner use +// one immutable allocation generation. GroupSels is copied into JoinMap at +// publication, so its physical slices retain this provenance until the last +// consumer frees the map. +func (sels *GroupSels) InitWithAllocation( + n int, + mp *mpool.MPool, + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, + site mpool.AllocationSite, +) error { + if n < 0 || n > math.MaxInt/2 { + return mpool.ErrAllocationAccountInvalid + } + if sels.tmp != nil || sels.vals != nil || sels.offsets != nil { + return mpool.ErrAllocationAccountInvariant + } var err error - sels.tmp, err = mpool.MakeSlice[int32](n*2, mp, false) + if account == nil { + if owner != 0 || site != 0 { + return mpool.ErrAllocationAccountInvalid + } + sels.tmp, err = mpool.MakeSlice[int32](n*2, mp, false) + } else { + sels.tmp, err = mpool.MakeSliceAccounted[int32]( + n*2, + mp, + account, + owner, + site, + ) + } if err != nil { return err } + sels.account = account + sels.owner = owner + sels.site = site sels.tmp = sels.tmp[:0] return nil } +func (sels *GroupSels) makeSlice(n int, mp *mpool.MPool) ([]int32, error) { + if sels.account == nil { + return mpool.MakeSlice[int32](n, mp, false) + } + return mpool.MakeSliceAccounted[int32]( + n, + mp, + sels.account, + sels.owner, + sels.site, + ) +} + func (sels *GroupSels) Free(mp *mpool.MPool) { if mp != nil { freeSlice(mp, sels.vals) @@ -64,6 +118,9 @@ func (sels *GroupSels) Free(mp *mpool.MPool) { sels.vals = nil sels.offsets = nil sels.tmp = nil + sels.account = nil + sels.owner = 0 + sels.site = 0 } func (sels *GroupSels) Size() int64 { @@ -107,8 +164,7 @@ func (sels *GroupSels) Finalize(groupCount int, inputRowCount int, mp *mpool.MPo } } // groupCount+2: +1 for sentinel, +1 for 1-based callers (dedup UPDATE uses keys 1..groupCount) - var err error - sels.offsets, err = mpool.MakeSlice[int32](groupCount+2, mp, false) + offsets, err := sels.makeSlice(groupCount+2, mp) if err != nil { return err } @@ -116,26 +172,29 @@ func (sels *GroupSels) Finalize(groupCount int, inputRowCount int, mp *mpool.MPo // count occurrences per group for i := 0; i < len(sels.tmp); i += 2 { k := sels.tmp[i] - sels.offsets[k+1]++ + offsets[k+1]++ } // prefix sum - for i := int32(1); i < int32(len(sels.offsets)); i++ { - sels.offsets[i] += sels.offsets[i-1] + for i := int32(1); i < int32(len(offsets)); i++ { + offsets[i] += offsets[i-1] } // scatter vals using offsets as write cursors, then recover - sels.vals, err = mpool.MakeSlice[int32](n, mp, false) + vals, err := sels.makeSlice(n, mp) if err != nil { + freeSlice(mp, offsets) return err } for i := 0; i < len(sels.tmp); i += 2 { k := sels.tmp[i] v := sels.tmp[i+1] - sels.vals[sels.offsets[k]] = v - sels.offsets[k]++ + vals[offsets[k]] = v + offsets[k]++ } // recover offsets: shift right by one - copy(sels.offsets[1:], sels.offsets[:len(sels.offsets)-1]) - sels.offsets[0] = 0 + copy(offsets[1:], offsets[:len(offsets)-1]) + offsets[0] = 0 + sels.vals = vals + sels.offsets = offsets freeSlice(mp, sels.tmp) sels.tmp = nil return nil diff --git a/pkg/vm/process/hashbuild_budget.go b/pkg/vm/process/hashbuild_budget.go index dc0db22802451..f8283a30ff50f 100644 --- a/pkg/vm/process/hashbuild_budget.go +++ b/pkg/vm/process/hashbuild_budget.go @@ -34,6 +34,11 @@ const hashBuildMinimumReserve = uint64(4 << 30) const ( hashBuildAllocationGenerationSlots = uint32(131_072) hashBuildMinimumCellBlockBytes = uint64(16 << 10) + // Three slots close the cell/descriptor replacement transaction. The + // copied-batch activation adds up to three vector buffers per minimum-width + // 8,192-row destination (data, nulls, grouping), so six slots per 16 KiB of + // aggregate capacity is the first combined owner bound. + hashBuildAllocationSlotsPerBlock = uint64(6) ) const ( @@ -1136,8 +1141,9 @@ func (g *HashBuildBudgetGeneration) Closed() bool { // AllocationAccountRegistry returns the bounded CN-local registry shared by // every activated HashBuild generation under this aggregate budget. The slot -// formula covers every minimum-size cell block, its published descriptor, and -// one private replacement descriptor/block transaction. +// formula covers every minimum-size cell block, its published descriptor, one +// private replacement transaction, and the copied-batch vector buffers added +// by the second activation. func (g *HashBuildBudgetGeneration) AllocationAccountRegistry() ( *mpool.AllocationAccountRegistry, error, @@ -1155,14 +1161,14 @@ func (g *HashBuildBudgetGeneration) AllocationAccountRegistry() ( if blocks == 0 { blocks = 1 } - if blocks > math.MaxUint64/3 { + if blocks > math.MaxUint64/hashBuildAllocationSlotsPerBlock { b.allocationRegistryErr = ErrHashBuildBudgetInvalid return } b.allocationRegistry, b.allocationRegistryErr = mpool.NewAllocationAccountRegistry( hashBuildAllocationGenerationSlots, - blocks*3, + blocks*hashBuildAllocationSlotsPerBlock, ) }) return b.allocationRegistry, b.allocationRegistryErr diff --git a/pkg/vm/process/hashbuild_budget_test.go b/pkg/vm/process/hashbuild_budget_test.go index a3e9e8f7ee027..e06cb37be6b47 100644 --- a/pkg/vm/process/hashbuild_budget_test.go +++ b/pkg/vm/process/hashbuild_budget_test.go @@ -238,8 +238,12 @@ func TestHashBuildAllocationAccountRegistryUsesBoundedFormula(t *testing.T) { if registry.GenerationCapacity() != hashBuildAllocationGenerationSlots { t.Fatalf("generation slots = %d", registry.GenerationCapacity()) } - if registry.MaxAllocationMetadata() != 3 { - t.Fatalf("allocation slots = %d, want 3", registry.MaxAllocationMetadata()) + if registry.MaxAllocationMetadata() != hashBuildAllocationSlotsPerBlock { + t.Fatalf( + "allocation slots = %d, want %d", + registry.MaxAllocationMetadata(), + hashBuildAllocationSlotsPerBlock, + ) } second, err := budget.OpenGeneration(2) if err != nil { diff --git a/pkg/vm/process/process.go b/pkg/vm/process/process.go index 65f83c90ff6d7..65fa44429e540 100644 --- a/pkg/vm/process/process.go +++ b/pkg/vm/process/process.go @@ -225,18 +225,46 @@ func (proc *Process) AllocVectorOfRows(typ types.Type, nele int, nsp *nulls.Null } func (proc *Process) NewBatchFromSrc(src *batch.Batch, preAllocSize int) (*batch.Batch, error) { + return proc.NewBatchFromSrcWithAllocation(src, preAllocSize, nil) +} + +// NewBatchFromSrcWithAllocation creates an empty off-heap destination whose +// first vector growth uses the supplied immutable allocation provenance. +func (proc *Process) NewBatchFromSrcWithAllocation( + src *batch.Batch, + preAllocSize int, + selection *vector.AllocationAccountSelection, +) (_ *batch.Batch, retErr error) { + if proc == nil || src == nil || preAllocSize < 0 { + return nil, mpool.ErrAllocationAccountInvalid + } bat := batch.NewOffHeapWithSize(len(src.Vecs)) + defer func() { + if retErr != nil { + bat.Clean(proc.Mp()) + } + }() bat.SetAttributes(src.Attrs) bat.Recursive = src.Recursive for i := range bat.Vecs { - v := vector.NewOffHeapVecWithType(*src.Vecs[i].GetType()) + if src.Vecs[i] == nil { + return nil, mpool.ErrAllocationAccountInvalid + } + bat.Vecs[i] = vector.NewOffHeapVecWithType(*src.Vecs[i].GetType()) + } + if selection != nil { + if err := bat.SetAllocationAccount(selection); err != nil { + return nil, err + } + } + for i := range bat.Vecs { + v := bat.Vecs[i] if v.Capacity() < preAllocSize { err := v.PreExtend(preAllocSize, proc.Mp()) if err != nil { return nil, err } } - bat.Vecs[i] = v } return bat, nil } From 8a71c16252b48ea764185cfd47877763c5fcefe6 Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 21:58:58 +0800 Subject: [PATCH 17/61] executor: activate accounted hash-build expressions --- pkg/common/mpool/mpool.go | 14 +- pkg/container/vector/allocation_account.go | 5 + .../colexec/hashbuild/expression_memory.go | 102 ++++ .../hashbuild/expression_memory_test.go | 457 ++++++++++++++++++ pkg/sql/colexec/hashbuild/hashmap.go | 27 +- pkg/sql/colexec/hashbuild/hashmap_test.go | 120 +++++ pkg/sql/colexec/hashbuild/spill.go | 32 +- pkg/sql/colexec/hashbuild/types.go | 15 +- pkg/sql/plan/function/func_builtin.go | 30 +- pkg/sql/plan/function/func_cast.go | 6 +- 10 files changed, 780 insertions(+), 28 deletions(-) diff --git a/pkg/common/mpool/mpool.go b/pkg/common/mpool/mpool.go index 4e88b54b1f070..0981aae300bd3 100644 --- a/pkg/common/mpool/mpool.go +++ b/pkg/common/mpool/mpool.go @@ -922,14 +922,24 @@ func (mp *MPool) allocAccounted( }() if err = request.account.acquire(uint64(sz)); err != nil { - return nil, err + return nil, fmt.Errorf( + "allocation owner=%d site=%d: %w", + request.owner, + request.site, + err, + ) } accountHeld = true if err = request.reach(allocationAfterAccount); err != nil { return nil, err } if err = request.account.registry.reserveMetadata(); err != nil { - return nil, err + return nil, fmt.Errorf( + "allocation owner=%d site=%d: %w", + request.owner, + request.site, + err, + ) } metadataHeld = true if err = request.reach(allocationAfterMetadata); err != nil { diff --git a/pkg/container/vector/allocation_account.go b/pkg/container/vector/allocation_account.go index 0daaf8151d543..057c0910e57d8 100644 --- a/pkg/container/vector/allocation_account.go +++ b/pkg/container/vector/allocation_account.go @@ -326,6 +326,11 @@ func (v *Vector) ensureBitmapCapacity(rows int, mp *mpool.MPool) error { if rows > 0 { rows++ } + requiredWords := (rows + 63) / 64 + if requiredWords <= v.nsp.GetBitmap().ExternalStorageCapacity() && + requiredWords <= v.gsp.GetBitmap().ExternalStorageCapacity() { + return nil + } nulls, err := v.allocateBitmapGrowth( v.nsp.GetBitmap(), rows, diff --git a/pkg/sql/colexec/hashbuild/expression_memory.go b/pkg/sql/colexec/hashbuild/expression_memory.go index 99dd1cd2134ce..03a24a1e806b7 100644 --- a/pkg/sql/colexec/hashbuild/expression_memory.go +++ b/pkg/sql/colexec/hashbuild/expression_memory.go @@ -23,6 +23,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/plan/function" "github.com/matrixorigin/matrixone/pkg/vm/process" ) @@ -54,6 +55,107 @@ type ExpressionMemoryLease struct { released bool } +// NewAllocationAccountedExpressionExecutors constructs only expression trees +// whose complete retained and call-scoped allocation ledger is closed. The +// exact MPool leases are the sole capacity charge; unsupported function +// families continue through NewBudgetedExpressionExecutors until their own +// scratch owner is migrated. +func NewAllocationAccountedExpressionExecutors( + proc *process.Process, + exprs []*plan.Expr, + allocation *colexec.ExpressionAllocationAccount, +) ([]colexec.ExpressionExecutor, error) { + if allocation == nil || !expressionSetAllocationClosed(exprs) { + return nil, process.ErrHashBuildBudgetInvalid + } + return colexec.NewExpressionExecutorsFromPlanExpressionsWithAllocation( + proc, + exprs, + allocation, + ) +} + +func expressionSetAllocationClosed(exprs []*plan.Expr) bool { + for _, expr := range exprs { + if !expressionAllocationClosed(expr) { + return false + } + } + return true +} + +func expressionAllocationClosed(expr *plan.Expr) bool { + if expr == nil { + return false + } + switch node := expr.Expr.(type) { + case *plan.Expr_Col, *plan.Expr_Lit, *plan.Expr_T, + *plan.Expr_P, *plan.Expr_V, *plan.Expr_Vec, *plan.Expr_Fold: + return true + case *plan.Expr_F: + if node.F == nil || node.F.Func == nil { + return false + } + // Keep this as an implementation audit list, not a semantic function + // list. CONCAT writes directly into admitted result storage, CASE owns + // its row selections through ExpressionAllocationAccount, and varchar + // equality has no row-scaled scratch. Integer string casts use a + // stack-backed formatter; inserted casts of literals are plan-bounded. + functionID, _ := function.DecodeOverloadID(node.F.Func.Obj) + switch functionID { + case function.CONCAT, function.CASE: + case function.EQUAL: + if !closedHashBuildEqual(node.F.Args) { + return false + } + case function.CAST: + if !closedHashBuildCast(expr, node.F.Args) { + return false + } + default: + return false + } + for _, arg := range node.F.Args { + if !expressionAllocationClosed(arg) { + return false + } + } + return true + default: + return false + } +} + +func closedHashBuildEqual(args []*plan.Expr) bool { + if len(args) != 2 || args[0] == nil || args[1] == nil { + return false + } + for _, arg := range args { + oid := types.T(arg.Typ.Id) + if oid != types.T_char && oid != types.T_varchar { + return false + } + } + return true +} + +func closedHashBuildCast(result *plan.Expr, args []*plan.Expr) bool { + if result == nil || len(args) == 0 || args[0] == nil { + return false + } + source := types.T(args[0].Typ.Id) + target := types.T(result.Typ.Id) + if !source.ToType().IsIntOrUint() { + if (source == types.T_char || source == types.T_varchar) && + (target == types.T_char || target == types.T_varchar) { + _, literal := args[0].Expr.(*plan.Expr_Lit) + return literal + } + return false + } + return target == types.T_char || target == types.T_varchar +} + // NewBudgetedExpressionExecutors admits the mpool-backed capacity owned by // constant children before constructing them. The returned lease adopts those // reservations, so construction and later evaluation have one continuous diff --git a/pkg/sql/colexec/hashbuild/expression_memory_test.go b/pkg/sql/colexec/hashbuild/expression_memory_test.go index d5c5577e71740..b095f6186165c 100644 --- a/pkg/sql/colexec/hashbuild/expression_memory_test.go +++ b/pkg/sql/colexec/hashbuild/expression_memory_test.go @@ -58,6 +58,463 @@ func makeExpressionLeaseTestBatch(proc *process.Process, rows int) *batch.Batch return bat } +func makeIssue26454ConcatKey(t testing.TB, proc *process.Process) *plan.Expr { + t.Helper() + cast := func(colPos int32) *plan.Expr { + col := &plan.Expr{ + Typ: plan.Type{Id: int32(types.T_int32)}, + Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: colPos}}, + } + targetType := plan.Type{ + Id: int32(types.T_varchar), + Width: types.MaxVarcharLen, + } + expr, err := plan2.BindFuncExprImplByPlanExpr( + proc.Ctx, + "cast", + []*plan.Expr{ + col, + { + Typ: targetType, + Expr: &plan.Expr_T{T: &plan.TargetType{}}, + }, + }, + ) + require.NoError(t, err) + return expr + } + expr, err := plan2.BindFuncExprImplByPlanExpr( + proc.Ctx, + "concat", + []*plan.Expr{ + cast(0), + plan2.MakePlan2StringConstExprWithType("-"), + cast(1), + }, + ) + require.NoError(t, err) + return expr +} + +func makeIssue26454CaseKey(t testing.TB, proc *process.Process) *plan.Expr { + t.Helper() + column := &plan.Expr{ + Typ: plan.Type{ + Id: int32(types.T_varchar), + Width: types.MaxVarcharLen, + }, + Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, + } + condition, err := plan2.BindFuncExprImplByPlanExpr( + proc.Ctx, + "=", + []*plan.Expr{ + column, + plan2.MakePlan2StringConstExprWithType("ATM_CON"), + }, + ) + require.NoError(t, err) + expr, err := plan2.BindFuncExprImplByPlanExpr( + proc.Ctx, + "case", + []*plan.Expr{ + condition, + plan2.MakePlan2StringConstExprWithType("CON_CONTRACT_HEADERS"), + plan2.MakePlan2StringConstExprWithType("CON_CONTRACT_DOC"), + }, + ) + require.NoError(t, err) + return expr +} + +func TestAllocationAccountedExpressionIssue26454AndOneByteShort(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + expr := makeIssue26454ConcatKey(t, proc) + require.True(t, expressionSetAllocationClosed([]*plan.Expr{expr})) + require.False(t, expressionSetAllocationClosed( + []*plan.Expr{makeExpressionLeaseTestExpr(t, proc)}, + )) + input := batch.NewWithSize(2) + input.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp()) + input.Vecs[1] = testutil.MakeInt32Vector([]int32{3, 4}, nil, proc.Mp()) + input.SetRowCount(2) + defer input.Clean(proc.Mp()) + + run := func(limit uint64, verify bool) (uint64, error) { + budget := process.MustNewHashBuildBudget(1<<20, 1<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + allocation, err := colexec.NewExpressionAllocationAccount( + account, + HashBuildAllocationOwner, + ) + require.NoError(t, err) + executors, runErr := NewAllocationAccountedExpressionExecutors( + proc, + []*plan.Expr{expr}, + allocation, + ) + if runErr == nil { + var result *vector.Vector + result, runErr = executors[0].Eval( + proc, + []*batch.Batch{input}, + nil, + ) + if runErr == nil && verify { + require.Equal(t, []string{"1-3", "2-4"}, + vector.InefficientMustStrCol(result)) + } + } + peak := account.Snapshot().Peak + freeExpressionLeaseTestExecutors(executors) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + _, _, terminalErr := registry.CompleteTerminal(account) + require.NoError(t, terminalErr) + return peak, runErr + } + + peak, err := run(1<<20, true) + require.NoError(t, err) + require.Positive(t, peak) + _, err = run(peak-1, false) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Contains(t, err.Error(), "allocation owner=1 site=") +} + +func TestAllocationAccountedExpressionIssue26454CaseKey(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + expr := makeIssue26454CaseKey(t, proc) + require.True(t, expressionSetAllocationClosed([]*plan.Expr{expr})) + input := batch.NewWithSize(1) + input.Vecs[0] = testutil.MakeVarcharVector( + []string{"ATM_CON", "OTHER"}, + nil, + proc.Mp(), + ) + input.SetRowCount(2) + defer input.Clean(proc.Mp()) + + budget := process.MustNewHashBuildBudget(1<<20, 1<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.OpenWithController(1<<20, generation) + require.NoError(t, err) + allocation, err := colexec.NewExpressionAllocationAccount( + account, + HashBuildAllocationOwner, + ) + require.NoError(t, err) + executors, err := NewAllocationAccountedExpressionExecutors( + proc, + []*plan.Expr{expr}, + allocation, + ) + require.NoError(t, err) + result, err := executors[0].Eval(proc, []*batch.Batch{input}, nil) + require.NoError(t, err) + require.Equal(t, + []string{"CON_CONTRACT_HEADERS", "CON_CONTRACT_DOC"}, + vector.InefficientMustStrCol(result), + ) + require.Positive(t, account.Snapshot().Used) + freeExpressionLeaseTestExecutors(executors) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestAllocationAccountedExpressionRealValueOverCapRollsBack(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + column := &plan.Expr{ + Typ: plan.Type{ + Id: int32(types.T_varchar), + Width: types.MaxVarcharLen, + }, + Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, + } + expr, err := plan2.BindFuncExprImplByPlanExpr( + proc.Ctx, + "concat", + []*plan.Expr{column, plan2.MakePlan2StringConstExprWithType("-suffix")}, + ) + require.NoError(t, err) + require.True(t, expressionSetAllocationClosed([]*plan.Expr{expr})) + + const capBytes = uint64(4 << 10) + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.OpenWithController(capBytes, generation) + require.NoError(t, err) + allocation, err := colexec.NewExpressionAllocationAccount( + account, + HashBuildAllocationOwner, + ) + require.NoError(t, err) + executors, err := NewAllocationAccountedExpressionExecutors( + proc, + []*plan.Expr{expr}, + allocation, + ) + require.NoError(t, err) + + eval := func(value string) (*vector.Vector, error) { + input := batch.NewWithSize(1) + input.Vecs[0] = testutil.MakeVarcharVector([]string{value}, nil, proc.Mp()) + input.SetRowCount(1) + defer input.Clean(proc.Mp()) + return executors[0].Eval(proc, []*batch.Batch{input}, nil) + } + _, err = eval(strings.Repeat("x", 8<<10)) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Contains(t, err.Error(), "allocation owner=1 site=") + result, err := eval("ok") + require.NoError(t, err) + require.Equal(t, []string{"ok-suffix"}, vector.InefficientMustStrCol(result)) + + freeExpressionLeaseTestExecutors(executors) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestHashmapBuilderFallsBackOnlyForUnclosedExpressionScratch(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + for _, tc := range []struct { + name string + expr *plan.Expr + accounted bool + }{ + {name: "closed concat cast", expr: makeIssue26454ConcatKey(t, proc), accounted: true}, + {name: "closed case equality", expr: makeIssue26454CaseKey(t, proc), accounted: true}, + {name: "unclosed modulo", expr: makeExpressionLeaseTestExpr(t, proc)}, + } { + t.Run(tc.name, func(t *testing.T) { + budget := process.MustNewHashBuildBudget(16<<20, 16<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.OpenWithController(16<<20, generation) + require.NoError(t, err) + var op HashBuild + op.NeedHashMap = true + require.NoError(t, op.SetAllocationAccount(account)) + hb := &op.ctr.hashmapBuilder + hb.setBudget(generation) + require.NoError(t, hb.Prepare( + []*plan.Expr{tc.expr}, + -1, + -1, + nil, + proc, + )) + if tc.accounted { + require.Nil(t, hb.expressionLease) + require.Equal(t, generation.Used(), generation.Snapshot().AllocationUsed) + } else { + require.NotNil(t, hb.expressionLease) + } + hb.FreeExecutors() + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.NoError(t, op.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + }) + } +} + +func BenchmarkIssue26454ExpressionAccounting(b *testing.B) { + const capBytes = uint64(8 << 30) + proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) + defer proc.Free() + expr := makeIssue26454ConcatKey(b, proc) + input := testutil.NewBatch( + []types.Type{types.T_int32.ToType(), types.T_int32.ToType()}, + true, + colexec.DefaultBatchSize, + proc.Mp(), + ) + defer input.Clean(proc.Mp()) + + b.Run("legacy", func(b *testing.B) { + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err := budget.OpenGeneration(1) + if err != nil { + b.Fatal(err) + } + executors, lease, err := NewBudgetedExpressionExecutors( + proc, + generation, + []*plan.Expr{expr}, + false, + ) + if err != nil { + b.Fatal(err) + } + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if err = lease.Eval( + proc, + []*batch.Batch{input}, + input.RowCount(), + func(_ int, _ *vector.Vector) error { return nil }, + ); err != nil { + b.Fatal(err) + } + } + b.StopTimer() + freeExpressionLeaseTestExecutors(executors) + lease.Release() + if generation.Used() != 0 { + b.Fatalf("generation used = %d", generation.Used()) + } + }) + + b.Run("accounted", func(b *testing.B) { + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err := budget.OpenGeneration(1) + if err != nil { + b.Fatal(err) + } + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + if err != nil { + b.Fatal(err) + } + account, err := registry.OpenWithController(capBytes, generation) + if err != nil { + b.Fatal(err) + } + allocation, err := colexec.NewExpressionAllocationAccount( + account, + HashBuildAllocationOwner, + ) + if err != nil { + b.Fatal(err) + } + executors, err := NewAllocationAccountedExpressionExecutors( + proc, + []*plan.Expr{expr}, + allocation, + ) + if err != nil { + b.Fatal(err) + } + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if _, err = executors[0].Eval( + proc, + []*batch.Batch{input}, + nil, + ); err != nil { + b.Fatal(err) + } + } + b.StopTimer() + freeExpressionLeaseTestExecutors(executors) + if account.Snapshot().Used != 0 || generation.Used() != 0 { + b.Fatalf( + "live account=%d generation=%d", + account.Snapshot().Used, + generation.Used(), + ) + } + if _, _, err = registry.CompleteTerminal(account); err != nil { + b.Fatal(err) + } + }) +} + +func BenchmarkIssue26454CaseExpressionAccounting(b *testing.B) { + const capBytes = uint64(64 << 20) + proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) + defer proc.Free() + expr := makeIssue26454CaseKey(b, proc) + values := make([]string, colexec.DefaultBatchSize) + for i := range values { + if i%2 == 0 { + values[i] = "ATM_CON" + } else { + values[i] = "OTHER" + } + } + input := batch.NewWithSize(1) + input.Vecs[0] = testutil.MakeVarcharVector(values, nil, proc.Mp()) + input.SetRowCount(len(values)) + defer input.Clean(proc.Mp()) + + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err := budget.OpenGeneration(1) + if err != nil { + b.Fatal(err) + } + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + if err != nil { + b.Fatal(err) + } + account, err := registry.OpenWithController(capBytes, generation) + if err != nil { + b.Fatal(err) + } + allocation, err := colexec.NewExpressionAllocationAccount( + account, + HashBuildAllocationOwner, + ) + if err != nil { + b.Fatal(err) + } + executors, err := NewAllocationAccountedExpressionExecutors( + proc, + []*plan.Expr{expr}, + allocation, + ) + if err != nil { + b.Fatal(err) + } + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if _, err = executors[0].Eval( + proc, + []*batch.Batch{input}, + nil, + ); err != nil { + b.Fatal(err) + } + } + b.StopTimer() + freeExpressionLeaseTestExecutors(executors) + if account.Snapshot().Used != 0 || generation.Used() != 0 { + b.Fatalf( + "live account=%d generation=%d", + account.Snapshot().Used, + generation.Used(), + ) + } + if _, _, err = registry.CompleteTerminal(account); err != nil { + b.Fatal(err) + } +} + func makeMaxArrayLeaseTestVector[T types.ArrayElement]( t *testing.T, proc *process.Process, diff --git a/pkg/sql/colexec/hashbuild/hashmap.go b/pkg/sql/colexec/hashbuild/hashmap.go index 7f72692e241cb..9ca87ece559ca 100644 --- a/pkg/sql/colexec/hashbuild/hashmap.go +++ b/pkg/sql/colexec/hashbuild/hashmap.go @@ -92,6 +92,7 @@ type HashmapBuilder struct { mapAllocationAccount *mpool.AllocationAccount mapAllocation *hashtable.AllocationAccountSelection batchAllocation *vector.AllocationAccountSelection + expressionAllocation *colexec.ExpressionAllocationAccount } func (hb *HashmapBuilder) GetSize() int64 { @@ -185,12 +186,26 @@ func (hb *HashmapBuilder) Prepare( } keyWidth += width } - executors, expressionLease, err := NewBudgetedExpressionExecutors( - proc, - hb.budget, - keyCols, - needDupVec, + var ( + executors []colexec.ExpressionExecutor + expressionLease *ExpressionMemoryLease + err error ) + if hb.expressionAllocation != nil && + expressionSetAllocationClosed(keyCols) { + executors, err = NewAllocationAccountedExpressionExecutors( + proc, + keyCols, + hb.expressionAllocation, + ) + } else { + executors, expressionLease, err = NewBudgetedExpressionExecutors( + proc, + hb.budget, + keyCols, + needDupVec, + ) + } if err != nil { return err } @@ -247,6 +262,7 @@ func (hb *HashmapBuilder) Reset(proc *process.Process, hashTableHasNotSent bool) hb.mapAllocationAccount = nil hb.mapAllocation = nil hb.batchAllocation = nil + hb.expressionAllocation = nil } func (hb *HashmapBuilder) Free(proc *process.Process) { @@ -271,6 +287,7 @@ func (hb *HashmapBuilder) Free(proc *process.Process) { hb.mapAllocationAccount = nil hb.mapAllocation = nil hb.batchAllocation = nil + hb.expressionAllocation = nil } func (hb *HashmapBuilder) FreeExecutors() { diff --git a/pkg/sql/colexec/hashbuild/hashmap_test.go b/pkg/sql/colexec/hashbuild/hashmap_test.go index 3366ce427e66c..9f2409c69a213 100644 --- a/pkg/sql/colexec/hashbuild/hashmap_test.go +++ b/pkg/sql/colexec/hashbuild/hashmap_test.go @@ -1422,6 +1422,126 @@ func TestSpillExpressionHashKeyUsesBoundedAdmission(t *testing.T) { require.Zero(t, generation.Used()) } +func TestSpillExpressionUsesExactAccountForClosedKey(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + budget := process.MustNewHashBuildBudget(16<<20, 16<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.OpenWithController(16<<20, generation) + require.NoError(t, err) + var op HashBuild + op.NeedHashMap = true + require.NoError(t, op.SetAllocationAccount(account)) + ctr := &op.ctr + ctr.hashmapBuilder.setBudget(generation) + expr := makeIssue26454ConcatKey(t, proc) + executors, err := ctr.initSpillExprExecs(proc, []*plan.Expr{expr}) + require.NoError(t, err) + require.True(t, ctr.spillExprAccounted) + require.Nil(t, ctr.spillExprLease) + input := batch.NewWithSize(2) + input.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp()) + input.Vecs[1] = testutil.MakeInt32Vector([]int32{3, 4}, nil, proc.Mp()) + input.SetRowCount(2) + defer input.Clean(proc.Mp()) + result, err := executors[0].Eval(proc, []*batch.Batch{input}, nil) + require.NoError(t, err) + require.Equal(t, []string{"1-3", "2-4"}, vector.InefficientMustStrCol(result)) + require.Equal(t, generation.Used(), generation.Snapshot().AllocationUsed) + require.Positive(t, account.Snapshot().Used) + + ctr.freeSpillExprExecs() + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.NoError(t, op.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestIssue26454ExpressionKeyBuildUsesActualCapacity(t *testing.T) { + const capBytes = uint64(16 << 20) + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + for _, tc := range []struct { + name string + expr *plan.Expr + input func() *batch.Batch + }{ + { + name: "concat cast key", + expr: makeIssue26454ConcatKey(t, proc), + input: func() *batch.Batch { + return testutil.NewBatch( + []types.Type{types.T_int32.ToType(), types.T_int32.ToType()}, + true, + 10_000, + proc.Mp(), + ) + }, + }, + { + name: "case equality key", + expr: makeIssue26454CaseKey(t, proc), + input: func() *batch.Batch { + values := make([]string, 10_000) + for i := range values { + if i%2 == 0 { + values[i] = "ATM_CON" + } else { + values[i] = "OTHER" + } + } + bat := batch.NewWithSize(1) + bat.Vecs[0] = testutil.MakeVarcharVector(values, nil, proc.Mp()) + bat.SetRowCount(len(values)) + return bat + }, + }, + } { + t.Run(tc.name, func(t *testing.T) { + legacyPeak, err := expressionVectorPeak(proc, tc.expr, 10_000, false) + require.NoError(t, err) + require.Greater(t, legacyPeak, capBytes) + + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 128) + require.NoError(t, err) + account, err := registry.OpenWithController(capBytes, generation) + require.NoError(t, err) + var op HashBuild + op.NeedHashMap = true + require.NoError(t, op.SetAllocationAccount(account)) + hb := &op.ctr.hashmapBuilder + hb.setBudget(generation) + require.NoError(t, hb.Prepare([]*plan.Expr{tc.expr}, -1, -1, nil, proc)) + require.Nil(t, hb.expressionLease) + input := tc.input() + require.NoError(t, hb.copyBuildBatch(input, proc)) + hb.InputBatchRowCount = input.RowCount() + input.Clean(proc.Mp()) + require.NoError(t, hb.BuildHashmap(false, false, false, proc)) + require.LessOrEqual(t, generation.Used(), capBytes) + + jm := hb.GetJoinMap(proc.Mp()) + require.NotNil(t, jm) + jm.IncRef(1) + hb.Reset(proc, false) + jm.Free() + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + terminal, first, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.True(t, first) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) + }) + } +} + func TestExpressionHashKeyReservesDeclaredPeakBeforeEval(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() diff --git a/pkg/sql/colexec/hashbuild/spill.go b/pkg/sql/colexec/hashbuild/spill.go index c92bd356d3f51..320ad025e4faa 100644 --- a/pkg/sql/colexec/hashbuild/spill.go +++ b/pkg/sql/colexec/hashbuild/spill.go @@ -833,20 +833,37 @@ func (ctr *container) initSpillExprExecs(proc *process.Process, conditions []*pl return nil, &process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorInvalid, Message: "nil shuffle spill key"} } } - if len(ctr.spillExprExecs) != len(conditions) { - execs, lease, err := NewBudgetedExpressionExecutors( - proc, - ctr.hashmapBuilder.budget, - conditions, - false, + wantAccounted := ctr.hashmapBuilder.expressionAllocation != nil && + expressionSetAllocationClosed(conditions) + if len(ctr.spillExprExecs) != len(conditions) || + ctr.spillExprAccounted != wantAccounted { + var ( + execs []colexec.ExpressionExecutor + lease *ExpressionMemoryLease + err error ) + if wantAccounted { + execs, err = NewAllocationAccountedExpressionExecutors( + proc, + conditions, + ctr.hashmapBuilder.expressionAllocation, + ) + } else { + execs, lease, err = NewBudgetedExpressionExecutors( + proc, + ctr.hashmapBuilder.budget, + conditions, + false, + ) + } if err != nil { return nil, err } ctr.freeSpillExprExecs() ctr.spillExprExecs = execs ctr.spillExprLease = lease - } else if ctr.spillExprLease == nil { + ctr.spillExprAccounted = wantAccounted + } else if !ctr.spillExprAccounted && ctr.spillExprLease == nil { lease, err := NewExpressionMemoryLease( ctr.hashmapBuilder.budget, conditions, @@ -869,6 +886,7 @@ func (ctr *container) freeSpillExprExecs() { } } ctr.spillExprExecs = nil + ctr.spillExprAccounted = false if ctr.spillExprLease != nil { ctr.spillExprLease.Release() ctr.spillExprLease = nil diff --git a/pkg/sql/colexec/hashbuild/types.go b/pkg/sql/colexec/hashbuild/types.go index 83139dca12c30..59d24c58aeffe 100644 --- a/pkg/sql/colexec/hashbuild/types.go +++ b/pkg/sql/colexec/hashbuild/types.go @@ -108,6 +108,9 @@ type container struct { // cached expression executors for spill (reused across batches) spillExprExecs []colexec.ExpressionExecutor spillExprLease *ExpressionMemoryLease + // spillExprAccounted distinguishes an exact executor set from a legacy set + // whose retained lease has not yet been installed. + spillExprAccounted bool } // spillFileBundle is deliberately owned by hashbuild. Build converts each @@ -318,9 +321,17 @@ func (hashBuild *HashBuild) SetAllocationAccount( if err != nil { return err } + expressionAllocation, err := colexec.NewExpressionAllocationAccount( + account, + HashBuildAllocationOwner, + ) + if err != nil { + return err + } builder.mapAllocationAccount = account builder.mapAllocation = selection builder.batchAllocation = batchSelection + builder.expressionAllocation = expressionAllocation return nil } @@ -335,12 +346,14 @@ func (hashBuild *HashBuild) ClearAllocationAccount( return mpool.ErrAllocationAccountMismatch } if builder.IntHashMap != nil || builder.StrHashMap != nil || - len(builder.Batches.Buf) != 0 || builder.Sels.Size() != 0 { + len(builder.Batches.Buf) != 0 || builder.Sels.Size() != 0 || + len(builder.executors) != 0 || len(hashBuild.ctr.spillExprExecs) != 0 { return mpool.ErrAllocationAccountInvariant } builder.mapAllocationAccount = nil builder.mapAllocation = nil builder.batchAllocation = nil + builder.expressionAllocation = nil return nil } diff --git a/pkg/sql/plan/function/func_builtin.go b/pkg/sql/plan/function/func_builtin.go index da66af88d3aa7..ace53043744e0 100644 --- a/pkg/sql/plan/function/func_builtin.go +++ b/pkg/sql/plan/function/func_builtin.go @@ -799,25 +799,33 @@ func builtInConcat(parameters []*vector.Vector, result vector.FunctionResultWrap } for i := uint64(0); i < uint64(length); i++ { - var vs string - apv := true - + total := 0 + null := false for _, p := range ps { - v, null := p.GetStrValue(i) - if null { + v, isNull := p.GetStrValue(i) + if isNull { if err := rs.AppendBytes(nil, true); err != nil { return err } - apv = false + null = true break - } else { - vs += string(v) } + if len(v) > math.MaxInt-total { + return moerr.NewInvalidInputNoCtx("CONCAT result is too large") + } + total += len(v) } - if apv { - if err := rs.AppendBytes([]byte(vs), false); err != nil { - return err + if null { + continue + } + if err := rs.AppendBytesWithFill(total, func(dst []byte) { + offset := 0 + for _, p := range ps { + v, _ := p.GetStrValue(i) + offset += copy(dst[offset:], v) } + }); err != nil { + return err } } return nil diff --git a/pkg/sql/plan/function/func_cast.go b/pkg/sql/plan/function/func_cast.go index 30fc6d0a9d8eb..ca8bfe3a70853 100644 --- a/pkg/sql/plan/function/func_cast.go +++ b/pkg/sql/plan/function/func_cast.go @@ -3613,7 +3613,8 @@ func signedToStr[T constraints.Integer]( return err } } else { - result := []byte(strconv.FormatInt(int64(v), 10)) + var scratch [20]byte + result := strconv.AppendInt(scratch[:0], int64(v), 10) if toType.Oid == types.T_binary || toType.Oid == types.T_varbinary { if int32(len(result)) > toType.Width { return moerr.NewDataTruncatedNoCtx("Signed", " truncated for binary/varbinary") @@ -3662,7 +3663,8 @@ func unsignedToStr[T constraints.Unsigned]( return err } } else { - result := []byte(strconv.FormatUint(uint64(v), 10)) + var scratch [20]byte + result := strconv.AppendUint(scratch[:0], uint64(v), 10) if toType.Oid == types.T_binary || toType.Oid == types.T_varbinary { if int32(len(result)) > toType.Width { return moerr.NewDataTruncatedNoCtx("Unsigned", "truncated for binary/varbinary") From ca3b6e78b1d5106c4cf390334f4cfb8de40bc9e8 Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 22:26:23 +0800 Subject: [PATCH 18/61] executor: account spill and runtime-filter memory --- pkg/sql/colexec/dedupjoin/join.go | 12 +- pkg/sql/colexec/dedupjoin/types.go | 38 ++ pkg/sql/colexec/hashbuild/budget.go | 48 +- pkg/sql/colexec/hashbuild/build.go | 66 ++- pkg/sql/colexec/hashbuild/build_test.go | 121 +++- .../colexec/hashbuild/expression_memory.go | 7 + pkg/sql/colexec/hashbuild/hashmap.go | 31 +- pkg/sql/colexec/hashbuild/hashmap_test.go | 54 ++ pkg/sql/colexec/hashbuild/spill.go | 558 ++++++++++++++++-- pkg/sql/colexec/hashbuild/spill_test.go | 70 +++ pkg/sql/colexec/hashbuild/types.go | 112 +++- pkg/sql/colexec/hashjoin/join.go | 12 +- pkg/sql/colexec/hashjoin/types.go | 38 ++ pkg/sql/colexec/rightdedupjoin/join.go | 12 +- pkg/sql/colexec/rightdedupjoin/types.go | 3 +- .../spillutil/allocation_account_test.go | 223 +++++++ pkg/sql/colexec/spillutil/join_spill.go | 133 ++++- 17 files changed, 1443 insertions(+), 95 deletions(-) diff --git a/pkg/sql/colexec/dedupjoin/join.go b/pkg/sql/colexec/dedupjoin/join.go index c7a60ba00ef3e..6f30955c13dd9 100644 --- a/pkg/sql/colexec/dedupjoin/join.go +++ b/pkg/sql/colexec/dedupjoin/join.go @@ -327,7 +327,7 @@ func (dedupJoin *DedupJoin) build(analyzer process.Analyzer, proc *process.Proce return leaseErr } ctr.probeExpressionLease = probeExpressionLease - engine := spillutil.NewSpillEngine(spillutil.SpillEngineConfig{ + engine, engineErr := spillutil.NewSpillEngineForAccount(spillutil.SpillEngineConfig{ BuildKeyExprs: dedupJoin.Conditions[1], ProbeKeyExprs: dedupJoin.Conditions[0], SpillThreshold: ctr.spillThreshold, @@ -344,7 +344,15 @@ func (dedupJoin *DedupJoin) build(analyzer process.Analyzer, proc *process.Proce DedupDeleteKeepColIdxList: dedupJoin.DedupDeleteKeepColIdxList, Budget: budget, ProbeExpressionLease: probeExpressionLease, - }) + }, dedupJoin.allocationAccount, hashbuild.HashBuildAllocationOwner) + if engineErr != nil { + _ = payload.Close() + ctr.mp.Free() + ctr.mp = nil + ctr.cleanEvalVectors() + ctr.releaseProbeExpressionLease() + return engineErr + } if len(payload.Files) > 0 { engine.InitFromSpilledFiles(payload.Files) } else { diff --git a/pkg/sql/colexec/dedupjoin/types.go b/pkg/sql/colexec/dedupjoin/types.go index e307bfbef3f02..a8dad265652e4 100644 --- a/pkg/sql/colexec/dedupjoin/types.go +++ b/pkg/sql/colexec/dedupjoin/types.go @@ -21,6 +21,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/hashmap" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/reuse" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" @@ -298,10 +299,45 @@ type DedupJoin struct { // main-table scan path; empty for regular INSERT/UPDATE. OldColCapturePlaceholderIdxList []int32 OldColCaptureProbeIdxList []int32 + allocationAccount *mpool.AllocationAccount vm.OperatorBase } +func (dedupJoin *DedupJoin) AllocationAccountEnabled() bool { + return dedupJoin != nil +} + +func (dedupJoin *DedupJoin) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if account == nil || account.Handle() == 0 { + return mpool.ErrAllocationAccountInvalid + } + if dedupJoin.allocationAccount != nil && + dedupJoin.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + dedupJoin.allocationAccount = account + return nil +} + +func (dedupJoin *DedupJoin) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if dedupJoin.allocationAccount == nil { + return nil + } + if dedupJoin.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if dedupJoin.ctr.mp != nil || dedupJoin.ctr.spillEngine != nil { + return mpool.ErrAllocationAccountInvariant + } + dedupJoin.allocationAccount = nil + return nil +} + func (dedupJoin *DedupJoin) GetOperatorBase() *vm.OperatorBase { return &dedupJoin.OperatorBase } @@ -386,6 +422,7 @@ func (dedupJoin *DedupJoin) Reset(proc *process.Process, pipelineFailed bool, er ctr.roundStatusPublished = false ctr.state = Build ctr.lastPos = 0 + dedupJoin.allocationAccount = nil } func (dedupJoin *DedupJoin) Free(proc *process.Process, pipelineFailed bool, err error) { @@ -405,6 +442,7 @@ func (dedupJoin *DedupJoin) Free(proc *process.Process, pipelineFailed bool, err } ctr.cleanEvalVectors() ctr.releaseProbeExpressionLease() + dedupJoin.allocationAccount = nil } func (dedupJoin *DedupJoin) ExecProjection(proc *process.Process, input *batch.Batch) (*batch.Batch, error) { diff --git a/pkg/sql/colexec/hashbuild/budget.go b/pkg/sql/colexec/hashbuild/budget.go index 59dfc9fd59b9e..ae829052a1be0 100644 --- a/pkg/sql/colexec/hashbuild/budget.go +++ b/pkg/sql/colexec/hashbuild/budget.go @@ -635,7 +635,7 @@ func (hb *HashmapBuilder) buildAuxBytesWithUniqueProjection( // separately. Charging multiple whole-batch copies here double-counts those // owners and can reject a build before any auxiliary allocation occurs. bytes := batchesAllocated(hb.Batches.Buf) - if needUniqueVec { + if needUniqueVec && hb.uniqueKeyAllocation == nil { growthSlack := bytes / 4 if bytes%4 != 0 { growthSlack++ @@ -870,7 +870,7 @@ func unionBatchAreaBytes( } func (hb *HashmapBuilder) reserveUniqueAppendOverlap(dst *vector.Vector, rows, areaBytes int) (*process.HashBuildReservation, error) { - if hb.budget == nil { + if hb.budget == nil || hb.uniqueKeyAllocation != nil { return nil, nil } if dst == nil || rows < 0 || areaBytes < 0 { @@ -946,8 +946,48 @@ func (hb *HashmapBuilder) reserveUniqueAppendOverlap(dst *vector.Vector, rows, a return hb.budget.Reserve(overlap) } -func (hb *HashmapBuilder) marshalRuntimeFilterVector(vec *vector.Vector) ([]byte, func(), error) { - return runtimefilter.MarshalExactFilterVector(vec, hb.budget) +func (hb *HashmapBuilder) marshalRuntimeFilterVector( + vec *vector.Vector, + mp *mpool.MPool, +) ([]byte, func(), error) { + if vec == nil || vec.GetNulls().Any() { + return nil, nil, process.ErrHashBuildBudgetInvalid + } + if hb.mapAllocationAccount == nil { + return runtimefilter.MarshalExactFilterVector(vec, hb.budget) + } + if mp == nil { + return nil, nil, mpool.ErrAllocationAccountInvalid + } + size, err := vec.MarshalBinarySize() + if err != nil { + return nil, nil, err + } + buf, err := mpool.NewAccountedBuffer( + mp, + hb.mapAllocationAccount, + HashBuildAllocationOwner, + HashBuildAllocationSiteRuntimeFilterPayload, + ) + if err != nil { + return nil, nil, err + } + if err = buf.EnsureCapacity(size); err != nil { + buf.Free() + if mpool.IsRetryableAllocationCapacity(err) { + err = runtimefilter.MarkOptionalAllocationError(err) + } + return nil, nil, err + } + if err = vec.MarshalBinaryTo(buf); err != nil { + buf.Free() + return nil, nil, err + } + if buf.Len() != size { + buf.Free() + return nil, nil, process.ErrHashBuildBudgetInvalid + } + return buf.Bytes(), buf.Free, nil } func (hb *HashmapBuilder) releaseBatchReservations() { diff --git a/pkg/sql/colexec/hashbuild/build.go b/pkg/sql/colexec/hashbuild/build.go index bd919892ed6fc..467defadf1029 100644 --- a/pkg/sql/colexec/hashbuild/build.go +++ b/pkg/sql/colexec/hashbuild/build.go @@ -318,6 +318,49 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz // particular, a rejected retained-copy admission below must not add the // same upstream batch a second time when it is spilled directly. ctr.hashmapBuilder.InputBatchRowCount += result.Batch.RowCount() + if hashBuild.IsShuffle { + // First prove that the current upstream batch can always be spilled + // directly. This uses its actual materialization semantics and never + // projects a hypothetical retained batch. + var directProofErr error + if !spillMode || ctr.spillBatchAllocation == nil { + directProofErr = ctr.ensureDirectSpillScratchReservation( + result.Batch, + analyzer, + ) + } + if directProofErr != nil { + // Existing retained batches were admitted with a future-drain + // proof. Drain them under that lease, then retry the direct proof + // after their source reservations have been released. + if spillMode || + !errors.Is(directProofErr, process.ErrHashBuildBudgetAdmission) || + len(ctr.hashmapBuilder.Batches.Buf) == 0 { + return directProofErr + } + if err := startSpill(); err != nil { + return err + } + if ctr.spillBatchAllocation == nil { + if err := ctr.ensureDirectSpillScratchReservation(result.Batch, analyzer); err != nil { + return err + } + } + } + if !spillMode { + // A batch may become retained only after its future spill scratch + // is admitted. If that proof does not fit, do not copy it: switch + // to the already-proven direct-spill path. + if err := ctr.ensureRetainedSpillScratchReservation(result.Batch, analyzer); err != nil { + if !errors.Is(err, process.ErrHashBuildBudgetAdmission) { + return err + } + if err := startSpill(); err != nil { + return err + } + } + } + } // If in spill mode, spill this batch directly to open files. if spillMode { err := ctr.spillBatchBounded(proc, result.Batch, spillFiles, ctr.spillExprExecs, analyzer, false) @@ -749,7 +792,7 @@ func (hashBuild *HashBuild) handleRuntimeFilter( // build a membership filter, based on its own threshold. runtimeFilter.Typ = message.RuntimeFilter_UNIQUEJOINKEYS - data, release, err := ctr.hashmapBuilder.marshalRuntimeFilterVector(keyVec) + data, release, err := ctr.hashmapBuilder.marshalRuntimeFilterVector(keyVec, proc.Mp()) if err != nil { if hashBuild.fallbackOptionalRuntimeFilter(err, &runtimeFilter, spec, proc) { return nil @@ -850,7 +893,10 @@ func (hashBuild *HashBuild) handleRuntimeFilter( } keyVec.GetNulls().Reset() keyVec.InplaceSort() - data, release, err := ctr.hashmapBuilder.marshalRuntimeFilterVector(keyVec) + data, release, err := ctr.hashmapBuilder.marshalRuntimeFilterVector( + keyVec, + proc.Mp(), + ) if err != nil { if hashBuild.fallbackOptionalRuntimeFilter(err, &runtimeFilter, spec, proc) { return nil @@ -1071,7 +1117,10 @@ func (hashBuild *HashBuild) materializeSerializedRuntimeFilter( } payload.InplaceSort() data, release, err = - hashBuild.ctr.hashmapBuilder.marshalRuntimeFilterVector(payload) + hashBuild.ctr.hashmapBuilder.marshalRuntimeFilterVector( + payload, + proc.Mp(), + ) if err != nil { if release != nil { release() @@ -1209,6 +1258,17 @@ func (hashBuild *HashBuild) fallbackOptionalRuntimeFilter( } else { stats.AddExtraStat( "HashBuildRuntimeFilterAllocationFallbacks", 1) + if account := hashBuild.ctr.hashmapBuilder.mapAllocationAccount; account != nil { + snapshot := account.Snapshot() + stats.SetMaxExtraStat( + "HashBuildRuntimeFilterBudgetFallbackUsedBytes", + hashBuildStatInt64(snapshot.Used), + ) + stats.SetMaxExtraStat( + "HashBuildRuntimeFilterBudgetFallbackCapBytes", + hashBuildStatInt64(snapshot.Limit), + ) + } } } *runtimeFilter = message.RuntimeFilterMessage{ diff --git a/pkg/sql/colexec/hashbuild/build_test.go b/pkg/sql/colexec/hashbuild/build_test.go index b197d164d6a46..d34ec4f562a5e 100644 --- a/pkg/sql/colexec/hashbuild/build_test.go +++ b/pkg/sql/colexec/hashbuild/build_test.go @@ -2120,7 +2120,7 @@ func TestRuntimeFilterMarshalUsesSinglePayloadBudget(t *testing.T) { require.NoError(t, err) tc.arg.ctr.hashmapBuilder.setBudget(generation) - data, release, err := tc.arg.ctr.hashmapBuilder.marshalRuntimeFilterVector(vec) + data, release, err := tc.arg.ctr.hashmapBuilder.marshalRuntimeFilterVector(vec, tc.proc.Mp()) require.NoError(t, err) require.NotEmpty(t, data) require.Equal(t, projected, generation.Peak()) @@ -2152,7 +2152,7 @@ func TestRuntimeFilterMarshalSinglePayloadCoversVarlenaPeak(t *testing.T) { require.NoError(t, err) tc.arg.ctr.hashmapBuilder.setBudget(generation) - data, release, err := tc.arg.ctr.hashmapBuilder.marshalRuntimeFilterVector(vec) + data, release, err := tc.arg.ctr.hashmapBuilder.marshalRuntimeFilterVector(vec, tc.proc.Mp()) require.NoError(t, err) require.NotEmpty(t, data) require.Equal(t, projected, generation.Peak()) @@ -2166,6 +2166,123 @@ func TestRuntimeFilterMarshalSinglePayloadCoversVarlenaPeak(t *testing.T) { require.Zero(t, tc.proc.Mp().CurrNB()) } +func TestRuntimeFilterMarshalAccountedPayloadMessageLifecycle(t *testing.T) { + tc := newTestCase(t, []bool{true}, []types.Type{types.T_varchar.ToType()}, + []*plan.Expr{newExpr(0, types.T_varchar.ToType())}) + vec := testutil.MakeVarcharVector( + []string{strings.Repeat("x", 4<<10), strings.Repeat("y", 8<<10)}, + []uint64{1}, + tc.proc.Mp(), + ) + + const limit = uint64(1 << 20) + budget := process.MustNewHashBuildBudget(limit, limit) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + tc.arg.NeedHashMap = true + require.NoError(t, tc.arg.SetAllocationAccount(account)) + tc.arg.ctr.hashmapBuilder.setBudget(generation) + + data, release, err := tc.arg.ctr.hashmapBuilder.marshalRuntimeFilterVector( + vec, + tc.proc.Mp(), + ) + require.NoError(t, err) + require.NotEmpty(t, data) + require.NotNil(t, release) + snapshot := account.Snapshot() + require.Positive(t, snapshot.Used) + require.Equal(t, snapshot.Used, generation.Snapshot().AllocationUsed) + require.Equal(t, snapshot.Used, generation.Used()) + + spec := &plan.RuntimeFilterSpec{Tag: 103} + runtimeFilter := message.RuntimeFilterMessage{ + Tag: spec.Tag, + Typ: message.RuntimeFilter_IN, + Card: 2, + Data: data, + } + runtimeFilter.SetMemoryRelease(release) + message.SendRuntimeFilter(runtimeFilter, spec, tc.proc.GetMessageBoard()) + require.True(t, tc.proc.GetMessageBoard().CloseAndDrain()) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + vec.Free(tc.proc.Mp()) + + require.NoError(t, tc.arg.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + generation.Close() + tc.proc.Free() + require.Zero(t, tc.proc.Mp().CurrNB()) +} + +func TestRuntimeFilterMarshalAccountedOneByteShortFallsBackToPass(t *testing.T) { + tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, + []*plan.Expr{newExpr(0, types.T_int32.ToType())}) + vec := testutil.MakeInt32Vector([]int32{1, 2, 3, 4}, nil, tc.proc.Mp()) + size, err := vec.MarshalBinarySize() + require.NoError(t, err) + capacity, ok := mpool.GrowCapacity(0, int64(size)) + require.True(t, ok) + require.Positive(t, capacity) + vec.Free(tc.proc.Mp()) + + limit := uint64(capacity - 1) + budget := process.MustNewHashBuildBudget(limit, limit) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + tc.arg.NeedHashMap = true + require.NoError(t, tc.arg.SetAllocationAccount(account)) + tc.arg.ctr.hashmapBuilder.setBudget(generation) + tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{ + Tag: 104, + UpperLimit: 100, + Expr: newExpr(0, types.T_int32.ToType()), + } + tc.arg.OpAnalyzer = process.NewAnalyzer(0, false, false, "hash build") + tc.arg.ctr.hashmapBuilder.InputBatchRowCount = 4 + tc.arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{ + testutil.MakeInt32Vector([]int32{1, 2, 3, 4}, nil, tc.proc.Mp()), + } + + require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.Equal(t, int64(1), + tc.arg.OpAnalyzer.GetOpStats().ExtraStats["HashBuildRuntimeFilterBudgetFallbacks"]) + + receiver := message.NewMessageReceiver( + []int32{tc.arg.RuntimeFilterSpec.Tag}, + message.AddrBroadCastOnCurrentCN(), + tc.proc.GetMessageBoard(), + ) + msgs, done, err := receiver.ReceiveMessage(false, tc.proc.Ctx) + require.NoError(t, err) + require.False(t, done) + require.Len(t, msgs, 1) + runtimeFilter, ok := msgs[0].(message.RuntimeFilterMessage) + require.True(t, ok) + require.Equal(t, int32(message.RuntimeFilter_PASS), runtimeFilter.Typ) + require.Empty(t, runtimeFilter.Data) + + require.True(t, tc.proc.GetMessageBoard().CloseAndDrain()) + require.NoError(t, tc.arg.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + generation.Close() + tc.proc.Free() + require.Zero(t, tc.proc.Mp().CurrNB()) +} + func TestRuntimeFilterMarshalClosedBudgetRemainsFatal(t *testing.T) { tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) diff --git a/pkg/sql/colexec/hashbuild/expression_memory.go b/pkg/sql/colexec/hashbuild/expression_memory.go index 03a24a1e806b7..a9664ca57a2e0 100644 --- a/pkg/sql/colexec/hashbuild/expression_memory.go +++ b/pkg/sql/colexec/hashbuild/expression_memory.go @@ -84,6 +84,13 @@ func expressionSetAllocationClosed(exprs []*plan.Expr) bool { return true } +// AllocationAccountedExpressionSetSupported reports whether every execution +// path in the expression set has a closed physical-allocation ledger. Spill +// rebuild uses this gate before selecting the shared exact account. +func AllocationAccountedExpressionSetSupported(exprs []*plan.Expr) bool { + return expressionSetAllocationClosed(exprs) +} + func expressionAllocationClosed(expr *plan.Expr) bool { if expr == nil { return false diff --git a/pkg/sql/colexec/hashbuild/hashmap.go b/pkg/sql/colexec/hashbuild/hashmap.go index 9ca87ece559ca..851cafcbf2291 100644 --- a/pkg/sql/colexec/hashbuild/hashmap.go +++ b/pkg/sql/colexec/hashbuild/hashmap.go @@ -92,6 +92,7 @@ type HashmapBuilder struct { mapAllocationAccount *mpool.AllocationAccount mapAllocation *hashtable.AllocationAccountSelection batchAllocation *vector.AllocationAccountSelection + uniqueKeyAllocation *vector.AllocationAccountSelection expressionAllocation *colexec.ExpressionAllocationAccount } @@ -262,6 +263,7 @@ func (hb *HashmapBuilder) Reset(proc *process.Process, hashTableHasNotSent bool) hb.mapAllocationAccount = nil hb.mapAllocation = nil hb.batchAllocation = nil + hb.uniqueKeyAllocation = nil hb.expressionAllocation = nil } @@ -287,6 +289,7 @@ func (hb *HashmapBuilder) Free(proc *process.Process) { hb.mapAllocationAccount = nil hb.mapAllocation = nil hb.batchAllocation = nil + hb.uniqueKeyAllocation = nil hb.expressionAllocation = nil } @@ -802,6 +805,7 @@ func (hb *HashmapBuilder) buildHashmap( ignoreCandidateOldKey = make([]*vector.Vector, 1) } +buildUnits: for i := 0; i < hb.InputBatchRowCount; i += hashmap.UnitLimit { if i%(hashmap.UnitLimit*32) == 0 { if err := checkHashBuildCanceled(proc); err != nil { @@ -949,9 +953,30 @@ func (hb *HashmapBuilder) buildHashmap( if len(hb.UniqueJoinKeys) == 0 { hb.UniqueJoinKeys = make([]*vector.Vector, len(hb.executors)) for j, vec := range hb.curVecs { - if hb.collectUniqueKeySlot(j) { - hb.UniqueJoinKeys[j] = - vector.NewOffHeapVecWithType(*vec.GetType()) + if !hb.collectUniqueKeySlot(j) { + continue + } + if hb.uniqueKeyAllocation == nil { + hb.UniqueJoinKeys[j] = vector.NewOffHeapVecWithType(*vec.GetType()) + continue + } + hb.UniqueJoinKeys[j], err = vector.NewOffHeapVecWithTypeAndAllocation( + *vec.GetType(), + hb.uniqueKeyAllocation, + ) + if err != nil { + cause := err + if mpool.IsRetryableAllocationCapacity(err) { + cause = runtimefilter.MarkOptionalAllocationError(err) + } + if fatalErr := hb.fallbackOptionalRuntimeFilterCollection( + proc, + cause, + ); fatalErr != nil { + return fatalErr + } + needUniqueVec = false + continue buildUnits } } } diff --git a/pkg/sql/colexec/hashbuild/hashmap_test.go b/pkg/sql/colexec/hashbuild/hashmap_test.go index 9f2409c69a213..4dc5e128ee731 100644 --- a/pkg/sql/colexec/hashbuild/hashmap_test.go +++ b/pkg/sql/colexec/hashbuild/hashmap_test.go @@ -1291,6 +1291,60 @@ func TestUniqueAppendBudgetIncludesDeadAreaCopiedByUnionBatch(t *testing.T) { generation.Close() } +func TestAccountedRuntimeFilterUniqueKeysDegradeWithoutFailingHashBuild(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + values := make([]string, 1_024) + for i := range values { + values[i] = strconv.Itoa(i) + strings.Repeat(string(rune('a'+i%26)), 4<<10) + } + input := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.MakeVarcharVector(values, nil, proc.Mp()), + }, nil) + defer input.Clean(proc.Mp()) + exprs := []*plan.Expr{newExpr(0, types.T_varchar.ToType())} + + run := func(limit uint64, needUnique bool) (mpool.AllocationAccountSnapshot, bool) { + budget := process.MustNewHashBuildBudget(64<<20, 64<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 128) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + builder := &HashmapBuilder{} + builder.SetBudget(generation) + require.NoError(t, builder.SetAllocationAccount(account)) + require.NoError(t, builder.Prepare(exprs, -1, -1, nil, proc)) + require.NoError(t, builder.CopyBuildBatch(input, proc)) + builder.InputBatchRowCount = input.RowCount() + require.NoError(t, builder.BuildHashmap(false, false, needUnique, proc)) + snapshot := account.Snapshot() + fallback, _ := builder.runtimeFilterFallbackState() + if needUnique { + require.True(t, fallback) + require.Empty(t, builder.UniqueJoinKeys) + require.NotNil(t, builder.StrHashMap) + require.Equal(t, uint64(input.RowCount()), builder.StrHashMap.GroupCount()) + } + builder.Free(proc) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + return snapshot, fallback + } + + baseline, fallback := run(64<<20, false) + require.False(t, fallback) + require.Positive(t, baseline.Peak) + // The exact baseline peak is sufficient for the required hash build, but + // not for a second, optional copy of the 4 MiB runtime-filter key payload. + constrained, fallback := run(baseline.Peak, true) + require.True(t, fallback) + require.LessOrEqual(t, constrained.Peak, baseline.Peak) +} + func TestCleanCopiedBatchReleasesCoalescedIngressReservations(t *testing.T) { const budgetCap = uint64(4 << 20) budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) diff --git a/pkg/sql/colexec/hashbuild/spill.go b/pkg/sql/colexec/hashbuild/spill.go index 320ad025e4faa..fbed914771f9f 100644 --- a/pkg/sql/colexec/hashbuild/spill.go +++ b/pkg/sql/colexec/hashbuild/spill.go @@ -23,6 +23,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/hashmap/keycodec" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" @@ -42,6 +43,19 @@ const ( spillWriteCoalesceSize = 64 << 10 ) +type spillMaterializationMode uint8 + +const ( + // spillDirectMaterialization models UnionInt32 on the current upstream + // batch. A const varlen source copies its out-of-line payload once and + // broadcasts the resulting descriptor. + spillDirectMaterialization spillMaterializationMode = iota + // spillRetainedMaterialization models the compact non-const batch produced + // by CopyIntoBatches. A later UnionInt32 treats every retained row as an + // independent value, even when the ingress vector was const. + spillRetainedMaterialization +) + func spillCheckedAdd(total, value uint64) (uint64, error) { if total > math.MaxUint64-value { return 0, process.ErrHashBuildBudgetInvalid @@ -87,9 +101,13 @@ func spillCapacityReplacementOverlap(rows, keys, hashCap, rowIDCap, keyCap int) // spillMaterializedBytes models the batch that spillBatchBounded creates with // UnionInt32. It follows vector materialization semantics instead of retained // capacity or stale logical length: fixed-width descriptors are per output -// row, null payload is skipped, and const varlen payload is copied once. -func spillMaterializedBytes(bat *batch.Batch) (uint64, error) { - if bat == nil || bat.RowCount() <= 0 { +// row, null payload is skipped, and direct const varlen payload is copied once. +func spillMaterializedBytesFor( + bat *batch.Batch, + targetRows uint64, + mode spillMaterializationMode, +) (uint64, error) { + if bat == nil || bat.RowCount() <= 0 || targetRows == 0 { return 0, nil } liveRows := uint64(bat.RowCount()) @@ -102,7 +120,7 @@ func spillMaterializedBytes(bat *batch.Batch) (uint64, error) { if typeSize < 0 { return 0, process.ErrHashBuildBudgetInvalid } - descriptors, err := spillCheckedMul(liveRows, uint64(typeSize)) + descriptors, err := spillCheckedMul(targetRows, uint64(typeSize)) if err != nil { return 0, err } @@ -136,13 +154,40 @@ func spillMaterializedBytes(bat *batch.Batch) (uint64, error) { } } - if materialized, err = spillCheckedAdd(materialized, livePayload); err != nil { + projectedPayload := livePayload + if !(mode == spillDirectMaterialization && vec.IsConst()) { + // A retained CopyIntoBatches destination is non-const. Repeating + // the complete live sample is a conservative bound for any compact + // target batch assembled from ingress batches whose individual + // high-water estimates were admitted before copying. + roundedRows, err := spillCheckedAdd(targetRows, valueRows-1) + if err != nil { + return 0, err + } + repeats := roundedRows / valueRows + if projectedPayload, err = spillCheckedMul(livePayload, repeats); err != nil { + return 0, err + } + } + + if materialized, err = spillCheckedAdd(materialized, projectedPayload); err != nil { return 0, err } } return materialized, nil } +func spillMaterializedBytes(bat *batch.Batch) (uint64, error) { + if bat == nil || bat.RowCount() <= 0 { + return 0, nil + } + return spillMaterializedBytesFor( + bat, + uint64(bat.RowCount()), + spillDirectMaterialization, + ) +} + func spillMarshalSlack(columns uint64) (uint64, error) { const ( fixedSlack = uint64(64 << 10) @@ -207,7 +252,11 @@ func spillBudgetBytes(bat *batch.Batch) (uint64, error) { return 0, nil } rows := uint64(bat.RowCount()) - selected, err := spillMaterializedBytes(bat) + selected, err := spillMaterializedBytesFor( + bat, + rows, + spillDirectMaterialization, + ) if err != nil { return 0, err } @@ -239,6 +288,129 @@ func spillScratchBudgetBytes(bat *batch.Batch, sourceAlreadyCharged bool) (uint6 return need - source, nil } +// spillRetainedBudgetBytes is the future-drain proof required before +// CopyIntoBatches may retain a small input. The destination loses constness, +// so its selected payload follows retained rather than direct semantics. +func spillRetainedBudgetBytes(bat *batch.Batch) (uint64, error) { + if bat == nil || bat.RowCount() <= 0 { + return 0, nil + } + rows := uint64(bat.RowCount()) + targetRows := rows + if rows < uint64(colexec.DefaultBatchSize) { + targetRows = uint64(colexec.DefaultBatchSize) + } + selected, err := spillMaterializedBytesFor( + bat, + targetRows, + spillRetainedMaterialization, + ) + if err != nil { + return 0, err + } + metadata, ok := retainedMetadataAllowance(bat) + if !ok || metadata > math.MaxUint64/targetRows { + return 0, process.ErrHashBuildBudgetInvalid + } + projectedMetadata := metadata + if targetRows > rows { + projectedMetadata, err = spillCheckedMul(metadata, targetRows) + if err != nil { + return 0, err + } + projectedMetadata, err = spillCheckedAdd(projectedMetadata, rows-1) + if err != nil { + return 0, err + } + projectedMetadata /= rows + } + if selected, err = spillCheckedAdd(selected, projectedMetadata); err != nil { + return 0, err + } + materializationSlack, err := spillMaterializationSlack(uint64(len(bat.Vecs))) + if err != nil { + return 0, err + } + if selected, err = spillCheckedAdd(selected, materializationSlack); err != nil { + return 0, err + } + // The retained source itself is covered by batchReservations. + return spillPeakBudgetFor( + targetRows, + 0, + selected, + uint64(len(bat.Vecs)), + ) +} + +func (ctr *container) ensureSpillScratchReservationBytes( + need uint64, + analyzer process.Analyzer, +) error { + if ctr.hashmapBuilder.budget == nil || need == 0 { + return nil + } + var err error + if ctr.spillScratchReservation == nil { + ctr.spillScratchReservation, err = + ctr.hashmapBuilder.budget.Reserve(need) + if err == nil { + analyzer.GetOpStats().SetMaxExtraStat( + "HashBuildEmergencyScratchBytes", + hashBuildStatInt64(need), + ) + ctr.spillScratchEmergency = true + ctr.spillScratchBase = need + } + return err + } + if ctr.spillScratchBase >= need { + ctr.spillScratchEmergency = true + return nil + } + grow := need - ctr.spillScratchBase + if err := ctr.spillScratchReservation.Grow(grow); err != nil { + analyzer.GetOpStats().AddExtraStat( + "HashBuildEmergencyScratchGrowRejects", + 1, + ) + return err + } + analyzer.GetOpStats().AddExtraStat( + "HashBuildEmergencyScratchGrowCount", + 1, + ) + analyzer.GetOpStats().AddExtraStat( + "HashBuildEmergencyScratchGrowBytes", + hashBuildStatInt64(grow), + ) + ctr.spillScratchBase = need + ctr.spillScratchEmergency = true + return nil +} + +func (ctr *container) ensureDirectSpillScratchReservation( + bat *batch.Batch, + analyzer process.Analyzer, +) error { + need, err := spillBudgetBytes(bat) + if err != nil { + return err + } + return ctr.ensureSpillScratchReservationBytes(need, analyzer) +} + +func (ctr *container) ensureRetainedSpillScratchReservation( + bat *batch.Batch, + analyzer process.Analyzer, +) error { + need, err := spillRetainedBudgetBytes(bat) + if err != nil { + return err + } + return ctr.ensureSpillScratchReservationBytes(need, analyzer) +} + func (ctr *container) growSpillScratchTransient( required uint64, analyzer process.Analyzer, @@ -271,10 +443,29 @@ func (ctr *container) releaseSpillScratchReservation() { ctr.spillScratchReservation.Release() ctr.spillScratchReservation = nil } + ctr.spillScratchEmergency = false ctr.spillScratchBase = 0 } func (ctr *container) dropSpillScratchBuffers() { + if ctr.spillBatchAllocation != nil { + if cap(ctr.spillHashValues) > 0 { + mpool.FreeSlice(ctr.spillAllocationMP, ctr.spillHashValues) + } + if cap(ctr.spillBucketRowIds) > 0 { + mpool.FreeSlice(ctr.spillAllocationMP, ctr.spillBucketRowIds) + } + if ctr.spillAccountedWrite != nil { + ctr.spillAccountedWrite.Free() + ctr.spillAccountedWrite = nil + } + for i := range ctr.spillAccountedBuckets { + if ctr.spillAccountedBuckets[i] != nil { + ctr.spillAccountedBuckets[i].Free() + ctr.spillAccountedBuckets[i] = nil + } + } + } for bucket := range ctr.spillBucketWriteBufs { ctr.spillBucketWriteBufs[bucket] = bytes.Buffer{} ctr.spillBucketWriteRows[bucket] = 0 @@ -290,6 +481,48 @@ func (ctr *container) dropSpillScratchBuffers() { ctr.spillSelection = nil ctr.spillKeyVecs = nil ctr.spillWriteBuf = bytes.Buffer{} + ctr.spillAllocationMP = nil +} + +func growHashBuildSpillSlice[T any]( + values []T, + length int, + mp *mpool.MPool, + account *mpool.AllocationAccount, + site mpool.AllocationSite, +) ([]T, error) { + if length < 0 || account == nil { + return nil, mpool.ErrAllocationAccountInvalid + } + if length <= cap(values) { + return values[:length], nil + } + capacity := cap(values) + if capacity == 0 { + capacity = 1 + } + for capacity < length { + if capacity > math.MaxInt/2 { + capacity = length + break + } + capacity *= 2 + } + next, err := mpool.MakeSliceAccounted[T]( + capacity, + mp, + account, + HashBuildAllocationOwner, + site, + ) + if err != nil { + return nil, err + } + copy(next, values) + if cap(values) > 0 { + mpool.FreeSlice(mp, values) + } + return next[:length], nil } func spillMarshalGrowBytes(bat *batch.Batch) (uint64, error) { @@ -347,6 +580,49 @@ func marshalSpillRecord(bat *batch.Batch, buf *bytes.Buffer) (int64, error) { return cnt, nil } +func marshalSpillRecordAccounted( + bat *batch.Batch, + buf *mpool.AccountedBuffer, +) (int64, error) { + if bat == nil || bat.RowCount() == 0 || buf == nil { + return 0, nil + } + cnt := int64(bat.RowCount()) + buf.Reset() + batchSize, err := bat.MarshalBinarySize() + if err != nil || batchSize > math.MaxInt-24 { + if err != nil { + return 0, err + } + return 0, process.ErrHashBuildBudgetInvalid + } + if err := buf.EnsureCapacity(batchSize + 24); err != nil { + return 0, err + } + if _, err := buf.Write(types.EncodeInt64(&cnt)); err != nil { + return 0, err + } + batchSizePos := buf.Len() + var zero int64 + if _, err := buf.Write(types.EncodeInt64(&zero)); err != nil { + return 0, err + } + batchStart := buf.Len() + if err := bat.MarshalBinaryTo(buf); err != nil { + return 0, err + } + serializedSize := int64(buf.Len() - batchStart) + copy( + buf.Bytes()[batchSizePos:batchSizePos+8], + types.EncodeInt64(&serializedSize), + ) + magic := uint64(spillMagic) + if _, err := buf.Write(types.EncodeUint64(&magic)); err != nil { + return 0, err + } + return cnt, nil +} + func (ctr *container) writeSpillPayload( proc *process.Process, file *os.File, @@ -478,15 +754,48 @@ func (ctr *container) spillBatchBounded(proc *process.Process, bat *batch.Batch, if err := checkHashBuildCanceled(proc); err != nil { return err } - need, err := spillScratchBudgetBytes(bat, sourceAlreadyCharged) - if err != nil { - return err + exact := ctr.spillBatchAllocation != nil + var ( + need uint64 + externalSource *process.HashBuildReservation + err error + ) + if exact { + if ctr.hashmapBuilder.mapAllocationAccount == nil { + return mpool.ErrAllocationAccountInvalid + } + if ctr.spillAllocationMP != nil && ctr.spillAllocationMP != proc.Mp() { + return mpool.ErrAllocationAccountInvalid + } + ctr.spillAllocationMP = proc.Mp() + // A pre-spill token is headroom, not physical ownership. Release it + // immediately before the exact scratch allocations consume that space. + ctr.releaseSpillScratchReservation() + if ctr.hashmapBuilder.budget != nil && !sourceAlreadyCharged { + externalBytes := uint64(bat.Allocated()) + if size := uint64(bat.Size()); size > externalBytes { + externalBytes = size + } + if externalBytes > 0 { + externalSource, err = ctr.hashmapBuilder.budget.Reserve(externalBytes) + if err != nil { + return err + } + defer externalSource.Release() + } + } + } else { + need, err = spillScratchBudgetBytes(bat, sourceAlreadyCharged) + if err != nil { + return err + } } - // Scratch belongs to the execution generation, not to one batch. Establish - // and grow the lease lazily before allocating spill buffers. Keep it live - // while capacities are retained and release it from Reset/Free/build cleanup - // exactly once. - if ctr.hashmapBuilder.budget != nil { + // Scratch belongs to the execution generation, not to one batch. Build + // normally pre-admits the emergency lease before calling us; direct callers + // (including recovery/error paths and unit tests) establish the same lease + // here. Keep it live while capacities are retained and release it from + // Reset/Free/build cleanup exactly once. + if ctr.hashmapBuilder.budget != nil && !exact { if ctr.spillScratchReservation == nil { ctr.spillScratchReservation, err = ctr.hashmapBuilder.budget.Reserve(need) if err != nil { @@ -515,34 +824,63 @@ func (ctr *container) spillBatchBounded(proc *process.Process, bat *batch.Batch, } rows := bat.RowCount() - replacementOverlap, err := spillCapacityReplacementOverlap( - rows, - len(executors), - cap(ctr.spillHashValues), - cap(ctr.spillBucketRowIds), - cap(ctr.spillKeyVecs), - ) - if err != nil { - return err - } - replacementPeak, err := spillCheckedAdd(need, replacementOverlap) - if err != nil { - return err - } - oldScratchSize, grewScratch, err := ctr.growSpillScratchTransient(replacementPeak, analyzer) - if err != nil { - return err + var oldScratchSize uint64 + var grewScratch bool + if !exact { + replacementOverlap, overlapErr := spillCapacityReplacementOverlap( + rows, + len(executors), + cap(ctr.spillHashValues), + cap(ctr.spillBucketRowIds), + cap(ctr.spillKeyVecs), + ) + if overlapErr != nil { + return overlapErr + } + replacementPeak, addErr := spillCheckedAdd(need, replacementOverlap) + if addErr != nil { + return addErr + } + oldScratchSize, grewScratch, err = ctr.growSpillScratchTransient( + replacementPeak, + analyzer, + ) + if err != nil { + return err + } } if cap(ctr.spillKeyVecs) < len(executors) { ctr.spillKeyVecs = make([]*vector.Vector, len(executors)) } - if cap(ctr.spillHashValues) < rows { + if exact { + ctr.spillHashValues, err = growHashBuildSpillSlice( + ctr.spillHashValues, + rows, + proc.Mp(), + ctr.hashmapBuilder.mapAllocationAccount, + HashBuildSpillAllocationSiteHashValues, + ) + } else if cap(ctr.spillHashValues) < rows { ctr.spillHashValues = make([]uint64, rows) } - if cap(ctr.spillBucketRowIds) < rows { + if err != nil { + return err + } + if exact { + ctr.spillBucketRowIds, err = growHashBuildSpillSlice( + ctr.spillBucketRowIds, + rows, + proc.Mp(), + ctr.hashmapBuilder.mapAllocationAccount, + HashBuildSpillAllocationSiteRowIDs, + ) + } else if cap(ctr.spillBucketRowIds) < rows { ctr.spillBucketRowIds = make([]int32, rows) } + if err != nil { + return err + } if err := ctr.restoreSpillScratchTransient(oldScratchSize, grewScratch); err != nil { return err } @@ -630,12 +968,30 @@ func (ctr *container) spillBatchBounded(proc *process.Process, bat *batch.Batch, } if selected == nil { selected = batch.NewOffHeapWithSize(len(bat.Vecs)) + if exact { + if err := selected.SetAllocationAccount( + ctr.spillBatchAllocation, + ); err != nil { + return err + } + } selected.Attrs = bat.Attrs for i, vec := range bat.Vecs { if vec == nil { return process.ErrHashBuildBudgetInvalid } - selected.Vecs[i] = vector.NewOffHeapVecWithType(*vec.GetType()) + if exact { + selected.Vecs[i], err = + vector.NewOffHeapVecWithTypeAndAllocation( + *vec.GetType(), + ctr.spillBatchAllocation, + ) + if err != nil { + return err + } + } else { + selected.Vecs[i] = vector.NewOffHeapVecWithType(*vec.GetType()) + } } } selected.CleanOnlyData() @@ -685,6 +1041,15 @@ func (ctr *container) appendSpillRecord( if bucket < 0 || bucket >= spillNumBuckets { return process.ErrHashBuildBudgetInvalid } + if ctr.spillBatchAllocation != nil { + return ctr.appendAccountedSpillRecord( + proc, + file, + bucket, + bat, + analyzer, + ) + } grow, err := spillMarshalGrowBytes(bat) if err != nil { return err @@ -734,6 +1099,90 @@ func (ctr *container) appendSpillRecord( return nil } +func (ctr *container) appendAccountedSpillRecord( + proc *process.Process, + file *os.File, + bucket int, + bat *batch.Batch, + analyzer process.Analyzer, +) error { + if ctr.spillAllocationMP != proc.Mp() || + ctr.hashmapBuilder.mapAllocationAccount == nil { + return mpool.ErrAllocationAccountInvalid + } + if ctr.spillAccountedWrite == nil { + var err error + ctr.spillAccountedWrite, err = mpool.NewAccountedBuffer( + proc.Mp(), + ctr.hashmapBuilder.mapAllocationAccount, + HashBuildAllocationOwner, + HashBuildSpillAllocationSiteMarshalBuffer, + ) + if err != nil { + return err + } + } + cnt, err := marshalSpillRecordAccounted(bat, ctr.spillAccountedWrite) + if err != nil { + return err + } + payload := ctr.spillAccountedWrite.Bytes() + buffer := ctr.spillAccountedBuckets[bucket] + if buffer != nil && buffer.Len() > 0 && + buffer.Len()+len(payload) > spillWriteCoalesceSize { + if err := ctr.flushPendingSpillBucket( + proc, + file, + bucket, + analyzer, + ); err != nil { + return err + } + } + if len(payload) > spillWriteCoalesceSize { + return ctr.writeSpillPayload(proc, file, payload, cnt, analyzer) + } + if buffer == nil { + buffer, err = mpool.NewAccountedBuffer( + proc.Mp(), + ctr.hashmapBuilder.mapAllocationAccount, + HashBuildAllocationOwner, + HashBuildSpillAllocationSiteCoalesceBuffer, + ) + if err != nil { + return err + } + ctr.spillAccountedBuckets[bucket] = buffer + } + if buffer.Len() == 0 && buffer.Cap() < spillWriteCoalesceSize { + if err := buffer.EnsureCapacity(spillWriteCoalesceSize); err != nil { + if mpool.IsRetryableAllocationCapacity(err) { + return ctr.writeSpillPayload( + proc, + file, + payload, + cnt, + analyzer, + ) + } + return err + } + } + if _, err := buffer.Write(payload); err != nil { + return err + } + ctr.spillBucketWriteRows[bucket] += cnt + if buffer.Len() >= spillWriteCoalesceSize { + return ctr.flushPendingSpillBucket( + proc, + file, + bucket, + analyzer, + ) + } + return nil +} + func (ctr *container) ensureSpillCoalesceCapacity(buf *bytes.Buffer, analyzer process.Analyzer) bool { if buf == nil || buf.Cap() >= spillWriteCoalesceSize { return true @@ -764,16 +1213,29 @@ func (ctr *container) flushPendingSpillBucket( if bucket < 0 || bucket >= spillNumBuckets { return process.ErrHashBuildBudgetInvalid } - buf := &ctr.spillBucketWriteBufs[bucket] - if buf.Len() == 0 { - return nil - } rows := ctr.spillBucketWriteRows[bucket] - payload := buf.Bytes() + var payload []byte + if ctr.spillBatchAllocation != nil { + buffer := ctr.spillAccountedBuckets[bucket] + if buffer == nil || buffer.Len() == 0 { + return nil + } + payload = buffer.Bytes() + } else { + buf := &ctr.spillBucketWriteBufs[bucket] + if buf.Len() == 0 { + return nil + } + payload = buf.Bytes() + } err := ctr.writeSpillPayload(proc, file, payload, rows, analyzer) // Clear even on a failed/partial write. A caller's enclosing failure path // owns cleanup, and retrying the same bytes could duplicate records. - buf.Reset() + if ctr.spillBatchAllocation != nil { + ctr.spillAccountedBuckets[bucket].Reset() + } else { + ctr.spillBucketWriteBufs[bucket].Reset() + } ctr.spillBucketWriteRows[bucket] = 0 return err } @@ -785,17 +1247,28 @@ func (ctr *container) flushPendingSpillBucket( func (ctr *container) flushSpillBuffers(proc *process.Process, files []*os.File, analyzer process.Analyzer) error { var firstErr error for bucket := 0; bucket < spillNumBuckets; bucket++ { - if ctr.spillBucketWriteBufs[bucket].Len() == 0 { + pending := ctr.spillBucketWriteBufs[bucket].Len() + if ctr.spillBatchAllocation != nil && + ctr.spillAccountedBuckets[bucket] != nil { + pending = ctr.spillAccountedBuckets[bucket].Len() + } + if pending == 0 { continue } if firstErr != nil { ctr.spillBucketWriteBufs[bucket].Reset() + if ctr.spillAccountedBuckets[bucket] != nil { + ctr.spillAccountedBuckets[bucket].Reset() + } ctr.spillBucketWriteRows[bucket] = 0 continue } if err := checkHashBuildCanceled(proc); err != nil { firstErr = err ctr.spillBucketWriteBufs[bucket].Reset() + if ctr.spillAccountedBuckets[bucket] != nil { + ctr.spillAccountedBuckets[bucket].Reset() + } ctr.spillBucketWriteRows[bucket] = 0 continue } @@ -806,6 +1279,9 @@ func (ctr *container) flushSpillBuffers(proc *process.Process, files []*os.File, if file == nil { firstErr = process.ErrHashBuildBudgetInvalid ctr.spillBucketWriteBufs[bucket].Reset() + if ctr.spillAccountedBuckets[bucket] != nil { + ctr.spillAccountedBuckets[bucket].Reset() + } ctr.spillBucketWriteRows[bucket] = 0 continue } diff --git a/pkg/sql/colexec/hashbuild/spill_test.go b/pkg/sql/colexec/hashbuild/spill_test.go index d8770b2d6e1c6..97837eab1c52f 100644 --- a/pkg/sql/colexec/hashbuild/spill_test.go +++ b/pkg/sql/colexec/hashbuild/spill_test.go @@ -1606,3 +1606,73 @@ func TestCleanupSpillFiles(t *testing.T) { require.Error(t, err, "file should be closed") } } + +func TestAccountedInitialSpillConvertsHeadroomToPhysicalOwnership(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + const limit = uint64(8 << 20) + budget := process.MustNewHashBuildBudget(limit, limit) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 128) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + var op HashBuild + op.NeedHashMap = true + require.NoError(t, op.SetAllocationAccount(account)) + ctr := &op.ctr + ctr.hashmapBuilder.setBudget(generation) + ctr.spillUUID = "accounted-initial-spill" + exprs := []*plan.Expr{newExpr(0, types.T_int64.ToType())} + executors, err := ctr.initSpillExprExecs(proc, exprs) + require.NoError(t, err) + require.True(t, ctr.spillExprAccounted) + input := testutil.NewBatch( + []types.Type{types.T_int64.ToType(), types.T_varchar.ToType()}, + true, + 1_024, + proc.Mp(), + ) + defer input.Clean(proc.Mp()) + headroom, err := spillBudgetBytes(input) + require.NoError(t, err) + ctr.spillScratchReservation, err = generation.Reserve(headroom) + require.NoError(t, err) + ctr.spillScratchBase = headroom + ctr.spillScratchEmergency = true + files := make([]*os.File, spillNumBuckets) + analyzer := process.NewAnalyzer(0, false, false, "test") + require.NoError(t, ctr.spillBatchBounded( + proc, + input, + files, + executors, + analyzer, + false, + )) + require.Nil(t, ctr.spillScratchReservation) + snapshot := generation.Snapshot() + require.Equal(t, snapshot.AllocationUsed, snapshot.Used, + "external input ownership is transient and headroom is not stacked") + require.Positive(t, snapshot.AllocationUsed) + require.NotNil(t, ctr.spillAccountedWrite) + require.NoError(t, ctr.flushSpillBuffers(proc, files, analyzer)) + + ctr.dropSpillScratchBuffers() + ctr.freeSpillExprExecs() + for _, file := range files { + if file != nil { + _ = file.Close() + } + } + if ctr.spillBundle != nil { + ctr.spillBundle.release() + ctr.spillBundle = nil + } + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.NoError(t, op.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} diff --git a/pkg/sql/colexec/hashbuild/types.go b/pkg/sql/colexec/hashbuild/types.go index 59d24c58aeffe..9c0c2f66d7b95 100644 --- a/pkg/sql/colexec/hashbuild/types.go +++ b/pkg/sql/colexec/hashbuild/types.go @@ -50,6 +50,17 @@ const ( const HashBuildAllocationOwner mpool.AllocationOwner = 1 +const ( + HashBuildSpillAllocationSiteSelectedData mpool.AllocationSite = iota + 64 + HashBuildSpillAllocationSiteSelectedArea + HashBuildSpillAllocationSiteSelectedNulls + HashBuildSpillAllocationSiteSelectedGrouping + HashBuildSpillAllocationSiteHashValues + HashBuildSpillAllocationSiteRowIDs + HashBuildSpillAllocationSiteMarshalBuffer + HashBuildSpillAllocationSiteCoalesceBuffer +) + const ( HashBuildAllocationSiteHashCell mpool.AllocationSite = iota + 24 HashBuildAllocationSiteHashDescriptor @@ -60,6 +71,18 @@ const ( HashBuildAllocationSiteGroupSels ) +// Runtime-filter keys and their published wire payload have lifetimes that +// differ from copied build batches: keys die after publication while the +// payload lives on the message board until every receiver destroys it. Keep +// their sites distinct from both the builder and SpillEngine ranges. +const ( + HashBuildAllocationSiteUniqueKeyData mpool.AllocationSite = iota + 44 + HashBuildAllocationSiteUniqueKeyArea + HashBuildAllocationSiteUniqueKeyNulls + HashBuildAllocationSiteUniqueKeyGrouping + HashBuildAllocationSiteRuntimeFilterPayload +) + type container struct { state int runtimeFilterIn bool @@ -94,13 +117,23 @@ type container struct { // spillBucketWriteBufs coalesce serialized records across source batches. // Each buffer is bounded by spillWriteCoalesceSize (plus bytes.Buffer's // bounded growth slack), so fanout does not imply fanout-sized vectors. - spillBucketWriteBufs [spillNumBuckets]bytes.Buffer - spillBucketWriteRows [spillNumBuckets]int64 - spillKeyVecs []*vector.Vector - // spillScratchReservation is a query/CN-charged lease retained while spill - // buffers are reusable. It is established lazily before the first scratch - // allocation and released with the execution generation. + spillBucketWriteBufs [spillNumBuckets]bytes.Buffer + spillBucketWriteRows [spillNumBuckets]int64 + spillKeyVecs []*vector.Vector + spillBatchAllocation *vector.AllocationAccountSelection + spillAllocationMP *mpool.MPool + spillAccountedWrite *mpool.AccountedBuffer + spillAccountedBuckets [spillNumBuckets]*mpool.AccountedBuffer + // spillScratchReservation is a query/CN-charged emergency lease retained + // while Shuffle build batches accumulate. It prevents retained copies from + // consuming the scratch required to recover from hard-budget rejection. spillScratchReservation *process.HashBuildReservation + // spillScratchEmergency marks a lease pre-admitted by + // ensureDirectSpillScratchReservation or + // ensureRetainedSpillScratchReservation. An uncharged upstream batch may + // not grow beyond this lease. A retained batch may grow it because its + // source memory remains charged separately while the batch is drained. + spillScratchEmergency bool // spillScratchBase is the retained scratch floor. Coalesce-buffer growth is // charged on top and must never be mistaken for this floor. spillScratchBase uint64 @@ -294,7 +327,32 @@ func (hashBuild *HashBuild) AllocationAccountEnabled() bool { func (hashBuild *HashBuild) SetAllocationAccount( account *mpool.AllocationAccount, ) error { - builder := &hashBuild.ctr.hashmapBuilder + selection, err := vector.NewAllocationAccountSelectionWithBitmaps( + account, + HashBuildAllocationOwner, + HashBuildSpillAllocationSiteSelectedData, + HashBuildSpillAllocationSiteSelectedArea, + HashBuildSpillAllocationSiteSelectedNulls, + HashBuildSpillAllocationSiteSelectedGrouping, + ) + if err != nil { + return err + } + if err = hashBuild.ctr.hashmapBuilder.SetAllocationAccount(account); err != nil { + return err + } + hashBuild.ctr.spillBatchAllocation = selection + return nil +} + +// SetAllocationAccount activates the physical allocation provenance shared by +// the producer HashBuild and SpillEngine rebuild builders. A builder is always +// single-generation and clears the selection only after all owned resources +// have either been freed or transferred to a JoinMap. +func (hb *HashmapBuilder) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + builder := hb if builder.mapAllocationAccount != nil { if builder.mapAllocationAccount == account { return nil @@ -321,6 +379,17 @@ func (hashBuild *HashBuild) SetAllocationAccount( if err != nil { return err } + uniqueKeySelection, err := vector.NewAllocationAccountSelectionWithBitmaps( + account, + HashBuildAllocationOwner, + HashBuildAllocationSiteUniqueKeyData, + HashBuildAllocationSiteUniqueKeyArea, + HashBuildAllocationSiteUniqueKeyNulls, + HashBuildAllocationSiteUniqueKeyGrouping, + ) + if err != nil { + return err + } expressionAllocation, err := colexec.NewExpressionAllocationAccount( account, HashBuildAllocationOwner, @@ -331,6 +400,7 @@ func (hashBuild *HashBuild) SetAllocationAccount( builder.mapAllocationAccount = account builder.mapAllocation = selection builder.batchAllocation = batchSelection + builder.uniqueKeyAllocation = uniqueKeySelection builder.expressionAllocation = expressionAllocation return nil } @@ -339,6 +409,31 @@ func (hashBuild *HashBuild) ClearAllocationAccount( account *mpool.AllocationAccount, ) error { builder := &hashBuild.ctr.hashmapBuilder + if len(hashBuild.ctr.spillExprExecs) != 0 { + return mpool.ErrAllocationAccountInvariant + } + if hashBuild.ctr.spillAllocationMP != nil || + hashBuild.ctr.spillAccountedWrite != nil { + return mpool.ErrAllocationAccountInvariant + } + for _, buffer := range hashBuild.ctr.spillAccountedBuckets { + if buffer != nil { + return mpool.ErrAllocationAccountInvariant + } + } + if err := builder.ClearAllocationAccount(account); err != nil { + return err + } + hashBuild.ctr.spillBatchAllocation = nil + return nil +} + +// ClearAllocationAccount verifies that no builder-owned object can allocate +// through the generation before dropping its selections. +func (hb *HashmapBuilder) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + builder := hb if builder.mapAllocationAccount == nil { return nil } @@ -347,12 +442,13 @@ func (hashBuild *HashBuild) ClearAllocationAccount( } if builder.IntHashMap != nil || builder.StrHashMap != nil || len(builder.Batches.Buf) != 0 || builder.Sels.Size() != 0 || - len(builder.executors) != 0 || len(hashBuild.ctr.spillExprExecs) != 0 { + len(builder.executors) != 0 { return mpool.ErrAllocationAccountInvariant } builder.mapAllocationAccount = nil builder.mapAllocation = nil builder.batchAllocation = nil + builder.uniqueKeyAllocation = nil builder.expressionAllocation = nil return nil } diff --git a/pkg/sql/colexec/hashjoin/join.go b/pkg/sql/colexec/hashjoin/join.go index fc932967c3060..1941151abcfd9 100644 --- a/pkg/sql/colexec/hashjoin/join.go +++ b/pkg/sql/colexec/hashjoin/join.go @@ -348,7 +348,7 @@ func (hashJoin *HashJoin) build(analyzer process.Analyzer, proc *process.Process return leaseErr } ctr.probeExpressionLease = probeExpressionLease - engine := spillutil.NewSpillEngine(spillutil.SpillEngineConfig{ + engine, engineErr := spillutil.NewSpillEngineForAccount(spillutil.SpillEngineConfig{ BuildKeyExprs: hashJoin.EqConds[1], ProbeKeyExprs: hashJoin.EqConds[0], SpillThreshold: ctr.spillThreshold, @@ -359,7 +359,15 @@ func (hashJoin *HashJoin) build(analyzer process.Analyzer, proc *process.Process NeedBatches: hashJoin.NeedBuildBatches(), Budget: budget, ProbeExpressionLease: probeExpressionLease, - }) + }, hashJoin.allocationAccount, hashbuild.HashBuildAllocationOwner) + if engineErr != nil { + _ = payload.Close() + ctr.mp.Free() + ctr.mp = nil + ctr.cleanEqCondExecutors() + ctr.releaseProbeExpressionLease() + return engineErr + } if len(payload.Files) > 0 { engine.InitFromSpilledFiles(payload.Files) } else { diff --git a/pkg/sql/colexec/hashjoin/types.go b/pkg/sql/colexec/hashjoin/types.go index c0330797c94c7..c8b4a6ae71856 100644 --- a/pkg/sql/colexec/hashjoin/types.go +++ b/pkg/sql/colexec/hashjoin/types.go @@ -17,6 +17,7 @@ package hashjoin import ( "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/hashmap" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/reuse" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" @@ -138,10 +139,45 @@ type HashJoin struct { RuntimeFilterSpecs []*plan.RuntimeFilterSpec JoinMapTag int32 SpillThreshold int64 + allocationAccount *mpool.AllocationAccount vm.OperatorBase } +func (hashJoin *HashJoin) AllocationAccountEnabled() bool { + return hashJoin != nil +} + +func (hashJoin *HashJoin) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if account == nil || account.Handle() == 0 { + return mpool.ErrAllocationAccountInvalid + } + if hashJoin.allocationAccount != nil && + hashJoin.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + hashJoin.allocationAccount = account + return nil +} + +func (hashJoin *HashJoin) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if hashJoin.allocationAccount == nil { + return nil + } + if hashJoin.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if hashJoin.ctr.mp != nil || hashJoin.ctr.spillEngine != nil { + return mpool.ErrAllocationAccountInvariant + } + hashJoin.allocationAccount = nil + return nil +} + func (hashJoin *HashJoin) GetOperatorBase() *vm.OperatorBase { return &hashJoin.OperatorBase } @@ -216,6 +252,7 @@ func (hashJoin *HashJoin) Reset(proc *process.Process, pipelineFailed bool, err ctr.state = Build ctr.probeState = psNextBatch ctr.lastIdx = 0 + hashJoin.allocationAccount = nil if hashJoin.OpAnalyzer != nil { hashJoin.OpAnalyzer.Alloc(ctr.maxAllocSize) @@ -232,6 +269,7 @@ func (hashJoin *HashJoin) Free(proc *process.Process, pipelineFailed bool, err e ctr.releaseProbeExpressionLease() ctr.cleanHashMap() ctr.cleanNonEqCondExecutor() + hashJoin.allocationAccount = nil } func (ctr *container) resetNonEqCondExecutor() { diff --git a/pkg/sql/colexec/rightdedupjoin/join.go b/pkg/sql/colexec/rightdedupjoin/join.go index 52df5361a3527..19186ff3de60d 100644 --- a/pkg/sql/colexec/rightdedupjoin/join.go +++ b/pkg/sql/colexec/rightdedupjoin/join.go @@ -227,7 +227,7 @@ func (rightDedupJoin *RightDedupJoin) build(analyzer process.Analyzer, proc *pro return leaseErr } ctr.probeExpressionLease = probeExpressionLease - engine := spillutil.NewSpillEngine(spillutil.SpillEngineConfig{ + engine, engineErr := spillutil.NewSpillEngineForAccount(spillutil.SpillEngineConfig{ BuildKeyExprs: rightDedupJoin.Conditions[1], ProbeKeyExprs: rightDedupJoin.Conditions[0], SpillThreshold: ctr.spillThreshold, @@ -235,7 +235,15 @@ func (rightDedupJoin *RightDedupJoin) build(analyzer process.Analyzer, proc *pro MergeProbeBatches: true, Budget: budget, ProbeExpressionLease: probeExpressionLease, - }) + }, rightDedupJoin.allocationAccount, hashbuild.HashBuildAllocationOwner) + if engineErr != nil { + _ = payload.Close() + ctr.mp.Free() + ctr.mp = nil + ctr.cleanEvalVectors() + ctr.releaseProbeExpressionLease() + return engineErr + } if len(payload.Files) > 0 { engine.InitFromSpilledFiles(payload.Files) } else { diff --git a/pkg/sql/colexec/rightdedupjoin/types.go b/pkg/sql/colexec/rightdedupjoin/types.go index 5489c1e98e602..aef8d1b2e931b 100644 --- a/pkg/sql/colexec/rightdedupjoin/types.go +++ b/pkg/sql/colexec/rightdedupjoin/types.go @@ -123,7 +123,8 @@ func (rightDedupJoin *RightDedupJoin) ClearAllocationAccount( if rightDedupJoin.allocationAccount != account { return mpool.ErrAllocationAccountMismatch } - if rightDedupJoin.ctr.mp != nil { + if rightDedupJoin.ctr.mp != nil || + rightDedupJoin.ctr.spillEngine != nil { return mpool.ErrAllocationAccountInvariant } rightDedupJoin.allocationAccount = nil diff --git a/pkg/sql/colexec/spillutil/allocation_account_test.go b/pkg/sql/colexec/spillutil/allocation_account_test.go index cbc59c88bd7bf..7b5e9d22bd2de 100644 --- a/pkg/sql/colexec/spillutil/allocation_account_test.go +++ b/pkg/sql/colexec/spillutil/allocation_account_test.go @@ -18,12 +18,14 @@ import ( "bytes" "os" "path/filepath" + "strings" "testing" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" "github.com/matrixorigin/matrixone/pkg/testutil" "github.com/matrixorigin/matrixone/pkg/vm/process" "github.com/stretchr/testify/require" @@ -175,6 +177,90 @@ func TestSpillAllocationAccountDecodedBatchLifecycle(t *testing.T) { finalizeTestSpillAllocationAccount(t, state) } +func TestSpillAllocationAccountDecodedReuseRetriesFromCleanRecord(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-decoded-retry"), + ) + defer proc.Free() + makeSource := func(width int) *batch.Batch { + values := make([]string, 1_024) + for i := range values { + values[i] = strings.Repeat("x", width) + } + return testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.NewVector( + len(values), + types.T_varchar.ToType(), + proc.Mp(), + false, + values, + ), + }, nil) + } + first := makeSource(512) + second := makeSource(2_048) + defer first.Clean(proc.Mp()) + defer second.Clean(proc.Mp()) + + measure := newTestSpillAllocationAccount(t, 64<<20, 128) + reader := BucketReader{ + fd: writeSpillAllocationTestRecords(t, first, second), + allocation: measure.allocation, + } + reuse := batch.NewOffHeapWithSize(0) + _, err := reader.ReadBatch(proc, reuse) + require.NoError(t, err) + firstUsed := measure.account.Snapshot().Used + require.Positive(t, firstUsed) + reuse.Clean(proc.Mp()) + require.NoError(t, reuse.SetAllocationAccount(measure.allocation.decoded)) + _, err = reader.ReadBatch(proc, reuse) + require.NoError(t, err) + secondUsed := measure.account.Snapshot().Used + require.Positive(t, secondUsed) + reuse.Clean(proc.Mp()) + reader.Close() + finalizeTestSpillAllocationAccount(t, measure) + + limit := max(firstUsed, secondUsed) + 128<<10 + require.Less(t, limit, firstUsed+secondUsed) + budget := process.MustNewHashBuildBudget(limit, limit) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 128) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + allocation, err := NewSpillAllocationAccount(account, 2) + require.NoError(t, err) + reader = BucketReader{allocation: allocation} + require.NoError(t, reader.EnsureBuffer(generation)) + reader.ResetForFd(writeSpillAllocationTestRecords(t, first, second)) + reuse = batch.NewOffHeapWithSize(0) + _, err = reader.ReadBatch(proc, reuse) + require.NoError(t, err) + rejects := generation.RejectCount() + _, err = reader.ReadBatch(proc, reuse) + require.NoError(t, err) + require.Equal(t, rejects, generation.RejectCount(), + "the local account rejects the overlap before the shared controller") + require.Equal(t, uint64(1), reader.cleanRetries, + "replacement overlap must exercise the clean-record retry") + require.Equal(t, + generation.Snapshot().AllocationUsed+uint64(64<<10), + generation.Used(), + "decoded payloads have no duplicate hard reservation", + ) + reuse.Clean(proc.Mp()) + reader.Close() + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + func TestSpillAllocationAccountScatterScratchLifecycle(t *testing.T) { proc := testutil.NewProcessWithMPool( t, @@ -238,6 +324,69 @@ func TestSpillAllocationAccountScatterScratchLifecycle(t *testing.T) { finalizeTestSpillAllocationAccount(t, state) } +func TestSpillAllocationAccountScatterChargesOnlyExternalSource(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-scatter-source"), + ) + defer proc.Free() + const limit = uint64(8 << 20) + budget := process.MustNewHashBuildBudget(limit, limit) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + allocation, err := NewSpillAllocationAccount(account, 2) + require.NoError(t, err) + engine, err := NewSpillEngineWithAllocation( + SpillEngineConfig{Budget: generation}, + allocation, + ) + require.NoError(t, err) + source := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.NewVector( + 8, + types.T_int64.ToType(), + proc.Mp(), + false, + []int64{1, 2, 3, 4, 5, 6, 7, 8}, + ), + }, nil) + defer source.Clean(proc.Mp()) + writers := MakeBucketWriters("spill_allocation_scatter_source") + defer func() { + for i := range writers { + writers[i].Close() + } + }() + require.NoError(t, engine.scatterBatchBounded( + proc, + source, + source.Vecs, + writers, + 0, + false, + process.NewAnalyzer(0, false, false, "test"), + )) + snapshot := generation.Snapshot() + require.Nil(t, engine.scatterScratchReservation) + require.Equal(t, snapshot.AllocationUsed, snapshot.Used, + "the upstream source token is transient and private spill bytes are exact") + require.Positive(t, snapshot.ReserveCount, + "the unaccounted upstream source remains part of the peak") + require.Greater(t, snapshot.PeakUsed, snapshot.Used) + + engine.releaseScatterScratch() + engine.Cleanup(proc) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + func TestSpillAllocationAccountMarshalBufferLifecycle(t *testing.T) { proc := testutil.NewProcessWithMPool( t, @@ -377,3 +526,77 @@ func TestSpillAllocationAccountScatterFailureCleanup(t *testing.T) { engine.Cleanup(proc) finalizeTestSpillAllocationAccount(t, state) } + +func TestSpillAllocationAccountRebuildAndRecursiveSpillLifecycle(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-rebuild"), + ) + defer proc.Free() + const limit = uint64(64 << 20) + budget := process.MustNewHashBuildBudget(limit, limit) + generation, err := budget.OpenGenerationWithSpillCaps( + 1, + limit, + 1<<30, + 4_096, + ) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + engine, err := NewSpillEngineForAccount( + SpillEngineConfig{ + BuildKeyExprs: makeTestKeyExpr(), + Budget: generation, + SpillThreshold: 100, + NeedsBuildForEmptyProbe: true, + }, + account, + hashbuild.HashBuildAllocationOwner, + ) + require.NoError(t, err) + + values := make([]int32, 5_000) + for i := range values { + values[i] = int32(i) + } + source := makeInt32Batch(proc, values) + fd := writeBuildFile(proc, "accounted_recursive_build", source) + source.Clean(proc.Mp()) + engine.InitFromSpilledMap([]*os.File{fd}) + analyzer := process.NewAnalyzer(0, false, false, "test") + + respills := 0 + ready := 0 + for steps := 0; engine.HasMoreBuckets(); steps++ { + require.Less(t, steps, 4_096, "recursive spill queue made no progress") + jm, result, rebuildErr := engine.RebuildHashmap(proc, analyzer) + require.NoError(t, rebuildErr) + switch result { + case BucketReSpilled: + respills++ + case BucketReady: + ready++ + require.NotNil(t, jm) + jm.Free() + case BucketSkip, BucketEmptyBuild: + require.Nil(t, jm) + default: + require.NotEqual(t, BucketQueueEmpty, result) + } + } + require.Positive(t, respills) + require.Positive(t, ready) + require.Positive(t, account.Snapshot().Peak) + + engine.Cleanup(proc) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.Zero(t, generation.SpillDiskUsed()) + require.Zero(t, generation.SpillFDUsed()) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} diff --git a/pkg/sql/colexec/spillutil/join_spill.go b/pkg/sql/colexec/spillutil/join_spill.go index b7d0d210c4179..330aa0f7aed50 100644 --- a/pkg/sql/colexec/spillutil/join_spill.go +++ b/pkg/sql/colexec/spillutil/join_spill.go @@ -89,6 +89,7 @@ type BucketReader struct { spillFile *message.SpillFile mergeRecords bool allocation *SpillAllocationAccount + cleanRetries uint64 } func (r *BucketReader) ReadBatch(proc *process.Process, reuseBat *batch.Batch) (*batch.Batch, error) { @@ -191,7 +192,7 @@ func (r *BucketReader) ReadBatch(proc *process.Process, reuseBat *batch.Batch) ( return nil, r.mergeReadError(proc, reuseBat, next, nextToken, err) } var mergeToken *process.HashBuildReservation - if r.budget != nil { + if r.budget != nil && r.allocation == nil { // Keep the current destination (O) and the source record (N) live // while admitting the final destination (D). UnionBatch may retain // rounded capacities larger than O+N, so reserving O+N here is not a @@ -541,7 +542,7 @@ func (r *BucketReader) readBatchRecord( if err := checkSpillCanceled(proc); err != nil { return nil, token, charge, err } - if r.budget != nil { + if r.budget != nil && r.allocation == nil { payload := uint64(batchSize) if payload > uint64(maxIntValue())-(64<<10) { return nil, token, charge, process.ErrHashBuildBudgetInvalid @@ -571,13 +572,6 @@ func (r *BucketReader) readBatchRecord( // A caller-provided reuse batch has no budget ownership on the first // read. Drop it before admitting the decoded payload. reuseBat.Clean(proc.Mp()) - if r.allocation != nil { - if err := reuseBat.SetAllocationAccount( - r.allocation.decoded, - ); err != nil { - return nil, nil, 0, err - } - } var err error token, err = r.budget.Reserve(projected) if err != nil { @@ -602,14 +596,6 @@ func (r *BucketReader) readBatchRecord( } if !retainedOK || !peakOK || growErr != nil { reuseBat.Clean(proc.Mp()) - if r.allocation != nil { - if err := reuseBat.SetAllocationAccount( - r.allocation.decoded, - ); err != nil { - token.Release() - return nil, nil, 0, err - } - } token.Release() token = nil var err error @@ -624,14 +610,45 @@ func (r *BucketReader) readBatchRecord( } } - reuseBat.CleanOnlyData() - if err := checkSpillCanceled(proc); err != nil { - return nil, token, charge, err + var payloadOffset int64 = -1 + if r.allocation != nil { + physical, seekErr := r.fd.Seek(0, io.SeekCurrent) + if seekErr != nil { + return nil, token, charge, seekErr + } + payloadOffset = physical - int64(r.reader.Buffered()) + if payloadOffset < 0 { + return nil, token, charge, process.ErrHashBuildBudgetInvalid + } } - - limitReader := io.LimitedReader{R: r.reader, N: batchSize} - if err := reuseBat.UnmarshalFromReader(&limitReader, proc.Mp()); err != nil { - return nil, token, charge, err + decode := func() (io.LimitedReader, error) { + reuseBat.CleanOnlyData() + if err := checkSpillCanceled(proc); err != nil { + return io.LimitedReader{}, err + } + limited := io.LimitedReader{R: r.reader, N: batchSize} + err := reuseBat.UnmarshalFromReader(&limited, proc.Mp()) + return limited, err + } + limitReader, decodeErr := decode() + if decodeErr != nil && r.allocation != nil && + mpool.IsRetryableAllocationCapacity(decodeErr) { + // Reuse growth owns O while allocating N. If that exact overlap does + // not fit, rewind the seekable spill record, release O, and retry the + // same minimum payload from a clean batch. No multiplier predicts N. + reuseBat.Clean(proc.Mp()) + if err := reuseBat.SetAllocationAccount(r.allocation.decoded); err != nil { + return nil, token, charge, err + } + if _, err := r.fd.Seek(payloadOffset, io.SeekStart); err != nil { + return nil, token, charge, err + } + r.reader.Reset(r.fd) + r.cleanRetries++ + limitReader, decodeErr = decode() + } + if decodeErr != nil { + return nil, token, charge, decodeErr } if err := checkSpillCanceled(proc); err != nil { return nil, token, charge, err @@ -1384,6 +1401,23 @@ func spillStatInt64(v uint64) int64 { return int64(v) } +func externalScatterSourceBytes( + bat *batch.Batch, + sourceAlreadyCharged bool, +) (uint64, error) { + if bat == nil || bat.RowCount() < 0 { + return 0, process.ErrHashBuildBudgetInvalid + } + if sourceAlreadyCharged { + return 0, nil + } + allocated := uint64(bat.Allocated()) + if size := uint64(bat.Size()); size > allocated { + allocated = size + } + return allocated, nil +} + func (e *SpillEngine) scatterRetainedBytes() (uint64, bool) { actual := uint64(0) add := func(v uint64) bool { @@ -1468,6 +1502,7 @@ func (e *SpillEngine) scatterBatchBounded( } rows := bat.RowCount() var selected *batch.Batch + var externalSource *process.HashBuildReservation defer func() { if selected != nil { selected.Clean(proc.Mp()) @@ -1480,8 +1515,11 @@ func (e *SpillEngine) scatterBatchBounded( if retErr != nil { e.discardScatterBuffers() } + if externalSource != nil { + externalSource.Release() + } }() - if e.cfg.Budget != nil { + if e.cfg.Budget != nil && e.allocation == nil { // Start with retained capacities already owned by this token, add only // row/hash capacity growth, then add each per-batch transient once. retained, ok := e.scatterRetainedBytes() @@ -1514,6 +1552,20 @@ func (e *SpillEngine) scatterBatchBounded( if err != nil { return err } + } else if e.cfg.Budget != nil { + externalBytes, err := externalScatterSourceBytes( + bat, + sourceAlreadyCharged, + ) + if err != nil { + return err + } + if externalBytes > 0 { + externalSource, err = e.cfg.Budget.Reserve(externalBytes) + if err != nil { + return err + } + } } if e.allocation != nil { @@ -2022,6 +2074,24 @@ func NewSpillEngineWithAllocation( return newSpillEngine(cfg, allocation), nil } +// NewSpillEngineForAccount activates exact spill ownership when an execution +// attempt supplied an account and preserves the legacy constructor for direct +// operator tests and callers outside the statement lifecycle. +func NewSpillEngineForAccount( + cfg SpillEngineConfig, + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, +) (*SpillEngine, error) { + if account == nil { + return NewSpillEngine(cfg), nil + } + allocation, err := NewSpillAllocationAccount(account, owner) + if err != nil { + return nil, err + } + return NewSpillEngineWithAllocation(cfg, allocation) +} + func newSpillEngine( cfg SpillEngineConfig, allocation *SpillAllocationAccount, @@ -2341,6 +2411,12 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana builder := &hashbuild.HashmapBuilder{} builder.SetBudget(e.cfg.Budget) + if e.allocation != nil { + if err := builder.SetAllocationAccount(e.allocation.account); err != nil { + builder.Free(proc) + return nil, BucketSkip, err + } + } builder.IsDedup = e.cfg.IsDedup builder.OnDuplicateAction = e.cfg.OnDuplicateAction builder.DedupBuildKeepLast = e.cfg.DedupBuildKeepLast @@ -2598,7 +2674,10 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal var execs []colexec.ExpressionExecutor var lease *hashbuild.ExpressionMemoryLease var err error - if e.allocation == nil { + if e.allocation == nil || + !hashbuild.AllocationAccountedExpressionSetSupported( + e.cfg.BuildKeyExprs, + ) { execs, lease, err = hashbuild.NewBudgetedExpressionExecutors( proc, @@ -2608,7 +2687,7 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal ) } else { execs, err = - colexec.NewExpressionExecutorsFromPlanExpressionsWithAllocation( + hashbuild.NewAllocationAccountedExpressionExecutors( proc, e.cfg.BuildKeyExprs, e.allocation.expression, From 28312732bcf8aa07e1bcfce88bda037c83dc5630 Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 23:37:34 +0800 Subject: [PATCH 19/61] executor: unify allocation pressure recovery --- pkg/common/mpool/allocation_account.go | 49 +- .../mpool/allocation_account_mpool_test.go | 26 + pkg/common/mpool/mpool.go | 26 +- .../dedupjoin/expression_memory_test.go | 34 + pkg/sql/colexec/dedupjoin/join.go | 35 +- .../colexec/dedupjoin/terminal_budget_test.go | 2 +- pkg/sql/colexec/dedupjoin/types.go | 9 +- pkg/sql/colexec/hashbuild/budget.go | 6 + pkg/sql/colexec/hashbuild/build.go | 22 +- pkg/sql/colexec/hashbuild/build_test.go | 71 ++ pkg/sql/colexec/hashbuild/dedup_memory.go | 100 +++ pkg/sql/colexec/hashbuild/errors.go | 8 +- pkg/sql/colexec/hashbuild/errors_test.go | 16 +- .../colexec/hashbuild/expression_memory.go | 20 + pkg/sql/colexec/hashbuild/hashmap.go | 146 +++- pkg/sql/colexec/hashbuild/hashmap_test.go | 115 ++- pkg/sql/colexec/hashbuild/pressure.go | 207 +++++ pkg/sql/colexec/hashbuild/pressure_test.go | 69 ++ pkg/sql/colexec/hashbuild/spill.go | 714 ++++++++++++++++-- pkg/sql/colexec/hashbuild/spill_test.go | 152 ++++ pkg/sql/colexec/hashbuild/types.go | 20 +- .../hashjoin/expression_memory_test.go | 34 + pkg/sql/colexec/hashjoin/join.go | 27 +- .../colexec/hashjoin/terminal_budget_test.go | 2 +- pkg/sql/colexec/hashjoin/types.go | 11 +- pkg/sql/colexec/join_util.go | 72 +- pkg/sql/colexec/join_util_test.go | 57 ++ .../rightdedupjoin/expression_memory_test.go | 34 + pkg/sql/colexec/rightdedupjoin/join.go | 35 +- pkg/sql/colexec/rightdedupjoin/types.go | 11 +- .../spillutil/allocation_account_test.go | 122 ++- pkg/sql/colexec/spillutil/join_spill.go | 600 ++++++++++++--- pkg/sql/colexec/spillutil/join_spill_test.go | 40 +- pkg/vm/process/hashbuild_budget.go | 99 ++- pkg/vm/process/hashbuild_budget_test.go | 92 ++- 35 files changed, 2723 insertions(+), 360 deletions(-) create mode 100644 pkg/sql/colexec/hashbuild/dedup_memory.go create mode 100644 pkg/sql/colexec/hashbuild/pressure.go create mode 100644 pkg/sql/colexec/hashbuild/pressure_test.go diff --git a/pkg/common/mpool/allocation_account.go b/pkg/common/mpool/allocation_account.go index 25584c257b437..8923056f74389 100644 --- a/pkg/common/mpool/allocation_account.go +++ b/pkg/common/mpool/allocation_account.go @@ -142,7 +142,10 @@ func IsRetryableAllocationCapacity(err error) bool { // the terminal boundary even if a later physical Free drains the tombstone. type AllocationAccountTerminalSnapshot struct { AllocationAccountSnapshot - State AllocationAccountTerminalState + State AllocationAccountTerminalState + LiveOwner AllocationOwner + LiveSite AllocationSite + LiveAllocations uint64 } // AllocationAccountCheckpoint records the physical live-byte boundary before @@ -552,6 +555,8 @@ func (r *AllocationAccountRegistry) CompleteTerminalWithError( } snapshot.State = AllocationAccountTerminalInvariantFailure + snapshot.LiveOwner, snapshot.LiveSite, snapshot.LiveAllocations = + allocationAccountLiveDiagnostic(account) entry.terminal = &snapshot entry.tombstone = true r.tombstones++ @@ -572,15 +577,55 @@ func newAllocationTerminalInvariantError( snapshot AllocationAccountTerminalSnapshot, ) error { return fmt.Errorf( - "%w: handle=%d used=%d peak=%d limit=%d", + "%w: handle=%d used=%d peak=%d limit=%d owner=%d site=%d live-allocations=%d", ErrAllocationAccountInvariant, snapshot.Handle, snapshot.Used, snapshot.Peak, snapshot.Limit, + snapshot.LiveOwner, + snapshot.LiveSite, + snapshot.LiveAllocations, ) } +// allocationAccountLiveDiagnostic is a terminal-only scan. It does not add a +// per-allocation hot-path counter: provenance already lives in the pointer +// metadata required for physical Free. The first live owner/site plus the +// exact live allocation count makes a nonzero terminal snapshot actionable. +func allocationAccountLiveDiagnostic( + account *AllocationAccount, +) (AllocationOwner, AllocationSite, uint64) { + if account == nil { + return 0, 0, 0 + } + var owner AllocationOwner + var site AllocationSite + var count uint64 + record := func(lease allocationLease) { + if lease.account != account { + return + } + if count == 0 { + owner = lease.owner + site = lease.site + } + count++ + } + for i := range globalPtrShards { + shard := &globalPtrShards[i] + shard.mu.Lock() + for _, lease := range shard.leases { + record(lease) + } + shard.mu.Unlock() + } + // noLock pools intentionally provide no synchronization for their local + // maps. Do not race unrelated single-threaded pools merely to enrich a + // terminal error; production query pools use the sharded metadata above. + return owner, site, count +} + func (r *AllocationAccountRegistry) removeSlotLocked( slot uint32, account *AllocationAccount, diff --git a/pkg/common/mpool/allocation_account_mpool_test.go b/pkg/common/mpool/allocation_account_mpool_test.go index cea0b62194d8b..df20d08fd8d0f 100644 --- a/pkg/common/mpool/allocation_account_mpool_test.go +++ b/pkg/common/mpool/allocation_account_mpool_test.go @@ -123,6 +123,30 @@ func TestMPoolAccountedAllocGrowFree(t *testing.T) { finalizeTestAllocationAccount(t, registry, account) } +func TestMPoolTerminalLeakDiagnosticUsesPublishedProvenance(t *testing.T) { + registry, account := newTestAllocationAccount(t, 64, 1) + mp := MustNew("accounted-terminal-diagnostic") + defer DeleteMPool(mp) + buffer, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + snapshot, first, err := registry.CompleteTerminal(account) + require.True(t, first) + require.ErrorIs(t, err, ErrAllocationAccountInvariant) + require.Equal(t, testAllocationOwner, snapshot.LiveOwner) + require.Equal(t, testAllocationSite, snapshot.LiveSite) + require.Equal(t, uint64(1), snapshot.LiveAllocations) + require.Contains(t, err.Error(), "owner=1 site=1 live-allocations=1") + mp.Free(buffer) + require.False(t, registry.AdmissionSuspended()) + _, ok := registry.Resolve(snapshot.Handle) + require.False(t, ok) +} + func TestMPoolMakeSliceAccounted(t *testing.T) { registry, account := newTestAllocationAccount(t, 64, 2) mp := MustNew("accounted-typed-slice") @@ -200,6 +224,7 @@ func TestMPoolAccountedRollback(t *testing.T) { testAllocationSite, ) require.ErrorIs(t, err, ErrAllocationAccountCapacity) + require.Contains(t, err.Error(), "owner=1 site=1") require.Zero(t, account.Snapshot().Used) require.Zero(t, registry.LiveAllocationMetadata()) finalizeTestAllocationAccount(t, registry, account) @@ -217,6 +242,7 @@ func TestMPoolAccountedRollback(t *testing.T) { testAllocationSite, ) require.ErrorIs(t, err, ErrAllocationMetadataSlots) + require.Contains(t, err.Error(), "owner=1 site=1") require.Zero(t, account.Snapshot().Used) require.Zero(t, registry.LiveAllocationMetadata()) finalizeTestAllocationAccount(t, registry, account) diff --git a/pkg/common/mpool/mpool.go b/pkg/common/mpool/mpool.go index 0981aae300bd3..288189f2ac31e 100644 --- a/pkg/common/mpool/mpool.go +++ b/pkg/common/mpool/mpool.go @@ -809,7 +809,17 @@ func (mp *MPool) allocAccountedWithDetailK( detailk string, sz int64, request allocationAccountRequest, -) ([]byte, error) { +) (result []byte, retErr error) { + defer func() { + if retErr != nil { + retErr = fmt.Errorf( + "allocation owner=%d site=%d: %w", + request.owner, + request.site, + retErr, + ) + } + }() if err := request.validate(); err != nil { return nil, err } @@ -922,24 +932,14 @@ func (mp *MPool) allocAccounted( }() if err = request.account.acquire(uint64(sz)); err != nil { - return nil, fmt.Errorf( - "allocation owner=%d site=%d: %w", - request.owner, - request.site, - err, - ) + return nil, err } accountHeld = true if err = request.reach(allocationAfterAccount); err != nil { return nil, err } if err = request.account.registry.reserveMetadata(); err != nil { - return nil, fmt.Errorf( - "allocation owner=%d site=%d: %w", - request.owner, - request.site, - err, - ) + return nil, err } metadataHeld = true if err = request.reach(allocationAfterMetadata); err != nil { diff --git a/pkg/sql/colexec/dedupjoin/expression_memory_test.go b/pkg/sql/colexec/dedupjoin/expression_memory_test.go index 4265a7313bb4b..7f053f871558f 100644 --- a/pkg/sql/colexec/dedupjoin/expression_memory_test.go +++ b/pkg/sql/colexec/dedupjoin/expression_memory_test.go @@ -62,3 +62,37 @@ func TestDedupJoinResetReleasesProbeExpressionLease(t *testing.T) { require.Nil(t, arg.ctr.vecs) require.Nil(t, arg.ctr.probeExpressionLease) } + +func TestDedupJoinResetReleasesAccountedProbeExpressions(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + const capBytes = uint64(1 << 20) + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.OpenWithController(capBytes, generation) + require.NoError(t, err) + expr := &plan.Expr{Typ: plan.Type{Id: int32(types.T_int32)}, Expr: &plan.Expr_Lit{Lit: &plan.Literal{Value: &plan.Literal_I32Val{I32Val: 1}}}} + executors, err := hashbuild.NewAllocationAccountedExpressionExecutorsForAccount( + proc, []*plan.Expr{expr}, account, hashbuild.HashBuildAllocationOwner) + require.NoError(t, err) + arg := &DedupJoin{allocationAccount: account} + arg.ctr.evecs = []evalVector{{executor: executors[0]}} + arg.ctr.vecs = make([]*vector.Vector, len(executors)) + arg.ctr.probeExpressionsAccounted = true + input := batch.NewWithSize(0) + input.SetRowCount(4) + require.NoError(t, arg.ctr.evalJoinConditionBudgeted(input, proc)) + require.Positive(t, account.Snapshot().Used) + + arg.Reset(proc, false, nil) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.False(t, arg.ctr.probeExpressionsAccounted) + require.Nil(t, arg.ctr.evecs) + terminal, _, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) +} diff --git a/pkg/sql/colexec/dedupjoin/join.go b/pkg/sql/colexec/dedupjoin/join.go index 6f30955c13dd9..d347b161aa5c1 100644 --- a/pkg/sql/colexec/dedupjoin/join.go +++ b/pkg/sql/colexec/dedupjoin/join.go @@ -312,12 +312,35 @@ func (dedupJoin *DedupJoin) build(analyzer process.Analyzer, proc *process.Proce if takeErr != nil { return takeErr } - probeExecutors := make([]colexec.ExpressionExecutor, len(ctr.evecs)) - for i := range ctr.evecs { - probeExecutors[i] = ctr.evecs[i].executor + var probeExpressionLease *hashbuild.ExpressionMemoryLease + var leaseErr error + if dedupJoin.allocationAccount != nil && + hashbuild.AllocationAccountedExpressionSetSupported(dedupJoin.Conditions[0]) { + ctr.cleanEvalVectors() + var probeExecutors []colexec.ExpressionExecutor + probeExecutors, leaseErr = + hashbuild.NewAllocationAccountedExpressionExecutorsForAccount( + proc, + dedupJoin.Conditions[0], + dedupJoin.allocationAccount, + hashbuild.HashBuildAllocationOwner, + ) + if leaseErr == nil { + ctr.evecs = make([]evalVector, len(probeExecutors)) + ctr.vecs = make([]*vector.Vector, len(probeExecutors)) + for i := range probeExecutors { + ctr.evecs[i].executor = probeExecutors[i] + } + ctr.probeExpressionsAccounted = true + } + } else { + probeExecutors := make([]colexec.ExpressionExecutor, len(ctr.evecs)) + for i := range ctr.evecs { + probeExecutors[i] = ctr.evecs[i].executor + } + probeExpressionLease, leaseErr = hashbuild.NewExpressionMemoryLease( + budget, dedupJoin.Conditions[0], probeExecutors, false) } - probeExpressionLease, leaseErr := hashbuild.NewExpressionMemoryLease( - budget, dedupJoin.Conditions[0], probeExecutors, false) if leaseErr != nil { _ = payload.Close() ctr.mp.Free() @@ -365,7 +388,7 @@ func (dedupJoin *DedupJoin) build(analyzer process.Analyzer, proc *process.Proce }, analyzer, func(bat *batch.Batch) ([]*vector.Vector, error) { - if err := ctr.evalJoinCondition(bat, proc); err != nil { + if err := ctr.evalJoinConditionBudgeted(bat, proc); err != nil { return nil, err } return ctr.vecs, nil diff --git a/pkg/sql/colexec/dedupjoin/terminal_budget_test.go b/pkg/sql/colexec/dedupjoin/terminal_budget_test.go index 5a92ea6e3fe93..ec706bf726166 100644 --- a/pkg/sql/colexec/dedupjoin/terminal_budget_test.go +++ b/pkg/sql/colexec/dedupjoin/terminal_budget_test.go @@ -42,7 +42,7 @@ func TestDedupJoinCallConvertsTerminalBudgetAdmission(t *testing.T) { admission := &process.HashBuildBudgetError{ Kind: process.HashBuildBudgetErrorAdmission, - Resource: process.HashBuildBudgetResourceMemory, + Component: process.HashBuildBudgetComponentMemory, Requested: 2, Used: 1, Cap: 1, diff --git a/pkg/sql/colexec/dedupjoin/types.go b/pkg/sql/colexec/dedupjoin/types.go index a8dad265652e4..272c7686cb99c 100644 --- a/pkg/sql/colexec/dedupjoin/types.go +++ b/pkg/sql/colexec/dedupjoin/types.go @@ -261,7 +261,8 @@ type container struct { // Non-nil only for spilled joins, where probe expressions are part of the // shared HashBuild/spill working set. Resident probe expressions remain // under normal process/mpool accounting; this is not a general query budget. - probeExpressionLease *hashbuild.ExpressionMemoryLease + probeExpressionLease *hashbuild.ExpressionMemoryLease + probeExpressionsAccounted bool } type DedupJoin struct { @@ -331,7 +332,8 @@ func (dedupJoin *DedupJoin) ClearAllocationAccount( if dedupJoin.allocationAccount != account { return mpool.ErrAllocationAccountMismatch } - if dedupJoin.ctr.mp != nil || dedupJoin.ctr.spillEngine != nil { + if dedupJoin.ctr.mp != nil || dedupJoin.ctr.spillEngine != nil || + dedupJoin.ctr.probeExpressionsAccounted { return mpool.ErrAllocationAccountInvariant } dedupJoin.allocationAccount = nil @@ -413,7 +415,7 @@ func (dedupJoin *DedupJoin) Reset(proc *process.Process, pipelineFailed bool, er ctr.spillEngine.Cleanup(proc) ctr.spillEngine = nil } - if ctr.probeExpressionLease != nil { + if ctr.probeExpressionLease != nil || ctr.probeExpressionsAccounted { ctr.cleanEvalVectors() ctr.releaseProbeExpressionLease() } else { @@ -529,6 +531,7 @@ func (ctr *container) cleanEvalVectors() { } ctr.evecs = nil ctr.vecs = nil + ctr.probeExpressionsAccounted = false } func (ctr *container) resetEvalVectors() { diff --git a/pkg/sql/colexec/hashbuild/budget.go b/pkg/sql/colexec/hashbuild/budget.go index ae829052a1be0..9d61b37ca40a8 100644 --- a/pkg/sql/colexec/hashbuild/budget.go +++ b/pkg/sql/colexec/hashbuild/budget.go @@ -676,6 +676,12 @@ func (hb *HashmapBuilder) reserveBuildAux( if hb.budget == nil { return nil } + if hb.batchAllocation != nil { + // Exact mode charges every data-scaled auxiliary owner at its physical + // allocation boundary. A second estimator reservation would double count + // the same memory and reintroduce false admission failures. + return nil + } if hb.auxReservation != nil { // BuildHashmap can be retried on the same retained batches with a // different optional-runtime-filter decision. Reconcile the existing diff --git a/pkg/sql/colexec/hashbuild/build.go b/pkg/sql/colexec/hashbuild/build.go index 467defadf1029..9a44c76c9f380 100644 --- a/pkg/sql/colexec/hashbuild/build.go +++ b/pkg/sql/colexec/hashbuild/build.go @@ -215,6 +215,7 @@ func (hashBuild *HashBuild) finalizeBuildFailure(proc *process.Process, err erro func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyzer) error { ctr := &hashBuild.ctr + ctr.spillConditions = hashBuild.Conditions spillMode := false var spillFiles []*os.File bundleTransferred := false @@ -231,6 +232,7 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz ctr.spillBundle = nil } ctr.freeSpillExprExecs() + ctr.spillConditions = nil // Build-key executors are producer scratch. No consumer reads them after // build() returns, so release their retained vectors and expression lease // here instead of holding both until pipeline Reset. @@ -258,6 +260,12 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz hashBuild.JoinMapRefCnt, ) } + if ctr.spillBatchAllocation != nil { + // Exact mode converts the one-unit forward-progress token into the + // expression and scatter allocations that follow. Keeping both live + // would stack the same capacity and reject the recovery path itself. + ctr.releaseSpillScratchReservation() + } execs, err := ctr.initSpillExprExecs(proc, hashBuild.Conditions) if err != nil { return err @@ -281,7 +289,7 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz } continue } - if err := ctr.spillBatchBounded(proc, bat, spillFiles, execs, analyzer, true); err != nil { + if err := ctr.spillBatchWithPressure(proc, bat, spillFiles, execs, analyzer, true); err != nil { return err } if err := ctr.hashmapBuilder.CleanCopiedBatchAt(0, proc); err != nil { @@ -334,7 +342,7 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz // proof. Drain them under that lease, then retry the direct proof // after their source reservations have been released. if spillMode || - !errors.Is(directProofErr, process.ErrHashBuildBudgetAdmission) || + !IsRetryableMemoryCapacity(directProofErr) || len(ctr.hashmapBuilder.Batches.Buf) == 0 { return directProofErr } @@ -352,7 +360,7 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz // is admitted. If that proof does not fit, do not copy it: switch // to the already-proven direct-spill path. if err := ctr.ensureRetainedSpillScratchReservation(result.Batch, analyzer); err != nil { - if !errors.Is(err, process.ErrHashBuildBudgetAdmission) { + if !IsRetryableMemoryCapacity(err) { return err } if err := startSpill(); err != nil { @@ -363,7 +371,7 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz } // If in spill mode, spill this batch directly to open files. if spillMode { - err := ctr.spillBatchBounded(proc, result.Batch, spillFiles, ctr.spillExprExecs, analyzer, false) + err := ctr.spillBatchWithPressure(proc, result.Batch, spillFiles, ctr.spillExprExecs, analyzer, false) if err != nil { return err } @@ -387,14 +395,14 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz retainedMemBefore := ctr.hashmapBuilder.Batches.MemSize err = ctr.hashmapBuilder.copyBuildBatch(result.Batch, proc) if err != nil { - if hashBuild.IsShuffle && errors.Is(err, process.ErrHashBuildBudgetAdmission) { + if hashBuild.IsShuffle && IsRetryableMemoryCapacity(err) { // The source batch is still owned by the upstream operator. Do // not retry CopyIntoBatches (or increment row count again); enter // spill recovery and write this batch directly. if err := startSpill(); err != nil { return err } - if err := ctr.spillBatchBounded(proc, result.Batch, spillFiles, ctr.spillExprExecs, analyzer, false); err != nil { + if err := ctr.spillBatchWithPressure(proc, result.Batch, spillFiles, ctr.spillExprExecs, analyzer, false); err != nil { return err } continue @@ -480,7 +488,7 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz "HashBuildRuntimeFilterCollectionFallbacks", 1) } if err != nil { - if !hashBuild.IsShuffle || !errors.Is(err, process.ErrHashBuildBudgetAdmission) { + if !hashBuild.IsShuffle || !IsRetryableMemoryCapacity(err) { return err } if !rebuildSafe { diff --git a/pkg/sql/colexec/hashbuild/build_test.go b/pkg/sql/colexec/hashbuild/build_test.go index d34ec4f562a5e..cb8695660b846 100644 --- a/pkg/sql/colexec/hashbuild/build_test.go +++ b/pkg/sql/colexec/hashbuild/build_test.go @@ -2906,6 +2906,77 @@ func TestShuffleHashBuildSpillsExpressionKey(t *testing.T) { bindProc.Free() } +func TestShuffleHashBuildAccountedSpillLifecycle(t *testing.T) { + tc := newTestCase( + t, + []bool{false}, + []types.Type{types.T_int64.ToType()}, + []*plan.Expr{newExpr(0, types.T_int64.ToType())}, + ) + tc.arg.IsShuffle = true + tc.arg.ShuffleIdx = 0 + tc.arg.SpillThreshold = 1 + tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{ + Tag: tc.arg.JoinMapTag + 4_500, + } + tc.arg.SetChildren([]vm.Operator{tc.marg}) + const limit = uint64(8 << 20) + budget := process.MustNewHashBuildBudget(limit, limit) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 256) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + require.NoError(t, tc.arg.SetAllocationAccount(account)) + require.NoError(t, tc.marg.Prepare(tc.proc)) + require.NoError(t, tc.arg.Prepare(tc.proc)) + tc.arg.ctr.hashmapBuilder.setBudget(generation) + + build := newBatch(tc.types, tc.proc, colexec.DefaultBatchSize) + tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly( + build, + nil, + tc.proc.Mp(), + ) + tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly( + nil, + nil, + tc.proc.Mp(), + ) + _, err = vm.Exec(tc.arg, tc.proc) + require.NoError(t, err) + result, err := message.ReceiveJoinMapResult( + tc.arg.JoinMapTag, + true, + tc.arg.ShuffleIdx, + tc.proc.GetMessageBoard(), + tc.proc.Ctx, + ) + require.NoError(t, err) + require.True(t, result.IsSuccess()) + jm := result.JoinMap() + require.NotNil(t, jm) + require.True(t, jm.IsSpilled()) + require.Equal(t, int64(colexec.DefaultBatchSize), jm.GetRowCount()) + payload, err := jm.TakeSpillBuildPayload() + require.NoError(t, err) + require.NoError(t, payload.Close()) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.Zero(t, generation.SpillDiskUsed()) + require.Zero(t, generation.SpillFDUsed()) + + tc.arg.Reset(tc.proc, false, nil) + tc.marg.Reset(tc.proc, false, nil) + require.NoError(t, tc.arg.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + tc.arg.Free(tc.proc, false, nil) + tc.proc.Free() + require.Zero(t, tc.proc.Mp().CurrNB()) +} + func TestShuffleHashBuildResizeRejectReleasesPartialMapAndSpills(t *testing.T) { tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) tc.arg.IsShuffle = true diff --git a/pkg/sql/colexec/hashbuild/dedup_memory.go b/pkg/sql/colexec/hashbuild/dedup_memory.go new file mode 100644 index 0000000000000..67634b5d88d5f --- /dev/null +++ b/pkg/sql/colexec/hashbuild/dedup_memory.go @@ -0,0 +1,100 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package hashbuild + +import ( + "math" + + "github.com/matrixorigin/matrixone/pkg/common/bitmap" + "github.com/matrixorigin/matrixone/pkg/common/mpool" +) + +// makeDedupSlice keeps the legacy allocation unchanged and moves only the +// activated, data-scaled owner to exact off-heap accounting. +func makeDedupSlice[T any]( + hb *HashmapBuilder, + n int, + mp *mpool.MPool, + site mpool.AllocationSite, +) ([]T, error) { + if n < 0 { + return nil, mpool.ErrAllocationAccountInvalid + } + if hb.mapAllocationAccount == nil { + return make([]T, n), nil + } + return mpool.MakeSliceAccounted[T]( + n, + mp, + hb.mapAllocationAccount, + HashBuildAllocationOwner, + site, + ) +} + +func freeDedupSlice[T any](hb *HashmapBuilder, values []T, mp *mpool.MPool) { + if hb.mapAllocationAccount != nil && cap(values) > 0 { + mpool.FreeSlice(mp, values) + } +} + +func (hb *HashmapBuilder) newDedupBitmap( + rows int, + mp *mpool.MPool, + site mpool.AllocationSite, +) (*bitmap.Bitmap, error) { + if rows < 0 || rows > math.MaxInt-63 { + return nil, mpool.ErrAllocationAccountInvalid + } + bm := &bitmap.Bitmap{} + if hb.mapAllocationAccount == nil { + bm.InitWithSize(int64(rows)) + return bm, nil + } + words := (rows + 63) / 64 + storage, err := mpool.MakeSliceAccounted[uint64]( + words, + mp, + hb.mapAllocationAccount, + HashBuildAllocationOwner, + site, + ) + if err != nil { + return nil, err + } + bm.InstallExternalStorage(storage) + bm.InitWithSize(int64(rows)) + return bm, nil +} + +func releaseDedupBitmap(bm *bitmap.Bitmap, mp *mpool.MPool) { + if bm == nil || !bm.HasExternalStorage() { + return + } + storage := bm.ReleaseExternalStorage() + if cap(storage) > 0 { + mpool.FreeSlice(mp, storage) + } +} + +func (hb *HashmapBuilder) freeIgnoreRows(mp *mpool.MPool) { + releaseDedupBitmap(hb.IgnoreRows, mp) + hb.IgnoreRows = nil +} + +func (hb *HashmapBuilder) freeDelRows(mp *mpool.MPool) { + releaseDedupBitmap(hb.DelRows, mp) + hb.DelRows = nil +} diff --git a/pkg/sql/colexec/hashbuild/errors.go b/pkg/sql/colexec/hashbuild/errors.go index 2a1b2bf46b6aa..39f432e147eea 100644 --- a/pkg/sql/colexec/hashbuild/errors.go +++ b/pkg/sql/colexec/hashbuild/errors.go @@ -48,14 +48,14 @@ func TerminalBudgetError(ctx context.Context, err error) error { reason := terminalBudgetReason(budgetErr.Message) var resource, action string - switch budgetErr.Resource { - case process.HashBuildBudgetResourceMemory: + switch budgetErr.Component { + case process.HashBuildBudgetComponentMemory: resource = "memory" action = "reduce join build width or query concurrency, increase processLimitationSize, or lower join_spill_mem for an eligible shuffle join; automatic spill can still exhaust recovery headroom for wide or skewed partitions" - case process.HashBuildBudgetResourceSpillDisk: + case process.HashBuildBudgetComponentSpillDisk: resource = "spill disk" action = "free spill storage or increase processLimitationSpillSize" - case process.HashBuildBudgetResourceSpillFD: + case process.HashBuildBudgetComponentSpillFD: resource = "spill file descriptor" action = "reduce concurrent spill work or raise the CN open-file limit" default: diff --git a/pkg/sql/colexec/hashbuild/errors_test.go b/pkg/sql/colexec/hashbuild/errors_test.go index b6d890a976cbb..d4735e2734c90 100644 --- a/pkg/sql/colexec/hashbuild/errors_test.go +++ b/pkg/sql/colexec/hashbuild/errors_test.go @@ -32,18 +32,18 @@ func TestTerminalBudgetError(t *testing.T) { }) for _, tc := range []struct { - name string - resource process.HashBuildBudgetResource - want []string + name string + component process.HashBuildBudgetComponent + want []string }{ - {"memory", process.HashBuildBudgetResourceMemory, []string{"memory", "requested=3", "used=5", "limit=7", "build width", "processLimitationSize", "join_spill_mem", "recovery headroom"}}, - {"spill disk", process.HashBuildBudgetResourceSpillDisk, []string{"spill disk", "requested=3", "used=5", "limit=7", "processLimitationSpillSize"}}, - {"spill fd", process.HashBuildBudgetResourceSpillFD, []string{"spill file descriptor", "requested=3", "used=5", "limit=7", "open-file limit"}}, + {"memory", process.HashBuildBudgetComponentMemory, []string{"memory", "requested=3", "used=5", "limit=7", "build width", "processLimitationSize", "join_spill_mem", "recovery headroom"}}, + {"spill disk", process.HashBuildBudgetComponentSpillDisk, []string{"spill disk", "requested=3", "used=5", "limit=7", "processLimitationSpillSize"}}, + {"spill fd", process.HashBuildBudgetComponentSpillFD, []string{"spill file descriptor", "requested=3", "used=5", "limit=7", "open-file limit"}}, } { t.Run(tc.name, func(t *testing.T) { err := TerminalBudgetError(context.Background(), &process.HashBuildBudgetError{ Kind: process.HashBuildBudgetErrorAdmission, - Resource: tc.resource, + Component: tc.component, Requested: 3, Used: 5, Cap: 7, @@ -72,7 +72,7 @@ func TestTerminalBudgetError(t *testing.T) { t.Run("resource admission keeps spill depth context", func(t *testing.T) { err := TerminalBudgetError(context.Background(), &process.HashBuildBudgetError{ Kind: process.HashBuildBudgetErrorAdmission, - Resource: process.HashBuildBudgetResourceMemory, + Component: process.HashBuildBudgetComponentMemory, Requested: 3, Used: 5, Cap: 7, diff --git a/pkg/sql/colexec/hashbuild/expression_memory.go b/pkg/sql/colexec/hashbuild/expression_memory.go index a9664ca57a2e0..07f39df0c6cd5 100644 --- a/pkg/sql/colexec/hashbuild/expression_memory.go +++ b/pkg/sql/colexec/hashbuild/expression_memory.go @@ -75,6 +75,26 @@ func NewAllocationAccountedExpressionExecutors( ) } +// NewAllocationAccountedExpressionExecutorsForAccount is the join-consumer +// entry point. It derives the same owner-scoped expression provenance used by +// HashBuild, so probe and rebuild expressions share one exact generation. +func NewAllocationAccountedExpressionExecutorsForAccount( + proc *process.Process, + exprs []*plan.Expr, + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, +) ([]colexec.ExpressionExecutor, error) { + allocation, err := colexec.NewExpressionAllocationAccount(account, owner) + if err != nil { + return nil, err + } + return NewAllocationAccountedExpressionExecutors( + proc, + exprs, + allocation, + ) +} + func expressionSetAllocationClosed(exprs []*plan.Expr) bool { for _, expr := range exprs { if !expressionAllocationClosed(expr) { diff --git a/pkg/sql/colexec/hashbuild/hashmap.go b/pkg/sql/colexec/hashbuild/hashmap.go index 851cafcbf2291..f79c0d5eaeba0 100644 --- a/pkg/sql/colexec/hashbuild/hashmap.go +++ b/pkg/sql/colexec/hashbuild/hashmap.go @@ -124,7 +124,8 @@ func (hb *HashmapBuilder) GetJoinMap(mp *mpool.MPool) *message.JoinMap { } sels := hb.Sels hb.Sels = message.GroupSels{} - jm := message.NewJoinMap(sels, hb.IntHashMap, hb.StrHashMap, hb.DelRows, hb.Batches.Buf, mp) + jmDelRows := hb.DelRows + jm := message.NewJoinMap(sels, hb.IntHashMap, hb.StrHashMap, jmDelRows, hb.Batches.Buf, mp) jm.SetHasNullKey(hb.HasNullKey) hb.IntHashMap = nil hb.StrHashMap = nil @@ -134,11 +135,12 @@ func (hb *HashmapBuilder) GetJoinMap(mp *mpool.MPool) *message.JoinMap { // Drop budgeted cached backing before transferring the encompassing aux // reservation to a consumer that may free it immediately after publication. hb.detachAndPruneCachedIterators() - hb.IgnoreRows = nil + hb.freeIgnoreRows(mp) hb.uniqueSels = nil hb.curVecs = nil release := hb.detachReservations() jm.SetMemoryRelease(func() { + releaseDedupBitmap(jmDelRows, mp) release() }) return jm @@ -247,8 +249,8 @@ func (hb *HashmapBuilder) Reset(proc *process.Process, hashTableHasNotSent bool) hb.Batches.Reset() hb.IntHashMap = nil hb.StrHashMap = nil - hb.IgnoreRows = nil - hb.DelRows = nil + hb.freeIgnoreRows(proc.Mp()) + hb.freeDelRows(proc.Mp()) for i := range hb.UniqueJoinKeys { if hb.UniqueJoinKeys[i] != nil { hb.UniqueJoinKeys[i].Free(proc.Mp()) @@ -273,6 +275,8 @@ func (hb *HashmapBuilder) Free(proc *process.Process) { hb.cachedStrIterator = nil hb.FreeHashMapAndBatches(proc) hb.FreeTemporaryVectors(proc) + hb.freeIgnoreRows(proc.Mp()) + hb.freeDelRows(proc.Mp()) hb.needDupVec = false hb.HasNullKey = false hb.Batches.Reset() @@ -326,6 +330,8 @@ func (hb *HashmapBuilder) FreeHashMapAndBatches(proc *process.Process) { } hb.Sels.Free(proc.Mp()) hb.Batches.Clean(proc.Mp()) + hb.freeIgnoreRows(proc.Mp()) + hb.freeDelRows(proc.Mp()) hb.releaseReservations() } @@ -774,8 +780,24 @@ func (hb *HashmapBuilder) buildHashmap( } if hb.IsDedup && (hb.OnDuplicateAction == plan.Node_IGNORE || dedupBuildKeepLast) { - hb.IgnoreRows = &bitmap.Bitmap{} - hb.IgnoreRows.InitWithSize(int64(hb.InputBatchRowCount)) + hb.IgnoreRows, err = hb.newDedupBitmap( + hb.InputBatchRowCount, + proc.Mp(), + HashBuildAllocationSiteDedupIgnoreBitmap, + ) + if err != nil { + return err + } + } + if hb.delColIdx != -1 && hb.DelRows == nil { + hb.DelRows, err = hb.newDedupBitmap( + hb.InputBatchRowCount, + proc.Mp(), + HashBuildAllocationSiteDedupDeleteBitmap, + ) + if err != nil { + return err + } } var ( @@ -790,15 +812,48 @@ func (hb *HashmapBuilder) buildHashmap( ignoreCandidateOwnsKey []bool ignoreCandidateOldKey []*vector.Vector ) + cleanupDedupScratch := func() { + freeDedupSlice(hb, lastRows, proc.Mp()) + lastRows = nil + freeDedupSlice(hb, ignoreSurvivorRows, proc.Mp()) + ignoreSurvivorRows = nil + freeDedupSlice(hb, ignoreSurvivorOwnsKey, proc.Mp()) + ignoreSurvivorOwnsKey = nil + } + defer cleanupDedupScratch() if dedupBuildKeepLast { - lastRows = make([]int64, hb.InputBatchRowCount+1) + lastRows, err = makeDedupSlice[int64]( + hb, + hb.InputBatchRowCount+1, + proc.Mp(), + HashBuildAllocationSiteDedupLastRows, + ) + if err != nil { + return err + } for i := range lastRows { lastRows[i] = -1 } } if hb.IsDedup && hb.OnDuplicateAction == plan.Node_IGNORE && hb.delColIdx >= 0 { - ignoreSurvivorRows = make([]int64, hb.InputBatchRowCount+1) - ignoreSurvivorOwnsKey = make([]bool, hb.InputBatchRowCount+1) + ignoreSurvivorRows, err = makeDedupSlice[int64]( + hb, + hb.InputBatchRowCount+1, + proc.Mp(), + HashBuildAllocationSiteDedupSurvivorRows, + ) + if err != nil { + return err + } + ignoreSurvivorOwnsKey, err = makeDedupSlice[bool]( + hb, + hb.InputBatchRowCount+1, + proc.Mp(), + HashBuildAllocationSiteDedupSurvivorOwnsKey, + ) + if err != nil { + return err + } ignoreBuildGroups = make([]uint64, hashmap.UnitLimit) ignoreBuildZvals = make([]int64, hashmap.UnitLimit) ignoreCandidateOwnsKey = make([]bool, hashmap.UnitLimit) @@ -1096,6 +1151,7 @@ buildUnits: hb.InputBatchRowCount = totalRowCount } hb.hashMapRowCount = hb.InputBatchRowCount + cleanupDedupScratch() hb.resetHashStateForRebuild(proc) needUniqueVec, err = hb.prepareCanonicalRuntimeFilterCollection( runtimeFilterRequested) @@ -1124,7 +1180,8 @@ buildUnits: } hb.InputBatchRowCount = hb.Batches.RowCount() hb.hashMapRowCount = hb.InputBatchRowCount - hb.DelRows = nil + cleanupDedupScratch() + hb.freeDelRows(proc.Mp()) hb.resetHashStateForRebuild(proc) needUniqueVec, err = hb.prepareCanonicalRuntimeFilterCollection( runtimeFilterRequested) @@ -1136,8 +1193,18 @@ buildUnits: if hb.delColIdx != -1 { if hb.DelRows == nil { - hb.DelRows = &bitmap.Bitmap{} - hb.DelRows.InitWithSize(int64(max(cardinality, uint64(hb.Batches.RowCount())))) + delRows := max(cardinality, uint64(hb.Batches.RowCount())) + if delRows > uint64(math.MaxInt) { + return process.ErrHashBuildBudgetInvalid + } + hb.DelRows, err = hb.newDedupBitmap( + int(delRows), + proc.Mp(), + HashBuildAllocationSiteDedupDeleteBitmap, + ) + if err != nil { + return err + } } // Scan every build row, including the delete-only rows appended by @@ -1240,7 +1307,7 @@ func (hb *HashmapBuilder) resetHashStateForRebuild(proc *process.Process) { hb.executors[i].ResetForNextQuery() } } - hb.IgnoreRows = nil + hb.freeIgnoreRows(proc.Mp()) } // FreeHashMapOnly discards a partial hash build while preserving the copied @@ -1249,7 +1316,7 @@ func (hb *HashmapBuilder) resetHashStateForRebuild(proc *process.Process) { // bounded spill recovery. func (hb *HashmapBuilder) FreeHashMapOnly(proc *process.Process) { hb.resetHashStateForRebuild(proc) - hb.DelRows = nil + hb.freeDelRows(proc.Mp()) if hb.auxReservation != nil { hb.auxReservation.Release() hb.auxReservation = nil @@ -1261,11 +1328,19 @@ func (hb *HashmapBuilder) keepDiscardedRowsForDelete(proc *process.Process) erro return hb.Batches.Shrink(hb.IgnoreRows, proc) } - activeRows := hb.IgnoreRows.Clone() - activeRows.Negate() - activeCount := activeRows.Count() + activeCount := int(hb.IgnoreRows.Len()) - hb.IgnoreRows.Count() - discardedWithDeletes := make([]int32, 0, hb.IgnoreRows.Count()) + discardedStorage, err := makeDedupSlice[int32]( + hb, + hb.IgnoreRows.Count(), + proc.Mp(), + HashBuildAllocationSiteDedupDiscardedRows, + ) + if err != nil { + return err + } + defer freeDedupSlice(hb, discardedStorage, proc.Mp()) + discardedWithDeletes := discardedStorage[:0] itr := hb.IgnoreRows.Iterator() for itr.HasNext() { row := itr.Next() @@ -1308,8 +1383,19 @@ func (hb *HashmapBuilder) keepDiscardedRowsForDelete(proc *process.Process) erro return err } - hb.DelRows = &bitmap.Bitmap{} - hb.DelRows.InitWithSize(int64(activeCount + len(discardedWithDeletes))) + newRows := activeCount + len(discardedWithDeletes) + if hb.DelRows == nil { + hb.DelRows, err = hb.newDedupBitmap( + newRows, + proc.Mp(), + HashBuildAllocationSiteDedupDeleteBitmap, + ) + if err != nil { + return err + } + } else { + hb.DelRows.InitWithSize(int64(newRows)) + } for i := range discardedWithDeletes { hb.DelRows.Add(uint64(activeCount + i)) } @@ -1323,9 +1409,27 @@ func (hb *HashmapBuilder) makeDeleteOnlyBatch(rows []int32, proc *process.Proces } bat := batch.NewOffHeapWithSize(len(hb.Batches.Buf[0].Vecs)) + if hb.mapAllocationAccount != nil { + selection, err := vector.NewAllocationAccountSelectionWithBitmaps( + hb.mapAllocationAccount, + HashBuildAllocationOwner, + HashBuildAllocationSiteDedupDeleteOnlyData, + HashBuildAllocationSiteDedupDeleteOnlyArea, + HashBuildAllocationSiteDedupDeleteOnlyNulls, + HashBuildAllocationSiteDedupDeleteOnlyGrouping, + ) + if err != nil { + bat.Clean(proc.Mp()) + return nil, err + } + if err = bat.SetAllocationAccount(selection); err != nil { + bat.Clean(proc.Mp()) + return nil, err + } + } bat.Attrs = hb.Batches.Buf[0].Attrs for colIdx, vec := range hb.Batches.Buf[0].Vecs { - bat.Vecs[colIdx] = vector.NewOffHeapVecWithType(*vec.GetType()) + bat.SetVector(int32(colIdx), vector.NewOffHeapVecWithType(*vec.GetType())) } cleanOnErr := true diff --git a/pkg/sql/colexec/hashbuild/hashmap_test.go b/pkg/sql/colexec/hashbuild/hashmap_test.go index 4dc5e128ee731..200ced7c14b9e 100644 --- a/pkg/sql/colexec/hashbuild/hashmap_test.go +++ b/pkg/sql/colexec/hashbuild/hashmap_test.go @@ -1147,10 +1147,9 @@ func TestReserveBuildAuxChargesOneRetainedCopy(t *testing.T) { hb.Batches.Buf = nil } -func TestReserveBuildAuxDoesNotStackAccountedGroupSels(t *testing.T) { +func TestReserveBuildAuxExactDoesNotDuplicatePhysicalOrHeadroomOwners(t *testing.T) { const rows = 10_000 - const iteratorScratch = uint64(640 << 10) - want := uint64(rows)*48 + iteratorScratch + const want = uint64(1) budget := process.MustNewHashBuildBudget(want, want) generation, err := budget.OpenGeneration(1) require.NoError(t, err) @@ -1166,7 +1165,7 @@ func TestReserveBuildAuxDoesNotStackAccountedGroupSels(t *testing.T) { hb.setBudget(generation) require.NoError(t, hb.reserveBuildAux(false, true)) - require.Equal(t, want, generation.Used()) + require.Zero(t, generation.Used()) require.Zero(t, generation.Snapshot().AllocationUsed) hb.releaseReservations() require.Zero(t, generation.Used()) @@ -2129,6 +2128,114 @@ func TestDedupBuildKeepLastPreservesDeleteOnlyRows(t *testing.T) { require.Equal(t, int32(100), markers[2]) } +func TestAccountedDedupScratchAndDeleteBitmapFollowJoinMapLifetime(t *testing.T) { + const capBytes = uint64(64 << 20) + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 256) + require.NoError(t, err) + account, err := registry.OpenWithController(capBytes, generation) + require.NoError(t, err) + + var op HashBuild + op.NeedHashMap = true + require.NoError(t, op.SetAllocationAccount(account)) + hb := &op.ctr.hashmapBuilder + hb.setBudget(generation) + hb.IsDedup = true + hb.DedupBuildKeepLast = true + hb.OnDuplicateAction = plan.Node_FAIL + hb.DedupColName = "id" + hb.DedupColTypes = []plan.Type{newExpr(0, types.T_int32.ToType()).Typ} + require.NoError(t, hb.Prepare( + []*plan.Expr{newExpr(0, types.T_int32.ToType())}, + -1, + 2, + []int32{2}, + proc, + )) + input := makeIntKeyValueBatchWithMarker( + proc, + []int32{1, 1, 2}, + []int32{10, 20, 30}, + []int32{100, 0, 0}, + []uint64{1, 2}, + ) + require.NoError(t, hb.copyBuildBatch(input, proc)) + hb.InputBatchRowCount = input.RowCount() + input.Clean(proc.Mp()) + + require.NoError(t, hb.BuildHashmap(false, false, false, proc)) + require.NotNil(t, hb.DelRows) + require.True(t, hb.DelRows.HasExternalStorage()) + require.True(t, hb.DelRows.Contains(2)) + require.Equal(t, generation.Used(), generation.Snapshot().AllocationUsed) + require.Positive(t, account.Snapshot().Used) + + jm := hb.GetJoinMap(proc.Mp()) + require.NotNil(t, jm) + jm.IncRef(1) + hb.Reset(proc, false) + // DelRows remains physically owned by the consumer together with the map. + require.Positive(t, account.Snapshot().Used) + require.True(t, jm.IsDeleted(2)) + jm.Free() + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.NoError(t, op.ClearAllocationAccount(account)) + terminal, first, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.True(t, first) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) +} + +func TestAccountedDedupBitmapExactBoundaryRollsBack(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + for _, tc := range []struct { + name string + cap uint64 + wantErr bool + }{ + {name: "one byte short", cap: 7, wantErr: true}, + {name: "exact", cap: 8}, + } { + t.Run(tc.name, func(t *testing.T) { + budget := process.MustNewHashBuildBudget(tc.cap, tc.cap) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + account, err := registry.OpenWithController(tc.cap, generation) + require.NoError(t, err) + var hb HashmapBuilder + hb.mapAllocationAccount = account + bm, err := hb.newDedupBitmap( + 64, + proc.Mp(), + HashBuildAllocationSiteDedupIgnoreBitmap, + ) + if tc.wantErr { + require.Error(t, err) + require.True(t, IsRetryableMemoryCapacity(err)) + require.Nil(t, bm) + require.Zero(t, account.Snapshot().Used) + } else { + require.NoError(t, err) + require.Equal(t, uint64(8), account.Snapshot().Used) + releaseDedupBitmap(bm, proc.Mp()) + require.Zero(t, account.Snapshot().Used) + } + terminal, _, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) + }) + } +} + // TestDedupBuildKeepLastMarksConflictBucketForDiscardedFanout reproduces the // REPLACE multi-UK fan-out case (issue #24428) at the hashbuild layer: one new // row (same new PK) fans out to several build rows that carry DIFFERENT old diff --git a/pkg/sql/colexec/hashbuild/pressure.go b/pkg/sql/colexec/hashbuild/pressure.go new file mode 100644 index 0000000000000..1fbe2d4ff3244 --- /dev/null +++ b/pkg/sql/colexec/hashbuild/pressure.go @@ -0,0 +1,207 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package hashbuild + +import ( + "errors" + "fmt" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/vm/process" +) + +// MemoryPressureReason is the single classification used by HashBuild and all +// spilled join consumers. Only Capacity may enter reclaim/spill/reduce/degrade +// control flow. Sealed and invariant failures are lifecycle bugs and must +// remain terminal. HashBuildBudgetError now exposes disjoint lifecycle and +// capacity identities; this classifier is the one control-flow boundary for +// physical-account and legacy-budget failures. +type MemoryPressureReason uint8 + +const ( + MemoryPressureNone MemoryPressureReason = iota + MemoryPressureCapacity + MemoryPressureSealed + MemoryPressureMismatch + MemoryPressureAllocatorLimit + MemoryPressureInvariant + MemoryPressureInvalid + MemoryPressureMinimumUnit + MemoryPressureSpillDiskLimit + MemoryPressureSpillFDLimit +) + +func MemoryPressureReasonOf(err error) MemoryPressureReason { + if err == nil { + return MemoryPressureNone + } + var minimum *MinimumAllocationPressureError + if errors.As(err, &minimum) { + return MemoryPressureMinimumUnit + } + + var budgetErr *process.HashBuildBudgetError + if errors.As(err, &budgetErr) { + switch budgetErr.Kind { + case process.HashBuildBudgetErrorAdmission: + switch budgetErr.Component { + case process.HashBuildBudgetComponentSpillDisk: + return MemoryPressureSpillDiskLimit + case process.HashBuildBudgetComponentSpillFD: + return MemoryPressureSpillFDLimit + } + return MemoryPressureCapacity + case process.HashBuildBudgetErrorClosed: + return MemoryPressureSealed + case process.HashBuildBudgetErrorInvalid, + process.HashBuildBudgetErrorCeilingMissing: + return MemoryPressureInvalid + default: + return MemoryPressureInvalid + } + } + + switch mpool.AllocationFailureReasonOf(err) { + case mpool.AllocationFailureCapacity: + return MemoryPressureCapacity + case mpool.AllocationFailureSealed, + mpool.AllocationFailureSuspended: + return MemoryPressureSealed + case mpool.AllocationFailureMismatch: + return MemoryPressureMismatch + case mpool.AllocationFailureAllocatorLimit: + return MemoryPressureAllocatorLimit + case mpool.AllocationFailureInvariant: + return MemoryPressureInvariant + } + + // A few compatibility call sites return the bare sentinel. Check closed + // first; structured errors above never rely on the legacy Is alias. + if errors.Is(err, process.ErrHashBuildBudgetClosed) { + return MemoryPressureSealed + } + if errors.Is(err, process.ErrHashBuildBudgetAdmission) || + errors.Is(err, process.ErrHashBuildBudgetRejected) { + return MemoryPressureCapacity + } + if errors.Is(err, process.ErrHashBuildBudgetInvalid) || + errors.Is(err, process.ErrHashBuildCeilingMissing) { + return MemoryPressureInvalid + } + return MemoryPressureNone +} + +func IsRetryableMemoryCapacity(err error) bool { + return MemoryPressureReasonOf(err) == MemoryPressureCapacity +} + +// MinimumAllocationPressureError means the operation has already reclaimed +// optional storage and reduced itself to one indivisible input unit. It does +// not unwrap the last capacity error: callers must not mistake the terminal +// boundary for another retryable admission failure. +type MinimumAllocationPressureError struct { + Owner string + Site string + Response string + Used uint64 + Limit uint64 +} + +func (e *MinimumAllocationPressureError) Error() string { + if e == nil { + return "minimum allocation cannot be admitted" + } + return fmt.Sprintf( + "minimum allocation cannot be admitted: owner=%s site=%s response=%s used=%d limit=%d", + e.Owner, + e.Site, + e.Response, + e.Used, + e.Limit, + ) +} + +func NewMinimumAllocationPressureError( + owner string, + site string, + account *mpool.AllocationAccount, +) error { + err := &MinimumAllocationPressureError{ + Owner: owner, + Site: site, + Response: "reclaim/reduce/degrade exhausted", + } + if account != nil { + snapshot := account.Snapshot() + err.Used = snapshot.Used + err.Limit = snapshot.Limit + } + return err +} + +// PressureProgress is the monotonic proof required before retrying one +// logical operation. A retry is legal only after memory was reclaimed, spill +// state advanced, the input unit shrank, or optional work was disabled. +type PressureProgress struct { + Used uint64 + SpillEpoch uint64 + InputUnits int + OptionalDisabled bool +} + +type PressureRetryGuard struct { + previous PressureProgress + attempts int + limit int +} + +func NewPressureRetryGuard(initial PressureProgress, limit int) *PressureRetryGuard { + if limit <= 0 { + limit = 64 + } + return &PressureRetryGuard{previous: initial, limit: limit} +} + +func (g *PressureRetryGuard) Advance(next PressureProgress) error { + if g == nil || next.InputUnits < 0 { + return process.ErrHashBuildBudgetInvalid + } + if g.attempts >= g.limit { + return fmt.Errorf( + "%w: memory-pressure retry limit exceeded", + process.ErrHashBuildBudgetInvalid, + ) + } + progress := next.Used < g.previous.Used || + next.SpillEpoch > g.previous.SpillEpoch || + (g.previous.InputUnits > 0 && next.InputUnits < g.previous.InputUnits) || + (!g.previous.OptionalDisabled && next.OptionalDisabled) + if !progress { + return fmt.Errorf( + "%w: memory-pressure retry made no progress", + process.ErrHashBuildBudgetInvalid, + ) + } + g.previous = next + g.attempts++ + return nil +} + +func (g *PressureRetryGuard) Attempts() int { + if g == nil { + return 0 + } + return g.attempts +} diff --git a/pkg/sql/colexec/hashbuild/pressure_test.go b/pkg/sql/colexec/hashbuild/pressure_test.go new file mode 100644 index 0000000000000..15127329c31a3 --- /dev/null +++ b/pkg/sql/colexec/hashbuild/pressure_test.go @@ -0,0 +1,69 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package hashbuild + +import ( + "fmt" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/vm/process" + "github.com/stretchr/testify/require" +) + +func TestMemoryPressureReasonSeparatesCapacityFromLifecycle(t *testing.T) { + tests := []struct { + err error + reason MemoryPressureReason + }{ + {nil, MemoryPressureNone}, + {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorAdmission}, MemoryPressureCapacity}, + {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorAdmission, Component: process.HashBuildBudgetComponentMemory}, MemoryPressureCapacity}, + {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorAdmission, Component: process.HashBuildBudgetComponentSpillDisk}, MemoryPressureSpillDiskLimit}, + {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorAdmission, Component: process.HashBuildBudgetComponentSpillFD}, MemoryPressureSpillFDLimit}, + {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorClosed}, MemoryPressureSealed}, + {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorInvalid}, MemoryPressureInvalid}, + {fmt.Errorf("wrapped: %w", process.ErrHashBuildBudgetAdmission), MemoryPressureCapacity}, + {mpool.ErrAllocationAccountCapacity, MemoryPressureCapacity}, + {mpool.ErrAllocationMetadataSlots, MemoryPressureCapacity}, + {mpool.ErrAllocationAccountSealed, MemoryPressureSealed}, + {mpool.ErrAllocationAccountMismatch, MemoryPressureMismatch}, + {mpool.ErrAllocationAllocatorLimit, MemoryPressureAllocatorLimit}, + {mpool.ErrAllocationAccountInvariant, MemoryPressureInvariant}, + {NewMinimumAllocationPressureError("hashbuild", "spill", nil), MemoryPressureMinimumUnit}, + } + for _, test := range tests { + require.Equal(t, test.reason, MemoryPressureReasonOf(test.err)) + require.Equal(t, test.reason == MemoryPressureCapacity, IsRetryableMemoryCapacity(test.err)) + } +} + +func TestPressureRetryGuardRequiresMonotonicProgress(t *testing.T) { + initial := PressureProgress{Used: 100, SpillEpoch: 1, InputUnits: 16} + for _, next := range []PressureProgress{ + {Used: 99, SpillEpoch: 1, InputUnits: 16}, + {Used: 99, SpillEpoch: 2, InputUnits: 16}, + {Used: 99, SpillEpoch: 2, InputUnits: 8}, + {Used: 99, SpillEpoch: 2, InputUnits: 8, OptionalDisabled: true}, + } { + guard := NewPressureRetryGuard(initial, 1) + require.NoError(t, guard.Advance(next)) + require.Equal(t, 1, guard.Attempts()) + require.Error(t, guard.Advance(next), "the retry limit remains fail-closed") + } + guard := NewPressureRetryGuard(initial, 4) + require.Error(t, guard.Advance(initial)) + require.Zero(t, guard.Attempts()) +} diff --git a/pkg/sql/colexec/hashbuild/spill.go b/pkg/sql/colexec/hashbuild/spill.go index fbed914771f9f..cd0e2a84a0689 100644 --- a/pkg/sql/colexec/hashbuild/spill.go +++ b/pkg/sql/colexec/hashbuild/spill.go @@ -30,6 +30,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/fileservice" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/plan/function" "github.com/matrixorigin/matrixone/pkg/vm/process" ) @@ -389,28 +390,6 @@ func (ctr *container) ensureSpillScratchReservationBytes( return nil } -func (ctr *container) ensureDirectSpillScratchReservation( - bat *batch.Batch, - analyzer process.Analyzer, -) error { - need, err := spillBudgetBytes(bat) - if err != nil { - return err - } - return ctr.ensureSpillScratchReservationBytes(need, analyzer) -} - -func (ctr *container) ensureRetainedSpillScratchReservation( - bat *batch.Batch, - analyzer process.Analyzer, -) error { - need, err := spillRetainedBudgetBytes(bat) - if err != nil { - return err - } - return ctr.ensureSpillScratchReservationBytes(need, analyzer) -} - func (ctr *container) growSpillScratchTransient( required uint64, analyzer process.Analyzer, @@ -438,6 +417,381 @@ func (ctr *container) restoreSpillScratchTransient(oldSize uint64, grew bool) er return err } +func (ctr *container) ensureDirectSpillScratchReservation(bat *batch.Batch, analyzer process.Analyzer) error { + if ctr.spillBatchAllocation != nil { + // A borrowed upstream batch is already live and cannot be reclaimed by + // reserving another logical token. Exact scatter/expression allocations + // admit their physical capacities and adapt the unpublished input. The + // retained-copy path below still keeps a one-unit future-progress token + // because choosing to retain is under HashBuild's control. + return nil + } + var ( + need uint64 + err error + ) + need, err = spillBudgetBytes(bat) + if err != nil { + return err + } + return ctr.ensureSpillScratchReservationBytes(need, analyzer) +} + +func (ctr *container) ensureRetainedSpillScratchReservation(bat *batch.Batch, analyzer process.Analyzer) error { + var ( + need uint64 + err error + ) + if ctr.spillBatchAllocation != nil { + need, err = spillMinimumUnitBudgetBytes(bat, ctr.spillConditions) + } else { + need, err = spillRetainedBudgetBytes(bat) + } + if err != nil { + return err + } + return ctr.ensureSpillScratchReservationBytes(need, analyzer) +} + +// spillMinimumUnitBudgetBytes keeps only the headroom for one physical spill +// unit. It derives capacities from actual input values and the closed +// expression family; it neither scales the whole batch nor applies a safety +// multiplier. The token is converted into exact allocations on spill entry. +func spillMinimumUnitBudgetBytes( + bat *batch.Batch, + exprs []*plan.Expr, +) (uint64, error) { + if bat == nil || bat.RowCount() <= 0 { + return 0, nil + } + selected, wire, err := spillMinimumSelectedAndWireBytes(bat) + if err != nil { + return 0, err + } + if wire > math.MaxUint64-24 { + return 0, process.ErrHashBuildBudgetInvalid + } + marshal, err := initialAllocationCapacity(wire + 24) + if err != nil { + return 0, err + } + expression, err := spillMinimumExpressionBytes(exprs, bat) + if err != nil { + return 0, err + } + total := uint64(12) // one hash plus one row id + for _, value := range []uint64{selected, marshal, expression} { + if total > math.MaxUint64-value { + return 0, process.ErrHashBuildBudgetInvalid + } + total += value + } + return total, nil +} + +func spillMinimumSelectedAndWireBytes( + bat *batch.Batch, +) (selected uint64, wire uint64, err error) { + // Batch framing plus Attr/ExtraBuf length prefixes. + wire = 8 + 4 + 4 + 4 + 4 + 4 + uint64(len(bat.ExtraBuf)) + for _, attr := range bat.Attrs { + if wire > math.MaxUint64-4-uint64(len(attr)) { + return 0, 0, process.ErrHashBuildBudgetInvalid + } + wire += 4 + uint64(len(attr)) + } + for _, vec := range bat.Vecs { + if vec == nil || vec.GetType().TypeSize() < 0 { + return 0, 0, process.ErrHashBuildBudgetInvalid + } + data, capErr := initialAllocationCapacity( + uint64(vec.GetType().TypeSize()), + ) + if capErr != nil { + return 0, 0, capErr + } + areaPayload, payloadErr := maxVectorValueBytes(vec) + if payloadErr != nil { + return 0, 0, payloadErr + } + var area uint64 + if areaPayload > types.VarlenaInlineSize { + area, capErr = initialAllocationCapacity(areaPayload) + if capErr != nil { + return 0, 0, capErr + } + } + // Accounted vectors install one null and one grouping word before + // extending their first row. + physical := data + area + 16 + if selected > math.MaxUint64-physical { + return 0, 0, process.ErrHashBuildBudgetInvalid + } + selected += physical + // Physical capacities upper-bound one-row logical data, area, and null + // payload; only the fixed wire framing is added separately. + const vectorFraming = uint64(4 + 1 + types.TSize + 4 + 4 + 4 + 4 + 1) + if wire > math.MaxUint64-vectorFraming-physical { + return 0, 0, process.ErrHashBuildBudgetInvalid + } + wire += vectorFraming + physical + } + return selected, wire, nil +} + +func maxVectorValueBytes(vec *vector.Vector) (uint64, error) { + if vec == nil || !vec.GetType().IsVarlen() || vec.IsConstNull() { + return 0, nil + } + values, _ := vector.MustVarlenaRawData(vec) + rows := vec.Length() + if vec.IsConst() && rows > 0 { + rows = 1 + } + if rows > len(values) { + return 0, process.ErrHashBuildBudgetInvalid + } + var maximum uint64 + for row := 0; row < rows; row++ { + if vec.GetNulls().Contains(uint64(row)) { + continue + } + var length uint64 + if values[row].IsSmall() { + length = uint64(len(values[row].GetByteSlice(nil))) + } else { + _, valueLen := values[row].OffsetLen() + length = uint64(valueLen) + } + if length > maximum { + maximum = length + } + } + return maximum, nil +} + +func spillMinimumExpressionBytes( + exprs []*plan.Expr, + bat *batch.Batch, +) (uint64, error) { + if len(exprs) == 0 { + return 0, nil + } + if !AllocationAccountedExpressionSetSupported(exprs) { + return 0, process.ErrHashBuildBudgetInvalid + } + var total uint64 + for _, expr := range exprs { + bytes, err := spillMinimumExpressionTreeBytes(expr, bat) + if err != nil || total > math.MaxUint64-bytes { + return 0, process.ErrHashBuildBudgetInvalid + } + total += bytes + } + return total, nil +} + +func spillMinimumExpressionTreeBytes( + expr *plan.Expr, + bat *batch.Batch, +) (uint64, error) { + if expr == nil { + return 0, process.ErrHashBuildBudgetInvalid + } + switch node := expr.Expr.(type) { + case *plan.Expr_Col: + return 0, nil + case *plan.Expr_Lit: + return expressionInitialOwnedBytes(expr) + case *plan.Expr_F: + if node.F == nil || node.F.Func == nil { + return 0, process.ErrHashBuildBudgetInvalid + } + var total uint64 + for _, arg := range node.F.Args { + child, err := spillMinimumExpressionTreeBytes(arg, bat) + if err != nil || total > math.MaxUint64-child { + return 0, process.ErrHashBuildBudgetInvalid + } + total += child + } + result, err := spillMinimumExpressionResultBytes(expr, node.F.Args, bat) + if err != nil || total > math.MaxUint64-result { + return 0, process.ErrHashBuildBudgetInvalid + } + total += result + functionID, _ := function.DecodeOverloadID(node.F.Func.Obj) + if functionID == function.CASE { + if total > math.MaxUint64-8 { + return 0, process.ErrHashBuildBudgetInvalid + } + total += 8 + } + return total, nil + default: + return expressionInitialOwnedBytes(expr) + } +} + +func spillMinimumExpressionResultBytes( + expr *plan.Expr, + args []*plan.Expr, + bat *batch.Batch, +) (uint64, error) { + oid := types.T(expr.Typ.Id) + typ := oid.ToType() + if typ.TypeSize() < 0 { + return 0, process.ErrHashBuildBudgetInvalid + } + data, err := initialAllocationCapacity(uint64(typ.TypeSize())) + if err != nil { + return 0, err + } + result := data + 16 + if !typ.IsVarlen() { + return result, nil + } + payload, err := spillExpressionPayloadBytes(expr, args, bat) + if err != nil { + return 0, err + } + if payload > types.VarlenaInlineSize { + area, err := initialAllocationCapacity(payload) + if err != nil || result > math.MaxUint64-area { + return 0, process.ErrHashBuildBudgetInvalid + } + result += area + } + return result, nil +} + +func spillExpressionPayloadBytes( + expr *plan.Expr, + args []*plan.Expr, + bat *batch.Batch, +) (uint64, error) { + if bat == nil || bat.RowCount() < 0 { + return 0, process.ErrHashBuildBudgetInvalid + } + var maximum uint64 + for row := 0; row < bat.RowCount(); row++ { + value, err := spillExpressionPayloadBytesAt(expr, args, bat, row) + if err != nil { + return 0, err + } + if value > maximum { + maximum = value + } + } + return maximum, nil +} + +func spillExpressionPayloadBytesAt( + expr *plan.Expr, + args []*plan.Expr, + bat *batch.Batch, + row int, +) (uint64, error) { + node, ok := expr.Expr.(*plan.Expr_F) + if !ok || node.F == nil || node.F.Func == nil { + return 0, nil + } + functionID, _ := function.DecodeOverloadID(node.F.Func.Obj) + switch functionID { + case function.CONCAT: + var total uint64 + for _, arg := range args { + value, err := spillExpressionArgPayloadBytesAt(arg, bat, row) + if err != nil || total > math.MaxUint64-value { + return 0, process.ErrHashBuildBudgetInvalid + } + total += value + } + return total, nil + case function.CASE: + // CASE's exact selected branch is evaluated later. The largest varlen + // branch value at this same row is a safe one-row bound without combining + // maxima taken from different rows. + var maximum uint64 + for _, arg := range args { + if arg == nil || !types.T(arg.Typ.Id).ToType().IsVarlen() { + continue + } + value, err := spillExpressionArgPayloadBytesAt(arg, bat, row) + if err != nil { + return 0, err + } + if value > maximum { + maximum = value + } + } + return maximum, nil + case function.CAST: + if len(args) == 0 || args[0] == nil { + return 0, process.ErrHashBuildBudgetInvalid + } + if types.T(args[0].Typ.Id).ToType().IsIntOrUint() { + return 20, nil + } + return spillExpressionArgPayloadBytesAt(args[0], bat, row) + default: + return 0, nil + } +} + +func spillExpressionArgPayloadBytesAt( + expr *plan.Expr, + bat *batch.Batch, + row int, +) (uint64, error) { + if expr == nil || bat == nil || row < 0 || row >= bat.RowCount() { + return 0, process.ErrHashBuildBudgetInvalid + } + switch node := expr.Expr.(type) { + case *plan.Expr_Col: + if node.Col == nil || node.Col.ColPos < 0 || + int(node.Col.ColPos) >= len(bat.Vecs) { + return 0, process.ErrHashBuildBudgetInvalid + } + return vectorValueBytesAt(bat.Vecs[node.Col.ColPos], row) + case *plan.Expr_Lit: + if node.Lit == nil || node.Lit.GetIsnull() { + return 0, nil + } + return uint64(len(node.Lit.GetSval())), nil + case *plan.Expr_F: + return spillExpressionPayloadBytesAt(expr, node.F.GetArgs(), bat, row) + default: + return 0, nil + } +} + +func vectorValueBytesAt(vec *vector.Vector, row int) (uint64, error) { + if vec == nil || row < 0 || row >= vec.Length() || + !vec.GetType().IsVarlen() || vec.IsConstNull() { + if vec == nil || row < 0 || row >= vec.Length() { + return 0, process.ErrHashBuildBudgetInvalid + } + return 0, nil + } + index := row + if vec.IsConst() { + index = 0 + } + if vec.GetNulls().Contains(uint64(index)) { + return 0, nil + } + values, _ := vector.MustVarlenaRawData(vec) + if index >= len(values) { + return 0, process.ErrHashBuildBudgetInvalid + } + if values[index].IsSmall() { + return uint64(len(values[index].GetByteSlice(nil))), nil + } + _, length := values[index].OffsetLen() + return uint64(length), nil +} + func (ctr *container) releaseSpillScratchReservation() { if ctr.spillScratchReservation != nil { ctr.spillScratchReservation.Release() @@ -482,6 +836,7 @@ func (ctr *container) dropSpillScratchBuffers() { ctr.spillKeyVecs = nil ctr.spillWriteBuf = bytes.Buffer{} ctr.spillAllocationMP = nil + ctr.spillCoalesceDisabled = false } func growHashBuildSpillSlice[T any]( @@ -756,9 +1111,8 @@ func (ctr *container) spillBatchBounded(proc *process.Process, bat *batch.Batch, } exact := ctr.spillBatchAllocation != nil var ( - need uint64 - externalSource *process.HashBuildReservation - err error + need uint64 + err error ) if exact { if ctr.hashmapBuilder.mapAllocationAccount == nil { @@ -771,18 +1125,19 @@ func (ctr *container) spillBatchBounded(proc *process.Process, bat *batch.Batch, // A pre-spill token is headroom, not physical ownership. Release it // immediately before the exact scratch allocations consume that space. ctr.releaseSpillScratchReservation() - if ctr.hashmapBuilder.budget != nil && !sourceAlreadyCharged { - externalBytes := uint64(bat.Allocated()) - if size := uint64(bat.Size()); size > externalBytes { + if !sourceAlreadyCharged { + // The upstream batch is borrowed, already physically live, and cannot + // be made smaller by rejecting a new logical token. Record it as + // observation only; every new HashBuild-owned byte below is admitted + // by the exact account and the process MPool remains the global guard. + externalBytes := bat.Allocated() + if size := bat.Size(); size > externalBytes { externalBytes = size } - if externalBytes > 0 { - externalSource, err = ctr.hashmapBuilder.budget.Reserve(externalBytes) - if err != nil { - return err - } - defer externalSource.Release() - } + analyzer.GetOpStats().SetMaxExtraStat( + "HashBuildSpillBorrowedSourceBytes", + int64(externalBytes), + ) } } else { need, err = spillScratchBudgetBytes(bat, sourceAlreadyCharged) @@ -917,7 +1272,9 @@ func (ctr *container) spillBatchBounded(proc *process.Process, bat *batch.Batch, // Eval may leave newly allocated child/result vectors cached in the // executor tree. Destroy that tree while both the previous and // candidate reservations are still charged. - ctr.freeSpillExprExecs() + if !exact { + ctr.freeSpillExprExecs() + } return err } if err := checkHashBuildCanceled(proc); err != nil { @@ -994,34 +1351,274 @@ func (ctr *container) spillBatchBounded(proc *process.Process, bat *batch.Batch, } } } - selected.CleanOnlyData() - sels := ctr.spillBucketRowIds[start:end] - n := int(end - start) - var spillErr error - for i, vec := range bat.Vecs { - if vec == nil { - spillErr = process.ErrHashBuildBudgetInvalid - break + cursor := start + for cursor < end { + attemptEnd := end + reclaimedMinimum := false + for { + selected.CleanOnlyData() + sels := ctr.spillBucketRowIds[cursor:attemptEnd] + n := int(attemptEnd - cursor) + var spillErr error + for i, vec := range bat.Vecs { + if vec == nil { + spillErr = process.ErrHashBuildBudgetInvalid + break + } + if spillErr = selected.Vecs[i].PreExtend(n, proc.Mp()); spillErr != nil { + break + } + if spillErr = selected.Vecs[i].UnionInt32(vec, sels, proc.Mp()); spillErr != nil { + break + } + } + if spillErr == nil { + selected.SetRowCount(n) + var file *os.File + file, spillErr = ctr.ensureSpillFile(proc, files, int(bucket)) + if spillErr == nil { + spillErr = ctr.appendSpillRecord( + proc, + file, + int(bucket), + selected, + need, + analyzer, + ) + } + } + selected.CleanOnlyData() + if spillErr == nil { + cursor = attemptEnd + break + } + if !exact || !IsRetryableMemoryCapacity(spillErr) { + return spillErr + } + if err := checkHashBuildCanceled(proc); err != nil { + return err + } + if n > 1 { + attemptEnd = cursor + int32((n+1)/2) + analyzer.GetOpStats().AddExtraStat( + "HashBuildSpillBatchReductions", + 1, + ) + continue + } + if !reclaimedMinimum { + before := ctr.hashmapBuilder.mapAllocationAccount.Snapshot().Used + if err := ctr.reclaimOptionalSpillBuffers( + proc, + files, + analyzer, + ); err != nil { + return err + } + reclaimedMinimum = true + after := ctr.hashmapBuilder.mapAllocationAccount.Snapshot().Used + if after >= before { + analyzer.GetOpStats().AddExtraStat( + "HashBuildSpillMinimumUnitErrors", + 1, + ) + return NewMinimumAllocationPressureError( + "hashbuild", + "spill-selected-or-codec", + ctr.hashmapBuilder.mapAllocationAccount, + ) + } + analyzer.GetOpStats().AddExtraStat( + "HashBuildSpillOptionalReclaims", + 1, + ) + continue + } + analyzer.GetOpStats().AddExtraStat( + "HashBuildSpillMinimumUnitErrors", + 1, + ) + return NewMinimumAllocationPressureError( + "hashbuild", + "spill-selected-or-codec", + ctr.hashmapBuilder.mapAllocationAccount, + ) } - if spillErr = selected.Vecs[i].PreExtend(n, proc.Mp()); spillErr != nil { - break + } + } + return nil +} + +// reclaimOptionalSpillBuffers publishes already completed coalesced records, +// then drops codec/coalesce capacity. The current selected record has not been +// published when this is called, so retrying that one record is idempotent. +func (ctr *container) reclaimOptionalSpillBuffers( + proc *process.Process, + files []*os.File, + analyzer process.Analyzer, +) error { + for bucket, buffer := range ctr.spillAccountedBuckets { + if buffer == nil { + continue + } + if buffer.Len() > 0 { + if bucket >= len(files) || files[bucket] == nil { + return process.ErrHashBuildBudgetInvalid } - if spillErr = selected.Vecs[i].UnionInt32(vec, sels[:n], proc.Mp()); spillErr != nil { - break + if err := ctr.flushPendingSpillBucket( + proc, + files[bucket], + bucket, + analyzer, + ); err != nil { + return err } } - if spillErr == nil { - selected.SetRowCount(n) - var file *os.File - file, spillErr = ctr.ensureSpillFile(proc, files, int(bucket)) - if spillErr == nil { - spillErr = ctr.appendSpillRecord(proc, file, int(bucket), selected, need, analyzer) + buffer.Free() + ctr.spillAccountedBuckets[bucket] = nil + } + ctr.spillCoalesceDisabled = true + if ctr.spillAccountedWrite != nil { + ctr.spillAccountedWrite.Free() + ctr.spillAccountedWrite = nil + } + return nil +} + +func (ctr *container) releaseSpillComputeScratch() { + if ctr.spillBatchAllocation == nil || ctr.spillAllocationMP == nil { + return + } + if cap(ctr.spillHashValues) > 0 { + mpool.FreeSlice(ctr.spillAllocationMP, ctr.spillHashValues) + } + if cap(ctr.spillBucketRowIds) > 0 { + mpool.FreeSlice(ctr.spillAllocationMP, ctr.spillBucketRowIds) + } + ctr.spillHashValues = nil + ctr.spillBucketRowIds = nil + ctr.spillSelection = nil +} + +// spillBatchWithPressure retries only the unpublished prefix of an exact +// spill operation. Hash/expression capacity failures happen before any bucket +// write; selected/codec failures are handled transactionally inside +// spillBatchBounded. Each retry halves the input or reclaims memory, and a +// one-row failure becomes a controlled minimum-unit error. +func (ctr *container) spillBatchWithPressure( + proc *process.Process, + bat *batch.Batch, + files []*os.File, + executors []colexec.ExpressionExecutor, + analyzer process.Analyzer, + sourceAlreadyCharged bool, +) error { + if ctr.spillBatchAllocation == nil || bat == nil || bat.RowCount() == 0 { + return ctr.spillBatchBounded( + proc, + bat, + files, + executors, + analyzer, + sourceAlreadyCharged, + ) + } + rows := bat.RowCount() + chunk := rows + minimumRetried := false + guard := NewPressureRetryGuard(PressureProgress{ + Used: ctr.hashmapBuilder.mapAllocationAccount.Snapshot().Used, + InputUnits: chunk, + OptionalDisabled: ctr.spillCoalesceDisabled, + }, 64) + for start := 0; start < rows; { + end := rows + if chunk < rows-start { + end = start + chunk + } + current := bat + if start != 0 || end != rows { + var err error + current, err = bat.Window(start, end) + if err != nil { + return err + } + } + err := ctr.spillBatchBounded( + proc, + current, + files, + executors, + analyzer, + sourceAlreadyCharged, + ) + if current != bat { + current.Clean(proc.Mp()) + } + if err == nil { + start = end + minimumRetried = false + nextUnits := chunk + if remaining := rows - start; remaining < nextUnits { + nextUnits = remaining } + guard = NewPressureRetryGuard(PressureProgress{ + Used: ctr.hashmapBuilder.mapAllocationAccount.Snapshot().Used, + InputUnits: nextUnits, + OptionalDisabled: ctr.spillCoalesceDisabled, + }, 64) + continue + } + if !IsRetryableMemoryCapacity(err) { + return err } - selected.CleanOnlyData() - if spillErr != nil { - return spillErr + if cancelErr := checkHashBuildCanceled(proc); cancelErr != nil { + return cancelErr } + ctr.releaseSpillComputeScratch() + attempted := end - start + if attempted <= 1 { + if !minimumRetried { + if reclaimErr := ctr.reclaimOptionalSpillBuffers( + proc, + files, + analyzer, + ); reclaimErr != nil { + return reclaimErr + } + next := PressureProgress{ + Used: ctr.hashmapBuilder.mapAllocationAccount.Snapshot().Used, + InputUnits: attempted, + OptionalDisabled: ctr.spillCoalesceDisabled, + } + if guard.Advance(next) != nil { + return NewMinimumAllocationPressureError( + "hashbuild", + "spill-hash-or-expression", + ctr.hashmapBuilder.mapAllocationAccount, + ) + } + minimumRetried = true + analyzer.GetOpStats().AddExtraStat( + "HashBuildSpillMinimumRetries", + 1, + ) + continue + } + return NewMinimumAllocationPressureError( + "hashbuild", + "spill-hash-or-expression", + ctr.hashmapBuilder.mapAllocationAccount, + ) + } + chunk = (attempted + 1) / 2 + if err := guard.Advance(PressureProgress{ + Used: ctr.hashmapBuilder.mapAllocationAccount.Snapshot().Used, + InputUnits: chunk, + }); err != nil { + return err + } + analyzer.GetOpStats().AddExtraStat("HashBuildSpillInputReductions", 1) } return nil } @@ -1127,6 +1724,9 @@ func (ctr *container) appendAccountedSpillRecord( return err } payload := ctr.spillAccountedWrite.Bytes() + if ctr.spillCoalesceDisabled { + return ctr.writeSpillPayload(proc, file, payload, cnt, analyzer) + } buffer := ctr.spillAccountedBuckets[bucket] if buffer != nil && buffer.Len() > 0 && buffer.Len()+len(payload) > spillWriteCoalesceSize { diff --git a/pkg/sql/colexec/hashbuild/spill_test.go b/pkg/sql/colexec/hashbuild/spill_test.go index 97837eab1c52f..1a7a68aacf41f 100644 --- a/pkg/sql/colexec/hashbuild/spill_test.go +++ b/pkg/sql/colexec/hashbuild/spill_test.go @@ -21,6 +21,7 @@ import ( "io" "math" "os" + "strings" "testing" "github.com/matrixorigin/matrixone/pkg/common/mpool" @@ -29,6 +30,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" + plan2 "github.com/matrixorigin/matrixone/pkg/sql/plan" "github.com/matrixorigin/matrixone/pkg/testutil" "github.com/matrixorigin/matrixone/pkg/vm/process" "github.com/stretchr/testify/require" @@ -1676,3 +1678,153 @@ func TestAccountedInitialSpillConvertsHeadroomToPhysicalOwnership(t *testing.T) _, _, err = registry.CompleteTerminal(account) require.NoError(t, err) } + +func TestAccountedMinimumSpillHeadroomIsOneUnitNotWholeBatch(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + makeBatch := func(rows int) *batch.Batch { + values := make([]string, rows) + for i := range values { + values[i] = strings.Repeat("x", 4<<10) + } + bat := batch.NewWithSize(1) + bat.Vecs[0] = testutil.MakeVarcharVector(values, nil, proc.Mp()) + bat.SetRowCount(rows) + return bat + } + one := makeBatch(1) + many := makeBatch(128) + defer one.Clean(proc.Mp()) + defer many.Clean(proc.Mp()) + exprs := []*plan.Expr{newExpr(0, types.T_varchar.ToType())} + oneUnit, err := spillMinimumUnitBudgetBytes(one, exprs) + require.NoError(t, err) + manyUnits, err := spillMinimumUnitBudgetBytes(many, exprs) + require.NoError(t, err) + require.Equal(t, oneUnit, manyUnits) + require.Positive(t, oneUnit) + legacyWholeBatch, err := spillBudgetBytes(many) + require.NoError(t, err) + require.Greater(t, legacyWholeBatch, manyUnits) +} + +func TestAccountedConcatSpillHeadroomUsesOnePhysicalRow(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + column := func(pos int32) *plan.Expr { + return &plan.Expr{ + Typ: plan.Type{Id: int32(types.T_varchar), Width: types.MaxVarcharLen}, + Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: pos}}, + } + } + expr, err := plan2.BindFuncExprImplByPlanExpr( + proc.Ctx, + "concat", + []*plan.Expr{column(0), column(1)}, + ) + require.NoError(t, err) + bat := batch.NewWithSize(2) + bat.Vecs[0] = testutil.MakeVarcharVector( + []string{strings.Repeat("a", 4<<10), "a"}, nil, proc.Mp()) + bat.Vecs[1] = testutil.MakeVarcharVector( + []string{"b", strings.Repeat("b", 4<<10)}, nil, proc.Mp()) + bat.SetRowCount(2) + defer bat.Clean(proc.Mp()) + + payload, err := spillExpressionPayloadBytes( + expr, + expr.GetF().GetArgs(), + bat, + ) + require.NoError(t, err) + require.Equal(t, uint64((4<<10)+1), payload) + headroom, err := spillMinimumUnitBudgetBytes(bat, []*plan.Expr{expr}) + require.NoError(t, err) + require.Positive(t, headroom) +} + +func TestAccountedInitialSpillReducesUnpublishedInputAndPreservesRows(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + const limit = uint64(80 << 10) + budget := process.MustNewHashBuildBudget(limit, limit) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 128) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + var op HashBuild + op.NeedHashMap = true + require.NoError(t, op.SetAllocationAccount(account)) + ctr := &op.ctr + ctr.hashmapBuilder.setBudget(generation) + ctr.spillUUID = "accounted-adaptive-spill" + exprs := []*plan.Expr{newExpr(0, types.T_int64.ToType())} + executors, err := ctr.initSpillExprExecs(proc, exprs) + require.NoError(t, err) + values := make([]int64, colexec.DefaultBatchSize) + for i := range values { + values[i] = int64(i) + } + input := batch.NewWithSize(1) + input.Vecs[0] = testutil.MakeInt64Vector(values, nil, proc.Mp()) + input.SetRowCount(len(values)) + defer input.Clean(proc.Mp()) + files := make([]*os.File, spillNumBuckets) + analyzer := process.NewAnalyzer(0, false, false, "test") + require.NoError(t, ctr.spillBatchWithPressure( + proc, + input, + files, + executors, + analyzer, + false, + )) + require.Positive(t, + analyzer.GetOpStats().ExtraStats["HashBuildSpillInputReductions"]) + require.NoError(t, ctr.flushSpillBuffers(proc, files, analyzer)) + + var totalRows int64 + for _, file := range files { + if file == nil { + continue + } + _, err = file.Seek(0, io.SeekStart) + require.NoError(t, err) + reader := bufio.NewReader(file) + for { + var header [16]byte + _, err = io.ReadFull(reader, header[:]) + if err == io.EOF { + break + } + require.NoError(t, err) + rows := types.DecodeInt64(header[:8]) + payload := types.DecodeInt64(header[8:]) + require.NoError(t, func() error { + _, copyErr := io.CopyN(io.Discard, reader, payload+8) + return copyErr + }()) + totalRows += rows + } + } + require.Equal(t, int64(len(values)), totalRows) + + ctr.dropSpillScratchBuffers() + ctr.freeSpillExprExecs() + for _, file := range files { + if file != nil { + _ = file.Close() + } + } + if ctr.spillBundle != nil { + ctr.spillBundle.release() + ctr.spillBundle = nil + } + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.NoError(t, op.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} diff --git a/pkg/sql/colexec/hashbuild/types.go b/pkg/sql/colexec/hashbuild/types.go index 9c0c2f66d7b95..d5e6ae43de003 100644 --- a/pkg/sql/colexec/hashbuild/types.go +++ b/pkg/sql/colexec/hashbuild/types.go @@ -81,6 +81,16 @@ const ( HashBuildAllocationSiteUniqueKeyNulls HashBuildAllocationSiteUniqueKeyGrouping HashBuildAllocationSiteRuntimeFilterPayload + HashBuildAllocationSiteDedupIgnoreBitmap + HashBuildAllocationSiteDedupDeleteBitmap + HashBuildAllocationSiteDedupLastRows + HashBuildAllocationSiteDedupSurvivorRows + HashBuildAllocationSiteDedupSurvivorOwnsKey + HashBuildAllocationSiteDedupDiscardedRows + HashBuildAllocationSiteDedupDeleteOnlyData + HashBuildAllocationSiteDedupDeleteOnlyArea + HashBuildAllocationSiteDedupDeleteOnlyNulls + HashBuildAllocationSiteDedupDeleteOnlyGrouping ) type container struct { @@ -124,6 +134,7 @@ type container struct { spillAllocationMP *mpool.MPool spillAccountedWrite *mpool.AccountedBuffer spillAccountedBuckets [spillNumBuckets]*mpool.AccountedBuffer + spillCoalesceDisabled bool // spillScratchReservation is a query/CN-charged emergency lease retained // while Shuffle build batches accumulate. It prevents retained copies from // consuming the scratch required to recover from hard-budget rejection. @@ -139,8 +150,9 @@ type container struct { spillScratchBase uint64 // cached expression executors for spill (reused across batches) - spillExprExecs []colexec.ExpressionExecutor - spillExprLease *ExpressionMemoryLease + spillExprExecs []colexec.ExpressionExecutor + spillExprLease *ExpressionMemoryLease + spillConditions []*plan.Expr // spillExprAccounted distinguishes an exact executor set from a legacy set // whose retained lease has not yet been installed. spillExprAccounted bool @@ -442,7 +454,9 @@ func (hb *HashmapBuilder) ClearAllocationAccount( } if builder.IntHashMap != nil || builder.StrHashMap != nil || len(builder.Batches.Buf) != 0 || builder.Sels.Size() != 0 || - len(builder.executors) != 0 { + len(builder.executors) != 0 || len(builder.curVecs) != 0 || + len(builder.UniqueJoinKeys) != 0 || + builder.IgnoreRows != nil || builder.DelRows != nil { return mpool.ErrAllocationAccountInvariant } builder.mapAllocationAccount = nil diff --git a/pkg/sql/colexec/hashjoin/expression_memory_test.go b/pkg/sql/colexec/hashjoin/expression_memory_test.go index 747d35fa2341a..664cd625c9f6a 100644 --- a/pkg/sql/colexec/hashjoin/expression_memory_test.go +++ b/pkg/sql/colexec/hashjoin/expression_memory_test.go @@ -62,3 +62,37 @@ func TestHashJoinResetReleasesProbeExpressionLease(t *testing.T) { require.Nil(t, arg.ctr.eqCondVecs) require.Nil(t, arg.ctr.probeExpressionLease) } + +func TestHashJoinResetReleasesAccountedProbeExpressions(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + const capBytes = uint64(1 << 20) + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.OpenWithController(capBytes, generation) + require.NoError(t, err) + expr := &plan.Expr{Typ: plan.Type{Id: int32(types.T_int32)}, Expr: &plan.Expr_Lit{Lit: &plan.Literal{Value: &plan.Literal_I32Val{I32Val: 1}}}} + executors, err := hashbuild.NewAllocationAccountedExpressionExecutorsForAccount( + proc, []*plan.Expr{expr}, account, hashbuild.HashBuildAllocationOwner) + require.NoError(t, err) + arg := &HashJoin{allocationAccount: account} + arg.ctr.eqCondExecs = executors + arg.ctr.eqCondVecs = make([]*vector.Vector, len(executors)) + arg.ctr.probeExpressionsAccounted = true + input := batch.NewWithSize(0) + input.SetRowCount(4) + require.NoError(t, arg.ctr.evalJoinConditionBudgeted(input, proc)) + require.Positive(t, account.Snapshot().Used) + + arg.Reset(proc, false, nil) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.False(t, arg.ctr.probeExpressionsAccounted) + require.Nil(t, arg.ctr.eqCondExecs) + terminal, _, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) +} diff --git a/pkg/sql/colexec/hashjoin/join.go b/pkg/sql/colexec/hashjoin/join.go index 1941151abcfd9..d9f117d1b313f 100644 --- a/pkg/sql/colexec/hashjoin/join.go +++ b/pkg/sql/colexec/hashjoin/join.go @@ -337,8 +337,29 @@ func (hashJoin *HashJoin) build(analyzer process.Analyzer, proc *process.Process if takeErr != nil { return takeErr } - probeExpressionLease, leaseErr := hashbuild.NewExpressionMemoryLease( - budget, hashJoin.EqConds[0], ctr.eqCondExecs, false) + var probeExpressionLease *hashbuild.ExpressionMemoryLease + var leaseErr error + if hashJoin.allocationAccount != nil && + hashbuild.AllocationAccountedExpressionSetSupported(hashJoin.EqConds[0]) { + ctr.cleanEqCondExecutors() + ctr.eqCondExecs, leaseErr = + hashbuild.NewAllocationAccountedExpressionExecutorsForAccount( + proc, + hashJoin.EqConds[0], + hashJoin.allocationAccount, + hashbuild.HashBuildAllocationOwner, + ) + if leaseErr == nil { + ctr.eqCondVecs = make( + []*vector.Vector, + len(hashJoin.EqConds[0]), + ) + ctr.probeExpressionsAccounted = true + } + } else { + probeExpressionLease, leaseErr = hashbuild.NewExpressionMemoryLease( + budget, hashJoin.EqConds[0], ctr.eqCondExecs, false) + } if leaseErr != nil { _ = payload.Close() ctr.mp.Free() @@ -380,7 +401,7 @@ func (hashJoin *HashJoin) build(analyzer process.Analyzer, proc *process.Process }, analyzer, func(bat *batch.Batch) ([]*vector.Vector, error) { - if err := ctr.evalJoinCondition(bat, proc); err != nil { + if err := ctr.evalJoinConditionBudgeted(bat, proc); err != nil { return nil, err } return ctr.eqCondVecs, nil diff --git a/pkg/sql/colexec/hashjoin/terminal_budget_test.go b/pkg/sql/colexec/hashjoin/terminal_budget_test.go index 3ee3c2b8441c4..8763dd9348ca5 100644 --- a/pkg/sql/colexec/hashjoin/terminal_budget_test.go +++ b/pkg/sql/colexec/hashjoin/terminal_budget_test.go @@ -42,7 +42,7 @@ func TestHashJoinCallConvertsTerminalBudgetAdmission(t *testing.T) { admission := &process.HashBuildBudgetError{ Kind: process.HashBuildBudgetErrorAdmission, - Resource: process.HashBuildBudgetResourceMemory, + Component: process.HashBuildBudgetComponentMemory, Requested: 2, Used: 1, Cap: 1, diff --git a/pkg/sql/colexec/hashjoin/types.go b/pkg/sql/colexec/hashjoin/types.go index c8b4a6ae71856..03e5697c0fed3 100644 --- a/pkg/sql/colexec/hashjoin/types.go +++ b/pkg/sql/colexec/hashjoin/types.go @@ -110,8 +110,9 @@ type container struct { // Non-nil only for spilled joins, where probe expressions are part of the // shared HashBuild/spill working set. Resident probe expressions remain // under normal process/mpool accounting; this is not a general query budget. - probeExpressionLease *hashbuild.ExpressionMemoryLease - probeBucketActive bool // true while reading probe batches from a bucket + probeExpressionLease *hashbuild.ExpressionMemoryLease + probeExpressionsAccounted bool + probeBucketActive bool // true while reading probe batches from a bucket } type HashJoin struct { @@ -171,7 +172,8 @@ func (hashJoin *HashJoin) ClearAllocationAccount( if hashJoin.allocationAccount != account { return mpool.ErrAllocationAccountMismatch } - if hashJoin.ctr.mp != nil || hashJoin.ctr.spillEngine != nil { + if hashJoin.ctr.mp != nil || hashJoin.ctr.spillEngine != nil || + hashJoin.ctr.probeExpressionsAccounted { return mpool.ErrAllocationAccountInvariant } hashJoin.allocationAccount = nil @@ -234,7 +236,7 @@ func (hashJoin *HashJoin) Reset(proc *process.Process, pipelineFailed bool, err // SpillEngine borrows the probe executor lease. End that borrow before the // join frees the executors and releases their reservation. ctr.cleanBucketBatches(proc) - if ctr.probeExpressionLease != nil { + if ctr.probeExpressionLease != nil || ctr.probeExpressionsAccounted { ctr.cleanEqCondExecutors() ctr.releaseProbeExpressionLease() } else { @@ -325,6 +327,7 @@ func (ctr *container) cleanEqCondExecutors() { } ctr.eqCondExecs = nil ctr.eqCondVecs = nil + ctr.probeExpressionsAccounted = false } func (ctr *container) resetEqCondExecutors() { diff --git a/pkg/sql/colexec/join_util.go b/pkg/sql/colexec/join_util.go index 73c973e5ba69f..7f951f7cb6b16 100644 --- a/pkg/sql/colexec/join_util.go +++ b/pkg/sql/colexec/join_util.go @@ -153,45 +153,63 @@ func (bs *Batches) Shrink(ignoreRow *bitmap.Bitmap, proc *process.Process) error if ignoreRow.Count() == 0 { return nil } - - ignoreRow.Negate() - count := int64(ignoreRow.Count()) - sels := make([]int32, 0, count) - itr := ignoreRow.Iterator() - for itr.HasNext() { - r := itr.Next() - sels = append(sels, int32(r)) + if len(bs.Buf) == 0 || bs.Buf[0] == nil { + return mpool.ErrAllocationAccountInvalid } - n := (len(sels)-1)/DefaultBatchSize + 1 + ignoreRow.Negate() + // Build the replacement privately and stream the active row IDs directly + // from the bitmap. The old implementation materialized one Go int32 per + // row and silently dropped the copied-batch allocation provenance. + count := ignoreRow.Count() + n := (count + DefaultBatchSize - 1) / DefaultBatchSize + if n == 0 { + n = 1 + } + selection := bs.Buf[0].AllocationAccountSelection() newBuf := make([]*batch.Batch, n) + cleanup := true + defer func() { + if cleanup { + for _, bat := range newBuf { + if bat != nil { + bat.Clean(proc.Mp()) + } + } + // Preserve the caller's ignore-row checkpoint on failure. + ignoreRow.Negate() + } + }() for i := range newBuf { - newBuf[i] = batch.NewOffHeapWithSize(len(bs.Buf[i].Vecs)) - for j, vec := range bs.Buf[0].Vecs { - newBuf[i].Vecs[j] = vector.NewOffHeapVecWithType(*vec.GetType()) + newBuf[i] = batch.NewOffHeapWithSize(len(bs.Buf[0].Vecs)) + if err := newBuf[i].SetAllocationAccount(selection); err != nil { + return err } - var newsels []int32 - if (i+1)*DefaultBatchSize <= len(sels) { - newsels = sels[i*DefaultBatchSize : (i+1)*DefaultBatchSize] - } else { - newsels = sels[i*DefaultBatchSize:] + for j, vec := range bs.Buf[0].Vecs { + newBuf[i].SetVector(int32(j), vector.NewOffHeapVecWithType(*vec.GetType())) } - for _, sel := range newsels { - idx1, idx2 := sel/DefaultBatchSize, sel%DefaultBatchSize - for j, vec := range bs.Buf[idx1].Vecs { - if err := newBuf[i].Vecs[j].UnionOne(vec, int64(idx2), proc.Mp()); err != nil { - for k := 0; k <= i; k++ { - newBuf[k].Clean(proc.Mp()) - } - return err - } + } + itr := ignoreRow.Iterator() + outRow := 0 + for itr.HasNext() { + sel := int(itr.Next()) + srcBatch, srcRow := sel/DefaultBatchSize, sel%DefaultBatchSize + dstBatch := outRow / DefaultBatchSize + for j, vec := range bs.Buf[srcBatch].Vecs { + if err := newBuf[dstBatch].Vecs[j].UnionOne(vec, int64(srcRow), proc.Mp()); err != nil { + return err } } - newBuf[i].SetRowCount(len(newsels)) + newBuf[dstBatch].AddRowCount(1) + outRow++ } bs.Clean(proc.Mp()) bs.Buf = newBuf + for _, bat := range newBuf { + bs.MemSize += int64(bat.Size()) + } + cleanup = false return nil } diff --git a/pkg/sql/colexec/join_util_test.go b/pkg/sql/colexec/join_util_test.go index 8d4dadecd30af..a2b96845b29ee 100644 --- a/pkg/sql/colexec/join_util_test.go +++ b/pkg/sql/colexec/join_util_test.go @@ -22,6 +22,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/testutil" ) @@ -76,3 +77,59 @@ func TestBatches(t *testing.T) { batches.Clean(proc.Mp()) require.Equal(t, int64(0), proc.Mp().CurrNB()) } + +func TestBatchesShrinkPreservesAllocationAndRollback(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + input := testutil.NewBatch( + []types.Type{types.T_int32.ToType()}, + true, + DefaultBatchSize, + proc.Mp(), + ) + defer input.Clean(proc.Mp()) + + measure := func(limit uint64, shrink bool) (uint64, error) { + registry, err := mpool.NewAllocationAccountRegistry(1, 32) + require.NoError(t, err) + account, err := registry.Open(limit) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelectionWithBitmaps( + account, 1, 1, 2, 3, 4) + require.NoError(t, err) + var batches Batches + require.NoError(t, batches.CopyIntoBatchesWithAllocation(input, proc, selection)) + before := account.Snapshot().Used + require.Positive(t, before) + var shrinkErr error + if shrink { + ignore := &bitmap.Bitmap{} + ignore.InitWithSize(DefaultBatchSize) + ignore.Add(0) + shrinkErr = batches.Shrink(ignore, proc) + if shrinkErr == nil { + require.Equal(t, DefaultBatchSize-1, batches.RowCount()) + for _, bat := range batches.Buf { + require.Same(t, selection, bat.AllocationAccountSelection()) + } + } else { + require.Equal(t, 1, ignore.Count(), "failed shrink restores ignore-row checkpoint") + require.Equal(t, DefaultBatchSize, batches.RowCount()) + require.Equal(t, before, account.Snapshot().Used) + } + } + batches.Clean(proc.Mp()) + require.Zero(t, account.Snapshot().Used) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + return before, shrinkErr + } + + used, err := measure(1<<20, false) + require.NoError(t, err) + _, err = measure(used, true) + require.Error(t, err) + require.True(t, mpool.IsRetryableAllocationCapacity(err)) + _, err = measure(1<<20, true) + require.NoError(t, err) +} diff --git a/pkg/sql/colexec/rightdedupjoin/expression_memory_test.go b/pkg/sql/colexec/rightdedupjoin/expression_memory_test.go index 971f251b0991f..3b30a749fcea2 100644 --- a/pkg/sql/colexec/rightdedupjoin/expression_memory_test.go +++ b/pkg/sql/colexec/rightdedupjoin/expression_memory_test.go @@ -62,3 +62,37 @@ func TestRightDedupJoinResetReleasesProbeExpressionLease(t *testing.T) { require.Nil(t, arg.ctr.vecs) require.Nil(t, arg.ctr.probeExpressionLease) } + +func TestRightDedupJoinResetReleasesAccountedProbeExpressions(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + const capBytes = uint64(1 << 20) + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.OpenWithController(capBytes, generation) + require.NoError(t, err) + expr := &plan.Expr{Typ: plan.Type{Id: int32(types.T_int32)}, Expr: &plan.Expr_Lit{Lit: &plan.Literal{Value: &plan.Literal_I32Val{I32Val: 1}}}} + executors, err := hashbuild.NewAllocationAccountedExpressionExecutorsForAccount( + proc, []*plan.Expr{expr}, account, hashbuild.HashBuildAllocationOwner) + require.NoError(t, err) + arg := &RightDedupJoin{allocationAccount: account} + arg.ctr.evecs = []evalVector{{executor: executors[0]}} + arg.ctr.vecs = make([]*vector.Vector, len(executors)) + arg.ctr.probeExpressionsAccounted = true + input := batch.NewWithSize(0) + input.SetRowCount(4) + require.NoError(t, arg.ctr.evalJoinConditionBudgeted(input, proc)) + require.Positive(t, account.Snapshot().Used) + + arg.Reset(proc, false, nil) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.False(t, arg.ctr.probeExpressionsAccounted) + require.Nil(t, arg.ctr.evecs) + terminal, _, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) +} diff --git a/pkg/sql/colexec/rightdedupjoin/join.go b/pkg/sql/colexec/rightdedupjoin/join.go index 19186ff3de60d..0521a15ec8f54 100644 --- a/pkg/sql/colexec/rightdedupjoin/join.go +++ b/pkg/sql/colexec/rightdedupjoin/join.go @@ -212,12 +212,35 @@ func (rightDedupJoin *RightDedupJoin) build(analyzer process.Analyzer, proc *pro if takeErr != nil { return takeErr } - probeExecutors := make([]colexec.ExpressionExecutor, len(ctr.evecs)) - for i := range ctr.evecs { - probeExecutors[i] = ctr.evecs[i].executor + var probeExpressionLease *hashbuild.ExpressionMemoryLease + var leaseErr error + if rightDedupJoin.allocationAccount != nil && + hashbuild.AllocationAccountedExpressionSetSupported(rightDedupJoin.Conditions[0]) { + ctr.cleanEvalVectors() + var probeExecutors []colexec.ExpressionExecutor + probeExecutors, leaseErr = + hashbuild.NewAllocationAccountedExpressionExecutorsForAccount( + proc, + rightDedupJoin.Conditions[0], + rightDedupJoin.allocationAccount, + hashbuild.HashBuildAllocationOwner, + ) + if leaseErr == nil { + ctr.evecs = make([]evalVector, len(probeExecutors)) + ctr.vecs = make([]*vector.Vector, len(probeExecutors)) + for i := range probeExecutors { + ctr.evecs[i].executor = probeExecutors[i] + } + ctr.probeExpressionsAccounted = true + } + } else { + probeExecutors := make([]colexec.ExpressionExecutor, len(ctr.evecs)) + for i := range ctr.evecs { + probeExecutors[i] = ctr.evecs[i].executor + } + probeExpressionLease, leaseErr = hashbuild.NewExpressionMemoryLease( + budget, rightDedupJoin.Conditions[0], probeExecutors, false) } - probeExpressionLease, leaseErr := hashbuild.NewExpressionMemoryLease( - budget, rightDedupJoin.Conditions[0], probeExecutors, false) if leaseErr != nil { _ = payload.Close() ctr.mp.Free() @@ -256,7 +279,7 @@ func (rightDedupJoin *RightDedupJoin) build(analyzer process.Analyzer, proc *pro }, analyzer, func(bat *batch.Batch) ([]*vector.Vector, error) { - if err := ctr.evalJoinCondition(bat, proc); err != nil { + if err := ctr.evalJoinConditionBudgeted(bat, proc); err != nil { return nil, err } return ctr.vecs, nil diff --git a/pkg/sql/colexec/rightdedupjoin/types.go b/pkg/sql/colexec/rightdedupjoin/types.go index aef8d1b2e931b..490208b5a6356 100644 --- a/pkg/sql/colexec/rightdedupjoin/types.go +++ b/pkg/sql/colexec/rightdedupjoin/types.go @@ -68,8 +68,9 @@ type container struct { // Non-nil only for spilled joins, where probe expressions are part of the // shared HashBuild/spill working set. Resident probe expressions remain // under normal process/mpool accounting; this is not a general query budget. - probeExpressionLease *hashbuild.ExpressionMemoryLease - resultBatch *batch.Batch + probeExpressionLease *hashbuild.ExpressionMemoryLease + probeExpressionsAccounted bool + resultBatch *batch.Batch } type RightDedupJoin struct { @@ -124,7 +125,8 @@ func (rightDedupJoin *RightDedupJoin) ClearAllocationAccount( return mpool.ErrAllocationAccountMismatch } if rightDedupJoin.ctr.mp != nil || - rightDedupJoin.ctr.spillEngine != nil { + rightDedupJoin.ctr.spillEngine != nil || + rightDedupJoin.ctr.probeExpressionsAccounted { return mpool.ErrAllocationAccountInvariant } rightDedupJoin.allocationAccount = nil @@ -180,7 +182,7 @@ func (rightDedupJoin *RightDedupJoin) Reset(proc *process.Process, pipelineFaile ctr.spillEngine.Cleanup(proc) ctr.spillEngine = nil } - if ctr.probeExpressionLease != nil { + if ctr.probeExpressionLease != nil || ctr.probeExpressionsAccounted { ctr.cleanEvalVectors() ctr.releaseProbeExpressionLease() } else { @@ -261,6 +263,7 @@ func (ctr *container) cleanEvalVectors() { } ctr.evecs = nil ctr.vecs = nil + ctr.probeExpressionsAccounted = false } func (ctr *container) resetEvalVectors() { diff --git a/pkg/sql/colexec/spillutil/allocation_account_test.go b/pkg/sql/colexec/spillutil/allocation_account_test.go index 7b5e9d22bd2de..7a37bf1a9285d 100644 --- a/pkg/sql/colexec/spillutil/allocation_account_test.go +++ b/pkg/sql/colexec/spillutil/allocation_account_test.go @@ -324,7 +324,7 @@ func TestSpillAllocationAccountScatterScratchLifecycle(t *testing.T) { finalizeTestSpillAllocationAccount(t, state) } -func TestSpillAllocationAccountScatterChargesOnlyExternalSource(t *testing.T) { +func TestSpillAllocationAccountScatterDoesNotReadmitBorrowedSource(t *testing.T) { proc := testutil.NewProcessWithMPool( t, "", @@ -374,9 +374,7 @@ func TestSpillAllocationAccountScatterChargesOnlyExternalSource(t *testing.T) { snapshot := generation.Snapshot() require.Nil(t, engine.scatterScratchReservation) require.Equal(t, snapshot.AllocationUsed, snapshot.Used, - "the upstream source token is transient and private spill bytes are exact") - require.Positive(t, snapshot.ReserveCount, - "the unaccounted upstream source remains part of the peak") + "borrowed input is already live; only new private spill bytes are admitted") require.Greater(t, snapshot.PeakUsed, snapshot.Used) engine.releaseScatterScratch() @@ -518,7 +516,11 @@ func TestSpillAllocationAccountScatterFailureCleanup(t *testing.T) { false, process.NewAnalyzer(0, false, false, "test"), ) - require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Equal( + t, + hashbuild.MemoryPressureMinimumUnit, + hashbuild.MemoryPressureReasonOf(err), + ) require.Equal(t, uint64(rows*(8+4)), state.account.Snapshot().Used) engine.releaseScatterScratch() @@ -527,6 +529,116 @@ func TestSpillAllocationAccountScatterFailureCleanup(t *testing.T) { finalizeTestSpillAllocationAccount(t, state) } +func TestSpillAllocationAccountScatterReducesUnpublishedInput(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-scatter-reduce"), + ) + defer proc.Free() + state := newTestSpillAllocationAccount(t, 80<<10, 128) + engine, err := NewSpillEngineWithAllocation( + SpillEngineConfig{}, + state.allocation, + ) + require.NoError(t, err) + values := make([]int64, 8_192) + for i := range values { + values[i] = int64(i) + } + source := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.MakeInt64Vector(values, nil, proc.Mp()), + }, nil) + defer source.Clean(proc.Mp()) + writers := MakeBucketWriters("spill_allocation_scatter_reduce") + defer func() { + for i := range writers { + writers[i].Close() + } + }() + analyzer := process.NewAnalyzer(0, false, false, "test") + require.NoError(t, engine.scatterBatchWithPressure( + proc, + source, + source.Vecs, + writers, + 0, + false, + analyzer, + )) + require.Positive(t, + analyzer.GetOpStats().ExtraStats["JoinSpillInputReductions"]) + require.NoError(t, engine.flushScatterBuffers(proc, writers, analyzer)) + var rows int64 + for i := range writers { + rows += writers[i].Rows + } + require.Equal(t, int64(len(values)), rows) + + engine.releaseScatterScratch() + engine.Cleanup(proc) + require.Zero(t, state.account.Snapshot().Used) + finalizeTestSpillAllocationAccount(t, state) +} + +func TestSpillAllocationAccountExpressionPressureReducesBeforePublication(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-expression-reduce"), + ) + defer proc.Free() + state := newTestSpillAllocationAccount(t, 2<<20, 4_096) + engine, err := NewSpillEngineWithAllocation( + SpillEngineConfig{}, + state.allocation, + ) + require.NoError(t, err) + values := make([]int64, 257) + for i := range values { + values[i] = int64(i) + } + source := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.MakeInt64Vector(values, nil, proc.Mp()), + }, nil) + defer source.Clean(proc.Mp()) + writers := MakeBucketWriters("spill_allocation_expression_reduce") + defer func() { + for i := range writers { + writers[i].Close() + } + }() + analyzer := process.NewAnalyzer(0, false, false, "test") + require.NoError(t, engine.scatterEvaluatedBatchWithPressure( + proc, + source, + writers, + 0, + false, + analyzer, + func(current *batch.Batch) ([]*vector.Vector, error) { + if current.RowCount() > 32 { + return nil, mpool.ErrAllocationAccountCapacity + } + return current.Vecs, nil + }, + )) + require.Positive(t, + analyzer.GetOpStats().ExtraStats["JoinSpillExpressionInputReductions"]) + require.NoError(t, engine.flushScatterBuffers(proc, writers, analyzer)) + var rows int64 + for i := range writers { + rows += writers[i].Rows + } + require.Equal(t, int64(len(values)), rows, + "evaluation retries must not duplicate or omit published rows") + + engine.releaseScatterScratch() + engine.Cleanup(proc) + require.Zero(t, state.account.Snapshot().Used) + finalizeTestSpillAllocationAccount(t, state) +} + func TestSpillAllocationAccountRebuildAndRecursiveSpillLifecycle(t *testing.T) { proc := testutil.NewProcessWithMPool( t, diff --git a/pkg/sql/colexec/spillutil/join_spill.go b/pkg/sql/colexec/spillutil/join_spill.go index 330aa0f7aed50..8d1cf60bbf2b1 100644 --- a/pkg/sql/colexec/spillutil/join_spill.go +++ b/pkg/sql/colexec/spillutil/join_spill.go @@ -591,7 +591,7 @@ func (r *BucketReader) readBatchRecord( growErr = token.Grow(peak - charge) } if growErr != nil && - !errors.Is(growErr, process.ErrHashBuildBudgetAdmission) { + !hashbuild.IsRetryableMemoryCapacity(growErr) { return nil, token, charge, growErr } if !retainedOK || !peakOK || growErr != nil { @@ -1234,20 +1234,6 @@ func scatterImpl( return nil } -// scatterBatch scatters bat using the engine's reusable hash/row-id buffers. -func (e *SpillEngine) scatterBatch( - proc *process.Process, - bat *batch.Batch, - keyVecs []*vector.Vector, - writers []BucketWriter, - buffers []*batch.Batch, - partitionLevel uint64, - sourceAlreadyCharged bool, - analyzer process.Analyzer, -) error { - return e.scatterBatchBounded(proc, bat, keyVecs, writers, partitionLevel, sourceAlreadyCharged, analyzer) -} - func scatterTransientBudgetBytes(bat *batch.Batch, sourceAlreadyCharged bool) (uint64, error) { if bat == nil || bat.RowCount() < 0 { return 0, process.ErrHashBuildBudgetInvalid @@ -1502,7 +1488,6 @@ func (e *SpillEngine) scatterBatchBounded( } rows := bat.RowCount() var selected *batch.Batch - var externalSource *process.HashBuildReservation defer func() { if selected != nil { selected.Clean(proc.Mp()) @@ -1512,12 +1497,10 @@ func (e *SpillEngine) scatterBatchBounded( if reconcileErr != nil && retErr == nil { retErr = reconcileErr } - if retErr != nil { + if retErr != nil && + !(e.allocation != nil && hashbuild.IsRetryableMemoryCapacity(retErr)) { e.discardScatterBuffers() } - if externalSource != nil { - externalSource.Release() - } }() if e.cfg.Budget != nil && e.allocation == nil { // Start with retained capacities already owned by this token, add only @@ -1552,20 +1535,18 @@ func (e *SpillEngine) scatterBatchBounded( if err != nil { return err } - } else if e.cfg.Budget != nil { - externalBytes, err := externalScatterSourceBytes( - bat, - sourceAlreadyCharged, - ) - if err != nil { - return err - } - if externalBytes > 0 { - externalSource, err = e.cfg.Budget.Reserve(externalBytes) - if err != nil { - return err - } + } else if e.cfg.Budget != nil && !sourceAlreadyCharged { + // The child batch is borrowed and already physically live. Rejecting a + // new logical token cannot reclaim it, so observe it while exact-account + // admission governs every new scatter allocation. + externalBytes := bat.Allocated() + if size := bat.Size(); size > externalBytes { + externalBytes = size } + analyzer.GetOpStats().SetMaxExtraStat( + "JoinSpillBorrowedSourceBytes", + int64(externalBytes), + ) } if e.allocation != nil { @@ -1622,7 +1603,6 @@ func (e *SpillEngine) scatterBatchBounded( if start == end || writers[bucketID].Name == "" { continue } - sels := e.scatterBucketRowIds[start:end] if selected == nil { var selection *vector.AllocationAccountSelection if e.allocation != nil { @@ -1642,18 +1622,419 @@ func (e *SpillEngine) scatterBatchBounded( } } } - selected.CleanOnlyData() - for j, vec := range bat.Vecs { - if err := selected.Vecs[j].UnionInt32(vec, sels, proc.Mp()); err != nil { + cursor := start + for cursor < end { + attemptEnd := end + reclaimedMinimum := false + for { + selected.CleanOnlyData() + sels := e.scatterBucketRowIds[cursor:attemptEnd] + var scatterErr error + for j, vec := range bat.Vecs { + if scatterErr = selected.Vecs[j].UnionInt32( + vec, + sels, + proc.Mp(), + ); scatterErr != nil { + break + } + } + if scatterErr == nil { + selected.SetRowCount(len(sels)) + scatterErr = e.appendScatterRecord( + proc, + selected, + &writers[bucketID], + bucketID, + analyzer, + ) + } selected.CleanOnlyData() + if scatterErr == nil { + cursor = attemptEnd + break + } + if e.allocation == nil || + !hashbuild.IsRetryableMemoryCapacity(scatterErr) { + return scatterErr + } + if err := checkSpillCanceled(proc); err != nil { + return err + } + n := int(attemptEnd - cursor) + if n > 1 { + attemptEnd = cursor + int32((n+1)/2) + analyzer.GetOpStats().AddExtraStat( + "JoinSpillBatchReductions", + 1, + ) + continue + } + if !reclaimedMinimum { + before := e.allocation.account.Snapshot().Used + if err := e.reclaimOptionalScatterBuffers( + proc, + writers, + analyzer, + ); err != nil { + return err + } + reclaimedMinimum = true + after := e.allocation.account.Snapshot().Used + if after >= before { + return hashbuild.NewMinimumAllocationPressureError( + "join-spill", + "scatter-selected-or-codec", + e.allocation.account, + ) + } + analyzer.GetOpStats().AddExtraStat( + "JoinSpillOptionalReclaims", + 1, + ) + continue + } + return hashbuild.NewMinimumAllocationPressureError( + "join-spill", + "scatter-selected-or-codec", + e.allocation.account, + ) + } + } + } + return nil +} + +func (e *SpillEngine) reclaimOptionalScatterBuffers( + proc *process.Process, + writers []BucketWriter, + analyzer process.Analyzer, +) error { + for bucket, buffer := range e.scatterAccountedWriteBuffers { + if buffer == nil { + continue + } + if buffer.Len() > 0 { + if bucket >= len(writers) { + return process.ErrHashBuildBudgetInvalid + } + if err := e.flushPendingScatterBucket( + proc, + &writers[bucket], + bucket, + analyzer, + ); err != nil { + return err + } + } + buffer.Free() + e.scatterAccountedWriteBuffers[bucket] = nil + } + e.scatterCoalesceDisabled = true + if e.scatterAccountedWriteBuf != nil { + e.scatterAccountedWriteBuf.Free() + e.scatterAccountedWriteBuf = nil + } + return nil +} + +func (e *SpillEngine) releaseScatterComputeScratch() { + if e.allocation == nil || e.allocationMP == nil { + return + } + freeSpillSlice(e.scatterHashValues, e.allocationMP, e.allocation) + freeSpillSlice(e.scatterBucketRowIds, e.allocationMP, e.allocation) + e.scatterHashValues = nil + e.scatterBucketRowIds = nil +} + +func (e *SpillEngine) scatterBatchWithPressure( + proc *process.Process, + bat *batch.Batch, + keyVecs []*vector.Vector, + writers []BucketWriter, + partitionLevel uint64, + sourceAlreadyCharged bool, + analyzer process.Analyzer, +) error { + if e.allocation == nil || bat == nil || bat.RowCount() == 0 { + return e.scatterBatchBounded( + proc, + bat, + keyVecs, + writers, + partitionLevel, + sourceAlreadyCharged, + analyzer, + ) + } + rows := bat.RowCount() + chunk := rows + minimumRetried := false + guard := hashbuild.NewPressureRetryGuard(hashbuild.PressureProgress{ + Used: e.allocation.account.Snapshot().Used, + InputUnits: rows, + OptionalDisabled: e.scatterCoalesceDisabled, + }, 64) + for start := 0; start < rows; { + end := rows + if chunk < rows-start { + end = start + chunk + } + current := bat + currentKeys := keyVecs + if start != 0 || end != rows { + var err error + current, err = bat.Window(start, end) + if err != nil { return err } + currentKeys = make([]*vector.Vector, len(keyVecs)) + for i, key := range keyVecs { + currentKeys[i], err = key.Window(start, end) + if err != nil { + for j := 0; j < i; j++ { + currentKeys[j].Free(proc.Mp()) + } + current.Clean(proc.Mp()) + return err + } + } + } + err := e.scatterBatchBounded( + proc, + current, + currentKeys, + writers, + partitionLevel, + sourceAlreadyCharged, + analyzer, + ) + if current != bat { + for _, key := range currentKeys { + key.Free(proc.Mp()) + } + current.Clean(proc.Mp()) + } + if err == nil { + start = end + minimumRetried = false + nextUnits := chunk + if remaining := rows - start; remaining < nextUnits { + nextUnits = remaining + } + guard = hashbuild.NewPressureRetryGuard(hashbuild.PressureProgress{ + Used: e.allocation.account.Snapshot().Used, + InputUnits: nextUnits, + OptionalDisabled: e.scatterCoalesceDisabled, + }, 64) + continue + } + if !hashbuild.IsRetryableMemoryCapacity(err) { + return err + } + if cancelErr := checkSpillCanceled(proc); cancelErr != nil { + return cancelErr + } + e.releaseScatterComputeScratch() + attempted := end - start + if attempted <= 1 { + if !minimumRetried { + if reclaimErr := e.reclaimOptionalScatterBuffers( + proc, + writers, + analyzer, + ); reclaimErr != nil { + return reclaimErr + } + next := hashbuild.PressureProgress{ + Used: e.allocation.account.Snapshot().Used, + InputUnits: attempted, + OptionalDisabled: e.scatterCoalesceDisabled, + } + if guard.Advance(next) != nil { + return hashbuild.NewMinimumAllocationPressureError( + "join-spill", + "scatter-hash", + e.allocation.account, + ) + } + minimumRetried = true + analyzer.GetOpStats().AddExtraStat( + "JoinSpillMinimumRetries", + 1, + ) + continue + } + return hashbuild.NewMinimumAllocationPressureError( + "join-spill", + "scatter-hash", + e.allocation.account, + ) } - selected.SetRowCount(len(sels)) - if err := e.appendScatterRecord(proc, selected, &writers[bucketID], bucketID, analyzer); err != nil { - selected.CleanOnlyData() + chunk = (attempted + 1) / 2 + if err := guard.Advance(hashbuild.PressureProgress{ + Used: e.allocation.account.Snapshot().Used, + InputUnits: chunk, + }); err != nil { return err } + analyzer.GetOpStats().AddExtraStat("JoinSpillInputReductions", 1) + } + return nil +} + +// scatterEvaluatedBatchWithPressure extends the same unpublished-input +// checkpoint across key evaluation and scatter. Exact expression executors may +// retain successfully admitted capacities after a later child/result growth +// fails; evaluating a smaller immutable window can then reuse those capacities +// without replaying any bucket record. scatterBatchWithPressure owns the +// transactional boundary after evaluation, so a capacity error returned here +// has not published the current window. +func (e *SpillEngine) scatterEvaluatedBatchWithPressure( + proc *process.Process, + bat *batch.Batch, + writers []BucketWriter, + partitionLevel uint64, + sourceAlreadyCharged bool, + analyzer process.Analyzer, + eval func(*batch.Batch) ([]*vector.Vector, error), +) error { + if bat == nil || bat.RowCount() == 0 { + return nil + } + if eval == nil { + return process.ErrHashBuildBudgetInvalid + } + if e.allocation == nil { + keyVecs, err := eval(bat) + if err != nil { + return err + } + return e.scatterBatchWithPressure( + proc, + bat, + keyVecs, + writers, + partitionLevel, + sourceAlreadyCharged, + analyzer, + ) + } + + rows := bat.RowCount() + chunk := rows + minimumRetried := false + guard := hashbuild.NewPressureRetryGuard(hashbuild.PressureProgress{ + Used: e.allocation.account.Snapshot().Used, + InputUnits: chunk, + OptionalDisabled: e.scatterCoalesceDisabled, + }, 64) + for start := 0; start < rows; { + end := rows + if chunk < rows-start { + end = start + chunk + } + current := bat + if start != 0 || end != rows { + var err error + current, err = bat.Window(start, end) + if err != nil { + return err + } + } + + var err error + if err = checkSpillCanceled(proc); err == nil { + var keyVecs []*vector.Vector + keyVecs, err = eval(current) + if err == nil { + err = checkSpillCanceled(proc) + } + if err == nil { + err = e.scatterBatchWithPressure( + proc, + current, + keyVecs, + writers, + partitionLevel, + sourceAlreadyCharged, + analyzer, + ) + } + } + if current != bat { + current.Clean(proc.Mp()) + } + if err == nil { + start = end + minimumRetried = false + nextUnits := chunk + if remaining := rows - start; remaining < nextUnits { + nextUnits = remaining + } + guard = hashbuild.NewPressureRetryGuard(hashbuild.PressureProgress{ + Used: e.allocation.account.Snapshot().Used, + InputUnits: nextUnits, + OptionalDisabled: e.scatterCoalesceDisabled, + }, 64) + continue + } + if !hashbuild.IsRetryableMemoryCapacity(err) { + return err + } + if cancelErr := checkSpillCanceled(proc); cancelErr != nil { + return cancelErr + } + + e.releaseScatterComputeScratch() + attempted := end - start + if attempted > 1 { + chunk = (attempted + 1) / 2 + if err := guard.Advance(hashbuild.PressureProgress{ + Used: e.allocation.account.Snapshot().Used, + InputUnits: chunk, + OptionalDisabled: e.scatterCoalesceDisabled, + }); err != nil { + return err + } + analyzer.GetOpStats().AddExtraStat( + "JoinSpillExpressionInputReductions", + 1, + ) + continue + } + + if minimumRetried { + return hashbuild.NewMinimumAllocationPressureError( + "join-spill", + "scatter-expression", + e.allocation.account, + ) + } + if reclaimErr := e.reclaimOptionalScatterBuffers( + proc, + writers, + analyzer, + ); reclaimErr != nil { + return reclaimErr + } + if err := guard.Advance(hashbuild.PressureProgress{ + Used: e.allocation.account.Snapshot().Used, + InputUnits: attempted, + OptionalDisabled: e.scatterCoalesceDisabled, + }); err != nil { + return hashbuild.NewMinimumAllocationPressureError( + "join-spill", + "scatter-expression", + e.allocation.account, + ) + } + minimumRetried = true + analyzer.GetOpStats().AddExtraStat( + "JoinSpillExpressionMinimumRetries", + 1, + ) } return nil } @@ -1731,6 +2112,9 @@ func (e *SpillEngine) appendAccountedScatterRecord( return err } payload := e.scatterAccountedWriteBuf.Bytes() + if e.scatterCoalesceDisabled { + return writeBucketPayload(proc, payload, rows, writer, analyzer) + } buf := e.scatterAccountedWriteBuffers[bucket] if buf != nil && buf.Len() > 0 && buf.Len()+len(payload) > spillWriteCoalesceSize { @@ -1899,6 +2283,7 @@ func (e *SpillEngine) releaseScatterScratch() { e.scatterWriteRows[i] = 0 } e.allocationMP = nil + e.scatterCoalesceDisabled = false if e.scatterScratchReservation != nil { e.scatterScratchReservation.Release() e.scatterScratchReservation = nil @@ -2041,6 +2426,7 @@ type SpillEngine struct { scatterWriteBuffers [SpillNumBuckets]bytes.Buffer scatterAccountedWriteBuf *mpool.AccountedBuffer scatterAccountedWriteBuffers [SpillNumBuckets]*mpool.AccountedBuffer + scatterCoalesceDisabled bool scatterWriteRows [SpillNumBuckets]int64 // The lease follows the reusable scratch capacities for the engine // lifetime. It is released only by Cleanup, after all backing arrays have @@ -2244,14 +2630,17 @@ func (e *SpillEngine) ScatterProbeTable( if bat.IsEmpty() { continue } - keyVecs, err := e.evalProbeKeys(proc, bat, evalKeysFn) - if err != nil { - return err - } - if err := checkSpillCanceled(proc); err != nil { - return err - } - if err := e.scatterBatch(proc, bat, keyVecs, writers, nil, 0, false, analyzer); err != nil { + if err := e.scatterEvaluatedBatchWithPressure( + proc, + bat, + writers, + 0, + false, + analyzer, + func(current *batch.Batch) ([]*vector.Vector, error) { + return e.evalProbeKeys(proc, current, evalKeysFn) + }, + ); err != nil { return err } } @@ -2516,7 +2905,7 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana builder.FreeHashMapAndBatches(proc) builder.Free(proc) if isBudgetAdmission(err) { - return nil, BucketSkip, noProgressError(bucket.Depth, err) + return nil, BucketSkip, noProgressError(proc, bucket.Depth) } return nil, BucketSkip, err } @@ -2568,7 +2957,7 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana builder.FreeHashMapAndBatches(proc) builder.Free(proc) if isBudgetAdmission(err) { - return nil, BucketSkip, noProgressError(bucket.Depth, err) + return nil, BucketSkip, noProgressError(proc, bucket.Depth) } return nil, BucketSkip, err } @@ -2718,13 +3107,9 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal evalAndScatter := func( bat *batch.Batch, writers []BucketWriter, - buffers []*batch.Batch, execs []colexec.ExpressionExecutor, sourceAlreadyCharged bool, ) error { - if err := checkSpillCanceled(proc); err != nil { - return err - } if cap(e.keyVecs) < len(execs) { e.keyVecs = make([]*vector.Vector, len(execs)) } @@ -2734,24 +3119,39 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal keyVecs[i] = nil } }() - err := e.buildExprLease.Run(proc, bat.RowCount(), func(i int) error { - vec, evalErr := execs[i].Eval(proc, []*batch.Batch{bat}, nil) - if evalErr != nil { - return evalErr - } - keyVecs[i] = vec - return nil - }) - if err != nil { - // Eval may leave newly allocated child/result vectors cached. - // Destroy the owned executor tree before releasing its lease. - e.freeKeyExecs() - return err - } - if err := checkSpillCanceled(proc); err != nil { - return err - } - return e.scatterBatch(proc, bat, keyVecs, writers, nil, partitionLevel, sourceAlreadyCharged, analyzer) + return e.scatterEvaluatedBatchWithPressure( + proc, + bat, + writers, + partitionLevel, + sourceAlreadyCharged, + analyzer, + func(current *batch.Batch) ([]*vector.Vector, error) { + for i := range keyVecs { + keyVecs[i] = nil + } + err := e.buildExprLease.Run(proc, current.RowCount(), func(i int) error { + vec, evalErr := execs[i].Eval(proc, []*batch.Batch{current}, nil) + if evalErr != nil { + return evalErr + } + keyVecs[i] = vec + return nil + }) + if err != nil { + // Exact capacity pressure keeps the executor tree as the + // rollback checkpoint: admitted child/result capacities may + // make a smaller immutable window fit. Every other failure is + // terminal and can destroy the private tree immediately. + if e.allocation == nil || + !hashbuild.IsRetryableMemoryCapacity(err) { + e.freeKeyExecs() + } + return nil, err + } + return keyVecs, nil + }, + ) } var buildRows int64 @@ -2759,7 +3159,7 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal if b != nil { buildRows += int64(b.RowCount()) } - if err := evalAndScatter(b, buildWriters, nil, e.keyExecs, true); err != nil { + if err := evalAndScatter(b, buildWriters, e.keyExecs, true); err != nil { return err } return nil @@ -2769,7 +3169,7 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal if pending != nil && pending.RowCount() > 0 { // pending is the current BucketReader batch whose copy admission failed; // the reader keeps its batch token live until the next ReadBatch. - if err := evalAndScatter(pending, buildWriters, nil, e.keyExecs, true); err != nil { + if err := evalAndScatter(pending, buildWriters, e.keyExecs, true); err != nil { return nil, err } } @@ -2789,7 +3189,7 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal return nil, err } buildRows += int64(bat.RowCount()) - if err := evalAndScatter(bat, buildWriters, nil, e.keyExecs, true); err != nil { + if err := evalAndScatter(bat, buildWriters, e.keyExecs, true); err != nil { return nil, err } } @@ -2835,7 +3235,7 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal if err != nil { return nil, err } - if err := scatterProbe(proc, e, bat, probeWriters, nil, partitionLevel, analyzer); err != nil { + if err := scatterProbe(proc, e, bat, probeWriters, partitionLevel, analyzer); err != nil { return nil, err } } @@ -2862,8 +3262,8 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal if enqueue { if len(e.buckets)-1+len(subBuckets)+1 > e.cfg.MaxQueue { return nil, &process.HashBuildBudgetError{ - Kind: process.HashBuildBudgetErrorAdmission, - Message: fmt.Sprintf("join spill queue limit exceeded (limit=%d); reduce join-key skew or increase processLimitationSize", e.cfg.MaxQueue), + Kind: process.HashBuildBudgetErrorInvalid, + Message: fmt.Sprintf("spill queue limit exceeded: limit=%d", e.cfg.MaxQueue), } } buildFile, err := buildWriters[i].handOffSpillFile() @@ -2998,12 +3398,18 @@ func (e *SpillEngine) AdvanceToNextBucket( // scatterProbe evaluates probe-side keys (EqConds[0]) for probe re-scatter. // It uses the borrowed probe lease, not build-side keyExecs; probeKeyEval is // retained only as the unbudgeted fallback. -func scatterProbe(proc *process.Process, e *SpillEngine, bat *batch.Batch, writers []BucketWriter, buffers []*batch.Batch, seed uint64, analyzer process.Analyzer) error { - keyVecs, err := e.evalProbeKeys(proc, bat, e.probeKeyEval) - if err != nil { - return err - } - return e.scatterBatch(proc, bat, keyVecs, writers, buffers, seed, true, analyzer) +func scatterProbe(proc *process.Process, e *SpillEngine, bat *batch.Batch, writers []BucketWriter, seed uint64, analyzer process.Analyzer) error { + return e.scatterEvaluatedBatchWithPressure( + proc, + bat, + writers, + seed, + true, + analyzer, + func(current *batch.Batch) ([]*vector.Vector, error) { + return e.evalProbeKeys(proc, current, e.probeKeyEval) + }, + ) } func (e *SpillEngine) evalProbeKeys( @@ -3056,26 +3462,16 @@ func (e *SpillEngine) freeKeyExecs() { } func isBudgetAdmission(err error) bool { - return err != nil && - errors.Is(err, process.ErrHashBuildBudgetAdmission) + return hashbuild.IsRetryableMemoryCapacity(err) } -func noProgressError(depth int, cause error) error { - budgetErr := &process.HashBuildBudgetError{ - Kind: process.HashBuildBudgetErrorAdmission, - Message: fmt.Sprintf("join spill cannot make progress at depth %d; reduce join-key skew or increase processLimitationSize", depth), - } - if cause != nil { - var budgetCause *process.HashBuildBudgetError - if errors.As(cause, &budgetCause) && budgetCause.Kind == process.HashBuildBudgetErrorAdmission { - budgetErr.Resource = budgetCause.Resource - budgetErr.Requested = budgetCause.Requested - budgetErr.Used = budgetCause.Used - budgetErr.Cap = budgetCause.Cap - budgetErr.Message = fmt.Sprintf("join spill cannot make progress at depth %d", depth) - } - } - return budgetErr +func noProgressError(proc *process.Process, depth int) error { + _ = proc + return hashbuild.NewMinimumAllocationPressureError( + "join-spill", + fmt.Sprintf("partition-depth-%d", depth), + nil, + ) } // Cleanup releases all engine resources. diff --git a/pkg/sql/colexec/spillutil/join_spill_test.go b/pkg/sql/colexec/spillutil/join_spill_test.go index 5473ed6437e18..aa41b34632588 100644 --- a/pkg/sql/colexec/spillutil/join_spill_test.go +++ b/pkg/sql/colexec/spillutil/join_spill_test.go @@ -1941,33 +1941,17 @@ func TestSpillEngineInitFromOwnedFilesAndErrorClassification(t *testing.T) { require.False(t, isBudgetAdmission(io.EOF)) require.True(t, isBudgetAdmission(process.ErrHashBuildBudgetAdmission)) require.False(t, isBudgetAdmission(process.ErrHashBuildBudgetClosed)) - require.True(t, isBudgetAdmission(&process.HashBuildBudgetError{ - Kind: process.HashBuildBudgetErrorAdmission, - })) require.False(t, isBudgetAdmission(&process.HashBuildBudgetError{ - Kind: process.HashBuildBudgetErrorClosed, + Kind: process.HashBuildBudgetErrorAdmission, + Component: process.HashBuildBudgetComponentSpillDisk, })) - logicalNoProgress := noProgressError(3, nil) - require.ErrorIs(t, logicalNoProgress, process.ErrHashBuildBudgetAdmission) - require.Contains(t, logicalNoProgress.Error(), "depth 3") - require.Contains(t, logicalNoProgress.Error(), "reduce join-key skew") - require.NotContains(t, logicalNoProgress.Error(), process.ErrHashBuildBudgetAdmission.Error()) - - memoryAdmission := &process.HashBuildBudgetError{ + require.False(t, isBudgetAdmission(&process.HashBuildBudgetError{ Kind: process.HashBuildBudgetErrorAdmission, - Resource: process.HashBuildBudgetResourceMemory, - Requested: 11, - Used: 13, - Cap: 17, - } - budgetNoProgress := noProgressError(4, memoryAdmission) - var budgetErr *process.HashBuildBudgetError - require.ErrorAs(t, budgetNoProgress, &budgetErr) - require.Equal(t, process.HashBuildBudgetResourceMemory, budgetErr.Resource) - require.Equal(t, uint64(11), budgetErr.Requested) - require.Equal(t, uint64(13), budgetErr.Used) - require.Equal(t, uint64(17), budgetErr.Cap) - require.Equal(t, "join spill cannot make progress at depth 4", budgetErr.Message) + Component: process.HashBuildBudgetComponentSpillFD, + })) + require.Equal(t, + hashbuild.MemoryPressureMinimumUnit, + hashbuild.MemoryPressureReasonOf(noProgressError(nil, 3))) require.NoError(t, owned.Close()) } @@ -2895,12 +2879,11 @@ func TestSpillEntryPointsRejectPreCanceledProcessWithoutOwnershipTransfer(t *tes t.Cleanup(cleanEngine) scatterWriters := []BucketWriter{{Name: "must_not_be_created"}} t.Cleanup(scatterWriters[0].Close) - err = engine.scatterBatch( + err = engine.scatterBatchWithPressure( proc, input, []*vector.Vector{input.Vecs[0]}, scatterWriters, - nil, 0, false, process.NewAnalyzer(0, false, false, "test"), @@ -4066,16 +4049,15 @@ func TestScatterProbeFunctionUsesStoredEval(t *testing.T) { } writers := MakeBucketWriters("test_scatter_func") - buffers := make([]*batch.Batch, len(writers)) bat := makeInt32Batch(proc, []int32{5, 15, 25}) - err := scatterProbe(proc, engine, bat, writers, buffers, 1, nil) + err := scatterProbe(proc, engine, bat, writers, 1, nil) require.NoError(t, err) require.True(t, evalCalled, "probeKeyEval must be used for scatterProbe") wantErr := errors.New("probe key evaluation failed") engine.probeKeyEval = func(*batch.Batch) ([]*vector.Vector, error) { return nil, wantErr } - require.ErrorIs(t, scatterProbe(proc, engine, bat, writers, buffers, 1, nil), wantErr) + require.ErrorIs(t, scatterProbe(proc, engine, bat, writers, 1, nil), wantErr) for i := range writers { writers[i].Close() diff --git a/pkg/vm/process/hashbuild_budget.go b/pkg/vm/process/hashbuild_budget.go index f8283a30ff50f..4734ceffc1c15 100644 --- a/pkg/vm/process/hashbuild_budget.go +++ b/pkg/vm/process/hashbuild_budget.go @@ -33,12 +33,16 @@ const hashBuildMinimumReserve = uint64(4 << 30) const ( hashBuildAllocationGenerationSlots = uint32(131_072) - hashBuildMinimumCellBlockBytes = uint64(16 << 10) - // Three slots close the cell/descriptor replacement transaction. The - // copied-batch activation adds up to three vector buffers per minimum-width - // 8,192-row destination (data, nulls, grouping), so six slots per 16 KiB of - // aggregate capacity is the first combined owner bound. - hashBuildAllocationSlotsPerBlock = uint64(6) + // Account metadata lives outside the MPool payload cap. Bound its worst-case + // Go-heap footprint to half of hashBuildMinimumReserve: 128 bytes covers one + // sparse pointer/lease-map entry, and 16,777,216 live entries consume at most + // 2 GiB by construction. Small aggregate caps use the tighter byte + // conservation bound because every physical allocation owns at least one + // byte. Slot exhaustion is real metadata capacity pressure, not an + // estimator rejection. + hashBuildAllocationMetadataBytesPerSlot = uint64(128) + hashBuildAllocationMetadataHeadroom = hashBuildMinimumReserve / 2 + hashBuildAllocationMetadataMaxSlots = hashBuildAllocationMetadataHeadroom / hashBuildAllocationMetadataBytesPerSlot ) const ( @@ -146,16 +150,18 @@ const ( HashBuildBudgetErrorCeilingMissing ) -// HashBuildBudgetResource identifies the finite resource whose admission -// failed. It is separate from Kind: every resource can reject capacity, while -// lifecycle and accounting failures remain resource-independent. -type HashBuildBudgetResource uint8 +// HashBuildBudgetComponent identifies the independently bounded resource that +// rejected an admission. The zero value remains the memory component for +// compatibility with older callers that construct HashBuildBudgetError +// directly. A spill-disk or spill-FD rejection must never enter the memory +// reclaim/reduce loop: reducing an in-memory batch cannot create either +// resource and may replay already-published spill records. +type HashBuildBudgetComponent uint8 const ( - HashBuildBudgetResourceUnknown HashBuildBudgetResource = iota - HashBuildBudgetResourceMemory - HashBuildBudgetResourceSpillDisk - HashBuildBudgetResourceSpillFD + HashBuildBudgetComponentMemory HashBuildBudgetComponent = iota + HashBuildBudgetComponentSpillDisk + HashBuildBudgetComponentSpillFD ) // HashBuildBudgetError carries bounded, observational details for an @@ -164,7 +170,7 @@ const ( // to inspect; they are never produced by overflowing arithmetic. type HashBuildBudgetError struct { Kind HashBuildBudgetErrorKind - Resource HashBuildBudgetResource + Component HashBuildBudgetComponent Requested uint64 Used uint64 Cap uint64 @@ -422,10 +428,10 @@ func (b *HashBuildBudget) SetSpillCaps(diskBytes, fds uint64) error { } effectiveFDCap := clampSpillFDCap(fds, processLimit, limitKnown) if b.spillDiskUsed > diskBytes { - return newAdmissionError(HashBuildBudgetResourceSpillDisk, 0, b.spillDiskUsed, diskBytes) + return newComponentAdmissionError(HashBuildBudgetComponentSpillDisk, 0, b.spillDiskUsed, diskBytes) } if b.spillFDUsed > effectiveFDCap { - return newAdmissionError(HashBuildBudgetResourceSpillFD, 0, b.spillFDUsed, effectiveFDCap) + return newComponentAdmissionError(HashBuildBudgetComponentSpillFD, 0, b.spillFDUsed, effectiveFDCap) } b.spillDiskCap = diskBytes b.spillFDConfiguredCap = fds @@ -1141,9 +1147,12 @@ func (g *HashBuildBudgetGeneration) Closed() bool { // AllocationAccountRegistry returns the bounded CN-local registry shared by // every activated HashBuild generation under this aggregate budget. The slot -// formula covers every minimum-size cell block, its published descriptor, one -// private replacement transaction, and the copied-batch vector buffers added -// by the second activation. +// bound follows a conservation fact rather than a per-operator multiplier: +// every live allocation owns at least one byte and all accounts share the +// aggregate byte cap, so live metadata cannot exceed aggregate capacity. A +// second fixed bound reserves at most 2 GiB of the existing 4 GiB CN +// headroom at 128 bytes per metadata entry. The registry stores only the +// resulting scalar limit; it does not preallocate one object per slot. func (g *HashBuildBudgetGeneration) AllocationAccountRegistry() ( *mpool.AllocationAccountRegistry, error, @@ -1154,21 +1163,15 @@ func (g *HashBuildBudgetGeneration) AllocationAccountRegistry() ( b := g.budget b.allocationRegistryOnce.Do(func() { capBytes := b.AggregateCap() - blocks := capBytes / hashBuildMinimumCellBlockBytes - if capBytes%hashBuildMinimumCellBlockBytes != 0 { - blocks++ - } - if blocks == 0 { - blocks = 1 - } - if blocks > math.MaxUint64/hashBuildAllocationSlotsPerBlock { + if capBytes == 0 { b.allocationRegistryErr = ErrHashBuildBudgetInvalid return } + allocationSlots := min(capBytes, hashBuildAllocationMetadataMaxSlots) b.allocationRegistry, b.allocationRegistryErr = mpool.NewAllocationAccountRegistry( hashBuildAllocationGenerationSlots, - blocks*hashBuildAllocationSlotsPerBlock, + allocationSlots, ) }) return b.allocationRegistry, b.allocationRegistryErr @@ -1355,7 +1358,7 @@ func (g *HashBuildBudgetGeneration) reserveLocked( g.rejectCount++ observeHashBuildBudget("memory", "reject", "cn", size) } - return nil, newAdmissionError(HashBuildBudgetResourceMemory, size, b.aggregateUsed, b.aggregateCap), true + return nil, newAdmissionError(size, b.aggregateUsed, b.aggregateCap), true } b.aggregateUsed += size if g.used > g.cap || size > g.cap-g.used { @@ -1363,7 +1366,7 @@ func (g *HashBuildBudgetGeneration) reserveLocked( b.aggregateUsed -= size g.rejectCount++ observeHashBuildBudget("memory", "reject", "query", size) - return nil, newAdmissionError(HashBuildBudgetResourceMemory, size, g.used, g.cap), false + return nil, newAdmissionError(size, g.used, g.cap), false } g.used += size g.reserveCount++ @@ -1490,12 +1493,12 @@ func (r *HashBuildReservation) growLocked(additional uint64, recordAggregateReje g.rejectCount++ observeHashBuildBudget("memory", "reject", "cn", additional) } - return newAdmissionError(HashBuildBudgetResourceMemory, additional, b.aggregateUsed, b.aggregateCap), true + return newAdmissionError(additional, b.aggregateUsed, b.aggregateCap), true } if g.used > g.cap || additional > g.cap-g.used { g.rejectCount++ observeHashBuildBudget("memory", "reject", "query", additional) - return newAdmissionError(HashBuildBudgetResourceMemory, additional, g.used, g.cap), false + return newAdmissionError(additional, g.used, g.cap), false } if r.core.size > math.MaxUint64-additional { return &HashBuildBudgetError{Kind: HashBuildBudgetErrorInvalid, Requested: additional, Message: "hash build reservation size overflow"}, false @@ -1510,10 +1513,24 @@ func (r *HashBuildReservation) growLocked(additional uint64, recordAggregateReje return nil, false } -func newAdmissionError(resource HashBuildBudgetResource, requested, used, cap uint64) error { +func newAdmissionError(requested, used, cap uint64) error { + return newComponentAdmissionError( + HashBuildBudgetComponentMemory, + requested, + used, + cap, + ) +} + +func newComponentAdmissionError( + component HashBuildBudgetComponent, + requested uint64, + used uint64, + cap uint64, +) error { return &HashBuildBudgetError{ Kind: HashBuildBudgetErrorAdmission, - Resource: resource, + Component: component, Requested: requested, Used: used, Cap: cap, @@ -1721,12 +1738,12 @@ func (g *HashBuildBudgetGeneration) ReserveSpillDisk(size uint64) (*HashBuildSpi if b.spillDiskUsed > b.spillDiskCap || size > b.spillDiskCap-b.spillDiskUsed { g.rejectCount++ observeHashBuildBudget("spill_disk", "reject", "cn", size) - return nil, newAdmissionError(HashBuildBudgetResourceSpillDisk, size, b.spillDiskUsed, b.spillDiskCap) + return nil, newComponentAdmissionError(HashBuildBudgetComponentSpillDisk, size, b.spillDiskUsed, b.spillDiskCap) } if g.spillDiskUsed > g.spillDiskCap || size > g.spillDiskCap-g.spillDiskUsed { g.rejectCount++ observeHashBuildBudget("spill_disk", "reject", "query", size) - return nil, newAdmissionError(HashBuildBudgetResourceSpillDisk, size, g.spillDiskUsed, g.spillDiskCap) + return nil, newComponentAdmissionError(HashBuildBudgetComponentSpillDisk, size, g.spillDiskUsed, g.spillDiskCap) } b.spillDiskUsed += size g.spillDiskUsed += size @@ -1764,12 +1781,12 @@ func (r *HashBuildSpillDiskReservation) Grow(additional uint64) error { if b.spillDiskUsed > b.spillDiskCap || additional > b.spillDiskCap-b.spillDiskUsed { g.rejectCount++ observeHashBuildBudget("spill_disk", "reject", "cn", additional) - return newAdmissionError(HashBuildBudgetResourceSpillDisk, additional, b.spillDiskUsed, b.spillDiskCap) + return newComponentAdmissionError(HashBuildBudgetComponentSpillDisk, additional, b.spillDiskUsed, b.spillDiskCap) } if g.spillDiskUsed > g.spillDiskCap || additional > g.spillDiskCap-g.spillDiskUsed { g.rejectCount++ observeHashBuildBudget("spill_disk", "reject", "query", additional) - return newAdmissionError(HashBuildBudgetResourceSpillDisk, additional, g.spillDiskUsed, g.spillDiskCap) + return newComponentAdmissionError(HashBuildBudgetComponentSpillDisk, additional, g.spillDiskUsed, g.spillDiskCap) } if r.core.size > math.MaxUint64-additional { return &HashBuildBudgetError{Kind: HashBuildBudgetErrorInvalid, Requested: additional, Message: "spill disk reservation size overflow"} @@ -1807,12 +1824,12 @@ func (g *HashBuildBudgetGeneration) ReserveSpillFD(size uint64) (*HashBuildSpill if b.spillFDUsed > b.spillFDCap || size > b.spillFDCap-b.spillFDUsed { g.rejectCount++ observeHashBuildBudget("spill_fd", "reject", "cn", size) - return nil, newAdmissionError(HashBuildBudgetResourceSpillFD, size, b.spillFDUsed, b.spillFDCap) + return nil, newComponentAdmissionError(HashBuildBudgetComponentSpillFD, size, b.spillFDUsed, b.spillFDCap) } if g.spillFDUsed > g.spillFDCap || size > g.spillFDCap-g.spillFDUsed { g.rejectCount++ observeHashBuildBudget("spill_fd", "reject", "query", size) - return nil, newAdmissionError(HashBuildBudgetResourceSpillFD, size, g.spillFDUsed, g.spillFDCap) + return nil, newComponentAdmissionError(HashBuildBudgetComponentSpillFD, size, g.spillFDUsed, g.spillFDCap) } b.spillFDUsed += size g.spillFDUsed += size diff --git a/pkg/vm/process/hashbuild_budget_test.go b/pkg/vm/process/hashbuild_budget_test.go index e06cb37be6b47..f512c6ed442f8 100644 --- a/pkg/vm/process/hashbuild_budget_test.go +++ b/pkg/vm/process/hashbuild_budget_test.go @@ -73,12 +73,12 @@ func TestHashBuildBudgetAdmissionIdentifiesResource(t *testing.T) { tests := []struct { name string - want HashBuildBudgetResource + want HashBuildBudgetComponent call func() error }{ { name: "memory", - want: HashBuildBudgetResourceMemory, + want: HashBuildBudgetComponentMemory, call: func() error { _, reserveErr := g.Reserve(11) return reserveErr @@ -86,7 +86,7 @@ func TestHashBuildBudgetAdmissionIdentifiesResource(t *testing.T) { }, { name: "spill disk", - want: HashBuildBudgetResourceSpillDisk, + want: HashBuildBudgetComponentSpillDisk, call: func() error { _, reserveErr := g.ReserveSpillDisk(6) return reserveErr @@ -94,7 +94,7 @@ func TestHashBuildBudgetAdmissionIdentifiesResource(t *testing.T) { }, { name: "spill fd", - want: HashBuildBudgetResourceSpillFD, + want: HashBuildBudgetComponentSpillFD, call: func() error { _, reserveErr := g.ReserveSpillFD(2) return reserveErr @@ -107,9 +107,9 @@ func TestHashBuildBudgetAdmissionIdentifiesResource(t *testing.T) { if err := test.call(); !errors.As(err, &budgetErr) { t.Fatalf("error=%v, want typed admission", err) } - if budgetErr.Kind != HashBuildBudgetErrorAdmission || budgetErr.Resource != test.want { + if budgetErr.Kind != HashBuildBudgetErrorAdmission || budgetErr.Component != test.want { t.Fatalf("admission kind/resource=(%d,%d), want=(%d,%d)", - budgetErr.Kind, budgetErr.Resource, HashBuildBudgetErrorAdmission, test.want) + budgetErr.Kind, budgetErr.Component, HashBuildBudgetErrorAdmission, test.want) } }) } @@ -225,7 +225,7 @@ func TestHashBuildBudgetAllocationAccountAdapter(t *testing.T) { } } -func TestHashBuildAllocationAccountRegistryUsesBoundedFormula(t *testing.T) { +func TestHashBuildAllocationAccountRegistryUsesByteConservationBound(t *testing.T) { budget := MustNewHashBuildBudget(16<<10, 16<<10) first, err := budget.OpenGeneration(1) if err != nil { @@ -238,11 +238,11 @@ func TestHashBuildAllocationAccountRegistryUsesBoundedFormula(t *testing.T) { if registry.GenerationCapacity() != hashBuildAllocationGenerationSlots { t.Fatalf("generation slots = %d", registry.GenerationCapacity()) } - if registry.MaxAllocationMetadata() != hashBuildAllocationSlotsPerBlock { + if registry.MaxAllocationMetadata() != 16<<10 { t.Fatalf( "allocation slots = %d, want %d", registry.MaxAllocationMetadata(), - hashBuildAllocationSlotsPerBlock, + uint64(16<<10), ) } second, err := budget.OpenGeneration(2) @@ -258,6 +258,24 @@ func TestHashBuildAllocationAccountRegistryUsesBoundedFormula(t *testing.T) { } } +func TestHashBuildAllocationAccountRegistryCapsMetadataHeadroom(t *testing.T) { + budget := MustNewHashBuildBudget( + hashBuildAllocationMetadataMaxSlots+1, + hashBuildAllocationMetadataMaxSlots+1, + ) + generation, err := budget.OpenGeneration(1) + if err != nil { + t.Fatal(err) + } + registry, err := generation.AllocationAccountRegistry() + if err != nil { + t.Fatal(err) + } + if got := registry.MaxAllocationMetadata(); got != hashBuildAllocationMetadataMaxSlots { + t.Fatalf("allocation slots = %d, want %d", got, hashBuildAllocationMetadataMaxSlots) + } +} + func TestHashBuildBudgetQueryRejectRollsBackCN(t *testing.T) { b := MustNewHashBuildBudget(10, 4) g1, _ := b.OpenGeneration(1) @@ -1195,6 +1213,9 @@ func TestHashBuildBudgetCompatibilityAndObservabilitySurface(t *testing.T) { !errors.Is(unknown, ErrHashBuildBudgetInvalid) { t.Fatal("unknown error kind must remain a fatal invalid error") } + if unknown.Is(ErrHashBuildBudgetInvalid) { + t.Fatal("unknown error kind matched a sentinel") + } message := &HashBuildBudgetError{Message: "explicit"} if message.Error() != "explicit" { t.Fatalf("explicit message = %q", message.Error()) @@ -1472,6 +1493,59 @@ func TestHashBuildBudgetCompatibilityUnhappyPaths(t *testing.T) { } } +func TestHashBuildBudgetAdmissionNamesIndependentComponent(t *testing.T) { + b, err := NewHashBuildBudgetWithSpillCaps(8, 8, 1, 1) + if err != nil { + t.Fatal(err) + } + g, err := b.OpenGenerationWithSpillCaps(1, 8, 1, 1) + if err != nil { + t.Fatal(err) + } + defer g.Close() + + assertComponent := func(err error, want HashBuildBudgetComponent) { + t.Helper() + var budgetErr *HashBuildBudgetError + if !errors.As(err, &budgetErr) || budgetErr.Component != want { + t.Fatalf("admission component: err=%v got=%v want=%v", err, budgetErr, want) + } + } + + memory, err := g.Reserve(8) + if err != nil { + t.Fatal(err) + } + defer memory.Release() + _, err = g.Reserve(1) + assertComponent(err, HashBuildBudgetComponentMemory) + + disk, err := g.ReserveSpillDisk(1) + if err != nil { + t.Fatal(err) + } + defer disk.Release() + _, err = g.ReserveSpillDisk(1) + assertComponent(err, HashBuildBudgetComponentSpillDisk) + + fd, err := g.ReserveSpillFD(1) + if err != nil { + t.Fatal(err) + } + defer fd.Release() + _, err = g.ReserveSpillFD(1) + assertComponent(err, HashBuildBudgetComponentSpillFD) + + err = b.SetSpillCaps(0, 1) + if err != nil { + t.Fatal(err) + } + err = b.SetSpillCaps(1, 0) + if err != nil { + t.Fatal(err) + } +} + func TestGetHashBuildBudgetInitializesAndReusesCNAggregate(t *testing.T) { const localService = "__process_local_cn__" hashBuildCNBudgets.Delete(localService) From e27a3324fe3e117eb8ac74a9bb50269de821362d Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 23:45:01 +0800 Subject: [PATCH 20/61] test: close allocation accounting performance matrix --- pkg/common/mpool/mpool.go | 36 ++--- pkg/sql/colexec/hashbuild/hashmap_test.go | 101 ++++++++++++++ .../spillutil/allocation_account_test.go | 94 +++++++++++++ pkg/vm/process/hashbuild_budget_test.go | 123 ++++++++++++++++++ 4 files changed, 339 insertions(+), 15 deletions(-) diff --git a/pkg/common/mpool/mpool.go b/pkg/common/mpool/mpool.go index 288189f2ac31e..e3d5da9bf8284 100644 --- a/pkg/common/mpool/mpool.go +++ b/pkg/common/mpool/mpool.go @@ -809,34 +809,40 @@ func (mp *MPool) allocAccountedWithDetailK( detailk string, sz int64, request allocationAccountRequest, -) (result []byte, retErr error) { - defer func() { - if retErr != nil { - retErr = fmt.Errorf( - "allocation owner=%d site=%d: %w", - request.owner, - request.site, - retErr, - ) - } - }() +) ([]byte, error) { if err := request.validate(); err != nil { - return nil, err + return nil, allocationAccountSiteError(request, err) } // reject unexpected alloc size. if sz < 0 || sz > maxAllocationSize() { logutil.Errorf("mpool memory allocation exceed limit with requested size %d: %s", sz, string(debug.Stack())) - return nil, fmt.Errorf( + return nil, allocationAccountSiteError(request, fmt.Errorf( "%w: requested=%d maximum=%d", ErrAllocationAllocatorLimit, sz, maxAllocationSize(), - ) + )) } if sz == 0 { return nil, nil } - return mp.allocAccounted(detailk, sz, request) + result, err := mp.allocAccounted(detailk, sz, request) + if err != nil { + return nil, allocationAccountSiteError(request, err) + } + return result, nil +} + +func allocationAccountSiteError( + request allocationAccountRequest, + err error, +) error { + return fmt.Errorf( + "allocation owner=%d site=%d: %w", + request.owner, + request.site, + err, + ) } func (mp *MPool) alloc( diff --git a/pkg/sql/colexec/hashbuild/hashmap_test.go b/pkg/sql/colexec/hashbuild/hashmap_test.go index 200ced7c14b9e..e120b56006e13 100644 --- a/pkg/sql/colexec/hashbuild/hashmap_test.go +++ b/pkg/sql/colexec/hashbuild/hashmap_test.go @@ -17,6 +17,7 @@ package hashbuild import ( "context" "errors" + "fmt" "math" "reflect" "strconv" @@ -2791,6 +2792,106 @@ func BenchmarkCopyBuildBatchAccounting(b *testing.B) { } } +// BenchmarkResidentHashBuildAccounting compares the complete resident owner +// closure, not only the primitive allocator: copied batches, key expression, +// hash cells/descriptors, and terminal release all run on every iteration. +// The 32-row case models high-frequency TP statements; 8,192 rows exercises a +// full physical batch without entering spill. +func BenchmarkResidentHashBuildAccounting(b *testing.B) { + const capBytes = uint64(512 << 20) + for _, rows := range []int{32, colexec.DefaultBatchSize} { + for _, stringKey := range []bool{false, true} { + kind := "int" + if stringKey { + kind = "varchar" + } + for _, accounted := range []bool{false, true} { + mode := "legacy" + if accounted { + mode = "accounted" + } + b.Run(fmt.Sprintf("%s/%s/rows-%d", mode, kind, rows), func(b *testing.B) { + proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) + defer proc.Free() + var input *batch.Batch + var keyType types.Type + if stringKey { + input = makeStrBatch(b, rows, proc) + keyType = types.T_varchar.ToType() + } else { + input = makeIntBatch(b, rows, proc) + keyType = types.T_int32.ToType() + } + defer input.Clean(proc.Mp()) + + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err := budget.OpenGeneration(1) + if err != nil { + b.Fatal(err) + } + var ( + registry *mpool.AllocationAccountRegistry + account *mpool.AllocationAccount + ) + if accounted { + registry, err = mpool.NewAllocationAccountRegistry(1, 4_096) + if err != nil { + b.Fatal(err) + } + account, err = registry.OpenWithController(capBytes, generation) + if err != nil { + b.Fatal(err) + } + } + + b.ReportAllocs() + b.SetBytes(int64(input.Size())) + b.ResetTimer() + for range b.N { + hb := &HashmapBuilder{} + hb.SetBudget(generation) + if accounted { + if err = hb.SetAllocationAccount(account); err != nil { + b.Fatal(err) + } + } + if err = hb.Prepare( + []*plan.Expr{newExpr(0, keyType)}, + -1, + -1, + nil, + proc, + ); err != nil { + b.Fatal(err) + } + hb.InputBatchRowCount = input.RowCount() + if err = hb.CopyBuildBatch(input, proc); err != nil { + b.Fatal(err) + } + if err = hb.BuildHashmap(false, false, false, proc); err != nil { + b.Fatal(err) + } + hb.Free(proc) + } + b.StopTimer() + if generation.Used() != 0 { + b.Fatalf("generation used = %d", generation.Used()) + } + if accounted { + if account.Snapshot().Used != 0 { + b.Fatalf("account used = %d", account.Snapshot().Used) + } + if _, _, err = registry.CompleteTerminal(account); err != nil { + b.Fatal(err) + } + } + generation.Close() + }) + } + } + } +} + func TestExtractRestoreCachedIterators(t *testing.T) { var hb HashmapBuilder mp := mpool.MustNewZero() diff --git a/pkg/sql/colexec/spillutil/allocation_account_test.go b/pkg/sql/colexec/spillutil/allocation_account_test.go index 7a37bf1a9285d..8ededc0917478 100644 --- a/pkg/sql/colexec/spillutil/allocation_account_test.go +++ b/pkg/sql/colexec/spillutil/allocation_account_test.go @@ -639,6 +639,100 @@ func TestSpillAllocationAccountExpressionPressureReducesBeforePublication(t *tes finalizeTestSpillAllocationAccount(t, state) } +// BenchmarkSpillScatterAccounting measures the steady streaming closure with +// the same selected-vector, hash/row-ID, marshal, and coalesce owners used by +// initial and recursive join spill. /dev/null keeps the benchmark bounded and +// retains the write syscall without turning repeated measurements into a disk +// capacity test. +func BenchmarkSpillScatterAccounting(b *testing.B) { + for _, accounted := range []bool{false, true} { + mode := "legacy" + if accounted { + mode = "accounted" + } + b.Run(mode, func(b *testing.B) { + proc := testutil.NewProcessWithMPool( + b, + "", + mpool.MustNewZero(), + ) + defer proc.Free() + values := make([]int64, 4_096) + for i := range values { + values[i] = int64(i) + } + source := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.MakeInt64Vector(values, nil, proc.Mp()), + }, nil) + defer source.Clean(proc.Mp()) + + var ( + engine *SpillEngine + state testSpillAllocationAccount + err error + ) + if accounted { + state = newTestSpillAllocationAccount(b, 64<<20, 4_096) + engine, err = NewSpillEngineWithAllocation( + SpillEngineConfig{}, + state.allocation, + ) + if err != nil { + b.Fatal(err) + } + } else { + engine = NewSpillEngine(SpillEngineConfig{}) + } + writers := make([]BucketWriter, SpillNumBuckets) + for i := range writers { + writers[i].Name = "benchmark-discard" + writers[i].Fd, err = os.OpenFile(os.DevNull, os.O_WRONLY, 0) + if err != nil { + b.Fatal(err) + } + } + defer func() { + for i := range writers { + writers[i].Close() + } + }() + analyzer := process.NewAnalyzer(0, false, false, "benchmark") + + b.ReportAllocs() + b.SetBytes(int64(source.Size())) + b.ResetTimer() + for range b.N { + if err = engine.scatterBatchWithPressure( + proc, + source, + source.Vecs, + writers, + 0, + false, + analyzer, + ); err != nil { + b.Fatal(err) + } + if err = engine.flushScatterBuffers(proc, writers, analyzer); err != nil { + b.Fatal(err) + } + for i := range writers { + writers[i].Rows = 0 + writers[i].Bytes = 0 + } + } + b.StopTimer() + engine.Cleanup(proc) + if accounted { + if state.account.Snapshot().Used != 0 { + b.Fatalf("account used = %d", state.account.Snapshot().Used) + } + finalizeTestSpillAllocationAccount(b, state) + } + }) + } +} + func TestSpillAllocationAccountRebuildAndRecursiveSpillLifecycle(t *testing.T) { proc := testutil.NewProcessWithMPool( t, diff --git a/pkg/vm/process/hashbuild_budget_test.go b/pkg/vm/process/hashbuild_budget_test.go index f512c6ed442f8..3d250e54e9f79 100644 --- a/pkg/vm/process/hashbuild_budget_test.go +++ b/pkg/vm/process/hashbuild_budget_test.go @@ -20,6 +20,7 @@ import ( "os" "os/exec" "runtime" + "sort" "strconv" "sync" "sync/atomic" @@ -2111,6 +2112,128 @@ func BenchmarkHashBuildBudgetAllocationAccount(b *testing.B) { } } +// BenchmarkHashBuildAllocationAttemptLifecycle covers the production control +// plane around high-frequency statements: concurrent generation open, account +// publication, one physical owner allocation, release, terminal snapshot, and +// slot reuse. The explicit quantiles make the contention tail visible when the +// benchmark is run with -cpu=1,8. +func BenchmarkHashBuildAllocationAttemptLifecycle(b *testing.B) { + const ( + aggregateCap = uint64(1 << 50) + attemptCap = uint64(1 << 20) + ) + budget := MustNewHashBuildBudget(aggregateCap, attemptCap) + registry, err := commonmpool.NewAllocationAccountRegistry(4_096, 4_096) + if err != nil { + b.Fatal(err) + } + mp := commonmpool.MustNew("hash-build-attempt-lifecycle-benchmark") + defer commonmpool.DeleteMPool(mp) + latencies := make([]int64, b.N) + var ( + nextID atomic.Uint64 + sample atomic.Uint64 + ) + + b.ReportAllocs() + b.SetBytes(4 << 10) + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + started := time.Now() + generation, openErr := budget.OpenGenerationWithCap( + nextID.Add(1), + attemptCap, + ) + if openErr != nil { + b.Errorf("open generation: %v", openErr) + return + } + account, openErr := registry.OpenWithController(attemptCap, generation) + if openErr != nil { + generation.Close() + b.Errorf("open account: %v", openErr) + return + } + buffer, allocErr := mp.AllocAccounted(4<<10, account, 1, 1) + if allocErr == nil { + mp.Free(buffer) + } + _, _, terminalErr := registry.CompleteTerminal(account) + generation.Close() + if allocErr != nil || terminalErr != nil { + b.Errorf("attempt alloc=%v terminal=%v", allocErr, terminalErr) + return + } + index := sample.Add(1) - 1 + latencies[index] = time.Since(started).Nanoseconds() + } + }) + b.StopTimer() + count := int(sample.Load()) + if count != b.N { + b.Fatalf("completed attempts = %d, want %d", count, b.N) + } + sort.Slice(latencies, func(i, j int) bool { return latencies[i] < latencies[j] }) + b.ReportMetric(float64(latencies[(count-1)*50/100]), "p50-ns/op") + b.ReportMetric(float64(latencies[(count-1)*99/100]), "p99-ns/op") + if budget.AggregateUsed() != 0 || registry.LiveAllocationMetadata() != 0 { + b.Fatalf( + "terminal leak: budget=%d metadata=%d", + budget.AggregateUsed(), + registry.LiveAllocationMetadata(), + ) + } +} + +// BenchmarkHashBuildAllocationReleaseStorm isolates concurrent physical +// alloc/free against one live generation, the shape produced when broadcast +// consumers and spill buffers drain together. +func BenchmarkHashBuildAllocationReleaseStorm(b *testing.B) { + const capacity = uint64(1 << 50) + budget := MustNewHashBuildBudget(capacity, capacity) + generation, err := budget.OpenGeneration(1) + if err != nil { + b.Fatal(err) + } + registry, err := commonmpool.NewAllocationAccountRegistry(1, 65_536) + if err != nil { + b.Fatal(err) + } + account, err := registry.OpenWithController(capacity, generation) + if err != nil { + b.Fatal(err) + } + mp := commonmpool.MustNew("hash-build-release-storm-benchmark") + defer commonmpool.DeleteMPool(mp) + + b.ReportAllocs() + b.SetBytes(4 << 10) + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + buffer, allocErr := mp.AllocAccounted(4<<10, account, 1, 1) + if allocErr != nil { + b.Errorf("allocate: %v", allocErr) + return + } + mp.Free(buffer) + } + }) + b.StopTimer() + if account.Snapshot().Used != 0 || generation.Used() != 0 { + b.Fatalf( + "release storm leak: account=%d generation=%d", + account.Snapshot().Used, + generation.Used(), + ) + } + if _, _, err = registry.CompleteTerminal(account); err != nil { + b.Fatal(err) + } + generation.Close() +} + func TestResolveHashBuildCeiling(t *testing.T) { const gib = uint64(1 << 30) got, err := ResolveHashBuildCeiling(HashBuildCeilingInputs{ From 1f6cc299a765a7fa0d2694557d7903b10156d443 Mon Sep 17 00:00:00 2001 From: aptend Date: Thu, 30 Jul 2026 18:20:47 +0800 Subject: [PATCH 21/61] [RFC] allocation-accounted memory admission --- ...ocation_accounted_memory_admission_impl.md | 807 +++++++++++++++ .../26459_allocation_accounting_bench.txt | 156 +++ ...0_allocation_accounted_memory_admission.md | 949 ++++++++++++++++++ 3 files changed, 1912 insertions(+) create mode 100644 docs/design/allocation_accounted_memory_admission_impl.md create mode 100644 docs/design/evidence/26459_allocation_accounting_bench.txt create mode 100644 docs/rfcs/00000000_allocation_accounted_memory_admission.md diff --git a/docs/design/allocation_accounted_memory_admission_impl.md b/docs/design/allocation_accounted_memory_admission_impl.md new file mode 100644 index 0000000000000..1bb24759d21d8 --- /dev/null +++ b/docs/design/allocation_accounted_memory_admission_impl.md @@ -0,0 +1,807 @@ +# Allocation-Accounted Memory Admission: Implementation Plan + +- Status: draft +- Tracking issue: + [#26459](https://github.com/matrixorigin/matrixone/issues/26459) +- Architecture: + [Allocation-Accounted Memory Admission RFC](../rfcs/00000000_allocation_accounted_memory_admission.md) +- Baseline at plan creation: `main` at `38ce3a774` +- Rebased implementation baseline: `main` at `43c896462` +- Merged prerequisite: #26455 at `93e8b22d2` +- Independent design review: completed against RFC commit `a7d54cb5f` +- Activation status: blocked until PRs 1--4 and the selected owner's + allocation-site/Go-heap gates close + +## 1. Purpose and rules + +The RFC owns architecture and invariants. This file contains only implementation +decisions, allocation-site status, PR scopes, and evidence for #26459. + +Every implementation PR follows these rules: + +1. Migrate a complete owner closure: alloc, grow, reuse, handoff, Reset, Free, + and failure rollback. +2. One site is legacy, allocation-accounted, synchronously reclaimable named + scratch, or small statically bounded headroom metadata. +3. Enable exact accounting and remove the same owner's legacy hard charge in + one PR. +4. Do not call the account per row or on within-capacity reuse. +5. Do not store a mutable current account on Process or MPool. +6. Spill and cleanup memory remain accounted and preserve bounded progress. +7. Every merged PR is independently safe; a later PR cannot repair an unsafe + interval. +8. The final state has no permanent legacy/exact behavior switch. + +## 2. Starting point and allocation-site ledger + +### Allocation and container boundaries + +- `pkg/common/mpool/mpool.go`: `memHdr`, `Alloc`, `Grow`, `Grow2`, `Free`, and + `GrowCapacity`. `memHdr` is currently fixed at 16 bytes. Cross-pool Free + already delegates to the original MPool. +- `pkg/container/vector/vector.go`: data and varlen area are separate physical + buffers. `Reset*` retains capacity; `Free` releases owned buffers; `cantFree*` + identifies non-owning views. +- `pkg/container/batch/batch.go`: Clone/Dup/Union allocate destination buffers; + `CleanOnlyData` retains capacity and `Clean` frees vectors. +- `pkg/sql/colexec/evalExpression*.go` and + `pkg/container/vector/functionTools.go`: FunctionResult and expression + executors retain result capacity across Reset and release it at Free. + +### HashBuild and spill boundaries + +- `pkg/sql/colexec/hashbuild/{budget,hashmap,spill,types}.go`: copied batches, + expression estimates, hash-table tokens, auxiliary buffers, runtime-filter + marshal, and spill scratch. +- `pkg/common/hashmap`, `pkg/container/hashtable`: physical hash-table blocks + and resize plans. +- `pkg/vm/message/joinMapMsg.go`: reference-counted producer-to-consumer + handoff; final Free may happen after HashBuild Reset. +- `pkg/sql/colexec/spillutil/join_spill.go`: decoded batches, retained reuse, + scatter/row-ID/encode/decode buffers, and BucketReader/Writer cleanup. +- HashJoin, DedupJoin, and RightDedupJoin share the pressure and handoff + contract. + +### Merged #26455 boundary + +#26455 is the immediate correctness prerequisite, not the final accounting +mechanism: + +- it gives each HashBuild-owned expression root a generation-scoped retained + lease and fixes reuse/reset lifetime mismatches; +- it still reconciles executor-owned capacity at the operator layer and keeps + estimator-derived admission for uncovered growth overlap; +- it does not attach provenance to each MPool allocation, cover arbitrary + built-in Go-heap scratch, or provide statement terminal finalization; +- an activation PR removes the matching #26455 retained lease at the same time + exact allocation provenance becomes complete for that owner; +- exact MPool charges must never be stacked on the #26455 charge for the same + physical capacity. + +Ledger states: + +- `L`: legacy prediction/reservation; +- `D`: exact primitive exists but this production owner is dormant; +- `A`: allocation-accounted; +- `S`: named bounded explicit scratch; +- `H`: small, statically bounded Go/runtime metadata covered by CN headroom; +- `R`: removed. + +The independent review rejected owner-class rows as proof of closure. The +working ledger is allocation-site based: + +| Allocation site | Allocator/mode and size | Terminal owner | Initial | Target/blocker | +| --- | --- | --- | ---: | --- | +| `mpool.memHdr` and account-ID side map | Go maps; one pointer record plus optional account record per live allocation | pointer removal at physical deallocation | L | H after per-entry and maximum-live-count proof | +| `Vector.data` | MPool; capacity from `Grow`, on/off-heap follows `v.offHeap` | owning `Vector.Free` | L | A only when off-heap | +| `Vector.area` | MPool; independent varlen payload capacity | owning `Vector.Free` | L | A only when off-heap | +| `Vector.nsp/gsp` bitmap data | Go `[]uint64`; `ceil(rows/64)*8`, retained by `Clear` | bitmap `Reset` from `Vector.Free` | L | move off-heap; blocks Vector-dependent activation | +| `FunctionResult.vec` data/area | off-heap Vector; rows and appended payload | executor `Free` | L | A | +| `FunctionResult.convenientParam` | Go slice; expression arity, not rows | executor `Free`/reuse | L | H after a proved arity bound | +| decimal parameter conversion | Go `[]T`; `rows*sizeof(T)` in `GenerateFunctionFixedTypeParameter` | evaluation wrapper/GC | L | move off-heap; blocks expression activation | +| IFF/CASE/COALESCE selection arrays | Go `[]bool`; one or two arrays of `rows`, retained by executor | executor `Free`/reuse | L | move off-heap; blocks expression activation | +| selected row IDs | Go `[]int64`; up to `rows`, retained by executor | executor `Free`/reuse | L | move off-heap; blocks expression activation | +| selected parameter/result vectors | off-heap Vector capacities | executor `Free` | L | A | +| hash-table initial cell block | off-heap `mpool.MakeSlice(..., true)`; 16 KiB int / 32 KiB string | hash map / `JoinMap.FreeMemory` | L | A in first activation | +| hash-table replacement/appended cell blocks | off-heap blocks, at most 4 MiB each; old+new overlap is physically visible | hash map / `JoinMap.FreeMemory` | L | A in first activation | +| hash-table `cells`/`newBlocks` descriptors | Go `[][]Cell`; 24 bytes per block header plus geometric resize backing arrays; GC is not treated as synchronous Free | hash map lifetime / GC | L | replace with an owning off-heap descriptor buffer and account its initial/replacement capacity; blocks first activation | +| hash-table `ResizePlan` and callback | fixed-size Go values/closures, one per table/resize | resize return / hash map Free | L | H; remove legacy reservation owner after cell activation | +| `GroupSels.{tmp,vals,offsets}` | on-heap `mpool.MakeSlice(..., false)`; O(build rows/groups) | builder or `JoinMap.FreeMemory` | L | switch off-heap; blocks auxiliary/copied-batch activation | +| copied build-batch vector buffers | MPool Vector data/area | builder or `JoinMap.FreeMemory` | L | A after per-buffer provenance | +| spill marshal/coalesce buffers | Go `bytes.Buffer`; O(serialized batch), retained by phase | spill cleanup | L | off-heap writer; blocks spill activation | +| spill hash values | Go `[]uint64`; `8*rows`, retained by phase | spill cleanup | L | move off-heap; blocks spill activation | +| spill row IDs | Go `[]int32`; `4*rows`, retained by phase | spill cleanup | L | move off-heap; blocks spill activation | +| spill counts/offsets/positions | Go `[]int32`; O(bucket count), bucket count finite | spill cleanup | L | H after bound is asserted | +| selected spill bucket vectors | off-heap Vector capacities | selected batch cleanup | L | A | +| BucketReader decoded vectors | MPool Vector data/area | `BucketReader.Close` | L | A after mode/provenance audit | +| runtime-filter serialized payload | Go buffer/message payload; O(filter rows) | message release | L | off-heap or PASS degradation; blocks runtime-filter activation | +| spill disk and FD | disk/FD ledgers | file removal/close | A | A | + +This is the known-site ledger, not yet the completion ledger for every later +owner. The hash cell/descriptor first-activation inventory is closed here. +PR 3 must generate and review the remaining built-in/function-specific `make`, +`append`, `bytes.Buffer`, builder, and codec sites before the corresponding +expression or spill closure can activate. No row named “other” or “unbounded +scratch” can declare closure. + +## 3. Decisions required before production integration + +### A. Allocation metadata representation + +PR 0 prototype results on linux/amd64, Go 1.26.4, i7-11700: + +| Representation | Map construction bytes/base entry | insert+delete median (five-run range) | +| --- | ---: | ---: | +| current 16-byte `memHdr` | 55.94 | 40.13 ns (39.97--40.40) | +| inline 24-byte header, all allocations | 83.97 | 40.21 ns (40.00--40.67) | +| side map, 1% accounted | 56.24 | 40.17 ns (40.16--40.43) | +| side map, 10% accounted | 58.29 | 40.81 ns (40.77--41.19) | +| side map, 100% accounted | 93.76 | 49.39 ns (48.39--49.89) | + +A bounded fixed registry prototype used 40.04 bytes per slot, including its +account state. Concurrent lookup across 1,024 generation-tagged handles +measured 0.4278 ns/op aggregate median (0.2876--0.4463) at `GOMAXPROCS=8`, +versus 5.554 ns/op (5.225--5.682) for the `sync.Map` comparison; both reported +zero allocations. The exact command, environment, and all five samples are in +[the raw benchmark record](evidence/26459_allocation_accounting_bench.txt). + +The same artifact now records real unaccounted baselines at `GOMAXPROCS=8`: + +| Existing path | Median | Five-run range | +| --- | ---: | ---: | +| sharded MPool alloc/free, 64 B | 245.6 ns | 242.4--249.6 ns | +| sharded MPool alloc/free, 4 KiB | 291.8 ns | 290.4--295.6 ns | +| sharded MPool alloc/free, 64 KiB | 999.2 ns | 991.2--1,003 ns | +| sharded MPool grow, 64 B to 64 KiB | 1,267 ns | 1,244--1,279 ns | +| `noLock` MPool alloc/free, 64 B | 222.7 ns | 222.5--226.0 ns | +| parallel sharded MPool alloc/free, 64 B | 209.4 ns | 176.6--213.2 ns | +| fixed Vector pre-extend/free, 8,192 rows | 1,042 ns | 1,028--1,089 ns | +| varlen Vector data+area pre-extend/free | 28,011 ns | 27,634--28,170 ns | +| fixed Vector Reset/capacity reuse | 2.810 ns | 2.691--2.852 ns | + +These measurements establish the pre-integration comparison baseline. The +prototype-map results are still not final accounted-MPool results, so they do +not by themselves freeze the representation. + +A test-only account-aware wrapper around the real MPool then measured: + +| Prototype path | Median | Five-run range | Delta from sharded baseline | +| --- | ---: | ---: | ---: | +| alloc/free, 64 B | 327.0 ns | 321.4--329.7 ns | +33.1% | +| alloc/free, 4 KiB | 367.6 ns | 361.9--381.9 ns | +26.0% | +| alloc/free, 64 KiB | 1,085 ns | 1,072--1,106 ns | +8.6% | +| grow, 64 B to 64 KiB | 1,400 ns | 1,390--1,418 ns | +10.5% | +| parallel alloc/free, 64 B | 227.2 ns | 203.1--233.5 ns | +8.5% | + +Every sample reported zero Go allocations. This wrapper intentionally takes a +second metadata-shard lock after MPool has already published its pointer +header. The 26--33% small-allocation cost rejects that shape for production; +PR 1 must publish the optional account side record under MPool's existing +pointer-shard transaction and re-run the comparison. The wrapper remains a +conservative upper bound and validates real allocation/growth rollback. + +Provisional choice: + +- keep `memHdr` at 16 bytes and replace the final `offHeap` byte with a + three-state allocation kind: on-heap, unaccounted off-heap, or accounted + off-heap; +- store a compact 16-byte `pointer -> account pointer + owner/site` lease only + for accounted allocations in the same pointer shard, or in the same + pool-local metadata store for `noLock` pools. The direct account pointer + keeps the original generation alive through late physical `Free` without a + process-global registry lookup; +- publish/remove the pointer header and optional account ID in one metadata + transaction under the same lock; a metadata failure rolls back both before + allocation publication; +- use a finite registry of reusable slots and encode `(slot, generation)` in + the handle; reuse a slot only after attempt-owned seal and exact zero, and + retire it rather than allowing the generation counter to wrap; +- reserve one finite CN-local allocation-metadata slot before publishing an + accounted allocation or replacement and return it on physical Free; +- size the registry and side-record backing stores from explicit CN headroom, + rather than allowing maps to grow without a hard count bound. + +The first hash-table activation fixes the initial sizing policy: + +- reserve 131,072 generation slots. The production registry plus a live + 64-byte account measures 80.09 bytes/slot; budgeting 128 bytes/slot reserves + 16 MiB. This default exceeds the + frontend `max_connections` system variable's declared upper bound of 100,000 + and leaves 31,072 slots for remote/internal attempts. Only an attempt whose + physical plan contains an activated owner opens a slot. The slot limit is + itself the hard supported activated-attempt concurrency; a deployment that + intends to support more must raise it and pass the same headroom check before + enabling the owner; +- every integer/string cell block is at least 16 KiB. After the outer + descriptor is made an owning accounted off-heap allocation, every live + table has at most one live descriptor buffer per live table version. + Therefore the first activation uses + `3 * ceil(MPoolGlobalCap / 16 KiB)` allocation-metadata slots. Two units + cover every live cell block plus its table's published descriptor; the third + covers one unpublished replacement descriptor per live table before the + matching cell allocation either publishes or rolls back; +- the production base pointer map plus a 16-byte all-accounted lease map + measures 111.87 bytes/entry, of which 55.93 bytes/entry is incremental over + the existing pointer map. Budget 128 bytes/allocation-metadata slot to cover + sparse shards and Go-map growth overlap, and 128 bytes/generation slot for + the final account fields. At MPool's 1 GiB minimum global cap this is 24 MiB + plus 16 MiB; at + larger caps the allocation component is 2.34375% of the MPool cap and the + fixed registry fraction decreases. PR 1 must measure construction, sparse + occupancy, and grow/evacuate high water; exceeding these conservative + constants blocks merge rather than silently consuming payload headroom; +- startup must reserve + `allocationSlots * 128 + generationSlots * 128` bytes outside the MPool + payload cap. If the host/container limit cannot supply it, the activated + owner is refused at startup rather than running with an unproved headroom + assumption. + +Later activations with smaller allocations must derive a new simultaneous +allocation bound and resize this headroom before they can become `A`; they +cannot inherit the hash-cell formula merely because the generic API exists. + +This choice is not frozen until real MPool alloc/grow/free, cross-pool Free, +deleted-pool fallback, real accounted ratios, and P50/P99 concurrent latency +match the prototype result. + +Metadata is `H`, not silently included in payload capacity. Safety comes from +finite generation and allocation-metadata slot limits; slot exhaustion is a +typed exact-pressure result. Each activation PR must also show that its +supported simultaneous allocations and generations fit the configured limits, +using measured pointer/side/registry bytes per entry and resulting aggregate +CN headroom. This prevents a safe-but-impractical false metadata-pressure +regression. + +### B. Account-aware API shape + +The provisional API rules are: + +- only an explicit first off-heap allocation accepts an account, owner, and + site; +- ordinary `Grow` inherits account provenance from allocation metadata and + takes no replacement account argument; +- account-A memory cannot grow under account B; +- an accounted on-heap allocation is rejected; +- unaccounted-to-accounted conversion allocates a new destination; +- ordinary unaccounted callers keep current behavior; +- helper delegation cannot silently drop the account. + +One `AllocationCapacity` rule must cover initial allocation, growth, runtime +rounding, and `CapLimit-kMemHdrSz`. `recordPtrHdr` failure, allocator panic, +cross-pool Free, deleted-pool fallback, and pool teardown are explicit +transaction branches. + +### C. Generation owner and terminal snapshot + +One execution attempt of `Compile` on each CN owns the generation. The +statement `ResourceRoot` aggregates its immutable terminal snapshot, but does +not own allocation release. HashBuild Reset is also not the generation owner. + +The attempt opens before `prePipelineInitializer`/operator Prepare and owns: + +```text +all local scopes, remote notifiers, and message consumers quiescent + -> close and drain that attempt's MessageBoard + -> seal new admission + -> release remaining live leases + -> used=0 + -> export one immutable snapshot + -> remove registry entry +``` + +`Scope.Run` defers pipeline cleanup, and `Scope.MergeRun` joins pre-scopes and +remote notifier goroutines before returning. Therefore the local hook belongs +in a deferred attempt finalizer around `Compile.runOnce`, after its result and +before retry transition or attempt publication. A retry finalizes the failed +attempt before `buildRetryCompile` opens the next generation. + +The remote hook belongs after `Scope.MergeRun` and in the existing +`runCompile.clear` terminal defer, which releases operators, resets the +MessageBoard, and snapshots the remote MPool before replying. PR 4 must add an +explicit MessageBoard close-and-drain operation: ordinary multi-CN `Reset` +only removes the board from `StmtIDToBoard` because producers may still access +it, whereas the attempt hook runs after the existing sender/receiver cleanup +barriers have proved quiescence. + +The same deferred finalizer covers success, error, panic, cancellation, retry, +broadcast, remote execution, and a JoinMap freed after producer Reset. +`SetStmtProfile` turnover, frontend `StatementInfo.EndStatement`, and +`HashBuildBudgetGeneration.Close` are observability or operator boundaries, +not acceptable release substitutes. + +If terminal cleanup ends nonzero, the attempt coordinator exports one immutable +invariant-failure snapshot and retains a release-capable tombstone. That CN +admits no new accounted generation until all such tombstones drain to zero, so +registry growth is bounded by generations already active at detection. A +deadline escalates owner/site diagnostics and allows controlled CN restart; it +never deletes live provenance. + +Generation open and suspension publication use one CN-local linearization +gate. An open that linearizes after suspension cannot publish. + +PR 4 derives this terminal matrix: + +| Attempt path | Required terminal ordering | Oracle | +| --- | --- | --- | +| local success | scopes join -> board close/drain -> seal -> zero -> publish | one valid snapshot, no queued message | +| local error/cancel/panic | cancellation -> every started scope cleanup/join -> board drain -> seal | one terminal snapshot; no goroutine or allocation survives | +| failure before `runOnce` | opened generation -> initializer rollback -> board drain -> seal | zero or one named invariant failure, never an abandoned open slot | +| retry | attempt N fully finalizes -> attempt N+1 opens | old handles are stale; no cross-attempt publication | +| remote execution | remote `MergeRun` joins -> `runCompile.clear`/board drain -> snapshot -> response | parent receives one immutable child snapshot | +| broadcast/late JoinMap Free | producer Reset -> every consumer cleanup -> queued refs drain | physical final Free releases the original generation exactly once | +| prepared reuse | attempt finalizes and replaces its board -> cached pipeline Reset -> next attempt opens | no retained accounted capacity crosses statement generations | +| nonzero terminal | seal -> failure snapshot -> tombstone/suspend -> late Free | no new open until every tombstone reaches zero | + +### D. Go-heap boundary + +`MPool.Alloc(..., false)` records requested bytes but `Free` does not reclaim +them synchronously. Therefore: + +- all data/row/payload-scaled controlled allocations move off-heap before + activation; +- small Go metadata may be `H` only with a static bound and separate CN + headroom; +- no data-scaled Go slice may be relabeled `S` to bypass migration. + +### E. Pressure and operation rollback + +Capacity pressure, sealed generation, account mismatch, allocator-size limit, +and invariant corruption are distinct typed results. Only capacity pressure is +recoverable. + +Each retryable owner records an operation checkpoint and cleanup/restart rule. +For example, `Vector.PreExtendWithArea` may grow data and then fail area growth; +that retained growth is valid accounting state but not proof that re-running +the logical operation is idempotent. Before a shared controller exists, an +exact rejection is a controlled terminal pressure error. + +The PR 0 reference model validates the provisional shared rule: + +- replacements remain private while the old allocation stays published and + charged; +- all new allocations and replacements commit as one logical operation; +- cancellation or later allocation failure frees private allocations and + restores the checkpoint before retry; +- a second attempt cannot begin while the failed operation remains active; +- retry may reduce the requested capacity, but cannot duplicate publication; +- an owner that cannot preserve or reconstruct the checkpoint is not + retryable and returns the typed pressure error after cleanup. + +For each spill owner choose an already allocated reusable buffer, a finite +progress sub-cap, or a smaller chunk. Normal work and progress allocations stay +under the same total query/CN cap; no uncharged emergency scratch is allowed. + +## 4. Pull request sequence + +### PR 0: model and measured design decisions + +Scope: + +- close decisions A--E with prototypes and benchmarks; +- finalize bounded owner/site enums; +- implement a test-only reference state machine; +- record MPool/vector allocation baselines; +- complete the first activation owner inventory and record the generation + method for later expression/spill inventories without changing production + behavior. + +Current evidence: + +- the reproducible test-only artifact is + `experiment/26459-allocation-accounting-validation` at `cde44cd099`: + `pkg/common/mpool/allocation_account_validation_test.go` and + `pkg/common/mpool/allocation_account_benchmark_test.go` and + `pkg/container/vector/allocation_account_validation_test.go`; +- the test-only model passes alloc, within-capacity reuse, old+new growth, + injected unpublished failures, views, deep copy, Reset, multi-allocation + checkpoint/commit/rollback, cancellation before and after allocation, + smaller retry without duplicate publication, bounded generation and + allocation-metadata slots, stale slot generations, exact metadata overlap + on replacement, generation-counter exhaustion without wrap, accounted + on-heap rejection, sealed-error precedence, handoff, cross-pool Free, + sealed-vs-capacity errors, zero finalization, nonzero + tombstone/suspension/drain, normal and `noLock` pool teardown, + open-vs-suspend linearization, stale generation, and 20,000 deterministic + randomized operations; +- the metadata and contention microbenchmarks in section 3 passed five runs; +- real unaccounted MPool alloc/free/grow and Vector allocate/reuse baselines in + section 3 passed five runs; +- the account-aware real-MPool wrapper passes allocation, reuse, old+new growth, + rollback, and final-zero validation; its five-run result rejects a second + side-metadata lock for production; +- at `GOMAXPROCS=8`, a serialized mutex acquire/release prototype measured + 80.73 ns/op median versus 21.11 ns/op for a two-operation atomic prototype, + so real aggregate-account contention remains a mandatory design benchmark; +- production behavior is unchanged. + +Gate: + +- the existing MPool/Vector baseline and rejected separate-lock shape have + reproducible five-run measurements; the selected same-shard transaction is + an explicit PR 1 merge gate; +- the model covers alloc, grow, failure, view/copy, Reset, operation rollback, + cancellation/retry, finite metadata slots, handoff, cross-pool Free, seal, + and stale generation; +- per-entry metadata and maximum simultaneous allocation/generation counts + prove finite aggregate CN headroom; +- the site ledger is complete for cell and descriptor initial allocation, + replacement, segmented growth, rollback, and terminal Free of the hash-table + first activation; +- generation owner, Go-heap classification, typed errors, and retry checkpoint + decisions are closed; +- production owners remain `L`. + +PR 0 design gates are closed. The metadata representation remains provisional +until PR 1's integrated benchmark passes, and no production owner may switch +from `L` to `A` before PRs 1--4 close. The separate-lock prototype is a +recorded rejected design, not an implementation candidate. + +### PR 1: generic account and MPool allocation transaction + +Scope: + +- low-level account contract below SQL/process; +- compact account-ID registry and dormant `HashBuildBudgetGeneration` adapter; +- finite generation/allocation-metadata slot limits and their typed pressure + results; +- account-aware alloc, Grow/Grow2, Free, and immutable snapshots. + +Required behavior: + +- reject accounted on-heap allocation; +- reserve the complete new capacity before allocation; +- for growth keep old and complete new capacity live until publication; +- roll back on admission, MPool/global-cap, metadata, or allocation failure; +- use one allocation-capacity rule for initial and growth boundaries; +- release through normal and cross-pool physical Free, deleted-owner-pool + fallback, and `noLock` teardown that physically deallocates; +- retain metadata and charge when normal-pool teardown only unregisters the + pool, and report live accounted allocations there as an invariant; +- reject account mismatch and stale handles. + +Gate: + +- exact/one-byte-short, allocator rounding, `CapLimit-kMemHdrSz`, and + `GrowCapacity` boundaries; +- old+new overlap; +- injected failure or panic at every unpublished step, including metadata; +- atomic header/account-ID publication and removal for sharded and `noLock` + pool metadata; +- normal-pool unregister plus late Free, `noLock` physical teardown, and no + premature release in either case; +- concurrent acquire/release, double Free, seal, and final zero; +- measured unaccounted/accounted alloc/free/grow overhead and concurrent + generation P50/P99 latency; +- no production owner selects an account. + +The current PR 1 candidate is +`feature/26459-allocation-account` at generic commit `766e1501c3` and +HashBuild-adapter commit `0655af4443`. It remains dormant. Its same-lock +pointer/lease transaction, finite registry, stale-handle checks, old+new +growth, deleted-pool/noLock lifetime rules, and tokenless HashBuild adapter are +implemented. Returned-error and panic rollback is injected after account, +metadata, global stats, pool stats, physical allocation, and header +publication for both sharded and `noLock` metadata. Normal package tests, +package vet, full package race tests, and the focused lifecycle/rollback race +matrix at 100 repetitions pass. Representation, latency, and integrated +benchmarks are recorded in the evidence artifact. No production owner selects +an account, and no legacy hard charge has been removed. + +### PR 2: dormant Vector and Batch propagation + +Scope: + +- optional account selection for owning off-heap Vector buffers; +- data and area growth; +- Batch Clone/Dup/Union destination propagation. + +Required behavior: + +- Reset retains charge; Free releases it; +- within-capacity append performs no account operation; +- aliases/views/const/shared area do not create another charge; +- deep copies use the destination account; +- on-heap null/group bitmaps remain explicit ledger blockers, not silently + included in the Vector charge; +- HashBuild production remains legacy until a later owner migration. + +Gate: + +- randomized fixed/varlen append; +- Reset/reuse/Free, views, partial selection, copy rollback, cross-pool Free; +- package race tests and vector benchmarks; +- Vector/Batch ledger rows become `D`. + +### PR 3: allocation-site closure and dormant propagation + +Scope: + +- complete the generated expression/built-in and spill allocation-site ledger; +- propagate dormant accounts through FunctionResult, expression result, + selected result, decoded Vector, and off-heap scratch constructors; +- replace data-scaled Go vector null/group bitmaps, + selection/conversion/hash/row-ID/serialization buffers with off-heap owners + or direct output; +- do not enable production accounting or remove a legacy hard gate. + +Gate: + +- every reachable `make`, capacity-growing `append`, `bytes.Buffer`, MPool, and + nested executor result has allocator, bound, terminal owner, and test; +- fixed/varlen/const/null, CAST, CONCAT, CASE, nested and selected paths close; +- repeated Eval/Reset/Free and construction failure reach the same terminal + owners; +- all migrated rows become `D`; production rows remain `L`. + +### PR 4: statement lifecycle and minimum pressure foundation + +Scope: + +- add the attempt-owned post-pipeline/MessageBoard-close seal/finalize hook; +- export one immutable valid or invariant-failure generation snapshot; +- retain release-capable tombstones and suspend new accounted generations + after nonzero terminal cleanup; +- introduce non-overlapping capacity, sealed, mismatch, allocator-limit, and + invariant error reasons; +- define operation checkpoint/rollback helpers needed by later retry; +- keep owner accounting dormant. + +Gate: + +- success, execution error, cancellation, retry, local/remote scope, broadcast, + and message-board teardown each seal exactly once; +- JoinMap/spill payload released after producer Reset still releases the + original generation; +- zero finalization exports once and removes the registry entry; +- nonzero terminal cleanup exports one failure snapshot, retains a + release-capable tombstone, stops new accounted generations on that CN, and + removes the tombstone only after late Free reaches zero; +- concurrent terminal failures are bounded by generations active at first + detection, and stale IDs cannot resolve after removal; +- a race test proves every successfully published generation linearized before + suspension and every later open is rejected; +- closed generation never enters reclaim/spill/retry; +- no production legacy charge is removed. + +### PR 5: hash-table cell-block activation + +Scope: + +- activate only the integer/string hash-table cell blocks; +- replace the Go `[][]Cell` outer backing store with an owning off-heap + descriptor buffer whose initial and replacement capacities use the same + account; +- attach the account to initial allocation, full replacement, segmented + appended blocks, and consumer-side empty-map growth; +- remove only the matching `hashMapReservationOwner`/`ResizeReservation` + budget charge; +- retain legacy copied-batch, expression, auxiliary, and spill charges. + +Gate: + +- int/string initial, no-op, full replacement, segmented reuse, stale plan, + injected allocation failure, and terminal Free are covered; +- old+new replacement and old+appended-block peaks match live cell allocation + plus descriptor allocation capacity; +- descriptor replacement rollback is atomic with cell-block publication, and + no data-scaled Go backing array remains; +- the first-activation metadata-slot formula holds at exact and one-slot-short + boundaries; +- consumer growth after HashBuild handoff keeps original provenance; +- generation final snapshot reaches zero; +- #25782 high-cardinality no-OOM regression and hash-table performance pass. + +### PR 6: copied batches and JoinMap activation + +Scope: + +- copied build-batch destinations; +- HashmapBuilder-to-JoinMap ownership handoff. + +Required behavior: + +- replace projected/reconciled batch tokens with physical leases; +- preserve provenance through producer Reset and broadcast consumers; +- let physical Free replace budget-only `SetMemoryRelease` callbacks. + +Gate: + +- empty/single/large build and both resize modes; +- failed publish, cancellation, multiple consumers, duplicate cleanup; +- old-generation allocation freed after a new generation opens; +- #25782 high-cardinality and #26413 external-table self-join regressions pass + before activation merges; +- copied-batch/JoinMap site rows become `A` or proved `H`. + +### PR 7: expression owner activation + +Scope: + +- activate the complete HashBuild-owned expression site closure; +- remove exactly the corresponding + `expressionVectorPeak`/`expressionTypePeak` hard charge; +- return controlled terminal pressure until a proved retry checkpoint supports + a smaller-batch retry; +- do not stack exact leases on #26455's operator-held lifetime charge. + +Gate: + +- every activated expression site is `A` or proved `H`; no data-scaled Go + allocation remains; +- one-byte-short and real single-value-over-cap diagnostics name actual + capacity; +- partial growth/evaluation failure publishes no duplicate rows; +- generation final snapshot reaches zero; +- #26454 workload and expression performance regressions pass. + +### PR 8 family: spill and runtime-filter closures + +Split into independently safe owner closures: + +1. decoded batches and retained reader reuse; +2. scatter/hash/row-ID/codec/coalesce buffers and forward-progress policy; +3. runtime-filter payload transfer and PASS degradation. + +Each sub-PR removes only its matching hard reservation. Required gates include +first spill, recursive spill, skew, empty bucket, EOF, failure/cancel at every +I/O/publication edge, minimum progress over cap, message destruction, and final +memory/disk/FD zero. The scatter/codec closure must pass #26174 fulltext INSERT; +the decoded/retained-reader closure must pass #26192 LOAD DATA. A runtime-filter +closure adds its own build/probe and PASS-degradation workload before merging. + +### PR 9: unified join pressure controller and remaining legacy deletion + +Implement across HashBuild, HashJoin, DedupJoin, and RightDedupJoin: + +```text +exact capacity rejection + -> rollback to operation checkpoint + -> reclaim -> retry + -> spill/re-spill -> retry + -> reduce batch -> retry + -> degrade optional owner + -> controlled minimum-unit error +``` + +Gate: + +- retry only after usage decreases, spill advances, input shrinks, or optional + work is disabled; +- partially grown retained buffers and output publication are idempotent; +- no infinite pressure loop; +- sealed/lifecycle errors never retry; +- cancellation/downstream failure is covered in every state; +- remaining multiplier hard gates and duplicate token owners are deleted; +- every site row is `A`, justified `H`, `S`, or `R`. + +### PR 10: workload, performance, and cleanup + +This re-runs all incident workloads together and supplies long-run and +comparative confirmation; it is not the first incident-level validation of an +earlier activation. + +Required workloads: + +- #25782 high-cardinality case; +- #26174 fulltext INSERT; +- #26192 LOAD DATA; +- #26413 Hive external-table self-join; +- #26454 string-expression join; +- TPCH 100G non-spill and TPCH 1T spill. + +Required proof: + +- no CN OOM/restart and no cap increase/spill-disable workaround; +- account never exceeds cap and every generation returns to zero; +- only a real minimum allocation can produce terminal pressure; +- diagnostics name owner/site and attempted response; +- resident and spill performance are compared separately with profiles; +- concurrent-generation P50/P99, release storm, and high-frequency TP + allocation results meet the recorded gate; +- temporary migration helpers are removed. + +## 5. Verification and conservation model + +Every semantic PR runs build, vet, complete package tests, focused adaptive race +stress, and dependent package tests for its closure. Likely packages are: + +```text +pkg/common/mpool +pkg/container/vector +pkg/container/batch +pkg/sql/colexec +pkg/sql/colexec/{hashbuild,hashjoin,dedupjoin,rightdedupjoin,spillutil} +pkg/vm/{message,process} +``` + +Before direct tests that can reach usearch, build `thirdparties` and use the +repository CGO include/library/rpath environment. Compilation or one successful +SQL run is not completion evidence. + +The completed PR 0 reference model must track allocation ID, account ID, +allocator mode, capacity, pool, logical owner, and +unpublished/live/freed/tombstone state. Its required generated operations +include allocation, within-capacity reuse, replacement growth, injected +failure, Free, cross-pool Free, view/copy, handoff, Reset, seal, finalization, +stale ID, operation checkpoint, cancellation, and retry. + +The current artifact generates allocation, growth, handoff, and Free, and +separately tests failure injection, cross-pool Free, seal, zero/nonzero +terminal paths, stale generation-tagged slots, bounded registry and allocation +metadata, teardown, open/suspend linearization, view/copy, Reset, +multi-allocation checkpoints, cancellation, rollback, and smaller retry +without duplicate publication. + +After every operation: + +```text +account.used + = sum(live accounted allocation capacities) + + sum(live named scratch) +``` + +Failed unpublished work leaves allocator/account state unchanged; no allocation +releases twice; sealed generations admit nothing new; final cleanup reaches +zero. Accounted on-heap allocation is rejected. A logical retry may retain +already published reusable capacity but cannot retain partially published +rows. Deterministic CI seeds print on failure; longer randomized runs may run +nightly. + +## 6. Resource Accounting integration + +Admission and SQL Resource Accounting share allocator facts, not mutable +control state. Each query-CN generation has a stable identity and its +`Compile` attempt coordinator exports exactly one immutable controlled-domain +snapshot: + +- cap, exact peak, and final live bytes; +- alloc/grow/free and real pressure counts; +- reclaim/spill/batch-reduction/degrade responses; +- invariant quality. + +Initial integration keeps `statement_info.stats[2]` unchanged and exposes the +controlled-domain snapshot in physical-plan/operator diagnostics. Resource +Accounting may mark a missing/inconsistent snapshot and aggregate terminal +facts, but its summaries never feed hard admission. Multiple operators may +reference one generation; their diagnostics must not be summed as separate +physical domains. The controlled-domain and MPool-domain peaks are not declared +equal until allocator mode and owner coverage match. + +No per-allocation log or Prometheus series is added. Owner/site values are +bounded enums and counters are emitted at generation/operator completion or +terminal pressure. + +## 7. Rollout, rollback, and completion + +Rollout: + +- rebase every PR on current `main`; +- keep generic primitives dormant until one complete owner closure migrates; +- enable exact accounting and remove the same legacy charge atomically; +- update the issue ledger only after evidence exists. + +Rollback reverts an owner's exact enablement, legacy removal, pressure response, +and tests together. Never roll back only the accounting or pressure half. If a +missing owner is found, complete the closure before rollout or revert that +owner; do not add another permanent estimator multiplier. + +Done means: + +- every site-ledger row is `A`, justified `H`, `S`, or `R`; +- every `H` row has a per-entry cost, maximum-live-count proof, and aggregate + CN headroom; +- allocation/growth failure is atomic; +- attempt-owned seal/finalize and exactly-once snapshot export are proven; +- nonzero terminal generations retain release provenance without unbounded + registry growth; +- Reset, Free, broadcast, and cross-generation handoff preserve provenance; +- data-scaled controlled Go allocations have moved off-heap; +- operation-level retry checkpoints prevent partial output replay; +- bounded forward progress is accounted; +- predictions cannot terminally reject SQL; +- no data-scaled HashBuild-owned allocation remains outside coverage; +- the original no-OOM case and all listed false-budget regressions pass; +- performance gates pass with evidence; +- Resource Accounting receives immutable diagnostics without entering the + admission decision; +- legacy hard estimators and duplicate token owners are deleted. diff --git a/docs/design/evidence/26459_allocation_accounting_bench.txt b/docs/design/evidence/26459_allocation_accounting_bench.txt new file mode 100644 index 0000000000000..4ed578af4ed7d --- /dev/null +++ b/docs/design/evidence/26459_allocation_accounting_bench.txt @@ -0,0 +1,156 @@ +Artifact: + +branch: experiment/26459-allocation-accounting-validation +commit: cde44cd099 (rebased validation artifact) +go: go version go1.26.4 linux/amd64 + +Setup: + +make thirdparties +make cgo +export CGO_ENABLED=1 + +Metadata command: + +go test -v ./pkg/common/mpool -run '^TestAllocationMetadataRepresentationCosts$' -count=1 + +Metadata results: + +baseline: 13985992 bytes total, 55.94 bytes/base-entry +inline-all: 20992896 bytes total, 83.97 bytes/base-entry +side-1pct: 14059896 bytes total, 56.24 bytes/base-entry +side-10pct: 14571520 bytes total, 58.29 bytes/base-entry +side-all: 23438864 bytes total, 93.76 bytes/base-entry +fixed registry: 10010736 bytes total, 40.04 bytes/slot + +Benchmark command: + +GOMAXPROCS=8 CGO_ENABLED=1 .agents/skills/mo-dev/scripts/mo-cgo-test -run '^$' -bench='BenchmarkAllocation(MetadataInsertFree|ConcurrentRegistryResolve|FixedRegistryResolve|AccountAcquireRelease)$' -benchmem -benchtime=500ms -count=5 ./pkg/common/mpool + +Environment: + +goos: linux +goarch: amd64 +cpu: 11th Gen Intel(R) Core(TM) i7-11700 @ 2.50GHz +GOMAXPROCS: 8 +CPU pinning: none + +Results: + +BenchmarkAllocationMetadataInsertFree/baseline-8 40.13 39.99 40.40 40.37 39.97 ns/op +BenchmarkAllocationMetadataInsertFree/inline-all-8 40.67 40.36 40.21 40.00 40.17 ns/op +BenchmarkAllocationMetadataInsertFree/side-1pct-8 40.17 40.16 40.42 40.17 40.43 ns/op +BenchmarkAllocationMetadataInsertFree/side-10pct-8 40.81 40.97 40.77 40.79 41.19 ns/op +BenchmarkAllocationMetadataInsertFree/side-all-8 49.89 48.39 48.84 49.39 49.42 ns/op +BenchmarkAllocationConcurrentRegistryResolve-8 5.682 5.564 5.554 5.287 5.225 ns/op +BenchmarkAllocationFixedRegistryResolve-8 0.4278 0.4461 0.4463 0.2876 0.3265 ns/op +BenchmarkAllocationAccountAcquireRelease/locked-8 73.71 81.29 72.94 88.61 80.73 ns/op +BenchmarkAllocationAccountAcquireRelease/atomic-8 21.68 21.31 20.77 20.88 21.11 ns/op + +Every sample reported 0 B/op and 0 allocs/op. + +Real MPool and Vector baseline command: + +GOMAXPROCS=8 CGO_ENABLED=1 .agents/skills/mo-dev/scripts/mo-cgo-test -run '^$' -bench '^BenchmarkAllocationAccounting(MPool|Vector)Baseline$' -benchmem -benchtime=200ms -count=5 ./pkg/common/mpool ./pkg/container/vector + +Real baseline results: + +BenchmarkAllocationAccountingMPoolBaseline/alloc-free/sharded/64-8 249.6 244.3 248.0 242.4 245.6 ns/op +BenchmarkAllocationAccountingMPoolBaseline/alloc-free/sharded/4096-8 291.8 291.0 292.5 290.4 295.6 ns/op +BenchmarkAllocationAccountingMPoolBaseline/alloc-free/sharded/65536-8 999.2 1003 991.2 1002 993.9 ns/op +BenchmarkAllocationAccountingMPoolBaseline/grow-replacement/sharded-8 1244 1267 1279 1254 1277 ns/op +BenchmarkAllocationAccountingMPoolBaseline/alloc-free/no-lock/64-8 222.7 226.0 223.2 222.5 222.6 ns/op +BenchmarkAllocationAccountingMPoolBaseline/alloc-free/no-lock/4096-8 280.5 273.3 276.5 271.6 276.5 ns/op +BenchmarkAllocationAccountingMPoolBaseline/alloc-free/no-lock/65536-8 984.2 979.6 979.9 994.1 980.9 ns/op +BenchmarkAllocationAccountingMPoolBaseline/grow-replacement/no-lock-8 1219 1199 1198 1248 1209 ns/op +BenchmarkAllocationAccountingMPoolBaseline/parallel-alloc-free/sharded/64-8 176.6 213.2 211.1 180.5 209.4 ns/op +BenchmarkAllocationAccountingVectorBaseline/fixed-preextend-free-8 1089 1028 1042 1043 1037 ns/op +BenchmarkAllocationAccountingVectorBaseline/varlen-preextend-free-8 28130 28011 27634 28170 27778 ns/op +BenchmarkAllocationAccountingVectorBaseline/fixed-reset-reuse-8 2.810 2.691 2.852 2.701 2.835 ns/op + +Every MPool sample reported 0 B/op and 0 allocs/op. Fixed Vector +pre-extend/free and Reset/reuse also reported 0 B/op and 0 allocs/op. Varlen +Vector pre-extend/free reported 48 B/op and 2 allocs/op. + +Account-aware real-MPool wrapper command: + +GOMAXPROCS=8 CGO_ENABLED=1 .agents/skills/mo-dev/scripts/mo-cgo-test -run '^$' -bench '^BenchmarkAllocationAccountedMPoolPrototype$' -benchmem -benchtime=200ms -count=5 ./pkg/common/mpool + +Account-aware wrapper results: + +BenchmarkAllocationAccountedMPoolPrototype/alloc-free/64-8 329.7 328.4 323.7 327.0 321.4 ns/op +BenchmarkAllocationAccountedMPoolPrototype/alloc-free/4096-8 367.6 361.9 381.9 366.5 368.4 ns/op +BenchmarkAllocationAccountedMPoolPrototype/alloc-free/65536-8 1106 1072 1085 1080 1091 ns/op +BenchmarkAllocationAccountedMPoolPrototype/grow-replacement-8 1418 1404 1390 1400 1393 ns/op +BenchmarkAllocationAccountedMPoolPrototype/parallel-alloc-free/64-8 203.1 227.2 205.6 233.5 230.2 ns/op + +Every sample reported 0 B/op and 0 allocs/op. This prototype takes a second +side-metadata lock and is intentionally an upper-bound comparison, not the +accepted production transaction shape. + +Final validation commands: + +go test -count=1 ./pkg/common/mpool ./pkg/container/vector +go vet ./pkg/common/mpool ./pkg/container/vector +go test -race -count=100 -run '^(TestAllocationAccountedMPoolPrototype|TestAllocationAccountingReferenceModel|TestAllocationAccountingViewCopyReset|TestAllocationAccountingMetadataSlotBounds|TestAllocationAccountingOperationRollbackRetry|TestAllocationAccountingTerminalTombstone|TestAllocationAccountingOpenSuspendLinearization|TestAllocationAccountingPoolTeardown)$' ./pkg/common/mpool +go test -race -count=23 -run '^TestAllocationAccountingReferenceModelRandomized$' ./pkg/common/mpool +go test -race -count=1 ./pkg/common/mpool ./pkg/container/vector + +Production PR 1 representation measurement (250,000 entries): + +branch: feature/26459-allocation-account +generic MPool commit: 766e1501c3 +HashBuild adapter commit: 0655af4443 + +base pointer map: 13,985,992 bytes total, 55.94 bytes/slot +base + 16-byte lease map: 27,966,424 bytes total, 111.87 bytes/slot +fixed registry backing: 4,022,272 bytes total, 16.09 bytes/slot +registry + live 64-byte accounts: 20,022,400 bytes total, 80.09 bytes/slot + +The conservative constants remain 128 bytes/allocation slot and are raised to +128 bytes/generation slot. The first-activation sizing test therefore checks +that a 1 GiB MPool cap yields 196,608 allocation-metadata slots and 40 MiB +total conservative metadata headroom, including 131,072 generation slots. + +Production PR 1 benchmark command: + +GOMAXPROCS=8 go test -run '^$' -bench '^BenchmarkMPoolAccountedAllocation$' \ + -benchmem -benchtime=500ms -count=5 ./pkg/common/mpool + +Clean-main medians from the paired baseline worktree: + +alloc/free 64: 244.6 ns/op +alloc/free 4096: 297.0 ns/op +alloc/free 16384: 333.4 ns/op +alloc/free 65536: 1023 ns/op +grow replacement: 1274 ns/op + +Production PR 1 unaccounted medians: + +alloc/free 64: 252.7 ns/op (+3.3%) +alloc/free 4096: 303.2 ns/op (+2.1%) +alloc/free 16384: 339.2 ns/op (+1.7%) +alloc/free 65536: 1034 ns/op (+1.1%) +grow replacement: 1306 ns/op (+2.5%) + +Production PR 1 accounted medians: + +alloc/free 64: 352.4 ns/op (+44.1% versus clean main) +alloc/free 4096: 400.4 ns/op (+34.8% versus clean main) +alloc/free 16384: 435.1 ns/op (+30.5% versus clean main) +alloc/free 65536: 1115 ns/op (+9.0% versus clean main) +grow replacement: 1492 ns/op (+17.1% versus clean main) +parallel accounted alloc/free 65536: 300.9 ns/op + +The dormant HashBuild controller adapter, including query/CN policy locking +and metrics, measured 1260 ns/op for 65536-byte alloc/free. Every production +PR 1 sample reported 0 B/op and 0 allocs/op. The first activated hash cell +allocation is at least 16 KiB; activation PRs still require workload-level +resident/spill validation because the small-allocation microbenchmark is not a +per-row path. + +An eight-generation, GOMAXPROCS=8 acquire/release latency harness sampled +200,000 operations per run. Across five runs P50 was 48 ns in every run; P99 +was 100, 98, 97, 68, and 73 ns (97 ns median). The measurement includes the +`time.Now`/`time.Since` sampling cost and is therefore a conservative +end-to-end latency observation, not a cycle-level atomic benchmark. diff --git a/docs/rfcs/00000000_allocation_accounted_memory_admission.md b/docs/rfcs/00000000_allocation_accounted_memory_admission.md new file mode 100644 index 0000000000000..4a2be8adabd2f --- /dev/null +++ b/docs/rfcs/00000000_allocation_accounted_memory_admission.md @@ -0,0 +1,949 @@ +- Status: draft +- Start Date: 2026-07-30 +- Authors: aptend +- Implementation PR: TBD +- Issue for this RFC: + [#26459](https://github.com/matrixorigin/matrixone/issues/26459) +- Implementation plan: + [Allocation-Accounted Memory Admission Implementation Plan](../design/allocation_accounted_memory_admission_impl.md) + +# Allocation-Accounted Memory Admission for Spillable SQL Execution + +## Summary + +MatrixOne currently protects HashBuild with finite query and CN budgets and +bounded spill. That protection is necessary: before it existed, a large join +could continue allocating until the CN was killed by OOM. The remaining +problem is that several hard admission decisions are based on predictions of +an operation's future memory rather than the capacity of memory that is +actually allocated and retained. + +Predictions such as SQL type maximums, payload multipliers, recursive +expression-tree peaks, or logical batch sizes are useful for deciding to spill +early. They are not reliable enough to decide that a valid statement cannot +run. The same prediction can over-count aliases and retained buffers, or +under-count an allocate-copy-free overlap. + +This RFC replaces estimator-driven hard rejection with allocation-accounted +ownership: + +```text +hard admission = capacity of real live allocations + + named, bounded non-allocator scratch + +prediction = scheduling hint only +``` + +An account-aware allocation reserves the exact physical capacity before it is +allocated. The resulting charge follows the allocation across reuse and owner +handoff, and is released by the same physical `Free`. Growth accounts for the +real replacement overlap by keeping the old allocation charged while +admitting the complete replacement capacity. + +When a real allocation cannot be admitted, execution treats it as typed memory +pressure: reclaim retained state, spill, retry, reduce the processing batch +where possible, degrade optional structures, and only then return a controlled +error for a minimum indivisible allocation that cannot fit. + +The first consumer is HashBuild and the joins that share its spill lifecycle. +The accounting primitive is deliberately defined below the SQL operator layer +so other spillable operators can adopt the same model later. + +## Motivation + +### Confirmed failures + +The incidents below include both false rejection and real under-accounting. +They are opposite results of using predictions as allocation facts. + +| Incident | Observed behavior | Accounting mismatch | +| --- | --- | --- | +| #25782 / #25837 | HashBuild could OOM a CN instead of spilling | hash-table growth was not admitted at its physical allocation boundary | +| #26174 / #26178 | fulltext INSERT requested 18.72 GiB with about 1.29 GiB used | current ingress, retained tail, const materialization, and future drain were charged as if simultaneously owned | +| #26192 / #26231 / #26318 | LOAD DATA was rejected | runtime-filter and payload multipliers duplicated already retained owners | +| #26413 / #26438 | a Hive external-table self-join was rejected | a 50K-row logical ingress estimate did not match the real 8192-row copy segmentation and allocator rounding | +| #26454 | a string-expression join requested exactly 551,368,048,640 bytes while observed memory was about 6--7 GiB | TEXT maximum size was multiplied by row count and recursively summed through CAST/CONCAT nodes | +| #26433 / #26455 | expression result capacity could be under-counted or double-charged across reuse | the budget lease lifetime did not match retained `ExpressionExecutor` capacity | +| #26186 | a spill transition could under-count ingress overlap | independent estimates omitted a simultaneously live ownership state | + +Fixing one multiplier or expression kind does not close the defect class. A +more conservative estimate prevents one OOM shape but rejects more valid +queries. Relaxing the estimate restores those queries but can miss a different +physical overlap. + +### Root problem + +Three concerns are currently mixed: + +1. **Ownership accounting**: which allocations are live, which finite account + owns them, and when their charges end. +2. **Operation prediction**: how much an upcoming expression, batch copy, + marshal, or spill transition might allocate. +3. **Pressure response**: what execution does when the next allocation cannot + fit. + +Only the first concern can support a hard capacity invariant. Prediction and +pressure response remain necessary, but they cannot be the source of truth for +live memory. + +### Relationship to SQL Resource Accounting + +This RFC is separate from +[`SQL Resource Accounting`](./00000000_sql_resource_accounting.md). + +SQL Resource Accounting defines observational facts used by statement trace, +CU, and physical-plan diagnostics. Its memory fields describe MPool domain +peaks after execution; it explicitly does not provide allocation-site +attribution or admission control. + +This RFC defines an execution-time control mechanism: + +| Concern | SQL Resource Accounting | This RFC | +| --- | --- | --- | +| Primary purpose | observe and persist usage | prevent an allocation from exceeding a finite execution account | +| Time of decision | execution summary / statement completion | immediately before allocation or growth | +| Unit | domain usage and peak | live physical allocation capacity | +| Missing data | quality flag | the path is not considered fully protected | +| Failure behavior | report incomplete facts | reclaim, spill, retry, or controlled pressure error | + +The two systems may share metrics and consistency checks, but one must not be +derived from the other. A terminal MPool peak cannot authorize an allocation +that has already happened, and a HashBuild account is not a complete statement +memory measurement. + +## Goals + +1. Make every hard memory charge correspond to a live physical allocation or + a named, bounded scratch owner. +2. Admit memory before allocation and leave allocator/account state unchanged + on failure. +3. Account for the real capacity chosen by MPool, including replacement + overlap during `Grow` and `Grow2`. +4. Make allocation provenance survive vector reuse, cross-owner handoff, and + cross-MPool `Free`. +5. Make Reset, Free, cancellation, retry, and generation turnover obey one + ownership contract. +6. Treat real capacity rejection as recoverable pressure where execution can + make progress. +7. Prevent a prediction alone from rejecting a query. +8. Add no per-row reservation or shared budget lock in the steady-state reuse + path. +9. Make the primitive reusable by spillable operators outside HashBuild. +10. Retain the no-OOM safety objective introduced for #25782. + +## Non-goals + +This RFC does not: + +- make the account equal to CN RSS or total Go runtime memory; +- remove process/CN headroom for caches, RPC, logs, goroutine stacks, or other + memory outside the controlled allocation domain; +- guarantee that every SQL statement completes under every finite cap; +- make an indivisible value smaller than its real representation; +- change SQL syntax, catalog metadata, or a persisted/wire format; +- replace spill disk and file-descriptor accounting; +- use type-specific exemptions for TEXT, CONCAT, CASE, LOAD DATA, or a + particular benchmark; +- raise the existing cap to hide false estimates; +- add a mutable process-wide "current memory account"; +- require all MPool users to become accounted in the first implementation. + +## Terminology + +### Allocation account + +A finite ledger that admits and releases bytes for one execution ownership +domain. `HashBuildBudgetGeneration` is the initial policy implementation, but +the allocator contract is not HashBuild-specific. + +### Accounted allocation + +An allocator-owned allocation whose metadata records: + +- its actual allocated capacity; +- an opaque reference to the account charge; +- bounded diagnostic classification such as owner class and allocation site. + +The charge belongs to the allocation, not to the operator field that currently +references it. + +The first implementation makes data-sized accounted allocations off-heap. +`MPool.Alloc(..., false)` uses a Go allocation: its requested size is not the +runtime size class, and removing MPool metadata does not make the GC reclaim it. +Such memory cannot be described as exact physical ownership at `Free`. + +### Allocation lease + +An exactly-once charge returned after successful admission. It is attached to +the allocation before the allocation is published to its caller. `Free` +releases it. Copying a lease value must not create a second release owner. + +### Explicit scratch lease + +A lease for bounded memory that cannot yet be allocated through MPool. It must +have one named owner, a finite size, and an explicit release point. It is an +exception used during migration, not a substitute for accounting ordinary +buffers. A data- or row-scaled Go allocation is not accepted merely by adding +such a lease: it must move off-heap or retain a conservative charge through a +proved GC-reclamation boundary. The initial implementation permits only small, +statically bounded Go metadata under explicit CN headroom. + +### Prediction + +A non-authoritative estimate used to select an execution strategy, start spill +before a hard limit, or choose an initial batch size. A prediction does not +create a retained allocation lease and cannot directly produce a terminal +budget error. + +### Generation + +One query-CN execution ownership epoch. The generation fixes the account +identity used by allocations that can outlive one operator call or move +between producer and consumer operators. + +## Required invariants + +### I1. Conservation + +For a live account at every observable transition: + +```text +account.used + = sum(capacity of live off-heap allocations charged to the account) + + sum(size of live named explicit scratch leases) +``` + +Predictions, logical vector length, source batch size, and SQL type maximum do +not appear in this equation. Small allocation/runtime metadata classified as +headroom is outside `account.used`; I9 separately requires a finite aggregate +bound for it. + +### I2. Admission precedes allocation + +The complete capacity of a new physical allocation is admitted before MPool or +another allocator changes state. + +### I3. Failure atomicity + +If admission or allocation fails: + +- no new allocation is published; +- the old allocation remains valid; +- any provisional lease is released; +- account usage and MPool usage return to their pre-operation values. + +### I4. Reclaimable ownership defines charge lifetime + +- shrinking a logical length does not release capacity; +- Reset retains both reusable capacity and its charge; +- reuse within existing capacity performs no new admission; +- off-heap `Free` releases physical memory and the charge together; +- a Go reference becoming unreachable is not treated as physical release; +- replacing a buffer transfers publication only after the replacement is + complete. + +### I5. One physical allocation has one charge + +Aliases, vector windows, const views, and shared areas do not create a second +charge. A deep copy creates a new allocation and therefore a new charge. + +### I6. Provenance survives handoff + +The account identity and charge remain correct when: + +- a build-side buffer is published to a `JoinMap`; +- a vector or batch moves to another operator; +- an allocation is freed through a different MPool; +- the producing operator has already Reset; +- a retained buffer is reused in a later call of the same execution + generation. + +### I7. Generation closure is not implicit memory release + +Sealing a generation prevents new admission but does not pretend that live +allocations disappeared. Existing allocation leases remain releasable. Normal +owner cleanup must bring usage to zero; a nonzero terminal value is an +invariant failure and a leak signal. + +### I8. Prediction cannot hard-reject + +An estimate may trigger early reclaim or spill. A terminal memory pressure +error must name a real allocation or explicit scratch request that failed +admission after applicable pressure responses. + +### I9. Metadata headroom is finite + +Every accounted owner has a proved maximum number of simultaneously live +allocations. Pointer headers, account-ID side records, registry entries, and +other per-allocation Go metadata have measured per-entry bounds. Their +aggregate bound is reserved as CN headroom: + +```text +metadata headroom + >= maximum live allocation count * measured metadata bytes per allocation + + maximum live generation count * measured registry bytes per generation +``` + +The implementation enforces both counts with finite CN-local metadata slots. +Opening a generation consumes a generation slot; publishing an accounted +allocation or replacement consumes an allocation-metadata slot; physical Free +returns that slot. Replacement growth temporarily consumes slots for both old +and new allocations. Slot exhaustion is exact metadata pressure, not a +payload estimate. An owner with no supported bound cannot be activated merely +because each payload allocation is charged. + +The first hash-table activation's concrete slot counts, conservative +per-entry bytes, and startup headroom formula are fixed in the implementation +plan. A later owner with smaller allocations must re-derive and provision its +own count before activation; the generic API does not silently widen the +proved domain. + +## Technical design + +### 1. Separate accounting mechanism from cap policy + +The low-level allocator must not import the SQL operator or `process` package. +It depends on a small generic contract, conceptually: + +```go +// Names are illustrative; this RFC does not freeze the Go API. +type AllocationAccount interface { + Acquire(AllocationRequest) (AllocationLease, error) +} + +type AllocationRequest struct { + Capacity uint64 + Class AllocationClass // bounded enum + Site AllocationSite // bounded enum +} + +type AllocationLease interface { + Capacity() uint64 + Release() +} +``` + +The contract is synchronous and does not wait for memory or reclamation; it may +briefly contend on account synchronization. It is called only when a physical +allocation or growth is required, not for every row or append. +It returns non-overlapping typed reasons: finite capacity pressure, +sealed generation, account mismatch, allocator size limit, and invariant +corruption. Only finite capacity pressure may enter reclaim/spill/retry. + +The policy layer remains responsible for: + +- query and CN caps; +- cap refresh; +- concurrency and linearization; +- spill disk/FD ledgers; +- metrics; +- typed pressure errors. + +The allocator layer is responsible for: + +- requesting the actual allocation capacity; +- failure rollback; +- associating the lease with allocation metadata; +- releasing the lease exactly once with physical memory. + +### 2. Allocation metadata and provenance + +MPool already tracks each allocation in pointer metadata so it can identify +the original pool, allocation size, off-heap status, double free, and +cross-pool free. Account provenance belongs at this same boundary. + +An accounted allocation adds an optional opaque charge handle to that +metadata. Unaccounted allocations retain current behavior. The handle's exact +representation must be benchmarked: adding a Go interface to every metadata +entry is not assumed acceptable. A compact pointer or account-local lease +record is preferred if it preserves exactly-once release and diagnostic +identity. + +The metadata is authoritative. Operator-maintained byte totals may remain as +diagnostics during migration, but they cannot independently release or +reconstruct the charge. + +Cross-MPool `Free` already delegates physical release to the original MPool. +The same terminal path releases the account charge, so the freeing caller does +not need to recover the producing operator or generation. + +An existing allocation's metadata also decides the account used by growth. +Growing an account-A allocation under account B is an invariant error. +Unaccounted-to-accounted conversion is never implicit: a migrated owner creates +an accounted destination and copies from the unaccounted source, or creates its +own buffers as accounted from the beginning. This prevents a caller from +changing only a vector field while the retained physical buffer still has +different provenance. + +For the compact-handle design, the registry is a finite set of reusable slots. +An opaque handle contains the slot and its generation counter. A slot is +reused only after sealing and exact zero; incrementing its generation makes +every older handle stale. Generation counters never wrap: a slot whose counter +is exhausted is retired. A missing or generation-mismatched registry entry +during `Free` is an invariant failure, not permission to drop the charge. + +Pointer headers, account-ID side records, and registry entries are `H` +metadata, not accounted payload. Their backing stores are sized from finite +CN-local slot limits and charged to explicit CN headroom. Before activating an +owner, supported live-allocation and generation counts must fit those limits; +the limits remain the safety backstop if an ownership assumption is wrong. + +### 3. New allocation protocol + +Initial allocation and growth share one authoritative +`AllocationCapacity(request, allocatorMode)` calculation. It includes +allocator rounding and the actual maximum accepted by `Alloc`; callers do not +infer capacity from a logical size or from `GrowCapacity` alone. + +For an account-aware allocation of capacity `C`: + +```text +calculate actual allocator capacity C + -> acquire lease(C) + -> reserve one allocation-metadata slot + -> allocate C + -> on failure: return metadata slot, release lease(C), return error + -> attach lease to allocation metadata + -> publish allocation to caller +``` + +If metadata attachment can fail, it is part of the unpublished transaction: +free the new allocation, release the lease, and return an error. + +MPool cap failure, global cap failure, and underlying allocator failure all +follow the same rollback contract. + +Metadata publication is checked. An allocator panic before publication must +run provisional-lease cleanup. Cross-pool `Free` and deleted-owner-pool +fallback release the lease when they physically deallocate the allocation. + +Pool teardown is not itself a universal release event. A `noLock` teardown +that physically deallocates pool-local allocations releases each matching +lease. Unregistering a normal pool retains global pointer metadata and the +charge until later physical `Free`; live accounted allocations at teardown are +reported as an invariant violation, never bulk-released. + +### 4. Growth and replacement protocol + +MPool `Grow` and `Grow2` currently allocate a replacement, copy the old bytes, +then free the old allocation. Hard accounting must represent that real +overlap. + +For old capacity `O` and required logical size `R`: + +1. calculate `N = AllocationCapacity(GrowCapacity(O, R), allocatorMode)`; +2. keep the old allocation and its `O` charge live; +3. acquire a complete `N` lease, not `N - O`; +4. allocate `N`; +5. copy old data and any second source; +6. attach the new lease and publish the replacement; +7. free the old allocation, releasing its `O` lease. + +Peak account usage during replacement is therefore: + +```text +other live allocations + O + N +``` + +After publication and old-buffer release it is: + +```text +other live allocations + N +``` + +Reserving only the delta would under-count the actual allocate-copy-free peak. +Using a multiplier would be an estimate of the same fact even though the +allocator already knows `O` and `N`. + +If `R <= O`, no physical growth occurs and no account call is made. + +#### Growth failure table + +| Failure point | Old buffer | New buffer | Account result | +| --- | --- | --- | --- | +| capacity calculation | unchanged | none | unchanged | +| new lease admission | unchanged | none | unchanged | +| physical allocation | unchanged | freed/not published | new lease released | +| copy before publication | unchanged | freed/not published | new lease released | +| publication succeeds | released afterward | live | old lease released; new lease retained | + +The implementation must preserve this table under panic-safe cleanup where +MPool currently permits a recoverable error. + +### 5. Vector and batch propagation + +An allocation's metadata owns the retained charge. A vector additionally +needs an optional account selection for the first allocation of a currently +nil buffer. Subsequent growth inherits or verifies the existing allocation's +account. + +The following rules apply: + +| Operation | Rule | +| --- | --- | +| append within capacity | no admission | +| append causing data growth | admit exact replacement capacity | +| append causing varlen-area growth | independently admit exact replacement capacity | +| logical reset / set length to zero | retain allocation and charge | +| vector Free | free every owning buffer; each allocation releases its own charge | +| window/view/const alias | no new allocation and no new charge | +| deep Dup/copy | destination allocations use the destination account | +| batch handoff | allocation metadata preserves charges; no sum-and-rereserve | +| cross-pool Free | original allocation metadata releases the original charge | + +Data and varlen area are separate physical allocations and therefore separate +charges. Null bitmap and other auxiliary buffers follow their actual +allocation ownership rather than a synthetic per-vector total. + +The first production migration accounts only off-heap vector buffers. +Row-scaled null/group bitmaps currently allocate Go `[]uint64`; they must move +to an off-heap owner or remain an explicit activation blocker. Switching a +vector field to an account while its existing backing remains on-heap is not a +valid migration. + +A shared area must retain one physical release owner. If current vector sharing +permits multiple logical owners, account integration must use the same +reference/ownership mechanism that prevents physical double free; it must not +introduce a second budget-only reference count. + +### 6. Expression execution + +`FunctionResult` and `ExpressionExecutor` must create result vectors with the +execution account selected for the owning HashBuild path. + +Fixed-width results allocate from actual row count and element width through +normal vector growth. Varlen results allocate from actual appended payload. +Neither uses the maximum SQL type width multiplied by row count. + +Intermediate expression results are charged only while their buffers are +physically live. Reuse keeps the charge. Reset does not release it unless Reset +also frees the buffer. + +Expression implementations that create unbounded temporary Go strings or byte +slices bypass MPool and must be changed by one of these methods: + +1. write directly into an account-aware result buffer; +2. use an account-aware MPool scratch buffer; +3. for small metadata only, use a named explicit scratch lease with a proved + finite bound and one cleanup owner. + +For example, CONCAT should not construct an unaccounted complete Go string and +then copy it into an accounted result vector. The temporary and final buffers +can be simultaneously live, so omitting the temporary would violate I1. + +Recursive expression peak calculation may remain temporarily as an early-spill +hint. Before deleting it as a hard gate, an allocation-site ledger must cover +every reachable data-scaled allocation in generic evaluation and built-ins, +including selection arrays, conversion slices, nested result vectors, and +function-specific scratch. Migrating only `FunctionResult` is not closure. + +### 7. Non-vector HashBuild allocations + +The migration inventory must include all memory whose lifetime is owned by the +HashBuild execution domain: + +- copied build and probe batches; +- integer and string hash-table blocks; +- selection lists and group mappings; +- join-map auxiliary storage; +- expression keys and intermediate results; +- spill scatter buffers; +- spill/re-spill read and decode buffers; +- runtime-filter buffers; +- marshal/unmarshal scratch; +- retained emergency scratch. + +An ownership closure is not complete merely because its largest vector is +accounted. Every data-scaled allocation reachable from an activated owner must +use account-aware, synchronously reclaimable allocation. Unactivated owners +retain their complete legacy charge until their own closure; only small, +statically bounded metadata may remain under named headroom. + +Hash-table callbacks that already reserve from `ResizePlan` are an intermediate +bridge. The final charge should be owned by the physical hash-table allocation +metadata rather than a parallel slice of reservation tokens in the operator. + +### 8. Pressure response state machine + +A failed exact allocation admission returns typed memory pressure, not +`ErrHashBuildBudgetInvalid`. + +The owner handling the request proceeds through a bounded state machine: + +```text +need allocation + -> exact admission succeeds + -> allocate and continue + -> exact admission rejected + -> release reclaimable retained capacity + -> retry exact allocation + -> start/advance spill + -> retry exact allocation + -> reduce processing batch where semantics permit + -> retry exact allocation + -> degrade optional structure where permitted + -> retry/continue + -> return controlled minimum-unit pressure error +``` + +Not every allocation supports every response. The request carries an owner +class that maps to a policy: + +| Owner class | Permitted response | +| --- | --- | +| retained reusable result | release retained capacity, then retry | +| spillable build/probe input | spill/re-spill, then retry | +| splittable expression batch | reduce batch, then retry | +| runtime filter | degrade to PASS when correctness is unchanged | +| indivisible single value / minimum hash block | controlled error with actual capacity | +| invariant corruption | fail immediately as invalid state | + +The state machine records which responses were attempted so it cannot loop +without progress. A retry is justified only after account usage decreased, +the input unit became smaller, spill state advanced, or an optional owner was +disabled. + +`ErrHashBuildBudgetInvalid` remains reserved for arithmetic overflow, corrupted +ownership, double release, account mismatch, or an impossible lifecycle +transition. Ordinary finite pressure is not an invariant failure. + +A sealed or stale generation is a lifecycle result distinct from finite +pressure. It cannot enter the retry state machine. + +#### Retry transaction boundary + +Allocation rollback alone does not make an operator operation retryable. A +vector or expression may successfully grow one retained buffer, fail on the +next allocation, and leave reusable capacity or partially written output. +Each retry-capable owner therefore defines: + +- the unpublished operation checkpoint; +- which retained capacity may survive failure; +- how row/output publication is rolled back; +- the Reset/Free actions required before retry; +- how a smaller batch resumes without duplicate work. + +Until the shared controller and these checkpoints exist, an exact rejection +returns a controlled terminal pressure error. An earlier migration must not +claim spill/reduce/retry merely because it has allocation-level rollback. + +#### Forward-progress memory + +Spill and reclaim paths do not bypass admission. Encoding, decoding, scatter, +and IO buffers are allocations too. Allowing normal work to consume the last +byte and then allocating uncharged "emergency" scratch would reproduce the +original safety hole at a different site. + +Before retaining a unit of work, an operator must preserve one bounded way to +make progress: + +1. reuse an already allocated and accounted spill buffer; +2. keep a finite progress sub-cap that normal work cannot consume and charge + actual spill/cleanup allocations against it; or +3. reduce the retained or spill chunk until the minimum real progress + allocation fits. + +Progress headroom is cap policy, not a fabricated live-memory charge: +`account.used` still contains only actual allocations and explicit live +scratch. If a sub-cap is used, normal and progress allocations remain bounded +by the same total query/CN cap, with the progress portion unavailable to normal +growth. + +The minimum progress unit is derived from the concrete buffer/chunk layout, not +from SQL type maximums. If retained state plus that minimum real unit cannot +fit, admission must stop earlier or return a controlled pressure error; it +must not wait until spill itself is unable to start. + +### 9. Concurrency and generation lifecycle + +Child pipelines may share a `BaseProcess` and execute concurrently. Therefore +neither Process nor MPool may contain a mutable "current account" used +implicitly by all allocations. + +Account selection is explicit at the owner/allocation boundary and immutable +for a published allocation. Concurrent allocations linearize in the account's +existing query/CN admission operation. Reuse within capacity does not enter +that lock. + +Normal generation lifecycle is: + +```text +open + -> allocations and explicit scratch may acquire leases + -> seal: no new leases + -> owner cleanup frees all live allocations/scratch + -> used reaches zero + -> finalize +``` + +One `Compile` execution attempt on each CN is the sole seal/finalize owner for +that CN's generation. HashBuild Reset cannot seal: `JoinMap`, spill payload, +broadcast, remote scope, and message-board consumers may outlive the producing +operator and still own or create controlled work. The attempt coordinator +seals only after all scopes and remote notifiers it owns have stopped +publishing work and its MessageBoard consumers have reached a terminal state. +Cleanup may then release old leases; exact zero produces one immutable +terminal snapshot and permits registry-slot reuse. The statement +`ResourceRoot` aggregates these immutable CN-attempt snapshots but does not own +allocation release. + +A nonzero terminal generation does not disappear and does not wait forever: + +```text +seal after execution/message quiescence + -> run terminal owner cleanup + -> zero: + export one valid immutable snapshot + remove the registry entry + -> nonzero at terminal-cleanup completion: + export one immutable invariant-failure snapshot + retain a release-capable tombstone until late Free reaches zero + suspend admission of new accounted generations on that CN +``` + +Suspension bounds tombstone growth to generations already active when the +first invariant failure is detected. It is lifted only after every tombstone +reaches zero; an operational deadline escalates with owner/site diagnostics +and permits a controlled CN restart rather than deleting provenance. Late +release may update bounded health counters, but it cannot rewrite or duplicate +the exported snapshot. + +Generation open and nonzero-terminal suspension linearize through the same +CN-local generation gate. Opening publishes a generation only if the +suspension check succeeds in that transaction. Once suspension publication +linearizes, no later open may publish; opens that linearized earlier are the +finite active set allowed to finish or become tombstones. + +Current `SetStmtProfile` turnover, frontend `StatementInfo.EndStatement`, and +`HashBuildBudgetGeneration.Close` do not prove this per-CN quiescence or +validate zero. Production activation is blocked until the `Compile` attempt +owns an explicit post-pipeline/MessageBoard-close transition for success, +failure, panic, cancellation, retry, broadcast, prepared reuse, and remote +execution. A forced close must not silently zero accounting while allocations +remain live. + +Ownership transfer does not change generations. If a transfer would cross to a +different generation, it must either: + +- retain the original generation until the allocation is freed; or +- perform one explicit atomic charge transfer before publication. + +The initial implementation should prefer retaining original provenance; charge +transfer adds a second failure and rollback boundary and is unnecessary for +normal HashBuild producer-to-consumer handoff. + +### 10. Observability + +Admission exposes bounded owner/site, actual capacity, used/cap, attempted +pressure response, and terminal result. Exact allocation events, prediction +hints, pressure responses, and invariant failures remain distinguishable. + +Owner/site values are bounded enums. Metrics and logs aggregate at +operator/generation or terminal-pressure boundaries; there is no +per-allocation log or unbounded SQL/stack label. + +The controlled-domain snapshot has a stable generation identity and is +exported exactly once by its CN attempt coordinator. The statement resource +root may aggregate those snapshots. It is separate from SQL Resource +Accounting's current off-heap MPool domain; consumers must not sum duplicate +operator references to one generation or claim the domains match before owner +coverage does. + +## Migration plan + +Migration is incremental by complete physical owner. Exact accounting and +removal of that owner's legacy hard charge happen atomically; one buffer is +never charged by both models. + +The implementation order is: + +1. measured metadata/API decisions and a reference model; +2. generic MPool allocation transaction; +3. dormant Vector/Batch propagation; +4. allocation-site closure and dormant expression/spill propagation; +5. statement generation lifecycle, typed pressure, and retry checkpoints; +6. hash-table cell/descriptor activation with only its legacy charge removed; +7. copied-batch and JoinMap activation; +8. expression-owner activation; +9. spill and runtime-filter closures; +10. unified join pressure control and remaining legacy estimator deletion; +11. workload and performance acceptance. + +The allocation-site ledger, PR gates, rollback rules, and test commands live in +the +[implementation plan](../design/allocation_accounted_memory_admission_impl.md). + +## Rollout and compatibility + +The change is internal to one CN binary and does not change SQL, catalog, disk, +or RPC formats. Unmigrated MPool users retain current behavior. A HashBuild +owner enables exact accounting only after its full alloc-to-Free closure is +covered and its legacy hard charge is removed. The final design has no +permanent legacy/exact switch. + +## Testing strategy + +Testing is derived from the invariants: + +- a randomized reference model checks conservation after alloc, grow, failure, + reuse, handoff, Free, seal, and cancellation; +- boundary and fault tests cover exact cap, allocator rounding, unpublished + rollback, cross-pool Free, and generation turnover; +- container/operator tests cover aliases, varlen data, Reset, broadcast, + spill/re-spill, pressure progress, and optional degradation; +- workload regressions cover #25782, #26174, #26192, #26413, #26454, and TPCH + spill/non-spill paths; +- performance gates verify no per-row account operation, no budget lock on + within-capacity reuse, bounded metadata cost, concurrent-generation P50/P99 + admission latency, release storms, and separately measured resident/spill + behavior. + +Exact matrices and per-PR gates are maintained in the +[implementation plan](../design/allocation_accounted_memory_admission_impl.md). + +## Drawbacks + +### Allocator and container changes are invasive + +Correct ownership crosses MPool, vector, batch, expression, hash table, and +operator handoff boundaries. A partial implementation can create a more +convincing but still incomplete safety claim. + +Mitigation: migrate by complete owner class, maintain the ownership ledger, and +gate each phase on conservation properties rather than workload success alone. + +### Allocation metadata has a cost + +An account handle can increase pointer-map memory and alloc/free work even when +only some allocations are accounted. + +Mitigation: keep it optional and compact, benchmark representation choices, +and avoid a Go interface stored inline in every metadata record unless +measurement supports it. + +### Exact replacement overlap can reject earlier than steady-state size + +If a 6 GiB buffer grows to 8 GiB, the current allocator may need 14 GiB live +during copy even though the final buffer is 8 GiB. Charging only 2 GiB would +look friendlier but would not protect the real peak. + +Mitigation: reclaim or spill before growth, reduce the processing unit, or +introduce a genuinely lower-overlap allocator operation. Do not hide the peak +with delta accounting. + +### Go-heap reclamation is not allocator-controlled + +Dropping an MPool pointer record for `Alloc(..., false)` does not synchronously +return its backing bytes. Treating that event as exact physical release would +allow new work while the old Go object remains resident. + +Mitigation: keep data-scaled controlled owners off-heap. Audit and migrate +row/payload-scaled Go slices before activation; reserve separate CN headroom +only for small, proved-bounded runtime metadata. + +### This is not complete RSS accounting + +Go runtime, goroutine stacks, caches, and unrelated subsystems remain outside +the HashBuild account. + +Mitigation: preserve explicit CN headroom and state coverage boundaries. A +future broader account may reuse the primitive but requires its own ownership +inventory. + +## Rationale and alternatives + +### Tune multipliers and type rules + +Rejected. The confirmed incidents demonstrate that no fixed multiplier +represents aliasing, reuse, variable payloads, segmentation, and replacement +overlap simultaneously. + +### Use predictions only to choose spill, then remove hard budgets + +Rejected. Prediction errors in the other direction can again allow #25782 to +OOM the CN. Real allocations still need finite admission. + +### Charge only `newCapacity - oldCapacity` on growth + +Rejected for the current allocate-copy-free implementation. It under-counts +the period where both allocations are live. + +### Use MPool current bytes as the HashBuild budget + +Rejected. A shared MPool contains other owners and does not preserve +HashBuild/query provenance across child pipelines. Sampling after allocation +also cannot provide pre-allocation safety. + +### Store a mutable current account on Process or MPool + +Rejected. Concurrent child pipelines can share the same base process and MPool. +The wrong goroutine could charge or free against another generation. + +### Poll RSS and spill near the cgroup limit + +Rejected as the primary control. RSS is delayed, includes unrelated memory, +and cannot make a specific allocation failure-atomic. It remains useful as a +coarse pressure signal and validation metric. + +### Wrap MPool only at operator call sites + +Insufficient by itself. A wrapper can select the account for initial +allocation, but provenance must still survive vector growth, physical handoff, +and cross-pool Free. The terminal charge belongs in allocation metadata. + +### Raise or disable the cap + +Rejected. It hides false rejection while weakening the original no-OOM +requirement. + +## Unresolved questions + +1. Does the provisional 16-byte-header plus side-map representation retain its + advantage in real MPool, cross-pool, and high-concurrency benchmarks? +2. What explicit MessageBoard close-and-drain primitive and tests prove + quiescence at the selected local and remote `Compile` attempt hooks? +3. Which remaining data-scaled Go-heap sites can write directly to off-heap + output, and which need a new off-heap scratch abstraction? +4. What is the minimum semantically safe batch and operation checkpoint for + each expression and spill phase? +5. What cap/headroom policy is appropriate once exact HashBuild ownership + replaces conservative estimates? This is policy work and must not change the + accounting invariant. +6. What measured metadata and hot-path overhead is acceptable for enabling the + primitive beyond HashBuild? + +These questions affect implementation shape, not the core decision that hard +admission must be tied to owned physical capacity. + +## Acceptance criteria + +The RFC is implemented only when: + +- I1--I9 hold under property tests, fault injection, cancellation, and race + execution; +- account-aware alloc/grow/free use the same real capacity calculation as + MPool; +- every HashBuild-owned data-scaled allocation is off-heap and accounted; +- every excluded Go/runtime metadata allocation is statically bounded and + covered by explicit CN headroom; +- live-allocation and generation counts prove the aggregate I9 metadata + headroom; +- each CN `Compile` attempt seals, validates exact zero, and exports one + generation snapshot after execution/MessageBoard quiescence; +- a nonzero terminal generation exports one failure snapshot, preserves + release provenance, and cannot accumulate unbounded tombstones; +- Reset retains reusable allocation charges and Free releases them exactly + once; +- aliases and handoffs do not duplicate charges; +- estimator-only false rejection is structurally impossible on migrated + paths; +- exact pressure triggers bounded reclaim/spill/retry/reduce behavior only + across proved operation checkpoints; +- real minimum-unit over-cap errors report actual allocation capacity and + owner/site; +- #25782 cannot exceed the finite account or OOM the CN; +- #26174, #26192, #26413, and #26454 pass durable workload regressions; +- TPCH spill and non-spill performance gates pass; +- superseded hard estimators and parallel reservation owners are removed; +- the implementation documentation states the remaining memory outside the + account and preserves corresponding CN headroom. From a3aed96f2f784330e02e049b0b3c58cb7c72adf7 Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 13:41:10 +0800 Subject: [PATCH 22/61] docs: record dormant vector accounting validation --- ...ocation_accounted_memory_admission_impl.md | 53 +++++++++++++++++-- 1 file changed, 49 insertions(+), 4 deletions(-) diff --git a/docs/design/allocation_accounted_memory_admission_impl.md b/docs/design/allocation_accounted_memory_admission_impl.md index 1bb24759d21d8..99a8ea21edd92 100644 --- a/docs/design/allocation_accounted_memory_admission_impl.md +++ b/docs/design/allocation_accounted_memory_admission_impl.md @@ -90,11 +90,11 @@ Ledger states: The independent review rejected owner-class rows as proof of closure. The working ledger is allocation-site based: -| Allocation site | Allocator/mode and size | Terminal owner | Initial | Target/blocker | +| Allocation site | Allocator/mode and size | Terminal owner | Current | Target/blocker | | --- | --- | --- | ---: | --- | -| `mpool.memHdr` and account-ID side map | Go maps; one pointer record plus optional account record per live allocation | pointer removal at physical deallocation | L | H after per-entry and maximum-live-count proof | -| `Vector.data` | MPool; capacity from `Grow`, on/off-heap follows `v.offHeap` | owning `Vector.Free` | L | A only when off-heap | -| `Vector.area` | MPool; independent varlen payload capacity | owning `Vector.Free` | L | A only when off-heap | +| `mpool.memHdr` and account-ID side map | Go maps; one pointer record plus optional account record per live allocation | pointer removal at physical deallocation | H | bounded by the measured finite registry/allocation-slot policy | +| `Vector.data` | MPool; capacity from `Grow`, on/off-heap follows `v.offHeap` | owning `Vector.Free` | D | A only when off-heap | +| `Vector.area` | MPool; independent varlen payload capacity | owning `Vector.Free` | D | A only when off-heap | | `Vector.nsp/gsp` bitmap data | Go `[]uint64`; `ceil(rows/64)*8`, retained by `Clear` | bitmap `Reset` from `Vector.Free` | L | move off-heap; blocks Vector-dependent activation | | `FunctionResult.vec` data/area | off-heap Vector; rows and appended payload | executor `Free` | L | A | | `FunctionResult.convenientParam` | Go slice; expression arity, not rows | executor `Free`/reuse | L | H after a proved arity bound | @@ -114,6 +114,7 @@ working ledger is allocation-site based: | spill counts/offsets/positions | Go `[]int32`; O(bucket count), bucket count finite | spill cleanup | L | H after bound is asserted | | selected spill bucket vectors | off-heap Vector capacities | selected batch cleanup | L | A | | BucketReader decoded vectors | MPool Vector data/area | `BucketReader.Close` | L | A after mode/provenance audit | +| `pSpool` cached Vector data/area | raw MPool slices retained and reassigned independently of their Vector | `spoolBuffer.clean` or the receiving Vector's `Free` | L | persist generation/selection provenance for a missing data or area allocation; blocks pipeline activation | | runtime-filter serialized payload | Go buffer/message payload; O(filter rows) | message release | L | off-heap or PASS degradation; blocks runtime-filter activation | | spill disk and FD | disk/FD ledgers | file removal/close | A | A | @@ -124,6 +125,15 @@ PR 3 must generate and review the remaining built-in/function-specific `make`, expression or spill closure can activate. No row named “other” or “unbounded scratch” can declare closure. +Batch destination propagation is now `D`: Clone, Dup, selected-column copy, +Union destinations, windows, reader decode, Clean, and FreeColumns preserve +the immutable destination selection without creating a synthetic batch-level +charge. `pSpool` is deliberately still `L`; its raw buffer cache can retain an +allocation after detaching it from a Vector, so merely copying the Batch +selection would disagree with the original account still recorded in the +MPool lease. `Vector.SetTypeAndFixData` also remains a PR 3 closure blocker +because its legacy API cannot currently return a failed growth admission. + ## 3. Decisions required before production integration ### A. Allocation metadata representation @@ -512,6 +522,41 @@ Gate: - package race tests and vector benchmarks; - Vector/Batch ledger rows become `D`. +The current PR 2 candidate is +`feature/26459-vector-propagation` at commit `dbfee20ecc`. It remains dormant. +It adds one immutable shared selection pointer to Vector and Batch, accounts +the first owned off-heap data/area allocation, lets later Grow/Grow2 inherit +the physical MPool lease, and rejects implicit conversion to on-heap or +no-copy aliases. Reset retains the selection and charge; Free clears the +selection after the physical allocations release their leases. Views carry no +selection, while Batch windows retain only the destination context needed for +a later deep copy. + +The implementation also closes two error edges found during self-review: +reader growth publishes the replacement buffer before a short read can return, +so cleanup never retains a freed old pointer, and no-copy Batch decode +explicitly detaches an empty Vector selection while retaining the Batch +destination context. + +Fresh local evidence on linux/amd64, Go 1.26.4, i7-11700: + +| Vector operation, `GOMAXPROCS=8` | Legacy median | Accounted median | Difference | +| --- | ---: | ---: | ---: | +| fixed pre-extend/free, 8,192 rows | 1,026 ns | 1,116 ns | +8.8% | +| varlen data+1 MiB area pre-extend/free | 76,766 ns | 77,650 ns | +1.2% | +| accounted fixed Reset/reuse | n/a | 1.524 ns | no account operation | + +Fixed paths remain 0 B/op and 0 allocs/op. Both varlen paths report the same +48 B/op and 2 allocs/op, so accounting adds no Go allocation. Randomized +fixed/varlen append, separate data/area charge, within-capacity reuse, Reset, +Free, views, partial selection, cross-owner copies, metadata rollback, +cross-pool Free, sealed accounts, shuffle replacement, copy/reader decode, and +Batch Clone/Dup/Union/FreeColumns pass. Every new and directly affected test +passed an exact `-race -count=100` run; both owning packages passed complete +race runs, build, vet, coverage, and dependent HashBuild/SQL/engine package +tests. No production owner selects an account and no legacy hard gate is +removed. + ### PR 3: allocation-site closure and dormant propagation Scope: From b8ae526349ce931f549e1057501c598c10d740de Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 14:40:52 +0800 Subject: [PATCH 23/61] docs: record expression and spill accounting progress --- ...ocation_accounted_memory_admission_impl.md | 79 +++++++++++++++++-- 1 file changed, 71 insertions(+), 8 deletions(-) diff --git a/docs/design/allocation_accounted_memory_admission_impl.md b/docs/design/allocation_accounted_memory_admission_impl.md index 99a8ea21edd92..a24d825659f0c 100644 --- a/docs/design/allocation_accounted_memory_admission_impl.md +++ b/docs/design/allocation_accounted_memory_admission_impl.md @@ -96,12 +96,12 @@ working ledger is allocation-site based: | `Vector.data` | MPool; capacity from `Grow`, on/off-heap follows `v.offHeap` | owning `Vector.Free` | D | A only when off-heap | | `Vector.area` | MPool; independent varlen payload capacity | owning `Vector.Free` | D | A only when off-heap | | `Vector.nsp/gsp` bitmap data | Go `[]uint64`; `ceil(rows/64)*8`, retained by `Clear` | bitmap `Reset` from `Vector.Free` | L | move off-heap; blocks Vector-dependent activation | -| `FunctionResult.vec` data/area | off-heap Vector; rows and appended payload | executor `Free` | L | A | +| `FunctionResult.vec` data/area | off-heap Vector; rows and appended payload | executor `Free` | D | A after expression ledger closure | | `FunctionResult.convenientParam` | Go slice; expression arity, not rows | executor `Free`/reuse | L | H after a proved arity bound | | decimal parameter conversion | Go `[]T`; `rows*sizeof(T)` in `GenerateFunctionFixedTypeParameter` | evaluation wrapper/GC | L | move off-heap; blocks expression activation | -| IFF/CASE/COALESCE selection arrays | Go `[]bool`; one or two arrays of `rows`, retained by executor | executor `Free`/reuse | L | move off-heap; blocks expression activation | -| selected row IDs | Go `[]int64`; up to `rows`, retained by executor | executor `Free`/reuse | L | move off-heap; blocks expression activation | -| selected parameter/result vectors | off-heap Vector capacities | executor `Free` | L | A | +| IFF/CASE/COALESCE selection arrays | allocation-accounted off-heap `[]bool`; one or two arrays of `rows`, retained by executor | executor `Free`/reuse | D | A after expression ledger closure | +| selected row IDs | allocation-accounted off-heap `[]int64`; capacity up to `rows`, retained by executor | executor `Free`/reuse | D | A after expression ledger closure | +| selected parameter/result vectors | allocation-accounted off-heap Vector capacities | executor `Free` | D | A after expression ledger closure | | hash-table initial cell block | off-heap `mpool.MakeSlice(..., true)`; 16 KiB int / 32 KiB string | hash map / `JoinMap.FreeMemory` | L | A in first activation | | hash-table replacement/appended cell blocks | off-heap blocks, at most 4 MiB each; old+new overlap is physically visible | hash map / `JoinMap.FreeMemory` | L | A in first activation | | hash-table `cells`/`newBlocks` descriptors | Go `[][]Cell`; 24 bytes per block header plus geometric resize backing arrays; GC is not treated as synchronous Free | hash map lifetime / GC | L | replace with an owning off-heap descriptor buffer and account its initial/replacement capacity; blocks first activation | @@ -109,11 +109,11 @@ working ledger is allocation-site based: | `GroupSels.{tmp,vals,offsets}` | on-heap `mpool.MakeSlice(..., false)`; O(build rows/groups) | builder or `JoinMap.FreeMemory` | L | switch off-heap; blocks auxiliary/copied-batch activation | | copied build-batch vector buffers | MPool Vector data/area | builder or `JoinMap.FreeMemory` | L | A after per-buffer provenance | | spill marshal/coalesce buffers | Go `bytes.Buffer`; O(serialized batch), retained by phase | spill cleanup | L | off-heap writer; blocks spill activation | -| spill hash values | Go `[]uint64`; `8*rows`, retained by phase | spill cleanup | L | move off-heap; blocks spill activation | -| spill row IDs | Go `[]int32`; `4*rows`, retained by phase | spill cleanup | L | move off-heap; blocks spill activation | +| spill hash values | allocation-accounted off-heap `[]uint64`; geometric capacity, `8*cap` | spill phase cleanup | D | A after spill ledger closure | +| spill row IDs | allocation-accounted off-heap `[]int32`; geometric capacity, `4*cap` | spill phase cleanup | D | A after spill ledger closure | | spill counts/offsets/positions | Go `[]int32`; O(bucket count), bucket count finite | spill cleanup | L | H after bound is asserted | -| selected spill bucket vectors | off-heap Vector capacities | selected batch cleanup | L | A | -| BucketReader decoded vectors | MPool Vector data/area | `BucketReader.Close` | L | A after mode/provenance audit | +| selected spill bucket vectors | allocation-accounted off-heap Vector capacities | per-call selected-batch cleanup | D | A after spill ledger closure | +| BucketReader decoded vectors | allocation-accounted MPool Vector data/area | reusable batch cleanup / `BucketReader.Close` | D | A after spill ledger closure | | `pSpool` cached Vector data/area | raw MPool slices retained and reassigned independently of their Vector | `spoolBuffer.clean` or the receiving Vector's `Free` | L | persist generation/selection provenance for a missing data or area allocation; blocks pipeline activation | | runtime-filter serialized payload | Go buffer/message payload; O(filter rows) | message release | L | off-heap or PASS degradation; blocks runtime-filter activation | | spill disk and FD | disk/FD ledgers | file removal/close | A | A | @@ -578,6 +578,69 @@ Gate: owners; - all migrated rows become `D`; production rows remain `L`. +The current dormant PR 3 candidate is +`feature/26459-expression-propagation` at expression commit `03296c5246` and +spill commit `0b4a7b5bb5`, stacked on PR 2 commit `dbfee20ecc`. + +Its propagation call chains are: + +```text +NewExpressionExecutorWithAllocation + -> recursive expression construction + -> constant/result/scratch AllocationAccountSelection + -> FunctionResult or selected/decode Vector growth + -> MPool AllocAccounted/Grow/Free + +NewSpillEngineWithAllocation + -> BucketReader decoded/reused Batch selection + -> scatter hash/row typed slices and selected Batch selection + -> MPool AllocAccounted/Grow/Free +``` + +The candidate adds no production caller of either dormant constructor and +removes no legacy reservation. It covers fixed, varlen, NULL, decoded-vector, +nested `CASE(CONCAT(CAST))`, folded and non-folded result transfer, partial +selection, repeated Reset/reuse, construction rollback, zero-length retained +scratch growth, decoded-record merge/error cleanup, scatter selected-vector +peak, and capacity-failure cleanup. Typed slices grow geometrically only on +the accounted path; the old and replacement capacities are simultaneously +charged until publication, and terminal cleanup frees a zero-length view by +its retained capacity. + +The first syntax inventory over non-test expression/builtin and spill sources +found 484 candidate lines. Because one line can match more than one category, +the overlapping counts are 233 `make([]...)`, 205 capacity-growing or +potentially growing `append`, 44 `bytes.Buffer`/`NewBuffer`, and 5 +`strings.Builder` sites. This is a review queue, not proof that every match is +data-scaled or reachable. The currently closed rows are the ones marked `D` +in the ledger above. Activation remains blocked by: + +- Vector null/group bitmap backing; +- decimal parameter conversion slices; +- spill marshal and coalesce `bytes.Buffer` backing; +- data-scaled function-specific Go-heap scratch identified by the remaining + built-in scan; +- `pSpool` raw-buffer provenance and `Vector.SetTypeAndFixData`. + +Fresh local validation on linux/amd64 with CGO and the repository-built +third-party artifacts passes: + +- complete normal tests for MPool, Vector, Batch, SQL util, expression, + spillutil, HashBuild, HashJoin, DedupJoin, RightDedupJoin, and Process; +- build and vet for the same package closure; +- exact `-race -count=100` runs for every new test plus directly affected flow + control, BucketReader merge, scatter lifecycle, and re-spill tests; +- complete race runs for MPool, Vector, Batch, SQL util, expression, and + spillutil; +- package coverage of 73.3% MPool, 47.8% Vector, 74.4% Batch, 26.6% SQL util, + 64.3% expression, and 79.7% spillutil. + +The existing constant-flow-control benchmark remains 0 B/op and +0 allocs/op. Five-run medians at `GOMAXPROCS=8` were 4.480 ns/op on the PR 2 +base and 4.464 ns/op on the PR 3 candidate; this focused benchmark shows no +measurable legacy fast-path regression, but it is not an activation-level +performance result. + ### PR 4: statement lifecycle and minimum pressure foundation Scope: From f88bfcf7d40808e7de85483e1bbd1b672f93d6c8 Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 15:23:01 +0800 Subject: [PATCH 24/61] docs: record spill and spool accounting closure --- ...ocation_accounted_memory_admission_impl.md | 65 ++++++++++++++----- 1 file changed, 48 insertions(+), 17 deletions(-) diff --git a/docs/design/allocation_accounted_memory_admission_impl.md b/docs/design/allocation_accounted_memory_admission_impl.md index a24d825659f0c..93a9757c250ef 100644 --- a/docs/design/allocation_accounted_memory_admission_impl.md +++ b/docs/design/allocation_accounted_memory_admission_impl.md @@ -108,13 +108,13 @@ working ledger is allocation-site based: | hash-table `ResizePlan` and callback | fixed-size Go values/closures, one per table/resize | resize return / hash map Free | L | H; remove legacy reservation owner after cell activation | | `GroupSels.{tmp,vals,offsets}` | on-heap `mpool.MakeSlice(..., false)`; O(build rows/groups) | builder or `JoinMap.FreeMemory` | L | switch off-heap; blocks auxiliary/copied-batch activation | | copied build-batch vector buffers | MPool Vector data/area | builder or `JoinMap.FreeMemory` | L | A after per-buffer provenance | -| spill marshal/coalesce buffers | Go `bytes.Buffer`; O(serialized batch), retained by phase | spill cleanup | L | off-heap writer; blocks spill activation | +| spill marshal/coalesce buffers | allocation-accounted off-heap streaming buffers; exact serialized size and bounded 64 KiB per-bucket coalesce capacity | spill phase cleanup | D | A after spill ledger closure | | spill hash values | allocation-accounted off-heap `[]uint64`; geometric capacity, `8*cap` | spill phase cleanup | D | A after spill ledger closure | | spill row IDs | allocation-accounted off-heap `[]int32`; geometric capacity, `4*cap` | spill phase cleanup | D | A after spill ledger closure | | spill counts/offsets/positions | Go `[]int32`; O(bucket count), bucket count finite | spill cleanup | L | H after bound is asserted | | selected spill bucket vectors | allocation-accounted off-heap Vector capacities | per-call selected-batch cleanup | D | A after spill ledger closure | | BucketReader decoded vectors | allocation-accounted MPool Vector data/area | reusable batch cleanup / `BucketReader.Close` | D | A after spill ledger closure | -| `pSpool` cached Vector data/area | raw MPool slices retained and reassigned independently of their Vector | `spoolBuffer.clean` or the receiving Vector's `Free` | L | persist generation/selection provenance for a missing data or area allocation; blocks pipeline activation | +| `pSpool` cached Vector data/area | provenance-bearing detached MPool buffers; accounted data/area sites cannot cross and legacy buffers keep a guarded fast path | `spoolBuffer.clean` or the receiving Vector's `Free` | D | A with the pipeline owner activation | | runtime-filter serialized payload | Go buffer/message payload; O(filter rows) | message release | L | off-heap or PASS degradation; blocks runtime-filter activation | | spill disk and FD | disk/FD ledgers | file removal/close | A | A | @@ -128,11 +128,13 @@ scratch” can declare closure. Batch destination propagation is now `D`: Clone, Dup, selected-column copy, Union destinations, windows, reader decode, Clean, and FreeColumns preserve the immutable destination selection without creating a synthetic batch-level -charge. `pSpool` is deliberately still `L`; its raw buffer cache can retain an -allocation after detaching it from a Vector, so merely copying the Batch -selection would disagree with the original account still recorded in the -MPool lease. `Vector.SetTypeAndFixData` also remains a PR 3 closure blocker -because its legacy API cannot currently return a failed growth admission. +charge. `pSpool` now transfers a detached buffer together with its immutable +selection and data/area site, reuses it only for the same provenance, and +returns every cache ID on construction failure. Its allocation-unaccounted +production path retains the old raw-slice representation behind guards that +reject an accounted Vector. `Vector.SetTypeAndFixData` now returns a failed +growth admission and rolls type and length back without losing the original +backing. ## 3. Decisions required before production integration @@ -579,8 +581,11 @@ Gate: - all migrated rows become `D`; production rows remain `L`. The current dormant PR 3 candidate is -`feature/26459-expression-propagation` at expression commit `03296c5246` and -spill commit `0b4a7b5bb5`, stacked on PR 2 commit `dbfee20ecc`. +`feature/26459-expression-propagation` at commit `1f82b218c2`, stacked on PR 2 +commit `dbfee20ecc`. Its closure commits are `03296c5246` (expressions), +`0b4a7b5bb5` (spill vectors/scratch), `a0754a6beb` (streaming buffers), +`9f88a0e83e` (spill serialization), and `1f82b218c2` (spool/type-change +ownership). Its propagation call chains are: @@ -594,7 +599,13 @@ NewExpressionExecutorWithAllocation NewSpillEngineWithAllocation -> BucketReader decoded/reused Batch selection -> scatter hash/row typed slices and selected Batch selection + -> exact streaming record and optional coalesce buffers -> MPool AllocAccounted/Grow/Free + +Pipeline spool accounted copy + -> detach Vector data/area with immutable provenance + -> cache only by matching selection and data/area site + -> attach to the next owning Vector or free at spool cleanup ``` The candidate adds no production caller of either dormant constructor and @@ -607,6 +618,21 @@ the accounted path; the old and replacement capacities are simultaneously charged until publication, and terminal cleanup frees a zero-length view by its retained capacity. +The spill record path now streams Bitmap, Nulls, Vector, and Batch wire formats +directly into one allocation-accounted off-heap buffer. It computes the exact +wire size before admission, retains one record buffer for the phase, bounds +each bucket's coalesce buffer at 64 KiB, and degrades optional coalescing to a +direct write when payload or metadata capacity rejects it. Wire round trips, +legacy byte equivalence, retained-buffer reuse, coalesce fallback, write +failure, and phase cleanup are covered. + +The pipeline spool now preserves Batch and Vector selection through cached +copies, keeps data and area allocation sites distinct, refuses cross-account +reuse, returns the cache slot after allocation failure, and fixes a legacy +non-last cache removal that previously dropped a still-owned buffer +descriptor. `SetTypeAndFixData` publishes a fixed-width type change only after +growth succeeds and propagates its error through all four callers. + The first syntax inventory over non-test expression/builtin and spill sources found 484 candidate lines. Because one line can match more than one category, the overlapping counts are 233 `make([]...)`, 205 capacity-growing or @@ -617,23 +643,23 @@ in the ledger above. Activation remains blocked by: - Vector null/group bitmap backing; - decimal parameter conversion slices; -- spill marshal and coalesce `bytes.Buffer` backing; - data-scaled function-specific Go-heap scratch identified by the remaining - built-in scan; -- `pSpool` raw-buffer provenance and `Vector.SetTypeAndFixData`. + built-in scan. Fresh local validation on linux/amd64 with CGO and the repository-built third-party artifacts passes: - complete normal tests for MPool, Vector, Batch, SQL util, expression, - spillutil, HashBuild, HashJoin, DedupJoin, RightDedupJoin, and Process; + pSpool, spillutil, HashBuild, HashJoin, DedupJoin, RightDedupJoin, + Connector, Dispatch, Function, and Process; - build and vet for the same package closure; - exact `-race -count=100` runs for every new test plus directly affected flow control, BucketReader merge, scatter lifecycle, and re-spill tests; -- complete race runs for MPool, Vector, Batch, SQL util, expression, and - spillutil; -- package coverage of 73.3% MPool, 47.8% Vector, 74.4% Batch, 26.6% SQL util, - 64.3% expression, and 79.7% spillutil. +- complete race runs for MPool, Vector, Batch, SQL util, expression, spillutil, + pSpool, Connector, and Dispatch; +- package coverage after the new closures of 74.0% MPool, 77.0% Bitmap, 41.2% + Nulls, 48.3% Vector, 73.4% Batch, 82.1% pSpool, 78.8% spillutil, and 54.1% + Function. The existing constant-flow-control benchmark remains 0 B/op and 0 allocs/op. Five-run medians at `GOMAXPROCS=8` were 4.480 ns/op on the PR 2 @@ -641,6 +667,11 @@ base and 4.464 ns/op on the PR 3 candidate; this focused benchmark shows no measurable legacy fast-path regression, but it is not an activation-level performance result. +The legacy 8,192-row pipeline-spool copy/reuse benchmark remains 200 B/op and +2 allocs/op. An interleaved five-run comparison at `GOMAXPROCS=8` measured a +1,542 ns/op median at commit `9f88a0e83e` and 1,540 ns/op at `1f82b218c2`; +the guarded allocation-unaccounted fast path has no measurable regression. + ### PR 4: statement lifecycle and minimum pressure foundation Scope: From 8830b4b780d6cadc5a1d18b6e1c9615d26edb375 Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 16:27:24 +0800 Subject: [PATCH 25/61] docs: record bitmap and conversion accounting closure --- ...ocation_accounted_memory_admission_impl.md | 76 +++++++++++++++---- 1 file changed, 62 insertions(+), 14 deletions(-) diff --git a/docs/design/allocation_accounted_memory_admission_impl.md b/docs/design/allocation_accounted_memory_admission_impl.md index 93a9757c250ef..6398721b2521a 100644 --- a/docs/design/allocation_accounted_memory_admission_impl.md +++ b/docs/design/allocation_accounted_memory_admission_impl.md @@ -95,10 +95,10 @@ working ledger is allocation-site based: | `mpool.memHdr` and account-ID side map | Go maps; one pointer record plus optional account record per live allocation | pointer removal at physical deallocation | H | bounded by the measured finite registry/allocation-slot policy | | `Vector.data` | MPool; capacity from `Grow`, on/off-heap follows `v.offHeap` | owning `Vector.Free` | D | A only when off-heap | | `Vector.area` | MPool; independent varlen payload capacity | owning `Vector.Free` | D | A only when off-heap | -| `Vector.nsp/gsp` bitmap data | Go `[]uint64`; `ceil(rows/64)*8`, retained by `Clear` | bitmap `Reset` from `Vector.Free` | L | move off-heap; blocks Vector-dependent activation | +| `Vector.nsp/gsp` bitmap data | allocation-accounted off-heap `[]uint64`; independent geometric capacity, paired replacement admission | owning `Vector.Free`; `Reset` retains and clears only published words | D | A after the selected Vector owner closes | | `FunctionResult.vec` data/area | off-heap Vector; rows and appended payload | executor `Free` | D | A after expression ledger closure | | `FunctionResult.convenientParam` | Go slice; expression arity, not rows | executor `Free`/reuse | L | H after a proved arity bound | -| decimal parameter conversion | Go `[]T`; `rows*sizeof(T)` in `GenerateFunctionFixedTypeParameter` | evaluation wrapper/GC | L | move off-heap; blocks expression activation | +| decimal parameter conversion | retained allocation-accounted off-heap buffer; `rows*sizeof(Decimal128)` for decimal64/float32/float64 promotion | `FunctionResult.Free`; Reset/evaluation reuse capacity | D | A after expression ledger closure | | IFF/CASE/COALESCE selection arrays | allocation-accounted off-heap `[]bool`; one or two arrays of `rows`, retained by executor | executor `Free`/reuse | D | A after expression ledger closure | | selected row IDs | allocation-accounted off-heap `[]int64`; capacity up to `rows`, retained by executor | executor `Free`/reuse | D | A after expression ledger closure | | selected parameter/result vectors | allocation-accounted off-heap Vector capacities | executor `Free` | D | A after expression ledger closure | @@ -136,6 +136,16 @@ reject an accounted Vector. `Vector.SetTypeAndFixData` now returns a failed growth admission and rolls type and length back without losing the original backing. +Bitmap-aware selections are also `D`. Null and grouping backing remain +independent physical allocations, are included in `Vector.Allocated`, and use +the same immutable account/owner with distinct sites. The legacy `Bitmap` +footprint is unchanged: a tagged nonnegative/bitwise-complemented length +records backing ownership without adding a field. Accounted Vector growth +admits both bitmap replacements before either publishes, raw unadmitted bitmap +growth fails instead of escaping to the Go heap, Reset retains capacity while +clearing only represented words, copy/reader decode fills admitted backing +directly, and Free is the one terminal release owner. + ## 3. Decisions required before production integration ### A. Allocation metadata representation @@ -513,8 +523,8 @@ Required behavior: - within-capacity append performs no account operation; - aliases/views/const/shared area do not create another charge; - deep copies use the destination account; -- on-heap null/group bitmaps remain explicit ledger blockers, not silently - included in the Vector charge; +- at the PR 2 boundary, on-heap null/group bitmaps remain explicit later-PR + blockers and are not silently included in the Vector charge; - HashBuild production remains legacy until a later owner migration. Gate: @@ -581,19 +591,20 @@ Gate: - all migrated rows become `D`; production rows remain `L`. The current dormant PR 3 candidate is -`feature/26459-expression-propagation` at commit `1f82b218c2`, stacked on PR 2 +`feature/26459-expression-propagation` at commit `a63c68b07b`, stacked on PR 2 commit `dbfee20ecc`. Its closure commits are `03296c5246` (expressions), `0b4a7b5bb5` (spill vectors/scratch), `a0754a6beb` (streaming buffers), `9f88a0e83e` (spill serialization), and `1f82b218c2` (spool/type-change -ownership). +ownership), followed by `a63c68b07b` (Vector bitmaps and decimal conversion +scratch). Its propagation call chains are: ```text NewExpressionExecutorWithAllocation -> recursive expression construction - -> constant/result/scratch AllocationAccountSelection - -> FunctionResult or selected/decode Vector growth + -> constant/result/scratch bitmap-aware AllocationAccountSelection + -> FunctionResult result/parameter scratch or selected/decode Vector growth -> MPool AllocAccounted/Grow/Free NewSpillEngineWithAllocation @@ -618,6 +629,22 @@ the accounted path; the old and replacement capacities are simultaneously charged until publication, and terminal cleanup frees a zero-length view by its retained capacity. +Vector null/group bitmap growth now follows the same rule. Both replacement +buffers are admitted before publication, so rejection of the second buffer +rolls the first unpublished buffer back and preserves both old owners. +`pSpool` releases bitmap backing after data/area detach, recreates it under the +destination provenance, and preserves grouping semantics including constant +vectors. Accounted window/duplicate/decode paths pre-admit bitmap coverage +before legacy raw bitmap APIs can mutate it. + +Decimal128 promotion from decimal64, float32, and float64 now writes into one +retained allocation-accounted parameter buffer per argument. Const promotion +uses the scalar wrapper and allocates no row scratch; normal/null promotion +reuses the admitted buffer across evaluations. Capacity, sealed-account, and +metadata failures return through the function executor instead of panicking or +falling back to a Go slice. The allocation selection and parameter buffer must +share the same account and owner. + The spill record path now streams Bitmap, Nulls, Vector, and Batch wire formats directly into one allocation-accounted off-heap buffer. It computes the exact wire size before admission, retains one record buffer for the phase, bounds @@ -639,12 +666,11 @@ the overlapping counts are 233 `make([]...)`, 205 capacity-growing or potentially growing `append`, 44 `bytes.Buffer`/`NewBuffer`, and 5 `strings.Builder` sites. This is a review queue, not proof that every match is data-scaled or reachable. The currently closed rows are the ones marked `D` -in the ledger above. Activation remains blocked by: - -- Vector null/group bitmap backing; -- decimal parameter conversion slices; -- data-scaled function-specific Go-heap scratch identified by the remaining - built-in scan. +in the ledger above. Vector null/group backing and decimal parameter conversion +no longer block activation. Expression activation remains blocked by +data-scaled function-specific Go-heap scratch identified by the remaining +built-in scan; the dormant constructor is not a license to enable partial +accounting. Fresh local validation on linux/amd64 with CGO and the repository-built third-party artifacts passes: @@ -661,6 +687,13 @@ third-party artifacts passes: Nulls, 48.3% Vector, 73.4% Batch, 82.1% pSpool, 78.8% spillutil, and 54.1% Function. +For `a63c68b07b`, every new bitmap, buffer, Vector, decimal-conversion, spool, +spill, and directly affected aggregate test passed an exact +`-race -count=100` run after the final edit. Bitmap, MPool, Vector, pSpool, +spillutil, aggregate, expression, and Function packages then passed complete +`-race -count=1` runs; the full affected normal-test and vet closure also +passed with repository-built CGO third parties. + The existing constant-flow-control benchmark remains 0 B/op and 0 allocs/op. Five-run medians at `GOMAXPROCS=8` were 4.480 ns/op on the PR 2 base and 4.464 ns/op on the PR 3 candidate; this focused benchmark shows no @@ -672,6 +705,21 @@ The legacy 8,192-row pipeline-spool copy/reuse benchmark remains 200 B/op and 1,542 ns/op median at commit `9f88a0e83e` and 1,540 ns/op at `1f82b218c2`; the guarded allocation-unaccounted fast path has no measurable regression. +At `a63c68b07b`, the same legacy pipeline-spool benchmark remains 200 B/op and +2 allocs/op with a 1,545 ns/op five-run median. The constant flow-control +benchmark remains 0 B/op and 0 allocs/op with a 4.259 ns/op median. The tagged +bitmap ownership representation is what preserves the legacy allocation +footprint. + +An 8,192-row first pre-extend/free costs 1,059 ns/op for the legacy Vector, +1,159 ns/op for data-only dormant accounting, and 1,986 ns/op when the two +independent bitmap allocations are also admitted; all three remain 0 B/op and +0 allocs/op. This one-time physical-allocation cost is explicit rather than +hidden. Capacity reuse is the steady-state path: empty Reset measured +2.431 ns/op for data-only accounting and 2.650 ns/op with bitmap ownership, +again with zero Go allocations. Activation performance tests must retain this +distinction instead of presenting the first-allocation cost as a per-row cost. + ### PR 4: statement lifecycle and minimum pressure foundation Scope: From 7f58b954890cfe0b3ed2d3a3e64d1c024ca86ff5 Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 16:31:54 +0800 Subject: [PATCH 26/61] docs: record direct-output scratch removal --- .../allocation_accounted_memory_admission_impl.md | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/docs/design/allocation_accounted_memory_admission_impl.md b/docs/design/allocation_accounted_memory_admission_impl.md index 6398721b2521a..02483dcbb9789 100644 --- a/docs/design/allocation_accounted_memory_admission_impl.md +++ b/docs/design/allocation_accounted_memory_admission_impl.md @@ -591,12 +591,12 @@ Gate: - all migrated rows become `D`; production rows remain `L`. The current dormant PR 3 candidate is -`feature/26459-expression-propagation` at commit `a63c68b07b`, stacked on PR 2 +`feature/26459-expression-propagation` at commit `0fd87ffebe`, stacked on PR 2 commit `dbfee20ecc`. Its closure commits are `03296c5246` (expressions), `0b4a7b5bb5` (spill vectors/scratch), `a0754a6beb` (streaming buffers), `9f88a0e83e` (spill serialization), and `1f82b218c2` (spool/type-change ownership), followed by `a63c68b07b` (Vector bitmaps and decimal conversion -scratch). +scratch) and `0fd87ffebe` (direct-output Field/VALUES paths). Its propagation call chains are: @@ -672,6 +672,14 @@ data-scaled function-specific Go-heap scratch identified by the remaining built-in scan; the dormant constructor is not a license to enable partial accounting. +The first follow-up scan also removed two allocations instead of moving them: +`FIELD` now writes directly into its pre-extended result and retains +arity-bounded parameter wrappers, while `VALUES` uses contiguous `UnionBatch` +instead of constructing one row ID per input row. Their exact tests passed +`-race -count=100` and the complete Function package passed normal, vet, and +race runs. Remaining candidates must first make this same +direct-output/streaming test before introducing a scratch owner. + Fresh local validation on linux/amd64 with CGO and the repository-built third-party artifacts passes: From 6af8b82d41a052336039ad3aa24548764fd1a34e Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 17:40:37 +0800 Subject: [PATCH 27/61] docs: record direct function scratch closure --- ...ocation_accounted_memory_admission_impl.md | 65 ++++++++++++++++++- 1 file changed, 63 insertions(+), 2 deletions(-) diff --git a/docs/design/allocation_accounted_memory_admission_impl.md b/docs/design/allocation_accounted_memory_admission_impl.md index 02483dcbb9789..45218a19c825a 100644 --- a/docs/design/allocation_accounted_memory_admission_impl.md +++ b/docs/design/allocation_accounted_memory_admission_impl.md @@ -102,6 +102,14 @@ working ledger is allocation-site based: | IFF/CASE/COALESCE selection arrays | allocation-accounted off-heap `[]bool`; one or two arrays of `rows`, retained by executor | executor `Free`/reuse | D | A after expression ledger closure | | selected row IDs | allocation-accounted off-heap `[]int64`; capacity up to `rows`, retained by executor | executor `Free`/reuse | D | A after expression ledger closure | | selected parameter/result vectors | allocation-accounted off-heap Vector capacities | executor `Free` | D | A after expression ledger closure | +| `JSON_ROW` per-row encoders and buffers | one Go `bytes.Buffer` per output row, each retaining its row payload | function operator reuse / GC | R | one reusable row encoder now streams directly to the accounted result; parameter closures are arity-bounded and cleared after every call | +| float array-distance row descriptors and result scratch | Go `[][]T` with one descriptor per row plus `[]float32` with one value per row | call return / GC | R | caller row accessor plus the upper half of the admitted `[]float64` result backing | +| GPU array-distance flattened inputs | C allocator; `query bytes + rows*dimension*4` for the SQL float32 GPU path | GPU job Wait or launch rollback | L | external allocation remains data-scaled and blocks expression activation until it is admitted or the activated owner excludes GPU dispatch | +| `NORMALIZE_L2` output scratch | pooled or per-row Go arrays with one output element per input element | pool / call return and GC | R | normalize directly into admitted varlen result storage | +| AES/ENCODE/DECODE output scratch | Go payload slices, one plaintext/ciphertext-sized allocation per row; AES padding could append into spare input capacity | call return / GC | R | validate size/padding first and write directly into admitted result storage | +| JQ object sort keys and gojq value graph | Go slices/maps/interfaces proportional to input JSON structure | row completion / GC | L | requires a separately bounded or allocation-accounted JSON execution owner | +| JSON typed-array/value builders and pretty-print buffer | Go `[]any`, ByteJson payloads, and `bytes.Buffer` proportional to row input/output | row completion / GC | L | direct ByteJson/result builders or an admitted reusable JSON scratch owner | +| geometry parse/overlay scratch | Go point/ring/interval/match slices proportional to geometry payload | row completion / GC | L | admitted per-row geometry workspace or streaming algorithm; remains an activation blocker | | hash-table initial cell block | off-heap `mpool.MakeSlice(..., true)`; 16 KiB int / 32 KiB string | hash map / `JoinMap.FreeMemory` | L | A in first activation | | hash-table replacement/appended cell blocks | off-heap blocks, at most 4 MiB each; old+new overlap is physically visible | hash map / `JoinMap.FreeMemory` | L | A in first activation | | hash-table `cells`/`newBlocks` descriptors | Go `[][]Cell`; 24 bytes per block header plus geometric resize backing arrays; GC is not treated as synchronous Free | hash map lifetime / GC | L | replace with an owning off-heap descriptor buffer and account its initial/replacement capacity; blocks first activation | @@ -591,12 +599,13 @@ Gate: - all migrated rows become `D`; production rows remain `L`. The current dormant PR 3 candidate is -`feature/26459-expression-propagation` at commit `0fd87ffebe`, stacked on PR 2 +`feature/26459-expression-propagation` at commit `ad71c4fb50`, stacked on PR 2 commit `dbfee20ecc`. Its closure commits are `03296c5246` (expressions), `0b4a7b5bb5` (spill vectors/scratch), `a0754a6beb` (streaming buffers), `9f88a0e83e` (spill serialization), and `1f82b218c2` (spool/type-change ownership), followed by `a63c68b07b` (Vector bitmaps and decimal conversion -scratch) and `0fd87ffebe` (direct-output Field/VALUES paths). +scratch), `0fd87ffebe` (direct-output Field/VALUES paths), and `ad71c4fb50` +(direct-output function scratch). Its propagation call chains are: @@ -680,6 +689,58 @@ instead of constructing one row ID per input row. Their exact tests passed race runs. Remaining candidates must first make this same direct-output/streaming test before introducing a scratch owner. +The second follow-up removes rather than accounts four more scratch families: + +- `JSON_ROW` no longer owns one encoder and buffer per row. It prepares + arity-bounded typed column closures once, streams one complete row through a + single reusable encoder, appends it to the admitted result, and clears every + closure on success, error, or panic. The unsigned path now preserves the full + `uint64` range. +- float array distance no longer materializes one `[]T` descriptor per input + row or a second result slice. The metric layer accepts a synchronous row + accessor and caller-owned output; SQL uses the upper half of its already + admitted `[]float64` result bytes as `[]float32` scratch and converts forward + only after Wait. GPU launch still copies flattened inputs into the existing C + allocator and therefore remains a named activation blocker. +- `NORMALIZE_L2` writes float32, float64, BF16, Float16, int8, and uint8 results + directly into admitted varlen storage. The old float pools and per-row + widened/narrowed arrays are deleted. +- AES ECB/CBC and legacy ENCODE/DECODE write directly into admitted result + storage. AES validates the final decrypted block before result publication, + and padding is assembled in one stack block, removing both payload copies + and the old possibility that `append` modified spare input-vector capacity. + +`FunctionResult.AppendBytesWithFill` is the shared publication primitive. It +admits result data/area capacity before exposing a row, publishes length only +after the callback finishes, and restores the unpublished varlena header and +area length on panic. Capacity acquired before a callback panic remains owned +and charged to the result until normal reuse or `Free`; no allocation is +released without its terminal owner. + +Fresh five-run benchmarks on linux/amd64, Go 1.26.4, i7-11700: + +| 8,192-row function path | Before median | `ad71c4fb50` median | Before allocation | `ad71c4fb50` allocation | +| --- | ---: | ---: | ---: | ---: | +| float32 L2-squared batch distance, 128 dimensions | about 483 us | 429.5 us | 229,488 B/op, 6 allocs/op | 112 B/op, 4 allocs/op | +| `JSON_ROW`, fresh operator | about 838 us | 449.6 us | 1,704,081 B/op, 16,387 allocs/op | 440 B/op, 7 allocs/op | +| `JSON_ROW`, reused operator | about 474.7 us | 450.9 us | about 831 B/op, 8 allocs/op | 200 B/op, 4 allocs/op | + +The direct-output candidate passed each new or directly affected behavioral +test separately under `-race -count=100` with the default 30-second adaptive +budget (the measured exact-test durations were 0 or 0.01 seconds, so the +100-run cap applied). Vector, metric, Function, and colexec passed complete +normal and race runs; the same package closure passed build and vet. The +unsafe result alias has its own functional test, and callback panic rollback +proves result length, area length, reuse, and final account zero. + +The GPU source path passed manual ownership review: dimension or allocation +failure releases every C buffer locally, successful launch transfers both +buffers to the job, and Wait deallocates them after the asynchronous kernel +terminates. This host has neither `nvcc` nor a CUDA conda environment, so the +GPU-tag build and tests were not run locally. That validation gap and the +data-scaled C allocation both remain explicit merge/activation gates; CPU +success is not treated as GPU evidence. + Fresh local validation on linux/amd64 with CGO and the repository-built third-party artifacts passes: From 11f8fc3c8f5e81ba441f087c1c2dc9b6731c30d9 Mon Sep 17 00:00:00 2001 From: aptend Date: Fri, 31 Jul 2026 19:31:34 +0800 Subject: [PATCH 28/61] docs: update allocation closure implementation evidence --- ...ocation_accounted_memory_admission_impl.md | 164 ++++++++++++------ 1 file changed, 112 insertions(+), 52 deletions(-) diff --git a/docs/design/allocation_accounted_memory_admission_impl.md b/docs/design/allocation_accounted_memory_admission_impl.md index 45218a19c825a..76549cb12e50f 100644 --- a/docs/design/allocation_accounted_memory_admission_impl.md +++ b/docs/design/allocation_accounted_memory_admission_impl.md @@ -6,7 +6,8 @@ - Architecture: [Allocation-Accounted Memory Admission RFC](../rfcs/00000000_allocation_accounted_memory_admission.md) - Baseline at plan creation: `main` at `38ce3a774` -- Rebased implementation baseline: `main` at `43c896462` +- Rebased implementation baseline: `main` at `60e36bef64` +- Current dormant PR 3 head: `fd2fcc953b` - Merged prerequisite: #26455 at `93e8b22d2` - Independent design review: completed against RFC commit `a7d54cb5f` - Activation status: blocked until PRs 1--4 and the selected owner's @@ -102,14 +103,21 @@ working ledger is allocation-site based: | IFF/CASE/COALESCE selection arrays | allocation-accounted off-heap `[]bool`; one or two arrays of `rows`, retained by executor | executor `Free`/reuse | D | A after expression ledger closure | | selected row IDs | allocation-accounted off-heap `[]int64`; capacity up to `rows`, retained by executor | executor `Free`/reuse | D | A after expression ledger closure | | selected parameter/result vectors | allocation-accounted off-heap Vector capacities | executor `Free` | D | A after expression ledger closure | -| `JSON_ROW` per-row encoders and buffers | one Go `bytes.Buffer` per output row, each retaining its row payload | function operator reuse / GC | R | one reusable row encoder now streams directly to the accounted result; parameter closures are arity-bounded and cleared after every call | +| `JSON_ROW` output | one retained allocation-accounted function scratch buffer, then copy into the admitted result; both physical capacities are charged | `FunctionResult.Free`; row publication copies synchronously | D | A after expression-owner activation; arity-bounded column closures are cleared after every call | | float array-distance row descriptors and result scratch | Go `[][]T` with one descriptor per row plus `[]float32` with one value per row | call return / GC | R | caller row accessor plus the upper half of the admitted `[]float64` result backing | -| GPU array-distance flattened inputs | C allocator; `query bytes + rows*dimension*4` for the SQL float32 GPU path | GPU job Wait or launch rollback | L | external allocation remains data-scaled and blocks expression activation until it is admitted or the activated owner excludes GPU dispatch | +| GPU array-distance flattened inputs | allocation-accounted caller scratch; `query bytes + rows*dimension*4` for the SQL float32 GPU path; legacy callers retain the C allocator | GPU job retains the caller slice until Wait; launch failure rolls back before return | D | A only after the GPU-tag build/tests pass; the legacy API remains unaccounted and is not an activated owner | | `NORMALIZE_L2` output scratch | pooled or per-row Go arrays with one output element per input element | pool / call return and GC | R | normalize directly into admitted varlen result storage | | AES/ENCODE/DECODE output scratch | Go payload slices, one plaintext/ciphertext-sized allocation per row; AES padding could append into spare input capacity | call return / GC | R | validate size/padding first and write directly into admitted result storage | -| JQ object sort keys and gojq value graph | Go slices/maps/interfaces proportional to input JSON structure | row completion / GC | L | requires a separately bounded or allocation-accounted JSON execution owner | -| JSON typed-array/value builders and pretty-print buffer | Go `[]any`, ByteJson payloads, and `bytes.Buffer` proportional to row input/output | row completion / GC | L | direct ByteJson/result builders or an admitted reusable JSON scratch owner | +| JQ visible output | one retained allocation-accounted function scratch buffer, then copy into the admitted result | `FunctionResult.Free`; row publication copies synchronously | D | A after expression-owner activation; `TRY_JQ` propagates infrastructure/account rejection instead of converting it to SQL NULL | +| JQ object sort keys and gojq value graph | Go slices/maps/interfaces proportional to input JSON structure | row completion / GC | L | requires a separately bounded or allocation-accounted JSON execution owner; blocks JQ activation | +| JSON_ARRAY/OBJECT/KEYS/PRETTY output and object-key scratch | direct storage-compatible ByteJSON/result builders; JSON_OBJECT keys use retained allocation-accounted function scratch | result or `FunctionResult.Free` | D | A after expression-owner activation | +| JSON parse/path/modify/merge/schema execution | Go strings, paths, maps/slices, ByteJson modification payloads, and schema graphs proportional to row input | row completion / GC | L | requires an admitted JSON execution arena or streaming algorithms; blocks these functions from activation | +| HASH/IN/PREFIX_IN/narrow-array conversion scratch | retained allocation-accounted function scratch, bounded per hash chunk or exact tuple/array bytes | `FunctionResult.Free` | D | A after expression-owner activation | +| codec output (HEX/base64/MD5/COMPRESS/UNCOMPRESS/random/quote/date formatting) | writes directly into admitted result backing; flate's fixed codec state remains Go/runtime memory | result Free; codec state ends at call return | D | A after result activation; flate state needs an explicit fixed-headroom measurement before H classification | | geometry parse/overlay scratch | Go point/ring/interval/match slices proportional to geometry payload | row completion / GC | L | admitted per-row geometry workspace or streaming algorithm; remains an activation blocker | +| regexp compile/match/output scratch | Go regexp program and match/output buffers proportional to pattern/input | operator cache / row completion / GC | L | bounded regexp owner or exclusion from expression activation | +| H3/S2 neighborhood scratch | Go slices; some S2 paths are statically bounded, H3 grid-disk output scales with radius | row completion / GC | L | split proved fixed bounds from data-scaled paths before activation | +| JSON cast visible serialization | `MarshalJSON` payload slices proportional to the JSON value | row completion / GC | L | direct visible writer or exclusion from expression activation | | hash-table initial cell block | off-heap `mpool.MakeSlice(..., true)`; 16 KiB int / 32 KiB string | hash map / `JoinMap.FreeMemory` | L | A in first activation | | hash-table replacement/appended cell blocks | off-heap blocks, at most 4 MiB each; old+new overlap is physically visible | hash map / `JoinMap.FreeMemory` | L | A in first activation | | hash-table `cells`/`newBlocks` descriptors | Go `[][]Cell`; 24 bytes per block header plus geometric resize backing arrays; GC is not treated as synchronous Free | hash map lifetime / GC | L | replace with an owning off-heap descriptor buffer and account its initial/replacement capacity; blocks first activation | @@ -599,13 +607,21 @@ Gate: - all migrated rows become `D`; production rows remain `L`. The current dormant PR 3 candidate is -`feature/26459-expression-propagation` at commit `ad71c4fb50`, stacked on PR 2 -commit `dbfee20ecc`. Its closure commits are `03296c5246` (expressions), -`0b4a7b5bb5` (spill vectors/scratch), `a0754a6beb` (streaming buffers), -`9f88a0e83e` (spill serialization), and `1f82b218c2` (spool/type-change -ownership), followed by `a63c68b07b` (Vector bitmaps and decimal conversion -scratch), `0fd87ffebe` (direct-output Field/VALUES paths), and `ad71c4fb50` -(direct-output function scratch). +`feature/26459-expression-propagation` at commit `fd2fcc953b`, rebased on +`main` commit `60e36bef64`. Its ordered closure commits are: + +- `06d39287e4`: generic allocation-accounted MPool ownership; +- `968c68df4c`: dormant HashBuild budget bridge; +- `941212edf0`: Vector propagation; +- `61a6e7a0e2`: expression propagation; +- `b03ec24c2f`, `309d070c71`, and `556d64d5ba`: spill scratch, streaming, + and serialization; +- `86e041ff53`: retained Vector/spool ownership gaps; +- `2b05fdac79`: Vector bitmap and conversion scratch; +- `24b5f6fef7` and `57940c8f0f`: direct-output row/function scratch removal; +- `674f78dea9`: general retained function-scratch provenance; +- `fd2fcc953b`: direct codecs, HASH/IN/PREFIX, GPU SQL flattening, and + ByteJSON/JQ/JSON output closure. Its propagation call chains are: @@ -613,7 +629,7 @@ Its propagation call chains are: NewExpressionExecutorWithAllocation -> recursive expression construction -> constant/result/scratch bitmap-aware AllocationAccountSelection - -> FunctionResult result/parameter scratch or selected/decode Vector growth + -> FunctionResult result/parameter/function scratch or selected/decode Vector growth -> MPool AllocAccounted/Grow/Free NewSpillEngineWithAllocation @@ -670,16 +686,21 @@ descriptor. `SetTypeAndFixData` publishes a fixed-width type change only after growth succeeds and propagates its error through all four callers. The first syntax inventory over non-test expression/builtin and spill sources -found 484 candidate lines. Because one line can match more than one category, -the overlapping counts are 233 `make([]...)`, 205 capacity-growing or -potentially growing `append`, 44 `bytes.Buffer`/`NewBuffer`, and 5 -`strings.Builder` sites. This is a review queue, not proof that every match is -data-scaled or reachable. The currently closed rows are the ones marked `D` -in the ledger above. Vector null/group backing and decimal parameter conversion -no longer block activation. Expression activation remains blocked by -data-scaled function-specific Go-heap scratch identified by the remaining -built-in scan; the dormant constructor is not a license to enable partial -accounting. +found 484 candidate lines. The latest focused function scan still reports +candidate syntax, not live-byte proof: the largest `make([]...)` concentrations +are `func_unary.go` and `func_binary.go` (28 each) and +`func_builtin_json.go` (26); the largest potentially growing `append` groups +are `func_unary.go` (37), `func_builtin.go` (36), `func_cast.go` (16), and +`func_builtin_json.go`/`func_binary.go` (13 each). Buffer/builder and JSON +serialization scans are recorded in the site rows above. Admin/CTL/UDF code, +arity-only slices, fixed stack buffers, and legacy-only branches are not +silently counted as controlled payload, but each must be explicitly excluded +or bounded before activation. Vector null/group backing, decimal conversion, +direct codecs, HASH/IN/PREFIX scratch, SQL GPU flattening, and visible +JSON/JQ output no longer create data-scaled unaccounted payload on the dormant +path. JQ graphs, JSON parse/modify/schema, geometry, regexp, scalable H3, and +JSON cast serialization remain explicit activation blockers; the dormant +constructor is not a license to enable partial accounting. The first follow-up scan also removed two allocations instead of moving them: `FIELD` now writes directly into its pre-extended result and retains @@ -693,15 +714,20 @@ The second follow-up removes rather than accounts four more scratch families: - `JSON_ROW` no longer owns one encoder and buffer per row. It prepares arity-bounded typed column closures once, streams one complete row through a - single reusable encoder, appends it to the admitted result, and clears every - closure on success, error, or panic. The unsigned path now preserves the full - `uint64` range. + single reusable encoder, and clears every closure on success, error, or + panic. Legacy execution reuses one `bytes.Buffer`; dormant exact execution + uses one retained allocation-accounted function buffer and then copies the + completed row into the admitted result. The scratch and result capacities + are both physical and therefore both charged. The unsigned path preserves + the full `uint64` range. - float array distance no longer materializes one `[]T` descriptor per input row or a second result slice. The metric layer accepts a synchronous row accessor and caller-owned output; SQL uses the upper half of its already admitted `[]float64` result bytes as `[]float32` scratch and converts forward - only after Wait. GPU launch still copies flattened inputs into the existing C - allocator and therefore remains a named activation blocker. + only after Wait. The SQL float32 GPU path now flattens query and row inputs + into caller-owned allocation-accounted function scratch, and the GPU job + retains that slice until Wait. The public legacy GPU API keeps its C-allocator + behavior and is not part of the activated owner. - `NORMALIZE_L2` writes float32, float64, BF16, Float16, int8, and uint8 results directly into admitted varlen storage. The old float pools and per-row widened/narrowed arrays are deleted. @@ -710,36 +736,70 @@ The second follow-up removes rather than accounts four more scratch families: and padding is assembled in one stack block, removing both payload copies and the old possibility that `append` modified spare input-vector capacity. -`FunctionResult.AppendBytesWithFill` is the shared publication primitive. It -admits result data/area capacity before exposing a row, publishes length only -after the callback finishes, and restores the unpublished varlena header and -area length on panic. Capacity acquired before a callback panic remains owned -and charged to the result until normal reuse or `Free`; no allocation is -released without its terminal owner. +The final dormant closure adds one retained `FunctionResult` scratch owner with +a site distinct from decimal parameter conversion. Detection is allocation-free; +the first nonzero resize admits physical capacity, geometric growth charges old +and replacement capacity until publication, Reset retains it, and +`FunctionResult.Free` is the sole terminal release. It is used for exact HASH +key chunks, fixed/string IN tuples, PREFIX_IN entries, narrow array conversion, +JSON_OBJECT keys, JQ/JSON_ROW visible output, JSON modify value +pre-materialization, and SQL float32 GPU flattening. One-byte-short tests cover +the generic allocator boundaries, and capacity-rejection tests cover each +new scalable function family. `TRY_JQ` suppresses jq/domain errors but never an output +allocation failure. + +Storage-compatible ByteJSON encoders now write scalar, array, object, +object-key-array, typed, opaque, bit, and nested values directly into Vector +area backing. JSON_ARRAY, JSON_OBJECT, JSON_KEYS, JSON_PRETTY, JSON_QUOTE, and +JSON_ROW no longer build a second output-sized ByteJSON or visible-text slice. +JSON_EXTRACT no longer retains `rows * path-count` path arrays: nonconstant +paths are parsed one row at a time into arity-sized reusable storage. This does +not close the parser/modifier/schema graphs themselves; those remain `L`. + +HEX/UNHEX, MD5, base64, vector-base64, COMPRESS/UNCOMPRESS, RANDOM_BYTES, +CHAR, MAKE_SET, EXPORT_SET, bitwise strings, array casts, quote, date/time +formatting, and FROM_UNIXTIME now publish directly into admitted result +capacity. The legacy production path for two-pass-capable formatting remains +one-pass; exact two-pass sizing is selected only by the dormant allocation +owner, so this PR does not impose duplicate formatting work before activation. + +`FunctionResult.AppendBytesWithBuilder` is the shared publication primitive; +`AppendBytesWithFill` is its exact-size convenience wrapper. It admits result +data/area capacity before exposing a row, accepts an actual length no larger +than the admitted capacity, collapses a short result into inline storage when +possible, and publishes length only after successful completion. Error, +invalid-length, and panic paths restore the unpublished varlena header and area +length. Capacity acquired before rollback remains owned and charged to the +result until normal reuse or `Free`; no allocation is released without its +terminal owner. Fresh five-run benchmarks on linux/amd64, Go 1.26.4, i7-11700: -| 8,192-row function path | Before median | `ad71c4fb50` median | Before allocation | `ad71c4fb50` allocation | +| 8,192-row function path | Before median | `fd2fcc953b` median | Before allocation | `fd2fcc953b` allocation | | --- | ---: | ---: | ---: | ---: | -| float32 L2-squared batch distance, 128 dimensions | about 483 us | 429.5 us | 229,488 B/op, 6 allocs/op | 112 B/op, 4 allocs/op | -| `JSON_ROW`, fresh operator | about 838 us | 449.6 us | 1,704,081 B/op, 16,387 allocs/op | 440 B/op, 7 allocs/op | -| `JSON_ROW`, reused operator | about 474.7 us | 450.9 us | about 831 B/op, 8 allocs/op | 200 B/op, 4 allocs/op | +| float32 L2-squared batch distance, 128 dimensions | about 483 us | 437.8 us | 229,488 B/op, 6 allocs/op | 112 B/op, 4 allocs/op | +| `JSON_ROW`, fresh operator | about 838 us | 472.7 us | 1,704,081 B/op, 16,387 allocs/op | 520 B/op, 8 allocs/op | +| `JSON_ROW`, reused operator | about 474.7 us | 471.4 us | about 831 B/op, 8 allocs/op | 264 B/op, 5 allocs/op | The direct-output candidate passed each new or directly affected behavioral -test separately under `-race -count=100` with the default 30-second adaptive -budget (the measured exact-test durations were 0 or 0.01 seconds, so the -100-run cap applied). Vector, metric, Function, and colexec passed complete -normal and race runs; the same package closure passed build and vet. The -unsafe result alias has its own functional test, and callback panic rollback -proves result length, area length, reuse, and final account zero. - -The GPU source path passed manual ownership review: dimension or allocation -failure releases every C buffer locally, successful launch transfers both -buffers to the job, and Wait deallocates them after the asynchronous kernel -terminates. This host has neither `nvcc` nor a CUDA conda environment, so the -GPU-tag build and tests were not run locally. That validation gap and the -data-scaled C allocation both remain explicit merge/activation gates; CPU -success is not treated as GPU evidence. +test separately under race with a 30-second adaptive budget. Every measured +test used `-count=100` except the 0.35-second incompressible DEFLATE-bound test, +which used `-count=85`; no empty regex was accepted as stress evidence. +ByteJSON, Vector, metric, Function, and colexec passed complete normal and race +runs; the same package closure passed vet and `golangci-lint` with zero issues. +Coverage for that closure is 78.3% ByteJSON, 50.9% Vector, 90.3% metric, 54.8% +Function, and 64.3% colexec. The unsafe result alias has its own functional +test, and error/invalid-length/panic rollback proves result length, area +length, reuse, and final account zero. + +The GPU source path passed manual ownership review: exact SQL execution sizes +and fills caller-owned scratch before launch, a successful launch transfers a +retained slice reference to the job, Wait drops it only after the asynchronous +kernel terminates, and dimension/launch failure publishes no job owner. The +legacy nil-scratch API keeps its existing C-buffer allocate/rollback/Wait +contract. This host has neither `nvcc` nor a CUDA conda environment, so the +GPU-tag build and tests were not run locally. That validation gap remains an +explicit merge/activation gate; CPU success is not treated as GPU evidence. Fresh local validation on linux/amd64 with CGO and the repository-built third-party artifacts passes: From ba0e217df92be27ede5d6a2f4e86cc8c887b39ee Mon Sep 17 00:00:00 2001 From: aptend Date: Sat, 1 Aug 2026 00:17:41 +0800 Subject: [PATCH 29/61] executor: make allocation activation owner-atomic --- pkg/container/vector/allocation_account.go | 2 +- .../dedupjoin/expression_memory_test.go | 10 ++ pkg/sql/colexec/dedupjoin/types.go | 14 +- pkg/sql/colexec/evalExpression.go | 6 +- pkg/sql/colexec/eval_expression_allocation.go | 4 +- pkg/sql/colexec/hashbuild/build_test.go | 2 +- .../colexec/hashbuild/expression_memory.go | 3 + .../hashbuild/expression_memory_test.go | 22 +++ pkg/sql/colexec/hashbuild/spill.go | 15 -- pkg/sql/colexec/hashbuild/spill_test.go | 84 +++------- pkg/sql/colexec/hashbuild/types.go | 14 +- .../hashjoin/expression_memory_test.go | 10 ++ pkg/sql/colexec/hashjoin/types.go | 14 +- .../rightdedupjoin/expression_memory_test.go | 10 ++ pkg/sql/colexec/rightdedupjoin/types.go | 16 +- .../colexec/spillutil/allocation_account.go | 2 +- pkg/sql/colexec/spillutil/join_spill.go | 149 ++---------------- pkg/sql/colexec/spillutil/join_spill_test.go | 144 +++++++++++++++++ .../compile/allocation_account_lifecycle.go | 39 +++-- .../allocation_account_lifecycle_test.go | 35 ++++ pkg/sql/compile/compile.go | 6 +- pkg/sql/compile/compile2.go | 18 +-- 22 files changed, 364 insertions(+), 255 deletions(-) diff --git a/pkg/container/vector/allocation_account.go b/pkg/container/vector/allocation_account.go index 057c0910e57d8..43ee95db45df7 100644 --- a/pkg/container/vector/allocation_account.go +++ b/pkg/container/vector/allocation_account.go @@ -139,7 +139,7 @@ func NewOffHeapVecWithTypeAndAllocation( return vec, nil } -// NewConstNullWithAllocation constructs a constant NULL Vector with dormant +// NewConstNullWithAllocation constructs a constant NULL Vector with explicit // allocation provenance for any future owned backing. func NewConstNullWithAllocation( typ types.Type, diff --git a/pkg/sql/colexec/dedupjoin/expression_memory_test.go b/pkg/sql/colexec/dedupjoin/expression_memory_test.go index 7f053f871558f..535a32856f68f 100644 --- a/pkg/sql/colexec/dedupjoin/expression_memory_test.go +++ b/pkg/sql/colexec/dedupjoin/expression_memory_test.go @@ -96,3 +96,13 @@ func TestDedupJoinResetReleasesAccountedProbeExpressions(t *testing.T) { require.NoError(t, err) require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) } + +func TestDedupJoinAllocationActivationRequiresBothKeySides(t *testing.T) { + col := &plan.Expr{Typ: plan.Type{Id: int32(types.T_int64)}, Expr: &plan.Expr_Col{Col: &plan.ColRef{}}} + arg := &DedupJoin{Conditions: [][]*plan.Expr{{col}, {col}}} + require.True(t, arg.AllocationAccountEnabled()) + require.False(t, arg.AllocationAccountActivationBlocked()) + arg.Conditions[0] = []*plan.Expr{nil} + require.False(t, arg.AllocationAccountEnabled()) + require.True(t, arg.AllocationAccountActivationBlocked()) +} diff --git a/pkg/sql/colexec/dedupjoin/types.go b/pkg/sql/colexec/dedupjoin/types.go index 272c7686cb99c..f1310fbe6cf10 100644 --- a/pkg/sql/colexec/dedupjoin/types.go +++ b/pkg/sql/colexec/dedupjoin/types.go @@ -306,7 +306,19 @@ type DedupJoin struct { } func (dedupJoin *DedupJoin) AllocationAccountEnabled() bool { - return dedupJoin != nil + return dedupJoin != nil && dedupJoin.allocationAccountExpressionOwnerClosed() +} + +func (dedupJoin *DedupJoin) AllocationAccountActivationBlocked() bool { + return dedupJoin != nil && !dedupJoin.allocationAccountExpressionOwnerClosed() +} + +func (dedupJoin *DedupJoin) allocationAccountExpressionOwnerClosed() bool { + if dedupJoin == nil || len(dedupJoin.Conditions) != 2 { + return false + } + return hashbuild.AllocationAccountedExpressionSetSupported(dedupJoin.Conditions[0]) && + hashbuild.AllocationAccountedExpressionSetSupported(dedupJoin.Conditions[1]) } func (dedupJoin *DedupJoin) SetAllocationAccount( diff --git a/pkg/sql/colexec/evalExpression.go b/pkg/sql/colexec/evalExpression.go index 1ae0acb4f2ff6..3415c205bbbf4 100644 --- a/pkg/sql/colexec/evalExpression.go +++ b/pkg/sql/colexec/evalExpression.go @@ -95,9 +95,9 @@ func NewExpressionExecutorsFromPlanExpressions(proc *process.Process, planExprs return newExpressionExecutorsFromPlanExpressions(proc, planExprs, nil) } -// NewExpressionExecutorsFromPlanExpressionsWithAllocation constructs dormant -// allocation-accounted expression trees. No production caller uses this path -// until the complete expression allocation-site ledger is closed. +// NewExpressionExecutorsFromPlanExpressionsWithAllocation constructs +// allocation-accounted expression trees for callers that have already proved +// the complete expression allocation-site ledger closed. func NewExpressionExecutorsFromPlanExpressionsWithAllocation( proc *process.Process, planExprs []*plan.Expr, diff --git a/pkg/sql/colexec/eval_expression_allocation.go b/pkg/sql/colexec/eval_expression_allocation.go index aa65a8fbabd2e..ca31d0830658f 100644 --- a/pkg/sql/colexec/eval_expression_allocation.go +++ b/pkg/sql/colexec/eval_expression_allocation.go @@ -23,8 +23,8 @@ import ( ) // Expression allocation sites are stable diagnostics within the owner chosen -// by the caller. The API is dormant: legacy expression constructors do not -// create or select an account. +// by the caller. Legacy expression constructors still do not create or select +// an account; an activated owner must opt in explicitly. const ( ExpressionAllocationSiteConstantData mpool.AllocationSite = iota + 1 ExpressionAllocationSiteConstantArea diff --git a/pkg/sql/colexec/hashbuild/build_test.go b/pkg/sql/colexec/hashbuild/build_test.go index cb8695660b846..748d811dc4149 100644 --- a/pkg/sql/colexec/hashbuild/build_test.go +++ b/pkg/sql/colexec/hashbuild/build_test.go @@ -3155,7 +3155,7 @@ func TestShuffleHashBuildSpillFailureReleasesEmergencyResources(t *testing.T) { require.NotErrorIs(t, buildErr, process.ErrHashBuildBudgetAdmission) require.Nil(t, tc.arg.ctr.spillScratchReservation) require.Zero(t, cap(tc.arg.ctr.spillHashValues)) - require.Zero(t, cap(tc.arg.ctr.spillSelection)) + require.Zero(t, cap(tc.arg.ctr.spillBucketRowIds)) require.Zero(t, cap(tc.arg.ctr.spillKeyVecs)) require.Zero(t, tc.arg.ctr.spillWriteBuf.Cap()) diff --git a/pkg/sql/colexec/hashbuild/expression_memory.go b/pkg/sql/colexec/hashbuild/expression_memory.go index 07f39df0c6cd5..283b08c85c478 100644 --- a/pkg/sql/colexec/hashbuild/expression_memory.go +++ b/pkg/sql/colexec/hashbuild/expression_memory.go @@ -96,6 +96,9 @@ func NewAllocationAccountedExpressionExecutorsForAccount( } func expressionSetAllocationClosed(exprs []*plan.Expr) bool { + if len(exprs) == 0 { + return false + } for _, expr := range exprs { if !expressionAllocationClosed(expr) { return false diff --git a/pkg/sql/colexec/hashbuild/expression_memory_test.go b/pkg/sql/colexec/hashbuild/expression_memory_test.go index b095f6186165c..13e45f54a2ae8 100644 --- a/pkg/sql/colexec/hashbuild/expression_memory_test.go +++ b/pkg/sql/colexec/hashbuild/expression_memory_test.go @@ -341,6 +341,28 @@ func TestHashmapBuilderFallsBackOnlyForUnclosedExpressionScratch(t *testing.T) { } } +func TestHashBuildAllocationActivationRequiresClosedExpressionOwner(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + + closed := &HashBuild{ + NeedHashMap: true, + Conditions: []*plan.Expr{makeIssue26454ConcatKey(t, proc)}, + } + require.True(t, closed.AllocationAccountEnabled()) + + unclosed := &HashBuild{ + NeedHashMap: true, + Conditions: []*plan.Expr{makeExpressionLeaseTestExpr(t, proc)}, + } + require.False(t, unclosed.AllocationAccountEnabled()) + + withoutMap := &HashBuild{ + Conditions: []*plan.Expr{makeIssue26454ConcatKey(t, proc)}, + } + require.False(t, withoutMap.AllocationAccountEnabled()) +} + func BenchmarkIssue26454ExpressionAccounting(b *testing.B) { const capBytes = uint64(8 << 30) proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) diff --git a/pkg/sql/colexec/hashbuild/spill.go b/pkg/sql/colexec/hashbuild/spill.go index cd0e2a84a0689..f7a2fb0dd538a 100644 --- a/pkg/sql/colexec/hashbuild/spill.go +++ b/pkg/sql/colexec/hashbuild/spill.go @@ -832,7 +832,6 @@ func (ctr *container) dropSpillScratchBuffers() { for i := range ctr.spillBucketOffsets { ctr.spillBucketOffsets[i] = 0 } - ctr.spillSelection = nil ctr.spillKeyVecs = nil ctr.spillWriteBuf = bytes.Buffer{} ctr.spillAllocationMP = nil @@ -1288,10 +1287,6 @@ func (ctr *container) spillBatchBounded(proc *process.Process, bat *batch.Batch, if err := checkHashBuildCanceled(proc); err != nil { return err } - // Keep the legacy spillSelection field as an alias for callers/tests that - // inspect it. It intentionally points at the same backing array: no second - // row-id allocation is made. - ctr.spillSelection = ctr.spillBucketRowIds counts := ctr.spillBucketCounts[:] for i := range counts { counts[i] = 0 @@ -1497,7 +1492,6 @@ func (ctr *container) releaseSpillComputeScratch() { } ctr.spillHashValues = nil ctr.spillBucketRowIds = nil - ctr.spillSelection = nil } // spillBatchWithPressure retries only the unpublished prefix of an exact @@ -1892,15 +1886,6 @@ func (ctr *container) flushSpillBuffers(proc *process.Process, files []*os.File, return firstErr } -func (ctr *container) appendBuildBatchToSpillFiles(proc *process.Process, bat *batch.Batch, files []*os.File, buffers []*batch.Batch, executors []colexec.ExpressionExecutor, analyzer process.Analyzer) error { - // buffers is retained in the signature for source compatibility with older - // unit tests and callers. The implementation intentionally ignores it: - // every non-empty bucket is selected and flushed before the next bucket is - // materialized, so no persistent fanout-sized vector set can grow. - _ = buffers - return ctr.spillBatchBounded(proc, bat, files, executors, analyzer, false) -} - // initSpillExprExecs initializes or validates spill expression executors. // Returns the executors slice ready for use. Called once when entering spill mode. func (ctr *container) initSpillExprExecs(proc *process.Process, conditions []*plan.Expr) ([]colexec.ExpressionExecutor, error) { diff --git a/pkg/sql/colexec/hashbuild/spill_test.go b/pkg/sql/colexec/hashbuild/spill_test.go index 1a7a68aacf41f..4678d02ab669a 100644 --- a/pkg/sql/colexec/hashbuild/spill_test.go +++ b/pkg/sql/colexec/hashbuild/spill_test.go @@ -378,24 +378,13 @@ func TestAppendBatchToSpillFilesPartitioning(t *testing.T) { }, } - buffers := make([]*batch.Batch, spillNumBuckets) - analyzer := process.NewAnalyzer(0, false, false, "test") ctr := &container{spillUUID: t.Name()} _, err := ctr.initSpillExprExecs(proc, conditions) require.NoError(t, err) - err = ctr.appendBuildBatchToSpillFiles(proc, bat, files, buffers, ctr.spillExprExecs, analyzer) + err = ctr.spillBatchBounded(proc, bat, files, ctr.spillExprExecs, analyzer, false) require.NoError(t, err) - // Flush remaining buffers (lazy file creation via ensureSpillFile) - for i, buf := range buffers { - if buf != nil && buf.RowCount() > 0 { - file, err := ctr.ensureSpillFile(proc, files, i) - require.NoError(t, err) - _, err = ctr.flushBucketBuffer(proc, buf, file, analyzer) - require.NoError(t, err) - } - } } func TestEmptyBatchSpill(t *testing.T) { @@ -424,13 +413,11 @@ func TestEmptyBatchSpill(t *testing.T) { }, } - buffers := make([]*batch.Batch, spillNumBuckets) - analyzer := process.NewAnalyzer(0, false, false, "test") ctr := &container{spillUUID: t.Name()} _, err := ctr.initSpillExprExecs(proc, conditions) require.NoError(t, err) - err = ctr.appendBuildBatchToSpillFiles(proc, bat, files, buffers, ctr.spillExprExecs, analyzer) + err = ctr.spillBatchBounded(proc, bat, files, ctr.spillExprExecs, analyzer, false) require.NoError(t, err) } @@ -467,24 +454,14 @@ func TestAppendBuildBatchMultipleFlushes(t *testing.T) { }, } - buffers := make([]*batch.Batch, spillNumBuckets) analyzer := process.NewAnalyzer(0, false, false, "test") ctr := &container{spillUUID: t.Name()} _, err := ctr.initSpillExprExecs(proc, conditions) require.NoError(t, err) - err = ctr.appendBuildBatchToSpillFiles(proc, bat, files, buffers, ctr.spillExprExecs, analyzer) + err = ctr.spillBatchBounded(proc, bat, files, ctr.spillExprExecs, analyzer, false) require.NoError(t, err) - // Flush remaining (lazy file creation via ensureSpillFile) - for i, buf := range buffers { - if buf != nil && buf.RowCount() > 0 { - file, err := ctr.ensureSpillFile(proc, files, i) - require.NoError(t, err) - _, err = ctr.flushBucketBuffer(proc, buf, file, analyzer) - require.NoError(t, err) - } - } } func TestAppendBuildBatchWithNulls(t *testing.T) { @@ -513,24 +490,14 @@ func TestAppendBuildBatchWithNulls(t *testing.T) { }, } - buffers := make([]*batch.Batch, spillNumBuckets) analyzer := process.NewAnalyzer(0, false, false, "test") ctr := &container{spillUUID: t.Name()} _, err := ctr.initSpillExprExecs(proc, conditions) require.NoError(t, err) - err = ctr.appendBuildBatchToSpillFiles(proc, bat, files, buffers, ctr.spillExprExecs, analyzer) + err = ctr.spillBatchBounded(proc, bat, files, ctr.spillExprExecs, analyzer, false) require.NoError(t, err) - // Flush remaining (lazy file creation via ensureSpillFile) - for i, buf := range buffers { - if buf != nil && buf.RowCount() > 0 { - file, err := ctr.ensureSpillFile(proc, files, i) - require.NoError(t, err) - _, err = ctr.flushBucketBuffer(proc, buf, file, analyzer) - require.NoError(t, err) - } - } } func TestAppendBuildBatchMultiColumn(t *testing.T) { @@ -566,24 +533,14 @@ func TestAppendBuildBatchMultiColumn(t *testing.T) { }, } - buffers := make([]*batch.Batch, spillNumBuckets) analyzer := process.NewAnalyzer(0, false, false, "test") ctr := &container{spillUUID: t.Name()} _, err := ctr.initSpillExprExecs(proc, conditions) require.NoError(t, err) - err = ctr.appendBuildBatchToSpillFiles(proc, bat, files, buffers, ctr.spillExprExecs, analyzer) + err = ctr.spillBatchBounded(proc, bat, files, ctr.spillExprExecs, analyzer, false) require.NoError(t, err) - // Flush remaining (lazy file creation via ensureSpillFile) - for i, buf := range buffers { - if buf != nil && buf.RowCount() > 0 { - file, err := ctr.ensureSpillFile(proc, files, i) - require.NoError(t, err) - _, err = ctr.flushBucketBuffer(proc, buf, file, analyzer) - require.NoError(t, err) - } - } } func TestShouldSpillBatchesRowThreshold(t *testing.T) { @@ -692,26 +649,24 @@ func TestAppendBuildBatchSingleBucket(t *testing.T) { }, } - buffers := make([]*batch.Batch, spillNumBuckets) analyzer := process.NewAnalyzer(0, false, false, "test") ctr := &container{spillUUID: t.Name()} _, err := ctr.initSpillExprExecs(proc, conditions) require.NoError(t, err) - err = ctr.appendBuildBatchToSpillFiles(proc, bat, files, buffers, ctr.spillExprExecs, analyzer) + err = ctr.spillBatchBounded(proc, bat, files, ctr.spillExprExecs, analyzer, false) require.NoError(t, err) - // Most buffers should be nil - nilCount := 0 - for _, buf := range buffers { - if buf == nil { - nilCount++ + fileCount := 0 + for _, file := range files { + if file != nil { + fileCount++ } } - require.Greater(t, nilCount, spillNumBuckets-5) + require.Equal(t, 1, fileCount) } -func TestBufferReuse(t *testing.T) { +func TestSpillScratchReuse(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() @@ -733,7 +688,6 @@ func TestBufferReuse(t *testing.T) { }, } - buffers := make([]*batch.Batch, spillNumBuckets) analyzer := process.NewAnalyzer(0, false, false, "test") ctr := &container{spillUUID: t.Name()} @@ -745,16 +699,22 @@ func TestBufferReuse(t *testing.T) { bat1.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp()) bat1.SetRowCount(2) - err = ctr.appendBuildBatchToSpillFiles(proc, bat1, files, buffers, ctr.spillExprExecs, analyzer) + err = ctr.spillBatchBounded(proc, bat1, files, ctr.spillExprExecs, analyzer, false) require.NoError(t, err) + hashCapacity := cap(ctr.spillHashValues) + rowIDCapacity := cap(ctr.spillBucketRowIds) + require.Positive(t, hashCapacity) + require.Positive(t, rowIDCapacity) - // Second batch - buffers should be reused + // An equal-size batch reuses the retained hash and row-id scratch. bat2 := batch.NewWithSize(1) bat2.Vecs[0] = testutil.MakeInt32Vector([]int32{3, 4}, nil, proc.Mp()) bat2.SetRowCount(2) - err = ctr.appendBuildBatchToSpillFiles(proc, bat2, files, buffers, ctr.spillExprExecs, analyzer) + err = ctr.spillBatchBounded(proc, bat2, files, ctr.spillExprExecs, analyzer, false) require.NoError(t, err) + require.Equal(t, hashCapacity, cap(ctr.spillHashValues)) + require.Equal(t, rowIDCapacity, cap(ctr.spillBucketRowIds)) } func TestSpillExpressionLeaseRetainsLargeBatchHighWater(t *testing.T) { @@ -829,7 +789,7 @@ func TestSpillWriteCoalescesAcrossBatches(t *testing.T) { bat.SetRowCount(3) defer bat.Clean(proc.Mp()) for i := 0; i < 2; i++ { - require.NoError(t, ctr.appendBuildBatchToSpillFiles(proc, bat, files, nil, ctr.spillExprExecs, analyzer)) + require.NoError(t, ctr.spillBatchBounded(proc, bat, files, ctr.spillExprExecs, analyzer, false)) } var pending int for i := range ctr.spillBucketWriteBufs { diff --git a/pkg/sql/colexec/hashbuild/types.go b/pkg/sql/colexec/hashbuild/types.go index d5e6ae43de003..8be3e85b30ce9 100644 --- a/pkg/sql/colexec/hashbuild/types.go +++ b/pkg/sql/colexec/hashbuild/types.go @@ -122,7 +122,6 @@ type container struct { spillBucketRowIds []int32 spillBucketCounts [spillNumBuckets]int32 spillBucketOffsets [spillNumBuckets + 1]int32 - spillSelection []int32 spillWriteBuf bytes.Buffer // spillBucketWriteBufs coalesce serialized records across source batches. // Each buffer is bounded by spillWriteCoalesceSize (plus bytes.Buffer's @@ -330,7 +329,18 @@ func (hashBuild *HashBuild) GetOperatorBase() *vm.OperatorBase { } func (hashBuild *HashBuild) AllocationAccountEnabled() bool { - return hashBuild != nil && hashBuild.NeedHashMap + // Activate one complete physical owner closure. An expression family whose + // call-scoped allocation ledger is not closed must keep the whole HashBuild + // on the legacy path; mixing an exact map/batch owner with an estimator-gated + // expression would reintroduce the false-rejection mechanism that this + // activation removes. + return hashBuild != nil && hashBuild.NeedHashMap && + expressionSetAllocationClosed(hashBuild.Conditions) +} + +func (hashBuild *HashBuild) AllocationAccountActivationBlocked() bool { + return hashBuild != nil && hashBuild.NeedHashMap && + !expressionSetAllocationClosed(hashBuild.Conditions) } // SetAllocationAccount selects immutable provenance for the hash-table owner diff --git a/pkg/sql/colexec/hashjoin/expression_memory_test.go b/pkg/sql/colexec/hashjoin/expression_memory_test.go index 664cd625c9f6a..2baa2fd96f5cd 100644 --- a/pkg/sql/colexec/hashjoin/expression_memory_test.go +++ b/pkg/sql/colexec/hashjoin/expression_memory_test.go @@ -96,3 +96,13 @@ func TestHashJoinResetReleasesAccountedProbeExpressions(t *testing.T) { require.NoError(t, err) require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) } + +func TestHashJoinAllocationActivationRequiresBothKeySides(t *testing.T) { + col := &plan.Expr{Typ: plan.Type{Id: int32(types.T_int64)}, Expr: &plan.Expr_Col{Col: &plan.ColRef{}}} + arg := &HashJoin{EqConds: [][]*plan.Expr{{col}, {col}}} + require.True(t, arg.AllocationAccountEnabled()) + require.False(t, arg.AllocationAccountActivationBlocked()) + arg.EqConds[1] = []*plan.Expr{nil} + require.False(t, arg.AllocationAccountEnabled()) + require.True(t, arg.AllocationAccountActivationBlocked()) +} diff --git a/pkg/sql/colexec/hashjoin/types.go b/pkg/sql/colexec/hashjoin/types.go index 03e5697c0fed3..5dca67aef0183 100644 --- a/pkg/sql/colexec/hashjoin/types.go +++ b/pkg/sql/colexec/hashjoin/types.go @@ -146,7 +146,19 @@ type HashJoin struct { } func (hashJoin *HashJoin) AllocationAccountEnabled() bool { - return hashJoin != nil + return hashJoin != nil && hashJoin.allocationAccountExpressionOwnerClosed() +} + +func (hashJoin *HashJoin) AllocationAccountActivationBlocked() bool { + return hashJoin != nil && !hashJoin.allocationAccountExpressionOwnerClosed() +} + +func (hashJoin *HashJoin) allocationAccountExpressionOwnerClosed() bool { + if hashJoin == nil || len(hashJoin.EqConds) != 2 { + return false + } + return hashbuild.AllocationAccountedExpressionSetSupported(hashJoin.EqConds[0]) && + hashbuild.AllocationAccountedExpressionSetSupported(hashJoin.EqConds[1]) } func (hashJoin *HashJoin) SetAllocationAccount( diff --git a/pkg/sql/colexec/rightdedupjoin/expression_memory_test.go b/pkg/sql/colexec/rightdedupjoin/expression_memory_test.go index 3b30a749fcea2..a7e9f00be9e4a 100644 --- a/pkg/sql/colexec/rightdedupjoin/expression_memory_test.go +++ b/pkg/sql/colexec/rightdedupjoin/expression_memory_test.go @@ -96,3 +96,13 @@ func TestRightDedupJoinResetReleasesAccountedProbeExpressions(t *testing.T) { require.NoError(t, err) require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) } + +func TestRightDedupJoinAllocationActivationRequiresBothKeySides(t *testing.T) { + col := &plan.Expr{Typ: plan.Type{Id: int32(types.T_int64)}, Expr: &plan.Expr_Col{Col: &plan.ColRef{}}} + arg := &RightDedupJoin{Conditions: [][]*plan.Expr{{col}, {col}}} + require.True(t, arg.AllocationAccountEnabled()) + require.False(t, arg.AllocationAccountActivationBlocked()) + arg.Conditions[1] = []*plan.Expr{nil} + require.False(t, arg.AllocationAccountEnabled()) + require.True(t, arg.AllocationAccountActivationBlocked()) +} diff --git a/pkg/sql/colexec/rightdedupjoin/types.go b/pkg/sql/colexec/rightdedupjoin/types.go index 490208b5a6356..d7caa01c3041b 100644 --- a/pkg/sql/colexec/rightdedupjoin/types.go +++ b/pkg/sql/colexec/rightdedupjoin/types.go @@ -98,7 +98,21 @@ type RightDedupJoin struct { } func (rightDedupJoin *RightDedupJoin) AllocationAccountEnabled() bool { - return rightDedupJoin != nil + return rightDedupJoin != nil && + rightDedupJoin.allocationAccountExpressionOwnerClosed() +} + +func (rightDedupJoin *RightDedupJoin) AllocationAccountActivationBlocked() bool { + return rightDedupJoin != nil && + !rightDedupJoin.allocationAccountExpressionOwnerClosed() +} + +func (rightDedupJoin *RightDedupJoin) allocationAccountExpressionOwnerClosed() bool { + if rightDedupJoin == nil || len(rightDedupJoin.Conditions) != 2 { + return false + } + return hashbuild.AllocationAccountedExpressionSetSupported(rightDedupJoin.Conditions[0]) && + hashbuild.AllocationAccountedExpressionSetSupported(rightDedupJoin.Conditions[1]) } func (rightDedupJoin *RightDedupJoin) SetAllocationAccount( diff --git a/pkg/sql/colexec/spillutil/allocation_account.go b/pkg/sql/colexec/spillutil/allocation_account.go index 1be4e02d81776..451c51bb9ae82 100644 --- a/pkg/sql/colexec/spillutil/allocation_account.go +++ b/pkg/sql/colexec/spillutil/allocation_account.go @@ -41,7 +41,7 @@ const ( SpillAllocationSiteSelectedGrouping ) -// SpillAllocationAccount is the dormant allocation provenance for one spill +// SpillAllocationAccount is the allocation provenance for one spill // engine. type SpillAllocationAccount struct { account *mpool.AllocationAccount diff --git a/pkg/sql/colexec/spillutil/join_spill.go b/pkg/sql/colexec/spillutil/join_spill.go index 8d1cf60bbf2b1..0bde01a41f12b 100644 --- a/pkg/sql/colexec/spillutil/join_spill.go +++ b/pkg/sql/colexec/spillutil/join_spill.go @@ -879,22 +879,6 @@ func MakeBucketWriters(prefix string) []BucketWriter { return writers } -// FlushBucketBatch writes bat to w, creating the spill file on first write. -// If analyzer is non-nil, spill bytes/rows are tracked. -func FlushBucketBatch(proc *process.Process, bat *batch.Batch, w *BucketWriter, bucketBuf *bytes.Buffer, analyzer process.Analyzer) error { - if bat == nil || bat.RowCount() == 0 { - return nil - } - // Serialize before creating the file. This admits marshal scratch and the - // exact disk extent before CreateAndRemoveFile/write, so a rejected write - // leaves both the writer and source batch intact. - cnt := int64(bat.RowCount()) - if err := marshalSpillRecord(bat, bucketBuf); err != nil { - return err - } - return writeBucketPayload(proc, bucketBuf.Bytes(), cnt, w, analyzer) -} - type spillRecordBuffer interface { io.Writer Bytes() []byte @@ -1094,7 +1078,10 @@ func ComputeXXHash(keyVecs []*vector.Vector, hashValues []uint64, seed uint64) { // id array in two linear passes. This replaces the historical bucket-by-bucket // scan of hashValues (which revisited every row once for each bucket). func classifyRows(hashValues []uint64, bucketCount int, shift uint64, rowIDs []int32, counts []int32, offsets []int32) error { - if bucketCount <= 0 || bucketCount&(bucketCount-1) != 0 || shift >= 64 || len(rowIDs) < len(hashValues) || len(counts) < bucketCount || len(offsets) < bucketCount+1 { + if bucketCount <= 0 || bucketCount > SpillNumBuckets || + bucketCount&(bucketCount-1) != 0 || shift >= 64 || + len(rowIDs) < len(hashValues) || len(counts) < bucketCount || + len(offsets) < bucketCount+1 { return process.ErrHashBuildBudgetInvalid } for i := 0; i < bucketCount; i++ { @@ -1109,128 +1096,13 @@ func classifyRows(hashValues []uint64, bucketCount int, shift uint64, rowIDs []i offsets[i+1] = offsets[i] + counts[i] } var writePos [SpillNumBuckets]int32 - if bucketCount <= len(writePos) { - copy(writePos[:bucketCount], offsets[:bucketCount]) - for row, hash := range hashValues { - bucket := int((hash >> shift) & mask) - pos := writePos[bucket] - rowIDs[pos] = int32(row) - writePos[bucket] = pos + 1 - } - return nil - } - // SpillNumBuckets is the production fanout. Keep the helper correct for - // callers using another power-of-two fanout without allocating a second - // row-id structure. - positions := make([]int32, bucketCount) - copy(positions, offsets[:bucketCount]) + copy(writePos[:bucketCount], offsets[:bucketCount]) for row, hash := range hashValues { bucket := int((hash >> shift) & mask) - pos := positions[bucket] + pos := writePos[bucket] rowIDs[pos] = int32(row) - positions[bucket] = pos + 1 - } - return nil -} - -// scatterImpl is the internal implementation that accepts reusable buffers. -func scatterImpl( - proc *process.Process, - bat *batch.Batch, - keyVecs []*vector.Vector, - writers []BucketWriter, - buffers []*batch.Batch, - seed uint64, - bucketBuf *bytes.Buffer, - analyzer process.Analyzer, - reuseHashValues *[]uint64, - reuseBucketRowIds *[][]int32, -) error { - rowCount := bat.RowCount() - if rowCount == 0 { - return nil - } - - var hashValues []uint64 - if reuseHashValues != nil && cap(*reuseHashValues) >= rowCount { - hashValues = (*reuseHashValues)[:rowCount] - } else { - hashValues = make([]uint64, rowCount) - if reuseHashValues != nil { - *reuseHashValues = hashValues - } - } - ComputeXXHash(keyVecs, hashValues, seed) - - if len(writers) == 0 || len(writers)&(len(writers)-1) != 0 { - return process.ErrHashBuildBudgetInvalid - } - // Build one contiguous row-id array, then expose each bucket as a slice of - // that array for compatibility with the buffered path. - var bucketRowIds [][]int32 - if reuseBucketRowIds != nil { - bucketRowIds = *reuseBucketRowIds - if cap(bucketRowIds) < len(writers) { - bucketRowIds = make([][]int32, len(writers)) - *reuseBucketRowIds = bucketRowIds - } else { - bucketRowIds = bucketRowIds[:len(writers)] - } - } else { - bucketRowIds = make([][]int32, len(writers)) - } - var rowIDs []int32 - if len(bucketRowIds) > 0 && cap(bucketRowIds[0]) >= rowCount { - rowIDs = bucketRowIds[0][:rowCount] - } else { - rowIDs = make([]int32, rowCount) - } - var countsFixed [SpillNumBuckets]int32 - var offsetsFixed [SpillNumBuckets + 1]int32 - counts := countsFixed[:len(writers)] - offsets := offsetsFixed[:len(writers)+1] - if len(writers) > SpillNumBuckets { - counts = make([]int32, len(writers)) - offsets = make([]int32, len(writers)+1) - } - if err := classifyRows(hashValues, len(writers), 0, rowIDs, counts, offsets); err != nil { - return err + writePos[bucket] = pos + 1 } - for i := range bucketRowIds { - bucketRowIds[i] = rowIDs[offsets[i]:offsets[i+1]] - } - - // Only iterate non-empty buckets. - for bucketId, sels := range bucketRowIds { - if len(sels) == 0 { - continue - } - if writers[bucketId].Name == "" { - continue // disabled bucket — discard rows - } - buf := buffers[bucketId] - if buf == nil { - buf = batch.NewOffHeapWithSize(len(bat.Vecs)) - for j, vec := range bat.Vecs { - buf.Vecs[j] = vector.NewOffHeapVecWithType(*vec.GetType()) - buf.Vecs[j].PreExtend(8192, proc.Mp()) - } - buffers[bucketId] = buf - } - for j, vec := range bat.Vecs { - if err := buf.Vecs[j].UnionInt32(vec, sels, proc.Mp()); err != nil { - return err - } - } - buf.SetRowCount(buf.RowCount() + len(sels)) - if buf.RowCount() >= 8192 { - if err := FlushBucketBatch(proc, buf, &writers[bucketId], bucketBuf, analyzer); err != nil { - return err - } - buf.CleanOnlyData() - } - } - return nil } @@ -2448,8 +2320,8 @@ func NewSpillEngine(cfg SpillEngineConfig) *SpillEngine { return newSpillEngine(cfg, nil) } -// NewSpillEngineWithAllocation constructs the dormant allocation-accounted -// spill path. Legacy production callers continue to use NewSpillEngine. +// NewSpillEngineWithAllocation constructs the allocation-accounted spill path. +// NewSpillEngine remains available to callers outside a statement account. func NewSpillEngineWithAllocation( cfg SpillEngineConfig, allocation *SpillAllocationAccount, @@ -3209,7 +3081,8 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal e.probeReadBatch = readBatch } - // Scatter probe file. Reuse reader's 4 MiB buffer from the build pass. + // Scatter the probe file through the same admitted 64 KiB reader buffer + // used for the build pass. if bucket.ProbeFd != nil { if err := reader.EnsureBuffer(e.cfg.Budget); err != nil { return nil, err diff --git a/pkg/sql/colexec/spillutil/join_spill_test.go b/pkg/sql/colexec/spillutil/join_spill_test.go index aa41b34632588..7b6430b72cfda 100644 --- a/pkg/sql/colexec/spillutil/join_spill_test.go +++ b/pkg/sql/colexec/spillutil/join_spill_test.go @@ -67,6 +67,137 @@ func (r *boundaryCancelReader) Read(p []byte) (int, error) { return n, err } +// FlushBucketBatch writes one legacy-format spill record. Production +// scatter goes through SpillEngine.scatterBatchBounded; tests that construct a +// reader fixture directly need only this codec/writer composition. +func FlushBucketBatch( + proc *process.Process, + bat *batch.Batch, + w *BucketWriter, + bucketBuf *bytes.Buffer, + analyzer process.Analyzer, +) error { + if bat == nil || bat.RowCount() == 0 { + return nil + } + if err := marshalSpillRecord(bat, bucketBuf); err != nil { + return err + } + return writeBucketPayload( + proc, + bucketBuf.Bytes(), + int64(bat.RowCount()), + w, + analyzer, + ) +} + +// scatterImpl retains the former buffered implementation solely as a +// compatibility oracle. It is intentionally test-only: production owns one +// selected batch at a time in SpillEngine.scatterBatchBounded. +func scatterImpl( + proc *process.Process, + bat *batch.Batch, + keyVecs []*vector.Vector, + writers []BucketWriter, + buffers []*batch.Batch, + seed uint64, + bucketBuf *bytes.Buffer, + analyzer process.Analyzer, + reuseHashValues *[]uint64, + reuseBucketRowIds *[][]int32, +) error { + rowCount := bat.RowCount() + if rowCount == 0 { + return nil + } + + var hashValues []uint64 + if reuseHashValues != nil && cap(*reuseHashValues) >= rowCount { + hashValues = (*reuseHashValues)[:rowCount] + } else { + hashValues = make([]uint64, rowCount) + if reuseHashValues != nil { + *reuseHashValues = hashValues + } + } + ComputeXXHash(keyVecs, hashValues, seed) + + if len(writers) == 0 || len(writers) > SpillNumBuckets || + len(writers)&(len(writers)-1) != 0 { + return process.ErrHashBuildBudgetInvalid + } + var bucketRowIds [][]int32 + if reuseBucketRowIds != nil { + bucketRowIds = *reuseBucketRowIds + if cap(bucketRowIds) < len(writers) { + bucketRowIds = make([][]int32, len(writers)) + *reuseBucketRowIds = bucketRowIds + } else { + bucketRowIds = bucketRowIds[:len(writers)] + } + } else { + bucketRowIds = make([][]int32, len(writers)) + } + var rowIDs []int32 + if len(bucketRowIds) > 0 && cap(bucketRowIds[0]) >= rowCount { + rowIDs = bucketRowIds[0][:rowCount] + } else { + rowIDs = make([]int32, rowCount) + } + var counts [SpillNumBuckets]int32 + var offsets [SpillNumBuckets + 1]int32 + if err := classifyRows( + hashValues, + len(writers), + 0, + rowIDs, + counts[:], + offsets[:], + ); err != nil { + return err + } + for i := range bucketRowIds { + bucketRowIds[i] = rowIDs[offsets[i]:offsets[i+1]] + } + + for bucketID, sels := range bucketRowIds { + if len(sels) == 0 || writers[bucketID].Name == "" { + continue + } + buf := buffers[bucketID] + if buf == nil { + buf = batch.NewOffHeapWithSize(len(bat.Vecs)) + for i, vec := range bat.Vecs { + buf.Vecs[i] = vector.NewOffHeapVecWithType(*vec.GetType()) + if err := buf.Vecs[i].PreExtend(8192, proc.Mp()); err != nil { + return err + } + } + buffers[bucketID] = buf + } + for i, vec := range bat.Vecs { + if err := buf.Vecs[i].UnionInt32(vec, sels, proc.Mp()); err != nil { + return err + } + } + buf.SetRowCount(buf.RowCount() + len(sels)) + if buf.RowCount() >= 8192 { + if err := FlushBucketBatch( + proc, + buf, + &writers[bucketID], + bucketBuf, + analyzer, + ); err != nil { + return err + } + buf.CleanOnlyData() + } + } + return nil +} + func TestTakeSpillBuildPayloadRejectsWrongBudgetRef(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() @@ -198,6 +329,19 @@ func TestClassifyRowsConservesRows(t *testing.T) { require.Equal(t, int32(len(hashes)), offsets[SpillNumBuckets]) } +func TestClassifyRowsRejectsNonProductionFanout(t *testing.T) { + const fanout = SpillNumBuckets * 2 + err := classifyRows( + []uint64{0}, + fanout, + 0, + make([]int32, 1), + make([]int32, fanout), + make([]int32, fanout+1), + ) + require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) +} + func legacyClassifyRows(hashes []uint64, rowIDs []int32) { pos := 0 for bucket := uint64(0); bucket < SpillNumBuckets; bucket++ { diff --git a/pkg/sql/compile/allocation_account_lifecycle.go b/pkg/sql/compile/allocation_account_lifecycle.go index 0d797ca709a5c..2de7d883f4986 100644 --- a/pkg/sql/compile/allocation_account_lifecycle.go +++ b/pkg/sql/compile/allocation_account_lifecycle.go @@ -30,6 +30,15 @@ type executionAllocationAccountOwner interface { ClearAllocationAccount(*mpool.AllocationAccount) error } +// executionAllocationAccountBlocker marks an operator whose physical owner is +// known but whose allocation-site closure is not yet complete. Automatic +// activation is statement-atomic: one blocker keeps every participating +// operator on the legacy path instead of creating a mixed exact/estimated +// generation. +type executionAllocationAccountBlocker interface { + AllocationAccountActivationBlocked() bool +} + // statementAllocationAttempt owns one local execution generation. The // MessageBoard pointer is captured at open so prepared/retry Reset cannot make // terminal cleanup drain a newer board. @@ -127,6 +136,10 @@ func configureAllocationAccountOwners( if err := vm.HandleAllOp( scope.RootOp, func(_ vm.Operator, op vm.Operator) error { + if blocker, ok := op.(executionAllocationAccountBlocker); ok && + blocker.AllocationAccountActivationBlocked() { + return mpool.ErrAllocationAccountInvariant + } if owner, ok := op.(executionAllocationAccountOwner); ok && owner.AllocationAccountEnabled() { if isConfigured(owner) { @@ -158,35 +171,31 @@ func configureAllocationAccountOwners( } func hasAllocationAccountOwner(scopes []*Scope) bool { - var inspect func(*Scope) bool - inspect = func(scope *Scope) bool { + found, blocked := false, false + var inspect func(*Scope) + inspect = func(scope *Scope) { if scope == nil { - return false + return } - found := false _ = vm.HandleAllOp(scope.RootOp, func(_ vm.Operator, op vm.Operator) error { + if blocker, ok := op.(executionAllocationAccountBlocker); ok && + blocker.AllocationAccountActivationBlocked() { + blocked = true + } if owner, ok := op.(executionAllocationAccountOwner); ok && owner.AllocationAccountEnabled() { found = true } return nil }) - if found { - return true - } for _, preScope := range scope.PreScopes { - if inspect(preScope) { - return true - } + inspect(preScope) } - return false } for _, scope := range scopes { - if inspect(scope) { - return true - } + inspect(scope) } - return false + return found && !blocked } // ensureAllocationAccountLifecycle activates accounting only when the physical diff --git a/pkg/sql/compile/allocation_account_lifecycle_test.go b/pkg/sql/compile/allocation_account_lifecycle_test.go index 883d9737bddf1..7a1f5b21e6edd 100644 --- a/pkg/sql/compile/allocation_account_lifecycle_test.go +++ b/pkg/sql/compile/allocation_account_lifecycle_test.go @@ -24,6 +24,7 @@ import ( "github.com/golang/mock/gomock" "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/pb/txn" "github.com/matrixorigin/matrixone/pkg/sql/colexec" @@ -45,6 +46,7 @@ type allocationLifecycleOwnerOperator struct { account *mpool.AllocationAccount failSet bool failClear bool + blocked bool clears int } @@ -52,6 +54,10 @@ func (op *allocationLifecycleOwnerOperator) AllocationAccountEnabled() bool { return true } +func (op *allocationLifecycleOwnerOperator) AllocationAccountActivationBlocked() bool { + return op.blocked +} + func (op *allocationLifecycleOwnerOperator) SetAllocationAccount( account *mpool.AllocationAccount, ) error { @@ -332,6 +338,31 @@ func TestStatementAllocationAttemptOwnerConfigurationRollsBack(t *testing.T) { require.NoError(t, err) } +func TestAllocationAccountActivationIsStatementAtomic(t *testing.T) { + eligible := &allocationLifecycleOwnerOperator{ + MockOperator: colexec.NewMockOperator(), + } + blocker := &allocationLifecycleOwnerOperator{ + MockOperator: colexec.NewMockOperator(), + blocked: true, + } + scopes := []*Scope{{RootOp: eligible}, {RootOp: blocker}} + require.False(t, hasAllocationAccountOwner(scopes), + "one unclosed physical owner must keep the whole statement legacy") + + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + configured, err := configureAllocationAccountOwners(scopes, account) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + require.Nil(t, configured) + require.Nil(t, eligible.account) + require.Equal(t, 1, eligible.clears) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + func TestStatementAllocationAttemptOwnerTeardownFailureExportsFailure(t *testing.T) { registry, err := mpool.NewAllocationAccountRegistry(1, 1) require.NoError(t, err) @@ -371,6 +402,10 @@ func TestCompileAutomaticallyActivatesCompleteHashTableOwner(t *testing.T) { } owner := hashbuild.NewArgument() owner.NeedHashMap = true + owner.Conditions = []*plan.Expr{{ + Typ: plan.Type{Id: int32(types.T_int64)}, + Expr: &plan.Expr_Col{Col: &plan.ColRef{}}, + }} c.scopes = []*Scope{{RootOp: owner}} require.NoError(t, c.ensureAllocationAccountLifecycle(func( diff --git a/pkg/sql/compile/compile.go b/pkg/sql/compile/compile.go index 8d7cc309fb5c0..ec2a07ea04d44 100644 --- a/pkg/sql/compile/compile.go +++ b/pkg/sql/compile/compile.go @@ -6997,9 +6997,9 @@ func (c *Compile) SetResourceAttemptOwnerEligible() { c.resourceAttemptOwnerEligible = true } -// ConfigureAllocationAccountLifecycle installs the dormant generation -// provider used by owner-activation PRs. A nil registry keeps production on -// the legacy path and opens no generation. +// ConfigureAllocationAccountLifecycle installs the generation provider used +// by allocation-accounted owners. A nil registry keeps production on the +// legacy path and opens no generation. func (c *Compile) ConfigureAllocationAccountLifecycle( registry *mpool.AllocationAccountRegistry, limit uint64, diff --git a/pkg/sql/compile/compile2.go b/pkg/sql/compile/compile2.go index c12e920192924..30bd84bdf65b8 100644 --- a/pkg/sql/compile/compile2.go +++ b/pkg/sql/compile/compile2.go @@ -467,15 +467,6 @@ func (c *Compile) Run(_ uint64) (queryResult *util2.RunResult, err error) { coordinatorPhaseStart = time.Time{} coordinatorPhaseBase = 0 } - if terminalErr := finishAllocationAttempt(); terminalErr != nil { - err = errors.Join(err, terminalErr) - resourceRecorder.finishAttempt( - uint64(retryTimes), attemptStart, attemptPreRunWall, attemptRemoteWait, stats, - attemptScopes, attemptAnal, c.addr, false, - ) - attemptOpen = false - return nil, err - } queryResult.AffectRows = runC.getAffectedRows() if c.uid != "mo_logger" && strings.Contains(strings.ToLower(c.sql), "insert") && @@ -496,6 +487,15 @@ func (c *Compile) Run(_ uint64) (queryResult *util2.RunResult, err error) { // outcome. The panic defer above remains the single terminal owner until // this call returns. c.AnalyzeExecPlan(runC, queryResult, stats, isExplainPhyPlan, option) + if terminalErr := finishAllocationAttempt(); terminalErr != nil { + err = errors.Join(err, terminalErr) + resourceRecorder.finishAttempt( + uint64(retryTimes), attemptStart, attemptPreRunWall, attemptRemoteWait, stats, + attemptScopes, attemptAnal, c.addr, false, + ) + attemptOpen = false + return nil, err + } resourceRecorder.finishAttempt( uint64(retryTimes), attemptStart, attemptPreRunWall, attemptRemoteWait, stats, From c3540319e94d55b2775cbcc0fd6b50b83bb89b55 Mon Sep 17 00:00:00 2001 From: aptend Date: Sat, 1 Aug 2026 00:18:00 +0800 Subject: [PATCH 30/61] docs: record allocation activation validation --- ...ocation_accounted_memory_admission_impl.md | 129 +++++++++++++----- .../evidence/26459_activation_validation.md | 123 +++++++++++++++++ .../26459_allocation_accounting_bench.txt | 28 ++++ ...0_allocation_accounted_memory_admission.md | 27 ++-- 4 files changed, 264 insertions(+), 43 deletions(-) create mode 100644 docs/design/evidence/26459_activation_validation.md diff --git a/docs/design/allocation_accounted_memory_admission_impl.md b/docs/design/allocation_accounted_memory_admission_impl.md index 76549cb12e50f..a6fa9bb45a312 100644 --- a/docs/design/allocation_accounted_memory_admission_impl.md +++ b/docs/design/allocation_accounted_memory_admission_impl.md @@ -1,17 +1,24 @@ # Allocation-Accounted Memory Admission: Implementation Plan -- Status: draft +- Status: implementation validation - Tracking issue: [#26459](https://github.com/matrixorigin/matrixone/issues/26459) - Architecture: [Allocation-Accounted Memory Admission RFC](../rfcs/00000000_allocation_accounted_memory_admission.md) - Baseline at plan creation: `main` at `38ce3a774` -- Rebased implementation baseline: `main` at `60e36bef64` -- Current dormant PR 3 head: `fd2fcc953b` +- Rebased implementation baseline: `main` at `5b9eeb54ec` +- Allocation-site closure baseline (PR 3): `f61b64d56c` +- Lifecycle/activation heads (PRs 4--10): `8633757dc1`, `5ae8eca00a`, + `8e4b689f45`, `c9ad0ea810`, `e072568998`, `eec23dcdc5`, and + `383fc6dce3` +- Owner-atomic activation and final cleanup: `656e254fe6` - Merged prerequisite: #26455 at `93e8b22d2` - Independent design review: completed against RFC commit `a7d54cb5f` -- Activation status: blocked until PRs 1--4 and the selected owner's - allocation-site/Go-heap gates close +- Activation status: enabled for the closed HashBuild expression owner set. + Build/probe key closure is checked statement-atomically. Unsupported + expression families keep all participating HashBuild/HashJoin/DedupJoin + owners in that local attempt on the legacy path; they are not partially + mixed into an exact generation. ## 1. Purpose and rules @@ -94,15 +101,15 @@ working ledger is allocation-site based: | Allocation site | Allocator/mode and size | Terminal owner | Current | Target/blocker | | --- | --- | --- | ---: | --- | | `mpool.memHdr` and account-ID side map | Go maps; one pointer record plus optional account record per live allocation | pointer removal at physical deallocation | H | bounded by the measured finite registry/allocation-slot policy | -| `Vector.data` | MPool; capacity from `Grow`, on/off-heap follows `v.offHeap` | owning `Vector.Free` | D | A only when off-heap | -| `Vector.area` | MPool; independent varlen payload capacity | owning `Vector.Free` | D | A only when off-heap | -| `Vector.nsp/gsp` bitmap data | allocation-accounted off-heap `[]uint64`; independent geometric capacity, paired replacement admission | owning `Vector.Free`; `Reset` retains and clears only published words | D | A after the selected Vector owner closes | -| `FunctionResult.vec` data/area | off-heap Vector; rows and appended payload | executor `Free` | D | A after expression ledger closure | -| `FunctionResult.convenientParam` | Go slice; expression arity, not rows | executor `Free`/reuse | L | H after a proved arity bound | -| decimal parameter conversion | retained allocation-accounted off-heap buffer; `rows*sizeof(Decimal128)` for decimal64/float32/float64 promotion | `FunctionResult.Free`; Reset/evaluation reuse capacity | D | A after expression ledger closure | -| IFF/CASE/COALESCE selection arrays | allocation-accounted off-heap `[]bool`; one or two arrays of `rows`, retained by executor | executor `Free`/reuse | D | A after expression ledger closure | -| selected row IDs | allocation-accounted off-heap `[]int64`; capacity up to `rows`, retained by executor | executor `Free`/reuse | D | A after expression ledger closure | -| selected parameter/result vectors | allocation-accounted off-heap Vector capacities | executor `Free` | D | A after expression ledger closure | +| `Vector.data` | MPool; capacity from `Grow`, on/off-heap follows `v.offHeap` | owning `Vector.Free` | A | activated HashBuild destinations use immutable off-heap provenance | +| `Vector.area` | MPool; independent varlen payload capacity | owning `Vector.Free` | A | activated HashBuild destinations use immutable off-heap provenance | +| `Vector.nsp/gsp` bitmap data | allocation-accounted off-heap `[]uint64`; independent geometric capacity, paired replacement admission | owning `Vector.Free`; `Reset` retains and clears only published words | A | selected HashBuild vectors and expression results are active | +| `FunctionResult.vec` data/area | off-heap Vector; rows and appended payload | executor `Free` | A | active for the closed HashBuild expression set | +| `FunctionResult.convenientParam` | Go slice; expression arity, not rows | executor `Free`/reuse | H | bounded by plan expression arity, not input rows or payload | +| decimal parameter conversion | retained allocation-accounted off-heap buffer; `rows*sizeof(Decimal128)` for decimal64/float32/float64 promotion | `FunctionResult.Free`; Reset/evaluation reuse capacity | A | active when present in a closed expression owner | +| IFF/CASE/COALESCE selection arrays | allocation-accounted off-heap `[]bool`; one or two arrays of `rows`, retained by executor | executor `Free`/reuse | A | CASE is active in the closed expression set | +| selected row IDs | allocation-accounted off-heap `[]int64`; capacity up to `rows`, retained by executor | executor `Free`/reuse | A | active when present in a closed expression owner | +| selected parameter/result vectors | allocation-accounted off-heap Vector capacities | executor `Free` | A | active when present in a closed expression owner | | `JSON_ROW` output | one retained allocation-accounted function scratch buffer, then copy into the admitted result; both physical capacities are charged | `FunctionResult.Free`; row publication copies synchronously | D | A after expression-owner activation; arity-bounded column closures are cleared after every call | | float array-distance row descriptors and result scratch | Go `[][]T` with one descriptor per row plus `[]float32` with one value per row | call return / GC | R | caller row accessor plus the upper half of the admitted `[]float64` result backing | | GPU array-distance flattened inputs | allocation-accounted caller scratch; `query bytes + rows*dimension*4` for the SQL float32 GPU path; legacy callers retain the C allocator | GPU job retains the caller slice until Wait; launch failure rolls back before return | D | A only after the GPU-tag build/tests pass; the legacy API remains unaccounted and is not an activated owner | @@ -118,28 +125,31 @@ working ledger is allocation-site based: | regexp compile/match/output scratch | Go regexp program and match/output buffers proportional to pattern/input | operator cache / row completion / GC | L | bounded regexp owner or exclusion from expression activation | | H3/S2 neighborhood scratch | Go slices; some S2 paths are statically bounded, H3 grid-disk output scales with radius | row completion / GC | L | split proved fixed bounds from data-scaled paths before activation | | JSON cast visible serialization | `MarshalJSON` payload slices proportional to the JSON value | row completion / GC | L | direct visible writer or exclusion from expression activation | -| hash-table initial cell block | off-heap `mpool.MakeSlice(..., true)`; 16 KiB int / 32 KiB string | hash map / `JoinMap.FreeMemory` | L | A in first activation | -| hash-table replacement/appended cell blocks | off-heap blocks, at most 4 MiB each; old+new overlap is physically visible | hash map / `JoinMap.FreeMemory` | L | A in first activation | -| hash-table `cells`/`newBlocks` descriptors | Go `[][]Cell`; 24 bytes per block header plus geometric resize backing arrays; GC is not treated as synchronous Free | hash map lifetime / GC | L | replace with an owning off-heap descriptor buffer and account its initial/replacement capacity; blocks first activation | -| hash-table `ResizePlan` and callback | fixed-size Go values/closures, one per table/resize | resize return / hash map Free | L | H; remove legacy reservation owner after cell activation | -| `GroupSels.{tmp,vals,offsets}` | on-heap `mpool.MakeSlice(..., false)`; O(build rows/groups) | builder or `JoinMap.FreeMemory` | L | switch off-heap; blocks auxiliary/copied-batch activation | -| copied build-batch vector buffers | MPool Vector data/area | builder or `JoinMap.FreeMemory` | L | A after per-buffer provenance | -| spill marshal/coalesce buffers | allocation-accounted off-heap streaming buffers; exact serialized size and bounded 64 KiB per-bucket coalesce capacity | spill phase cleanup | D | A after spill ledger closure | -| spill hash values | allocation-accounted off-heap `[]uint64`; geometric capacity, `8*cap` | spill phase cleanup | D | A after spill ledger closure | -| spill row IDs | allocation-accounted off-heap `[]int32`; geometric capacity, `4*cap` | spill phase cleanup | D | A after spill ledger closure | -| spill counts/offsets/positions | Go `[]int32`; O(bucket count), bucket count finite | spill cleanup | L | H after bound is asserted | -| selected spill bucket vectors | allocation-accounted off-heap Vector capacities | per-call selected-batch cleanup | D | A after spill ledger closure | -| BucketReader decoded vectors | allocation-accounted MPool Vector data/area | reusable batch cleanup / `BucketReader.Close` | D | A after spill ledger closure | -| `pSpool` cached Vector data/area | provenance-bearing detached MPool buffers; accounted data/area sites cannot cross and legacy buffers keep a guarded fast path | `spoolBuffer.clean` or the receiving Vector's `Free` | D | A with the pipeline owner activation | -| runtime-filter serialized payload | Go buffer/message payload; O(filter rows) | message release | L | off-heap or PASS degradation; blocks runtime-filter activation | +| hash-table initial cell block | off-heap `mpool.MakeSlice(..., true)`; 16 KiB int / 32 KiB string | hash map / `JoinMap.FreeMemory` | A | physical allocation is the sole charge | +| hash-table replacement/appended cell blocks | off-heap blocks, at most 4 MiB each; old+new overlap is physically visible | hash map / `JoinMap.FreeMemory` | A | replacement overlap is admitted before publication | +| hash-table `cells`/`newBlocks` descriptors | owning off-heap descriptor buffer; initial and replacement capacities are distinct physical allocations | hash map / `JoinMap.FreeMemory` | A | no row-scaled Go descriptor backing remains | +| hash-table `ResizePlan` and callback | fixed-size Go values/closures, one per table/resize | resize return / hash map Free | H | table-count bounded metadata; no row-scaled backing | +| `GroupSels.{tmp,vals,offsets}` | allocation-accounted off-heap buffers; O(build rows/groups) | builder or `JoinMap.FreeMemory` | A | provenance follows final JoinMap consumer | +| copied build-batch vector buffers | MPool Vector data/area | builder or `JoinMap.FreeMemory` | A | physical Vector leases replace projected batch tokens | +| spill marshal/coalesce buffers | allocation-accounted off-heap streaming buffers; exact serialized size and bounded 64 KiB per-bucket coalesce capacity | spill phase cleanup | A | optional coalesce degrades to direct write under pressure | +| spill hash values | allocation-accounted off-heap `[]uint64`; geometric capacity, `8*cap` | spill phase cleanup | A | exact capacity and rollback covered | +| spill row IDs | allocation-accounted off-heap `[]int32`; geometric capacity, `4*cap` | spill phase cleanup | A | exact capacity and rollback covered | +| spill counts/offsets/positions | Go `[]int32`; O(bucket count), bucket count finite | spill cleanup | H | bucket count is the fixed spill fanout | +| selected spill bucket vectors | allocation-accounted off-heap Vector capacities | per-call selected-batch cleanup | A | adaptive unpublished windows bound progress | +| BucketReader decoded vectors | allocation-accounted MPool Vector data/area | reusable batch cleanup / `BucketReader.Close` | A | replacement and retry start from a clean record | +| BucketReader input buffer | one 64 KiB Go buffer per production spill reader; direct legacy readers use a bounded 4 MiB default | `BucketReader.Close` / GC | H | reader count is bounded by the live SpillEngine set and the statement/CN hash cap retains fixed non-payload headroom | +| `pSpool` cached Vector data/area | provenance-bearing detached MPool buffers; accounted data/area sites cannot cross and legacy buffers keep a guarded fast path | `spoolBuffer.clean` or the receiving Vector's `Free` | A | activated vectors retain immutable provenance through cache reuse | +| runtime-filter serialized payload | one allocation-accounted payload; optional filter publication | message release | A | capacity pressure degrades to PASS before publication | | spill disk and FD | disk/FD ledgers | file removal/close | A | A | -This is the known-site ledger, not yet the completion ledger for every later -owner. The hash cell/descriptor first-activation inventory is closed here. -PR 3 must generate and review the remaining built-in/function-specific `make`, -`append`, `bytes.Buffer`, builder, and codec sites before the corresponding -expression or spill closure can activate. No row named “other” or “unbounded -scratch” can declare closure. +This is the known-site ledger for the activated HashBuild/join owner. Generic SQL +functions with unbounded JSON, geometry, regexp, H3/S2, or JSON-cast Go scratch +remain `L`; the build and probe activation gates reject the local statement's +entire owner set when one of those families appears. Consequently an activated +owner has no estimator-gated expression subtree, while an unsupported plan +remains wholly legacy and is an explicit later migration rather than a +partially exact owner. +No row named “other” or “unbounded scratch” can declare closure. Batch destination propagation is now `D`: Clone, Dup, selected-column copy, Union destinations, windows, reader decode, Clean, and FreeColumns preserve @@ -851,6 +861,11 @@ distinction instead of presenting the first-allocation cost as a per-row cost. ### PR 4: statement lifecycle and minimum pressure foundation +Implementation: `8633757dc1`. The attempt coordinator now owns board drain, +owner teardown, exactly-once terminal completion, immutable export, and the +release-capable tombstone/suspension path. Typed lifecycle failures are +disjoint from retryable capacity pressure. + Scope: - add the attempt-owned post-pipeline/MessageBoard-close seal/finalize hook; @@ -881,6 +896,11 @@ Gate: ### PR 5: hash-table cell-block activation +Implementation: `5ae8eca00a`. Integer and string cells plus their descriptor +backing are allocation-accounted through initial allocation, both resize +modes, producer/consumer handoff, and terminal Free. Activated maps do not +install the legacy resize reservation owner. + Scope: - activate only the integer/string hash-table cell blocks; @@ -909,6 +929,10 @@ Gate: ### PR 6: copied batches and JoinMap activation +Implementation: `8e4b689f45`. Copied build vectors, GroupSels, dedup scratch, +delete bitmaps, and final JoinMap ownership retain immutable provenance until +the last consumer releases them. + Scope: - copied build-batch destinations; @@ -931,6 +955,12 @@ Gate: ### PR 7: expression owner activation +Implementation: `c9ad0ea810`. COL/LIT/PARAM/VAR/VEC/FOLD plus audited CONCAT, +CASE, varchar EQUAL, and integer-to-string/literal string CAST trees use exact +result, bitmap, conversion, and function-scratch allocations. A HashBuild +containing any other expression family is not activated at all; it cannot mix +an exact map owner with a legacy expression estimator. + Scope: - activate the complete HashBuild-owned expression site closure; @@ -952,6 +982,12 @@ Gate: ### PR 8 family: spill and runtime-filter closures +Implementation: `e072568998`. Decoded reuse, selected vectors, hash/row IDs, +marshal/coalesce buffers, recursive rebuild, spill disk/FD, pSpool provenance, +and runtime-filter payload publication all have explicit allocation or +resource owners. Optional coalesce and runtime filters degrade without +publishing partial state. + Split into independently safe owner closures: 1. decoded batches and retained reader reuse; @@ -967,6 +1003,12 @@ closure adds its own build/probe and PASS-degradation workload before merging. ### PR 9: unified join pressure controller and remaining legacy deletion +Implementation: `eec23dcdc5`. HashBuild, SpillEngine, HashJoin, DedupJoin, and +RightDedupJoin share typed pressure classification and a monotonic retry guard. +Only memory-capacity failures retry; lifecycle, invariant, disk, and FD +failures are terminal. Unpublished input windows shrink, optional coalescing +is disabled once, and real minimum-unit failure is finite and diagnostic. + Implement across HashBuild, HashJoin, DedupJoin, and RightDedupJoin: ```text @@ -992,6 +1034,12 @@ Gate: ### PR 10: workload, performance, and cleanup +Local implementation and benchmark harness: `383fc6dce3`. The harness covers +resident int/varchar keys at 32 and 8,192 rows, complete spill scatter, +#26454's expression, copy ownership, lifecycle latency, and concurrent release +storms. Raw commands and medians are in +[`evidence/26459_activation_validation.md`](evidence/26459_activation_validation.md). + This re-runs all incident workloads together and supplies long-run and comparative confirmation; it is not the first incident-level validation of an earlier activation. @@ -1016,6 +1064,19 @@ Required proof: allocation results meet the recorded gate; - temporary migration helpers are removed. +Current validation status: + +- the complete affected package matrix passes after rebasing on + `5b9eeb54ec`; +- the complete affected package matrix, the same matrix under `-race`, a + 20-iteration focused lifecycle/pressure race stress, affected-package vet, + and `make build` pass at `656e254fe6`; +- local regression tests cover the five incident mechanisms without an + estimator-only rejection on the activated owner; +- TPCH 100G/1T candidate-versus-main runs remain the remote workload gate and + must be recorded in the linked evidence before this document moves from + `implementation validation` to `implemented`. + ## 5. Verification and conservation model Every semantic PR runs build, vet, complete package tests, focused adaptive race diff --git a/docs/design/evidence/26459_activation_validation.md b/docs/design/evidence/26459_activation_validation.md new file mode 100644 index 0000000000000..a1feab10b5c65 --- /dev/null +++ b/docs/design/evidence/26459_activation_validation.md @@ -0,0 +1,123 @@ +# #26459 Allocation-Accounted Activation Validation + +## Candidate + +- Rebased main: `5b9eeb54ec` +- PR 4 lifecycle: `8633757dc1` +- PR 5 hash-table activation: `5ae8eca00a` +- PR 6 retained batch/JoinMap activation: `8e4b689f45` +- PR 7 expression activation: `c9ad0ea810` +- PR 8 spill/runtime-filter activation: `e072568998` +- PR 9 unified pressure recovery: `eec23dcdc5` +- PR 10 benchmark harness: `383fc6dce3` +- Owner-atomic activation and cleanup: `656e254fe6` + +The activated expression owner is deliberately closed: COL, literal, param, +variable, vector/fold, CONCAT, CASE, varchar EQUAL, and the audited string CAST +forms. Build and probe keys are checked together. If any HashBuild, HashJoin, +DedupJoin, or RightDedupJoin key contains another function family, automatic +activation keeps every participating operator in that local statement attempt +on the legacy path. This prevents a partially exact owner from retaining an +estimator-only expression rejection. It does not claim that generic JSON, +regexp, geometry, or spatial execution has already migrated. + +## Local correctness matrix + +All direct Go tests use the repository CGO wrapper so `usearch` links against +the locally built `thirdparties` artifacts. + +```text +.agents/skills/mo-dev/scripts/mo-cgo-test -count=1 \ + ./pkg/common/mpool ./pkg/common/hashmap \ + ./pkg/container/hashtable ./pkg/container/vector ./pkg/container/batch \ + ./pkg/sql/colexec ./pkg/sql/colexec/hashbuild \ + ./pkg/sql/colexec/hashjoin ./pkg/sql/colexec/dedupjoin \ + ./pkg/sql/colexec/rightdedupjoin ./pkg/sql/colexec/spillutil \ + ./pkg/sql/compile ./pkg/util/resource ./pkg/vm/message ./pkg/vm/process + +result: PASS +``` + +The same package matrix passes with `-race -p=2 -count=1`. A focused +`-race -p=2 -count=20` stress also passes for concurrent account alloc/free, +seal/open linearization, statement terminal one-shot behavior, owner-atomic +activation, accounted spill/reduction, and all three join-consumer activation +gates. Affected-package `go vet` and `make build` pass on the same candidate. + +Incident-mechanism regression mapping: + +| Incident | Local regression proof | +| --- | --- | +| #25782 high-cardinality HashBuild/spill | `TestAccountedInitialSpillReducesUnpublishedInputAndPreservesRows`, `TestShuffleHashBuildAccountedSpillLifecycle`, hash cell/descriptor replacement tests | +| #26174 dedup/fulltext INSERT | `TestAccountedDedupScratchAndDeleteBitmapFollowJoinMapLifetime`, `TestAccountedDedupBitmapExactBoundaryRollsBack` | +| #26192 LOAD DATA decoded/rebuild path | `TestSpillAllocationAccountDecodedBatchLifecycle`, `TestSpillAllocationAccountDecodedReuseRetriesFromCleanRecord`, recursive rebuild lifecycle test | +| #26413 external/self-join lifetime | statement-attempt zero/late-Free/cancel/error/panic tests, `TestAccountedJoinMapLateFreeKeepsOriginalGeneration` | +| #26454 string expression false budget | `TestAllocationAccountedExpressionIssue26454AndOneByteShort`, `TestIssue26454ExpressionKeyBuildUsesActualCapacity`, adaptive expression-pressure tests | + +The activation boundary itself is covered by +`TestHashBuildAllocationActivationRequiresClosedExpressionOwner` and +`TestAllocationAccountActivationIsStatementAtomic`: a closed #26454 +expression activates, an unclosed modulo expression keeps the owner legacy, a +no-map operator does not open a generation, and one unclosed owner rolls back +all already configured owners. + +## Local performance evidence + +Host: linux/amd64, Intel i7-11700, Go 1.26.4. Unless stated otherwise, +GOMAXPROCS was 16. Medians are from five runs; resident tests used a 200 ms +benchtime and spill scatter used 300 ms. + +### Resident HashBuild + +| Key/rows | Legacy median | Accounted median | Delta | Allocation effect | +| --- | ---: | ---: | ---: | --- | +| int / 32 | 4,161 ns | 5,841 ns | +40.4% (+1.68 us) | 32 -> 27 allocs/op | +| varchar / 32 | 7,147 ns | 8,907 ns | +24.6% (+1.76 us) | 96 -> 91 allocs/op | +| int / 8,192 | 139,602 ns | 143,254 ns | +2.62% | 43,738 -> 10,401 B/op; 56 -> 29 allocs/op | +| varchar / 8,192 | 410,074 ns | 400,782 ns | -2.27% | 222,055 -> 24,899 B/op; 600 -> 574 allocs/op | + +The 32-row cases quantify the fixed first-allocation tax for high-frequency TP +work; they are not presented as a per-row tax. Full-batch resident HashBuild is +within 2.7% for int keys and improves the measured varchar case. + +### Spill and expression paths + +| Benchmark | Legacy median | Accounted median | Delta | +| --- | ---: | ---: | ---: | +| 4,096-row scatter including hash/select/marshal/coalesce/write | 59,073 ns | 60,571 ns | +2.54% | +| #26454 expression | 909,082 ns | 940,454 ns | +3.45% | +| copied build batch | 18,584 ns | 22,051 ns | +18.7% | + +The copied-batch result also reduced 230,032 B/op to 632 B/op and 10 to 6 +allocations/op. Scatter reduced 2,711 B/op to 2,392 B/op; syscall and codec work +remain included. + +### Generation and release concurrency + +| Benchmark | CPU | Median | Recorded latency | +| --- | ---: | ---: | --- | +| full allocation-attempt lifecycle | 1 | 843.5 ns/op | p50 730 ns; p99 2,011 ns | +| full allocation-attempt lifecycle | 8 | 762.7 ns/op | p50 1,995 ns; p99 73,349 ns | +| same-generation release storm | 1 | 558.4 ns/op | 0 B/op, 0 allocs/op | +| same-generation release storm | 8 | 470.0 ns/op | 0 B/op, 0 allocs/op | + +The lifecycle benchmark includes generation open, controller/account creation, +one exact 4 KiB allocation/free, terminal completion, and concurrent operation. + +## Remote workload gate + +The candidate must be compared with the exact rebased main on the same TKE +workflow and resource configuration. Required results are: + +| Workload | Required evidence | Status | +| --- | --- | --- | +| TPCH 100G | three rounds, no spill/OOM/restart, total and per-query comparison | pending | +| TPCH 1T | three rounds, spill succeeds, no OOM/restart, total and per-query comparison | pending | +| #26174 fulltext INSERT | workload pass under unchanged cap | pending | +| #26192 LOAD DATA | workload pass under unchanged cap | pending | +| #26413 external self-join | workload pass and zero terminal generation | pending | + +No cap increase, reduced dataset, disabled spill, or plan-specific bypass is an +acceptable pass. The candidate and main run URLs, exact SHAs, load times, +query totals, spill evidence, restart/OOM search, and profile comparison belong +in this table before PR 10 is accepted. diff --git a/docs/design/evidence/26459_allocation_accounting_bench.txt b/docs/design/evidence/26459_allocation_accounting_bench.txt index 4ed578af4ed7d..522caa2537fcf 100644 --- a/docs/design/evidence/26459_allocation_accounting_bench.txt +++ b/docs/design/evidence/26459_allocation_accounting_bench.txt @@ -154,3 +154,31 @@ An eight-generation, GOMAXPROCS=8 acquire/release latency harness sampled was 100, 98, 97, 68, and 73 ns (97 ns median). The measurement includes the `time.Now`/`time.Since` sampling cost and is therefore a conservative end-to-end latency observation, not a cycle-level atomic benchmark. + +Production activation benchmark (rebased main 5b9eeb54ec, linux/amd64, +Go 1.26.4, i7-11700, GOMAXPROCS=16; sequential package runs): + +BenchmarkResidentHashBuildAccounting/legacy/int/rows-32 median: 4161 ns/op +BenchmarkResidentHashBuildAccounting/accounted/int/rows-32 median: 5841 ns/op +BenchmarkResidentHashBuildAccounting/legacy/varchar/rows-32 median: 7147 ns/op +BenchmarkResidentHashBuildAccounting/accounted/varchar/rows-32 median: 8907 ns/op +BenchmarkResidentHashBuildAccounting/legacy/int/rows-8192 median: 139602 ns/op +BenchmarkResidentHashBuildAccounting/accounted/int/rows-8192 median: 143254 ns/op +BenchmarkResidentHashBuildAccounting/legacy/varchar/rows-8192 median: 410074 ns/op +BenchmarkResidentHashBuildAccounting/accounted/varchar/rows-8192 median: 400782 ns/op +BenchmarkIssue26454ExpressionAccounting/legacy median: 909082 ns/op +BenchmarkIssue26454ExpressionAccounting/accounted median: 940454 ns/op +BenchmarkCopyBuildBatchAccounting/legacy median: 18584 ns/op, 230032 B/op, 10 allocs/op +BenchmarkCopyBuildBatchAccounting/accounted median: 22051 ns/op, 632 B/op, 6 allocs/op +BenchmarkSpillScatterAccounting/legacy median: 59073 ns/op +BenchmarkSpillScatterAccounting/accounted median: 60571 ns/op + +BenchmarkHashBuildAllocationAttemptLifecycle cpu=1 median: 843.5 ns/op, +p50 730 ns/op, p99 2011 ns/op, 256 B/op, 3 allocs/op +BenchmarkHashBuildAllocationAttemptLifecycle cpu=8 median: 762.7 ns/op, +p50 1995 ns/op, p99 73349 ns/op, 256 B/op, 3 allocs/op +BenchmarkHashBuildAllocationReleaseStorm cpu=1 median: 558.4 ns/op, 0 B/op, 0 allocs/op +BenchmarkHashBuildAllocationReleaseStorm cpu=8 median: 470.0 ns/op, 0 B/op, 0 allocs/op + +Commands and workload gates are recorded in +docs/design/evidence/26459_activation_validation.md. diff --git a/docs/rfcs/00000000_allocation_accounted_memory_admission.md b/docs/rfcs/00000000_allocation_accounted_memory_admission.md index 4a2be8adabd2f..bc0afeb50aa91 100644 --- a/docs/rfcs/00000000_allocation_accounted_memory_admission.md +++ b/docs/rfcs/00000000_allocation_accounted_memory_admission.md @@ -1,7 +1,8 @@ -- Status: draft +- Status: implementation validation - Start Date: 2026-07-30 - Authors: aptend -- Implementation PR: TBD +- Implementation candidate: `feature/26459-statement-lifecycle` at + `656e254fe6` - Issue for this RFC: [#26459](https://github.com/matrixorigin/matrixone/issues/26459) - Implementation plan: @@ -49,6 +50,15 @@ The first consumer is HashBuild and the joins that share its spill lifecycle. The accounting primitive is deliberately defined below the SQL operator layer so other spillable operators can adopt the same model later. +The current activation is deliberately owner-closed. Build and probe trees +from the audited COL/literal/CONCAT/CASE/varchar-EQUAL/string-CAST set activate +exact accounting for the whole local HashBuild/join owner set. A tree +containing an unclosed generic function family keeps all participating owners +in that local attempt on the legacy path; exact map/batch ownership is never +mixed with an estimator-gated expression. The remaining generic-function +migration is therefore explicit rather than being misreported as RFC +completion. + ## Motivation ### Confirmed failures @@ -714,13 +724,12 @@ suspension check succeeds in that transaction. Once suspension publication linearizes, no later open may publish; opens that linearized earlier are the finite active set allowed to finish or become tombstones. -Current `SetStmtProfile` turnover, frontend `StatementInfo.EndStatement`, and -`HashBuildBudgetGeneration.Close` do not prove this per-CN quiescence or -validate zero. Production activation is blocked until the `Compile` attempt -owns an explicit post-pipeline/MessageBoard-close transition for success, -failure, panic, cancellation, retry, broadcast, prepared reuse, and remote -execution. A forced close must not silently zero accounting while allocations -remain live. +`SetStmtProfile` turnover, frontend `StatementInfo.EndStatement`, and +`HashBuildBudgetGeneration.Close` alone do not prove this per-CN quiescence or +validate zero. The implementation therefore gives the `Compile` attempt an +explicit post-pipeline/MessageBoard-close transition for success, failure, +panic, cancellation, retry, broadcast, prepared reuse, and remote execution. +A forced close never silently zeros accounting while allocations remain live. Ownership transfer does not change generations. If a transfer would cross to a different generation, it must either: From 3e0cb4fed2799c6a5a2dd2cd88283f5a80f0edd6 Mon Sep 17 00:00:00 2001 From: aptend Date: Sat, 1 Aug 2026 00:48:28 +0800 Subject: [PATCH 31/61] executor: close accounted vector ownership boundaries --- pkg/common/bitmap/bitmap.go | 176 +++++++++++++++++- pkg/common/bitmap/bitmap_test.go | 54 ++++++ pkg/common/mpool/allocation_account.go | 99 +++++++--- pkg/common/mpool/mpool.go | 14 +- .../bytejson/bytejson_text_writer.go | 15 +- pkg/container/nulls/nulls.go | 26 +++ pkg/container/nulls/nulls_test.go | 53 ++++++ pkg/container/vector/allocation_account.go | 24 +-- .../vector/allocation_account_test.go | 97 ++++++++++ pkg/container/vector/pSpoolTools.go | 21 +-- pkg/container/vector/vector.go | 137 +++++++++++--- pkg/container/vector/versions.go | 7 +- pkg/sql/colexec/aggexec/maxby.go | 10 +- pkg/sql/colexec/aggexec/maxby_test.go | 42 +++++ pkg/sql/colexec/hashbuild/pressure.go | 13 +- pkg/sql/plan/function/func_unary.go | 4 +- pkg/sql/util/copy_batch_test.go | 43 +++++ pkg/sql/util/util.go | 14 +- 18 files changed, 744 insertions(+), 105 deletions(-) diff --git a/pkg/common/bitmap/bitmap.go b/pkg/common/bitmap/bitmap.go index 06603cdac077b..6ba3fc586a91f 100644 --- a/pkg/common/bitmap/bitmap.go +++ b/pkg/common/bitmap/bitmap.go @@ -23,6 +23,7 @@ import ( "math/bits" "unsafe" + "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/container/types" ) @@ -461,6 +462,173 @@ func (n *Bitmap) TryExpandWithSize(size int) { } } +// RemapOrdered rewrites the bitmap in place for an ordered row selection. +// When negate is false, output row i comes from sels[i]. When negate is true, +// sels identifies rows to remove. The caller must provide strictly increasing, +// non-negative row indexes. Because every destination row is at or before its +// source row, one cached source word is sufficient to avoid allocating a +// second data-scaled bitmap. Selection rows beyond the bitmap's logical +// length are valid and read as clear: a null bitmap may be shorter than its +// owning vector when the vector's trailing rows are all non-null. +func (n *Bitmap) RemapOrdered(sels []int64, negate bool) { + if n == nil { + return + } + oldLength := n.logicalLen() + previous := int64(-1) + for _, sel := range sels { + if sel <= previous || sel < 0 { + panic("bitmap ordered remap requires strictly increasing non-negative rows") + } + previous = sel + } + logicalLength := int64(len(sels)) + if negate { + logicalLength = oldLength + } + n.prepareOrderedRemap(logicalLength) + + sourceWordIndex := int64(-1) + var sourceWord uint64 + readSource := func(row int64) bool { + if row >= oldLength { + return false + } + wordIndex := row >> 6 + if wordIndex != sourceWordIndex { + sourceWordIndex = wordIndex + sourceWord = n.data[wordIndex] + } + return sourceWord&(uint64(1)<> 6 + mask := uint64(1) << uint(row&63) + if value { + n.data[wordIndex] |= mask + } else { + n.data[wordIndex] &^= mask + } + } + + output := int64(0) + if !negate { + for _, sel := range sels { + writeDestination(output, readSource(sel)) + output++ + } + } else { + selIndex := 0 + for source := int64(0); source < oldLength; source++ { + if selIndex < len(sels) && source == sels[selIndex] { + selIndex++ + continue + } + writeDestination(output, readSource(source)) + output++ + } + } + n.finishOrderedRemap(output, logicalLength) +} + +// RemapMaskOrdered is RemapOrdered for an ordered bitmap selection. Selection +// bitmap iteration is monotonic, so the rewrite uses no row-scaled scratch. +func (n *Bitmap) RemapMaskOrdered(sels *Bitmap, negate bool) { + if n == nil || sels == nil { + return + } + oldLength := n.logicalLen() + logicalLength := int64(sels.Count()) + if negate { + logicalLength = oldLength + } + n.prepareOrderedRemap(logicalLength) + sourceWordIndex := int64(-1) + var sourceWord uint64 + readSource := func(row int64) bool { + if row >= oldLength { + return false + } + wordIndex := row >> 6 + if wordIndex != sourceWordIndex { + sourceWordIndex = wordIndex + sourceWord = n.data[wordIndex] + } + return sourceWord&(uint64(1)<> 6 + mask := uint64(1) << uint(row&63) + if value { + n.data[wordIndex] |= mask + } else { + n.data[wordIndex] &^= mask + } + } + + output := int64(0) + iterator := sels.Iterator() + if !negate { + for iterator.HasNext() { + source := int64(iterator.Next()) + writeDestination(output, readSource(source)) + output++ + } + } else { + var selected int64 = -1 + if iterator.HasNext() { + selected = int64(iterator.Next()) + } + for source := int64(0); source < oldLength; source++ { + if source == selected { + if iterator.HasNext() { + selected = int64(iterator.Next()) + } else { + selected = -1 + } + continue + } + writeDestination(output, readSource(source)) + output++ + } + } + n.finishOrderedRemap(output, logicalLength) +} + +func (n *Bitmap) prepareOrderedRemap(logicalLength int64) { + words := int((logicalLength + 63) / 64) + if words > cap(n.data) { + panic("bitmap external storage capacity exceeded") + } + if words > len(n.data) { + storage := n.data[:cap(n.data)] + clear(storage[len(n.data):words]) + n.data = storage[:words] + } +} + +func (n *Bitmap) finishOrderedRemap(written, logicalLength int64) { + words := int((logicalLength + 63) / 64) + if written < logicalLength { + word := int(written >> 6) + if tail := uint(written & 63); tail != 0 { + n.data[word] &= (uint64(1) << tail) - 1 + word++ + } + clear(n.data[word:words]) + } + if words > 0 && logicalLength&63 != 0 { + n.data[words-1] &= (uint64(1) << uint(logicalLength&63)) - 1 + } + clear(n.data[words:]) + n.data = n.data[:words] + n.setLogicalLen(logicalLength) + n.count = 0 + for _, word := range n.data { + n.count += int64(bits.OnesCount64(word)) + } +} + func (n *Bitmap) Filter(sels []int64) *Bitmap { var b Bitmap b.InitWithSize(n.logicalLen()) @@ -530,13 +698,13 @@ func DecodeMarshalHeader(data []byte) ( rawBitLength > math.MaxInt64 || rawDataSize > math.MaxInt || rawDataSize%8 != 0 { - return 0, 0, 0, fmt.Errorf("invalid bitmap wire header") + return 0, 0, 0, moerr.NewInvalidInputNoCtx("invalid bitmap wire header") } bitLength = int64(rawBitLength) dataSize = int(rawDataSize) if count > bitLength || uint64(dataSize/8) != (rawBitLength+63)/64 { - return 0, 0, 0, fmt.Errorf("invalid bitmap wire header") + return 0, 0, 0, moerr.NewInvalidInputNoCtx("invalid bitmap wire header") } return count, bitLength, dataSize, nil } @@ -548,7 +716,7 @@ func (n *Bitmap) PrepareExternalUnmarshal( totalSize int, ) ([]byte, error) { if !n.HasExternalStorage() { - return nil, fmt.Errorf("bitmap does not use external storage") + return nil, moerr.NewInvalidInputNoCtx("bitmap does not use external storage") } count, bitLength, dataSize, err := DecodeMarshalHeader(header) if err != nil { @@ -556,7 +724,7 @@ func (n *Bitmap) PrepareExternalUnmarshal( } if totalSize != MarshalHeaderSize+dataSize || dataSize/8 > cap(n.data) { - return nil, fmt.Errorf("invalid bitmap external storage capacity") + return nil, moerr.NewInvalidInputNoCtx("invalid bitmap external storage capacity") } storage := n.data[:cap(n.data)] clear(storage) diff --git a/pkg/common/bitmap/bitmap_test.go b/pkg/common/bitmap/bitmap_test.go index 39a55ac82366d..131e764d258dc 100644 --- a/pkg/common/bitmap/bitmap_test.go +++ b/pkg/common/bitmap/bitmap_test.go @@ -270,6 +270,60 @@ func TestBitmapExternalStorageLifecycle(t *testing.T) { require.True(t, value.Contains(128)) } +func TestBitmapRemapOrdered(t *testing.T) { + for _, external := range []bool{false, true} { + t.Run(fmt.Sprintf("external=%t", external), func(t *testing.T) { + var value Bitmap + if external { + value.InstallExternalStorage(make([]uint64, 3)) + } + value.InitWithSize(130) + value.AddMany([]uint64{0, 2, 63, 64, 65, 128, 129}) + + value.RemapOrdered([]int64{0, 2, 64, 65, 129}, false) + require.Equal(t, int64(5), value.Len()) + require.Equal(t, 5, value.Count()) + for row := uint64(0); row < 5; row++ { + require.True(t, value.Contains(row)) + } + + value.InitWithSize(130) + value.AddMany([]uint64{0, 2, 63, 64, 65, 128, 129}) + value.RemapOrdered([]int64{1, 63, 128}, true) + require.Equal(t, int64(130), value.Len()) + require.Equal(t, 5, value.Count()) + require.True(t, value.Contains(0)) + require.True(t, value.Contains(1)) + require.True(t, value.Contains(62)) + require.True(t, value.Contains(63)) + require.True(t, value.Contains(126)) + }) + } +} + +func TestBitmapRemapMaskOrdered(t *testing.T) { + var value Bitmap + value.InitWithSize(130) + value.AddMany([]uint64{1, 63, 64, 127, 129}) + + selection := newBm(130) + selection.AddMany([]uint64{1, 64, 129}) + value.RemapMaskOrdered(selection, false) + require.Equal(t, int64(3), value.Len()) + require.Equal(t, 3, value.Count()) + require.True(t, value.Contains(0)) + require.True(t, value.Contains(1)) + require.True(t, value.Contains(2)) + + value.InitWithSize(130) + value.AddMany([]uint64{1, 63, 64, 127, 129}) + value.RemapMaskOrdered(selection, true) + require.Equal(t, int64(130), value.Len()) + require.Equal(t, 2, value.Count()) + require.True(t, value.Contains(62)) + require.True(t, value.Contains(125)) +} + func TestBitmapExternalStorageUnmarshal(t *testing.T) { source := newBm(128) source.Add(1) diff --git a/pkg/common/mpool/allocation_account.go b/pkg/common/mpool/allocation_account.go index 8923056f74389..8a9d896c16ce1 100644 --- a/pkg/common/mpool/allocation_account.go +++ b/pkg/common/mpool/allocation_account.go @@ -38,23 +38,72 @@ const ( ) var ( - ErrAllocationAccountCapacity = errors.New("allocation account capacity exceeded") - ErrAllocationAccountSealed = errors.New("allocation account is sealed") - ErrAllocationAccountInvalid = errors.New("invalid allocation account") - ErrAllocationAccountStale = errors.New("stale allocation account handle") - ErrAllocationAccountMismatch = errors.New("allocation account ownership mismatch") - ErrAllocationAllocatorLimit = errors.New("allocation exceeds allocator size limit") - ErrAllocationAccountInvariant = errors.New( - "allocation account invariant failure", - ) - ErrAllocationAdmissionSuspended = errors.New( - "allocation account admission is suspended", - ) - ErrAllocationMetadataSlots = errors.New("allocation metadata slots exhausted") - ErrAllocationGenerationSlots = errors.New("allocation account generation slots exhausted") - ErrAllocationAccountLive = errors.New("allocation account still owns memory") + ErrAllocationAccountCapacity error = allocationAccountSentinel("allocation account capacity exceeded") + ErrAllocationAccountSealed error = allocationAccountSentinel("allocation account is sealed") + ErrAllocationAccountInvalid error = allocationAccountSentinel("invalid allocation account") + ErrAllocationAccountStale error = allocationAccountSentinel("stale allocation account handle") + ErrAllocationAccountMismatch error = allocationAccountSentinel("allocation account ownership mismatch") + ErrAllocationAllocatorLimit error = allocationAccountSentinel("allocation exceeds allocator size limit") + ErrAllocationAccountInvariant error = allocationAccountSentinel("allocation account invariant failure") + ErrAllocationAdmissionSuspended error = allocationAccountSentinel("allocation account admission is suspended") + ErrAllocationMetadataSlots error = allocationAccountSentinel("allocation metadata slots exhausted") + ErrAllocationGenerationSlots error = allocationAccountSentinel("allocation account generation slots exhausted") + ErrAllocationAccountLive error = allocationAccountSentinel("allocation account still owns memory") ) +type allocationAccountSentinel string + +func (e allocationAccountSentinel) Error() string { return string(e) } + +type allocationAccountDetailError struct { + cause error + detail string + detailFirst bool +} + +func (e *allocationAccountDetailError) Error() string { + if e == nil || e.cause == nil { + return "allocation account error" + } + if e.detail == "" { + return e.cause.Error() + } + if e.detailFirst { + return e.detail + ": " + e.cause.Error() + } + return e.cause.Error() + ": " + e.detail +} + +func prefixAllocationAccountError( + cause error, + format string, + args ...any, +) error { + return &allocationAccountDetailError{ + cause: cause, + detail: fmt.Sprintf(format, args...), + detailFirst: true, + } +} + +func (e *allocationAccountDetailError) Unwrap() error { + if e == nil { + return nil + } + return e.cause +} + +func wrapAllocationAccountError( + cause error, + format string, + args ...any, +) error { + return &allocationAccountDetailError{ + cause: cause, + detail: fmt.Sprintf(format, args...), + } +} + const ( allocationAccountSealedBit = uint64(1) << 63 allocationAccountUsedMask = allocationAccountSealedBit - 1 @@ -228,9 +277,9 @@ func (a *AllocationAccount) ValidateRollback( return ErrAllocationAccountInvalid } if checkpoint.Handle != a.handle { - return fmt.Errorf( - "%w: checkpoint=%d account=%d", + return wrapAllocationAccountError( ErrAllocationAccountMismatch, + "checkpoint=%d account=%d", checkpoint.Handle, a.handle, ) @@ -240,9 +289,9 @@ func (a *AllocationAccount) ValidateRollback( return ErrAllocationAccountSealed } if snapshot.Used != checkpoint.Used { - return fmt.Errorf( - "%w: checkpoint-used=%d current-used=%d", + return wrapAllocationAccountError( ErrAllocationAccountInvariant, + "checkpoint-used=%d current-used=%d", checkpoint.Used, snapshot.Used, ) @@ -340,9 +389,9 @@ func newAllocationAccountCapacityError( requested uint64, limit uint64, ) error { - return fmt.Errorf( - "%w: used=%d requested=%d limit=%d", + return wrapAllocationAccountError( ErrAllocationAccountCapacity, + "used=%d requested=%d limit=%d", used, requested, limit, @@ -530,9 +579,9 @@ func (r *AllocationAccountRegistry) CompleteTerminalWithError( return AllocationAccountTerminalSnapshot{ AllocationAccountSnapshot: current, State: AllocationAccountTerminalInvariantFailure, - }, false, fmt.Errorf( - "%w: terminal account is not quiescent", + }, false, wrapAllocationAccountError( ErrAllocationAccountInvariant, + "terminal account is not quiescent", ) } snapshot = AllocationAccountTerminalSnapshot{ @@ -576,9 +625,9 @@ func (r *AllocationAccountRegistry) CompleteTerminalWithError( func newAllocationTerminalInvariantError( snapshot AllocationAccountTerminalSnapshot, ) error { - return fmt.Errorf( - "%w: handle=%d used=%d peak=%d limit=%d owner=%d site=%d live-allocations=%d", + return wrapAllocationAccountError( ErrAllocationAccountInvariant, + "handle=%d used=%d peak=%d limit=%d owner=%d site=%d live-allocations=%d", snapshot.Handle, snapshot.Used, snapshot.Peak, diff --git a/pkg/common/mpool/mpool.go b/pkg/common/mpool/mpool.go index e3d5da9bf8284..ecc96bc4a0abc 100644 --- a/pkg/common/mpool/mpool.go +++ b/pkg/common/mpool/mpool.go @@ -816,9 +816,9 @@ func (mp *MPool) allocAccountedWithDetailK( // reject unexpected alloc size. if sz < 0 || sz > maxAllocationSize() { logutil.Errorf("mpool memory allocation exceed limit with requested size %d: %s", sz, string(debug.Stack())) - return nil, allocationAccountSiteError(request, fmt.Errorf( - "%w: requested=%d maximum=%d", + return nil, allocationAccountSiteError(request, wrapAllocationAccountError( ErrAllocationAllocatorLimit, + "requested=%d maximum=%d", sz, maxAllocationSize(), )) @@ -837,11 +837,11 @@ func allocationAccountSiteError( request allocationAccountRequest, err error, ) error { - return fmt.Errorf( - "allocation owner=%d site=%d: %w", + return prefixAllocationAccountError( + err, + "allocation owner=%d site=%d", request.owner, request.site, - err, ) } @@ -1369,9 +1369,9 @@ func MakeSliceAccounted[T any]( if elementSize == 0 || maxSize <= 0 || uint64(n) > uint64(maxSize)/uint64(elementSize) { - return nil, fmt.Errorf( - "%w: elements=%d element-size=%d maximum=%d", + return nil, wrapAllocationAccountError( ErrAllocationAllocatorLimit, + "elements=%d element-size=%d maximum=%d", n, elementSize, maxSize, diff --git a/pkg/container/bytejson/bytejson_text_writer.go b/pkg/container/bytejson/bytejson_text_writer.go index d91674766654f..79d10ca28a8f0 100644 --- a/pkg/container/bytejson/bytejson_text_writer.go +++ b/pkg/container/bytejson/bytejson_text_writer.go @@ -16,11 +16,12 @@ package bytejson import ( "encoding/base64" - "fmt" "io" "math" "strconv" "unicode/utf8" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" ) // WriteJSONText writes the visible JSON representation without allocating a @@ -72,7 +73,7 @@ func WriteJSONText(w io.Writer, value ByteJson) error { return writeBytes(w, strconv.AppendUint(buf[:0], value.GetUint64(), 10)) case TpCodeLiteral: if len(value.Data) == 0 { - return fmt.Errorf("invalid JSON literal") + return moerr.NewInvalidInputNoCtx("invalid JSON literal") } switch value.Data[0] { case LiteralNull: @@ -82,12 +83,12 @@ func WriteJSONText(w io.Writer, value ByteJson) error { case LiteralFalse: return writeString(w, "false") default: - return fmt.Errorf("invalid JSON literal %d", value.Data[0]) + return moerr.NewInvalidInputNoCtxf("invalid JSON literal %d", value.Data[0]) } case TpCodeFloat64: f := value.GetFloat64() if math.IsInf(f, 0) || math.IsNaN(f) { - return fmt.Errorf("invalid JSON float64 %f", f) + return moerr.NewInvalidInputNoCtxf("invalid JSON float64 %f", f) } format := byte('e') abs := math.Abs(f) @@ -125,7 +126,7 @@ func WriteJSONText(w io.Writer, value ByteJson) error { } return writeByte(w, '"') default: - return fmt.Errorf("invalid JSON type %d", value.Type) + return moerr.NewInvalidInputNoCtxf("invalid JSON type %d", value.Type) } } @@ -214,7 +215,7 @@ func WriteJSONString(w io.Writer, value []byte) error { } _, size := utf8.DecodeRune(value[offset:]) if size == 1 { - return fmt.Errorf("invalid UTF-8") + return moerr.NewInvalidInputNoCtx("invalid UTF-8") } offset += size } @@ -229,7 +230,7 @@ func writeNormalizedBase64(w io.Writer, encoded []byte) error { for offset := 0; offset < len(encoded); { n, next, ok := decodeBase64Chunk(encoded, offset, decoded[:]) if !ok { - return fmt.Errorf("invalid base64 JSON value") + return moerr.NewInvalidInputNoCtx("invalid base64 JSON value") } if err := writeRawBase64(w, decoded[:n]); err != nil { return err diff --git a/pkg/container/nulls/nulls.go b/pkg/container/nulls/nulls.go index fc16f895b6265..e01555d7bd531 100644 --- a/pkg/container/nulls/nulls.go +++ b/pkg/container/nulls/nulls.go @@ -281,6 +281,19 @@ func Filter(nsp *Nulls, sels []int64, negate bool) { } } +// FilterInPlaceOrdered preserves Filter semantics for Vector.Shrink's ordered +// selection contract without allocating a second row-scaled bitmap. +func FilterInPlaceOrdered(nsp *Nulls, sels []int64, negate bool) { + if nsp.np.EmptyByFlag() { + return + } + if !nsp.np.HasExternalStorage() { + Filter(nsp, sels, negate) + return + } + nsp.np.RemapOrdered(sels, negate) +} + func FilterByMask(nsp *Nulls, sels *bitmap.Bitmap, negate bool) { if nsp.np.EmptyByFlag() { return @@ -332,6 +345,19 @@ func FilterByMask(nsp *Nulls, sels *bitmap.Bitmap, negate bool) { } } +// FilterByMaskInPlace rewrites a null bitmap using the selection bitmap's +// naturally ordered iterator and therefore requires no row-scaled scratch. +func FilterByMaskInPlace(nsp *Nulls, sels *bitmap.Bitmap, negate bool) { + if nsp.np.EmptyByFlag() { + return + } + if !nsp.np.HasExternalStorage() { + FilterByMask(nsp, sels, negate) + return + } + nsp.np.RemapMaskOrdered(sels, negate) +} + // XXX This emptyFlag thing is broken -- it simply cannot be used concurrently. // Make any an alias of EmptyByFlag, otherwise there will be hell lots of race conditions. func (nsp *Nulls) Any() bool { diff --git a/pkg/container/nulls/nulls_test.go b/pkg/container/nulls/nulls_test.go index f4f4f895434e9..b3ef3cb4bc53e 100644 --- a/pkg/container/nulls/nulls_test.go +++ b/pkg/container/nulls/nulls_test.go @@ -16,8 +16,10 @@ package nulls import ( "bytes" + "fmt" "testing" + "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -230,6 +232,57 @@ func TestFilter(t *testing.T) { }) } +func TestFilterInPlaceOrderedMatchesFilter(t *testing.T) { + source := Build(130, 0, 2, 63, 64, 65, 128, 129) + for _, test := range []struct { + name string + sels []int64 + negate bool + }{ + {"select-across-words", []int64{0, 2, 64, 65, 129}, false}, + {"select-trailing-clear", []int64{1, 63, 130, 131}, false}, + {"remove-across-words", []int64{1, 63, 128}, true}, + {"remove-after-bitmap", []int64{130, 131}, true}, + } { + t.Run(test.name, func(t *testing.T) { + legacy := source.Clone() + Filter(legacy, test.sels, test.negate) + + var inPlace Nulls + inPlace.GetBitmap().InstallExternalStorage(make([]uint64, 4)) + inPlace.InitWith(source) + FilterInPlaceOrdered(&inPlace, test.sels, test.negate) + + require.True(t, inPlace.GetBitmap().IsSame(legacy.GetBitmap())) + require.Equal(t, legacy.GetBitmap().Len(), inPlace.GetBitmap().Len()) + require.Equal(t, legacy.Count(), inPlace.Count()) + }) + } +} + +func TestFilterByMaskInPlaceMatchesFilter(t *testing.T) { + source := Build(130, 0, 2, 63, 64, 65, 128, 129) + for _, negate := range []bool{false, true} { + t.Run(fmt.Sprintf("negate=%t", negate), func(t *testing.T) { + var selection bitmap.Bitmap + selection.InitWithSize(132) + selection.AddMany([]uint64{1, 63, 128, 130, 131}) + + legacy := source.Clone() + FilterByMask(legacy, &selection, negate) + + var inPlace Nulls + inPlace.GetBitmap().InstallExternalStorage(make([]uint64, 4)) + inPlace.InitWith(source) + FilterByMaskInPlace(&inPlace, &selection, negate) + + require.True(t, inPlace.GetBitmap().IsSame(legacy.GetBitmap())) + require.Equal(t, legacy.GetBitmap().Len(), inPlace.GetBitmap().Len()) + require.Equal(t, legacy.Count(), inPlace.Count()) + }) + } +} + func TestMerge(t *testing.T) { t.Run("merge test", func(t *testing.T) { var n, m Nulls diff --git a/pkg/container/vector/allocation_account.go b/pkg/container/vector/allocation_account.go index 43ee95db45df7..eaaebaa2410cd 100644 --- a/pkg/container/vector/allocation_account.go +++ b/pkg/container/vector/allocation_account.go @@ -15,7 +15,7 @@ package vector import ( - "fmt" + "errors" "io" "math" @@ -25,6 +25,13 @@ import ( "github.com/matrixorigin/matrixone/pkg/container/types" ) +func allocationAccountInvalid(message string) error { + return errors.Join( + mpool.ErrAllocationAccountInvalid, + moerr.NewInternalErrorNoCtx(message), + ) +} + // AllocationAccountSelection is an immutable choice for the first owned // off-heap allocations of a Vector. The physical MPool allocation metadata // remains the sole owner of the resulting charge. @@ -260,9 +267,8 @@ func (v *Vector) CanSetAllocationAccount( return err } if !v.offHeap { - return fmt.Errorf( - "%w: allocation-accounted vector must be off-heap", - mpool.ErrAllocationAccountInvalid, + return allocationAccountInvalid( + "allocation-accounted vector must be off-heap", ) } } @@ -270,10 +276,7 @@ func (v *Vector) CanSetAllocationAccount( return nil } if v.hasBackingStorage() { - return fmt.Errorf( - "%w: vector already has backing storage", - mpool.ErrAllocationAccountInvalid, - ) + return allocationAccountInvalid("vector already has backing storage") } return nil } @@ -426,9 +429,8 @@ func (v *Vector) allocOwned( return mp.Alloc(size, offHeap) } if !offHeap { - return nil, fmt.Errorf( - "%w: accounted allocation must be off-heap", - mpool.ErrAllocationAccountInvalid, + return nil, allocationAccountInvalid( + "accounted allocation must be off-heap", ) } site := v.allocationAccount.areaSite diff --git a/pkg/container/vector/allocation_account_test.go b/pkg/container/vector/allocation_account_test.go index 9594112db793a..ac132febac91f 100644 --- a/pkg/container/vector/allocation_account_test.go +++ b/pkg/container/vector/allocation_account_test.go @@ -313,6 +313,103 @@ func TestVectorAllocationAccountBitmapResetReuseAndFree(t *testing.T) { finalizeTestVectorAllocationAccount(t, state) } +func TestVectorAllocationAccountBitmapShrinkUsesNoScratch(t *testing.T) { + state := newTestVectorBitmapAllocationAccount(t, 8<<20, 16) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, vec.PreExtend(130, mp)) + for i := range 130 { + require.NoError(t, AppendFixed(vec, int64(i), false, mp)) + } + for _, row := range []uint64{0, 2, 63, 64, 129} { + vec.SetNull(row) + } + for _, row := range []uint64{1, 65, 128} { + vec.GetGrouping().Add(row) + } + before := state.account.Snapshot() + + vec.Shrink([]int64{0, 2, 64, 65, 129}, false) + require.Equal(t, []int64{0, 2, 64, 65, 129}, MustFixedColWithTypeCheck[int64](vec)) + for _, row := range []uint64{0, 1, 2, 4} { + require.True(t, vec.IsNull(row)) + } + require.Equal(t, 4, vec.GetNulls().Count()) + require.True(t, vec.GetGrouping().Contains(3)) + require.Equal(t, 1, vec.GetGrouping().Count()) + after := state.account.Snapshot() + require.Equal(t, before.Used, after.Used) + require.Equal(t, before.Peak, after.Peak) + + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountBitmapShuffleAccountsScratch(t *testing.T) { + state := newTestVectorBitmapAllocationAccount(t, 8<<20, 16) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, vec.PreExtend(130, mp)) + for i := range 130 { + require.NoError(t, AppendFixed(vec, int64(i), false, mp)) + } + for _, row := range []uint64{1, 64, 129} { + vec.SetNull(row) + } + for _, row := range []uint64{2, 63, 128} { + vec.GetGrouping().Add(row) + } + before := state.account.Snapshot() + + require.NoError(t, vec.Shuffle([]int64{129, 1, 64, 1, 2, 128, 63}, mp)) + require.Equal(t, []int64{129, 1, 64, 1, 2, 128, 63}, MustFixedColWithTypeCheck[int64](vec)) + for _, row := range []uint64{0, 1, 2, 3} { + require.True(t, vec.IsNull(row)) + } + require.Equal(t, 4, vec.GetNulls().Count()) + for _, row := range []uint64{4, 5, 6} { + require.True(t, vec.GetGrouping().Contains(row)) + } + require.Equal(t, 3, vec.GetGrouping().Count()) + after := state.account.Snapshot() + require.Greater(t, after.Peak, before.Peak) + require.Equal(t, uint64(3), state.registry.LiveAllocationMetadata()) + + var goScratch []byte + require.NoError(t, vec.ShuffleWithBuf([]int64{6, 5, 4, 3, 2, 1, 0}, mp, &goScratch)) + require.Nil(t, goScratch) + require.Equal(t, []int64{63, 128, 2, 1, 64, 1, 129}, MustFixedColWithTypeCheck[int64](vec)) + + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountBitmapShuffleFailurePreservesVector(t *testing.T) { + state := newTestVectorBitmapAllocationAccount(t, 8<<20, 4) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, vec.PreExtend(130, mp)) + for i := range 130 { + require.NoError(t, AppendFixed(vec, int64(i), false, mp)) + } + vec.SetNull(1) + vec.GetGrouping().Add(2) + before := state.account.Snapshot() + + err := vec.Shuffle([]int64{2, 1, 0}, mp) + require.ErrorIs(t, err, mpool.ErrAllocationMetadataSlots) + require.Equal(t, before.Used, state.account.Snapshot().Used) + require.Equal(t, uint64(3), state.registry.LiveAllocationMetadata()) + require.Equal(t, 130, vec.Length()) + require.True(t, vec.IsNull(1)) + require.True(t, vec.GetGrouping().Contains(2)) + require.Equal(t, int64(0), MustFixedColWithTypeCheck[int64](vec)[0]) + require.Equal(t, int64(129), MustFixedColWithTypeCheck[int64](vec)[129]) + + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + func TestVectorAllocationAccountBitmapGrowthFailurePreservesOwner(t *testing.T) { state := newTestVectorBitmapAllocationAccount(t, 1000, 8) mp := mpool.MustNewZero() diff --git a/pkg/container/vector/pSpoolTools.go b/pkg/container/vector/pSpoolTools.go index c62c482f337a9..d4b46d40583b7 100644 --- a/pkg/container/vector/pSpoolTools.go +++ b/pkg/container/vector/pSpoolTools.go @@ -14,11 +14,7 @@ package vector -import ( - "fmt" - - "github.com/matrixorigin/matrixone/pkg/common/mpool" -) +import "github.com/matrixorigin/matrixone/pkg/common/mpool" // DetachedBuffer transfers one owned Vector backing allocation through the // pipeline spool without losing its immutable allocation provenance. @@ -122,24 +118,21 @@ func (b *DetachedBuffer) AttachTo( kind DetachedBufferKind, ) error { if !b.CanAttachTo(v, kind) { - return fmt.Errorf( - "%w: detached vector buffer provenance mismatch", - mpool.ErrAllocationAccountInvalid, + return allocationAccountInvalid( + "detached vector buffer provenance mismatch", ) } if kind == DetachedAreaBuffer { if cap(v.area) != 0 { - return fmt.Errorf( - "%w: vector area already has backing storage", - mpool.ErrAllocationAccountInvalid, + return allocationAccountInvalid( + "vector area already has backing storage", ) } v.area = b.data } else { if cap(v.data) != 0 { - return fmt.Errorf( - "%w: vector data already has backing storage", - mpool.ErrAllocationAccountInvalid, + return allocationAccountInvalid( + "vector data already has backing storage", ) } v.data = b.data[:cap(b.data)] diff --git a/pkg/container/vector/vector.go b/pkg/container/vector/vector.go index 20691e3a43239..a309e012bfbac 100644 --- a/pkg/container/vector/vector.go +++ b/pkg/container/vector/vector.go @@ -920,9 +920,8 @@ func (v *Vector) UnmarshalBinaryTrusted(data []byte) error { func (v *Vector) unmarshalBinary(data []byte, validateValues bool) error { if v.allocationAccount != nil { - return fmt.Errorf( - "%w: cannot install aliases in an accounted vector", - mpool.ErrAllocationAccountInvalid, + return allocationAccountInvalid( + "cannot install aliases in an accounted vector", ) } read := func(size int) ([]byte, error) { @@ -1133,9 +1132,8 @@ func canonicalVectorTypeSize(typ types.Type) (int, error) { func (v *Vector) UnmarshalBinaryWithCopy(data []byte, mp *mpool.MPool) error { if v.allocationAccount != nil && v.hasBackingStorage() { - return fmt.Errorf( - "%w: cannot replace accounted vector storage without Free", - mpool.ErrAllocationAccountInvalid, + return allocationAccountInvalid( + "cannot replace accounted vector storage without Free", ) } var err error @@ -1348,9 +1346,8 @@ func (v *Vector) PreExtendWithArea(rows int, extraAreaSize int, mp *mpool.MPool) // Dup use to copy an identical vector func (v *Vector) Dup(mp *mpool.MPool) (*Vector, error) { if v.allocationAccount != nil { - return nil, fmt.Errorf( - "%w: accounted vector duplication requires an off-heap destination", - mpool.ErrAllocationAccountInvalid, + return nil, allocationAccountInvalid( + "accounted vector duplication requires an off-heap destination", ) } return v.dup(mp, false, v.offHeap, nil) @@ -1453,16 +1450,16 @@ func (v *Vector) dup( // unreferenced bytes in area are not propagated into batch memory accounting. func (v *Vector) CloneToFlatCompact(mp *mpool.MPool) (*Vector, error) { if v.allocationAccount != nil { - return nil, fmt.Errorf( - "%w: accounted compact clone requires a destination selection", - mpool.ErrAllocationAccountInvalid, + return nil, allocationAccountInvalid( + "accounted compact clone requires a destination selection", ) } return v.cloneToFlatCompact(mp, nil) } // CloneToFlatCompactWithAllocation creates an off-heap compact copy under the -// explicit destination selection. +// explicit destination selection. Passing nil creates an unaccounted +// destination and is reserved for a deliberate ownership boundary. func (v *Vector) CloneToFlatCompactWithAllocation( mp *mpool.MPool, selection *AllocationAccountSelection, @@ -1783,6 +1780,12 @@ func (v *Vector) ShuffleWithBuf(sels []int64, mp *mpool.MPool, buf *[]byte) (err if v.IsConst() { return nil } + // The reusable buffer is Go-heap storage and therefore has no physical + // allocation provenance. Allocation-accounted vectors must use Shuffle, + // whose replacement data and bitmap scratch are admitted to their owner. + if v.allocationAccount != nil { + return v.Shuffle(sels, mp) + } // Fall back to allocating Shuffle if the vector doesn't own its data // or the selection changes the element count. if v.cantFreeData || len(sels) != v.length { @@ -4431,8 +4434,8 @@ func shrinkFixed[T types.FixedSizeT](v *Vector, sels []int64, negate bool) { for i, sel := range sels { vs[i] = vs[sel] } - nulls.Filter(&v.gsp, sels, false) - nulls.Filter(&v.nsp, sels, false) + nulls.FilterInPlaceOrdered(&v.gsp, sels, false) + nulls.FilterInPlaceOrdered(&v.nsp, sels, false) v.length = len(sels) } else if len(sels) > 0 { for oldIdx, newIdx, selIdx, sel := 0, 0, 0, sels[0]; oldIdx < v.length; oldIdx++ { @@ -4451,8 +4454,8 @@ func shrinkFixed[T types.FixedSizeT](v *Vector, sels []int64, negate bool) { sel = sels[selIdx] } } - nulls.Filter(&v.gsp, sels, true) - nulls.Filter(&v.nsp, sels, true) + nulls.FilterInPlaceOrdered(&v.gsp, sels, true) + nulls.FilterInPlaceOrdered(&v.nsp, sels, true) v.length -= len(sels) } } @@ -4467,8 +4470,8 @@ func shrinkFixedByMask[T types.FixedSizeT](v *Vector, sels *bitmap.Bitmap, negat vs[idx] = vs[itr.Next()+offset] idx++ } - nulls.FilterByMask(&v.gsp, sels, false) - nulls.FilterByMask(&v.nsp, sels, false) + nulls.FilterByMaskInPlace(&v.gsp, sels, false) + nulls.FilterByMaskInPlace(&v.nsp, sels, false) v.length = length } else if length > 0 { sel := itr.Next() + offset @@ -4487,8 +4490,8 @@ func shrinkFixedByMask[T types.FixedSizeT](v *Vector, sels *bitmap.Bitmap, negat sel = itr.Next() + offset } } - nulls.FilterByMask(&v.gsp, sels, true) - nulls.FilterByMask(&v.nsp, sels, true) + nulls.FilterByMaskInPlace(&v.gsp, sels, true) + nulls.FilterByMaskInPlace(&v.nsp, sels, true) v.length -= length } } @@ -4504,12 +4507,14 @@ func shuffleFixedNoTypeCheck[T types.FixedSizeT](v *Vector, sels []int64, mp *mp if err != nil { return err } - v.data = data - ws := toSliceOfLengthNoTypeCheck[T](v, ns) + ws := util.UnsafeSliceCastToLength[T](data, ns) shuffle.FixedLengthShuffle(vs, ws, sels) - nulls.Filter(&v.gsp, sels, false) - nulls.Filter(&v.nsp, sels, false) + if err := v.remapShuffleBitmaps(sels, mp); err != nil { + mp.Free(data) + return err + } + v.data = data // XXX We should never allow "half-owned" vectors later. And unowned vector should be strictly read-only. if v.cantFreeData { v.cantFreeData = false @@ -4520,6 +4525,88 @@ func shuffleFixedNoTypeCheck[T types.FixedSizeT](v *Vector, sels []int64, mp *mp return nil } +type bitmapRemapScratch struct { + destination *bitmap.Bitmap + value bitmap.Bitmap + storage []uint64 +} + +func (s *bitmapRemapScratch) release(mp *mpool.MPool) { + if s == nil || cap(s.storage) == 0 { + return + } + s.value.ReleaseExternalStorage() + mpool.FreeSlice(mp, s.storage) + s.storage = nil +} + +// remapShuffleBitmaps preserves Shuffle's arbitrary-selection semantics. An +// allocation-accounted vector builds both results in admitted temporary +// storage before publishing either, so rejection cannot leave null and +// grouping ownership half-mutated. +func (v *Vector) remapShuffleBitmaps(sels []int64, mp *mpool.MPool) error { + if v.allocationAccount == nil || !v.allocationAccount.accountBitmaps { + nulls.Filter(&v.gsp, sels, false) + nulls.Filter(&v.nsp, sels, false) + return nil + } + + targets := [...]struct { + destination *bitmap.Bitmap + site mpool.AllocationSite + }{ + {v.gsp.GetBitmap(), v.allocationAccount.groupingSite}, + {v.nsp.GetBitmap(), v.allocationAccount.nullsSite}, + } + if targets[0].destination.EmptyByFlag() && + targets[1].destination.EmptyByFlag() { + return nil + } + if err := v.ensureBitmapCapacity(len(sels), mp); err != nil { + return err + } + + var scratch [2]bitmapRemapScratch + for i, target := range targets { + if target.destination.EmptyByFlag() { + continue + } + words := (len(sels) + 63) / 64 + storage, err := mpool.MakeSliceAccounted[uint64]( + words, + mp, + v.allocationAccount.account, + v.allocationAccount.owner, + target.site, + ) + if err != nil { + for j := range i { + scratch[j].release(mp) + } + return err + } + scratch[i].destination = target.destination + scratch[i].storage = storage + scratch[i].value.InstallExternalStorage(storage) + scratch[i].value.InitWithSize(int64(len(sels))) + for output, source := range sels { + if target.destination.Contains(uint64(source)) { + scratch[i].value.Add(uint64(output)) + } + } + } + + for i := range scratch { + if scratch[i].destination != nil { + scratch[i].destination.InitWith(&scratch[i].value) + } + } + for i := range scratch { + scratch[i].release(mp) + } + return nil +} + // shuffleFixedNoTypeCheckWithBuf permutes elements using a reusable scratch // buffer instead of allocating a new data buffer. Only valid when // len(sels) == v.length and !v.cantFreeData (caller checks). diff --git a/pkg/container/vector/versions.go b/pkg/container/vector/versions.go index b9fdcd7abf317..f5026b37b45a8 100644 --- a/pkg/container/vector/versions.go +++ b/pkg/container/vector/versions.go @@ -16,9 +16,7 @@ package vector import ( "bytes" - "fmt" - "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/types" ) @@ -73,9 +71,8 @@ func (v *Vector) MarshalBinaryWithBufferV1(buf *bytes.Buffer) error { func (v *Vector) UnmarshalBinaryV1(data []byte) error { if v.allocationAccount != nil { - return fmt.Errorf( - "%w: cannot install aliases in an accounted vector", - mpool.ErrAllocationAccountInvalid, + return allocationAccountInvalid( + "cannot install aliases in an accounted vector", ) } // read class diff --git a/pkg/sql/colexec/aggexec/maxby.go b/pkg/sql/colexec/aggexec/maxby.go index d39e17bd9cfb9..96dcd0f9270c8 100644 --- a/pkg/sql/colexec/aggexec/maxby.go +++ b/pkg/sql/colexec/aggexec/maxby.go @@ -286,7 +286,15 @@ func compactMaxByStateVector(vec *vector.Vector, mp *mpool.MPool) error { if areaCapacity <= 2*liveBytes+maxByVarlenaCompactionSlack { return nil } - compact, err := vec.CloneToFlatCompact(mp) + var ( + compact *vector.Vector + err error + ) + if selection := vec.AllocationAccountSelection(); selection != nil { + compact, err = vec.CloneToFlatCompactWithAllocation(mp, selection) + } else { + compact, err = vec.CloneToFlatCompact(mp) + } if err != nil { return err } diff --git a/pkg/sql/colexec/aggexec/maxby_test.go b/pkg/sql/colexec/aggexec/maxby_test.go index 78a2533e22b08..5876cbf268288 100644 --- a/pkg/sql/colexec/aggexec/maxby_test.go +++ b/pkg/sql/colexec/aggexec/maxby_test.go @@ -80,6 +80,48 @@ func TestMaxByCompactsReplacedVarlenaState(t *testing.T) { require.Less(t, state.Allocated(), 2<<20, "winner state must be bounded by live groups, not by replaced input rows") } +func TestCompactMaxByStateVectorPreservesAllocationOwner(t *testing.T) { + mp := mpool.MustNewZero() + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.Open(8 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection( + account, + mpool.AllocationOwner(1), + mpool.AllocationSite(1), + mpool.AllocationSite(2), + ) + require.NoError(t, err) + vec := vector.NewOffHeapVecWithType(types.T_varchar.ToType()) + require.NoError(t, vec.SetAllocationAccount(selection)) + defer func() { + vec.Free(mp) + snapshot := account.Seal() + require.Zero(t, snapshot.Used) + require.Zero(t, registry.LiveAllocationMetadata()) + _, err = registry.Finalize(account) + require.NoError(t, err) + require.Zero(t, mp.CurrNB()) + }() + + value := []byte(strings.Repeat("x", 4096)) + require.NoError(t, vector.AppendBytes(vec, value, false, mp)) + for i := 0; i < 400; i++ { + value[0] = byte(i) + require.NoError(t, vec.SetRawBytesAt(0, value, mp)) + } + before := account.Snapshot() + require.Greater(t, before.Used, uint64(maxByVarlenaCompactionSlack)) + + require.NoError(t, compactMaxByStateVector(vec, mp)) + require.Same(t, selection, vec.AllocationAccountSelection()) + require.Equal(t, value, vec.GetBytesAt(0)) + after := account.Snapshot() + require.Less(t, after.Used, before.Used) + require.GreaterOrEqual(t, after.Peak, before.Used) +} + func TestMaxByNullContractAndDeterministicMerge(t *testing.T) { mp := mpool.MustNewZero() params := []types.Type{types.T_varchar.ToType(), types.T_int64.ToType(), types.T_varchar.ToType()} diff --git a/pkg/sql/colexec/hashbuild/pressure.go b/pkg/sql/colexec/hashbuild/pressure.go index 1fbe2d4ff3244..ad1bab84acbc7 100644 --- a/pkg/sql/colexec/hashbuild/pressure.go +++ b/pkg/sql/colexec/hashbuild/pressure.go @@ -18,6 +18,7 @@ import ( "errors" "fmt" + "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/vm/process" ) @@ -179,9 +180,11 @@ func (g *PressureRetryGuard) Advance(next PressureProgress) error { return process.ErrHashBuildBudgetInvalid } if g.attempts >= g.limit { - return fmt.Errorf( - "%w: memory-pressure retry limit exceeded", + return errors.Join( process.ErrHashBuildBudgetInvalid, + moerr.NewInternalErrorNoCtx( + "memory-pressure retry limit exceeded", + ), ) } progress := next.Used < g.previous.Used || @@ -189,9 +192,11 @@ func (g *PressureRetryGuard) Advance(next PressureProgress) error { (g.previous.InputUnits > 0 && next.InputUnits < g.previous.InputUnits) || (!g.previous.OptionalDisabled && next.OptionalDisabled) if !progress { - return fmt.Errorf( - "%w: memory-pressure retry made no progress", + return errors.Join( process.ErrHashBuildBudgetInvalid, + moerr.NewInternalErrorNoCtx( + "memory-pressure retry made no progress", + ), ) } g.previous = next diff --git a/pkg/sql/plan/function/func_unary.go b/pkg/sql/plan/function/func_unary.go index 1bafbdda271b6..e15486326426c 100644 --- a/pkg/sql/plan/function/func_unary.go +++ b/pkg/sql/plan/function/func_unary.go @@ -6245,7 +6245,9 @@ func Uncompress(parameters []*vector.Vector, result vector.FunctionResultWrapper n, err := reader.Read(extra[:]) if err != io.EOF || n != 0 { if err == nil { - err = errors.New("decompressed length exceeds header") + err = moerr.NewInvalidInputNoCtx( + "decompressed length exceeds header", + ) } decodeErr = err return 0, err diff --git a/pkg/sql/util/copy_batch_test.go b/pkg/sql/util/copy_batch_test.go index 643a56073c103..6f36bb0a7ae9c 100644 --- a/pkg/sql/util/copy_batch_test.go +++ b/pkg/sql/util/copy_batch_test.go @@ -91,6 +91,49 @@ func TestCopyBatchCompactsAndFlattens(t *testing.T) { require.Equal(t, int64(0), mp.CurrNB()) } +func TestCopyBatchCrossesAllocationOwnershipBoundary(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelectionWithBitmaps( + account, + 1, + 1, + 2, + 3, + 4, + ) + require.NoError(t, err) + + src := batch.NewOffHeapWithSize(1) + vec := vector.NewOffHeapVecWithType(types.T_int64.ToType()) + require.NoError(t, vec.SetAllocationAccount(selection)) + require.NoError(t, vector.AppendFixed(vec, int64(42), false, mp)) + src.SetVector(0, vec) + src.SetRowCount(1) + sourceUsed := account.Snapshot().Used + require.Positive(t, sourceUsed) + + got, err := CopyBatch(src, proc) + require.NoError(t, err) + require.Nil(t, got.Vecs[0].AllocationAccountSelection()) + require.Equal(t, int64(42), vector.GetFixedAtNoTypeCheck[int64](got.Vecs[0], 0)) + require.Equal(t, sourceUsed, account.Snapshot().Used) + + src.Clean(mp) + require.Zero(t, account.Snapshot().Used) + require.Equal(t, int64(42), vector.GetFixedAtNoTypeCheck[int64](got.Vecs[0], 0)) + got.Clean(mp) + + account.Seal() + _, err = registry.Finalize(account) + require.NoError(t, err) + require.Equal(t, int64(0), mp.CurrNB()) +} + func BenchmarkCopyBatchCompact(b *testing.B) { proc := testutil.NewProcess(b) mp := proc.Mp() diff --git a/pkg/sql/util/util.go b/pkg/sql/util/util.go index f329edce53af4..00930be1ea470 100644 --- a/pkg/sql/util/util.go +++ b/pkg/sql/util/util.go @@ -21,6 +21,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/catalog" "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/defines" "github.com/matrixorigin/matrixone/pkg/sql/parsers/tree" "github.com/matrixorigin/matrixone/pkg/vm/process" @@ -44,7 +45,18 @@ func CopyBatch(bat *batch.Batch, proc *process.Process) (*batch.Batch, error) { rbat := batch.NewWithSize(len(bat.Vecs)) rbat.Attrs = append(rbat.Attrs, bat.Attrs...) for i, srcVec := range bat.Vecs { - vec, err := srcVec.CloneToFlatCompact(proc.Mp()) + var ( + vec *vector.Vector + err error + ) + if srcVec.AllocationAccountSelection() != nil { + // CopyBatch is an ownership boundary: the source keeps its physical + // account until Free, while the independent destination belongs to + // the generic downstream batch owner. + vec, err = srcVec.CloneToFlatCompactWithAllocation(proc.Mp(), nil) + } else { + vec, err = srcVec.CloneToFlatCompact(proc.Mp()) + } if err != nil { rbat.Clean(proc.Mp()) return nil, err From 9ed2cca77507b58e988ff3d9215edcbfea251c64 Mon Sep 17 00:00:00 2001 From: aptend Date: Sat, 1 Aug 2026 00:50:25 +0800 Subject: [PATCH 32/61] test: restore cross-package spill fixture writer --- pkg/sql/colexec/spillutil/join_spill.go | 26 ++++++++++++++++++++ pkg/sql/colexec/spillutil/join_spill_test.go | 25 ------------------- 2 files changed, 26 insertions(+), 25 deletions(-) diff --git a/pkg/sql/colexec/spillutil/join_spill.go b/pkg/sql/colexec/spillutil/join_spill.go index 0bde01a41f12b..55964a71a170c 100644 --- a/pkg/sql/colexec/spillutil/join_spill.go +++ b/pkg/sql/colexec/spillutil/join_spill.go @@ -879,6 +879,32 @@ func MakeBucketWriters(prefix string) []BucketWriter { return writers } +// FlushBucketBatch writes one framed batch to w. It remains the low-level +// fixture/compatibility boundary for callers that already own a spill writer; +// production scatter uses SpillEngine so its retained buffers and pressure +// retries stay under the statement allocation owner. +func FlushBucketBatch( + proc *process.Process, + bat *batch.Batch, + w *BucketWriter, + bucketBuf *bytes.Buffer, + analyzer process.Analyzer, +) error { + if bat == nil || bat.RowCount() == 0 { + return nil + } + if err := marshalSpillRecord(bat, bucketBuf); err != nil { + return err + } + return writeBucketPayload( + proc, + bucketBuf.Bytes(), + int64(bat.RowCount()), + w, + analyzer, + ) +} + type spillRecordBuffer interface { io.Writer Bytes() []byte diff --git a/pkg/sql/colexec/spillutil/join_spill_test.go b/pkg/sql/colexec/spillutil/join_spill_test.go index 7b6430b72cfda..0281c17a441d6 100644 --- a/pkg/sql/colexec/spillutil/join_spill_test.go +++ b/pkg/sql/colexec/spillutil/join_spill_test.go @@ -67,31 +67,6 @@ func (r *boundaryCancelReader) Read(p []byte) (int, error) { return n, err } -// FlushBucketBatch writes one legacy-format spill record. Production -// scatter goes through SpillEngine.scatterBatchBounded; tests that construct a -// reader fixture directly need only this codec/writer composition. -func FlushBucketBatch( - proc *process.Process, - bat *batch.Batch, - w *BucketWriter, - bucketBuf *bytes.Buffer, - analyzer process.Analyzer, -) error { - if bat == nil || bat.RowCount() == 0 { - return nil - } - if err := marshalSpillRecord(bat, bucketBuf); err != nil { - return err - } - return writeBucketPayload( - proc, - bucketBuf.Bytes(), - int64(bat.RowCount()), - w, - analyzer, - ) -} - // scatterImpl retains the former buffered implementation solely as a // compatibility oracle. It is intentionally test-only: production owns one // selected batch at a time in SpillEngine.scatterBatchBounded. From 3a065b7bb6a02d02aada7367d25f0f1ef895a87f Mon Sep 17 00:00:00 2001 From: aptend Date: Sat, 1 Aug 2026 01:08:29 +0800 Subject: [PATCH 33/61] executor: finalize allocation owners before release --- pkg/sql/colexec/spillutil/join_spill.go | 17 ------- .../allocation_account_lifecycle_test.go | 46 ++++++++++++++++--- pkg/sql/compile/compile.go | 9 ++-- 3 files changed, 46 insertions(+), 26 deletions(-) diff --git a/pkg/sql/colexec/spillutil/join_spill.go b/pkg/sql/colexec/spillutil/join_spill.go index 55964a71a170c..3d5532729360f 100644 --- a/pkg/sql/colexec/spillutil/join_spill.go +++ b/pkg/sql/colexec/spillutil/join_spill.go @@ -1285,23 +1285,6 @@ func spillStatInt64(v uint64) int64 { return int64(v) } -func externalScatterSourceBytes( - bat *batch.Batch, - sourceAlreadyCharged bool, -) (uint64, error) { - if bat == nil || bat.RowCount() < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - if sourceAlreadyCharged { - return 0, nil - } - allocated := uint64(bat.Allocated()) - if size := uint64(bat.Size()); size > allocated { - allocated = size - } - return allocated, nil -} - func (e *SpillEngine) scatterRetainedBytes() (uint64, bool) { actual := uint64(0) add := func(v uint64) bool { diff --git a/pkg/sql/compile/allocation_account_lifecycle_test.go b/pkg/sql/compile/allocation_account_lifecycle_test.go index 7a1f5b21e6edd..a478bbafb6c87 100644 --- a/pkg/sql/compile/allocation_account_lifecycle_test.go +++ b/pkg/sql/compile/allocation_account_lifecycle_test.go @@ -43,11 +43,18 @@ type allocationLifecycleErrorOperator struct { type allocationLifecycleOwnerOperator struct { *colexec.MockOperator - account *mpool.AllocationAccount - failSet bool - failClear bool - blocked bool - clears int + account *mpool.AllocationAccount + failSet bool + failClear bool + blocked bool + clears int + released bool + releaseSawLiveAccount bool +} + +func (op *allocationLifecycleOwnerOperator) Release() { + op.released = true + op.releaseSawLiveAccount = op.account != nil } func (op *allocationLifecycleOwnerOperator) AllocationAccountEnabled() bool { @@ -153,7 +160,6 @@ func TestStatementAllocationAttemptZeroTerminalExportsOnce(t *testing.T) { ) { exported = append(exported, snapshot) }) - attempt, err := c.beginAllocationAccountAttempt() require.NoError(t, err) require.NotNil(t, attempt) @@ -394,6 +400,34 @@ func TestStatementAllocationAttemptOwnerTeardownFailureExportsFailure(t *testing require.False(t, ok) } +func TestCompileClearFinalizesAllocationOwnerBeforeRelease(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + var exported []mpool.AllocationAccountTerminalSnapshot + c := newTestAllocationLifecycleCompile(t, registry, func( + snapshot mpool.AllocationAccountTerminalSnapshot, + ) { + exported = append(exported, snapshot) + }) + c.affectRows = &atomic.Uint64{} + owner := &allocationLifecycleOwnerOperator{ + MockOperator: colexec.NewMockOperator(), + } + c.scopes = []*Scope{{RootOp: owner}} + + _, err = c.beginAllocationAccountAttempt() + require.NoError(t, err) + c.clear() + + require.True(t, owner.released) + require.False(t, owner.releaseSawLiveAccount) + require.Nil(t, owner.account) + require.Equal(t, 1, owner.clears) + require.Len(t, exported, 1) + require.Equal(t, mpool.AllocationAccountTerminalValid, exported[0].State) + require.Zero(t, exported[0].Used) +} + func TestCompileAutomaticallyActivatesCompleteHashTableOwner(t *testing.T) { proc := testutil.NewProcess(t) c := &Compile{ diff --git a/pkg/sql/compile/compile.go b/pkg/sql/compile/compile.go index ec2a07ea04d44..a1397dd9af400 100644 --- a/pkg/sql/compile/compile.go +++ b/pkg/sql/compile/compile.go @@ -277,6 +277,12 @@ func (c *Compile) clear() { if c.anal != nil { c.anal.release() } + // The attempt owns references to allocation-aware operators. Finalize it + // before Scope.release returns those operators to reuse pools; otherwise a + // defensive cleanup path could clear an already-reset or reused owner. + if err := c.finishAllocationAccountAttempt(); err != nil { + logutil.Errorf("allocation account terminal cleanup failed: %v", err) + } for i := range c.scopes { c.scopes[i].release() } @@ -284,9 +290,6 @@ func (c *Compile) clear() { c.fuzzys[i].release() } - if err := c.finishAllocationAccountAttempt(); err != nil { - logutil.Errorf("allocation account terminal cleanup failed: %v", err) - } c.MessageBoard = c.MessageBoard.Reset() c.fuzzys = c.fuzzys[:0] c.scopes = c.scopes[:0] From 81426a00e6da3c755d15d2d7572b4688ca2f4f24 Mon Sep 17 00:00:00 2001 From: aptend Date: Sat, 1 Aug 2026 01:13:22 +0800 Subject: [PATCH 34/61] docs: refresh allocation activation closure evidence --- ...ocation_accounted_memory_admission_impl.md | 25 ++++++++++++++++--- .../evidence/26459_activation_validation.md | 23 ++++++++++++++--- 2 files changed, 40 insertions(+), 8 deletions(-) diff --git a/docs/design/allocation_accounted_memory_admission_impl.md b/docs/design/allocation_accounted_memory_admission_impl.md index a6fa9bb45a312..87491c884766b 100644 --- a/docs/design/allocation_accounted_memory_admission_impl.md +++ b/docs/design/allocation_accounted_memory_admission_impl.md @@ -6,12 +6,14 @@ - Architecture: [Allocation-Accounted Memory Admission RFC](../rfcs/00000000_allocation_accounted_memory_admission.md) - Baseline at plan creation: `main` at `38ce3a774` -- Rebased implementation baseline: `main` at `5b9eeb54ec` +- Rebased implementation baseline: `main` at `a9c13b02d8` - Allocation-site closure baseline (PR 3): `f61b64d56c` - Lifecycle/activation heads (PRs 4--10): `8633757dc1`, `5ae8eca00a`, `8e4b689f45`, `c9ad0ea810`, `e072568998`, `eec23dcdc5`, and `383fc6dce3` -- Owner-atomic activation and final cleanup: `656e254fe6` +- Owner-atomic activation and cleanup: `656e254fe6` +- Post-activation ownership closure: `73d20085c0`, `3f10c23f21`, and + `91839ea91f` - Merged prerequisite: #26455 at `93e8b22d2` - Independent design review: completed against RFC commit `a7d54cb5f` - Activation status: enabled for the closed HashBuild expression owner set. @@ -104,6 +106,7 @@ working ledger is allocation-site based: | `Vector.data` | MPool; capacity from `Grow`, on/off-heap follows `v.offHeap` | owning `Vector.Free` | A | activated HashBuild destinations use immutable off-heap provenance | | `Vector.area` | MPool; independent varlen payload capacity | owning `Vector.Free` | A | activated HashBuild destinations use immutable off-heap provenance | | `Vector.nsp/gsp` bitmap data | allocation-accounted off-heap `[]uint64`; independent geometric capacity, paired replacement admission | owning `Vector.Free`; `Reset` retains and clears only published words | A | selected HashBuild vectors and expression results are active | +| ordered/masked Vector filtering | old bitmap backing is remapped in place; arbitrary row reordering uses allocation-accounted scratch admitted before publication | owning `Vector.Free`; scratch ends at `Shuffle` return | A | rejection leaves the old Vector and bitmap owners unchanged | | `FunctionResult.vec` data/area | off-heap Vector; rows and appended payload | executor `Free` | A | active for the closed HashBuild expression set | | `FunctionResult.convenientParam` | Go slice; expression arity, not rows | executor `Free`/reuse | H | bounded by plan expression arity, not input rows or payload | | decimal parameter conversion | retained allocation-accounted off-heap buffer; `rows*sizeof(Decimal128)` for decimal64/float32/float64 promotion | `FunctionResult.Free`; Reset/evaluation reuse capacity | A | active when present in a closed expression owner | @@ -172,6 +175,16 @@ growth fails instead of escaping to the Go heap, Reset retains capacity while clearing only represented words, copy/reader decode fills admitted backing directly, and Free is the one terminal release owner. +Post-activation filtering closes the corresponding mutation boundary. Ordered +filters remap null/group bitmaps in place, while arbitrary permutations admit +the replacement scratch against the same immutable allocation owner before +publication. `ShuffleWithBuf` cannot fall back to an unaccounted Go slice for +an accounted Vector. Aggregate compactors such as `maxby` preserve the same +selection while rewriting retained state. Generic `CopyBatch` is instead an +explicit ownership exit: it creates a compact clone owned by the destination +MPool and does not leak a statement-generation account into transaction or +bootstrap lifetime. + ## 3. Decisions required before production integration ### A. Allocation metadata representation @@ -1067,10 +1080,14 @@ Required proof: Current validation status: - the complete affected package matrix passes after rebasing on - `5b9eeb54ec`; + `a9c13b02d8`; - the complete affected package matrix, the same matrix under `-race`, a 20-iteration focused lifecycle/pressure race stress, affected-package vet, - and `make build` pass at `656e254fe6`; + `make static-check`, and `make build` pass through `91839ea91f`; +- accounted ordered/masked filtering, arbitrary shuffle rollback, `maxby` + retained-state compaction, generic `CopyBatch` ownership exit, CN bootstrap, + cross-package spill fixtures, and owner finalization-before-release have + dedicated regressions; - local regression tests cover the five incident mechanisms without an estimator-only rejection on the activated owner; - TPCH 100G/1T candidate-versus-main runs remain the remote workload gate and diff --git a/docs/design/evidence/26459_activation_validation.md b/docs/design/evidence/26459_activation_validation.md index a1feab10b5c65..55b55cea57f06 100644 --- a/docs/design/evidence/26459_activation_validation.md +++ b/docs/design/evidence/26459_activation_validation.md @@ -2,7 +2,7 @@ ## Candidate -- Rebased main: `5b9eeb54ec` +- Rebased main: `a9c13b02d8` - PR 4 lifecycle: `8633757dc1` - PR 5 hash-table activation: `5ae8eca00a` - PR 6 retained batch/JoinMap activation: `8e4b689f45` @@ -11,6 +11,9 @@ - PR 9 unified pressure recovery: `eec23dcdc5` - PR 10 benchmark harness: `383fc6dce3` - Owner-atomic activation and cleanup: `656e254fe6` +- Vector/bitmap ownership boundary closure: `73d20085c0` +- Cross-package spill fixture compatibility: `3f10c23f21` +- Owner finalization-before-release: `91839ea91f` The activated expression owner is deliberately closed: COL, literal, param, variable, vector/fold, CONCAT, CASE, varchar EQUAL, and the audited string CAST @@ -28,9 +31,10 @@ the locally built `thirdparties` artifacts. ```text .agents/skills/mo-dev/scripts/mo-cgo-test -count=1 \ - ./pkg/common/mpool ./pkg/common/hashmap \ + ./pkg/common/mpool ./pkg/common/hashmap ./pkg/common/bitmap \ ./pkg/container/hashtable ./pkg/container/vector ./pkg/container/batch \ - ./pkg/sql/colexec ./pkg/sql/colexec/hashbuild \ + ./pkg/container/nulls ./pkg/container/bytejson ./pkg/sql/util \ + ./pkg/sql/colexec ./pkg/sql/colexec/aggexec ./pkg/sql/colexec/hashbuild \ ./pkg/sql/colexec/hashjoin ./pkg/sql/colexec/dedupjoin \ ./pkg/sql/colexec/rightdedupjoin ./pkg/sql/colexec/spillutil \ ./pkg/sql/compile ./pkg/util/resource ./pkg/vm/message ./pkg/vm/process @@ -38,12 +42,23 @@ the locally built `thirdparties` artifacts. result: PASS ``` -The same package matrix passes with `-race -p=2 -count=1`. A focused +The same package matrix passes with `-race -p=2 -count=1`. `make static-check` +and `make build` also pass on the rebased candidate. A focused `-race -p=2 -count=20` stress also passes for concurrent account alloc/free, seal/open linearization, statement terminal one-shot behavior, owner-atomic activation, accounted spill/reduction, and all three join-consumer activation gates. Affected-package `go vet` and `make build` pass on the same candidate. +The previous PR CI failures were reduced to two branch-local ownership +boundaries and fixed before rerunning the matrix: generic `CopyBatch` now +performs an explicit compact clone when leaving an accounted statement owner, +which restores CN bootstrap, and the spill fixture writer remains available to +cross-package hashjoin/dedup tests. `Compile.clear` finalizes and detaches the +statement allocation owner before releasing operators back to reuse pools. +Dedicated tests verify accounted bitmap remap/shuffle rollback, `maxby` +selection preservation, bootstrap copy, fixture compilation, and that an +operator's `Release` cannot observe a live generation account. + Incident-mechanism regression mapping: | Incident | Local regression proof | From c39f49a19f44258ebd426d62fce0f7a8c0ee121a Mon Sep 17 00:00:00 2001 From: aptend Date: Sat, 1 Aug 2026 01:40:00 +0800 Subject: [PATCH 35/61] executor: close transaction PK ownership boundary --- pkg/vm/engine/disttae/txn_table.go | 7 ++++ pkg/vm/engine/disttae/txn_test.go | 52 ++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+) diff --git a/pkg/vm/engine/disttae/txn_table.go b/pkg/vm/engine/disttae/txn_table.go index 7054bbcdbdec5..82591c6a71756 100644 --- a/pkg/vm/engine/disttae/txn_table.go +++ b/pkg/vm/engine/disttae/txn_table.go @@ -3598,6 +3598,13 @@ func (tbl *txnTable) GetExtraInfo() *api.SchemaExtra { // If v has no NULLs it returns Dup(v). The caller must Free the result. func dupVectorWithoutNulls(v *vector.Vector, mp *mpool.MPool) (*vector.Vector, error) { if !v.HasNull() { + if v.AllocationAccountSelection() != nil { + // PK validation runs during transaction commit, after the producing SQL + // attempt may already have sealed its allocation account. The sorted + // copy is a short-lived transaction-engine owner, not a continuation of + // the statement owner, so make that ownership exit explicit. + return v.DupOffHeapWithAllocation(mp, nil) + } return v.Dup(mp) } filtered := vector.NewVec(*v.GetType()) diff --git a/pkg/vm/engine/disttae/txn_test.go b/pkg/vm/engine/disttae/txn_test.go index e3112368243f0..5cf3b861c72bd 100644 --- a/pkg/vm/engine/disttae/txn_test.go +++ b/pkg/vm/engine/disttae/txn_test.go @@ -26,6 +26,7 @@ import ( "github.com/golang/mock/gomock" "github.com/matrixorigin/matrixone/pkg/catalog" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" @@ -1240,6 +1241,57 @@ func TestDupVectorWithoutNulls(t *testing.T) { }) } +func TestDupVectorWithoutNullsLeavesSealedStatementOwner(t *testing.T) { + proc := testutil.NewProc(t) + mp := proc.Mp() + + for _, tc := range []struct { + name string + withNulls bool + }{ + {name: "no nulls"}, + {name: "with nulls", withNulls: true}, + } { + t.Run(tc.name, func(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection( + account, + mpool.AllocationOwner(1), + mpool.AllocationSite(1), + mpool.AllocationSite(2), + ) + require.NoError(t, err) + source, err := vector.NewOffHeapVecWithTypeAndAllocation( + types.T_int64.ToType(), + selection, + ) + require.NoError(t, err) + require.NoError(t, vector.AppendFixed(source, int64(1), false, mp)) + if tc.withNulls { + require.NoError(t, vector.AppendFixed(source, int64(0), true, mp)) + } + require.NoError(t, vector.AppendFixed(source, int64(2), false, mp)) + + used := account.Seal().Used + require.NotZero(t, used) + out, err := dupVectorWithoutNulls(source, mp) + require.NoError(t, err) + require.Nil(t, out.AllocationAccountSelection()) + require.Equal(t, used, account.Snapshot().Used) + require.Equal(t, []int64{1, 2}, vector.MustFixedColWithTypeCheck[int64](out)) + + out.Free(mp) + source.Free(mp) + require.Zero(t, account.Snapshot().Used) + _, err = registry.Finalize(account) + require.NoError(t, err) + }) + } +} + func newInt64BatchForTest( t *testing.T, proc *process.Process, From f8af2d6d8903309130a75ff11108bffd35c435e2 Mon Sep 17 00:00:00 2001 From: aptend Date: Sat, 1 Aug 2026 01:53:45 +0800 Subject: [PATCH 36/61] executor: preserve mixed batch allocation provenance --- .../batch/allocation_account_test.go | 184 ++++++++++++++++++ pkg/container/batch/batch.go | 76 ++++++-- pkg/container/batch/batch_set.go | 121 +++++++++++- .../colexec/multi_update/s3writer_delegate.go | 9 +- pkg/sql/colexec/shuffle/shufflepool.go | 2 +- pkg/sql/colexec/shuffle/shufflepool_test.go | 51 +++++ pkg/vm/engine/disttae/txn_table.go | 7 +- 7 files changed, 421 insertions(+), 29 deletions(-) diff --git a/pkg/container/batch/allocation_account_test.go b/pkg/container/batch/allocation_account_test.go index ddc74d21be964..ad3ad0a2d4d65 100644 --- a/pkg/container/batch/allocation_account_test.go +++ b/pkg/container/batch/allocation_account_test.go @@ -162,6 +162,190 @@ func TestBatchAllocationAccountCloneDupAndWindow(t *testing.T) { finalizeTestBatchAllocationAccount(t, state) } +func newMixedBatchAllocationSource( + t *testing.T, + mp *mpool.MPool, + selection *vector.AllocationAccountSelection, + rows int, +) *Batch { + t.Helper() + bat := NewOffHeapWithSize(2) + bat.Attrs = []string{"accounted", "legacy"} + bat.Vecs[0] = vector.NewOffHeapVecWithType(types.T_int64.ToType()) + require.NoError(t, bat.Vecs[0].SetAllocationAccount(selection)) + bat.Vecs[1] = vector.NewOffHeapVecWithType(types.T_varchar.ToType()) + for i := 0; i < rows; i++ { + require.NoError(t, vector.AppendFixed(bat.Vecs[0], int64(i), false, mp)) + require.NoError(t, vector.AppendBytes(bat.Vecs[1], []byte("legacy"), false, mp)) + } + bat.SetRowCount(rows) + return bat +} + +func TestMixedBatchAllocationClonePreservesVectorProvenance(t *testing.T) { + state := newTestBatchAllocationAccount(t, 128) + mp := mpool.MustNewZero() + source := newMixedBatchAllocationSource(t, mp, state.selection, 8) + sourceUsed := state.account.Snapshot().Used + require.NotZero(t, sourceUsed) + require.Nil(t, source.AllocationAccountSelection()) + + _, err := source.Clone(mp, false) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + for _, clone := range []func() (*Batch, error){ + func() (*Batch, error) { return source.Clone(mp, true) }, + func() (*Batch, error) { return source.Dup(mp) }, + } { + got, err := clone() + require.NoError(t, err) + require.Nil(t, got.AllocationAccountSelection()) + require.Same(t, state.selection, got.Vecs[0].AllocationAccountSelection()) + require.Nil(t, got.Vecs[1].AllocationAccountSelection()) + got.Clean(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + } + + accounted, err := source.CloneSelectedColumns([]int{0}, []string{"accounted"}, mp) + require.NoError(t, err) + require.Same(t, state.selection, accounted.Vecs[0].AllocationAccountSelection()) + accounted.Clean(mp) + legacy, err := source.CloneSelectedColumns([]int{1}, []string{"legacy"}, mp) + require.NoError(t, err) + require.Nil(t, legacy.Vecs[0].AllocationAccountSelection()) + legacy.Clean(mp) + + source.FreeColumns(mp) + require.Zero(t, state.account.Snapshot().Used) + require.Same(t, state.selection, source.Vecs[0].AllocationAccountSelection()) + require.Nil(t, source.Vecs[1].AllocationAccountSelection()) + require.NoError(t, vector.AppendFixed(source.Vecs[0], int64(9), false, mp)) + require.NotZero(t, state.account.Snapshot().Used) + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestMixedBatchAllocationBatchSetPreservesVectorProvenance(t *testing.T) { + state := newTestBatchAllocationAccount(t, 128) + mp := mpool.MustNewZero() + set := NewBatchSet(4) + first := newMixedBatchAllocationSource(t, mp, state.selection, 2) + second := newMixedBatchAllocationSource(t, mp, state.selection, 6) + + consumed, err := set.Extend(mp, first, nil) + require.NoError(t, err) + require.False(t, consumed) + consumed, err = set.Extend(mp, second, nil) + require.NoError(t, err) + require.False(t, consumed) + require.Equal(t, 2, set.Length()) + require.Equal(t, 8, set.RowCount()) + for i := 0; i < set.Length(); i++ { + require.Same(t, state.selection, set.Get(i).Vecs[0].AllocationAccountSelection()) + require.Nil(t, set.Get(i).Vecs[1].AllocationAccountSelection()) + } + + first.Clean(mp) + second.Clean(mp) + set.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchSetStartsNewTailWhenVectorProvenanceChanges(t *testing.T) { + state := newTestBatchAllocationAccount(t, 128) + mp := mpool.MustNewZero() + set := NewBatchSet(4) + legacy := newBatchAllocationTestSource(t, mp, nil) + legacy.Shrink([]int64{0, 1}, false) + mixed := newMixedBatchAllocationSource(t, mp, state.selection, 3) + + _, err := set.Extend(mp, legacy, nil) + require.NoError(t, err) + ready := set.ReadyCount() + require.Equal(t, 1, set.ReadyDeltaFor(mixed, mixed.RowCount())) + _, err = set.Extend(mp, mixed, nil) + require.NoError(t, err) + require.Equal(t, 1, set.ReadyCount()-ready) + require.Equal(t, 2, set.Length()) + require.Equal(t, 2, set.Get(0).RowCount()) + require.Equal(t, 3, set.Get(1).RowCount()) + require.Nil(t, set.Get(0).Vecs[0].AllocationAccountSelection()) + require.Same(t, state.selection, set.Get(1).Vecs[0].AllocationAccountSelection()) + + legacyUnion := newBatchAllocationTestSource(t, mp, nil) + ready = set.ReadyCount() + require.Equal(t, 1, set.ReadyDeltaFor(legacyUnion, 1)) + _, err = set.Union(mp, legacyUnion, []int32{0}, nil) + require.NoError(t, err) + require.Equal(t, 1, set.ReadyCount()-ready) + require.Equal(t, 3, set.Length()) + require.Nil(t, set.Get(2).Vecs[0].AllocationAccountSelection()) + + pushed := newMixedBatchAllocationSource(t, mp, state.selection, 1) + require.NoError(t, set.Push(mp, pushed)) + require.Equal(t, 4, set.Length()) + require.Same(t, state.selection, set.Get(3).Vecs[0].AllocationAccountSelection()) + require.Equal(t, 1, set.Get(3).RowCount()) + + legacy.Clean(mp) + mixed.Clean(mp) + legacyUnion.Clean(mp) + set.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchSetPreservesUniformBatchAllocationContext(t *testing.T) { + state := newTestBatchAllocationAccount(t, 256) + mp := mpool.MustNewZero() + set := NewBatchSet(16) + first := newBatchAllocationTestSource(t, mp, state.selection) + first.Shrink([]int64{0, 1, 2, 3, 4, 5, 6, 7}, false) + second := newBatchAllocationTestSource(t, mp, state.selection) + second.Shrink([]int64{ + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, + }, false) + + _, err := set.Extend(mp, first, nil) + require.NoError(t, err) + _, err = set.Extend(mp, second, nil) + require.NoError(t, err) + require.Equal(t, 2, set.Length()) + for i := 0; i < set.Length(); i++ { + require.Same(t, state.selection, set.Get(i).AllocationAccountSelection()) + for _, vec := range set.Get(i).Vecs { + require.Same(t, state.selection, vec.AllocationAccountSelection()) + } + } + + reuse := NewWithSchema( + true, + first.Attrs, + []types.Type{types.T_int64.ToType(), types.T_varchar.ToType()}, + ) + require.NoError(t, reuse.SetAllocationAccount(state.selection)) + third := newBatchAllocationTestSource(t, mp, state.selection) + third.Shrink([]int64{0, 1, 2, 3, 4, 5, 6, 7}, false) + consumed, err := set.Extend(mp, third, reuse) + require.NoError(t, err) + require.True(t, consumed) + require.Equal(t, 3, set.Length()) + require.Same(t, state.selection, set.Get(2).AllocationAccountSelection()) + + set.Get(2).FreeColumns(mp) + require.Same(t, state.selection, set.Get(2).AllocationAccountSelection()) + for _, vec := range set.Get(2).Vecs { + require.Same(t, state.selection, vec.AllocationAccountSelection()) + } + + first.Clean(mp) + second.Clean(mp) + third.Clean(mp) + set.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + func TestBatchAllocationAccountDestinationCloneUnionAndReuse(t *testing.T) { state := newTestBatchAllocationAccount(t, 64) mp := mpool.MustNewZero() diff --git a/pkg/container/batch/batch.go b/pkg/container/batch/batch.go index a1c8cf0724d53..20fee1392b7f2 100644 --- a/pkg/container/batch/batch.go +++ b/pkg/container/batch/batch.go @@ -742,10 +742,9 @@ func (bat *Batch) CloneSelectedColumns( cloned.Vecs[idx] = vector.NewVec(typ) } } - if bat.allocationAccount != nil { - if err = cloned.SetAllocationAccount(bat.allocationAccount); err != nil { - return nil, err - } + if err = configureCloneAllocation(bat, cloned, selectCols); err != nil { + cloned.Clean(mp) + return nil, err } if err = bat.CloneSelectedColumnsTo(selectCols, cloned, mp); err != nil { cloned.Clean(mp) @@ -842,11 +841,13 @@ func (bat *Batch) CleanOnlyData() { func (bat *Batch) FreeColumns(m *mpool.MPool) { for _, vec := range bat.Vecs { if vec != nil { + selection := vec.AllocationAccountSelection() vec.Free(m) if bat.allocationAccount != nil { - if err := vec.SetAllocationAccount(bat.allocationAccount); err != nil { - panic(err) - } + selection = bat.allocationAccount + } + if err := vec.SetAllocationAccount(selection); err != nil { + panic(err) } } } @@ -871,19 +872,66 @@ func (bat *Batch) GetSchema() (attrs []string, attrTypes []types.Type) { return } -func (bat *Batch) Clone(mp *mpool.MPool, offHeap bool) (*Batch, error) { - if bat.allocationAccount != nil && !offHeap { - return nil, mpool.ErrAllocationAccountInvalid +func vectorAllocationSelectionsMatch(left, right *Batch) bool { + if left == nil || right == nil || len(left.Vecs) != len(right.Vecs) { + return false + } + if left.allocationAccount != right.allocationAccount { + return false + } + for i := range left.Vecs { + if left.Vecs[i] == nil || right.Vecs[i] == nil { + if left.Vecs[i] != right.Vecs[i] { + return false + } + continue + } + if left.Vecs[i].AllocationAccountSelection() != + right.Vecs[i].AllocationAccountSelection() { + return false + } + } + return true +} + +func configureCloneAllocation( + source, destination *Batch, + selectedColumns []int, +) error { + if source == nil || destination == nil { + return mpool.ErrAllocationAccountInvalid + } + if source.allocationAccount != nil { + return destination.SetAllocationAccount(source.allocationAccount) + } + for destinationIdx := range destination.Vecs { + sourceIdx := destinationIdx + if len(selectedColumns) > 0 { + sourceIdx = selectedColumns[destinationIdx] + } + selection := source.Vecs[sourceIdx].AllocationAccountSelection() + if selection == nil { + continue + } + if !destination.offHeap { + return mpool.ErrAllocationAccountInvalid + } + if err := destination.Vecs[destinationIdx].SetAllocationAccount(selection); err != nil { + return err + } } + return nil +} + +func (bat *Batch) Clone(mp *mpool.MPool, offHeap bool) (*Batch, error) { var ( cloned *Batch attrs, attrTypes = bat.GetSchema() ) cloned = NewWithSchema(offHeap, attrs, attrTypes) - if offHeap && bat.allocationAccount != nil { - if err := cloned.SetAllocationAccount(bat.allocationAccount); err != nil { - return nil, err - } + if err := configureCloneAllocation(bat, cloned, nil); err != nil { + cloned.Clean(mp) + return nil, err } cloned.Recursive = bat.Recursive err := bat.CloneTo(cloned, mp) diff --git a/pkg/container/batch/batch_set.go b/pkg/container/batch/batch_set.go index 7e6ad16ec55d6..6b39d95dbdb01 100644 --- a/pkg/container/batch/batch_set.go +++ b/pkg/container/batch/batch_set.go @@ -42,8 +42,9 @@ func (bs *BatchSet) Length() int { return len(bs.batches) } -// ReadyCount returns the number of full batches that can be consumed. The last -// batch remains a writable tail only while it is partial. +// ReadyCount returns the number of batches that can be consumed. All batches +// before the last are sealed, including a partial batch whose successor has a +// different allocation provenance. The last is ready only when it is full. func (bs *BatchSet) ReadyCount() int { if len(bs.batches) == 0 { return 0 @@ -69,6 +70,23 @@ func (bs *BatchSet) ReadyDelta(rowCount int) int { return (lastRows+rowCount)/bs.batchMaxRow - lastRows/bs.batchMaxRow } +// ReadyDeltaFor returns how many batches become consumable when rows copied +// from source are appended. A provenance change seals the existing partial +// tail because already allocated vectors cannot be relabeled. +func (bs *BatchSet) ReadyDeltaFor(source *Batch, rowCount int) int { + if rowCount <= 0 { + return 0 + } + if len(bs.batches) > 0 { + last := bs.batches[len(bs.batches)-1] + if last.RowCount() < bs.batchMaxRow && + !vectorAllocationSelectionsMatch(last, source) { + return 1 + rowCount/bs.batchMaxRow + } + } + return bs.ReadyDelta(rowCount) +} + func (bs *BatchSet) Get(idx int) *Batch { if idx >= len(bs.batches) { return nil @@ -120,9 +138,14 @@ func (bs *BatchSet) Push(mpool *mpool.MPool, inBatch *Batch) error { return nil } - defer func() { - inBatch.Clean(mpool) - }() + if !vectorAllocationSelectionsMatch(bs.batches[batLen-1], inBatch) { + // A Vector cannot change provenance after its first allocation. Seal the + // existing partial tail and preserve the incoming Batch as a new owner. + bs.batches = append(bs.batches, inBatch) + return nil + } + + defer inBatch.Clean(mpool) // fast path 2 if lastBatRowCount+inBatch.RowCount() <= bs.batchMaxRow { @@ -148,8 +171,12 @@ func (bs *BatchSet) Extend(mpool *mpool.MPool, inBatch *Batch, reuseBuf *Batch) // empty bats or last batch is full - can directly use fast path lastIdx := batLen - 1 + if batLen > 0 && bs.batches[lastIdx].rowCount < bs.batchMaxRow && + !vectorAllocationSelectionsMatch(bs.batches[lastIdx], inBatch) { + return bs.extendWithNewProvenance(mpool, inBatch, reuseBuf) + } if batLen == 0 || bs.batches[lastIdx].rowCount == bs.batchMaxRow { - if reuseBuf != nil && len(reuseBuf.Vecs) == len(inBatch.Vecs) { + if vectorAllocationSelectionsMatch(reuseBuf, inBatch) { reuseBuf.CleanOnlyData() reuseBuf, err = reuseBuf.AppendWithCopy(context.TODO(), mpool, inBatch) if err != nil { @@ -168,7 +195,7 @@ func (bs *BatchSet) Extend(mpool *mpool.MPool, inBatch *Batch, reuseBuf *Batch) // fast path 2: inBatch is full if inBatch.rowCount == bs.batchMaxRow { - if reuseBuf != nil && len(reuseBuf.Vecs) == len(inBatch.Vecs) { + if vectorAllocationSelectionsMatch(reuseBuf, inBatch) { reuseBuf.CleanOnlyData() reuseBuf, err = reuseBuf.AppendWithCopy(context.TODO(), mpool, inBatch) if err != nil { @@ -205,6 +232,13 @@ func (bs *BatchSet) Union(mpool *mpool.MPool, inBatch *Batch, sels []int32, reus if selsLen > inBatch.RowCount() { panic("sels len > inBatch.RowCount()") } + if bs.Length() > 0 { + last := bs.batches[bs.Length()-1] + if last.rowCount < bs.batchMaxRow && + !vectorAllocationSelectionsMatch(last, inBatch) { + return bs.unionWithNewProvenance(mpool, inBatch, sels, reuseBuf) + } + } consumed := false if bs.Length() == 0 { @@ -283,7 +317,7 @@ func (bs *BatchSet) Union(mpool *mpool.MPool, inBatch *Batch, sels []int32, reus } func (bs *BatchSet) getOrCreateBatch(inBatch *Batch, reuseBuf *Batch, mpool *mpool.MPool) (*Batch, error) { - if reuseBuf != nil && len(reuseBuf.Vecs) == len(inBatch.Vecs) { + if vectorAllocationSelectionsMatch(reuseBuf, inBatch) { reuseBuf.CleanOnlyData() return reuseBuf, nil } @@ -291,9 +325,80 @@ func (bs *BatchSet) getOrCreateBatch(inBatch *Batch, reuseBuf *Batch, mpool *mpo for i := range tmpBat.Vecs { tmpBat.Vecs[i] = vector.NewOffHeapVecWithType(*inBatch.Vecs[i].GetType()) } + if inBatch.allocationAccount != nil { + if err := tmpBat.SetAllocationAccount(inBatch.allocationAccount); err != nil { + tmpBat.Clean(mpool) + return nil, err + } + return tmpBat, nil + } + for i := range tmpBat.Vecs { + if selection := inBatch.Vecs[i].AllocationAccountSelection(); selection != nil { + if err := tmpBat.Vecs[i].SetAllocationAccount(selection); err != nil { + tmpBat.Clean(mpool) + return nil, err + } + } + } return tmpBat, nil } +func (bs *BatchSet) extendWithNewProvenance( + mp *mpool.MPool, + inBatch *Batch, + reuseBuf *Batch, +) (bool, error) { + consumed := false + for start := 0; start < inBatch.RowCount(); { + tmpBat, err := bs.getOrCreateBatch(inBatch, reuseBuf, mp) + if err != nil { + return consumed, err + } + if tmpBat == reuseBuf { + consumed = true + reuseBuf = nil + } + bs.batches = append(bs.batches, tmpBat) + count := min(bs.batchMaxRow, inBatch.RowCount()-start) + if err := tmpBat.UnionWindow(inBatch, start, count, mp); err != nil { + return consumed, err + } + start += count + } + return consumed, nil +} + +func (bs *BatchSet) unionWithNewProvenance( + mp *mpool.MPool, + inBatch *Batch, + sels []int32, + reuseBuf *Batch, +) (bool, error) { + consumed := false + for start := 0; start < len(sels); { + tmpBat, err := bs.getOrCreateBatch(inBatch, reuseBuf, mp) + if err != nil { + return consumed, err + } + if tmpBat == reuseBuf { + consumed = true + reuseBuf = nil + } + bs.batches = append(bs.batches, tmpBat) + count := min(bs.batchMaxRow, len(sels)-start) + for i := range tmpBat.Vecs { + if err := tmpBat.Vecs[i].UnionInt32( + inBatch.Vecs[i], sels[start:start+count], mp, + ); err != nil { + return consumed, err + } + } + tmpBat.rowCount = count + start += count + } + return consumed, nil +} + func (bs *BatchSet) RowCount() int { rowCount := 0 for _, bat := range bs.batches { diff --git a/pkg/sql/colexec/multi_update/s3writer_delegate.go b/pkg/sql/colexec/multi_update/s3writer_delegate.go index 962e43fb001fe..91d6a534776fa 100644 --- a/pkg/sql/colexec/multi_update/s3writer_delegate.go +++ b/pkg/sql/colexec/multi_update/s3writer_delegate.go @@ -39,6 +39,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/sql/colexec" "github.com/matrixorigin/matrixone/pkg/sql/colexec/deletion" plan2 "github.com/matrixorigin/matrixone/pkg/sql/plan" + sqlutil "github.com/matrixorigin/matrixone/pkg/sql/util" "github.com/matrixorigin/matrixone/pkg/vm/engine/disttae" "github.com/matrixorigin/matrixone/pkg/vm/engine/tae/containers" "github.com/matrixorigin/matrixone/pkg/vm/engine/tae/options" @@ -286,7 +287,9 @@ func (writer *s3WriterDelegate) append( if tableType == UpdateMainTable { if mainTableNullPkFilter { var checked *batch.Batch - if checked, err = projBat.Clone(mp, false); err != nil { + // This validation copy leaves any allocation-accounted join owner; + // it is short-lived and never published back into that owner. + if checked, err = sqlutil.CopyBatch(projBat, proc); err != nil { return } nulls := checked.Vecs[mainTablePkProjectIdx].GetNulls().GetBitmap().Clone() @@ -321,7 +324,9 @@ func (writer *s3WriterDelegate) append( // Clone because SelectColumns shares vectors, and ShrinkByMask // modifies in-place. var filtered *batch.Batch - if filtered, err = projBat.Clone(mp, false); err != nil { + // The sinker owns the filtered copy independently of the input + // pipeline, so cross the allocation ownership boundary explicitly. + if filtered, err = sqlutil.CopyBatch(projBat, proc); err != nil { return } nullIdx := writer.sortIndexes[i] diff --git a/pkg/sql/colexec/shuffle/shufflepool.go b/pkg/sql/colexec/shuffle/shufflepool.go index d038e2db6abd8..9fbec8594b4be 100644 --- a/pkg/sql/colexec/shuffle/shufflepool.go +++ b/pkg/sql/colexec/shuffle/shufflepool.go @@ -618,7 +618,7 @@ func (sp *ShufflePool) tryWrite( sp.batchLocks[bucket].Unlock() break } - readyDelta := sp.batchSets[bucket].ReadyDelta(len(chunk)) + readyDelta := sp.batchSets[bucket].ReadyDeltaFor(srcBatch, len(chunk)) wait, ok := sp.reserveReady(int32(bucket), readyDelta) if !ok { sp.batchLocks[bucket].Unlock() diff --git a/pkg/sql/colexec/shuffle/shufflepool_test.go b/pkg/sql/colexec/shuffle/shufflepool_test.go index ab8ea47204816..60db6768d0448 100644 --- a/pkg/sql/colexec/shuffle/shufflepool_test.go +++ b/pkg/sql/colexec/shuffle/shufflepool_test.go @@ -182,6 +182,57 @@ func TestShufflePoolBoundsReadyBatchesAndResumes(t *testing.T) { require.Equal(t, int64(0), proc.Mp().CurrNB()) } +func TestShufflePoolReservesReadyCreditForProvenanceChange(t *testing.T) { + mp := mpool.MustNewZero() + proc := testutil.NewProcessWithMPool(t, "", mp) + defer proc.Free() + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2) + require.NoError(t, err) + + sp := NewShufflePool(1, 1, true) + legacy := testutil.NewBatch([]types.Type{types.T_int64.ToType()}, false, 2, mp) + done, err := writeBatchToBucketForTest(sp, legacy, proc, 0) + require.NoError(t, err) + require.True(t, done) + require.Zero(t, sp.readyCount) + + accounted := batch.NewWithSchema( + true, + nil, + []types.Type{types.T_int64.ToType()}, + ) + require.NoError(t, accounted.SetAllocationAccount(selection)) + require.NoError(t, vector.AppendFixed(accounted.Vecs[0], int64(7), false, mp)) + accounted.SetRowCount(1) + done, err = writeBatchToBucketForTest(sp, accounted, proc, 0) + require.NoError(t, err) + require.True(t, done) + require.Equal(t, 1, sp.readyCount) + + ready := sp.getAnyFullBatch() + require.NotNil(t, ready) + require.Equal(t, 2, ready.RowCount()) + sp.discardBatch(ready, mp) + require.Zero(t, sp.readyCount) + tail := sp.getAnyLastBatch() + require.NotNil(t, tail) + require.Equal(t, 1, tail.RowCount()) + require.Same(t, selection, tail.AllocationAccountSelection()) + sp.discardBatch(tail, mp) + + legacy.Clean(mp) + accounted.Clean(mp) + sp.abort(mp) + require.Zero(t, account.Seal().Used) + _, err = registry.Finalize(account) + require.NoError(t, err) + require.Equal(t, int64(0), mp.CurrNB()) +} + func TestShufflePoolFixedBucketsHaveIndependentBackpressure(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() diff --git a/pkg/vm/engine/disttae/txn_table.go b/pkg/vm/engine/disttae/txn_table.go index 82591c6a71756..858db5834ba58 100644 --- a/pkg/vm/engine/disttae/txn_table.go +++ b/pkg/vm/engine/disttae/txn_table.go @@ -3599,10 +3599,9 @@ func (tbl *txnTable) GetExtraInfo() *api.SchemaExtra { func dupVectorWithoutNulls(v *vector.Vector, mp *mpool.MPool) (*vector.Vector, error) { if !v.HasNull() { if v.AllocationAccountSelection() != nil { - // PK validation runs during transaction commit, after the producing SQL - // attempt may already have sealed its allocation account. The sorted - // copy is a short-lived transaction-engine owner, not a continuation of - // the statement owner, so make that ownership exit explicit. + // PK validation borrows caller-owned data. Its locally sorted copy is a + // short-lived transaction-engine owner, not a continuation of the + // statement owner, so make that ownership exit explicit. return v.DupOffHeapWithAllocation(mp, nil) } return v.Dup(mp) From 8ca209c69bebfa1e395ec8cad1bd2ff7bf669a5b Mon Sep 17 00:00:00 2001 From: aptend Date: Sat, 1 Aug 2026 02:02:54 +0800 Subject: [PATCH 37/61] docs: record allocation provenance closure --- ...ocation_accounted_memory_admission_impl.md | 18 +++++----- .../evidence/26459_activation_validation.md | 36 ++++++++++++------- 2 files changed, 33 insertions(+), 21 deletions(-) diff --git a/docs/design/allocation_accounted_memory_admission_impl.md b/docs/design/allocation_accounted_memory_admission_impl.md index 87491c884766b..124d7c4741f1d 100644 --- a/docs/design/allocation_accounted_memory_admission_impl.md +++ b/docs/design/allocation_accounted_memory_admission_impl.md @@ -6,14 +6,14 @@ - Architecture: [Allocation-Accounted Memory Admission RFC](../rfcs/00000000_allocation_accounted_memory_admission.md) - Baseline at plan creation: `main` at `38ce3a774` -- Rebased implementation baseline: `main` at `a9c13b02d8` +- Latest implementation baseline: `main` at `26ed429b60` - Allocation-site closure baseline (PR 3): `f61b64d56c` - Lifecycle/activation heads (PRs 4--10): `8633757dc1`, `5ae8eca00a`, `8e4b689f45`, `c9ad0ea810`, `e072568998`, `eec23dcdc5`, and `383fc6dce3` - Owner-atomic activation and cleanup: `656e254fe6` -- Post-activation ownership closure: `73d20085c0`, `3f10c23f21`, and - `91839ea91f` +- Post-activation ownership closure: `73d20085c0`, `3f10c23f21`, + `91839ea91f`, `abb1637882`, and `75cd58dc7a` - Merged prerequisite: #26455 at `93e8b22d2` - Independent design review: completed against RFC commit `a7d54cb5f` - Activation status: enabled for the closed HashBuild expression owner set. @@ -1079,15 +1079,15 @@ Required proof: Current validation status: -- the complete affected package matrix passes after rebasing on - `a9c13b02d8`; +- the complete affected package matrix passes after merging `26ed429b60`; - the complete affected package matrix, the same matrix under `-race`, a 20-iteration focused lifecycle/pressure race stress, affected-package vet, - `make static-check`, and `make build` pass through `91839ea91f`; + `make static-check`, and `make build` pass through `75cd58dc7a`; - accounted ordered/masked filtering, arbitrary shuffle rollback, `maxby` - retained-state compaction, generic `CopyBatch` ownership exit, CN bootstrap, - cross-package spill fixtures, and owner finalization-before-release have - dedicated regressions; + retained-state compaction, generic/transaction ownership exits, mixed Batch + clone and BatchSet provenance, provenance-aware ShufflePool credits, CN + bootstrap, cross-package spill fixtures, and owner finalization-before-release + have dedicated regressions; - local regression tests cover the five incident mechanisms without an estimator-only rejection on the activated owner; - TPCH 100G/1T candidate-versus-main runs remain the remote workload gate and diff --git a/docs/design/evidence/26459_activation_validation.md b/docs/design/evidence/26459_activation_validation.md index 55b55cea57f06..9f4095b835745 100644 --- a/docs/design/evidence/26459_activation_validation.md +++ b/docs/design/evidence/26459_activation_validation.md @@ -2,7 +2,7 @@ ## Candidate -- Rebased main: `a9c13b02d8` +- Latest main baseline: `26ed429b60` - PR 4 lifecycle: `8633757dc1` - PR 5 hash-table activation: `5ae8eca00a` - PR 6 retained batch/JoinMap activation: `8e4b689f45` @@ -14,6 +14,8 @@ - Vector/bitmap ownership boundary closure: `73d20085c0` - Cross-package spill fixture compatibility: `3f10c23f21` - Owner finalization-before-release: `91839ea91f` +- Transaction PK ownership exit: `abb1637882` +- Mixed Batch provenance closure: `75cd58dc7a` The activated expression owner is deliberately closed: COL, literal, param, variable, vector/fold, CONCAT, CASE, varchar EQUAL, and the audited string CAST @@ -37,27 +39,37 @@ the locally built `thirdparties` artifacts. ./pkg/sql/colexec ./pkg/sql/colexec/aggexec ./pkg/sql/colexec/hashbuild \ ./pkg/sql/colexec/hashjoin ./pkg/sql/colexec/dedupjoin \ ./pkg/sql/colexec/rightdedupjoin ./pkg/sql/colexec/spillutil \ - ./pkg/sql/compile ./pkg/util/resource ./pkg/vm/message ./pkg/vm/process + ./pkg/sql/colexec/shuffle ./pkg/sql/colexec/multi_update \ + ./pkg/sql/compile ./pkg/util/resource ./pkg/vm/message ./pkg/vm/process \ + ./pkg/vm/engine/disttae result: PASS ``` The same package matrix passes with `-race -p=2 -count=1`. `make static-check` -and `make build` also pass on the rebased candidate. A focused +and `make build` also pass on the latest-main candidate. A focused `-race -p=2 -count=20` stress also passes for concurrent account alloc/free, seal/open linearization, statement terminal one-shot behavior, owner-atomic activation, accounted spill/reduction, and all three join-consumer activation gates. Affected-package `go vet` and `make build` pass on the same candidate. -The previous PR CI failures were reduced to two branch-local ownership -boundaries and fixed before rerunning the matrix: generic `CopyBatch` now -performs an explicit compact clone when leaving an accounted statement owner, -which restores CN bootstrap, and the spill fixture writer remains available to -cross-package hashjoin/dedup tests. `Compile.clear` finalizes and detaches the -statement allocation owner before releasing operators back to reuse pools. -Dedicated tests verify accounted bitmap remap/shuffle rollback, `maxby` -selection preservation, bootstrap copy, fixture compilation, and that an -operator's `Release` cannot observe a live generation account. +An isolated `etc/launch/launch.toml` bootstrap using the freshly built binary +also passes on `26ed429b60` plus `75cd58dc7a`: the CN becomes queryable, system +catalog initialization completes, and the service log contains no allocation +account failure, terminal cleanup failure, or panic. + +The previous PR CI failures were reduced to concrete branch-local ownership +boundaries and fixed before rerunning the matrix. Generic `CopyBatch` and +transaction PK validation now make explicit exits from a borrowed statement +owner. Mixed Batch clones preserve per-Vector provenance; BatchSet preserves +uniform destination context, starts a new tail when provenance changes, and +ShufflePool reserves the corresponding ready credit. MultiUpdate's short-lived +validation/filter copies use the explicit generic boundary. The spill fixture +writer remains available to cross-package hashjoin/dedup tests, and +`Compile.clear` finalizes and detaches the statement allocation owner before +releasing operators back to reuse pools. Dedicated tests cover these paths, +accounted bitmap remap/shuffle rollback, `maxby` selection preservation, CN +bootstrap, fixture compilation, and release after finalization. Incident-mechanism regression mapping: From 31771e0941738eee3f1951fb0976f1bc66f660ee Mon Sep 17 00:00:00 2001 From: aptend Date: Sat, 1 Aug 2026 07:12:09 +0800 Subject: [PATCH 38/61] executor: converge memory admission on physical ownership --- ...ocation_accounted_memory_admission_impl.md | 1384 +---- .../evidence/26459_activation_validation.md | 150 - .../26459_allocation_accounting_bench.txt | 287 +- .../design/evidence/26459_local_validation.md | 95 + ...0_allocation_accounted_memory_admission.md | 1060 +--- pkg/common/bitmap/bitmap.go | 119 +- pkg/common/bitmap/types.go | 6 +- pkg/common/hashmap/inthashmap_lazy_test.go | 8 +- pkg/common/hashmap/inthashmap_test.go | 66 +- pkg/common/hashmap/iterator.go | 131 +- pkg/common/hashmap/keycodec/keycodec.go | 56 + pkg/common/hashmap/keycodec/keycodec_test.go | 80 + pkg/common/hashmap/strhashmap.go | 318 +- pkg/common/hashmap/strhashmap_test.go | 262 +- pkg/common/hashmap/types.go | 46 +- pkg/common/mpool/accounted_buffer.go | 63 + pkg/common/mpool/allocation_account.go | 4 +- pkg/compare/arraycompare.go | 6 +- pkg/compare/compare.go | 12 +- pkg/compare/grouping_compare_test.go | 90 + pkg/compare/strcompare.go | 6 +- .../batch/allocation_account_test.go | 241 +- pkg/container/batch/batch.go | 326 +- pkg/container/batch/grouping_codec.go | 131 + .../bytejson/bytejson_composite_plan.go | 358 -- .../bytejson/bytejson_composite_plan_test.go | 125 - pkg/container/bytejson/bytejson_keys_plan.go | 85 - .../bytejson/bytejson_keys_plan_test.go | 68 - .../bytejson/bytejson_scalar_plan.go | 66 - .../bytejson/bytejson_text_writer.go | 296 - .../bytejson/bytejson_text_writer_test.go | 60 - .../hashtable/allocation_account_test.go | 2 +- pkg/container/pSpool/buffer.go | 19 +- pkg/container/pSpool/copy.go | 86 +- pkg/container/pSpool/sender_test.go | 22 +- pkg/container/vector/allocation_account.go | 170 +- .../vector/allocation_account_test.go | 517 +- pkg/container/vector/functionTools.go | 609 +-- .../vector/function_result_allocation_test.go | 409 -- pkg/container/vector/pSpoolTools.go | 40 +- pkg/container/vector/tools.go | 27 +- pkg/container/vector/vector.go | 956 +++- pkg/container/vector/vector_test.go | 15 + pkg/sql/colexec/aggexec/maxby_test.go | 2 + pkg/sql/colexec/allocation_state.go | 65 + .../dedupjoin/allocation_test_helpers_test.go | 38 + .../dedupjoin/expression_memory_test.go | 108 - pkg/sql/colexec/dedupjoin/join.go | 154 +- .../dedupjoin/join_finalize_optimize_test.go | 1 + pkg/sql/colexec/dedupjoin/join_test.go | 82 +- .../colexec/dedupjoin/key_contract_test.go | 1 + pkg/sql/colexec/dedupjoin/types.go | 97 +- pkg/sql/colexec/evalExpression.go | 589 +- pkg/sql/colexec/evalExpressionReset.go | 64 +- pkg/sql/colexec/eval_expression_allocation.go | 262 - .../eval_expression_allocation_test.go | 545 -- pkg/sql/colexec/fuzzyfilter/filter.go | 35 +- pkg/sql/colexec/fuzzyfilter/filter_test.go | 54 +- pkg/sql/colexec/fuzzyfilter/types.go | 63 +- pkg/sql/colexec/group/helper.go | 4 +- .../hashbuild/allocation_test_helpers_test.go | 71 + pkg/sql/colexec/hashbuild/budget.go | 973 +--- pkg/sql/colexec/hashbuild/build.go | 294 +- pkg/sql/colexec/hashbuild/build_test.go | 828 +-- pkg/sql/colexec/hashbuild/dedup_memory.go | 9 +- .../colexec/hashbuild/expression_memory.go | 594 +- .../hashbuild/expression_memory_test.go | 1313 ----- .../hashbuild/expression_test_helpers_test.go | 94 + pkg/sql/colexec/hashbuild/hashmap.go | 552 +- pkg/sql/colexec/hashbuild/hashmap_test.go | 1853 ++----- pkg/sql/colexec/hashbuild/pressure.go | 13 +- pkg/sql/colexec/hashbuild/pressure_test.go | 2 +- pkg/sql/colexec/hashbuild/spill.go | 1293 +---- pkg/sql/colexec/hashbuild/spill_test.go | 1911 +------ pkg/sql/colexec/hashbuild/types.go | 112 +- .../hashjoin/allocation_test_helpers_test.go | 38 + pkg/sql/colexec/hashjoin/bitmap_mailbox.go | 98 + .../colexec/hashjoin/bitmap_mailbox_test.go | 105 + .../hashjoin/expression_memory_test.go | 108 - pkg/sql/colexec/hashjoin/join.go | 137 +- pkg/sql/colexec/hashjoin/join_test.go | 30 +- pkg/sql/colexec/hashjoin/key_contract_test.go | 1 + pkg/sql/colexec/hashjoin/mark_spill_test.go | 11 +- pkg/sql/colexec/hashjoin/spill_diskv2_test.go | 33 +- .../hashjoin/spill_integration_test.go | 163 +- pkg/sql/colexec/hashjoin/types.go | 98 +- pkg/sql/colexec/indexbuild/build.go | 64 +- pkg/sql/colexec/indexbuild/build_test.go | 56 +- pkg/sql/colexec/indexbuild/types.go | 66 +- pkg/sql/colexec/intersect/intersect.go | 5 +- pkg/sql/colexec/intersectall/intersectall.go | 5 +- pkg/sql/colexec/join_util.go | 55 +- pkg/sql/colexec/join_util_test.go | 2 +- pkg/sql/colexec/loopjoin/join.go | 52 +- pkg/sql/colexec/loopjoin/join_test.go | 56 +- pkg/sql/colexec/loopjoin/types.go | 83 +- pkg/sql/colexec/product/product.go | 95 +- pkg/sql/colexec/product/product_test.go | 62 +- pkg/sql/colexec/product/types.go | 89 +- .../colexec/productl2/joinmap_account_test.go | 97 + pkg/sql/colexec/productl2/product_l2.go | 10 +- pkg/sql/colexec/receiver_operator.go | 32 - .../allocation_test_helpers_test.go | 38 + .../rightdedupjoin/expression_memory_test.go | 108 - pkg/sql/colexec/rightdedupjoin/join.go | 123 +- pkg/sql/colexec/rightdedupjoin/join_test.go | 30 +- .../rightdedupjoin/key_contract_test.go | 1 + pkg/sql/colexec/rightdedupjoin/types.go | 84 +- pkg/sql/colexec/runtimefilter/contract.go | 94 +- .../colexec/runtimefilter/contract_test.go | 14 +- pkg/sql/colexec/sample/sample.go | 186 +- pkg/sql/colexec/sample/sample_test.go | 335 ++ pkg/sql/colexec/sample/types.go | 15 +- pkg/sql/colexec/shuffle/shufflepool_test.go | 8 +- .../colexec/spillutil/allocation_account.go | 47 +- .../spillutil/allocation_account_test.go | 243 +- .../spillutil/exact_test_helpers_test.go | 91 + pkg/sql/colexec/spillutil/join_spill.go | 1855 ++----- pkg/sql/colexec/spillutil/join_spill_test.go | 4806 ++--------------- .../compile/allocation_account_lifecycle.go | 242 +- .../allocation_account_lifecycle_test.go | 215 +- pkg/sql/compile/compile.go | 37 +- pkg/sql/compile/compile_test.go | 23 - pkg/sql/compile/operator.go | 15 +- pkg/sql/compile/operator_test.go | 17 +- pkg/sql/compile/scope.go | 8 + pkg/sql/compile/types.go | 2 +- pkg/sql/plan/function/baseTemplate.go | 128 +- pkg/sql/plan/function/func_binary.go | 879 +-- pkg/sql/plan/function/func_binary_aes_test.go | 120 +- .../func_binary_array_distance_gpu_test.go | 16 +- .../func_binary_array_distance_test.go | 135 +- pkg/sql/plan/function/func_builtin.go | 169 +- pkg/sql/plan/function/func_builtin_jq.go | 689 +-- pkg/sql/plan/function/func_builtin_json.go | 912 ++-- .../function/func_builtin_json_row_test.go | 179 - pkg/sql/plan/function/func_cast.go | 64 +- pkg/sql/plan/function/func_compare.go | 56 +- pkg/sql/plan/function/func_prefix.go | 143 +- .../plan/function/func_string_complex_test.go | 4 +- pkg/sql/plan/function/func_unary.go | 623 +-- .../function/func_unary_codec_scratch_test.go | 161 - .../function_allocation_scratch_test.go | 763 --- pkg/sql/plan/function/operator_in.go | 278 +- pkg/sql/util/copy_batch_test.go | 2 +- pkg/sql/util/eval_expr_util.go | 34 +- pkg/sql/util/eval_expr_util_test.go | 49 - pkg/util/resource/summary.go | 2 +- pkg/vectorindex/metric/cpu.go | 50 - pkg/vectorindex/metric/gpu.go | 317 +- pkg/vectorindex/metric/pairwise.go | 59 +- pkg/vectorindex/metric/pairwise_test.go | 32 - pkg/vm/engine/disttae/txn_test.go | 2 + pkg/vm/message/group_sels_test.go | 31 +- pkg/vm/message/joinMapDependency_test.go | 13 + pkg/vm/message/joinMapMsg.go | 160 +- pkg/vm/message/message_test.go | 160 +- pkg/vm/process/cte_memory_budget.go | 7 +- pkg/vm/process/hashbuild_budget.go | 638 +-- pkg/vm/process/hashbuild_budget_test.go | 2309 +------- 160 files changed, 10680 insertions(+), 29745 deletions(-) delete mode 100644 docs/design/evidence/26459_activation_validation.md create mode 100644 docs/design/evidence/26459_local_validation.md create mode 100644 pkg/compare/grouping_compare_test.go create mode 100644 pkg/container/batch/grouping_codec.go delete mode 100644 pkg/container/bytejson/bytejson_composite_plan.go delete mode 100644 pkg/container/bytejson/bytejson_composite_plan_test.go delete mode 100644 pkg/container/bytejson/bytejson_keys_plan.go delete mode 100644 pkg/container/bytejson/bytejson_keys_plan_test.go delete mode 100644 pkg/container/bytejson/bytejson_scalar_plan.go delete mode 100644 pkg/container/bytejson/bytejson_text_writer.go delete mode 100644 pkg/container/bytejson/bytejson_text_writer_test.go delete mode 100644 pkg/container/vector/function_result_allocation_test.go create mode 100644 pkg/sql/colexec/allocation_state.go create mode 100644 pkg/sql/colexec/dedupjoin/allocation_test_helpers_test.go delete mode 100644 pkg/sql/colexec/dedupjoin/expression_memory_test.go delete mode 100644 pkg/sql/colexec/eval_expression_allocation.go delete mode 100644 pkg/sql/colexec/eval_expression_allocation_test.go create mode 100644 pkg/sql/colexec/hashbuild/allocation_test_helpers_test.go delete mode 100644 pkg/sql/colexec/hashbuild/expression_memory_test.go create mode 100644 pkg/sql/colexec/hashbuild/expression_test_helpers_test.go create mode 100644 pkg/sql/colexec/hashjoin/allocation_test_helpers_test.go create mode 100644 pkg/sql/colexec/hashjoin/bitmap_mailbox.go create mode 100644 pkg/sql/colexec/hashjoin/bitmap_mailbox_test.go delete mode 100644 pkg/sql/colexec/hashjoin/expression_memory_test.go create mode 100644 pkg/sql/colexec/productl2/joinmap_account_test.go delete mode 100644 pkg/sql/colexec/receiver_operator.go create mode 100644 pkg/sql/colexec/rightdedupjoin/allocation_test_helpers_test.go delete mode 100644 pkg/sql/colexec/rightdedupjoin/expression_memory_test.go create mode 100644 pkg/sql/colexec/spillutil/exact_test_helpers_test.go delete mode 100644 pkg/sql/plan/function/func_builtin_json_row_test.go delete mode 100644 pkg/sql/plan/function/func_unary_codec_scratch_test.go delete mode 100644 pkg/sql/plan/function/function_allocation_scratch_test.go diff --git a/docs/design/allocation_accounted_memory_admission_impl.md b/docs/design/allocation_accounted_memory_admission_impl.md index 124d7c4741f1d..89cd9873f685c 100644 --- a/docs/design/allocation_accounted_memory_admission_impl.md +++ b/docs/design/allocation_accounted_memory_admission_impl.md @@ -1,1201 +1,183 @@ -# Allocation-Accounted Memory Admission: Implementation Plan - -- Status: implementation validation -- Tracking issue: - [#26459](https://github.com/matrixorigin/matrixone/issues/26459) -- Architecture: - [Allocation-Accounted Memory Admission RFC](../rfcs/00000000_allocation_accounted_memory_admission.md) -- Baseline at plan creation: `main` at `38ce3a774` -- Latest implementation baseline: `main` at `26ed429b60` -- Allocation-site closure baseline (PR 3): `f61b64d56c` -- Lifecycle/activation heads (PRs 4--10): `8633757dc1`, `5ae8eca00a`, - `8e4b689f45`, `c9ad0ea810`, `e072568998`, `eec23dcdc5`, and - `383fc6dce3` -- Owner-atomic activation and cleanup: `656e254fe6` -- Post-activation ownership closure: `73d20085c0`, `3f10c23f21`, - `91839ea91f`, `abb1637882`, and `75cd58dc7a` -- Merged prerequisite: #26455 at `93e8b22d2` -- Independent design review: completed against RFC commit `a7d54cb5f` -- Activation status: enabled for the closed HashBuild expression owner set. - Build/probe key closure is checked statement-atomically. Unsupported - expression families keep all participating HashBuild/HashJoin/DedupJoin - owners in that local attempt on the legacy path; they are not partially - mixed into an exact generation. - -## 1. Purpose and rules - -The RFC owns architecture and invariants. This file contains only implementation -decisions, allocation-site status, PR scopes, and evidence for #26459. - -Every implementation PR follows these rules: - -1. Migrate a complete owner closure: alloc, grow, reuse, handoff, Reset, Free, - and failure rollback. -2. One site is legacy, allocation-accounted, synchronously reclaimable named - scratch, or small statically bounded headroom metadata. -3. Enable exact accounting and remove the same owner's legacy hard charge in - one PR. -4. Do not call the account per row or on within-capacity reuse. -5. Do not store a mutable current account on Process or MPool. -6. Spill and cleanup memory remain accounted and preserve bounded progress. -7. Every merged PR is independently safe; a later PR cannot repair an unsafe - interval. -8. The final state has no permanent legacy/exact behavior switch. - -## 2. Starting point and allocation-site ledger - -### Allocation and container boundaries - -- `pkg/common/mpool/mpool.go`: `memHdr`, `Alloc`, `Grow`, `Grow2`, `Free`, and - `GrowCapacity`. `memHdr` is currently fixed at 16 bytes. Cross-pool Free - already delegates to the original MPool. -- `pkg/container/vector/vector.go`: data and varlen area are separate physical - buffers. `Reset*` retains capacity; `Free` releases owned buffers; `cantFree*` - identifies non-owning views. -- `pkg/container/batch/batch.go`: Clone/Dup/Union allocate destination buffers; - `CleanOnlyData` retains capacity and `Clean` frees vectors. -- `pkg/sql/colexec/evalExpression*.go` and - `pkg/container/vector/functionTools.go`: FunctionResult and expression - executors retain result capacity across Reset and release it at Free. - -### HashBuild and spill boundaries - -- `pkg/sql/colexec/hashbuild/{budget,hashmap,spill,types}.go`: copied batches, - expression estimates, hash-table tokens, auxiliary buffers, runtime-filter - marshal, and spill scratch. -- `pkg/common/hashmap`, `pkg/container/hashtable`: physical hash-table blocks - and resize plans. -- `pkg/vm/message/joinMapMsg.go`: reference-counted producer-to-consumer - handoff; final Free may happen after HashBuild Reset. -- `pkg/sql/colexec/spillutil/join_spill.go`: decoded batches, retained reuse, - scatter/row-ID/encode/decode buffers, and BucketReader/Writer cleanup. -- HashJoin, DedupJoin, and RightDedupJoin share the pressure and handoff - contract. - -### Merged #26455 boundary - -#26455 is the immediate correctness prerequisite, not the final accounting -mechanism: - -- it gives each HashBuild-owned expression root a generation-scoped retained - lease and fixes reuse/reset lifetime mismatches; -- it still reconciles executor-owned capacity at the operator layer and keeps - estimator-derived admission for uncovered growth overlap; -- it does not attach provenance to each MPool allocation, cover arbitrary - built-in Go-heap scratch, or provide statement terminal finalization; -- an activation PR removes the matching #26455 retained lease at the same time - exact allocation provenance becomes complete for that owner; -- exact MPool charges must never be stacked on the #26455 charge for the same - physical capacity. - -Ledger states: - -- `L`: legacy prediction/reservation; -- `D`: exact primitive exists but this production owner is dormant; -- `A`: allocation-accounted; -- `S`: named bounded explicit scratch; -- `H`: small, statically bounded Go/runtime metadata covered by CN headroom; -- `R`: removed. - -The independent review rejected owner-class rows as proof of closure. The -working ledger is allocation-site based: - -| Allocation site | Allocator/mode and size | Terminal owner | Current | Target/blocker | -| --- | --- | --- | ---: | --- | -| `mpool.memHdr` and account-ID side map | Go maps; one pointer record plus optional account record per live allocation | pointer removal at physical deallocation | H | bounded by the measured finite registry/allocation-slot policy | -| `Vector.data` | MPool; capacity from `Grow`, on/off-heap follows `v.offHeap` | owning `Vector.Free` | A | activated HashBuild destinations use immutable off-heap provenance | -| `Vector.area` | MPool; independent varlen payload capacity | owning `Vector.Free` | A | activated HashBuild destinations use immutable off-heap provenance | -| `Vector.nsp/gsp` bitmap data | allocation-accounted off-heap `[]uint64`; independent geometric capacity, paired replacement admission | owning `Vector.Free`; `Reset` retains and clears only published words | A | selected HashBuild vectors and expression results are active | -| ordered/masked Vector filtering | old bitmap backing is remapped in place; arbitrary row reordering uses allocation-accounted scratch admitted before publication | owning `Vector.Free`; scratch ends at `Shuffle` return | A | rejection leaves the old Vector and bitmap owners unchanged | -| `FunctionResult.vec` data/area | off-heap Vector; rows and appended payload | executor `Free` | A | active for the closed HashBuild expression set | -| `FunctionResult.convenientParam` | Go slice; expression arity, not rows | executor `Free`/reuse | H | bounded by plan expression arity, not input rows or payload | -| decimal parameter conversion | retained allocation-accounted off-heap buffer; `rows*sizeof(Decimal128)` for decimal64/float32/float64 promotion | `FunctionResult.Free`; Reset/evaluation reuse capacity | A | active when present in a closed expression owner | -| IFF/CASE/COALESCE selection arrays | allocation-accounted off-heap `[]bool`; one or two arrays of `rows`, retained by executor | executor `Free`/reuse | A | CASE is active in the closed expression set | -| selected row IDs | allocation-accounted off-heap `[]int64`; capacity up to `rows`, retained by executor | executor `Free`/reuse | A | active when present in a closed expression owner | -| selected parameter/result vectors | allocation-accounted off-heap Vector capacities | executor `Free` | A | active when present in a closed expression owner | -| `JSON_ROW` output | one retained allocation-accounted function scratch buffer, then copy into the admitted result; both physical capacities are charged | `FunctionResult.Free`; row publication copies synchronously | D | A after expression-owner activation; arity-bounded column closures are cleared after every call | -| float array-distance row descriptors and result scratch | Go `[][]T` with one descriptor per row plus `[]float32` with one value per row | call return / GC | R | caller row accessor plus the upper half of the admitted `[]float64` result backing | -| GPU array-distance flattened inputs | allocation-accounted caller scratch; `query bytes + rows*dimension*4` for the SQL float32 GPU path; legacy callers retain the C allocator | GPU job retains the caller slice until Wait; launch failure rolls back before return | D | A only after the GPU-tag build/tests pass; the legacy API remains unaccounted and is not an activated owner | -| `NORMALIZE_L2` output scratch | pooled or per-row Go arrays with one output element per input element | pool / call return and GC | R | normalize directly into admitted varlen result storage | -| AES/ENCODE/DECODE output scratch | Go payload slices, one plaintext/ciphertext-sized allocation per row; AES padding could append into spare input capacity | call return / GC | R | validate size/padding first and write directly into admitted result storage | -| JQ visible output | one retained allocation-accounted function scratch buffer, then copy into the admitted result | `FunctionResult.Free`; row publication copies synchronously | D | A after expression-owner activation; `TRY_JQ` propagates infrastructure/account rejection instead of converting it to SQL NULL | -| JQ object sort keys and gojq value graph | Go slices/maps/interfaces proportional to input JSON structure | row completion / GC | L | requires a separately bounded or allocation-accounted JSON execution owner; blocks JQ activation | -| JSON_ARRAY/OBJECT/KEYS/PRETTY output and object-key scratch | direct storage-compatible ByteJSON/result builders; JSON_OBJECT keys use retained allocation-accounted function scratch | result or `FunctionResult.Free` | D | A after expression-owner activation | -| JSON parse/path/modify/merge/schema execution | Go strings, paths, maps/slices, ByteJson modification payloads, and schema graphs proportional to row input | row completion / GC | L | requires an admitted JSON execution arena or streaming algorithms; blocks these functions from activation | -| HASH/IN/PREFIX_IN/narrow-array conversion scratch | retained allocation-accounted function scratch, bounded per hash chunk or exact tuple/array bytes | `FunctionResult.Free` | D | A after expression-owner activation | -| codec output (HEX/base64/MD5/COMPRESS/UNCOMPRESS/random/quote/date formatting) | writes directly into admitted result backing; flate's fixed codec state remains Go/runtime memory | result Free; codec state ends at call return | D | A after result activation; flate state needs an explicit fixed-headroom measurement before H classification | -| geometry parse/overlay scratch | Go point/ring/interval/match slices proportional to geometry payload | row completion / GC | L | admitted per-row geometry workspace or streaming algorithm; remains an activation blocker | -| regexp compile/match/output scratch | Go regexp program and match/output buffers proportional to pattern/input | operator cache / row completion / GC | L | bounded regexp owner or exclusion from expression activation | -| H3/S2 neighborhood scratch | Go slices; some S2 paths are statically bounded, H3 grid-disk output scales with radius | row completion / GC | L | split proved fixed bounds from data-scaled paths before activation | -| JSON cast visible serialization | `MarshalJSON` payload slices proportional to the JSON value | row completion / GC | L | direct visible writer or exclusion from expression activation | -| hash-table initial cell block | off-heap `mpool.MakeSlice(..., true)`; 16 KiB int / 32 KiB string | hash map / `JoinMap.FreeMemory` | A | physical allocation is the sole charge | -| hash-table replacement/appended cell blocks | off-heap blocks, at most 4 MiB each; old+new overlap is physically visible | hash map / `JoinMap.FreeMemory` | A | replacement overlap is admitted before publication | -| hash-table `cells`/`newBlocks` descriptors | owning off-heap descriptor buffer; initial and replacement capacities are distinct physical allocations | hash map / `JoinMap.FreeMemory` | A | no row-scaled Go descriptor backing remains | -| hash-table `ResizePlan` and callback | fixed-size Go values/closures, one per table/resize | resize return / hash map Free | H | table-count bounded metadata; no row-scaled backing | -| `GroupSels.{tmp,vals,offsets}` | allocation-accounted off-heap buffers; O(build rows/groups) | builder or `JoinMap.FreeMemory` | A | provenance follows final JoinMap consumer | -| copied build-batch vector buffers | MPool Vector data/area | builder or `JoinMap.FreeMemory` | A | physical Vector leases replace projected batch tokens | -| spill marshal/coalesce buffers | allocation-accounted off-heap streaming buffers; exact serialized size and bounded 64 KiB per-bucket coalesce capacity | spill phase cleanup | A | optional coalesce degrades to direct write under pressure | -| spill hash values | allocation-accounted off-heap `[]uint64`; geometric capacity, `8*cap` | spill phase cleanup | A | exact capacity and rollback covered | -| spill row IDs | allocation-accounted off-heap `[]int32`; geometric capacity, `4*cap` | spill phase cleanup | A | exact capacity and rollback covered | -| spill counts/offsets/positions | Go `[]int32`; O(bucket count), bucket count finite | spill cleanup | H | bucket count is the fixed spill fanout | -| selected spill bucket vectors | allocation-accounted off-heap Vector capacities | per-call selected-batch cleanup | A | adaptive unpublished windows bound progress | -| BucketReader decoded vectors | allocation-accounted MPool Vector data/area | reusable batch cleanup / `BucketReader.Close` | A | replacement and retry start from a clean record | -| BucketReader input buffer | one 64 KiB Go buffer per production spill reader; direct legacy readers use a bounded 4 MiB default | `BucketReader.Close` / GC | H | reader count is bounded by the live SpillEngine set and the statement/CN hash cap retains fixed non-payload headroom | -| `pSpool` cached Vector data/area | provenance-bearing detached MPool buffers; accounted data/area sites cannot cross and legacy buffers keep a guarded fast path | `spoolBuffer.clean` or the receiving Vector's `Free` | A | activated vectors retain immutable provenance through cache reuse | -| runtime-filter serialized payload | one allocation-accounted payload; optional filter publication | message release | A | capacity pressure degrades to PASS before publication | -| spill disk and FD | disk/FD ledgers | file removal/close | A | A | - -This is the known-site ledger for the activated HashBuild/join owner. Generic SQL -functions with unbounded JSON, geometry, regexp, H3/S2, or JSON-cast Go scratch -remain `L`; the build and probe activation gates reject the local statement's -entire owner set when one of those families appears. Consequently an activated -owner has no estimator-gated expression subtree, while an unsupported plan -remains wholly legacy and is an explicit later migration rather than a -partially exact owner. -No row named “other” or “unbounded scratch” can declare closure. - -Batch destination propagation is now `D`: Clone, Dup, selected-column copy, -Union destinations, windows, reader decode, Clean, and FreeColumns preserve -the immutable destination selection without creating a synthetic batch-level -charge. `pSpool` now transfers a detached buffer together with its immutable -selection and data/area site, reuses it only for the same provenance, and -returns every cache ID on construction failure. Its allocation-unaccounted -production path retains the old raw-slice representation behind guards that -reject an accounted Vector. `Vector.SetTypeAndFixData` now returns a failed -growth admission and rolls type and length back without losing the original -backing. - -Bitmap-aware selections are also `D`. Null and grouping backing remain -independent physical allocations, are included in `Vector.Allocated`, and use -the same immutable account/owner with distinct sites. The legacy `Bitmap` -footprint is unchanged: a tagged nonnegative/bitwise-complemented length -records backing ownership without adding a field. Accounted Vector growth -admits both bitmap replacements before either publishes, raw unadmitted bitmap -growth fails instead of escaping to the Go heap, Reset retains capacity while -clearing only represented words, copy/reader decode fills admitted backing -directly, and Free is the one terminal release owner. - -Post-activation filtering closes the corresponding mutation boundary. Ordered -filters remap null/group bitmaps in place, while arbitrary permutations admit -the replacement scratch against the same immutable allocation owner before -publication. `ShuffleWithBuf` cannot fall back to an unaccounted Go slice for -an accounted Vector. Aggregate compactors such as `maxby` preserve the same -selection while rewriting retained state. Generic `CopyBatch` is instead an -explicit ownership exit: it creates a compact clone owned by the destination -MPool and does not leak a statement-generation account into transaction or -bootstrap lifetime. - -## 3. Decisions required before production integration - -### A. Allocation metadata representation - -PR 0 prototype results on linux/amd64, Go 1.26.4, i7-11700: - -| Representation | Map construction bytes/base entry | insert+delete median (five-run range) | -| --- | ---: | ---: | -| current 16-byte `memHdr` | 55.94 | 40.13 ns (39.97--40.40) | -| inline 24-byte header, all allocations | 83.97 | 40.21 ns (40.00--40.67) | -| side map, 1% accounted | 56.24 | 40.17 ns (40.16--40.43) | -| side map, 10% accounted | 58.29 | 40.81 ns (40.77--41.19) | -| side map, 100% accounted | 93.76 | 49.39 ns (48.39--49.89) | - -A bounded fixed registry prototype used 40.04 bytes per slot, including its -account state. Concurrent lookup across 1,024 generation-tagged handles -measured 0.4278 ns/op aggregate median (0.2876--0.4463) at `GOMAXPROCS=8`, -versus 5.554 ns/op (5.225--5.682) for the `sync.Map` comparison; both reported -zero allocations. The exact command, environment, and all five samples are in -[the raw benchmark record](evidence/26459_allocation_accounting_bench.txt). - -The same artifact now records real unaccounted baselines at `GOMAXPROCS=8`: - -| Existing path | Median | Five-run range | -| --- | ---: | ---: | -| sharded MPool alloc/free, 64 B | 245.6 ns | 242.4--249.6 ns | -| sharded MPool alloc/free, 4 KiB | 291.8 ns | 290.4--295.6 ns | -| sharded MPool alloc/free, 64 KiB | 999.2 ns | 991.2--1,003 ns | -| sharded MPool grow, 64 B to 64 KiB | 1,267 ns | 1,244--1,279 ns | -| `noLock` MPool alloc/free, 64 B | 222.7 ns | 222.5--226.0 ns | -| parallel sharded MPool alloc/free, 64 B | 209.4 ns | 176.6--213.2 ns | -| fixed Vector pre-extend/free, 8,192 rows | 1,042 ns | 1,028--1,089 ns | -| varlen Vector data+area pre-extend/free | 28,011 ns | 27,634--28,170 ns | -| fixed Vector Reset/capacity reuse | 2.810 ns | 2.691--2.852 ns | - -These measurements establish the pre-integration comparison baseline. The -prototype-map results are still not final accounted-MPool results, so they do -not by themselves freeze the representation. - -A test-only account-aware wrapper around the real MPool then measured: - -| Prototype path | Median | Five-run range | Delta from sharded baseline | -| --- | ---: | ---: | ---: | -| alloc/free, 64 B | 327.0 ns | 321.4--329.7 ns | +33.1% | -| alloc/free, 4 KiB | 367.6 ns | 361.9--381.9 ns | +26.0% | -| alloc/free, 64 KiB | 1,085 ns | 1,072--1,106 ns | +8.6% | -| grow, 64 B to 64 KiB | 1,400 ns | 1,390--1,418 ns | +10.5% | -| parallel alloc/free, 64 B | 227.2 ns | 203.1--233.5 ns | +8.5% | - -Every sample reported zero Go allocations. This wrapper intentionally takes a -second metadata-shard lock after MPool has already published its pointer -header. The 26--33% small-allocation cost rejects that shape for production; -PR 1 must publish the optional account side record under MPool's existing -pointer-shard transaction and re-run the comparison. The wrapper remains a -conservative upper bound and validates real allocation/growth rollback. - -Provisional choice: - -- keep `memHdr` at 16 bytes and replace the final `offHeap` byte with a - three-state allocation kind: on-heap, unaccounted off-heap, or accounted - off-heap; -- store a compact 16-byte `pointer -> account pointer + owner/site` lease only - for accounted allocations in the same pointer shard, or in the same - pool-local metadata store for `noLock` pools. The direct account pointer - keeps the original generation alive through late physical `Free` without a - process-global registry lookup; -- publish/remove the pointer header and optional account ID in one metadata - transaction under the same lock; a metadata failure rolls back both before - allocation publication; -- use a finite registry of reusable slots and encode `(slot, generation)` in - the handle; reuse a slot only after attempt-owned seal and exact zero, and - retire it rather than allowing the generation counter to wrap; -- reserve one finite CN-local allocation-metadata slot before publishing an - accounted allocation or replacement and return it on physical Free; -- size the registry and side-record backing stores from explicit CN headroom, - rather than allowing maps to grow without a hard count bound. - -The first hash-table activation fixes the initial sizing policy: - -- reserve 131,072 generation slots. The production registry plus a live - 64-byte account measures 80.09 bytes/slot; budgeting 128 bytes/slot reserves - 16 MiB. This default exceeds the - frontend `max_connections` system variable's declared upper bound of 100,000 - and leaves 31,072 slots for remote/internal attempts. Only an attempt whose - physical plan contains an activated owner opens a slot. The slot limit is - itself the hard supported activated-attempt concurrency; a deployment that - intends to support more must raise it and pass the same headroom check before - enabling the owner; -- every integer/string cell block is at least 16 KiB. After the outer - descriptor is made an owning accounted off-heap allocation, every live - table has at most one live descriptor buffer per live table version. - Therefore the first activation uses - `3 * ceil(MPoolGlobalCap / 16 KiB)` allocation-metadata slots. Two units - cover every live cell block plus its table's published descriptor; the third - covers one unpublished replacement descriptor per live table before the - matching cell allocation either publishes or rolls back; -- the production base pointer map plus a 16-byte all-accounted lease map - measures 111.87 bytes/entry, of which 55.93 bytes/entry is incremental over - the existing pointer map. Budget 128 bytes/allocation-metadata slot to cover - sparse shards and Go-map growth overlap, and 128 bytes/generation slot for - the final account fields. At MPool's 1 GiB minimum global cap this is 24 MiB - plus 16 MiB; at - larger caps the allocation component is 2.34375% of the MPool cap and the - fixed registry fraction decreases. PR 1 must measure construction, sparse - occupancy, and grow/evacuate high water; exceeding these conservative - constants blocks merge rather than silently consuming payload headroom; -- startup must reserve - `allocationSlots * 128 + generationSlots * 128` bytes outside the MPool - payload cap. If the host/container limit cannot supply it, the activated - owner is refused at startup rather than running with an unproved headroom - assumption. - -Later activations with smaller allocations must derive a new simultaneous -allocation bound and resize this headroom before they can become `A`; they -cannot inherit the hash-cell formula merely because the generic API exists. - -This choice is not frozen until real MPool alloc/grow/free, cross-pool Free, -deleted-pool fallback, real accounted ratios, and P50/P99 concurrent latency -match the prototype result. - -Metadata is `H`, not silently included in payload capacity. Safety comes from -finite generation and allocation-metadata slot limits; slot exhaustion is a -typed exact-pressure result. Each activation PR must also show that its -supported simultaneous allocations and generations fit the configured limits, -using measured pointer/side/registry bytes per entry and resulting aggregate -CN headroom. This prevents a safe-but-impractical false metadata-pressure -regression. - -### B. Account-aware API shape - -The provisional API rules are: - -- only an explicit first off-heap allocation accepts an account, owner, and - site; -- ordinary `Grow` inherits account provenance from allocation metadata and - takes no replacement account argument; -- account-A memory cannot grow under account B; -- an accounted on-heap allocation is rejected; -- unaccounted-to-accounted conversion allocates a new destination; -- ordinary unaccounted callers keep current behavior; -- helper delegation cannot silently drop the account. - -One `AllocationCapacity` rule must cover initial allocation, growth, runtime -rounding, and `CapLimit-kMemHdrSz`. `recordPtrHdr` failure, allocator panic, -cross-pool Free, deleted-pool fallback, and pool teardown are explicit -transaction branches. - -### C. Generation owner and terminal snapshot - -One execution attempt of `Compile` on each CN owns the generation. The -statement `ResourceRoot` aggregates its immutable terminal snapshot, but does -not own allocation release. HashBuild Reset is also not the generation owner. - -The attempt opens before `prePipelineInitializer`/operator Prepare and owns: - -```text -all local scopes, remote notifiers, and message consumers quiescent - -> close and drain that attempt's MessageBoard - -> seal new admission - -> release remaining live leases - -> used=0 - -> export one immutable snapshot - -> remove registry entry -``` - -`Scope.Run` defers pipeline cleanup, and `Scope.MergeRun` joins pre-scopes and -remote notifier goroutines before returning. Therefore the local hook belongs -in a deferred attempt finalizer around `Compile.runOnce`, after its result and -before retry transition or attempt publication. A retry finalizes the failed -attempt before `buildRetryCompile` opens the next generation. - -The remote hook belongs after `Scope.MergeRun` and in the existing -`runCompile.clear` terminal defer, which releases operators, resets the -MessageBoard, and snapshots the remote MPool before replying. PR 4 must add an -explicit MessageBoard close-and-drain operation: ordinary multi-CN `Reset` -only removes the board from `StmtIDToBoard` because producers may still access -it, whereas the attempt hook runs after the existing sender/receiver cleanup -barriers have proved quiescence. - -The same deferred finalizer covers success, error, panic, cancellation, retry, -broadcast, remote execution, and a JoinMap freed after producer Reset. -`SetStmtProfile` turnover, frontend `StatementInfo.EndStatement`, and -`HashBuildBudgetGeneration.Close` are observability or operator boundaries, -not acceptable release substitutes. - -If terminal cleanup ends nonzero, the attempt coordinator exports one immutable -invariant-failure snapshot and retains a release-capable tombstone. That CN -admits no new accounted generation until all such tombstones drain to zero, so -registry growth is bounded by generations already active at detection. A -deadline escalates owner/site diagnostics and allows controlled CN restart; it -never deletes live provenance. - -Generation open and suspension publication use one CN-local linearization -gate. An open that linearizes after suspension cannot publish. - -PR 4 derives this terminal matrix: - -| Attempt path | Required terminal ordering | Oracle | -| --- | --- | --- | -| local success | scopes join -> board close/drain -> seal -> zero -> publish | one valid snapshot, no queued message | -| local error/cancel/panic | cancellation -> every started scope cleanup/join -> board drain -> seal | one terminal snapshot; no goroutine or allocation survives | -| failure before `runOnce` | opened generation -> initializer rollback -> board drain -> seal | zero or one named invariant failure, never an abandoned open slot | -| retry | attempt N fully finalizes -> attempt N+1 opens | old handles are stale; no cross-attempt publication | -| remote execution | remote `MergeRun` joins -> `runCompile.clear`/board drain -> snapshot -> response | parent receives one immutable child snapshot | -| broadcast/late JoinMap Free | producer Reset -> every consumer cleanup -> queued refs drain | physical final Free releases the original generation exactly once | -| prepared reuse | attempt finalizes and replaces its board -> cached pipeline Reset -> next attempt opens | no retained accounted capacity crosses statement generations | -| nonzero terminal | seal -> failure snapshot -> tombstone/suspend -> late Free | no new open until every tombstone reaches zero | - -### D. Go-heap boundary - -`MPool.Alloc(..., false)` records requested bytes but `Free` does not reclaim -them synchronously. Therefore: - -- all data/row/payload-scaled controlled allocations move off-heap before - activation; -- small Go metadata may be `H` only with a static bound and separate CN - headroom; -- no data-scaled Go slice may be relabeled `S` to bypass migration. - -### E. Pressure and operation rollback - -Capacity pressure, sealed generation, account mismatch, allocator-size limit, -and invariant corruption are distinct typed results. Only capacity pressure is -recoverable. - -Each retryable owner records an operation checkpoint and cleanup/restart rule. -For example, `Vector.PreExtendWithArea` may grow data and then fail area growth; -that retained growth is valid accounting state but not proof that re-running -the logical operation is idempotent. Before a shared controller exists, an -exact rejection is a controlled terminal pressure error. - -The PR 0 reference model validates the provisional shared rule: - -- replacements remain private while the old allocation stays published and - charged; -- all new allocations and replacements commit as one logical operation; -- cancellation or later allocation failure frees private allocations and - restores the checkpoint before retry; -- a second attempt cannot begin while the failed operation remains active; -- retry may reduce the requested capacity, but cannot duplicate publication; -- an owner that cannot preserve or reconstruct the checkpoint is not - retryable and returns the typed pressure error after cleanup. - -For each spill owner choose an already allocated reusable buffer, a finite -progress sub-cap, or a smaller chunk. Normal work and progress allocations stay -under the same total query/CN cap; no uncharged emergency scratch is allowed. - -## 4. Pull request sequence - -### PR 0: model and measured design decisions - -Scope: - -- close decisions A--E with prototypes and benchmarks; -- finalize bounded owner/site enums; -- implement a test-only reference state machine; -- record MPool/vector allocation baselines; -- complete the first activation owner inventory and record the generation - method for later expression/spill inventories without changing production - behavior. - -Current evidence: - -- the reproducible test-only artifact is - `experiment/26459-allocation-accounting-validation` at `cde44cd099`: - `pkg/common/mpool/allocation_account_validation_test.go` and - `pkg/common/mpool/allocation_account_benchmark_test.go` and - `pkg/container/vector/allocation_account_validation_test.go`; -- the test-only model passes alloc, within-capacity reuse, old+new growth, - injected unpublished failures, views, deep copy, Reset, multi-allocation - checkpoint/commit/rollback, cancellation before and after allocation, - smaller retry without duplicate publication, bounded generation and - allocation-metadata slots, stale slot generations, exact metadata overlap - on replacement, generation-counter exhaustion without wrap, accounted - on-heap rejection, sealed-error precedence, handoff, cross-pool Free, - sealed-vs-capacity errors, zero finalization, nonzero - tombstone/suspension/drain, normal and `noLock` pool teardown, - open-vs-suspend linearization, stale generation, and 20,000 deterministic - randomized operations; -- the metadata and contention microbenchmarks in section 3 passed five runs; -- real unaccounted MPool alloc/free/grow and Vector allocate/reuse baselines in - section 3 passed five runs; -- the account-aware real-MPool wrapper passes allocation, reuse, old+new growth, - rollback, and final-zero validation; its five-run result rejects a second - side-metadata lock for production; -- at `GOMAXPROCS=8`, a serialized mutex acquire/release prototype measured - 80.73 ns/op median versus 21.11 ns/op for a two-operation atomic prototype, - so real aggregate-account contention remains a mandatory design benchmark; -- production behavior is unchanged. - -Gate: - -- the existing MPool/Vector baseline and rejected separate-lock shape have - reproducible five-run measurements; the selected same-shard transaction is - an explicit PR 1 merge gate; -- the model covers alloc, grow, failure, view/copy, Reset, operation rollback, - cancellation/retry, finite metadata slots, handoff, cross-pool Free, seal, - and stale generation; -- per-entry metadata and maximum simultaneous allocation/generation counts - prove finite aggregate CN headroom; -- the site ledger is complete for cell and descriptor initial allocation, - replacement, segmented growth, rollback, and terminal Free of the hash-table - first activation; -- generation owner, Go-heap classification, typed errors, and retry checkpoint - decisions are closed; -- production owners remain `L`. - -PR 0 design gates are closed. The metadata representation remains provisional -until PR 1's integrated benchmark passes, and no production owner may switch -from `L` to `A` before PRs 1--4 close. The separate-lock prototype is a -recorded rejected design, not an implementation candidate. - -### PR 1: generic account and MPool allocation transaction - -Scope: - -- low-level account contract below SQL/process; -- compact account-ID registry and dormant `HashBuildBudgetGeneration` adapter; -- finite generation/allocation-metadata slot limits and their typed pressure - results; -- account-aware alloc, Grow/Grow2, Free, and immutable snapshots. - -Required behavior: - -- reject accounted on-heap allocation; -- reserve the complete new capacity before allocation; -- for growth keep old and complete new capacity live until publication; -- roll back on admission, MPool/global-cap, metadata, or allocation failure; -- use one allocation-capacity rule for initial and growth boundaries; -- release through normal and cross-pool physical Free, deleted-owner-pool - fallback, and `noLock` teardown that physically deallocates; -- retain metadata and charge when normal-pool teardown only unregisters the - pool, and report live accounted allocations there as an invariant; -- reject account mismatch and stale handles. - -Gate: - -- exact/one-byte-short, allocator rounding, `CapLimit-kMemHdrSz`, and - `GrowCapacity` boundaries; -- old+new overlap; -- injected failure or panic at every unpublished step, including metadata; -- atomic header/account-ID publication and removal for sharded and `noLock` - pool metadata; -- normal-pool unregister plus late Free, `noLock` physical teardown, and no - premature release in either case; -- concurrent acquire/release, double Free, seal, and final zero; -- measured unaccounted/accounted alloc/free/grow overhead and concurrent - generation P50/P99 latency; -- no production owner selects an account. - -The current PR 1 candidate is -`feature/26459-allocation-account` at generic commit `766e1501c3` and -HashBuild-adapter commit `0655af4443`. It remains dormant. Its same-lock -pointer/lease transaction, finite registry, stale-handle checks, old+new -growth, deleted-pool/noLock lifetime rules, and tokenless HashBuild adapter are -implemented. Returned-error and panic rollback is injected after account, -metadata, global stats, pool stats, physical allocation, and header -publication for both sharded and `noLock` metadata. Normal package tests, -package vet, full package race tests, and the focused lifecycle/rollback race -matrix at 100 repetitions pass. Representation, latency, and integrated -benchmarks are recorded in the evidence artifact. No production owner selects -an account, and no legacy hard charge has been removed. - -### PR 2: dormant Vector and Batch propagation - -Scope: - -- optional account selection for owning off-heap Vector buffers; -- data and area growth; -- Batch Clone/Dup/Union destination propagation. - -Required behavior: - -- Reset retains charge; Free releases it; -- within-capacity append performs no account operation; -- aliases/views/const/shared area do not create another charge; -- deep copies use the destination account; -- at the PR 2 boundary, on-heap null/group bitmaps remain explicit later-PR - blockers and are not silently included in the Vector charge; -- HashBuild production remains legacy until a later owner migration. - -Gate: - -- randomized fixed/varlen append; -- Reset/reuse/Free, views, partial selection, copy rollback, cross-pool Free; -- package race tests and vector benchmarks; -- Vector/Batch ledger rows become `D`. - -The current PR 2 candidate is -`feature/26459-vector-propagation` at commit `dbfee20ecc`. It remains dormant. -It adds one immutable shared selection pointer to Vector and Batch, accounts -the first owned off-heap data/area allocation, lets later Grow/Grow2 inherit -the physical MPool lease, and rejects implicit conversion to on-heap or -no-copy aliases. Reset retains the selection and charge; Free clears the -selection after the physical allocations release their leases. Views carry no -selection, while Batch windows retain only the destination context needed for -a later deep copy. - -The implementation also closes two error edges found during self-review: -reader growth publishes the replacement buffer before a short read can return, -so cleanup never retains a freed old pointer, and no-copy Batch decode -explicitly detaches an empty Vector selection while retaining the Batch -destination context. - -Fresh local evidence on linux/amd64, Go 1.26.4, i7-11700: - -| Vector operation, `GOMAXPROCS=8` | Legacy median | Accounted median | Difference | -| --- | ---: | ---: | ---: | -| fixed pre-extend/free, 8,192 rows | 1,026 ns | 1,116 ns | +8.8% | -| varlen data+1 MiB area pre-extend/free | 76,766 ns | 77,650 ns | +1.2% | -| accounted fixed Reset/reuse | n/a | 1.524 ns | no account operation | - -Fixed paths remain 0 B/op and 0 allocs/op. Both varlen paths report the same -48 B/op and 2 allocs/op, so accounting adds no Go allocation. Randomized -fixed/varlen append, separate data/area charge, within-capacity reuse, Reset, -Free, views, partial selection, cross-owner copies, metadata rollback, -cross-pool Free, sealed accounts, shuffle replacement, copy/reader decode, and -Batch Clone/Dup/Union/FreeColumns pass. Every new and directly affected test -passed an exact `-race -count=100` run; both owning packages passed complete -race runs, build, vet, coverage, and dependent HashBuild/SQL/engine package -tests. No production owner selects an account and no legacy hard gate is -removed. - -### PR 3: allocation-site closure and dormant propagation - -Scope: - -- complete the generated expression/built-in and spill allocation-site ledger; -- propagate dormant accounts through FunctionResult, expression result, - selected result, decoded Vector, and off-heap scratch constructors; -- replace data-scaled Go vector null/group bitmaps, - selection/conversion/hash/row-ID/serialization buffers with off-heap owners - or direct output; -- do not enable production accounting or remove a legacy hard gate. - -Gate: - -- every reachable `make`, capacity-growing `append`, `bytes.Buffer`, MPool, and - nested executor result has allocator, bound, terminal owner, and test; -- fixed/varlen/const/null, CAST, CONCAT, CASE, nested and selected paths close; -- repeated Eval/Reset/Free and construction failure reach the same terminal - owners; -- all migrated rows become `D`; production rows remain `L`. - -The current dormant PR 3 candidate is -`feature/26459-expression-propagation` at commit `fd2fcc953b`, rebased on -`main` commit `60e36bef64`. Its ordered closure commits are: - -- `06d39287e4`: generic allocation-accounted MPool ownership; -- `968c68df4c`: dormant HashBuild budget bridge; -- `941212edf0`: Vector propagation; -- `61a6e7a0e2`: expression propagation; -- `b03ec24c2f`, `309d070c71`, and `556d64d5ba`: spill scratch, streaming, - and serialization; -- `86e041ff53`: retained Vector/spool ownership gaps; -- `2b05fdac79`: Vector bitmap and conversion scratch; -- `24b5f6fef7` and `57940c8f0f`: direct-output row/function scratch removal; -- `674f78dea9`: general retained function-scratch provenance; -- `fd2fcc953b`: direct codecs, HASH/IN/PREFIX, GPU SQL flattening, and - ByteJSON/JQ/JSON output closure. - -Its propagation call chains are: - -```text -NewExpressionExecutorWithAllocation - -> recursive expression construction - -> constant/result/scratch bitmap-aware AllocationAccountSelection - -> FunctionResult result/parameter/function scratch or selected/decode Vector growth - -> MPool AllocAccounted/Grow/Free - -NewSpillEngineWithAllocation - -> BucketReader decoded/reused Batch selection - -> scatter hash/row typed slices and selected Batch selection - -> exact streaming record and optional coalesce buffers - -> MPool AllocAccounted/Grow/Free - -Pipeline spool accounted copy - -> detach Vector data/area with immutable provenance - -> cache only by matching selection and data/area site - -> attach to the next owning Vector or free at spool cleanup -``` - -The candidate adds no production caller of either dormant constructor and -removes no legacy reservation. It covers fixed, varlen, NULL, decoded-vector, -nested `CASE(CONCAT(CAST))`, folded and non-folded result transfer, partial -selection, repeated Reset/reuse, construction rollback, zero-length retained -scratch growth, decoded-record merge/error cleanup, scatter selected-vector -peak, and capacity-failure cleanup. Typed slices grow geometrically only on -the accounted path; the old and replacement capacities are simultaneously -charged until publication, and terminal cleanup frees a zero-length view by -its retained capacity. - -Vector null/group bitmap growth now follows the same rule. Both replacement -buffers are admitted before publication, so rejection of the second buffer -rolls the first unpublished buffer back and preserves both old owners. -`pSpool` releases bitmap backing after data/area detach, recreates it under the -destination provenance, and preserves grouping semantics including constant -vectors. Accounted window/duplicate/decode paths pre-admit bitmap coverage -before legacy raw bitmap APIs can mutate it. - -Decimal128 promotion from decimal64, float32, and float64 now writes into one -retained allocation-accounted parameter buffer per argument. Const promotion -uses the scalar wrapper and allocates no row scratch; normal/null promotion -reuses the admitted buffer across evaluations. Capacity, sealed-account, and -metadata failures return through the function executor instead of panicking or -falling back to a Go slice. The allocation selection and parameter buffer must -share the same account and owner. - -The spill record path now streams Bitmap, Nulls, Vector, and Batch wire formats -directly into one allocation-accounted off-heap buffer. It computes the exact -wire size before admission, retains one record buffer for the phase, bounds -each bucket's coalesce buffer at 64 KiB, and degrades optional coalescing to a -direct write when payload or metadata capacity rejects it. Wire round trips, -legacy byte equivalence, retained-buffer reuse, coalesce fallback, write -failure, and phase cleanup are covered. - -The pipeline spool now preserves Batch and Vector selection through cached -copies, keeps data and area allocation sites distinct, refuses cross-account -reuse, returns the cache slot after allocation failure, and fixes a legacy -non-last cache removal that previously dropped a still-owned buffer -descriptor. `SetTypeAndFixData` publishes a fixed-width type change only after -growth succeeds and propagates its error through all four callers. - -The first syntax inventory over non-test expression/builtin and spill sources -found 484 candidate lines. The latest focused function scan still reports -candidate syntax, not live-byte proof: the largest `make([]...)` concentrations -are `func_unary.go` and `func_binary.go` (28 each) and -`func_builtin_json.go` (26); the largest potentially growing `append` groups -are `func_unary.go` (37), `func_builtin.go` (36), `func_cast.go` (16), and -`func_builtin_json.go`/`func_binary.go` (13 each). Buffer/builder and JSON -serialization scans are recorded in the site rows above. Admin/CTL/UDF code, -arity-only slices, fixed stack buffers, and legacy-only branches are not -silently counted as controlled payload, but each must be explicitly excluded -or bounded before activation. Vector null/group backing, decimal conversion, -direct codecs, HASH/IN/PREFIX scratch, SQL GPU flattening, and visible -JSON/JQ output no longer create data-scaled unaccounted payload on the dormant -path. JQ graphs, JSON parse/modify/schema, geometry, regexp, scalable H3, and -JSON cast serialization remain explicit activation blockers; the dormant -constructor is not a license to enable partial accounting. - -The first follow-up scan also removed two allocations instead of moving them: -`FIELD` now writes directly into its pre-extended result and retains -arity-bounded parameter wrappers, while `VALUES` uses contiguous `UnionBatch` -instead of constructing one row ID per input row. Their exact tests passed -`-race -count=100` and the complete Function package passed normal, vet, and -race runs. Remaining candidates must first make this same -direct-output/streaming test before introducing a scratch owner. - -The second follow-up removes rather than accounts four more scratch families: - -- `JSON_ROW` no longer owns one encoder and buffer per row. It prepares - arity-bounded typed column closures once, streams one complete row through a - single reusable encoder, and clears every closure on success, error, or - panic. Legacy execution reuses one `bytes.Buffer`; dormant exact execution - uses one retained allocation-accounted function buffer and then copies the - completed row into the admitted result. The scratch and result capacities - are both physical and therefore both charged. The unsigned path preserves - the full `uint64` range. -- float array distance no longer materializes one `[]T` descriptor per input - row or a second result slice. The metric layer accepts a synchronous row - accessor and caller-owned output; SQL uses the upper half of its already - admitted `[]float64` result bytes as `[]float32` scratch and converts forward - only after Wait. The SQL float32 GPU path now flattens query and row inputs - into caller-owned allocation-accounted function scratch, and the GPU job - retains that slice until Wait. The public legacy GPU API keeps its C-allocator - behavior and is not part of the activated owner. -- `NORMALIZE_L2` writes float32, float64, BF16, Float16, int8, and uint8 results - directly into admitted varlen storage. The old float pools and per-row - widened/narrowed arrays are deleted. -- AES ECB/CBC and legacy ENCODE/DECODE write directly into admitted result - storage. AES validates the final decrypted block before result publication, - and padding is assembled in one stack block, removing both payload copies - and the old possibility that `append` modified spare input-vector capacity. - -The final dormant closure adds one retained `FunctionResult` scratch owner with -a site distinct from decimal parameter conversion. Detection is allocation-free; -the first nonzero resize admits physical capacity, geometric growth charges old -and replacement capacity until publication, Reset retains it, and -`FunctionResult.Free` is the sole terminal release. It is used for exact HASH -key chunks, fixed/string IN tuples, PREFIX_IN entries, narrow array conversion, -JSON_OBJECT keys, JQ/JSON_ROW visible output, JSON modify value -pre-materialization, and SQL float32 GPU flattening. One-byte-short tests cover -the generic allocator boundaries, and capacity-rejection tests cover each -new scalable function family. `TRY_JQ` suppresses jq/domain errors but never an output -allocation failure. - -Storage-compatible ByteJSON encoders now write scalar, array, object, -object-key-array, typed, opaque, bit, and nested values directly into Vector -area backing. JSON_ARRAY, JSON_OBJECT, JSON_KEYS, JSON_PRETTY, JSON_QUOTE, and -JSON_ROW no longer build a second output-sized ByteJSON or visible-text slice. -JSON_EXTRACT no longer retains `rows * path-count` path arrays: nonconstant -paths are parsed one row at a time into arity-sized reusable storage. This does -not close the parser/modifier/schema graphs themselves; those remain `L`. - -HEX/UNHEX, MD5, base64, vector-base64, COMPRESS/UNCOMPRESS, RANDOM_BYTES, -CHAR, MAKE_SET, EXPORT_SET, bitwise strings, array casts, quote, date/time -formatting, and FROM_UNIXTIME now publish directly into admitted result -capacity. The legacy production path for two-pass-capable formatting remains -one-pass; exact two-pass sizing is selected only by the dormant allocation -owner, so this PR does not impose duplicate formatting work before activation. - -`FunctionResult.AppendBytesWithBuilder` is the shared publication primitive; -`AppendBytesWithFill` is its exact-size convenience wrapper. It admits result -data/area capacity before exposing a row, accepts an actual length no larger -than the admitted capacity, collapses a short result into inline storage when -possible, and publishes length only after successful completion. Error, -invalid-length, and panic paths restore the unpublished varlena header and area -length. Capacity acquired before rollback remains owned and charged to the -result until normal reuse or `Free`; no allocation is released without its -terminal owner. - -Fresh five-run benchmarks on linux/amd64, Go 1.26.4, i7-11700: - -| 8,192-row function path | Before median | `fd2fcc953b` median | Before allocation | `fd2fcc953b` allocation | -| --- | ---: | ---: | ---: | ---: | -| float32 L2-squared batch distance, 128 dimensions | about 483 us | 437.8 us | 229,488 B/op, 6 allocs/op | 112 B/op, 4 allocs/op | -| `JSON_ROW`, fresh operator | about 838 us | 472.7 us | 1,704,081 B/op, 16,387 allocs/op | 520 B/op, 8 allocs/op | -| `JSON_ROW`, reused operator | about 474.7 us | 471.4 us | about 831 B/op, 8 allocs/op | 264 B/op, 5 allocs/op | - -The direct-output candidate passed each new or directly affected behavioral -test separately under race with a 30-second adaptive budget. Every measured -test used `-count=100` except the 0.35-second incompressible DEFLATE-bound test, -which used `-count=85`; no empty regex was accepted as stress evidence. -ByteJSON, Vector, metric, Function, and colexec passed complete normal and race -runs; the same package closure passed vet and `golangci-lint` with zero issues. -Coverage for that closure is 78.3% ByteJSON, 50.9% Vector, 90.3% metric, 54.8% -Function, and 64.3% colexec. The unsafe result alias has its own functional -test, and error/invalid-length/panic rollback proves result length, area -length, reuse, and final account zero. - -The GPU source path passed manual ownership review: exact SQL execution sizes -and fills caller-owned scratch before launch, a successful launch transfers a -retained slice reference to the job, Wait drops it only after the asynchronous -kernel terminates, and dimension/launch failure publishes no job owner. The -legacy nil-scratch API keeps its existing C-buffer allocate/rollback/Wait -contract. This host has neither `nvcc` nor a CUDA conda environment, so the -GPU-tag build and tests were not run locally. That validation gap remains an -explicit merge/activation gate; CPU success is not treated as GPU evidence. - -Fresh local validation on linux/amd64 with CGO and the repository-built -third-party artifacts passes: - -- complete normal tests for MPool, Vector, Batch, SQL util, expression, - pSpool, spillutil, HashBuild, HashJoin, DedupJoin, RightDedupJoin, - Connector, Dispatch, Function, and Process; -- build and vet for the same package closure; -- exact `-race -count=100` runs for every new test plus directly affected flow - control, BucketReader merge, scatter lifecycle, and re-spill tests; -- complete race runs for MPool, Vector, Batch, SQL util, expression, spillutil, - pSpool, Connector, and Dispatch; -- package coverage after the new closures of 74.0% MPool, 77.0% Bitmap, 41.2% - Nulls, 48.3% Vector, 73.4% Batch, 82.1% pSpool, 78.8% spillutil, and 54.1% - Function. - -For `a63c68b07b`, every new bitmap, buffer, Vector, decimal-conversion, spool, -spill, and directly affected aggregate test passed an exact -`-race -count=100` run after the final edit. Bitmap, MPool, Vector, pSpool, -spillutil, aggregate, expression, and Function packages then passed complete -`-race -count=1` runs; the full affected normal-test and vet closure also -passed with repository-built CGO third parties. - -The existing constant-flow-control benchmark remains 0 B/op and -0 allocs/op. Five-run medians at `GOMAXPROCS=8` were 4.480 ns/op on the PR 2 -base and 4.464 ns/op on the PR 3 candidate; this focused benchmark shows no -measurable legacy fast-path regression, but it is not an activation-level -performance result. - -The legacy 8,192-row pipeline-spool copy/reuse benchmark remains 200 B/op and -2 allocs/op. An interleaved five-run comparison at `GOMAXPROCS=8` measured a -1,542 ns/op median at commit `9f88a0e83e` and 1,540 ns/op at `1f82b218c2`; -the guarded allocation-unaccounted fast path has no measurable regression. - -At `a63c68b07b`, the same legacy pipeline-spool benchmark remains 200 B/op and -2 allocs/op with a 1,545 ns/op five-run median. The constant flow-control -benchmark remains 0 B/op and 0 allocs/op with a 4.259 ns/op median. The tagged -bitmap ownership representation is what preserves the legacy allocation -footprint. - -An 8,192-row first pre-extend/free costs 1,059 ns/op for the legacy Vector, -1,159 ns/op for data-only dormant accounting, and 1,986 ns/op when the two -independent bitmap allocations are also admitted; all three remain 0 B/op and -0 allocs/op. This one-time physical-allocation cost is explicit rather than -hidden. Capacity reuse is the steady-state path: empty Reset measured -2.431 ns/op for data-only accounting and 2.650 ns/op with bitmap ownership, -again with zero Go allocations. Activation performance tests must retain this -distinction instead of presenting the first-allocation cost as a per-row cost. - -### PR 4: statement lifecycle and minimum pressure foundation - -Implementation: `8633757dc1`. The attempt coordinator now owns board drain, -owner teardown, exactly-once terminal completion, immutable export, and the -release-capable tombstone/suspension path. Typed lifecycle failures are -disjoint from retryable capacity pressure. - -Scope: - -- add the attempt-owned post-pipeline/MessageBoard-close seal/finalize hook; -- export one immutable valid or invariant-failure generation snapshot; -- retain release-capable tombstones and suspend new accounted generations - after nonzero terminal cleanup; -- introduce non-overlapping capacity, sealed, mismatch, allocator-limit, and - invariant error reasons; -- define operation checkpoint/rollback helpers needed by later retry; -- keep owner accounting dormant. - -Gate: - -- success, execution error, cancellation, retry, local/remote scope, broadcast, - and message-board teardown each seal exactly once; -- JoinMap/spill payload released after producer Reset still releases the - original generation; -- zero finalization exports once and removes the registry entry; -- nonzero terminal cleanup exports one failure snapshot, retains a - release-capable tombstone, stops new accounted generations on that CN, and - removes the tombstone only after late Free reaches zero; -- concurrent terminal failures are bounded by generations active at first - detection, and stale IDs cannot resolve after removal; -- a race test proves every successfully published generation linearized before - suspension and every later open is rejected; -- closed generation never enters reclaim/spill/retry; -- no production legacy charge is removed. - -### PR 5: hash-table cell-block activation - -Implementation: `5ae8eca00a`. Integer and string cells plus their descriptor -backing are allocation-accounted through initial allocation, both resize -modes, producer/consumer handoff, and terminal Free. Activated maps do not -install the legacy resize reservation owner. - -Scope: - -- activate only the integer/string hash-table cell blocks; -- replace the Go `[][]Cell` outer backing store with an owning off-heap - descriptor buffer whose initial and replacement capacities use the same - account; -- attach the account to initial allocation, full replacement, segmented - appended blocks, and consumer-side empty-map growth; -- remove only the matching `hashMapReservationOwner`/`ResizeReservation` - budget charge; -- retain legacy copied-batch, expression, auxiliary, and spill charges. - -Gate: - -- int/string initial, no-op, full replacement, segmented reuse, stale plan, - injected allocation failure, and terminal Free are covered; -- old+new replacement and old+appended-block peaks match live cell allocation - plus descriptor allocation capacity; -- descriptor replacement rollback is atomic with cell-block publication, and - no data-scaled Go backing array remains; -- the first-activation metadata-slot formula holds at exact and one-slot-short - boundaries; -- consumer growth after HashBuild handoff keeps original provenance; -- generation final snapshot reaches zero; -- #25782 high-cardinality no-OOM regression and hash-table performance pass. - -### PR 6: copied batches and JoinMap activation - -Implementation: `8e4b689f45`. Copied build vectors, GroupSels, dedup scratch, -delete bitmaps, and final JoinMap ownership retain immutable provenance until -the last consumer releases them. - -Scope: - -- copied build-batch destinations; -- HashmapBuilder-to-JoinMap ownership handoff. - -Required behavior: - -- replace projected/reconciled batch tokens with physical leases; -- preserve provenance through producer Reset and broadcast consumers; -- let physical Free replace budget-only `SetMemoryRelease` callbacks. - -Gate: - -- empty/single/large build and both resize modes; -- failed publish, cancellation, multiple consumers, duplicate cleanup; -- old-generation allocation freed after a new generation opens; -- #25782 high-cardinality and #26413 external-table self-join regressions pass - before activation merges; -- copied-batch/JoinMap site rows become `A` or proved `H`. - -### PR 7: expression owner activation - -Implementation: `c9ad0ea810`. COL/LIT/PARAM/VAR/VEC/FOLD plus audited CONCAT, -CASE, varchar EQUAL, and integer-to-string/literal string CAST trees use exact -result, bitmap, conversion, and function-scratch allocations. A HashBuild -containing any other expression family is not activated at all; it cannot mix -an exact map owner with a legacy expression estimator. - -Scope: - -- activate the complete HashBuild-owned expression site closure; -- remove exactly the corresponding - `expressionVectorPeak`/`expressionTypePeak` hard charge; -- return controlled terminal pressure until a proved retry checkpoint supports - a smaller-batch retry; -- do not stack exact leases on #26455's operator-held lifetime charge. - -Gate: - -- every activated expression site is `A` or proved `H`; no data-scaled Go - allocation remains; -- one-byte-short and real single-value-over-cap diagnostics name actual - capacity; -- partial growth/evaluation failure publishes no duplicate rows; -- generation final snapshot reaches zero; -- #26454 workload and expression performance regressions pass. - -### PR 8 family: spill and runtime-filter closures - -Implementation: `e072568998`. Decoded reuse, selected vectors, hash/row IDs, -marshal/coalesce buffers, recursive rebuild, spill disk/FD, pSpool provenance, -and runtime-filter payload publication all have explicit allocation or -resource owners. Optional coalesce and runtime filters degrade without -publishing partial state. - -Split into independently safe owner closures: - -1. decoded batches and retained reader reuse; -2. scatter/hash/row-ID/codec/coalesce buffers and forward-progress policy; -3. runtime-filter payload transfer and PASS degradation. - -Each sub-PR removes only its matching hard reservation. Required gates include -first spill, recursive spill, skew, empty bucket, EOF, failure/cancel at every -I/O/publication edge, minimum progress over cap, message destruction, and final -memory/disk/FD zero. The scatter/codec closure must pass #26174 fulltext INSERT; -the decoded/retained-reader closure must pass #26192 LOAD DATA. A runtime-filter -closure adds its own build/probe and PASS-degradation workload before merging. - -### PR 9: unified join pressure controller and remaining legacy deletion - -Implementation: `eec23dcdc5`. HashBuild, SpillEngine, HashJoin, DedupJoin, and -RightDedupJoin share typed pressure classification and a monotonic retry guard. -Only memory-capacity failures retry; lifecycle, invariant, disk, and FD -failures are terminal. Unpublished input windows shrink, optional coalescing -is disabled once, and real minimum-unit failure is finite and diagnostic. - -Implement across HashBuild, HashJoin, DedupJoin, and RightDedupJoin: - -```text -exact capacity rejection - -> rollback to operation checkpoint - -> reclaim -> retry - -> spill/re-spill -> retry - -> reduce batch -> retry - -> degrade optional owner - -> controlled minimum-unit error -``` - -Gate: - -- retry only after usage decreases, spill advances, input shrinks, or optional - work is disabled; -- partially grown retained buffers and output publication are idempotent; -- no infinite pressure loop; -- sealed/lifecycle errors never retry; -- cancellation/downstream failure is covered in every state; -- remaining multiplier hard gates and duplicate token owners are deleted; -- every site row is `A`, justified `H`, `S`, or `R`. - -### PR 10: workload, performance, and cleanup - -Local implementation and benchmark harness: `383fc6dce3`. The harness covers -resident int/varchar keys at 32 and 8,192 rows, complete spill scatter, -#26454's expression, copy ownership, lifecycle latency, and concurrent release -storms. Raw commands and medians are in -[`evidence/26459_activation_validation.md`](evidence/26459_activation_validation.md). - -This re-runs all incident workloads together and supplies long-run and -comparative confirmation; it is not the first incident-level validation of an -earlier activation. - -Required workloads: - -- #25782 high-cardinality case; -- #26174 fulltext INSERT; -- #26192 LOAD DATA; -- #26413 Hive external-table self-join; -- #26454 string-expression join; -- TPCH 100G non-spill and TPCH 1T spill. - -Required proof: - -- no CN OOM/restart and no cap increase/spill-disable workaround; -- account never exceeds cap and every generation returns to zero; -- only a real minimum allocation can produce terminal pressure; -- diagnostics name owner/site and attempted response; -- resident and spill performance are compared separately with profiles; -- concurrent-generation P50/P99, release storm, and high-frequency TP - allocation results meet the recorded gate; -- temporary migration helpers are removed. - -Current validation status: - -- the complete affected package matrix passes after merging `26ed429b60`; -- the complete affected package matrix, the same matrix under `-race`, a - 20-iteration focused lifecycle/pressure race stress, affected-package vet, - `make static-check`, and `make build` pass through `75cd58dc7a`; -- accounted ordered/masked filtering, arbitrary shuffle rollback, `maxby` - retained-state compaction, generic/transaction ownership exits, mixed Batch - clone and BatchSet provenance, provenance-aware ShufflePool credits, CN - bootstrap, cross-package spill fixtures, and owner finalization-before-release - have dedicated regressions; -- local regression tests cover the five incident mechanisms without an - estimator-only rejection on the activated owner; -- TPCH 100G/1T candidate-versus-main runs remain the remote workload gate and - must be recorded in the linked evidence before this document moves from - `implementation validation` to `implemented`. - -## 5. Verification and conservation model - -Every semantic PR runs build, vet, complete package tests, focused adaptive race -stress, and dependent package tests for its closure. Likely packages are: - -```text -pkg/common/mpool -pkg/container/vector -pkg/container/batch -pkg/sql/colexec -pkg/sql/colexec/{hashbuild,hashjoin,dedupjoin,rightdedupjoin,spillutil} -pkg/vm/{message,process} -``` - -Before direct tests that can reach usearch, build `thirdparties` and use the -repository CGO include/library/rpath environment. Compilation or one successful -SQL run is not completion evidence. - -The completed PR 0 reference model must track allocation ID, account ID, -allocator mode, capacity, pool, logical owner, and -unpublished/live/freed/tombstone state. Its required generated operations -include allocation, within-capacity reuse, replacement growth, injected -failure, Free, cross-pool Free, view/copy, handoff, Reset, seal, finalization, -stale ID, operation checkpoint, cancellation, and retry. - -The current artifact generates allocation, growth, handoff, and Free, and -separately tests failure injection, cross-pool Free, seal, zero/nonzero -terminal paths, stale generation-tagged slots, bounded registry and allocation -metadata, teardown, open/suspend linearization, view/copy, Reset, -multi-allocation checkpoints, cancellation, rollback, and smaller retry -without duplicate publication. - -After every operation: - -```text -account.used - = sum(live accounted allocation capacities) - + sum(live named scratch) -``` - -Failed unpublished work leaves allocator/account state unchanged; no allocation -releases twice; sealed generations admit nothing new; final cleanup reaches -zero. Accounted on-heap allocation is rejected. A logical retry may retain -already published reusable capacity but cannot retain partially published -rows. Deterministic CI seeds print on failure; longer randomized runs may run -nightly. - -## 6. Resource Accounting integration - -Admission and SQL Resource Accounting share allocator facts, not mutable -control state. Each query-CN generation has a stable identity and its -`Compile` attempt coordinator exports exactly one immutable controlled-domain -snapshot: - -- cap, exact peak, and final live bytes; -- alloc/grow/free and real pressure counts; -- reclaim/spill/batch-reduction/degrade responses; -- invariant quality. - -Initial integration keeps `statement_info.stats[2]` unchanged and exposes the -controlled-domain snapshot in physical-plan/operator diagnostics. Resource -Accounting may mark a missing/inconsistent snapshot and aggregate terminal -facts, but its summaries never feed hard admission. Multiple operators may -reference one generation; their diagnostics must not be summed as separate -physical domains. The controlled-domain and MPool-domain peaks are not declared -equal until allocator mode and owner coverage match. - -No per-allocation log or Prometheus series is added. Owner/site values are -bounded enums and counters are emitted at generation/operator completion or -terminal pressure. - -## 7. Rollout, rollback, and completion - -Rollout: - -- rebase every PR on current `main`; -- keep generic primitives dormant until one complete owner closure migrates; -- enable exact accounting and remove the same legacy charge atomically; -- update the issue ledger only after evidence exists. - -Rollback reverts an owner's exact enablement, legacy removal, pressure response, -and tests together. Never roll back only the accounting or pressure half. If a -missing owner is found, complete the closure before rollout or revert that -owner; do not add another permanent estimator multiplier. - -Done means: - -- every site-ledger row is `A`, justified `H`, `S`, or `R`; -- every `H` row has a per-entry cost, maximum-live-count proof, and aggregate - CN headroom; -- allocation/growth failure is atomic; -- attempt-owned seal/finalize and exactly-once snapshot export are proven; -- nonzero terminal generations retain release provenance without unbounded - registry growth; -- Reset, Free, broadcast, and cross-generation handoff preserve provenance; -- data-scaled controlled Go allocations have moved off-heap; -- operation-level retry checkpoints prevent partial output replay; -- bounded forward progress is accounted; -- predictions cannot terminally reject SQL; -- no data-scaled HashBuild-owned allocation remains outside coverage; -- the original no-OOM case and all listed false-budget regressions pass; -- performance gates pass with evidence; -- Resource Accounting receives immutable diagnostics without entering the - admission decision; -- legacy hard estimators and duplicate token owners are deleted. +# Allocation-accounted HashBuild memory admission + +Issue: #26459 + +## Goal + +Within the controlled domain, HashBuild and its join consumers must admit the +bytes they physically retain, not a predicted multiple of logical payload size. +A query may spill or degrade an optional optimization when the retained +allocation cannot be admitted, but must not fail because an estimator guessed +a larger, non-existent allocation. + +The implementation has one production path: + +1. the MPool performs the physical allocation; +2. the allocation carries immutable account, owner, and site provenance; +3. the account charges the query generation and the CN aggregate controller; +4. the same allocation releases the charge when MPool frees it. + +There is no feature switch, activation gate, estimated-memory reservation, or +parallel compatibility ledger. Cardinality estimates may still select a hash +table capacity, but they never create a separately releasable memory charge. + +## Scope + +The account covers retained physical storage owned by the HashBuild execution +family: + +- hash table cells, descriptors, iterator keys, and selection lists; +- copied build batches and retained unique keys; +- JoinMap-owned batches and grouping metadata; +- join matched/capture/result state; +- Product result state; +- runtime-filter payloads until ownership transfer; +- spill encode/decode/scatter buffers and rebuilt retained state. + +ExpressionExecutor results, caches, and library-internal Go heap objects are +not HashBuild-retained storage. They remain in the existing MPool/Go runtime +domain. When an expression result is copied into a retained batch, key, or +result vector, that destination allocation is charged. This boundary avoids a +misleading partial "exact expression memory" account for regexp, JSON, JQ, and +other libraries that do not expose allocator/free hooks. + +The account is therefore not advertised as total query RSS. General transient +expression admission is a separate problem and must not be represented by an +estimated charge inside this exact, terminal-zero ledger. + +ProductL2 scratch and native CPU/GPU index storage are also outside this first +controlled domain. ProductL2 still consumes an accounted JoinMap: those source +allocations keep their original provenance through transfer and are released +by JoinMap `Free`. Its additional search index and scratch remain under the +existing implementation until their CPU and GPU allocators expose one common +physical capacity contract. Adding only the visible Go buffers would claim +false exactness and create a partial second path. + +## Ownership model + +Each physical allocation has one owner and one release path. Provenance is +attached before the first owned allocation and cannot change while storage is +live. Views borrow storage and do not create another charge. Copies allocate +new storage in their destination account. + +The main transfer boundaries are: + +| Storage | Initial owner | Transfer | Terminal release | +| --- | --- | --- | --- | +| retained build batch | HashmapBuilder | JoinMap | JoinMap `Free` | +| hash cells/descriptors | HashmapBuilder | JoinMap | JoinMap `Free` | +| grouping selections | HashmapBuilder | JoinMap | JoinMap `Free` | +| spill file + disk/FD tokens | HashBuild | SpillBuildPayload | SpillEngine/file close | +| matched bitmap | parallel worker | BitmapMailbox/merger | merger or mailbox drain | +| Product build batches | producer JoinMap | Product | Product reset/free | +| runtime-filter payload | HashBuild | message board | message destruction | + +Transfers are move-only. A successful send clears the sender's ownership; a +failed send leaves ownership with the sender. Cancellation seals mailboxes and +drains queued accounted objects before terminal validation. + +## Execution lifecycle + +Compile opens one allocation generation for each local statement attempt. +Every operator implementing `SetAllocationAccount` / `ClearAllocationAccount` +is an owner in that generation. + +The sequence is: + +1. collect owners from physical scope templates; +2. open the account with the live HashBuild capacity controller; +3. configure all owners atomically, rolling back in reverse order on failure; +4. attach parallel scan/load clones created during `runOnce` to the same + generation before worker `Prepare` runs; +5. execute and drain the message board; +6. clear owners in reverse order; +7. seal and finalize the account; +8. export exactly one terminal snapshot. + +Prepared statements and retries create a new generation. Reset frees all +generation-bound state; it does not carry an executor, bitmap, mailbox payload, +or allocation selection into the next attempt. Runtime parallel clones are +also registered as owners and are cleared before their reuse-pool release. + +A valid terminal snapshot requires zero live bytes and zero live allocation +metadata. A mismatch suspends new admission until late physical frees drain the +tombstone; it is never converted into retryable capacity pressure. + +## Capacity and pressure + +The controller enforces both the query-generation cap and the CN aggregate +cap. Admission uses checked arithmetic and charges the physical MPool capacity +requested by the allocator. + +Pressure reasons are typed and disjoint: + +- memory capacity: reclaim, spill, reduce an unpublished input unit, or degrade + an optional runtime filter; +- account sealed/suspended: terminal lifecycle error; +- owner/site mismatch or allocator invariant: terminal correctness error; +- spill disk cap: spill-resource error, never a memory-reduction retry; +- spill FD cap: spill-resource error, never a memory-reduction retry; +- minimum input unit: terminal capacity error after monotonic progress is no + longer possible. + +Retry is allowed only when progress is observable: retained bytes decrease, +spill epoch advances, input units shrink, or optional state is disabled. This +prevents a capacity loop from replaying the same publication or I/O. + +Runtime filters are optional. If their retained payload cannot be admitted, +HashBuild publishes PASS and releases unpublished scratch. Required hash/join +state does not silently bypass admission. + +## Spill resources + +SpillEngine requires the same live budget generation as its producer. There is +no nil-budget file path. + +Memory, disk bytes, and open file descriptors are separate physical resources: + +- memory is charged by MPool allocations; +- each spill file owns one growable disk token; +- each open spill file owns one FD token; +- a file handoff moves both tokens with the file; +- close releases all three exactly once. + +Recursive spill validates schema, framing, row conservation, queue bounds, and +file metadata. Repartitioning keeps only bounded control arrays plus admitted +scatter buffers. Test fixtures use the same builder copy and budget paths as +production. + +## Grouping semantics + +Grouping sentinels are a distinct key domain from ordinary zero/empty values. +HashBuild selects grouping-aware key encoding whenever input contains grouping +bits. Hashing preserves this distinction for partitioning, equality preserves +it in resident maps, copies preserve the bitmap, and ordering treats the +sentinel as SQL NULL for NULLS FIRST/LAST behavior. + +Sample keeps its ordinary batched hash path unchanged. If grouping bits appear +after ordinary groups were already inserted, it lazily opens a grouping-aware +key domain and translates both maps' local IDs into one sample-pool ID space. +Its iterators are reused per batch, so alternating grouping bits do not create +one iterator allocation per row. + +## Non-goals + +- estimating or limiting total query RSS; +- charging regexp/JQ/JSON/library-internal Go heap as if it were exact; +- changing optimizer join selection; +- making spill as fast as a sufficient-memory in-memory join; +- using remote benchmark workflows as a correctness oracle. + +## Completion criteria + +The implementation is complete only when: + +- production has no estimated HashBuild memory reservation or activation gate; +- every retained HashBuild/join allocation in the controlled domain has + immutable provenance; +- runtime parallel clones join the current attempt before `Prepare`; +- every transfer has exactly one owner after success and on cancellation; +- memory, disk, and FD rejection remain distinct; +- prepared/retry generations terminate independently at zero; +- local unit, race, build, vet, lifecycle, spill, and performance checks pass; +- independent reviews report no blocker or major correctness/performance issue. diff --git a/docs/design/evidence/26459_activation_validation.md b/docs/design/evidence/26459_activation_validation.md deleted file mode 100644 index 9f4095b835745..0000000000000 --- a/docs/design/evidence/26459_activation_validation.md +++ /dev/null @@ -1,150 +0,0 @@ -# #26459 Allocation-Accounted Activation Validation - -## Candidate - -- Latest main baseline: `26ed429b60` -- PR 4 lifecycle: `8633757dc1` -- PR 5 hash-table activation: `5ae8eca00a` -- PR 6 retained batch/JoinMap activation: `8e4b689f45` -- PR 7 expression activation: `c9ad0ea810` -- PR 8 spill/runtime-filter activation: `e072568998` -- PR 9 unified pressure recovery: `eec23dcdc5` -- PR 10 benchmark harness: `383fc6dce3` -- Owner-atomic activation and cleanup: `656e254fe6` -- Vector/bitmap ownership boundary closure: `73d20085c0` -- Cross-package spill fixture compatibility: `3f10c23f21` -- Owner finalization-before-release: `91839ea91f` -- Transaction PK ownership exit: `abb1637882` -- Mixed Batch provenance closure: `75cd58dc7a` - -The activated expression owner is deliberately closed: COL, literal, param, -variable, vector/fold, CONCAT, CASE, varchar EQUAL, and the audited string CAST -forms. Build and probe keys are checked together. If any HashBuild, HashJoin, -DedupJoin, or RightDedupJoin key contains another function family, automatic -activation keeps every participating operator in that local statement attempt -on the legacy path. This prevents a partially exact owner from retaining an -estimator-only expression rejection. It does not claim that generic JSON, -regexp, geometry, or spatial execution has already migrated. - -## Local correctness matrix - -All direct Go tests use the repository CGO wrapper so `usearch` links against -the locally built `thirdparties` artifacts. - -```text -.agents/skills/mo-dev/scripts/mo-cgo-test -count=1 \ - ./pkg/common/mpool ./pkg/common/hashmap ./pkg/common/bitmap \ - ./pkg/container/hashtable ./pkg/container/vector ./pkg/container/batch \ - ./pkg/container/nulls ./pkg/container/bytejson ./pkg/sql/util \ - ./pkg/sql/colexec ./pkg/sql/colexec/aggexec ./pkg/sql/colexec/hashbuild \ - ./pkg/sql/colexec/hashjoin ./pkg/sql/colexec/dedupjoin \ - ./pkg/sql/colexec/rightdedupjoin ./pkg/sql/colexec/spillutil \ - ./pkg/sql/colexec/shuffle ./pkg/sql/colexec/multi_update \ - ./pkg/sql/compile ./pkg/util/resource ./pkg/vm/message ./pkg/vm/process \ - ./pkg/vm/engine/disttae - -result: PASS -``` - -The same package matrix passes with `-race -p=2 -count=1`. `make static-check` -and `make build` also pass on the latest-main candidate. A focused -`-race -p=2 -count=20` stress also passes for concurrent account alloc/free, -seal/open linearization, statement terminal one-shot behavior, owner-atomic -activation, accounted spill/reduction, and all three join-consumer activation -gates. Affected-package `go vet` and `make build` pass on the same candidate. - -An isolated `etc/launch/launch.toml` bootstrap using the freshly built binary -also passes on `26ed429b60` plus `75cd58dc7a`: the CN becomes queryable, system -catalog initialization completes, and the service log contains no allocation -account failure, terminal cleanup failure, or panic. - -The previous PR CI failures were reduced to concrete branch-local ownership -boundaries and fixed before rerunning the matrix. Generic `CopyBatch` and -transaction PK validation now make explicit exits from a borrowed statement -owner. Mixed Batch clones preserve per-Vector provenance; BatchSet preserves -uniform destination context, starts a new tail when provenance changes, and -ShufflePool reserves the corresponding ready credit. MultiUpdate's short-lived -validation/filter copies use the explicit generic boundary. The spill fixture -writer remains available to cross-package hashjoin/dedup tests, and -`Compile.clear` finalizes and detaches the statement allocation owner before -releasing operators back to reuse pools. Dedicated tests cover these paths, -accounted bitmap remap/shuffle rollback, `maxby` selection preservation, CN -bootstrap, fixture compilation, and release after finalization. - -Incident-mechanism regression mapping: - -| Incident | Local regression proof | -| --- | --- | -| #25782 high-cardinality HashBuild/spill | `TestAccountedInitialSpillReducesUnpublishedInputAndPreservesRows`, `TestShuffleHashBuildAccountedSpillLifecycle`, hash cell/descriptor replacement tests | -| #26174 dedup/fulltext INSERT | `TestAccountedDedupScratchAndDeleteBitmapFollowJoinMapLifetime`, `TestAccountedDedupBitmapExactBoundaryRollsBack` | -| #26192 LOAD DATA decoded/rebuild path | `TestSpillAllocationAccountDecodedBatchLifecycle`, `TestSpillAllocationAccountDecodedReuseRetriesFromCleanRecord`, recursive rebuild lifecycle test | -| #26413 external/self-join lifetime | statement-attempt zero/late-Free/cancel/error/panic tests, `TestAccountedJoinMapLateFreeKeepsOriginalGeneration` | -| #26454 string expression false budget | `TestAllocationAccountedExpressionIssue26454AndOneByteShort`, `TestIssue26454ExpressionKeyBuildUsesActualCapacity`, adaptive expression-pressure tests | - -The activation boundary itself is covered by -`TestHashBuildAllocationActivationRequiresClosedExpressionOwner` and -`TestAllocationAccountActivationIsStatementAtomic`: a closed #26454 -expression activates, an unclosed modulo expression keeps the owner legacy, a -no-map operator does not open a generation, and one unclosed owner rolls back -all already configured owners. - -## Local performance evidence - -Host: linux/amd64, Intel i7-11700, Go 1.26.4. Unless stated otherwise, -GOMAXPROCS was 16. Medians are from five runs; resident tests used a 200 ms -benchtime and spill scatter used 300 ms. - -### Resident HashBuild - -| Key/rows | Legacy median | Accounted median | Delta | Allocation effect | -| --- | ---: | ---: | ---: | --- | -| int / 32 | 4,161 ns | 5,841 ns | +40.4% (+1.68 us) | 32 -> 27 allocs/op | -| varchar / 32 | 7,147 ns | 8,907 ns | +24.6% (+1.76 us) | 96 -> 91 allocs/op | -| int / 8,192 | 139,602 ns | 143,254 ns | +2.62% | 43,738 -> 10,401 B/op; 56 -> 29 allocs/op | -| varchar / 8,192 | 410,074 ns | 400,782 ns | -2.27% | 222,055 -> 24,899 B/op; 600 -> 574 allocs/op | - -The 32-row cases quantify the fixed first-allocation tax for high-frequency TP -work; they are not presented as a per-row tax. Full-batch resident HashBuild is -within 2.7% for int keys and improves the measured varchar case. - -### Spill and expression paths - -| Benchmark | Legacy median | Accounted median | Delta | -| --- | ---: | ---: | ---: | -| 4,096-row scatter including hash/select/marshal/coalesce/write | 59,073 ns | 60,571 ns | +2.54% | -| #26454 expression | 909,082 ns | 940,454 ns | +3.45% | -| copied build batch | 18,584 ns | 22,051 ns | +18.7% | - -The copied-batch result also reduced 230,032 B/op to 632 B/op and 10 to 6 -allocations/op. Scatter reduced 2,711 B/op to 2,392 B/op; syscall and codec work -remain included. - -### Generation and release concurrency - -| Benchmark | CPU | Median | Recorded latency | -| --- | ---: | ---: | --- | -| full allocation-attempt lifecycle | 1 | 843.5 ns/op | p50 730 ns; p99 2,011 ns | -| full allocation-attempt lifecycle | 8 | 762.7 ns/op | p50 1,995 ns; p99 73,349 ns | -| same-generation release storm | 1 | 558.4 ns/op | 0 B/op, 0 allocs/op | -| same-generation release storm | 8 | 470.0 ns/op | 0 B/op, 0 allocs/op | - -The lifecycle benchmark includes generation open, controller/account creation, -one exact 4 KiB allocation/free, terminal completion, and concurrent operation. - -## Remote workload gate - -The candidate must be compared with the exact rebased main on the same TKE -workflow and resource configuration. Required results are: - -| Workload | Required evidence | Status | -| --- | --- | --- | -| TPCH 100G | three rounds, no spill/OOM/restart, total and per-query comparison | pending | -| TPCH 1T | three rounds, spill succeeds, no OOM/restart, total and per-query comparison | pending | -| #26174 fulltext INSERT | workload pass under unchanged cap | pending | -| #26192 LOAD DATA | workload pass under unchanged cap | pending | -| #26413 external self-join | workload pass and zero terminal generation | pending | - -No cap increase, reduced dataset, disabled spill, or plan-specific bypass is an -acceptable pass. The candidate and main run URLs, exact SHAs, load times, -query totals, spill evidence, restart/OOM search, and profile comparison belong -in this table before PR 10 is accepted. diff --git a/docs/design/evidence/26459_allocation_accounting_bench.txt b/docs/design/evidence/26459_allocation_accounting_bench.txt index 522caa2537fcf..cf1d5b7b50707 100644 --- a/docs/design/evidence/26459_allocation_accounting_bench.txt +++ b/docs/design/evidence/26459_allocation_accounting_bench.txt @@ -1,184 +1,141 @@ -Artifact: +Artifact +======== -branch: experiment/26459-allocation-accounting-validation -commit: cde44cd099 (rebased validation artifact) -go: go version go1.26.4 linux/amd64 +Date: 2026-08-01 +Branch: feature/26459-statement-lifecycle +Base: origin/main ecc389d420 +Go: go1.26.4 linux/amd64 +CPU: 11th Gen Intel Core i7-11700, GOMAXPROCS=16 -Setup: +All commands used the MatrixOne CGo wrapper: -make thirdparties -make cgo -export CGO_ENABLED=1 + .agents/skills/mo-dev/scripts/mo-cgo-test -Metadata command: +No remote workflow or container build was used. Values below are medians of +the displayed count unless noted otherwise. -go test -v ./pkg/common/mpool -run '^TestAllocationMetadataRepresentationCosts$' -count=1 -Metadata results: +Physical MPool admission +======================== -baseline: 13985992 bytes total, 55.94 bytes/base-entry -inline-all: 20992896 bytes total, 83.97 bytes/base-entry -side-1pct: 14059896 bytes total, 56.24 bytes/base-entry -side-10pct: 14571520 bytes total, 58.29 bytes/base-entry -side-all: 23438864 bytes total, 93.76 bytes/base-entry -fixed registry: 10010736 bytes total, 40.04 bytes/slot +Command: -Benchmark command: + mo-cgo-test ./pkg/common/mpool -run '^$' \ + -bench '^BenchmarkMPoolAccountedAllocation$' \ + -benchmem -benchtime=1s -count=5 -GOMAXPROCS=8 CGO_ENABLED=1 .agents/skills/mo-dev/scripts/mo-cgo-test -run '^$' -bench='BenchmarkAllocation(MetadataInsertFree|ConcurrentRegistryResolve|FixedRegistryResolve|AccountAcquireRelease)$' -benchmem -benchtime=500ms -count=5 ./pkg/common/mpool +Operation Unaccounted Accounted Delta +alloc/free 64 B 248.6 ns 357.5 ns +43.8% +alloc/free 4 KiB 294.6 ns 410.5 ns +39.3% +alloc/free 16 KiB 341.4 ns 449.3 ns +31.6% +alloc/free 64 KiB 1014 ns 1121 ns +10.6% +grow replacement 1292 ns 1492 ns +15.5% +parallel accounted alloc/free 64K 285.4 ns -Environment: +Every sample reported 0 B/op and 0 allocs/op. The fixed transaction cost is +largest for tiny allocations; retained HashBuild cells and batch buffers use +larger capacity changes rather than one admission per row. -goos: linux -goarch: amd64 -cpu: 11th Gen Intel(R) Core(TM) i7-11700 @ 2.50GHz -GOMAXPROCS: 8 -CPU pinning: none -Results: +Vector growth and reuse +======================= -BenchmarkAllocationMetadataInsertFree/baseline-8 40.13 39.99 40.40 40.37 39.97 ns/op -BenchmarkAllocationMetadataInsertFree/inline-all-8 40.67 40.36 40.21 40.00 40.17 ns/op -BenchmarkAllocationMetadataInsertFree/side-1pct-8 40.17 40.16 40.42 40.17 40.43 ns/op -BenchmarkAllocationMetadataInsertFree/side-10pct-8 40.81 40.97 40.77 40.79 41.19 ns/op -BenchmarkAllocationMetadataInsertFree/side-all-8 49.89 48.39 48.84 49.39 49.42 ns/op -BenchmarkAllocationConcurrentRegistryResolve-8 5.682 5.564 5.554 5.287 5.225 ns/op -BenchmarkAllocationFixedRegistryResolve-8 0.4278 0.4461 0.4463 0.2876 0.3265 ns/op -BenchmarkAllocationAccountAcquireRelease/locked-8 73.71 81.29 72.94 88.61 80.73 ns/op -BenchmarkAllocationAccountAcquireRelease/atomic-8 21.68 21.31 20.77 20.88 21.11 ns/op +Command: -Every sample reported 0 B/op and 0 allocs/op. + mo-cgo-test ./pkg/container/vector -run '^$' \ + -bench '^BenchmarkVectorAllocationAccount$' \ + -benchmem -benchtime=1s -count=3 -Real MPool and Vector baseline command: +Operation Unaccounted Accounted Allocation result +fixed preextend/free 1047 ns 1158 ns 0 B/op, 0 allocs/op +varlen preextend/free 77963 ns 78360 ns 48 B/op, 2 allocs/op +accounted fixed Reset reuse 2.63 ns 0 B/op, 0 allocs/op -GOMAXPROCS=8 CGO_ENABLED=1 .agents/skills/mo-dev/scripts/mo-cgo-test -run '^$' -bench '^BenchmarkAllocationAccounting(MPool|Vector)Baseline$' -benchmem -benchtime=200ms -count=5 ./pkg/common/mpool ./pkg/container/vector +Accounting adds no Go object to vector allocation or Reset reuse. The varlen +difference was 0.5%, inside local run noise. -Real baseline results: -BenchmarkAllocationAccountingMPoolBaseline/alloc-free/sharded/64-8 249.6 244.3 248.0 242.4 245.6 ns/op -BenchmarkAllocationAccountingMPoolBaseline/alloc-free/sharded/4096-8 291.8 291.0 292.5 290.4 295.6 ns/op -BenchmarkAllocationAccountingMPoolBaseline/alloc-free/sharded/65536-8 999.2 1003 991.2 1002 993.9 ns/op -BenchmarkAllocationAccountingMPoolBaseline/grow-replacement/sharded-8 1244 1267 1279 1254 1277 ns/op -BenchmarkAllocationAccountingMPoolBaseline/alloc-free/no-lock/64-8 222.7 226.0 223.2 222.5 222.6 ns/op -BenchmarkAllocationAccountingMPoolBaseline/alloc-free/no-lock/4096-8 280.5 273.3 276.5 271.6 276.5 ns/op -BenchmarkAllocationAccountingMPoolBaseline/alloc-free/no-lock/65536-8 984.2 979.6 979.9 994.1 980.9 ns/op -BenchmarkAllocationAccountingMPoolBaseline/grow-replacement/no-lock-8 1219 1199 1198 1248 1209 ns/op -BenchmarkAllocationAccountingMPoolBaseline/parallel-alloc-free/sharded/64-8 176.6 213.2 211.1 180.5 209.4 ns/op -BenchmarkAllocationAccountingVectorBaseline/fixed-preextend-free-8 1089 1028 1042 1043 1037 ns/op -BenchmarkAllocationAccountingVectorBaseline/varlen-preextend-free-8 28130 28011 27634 28170 27778 ns/op -BenchmarkAllocationAccountingVectorBaseline/fixed-reset-reuse-8 2.810 2.691 2.852 2.701 2.835 ns/op +Retained HashBuild closure +========================== -Every MPool sample reported 0 B/op and 0 allocs/op. Fixed Vector -pre-extend/free and Reset/reuse also reported 0 B/op and 0 allocs/op. Varlen -Vector pre-extend/free reported 48 B/op and 2 allocs/op. - -Account-aware real-MPool wrapper command: - -GOMAXPROCS=8 CGO_ENABLED=1 .agents/skills/mo-dev/scripts/mo-cgo-test -run '^$' -bench '^BenchmarkAllocationAccountedMPoolPrototype$' -benchmem -benchtime=200ms -count=5 ./pkg/common/mpool - -Account-aware wrapper results: - -BenchmarkAllocationAccountedMPoolPrototype/alloc-free/64-8 329.7 328.4 323.7 327.0 321.4 ns/op -BenchmarkAllocationAccountedMPoolPrototype/alloc-free/4096-8 367.6 361.9 381.9 366.5 368.4 ns/op -BenchmarkAllocationAccountedMPoolPrototype/alloc-free/65536-8 1106 1072 1085 1080 1091 ns/op -BenchmarkAllocationAccountedMPoolPrototype/grow-replacement-8 1418 1404 1390 1400 1393 ns/op -BenchmarkAllocationAccountedMPoolPrototype/parallel-alloc-free/64-8 203.1 227.2 205.6 233.5 230.2 ns/op - -Every sample reported 0 B/op and 0 allocs/op. This prototype takes a second -side-metadata lock and is intentionally an upper-bound comparison, not the -accepted production transaction shape. - -Final validation commands: - -go test -count=1 ./pkg/common/mpool ./pkg/container/vector -go vet ./pkg/common/mpool ./pkg/container/vector -go test -race -count=100 -run '^(TestAllocationAccountedMPoolPrototype|TestAllocationAccountingReferenceModel|TestAllocationAccountingViewCopyReset|TestAllocationAccountingMetadataSlotBounds|TestAllocationAccountingOperationRollbackRetry|TestAllocationAccountingTerminalTombstone|TestAllocationAccountingOpenSuspendLinearization|TestAllocationAccountingPoolTeardown)$' ./pkg/common/mpool -go test -race -count=23 -run '^TestAllocationAccountingReferenceModelRandomized$' ./pkg/common/mpool -go test -race -count=1 ./pkg/common/mpool ./pkg/container/vector - -Production PR 1 representation measurement (250,000 entries): - -branch: feature/26459-allocation-account -generic MPool commit: 766e1501c3 -HashBuild adapter commit: 0655af4443 - -base pointer map: 13,985,992 bytes total, 55.94 bytes/slot -base + 16-byte lease map: 27,966,424 bytes total, 111.87 bytes/slot -fixed registry backing: 4,022,272 bytes total, 16.09 bytes/slot -registry + live 64-byte accounts: 20,022,400 bytes total, 80.09 bytes/slot - -The conservative constants remain 128 bytes/allocation slot and are raised to -128 bytes/generation slot. The first-activation sizing test therefore checks -that a 1 GiB MPool cap yields 196,608 allocation-metadata slots and 40 MiB -total conservative metadata headroom, including 131,072 generation slots. - -Production PR 1 benchmark command: - -GOMAXPROCS=8 go test -run '^$' -bench '^BenchmarkMPoolAccountedAllocation$' \ - -benchmem -benchtime=500ms -count=5 ./pkg/common/mpool - -Clean-main medians from the paired baseline worktree: - -alloc/free 64: 244.6 ns/op -alloc/free 4096: 297.0 ns/op -alloc/free 16384: 333.4 ns/op -alloc/free 65536: 1023 ns/op -grow replacement: 1274 ns/op - -Production PR 1 unaccounted medians: - -alloc/free 64: 252.7 ns/op (+3.3%) -alloc/free 4096: 303.2 ns/op (+2.1%) -alloc/free 16384: 339.2 ns/op (+1.7%) -alloc/free 65536: 1034 ns/op (+1.1%) -grow replacement: 1306 ns/op (+2.5%) - -Production PR 1 accounted medians: - -alloc/free 64: 352.4 ns/op (+44.1% versus clean main) -alloc/free 4096: 400.4 ns/op (+34.8% versus clean main) -alloc/free 16384: 435.1 ns/op (+30.5% versus clean main) -alloc/free 65536: 1115 ns/op (+9.0% versus clean main) -grow replacement: 1492 ns/op (+17.1% versus clean main) -parallel accounted alloc/free 65536: 300.9 ns/op - -The dormant HashBuild controller adapter, including query/CN policy locking -and metrics, measured 1260 ns/op for 65536-byte alloc/free. Every production -PR 1 sample reported 0 B/op and 0 allocs/op. The first activated hash cell -allocation is at least 16 KiB; activation PRs still require workload-level -resident/spill validation because the small-allocation microbenchmark is not a -per-row path. - -An eight-generation, GOMAXPROCS=8 acquire/release latency harness sampled -200,000 operations per run. Across five runs P50 was 48 ns in every run; P99 -was 100, 98, 97, 68, and 73 ns (97 ns median). The measurement includes the -`time.Now`/`time.Since` sampling cost and is therefore a conservative -end-to-end latency observation, not a cycle-level atomic benchmark. - -Production activation benchmark (rebased main 5b9eeb54ec, linux/amd64, -Go 1.26.4, i7-11700, GOMAXPROCS=16; sequential package runs): - -BenchmarkResidentHashBuildAccounting/legacy/int/rows-32 median: 4161 ns/op -BenchmarkResidentHashBuildAccounting/accounted/int/rows-32 median: 5841 ns/op -BenchmarkResidentHashBuildAccounting/legacy/varchar/rows-32 median: 7147 ns/op -BenchmarkResidentHashBuildAccounting/accounted/varchar/rows-32 median: 8907 ns/op -BenchmarkResidentHashBuildAccounting/legacy/int/rows-8192 median: 139602 ns/op -BenchmarkResidentHashBuildAccounting/accounted/int/rows-8192 median: 143254 ns/op -BenchmarkResidentHashBuildAccounting/legacy/varchar/rows-8192 median: 410074 ns/op -BenchmarkResidentHashBuildAccounting/accounted/varchar/rows-8192 median: 400782 ns/op -BenchmarkIssue26454ExpressionAccounting/legacy median: 909082 ns/op -BenchmarkIssue26454ExpressionAccounting/accounted median: 940454 ns/op -BenchmarkCopyBuildBatchAccounting/legacy median: 18584 ns/op, 230032 B/op, 10 allocs/op -BenchmarkCopyBuildBatchAccounting/accounted median: 22051 ns/op, 632 B/op, 6 allocs/op -BenchmarkSpillScatterAccounting/legacy median: 59073 ns/op -BenchmarkSpillScatterAccounting/accounted median: 60571 ns/op - -BenchmarkHashBuildAllocationAttemptLifecycle cpu=1 median: 843.5 ns/op, -p50 730 ns/op, p99 2011 ns/op, 256 B/op, 3 allocs/op -BenchmarkHashBuildAllocationAttemptLifecycle cpu=8 median: 762.7 ns/op, -p50 1995 ns/op, p99 73349 ns/op, 256 B/op, 3 allocs/op -BenchmarkHashBuildAllocationReleaseStorm cpu=1 median: 558.4 ns/op, 0 B/op, 0 allocs/op -BenchmarkHashBuildAllocationReleaseStorm cpu=8 median: 470.0 ns/op, 0 B/op, 0 allocs/op - -Commands and workload gates are recorded in -docs/design/evidence/26459_activation_validation.md. +Command: + + mo-cgo-test ./pkg/sql/colexec/hashbuild -run '^$' \ + -bench '^(BenchmarkResidentHashBuildAccounting|BenchmarkCopyBuildBatchAccounting)$' \ + -benchmem -benchtime=1s -count=3 + +Benchmark Median B/op allocs/op +copy retained build batch 18919 ns 760 7 +resident int key, 32 rows 4428 ns 2320 24 +resident varchar key, 32 rows 7919 ns 20136 23 +resident int key, 8192 rows 137526 ns 10432 26 +resident varchar key, 8192 rows 411112 ns 20233 27 + +The retained data itself is MPool-backed; B/op here is bounded Go control +state. Allocations do not scale per input row. + + +Pipeline spool ownership transfer +================================= + +Command: + + mo-cgo-test ./pkg/container/pSpool -run '^$' \ + -bench '^BenchmarkCachedBatchReuse$' \ + -benchmem -benchtime=2s -count=5 + +Median: 1584 ns/op, 224 B/op, 3 allocs/op. + +The same DetachedBuffer path handles both storage outside the controlled +HashBuild domain and account-provenance storage. There is no second raw-byte +cache implementation. + + +Grouping-aware Sample correctness path +====================================== + +Command: + + mo-cgo-test ./pkg/sql/colexec/sample -run '^$' \ + -bench '^(BenchmarkSampleGroupedHashFastPath|BenchmarkSampleAlternatingGrouping)$' \ + -benchmem -benchtime=1s -count=5 + +Ordinary 256-row grouped path: 6075 ns/op, 9344 B/op, 6 allocs/op. +Alternating ordinary/GROUPING rows: 26869 ns/op, 19256 B/op, 12 allocs/op. + +A paired temporary reference containing the pre-fix ordinary hash loop was +measured in the same binary and then removed: reference median 5989 ns/op, +final median 6064 ns/op (+1.25%); both were 9344 B/op and 6 allocs/op. The +committed alternating-domain test also applies AllocsPerRun and proves the +iterator count remains constant per batch rather than per grouping run. + + +Spill scatter and serialization +=============================== + +Command: + + mo-cgo-test ./pkg/sql/colexec/spillutil -run '^$' \ + -bench '^BenchmarkSpillScatterAccounting$' \ + -benchmem -benchtime=1s -count=5 + +Final median: 62490 ns/op, 344 B/op, 4 allocs/op, about 524 MB/s. + +An allocation profile found that primitive marshal fallbacks took the address +of integer parameters even when AccountedBuffer's typed writer fast path was +selected. Before the fix the median was 67751 ns/op, 2904 B/op, and 516 +allocs/op. Encoding fallback integers into branch-local fixed arrays reduced +time by 7.8% and removed 512 Go allocations per scatter operation without +changing the wire format. + + +Acceptance +========== + +- physical admission/release adds zero Go allocations; +- vector and spool reuse retain bounded allocation counts; +- ordinary Sample grouping has unchanged allocation shape and only 1.25% + measured CPU delta; +- mixed GROUPING input uses a fixed number of iterators per batch; +- spill serialization has four bounded Go allocations per 4096-row scatter; +- no benchmark indicates a per-row accounting object or compatibility ledger. diff --git a/docs/design/evidence/26459_local_validation.md b/docs/design/evidence/26459_local_validation.md new file mode 100644 index 0000000000000..673d75aa60b5a --- /dev/null +++ b/docs/design/evidence/26459_local_validation.md @@ -0,0 +1,95 @@ +# #26459 local validation evidence + +This file records evidence for the single production path described in +`../allocation_accounted_memory_admission_impl.md`. Removed implementations +are not retained as validation dimensions. + +## Static closure checks + +- no production allocation-account enable switch; +- no HashBuild logical-size memory reservation token; +- no join/HashBuild expression account pretending to cover library Go heap; +- SpillEngine construction rejects a missing or closed budget generation; +- runtime scan/load clones are attached to the current attempt before worker + `Prepare`; +- terminal lifecycle validates zero bytes and zero live metadata; +- memory, spill disk, and spill FD admission errors have distinct components. + +## Required local test matrix + +All Go tests use the repository CGo wrapper so `usearch` headers, libraries, +link flags, and runtime paths match the MatrixOne build contract. + +```text +.agents/skills/mo-dev/scripts/mo-cgo-test -count=1 -timeout=240s \ + ./pkg/common/mpool ./pkg/common/bitmap ./pkg/common/hashmap/... \ + ./pkg/container/vector ./pkg/container/batch ./pkg/vm/message \ + ./pkg/vm/process ./pkg/sql/colexec/hashbuild \ + ./pkg/sql/colexec/hashjoin ./pkg/sql/colexec/dedupjoin \ + ./pkg/sql/colexec/rightdedupjoin ./pkg/sql/colexec/loopjoin \ + ./pkg/sql/colexec/product ./pkg/sql/colexec/productl2 \ + ./pkg/sql/colexec/spillutil ./pkg/sql/compile +``` + +Selected race coverage: + +```text +.agents/skills/mo-dev/scripts/mo-cgo-test -race -count=1 -timeout=300s \ + ./pkg/common/mpool ./pkg/sql/colexec/hashjoin \ + ./pkg/container/pSpool ./pkg/sql/colexec/productl2 \ + ./pkg/sql/colexec/spillutil ./pkg/sql/colexec/sample ./pkg/sql/compile +``` + +Static checks: + +```text +go vet +go build +``` + +## Fresh local result + +The final semantic edit was followed by a clean local run on 2026-08-01. + +- the complete package matrix above passed, with Sample, pSpool, compare, + hashtable, nulls, shuffle, SQL util, vector-index CPU packages, and disttae + added to the command; +- `-race -p=2 -count=1` passed for mpool, hashjoin, spillutil, Sample, and + compile; +- `go vet -mod=readonly` passed for every modified production package; +- `go build -mod=readonly` passed for every modified production package; +- `git diff --check` and `gofmt` were clean; +- allocation and performance results are recorded in + `26459_allocation_accounting_bench.txt`. + +No partial or still-running session is counted as a pass. + +## Behavioral coverage + +The local suite covers: + +- exact physical allocation/release and capacity rollback; +- owner/site mismatch and sealed-generation terminal errors; +- prepared and retry generation reuse; +- runtime parallel clone attachment; +- JoinMap and spill payload move-only ownership; +- cancellation-safe bitmap mailbox seal/drain; +- spill disk/FD admission and release; +- recursive spill row/schema/file validation; +- minimum-unit and monotonic pressure termination; +- optional runtime-filter degradation; +- grouping-aware copy, hash, equality, and ordering; +- late and alternating Sample grouping domains across row, percent, and merge + modes; +- Product cleanup and account terminal zero; +- accounted JoinMap release after an unaccounted ProductL2 consumer frees it. + +## Performance evidence + +Performance validation is local and allocation-focused. Benchmarks record +`ns/op`, `B/op`, and `allocs/op` for account acquire/release, vector growth, +hash-map build/lookup, and spill scatter. The acceptance rule is no +new per-row or per-allocation Go object in steady state and no material +regression outside measurement noise. + +Remote auto-test is deliberately not part of this validation cycle. diff --git a/docs/rfcs/00000000_allocation_accounted_memory_admission.md b/docs/rfcs/00000000_allocation_accounted_memory_admission.md index bc0afeb50aa91..2d5c6b341e719 100644 --- a/docs/rfcs/00000000_allocation_accounted_memory_admission.md +++ b/docs/rfcs/00000000_allocation_accounted_memory_admission.md @@ -1,958 +1,252 @@ - Status: implementation validation - Start Date: 2026-07-30 - Authors: aptend -- Implementation candidate: `feature/26459-statement-lifecycle` at - `656e254fe6` -- Issue for this RFC: - [#26459](https://github.com/matrixorigin/matrixone/issues/26459) +- Issue: [#26459](https://github.com/matrixorigin/matrixone/issues/26459) - Implementation plan: - [Allocation-Accounted Memory Admission Implementation Plan](../design/allocation_accounted_memory_admission_impl.md) + [Allocation-accounted HashBuild memory admission](../design/allocation_accounted_memory_admission_impl.md) -# Allocation-Accounted Memory Admission for Spillable SQL Execution +# Allocation-accounted memory admission for spillable SQL execution ## Summary -MatrixOne currently protects HashBuild with finite query and CN budgets and -bounded spill. That protection is necessary: before it existed, a large join -could continue allocating until the CN was killed by OOM. The remaining -problem is that several hard admission decisions are based on predictions of -an operation's future memory rather than the capacity of memory that is -actually allocated and retained. +HashBuild needs a finite query/CN memory boundary, but a predicted future size +is not a reliable hard-admission fact. SQL maximum widths, payload multipliers, +recursive expression sums, and logical batch sizes can reject valid work by +orders of magnitude. Relaxing those estimates can instead miss a real +allocate-copy-free overlap and allow OOM. -Predictions such as SQL type maximums, payload multipliers, recursive -expression-tree peaks, or logical batch sizes are useful for deciding to spill -early. They are not reliable enough to decide that a valid statement cannot -run. The same prediction can over-count aliases and retained buffers, or -under-count an allocate-copy-free overlap. - -This RFC replaces estimator-driven hard rejection with allocation-accounted -ownership: +This RFC makes the capacity of live physical allocations the source of truth: ```text -hard admission = capacity of real live allocations - + named, bounded non-allocator scratch - -prediction = scheduling hint only +hard HashBuild admission = live retained physical allocation capacity +prediction = planning and early-spill hint only ``` -An account-aware allocation reserves the exact physical capacity before it is -allocated. The resulting charge follows the allocation across reuse and owner -handoff, and is released by the same physical `Free`. Growth accounts for the -real replacement overlap by keeping the old allocation charged while -admitting the complete replacement capacity. - -When a real allocation cannot be admitted, execution treats it as typed memory -pressure: reclaim retained state, spill, retry, reduce the processing batch -where possible, degrade optional structures, and only then return a controlled -error for a minimum indivisible allocation that cannot fit. - -The first consumer is HashBuild and the joins that share its spill lifecycle. -The accounting primitive is deliberately defined below the SQL operator layer -so other spillable operators can adopt the same model later. - -The current activation is deliberately owner-closed. Build and probe trees -from the audited COL/literal/CONCAT/CASE/varchar-EQUAL/string-CAST set activate -exact accounting for the whole local HashBuild/join owner set. A tree -containing an unclosed generic function family keeps all participating owners -in that local attempt on the legacy path; exact map/batch ownership is never -mixed with an estimator-gated expression. The remaining generic-function -migration is therefore explicit rather than being misreported as RFC -completion. +MPool admits an accounted allocation before allocating it. Immutable account, +owner, and site provenance follows the allocation across reuse and ownership +handoff. The same physical Free releases the charge. Growth admits the actual +replacement capacity while the old allocation is still charged, so overlap is +represented without a multiplier. -## Motivation +The first consumer is retained HashBuild and join storage. The design has one +production path: no activation switch, estimated-memory reservation, or +compatibility ledger remains beside physical accounting. -### Confirmed failures +## Motivation -The incidents below include both false rejection and real under-accounting. -They are opposite results of using predictions as allocation facts. +The related incidents show both sides of the same modeling error: -| Incident | Observed behavior | Accounting mismatch | +| Incidents | Failure | Mismatch | | --- | --- | --- | -| #25782 / #25837 | HashBuild could OOM a CN instead of spilling | hash-table growth was not admitted at its physical allocation boundary | -| #26174 / #26178 | fulltext INSERT requested 18.72 GiB with about 1.29 GiB used | current ingress, retained tail, const materialization, and future drain were charged as if simultaneously owned | -| #26192 / #26231 / #26318 | LOAD DATA was rejected | runtime-filter and payload multipliers duplicated already retained owners | -| #26413 / #26438 | a Hive external-table self-join was rejected | a 50K-row logical ingress estimate did not match the real 8192-row copy segmentation and allocator rounding | -| #26454 | a string-expression join requested exactly 551,368,048,640 bytes while observed memory was about 6--7 GiB | TEXT maximum size was multiplied by row count and recursively summed through CAST/CONCAT nodes | -| #26433 / #26455 | expression result capacity could be under-counted or double-charged across reuse | the budget lease lifetime did not match retained `ExpressionExecutor` capacity | -| #26186 | a spill transition could under-count ingress overlap | independent estimates omitted a simultaneously live ownership state | - -Fixing one multiplier or expression kind does not close the defect class. A -more conservative estimate prevents one OOM shape but rejects more valid -queries. Relaxing the estimate restores those queries but can miss a different -physical overlap. - -### Root problem - -Three concerns are currently mixed: - -1. **Ownership accounting**: which allocations are live, which finite account - owns them, and when their charges end. -2. **Operation prediction**: how much an upcoming expression, batch copy, - marshal, or spill transition might allocate. -3. **Pressure response**: what execution does when the next allocation cannot - fit. +| #25782 / #25837 | CN OOM instead of bounded spill | hash-table growth was not admitted at allocation time | +| #26174 / #26178 | false fulltext INSERT rejection | non-simultaneous ingress/tail/drain states were summed | +| #26192 / #26231 / #26318 | false LOAD DATA rejection | runtime-filter/payload multipliers duplicated live owners | +| #26413 / #26438 | false external-table join rejection | logical ingress did not match physical batch segmentation | +| #26454 | request reported hundreds of GiB with single-digit-GiB usage | TEXT maximum width was multiplied through CAST/CONCAT | +| #26186 | spill overlap could be under-counted | independent estimates omitted a live transition state | -Only the first concern can support a hard capacity invariant. Prediction and -pressure response remain necessary, but they cannot be the source of truth for -live memory. +Fixing individual multipliers cannot close this class. A stricter estimate +rejects more valid queries; a looser estimate misses a different overlap. -### Relationship to SQL Resource Accounting +Three concerns must remain separate: -This RFC is separate from -[`SQL Resource Accounting`](./00000000_sql_resource_accounting.md). +1. ownership accounting: which retained allocations are live; +2. prediction: which strategy or capacity may be useful; +3. pressure response: spill, reclaim, reduce, degrade, or error. -SQL Resource Accounting defines observational facts used by statement trace, -CU, and physical-plan diagnostics. Its memory fields describe MPool domain -peaks after execution; it explicitly does not provide allocation-site -attribution or admission control. - -This RFC defines an execution-time control mechanism: - -| Concern | SQL Resource Accounting | This RFC | -| --- | --- | --- | -| Primary purpose | observe and persist usage | prevent an allocation from exceeding a finite execution account | -| Time of decision | execution summary / statement completion | immediately before allocation or growth | -| Unit | domain usage and peak | live physical allocation capacity | -| Missing data | quality flag | the path is not considered fully protected | -| Failure behavior | report incomplete facts | reclaim, spill, retry, or controlled pressure error | - -The two systems may share metrics and consistency checks, but one must not be -derived from the other. A terminal MPool peak cannot authorize an allocation -that has already happened, and a HashBuild account is not a complete statement -memory measurement. +Only ownership accounting can support a hard live-memory invariant. ## Goals -1. Make every hard memory charge correspond to a live physical allocation or - a named, bounded scratch owner. -2. Admit memory before allocation and leave allocator/account state unchanged - on failure. -3. Account for the real capacity chosen by MPool, including replacement - overlap during `Grow` and `Grow2`. -4. Make allocation provenance survive vector reuse, cross-owner handoff, and - cross-MPool `Free`. -5. Make Reset, Free, cancellation, retry, and generation turnover obey one - ownership contract. -6. Treat real capacity rejection as recoverable pressure where execution can - make progress. -7. Prevent a prediction alone from rejecting a query. -8. Add no per-row reservation or shared budget lock in the steady-state reuse - path. -9. Make the primitive reusable by spillable operators outside HashBuild. -10. Retain the no-OOM safety objective introduced for #25782. +1. Make each HashBuild hard memory charge correspond to live physical retained + storage. +2. Admit before allocation and leave allocator/account state unchanged on + rejection. +3. Represent real replacement overlap during growth. +4. Preserve provenance through vector reuse, copies, and owner transfer. +5. Give Reset, Free, cancellation, retry, and prepared reuse one lifecycle. +6. Treat only typed memory-capacity rejection as reclaimable pressure. +7. Prevent a prediction alone from rejecting a statement. +8. Add no per-row reservation or Go allocation to steady-state reuse paths. +9. Keep spill disk and FD accounting distinct from memory. +10. Preserve the no-OOM objective introduced for #25782. ## Non-goals This RFC does not: -- make the account equal to CN RSS or total Go runtime memory; -- remove process/CN headroom for caches, RPC, logs, goroutine stacks, or other - memory outside the controlled allocation domain; -- guarantee that every SQL statement completes under every finite cap; -- make an indivisible value smaller than its real representation; -- change SQL syntax, catalog metadata, or a persisted/wire format; -- replace spill disk and file-descriptor accounting; -- use type-specific exemptions for TEXT, CONCAT, CASE, LOAD DATA, or a - particular benchmark; -- raise the existing cap to hide false estimates; -- add a mutable process-wide "current memory account"; -- require all MPool users to become accounted in the first implementation. - -## Terminology - -### Allocation account - -A finite ledger that admits and releases bytes for one execution ownership -domain. `HashBuildBudgetGeneration` is the initial policy implementation, but -the allocator contract is not HashBuild-specific. - -### Accounted allocation - -An allocator-owned allocation whose metadata records: - -- its actual allocated capacity; -- an opaque reference to the account charge; -- bounded diagnostic classification such as owner class and allocation site. - -The charge belongs to the allocation, not to the operator field that currently -references it. - -The first implementation makes data-sized accounted allocations off-heap. -`MPool.Alloc(..., false)` uses a Go allocation: its requested size is not the -runtime size class, and removing MPool metadata does not make the GC reclaim it. -Such memory cannot be described as exact physical ownership at `Free`. - -### Allocation lease - -An exactly-once charge returned after successful admission. It is attached to -the allocation before the allocation is published to its caller. `Free` -releases it. Copying a lease value must not create a second release owner. - -### Explicit scratch lease - -A lease for bounded memory that cannot yet be allocated through MPool. It must -have one named owner, a finite size, and an explicit release point. It is an -exception used during migration, not a substitute for accounting ordinary -buffers. A data- or row-scaled Go allocation is not accepted merely by adding -such a lease: it must move off-heap or retain a conservative charge through a -proved GC-reclamation boundary. The initial implementation permits only small, -statically bounded Go metadata under explicit CN headroom. - -### Prediction - -A non-authoritative estimate used to select an execution strategy, start spill -before a hard limit, or choose an initial batch size. A prediction does not -create a retained allocation lease and cannot directly produce a terminal -budget error. - -### Generation - -One query-CN execution ownership epoch. The generation fixes the account -identity used by allocations that can outlive one operator call or move -between producer and consumer operators. - -## Required invariants - -### I1. Conservation - -For a live account at every observable transition: - -```text -account.used - = sum(capacity of live off-heap allocations charged to the account) - + sum(size of live named explicit scratch leases) -``` - -Predictions, logical vector length, source batch size, and SQL type maximum do -not appear in this equation. Small allocation/runtime metadata classified as -headroom is outside `account.used`; I9 separately requires a finite aggregate -bound for it. - -### I2. Admission precedes allocation - -The complete capacity of a new physical allocation is admitted before MPool or -another allocator changes state. - -### I3. Failure atomicity - -If admission or allocation fails: - -- no new allocation is published; -- the old allocation remains valid; -- any provisional lease is released; -- account usage and MPool usage return to their pre-operation values. - -### I4. Reclaimable ownership defines charge lifetime - -- shrinking a logical length does not release capacity; -- Reset retains both reusable capacity and its charge; -- reuse within existing capacity performs no new admission; -- off-heap `Free` releases physical memory and the charge together; -- a Go reference becoming unreachable is not treated as physical release; -- replacing a buffer transfers publication only after the replacement is - complete. - -### I5. One physical allocation has one charge - -Aliases, vector windows, const views, and shared areas do not create a second -charge. A deep copy creates a new allocation and therefore a new charge. - -### I6. Provenance survives handoff - -The account identity and charge remain correct when: - -- a build-side buffer is published to a `JoinMap`; -- a vector or batch moves to another operator; -- an allocation is freed through a different MPool; -- the producing operator has already Reset; -- a retained buffer is reused in a later call of the same execution - generation. - -### I7. Generation closure is not implicit memory release - -Sealing a generation prevents new admission but does not pretend that live -allocations disappeared. Existing allocation leases remain releasable. Normal -owner cleanup must bring usage to zero; a nonzero terminal value is an -invariant failure and a leak signal. - -### I8. Prediction cannot hard-reject - -An estimate may trigger early reclaim or spill. A terminal memory pressure -error must name a real allocation or explicit scratch request that failed -admission after applicable pressure responses. - -### I9. Metadata headroom is finite - -Every accounted owner has a proved maximum number of simultaneously live -allocations. Pointer headers, account-ID side records, registry entries, and -other per-allocation Go metadata have measured per-entry bounds. Their -aggregate bound is reserved as CN headroom: - -```text -metadata headroom - >= maximum live allocation count * measured metadata bytes per allocation - + maximum live generation count * measured registry bytes per generation -``` - -The implementation enforces both counts with finite CN-local metadata slots. -Opening a generation consumes a generation slot; publishing an accounted -allocation or replacement consumes an allocation-metadata slot; physical Free -returns that slot. Replacement growth temporarily consumes slots for both old -and new allocations. Slot exhaustion is exact metadata pressure, not a -payload estimate. An owner with no supported bound cannot be activated merely -because each payload allocation is charged. - -The first hash-table activation's concrete slot counts, conservative -per-entry bytes, and startup headroom formula are fixed in the implementation -plan. A later owner with smaller allocations must re-derive and provision its -own count before activation; the generic API does not silently widen the -proved domain. - -## Technical design - -### 1. Separate accounting mechanism from cap policy - -The low-level allocator must not import the SQL operator or `process` package. -It depends on a small generic contract, conceptually: - -```go -// Names are illustrative; this RFC does not freeze the Go API. -type AllocationAccount interface { - Acquire(AllocationRequest) (AllocationLease, error) -} - -type AllocationRequest struct { - Capacity uint64 - Class AllocationClass // bounded enum - Site AllocationSite // bounded enum -} - -type AllocationLease interface { - Capacity() uint64 - Release() -} -``` - -The contract is synchronous and does not wait for memory or reclamation; it may -briefly contend on account synchronization. It is called only when a physical -allocation or growth is required, not for every row or append. -It returns non-overlapping typed reasons: finite capacity pressure, -sealed generation, account mismatch, allocator size limit, and invariant -corruption. Only finite capacity pressure may enter reclaim/spill/retry. - -The policy layer remains responsible for: - -- query and CN caps; -- cap refresh; -- concurrency and linearization; -- spill disk/FD ledgers; -- metrics; -- typed pressure errors. - -The allocator layer is responsible for: - -- requesting the actual allocation capacity; -- failure rollback; -- associating the lease with allocation metadata; -- releasing the lease exactly once with physical memory. - -### 2. Allocation metadata and provenance - -MPool already tracks each allocation in pointer metadata so it can identify -the original pool, allocation size, off-heap status, double free, and -cross-pool free. Account provenance belongs at this same boundary. - -An accounted allocation adds an optional opaque charge handle to that -metadata. Unaccounted allocations retain current behavior. The handle's exact -representation must be benchmarked: adding a Go interface to every metadata -entry is not assumed acceptable. A compact pointer or account-local lease -record is preferred if it preserves exactly-once release and diagnostic -identity. - -The metadata is authoritative. Operator-maintained byte totals may remain as -diagnostics during migration, but they cannot independently release or -reconstruct the charge. - -Cross-MPool `Free` already delegates physical release to the original MPool. -The same terminal path releases the account charge, so the freeing caller does -not need to recover the producing operator or generation. - -An existing allocation's metadata also decides the account used by growth. -Growing an account-A allocation under account B is an invariant error. -Unaccounted-to-accounted conversion is never implicit: a migrated owner creates -an accounted destination and copies from the unaccounted source, or creates its -own buffers as accounted from the beginning. This prevents a caller from -changing only a vector field while the retained physical buffer still has -different provenance. - -For the compact-handle design, the registry is a finite set of reusable slots. -An opaque handle contains the slot and its generation counter. A slot is -reused only after sealing and exact zero; incrementing its generation makes -every older handle stale. Generation counters never wrap: a slot whose counter -is exhausted is retired. A missing or generation-mismatched registry entry -during `Free` is an invariant failure, not permission to drop the charge. - -Pointer headers, account-ID side records, and registry entries are `H` -metadata, not accounted payload. Their backing stores are sized from finite -CN-local slot limits and charged to explicit CN headroom. Before activating an -owner, supported live-allocation and generation counts must fit those limits; -the limits remain the safety backstop if an ownership assumption is wrong. - -### 3. New allocation protocol - -Initial allocation and growth share one authoritative -`AllocationCapacity(request, allocatorMode)` calculation. It includes -allocator rounding and the actual maximum accepted by `Alloc`; callers do not -infer capacity from a logical size or from `GrowCapacity` alone. - -For an account-aware allocation of capacity `C`: +- equate the HashBuild account with total query RSS; +- account all Go runtime/library allocations; +- guarantee completion under every finite cap; +- change optimizer join selection; +- make disk spill as fast as sufficient-memory execution; +- add type- or workload-specific exemptions; +- raise a cap to hide a false estimate; +- introduce a mutable process-wide current account. -```text -calculate actual allocator capacity C - -> acquire lease(C) - -> reserve one allocation-metadata slot - -> allocate C - -> on failure: return metadata slot, release lease(C), return error - -> attach lease to allocation metadata - -> publish allocation to caller -``` - -If metadata attachment can fail, it is part of the unpublished transaction: -free the new allocation, release the lease, and return an error. - -MPool cap failure, global cap failure, and underlying allocator failure all -follow the same rollback contract. - -Metadata publication is checked. An allocator panic before publication must -run provisional-lease cleanup. Cross-pool `Free` and deleted-owner-pool -fallback release the lease when they physically deallocate the allocation. - -Pool teardown is not itself a universal release event. A `noLock` teardown -that physically deallocates pool-local allocations releases each matching -lease. Unregistering a normal pool retains global pointer metadata and the -charge until later physical `Free`; live accounted allocations at teardown are -reported as an invariant violation, never bulk-released. - -### 4. Growth and replacement protocol - -MPool `Grow` and `Grow2` currently allocate a replacement, copy the old bytes, -then free the old allocation. Hard accounting must represent that real -overlap. - -For old capacity `O` and required logical size `R`: - -1. calculate `N = AllocationCapacity(GrowCapacity(O, R), allocatorMode)`; -2. keep the old allocation and its `O` charge live; -3. acquire a complete `N` lease, not `N - O`; -4. allocate `N`; -5. copy old data and any second source; -6. attach the new lease and publish the replacement; -7. free the old allocation, releasing its `O` lease. - -Peak account usage during replacement is therefore: - -```text -other live allocations + O + N -``` - -After publication and old-buffer release it is: - -```text -other live allocations + N -``` - -Reserving only the delta would under-count the actual allocate-copy-free peak. -Using a multiplier would be an estimate of the same fact even though the -allocator already knows `O` and `N`. - -If `R <= O`, no physical growth occurs and no account call is made. - -#### Growth failure table - -| Failure point | Old buffer | New buffer | Account result | -| --- | --- | --- | --- | -| capacity calculation | unchanged | none | unchanged | -| new lease admission | unchanged | none | unchanged | -| physical allocation | unchanged | freed/not published | new lease released | -| copy before publication | unchanged | freed/not published | new lease released | -| publication succeeds | released afterward | live | old lease released; new lease retained | - -The implementation must preserve this table under panic-safe cleanup where -MPool currently permits a recoverable error. - -### 5. Vector and batch propagation - -An allocation's metadata owns the retained charge. A vector additionally -needs an optional account selection for the first allocation of a currently -nil buffer. Subsequent growth inherits or verifies the existing allocation's -account. - -The following rules apply: - -| Operation | Rule | -| --- | --- | -| append within capacity | no admission | -| append causing data growth | admit exact replacement capacity | -| append causing varlen-area growth | independently admit exact replacement capacity | -| logical reset / set length to zero | retain allocation and charge | -| vector Free | free every owning buffer; each allocation releases its own charge | -| window/view/const alias | no new allocation and no new charge | -| deep Dup/copy | destination allocations use the destination account | -| batch handoff | allocation metadata preserves charges; no sum-and-rereserve | -| cross-pool Free | original allocation metadata releases the original charge | - -Data and varlen area are separate physical allocations and therefore separate -charges. Null bitmap and other auxiliary buffers follow their actual -allocation ownership rather than a synthetic per-vector total. - -The first production migration accounts only off-heap vector buffers. -Row-scaled null/group bitmaps currently allocate Go `[]uint64`; they must move -to an off-heap owner or remain an explicit activation blocker. Switching a -vector field to an account while its existing backing remains on-heap is not a -valid migration. - -A shared area must retain one physical release owner. If current vector sharing -permits multiple logical owners, account integration must use the same -reference/ownership mechanism that prevents physical double free; it must not -introduce a second budget-only reference count. - -### 6. Expression execution - -`FunctionResult` and `ExpressionExecutor` must create result vectors with the -execution account selected for the owning HashBuild path. - -Fixed-width results allocate from actual row count and element width through -normal vector growth. Varlen results allocate from actual appended payload. -Neither uses the maximum SQL type width multiplied by row count. - -Intermediate expression results are charged only while their buffers are -physically live. Reuse keeps the charge. Reset does not release it unless Reset -also frees the buffer. - -Expression implementations that create unbounded temporary Go strings or byte -slices bypass MPool and must be changed by one of these methods: - -1. write directly into an account-aware result buffer; -2. use an account-aware MPool scratch buffer; -3. for small metadata only, use a named explicit scratch lease with a proved - finite bound and one cleanup owner. - -For example, CONCAT should not construct an unaccounted complete Go string and -then copy it into an accounted result vector. The temporary and final buffers -can be simultaneously live, so omitting the temporary would violate I1. - -Recursive expression peak calculation may remain temporarily as an early-spill -hint. Before deleting it as a hard gate, an allocation-site ledger must cover -every reachable data-scaled allocation in generic evaluation and built-ins, -including selection arrays, conversion slices, nested result vectors, and -function-specific scratch. Migrating only `FunctionResult` is not closure. - -### 7. Non-vector HashBuild allocations - -The migration inventory must include all memory whose lifetime is owned by the -HashBuild execution domain: - -- copied build and probe batches; -- integer and string hash-table blocks; -- selection lists and group mappings; -- join-map auxiliary storage; -- expression keys and intermediate results; -- spill scatter buffers; -- spill/re-spill read and decode buffers; -- runtime-filter buffers; -- marshal/unmarshal scratch; -- retained emergency scratch. - -An ownership closure is not complete merely because its largest vector is -accounted. Every data-scaled allocation reachable from an activated owner must -use account-aware, synchronously reclaimable allocation. Unactivated owners -retain their complete legacy charge until their own closure; only small, -statically bounded metadata may remain under named headroom. - -Hash-table callbacks that already reserve from `ResizePlan` are an intermediate -bridge. The final charge should be owned by the physical hash-table allocation -metadata rather than a parallel slice of reservation tokens in the operator. - -### 8. Pressure response state machine - -A failed exact allocation admission returns typed memory pressure, not -`ErrHashBuildBudgetInvalid`. - -The owner handling the request proceeds through a bounded state machine: - -```text -need allocation - -> exact admission succeeds - -> allocate and continue - -> exact admission rejected - -> release reclaimable retained capacity - -> retry exact allocation - -> start/advance spill - -> retry exact allocation - -> reduce processing batch where semantics permit - -> retry exact allocation - -> degrade optional structure where permitted - -> retry/continue - -> return controlled minimum-unit pressure error -``` - -Not every allocation supports every response. The request carries an owner -class that maps to a policy: - -| Owner class | Permitted response | -| --- | --- | -| retained reusable result | release retained capacity, then retry | -| spillable build/probe input | spill/re-spill, then retry | -| splittable expression batch | reduce batch, then retry | -| runtime filter | degrade to PASS when correctness is unchanged | -| indivisible single value / minimum hash block | controlled error with actual capacity | -| invariant corruption | fail immediately as invalid state | - -The state machine records which responses were attempted so it cannot loop -without progress. A retry is justified only after account usage decreased, -the input unit became smaller, spill state advanced, or an optional owner was -disabled. - -`ErrHashBuildBudgetInvalid` remains reserved for arithmetic overflow, corrupted -ownership, double release, account mismatch, or an impossible lifecycle -transition. Ordinary finite pressure is not an invariant failure. - -A sealed or stale generation is a lifecycle result distinct from finite -pressure. It cannot enter the retry state machine. - -#### Retry transaction boundary - -Allocation rollback alone does not make an operator operation retryable. A -vector or expression may successfully grow one retained buffer, fail on the -next allocation, and leave reusable capacity or partially written output. -Each retry-capable owner therefore defines: - -- the unpublished operation checkpoint; -- which retained capacity may survive failure; -- how row/output publication is rolled back; -- the Reset/Free actions required before retry; -- how a smaller batch resumes without duplicate work. +## Controlled domain -Until the shared controller and these checkpoints exist, an exact rejection -returns a controlled terminal pressure error. An earlier migration must not -claim spill/reduce/retry merely because it has allocation-level rollback. +The first implementation accounts physical storage retained by HashBuild and +its join consumers: hash tables, copied build batches, JoinMap state, retained +keys, join bitmaps/capture/result state, Product result state, +runtime-filter payloads, and spill encode/decode/rebuild buffers. -#### Forward-progress memory +ExpressionExecutor temporary results, caches, and library-internal Go heap +remain in the existing MPool/Go runtime domain. They cannot truthfully be put +in an exact terminal-zero account while regexp, JSON, JQ, and similar libraries +do not expose allocator/free hooks. An expression value becomes accounted when +it is physically copied into retained HashBuild/join storage. -Spill and reclaim paths do not bypass admission. Encoding, decoding, scatter, -and IO buffers are allocations too. Allowing normal work to consume the last -byte and then allocating uncharged "emergency" scratch would reproduce the -original safety hole at a different site. +This is a static ownership boundary, not a runtime fallback. A future general +expression-memory design must either use allocator-aware implementations or a +separate explicitly non-exact transient policy. Estimates must not be inserted +into this exact ledger. -Before retaining a unit of work, an operator must preserve one bounded way to -make progress: +ProductL2's accounted input remains the producer-owned JoinMap. Its additional +CPU/GPU index and scratch storage is outside the first controlled domain because +the native GPU allocator does not expose an admission/free capacity contract. +Partially charging only Go-visible buffers would be another ledger, not exact +physical ownership. ProductL2 can join this domain only when both platform +implementations provide the same allocation contract. -1. reuse an already allocated and accounted spill buffer; -2. keep a finite progress sub-cap that normal work cannot consume and charge - actual spill/cleanup allocations against it; or -3. reduce the retained or spill chunk until the minimum real progress - allocation fits. +## Allocation contract -Progress headroom is cap policy, not a fabricated live-memory charge: -`account.used` still contains only actual allocations and explicit live -scratch. If a sub-cap is used, normal and progress allocations remain bounded -by the same total query/CN cap, with the progress portion unavailable to normal -growth. +An accounted allocation records: -The minimum progress unit is derived from the concrete buffer/chunk layout, not -from SQL type maximums. If retained state plus that minimum real unit cannot -fit, admission must stop earlier or return a controlled pressure error; it -must not wait until spill itself is unable to start. +- actual allocator capacity; +- allocation account and generation; +- bounded owner class and allocation site. -### 9. Concurrency and generation lifecycle +The contract is transactional: -Child pipelines may share a `BaseProcess` and execute concurrently. Therefore -neither Process nor MPool may contain a mutable "current account" used -implicitly by all allocations. +1. compute the allocator capacity using checked arithmetic; +2. acquire account/controller capacity; +3. allocate physical storage; +4. publish metadata and ownership; +5. on failure, undo the acquisition before returning; +6. on Free, deallocate and release exactly once. -Account selection is explicit at the owner/allocation boundary and immutable -for a published allocation. Concurrent allocations linearize in the account's -existing query/CN admission operation. Reuse within capacity does not enter -that lock. +For growth, old and replacement allocations are both charged until the copy +succeeds and the old allocation is physically freed. Views borrow storage and +do not charge again. Copies charge their destination allocation. -Normal generation lifecycle is: +The query generation and CN aggregate are checked at the same controller +boundary. There is no operator-owned memory reservation token parallel to the +allocator lease. -```text -open - -> allocations and explicit scratch may acquire leases - -> seal: no new leases - -> owner cleanup frees all live allocations/scratch - -> used reaches zero - -> finalize -``` - -One `Compile` execution attempt on each CN is the sole seal/finalize owner for -that CN's generation. HashBuild Reset cannot seal: `JoinMap`, spill payload, -broadcast, remote scope, and message-board consumers may outlive the producing -operator and still own or create controlled work. The attempt coordinator -seals only after all scopes and remote notifiers it owns have stopped -publishing work and its MessageBoard consumers have reached a terminal state. -Cleanup may then release old leases; exact zero produces one immutable -terminal snapshot and permits registry-slot reuse. The statement -`ResourceRoot` aggregates these immutable CN-attempt snapshots but does not own -allocation release. - -A nonzero terminal generation does not disappear and does not wait forever: - -```text -seal after execution/message quiescence - -> run terminal owner cleanup - -> zero: - export one valid immutable snapshot - remove the registry entry - -> nonzero at terminal-cleanup completion: - export one immutable invariant-failure snapshot - retain a release-capable tombstone until late Free reaches zero - suspend admission of new accounted generations on that CN -``` - -Suspension bounds tombstone growth to generations already active when the -first invariant failure is detected. It is lifted only after every tombstone -reaches zero; an operational deadline escalates with owner/site diagnostics -and permits a controlled CN restart rather than deleting provenance. Late -release may update bounded health counters, but it cannot rewrite or duplicate -the exported snapshot. - -Generation open and nonzero-terminal suspension linearize through the same -CN-local generation gate. Opening publishes a generation only if the -suspension check succeeds in that transaction. Once suspension publication -linearizes, no later open may publish; opens that linearized earlier are the -finite active set allowed to finish or become tombstones. - -`SetStmtProfile` turnover, frontend `StatementInfo.EndStatement`, and -`HashBuildBudgetGeneration.Close` alone do not prove this per-CN quiescence or -validate zero. The implementation therefore gives the `Compile` attempt an -explicit post-pipeline/MessageBoard-close transition for success, failure, -panic, cancellation, retry, broadcast, prepared reuse, and remote execution. -A forced close never silently zeros accounting while allocations remain live. - -Ownership transfer does not change generations. If a transfer would cross to a -different generation, it must either: - -- retain the original generation until the allocation is freed; or -- perform one explicit atomic charge transfer before publication. - -The initial implementation should prefer retaining original provenance; charge -transfer adds a second failure and rollback boundary and is unnecessary for -normal HashBuild producer-to-consumer handoff. - -### 10. Observability - -Admission exposes bounded owner/site, actual capacity, used/cap, attempted -pressure response, and terminal result. Exact allocation events, prediction -hints, pressure responses, and invariant failures remain distinguishable. - -Owner/site values are bounded enums. Metrics and logs aggregate at -operator/generation or terminal-pressure boundaries; there is no -per-allocation log or unbounded SQL/stack label. - -The controlled-domain snapshot has a stable generation identity and is -exported exactly once by its CN attempt coordinator. The statement resource -root may aggregate those snapshots. It is separate from SQL Resource -Accounting's current off-heap MPool domain; consumers must not sum duplicate -operator references to one generation or claim the domains match before owner -coverage does. - -## Migration plan - -Migration is incremental by complete physical owner. Exact accounting and -removal of that owner's legacy hard charge happen atomically; one buffer is -never charged by both models. +## Lifecycle -The implementation order is: +Compile opens one generation for each local statement attempt, configures all +physical-plan owners, and records them for reverse-order cleanup. Parallel +scan/load workers created during `runOnce` join the same generation before +worker `Prepare` starts. Configuration is atomic: a failure clears newly +configured owners in reverse order. -1. measured metadata/API decisions and a reference model; -2. generic MPool allocation transaction; -3. dormant Vector/Batch propagation; -4. allocation-site closure and dormant expression/spill propagation; -5. statement generation lifecycle, typed pressure, and retry checkpoints; -6. hash-table cell/descriptor activation with only its legacy charge removed; -7. copied-batch and JoinMap activation; -8. expression-owner activation; -9. spill and runtime-filter closures; -10. unified join pressure control and remaining legacy estimator deletion; -11. workload and performance acceptance. +At completion, the message board is drained, owners are cleared, and the +account is sealed. A valid terminal state requires zero live bytes and zero +live metadata. A late allocation or release mismatch is a lifecycle invariant, +not capacity pressure. Prepared statements and retries use new generations; +generation-bound state cannot survive Reset. -The allocation-site ledger, PR gates, rollback rules, and test commands live in -the -[implementation plan](../design/allocation_accounted_memory_admission_impl.md). +## Pressure protocol -## Rollout and compatibility +Typed reasons keep control flow honest: -The change is internal to one CN binary and does not change SQL, catalog, disk, -or RPC formats. Unmigrated MPool users retain current behavior. A HashBuild -owner enables exact accounting only after its full alloc-to-Free closure is -covered and its legacy hard charge is removed. The final design has no -permanent legacy/exact switch. +- memory capacity may reclaim, spill, reduce an unpublished input, or disable + an optional optimization; +- sealed, suspended, owner/site mismatch, or allocator invariant is terminal; +- spill disk and spill FD rejection are resource-specific terminal results; +- minimum-unit pressure ends retries when no smaller valid input exists. -## Testing strategy +A retry requires monotonic evidence: fewer live bytes, a new spill epoch, a +smaller input unit, or an optional structure disabled. Publication and spill +I/O cannot be replayed merely because the same capacity error recurred. -Testing is derived from the invariants: +Runtime filters are optional and degrade to PASS if their retained payload +cannot be admitted. Required join state never bypasses admission. -- a randomized reference model checks conservation after alloc, grow, failure, - reuse, handoff, Free, seal, and cancellation; -- boundary and fault tests cover exact cap, allocator rounding, unpublished - rollback, cross-pool Free, and generation turnover; -- container/operator tests cover aliases, varlen data, Reset, broadcast, - spill/re-spill, pressure progress, and optional degradation; -- workload regressions cover #25782, #26174, #26192, #26413, #26454, and TPCH - spill/non-spill paths; -- performance gates verify no per-row account operation, no budget lock on - within-capacity reuse, bounded metadata cost, concurrent-generation P50/P99 - admission latency, release storms, and separately measured resident/spill - behavior. +## Spill ownership -Exact matrices and per-PR gates are maintained in the -[implementation plan](../design/allocation_accounted_memory_admission_impl.md). +SpillEngine must receive the producer's live budget generation. Memory, disk, +and FD tokens describe different physical resources. One file owns one +growable disk token and one FD token; file handoff moves both. Close releases +them once. Recursive spill validates framing, schema, row conservation, file +metadata, and bounded queue progress. -## Drawbacks +## Relationship to SQL Resource Accounting -### Allocator and container changes are invasive +[`SQL Resource Accounting`](./00000000_sql_resource_accounting.md) observes and +persists statement resource facts. This RFC controls an allocation before it +happens. -Correct ownership crosses MPool, vector, batch, expression, hash table, and -operator handoff boundaries. A partial implementation can create a more -convincing but still incomplete safety claim. - -Mitigation: migrate by complete owner class, maintain the ownership ledger, and -gate each phase on conservation properties rather than workload success alone. - -### Allocation metadata has a cost - -An account handle can increase pointer-map memory and alloc/free work even when -only some allocations are accounted. - -Mitigation: keep it optional and compact, benchmark representation choices, -and avoid a Go interface stored inline in every metadata record unless -measurement supports it. - -### Exact replacement overlap can reject earlier than steady-state size - -If a 6 GiB buffer grows to 8 GiB, the current allocator may need 14 GiB live -during copy even though the final buffer is 8 GiB. Charging only 2 GiB would -look friendlier but would not protect the real peak. - -Mitigation: reclaim or spill before growth, reduce the processing unit, or -introduce a genuinely lower-overlap allocator operation. Do not hide the peak -with delta accounting. - -### Go-heap reclamation is not allocator-controlled - -Dropping an MPool pointer record for `Alloc(..., false)` does not synchronously -return its backing bytes. Treating that event as exact physical release would -allow new work while the old Go object remains resident. - -Mitigation: keep data-scaled controlled owners off-heap. Audit and migrate -row/payload-scaled Go slices before activation; reserve separate CN headroom -only for small, proved-bounded runtime metadata. - -### This is not complete RSS accounting - -Go runtime, goroutine stacks, caches, and unrelated subsystems remain outside -the HashBuild account. - -Mitigation: preserve explicit CN headroom and state coverage boundaries. A -future broader account may reuse the primitive but requires its own ownership -inventory. +| Concern | SQL Resource Accounting | This RFC | +| --- | --- | --- | +| Purpose | observation and diagnostics | admission and pressure control | +| Time | during/after execution | immediately before allocation | +| Unit | domain usage/peak | live physical retained capacity | +| Missing data | quality flag | storage is outside this controlled domain | +| Failure action | report | reclaim/spill/reduce/degrade/error | -## Rationale and alternatives - -### Tune multipliers and type rules +The systems may cross-check terminal facts, but neither authorizes the other. -Rejected. The confirmed incidents demonstrate that no fixed multiplier -represents aliasing, reuse, variable payloads, segmentation, and replacement -overlap simultaneously. - -### Use predictions only to choose spill, then remove hard budgets - -Rejected. Prediction errors in the other direction can again allow #25782 to -OOM the CN. Real allocations still need finite admission. +## Performance constraints -### Charge only `newCapacity - oldCapacity` on growth +- no per-row account object; +- no per-allocation Go object in steady state; +- fixed-size provenance in allocator metadata; +- controller calls only when physical capacity changes; +- views and Reset reuse do not re-admit unchanged storage; +- spill tokens scale with open files, not records; +- no extra expression-account bookkeeping in hot predicates. -Rejected for the current allocate-copy-free implementation. It under-counts -the period where both allocations are live. +Benchmarks must compare unaccounted MPool behavior, accounted acquire/release, +vector growth/reuse, hash build/lookup, and spill scatter. Correctness +requires exact terminal zero; performance acceptance requires no material +regression outside measurement noise for unaccounted paths and no new +data-scaled Go allocation in accounted paths. -### Use MPool current bytes as the HashBuild budget +## Alternatives rejected -Rejected. A shared MPool contains other owners and does not preserve -HashBuild/query provenance across child pipelines. Sampling after allocation -also cannot provide pre-allocation safety. +### Patch each multiplier -### Store a mutable current account on Process or MPool +It cannot close both false-positive and false-negative estimate errors. -Rejected. Concurrent child pipelines can share the same base process and MPool. -The wrong goroutine could charge or free against another generation. +### Reserve an estimated amount, reconcile later -### Poll RSS and spill near the cgroup limit +The estimate can reject before real allocation or under-count live overlap. It +also creates a second release owner beside MPool. -Rejected as the primary control. RSS is delayed, includes unrelated memory, -and cannot make a specific allocation failure-atomic. It remains useful as a -coarse pressure signal and validation metric. +### Sample RSS or heap counters -### Wrap MPool only at operator call sites +Sampling is observational, process-wide, and too late to authorize a specific +allocation. -Insufficient by itself. A wrapper can select the account for initial -allocation, but provenance must still survive vector growth, physical handoff, -and cross-pool Free. The terminal charge belongs in allocation metadata. +### Account only selected expression buffers -### Raise or disable the cap +This reports false exactness while opaque library Go heap remains untracked. +The implementation therefore accounts retained copies, not partial expression +internals. -Rejected. It hides false rejection while weakening the original no-OOM -requirement. +### Keep old and new production paths behind a switch -## Unresolved questions +Two admission semantics double ownership states, tests, and failure modes. The +final implementation removes the old path instead. -1. Does the provisional 16-byte-header plus side-map representation retain its - advantage in real MPool, cross-pool, and high-concurrency benchmarks? -2. What explicit MessageBoard close-and-drain primitive and tests prove - quiescence at the selected local and remote `Compile` attempt hooks? -3. Which remaining data-scaled Go-heap sites can write directly to off-heap - output, and which need a new off-heap scratch abstraction? -4. What is the minimum semantically safe batch and operation checkpoint for - each expression and spill phase? -5. What cap/headroom policy is appropriate once exact HashBuild ownership - replaces conservative estimates? This is policy work and must not change the - accounting invariant. -6. What measured metadata and hot-path overhead is acceptable for enabling the - primitive beyond HashBuild? +## Rollout and validation -These questions affect implementation shape, not the core decision that hard -admission must be tied to owned physical capacity. +The merge unit is one PR with reviewable commits, but the final branch must +present a single production path. Local validation includes unit tests, +prepared/retry and parallel-clone lifecycle tests, cancellation/transfer tests, +spill resource tests, race tests, build, vet, and allocation benchmarks. -## Acceptance criteria +Remote benchmark workflows are not required for correctness convergence in +this implementation cycle. -The RFC is implemented only when: +## Completion invariant -- I1--I9 hold under property tests, fault injection, cancellation, and race - execution; -- account-aware alloc/grow/free use the same real capacity calculation as - MPool; -- every HashBuild-owned data-scaled allocation is off-heap and accounted; -- every excluded Go/runtime metadata allocation is statically bounded and - covered by explicit CN headroom; -- live-allocation and generation counts prove the aggregate I9 metadata - headroom; -- each CN `Compile` attempt seals, validates exact zero, and exports one - generation snapshot after execution/MessageBoard quiescence; -- a nonzero terminal generation exports one failure snapshot, preserves - release provenance, and cannot accumulate unbounded tombstones; -- Reset retains reusable allocation charges and Free releases them exactly - once; -- aliases and handoffs do not duplicate charges; -- estimator-only false rejection is structurally impossible on migrated - paths; -- exact pressure triggers bounded reclaim/spill/retry/reduce behavior only - across proved operation checkpoints; -- real minimum-unit over-cap errors report actual allocation capacity and - owner/site; -- #25782 cannot exceed the finite account or OOM the CN; -- #26174, #26192, #26413, and #26454 pass durable workload regressions; -- TPCH spill and non-spill performance gates pass; -- superseded hard estimators and parallel reservation owners are removed; -- the implementation documentation states the remaining memory outside the - account and preserves corresponding CN headroom. +The RFC is complete when every retained HashBuild/join allocation in scope has +one physical owner, one admission path, and one terminal release; estimated +hard memory reservations and activation gates are absent; runtime clones share +the statement generation; all pressure/resource types remain disjoint; and +independent review plus fresh local validation reports no blocker or major +regression. diff --git a/pkg/common/bitmap/bitmap.go b/pkg/common/bitmap/bitmap.go index 6ba3fc586a91f..cc9a10c9f8150 100644 --- a/pkg/common/bitmap/bitmap.go +++ b/pkg/common/bitmap/bitmap.go @@ -243,7 +243,7 @@ func (n *Bitmap) InstallExternalStorage(storage []uint64) []uint64 { } // ReleaseExternalStorage detaches caller-owned storage and clears the bitmap. -// It returns nil for a legacy Go-owned bitmap. +// It returns nil for a bitmap that owns its Go-allocated backing. func (n *Bitmap) ReleaseExternalStorage() []uint64 { if !n.HasExternalStorage() { return nil @@ -644,6 +644,61 @@ func (n *Bitmap) Count() int { return int(n.count) } +// CountRange returns the number of set bits in [start, end). It never scans +// outside the bitmap's logical coverage and does not allocate. +func (n *Bitmap) CountRange(start, end uint64) int { + if n == nil || start >= end || start >= uint64(n.logicalLen()) { + return 0 + } + if end > uint64(n.logicalLen()) { + end = uint64(n.logicalLen()) + } + first := start >> 6 + last := (end - 1) >> 6 + if first == last { + mask := (^uint64(0) << (start & 63)) & + (^uint64(0) >> ((-end) & 63)) + return bits.OnesCount64(n.data[first] & mask) + } + count := bits.OnesCount64(n.data[first] & (^uint64(0) << (start & 63))) + for word := first + 1; word < last; word++ { + count += bits.OnesCount64(n.data[word]) + } + count += bits.OnesCount64(n.data[last] & (^uint64(0) >> ((-end) & 63))) + return count +} + +// AnySetNotIn reports whether [start, end) contains a bit set in n and not in +// other. It is used when one provenance bitmap (GROUPING) overrides another +// (SQL NULL) without expanding either bitmap row by row. +func (n *Bitmap) AnySetNotIn(other *Bitmap, start, end uint64) bool { + if n == nil || start >= end || start >= uint64(n.logicalLen()) { + return false + } + if end > uint64(n.logicalLen()) { + end = uint64(n.logicalLen()) + } + first := start >> 6 + last := (end - 1) >> 6 + for word := first; word <= last; word++ { + mask := ^uint64(0) + if word == first { + mask &= ^uint64(0) << (start & 63) + } + if word == last { + mask &= ^uint64(0) >> ((-end) & 63) + } + value := n.data[word] & mask + if other != nil && word < uint64(len(other.data)) { + value &^= other.data[word] + } + if value != 0 { + return true + } + } + return false +} + func (n *Bitmap) ToArray() []uint64 { rows := make([]uint64, 0, n.Count()) ToArray(n, &rows) @@ -681,6 +736,27 @@ func (n *Bitmap) MarshalSize() int { return MarshalHeaderSize + len(n.data)*8 } +// Validate checks the in-memory representation after streaming decode. +func (n *Bitmap) Validate() error { + if n == nil || n.logicalLen() < 0 || n.count < 0 || + n.count > n.logicalLen() || + len(n.data) != int((n.logicalLen()+63)/64) { + return moerr.NewInvalidInputNoCtx("invalid bitmap representation") + } + actual := int64(0) + for i, word := range n.data { + if i == len(n.data)-1 && n.logicalLen()%64 != 0 && + word>>uint(n.logicalLen()%64) != 0 { + return moerr.NewInvalidInputNoCtx("invalid bitmap trailing bits") + } + actual += int64(bits.OnesCount64(word)) + } + if actual != n.count { + return moerr.NewInvalidInputNoCtx("invalid bitmap count") + } + return nil +} + // DecodeMarshalHeader validates the fixed bitmap wire header. func DecodeMarshalHeader(data []byte) ( count int64, @@ -743,19 +819,40 @@ func (n *Bitmap) MarshalTo(w io.Writer) error { } bitLength := uint64(n.logicalLen()) dataLength := uint64(len(n.data) * 8) - for _, value := range [][]byte{ - types.EncodeInt64(&n.count), - types.EncodeUint64(&bitLength), - types.EncodeUint64(&dataLength), - types.EncodeSlice(n.data), - } { - written, err := w.Write(value) - if err != nil { + if typed, ok := w.(interface { + WriteInt64(int64) error + WriteUint64(uint64) error + }); ok { + if err := typed.WriteInt64(n.count); err != nil { return err } - if written != len(value) { - return io.ErrShortWrite + if err := typed.WriteUint64(bitLength); err != nil { + return err } + if err := typed.WriteUint64(dataLength); err != nil { + return err + } + return writeBitmapMarshalBytes(w, types.EncodeSlice(n.data)) + } + if err := writeBitmapMarshalBytes(w, types.EncodeInt64(&n.count)); err != nil { + return err + } + if err := writeBitmapMarshalBytes(w, types.EncodeUint64(&bitLength)); err != nil { + return err + } + if err := writeBitmapMarshalBytes(w, types.EncodeUint64(&dataLength)); err != nil { + return err + } + return writeBitmapMarshalBytes(w, types.EncodeSlice(n.data)) +} + +func writeBitmapMarshalBytes(w io.Writer, value []byte) error { + written, err := w.Write(value) + if err != nil { + return err + } + if written != len(value) { + return io.ErrShortWrite } return nil } diff --git a/pkg/common/bitmap/types.go b/pkg/common/bitmap/types.go index 883625d144f86..0e12c750d66ad 100644 --- a/pkg/common/bitmap/types.go +++ b/pkg/common/bitmap/types.go @@ -22,9 +22,9 @@ type Iterator interface { type Bitmap struct { count int64 //in version 1, we use emptyFlag with type int32 to indicate whether it is empty - // taggedLen stores the logical length directly for legacy backing and its - // bitwise complement for caller-owned backing. This keeps Bitmap's legacy - // footprint unchanged while making backing ownership explicit. + // taggedLen stores the logical length directly for Go-owned backing and its + // bitwise complement for caller-owned backing. This keeps Bitmap's footprint + // unchanged while making backing ownership explicit. taggedLen int64 data []uint64 } diff --git a/pkg/common/hashmap/inthashmap_lazy_test.go b/pkg/common/hashmap/inthashmap_lazy_test.go index be4131421bbd6..665f964abd6c4 100644 --- a/pkg/common/hashmap/inthashmap_lazy_test.go +++ b/pkg/common/hashmap/inthashmap_lazy_test.go @@ -52,7 +52,8 @@ func TestIntHashMapIteratorLazyBuffers(t *testing.T) { insertedVs := append([]uint64(nil), vs...) insertedZvs := append([]int64(nil), zvs...) - foundVs, foundZvs := itr.Find(0, count, vecs) + foundVs, foundZvs, err := itr.Find(0, count, vecs) + require.NoError(t, err) if count > 0 { require.Equal(t, insertedVs, foundVs) require.Equal(t, insertedZvs, foundZvs) @@ -224,7 +225,10 @@ func BenchmarkIntHashMapFindFloat32(b *testing.B) { b.SetBytes(int64(count * types.T_float32.TypeLen())) b.ResetTimer() for i := 0; i < b.N; i++ { - benchmarkIntValues, benchmarkIntZValues = itr.Find(0, count, vecs) + benchmarkIntValues, benchmarkIntZValues, err = itr.Find(0, count, vecs) + if err != nil { + b.Fatal(err) + } } }) } diff --git a/pkg/common/hashmap/inthashmap_test.go b/pkg/common/hashmap/inthashmap_test.go index 8f28a8e63cba2..d5a746e1ae229 100644 --- a/pkg/common/hashmap/inthashmap_test.go +++ b/pkg/common/hashmap/inthashmap_test.go @@ -24,6 +24,48 @@ import ( "github.com/stretchr/testify/require" ) +func TestIntHashMapProbeGroupingDoesNotMatchRawKey(t *testing.T) { + mp := mpool.MustNewZero() + hashMap, err := NewIntHashMap(false, mp) + require.NoError(t, err) + defer hashMap.Free() + raw := vector.NewVec(types.T_uint8.ToType()) + require.NoError(t, vector.AppendFixed(raw, uint8(0), false, mp)) + grouping := vector.NewRollupConst(types.T_uint8.ToType(), 1, mp) + defer raw.Free(mp) + defer grouping.Free(mp) + + iterator := hashMap.NewIterator() + _, _, err = iterator.Insert(0, 1, []*vector.Vector{raw}) + require.NoError(t, err) + values, zValues, err := iterator.Find(0, 1, []*vector.Vector{grouping}) + require.NoError(t, err) + require.Equal(t, []uint64{0}, values) + require.Equal(t, []int64{0}, zValues) +} + +func TestIntHashMapPartialGroupingRowsDoNotMatchRawKeys(t *testing.T) { + mp := mpool.MustNewZero() + hashMap, err := NewIntHashMap(false, mp) + require.NoError(t, err) + defer hashMap.Free() + build := vector.NewVec(types.T_int32.ToType()) + probe := vector.NewVec(types.T_int32.ToType()) + require.NoError(t, vector.AppendFixedList(build, []int32{7, 8}, nil, mp)) + require.NoError(t, vector.AppendFixedList(probe, []int32{7, 8}, nil, mp)) + probe.GetGrouping().Add(0) + defer build.Free(mp) + defer probe.Free(mp) + + iterator := hashMap.NewIterator() + _, _, err = iterator.Insert(0, 2, []*vector.Vector{build}) + require.NoError(t, err) + values, zValues, err := iterator.Find(0, 2, []*vector.Vector{probe}) + require.NoError(t, err) + require.Equal(t, []uint64{0, 2}, values) + require.Equal(t, []int64{0, 1}, zValues) +} + func TestIntHashMap_Iterator(t *testing.T) { { m := mpool.MustNewZero() @@ -42,7 +84,8 @@ func TestIntHashMap_Iterator(t *testing.T) { vs, _, err := itr.Insert(0, rowCount, vecs) require.NoError(t, err) require.Equal(t, []uint64{1, 1, 1, 2, 2, 2, 3, 3, 3, 4}, vs) - vs, _ = itr.Find(0, rowCount, vecs) + vs, _, err = itr.Find(0, rowCount, vecs) + require.NoError(t, err) require.Equal(t, []uint64{1, 1, 1, 2, 2, 2, 3, 3, 3, 4}, vs) for _, vec := range vecs { vec.Free(m) @@ -63,7 +106,8 @@ func TestIntHashMap_Iterator(t *testing.T) { vs, _, err := itr.Insert(0, Rows, vecs) require.NoError(t, err) require.Equal(t, []uint64{1, 2, 1, 3, 1, 4, 1, 5, 1, 6}, vs[:Rows]) - vs, _ = itr.Find(0, Rows, vecs) + vs, _, err = itr.Find(0, Rows, vecs) + require.NoError(t, err) require.Equal(t, []uint64{1, 2, 1, 3, 1, 4, 1, 5, 1, 6}, vs[:Rows]) for _, vec := range vecs { vec.Free(m) @@ -83,7 +127,8 @@ func TestIntHashMap_Iterator(t *testing.T) { vs, _, err := itr.Insert(0, Rows, vecs) require.NoError(t, err) require.Equal(t, []uint64{1, 2, 1, 3, 1, 4, 1, 5, 1, 6}, vs[:Rows]) - vs, _ = itr.Find(0, Rows, vecs) + vs, _, err = itr.Find(0, Rows, vecs) + require.NoError(t, err) require.Equal(t, []uint64{1, 2, 1, 3, 1, 4, 1, 5, 1, 6}, vs[:Rows]) for _, vec := range vecs { vec.Free(m) @@ -103,7 +148,8 @@ func TestIntHashMap_Iterator(t *testing.T) { vs, _, err := itr.Insert(0, Rows, vecs) require.NoError(t, err) require.Equal(t, []uint64{1, 2, 1, 3, 1, 4, 1, 5, 1, 6}, vs[:Rows]) - vs, _ = itr.Find(0, Rows, vecs) + vs, _, err = itr.Find(0, Rows, vecs) + require.NoError(t, err) require.Equal(t, []uint64{1, 2, 1, 3, 1, 4, 1, 5, 1, 6}, vs[:Rows]) for _, vec := range vecs { vec.Free(m) @@ -123,7 +169,8 @@ func TestIntHashMap_Iterator(t *testing.T) { vs, _, err := itr.Insert(0, Rows, vecs) require.NoError(t, err) require.Equal(t, []uint64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, vs[:Rows]) - vs, _ = itr.Find(0, Rows, vecs) + vs, _, err = itr.Find(0, Rows, vecs) + require.NoError(t, err) require.Equal(t, []uint64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, vs[:Rows]) for _, vec := range vecs { vec.Free(m) @@ -143,7 +190,8 @@ func TestIntHashMap_Iterator(t *testing.T) { vs, _, err := itr.Insert(0, Rows, vecs) require.NoError(t, err) require.Equal(t, []uint64{0, 1, 0, 2, 0, 3, 0, 4, 0, 5}, vs[:Rows]) - vs, _ = itr.Find(0, Rows, vecs) + vs, _, err = itr.Find(0, Rows, vecs) + require.NoError(t, err) require.Equal(t, []uint64{0, 1, 0, 2, 0, 3, 0, 4, 0, 5}, vs[:Rows]) for _, vec := range vecs { vec.Free(m) @@ -208,7 +256,8 @@ func TestIntHashMap_MarshalUnmarshal(t *testing.T) { require.Equal(t, expectedGroupCount, unmarshaledMp.GroupCount()) require.Equal(t, mp.HasNull(), unmarshaledMp.HasNull()) - foundVs, _ := unmarshaledMp.NewIterator().Find(0, rowCount, vecs) + foundVs, _, err := unmarshaledMp.NewIterator().Find(0, rowCount, vecs) + require.NoError(t, err) require.Equal(t, expectedMappedValue, foundVs) }) @@ -245,7 +294,8 @@ func TestIntHashMap_MarshalUnmarshal(t *testing.T) { require.Equal(t, expectedGroupCount, unmarshaledMp.GroupCount()) require.Equal(t, mp.HasNull(), unmarshaledMp.HasNull()) - foundVs, foundZvs := unmarshaledMp.NewIterator().Find(0, numElements, vecs) + foundVs, foundZvs, err := unmarshaledMp.NewIterator().Find(0, numElements, vecs) + require.NoError(t, err) for i := 0; i < numElements; i++ { require.Equal(t, originalVs[i], foundVs[i], "Mismatch at index %d for mapped value", i) require.Equal(t, originalZvs[i], foundZvs[i], "Mismatch at index %d for zValue", i) diff --git a/pkg/common/hashmap/iterator.go b/pkg/common/hashmap/iterator.go index 8b99d6405551f..cac9b7eb2e299 100644 --- a/pkg/common/hashmap/iterator.go +++ b/pkg/common/hashmap/iterator.go @@ -17,9 +17,47 @@ package hashmap import ( "unsafe" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/vector" ) +func validateIteratorVectors( + vecs []*vector.Vector, + start int, + count int, +) error { + if len(vecs) == 0 || start < 0 || count < 0 || count > UnitLimit { + return mpool.ErrAllocationAccountInvalid + } + for _, vec := range vecs { + if vec == nil || start > vec.Length() || count > vec.Length()-start { + return mpool.ErrAllocationAccountInvalid + } + } + return nil +} + +func hasGroupingInRange(vecs []*vector.Vector, start, count int) bool { + end := uint64(start + count) + for _, vec := range vecs { + if vec != nil && vec.GetGrouping().GetBitmap().CountRange( + uint64(start), end, + ) > 0 { + return true + } + } + return false +} + +func rowHasGrouping(vecs []*vector.Vector, row int) bool { + for _, vec := range vecs { + if vec.GetGrouping().Contains(uint64(row)) { + return true + } + } + return false +} + // MaxStrIteratorCapacity limits how many bytes of backing storage we keep when // reusing a string iterator. Avoids retaining oversized buffers after handling // very large strings. @@ -31,7 +69,12 @@ func IteratorChangeOwner(itr Iterator, m HashMap) { return } it := itr.(*strHashmapIterator) - it.mp = m.(*StrHashMap) + next := m.(*StrHashMap) + if it.mp != nil && + it.mp.iteratorAllocation != next.iteratorAllocation { + it.releaseScratch() + } + it.mp = next } // IteratorClearOwner detaches the iterator from its hashmap to allow the old @@ -41,6 +84,7 @@ func IteratorClearOwner(itr Iterator) { case *intHashMapIterator: it.mp = nil case *strHashmapIterator: + it.releaseAccountedScratch() it.mp = nil } } @@ -52,28 +96,58 @@ func StrIteratorCapacity(itr Iterator) int { if !ok || it == nil { return 0 } - total := 0 - for i := range it.keys { - total += cap(it.keys[i]) + return cap(it.keyBuffer) +} + +func (itr *strHashmapIterator) releaseScratch() { + if itr == nil { + return + } + if cap(itr.keyBuffer) > 0 && itr.mp != nil && itr.mp.mp != nil && + itr.mp.iteratorAllocation != nil { + itr.mp.mp.Free(itr.keyBuffer) } - return total + itr.keyBuffer = nil + clear(itr.keys) } -func (itr *strHashmapIterator) Find(start, count int, vecs []*vector.Vector) ([]uint64, []int64) { - for i := 0; i < count; i++ { - itr.keys[i] = itr.keys[i][:0] +func (itr *strHashmapIterator) releaseAccountedScratch() { + if itr == nil || cap(itr.keyBuffer) == 0 || itr.mp == nil || + itr.mp.iteratorAllocation == nil { + return + } + itr.mp.mp.Free(itr.keyBuffer) + itr.keyBuffer = nil + clear(itr.keys) +} + +func (itr *strHashmapIterator) Find(start, count int, vecs []*vector.Vector) ([]uint64, []int64, error) { + if err := itr.prepareHashKeys(vecs, start, count); err != nil { + return nil, nil, err } copy(itr.zValues[:count], OneInt64s[:count]) copy(itr.values[:count], zeroUint64[:count]) itr.encodeHashKeys(vecs, start, count) itr.mp.hashMap.FindStringBatch(itr.strHashStates, itr.keys[:count], itr.values) - return itr.values[:count], itr.zValues[:count] + if !itr.mp.hasNull && !itr.mp.groupingAware && + hasGroupingInRange(vecs, start, count) { + for i := 0; i < count; i++ { + if rowHasGrouping(vecs, start+i) { + itr.values[i] = 0 + itr.zValues[i] = 0 + } + } + } + return itr.values[:count], itr.zValues[:count], nil } // Insert a row from multiple columns into the hashmap, return true if it is new, otherwise false func (itr *strHashmapIterator) DetectDup(vecs []*vector.Vector, row int) (bool, error) { keys := itr.keys defer func() { keys[0] = keys[0][:0] }() + if err := itr.prepareHashKeys(vecs, row, 1); err != nil { + return false, err + } itr.encodeHashKeys(vecs, row, 1) if err := itr.mp.hashMap.InsertStringBatch(itr.strHashStates, keys[:1], itr.values[:1]); err != nil { return false, err @@ -88,6 +162,9 @@ func (itr *strHashmapIterator) DetectDup(vecs []*vector.Vector, row int) (bool, func (itr *strHashmapIterator) Insert(start, count int, vecs []*vector.Vector) ([]uint64, []int64, error) { var err error + if err = itr.prepareHashKeys(vecs, start, count); err != nil { + return nil, nil, err + } defer func() { for i := 0; i < count; i++ { itr.keys[i] = itr.keys[i][:0] @@ -110,10 +187,16 @@ func (itr *strHashmapIterator) Insert(start, count int, vecs []*vector.Vector) ( return vs, zvs, err } -func (itr *intHashMapIterator) Find(start, count int, vecs []*vector.Vector) ([]uint64, []int64) { +func (itr *intHashMapIterator) Find(start, count int, vecs []*vector.Vector) ([]uint64, []int64, error) { + if itr == nil || itr.mp == nil { + return nil, nil, mpool.ErrAllocationAccountInvalid + } + if err := validateIteratorVectors(vecs, start, count); err != nil { + return nil, nil, err + } itr.ensureCapacity(count) if count == 0 { - return itr.values, itr.zValues + return itr.values, itr.zValues, nil } for i := 0; i < count; i++ { itr.keys[i] = 0 @@ -124,15 +207,37 @@ func (itr *intHashMapIterator) Find(start, count int, vecs []*vector.Vector) ([] itr.encodeHashKeys(vecs, start, count) copy(itr.hashes[:count], zeroUint64[:count]) itr.mp.hashMap.FindBatch(count, itr.hashes[:count], unsafe.Pointer(&itr.keys[0]), itr.values[:count]) - return itr.values[:count], itr.zValues[:count] + if hasGroupingInRange(vecs, start, count) { + for i := 0; i < count; i++ { + if rowHasGrouping(vecs, start+i) { + itr.values[i] = 0 + itr.zValues[i] = 0 + } + } + } + return itr.values[:count], itr.zValues[:count], nil } func (itr *intHashMapIterator) DetectDup(vecs []*vector.Vector, row int) (bool, error) { - panic("not implemented yet!!!") + if itr == nil || itr.mp == nil { + return false, mpool.ErrAllocationAccountInvalid + } + before := itr.mp.rows + values, zValues, err := itr.Insert(row, 1, vecs) + if err != nil { + return false, err + } + return zValues[0] != 0 && values[0] > before, nil } func (itr *intHashMapIterator) Insert(start, count int, vecs []*vector.Vector) ([]uint64, []int64, error) { var err error + if itr == nil || itr.mp == nil { + return nil, nil, mpool.ErrAllocationAccountInvalid + } + if err = validateIteratorVectors(vecs, start, count); err != nil { + return nil, nil, err + } itr.ensureCapacity(count) if count == 0 { return itr.values, itr.zValues, nil diff --git a/pkg/common/hashmap/keycodec/keycodec.go b/pkg/common/hashmap/keycodec/keycodec.go index b259784429efc..0f72f997e3e83 100644 --- a/pkg/common/hashmap/keycodec/keycodec.go +++ b/pkg/common/hashmap/keycodec/keycodec.go @@ -25,6 +25,23 @@ import ( "github.com/matrixorigin/matrixone/pkg/container/vector" ) +// ValidVectors verifies the row-shape contract required by both resident hash +// maps and spill partitioning. Hashing a short or nil key must never silently +// leave a suffix at its previous seed value. +func ValidVectors(vecs []*vector.Vector, rows int) bool { + if rows < 0 || len(vecs) == 0 { + return false + } + for _, vec := range vecs { + if vec == nil || vec.Length() != rows { + return false + } + } + return true +} + +var groupingColumnHash = xxhash.Sum64([]byte{2}) + // Float32Codec holds the SQL comparison normalization for one FLOAT32 type. // Construct it once per vector so scale processing is not repeated per row. type Float32Codec struct { @@ -194,6 +211,10 @@ func ComputeXXHash(keyVecs []*vector.Vector, hashValues []uint64, seed uint64) { } for _, vec := range keyVecs { + if vec.GetGrouping().GetBitmap().CountRange(0, uint64(rowCount)) > 0 { + computeGroupingXXHash(vec, hashValues) + continue + } switch vec.GetType().Oid { case types.T_float32: computeFloat32XXHash(vec, hashValues) @@ -273,6 +294,41 @@ func computeFloat32XXHash(vec *vector.Vector, hashValues []uint64) { } } +func computeGroupingXXHash(vec *vector.Vector, hashValues []uint64) { + rowCount := len(hashValues) + grouping := vec.GetGrouping() + nulls := vec.GetNulls() + for i := 0; i < rowCount; i++ { + if grouping.Contains(uint64(i)) { + hashValues[i] = HashCombine(hashValues[i], groupingColumnHash) + continue + } + if vec.IsConstNull() || nulls.Contains(uint64(i)) { + hashValues[i] = HashCombine(hashValues[i], 0) + continue + } + row := i + if vec.IsConst() { + row = 0 + } + switch vec.GetType().Oid { + case types.T_float32: + values := vector.MustFixedColNoTypeCheck[float32](vec) + value := NewFloat32Codec(vec.GetType().Scale).CanonicalBytes(values[row]) + hashValues[i] = HashCombine(hashValues[i], xxhash.Sum64(value[:])) + continue + case types.T_float64: + values := vector.MustFixedColNoTypeCheck[float64](vec) + value := CanonicalFloat64Bytes(values[row]) + hashValues[i] = HashCombine(hashValues[i], xxhash.Sum64(value[:])) + continue + } + hashValues[i] = HashCombine( + hashValues[i], xxhash.Sum64(vec.GetRawBytesAt(row)), + ) + } +} + func computeFloat64XXHash(vec *vector.Vector, hashValues []uint64) { rowCount := len(hashValues) if vec.IsConst() { diff --git a/pkg/common/hashmap/keycodec/keycodec_test.go b/pkg/common/hashmap/keycodec/keycodec_test.go index 8fd2fa380c822..79cdcbf784196 100644 --- a/pkg/common/hashmap/keycodec/keycodec_test.go +++ b/pkg/common/hashmap/keycodec/keycodec_test.go @@ -178,3 +178,83 @@ func TestComputeXXHashCompositeScaledFloat32Contract(t *testing.T) { require.NotEqual(t, hashes[0], hashes[2], "the FLOAT32 codec must preserve prior column hash state") require.NotEqual(t, hashes[0], hashes[3], "a distinct canonical FLOAT32 value must change the composite hash") } + +func TestComputeXXHashCanonicalizesGroupingRows(t *testing.T) { + mp := mpool.MustNewZero() + defer func() { require.Zero(t, mp.CurrNB()) }() + + for _, test := range []struct { + name string + left any + right any + newVec func(any) *vector.Vector + }{ + { + name: "fixed", + left: []int64{11, 22, 33}, + right: []int64{101, 22, 303}, + newVec: func(values any) *vector.Vector { + vec := vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixedList( + vec, values.([]int64), nil, mp, + )) + return vec + }, + }, + { + name: "float64", + left: []float64{11, 22, 33}, + right: []float64{101, 22, 303}, + newVec: func(values any) *vector.Vector { + vec := vector.NewVec(types.T_float64.ToType()) + require.NoError(t, vector.AppendFixedList( + vec, values.([]float64), nil, mp, + )) + return vec + }, + }, + } { + t.Run(test.name, func(t *testing.T) { + left := test.newVec(test.left) + right := test.newVec(test.right) + ordinary := test.newVec(test.left) + defer left.Free(mp) + defer right.Free(mp) + defer ordinary.Free(mp) + left.GetGrouping().AddRange(0, 3) + right.GetGrouping().AddRange(0, 3) + + leftHashes := make([]uint64, 3) + rightHashes := make([]uint64, 3) + ComputeXXHash([]*vector.Vector{left}, leftHashes, 17) + ComputeXXHash([]*vector.Vector{right}, rightHashes, 17) + + require.Equal(t, leftHashes, rightHashes) + + left.GetGrouping().Reset() + left.GetGrouping().Add(0) + ComputeXXHash([]*vector.Vector{left}, leftHashes, 17) + ComputeXXHash([]*vector.Vector{ordinary}, rightHashes, 17) + require.NotEqual(t, rightHashes[0], leftHashes[0]) + require.Equal(t, rightHashes[1:], leftHashes[1:]) + }) + } +} + +func TestComputeXXHashDoesNotTreatStaleGroupingAsFull(t *testing.T) { + mp := mpool.MustNewZero() + left := vector.NewVec(types.T_int64.ToType()) + right := vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixed(left, int64(11), false, mp)) + require.NoError(t, vector.AppendFixed(right, int64(22), false, mp)) + left.GetGrouping().Add(5) + right.GetGrouping().Add(5) + leftHash := []uint64{0} + rightHash := []uint64{0} + ComputeXXHash([]*vector.Vector{left}, leftHash, 17) + ComputeXXHash([]*vector.Vector{right}, rightHash, 17) + require.NotEqual(t, leftHash, rightHash) + left.Free(mp) + right.Free(mp) + require.Zero(t, mp.CurrNB()) +} diff --git a/pkg/common/hashmap/strhashmap.go b/pkg/common/hashmap/strhashmap.go index 6fa09dbf1ba56..b4ddb84f68a62 100644 --- a/pkg/common/hashmap/strhashmap.go +++ b/pkg/common/hashmap/strhashmap.go @@ -45,14 +45,30 @@ func NewStrHashMapWithAllocation( hasNull bool, memPool *mpool.MPool, allocation *hashtable.AllocationAccountSelection, +) (*StrHashMap, error) { + return NewStrHashMapWithAllocations( + hasNull, + memPool, + allocation, + nil, + ) +} + +func NewStrHashMapWithAllocations( + hasNull bool, + memPool *mpool.MPool, + allocation *hashtable.AllocationAccountSelection, + iteratorAllocation *IteratorAllocation, ) (*StrHashMap, error) { mp := &hashtable.StringHashMap{} if err := mp.InitWithAllocation(memPool, allocation); err != nil { return nil, err } return &StrHashMap{ - hashMap: mp, - hasNull: hasNull, + hashMap: mp, + hasNull: hasNull, + mp: memPool, + iteratorAllocation: iteratorAllocation, }, nil } @@ -66,10 +82,188 @@ func (m *StrHashMap) NewIterator() Iterator { } } +func (itr *strHashmapIterator) prepareHashKeys( + vecs []*vector.Vector, + start int, + count int, +) error { + if itr == nil || itr.mp == nil || start < 0 || count < 0 || + count > UnitLimit { + return mpool.ErrAllocationAccountInvalid + } + if err := validateIteratorVectors(vecs, start, count); err != nil { + return err + } + for i := 0; i < count; i++ { + itr.keyLengths[i] = 0 + } + const maxInt = int(^uint(0) >> 1) + add := func(row int, size int) error { + if size < 0 || itr.keyLengths[row] > maxInt-size { + return mpool.ErrAllocationAccountInvalid + } + itr.keyLengths[row] += size + return nil + } + for _, vec := range vecs { + withDomain := itr.mp.hasNull || itr.mp.groupingAware + prefix := 0 + if withDomain { + prefix = 1 + } + if vec.IsGrouping() { + for i := 0; i < count; i++ { + if err := add(i, 1); err != nil { + return err + } + } + continue + } + if vec.IsConstNull() { + if itr.mp.hasNull { + for i := 0; i < count; i++ { + if err := add(i, 1); err != nil { + return err + } + } + } + continue + } + + // Most join keys are flat and non-null. Size them from the physical + // representation directly, avoiding repeated type/null/const dispatch + // before the encoder's required value pass. + hasGrouping := withDomain && vec.HasGrouping() + if !hasGrouping && !vec.GetNulls().Any() { + if vec.GetType().IsFixedLen() { + size := prefix + vec.GetType().TypeSize() + for i := 0; i < count; i++ { + if err := add(i, size); err != nil { + return err + } + } + continue + } + if vec.IsConst() { + size := prefix + 4 + len(vec.GetBytesAt(0)) + for i := 0; i < count; i++ { + if err := add(i, size); err != nil { + return err + } + } + continue + } + values, area := vector.MustVarlenaRawData(vec) + for i := 0; i < count; i++ { + value := values[start+i].ByteSlice() + if area != nil { + value = values[start+i].GetByteSlice(area) + } + if err := add(i, prefix+4+len(value)); err != nil { + return err + } + } + continue + } + + fixed := vec.GetType().IsFixedLen() + for i := 0; i < count; i++ { + row := start + i + if withDomain && vec.GetGrouping().Contains(uint64(row)) { + if err := add(i, 1); err != nil { + return err + } + continue + } + if vec.GetNulls().Contains(uint64(row)) { + if itr.mp.hasNull { + if err := add(i, 1); err != nil { + return err + } + } + continue + } + if fixed { + if err := add(i, prefix+vec.GetType().TypeSize()); err != nil { + return err + } + continue + } + valueRow := row + if vec.IsConst() { + valueRow = 0 + } + if err := add(i, prefix+4+len(vec.GetBytesAt(valueRow))); err != nil { + return err + } + } + } + + total := 0 + for i := 0; i < count; i++ { + if itr.keyLengths[i] < 16 { + itr.keyLengths[i] = 16 + } + if total > maxInt-itr.keyLengths[i] { + return mpool.ErrAllocationAccountInvalid + } + total += itr.keyLengths[i] + } + if cap(itr.keyBuffer) < total { + if allocation := itr.mp.iteratorAllocation; allocation != nil { + capacity, ok := mpool.GrowCapacity( + int64(cap(itr.keyBuffer)), int64(total), + ) + if !ok || int64(int(capacity)) != capacity { + return mpool.ErrAllocationAllocatorLimit + } + var next []byte + var err error + if cap(itr.keyBuffer) > 0 { + next, err = itr.mp.mp.Grow(itr.keyBuffer, int(capacity), true) + } else { + next, err = itr.mp.mp.AllocAccounted( + int(capacity), + allocation.account, + allocation.owner, + allocation.site, + ) + } + if err != nil { + return err + } + itr.keyBuffer = next + } else { + itr.keyBuffer = make([]byte, total) + } + } + itr.keyBuffer = itr.keyBuffer[:total] + storage := itr.keyBuffer + offset := 0 + for i := 0; i < count; i++ { + end := offset + itr.keyLengths[i] + itr.keys[i] = storage[offset:offset:end] + offset = end + } + return nil +} + func (m *StrHashMap) HasNull() bool { return m.hasNull } +// SetGroupingAware selects a collision-free key domain for maps that may see +// GROUPING rows. It must be set before the +// first insert. Ordinary columns receive a 0 domain byte and GROUPING columns +// receive 2, so no raw fixed-width value can alias the sentinel. +func (m *StrHashMap) SetGroupingAware() error { + if m == nil || m.rows != 0 { + return mpool.ErrAllocationAccountInvalid + } + m.groupingAware = true + return nil +} + func (m *StrHashMap) Free() { m.hashMap.Free() } @@ -104,6 +298,10 @@ func (m *StrHashMap) Size() int64 { func (itr *strHashmapIterator) encodeHashKeys(vecs []*vector.Vector, start, count int) { for _, vec := range vecs { + if itr.mp.groupingAware || itr.mp.hasNull { + fillGroupingAwareStr(itr, vec, count, start) + continue + } if vec.GetType().IsFixedLen() { switch vec.GetType().Oid { case types.T_float32: @@ -145,7 +343,6 @@ func fillFloat32GroupStr(itr *strHashmapIterator, vec *vector.Vector, n, start i } return } - values := vector.MustFixedColNoTypeCheck[float32](vec) codec := keycodec.NewFloat32Codec(vec.GetType().Scale) if vec.IsConst() { @@ -198,6 +395,78 @@ func fillFloat32GroupStr(itr *strHashmapIterator, vec *vector.Vector, n, start i } } +func fillGroupingAwareStr( + itr *strHashmapIterator, + vec *vector.Vector, + n int, + start int, +) { + keys := itr.keys + if vec.IsGrouping() { + for i := 0; i < n; i++ { + keys[i] = append(keys[i], byte(2)) + } + return + } + if vec.IsConstNull() { + for i := 0; i < n; i++ { + row := start + i + if vec.GetGrouping().Contains(uint64(row)) { + keys[i] = append(keys[i], byte(2)) + } else if itr.mp.hasNull { + keys[i] = append(keys[i], byte(1)) + } else { + itr.zValues[i] = 0 + } + } + return + } + float32Codec := keycodec.NewFloat32Codec(vec.GetType().Scale) + for i := 0; i < n; i++ { + row := start + i + if vec.GetGrouping().Contains(uint64(row)) { + keys[i] = append(keys[i], byte(2)) + continue + } + if vec.GetNulls().Contains(uint64(row)) { + if itr.mp.hasNull { + keys[i] = append(keys[i], byte(1)) + } else { + itr.zValues[i] = 0 + } + continue + } + keys[i] = append(keys[i], byte(0)) + valueRow := row + if vec.IsConst() { + valueRow = 0 + } + switch vec.GetType().Oid { + case types.T_float32: + values := vector.MustFixedColNoTypeCheck[float32](vec) + value := float32Codec.CanonicalBytes(values[valueRow]) + keys[i] = append(keys[i], value[:]...) + continue + case types.T_float64: + values := vector.MustFixedColNoTypeCheck[float64](vec) + value := keycodec.CanonicalFloat64Bytes(values[valueRow]) + keys[i] = append(keys[i], value[:]...) + continue + } + if vec.GetType().IsFixedLen() { + size := vec.GetType().TypeSize() + data := vec.GetData() + value := data[valueRow*size : (valueRow+1)*size] + keys[i] = append(keys[i], value...) + continue + } + value := vec.GetBytesAt(valueRow) + length := uint32(len(value)) + keys[i] = append(keys[i], util.UnsafeToBytes(&length)...) + keys[i] = append(keys[i], value...) + } +} + func fillFloat64GroupStr(itr *strHashmapIterator, vec *vector.Vector, n, start int) { keys := itr.keys if vec.IsGrouping() { @@ -273,7 +542,7 @@ func fillFloat64GroupStr(itr *strHashmapIterator, vec *vector.Vector, n, start i func fillStringGroupStrForConstVec(itr *strHashmapIterator, vec *vector.Vector, n int, start int) { keys := itr.keys bytes := vec.GetBytesAt(start) - length := uint16(len(bytes)) + length := uint32(len(bytes)) // can't be const null if itr.mp.hasNull { gsp := vec.GetGrouping() @@ -350,7 +619,7 @@ func fillStringGroupStr(itr *strHashmapIterator, vec *vector.Vector, lenV int, s // this is not null value keys[i] = append(keys[i], 0) // give the length - length := uint16(len(bytes)) + length := uint32(len(bytes)) keys[i] = append(keys[i], util.UnsafeToBytes(&length)...) // append the pure value bytes keys[i] = append(keys[i], bytes...) @@ -367,7 +636,7 @@ func fillStringGroupStr(itr *strHashmapIterator, vec *vector.Vector, lenV int, s // this is not null value keys[i] = append(keys[i], 0) // give the length - length := uint16(len(bytes)) + length := uint32(len(bytes)) keys[i] = append(keys[i], util.UnsafeToBytes(&length)...) // append the pure value bytes keys[i] = append(keys[i], bytes...) @@ -380,7 +649,7 @@ func fillStringGroupStr(itr *strHashmapIterator, vec *vector.Vector, lenV int, s bytes := va[i+start].ByteSlice() // for "a"，"bc" and "ab","c", we need to distinct // give the length - length := uint16(len(bytes)) + length := uint32(len(bytes)) keys[i] = append(keys[i], util.UnsafeToBytes(&length)...) // append the pure value bytes keys[i] = append(keys[i], bytes...) @@ -390,7 +659,7 @@ func fillStringGroupStr(itr *strHashmapIterator, vec *vector.Vector, lenV int, s bytes := va[i+start].GetByteSlice(area) // for "a"，"bc" and "ab","c", we need to distinct // give the length - length := uint16(len(bytes)) + length := uint32(len(bytes)) keys[i] = append(keys[i], util.UnsafeToBytes(&length)...) // append the pure value bytes keys[i] = append(keys[i], bytes...) @@ -416,7 +685,7 @@ func fillStringGroupStr(itr *strHashmapIterator, vec *vector.Vector, lenV int, s // this is not null value keys[i] = append(keys[i], 0) // give the length - length := uint16(len(bytes)) + length := uint32(len(bytes)) keys[i] = append(keys[i], util.UnsafeToBytes(&length)...) // append the pure value bytes keys[i] = append(keys[i], bytes...) @@ -429,7 +698,7 @@ func fillStringGroupStr(itr *strHashmapIterator, vec *vector.Vector, lenV int, s bytes := va[i+start].ByteSlice() // for "a"，"bc" and "ab","c", we need to distinct // give the length - length := uint16(len(bytes)) + length := uint32(len(bytes)) keys[i] = append(keys[i], util.UnsafeToBytes(&length)...) // append the pure value bytes keys[i] = append(keys[i], bytes...) @@ -450,7 +719,7 @@ func fillStringGroupStr(itr *strHashmapIterator, vec *vector.Vector, lenV int, s // this is not null value keys[i] = append(keys[i], 0) // give the length - length := uint16(len(bytes)) + length := uint32(len(bytes)) keys[i] = append(keys[i], util.UnsafeToBytes(&length)...) // append the pure value bytes keys[i] = append(keys[i], bytes...) @@ -463,7 +732,7 @@ func fillStringGroupStr(itr *strHashmapIterator, vec *vector.Vector, lenV int, s bytes := va[i+start].GetByteSlice(area) // for "a"，"bc" and "ab","c", we need to distinct // give the length - length := uint16(len(bytes)) + length := uint32(len(bytes)) keys[i] = append(keys[i], util.UnsafeToBytes(&length)...) // append the pure value bytes keys[i] = append(keys[i], bytes...) @@ -566,15 +835,17 @@ func (m *StrHashMap) UnmarshalBinary(data []byte, mp *mpool.MPool) error { func (m *StrHashMap) WriteTo(w io.Writer) (int64, error) { var n int64 - // Serialize hasNull (1 byte) + // The low two bits retain the key grammar. Historical payloads used only + // bit zero, so 0/1 remain backward-compatible. + flags := byte(0) if m.hasNull { - if _, err := w.Write([]byte{1}); err != nil { - return 0, err - } - } else { - if _, err := w.Write([]byte{0}); err != nil { - return 0, err - } + flags |= 1 + } + if m.groupingAware { + flags |= 2 + } + if _, err := w.Write([]byte{flags}); err != nil { + return 0, err } n++ @@ -606,7 +877,11 @@ func (m *StrHashMap) UnmarshalFrom(r io.Reader, mp *mpool.MPool) (int64, error) return 0, err } n += int64(rn) - m.hasNull = b[0] == 1 + if b[0]&^byte(3) != 0 { + return 0, mpool.ErrAllocationAccountInvalid + } + m.hasNull = b[0]&1 != 0 + m.groupingAware = b[0]&2 != 0 // Deserialize rows rowsData := make([]byte, 8) @@ -615,6 +890,7 @@ func (m *StrHashMap) UnmarshalFrom(r io.Reader, mp *mpool.MPool) (int64, error) } n += int64(rn) m.rows = types.DecodeUint64(rowsData) + m.mp = mp // Deserialize the underlying StringHashMap m.hashMap = &hashtable.StringHashMap{} diff --git a/pkg/common/hashmap/strhashmap_test.go b/pkg/common/hashmap/strhashmap_test.go index 1d9609893878b..25b74b3610d85 100644 --- a/pkg/common/hashmap/strhashmap_test.go +++ b/pkg/common/hashmap/strhashmap_test.go @@ -141,7 +141,8 @@ func runIntHashMapFloat32CompositeFloatLastCase( require.NoError(t, err) require.Equal(t, []uint64{1}, values) require.Equal(t, []int64{1}, zValues) - values, zValues = itr.Find(0, 1, probeKeys) + values, zValues, err = itr.Find(0, 1, probeKeys) + require.NoError(t, err) require.Equal(t, []uint64{1}, values) require.Equal(t, []int64{1}, zValues) } @@ -236,6 +237,247 @@ func runFloatHashMapContract(t *testing.T, makeVectors floatHashMapVectorFactory } } +func TestStringHashMapCanonicalizesFullyGroupedKeys(t *testing.T) { + mp := mpool.MustNewZero() + hashMap, err := NewStrHashMap(false, mp) + require.NoError(t, err) + require.NoError(t, hashMap.SetGroupingAware()) + defer hashMap.Free() + + build := vector.NewVec(types.T_int32.ToType()) + probe := vector.NewVec(types.T_int32.ToType()) + require.NoError(t, vector.AppendFixed(build, int32(111), false, mp)) + require.NoError(t, vector.AppendFixed(probe, int32(222), false, mp)) + build.GetGrouping().Add(0) + probe.GetGrouping().Add(0) + defer build.Free(mp) + defer probe.Free(mp) + + iterator := hashMap.NewIterator() + values, zValues, err := iterator.Insert(0, 1, []*vector.Vector{build}) + require.NoError(t, err) + require.Equal(t, []uint64{1}, values) + require.Equal(t, []int64{1}, zValues) + values, zValues, err = iterator.Find(0, 1, []*vector.Vector{probe}) + require.NoError(t, err) + require.Equal(t, []uint64{1}, values) + require.Equal(t, []int64{1}, zValues) +} + +func TestGroupingAwareStringHashMapSeparatesRawSentinelBytes(t *testing.T) { + mp := mpool.MustNewZero() + hashMap, err := NewStrHashMap(false, mp) + require.NoError(t, err) + require.NoError(t, hashMap.SetGroupingAware()) + defer hashMap.Free() + + raw := vector.NewVec(types.T_uint8.ToType()) + require.NoError(t, vector.AppendFixed(raw, uint8(2), false, mp)) + grouping := vector.NewRollupConst(types.T_uint8.ToType(), 1, mp) + defer raw.Free(mp) + defer grouping.Free(mp) + + iterator := hashMap.NewIterator() + values, zValues, err := iterator.Insert(0, 1, []*vector.Vector{raw}) + require.NoError(t, err) + require.Equal(t, []uint64{1}, values) + require.Equal(t, []int64{1}, zValues) + values, zValues, err = iterator.Insert(0, 1, []*vector.Vector{grouping}) + require.NoError(t, err) + require.Equal(t, []uint64{2}, values) + require.Equal(t, []int64{1}, zValues) + require.Equal(t, uint64(2), hashMap.GroupCount()) +} + +func TestNullableStringHashMapTreatsGroupingRowsAsSentinel(t *testing.T) { + mp := mpool.MustNewZero() + hashMap, err := NewStrHashMap(true, mp) + require.NoError(t, err) + defer hashMap.Free() + + partial := vector.NewVec(types.T_int32.ToType()) + require.NoError(t, vector.AppendFixedList( + partial, + []int32{7, 7, 9}, + nil, + mp, + )) + partial.GetGrouping().Add(0) + partial.GetGrouping().Add(2) + partial.GetNulls().Add(2) + defer partial.Free(mp) + + iterator := hashMap.NewIterator() + values, zValues, err := iterator.Insert(0, 3, []*vector.Vector{partial}) + require.NoError(t, err) + require.Equal(t, []uint64{1, 2, 1}, values) + require.Equal(t, []int64{1, 1, 1}, zValues) + + fullGrouping := vector.NewRollupConst(types.T_int32.ToType(), 1, mp) + defer fullGrouping.Free(mp) + values, zValues, err = iterator.Insert(0, 1, []*vector.Vector{fullGrouping}) + require.NoError(t, err) + require.Equal(t, []uint64{1}, values) + require.Equal(t, []int64{1}, zValues) + require.Equal(t, uint64(2), hashMap.GroupCount()) +} + +func TestGroupingAwareStringHashMapConstNullUsesRowwiseGrouping(t *testing.T) { + mp := mpool.MustNewZero() + for _, hasNull := range []bool{false, true} { + hashMap, err := NewStrHashMap(hasNull, mp) + require.NoError(t, err) + require.NoError(t, hashMap.SetGroupingAware()) + vec := vector.NewConstNull(types.T_int32.ToType(), 2, mp) + vec.GetGrouping().Add(0) + + values, zValues, err := hashMap.NewIterator().Insert( + 0, 2, []*vector.Vector{vec}, + ) + require.NoError(t, err) + if hasNull { + require.Equal(t, []uint64{1, 2}, values) + require.Equal(t, []int64{1, 1}, zValues) + } else { + require.Equal(t, []uint64{1, 0}, values) + require.Equal(t, []int64{1, 0}, zValues) + } + + vec.Free(mp) + hashMap.Free() + } + require.Zero(t, mp.CurrNB()) +} + +func TestStringHashMapProbeGroupingDoesNotMatchRawKey(t *testing.T) { + mp := mpool.MustNewZero() + hashMap, err := NewStrHashMap(false, mp) + require.NoError(t, err) + defer hashMap.Free() + raw := vector.NewVec(types.T_uint8.ToType()) + require.NoError(t, vector.AppendFixed(raw, uint8(2), false, mp)) + grouping := vector.NewRollupConst(types.T_uint8.ToType(), 1, mp) + defer raw.Free(mp) + defer grouping.Free(mp) + iterator := hashMap.NewIterator() + _, _, err = iterator.Insert(0, 1, []*vector.Vector{raw}) + require.NoError(t, err) + values, zValues, err := iterator.Find(0, 1, []*vector.Vector{grouping}) + require.NoError(t, err) + require.Equal(t, []uint64{0}, values) + require.Equal(t, []int64{0}, zValues) +} + +func TestGroupingAwareStringHashMapRoundTripRetainsKeyGrammar(t *testing.T) { + mp := mpool.MustNewZero() + original, err := NewStrHashMap(false, mp) + require.NoError(t, err) + require.NoError(t, original.SetGroupingAware()) + raw := vector.NewVec(types.T_uint8.ToType()) + require.NoError(t, vector.AppendFixed(raw, uint8(2), false, mp)) + defer raw.Free(mp) + _, _, err = original.NewIterator().Insert(0, 1, []*vector.Vector{raw}) + require.NoError(t, err) + require.ErrorIs(t, original.SetGroupingAware(), mpool.ErrAllocationAccountInvalid) + + encoded, err := original.MarshalBinary() + require.NoError(t, err) + original.Free() + restored := &StrHashMap{} + require.NoError(t, restored.UnmarshalBinary(encoded, mp)) + values, zValues, err := restored.NewIterator().Find(0, 1, []*vector.Vector{raw}) + require.NoError(t, err) + require.Equal(t, []uint64{1}, values) + require.Equal(t, []int64{1}, zValues) + restored.Free() +} + +func TestStringHashMapWideVarlenaLengthsDoNotCollide(t *testing.T) { + mp := mpool.MustNewZero() + hashMap, err := NewStrHashMap(false, mp) + require.NoError(t, err) + defer hashMap.Free() + + large := make([]byte, 1<<16) + left := vector.NewVec(types.T_binary.ToType()) + right := vector.NewVec(types.T_binary.ToType()) + require.NoError(t, vector.AppendBytesList(left, [][]byte{large, nil}, nil, mp)) + require.NoError(t, vector.AppendBytesList(right, [][]byte{nil, large}, nil, mp)) + defer left.Free(mp) + defer right.Free(mp) + + values, zValues, err := hashMap.NewIterator().Insert( + 0, 2, []*vector.Vector{left, right}, + ) + require.NoError(t, err) + require.Equal(t, []uint64{1, 2}, values) + require.Equal(t, []int64{1, 1}, zValues) +} + +func TestStringHashMapConstAndFlatVarlenaUseSameKey(t *testing.T) { + mp := mpool.MustNewZero() + for _, buildConst := range []bool{false, true} { + hashMap, err := NewStrHashMap(true, mp) + require.NoError(t, err) + build := vector.NewVec(types.T_varchar.ToType()) + probe := vector.NewVec(types.T_varchar.ToType()) + if buildConst { + build, err = vector.NewConstBytes(types.T_varchar.ToType(), []byte("abc"), 1, mp) + require.NoError(t, err) + require.NoError(t, vector.AppendBytes(probe, []byte("abc"), false, mp)) + } else { + require.NoError(t, vector.AppendBytes(build, []byte("abc"), false, mp)) + probe, err = vector.NewConstBytes(types.T_varchar.ToType(), []byte("abc"), 1, mp) + require.NoError(t, err) + } + iterator := hashMap.NewIterator() + _, _, err = iterator.Insert(0, 1, []*vector.Vector{build}) + require.NoError(t, err) + values, zValues, err := iterator.Find(0, 1, []*vector.Vector{probe}) + require.NoError(t, err) + require.Equal(t, []uint64{1}, values) + require.Equal(t, []int64{1}, zValues) + build.Free(mp) + probe.Free(mp) + hashMap.Free() + } + require.Zero(t, mp.CurrNB()) +} + +func TestHashMapIteratorsRejectMalformedRowShapes(t *testing.T) { + mp := mpool.MustNewZero() + for _, makeIterator := range []func() (HashMap, Iterator){ + func() (HashMap, Iterator) { + m, err := NewIntHashMap(false, mp) + require.NoError(t, err) + return m, m.NewIterator() + }, + func() (HashMap, Iterator) { + m, err := NewStrHashMap(false, mp) + require.NoError(t, err) + return m, m.NewIterator() + }, + } { + m, iterator := makeIterator() + short, err := vector.NewConstFixed(types.T_int32.ToType(), int32(1), 1, mp) + require.NoError(t, err) + for _, vecs := range [][]*vector.Vector{ + nil, + []*vector.Vector{nil}, + []*vector.Vector{short}, + } { + _, _, err = iterator.Insert(0, 2, vecs) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + _, _, err = iterator.Find(1, 1, vecs) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + } + require.Zero(t, m.GroupCount()) + short.Free(mp) + m.Free() + } + require.Zero(t, mp.CurrNB()) +} + func runFloatHashMapShape( t *testing.T, composite bool, @@ -280,7 +522,8 @@ func runFloatHashMapShape( require.Equal(t, shape.wantValues, liveHashMapValues(values, zValues)) require.Equal(t, shape.wantGroups, hashMap.GroupCount()) - values, zValues = itr.Find(shape.start, shape.count, probe) + values, zValues, err = itr.Find(shape.start, shape.count, probe) + require.NoError(t, err) require.Equal(t, shape.wantZValues, zValues) require.Equal(t, shape.wantValues, liveHashMapValues(values, zValues)) require.Equal(t, shape.wantGroups, hashMap.GroupCount()) @@ -454,7 +697,8 @@ func TestIterator(t *testing.T) { vs, _, err := itr.Insert(0, Rows, vecs) require.NoError(t, err) require.Equal(t, []uint64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, vs[:Rows]) - vs, _ = itr.Find(0, Rows, vecs) + vs, _, err = itr.Find(0, Rows, vecs) + require.NoError(t, err) require.Equal(t, []uint64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, vs[:Rows]) for _, vec := range vecs { vec.Free(m) @@ -479,7 +723,8 @@ func TestIterator(t *testing.T) { vs, _, err := itr.Insert(0, Rows, vecs) require.NoError(t, err) require.Equal(t, []uint64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, vs[:Rows]) - vs, _ = itr.Find(0, Rows, vecs) + vs, _, err = itr.Find(0, Rows, vecs) + require.NoError(t, err) require.Equal(t, []uint64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, vs[:Rows]) for _, vec := range vecs { vec.Free(m) @@ -504,7 +749,8 @@ func TestIterator(t *testing.T) { vs, _, err := itr.Insert(0, Rows, vecs) require.NoError(t, err) require.Equal(t, []uint64{1, 2, 1, 3, 1, 4, 1, 5, 1, 6}, vs[:Rows]) - vs, _ = itr.Find(0, Rows, vecs) + vs, _, err = itr.Find(0, Rows, vecs) + require.NoError(t, err) require.Equal(t, []uint64{1, 2, 1, 3, 1, 4, 1, 5, 1, 6}, vs[:Rows]) for _, vec := range vecs { vec.Free(m) @@ -835,7 +1081,8 @@ func TestStrHashMap_MarshalUnmarshal(t *testing.T) { require.Equal(t, expectedGroupCount, unmarshaledMp.GroupCount()) require.Equal(t, mp.HasNull(), unmarshaledMp.HasNull()) - foundVs, _ := unmarshaledMp.NewIterator().Find(0, rowCount, vecs) + foundVs, _, err := unmarshaledMp.NewIterator().Find(0, rowCount, vecs) + require.NoError(t, err) require.Equal(t, expectedMappedValue, foundVs) }) @@ -872,7 +1119,8 @@ func TestStrHashMap_MarshalUnmarshal(t *testing.T) { require.Equal(t, expectedGroupCount, unmarshaledMp.GroupCount()) require.Equal(t, mp.HasNull(), unmarshaledMp.HasNull()) - foundVs, foundZvs := unmarshaledMp.NewIterator().Find(0, numElements, vecs) + foundVs, foundZvs, err := unmarshaledMp.NewIterator().Find(0, numElements, vecs) + require.NoError(t, err) for i := 0; i < numElements; i++ { require.Equal(t, originalVs[i], foundVs[i], "Mismatch at index %d for mapped value", i) require.Equal(t, originalZvs[i], foundZvs[i], "Mismatch at index %d for zValue", i) diff --git a/pkg/common/hashmap/types.go b/pkg/common/hashmap/types.go index f84800bb74237..190eaaa17dff7 100644 --- a/pkg/common/hashmap/types.go +++ b/pkg/common/hashmap/types.go @@ -77,16 +77,46 @@ type Iterator interface { // Find vecs[start, start+count) in hashmap // vs : the number of rows corresponding to each value in the hash table (start with 1, and 0 means not found.) // zvs : if zvs[i] is 0 indicates the presence null, 1 indicates the absence of a null. - Find(start, count int, vecs []*vector.Vector) (vs []uint64, zvs []int64) + Find(start, count int, vecs []*vector.Vector) (vs []uint64, zvs []int64, err error) +} + +// IteratorAllocation selects exact physical provenance for data-scaled hash +// key encoding scratch. It is immutable and shared by iterators created from +// one map generation. +type IteratorAllocation struct { + account *mpool.AllocationAccount + owner mpool.AllocationOwner + site mpool.AllocationSite +} + +func NewIteratorAllocation( + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, + site mpool.AllocationSite, +) (*IteratorAllocation, error) { + allocation := &IteratorAllocation{ + account: account, + owner: owner, + site: site, + } + if account == nil || account.Handle() == 0 || + owner < mpool.AllocationOwnerMin || owner > mpool.AllocationOwnerMax || + site < mpool.AllocationSiteMin { + return nil, mpool.ErrAllocationAccountInvalid + } + return allocation, nil } // StrHashMap key is []byte, value is an uint64 value (starting from 1) // // each time a new key is inserted, the hashtable returns a last-value+1 or, if the old key is inserted, the value corresponding to that key type StrHashMap struct { - hasNull bool - rows uint64 - hashMap *hashtable.StringHashMap + hasNull bool + groupingAware bool + rows uint64 + hashMap *hashtable.StringHashMap + mp *mpool.MPool + iteratorAllocation *IteratorAllocation } // IntHashMap key is int64, value is an uint64 (start from 1) @@ -99,9 +129,11 @@ type IntHashMap struct { } type strHashmapIterator struct { - mp *StrHashMap - keys [][]byte - values []uint64 + mp *StrHashMap + keys [][]byte + values []uint64 + keyBuffer []byte + keyLengths [UnitLimit]int // zValues, 0 indicates the presence null, 1 indicates the absence of a null zValues []int64 strHashStates [][3]uint64 diff --git a/pkg/common/mpool/accounted_buffer.go b/pkg/common/mpool/accounted_buffer.go index ee59f29eb7314..c7a1f00550999 100644 --- a/pkg/common/mpool/accounted_buffer.go +++ b/pkg/common/mpool/accounted_buffer.go @@ -15,6 +15,7 @@ package mpool import ( + "encoding/binary" "math" ) @@ -158,6 +159,68 @@ func (b *AccountedBuffer) WriteString(value string) (int, error) { return len(value), nil } +func (b *AccountedBuffer) appendSpace(length int) ([]byte, error) { + if b == nil || length < 0 || length > math.MaxInt-len(b.data) { + return nil, ErrAllocationAccountInvalid + } + start := len(b.data) + if err := b.Resize(start + length); err != nil { + return nil, err + } + return b.data[start:], nil +} + +func (b *AccountedBuffer) WriteByte(value byte) error { + dst, err := b.appendSpace(1) + if err != nil { + return err + } + dst[0] = value + return nil +} + +func (b *AccountedBuffer) WriteUint32(value uint32) error { + dst, err := b.appendSpace(4) + if err != nil { + return err + } + binary.NativeEndian.PutUint32(dst, value) + return nil +} + +func (b *AccountedBuffer) WriteInt32(value int32) error { + return b.WriteUint32(uint32(value)) +} + +func (b *AccountedBuffer) WriteUint64(value uint64) error { + dst, err := b.appendSpace(8) + if err != nil { + return err + } + binary.NativeEndian.PutUint64(dst, value) + return nil +} + +func (b *AccountedBuffer) WriteInt64(value int64) error { + return b.WriteUint64(uint64(value)) +} + +func (b *AccountedBuffer) SetUint32(offset int, value uint32) error { + if b == nil || offset < 0 || offset > len(b.data)-4 { + return ErrAllocationAccountInvalid + } + binary.NativeEndian.PutUint32(b.data[offset:offset+4], value) + return nil +} + +func (b *AccountedBuffer) SetInt64(offset int, value int64) error { + if b == nil || offset < 0 || offset > len(b.data)-8 { + return ErrAllocationAccountInvalid + } + binary.NativeEndian.PutUint64(b.data[offset:offset+8], uint64(value)) + return nil +} + func (b *AccountedBuffer) Reset() { if b != nil { b.data = b.data[:0] diff --git a/pkg/common/mpool/allocation_account.go b/pkg/common/mpool/allocation_account.go index 8a9d896c16ce1..15e31063f1918 100644 --- a/pkg/common/mpool/allocation_account.go +++ b/pkg/common/mpool/allocation_account.go @@ -207,8 +207,8 @@ type AllocationAccountCheckpoint struct { } // AllocationCapacityController lets an account share a higher-level aggregate -// cap during migration. The controller owns cap policy only; physical MPool -// metadata remains the sole release owner. +// cap. The controller owns cap policy only; physical MPool metadata remains +// the sole release owner. type AllocationCapacityController interface { AcquireAllocationCapacity(uint64) error ReleaseAllocationCapacity(uint64) diff --git a/pkg/compare/arraycompare.go b/pkg/compare/arraycompare.go index 3ac0695297ba8..d6e9c4410ee2c 100644 --- a/pkg/compare/arraycompare.go +++ b/pkg/compare/arraycompare.go @@ -45,8 +45,10 @@ func (c arrayCompare) Copy(vecSrc, vecDst int, src, dst int64, proc *process.Pro } func (c arrayCompare) Compare(veci, vecj int, vi, vj int64) int { - n0 := c.isConstNull[veci] || c.vs[veci].GetNulls().Contains(uint64(vi)) - n1 := c.isConstNull[vecj] || c.vs[vecj].GetNulls().Contains(uint64(vj)) + n0 := c.isConstNull[veci] || c.vs[veci].GetNulls().Contains(uint64(vi)) || + c.vs[veci].GetGrouping().Contains(uint64(vi)) + n1 := c.isConstNull[vecj] || c.vs[vecj].GetNulls().Contains(uint64(vj)) || + c.vs[vecj].GetGrouping().Contains(uint64(vj)) cmp := nullsCompare(n0, n1, c.nullsLast) if cmp != 0 { return cmp - nullsCompareFlag diff --git a/pkg/compare/compare.go b/pkg/compare/compare.go index ce547f19da34b..e9c4539bbe57e 100644 --- a/pkg/compare/compare.go +++ b/pkg/compare/compare.go @@ -230,8 +230,10 @@ func (c *compare[T]) Set(idx int, vec *vector.Vector) { } func (c *compare[T]) Compare(veci, vecj int, vi, vj int64) int { - n0 := c.isConstNull[veci] || c.ns[veci].Contains(uint64(vi)) - n1 := c.isConstNull[vecj] || c.ns[vecj].Contains(uint64(vj)) + n0 := c.isConstNull[veci] || c.ns[veci].Contains(uint64(vi)) || + c.gs[veci].Contains(uint64(vi)) + n1 := c.isConstNull[vecj] || c.ns[vecj].Contains(uint64(vj)) || + c.gs[vecj].Contains(uint64(vj)) cmp := nullsCompare(n0, n1, c.nullsLast) if cmp != 0 { return cmp - nullsCompareFlag @@ -239,7 +241,11 @@ func (c *compare[T]) Compare(veci, vecj int, vi, vj int64) int { return c.cmp(c.xs[veci][vi], c.xs[vecj][vj]) } -func (c *compare[T]) Copy(vecSrc, vecDst int, src, dst int64, _ *process.Process) error { +func (c *compare[T]) Copy(vecSrc, vecDst int, src, dst int64, proc *process.Process) error { + if c.gs[vecSrc].Contains(uint64(src)) || + c.gs[vecDst].Contains(uint64(dst)) { + return c.vs[vecDst].Copy(c.vs[vecSrc], dst, src, proc.Mp()) + } if c.isConstNull[vecSrc] || c.ns[vecSrc].Contains(uint64(src)) { nulls.Add(c.ns[vecDst], uint64(dst)) } else { diff --git a/pkg/compare/grouping_compare_test.go b/pkg/compare/grouping_compare_test.go new file mode 100644 index 0000000000000..cd228246e4198 --- /dev/null +++ b/pkg/compare/grouping_compare_test.go @@ -0,0 +1,90 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package compare + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/stretchr/testify/require" +) + +func TestGroupingSentinelComparesAsNull(t *testing.T) { + mp := mpool.MustNewZero() + tests := []struct { + name string + typ types.Type + add func(*vector.Vector, bool) error + }{ + { + name: "fixed", + typ: types.T_int64.ToType(), + add: func(v *vector.Vector, isNull bool) error { + return vector.AppendFixed(v, int64(0), isNull, mp) + }, + }, + { + name: "varlen", + typ: types.T_varchar.ToType(), + add: func(v *vector.Vector, isNull bool) error { + return vector.AppendBytes(v, nil, isNull, mp) + }, + }, + { + name: "array", + typ: types.T_array_float32.ToType(), + add: func(v *vector.Vector, isNull bool) error { + return vector.AppendArray(v, []float32{0}, isNull, mp) + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + grouping := vector.NewVec(test.typ) + ordinary := vector.NewVec(test.typ) + nullValue := vector.NewVec(test.typ) + t.Cleanup(func() { + grouping.Free(mp) + ordinary.Free(mp) + nullValue.Free(mp) + }) + require.NoError(t, test.add(grouping, false)) + require.NoError(t, test.add(ordinary, false)) + require.NoError(t, test.add(nullValue, true)) + grouping.GetGrouping().Add(0) + + for _, nullsLast := range []bool{false, true} { + for _, desc := range []bool{false, true} { + cmp := New(test.typ, desc, nullsLast) + cmp.Set(0, grouping) + cmp.Set(1, ordinary) + if nullsLast { + require.Positive(t, cmp.Compare(0, 1, 0, 0)) + require.Negative(t, cmp.Compare(1, 0, 0, 0)) + } else { + require.Negative(t, cmp.Compare(0, 1, 0, 0)) + require.Positive(t, cmp.Compare(1, 0, 0, 0)) + } + + cmp.Set(1, nullValue) + require.Zero(t, cmp.Compare(0, 1, 0, 0)) + } + } + }) + } +} diff --git a/pkg/compare/strcompare.go b/pkg/compare/strcompare.go index f74ddc2a6375c..ba155c435e2c6 100644 --- a/pkg/compare/strcompare.go +++ b/pkg/compare/strcompare.go @@ -46,8 +46,10 @@ func (c *strCompare) Copy(vecSrc, vecDst int, src, dst int64, proc *process.Proc } func (c *strCompare) Compare(veci, vecj int, vi, vj int64) int { - n0 := c.isConstNull[veci] || c.vs[veci].GetNulls().Contains(uint64(vi)) - n1 := c.isConstNull[vecj] || c.vs[vecj].GetNulls().Contains(uint64(vj)) + n0 := c.isConstNull[veci] || c.vs[veci].GetNulls().Contains(uint64(vi)) || + c.vs[veci].GetGrouping().Contains(uint64(vi)) + n1 := c.isConstNull[vecj] || c.vs[vecj].GetNulls().Contains(uint64(vj)) || + c.vs[vecj].GetGrouping().Contains(uint64(vj)) cmp := nullsCompare(n0, n1, c.nullsLast) if cmp != 0 { return cmp - nullsCompareFlag diff --git a/pkg/container/batch/allocation_account_test.go b/pkg/container/batch/allocation_account_test.go index ad3ad0a2d4d65..a2861f63390f3 100644 --- a/pkg/container/batch/allocation_account_test.go +++ b/pkg/container/batch/allocation_account_test.go @@ -39,7 +39,7 @@ func newTestBatchAllocationAccount( require.NoError(t, err) account, err := registry.Open(16 << 20) require.NoError(t, err) - selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) require.NoError(t, err) return testBatchAllocationAccount{ registry: registry, @@ -155,6 +155,14 @@ func TestBatchAllocationAccountCloneDupAndWindow(t *testing.T) { for _, vec := range aliasDecoded.Vecs { require.Nil(t, vec.AllocationAccountSelection()) } + require.NoError( + t, + aliasDecoded.UnmarshalFromReader(bytes.NewReader(data), mp), + ) + for _, vec := range aliasDecoded.Vecs { + require.Same(t, state.selection, vec.AllocationAccountSelection()) + } + require.Equal(t, source.RowCount(), aliasDecoded.RowCount()) aliasDecoded.Clean(mp) require.Equal(t, sourceUsed, state.account.Snapshot().Used) @@ -162,6 +170,211 @@ func TestBatchAllocationAccountCloneDupAndWindow(t *testing.T) { finalizeTestBatchAllocationAccount(t, state) } +func TestBatchAccountedReaderAcceptsBitmapCapacityBeyondLogicalRows(t *testing.T) { + state := newTestBatchAllocationAccount(t, 64) + mp := mpool.MustNewZero() + source := newBatchAllocationTestSource(t, mp, nil) + source.Vecs[0].GetNulls().Add(31) + source.Shrink([]int64{0, 1, 2, 3, 4}, false) + + var encoded bytes.Buffer + require.NoError(t, source.MarshalBinaryTo(&encoded)) + decoded := NewOffHeapEmpty() + require.NoError(t, decoded.SetAllocationAccount(state.selection)) + require.NoError(t, decoded.UnmarshalFromReader(&encoded, mp)) + require.Equal(t, 5, decoded.RowCount()) + require.Equal( + t, + int64(0), + vector.GetFixedAtWithTypeCheck[int64](decoded.Vecs[0], 0), + ) + + decoded.Clean(mp) + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchAccountedReaderPreservesRowsWhenVectorCountChanges(t *testing.T) { + state := newTestBatchAllocationAccount(t, 64) + mp := mpool.MustNewZero() + source := newBatchAllocationTestSource(t, mp, nil) + + var encoded bytes.Buffer + require.NoError(t, source.MarshalBinaryTo(&encoded)) + decoded := NewOffHeapWithSize(1) + decoded.Vecs[0] = vector.NewOffHeapVecWithType(types.T_int64.ToType()) + require.NoError(t, decoded.SetAllocationAccount(state.selection)) + decoded.SetRowCount(7) + require.NoError(t, decoded.UnmarshalFromReader(&encoded, mp)) + require.Equal(t, source.RowCount(), decoded.RowCount()) + require.Len(t, decoded.Vecs, 2) + require.Equal( + t, + int64(31), + vector.GetFixedAtWithTypeCheck[int64](decoded.Vecs[0], 31), + ) + + decoded.Clean(mp) + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchGroupingCodecRoundTrip(t *testing.T) { + state := newTestBatchAllocationAccount(t, 64) + mp := mpool.MustNewZero() + source := newBatchAllocationTestSource(t, mp, nil) + source.Vecs[0].GetGrouping().Add(1, 7, 31) + source.Vecs[1].GetGrouping().Add(2, 9) + source.ExtraBuf = bytes.Repeat([]byte("x"), 1<<20) + + var encoded bytes.Buffer + spillSize, err := source.MarshalBinaryWithGroupingSize() + require.NoError(t, err) + stableSize, err := source.MarshalBinarySize() + require.NoError(t, err) + require.Greater(t, stableSize-spillSize, len(source.ExtraBuf)/2) + require.NoError(t, source.MarshalBinaryWithGroupingTo(&encoded)) + decoded := NewOffHeapEmpty() + require.NoError(t, decoded.SetAllocationAccount(state.selection)) + require.NoError(t, decoded.UnmarshalFromReaderWithGrouping(&encoded, mp)) + require.Empty(t, decoded.Attrs) + require.Empty(t, decoded.ExtraBuf) + for i := range source.Vecs { + require.True(t, decoded.Vecs[i].GetGrouping().IsSame(source.Vecs[i].GetGrouping())) + } + withoutGrouping := newBatchAllocationTestSource(t, mp, nil) + encoded.Reset() + require.NoError(t, withoutGrouping.MarshalBinaryWithGroupingTo(&encoded)) + require.NoError(t, decoded.UnmarshalFromReaderWithGrouping(&encoded, mp)) + for _, vec := range decoded.Vecs { + require.True(t, vec.GetGrouping().IsEmpty()) + } + + decoded.Clean(mp) + withoutGrouping.Clean(mp) + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchGroupingCodecRejectsStableMetadataBeforePayloadAllocation(t *testing.T) { + state := newTestBatchAllocationAccount(t, 32) + mp := mpool.MustNewZero() + + for _, test := range []struct { + name string + attrs []string + extra []byte + want string + }{ + { + name: "attributes", + attrs: []string{string(bytes.Repeat([]byte("a"), 1<<20))}, + want: "attributes are not allowed", + }, + { + name: "extra buffer", + extra: bytes.Repeat([]byte("x"), 1<<20), + want: "extra buffer is not allowed", + }, + } { + t.Run(test.name, func(t *testing.T) { + source := NewWithSize(0) + source.Attrs = test.attrs + source.ExtraBuf = test.extra + var encoded bytes.Buffer + require.NoError(t, source.MarshalBinaryTo(&encoded)) + + decoded := NewOffHeapEmpty() + require.NoError(t, decoded.SetAllocationAccount(state.selection)) + require.ErrorContains( + t, + decoded.UnmarshalFromReaderWithGrouping(&encoded, mp), + test.want, + ) + require.Zero(t, state.account.Snapshot().Used) + decoded.Clean(mp) + }) + } + + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchGroupingCodecRejectsMismatchedRowCount(t *testing.T) { + state := newTestBatchAllocationAccount(t, 32) + mp := mpool.MustNewZero() + source := NewWithSize(1) + source.Vecs[0] = vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixed(source.Vecs[0], int64(1), false, mp)) + source.SetRowCount(2) + + var encoded bytes.Buffer + require.NoError(t, source.MarshalBinaryWithGroupingTo(&encoded)) + decoded := NewOffHeapEmpty() + require.NoError(t, decoded.SetAllocationAccount(state.selection)) + require.ErrorContains( + t, + decoded.UnmarshalFromReaderWithGrouping(&encoded, mp), + "vector length does not match row count", + ) + + decoded.Clean(mp) + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchAccountedReaderRejectsInvalidLengthsBeforeAllocation(t *testing.T) { + state := newTestBatchAllocationAccount(t, 16) + mp := mpool.MustNewZero() + encode := func(values ...[]byte) []byte { + return bytes.Join(values, nil) + } + zeroRows := int64(0) + zeroCount := int32(0) + negative := int32(-1) + huge := int32(1<<20 + 1) + one := int32(1) + oversized := int32(1 << 30) + + tests := []struct { + name string + data []byte + }{ + { + name: "negative vector count", + data: encode(types.EncodeInt64(&zeroRows), types.EncodeInt32(&negative)), + }, + { + name: "huge vector count", + data: encode(types.EncodeInt64(&zeroRows), types.EncodeInt32(&huge)), + }, + { + name: "negative attribute count", + data: encode(types.EncodeInt64(&zeroRows), types.EncodeInt32(&zeroCount), types.EncodeInt32(&negative)), + }, + { + name: "oversized attribute payload", + data: encode( + types.EncodeInt64(&zeroRows), + types.EncodeInt32(&zeroCount), + types.EncodeInt32(&one), + types.EncodeInt32(&oversized), + ), + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + decoded := NewOffHeapEmpty() + require.NoError(t, decoded.SetAllocationAccount(state.selection)) + require.NotPanics(t, func() { + require.Error(t, decoded.UnmarshalFromReader(bytes.NewReader(test.data), mp)) + }) + decoded.Clean(mp) + require.Zero(t, state.account.Snapshot().Used) + }) + } + finalizeTestBatchAllocationAccount(t, state) +} + func newMixedBatchAllocationSource( t *testing.T, mp *mpool.MPool, @@ -170,13 +383,13 @@ func newMixedBatchAllocationSource( ) *Batch { t.Helper() bat := NewOffHeapWithSize(2) - bat.Attrs = []string{"accounted", "legacy"} + bat.Attrs = []string{"accounted", "unaccounted"} bat.Vecs[0] = vector.NewOffHeapVecWithType(types.T_int64.ToType()) require.NoError(t, bat.Vecs[0].SetAllocationAccount(selection)) bat.Vecs[1] = vector.NewOffHeapVecWithType(types.T_varchar.ToType()) for i := 0; i < rows; i++ { require.NoError(t, vector.AppendFixed(bat.Vecs[0], int64(i), false, mp)) - require.NoError(t, vector.AppendBytes(bat.Vecs[1], []byte("legacy"), false, mp)) + require.NoError(t, vector.AppendBytes(bat.Vecs[1], []byte("unaccounted"), false, mp)) } bat.SetRowCount(rows) return bat @@ -211,10 +424,10 @@ func TestMixedBatchAllocationClonePreservesVectorProvenance(t *testing.T) { require.NoError(t, err) require.Same(t, state.selection, accounted.Vecs[0].AllocationAccountSelection()) accounted.Clean(mp) - legacy, err := source.CloneSelectedColumns([]int{1}, []string{"legacy"}, mp) + unaccounted, err := source.CloneSelectedColumns([]int{1}, []string{"unaccounted"}, mp) require.NoError(t, err) - require.Nil(t, legacy.Vecs[0].AllocationAccountSelection()) - legacy.Clean(mp) + require.Nil(t, unaccounted.Vecs[0].AllocationAccountSelection()) + unaccounted.Clean(mp) source.FreeColumns(mp) require.Zero(t, state.account.Snapshot().Used) @@ -256,11 +469,11 @@ func TestBatchSetStartsNewTailWhenVectorProvenanceChanges(t *testing.T) { state := newTestBatchAllocationAccount(t, 128) mp := mpool.MustNewZero() set := NewBatchSet(4) - legacy := newBatchAllocationTestSource(t, mp, nil) - legacy.Shrink([]int64{0, 1}, false) + unaccounted := newBatchAllocationTestSource(t, mp, nil) + unaccounted.Shrink([]int64{0, 1}, false) mixed := newMixedBatchAllocationSource(t, mp, state.selection, 3) - _, err := set.Extend(mp, legacy, nil) + _, err := set.Extend(mp, unaccounted, nil) require.NoError(t, err) ready := set.ReadyCount() require.Equal(t, 1, set.ReadyDeltaFor(mixed, mixed.RowCount())) @@ -273,10 +486,10 @@ func TestBatchSetStartsNewTailWhenVectorProvenanceChanges(t *testing.T) { require.Nil(t, set.Get(0).Vecs[0].AllocationAccountSelection()) require.Same(t, state.selection, set.Get(1).Vecs[0].AllocationAccountSelection()) - legacyUnion := newBatchAllocationTestSource(t, mp, nil) + unaccountedUnion := newBatchAllocationTestSource(t, mp, nil) ready = set.ReadyCount() - require.Equal(t, 1, set.ReadyDeltaFor(legacyUnion, 1)) - _, err = set.Union(mp, legacyUnion, []int32{0}, nil) + require.Equal(t, 1, set.ReadyDeltaFor(unaccountedUnion, 1)) + _, err = set.Union(mp, unaccountedUnion, []int32{0}, nil) require.NoError(t, err) require.Equal(t, 1, set.ReadyCount()-ready) require.Equal(t, 3, set.Length()) @@ -288,9 +501,9 @@ func TestBatchSetStartsNewTailWhenVectorProvenanceChanges(t *testing.T) { require.Same(t, state.selection, set.Get(3).Vecs[0].AllocationAccountSelection()) require.Equal(t, 1, set.Get(3).RowCount()) - legacy.Clean(mp) + unaccounted.Clean(mp) mixed.Clean(mp) - legacyUnion.Clean(mp) + unaccountedUnion.Clean(mp) set.Clean(mp) finalizeTestBatchAllocationAccount(t, state) } diff --git a/pkg/container/batch/batch.go b/pkg/container/batch/batch.go index 20fee1392b7f2..880646822327b 100644 --- a/pkg/container/batch/batch.go +++ b/pkg/container/batch/batch.go @@ -17,8 +17,10 @@ package batch import ( "bytes" "context" + "encoding/binary" "fmt" "io" + "math" "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/moerr" @@ -122,6 +124,12 @@ func (bat *Batch) MarshalBinaryWithBuffer(w *bytes.Buffer, reset bool) ([]byte, } func (bat *Batch) MarshalBinarySize() (int, error) { + return bat.prepareMarshalBinary(nil) +} + +func (bat *Batch) prepareMarshalBinary( + plans []vector.MarshalBinaryPlan, +) (int, error) { if bat == nil { return 0, moerr.NewInvalidInputNoCtx("invalid batch for marshal") } @@ -141,22 +149,29 @@ func (bat *Batch) MarshalBinarySize() (int, error) { "batch field exceeds marshal format", ) } - for _, vec := range bat.Vecs { + if plans != nil && len(plans) < len(bat.Vecs) { + return 0, moerr.NewInvalidInputNoCtx("short batch marshal plan") + } + for i, vec := range bat.Vecs { if vec == nil { return 0, moerr.NewInvalidInputNoCtx( "cannot marshal a nil batch vector", ) } - size, err := vec.MarshalBinarySize() + plan, err := vec.PrepareMarshalBinary() if err != nil { return 0, err } + size := plan.Size() if uint64(size) > uint64(^uint32(0)) || !add(4+uint64(size)) { return 0, moerr.NewInvalidInputNoCtx( "batch vector exceeds marshal format", ) } + if plans != nil { + plans[i] = plan + } } for _, attr := range bat.Attrs { if uint64(len(attr)) > uint64(^uint32(0)>>1) || @@ -178,42 +193,52 @@ func (bat *Batch) MarshalBinaryTo(w io.Writer) error { if bat == nil || w == nil { return io.ErrClosedPipe } - if _, err := bat.MarshalBinarySize(); err != nil { + var inlinePlans [64]vector.MarshalBinaryPlan + var plans []vector.MarshalBinaryPlan + if len(bat.Vecs) <= len(inlinePlans) { + plans = inlinePlans[:len(bat.Vecs)] + } else { + plans = make([]vector.MarshalBinaryPlan, len(bat.Vecs)) + } + size, err := bat.prepareMarshalBinary(plans) + if err != nil { return err } - rl := int64(bat.rowCount) - if err := writeBatchMarshalBytes(w, types.EncodeInt64(&rl)); err != nil { + if sized, ok := w.(interface { + Len() int + EnsureCapacity(int) error + }); ok { + if sized.Len() > math.MaxInt-size { + return moerr.NewInvalidInputNoCtx("batch marshal size exceeds platform limit") + } + if err := sized.EnsureCapacity(sized.Len() + size); err != nil { + return err + } + } + if err := writeBatchMarshalInt64(w, int64(bat.rowCount)); err != nil { return err } l := int32(len(bat.Vecs)) - if err := writeBatchMarshalBytes(w, types.EncodeInt32(&l)); err != nil { + if err := writeBatchMarshalInt32(w, l); err != nil { return err } for i := 0; i < int(l); i++ { - size, err := bat.Vecs[i].MarshalBinarySize() - if err != nil { - return err - } - wireSize := uint32(size) - if err := writeBatchMarshalBytes( - w, - types.EncodeUint32(&wireSize), - ); err != nil { + if err := writeBatchMarshalUint32(w, uint32(plans[i].Size())); err != nil { return err } - if err := bat.Vecs[i].MarshalBinaryTo(w); err != nil { + if err := plans[i].MarshalTo(w); err != nil { return err } } l = int32(len(bat.Attrs)) - if err := writeBatchMarshalBytes(w, types.EncodeInt32(&l)); err != nil { + if err := writeBatchMarshalInt32(w, l); err != nil { return err } for i := 0; i < int(l); i++ { size := int32(len(bat.Attrs[i])) - if err := writeBatchMarshalBytes(w, types.EncodeInt32(&size)); err != nil { + if err := writeBatchMarshalInt32(w, size); err != nil { return err } n, err := io.WriteString(w, bat.Attrs[i]) @@ -226,20 +251,50 @@ func (bat *Batch) MarshalBinaryTo(w io.Writer) error { } extraSize := int32(len(bat.ExtraBuf)) - if err := writeBatchMarshalBytes(w, types.EncodeInt32(&extraSize)); err != nil { + if err := writeBatchMarshalInt32(w, extraSize); err != nil { return err } if err := writeBatchMarshalBytes(w, bat.ExtraBuf); err != nil { return err } - if err := writeBatchMarshalBytes( - w, - types.EncodeInt32(&bat.Recursive), - ); err != nil { + if err := writeBatchMarshalInt32(w, bat.Recursive); err != nil { return err } - return writeBatchMarshalBytes(w, types.EncodeInt32(&bat.ShuffleIDX)) + return writeBatchMarshalInt32(w, bat.ShuffleIDX) +} + +type batchPrimitiveWriter interface { + WriteUint32(uint32) error + WriteInt32(int32) error + WriteInt64(int64) error +} + +func writeBatchMarshalUint32(w io.Writer, value uint32) error { + if typed, ok := w.(batchPrimitiveWriter); ok { + return typed.WriteUint32(value) + } + var data [4]byte + binary.NativeEndian.PutUint32(data[:], value) + return writeBatchMarshalBytes(w, data[:]) +} + +func writeBatchMarshalInt32(w io.Writer, value int32) error { + if typed, ok := w.(batchPrimitiveWriter); ok { + return typed.WriteInt32(value) + } + var data [4]byte + binary.NativeEndian.PutUint32(data[:], uint32(value)) + return writeBatchMarshalBytes(w, data[:]) +} + +func writeBatchMarshalInt64(w io.Writer, value int64) error { + if typed, ok := w.(batchPrimitiveWriter); ok { + return typed.WriteInt64(value) + } + var data [8]byte + binary.NativeEndian.PutUint64(data[:], uint64(value)) + return writeBatchMarshalBytes(w, data[:]) } func writeBatchMarshalBytes(w io.Writer, value []byte) error { @@ -505,35 +560,35 @@ func (bat *Batch) UnmarshalBinaryWithAnyMp(data []byte, mp *mpool.MPool) (err er } func (bat *Batch) UnmarshalFromReader(r io.Reader, mp *mpool.MPool) (err error) { - allocationAccount := bat.allocationAccount + return bat.unmarshalFromReader(r, mp, true) +} + +func (bat *Batch) unmarshalFromReader( + r io.Reader, + mp *mpool.MPool, + allowMetadata bool, +) (err error) { + if bat == nil || r == nil { + return io.ErrClosedPipe + } i64, err := types.ReadInt64(r) if err != nil { return err } - bat.rowCount = int(i64) + if i64 < 0 || int64(int(i64)) != i64 { + return moerr.NewInvalidInputNoCtx("invalid batch row count") + } + decodedRowCount := int(i64) l, err := types.ReadInt32AsInt(r) if err != nil { return err } - if l != len(bat.Vecs) { - if len(bat.Vecs) > 0 { - bat.Clean(mp) - bat.allocationAccount = allocationAccount - } - bat.Vecs = make([]*vector.Vector, l) - for i := range bat.Vecs { - if bat.offHeap { - bat.Vecs[i] = vector.NewOffHeapVec() - if allocationAccount != nil { - if err := bat.Vecs[i].SetAllocationAccount(allocationAccount); err != nil { - return err - } - } - } else { - bat.Vecs[i] = vector.NewVecFromReuse() - } - } + if err = validateReaderElementCount(r, l, 4, "vector"); err != nil { + return err + } + if err = bat.prepareOwnedDecodeVectors(l, mp); err != nil { + return err } vecs := bat.Vecs @@ -558,21 +613,39 @@ func (bat *Batch) UnmarshalFromReader(r io.Reader, mp *mpool.MPool) (err error) if err != nil { return err } + if err = validateReaderElementCount(r, l, 4, "attribute"); err != nil { + return err + } + if !allowMetadata && l != 0 { + return moerr.NewInvalidInputNoCtx("spill batch attributes are not allowed") + } if l != len(bat.Attrs) { bat.Attrs = make([]string, l) } for i := 0; i < int(l); i++ { - _, bs, err := types.ReadSizeBytes(r) + bs, err := readBatchSizedBytes(r) if err != nil { return err } bat.Attrs[i] = string(bs) } - // ExtraBuf - if _, bat.ExtraBuf, err = types.ReadSizeBytes(r); err != nil { - return err + // ExtraBuf is a data-scaled Go-heap field in the stable Batch codec. Spill + // records do not use it and reject it before allocating its payload. + if allowMetadata { + if bat.ExtraBuf, err = readBatchSizedBytes(r); err != nil { + return err + } + } else { + extraSize, readErr := types.ReadInt32AsInt(r) + if readErr != nil { + return readErr + } + if extraSize != 0 { + return moerr.NewInvalidInputNoCtx("spill batch extra buffer is not allowed") + } + bat.ExtraBuf = nil } if bat.Recursive, err = types.ReadInt32(r); err != nil { @@ -581,6 +654,128 @@ func (bat *Batch) UnmarshalFromReader(r io.Reader, mp *mpool.MPool) (err error) if bat.ShuffleIDX, err = types.ReadInt32(r); err != nil { return err } + bat.rowCount = decodedRowCount + return nil +} + +func readBatchSizedBytes(r io.Reader) ([]byte, error) { + size, err := types.ReadInt32AsInt(r) + if err != nil { + return nil, err + } + if size < 0 { + return nil, moerr.NewInvalidInputNoCtx("negative batch buffer size") + } + if limited, ok := r.(*io.LimitedReader); ok && int64(size) > limited.N { + return nil, io.ErrUnexpectedEOF + } + if lengthAware, ok := r.(interface{ Len() int }); ok && size > lengthAware.Len() { + return nil, io.ErrUnexpectedEOF + } + if size == 0 { + return nil, nil + } + value := make([]byte, size) + if _, err = io.ReadFull(r, value); err != nil { + return nil, err + } + return value, nil +} + +func validateReaderElementCount( + r io.Reader, + count int, + minimumWireBytes int64, + field string, +) error { + const maxBatchWireFields = 1 << 20 + if count < 0 || count > maxBatchWireFields || minimumWireBytes <= 0 { + return moerr.NewInvalidInputNoCtx("invalid batch " + field + " count") + } + var remaining int64 = -1 + switch reader := r.(type) { + case *io.LimitedReader: + remaining = reader.N + case interface{ Len() int }: + remaining = int64(reader.Len()) + } + if remaining >= 0 && int64(count) > remaining/minimumWireBytes { + return moerr.NewInvalidInputNoCtx("invalid batch " + field + " count") + } + return nil +} + +// prepareOwnedDecodeVectors makes every destination an independent owner. +// Alias decoding deliberately installs borrowed vector buffers; those buffers +// must never be grown or relabeled by the owned streaming decoder. +func (bat *Batch) prepareOwnedDecodeVectors(count int, mp *mpool.MPool) error { + if count < 0 { + return moerr.NewInvalidInputNoCtx("invalid batch vector count") + } + allocationAccount := bat.allocationAccount + if count != len(bat.Vecs) { + if len(bat.Vecs) > 0 { + bat.Clean(mp) + bat.allocationAccount = allocationAccount + } + bat.Vecs = make([]*vector.Vector, count) + } + + const inlineReceivers = 16 + var inline [inlineReceivers]*vector.Vector + var used map[*vector.Vector]struct{} + for i, vec := range bat.Vecs { + selection := allocationAccount + if selection == nil && vec != nil { + selection = vec.AllocationAccountSelection() + } + if vec != nil { + exists := false + if i < inlineReceivers { + for j := 0; j < i; j++ { + if inline[j] == vec { + exists = true + break + } + } + inline[i] = vec + } else { + if used == nil { + used = make(map[*vector.Vector]struct{}, count) + for _, prior := range inline { + if prior != nil { + used[prior] = struct{}{} + } + } + } + _, exists = used[vec] + used[vec] = struct{}{} + } + if exists { + vec = nil + } + } + if vec == nil { + if bat.offHeap { + vec = vector.NewOffHeapVec() + } else { + vec = vector.NewVecFromReuse() + } + } else if vec.NeedDup() { + vec.Free(mp) + } + vec.SetOffHeap(bat.offHeap) + if vec.AllocationAccountSelection() != selection { + if err := vec.CanSetAllocationAccount(selection); err != nil { + vec.Free(mp) + vec.SetOffHeap(bat.offHeap) + } + if err := vec.SetAllocationAccount(selection); err != nil { + return err + } + } + bat.Vecs[i] = vec + } return nil } @@ -1105,3 +1300,38 @@ func (bat *Batch) Window(start, end int) (*Batch, error) { b.rowCount = end - start return b, nil } + +// WindowWithAllocation is the allocation-accounted counterpart of Window. +// Vector data and area remain borrowed; null/grouping range bitmaps are owned +// by selection and are released when the returned batch is cleaned. +func (bat *Batch) WindowWithAllocation( + start int, + end int, + mp *mpool.MPool, + selection *vector.AllocationAccountSelection, +) (*Batch, error) { + if bat == nil || mp == nil || selection == nil || + start < 0 || end < start || end > bat.RowCount() { + return nil, mpool.ErrAllocationAccountInvalid + } + b := NewOffHeapWithSize(len(bat.Vecs)) + b.Attrs = bat.Attrs + if err := b.SetAllocationAccount(selection); err != nil { + b.Clean(mp) + return nil, err + } + for i, vec := range bat.Vecs { + if vec == nil { + b.Clean(mp) + return nil, mpool.ErrAllocationAccountInvalid + } + var err error + b.Vecs[i], err = vec.WindowWithAllocation(start, end, mp, selection) + if err != nil { + b.Clean(mp) + return nil, err + } + } + b.rowCount = end - start + return b, nil +} diff --git a/pkg/container/batch/grouping_codec.go b/pkg/container/batch/grouping_codec.go new file mode 100644 index 0000000000000..7d54d0a4c66a7 --- /dev/null +++ b/pkg/container/batch/grouping_codec.go @@ -0,0 +1,131 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package batch + +import ( + "io" + "math" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" +) + +// MarshalBinaryWithGroupingSize returns the spill-only Batch wire size. Spill +// records retain vectors and grouping provenance but deliberately omit Attrs +// and ExtraBuf: joins do not consume either field, and ExtraBuf would otherwise +// create an unaccounted data-scaled Go-heap owner while decoding. +func (bat *Batch) MarshalBinaryWithGroupingSize() (int, error) { + metadataFree := *bat + metadataFree.Attrs = nil + metadataFree.ExtraBuf = nil + size, err := metadataFree.MarshalBinarySize() + if err != nil { + return 0, err + } + if len(bat.Vecs) > math.MaxInt32 || size > math.MaxInt-4 { + return 0, moerr.NewInvalidInputNoCtx("batch grouping payload exceeds marshal format") + } + size += 4 + for _, vec := range bat.Vecs { + groupingSize := vec.GroupingMarshalBinarySize() + if groupingSize < 0 || groupingSize > math.MaxInt32 || + size > math.MaxInt-4-groupingSize { + return 0, moerr.NewInvalidInputNoCtx("batch grouping payload exceeds marshal format") + } + size += 4 + groupingSize + } + return size, nil +} + +func (bat *Batch) MarshalBinaryWithGroupingTo(w io.Writer) error { + if bat == nil || w == nil { + return io.ErrClosedPipe + } + metadataFree := *bat + metadataFree.Attrs = nil + metadataFree.ExtraBuf = nil + if err := metadataFree.MarshalBinaryTo(w); err != nil { + return err + } + return bat.marshalGroupingTo(w) +} + +func (bat *Batch) marshalGroupingTo(w io.Writer) error { + if err := writeBatchMarshalInt32(w, int32(len(bat.Vecs))); err != nil { + return err + } + for _, vec := range bat.Vecs { + size := vec.GroupingMarshalBinarySize() + if size > math.MaxInt32 { + return moerr.NewInvalidInputNoCtx("vector grouping payload exceeds marshal format") + } + if err := writeBatchMarshalInt32(w, int32(size)); err != nil { + return err + } + if size > 0 { + if err := vec.MarshalGroupingTo(w); err != nil { + return err + } + } + } + return nil +} + +func (bat *Batch) UnmarshalFromReaderWithGrouping( + r io.Reader, + mp *mpool.MPool, +) error { + if err := bat.unmarshalFromReader(r, mp, false); err != nil { + return err + } + return bat.unmarshalGroupingFromReader(r, mp) +} + +func (bat *Batch) unmarshalGroupingFromReader( + r io.Reader, + mp *mpool.MPool, +) error { + if err := bat.CheckLength(); err != nil { + return moerr.NewInvalidInputNoCtx("spill batch vector length does not match row count") + } + count, err := types.ReadInt32AsInt(r) + if err != nil { + return err + } + if count != len(bat.Vecs) { + return moerr.NewInvalidInputNoCtx("batch grouping vector count mismatch") + } + for _, vec := range bat.Vecs { + size, err := types.ReadInt32AsInt(r) + if err != nil { + return err + } + if size < 0 { + return moerr.NewInvalidInputNoCtx("invalid vector grouping payload size") + } + if remaining, ok := r.(*io.LimitedReader); ok && int64(size) > remaining.N { + return io.ErrUnexpectedEOF + } + limited := &io.LimitedReader{R: r, N: int64(size)} + if err = vec.UnmarshalGroupingFromReader(limited, size, mp); err != nil { + return err + } + if limited.N != 0 { + return moerr.NewInvalidInputNoCtx("vector grouping payload was not fully consumed") + } + } + return nil +} diff --git a/pkg/container/bytejson/bytejson_composite_plan.go b/pkg/container/bytejson/bytejson_composite_plan.go deleted file mode 100644 index f37df516f7c62..0000000000000 --- a/pkg/container/bytejson/bytejson_composite_plan.go +++ /dev/null @@ -1,358 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package bytejson - -import ( - "bytes" - "encoding/base64" - "encoding/binary" - "math" - "slices" - - "github.com/matrixorigin/matrixone/pkg/common/moerr" -) - -type fixedDataEncoder struct { - typeCode TpCode - data [numberSize]byte - size uint32 -} - -func NewLiteralDataEncoder(literal byte) ByteJsonDataEncoder { - return &fixedDataEncoder{typeCode: TpCodeLiteral, data: [numberSize]byte{literal}, size: 1} -} - -func NewInt64DataEncoder(value int64) ByteJsonDataEncoder { - encoder := &fixedDataEncoder{typeCode: TpCodeInt64, size: numberSize} - endian.PutUint64(encoder.data[:], uint64(value)) - return encoder -} - -func NewUint64DataEncoder(value uint64) ByteJsonDataEncoder { - encoder := &fixedDataEncoder{typeCode: TpCodeUint64, size: numberSize} - endian.PutUint64(encoder.data[:], value) - return encoder -} - -func NewFloat64DataEncoder(value float64) ByteJsonDataEncoder { - encoder := &fixedDataEncoder{typeCode: TpCodeFloat64, size: numberSize} - endian.PutUint64(encoder.data[:], math.Float64bits(value)) - return encoder -} - -func (e *fixedDataEncoder) TypeCode() TpCode { return e.typeCode } -func (e *fixedDataEncoder) DataSize() uint32 { return e.size } -func (e *fixedDataEncoder) EncodeDataInto(dst []byte) (int, error) { - if e == nil || len(dst) != int(e.size) { - return 0, moerr.NewInvalidArgNoCtx("JSON scalar", "result size mismatch") - } - return copy(dst, e.data[:e.size]), nil -} - -type rawDataEncoder struct { - value ByteJson -} - -// NewRawDataEncoder references an already storage-compatible value. -func NewRawDataEncoder(value ByteJson) (ByteJsonDataEncoder, error) { - if value.requiresLegacyBinaryEncoding() { - return nil, moerr.NewInvalidArgNoCtx( - "JSON value", - "value is not storage compatible", - ) - } - if uint64(len(value.Data)) > math.MaxUint32 { - return nil, moerr.NewInvalidArgNoCtx("JSON value", "value is too large") - } - return &rawDataEncoder{value: value}, nil -} - -func (e *rawDataEncoder) TypeCode() TpCode { return e.value.Type } -func (e *rawDataEncoder) DataSize() uint32 { return uint32(len(e.value.Data)) } -func (e *rawDataEncoder) EncodeDataInto(dst []byte) (int, error) { - if e == nil || len(dst) != len(e.value.Data) { - return 0, moerr.NewInvalidArgNoCtx("JSON value", "result size mismatch") - } - return copy(dst, e.value.Data), nil -} - -type typedStringDataEncoder struct { - typeCode TpCode - value []byte - dataSize uint32 -} - -func NewTypedStringDataEncoder(tp TpCode, value []byte) (ByteJsonDataEncoder, error) { - switch tp { - case TpCodeString, TpCodeDecimal, TpCodeDate, TpCodeTime, TpCodeDatetime, TpCodeBlob: - default: - return nil, moerr.NewInvalidArgNoCtx("JSON string", "invalid type code") - } - dataSize, err := binaryStringDataSize(len(value)) - if err != nil { - return nil, err - } - return &typedStringDataEncoder{typeCode: tp, value: value, dataSize: dataSize}, nil -} - -func (e *typedStringDataEncoder) TypeCode() TpCode { return e.typeCode } -func (e *typedStringDataEncoder) DataSize() uint32 { return e.dataSize } -func (e *typedStringDataEncoder) EncodeDataInto(dst []byte) (int, error) { - if e == nil || len(dst) != int(e.dataSize) { - return 0, moerr.NewInvalidArgNoCtx("JSON string", "result size mismatch") - } - written := binary.PutUvarint(dst, uint64(len(e.value))) - written += copy(dst[written:], e.value) - return written, nil -} - -type binaryDataEncoder struct { - value []byte - prefix string - dataSize uint32 -} - -func NewOpaqueDataEncoder(value []byte) (ByteJsonDataEncoder, error) { - return newBinaryDataEncoder(value, "") -} - -func NewBitDataEncoder(value []byte) (ByteJsonDataEncoder, error) { - return newBinaryDataEncoder(value, persistedBitPrefix) -} - -func newBinaryDataEncoder(value []byte, prefix string) (ByteJsonDataEncoder, error) { - encodedLength := uint64(base64.StdEncoding.EncodedLen(len(value))) + uint64(len(prefix)) - if encodedLength > math.MaxInt { - return nil, moerr.NewInvalidArgNoCtx("JSON binary", "value is too large") - } - dataSize, err := binaryStringDataSize(int(encodedLength)) - if err != nil { - return nil, err - } - return &binaryDataEncoder{value: value, prefix: prefix, dataSize: dataSize}, nil -} - -func (e *binaryDataEncoder) TypeCode() TpCode { return TpCodeBlob } -func (e *binaryDataEncoder) DataSize() uint32 { return e.dataSize } -func (e *binaryDataEncoder) EncodeDataInto(dst []byte) (int, error) { - if e == nil || len(dst) != int(e.dataSize) { - return 0, moerr.NewInvalidArgNoCtx("JSON binary", "result size mismatch") - } - encodedLength := base64.StdEncoding.EncodedLen(len(e.value)) + len(e.prefix) - written := binary.PutUvarint(dst, uint64(encodedLength)) - written += copy(dst[written:], e.prefix) - base64.StdEncoding.Encode(dst[written:], e.value) - return written + base64.StdEncoding.EncodedLen(len(e.value)), nil -} - -func binaryStringDataSize(length int) (uint32, error) { - if length < 0 { - return 0, moerr.NewInvalidArgNoCtx("JSON string", "invalid length") - } - var lengthBuffer [binary.MaxVarintLen64]byte - lengthSize := binary.PutUvarint(lengthBuffer[:], uint64(length)) - total := uint64(lengthSize) + uint64(length) - if total > math.MaxUint32 { - return 0, moerr.NewInvalidArgNoCtx("JSON string", "value is too large") - } - return uint32(total), nil -} - -type ArrayDataEncoder struct { - values []ByteJsonDataEncoder - dataSize uint32 -} - -func NewArrayDataEncoder(values []ByteJsonDataEncoder) (*ArrayDataEncoder, error) { - total := uint64(headerSize) + uint64(len(values))*valEntrySize - for _, value := range values { - if value == nil { - return nil, moerr.NewInvalidArgNoCtx("JSON array", "nil value encoder") - } - if value.TypeCode() != TpCodeLiteral { - total += uint64(value.DataSize()) - } - if total > math.MaxUint32 { - return nil, moerr.NewInvalidArgNoCtx("JSON array", "result is too large") - } - } - return &ArrayDataEncoder{values: values, dataSize: uint32(total)}, nil -} - -func (e *ArrayDataEncoder) TypeCode() TpCode { return TpCodeArray } -func (e *ArrayDataEncoder) DataSize() uint32 { return e.dataSize } -func (e *ArrayDataEncoder) EncodeDataInto(dst []byte) (int, error) { - if e == nil || len(dst) != int(e.dataSize) { - return 0, moerr.NewInvalidArgNoCtx("JSON array", "result size mismatch") - } - headerEnd := headerSize + len(e.values)*valEntrySize - clear(dst[:headerEnd]) - endian.PutUint32(dst, uint32(len(e.values))) - endian.PutUint32(dst[docSizeOff:], e.dataSize) - payloadOffset := headerEnd - for idx, value := range e.values { - entryOffset := headerSize + idx*valEntrySize - dst[entryOffset] = byte(value.TypeCode()) - if value.TypeCode() == TpCodeLiteral { - if _, err := value.EncodeDataInto(dst[entryOffset+valTypeSize : entryOffset+valTypeSize+1]); err != nil { - return 0, err - } - continue - } - endian.PutUint32(dst[entryOffset+valTypeSize:], uint32(payloadOffset)) - size := int(value.DataSize()) - written, err := value.EncodeDataInto(dst[payloadOffset : payloadOffset+size]) - if err != nil { - return 0, err - } - if written != size { - return 0, moerr.NewInvalidArgNoCtx("JSON array", "value size mismatch") - } - payloadOffset += size - } - return payloadOffset, nil -} - -type IndexedFloatArrayDataEncoder struct { - count int - valueAt func(int) float64 - dataSize uint32 -} - -func NewIndexedFloatArrayDataEncoder( - count int, - valueAt func(int) float64, -) (*IndexedFloatArrayDataEncoder, error) { - if count < 0 || valueAt == nil { - return nil, moerr.NewInvalidArgNoCtx("JSON array", "invalid value accessor") - } - total := uint64(headerSize) + uint64(count)*(valEntrySize+numberSize) - if total > math.MaxUint32 { - return nil, moerr.NewInvalidArgNoCtx("JSON array", "result is too large") - } - return &IndexedFloatArrayDataEncoder{ - count: count, valueAt: valueAt, dataSize: uint32(total), - }, nil -} - -func (e *IndexedFloatArrayDataEncoder) TypeCode() TpCode { return TpCodeArray } -func (e *IndexedFloatArrayDataEncoder) DataSize() uint32 { return e.dataSize } -func (e *IndexedFloatArrayDataEncoder) EncodeDataInto(dst []byte) (int, error) { - if e == nil || len(dst) != int(e.dataSize) { - return 0, moerr.NewInvalidArgNoCtx("JSON array", "result size mismatch") - } - headerEnd := headerSize + e.count*valEntrySize - clear(dst[:headerEnd]) - endian.PutUint32(dst, uint32(e.count)) - endian.PutUint32(dst[docSizeOff:], e.dataSize) - payloadOffset := headerEnd - for idx := 0; idx < e.count; idx++ { - entryOffset := headerSize + idx*valEntrySize - dst[entryOffset] = byte(TpCodeFloat64) - endian.PutUint32(dst[entryOffset+valTypeSize:], uint32(payloadOffset)) - endian.PutUint64(dst[payloadOffset:], math.Float64bits(e.valueAt(idx))) - payloadOffset += numberSize - } - return payloadOffset, nil -} - -type ObjectDataEncoderEntry struct { - Key []byte - Value ByteJsonDataEncoder - order int -} - -type ObjectDataEncoder struct { - entries []ObjectDataEncoderEntry - dataSize uint32 -} - -func NewObjectDataEncoder(entries []ObjectDataEncoderEntry) (*ObjectDataEncoder, error) { - for idx := range entries { - entries[idx].order = idx - if entries[idx].Value == nil || len(entries[idx].Key) > math.MaxUint16 { - return nil, moerr.NewInvalidArgNoCtx("JSON object", "invalid entry") - } - } - slices.SortFunc(entries, func(left, right ObjectDataEncoderEntry) int { - if order := bytes.Compare(left.Key, right.Key); order != 0 { - return order - } - return left.order - right.order - }) - unique := entries[:0] - for _, entry := range entries { - if len(unique) > 0 && bytes.Equal(unique[len(unique)-1].Key, entry.Key) { - unique[len(unique)-1] = entry - } else { - unique = append(unique, entry) - } - } - entries = unique - total := uint64(headerSize) + uint64(len(entries))*(keyEntrySize+valEntrySize) - for _, entry := range entries { - total += uint64(len(entry.Key)) - if entry.Value.TypeCode() != TpCodeLiteral { - total += uint64(entry.Value.DataSize()) - } - if total > math.MaxUint32 { - return nil, moerr.NewInvalidArgNoCtx("JSON object", "result is too large") - } - } - return &ObjectDataEncoder{entries: entries, dataSize: uint32(total)}, nil -} - -func (e *ObjectDataEncoder) TypeCode() TpCode { return TpCodeObject } -func (e *ObjectDataEncoder) DataSize() uint32 { return e.dataSize } -func (e *ObjectDataEncoder) EncodeDataInto(dst []byte) (int, error) { - if e == nil || len(dst) != int(e.dataSize) { - return 0, moerr.NewInvalidArgNoCtx("JSON object", "result size mismatch") - } - count := len(e.entries) - keyEntryBegin := headerSize - valueEntryBegin := keyEntryBegin + count*keyEntrySize - payloadOffset := valueEntryBegin + count*valEntrySize - clear(dst[:payloadOffset]) - endian.PutUint32(dst, uint32(count)) - endian.PutUint32(dst[docSizeOff:], e.dataSize) - for idx, entry := range e.entries { - entryOffset := keyEntryBegin + idx*keyEntrySize - endian.PutUint32(dst[entryOffset:], uint32(payloadOffset)) - endian.PutUint16(dst[entryOffset+keyOriginOff:], uint16(len(entry.Key))) - payloadOffset += copy(dst[payloadOffset:], entry.Key) - } - for idx, entry := range e.entries { - value := entry.Value - entryOffset := valueEntryBegin + idx*valEntrySize - dst[entryOffset] = byte(value.TypeCode()) - if value.TypeCode() == TpCodeLiteral { - if _, err := value.EncodeDataInto(dst[entryOffset+valTypeSize : entryOffset+valTypeSize+1]); err != nil { - return 0, err - } - continue - } - endian.PutUint32(dst[entryOffset+valTypeSize:], uint32(payloadOffset)) - size := int(value.DataSize()) - written, err := value.EncodeDataInto(dst[payloadOffset : payloadOffset+size]) - if err != nil { - return 0, err - } - if written != size { - return 0, moerr.NewInvalidArgNoCtx("JSON object", "value size mismatch") - } - payloadOffset += size - } - return payloadOffset, nil -} diff --git a/pkg/container/bytejson/bytejson_composite_plan_test.go b/pkg/container/bytejson/bytejson_composite_plan_test.go deleted file mode 100644 index ee3f3cf8376b1..0000000000000 --- a/pkg/container/bytejson/bytejson_composite_plan_test.go +++ /dev/null @@ -1,125 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package bytejson - -import ( - "testing" - - "github.com/stretchr/testify/require" -) - -func encodedByteJSON(t *testing.T, encoder ByteJsonDataEncoder) ByteJson { - t.Helper() - data := make([]byte, encoder.DataSize()) - written, err := encoder.EncodeDataInto(data) - require.NoError(t, err) - require.Equal(t, len(data), written) - return ByteJson{Type: encoder.TypeCode(), Data: data} -} - -func TestArrayDataEncoderMatchesCreateByteJSON(t *testing.T) { - stringEncoder, err := NewStringDataEncoder([]byte("value")) - require.NoError(t, err) - values := []ByteJsonDataEncoder{ - NewLiteralDataEncoder(LiteralNull), - NewLiteralDataEncoder(LiteralTrue), - NewInt64DataEncoder(-7), - NewUint64DataEncoder(9), - NewFloat64DataEncoder(1.25), - stringEncoder, - } - encoder, err := NewArrayDataEncoder(values) - require.NoError(t, err) - want, err := CreateByteJSON([]any{nil, true, int64(-7), uint64(9), 1.25, "value"}) - require.NoError(t, err) - require.Equal(t, want, encodedByteJSON(t, encoder)) -} - -func TestIndexedFloatArrayDataEncoder(t *testing.T) { - values := []float64{1.5, -2, 3.25} - encoder, err := NewIndexedFloatArrayDataEncoder( - len(values), - func(idx int) float64 { return values[idx] }, - ) - require.NoError(t, err) - want, err := CreateByteJSON([]any{1.5, -2.0, 3.25}) - require.NoError(t, err) - require.Equal(t, want, encodedByteJSON(t, encoder)) -} - -func TestObjectDataEncoderSortsAndKeepsLastDuplicate(t *testing.T) { - first, err := NewStringDataEncoder([]byte("first")) - require.NoError(t, err) - last, err := NewStringDataEncoder([]byte("last")) - require.NoError(t, err) - encoder, err := NewObjectDataEncoder([]ObjectDataEncoderEntry{ - {Key: []byte("z"), Value: NewInt64DataEncoder(1)}, - {Key: []byte("a"), Value: first}, - {Key: []byte("a"), Value: last}, - }) - require.NoError(t, err) - want, err := CreateByteJSON(map[string]any{"a": "last", "z": int64(1)}) - require.NoError(t, err) - require.Equal(t, want, encodedByteJSON(t, encoder)) -} - -func TestBinaryDataEncodersAreStorageCompatible(t *testing.T) { - raw := []byte{0, 1, 2, 250, 251} - for _, constructor := range []func([]byte) (ByteJsonDataEncoder, error){ - NewOpaqueDataEncoder, - NewBitDataEncoder, - } { - encoder, err := constructor(raw) - require.NoError(t, err) - value := encodedByteJSON(t, encoder) - require.Equal(t, TpCodeBlob, value.Type) - require.False(t, value.requiresLegacyBinaryEncoding()) - } -} - -func TestCompositeDataEncodersPreserveNestedAndTypedValues(t *testing.T) { - nested, err := CreateByteJSON(map[string]any{ - "key": []any{int64(1), "value"}, - }) - require.NoError(t, err) - raw, err := NewRawDataEncoder(nested) - require.NoError(t, err) - date, err := NewTypedStringDataEncoder(TpCodeDate, []byte("2026-07-31")) - require.NoError(t, err) - decimal, err := NewTypedStringDataEncoder(TpCodeDecimal, []byte("123.450")) - require.NoError(t, err) - array, err := NewArrayDataEncoder([]ByteJsonDataEncoder{raw, date, decimal}) - require.NoError(t, err) - encoded := encodedByteJSON(t, array) - require.Equal(t, nested, encoded.GetArrayElem(0)) - require.Equal(t, TpCodeDate, encoded.GetArrayElem(1).Type) - require.Equal(t, []byte("2026-07-31"), encoded.GetArrayElem(1).GetString()) - require.Equal(t, TpCodeDecimal, encoded.GetArrayElem(2).Type) - require.Equal(t, []byte("123.450"), encoded.GetArrayElem(2).GetString()) -} - -func TestCompositeDataEncodersRejectInvalidPlansAndDestinations(t *testing.T) { - _, err := NewArrayDataEncoder([]ByteJsonDataEncoder{nil}) - require.Error(t, err) - _, err = NewObjectDataEncoder([]ObjectDataEncoderEntry{{Key: []byte("key")}}) - require.Error(t, err) - _, err = NewTypedStringDataEncoder(TpCodeObject, []byte("invalid")) - require.Error(t, err) - - encoder, err := NewArrayDataEncoder([]ByteJsonDataEncoder{NewInt64DataEncoder(1)}) - require.NoError(t, err) - _, err = encoder.EncodeDataInto(make([]byte, encoder.DataSize()-1)) - require.Error(t, err) -} diff --git a/pkg/container/bytejson/bytejson_keys_plan.go b/pkg/container/bytejson/bytejson_keys_plan.go deleted file mode 100644 index 95a313e4dc915..0000000000000 --- a/pkg/container/bytejson/bytejson_keys_plan.go +++ /dev/null @@ -1,85 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package bytejson - -import ( - "encoding/binary" - "math" - - "github.com/matrixorigin/matrixone/pkg/common/moerr" -) - -// ObjectKeysArrayEncoder exposes an object's sorted keys as a JSON array -// without materializing []any or a second binary-JSON payload. -type ObjectKeysArrayEncoder struct { - object ByteJson - dataSize uint32 -} - -func NewObjectKeysArrayEncoder(object ByteJson) (*ObjectKeysArrayEncoder, error) { - if object.Type != TpCodeObject { - return nil, moerr.NewInvalidArgNoCtx("json_keys", "JSON value is not an object") - } - count := object.GetElemCnt() - total := uint64(headerSize) + uint64(count)*uint64(valEntrySize) - var lengthBuffer [binary.MaxVarintLen64]byte - for idx := 0; idx < count; idx++ { - key := object.GetObjectKey(idx) - lengthSize := binary.PutUvarint(lengthBuffer[:], uint64(len(key))) - total += uint64(lengthSize) + uint64(len(key)) - if total > math.MaxUint32 { - return nil, moerr.NewInvalidArgNoCtx("json_keys", "JSON result is too large") - } - } - return &ObjectKeysArrayEncoder{ - object: object, - dataSize: uint32(total), - }, nil -} - -func (e *ObjectKeysArrayEncoder) TypeCode() TpCode { - return TpCodeArray -} - -func (e *ObjectKeysArrayEncoder) DataSize() uint32 { - if e == nil { - return 0 - } - return e.dataSize -} - -func (e *ObjectKeysArrayEncoder) EncodeDataInto(dst []byte) (int, error) { - if e == nil || uint64(len(dst)) != uint64(e.dataSize) { - return 0, moerr.NewInvalidArgNoCtx("json_keys", "JSON result size mismatch") - } - count := e.object.GetElemCnt() - headerEnd := headerSize + count*valEntrySize - clear(dst[:headerEnd]) - endian.PutUint32(dst, uint32(count)) - endian.PutUint32(dst[docSizeOff:], e.dataSize) - payloadOffset := headerEnd - for idx := 0; idx < count; idx++ { - entryOffset := headerSize + idx*valEntrySize - dst[entryOffset] = byte(TpCodeString) - endian.PutUint32(dst[entryOffset+valTypeSize:], uint32(payloadOffset)) - key := e.object.GetObjectKey(idx) - payloadOffset += binary.PutUvarint(dst[payloadOffset:], uint64(len(key))) - payloadOffset += copy(dst[payloadOffset:], key) - } - if payloadOffset != len(dst) { - return 0, moerr.NewInvalidArgNoCtx("json_keys", "JSON result size mismatch") - } - return payloadOffset, nil -} diff --git a/pkg/container/bytejson/bytejson_keys_plan_test.go b/pkg/container/bytejson/bytejson_keys_plan_test.go deleted file mode 100644 index 608925d2f1cdc..0000000000000 --- a/pkg/container/bytejson/bytejson_keys_plan_test.go +++ /dev/null @@ -1,68 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package bytejson - -import ( - "testing" - - "github.com/stretchr/testify/require" -) - -func TestObjectKeysArrayEncoder(t *testing.T) { - object, err := CreateByteJSON(map[string]any{ - "z": int64(1), - "a": true, - "longer": nil, - }) - require.NoError(t, err) - encoder, err := NewObjectKeysArrayEncoder(object) - require.NoError(t, err) - encoded := make([]byte, encoder.DataSize()) - written, err := encoder.EncodeDataInto(encoded) - require.NoError(t, err) - require.Equal(t, len(encoded), written) - result := ByteJson{Type: encoder.TypeCode(), Data: encoded} - visible, err := result.MarshalJSON() - require.NoError(t, err) - require.JSONEq(t, `["a", "longer", "z"]`, string(visible)) -} - -func TestObjectKeysArrayEncoderRejectsInvalidInputAndSize(t *testing.T) { - array, err := CreateByteJSON([]any{int64(1)}) - require.NoError(t, err) - _, err = NewObjectKeysArrayEncoder(array) - require.Error(t, err) - - object, err := CreateByteJSON(map[string]any{"key": int64(1)}) - require.NoError(t, err) - encoder, err := NewObjectKeysArrayEncoder(object) - require.NoError(t, err) - _, err = encoder.EncodeDataInto(make([]byte, encoder.DataSize()-1)) - require.Error(t, err) -} - -func TestStringDataEncoder(t *testing.T) { - encoder, err := NewStringDataEncoder([]byte("a value \" with unicode 世界")) - require.NoError(t, err) - encoded := make([]byte, encoder.DataSize()) - written, err := encoder.EncodeDataInto(encoded) - require.NoError(t, err) - require.Equal(t, len(encoded), written) - result := ByteJson{Type: encoder.TypeCode(), Data: encoded} - require.Equal(t, []byte("a value \" with unicode 世界"), result.GetString()) - - _, err = encoder.EncodeDataInto(encoded[:len(encoded)-1]) - require.Error(t, err) -} diff --git a/pkg/container/bytejson/bytejson_scalar_plan.go b/pkg/container/bytejson/bytejson_scalar_plan.go deleted file mode 100644 index 164edac97c0b6..0000000000000 --- a/pkg/container/bytejson/bytejson_scalar_plan.go +++ /dev/null @@ -1,66 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package bytejson - -import ( - "encoding/binary" - "math" - "unicode/utf8" - - "github.com/matrixorigin/matrixone/pkg/common/moerr" -) - -// StringDataEncoder writes a binary-JSON string from caller-owned bytes. -// The source must remain valid until EncodeDataInto returns. -type StringDataEncoder struct { - value []byte - dataSize uint32 -} - -func NewStringDataEncoder(value []byte) (*StringDataEncoder, error) { - if !utf8.Valid(value) { - return nil, moerr.NewInvalidArgNoCtx("JSON string", "invalid UTF-8") - } - var lengthBuffer [binary.MaxVarintLen64]byte - lengthSize := binary.PutUvarint(lengthBuffer[:], uint64(len(value))) - total := uint64(lengthSize) + uint64(len(value)) - if total > math.MaxUint32 { - return nil, moerr.NewInvalidArgNoCtx("JSON string", "value is too large") - } - return &StringDataEncoder{ - value: value, - dataSize: uint32(total), - }, nil -} - -func (e *StringDataEncoder) TypeCode() TpCode { - return TpCodeString -} - -func (e *StringDataEncoder) DataSize() uint32 { - if e == nil { - return 0 - } - return e.dataSize -} - -func (e *StringDataEncoder) EncodeDataInto(dst []byte) (int, error) { - if e == nil || uint64(len(dst)) != uint64(e.dataSize) { - return 0, moerr.NewInvalidArgNoCtx("JSON string", "result size mismatch") - } - written := binary.PutUvarint(dst, uint64(len(e.value))) - written += copy(dst[written:], e.value) - return written, nil -} diff --git a/pkg/container/bytejson/bytejson_text_writer.go b/pkg/container/bytejson/bytejson_text_writer.go deleted file mode 100644 index 79d10ca28a8f0..0000000000000 --- a/pkg/container/bytejson/bytejson_text_writer.go +++ /dev/null @@ -1,296 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package bytejson - -import ( - "encoding/base64" - "io" - "math" - "strconv" - "unicode/utf8" - - "github.com/matrixorigin/matrixone/pkg/common/moerr" -) - -// WriteJSONText writes the visible JSON representation without allocating a -// payload-sized intermediate slice. It is byte-for-byte equivalent to -// ByteJson.MarshalJSON. -func WriteJSONText(w io.Writer, value ByteJson) error { - switch value.Type { - case TpCodeArray: - if err := writeByte(w, '['); err != nil { - return err - } - for idx := 0; idx < value.GetElemCnt(); idx++ { - if idx > 0 { - if err := writeString(w, ", "); err != nil { - return err - } - } - if err := WriteJSONText(w, value.GetArrayElem(idx)); err != nil { - return err - } - } - return writeByte(w, ']') - case TpCodeObject: - if err := writeByte(w, '{'); err != nil { - return err - } - for idx := 0; idx < value.GetElemCnt(); idx++ { - if idx > 0 { - if err := writeString(w, ", "); err != nil { - return err - } - } - if err := WriteJSONString(w, value.GetObjectKey(idx)); err != nil { - return err - } - if err := writeString(w, ": "); err != nil { - return err - } - if err := WriteJSONText(w, value.GetObjectVal(idx)); err != nil { - return err - } - } - return writeByte(w, '}') - case TpCodeInt64: - var buf [32]byte - return writeBytes(w, strconv.AppendInt(buf[:0], value.GetInt64(), 10)) - case TpCodeUint64: - var buf [32]byte - return writeBytes(w, strconv.AppendUint(buf[:0], value.GetUint64(), 10)) - case TpCodeLiteral: - if len(value.Data) == 0 { - return moerr.NewInvalidInputNoCtx("invalid JSON literal") - } - switch value.Data[0] { - case LiteralNull: - return writeString(w, "null") - case LiteralTrue: - return writeString(w, "true") - case LiteralFalse: - return writeString(w, "false") - default: - return moerr.NewInvalidInputNoCtxf("invalid JSON literal %d", value.Data[0]) - } - case TpCodeFloat64: - f := value.GetFloat64() - if math.IsInf(f, 0) || math.IsNaN(f) { - return moerr.NewInvalidInputNoCtxf("invalid JSON float64 %f", f) - } - format := byte('e') - abs := math.Abs(f) - if abs == 0 || 1e-6 <= abs && abs < 1e21 { - format = 'f' - } - var buf [32]byte - return writeBytes(w, strconv.AppendFloat(buf[:0], f, format, -1, 64)) - case TpCodeString: - return WriteJSONString(w, value.GetString()) - case TpCodeDecimal: - return writeBytes(w, value.GetString()) - case TpCodeDate, TpCodeTime, TpCodeDatetime: - if err := writeByte(w, '"'); err != nil { - return err - } - if err := writeBytes(w, value.GetString()); err != nil { - return err - } - return writeByte(w, '"') - case TpCodeBlob: - if err := writeByte(w, '"'); err != nil { - return err - } - if err := writeBinaryJSONText(w, value); err != nil { - return err - } - return writeByte(w, '"') - case TpCodeOpaque, TpCodeBit: - if err := writeByte(w, '"'); err != nil { - return err - } - if err := writeBinaryJSONText(w, value); err != nil { - return err - } - return writeByte(w, '"') - default: - return moerr.NewInvalidInputNoCtxf("invalid JSON type %d", value.Type) - } -} - -// WriteJSONObjectKeyText applies JSON_OBJECT's key coercion to an existing -// binary-JSON value without allocating its visible representation. -func WriteJSONObjectKeyText(w io.Writer, value ByteJson) error { - switch value.Type { - case TpCodeString, TpCodeDate, TpCodeTime, TpCodeDatetime: - return writeBytes(w, value.GetString()) - case TpCodeBlob, TpCodeOpaque, TpCodeBit: - return writeBinaryJSONText(w, value) - default: - return WriteJSONText(w, value) - } -} - -// WriteJSONBase64Text writes the visible text of a binary JSON scalar without -// surrounding quotes. -func WriteJSONBase64Text(w io.Writer, value []byte) error { - return writeRawBase64(w, value) -} - -func writeBinaryJSONText(w io.Writer, value ByteJson) error { - if value.Type == TpCodeOpaque || value.Type == TpCodeBit { - return writeRawBase64(w, value.GetString()) - } - data := value.GetString() - if len(data) >= len(persistedBitPrefix) && - string(data[:len(persistedBitPrefix)]) == persistedBitPrefix { - encoded := data[len(persistedBitPrefix):] - if _, ok := base64DecodedLen(encoded); ok { - return writeNormalizedBase64(w, encoded) - } - } - return writeBytes(w, data) -} - -// WriteJSONString writes one JSON string without allocating an escaped copy. -func WriteJSONString(w io.Writer, value []byte) error { - if err := writeByte(w, '"'); err != nil { - return err - } - start := 0 - for offset := 0; offset < len(value); { - b := value[offset] - if b < utf8.RuneSelf { - if b >= ' ' && b != '"' && b != '\\' { - offset++ - continue - } - if err := writeBytes(w, value[start:offset]); err != nil { - return err - } - var escaped string - switch b { - case '"': - escaped = `\"` - case '\\': - escaped = `\\` - case '\b': - escaped = `\b` - case '\f': - escaped = `\f` - case '\n': - escaped = `\n` - case '\r': - escaped = `\r` - case '\t': - escaped = `\t` - default: - const hex = "0123456789abcdef" - var escapedControl = [6]byte{'\\', 'u', '0', '0', hex[b>>4], hex[b&0xf]} - if err := writeBytes(w, escapedControl[:]); err != nil { - return err - } - offset++ - start = offset - continue - } - if err := writeString(w, escaped); err != nil { - return err - } - offset++ - start = offset - continue - } - _, size := utf8.DecodeRune(value[offset:]) - if size == 1 { - return moerr.NewInvalidInputNoCtx("invalid UTF-8") - } - offset += size - } - if err := writeBytes(w, value[start:]); err != nil { - return err - } - return writeByte(w, '"') -} - -func writeNormalizedBase64(w io.Writer, encoded []byte) error { - var decoded [binaryJSONCompareDecodedChunkSize]byte - for offset := 0; offset < len(encoded); { - n, next, ok := decodeBase64Chunk(encoded, offset, decoded[:]) - if !ok { - return moerr.NewInvalidInputNoCtx("invalid base64 JSON value") - } - if err := writeRawBase64(w, decoded[:n]); err != nil { - return err - } - offset = next - } - return nil -} - -func writeRawBase64(w io.Writer, raw []byte) error { - const decodedChunk = 3 * 256 - const encodedChunk = 4 * 256 - var encoded [encodedChunk]byte - for len(raw) > 0 { - length := min(len(raw), decodedChunk) - if length < len(raw) { - length -= length % 3 - } - written := base64.StdEncoding.EncodedLen(length) - base64.StdEncoding.Encode(encoded[:written], raw[:length]) - if err := writeBytes(w, encoded[:written]); err != nil { - return err - } - raw = raw[length:] - } - return nil -} - -func writeBytes(w io.Writer, value []byte) error { - for len(value) > 0 { - written, err := w.Write(value) - if err != nil { - return err - } - if written <= 0 || written > len(value) { - return io.ErrShortWrite - } - value = value[written:] - } - return nil -} - -func writeString(w io.Writer, value string) error { - if stringWriter, ok := w.(io.StringWriter); ok { - written, err := stringWriter.WriteString(value) - if err != nil { - return err - } - if written != len(value) { - return io.ErrShortWrite - } - return nil - } - return writeBytes(w, []byte(value)) -} - -func writeByte(w io.Writer, value byte) error { - if byteWriter, ok := w.(io.ByteWriter); ok { - return byteWriter.WriteByte(value) - } - buffer := [1]byte{value} - return writeBytes(w, buffer[:]) -} diff --git a/pkg/container/bytejson/bytejson_text_writer_test.go b/pkg/container/bytejson/bytejson_text_writer_test.go deleted file mode 100644 index 5519cc477c1b5..0000000000000 --- a/pkg/container/bytejson/bytejson_text_writer_test.go +++ /dev/null @@ -1,60 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package bytejson - -import ( - "bytes" - "encoding/base64" - "testing" - - "github.com/stretchr/testify/require" -) - -func TestWriteJSONTextMatchesMarshalJSON(t *testing.T) { - nested, err := CreateByteJSON(map[string]any{ - "array": []any{nil, true, int64(-7), uint64(9), 1.25, "a\n\"中"}, - "empty": map[string]any{}, - }) - require.NoError(t, err) - raw := []byte{0, 1, 2, 3, 250, 251, 252} - values := []ByteJson{ - nested, - {Type: TpCodeOpaque, Data: appendBinaryString(nil, string(raw))}, - {Type: TpCodeBit, Data: appendBinaryString(nil, string(raw))}, - { - Type: TpCodeBlob, - Data: appendBinaryString( - nil, - persistedBitPrefix+base64.StdEncoding.EncodeToString(raw), - ), - }, - { - Type: TpCodeBlob, - Data: appendBinaryString(nil, persistedBitPrefix+"not-base64!"), - }, - } - for _, value := range values { - want, err := value.MarshalJSON() - require.NoError(t, err) - var got bytes.Buffer - require.NoError(t, WriteJSONText(&got, value)) - require.Equal(t, want, got.Bytes()) - } -} - -func TestWriteJSONStringRejectsInvalidUTF8(t *testing.T) { - var output bytes.Buffer - require.Error(t, WriteJSONString(&output, []byte{'a', 0xff})) -} diff --git a/pkg/container/hashtable/allocation_account_test.go b/pkg/container/hashtable/allocation_account_test.go index 178261c973ef4..f57f55aa6556e 100644 --- a/pkg/container/hashtable/allocation_account_test.go +++ b/pkg/container/hashtable/allocation_account_test.go @@ -329,7 +329,7 @@ func TestHashTableAccountedHighCardinalityResizeReturnsToZero(t *testing.T) { func BenchmarkHashTableResizeAccounting(b *testing.B) { const rows = 100_000 - b.Run("legacy", func(b *testing.B) { + b.Run("unaccounted", func(b *testing.B) { mp := mpool.MustNewZero() b.ReportAllocs() b.ResetTimer() diff --git a/pkg/container/pSpool/buffer.go b/pkg/container/pSpool/buffer.go index c152fc0d07cf1..b3f7ae4767935 100644 --- a/pkg/container/pSpool/buffer.go +++ b/pkg/container/pSpool/buffer.go @@ -59,20 +59,7 @@ func (b *spoolBuffer) putCacheID(mp *mpool.MPool, id uint32, bat *batch.Batch) { // 1. const vector size was too small, // 2. vector doesn't own its data and area, // we don't need to cache it. - if vec.IsConst() || vec.NeedDup() { - vec.Free(mp) - } - - if vec.AllocationAccountSelection() == nil { - data := vector.DetachLegacyVectorData(vec) - area := vector.DetachLegacyVectorArea(vec) - if data != nil { - b.bytesCache[id].bs = append(b.bytesCache[id].bs, data) - } - if area != nil { - b.bytesCache[id].bs = append(b.bytesCache[id].bs, area) - } - } else { + if !vec.IsConst() && !vec.NeedDup() { data := vector.DetachVectorData(vec) area := vector.DetachVectorArea(vec) if data.Capacity() != 0 { @@ -119,10 +106,6 @@ func (b *spoolBuffer) getCacheID() (uint32, *batch.Batch) { func (b *spoolBuffer) clean(mp *mpool.MPool) { for i := range b.bytesCache { - for j := range b.bytesCache[i].bs { - mp.Free(b.bytesCache[i].bs[j]) - } - b.bytesCache[i].bs = nil for j := range b.bytesCache[i].buffers { b.bytesCache[i].buffers[j].Free(mp) } diff --git a/pkg/container/pSpool/copy.go b/pkg/container/pSpool/copy.go index 81c7116aa30f7..14898f482bad9 100644 --- a/pkg/container/pSpool/copy.go +++ b/pkg/container/pSpool/copy.go @@ -36,9 +36,6 @@ type cachedBatch struct { } type oneBatchMemoryCache struct { - // bs keeps the allocation-unaccounted production fast path unchanged. - bs [][]byte - // buffers copy vector data and area while preserving allocation provenance. buffers []vector.DetachedBuffer } @@ -155,7 +152,7 @@ func (cb *cachedBatch) GetCopiedBatch( cb.CacheBatch(true, cacheID, dst) return nil, false, 0, mpool.ErrAllocationAccountInvalid } - if err = dst.Vecs[i].PreExtendBitmap( + if err = dst.Vecs[i].PreExtendGrouping( int(groupingRows), cb.mp, ); err != nil { @@ -188,21 +185,6 @@ func (cb *cachedBatch) GetCopiedBatch( func (mc *oneBatchMemoryCache) setSuitableDataAreaToVector( dataSize, areaSize int, vec *vector.Vector, -) error { - if vec.AllocationAccountSelection() == nil { - mc.setSuitableLegacyDataAreaToVector(dataSize, areaSize, vec) - return nil - } - return mc.setSuitableAccountedDataAreaToVector( - dataSize, - areaSize, - vec, - ) -} - -func (mc *oneBatchMemoryCache) setSuitableAccountedDataAreaToVector( - dataSize, areaSize int, - vec *vector.Vector, ) error { // return directly once cache was empty. if len(mc.buffers) == 0 { @@ -302,72 +284,6 @@ func (mc *oneBatchMemoryCache) setSuitableAccountedDataAreaToVector( return nil } -func (mc *oneBatchMemoryCache) setSuitableLegacyDataAreaToVector( - dataSize, areaSize int, - vec *vector.Vector, -) { - if len(mc.bs) == 0 { - return - } - - setDataFirst := dataSize >= areaSize - first, second := dataSize, areaSize - if !setDataFirst { - first, second = areaSize, dataSize - } - - if first > 0 { - if idx := mc.bestLegacyBuffer(first); idx >= 0 { - mem := mc.removeLegacyBuffer(idx) - if setDataFirst { - vector.AttachLegacyVectorData(vec, mem) - } else { - vector.AttachLegacyVectorArea(vec, mem) - } - } - } - if second > 0 { - if idx := mc.bestLegacyBuffer(second); idx >= 0 { - mem := mc.removeLegacyBuffer(idx) - if setDataFirst { - vector.AttachLegacyVectorArea(vec, mem) - } else { - vector.AttachLegacyVectorData(vec, mem) - } - } - } - if len(mc.bs) > 0 && cap(vec.GetData()) == 0 && dataSize > 0 { - vector.AttachLegacyVectorData(vec, mc.removeLegacyBuffer(len(mc.bs)-1)) - } - if len(mc.bs) > 0 && cap(vec.GetArea()) == 0 && areaSize > 0 { - vector.AttachLegacyVectorArea(vec, mc.removeLegacyBuffer(len(mc.bs)-1)) - } -} - -func (mc *oneBatchMemoryCache) bestLegacyBuffer(size int) int { - best := -1 - difference := math.MaxInt - for i, buffer := range mc.bs { - if current := cap(buffer) - size; current > 0 && - current < difference { - best = i - difference = current - } - } - return best -} - -func (mc *oneBatchMemoryCache) removeLegacyBuffer(idx int) []byte { - last := len(mc.bs) - 1 - buffer := mc.bs[idx] - if idx != last { - mc.bs[idx] = mc.bs[last] - } - mc.bs[last] = nil - mc.bs = mc.bs[:last] - return buffer -} - func (mc *oneBatchMemoryCache) lastAttachable( vec *vector.Vector, kind vector.DetachedBufferKind, diff --git a/pkg/container/pSpool/sender_test.go b/pkg/container/pSpool/sender_test.go index 516e06c64eb2b..ca05143aeed85 100644 --- a/pkg/container/pSpool/sender_test.go +++ b/pkg/container/pSpool/sender_test.go @@ -105,7 +105,7 @@ func TestCachedBatchPreservesAllocationProvenance(t *testing.T) { require.NoError(t, err) account, err := registry.Open(1 << 20) require.NoError(t, err) - selection, err := vector.NewAllocationAccountSelectionWithBitmaps( + selection, err := vector.NewAllocationAccountSelection( account, 1, 1, @@ -116,7 +116,7 @@ func TestCachedBatchPreservesAllocationProvenance(t *testing.T) { require.NoError(t, err) otherAccount, err := registry.Open(1 << 20) require.NoError(t, err) - otherSelection, err := vector.NewAllocationAccountSelectionWithBitmaps( + otherSelection, err := vector.NewAllocationAccountSelection( otherAccount, 1, 1, @@ -139,6 +139,7 @@ func TestCachedBatchPreservesAllocationProvenance(t *testing.T) { false, mp, )) + require.NoError(t, vec.PreExtendGrouping(1, mp)) vec.GetGrouping().Add(0) source.SetRowCount(1) return source @@ -215,6 +216,8 @@ func TestCachedBatchAllocationFailureReturnsCacheOwnership(t *testing.T) { 1, 1, 2, + 3, + 4, ) require.NoError(t, err) source := batch.NewOffHeapWithSize(1) @@ -243,13 +246,14 @@ func TestCachedBatchAllocationFailureReturnsCacheOwnership(t *testing.T) { require.NoError(t, err) } -func TestLegacyCacheNonLastSelectionRetainsOwnership(t *testing.T) { +func TestUnaccountedCacheNonLastSelectionRetainsOwnership(t *testing.T) { mp := mpool.MustNewZero() cache := oneBatchMemoryCache{} for _, size := range []int{64, 128, 256} { - buffer, err := mp.Alloc(size, true) - require.NoError(t, err) - cache.bs = append(cache.bs, buffer) + owner := vector.NewOffHeapVecWithType(types.T_int8.ToType()) + require.NoError(t, owner.PreExtend(size, mp)) + cache.buffers = append(cache.buffers, vector.DetachVectorData(owner)) + owner.Free(mp) } vec := vector.NewOffHeapVecWithType(types.T_int8.ToType()) require.NoError(t, cache.setSuitableDataAreaToVector( @@ -257,11 +261,11 @@ func TestLegacyCacheNonLastSelectionRetainsOwnership(t *testing.T) { 0, vec, )) - require.Len(t, cache.bs, 2) + require.Len(t, cache.buffers, 2) vec.Free(mp) - for i := range cache.bs { - mp.Free(cache.bs[i]) + for i := range cache.buffers { + cache.buffers[i].Free(mp) } require.Zero(t, mp.CurrNB()) } diff --git a/pkg/container/vector/allocation_account.go b/pkg/container/vector/allocation_account.go index eaaebaa2410cd..af7f6c878a9cf 100644 --- a/pkg/container/vector/allocation_account.go +++ b/pkg/container/vector/allocation_account.go @@ -39,54 +39,12 @@ func allocationAccountInvalid(message string) error { // A selection may be shared by all vectors owned by one Batch. Views do not // copy it: they share storage and therefore must not create a second charge. type AllocationAccountSelection struct { - account *mpool.AllocationAccount - owner mpool.AllocationOwner - dataSite mpool.AllocationSite - areaSite mpool.AllocationSite - nullsSite mpool.AllocationSite - groupingSite mpool.AllocationSite - accountBitmaps bool -} - -// FunctionAllocation is the immutable allocation provenance for row-scaled -// function-owned scratch. Parameter conversion and general function scratch -// use distinct sites so the physical charges remain diagnosable. -type FunctionAllocation struct { - account *mpool.AllocationAccount - owner mpool.AllocationOwner - parameterSite mpool.AllocationSite - scratchSite mpool.AllocationSite -} - -func NewFunctionAllocation( - account *mpool.AllocationAccount, - owner mpool.AllocationOwner, - parameterSite mpool.AllocationSite, - scratchSite mpool.AllocationSite, -) (*FunctionAllocation, error) { - allocation := &FunctionAllocation{ - account: account, - owner: owner, - parameterSite: parameterSite, - scratchSite: scratchSite, - } - if err := allocation.validate(); err != nil { - return nil, err - } - return allocation, nil -} - -func (a *FunctionAllocation) validate() error { - if a == nil || - a.account == nil || - a.account.Handle() == 0 || - a.owner < mpool.AllocationOwnerMin || - a.owner > mpool.AllocationOwnerMax || - a.parameterSite < mpool.AllocationSiteMin || - a.scratchSite < mpool.AllocationSiteMin { - return mpool.ErrAllocationAccountInvalid - } - return nil + account *mpool.AllocationAccount + owner mpool.AllocationOwner + dataSite mpool.AllocationSite + areaSite mpool.AllocationSite + nullsSite mpool.AllocationSite + groupingSite mpool.AllocationSite } func NewAllocationAccountSelection( @@ -94,37 +52,16 @@ func NewAllocationAccountSelection( owner mpool.AllocationOwner, dataSite mpool.AllocationSite, areaSite mpool.AllocationSite, -) (*AllocationAccountSelection, error) { - selection := &AllocationAccountSelection{ - account: account, - owner: owner, - dataSite: dataSite, - areaSite: areaSite, - } - if err := selection.validate(); err != nil { - return nil, err - } - return selection, nil -} - -// NewAllocationAccountSelectionWithBitmaps additionally selects physical -// allocation sites for the Vector null and grouping bitmap backing. -func NewAllocationAccountSelectionWithBitmaps( - account *mpool.AllocationAccount, - owner mpool.AllocationOwner, - dataSite mpool.AllocationSite, - areaSite mpool.AllocationSite, nullsSite mpool.AllocationSite, groupingSite mpool.AllocationSite, ) (*AllocationAccountSelection, error) { selection := &AllocationAccountSelection{ - account: account, - owner: owner, - dataSite: dataSite, - areaSite: areaSite, - nullsSite: nullsSite, - groupingSite: groupingSite, - accountBitmaps: true, + account: account, + owner: owner, + dataSite: dataSite, + areaSite: areaSite, + nullsSite: nullsSite, + groupingSite: groupingSite, } if err := selection.validate(); err != nil { return nil, err @@ -237,16 +174,16 @@ func (s *AllocationAccountSelection) validate() error { s.owner > mpool.AllocationOwnerMax || s.dataSite < mpool.AllocationSiteMin || s.areaSite < mpool.AllocationSiteMin || - (s.accountBitmaps && - (s.nullsSite < mpool.AllocationSiteMin || - s.groupingSite < mpool.AllocationSiteMin)) { + s.nullsSite < mpool.AllocationSiteMin || + s.groupingSite < mpool.AllocationSiteMin { return mpool.ErrAllocationAccountInvalid } return nil } // AllocationAccountSelection returns the immutable selection used by this -// vector's future owned allocations. It is nil for legacy vectors and views. +// vector's future owned allocations. It is nil for unaccounted vectors and +// views outside the retained HashBuild domain. func (v *Vector) AllocationAccountSelection() *AllocationAccountSelection { if v == nil { return nil @@ -302,14 +239,12 @@ func (v *Vector) SetAllocationAccount( if v.allocationAccount == selection { return nil } - if v.allocationAccount != nil && - v.allocationAccount.accountBitmaps && - (selection == nil || !selection.accountBitmaps) { + if v.allocationAccount != nil && selection == nil { v.nsp.GetBitmap().ReleaseExternalStorage() v.gsp.GetBitmap().ReleaseExternalStorage() } v.allocationAccount = selection - if selection != nil && selection.accountBitmaps { + if selection != nil { v.nsp.GetBitmap().InstallExternalStorage(nil) v.gsp.GetBitmap().InstallExternalStorage(nil) } @@ -317,7 +252,7 @@ func (v *Vector) SetAllocationAccount( } func (v *Vector) ensureBitmapCapacity(rows int, mp *mpool.MPool) error { - if v.allocationAccount == nil || !v.allocationAccount.accountBitmaps { + if v.allocationAccount == nil { return nil } if rows < 0 || rows > math.MaxInt-64 || mp == nil { @@ -364,6 +299,53 @@ func (v *Vector) ensureBitmapCapacity(rows int, mp *mpool.MPool) error { return nil } +func (v *Vector) ensureNullCapacity(rows int, mp *mpool.MPool) error { + if v.allocationAccount == nil { + return nil + } + return v.ensureSingleBitmapCapacity( + v.nsp.GetBitmap(), + rows, + mp, + v.allocationAccount.nullsSite, + ) +} + +func (v *Vector) ensureGroupingCapacity(rows int, mp *mpool.MPool) error { + if v.allocationAccount == nil { + return nil + } + return v.ensureSingleBitmapCapacity( + v.gsp.GetBitmap(), + rows, + mp, + v.allocationAccount.groupingSite, + ) +} + +func (v *Vector) ensureSingleBitmapCapacity( + value *bitmap.Bitmap, + rows int, + mp *mpool.MPool, + site mpool.AllocationSite, +) error { + if rows < 0 || rows > math.MaxInt-64 || mp == nil { + return mpool.ErrAllocationAccountInvalid + } + if rows > 0 { + rows++ + } + storage, err := v.allocateBitmapGrowth(value, rows, mp, site) + if err != nil { + return err + } + if cap(storage) > 0 { + previous := value.InstallExternalStorage(storage) + mpool.FreeSlice(mp, previous) + } + return nil +} + func (v *Vector) allocateBitmapGrowth( value *bitmap.Bitmap, rows int, @@ -539,6 +521,9 @@ func (v *Vector) readSizeBytes( "negative vector buffer size", ) } + if err := validateStreamingReadSize(r, int64(size)); err != nil { + return size, nil, err + } var buf []byte if data { buf, err = v.growData(mp, int(size)) @@ -560,3 +545,20 @@ func (v *Vector) readSizeBytes( } return size, buf, nil } + +func validateStreamingReadSize(r io.Reader, size int64) error { + if size < 0 { + return moerr.NewInvalidInputNoCtx("negative vector buffer size") + } + var remaining int64 = -1 + switch reader := r.(type) { + case *io.LimitedReader: + remaining = reader.N + case interface{ Len() int }: + remaining = int64(reader.Len()) + } + if remaining >= 0 && size > remaining { + return io.ErrUnexpectedEOF + } + return nil +} diff --git a/pkg/container/vector/allocation_account_test.go b/pkg/container/vector/allocation_account_test.go index ac132febac91f..793e8b8a7ca98 100644 --- a/pkg/container/vector/allocation_account_test.go +++ b/pkg/container/vector/allocation_account_test.go @@ -17,6 +17,7 @@ package vector import ( "bytes" "errors" + "math" "math/rand" "testing" @@ -26,80 +27,17 @@ import ( ) const ( - testVectorAllocationOwner mpool.AllocationOwner = 1 - testVectorDataAllocationSite mpool.AllocationSite = 1 - testVectorAreaAllocationSite mpool.AllocationSite = 2 - testVectorNullAllocationSite mpool.AllocationSite = 3 - testVectorGroupAllocationSite mpool.AllocationSite = 4 - testVectorParamAllocationSite mpool.AllocationSite = 5 - testVectorScratchAllocationSite mpool.AllocationSite = 6 + testVectorAllocationOwner mpool.AllocationOwner = 1 + testVectorDataAllocationSite mpool.AllocationSite = 1 + testVectorAreaAllocationSite mpool.AllocationSite = 2 + testVectorNullAllocationSite mpool.AllocationSite = 3 + testVectorGroupAllocationSite mpool.AllocationSite = 4 ) type testVectorAllocationAccount struct { registry *mpool.AllocationAccountRegistry account *mpool.AllocationAccount selection *AllocationAccountSelection - function *FunctionAllocation -} - -func newTestVectorFunctionAllocationAccount( - t testing.TB, - limit uint64, - allocationSlots uint64, -) testVectorAllocationAccount { - t.Helper() - registry, err := mpool.NewAllocationAccountRegistry(1, allocationSlots) - require.NoError(t, err) - account, err := registry.Open(limit) - require.NoError(t, err) - selection, err := NewAllocationAccountSelectionWithBitmaps( - account, - testVectorAllocationOwner, - testVectorDataAllocationSite, - testVectorAreaAllocationSite, - testVectorNullAllocationSite, - testVectorGroupAllocationSite, - ) - require.NoError(t, err) - function, err := NewFunctionAllocation( - account, - testVectorAllocationOwner, - testVectorParamAllocationSite, - testVectorScratchAllocationSite, - ) - require.NoError(t, err) - return testVectorAllocationAccount{ - registry: registry, - account: account, - selection: selection, - function: function, - } -} - -func newTestVectorBitmapAllocationAccount( - t testing.TB, - limit uint64, - allocationSlots uint64, -) testVectorAllocationAccount { - t.Helper() - registry, err := mpool.NewAllocationAccountRegistry(1, allocationSlots) - require.NoError(t, err) - account, err := registry.Open(limit) - require.NoError(t, err) - selection, err := NewAllocationAccountSelectionWithBitmaps( - account, - testVectorAllocationOwner, - testVectorDataAllocationSite, - testVectorAreaAllocationSite, - testVectorNullAllocationSite, - testVectorGroupAllocationSite, - ) - require.NoError(t, err) - return testVectorAllocationAccount{ - registry: registry, - account: account, - selection: selection, - } } func newTestVectorAllocationAccount( @@ -117,6 +55,8 @@ func newTestVectorAllocationAccount( testVectorAllocationOwner, testVectorDataAllocationSite, testVectorAreaAllocationSite, + testVectorNullAllocationSite, + testVectorGroupAllocationSite, ) require.NoError(t, err) return testVectorAllocationAccount{ @@ -158,6 +98,8 @@ func TestVectorAllocationAccountConfiguration(t *testing.T) { testVectorAllocationOwner, testVectorDataAllocationSite, testVectorAreaAllocationSite, + testVectorNullAllocationSite, + testVectorGroupAllocationSite, ) require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) @@ -255,29 +197,13 @@ func TestVectorAllocationAccountVarlenaDataAndArea(t *testing.T) { finalizeTestVectorAllocationAccount(t, state) } -func TestVectorAllocationAccountLeavesGoBitmapsUnaccounted(t *testing.T) { - state := newTestVectorAllocationAccount(t, 1<<20, 8) - mp := mpool.MustNewZero() - vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) - require.NoError(t, vec.PreExtend(128, mp)) - before := state.account.Snapshot().Used - - // Null/group bitmaps still use Go []uint64. They are an explicit activation - // blocker and must not be mislabeled as part of the off-heap vector charge. - vec.SetAllNulls(32 * 1024) - vec.GetGrouping().AddRange(0, 32*1024) - require.Equal(t, before, state.account.Snapshot().Used) - - vec.Free(mp) - finalizeTestVectorAllocationAccount(t, state) -} - func TestVectorAllocationAccountBitmapResetReuseAndFree(t *testing.T) { - state := newTestVectorBitmapAllocationAccount(t, 8<<20, 16) + state := newTestVectorAllocationAccount(t, 8<<20, 16) mp := mpool.MustNewZero() vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) require.NoError(t, vec.PreExtend(32*1024, mp)) + require.NoError(t, vec.PreExtendBitmap(32*1024, mp)) vec.SetLength(32 * 1024) vec.SetAllNulls(32 * 1024) vec.GetGrouping().AddRange(0, 32*1024) @@ -304,6 +230,7 @@ func TestVectorAllocationAccountBitmapResetReuseAndFree(t *testing.T) { vec.ResetWithSameType() require.NoError(t, vec.PreExtend(64*1024, mp)) + require.NoError(t, vec.PreExtendBitmap(64*1024, mp)) grown := state.account.Snapshot() require.Greater(t, grown.Used, initial.Used) require.Greater(t, grown.Peak, grown.Used) @@ -314,10 +241,11 @@ func TestVectorAllocationAccountBitmapResetReuseAndFree(t *testing.T) { } func TestVectorAllocationAccountBitmapShrinkUsesNoScratch(t *testing.T) { - state := newTestVectorBitmapAllocationAccount(t, 8<<20, 16) + state := newTestVectorAllocationAccount(t, 8<<20, 16) mp := mpool.MustNewZero() vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) require.NoError(t, vec.PreExtend(130, mp)) + require.NoError(t, vec.PreExtendBitmap(130, mp)) for i := range 130 { require.NoError(t, AppendFixed(vec, int64(i), false, mp)) } @@ -346,10 +274,11 @@ func TestVectorAllocationAccountBitmapShrinkUsesNoScratch(t *testing.T) { } func TestVectorAllocationAccountBitmapShuffleAccountsScratch(t *testing.T) { - state := newTestVectorBitmapAllocationAccount(t, 8<<20, 16) + state := newTestVectorAllocationAccount(t, 8<<20, 16) mp := mpool.MustNewZero() vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) require.NoError(t, vec.PreExtend(130, mp)) + require.NoError(t, vec.PreExtendBitmap(130, mp)) for i := range 130 { require.NoError(t, AppendFixed(vec, int64(i), false, mp)) } @@ -385,10 +314,11 @@ func TestVectorAllocationAccountBitmapShuffleAccountsScratch(t *testing.T) { } func TestVectorAllocationAccountBitmapShuffleFailurePreservesVector(t *testing.T) { - state := newTestVectorBitmapAllocationAccount(t, 8<<20, 4) + state := newTestVectorAllocationAccount(t, 8<<20, 4) mp := mpool.MustNewZero() vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) require.NoError(t, vec.PreExtend(130, mp)) + require.NoError(t, vec.PreExtendBitmap(130, mp)) for i := range 130 { require.NoError(t, AppendFixed(vec, int64(i), false, mp)) } @@ -411,10 +341,11 @@ func TestVectorAllocationAccountBitmapShuffleFailurePreservesVector(t *testing.T } func TestVectorAllocationAccountBitmapGrowthFailurePreservesOwner(t *testing.T) { - state := newTestVectorBitmapAllocationAccount(t, 1000, 8) + state := newTestVectorAllocationAccount(t, 1000, 8) mp := mpool.MustNewZero() vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) require.NoError(t, vec.PreExtend(64, mp)) + require.NoError(t, vec.PreExtendBitmap(64, mp)) vec.SetLength(64) vec.SetNull(7) vec.GetGrouping().Add(9) @@ -425,7 +356,7 @@ func TestVectorAllocationAccountBitmapGrowthFailurePreservesOwner(t *testing.T) groupCapacity := vec.gsp.GetBitmap().ExternalStorageCapacity() // The null replacement fits by itself, but admitting the grouping // replacement would exceed the account. Neither replacement is published. - err := vec.PreExtend(2*1024, mp) + err := vec.PreExtendBitmap(2*1024, mp) require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) require.Equal(t, used, state.account.Snapshot().Used) require.Equal(t, dataCapacity, cap(vec.data)) @@ -439,16 +370,16 @@ func TestVectorAllocationAccountBitmapGrowthFailurePreservesOwner(t *testing.T) } func TestVectorAllocationAccountBitmapRejectsUnadmittedRawGrowth(t *testing.T) { - state := newTestVectorBitmapAllocationAccount(t, 1<<20, 8) + state := newTestVectorAllocationAccount(t, 1<<20, 8) mp := mpool.MustNewZero() - legacy := NewOffHeapVecWithType(types.T_int64.ToType()) - legacy.GetNulls().Add(0) + unaccounted := NewOffHeapVecWithType(types.T_int64.ToType()) + unaccounted.GetNulls().Add(0) require.ErrorIs( t, - legacy.SetAllocationAccount(state.selection), + unaccounted.SetAllocationAccount(state.selection), mpool.ErrAllocationAccountInvalid, ) - legacy.Free(mp) + unaccounted.Free(mp) vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) require.Panics(t, func() { @@ -461,7 +392,7 @@ func TestVectorAllocationAccountBitmapRejectsUnadmittedRawGrowth(t *testing.T) { } func TestVectorAllocationAccountBitmapCopyDecode(t *testing.T) { - state := newTestVectorBitmapAllocationAccount(t, 1<<20, 32) + state := newTestVectorAllocationAccount(t, 1<<20, 32) mp := mpool.MustNewZero() source := NewOffHeapVecWithType(types.T_int64.ToType()) for i := 0; i < 128; i++ { @@ -579,6 +510,103 @@ func TestVectorAllocationAccountViewAndDeepCopy(t *testing.T) { finalizeTestVectorAllocationAccount(t, stateB) } +func TestWindowPreservesGroupingProvenance(t *testing.T) { + mp := mpool.MustNewZero() + for _, typ := range []types.Type{types.T_int64.ToType(), types.T_varchar.ToType()} { + source := NewVec(typ) + if typ.IsVarlen() { + for _, value := range []string{"zero", "one", "two", "three"} { + require.NoError(t, AppendBytes(source, []byte(value), false, mp)) + } + } else { + require.NoError(t, AppendFixedList(source, []int64{0, 1, 2, 3}, nil, mp)) + } + source.GetGrouping().Add(1, 3) + + window, err := source.Window(1, 4) + require.NoError(t, err) + require.True(t, window.GetGrouping().Contains(0)) + require.False(t, window.GetGrouping().Contains(1)) + require.True(t, window.GetGrouping().Contains(2)) + window.Free(mp) + + clone, err := source.CloneWindow(1, 4, mp) + require.NoError(t, err) + require.True(t, clone.GetGrouping().Contains(0)) + require.False(t, clone.GetGrouping().Contains(1)) + require.True(t, clone.GetGrouping().Contains(2)) + clone.Free(mp) + source.Free(mp) + } + + rollup := NewRollupConst(types.T_int64.ToType(), 4, mp) + window, err := rollup.Window(1, 3) + require.NoError(t, err) + require.True(t, window.IsGrouping()) + window.Free(mp) + clone, err := rollup.CloneWindow(1, 3, mp) + require.NoError(t, err) + require.True(t, clone.IsGrouping()) + clone.Free(mp) + rollup.Free(mp) + require.Zero(t, mp.CurrNB()) +} + +func TestAccountedWindowOwnsRangeBitmaps(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 16) + mp := mpool.MustNewZero() + source := NewVec(types.T_int64.ToType()) + require.NoError(t, AppendFixedList(source, []int64{0, 1, 2, 3}, []bool{false, true, false, true}, mp)) + source.GetGrouping().Add(1, 2) + + window, err := source.WindowWithAllocation(1, 4, mp, state.selection) + require.NoError(t, err) + require.True(t, window.GetNulls().GetBitmap().HasExternalStorage()) + require.True(t, window.GetGrouping().GetBitmap().HasExternalStorage()) + require.True(t, window.GetGrouping().Contains(0)) + require.True(t, window.GetGrouping().Contains(1)) + require.NotZero(t, state.account.Snapshot().Used) + window.Free(mp) + require.Zero(t, state.account.Snapshot().Used) + + source.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestIsGroupingRejectsOutOfRangeBits(t *testing.T) { + for _, typ := range []types.Type{types.T_int64.ToType(), types.T_varchar.ToType()} { + vec := NewVec(typ) + vec.SetLength(1) + vec.GetGrouping().Add(5) + require.False(t, vec.IsGrouping()) + vec.Free(nil) + } +} + +func TestConstSetPreservesSelectedGrouping(t *testing.T) { + mp := mpool.MustNewZero() + for _, typ := range []types.Type{types.T_int64.ToType(), types.T_varchar.ToType()} { + source := NewVec(typ) + if typ.IsVarlen() { + require.NoError(t, AppendBytesList(source, [][]byte{[]byte("ordinary"), []byte("grouping")}, nil, mp)) + } else { + require.NoError(t, AppendFixedList(source, []int64{1, 2}, nil, mp)) + } + source.GetGrouping().Add(1) + destination := NewVec(typ) + set := GetConstSetFunction(typ, mp) + + require.NoError(t, set(destination, source, 1, 4)) + require.True(t, destination.IsGrouping()) + require.NoError(t, set(destination, source, 0, 4)) + require.False(t, destination.HasGrouping()) + + destination.Free(mp) + source.Free(mp) + } + require.Zero(t, mp.CurrNB()) +} + func TestVectorAllocationAccountCopyRollback(t *testing.T) { state := newTestVectorAllocationAccount(t, 1<<20, 1) mp := mpool.MustNewZero() @@ -673,6 +701,76 @@ func TestVectorAllocationAccountRandomizedAppendAndSelection(t *testing.T) { finalizeTestVectorAllocationAccount(t, state) } +func TestVectorAccountedUnionPreservesGroupingWithoutNulls(t *testing.T) { + state := newTestVectorAllocationAccount(t, 8<<20, 64) + mp := mpool.MustNewZero() + + for _, typ := range []types.Type{ + types.T_int32.ToType(), + types.T_varchar.ToType(), + } { + t.Run(typ.String(), func(t *testing.T) { + source := NewOffHeapVecWithType(typ) + for i := range 6 { + if typ.IsVarlen() { + require.NoError(t, AppendBytes(source, []byte{byte('a' + i)}, false, mp)) + } else { + require.NoError(t, AppendFixed(source, int32(i), false, mp)) + } + } + source.GetGrouping().Add(1, 4) + + tests := []struct { + name string + run func(*Vector) error + want []bool + }{ + { + name: "union", + run: func(dst *Vector) error { + return dst.Union(source, []int64{4, 0, 1}, mp) + }, + want: []bool{true, false, true}, + }, + { + name: "union int32", + run: func(dst *Vector) error { + return dst.UnionInt32(source, []int32{4, 0, 1}, mp) + }, + want: []bool{true, false, true}, + }, + { + name: "union batch", + run: func(dst *Vector) error { + return dst.UnionBatch(source, 1, 4, nil, mp) + }, + want: []bool{true, false, false, true}, + }, + { + name: "union batch flags", + run: func(dst *Vector) error { + return dst.UnionBatch(source, 1, 4, []uint8{1, 0, 1, 1}, mp) + }, + want: []bool{true, false, true}, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + dst := newAccountedTestVector(t, typ, state.selection) + require.NoError(t, test.run(dst)) + for row, want := range test.want { + require.Equal(t, want, dst.GetGrouping().Contains(uint64(row))) + } + require.True(t, dst.GetGrouping().GetBitmap().HasExternalStorage()) + dst.Free(mp) + }) + } + source.Free(mp) + }) + } + finalizeTestVectorAllocationAccount(t, state) +} + func TestVectorAllocationAccountDecodeCopyAndReader(t *testing.T) { state := newTestVectorAllocationAccount(t, 1<<20, 16) mp := mpool.MustNewZero() @@ -722,20 +820,112 @@ func TestVectorAllocationAccountDecodeCopyAndReader(t *testing.T) { mp, ), ) - require.NotZero(t, state.account.Snapshot().Used) + require.Zero(t, state.account.Snapshot().Used) short.Free(mp) source.Free(mp) finalizeTestVectorAllocationAccount(t, state) } +func TestVectorAccountedReaderRejectsMalformedWire(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 16) + mp := mpool.MustNewZero() + source := NewVec(types.T_int64.ToType()) + require.NoError(t, AppendFixed(source, int64(7), true, mp)) + encoded, err := source.MarshalBinary() + require.NoError(t, err) + + lengthOffset := 1 + types.TSize + dataLengthOffset := lengthOffset + 4 + dataLength := int(types.DecodeUint32(encoded[dataLengthOffset : dataLengthOffset+4])) + areaLengthOffset := dataLengthOffset + 4 + dataLength + areaLength := int(types.DecodeUint32(encoded[areaLengthOffset : areaLengthOffset+4])) + nullLengthOffset := areaLengthOffset + 4 + areaLength + nullOffset := nullLengthOffset + 4 + + tests := []struct { + name string + mutate func([]byte) + }{ + { + name: "invalid class", + mutate: func(data []byte) { + data[0] = 0xff + }, + }, + { + name: "negative length", + mutate: func(data []byte) { + value := uint32(math.MaxUint32) + copy(data[lengthOffset:lengthOffset+4], types.EncodeUint32(&value)) + }, + }, + { + name: "mismatched data length", + mutate: func(data []byte) { + value := uint32(2) + copy(data[lengthOffset:lengthOffset+4], types.EncodeUint32(&value)) + }, + }, + { + name: "oversized data payload", + mutate: func(data []byte) { + value := uint32(1 << 30) + copy(data[dataLengthOffset:dataLengthOffset+4], types.EncodeUint32(&value)) + }, + }, + { + name: "invalid null bitmap count", + mutate: func(data []byte) { + value := int64(2) + copy(data[nullOffset:nullOffset+8], types.EncodeInt64(&value)) + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + data := append([]byte(nil), encoded...) + test.mutate(data) + decoded := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NotPanics(t, func() { + require.Error(t, decoded.UnmarshalWithReader(bytes.NewReader(data), mp)) + }) + decoded.Free(mp) + require.Zero(t, state.account.Snapshot().Used) + + copied := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NotPanics(t, func() { + require.Error(t, copied.UnmarshalBinaryWithCopy(data, mp)) + }) + copied.Free(mp) + require.Zero(t, state.account.Snapshot().Used) + }) + } + for end := range encoded { + accounted := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NotPanics(t, func() { + require.Error(t, accounted.UnmarshalBinaryWithCopy(encoded[:end], mp)) + }) + accounted.Free(mp) + require.Zero(t, state.account.Snapshot().Used) + + unaccounted := NewOffHeapVecWithType(types.T_int64.ToType()) + require.NotPanics(t, func() { + require.Error(t, unaccounted.UnmarshalBinaryWithCopy(encoded[:end], mp)) + }) + unaccounted.Free(mp) + } + + source.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + func BenchmarkVectorAllocationAccount(b *testing.B) { const rows = 8192 mp := mpool.MustNewZero() state := newTestVectorAllocationAccount(b, 1<<40, 64) - bitmapState := newTestVectorBitmapAllocationAccount(b, 1<<40, 64) - b.Run("legacy-fixed-preextend-free", func(b *testing.B) { + b.Run("unaccounted-fixed-preextend-free", func(b *testing.B) { b.ReportAllocs() for i := 0; i < b.N; i++ { vec := NewOffHeapVecWithType(types.T_int64.ToType()) @@ -758,20 +948,7 @@ func BenchmarkVectorAllocationAccount(b *testing.B) { vec.Free(mp) } }) - b.Run("accounted-bitmap-fixed-preextend-free", func(b *testing.B) { - b.ReportAllocs() - for i := 0; i < b.N; i++ { - vec := NewOffHeapVecWithType(types.T_int64.ToType()) - if err := vec.SetAllocationAccount(bitmapState.selection); err != nil { - b.Fatal(err) - } - if err := vec.PreExtend(rows, mp); err != nil { - b.Fatal(err) - } - vec.Free(mp) - } - }) - b.Run("legacy-varlen-preextend-free", func(b *testing.B) { + b.Run("unaccounted-varlen-preextend-free", func(b *testing.B) { b.ReportAllocs() for i := 0; i < b.N; i++ { vec := NewOffHeapVecWithType(types.T_varchar.ToType()) @@ -807,26 +984,7 @@ func BenchmarkVectorAllocationAccount(b *testing.B) { b.StopTimer() vec.Free(mp) }) - b.Run("accounted-bitmap-fixed-reset-reuse", func(b *testing.B) { - vec := newAccountedTestVector( - b, - types.T_int64.ToType(), - bitmapState.selection, - ) - if err := vec.PreExtend(rows, mp); err != nil { - b.Fatal(err) - } - b.ReportAllocs() - b.ResetTimer() - for i := 0; i < b.N; i++ { - vec.ResetWithSameType() - } - b.StopTimer() - vec.Free(mp) - }) - finalizeTestVectorAllocationAccount(b, state) - finalizeTestVectorAllocationAccount(b, bitmapState) } func TestVectorAllocationAccountErrorsAreTyped(t *testing.T) { @@ -855,12 +1013,6 @@ func TestDetachedBufferPreservesAllocationProvenance(t *testing.T) { )) used := state.account.Snapshot().Used require.Positive(t, used) - require.Panics(t, func() { - DetachLegacyVectorData(source) - }) - require.Panics(t, func() { - DetachLegacyVectorArea(source) - }) data := DetachVectorData(source) area := DetachVectorArea(source) @@ -912,22 +1064,27 @@ func TestSetTypeAndFixDataAllocationFailureIsAtomic(t *testing.T) { finalizeTestVectorAllocationAccount(t, state) } -func TestDetachedLegacyBufferAndTypeChange(t *testing.T) { +func TestDetachedUnaccountedBufferAndTypeChange(t *testing.T) { mp := mpool.MustNewZero() source := NewOffHeapVecWithType(types.T_varchar.ToType()) require.NoError(t, AppendBytes( source, - []byte("legacy detached allocation payload"), + []byte("unaccounted detached allocation payload"), false, mp, )) - data := DetachLegacyVectorData(source) - area := DetachLegacyVectorArea(source) + data := DetachVectorData(source) + area := DetachVectorArea(source) source.Free(mp) destination := NewOffHeapVecWithType(types.T_varchar.ToType()) - AttachLegacyVectorData(destination, data) - AttachLegacyVectorArea(destination, area) + // Unaccounted buffers may serve either backing because there is no + // allocation-site provenance to preserve. + require.True(t, data.CanAttachTo(destination, DetachedAreaBuffer)) + require.NoError(t, area.AttachTo(destination, DetachedDataBuffer)) + require.NoError(t, data.AttachTo(destination, DetachedAreaBuffer)) + require.Zero(t, data.Capacity()) + require.Zero(t, area.Capacity()) destination.Free(mp) require.Zero(t, mp.CurrNB()) @@ -997,3 +1154,73 @@ func TestVectorAllocationAccountHelperBoundaries(t *testing.T) { vec.Free(mp) finalizeTestVectorAllocationAccount(t, state) } + +func TestUnionAllPreservesConstGrouping(t *testing.T) { + for _, typ := range []types.Type{ + types.T_int64.ToType(), + types.T_varchar.ToType(), + } { + t.Run(typ.String(), func(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 32) + mp := mpool.MustNewZero() + destination := newAccountedTestVector(t, typ, state.selection) + if typ.IsVarlen() { + require.NoError(t, AppendBytes(destination, []byte("prefix"), false, mp)) + } else { + require.NoError(t, AppendFixed(destination, int64(1), false, mp)) + } + + rollup := NewRollupConst(typ, 3, mp) + require.NoError(t, GetUnionAllFunction(typ, mp)(destination, rollup)) + require.Equal(t, 4, destination.Length()) + for row := uint64(1); row < 4; row++ { + require.True(t, destination.GetGrouping().Contains(row)) + } + rollup.Free(mp) + + var ordinary *Vector + var err error + if typ.IsVarlen() { + ordinary, err = NewConstBytes(typ, []byte("value"), 3, mp) + } else { + ordinary, err = NewConstFixed(typ, int64(2), 3, mp) + } + require.NoError(t, err) + ordinary.GetGrouping().Add(1) + require.NoError(t, GetUnionAllFunction(typ, mp)(destination, ordinary)) + require.Equal(t, 7, destination.Length()) + require.True(t, destination.GetGrouping().Contains(5)) + require.False(t, destination.GetGrouping().Contains(4)) + require.False(t, destination.GetGrouping().Contains(6)) + ordinary.Free(mp) + + require.True(t, + destination.GetGrouping().GetBitmap().HasExternalStorage()) + destination.Free(mp) + finalizeTestVectorAllocationAccount(t, state) + require.Zero(t, mp.CurrNB()) + }) + } +} + +func TestUnmarshalBinaryRejectsOwnedDestinationWithoutLosingBacking(t *testing.T) { + mp := mpool.MustNewZero() + source := NewOffHeapVecWithType(types.T_varchar.ToType()) + target := NewOffHeapVecWithType(types.T_varchar.ToType()) + require.NoError(t, AppendBytes(source, bytes.Repeat([]byte("s"), 64), false, mp)) + targetValue := bytes.Repeat([]byte("t"), 64) + require.NoError(t, AppendBytes(target, targetValue, false, mp)) + encoded, err := source.MarshalBinary() + require.NoError(t, err) + before := mp.CurrNB() + require.Positive(t, before) + + err = target.UnmarshalBinary(encoded) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + require.Equal(t, targetValue, target.GetBytesAt(0)) + require.Equal(t, before, mp.CurrNB()) + + source.Free(mp) + target.Free(mp) + require.Zero(t, mp.CurrNB()) +} diff --git a/pkg/container/vector/functionTools.go b/pkg/container/vector/functionTools.go index a0ff0af783531..9605a068a0ad4 100644 --- a/pkg/container/vector/functionTools.go +++ b/pkg/container/vector/functionTools.go @@ -16,12 +16,9 @@ package vector import ( "fmt" - "math" - "unsafe" "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/common/util" "github.com/matrixorigin/matrixone/pkg/container/bytejson" "github.com/matrixorigin/matrixone/pkg/container/nulls" "github.com/matrixorigin/matrixone/pkg/container/types" @@ -69,45 +66,106 @@ func GenerateFunctionFixedTypeParameter[T types.FixedSizeTExceptStrType](v *Vect } } + // Special handling for type conversions to decimal128 var cols []T - parameterType := *t + var convertedType types.Type var anyT T switch (any)(anyT).(type) { case types.Decimal128: - if needsDecimal128ParameterConversion[T](v) { - parameter, err := generateDecimal128ParameterWithScratch[T]( - nil, - 0, - v, - nil, - ) - if err != nil { - panic(err) + convertedType = types.T_decimal128.ToType() + convertedType.Width = 38 + if t.Oid == types.T_decimal64 { + convertedType.Scale = t.Scale + d64Cols := MustFixedColWithTypeCheck[types.Decimal64](v) + // Optimize: for const vector, only convert one element + if v.IsConst() { + d128 := functionUtil.ConvertD64ToD128(d64Cols[0]) + return &FunctionParameterScalar[T]{ + typ: convertedType, + sourceVector: v, + scalarValue: any(d128).(T), + } } - return parameter + cols = make([]T, len(d64Cols)) + for i, d64 := range d64Cols { + cols[i] = any(functionUtil.ConvertD64ToD128(d64)).(T) + } + } else if t.Oid == types.T_float64 { + convertedType.Scale = 16 + f64Cols := MustFixedColWithTypeCheck[float64](v) + // Optimize: for const vector, only convert one element + if v.IsConst() { + d128, err := types.Decimal128FromFloat64(f64Cols[0], 38, 16) + if err != nil { + // Conversion failed, use zero value (similar to MySQL behavior for invalid conversions) + d128 = types.Decimal128{B0_63: 0, B64_127: 0} + } + return &FunctionParameterScalar[T]{ + typ: convertedType, + sourceVector: v, + scalarValue: any(d128).(T), + } + } + cols = make([]T, len(f64Cols)) + for i, f64 := range f64Cols { + d128, err := types.Decimal128FromFloat64(f64, 38, 16) + if err != nil { + // Conversion failed, use zero value + d128 = types.Decimal128{B0_63: 0, B64_127: 0} + } + cols[i] = any(d128).(T) + } + } else if t.Oid == types.T_float32 { + convertedType.Scale = 7 + f32Cols := MustFixedColWithTypeCheck[float32](v) + // Optimize: for const vector, only convert one element + if v.IsConst() { + d128, err := types.Decimal128FromFloat64(float64(f32Cols[0]), 38, 7) + if err != nil { + // Conversion failed, use zero value + d128 = types.Decimal128{B0_63: 0, B64_127: 0} + } + return &FunctionParameterScalar[T]{ + typ: convertedType, + sourceVector: v, + scalarValue: any(d128).(T), + } + } + cols = make([]T, len(f32Cols)) + for i, f32 := range f32Cols { + d128, err := types.Decimal128FromFloat64(float64(f32), 38, 7) + if err != nil { + // Conversion failed, use zero value + d128 = types.Decimal128{B0_63: 0, B64_127: 0} + } + cols[i] = any(d128).(T) + } + } else { + convertedType = *t + cols = MustFixedColWithTypeCheck[T](v) } - cols = MustFixedColWithTypeCheck[T](v) default: + convertedType = *t cols = MustFixedColWithTypeCheck[T](v) } if v.IsConst() { return &FunctionParameterScalar[T]{ - typ: parameterType, + typ: convertedType, sourceVector: v, scalarValue: cols[0], } } if !v.nsp.IsEmpty() { return &FunctionParameterNormal[T]{ - typ: parameterType, + typ: convertedType, sourceVector: v, values: cols, nullMap: v.GetNulls().GetBitmap(), } } return &FunctionParameterWithoutNull[T]{ - typ: parameterType, + typ: convertedType, sourceVector: v, values: cols, } @@ -483,227 +541,19 @@ type reusableParameterWrapper interface{} type OptFunctionResultWrapper interface { UseOptFunctionParamFrame(paramCount int) getConvenientParamList() []reusableParameterWrapper - hasParameterScratch() bool - resizeParameterScratch(idx int, size int) ([]byte, error) - setFunctionAllocation(allocation *FunctionAllocation) - HasFunctionScratch() bool - ResizeFunctionScratch(size int) ([]byte, bool, error) } func OptGetParamFromWrapper[ParamType types.FixedSizeTExceptStrType]( - wrapper FunctionResultWrapper, - idx int, - src *Vector, -) (FunctionParameterWrapper[ParamType], error) { + wrapper FunctionResultWrapper, idx int, src *Vector) FunctionParameterWrapper[ParamType] { ws := wrapper.getConvenientParamList() - if needsDecimal128ParameterConversion[ParamType](src) { - if !wrapper.hasParameterScratch() { - fr := GenerateFunctionFixedTypeParameter[ParamType](src) - ws[idx] = fr - return fr, nil - } - fr, err := generateDecimal128ParameterWithScratch[ParamType]( - wrapper, - idx, - src, - ws[idx], - ) - if err != nil { - return nil, err - } - ws[idx] = fr - return fr, nil - } if fr, ok := ws[idx].(FunctionParameterWrapper[ParamType]); ok && ReuseFunctionFixedTypeParameter(src, fr) { - return fr, nil + return fr } fr := GenerateFunctionFixedTypeParameter[ParamType](src) ws[idx] = fr - return fr, nil -} - -func needsDecimal128ParameterConversion[T types.FixedSizeTExceptStrType]( - src *Vector, -) bool { - var value T - if _, ok := any(value).(types.Decimal128); !ok { - return false - } - switch src.GetType().Oid { - case types.T_decimal64, types.T_float32, types.T_float64: - return true - default: - return false - } -} - -func generateDecimal128ParameterWithScratch[ - T types.FixedSizeTExceptStrType, -]( - wrapper FunctionResultWrapper, - idx int, - src *Vector, - reuse reusableParameterWrapper, -) (FunctionParameterWrapper[T], error) { - if src.IsConstNull() { - if parameter, ok := reuse.(*FunctionParameterScalarNull[T]); ok { - parameter.typ = *src.GetType() - parameter.sourceVector = src - return parameter, nil - } - return &FunctionParameterScalarNull[T]{ - typ: *src.GetType(), - sourceVector: src, - }, nil - } - convertedType := types.T_decimal128.ToType() - convertedType.Width = 38 - switch src.GetType().Oid { - case types.T_decimal64: - convertedType.Scale = src.GetType().Scale - case types.T_float32: - convertedType.Scale = 7 - case types.T_float64: - convertedType.Scale = 16 - default: - return nil, mpool.ErrAllocationAccountInvalid - } - - convert := func(row int) types.Decimal128 { - switch src.GetType().Oid { - case types.T_decimal64: - values := MustFixedColWithTypeCheck[types.Decimal64](src) - return functionUtil.ConvertD64ToD128(values[row]) - case types.T_float32: - values := MustFixedColWithTypeCheck[float32](src) - value, err := types.Decimal128FromFloat64( - float64(values[row]), - 38, - 7, - ) - if err == nil { - return value - } - case types.T_float64: - values := MustFixedColWithTypeCheck[float64](src) - value, err := types.Decimal128FromFloat64(values[row], 38, 16) - if err == nil { - return value - } - } - return types.Decimal128{} - } - - if src.IsConst() { - value := any(convert(0)).(T) - if parameter, ok := reuse.(*FunctionParameterScalar[T]); ok { - parameter.typ = convertedType - parameter.sourceVector = src - parameter.scalarValue = value - return parameter, nil - } - return &FunctionParameterScalar[T]{ - typ: convertedType, - sourceVector: src, - scalarValue: value, - }, nil - } - - var values []T - if wrapper == nil { - values = make([]T, src.Length()) - } else { - var err error - values, err = parameterScratchSlice[T](wrapper, idx, src.Length()) - if err != nil { - return nil, err - } - } - switch src.GetType().Oid { - case types.T_decimal64: - source := MustFixedColWithTypeCheck[types.Decimal64](src) - for row := range values { - values[row] = any( - functionUtil.ConvertD64ToD128(source[row]), - ).(T) - } - case types.T_float32: - source := MustFixedColWithTypeCheck[float32](src) - for row := range values { - value, conversionErr := types.Decimal128FromFloat64( - float64(source[row]), - 38, - 7, - ) - if conversionErr != nil { - value = types.Decimal128{} - } - values[row] = any(value).(T) - } - case types.T_float64: - source := MustFixedColWithTypeCheck[float64](src) - for row := range values { - value, conversionErr := types.Decimal128FromFloat64( - source[row], - 38, - 16, - ) - if conversionErr != nil { - value = types.Decimal128{} - } - values[row] = any(value).(T) - } - } - if !src.nsp.IsEmpty() { - if parameter, ok := reuse.(*FunctionParameterNormal[T]); ok { - parameter.typ = convertedType - parameter.sourceVector = src - parameter.values = values - parameter.nullMap = src.GetNulls().GetBitmap() - return parameter, nil - } - return &FunctionParameterNormal[T]{ - typ: convertedType, - sourceVector: src, - values: values, - nullMap: src.GetNulls().GetBitmap(), - }, nil - } - if parameter, ok := reuse.(*FunctionParameterWithoutNull[T]); ok { - parameter.typ = convertedType - parameter.sourceVector = src - parameter.values = values - return parameter, nil - } - return &FunctionParameterWithoutNull[T]{ - typ: convertedType, - sourceVector: src, - values: values, - }, nil -} - -func parameterScratchSlice[T types.FixedSizeTExceptStrType]( - wrapper FunctionResultWrapper, - idx int, - length int, -) ([]T, error) { - var value T - elementSize := unsafe.Sizeof(value) - if length < 0 || - elementSize == 0 || - uint64(length) > uint64(math.MaxInt)/uint64(elementSize) { - return nil, mpool.ErrAllocationAccountInvalid - } - data, err := wrapper.resizeParameterScratch( - idx, - int(uint64(length)*uint64(elementSize)), - ) - if err != nil { - return nil, err - } - return util.UnsafeSliceCastToLength[T](data, length), nil + return fr } func OptGetBytesParamFromWrapper(wrapper FunctionResultWrapper, idx int, src *Vector) FunctionParameterWrapper[types.Varlena] { @@ -725,19 +575,15 @@ type FunctionResult[T types.FixedSizeT] struct { vec *Vector mp *mpool.MPool - allocationAccount *AllocationAccountSelection - functionAllocation *FunctionAllocation - isVarlena bool - cols []T - length uint64 + isVarlena bool + cols []T + length uint64 // convenientParam save parameter wrappers for easy getting row values. // // this field is for optimisation to reduce the allocation of FunctionParameterWrapper pointer. // there are still many built-in functions don't use it now, and will be fixed in the future. - convenientParam []reusableParameterWrapper - parameterScratch []*mpool.AccountedBuffer - functionScratch *mpool.AccountedBuffer + convenientParam []reusableParameterWrapper } func MustFunctionResult[T types.FixedSizeT](wrapper FunctionResultWrapper) *FunctionResult[T] { @@ -748,15 +594,11 @@ func MustFunctionResult[T types.FixedSizeT](wrapper FunctionResultWrapper) *Func } func newResultFunc[T types.FixedSizeT]( - resultType types.Type, - mp *mpool.MPool, - allocationAccount *AllocationAccountSelection, -) *FunctionResult[T] { + resultType types.Type, mp *mpool.MPool) *FunctionResult[T] { f := &FunctionResult[T]{ - typ: resultType, - mp: mp, - allocationAccount: allocationAccount, + typ: resultType, + mp: mp, } var tempT T @@ -771,100 +613,15 @@ func (fr *FunctionResult[T]) UseOptFunctionParamFrame(paramCount int) { if fr.convenientParam == nil { fr.convenientParam = make([]reusableParameterWrapper, paramCount) } - if fr.allocationAccount != nil && - fr.functionAllocation != nil && - fr.parameterScratch == nil { - fr.parameterScratch = make([]*mpool.AccountedBuffer, paramCount) - } } func (fr *FunctionResult[T]) getConvenientParamList() []reusableParameterWrapper { return fr.convenientParam } -func (fr *FunctionResult[T]) hasParameterScratch() bool { - return fr.functionAllocation != nil -} - -func (fr *FunctionResult[T]) setFunctionAllocation( - allocation *FunctionAllocation, -) { - fr.functionAllocation = allocation -} - -func (fr *FunctionResult[T]) HasFunctionScratch() bool { - return fr.functionAllocation != nil -} - -func (fr *FunctionResult[T]) resizeParameterScratch( - idx int, - size int, -) ([]byte, error) { - if !fr.hasParameterScratch() { - return nil, mpool.ErrAllocationAccountInvalid - } - if idx < 0 || idx >= len(fr.parameterScratch) { - return nil, mpool.ErrAllocationAccountInvalid - } - if fr.parameterScratch[idx] == nil { - buffer, err := mpool.NewAccountedBuffer( - fr.mp, - fr.functionAllocation.account, - fr.functionAllocation.owner, - fr.functionAllocation.parameterSite, - ) - if err != nil { - return nil, err - } - fr.parameterScratch[idx] = buffer - } - if err := fr.parameterScratch[idx].Resize(size); err != nil { - return nil, err - } - return fr.parameterScratch[idx].Bytes(), nil -} - -// ResizeFunctionScratch returns retained allocation-accounted off-heap -// scratch for data-scaled function internals. Legacy results report selected -// false so callers preserve their existing allocator path. -func (fr *FunctionResult[T]) ResizeFunctionScratch( - size int, -) ([]byte, bool, error) { - if fr.functionAllocation == nil { - return nil, false, nil - } - if fr.functionScratch == nil { - buffer, err := mpool.NewAccountedBuffer( - fr.mp, - fr.functionAllocation.account, - fr.functionAllocation.owner, - fr.functionAllocation.scratchSite, - ) - if err != nil { - return nil, true, err - } - fr.functionScratch = buffer - } - if err := fr.functionScratch.Resize(size); err != nil { - return nil, true, err - } - return fr.functionScratch.Bytes(), true, nil -} - func (fr *FunctionResult[T]) PreExtendAndReset(targetSize int) error { if fr.vec == nil { - var err error - if fr.allocationAccount == nil { - fr.vec = NewOffHeapVecWithType(fr.typ) - } else { - fr.vec, err = NewOffHeapVecWithTypeAndAllocation( - fr.typ, - fr.allocationAccount, - ) - if err != nil { - return err - } - } + fr.vec = NewOffHeapVecWithType(fr.typ) } oldLength := fr.vec.Length() @@ -914,90 +671,6 @@ func (fr *FunctionResult[T]) AppendBytes(val []byte, isnull bool) error { return nil } -// AppendBytesWithFill appends one non-null varlena value and lets fill write -// directly into the result Vector's admitted backing storage. The provided -// slice is valid only during fill. A panic rolls the unpublished row and area -// length back before propagating. Use AppendBytesWithBuilder when construction -// can return an error or a shorter value. -func (fr *FunctionResult[T]) AppendBytesWithFill( - size int, - fill func([]byte), -) error { - return fr.AppendBytesWithBuilder(size, func(dst []byte) (int, error) { - fill(dst) - return size, nil - }) -} - -// AppendBytesWithBuilder reserves capacity for one non-null varlena value and -// lets build return the number of bytes it initialized. This supports codecs -// whose exact output is known only after encoding without allocating a second -// payload buffer. Reserved capacity remains owned by the result across reuse; -// only the published area length is reduced to the actual value size. -func (fr *FunctionResult[T]) AppendBytesWithBuilder( - capacity int, - build func([]byte) (int, error), -) error { - if !fr.isVarlena || - fr.vec == nil || - fr.vec.IsConst() || - capacity < 0 || - build == nil { - return mpool.ErrAllocationAccountInvalid - } - oldAreaLen := len(fr.vec.area) - if uint64(oldAreaLen)+uint64(capacity) > uint64(math.MaxUint32) { - return mpool.ErrAllocationAccountInvalid - } - areaSize := capacity - if capacity <= types.VarlenaInlineSize { - areaSize = 0 - } - if err := fr.vec.PreExtendWithArea(1, areaSize, fr.mp); err != nil { - return err - } - - index := fr.vec.length - values := toSliceOfLengthNoTypeCheck[types.Varlena](fr.vec, index+1) - oldValue := values[index] - values[index] = types.Varlena{} - var target []byte - if capacity <= types.VarlenaInlineSize { - target = values[index][1 : 1+capacity] - } else { - fr.vec.area = fr.vec.area[:oldAreaLen+capacity] - target = fr.vec.area[oldAreaLen:] - } - - published := false - defer func() { - if !published { - fr.vec.area = fr.vec.area[:oldAreaLen] - values[index] = oldValue - } - }() - written, err := build(target) - if err != nil { - return err - } - if written < 0 || written > capacity { - return mpool.ErrAllocationAccountInvalid - } - if written <= types.VarlenaInlineSize { - if capacity > types.VarlenaInlineSize { - copy(values[index][1:1+written], target[:written]) - fr.vec.area = fr.vec.area[:oldAreaLen] - } - values[index][0] = byte(written) - } else { - fr.vec.area = fr.vec.area[:oldAreaLen+written] - values[index].SetOffsetLen(uint32(oldAreaLen), uint32(written)) - } - fr.vec.length++ - published = true - return nil -} - func (fr *FunctionResult[T]) AppendByteJson(bj bytejson.ByteJson, isnull bool) error { if !fr.vec.IsConst() { return AppendByteJson(fr.vec, bj, isnull, fr.mp) @@ -1084,121 +757,65 @@ func (fr *FunctionResult[T]) Free() { fr.vec.Free(fr.mp) fr.vec = nil } - for i := range fr.parameterScratch { - if fr.parameterScratch[i] != nil { - fr.parameterScratch[i].Free() - fr.parameterScratch[i] = nil - } - } - if fr.functionScratch != nil { - fr.functionScratch.Free() - fr.functionScratch = nil - } - fr.allocationAccount = nil - fr.functionAllocation = nil fr.convenientParam = nil - fr.parameterScratch = nil } func NewFunctionResultWrapper(typ types.Type, mp *mpool.MPool) FunctionResultWrapper { - return newFunctionResultWrapper(typ, mp, nil) -} - -// NewFunctionResultWrapperWithAllocation constructs a result owner whose -// lazily allocated Vector data and area use selection. Existing callers remain -// on the legacy path through NewFunctionResultWrapper. -func NewFunctionResultWrapperWithAllocation( - typ types.Type, - mp *mpool.MPool, - selection *AllocationAccountSelection, -) (FunctionResultWrapper, error) { - if err := selection.validate(); err != nil { - return nil, err - } - return newFunctionResultWrapper(typ, mp, selection), nil -} - -func NewFunctionResultWrapperWithFunctionAllocation( - typ types.Type, - mp *mpool.MPool, - selection *AllocationAccountSelection, - functionAllocation *FunctionAllocation, -) (FunctionResultWrapper, error) { - if err := selection.validate(); err != nil { - return nil, err - } - if err := functionAllocation.validate(); err != nil { - return nil, err - } - if selection.account != functionAllocation.account || - selection.owner != functionAllocation.owner { - return nil, mpool.ErrAllocationAccountInvalid - } - result := newFunctionResultWrapper(typ, mp, selection) - result.setFunctionAllocation(functionAllocation) - return result, nil -} - -func newFunctionResultWrapper( - typ types.Type, - mp *mpool.MPool, - selection *AllocationAccountSelection, -) FunctionResultWrapper { if typ.IsVarlen() { - return newResultFunc[types.Varlena](typ, mp, selection) + return newResultFunc[types.Varlena](typ, mp) } switch typ.Oid { case types.T_bool: - return newResultFunc[bool](typ, mp, selection) + return newResultFunc[bool](typ, mp) case types.T_bit: - return newResultFunc[uint64](typ, mp, selection) + return newResultFunc[uint64](typ, mp) case types.T_int8: - return newResultFunc[int8](typ, mp, selection) + return newResultFunc[int8](typ, mp) case types.T_int16: - return newResultFunc[int16](typ, mp, selection) + return newResultFunc[int16](typ, mp) case types.T_int32: - return newResultFunc[int32](typ, mp, selection) + return newResultFunc[int32](typ, mp) case types.T_int64: - return newResultFunc[int64](typ, mp, selection) + return newResultFunc[int64](typ, mp) case types.T_uint8: - return newResultFunc[uint8](typ, mp, selection) + return newResultFunc[uint8](typ, mp) case types.T_uint16: - return newResultFunc[uint16](typ, mp, selection) + return newResultFunc[uint16](typ, mp) case types.T_uint32: - return newResultFunc[uint32](typ, mp, selection) + return newResultFunc[uint32](typ, mp) case types.T_uint64: - return newResultFunc[uint64](typ, mp, selection) + return newResultFunc[uint64](typ, mp) case types.T_float32: - return newResultFunc[float32](typ, mp, selection) + return newResultFunc[float32](typ, mp) case types.T_float64: - return newResultFunc[float64](typ, mp, selection) + return newResultFunc[float64](typ, mp) case types.T_date: - return newResultFunc[types.Date](typ, mp, selection) + return newResultFunc[types.Date](typ, mp) case types.T_year: - return newResultFunc[types.MoYear](typ, mp, selection) + return newResultFunc[types.MoYear](typ, mp) case types.T_datetime: - return newResultFunc[types.Datetime](typ, mp, selection) + return newResultFunc[types.Datetime](typ, mp) case types.T_time: - return newResultFunc[types.Time](typ, mp, selection) + return newResultFunc[types.Time](typ, mp) case types.T_timestamp: - return newResultFunc[types.Timestamp](typ, mp, selection) + return newResultFunc[types.Timestamp](typ, mp) case types.T_decimal64: - return newResultFunc[types.Decimal64](typ, mp, selection) + return newResultFunc[types.Decimal64](typ, mp) case types.T_decimal128: - return newResultFunc[types.Decimal128](typ, mp, selection) + return newResultFunc[types.Decimal128](typ, mp) case types.T_decimal256: - return newResultFunc[types.Decimal256](typ, mp, selection) + return newResultFunc[types.Decimal256](typ, mp) case types.T_TS: - return newResultFunc[types.TS](typ, mp, selection) + return newResultFunc[types.TS](typ, mp) case types.T_Rowid: - return newResultFunc[types.Rowid](typ, mp, selection) + return newResultFunc[types.Rowid](typ, mp) case types.T_Blockid: - return newResultFunc[types.Blockid](typ, mp, selection) + return newResultFunc[types.Blockid](typ, mp) case types.T_uuid: - return newResultFunc[types.Uuid](typ, mp, selection) + return newResultFunc[types.Uuid](typ, mp) case types.T_enum: - return newResultFunc[types.Enum](typ, mp, selection) + return newResultFunc[types.Enum](typ, mp) } panic(fmt.Sprintf("unexpected type %s for function result", typ)) } diff --git a/pkg/container/vector/function_result_allocation_test.go b/pkg/container/vector/function_result_allocation_test.go deleted file mode 100644 index 50a987a93f72c..0000000000000 --- a/pkg/container/vector/function_result_allocation_test.go +++ /dev/null @@ -1,409 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package vector - -import ( - "errors" - "testing" - - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/stretchr/testify/require" -) - -func TestFunctionResultAllocationAccountLifecycle(t *testing.T) { - state := newTestVectorAllocationAccount(t, 1<<20, 8) - mp := mpool.MustNew("function-result-allocation") - defer mpool.DeleteMPool(mp) - - fixed, err := NewFunctionResultWrapperWithAllocation( - types.T_int64.ToType(), - mp, - state.selection, - ) - require.NoError(t, err) - require.NoError(t, fixed.PreExtendAndReset(8)) - fixedVector := fixed.GetResultVector() - require.Same( - t, - state.selection, - fixedVector.AllocationAccountSelection(), - ) - firstUsed := state.account.Snapshot().Used - require.Positive(t, firstUsed) - - require.NoError(t, fixed.PreExtendAndReset(2)) - require.Equal(t, firstUsed, state.account.Snapshot().Used) - - fixed.SetResultVector(nil) - require.NoError(t, fixed.PreExtendAndReset(16)) - transferredUsed := state.account.Snapshot().Used - require.Greater(t, transferredUsed, firstUsed) - fixed.Free() - require.Equal(t, firstUsed, state.account.Snapshot().Used) - fixedVector.Free(mp) - require.Zero(t, state.account.Snapshot().Used) - - varlen, err := NewFunctionResultWrapperWithAllocation( - types.T_varchar.ToType(), - mp, - state.selection, - ) - require.NoError(t, err) - require.NoError(t, varlen.PreExtendAndReset(2)) - result := MustFunctionResult[types.Varlena](varlen) - require.NoError(t, result.AppendBytes(make([]byte, 256), false)) - require.NoError(t, result.AppendBytes([]byte("small"), false)) - varlenUsed := state.account.Snapshot().Used - require.Positive(t, varlenUsed) - - require.NoError(t, varlen.PreExtendAndReset(1)) - require.Equal(t, varlenUsed, state.account.Snapshot().Used) - varlen.Free() - require.Zero(t, state.account.Snapshot().Used) - finalizeTestVectorAllocationAccount(t, state) -} - -func TestFunctionResultAllocationAccountFailure(t *testing.T) { - zeroMP := mpool.MustNewZero() - defer mpool.DeleteMPool(zeroMP) - require.ErrorIs( - t, - func() error { - _, err := NewFunctionResultWrapperWithAllocation( - types.T_int64.ToType(), - zeroMP, - nil, - ) - return err - }(), - mpool.ErrAllocationAccountInvalid, - ) - - state := newTestVectorAllocationAccount(t, 1<<20, 4) - _, err := NewFunctionResultWrapperWithFunctionAllocation( - types.T_int64.ToType(), - zeroMP, - state.selection, - nil, - ) - require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) - otherOwner, err := NewFunctionAllocation( - state.account, - testVectorAllocationOwner+1, - testVectorParamAllocationSite, - testVectorScratchAllocationSite, - ) - require.NoError(t, err) - _, err = NewFunctionResultWrapperWithFunctionAllocation( - types.T_int64.ToType(), - zeroMP, - state.selection, - otherOwner, - ) - require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) - _, err = NewFunctionAllocation( - nil, - testVectorAllocationOwner, - testVectorParamAllocationSite, - testVectorScratchAllocationSite, - ) - require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) - finalizeTestVectorAllocationAccount(t, state) - - state = newTestVectorAllocationAccount(t, 7, 1) - mp := mpool.MustNew("function-result-allocation-failure") - defer mpool.DeleteMPool(mp) - result, err := NewFunctionResultWrapperWithAllocation( - types.T_int64.ToType(), - mp, - state.selection, - ) - require.NoError(t, err) - - err = result.PreExtendAndReset(1) - require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) - require.Zero(t, state.account.Snapshot().Used) - require.Zero(t, state.registry.LiveAllocationMetadata()) - result.Free() - finalizeTestVectorAllocationAccount(t, state) -} - -func TestFunctionResultAppendBytesWithFillLifecycle(t *testing.T) { - state := newTestVectorAllocationAccount(t, 1<<20, 8) - mp := mpool.MustNew("function-result-fill") - defer mpool.DeleteMPool(mp) - wrapper, err := NewFunctionResultWrapperWithAllocation( - types.T_varchar.ToType(), - mp, - state.selection, - ) - require.NoError(t, err) - require.NoError(t, wrapper.PreExtendAndReset(3)) - result := MustFunctionResult[types.Varlena](wrapper) - - require.NoError(t, result.AppendBytesWithFill(5, func(dst []byte) { - copy(dst, "small") - })) - large := make([]byte, 256) - for idx := range large { - large[idx] = byte(idx) - } - require.NoError(t, result.AppendBytesWithFill(len(large), func(dst []byte) { - copy(dst, large) - })) - require.Equal(t, []byte("small"), wrapper.GetResultVector().GetBytesAt(0)) - require.Equal(t, large, wrapper.GetResultVector().GetBytesAt(1)) - - beforeLength := wrapper.GetResultVector().Length() - beforeAreaLength := len(wrapper.GetResultVector().GetArea()) - require.Panics(t, func() { - _ = result.AppendBytesWithFill(512, func(dst []byte) { - dst[0] = 1 - panic("injected fill failure") - }) - }) - require.Equal(t, beforeLength, wrapper.GetResultVector().Length()) - require.Equal(t, beforeAreaLength, len(wrapper.GetResultVector().GetArea())) - - fillErr := errors.New("injected fill error") - require.ErrorIs(t, result.AppendBytesWithBuilder(128, func(dst []byte) (int, error) { - dst[0] = 2 - return 0, fillErr - }), fillErr) - require.Equal(t, beforeLength, wrapper.GetResultVector().Length()) - require.Equal(t, beforeAreaLength, len(wrapper.GetResultVector().GetArea())) - require.ErrorIs(t, result.AppendBytesWithBuilder(128, func([]byte) (int, error) { - return 129, nil - }), mpool.ErrAllocationAccountInvalid) - require.Equal(t, beforeLength, wrapper.GetResultVector().Length()) - require.Equal(t, beforeAreaLength, len(wrapper.GetResultVector().GetArea())) - - require.NoError(t, result.AppendBytesWithBuilder(512, func(dst []byte) (int, error) { - copy(dst, "last") - return 4, nil - })) - require.Equal(t, []byte("last"), wrapper.GetResultVector().GetBytesAt(2)) - require.Equal(t, beforeAreaLength, len(wrapper.GetResultVector().GetArea())) - require.Positive(t, state.account.Snapshot().Used) - - wrapper.Free() - require.Zero(t, state.account.Snapshot().Used) - finalizeTestVectorAllocationAccount(t, state) -} - -func TestFunctionResultAllocationAccountDecimalParameterScratch(t *testing.T) { - state := newTestVectorFunctionAllocationAccount(t, 1<<20, 16) - mp := mpool.MustNew("function-parameter-allocation") - defer mpool.DeleteMPool(mp) - result, err := NewFunctionResultWrapperWithFunctionAllocation( - types.T_bool.ToType(), - mp, - state.selection, - state.function, - ) - require.NoError(t, err) - result.UseOptFunctionParamFrame(1) - - source := NewOffHeapVecWithType(types.T_decimal64.ToType()) - for i := int64(1); i <= 32; i++ { - require.NoError(t, AppendFixed( - source, - types.Decimal64(i), - i%7 == 0, - mp, - )) - } - parameter, err := OptGetParamFromWrapper[types.Decimal128]( - result, - 0, - source, - ) - require.NoError(t, err) - values := parameter.UnSafeGetAllValue() - require.Len(t, values, source.Length()) - require.Equal(t, types.Decimal128{B0_63: 1}, values[0]) - _, isNull := parameter.GetValue(6) - require.True(t, isNull) - first := state.account.Snapshot() - require.Positive(t, first.Used) - require.Equal(t, uint64(1), state.registry.LiveAllocationMetadata()) - - reused, err := OptGetParamFromWrapper[types.Decimal128]( - result, - 0, - source, - ) - require.NoError(t, err) - require.Same(t, parameter, reused) - require.Equal(t, first.Used, state.account.Snapshot().Used) - - float32Source := NewOffHeapVecWithType(types.T_float32.ToType()) - require.NoError(t, AppendFixed(float32Source, float32(1.25), false, mp)) - float32Parameter, err := OptGetParamFromWrapper[types.Decimal128]( - result, - 0, - float32Source, - ) - require.NoError(t, err) - expectedFloat32, err := types.Decimal128FromFloat64(1.25, 38, 7) - require.NoError(t, err) - require.Equal(t, expectedFloat32, float32Parameter.UnSafeGetAllValue()[0]) - require.Equal(t, int32(7), float32Parameter.GetType().Scale) - - float64Source := NewOffHeapVecWithType(types.T_float64.ToType()) - require.NoError(t, AppendFixed(float64Source, 2.5, false, mp)) - float64Parameter, err := OptGetParamFromWrapper[types.Decimal128]( - result, - 0, - float64Source, - ) - require.NoError(t, err) - expectedFloat64, err := types.Decimal128FromFloat64(2.5, 38, 16) - require.NoError(t, err) - require.Equal(t, expectedFloat64, float64Parameter.UnSafeGetAllValue()[0]) - require.Equal(t, int32(16), float64Parameter.GetType().Scale) - - constSource, err := NewConstFixed( - types.T_float64.ToType(), - 3.5, - 8, - mp, - ) - require.NoError(t, err) - beforeConst := state.account.Snapshot().Used - constParameter, err := OptGetParamFromWrapper[types.Decimal128]( - result, - 0, - constSource, - ) - require.NoError(t, err) - expectedConst, err := types.Decimal128FromFloat64(3.5, 38, 16) - require.NoError(t, err) - value, isNull := constParameter.GetValue(7) - require.False(t, isNull) - require.Equal(t, expectedConst, value) - require.Equal(t, beforeConst, state.account.Snapshot().Used) - - result.Free() - require.Zero(t, state.account.Snapshot().Used) - source.Free(mp) - float32Source.Free(mp) - float64Source.Free(mp) - constSource.Free(mp) - finalizeTestVectorAllocationAccount(t, state) -} - -func TestFunctionResultAllocationAccountDecimalParameterFailure(t *testing.T) { - state := newTestVectorFunctionAllocationAccount(t, 127, 4) - mp := mpool.MustNew("function-parameter-allocation-failure") - defer mpool.DeleteMPool(mp) - result, err := NewFunctionResultWrapperWithFunctionAllocation( - types.T_bool.ToType(), - mp, - state.selection, - state.function, - ) - require.NoError(t, err) - result.UseOptFunctionParamFrame(1) - - source := NewOffHeapVecWithType(types.T_decimal64.ToType()) - for i := 0; i < 8; i++ { - require.NoError(t, AppendFixed( - source, - types.Decimal64(i), - false, - mp, - )) - } - _, err = OptGetParamFromWrapper[types.Decimal128]( - result, - 0, - source, - ) - require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) - require.Zero(t, state.account.Snapshot().Used) - require.Zero(t, state.registry.LiveAllocationMetadata()) - - source.SetLength(7) - _, err = OptGetParamFromWrapper[types.Decimal128]( - result, - 0, - source, - ) - require.NoError(t, err) - require.Positive(t, state.account.Snapshot().Used) - - result.Free() - source.Free(mp) - finalizeTestVectorAllocationAccount(t, state) -} - -func TestFunctionResultAllocationAccountFunctionScratch(t *testing.T) { - state := newTestVectorFunctionAllocationAccount(t, 1<<20, 8) - mp := mpool.MustNew("function-scratch-allocation") - defer mpool.DeleteMPool(mp) - result, err := NewFunctionResultWrapperWithFunctionAllocation( - types.T_bool.ToType(), - mp, - state.selection, - state.function, - ) - require.NoError(t, err) - - scratch, selected, err := result.ResizeFunctionScratch(128) - require.NoError(t, err) - require.True(t, selected) - require.Len(t, scratch, 128) - used := state.account.Snapshot().Used - require.Positive(t, used) - - scratch, selected, err = result.ResizeFunctionScratch(64) - require.NoError(t, err) - require.True(t, selected) - require.Len(t, scratch, 64) - require.Equal(t, used, state.account.Snapshot().Used) - - result.Free() - require.Zero(t, state.account.Snapshot().Used) - finalizeTestVectorAllocationAccount(t, state) - - legacy := NewFunctionResultWrapper(types.T_bool.ToType(), mp) - scratch, selected, err = legacy.ResizeFunctionScratch(128) - require.NoError(t, err) - require.False(t, selected) - require.Nil(t, scratch) - legacy.Free() -} - -func TestFunctionResultAllocationAccountFunctionScratchFailure(t *testing.T) { - state := newTestVectorFunctionAllocationAccount(t, 127, 2) - mp := mpool.MustNew("function-scratch-allocation-failure") - defer mpool.DeleteMPool(mp) - result, err := NewFunctionResultWrapperWithFunctionAllocation( - types.T_bool.ToType(), - mp, - state.selection, - state.function, - ) - require.NoError(t, err) - - _, selected, err := result.ResizeFunctionScratch(128) - require.True(t, selected) - require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) - require.Zero(t, state.account.Snapshot().Used) - result.Free() - finalizeTestVectorAllocationAccount(t, state) -} diff --git a/pkg/container/vector/pSpoolTools.go b/pkg/container/vector/pSpoolTools.go index d4b46d40583b7..f2bdfecddf4b1 100644 --- a/pkg/container/vector/pSpoolTools.go +++ b/pkg/container/vector/pSpoolTools.go @@ -32,40 +32,6 @@ const ( DetachedAreaBuffer ) -// DetachLegacyVectorData is the allocation-unaccounted spool fast path. The -// explicit guard prevents raw ownership transfer from dropping provenance. -func DetachLegacyVectorData(v *Vector) []byte { - if v.allocationAccount != nil { - panic("cannot detach accounted vector data without provenance") - } - data := v.data - v.data = nil - return data -} - -func DetachLegacyVectorArea(v *Vector) []byte { - if v.allocationAccount != nil { - panic("cannot detach accounted vector area without provenance") - } - area := v.area - v.area = nil - return area -} - -func AttachLegacyVectorData(v *Vector, data []byte) { - if v.allocationAccount != nil || cap(v.data) != 0 { - panic("cannot attach legacy vector data") - } - v.data = data[:cap(data)] -} - -func AttachLegacyVectorArea(v *Vector, area []byte) { - if v.allocationAccount != nil || cap(v.area) != 0 { - panic("cannot attach legacy vector area") - } - v.area = area -} - func DetachVectorData(v *Vector) DetachedBuffer { if v == nil { return DetachedBuffer{} @@ -98,9 +64,9 @@ func (b *DetachedBuffer) Capacity() int { return cap(b.data) } -// CanAttachTo preserves data/area site provenance for accounted allocations. -// Legacy buffers have no site identity and retain the historical ability to -// serve either backing. +// CanAttachTo preserves data/area site provenance when an allocation has an +// account. Unaccounted storage has no site identity and can serve either +// backing without losing ownership information. func (b *DetachedBuffer) CanAttachTo( v *Vector, kind DetachedBufferKind, diff --git a/pkg/container/vector/tools.go b/pkg/container/vector/tools.go index ccce3160db174..6c0b0c1c6cba3 100644 --- a/pkg/container/vector/tools.go +++ b/pkg/container/vector/tools.go @@ -194,13 +194,23 @@ func MustVarlenaToInt64Slice(v *Vector) [][3]int64 { } func MustVarlenaRawData(v *Vector) (data []types.Varlena, area []byte) { - data = MustFixedColNoTypeCheck[types.Varlena](v) + data = ToSliceNoTypeCheck2[types.Varlena](v) area = v.area return } // XXX extend will extend the vector's Data to accommodate rows more entry. func extend(v *Vector, rows int, m *mpool.MPool) error { + return extendWithBitmaps(v, rows, m, false, false) +} + +func extendWithBitmaps( + v *Vector, + rows int, + m *mpool.MPool, + needNulls bool, + needGrouping bool, +) error { if rows <= 0 { // we will at least extent by 1. // This is a pure hack to @@ -208,8 +218,19 @@ func extend(v *Vector, rows int, m *mpool.MPool) error { } tgtLen := v.length + rows - if err := v.ensureBitmapCapacity(tgtLen, m); err != nil { - return err + switch { + case needNulls && needGrouping: + if err := v.ensureBitmapCapacity(tgtLen, m); err != nil { + return err + } + case needNulls: + if err := v.ensureNullCapacity(tgtLen, m); err != nil { + return err + } + case needGrouping: + if err := v.ensureGroupingCapacity(tgtLen, m); err != nil { + return err + } } tgtDataCap := tgtLen * v.typ.TypeSize() if tgtDataCap > cap(v.data) { diff --git a/pkg/container/vector/vector.go b/pkg/container/vector/vector.go index a309e012bfbac..e8cde3a52d936 100644 --- a/pkg/container/vector/vector.go +++ b/pkg/container/vector/vector.go @@ -16,6 +16,7 @@ package vector import ( "bytes" + "encoding/binary" "fmt" "io" "math" @@ -618,7 +619,9 @@ func (v *Vector) IsConst() bool { } func (v *Vector) IsGrouping() bool { - return v.length > 0 && v.length == v.gsp.Count() + return v.length > 0 && + v.length == v.gsp.Count() && + v.length == v.gsp.GetBitmap().CountRange(0, uint64(v.length)) } func (v *Vector) SetClass(class int) { @@ -777,19 +780,33 @@ func (v *Vector) MarshalBinaryWithBuffer(buf *bytes.Buffer) error { return v.MarshalBinaryTo(buf) } -func (v *Vector) MarshalBinarySize() (int, error) { +// MarshalBinaryPlan is a validated, allocation-free snapshot of one Vector's +// wire lengths. It lets batch writers size once and encode once. +type MarshalBinaryPlan struct { + vector *Vector + size int + dataLength uint32 + areaLength uint32 + nullLength uint32 +} + +func (p MarshalBinaryPlan) Size() int { + return p.size +} + +func (v *Vector) PrepareMarshalBinary() (MarshalBinaryPlan, error) { if v == nil || v.length < 0 { - return 0, moerr.NewInvalidInputNoCtx("invalid vector for marshal") + return MarshalBinaryPlan{}, moerr.NewInvalidInputNoCtx("invalid vector for marshal") } const maxWireBuffer = uint64(^uint32(0)) if uint64(v.length) > maxWireBuffer { - return 0, moerr.NewInvalidInputNoCtx( + return MarshalBinaryPlan{}, moerr.NewInvalidInputNoCtx( "vector length exceeds marshal format", ) } typeSize := v.typ.TypeSize() if typeSize < 0 { - return 0, moerr.NewInvalidInputNoCtx( + return MarshalBinaryPlan{}, moerr.NewInvalidInputNoCtx( "vector type has invalid marshal size", ) } @@ -797,7 +814,7 @@ func (v *Vector) MarshalBinarySize() (int, error) { if !v.IsConst() { if v.length != 0 && dataLength > ^uint64(0)/uint64(v.length) { - return 0, moerr.NewInvalidInputNoCtx( + return MarshalBinaryPlan{}, moerr.NewInvalidInputNoCtx( "vector data exceeds marshal format", ) } @@ -810,80 +827,114 @@ func (v *Vector) MarshalBinarySize() (int, error) { if dataLength > maxWireBuffer || areaLength > maxWireBuffer || nullLength > maxWireBuffer { - return 0, moerr.NewInvalidInputNoCtx( + return MarshalBinaryPlan{}, moerr.NewInvalidInputNoCtx( "vector buffer exceeds marshal format", ) } if dataLength > uint64(len(v.data)) { - return 0, moerr.NewInvalidInputNoCtx( + return MarshalBinaryPlan{}, moerr.NewInvalidInputNoCtx( "vector data is shorter than its marshal length", ) } total := uint64(1+types.TSize+4+4+4+4+1) + dataLength + areaLength + nullLength if total > uint64(^uint(0)>>1) { - return 0, moerr.NewInvalidInputNoCtx( + return MarshalBinaryPlan{}, moerr.NewInvalidInputNoCtx( "vector marshal size exceeds platform limit", ) } - return int(total), nil + return MarshalBinaryPlan{ + vector: v, + size: int(total), + dataLength: uint32(dataLength), + areaLength: uint32(areaLength), + nullLength: uint32(nullLength), + }, nil +} + +func (v *Vector) MarshalBinarySize() (int, error) { + plan, err := v.PrepareMarshalBinary() + return plan.Size(), err } func (v *Vector) MarshalBinaryTo(w io.Writer) error { - if w == nil { - return io.ErrClosedPipe - } - if _, err := v.MarshalBinarySize(); err != nil { + plan, err := v.PrepareMarshalBinary() + if err != nil { return err } - if err := writeVectorMarshalBytes(w, []byte{uint8(v.class)}); err != nil { + return plan.MarshalTo(w) +} + +func (p MarshalBinaryPlan) MarshalTo(w io.Writer) error { + v := p.vector + if v == nil || w == nil { + return io.ErrClosedPipe + } + if err := writeVectorMarshalByte(w, uint8(v.class)); err != nil { return err } if err := writeVectorMarshalBytes(w, types.EncodeType(&v.typ)); err != nil { return err } - length := uint32(v.length) - if err := writeVectorMarshalBytes(w, types.EncodeUint32(&length)); err != nil { + if err := writeVectorMarshalUint32(w, uint32(v.length)); err != nil { return err } - dataLen := uint32(v.typ.TypeSize()) - if !v.IsConst() { - dataLen *= uint32(v.length) - } else if v.IsConstNull() { - dataLen = 0 - } - if err := writeVectorMarshalBytes(w, types.EncodeUint32(&dataLen)); err != nil { + if err := writeVectorMarshalUint32(w, p.dataLength); err != nil { return err } - if dataLen > 0 { - if err := writeVectorMarshalBytes(w, v.data[:dataLen]); err != nil { + if p.dataLength > 0 { + if err := writeVectorMarshalBytes(w, v.data[:p.dataLength]); err != nil { return err } } - areaLen := uint32(len(v.area)) - if err := writeVectorMarshalBytes(w, types.EncodeUint32(&areaLen)); err != nil { + if err := writeVectorMarshalUint32(w, p.areaLength); err != nil { return err } - if areaLen > 0 { + if p.areaLength > 0 { if err := writeVectorMarshalBytes(w, v.area); err != nil { return err } } - nspLen := uint32(v.nsp.MarshalSize()) - if err := writeVectorMarshalBytes(w, types.EncodeUint32(&nspLen)); err != nil { + if err := writeVectorMarshalUint32(w, p.nullLength); err != nil { return err } - if nspLen > 0 { + if p.nullLength > 0 { if err := v.nsp.MarshalTo(w); err != nil { return err } } - return writeVectorMarshalBytes(w, types.EncodeBool(&v.sorted)) + if v.sorted { + return writeVectorMarshalByte(w, 1) + } + return writeVectorMarshalByte(w, 0) +} + +type vectorPrimitiveWriter interface { + WriteByte(byte) error + WriteUint32(uint32) error +} + +func writeVectorMarshalByte(w io.Writer, value byte) error { + if typed, ok := w.(vectorPrimitiveWriter); ok { + return typed.WriteByte(value) + } + var data [1]byte + data[0] = value + return writeVectorMarshalBytes(w, data[:]) +} + +func writeVectorMarshalUint32(w io.Writer, value uint32) error { + if typed, ok := w.(vectorPrimitiveWriter); ok { + return typed.WriteUint32(value) + } + var data [4]byte + binary.NativeEndian.PutUint32(data[:], value) + return writeVectorMarshalBytes(w, data[:]) } func writeVectorMarshalBytes(w io.Writer, value []byte) error { @@ -918,96 +969,157 @@ func (v *Vector) UnmarshalBinaryTrusted(data []byte) error { return v.unmarshalBinary(data, false) } -func (v *Vector) unmarshalBinary(data []byte, validateValues bool) error { - if v.allocationAccount != nil { - return allocationAccountInvalid( - "cannot install aliases in an accounted vector", - ) - } - read := func(size int) ([]byte, error) { - if size < 0 || size > len(data) { - return nil, io.ErrUnexpectedEOF - } - value := data[:size] - data = data[size:] - return value, nil - } - readUint32 := func() (uint32, error) { - value, err := read(4) - if err != nil { - return 0, err - } - return types.DecodeUint32(value), nil +type vectorBinaryLayout struct { + class byte + typ types.Type + length int + data []byte + area []byte + nulls []byte + sorted bool +} + +type vectorBinaryCursor struct { + data []byte + offset int +} + +func (c *vectorBinaryCursor) read(size int) ([]byte, error) { + if size < 0 || c.offset > len(c.data)-size { + return nil, io.ErrUnexpectedEOF } + value := c.data[c.offset : c.offset+size] + c.offset += size + return value, nil +} - class, err := read(1) +func (c *vectorBinaryCursor) readUint32() (uint32, error) { + value, err := c.read(4) if err != nil { - return err + return 0, err } - typ, err := read(types.TSize) + return types.DecodeUint32(value), nil +} + +func decodeVectorBinaryLayout( + data []byte, + validateValues bool, +) (vectorBinaryLayout, error) { + cursor := vectorBinaryCursor{data: data} + class, err := cursor.read(1) if err != nil { - return err + return vectorBinaryLayout{}, err } - length, err := readUint32() + typData, err := cursor.read(types.TSize) if err != nil { - return err + return vectorBinaryLayout{}, err } - dataLen, err := readUint32() - if err != nil { - return err + length, err := cursor.readUint32() + if err != nil || uint64(length) > uint64(math.MaxInt) { + if err != nil { + return vectorBinaryLayout{}, err + } + return vectorBinaryLayout{}, moerr.NewInvalidInputNoCtx("vector length exceeds platform limit") } - vecData, err := read(int(dataLen)) - if err != nil { - return err + readSized := func() ([]byte, error) { + size, err := cursor.readUint32() + if err != nil { + return nil, err + } + if uint64(size) > uint64(math.MaxInt) { + return nil, moerr.NewInvalidInputNoCtx("vector buffer exceeds platform limit") + } + return cursor.read(int(size)) } - areaLen, err := readUint32() + vectorData, err := readSized() if err != nil { - return err + return vectorBinaryLayout{}, err } - area, err := read(int(areaLen)) + area, err := readSized() if err != nil { - return err + return vectorBinaryLayout{}, err } - nspLen, err := readUint32() + nullData, err := readSized() if err != nil { - return err + return vectorBinaryLayout{}, err } - nspData, err := read(int(nspLen)) + sorted, err := cursor.read(1) if err != nil { - return err + return vectorBinaryLayout{}, err + } + if cursor.offset != len(cursor.data) { + return vectorBinaryLayout{}, moerr.NewInvalidInputNoCtx("trailing vector wire data") + } + if sorted[0] > 1 { + return vectorBinaryLayout{}, moerr.NewInvalidInputNoCtx("invalid vector sorted flag") + } + if err = validateVectorNullBitmap(nullData, validateValues); err != nil { + return vectorBinaryLayout{}, err + } + var decodedNulls nulls.Nulls + if len(nullData) > 0 { + if err = decodedNulls.ReadNoCopy(nullData); err != nil { + return vectorBinaryLayout{}, err + } + } + typ := types.DecodeType(typData) + if err = validateVectorBinary( + class[0], + typ, + length, + vectorData, + area, + &decodedNulls, + validateValues, + ); err != nil { + return vectorBinaryLayout{}, err + } + return vectorBinaryLayout{ + class: class[0], + typ: typ, + length: int(length), + data: vectorData, + area: area, + nulls: nullData, + sorted: sorted[0] != 0, + }, nil +} + +func (v *Vector) unmarshalBinary(data []byte, validateValues bool) error { + if v == nil { + return io.ErrClosedPipe + } + if v.allocationAccount != nil { + return allocationAccountInvalid( + "cannot install aliases in an accounted vector", + ) } - sorted, err := read(1) + layout, err := decodeVectorBinaryLayout(data, validateValues) if err != nil { return err } - - decodedType := types.DecodeType(typ) - if err := validateVectorNullBitmap(nspData, validateValues); err != nil { - return err + if v.hasBackingStorage() { + return allocationAccountInvalid( + "cannot replace owned vector storage with aliases", + ) } - var nsp nulls.Nulls - if len(nspData) > 0 { - if err := nsp.ReadNoCopy(nspData); err != nil { + var decodedNulls nulls.Nulls + if len(layout.nulls) > 0 { + if err = decodedNulls.ReadNoCopy(layout.nulls); err != nil { return err } } - if err := validateVectorBinary(class[0], decodedType, length, vecData, area, &nsp, validateValues); err != nil { - return err - } - v.class = int(class[0]) - v.typ = decodedType - v.length = int(length) - v.data = vecData - v.area = area - v.nsp = nsp - v.sorted = types.DecodeBool(sorted) - + v.class = int(layout.class) + v.typ = layout.typ + v.length = layout.length + v.data = layout.data + v.area = layout.area + v.nsp = decodedNulls + v.gsp.Reset() + v.sorted = layout.sorted v.cantFreeData = true v.cantFreeArea = true - // The decoded buffers alias the input byte slice. They have no physical - // MPool ownership and therefore cannot retain an allocation selection. v.allocationAccount = nil - return nil } @@ -1131,71 +1243,72 @@ func canonicalVectorTypeSize(typ types.Type) (int, error) { } func (v *Vector) UnmarshalBinaryWithCopy(data []byte, mp *mpool.MPool) error { - if v.allocationAccount != nil && v.hasBackingStorage() { + if v == nil || mp == nil { + return io.ErrClosedPipe + } + if v.hasBackingStorage() { return allocationAccountInvalid( - "cannot replace accounted vector storage without Free", + "cannot replace vector storage without Free", ) } - var err error - - // read class - v.class = int(data[0]) - data = data[1:] - - // read typ - v.typ = types.DecodeType(data[:types.TSize]) - data = data[types.TSize:] - - // read length - v.length = int(types.DecodeUint32(data[:4])) - data = data[4:] - if err = v.ensureBitmapCapacity(v.length, mp); err != nil { + layout, err := decodeVectorBinaryLayout(data, true) + if err != nil { return err } - - // read data - dataLen := int(types.DecodeUint32(data[:4])) - data = data[4:] - if dataLen > 0 { - v.data, err = v.allocData(mp, dataLen) + decoded := NewVec(layout.typ) + decoded.offHeap = v.offHeap + if v.allocationAccount != nil { + if err = decoded.SetAllocationAccount(v.allocationAccount); err != nil { + return err + } + } + committed := false + defer func() { + if !committed { + decoded.Free(mp) + } + }() + decoded.class = int(layout.class) + decoded.length = layout.length + if len(layout.data) > 0 { + decoded.data, err = decoded.allocData(mp, len(layout.data)) if err != nil { return err } - copy(v.data, data[:dataLen]) - data = data[dataLen:] + copy(decoded.data, layout.data) } - - // read area - areaLen := int(types.DecodeUint32(data[:4])) - data = data[4:] - if areaLen > 0 { - v.area, err = v.allocArea(mp, areaLen) + if len(layout.area) > 0 { + decoded.area, err = decoded.allocArea(mp, len(layout.area)) if err != nil { return err } - copy(v.area, data[:areaLen]) - data = data[areaLen:] + copy(decoded.area, layout.area) } - - // read nsp - nspLen := types.DecodeUint32(data[:4]) - data = data[4:] - if nspLen > 0 { - if err := v.nsp.Read(data[:nspLen]); err != nil { + if len(layout.nulls) > 0 { + if decoded.allocationAccount != nil { + _, bitLength, _, decodeErr := bitmap.DecodeMarshalHeader(layout.nulls) + if decodeErr != nil || bitLength > int64(math.MaxInt) { + return moerr.NewInvalidInputNoCtx("invalid vector null bitmap") + } + if err = decoded.ensureNullCapacity(int(bitLength), mp); err != nil { + return err + } + } + if err = decoded.nsp.Read(layout.nulls); err != nil { return err } - data = data[nspLen:] - } else { - v.nsp.Reset() } - - v.sorted = types.DecodeBool(data[:1]) - //data = data[1:] - + decoded.sorted = layout.sorted + *v = *decoded + committed = true return nil } func (v *Vector) UnmarshalWithReader(r io.Reader, mp *mpool.MPool) error { + if v == nil || r == nil { + return io.ErrClosedPipe + } + v.ResetWithSameType() var err error if v.class, err = types.ReadByteAsInt(r); err != nil { @@ -1209,8 +1322,11 @@ func (v *Vector) UnmarshalWithReader(r io.Reader, mp *mpool.MPool) error { if v.length, err = types.ReadInt32AsInt(r); err != nil { return err } - if err = v.ensureBitmapCapacity(v.length, mp); err != nil { - return err + if v.length < 0 { + return moerr.NewInvalidInputNoCtx("negative vector length") + } + if v.length > math.MaxUint32 { + return moerr.NewInvalidInputNoCtx("vector length exceeds marshal format") } // read data @@ -1231,7 +1347,7 @@ func (v *Vector) UnmarshalWithReader(r io.Reader, mp *mpool.MPool) error { v.area = areaBuf } - if err = v.readNullsWithReader(r); err != nil { + if err = v.readNullsWithReader(r, mp); err != nil { return err } @@ -1239,17 +1355,37 @@ func (v *Vector) UnmarshalWithReader(r io.Reader, mp *mpool.MPool) error { if err != nil { return err } - - return nil + return validateVectorBinary( + byte(v.class), + v.typ, + uint32(v.length), + v.data[:int(dataLen)], + v.area[:int(areaLen)], + &v.nsp, + true, + ) } -func (v *Vector) readNullsWithReader(r io.Reader) error { - if v.allocationAccount == nil || !v.allocationAccount.accountBitmaps { - nspLen, nspBuf, err := types.ReadSizeBytes(r) +func (v *Vector) readNullsWithReader(r io.Reader, mp *mpool.MPool) error { + if v.allocationAccount == nil { + nspLen, err := types.ReadInt32(r) if err != nil { return err } + if nspLen < 0 { + return moerr.NewInvalidInputNoCtx("negative vector null bitmap size") + } + if err = validateStreamingReadSize(r, int64(nspLen)); err != nil { + return err + } if nspLen > 0 { + nspBuf := make([]byte, nspLen) + if _, err = io.ReadFull(r, nspBuf); err != nil { + return err + } + if err := validateVectorNullBitmap(nspBuf, true); err != nil { + return err + } return v.nsp.Read(nspBuf) } v.nsp.Reset() @@ -1272,9 +1408,15 @@ func (v *Vector) readNullsWithReader(r io.Reader) error { return err } _, bitLength, _, err := bitmap.DecodeMarshalHeader(header[:]) - if err != nil || bitLength > int64(v.length)+1 { + if err != nil { return moerr.NewInvalidInputNoCtx("invalid vector null bitmap") } + if bitLength > int64(math.MaxInt) { + return moerr.NewInvalidInputNoCtx("vector null bitmap exceeds platform limit") + } + if err = v.ensureNullCapacity(int(bitLength), mp); err != nil { + return err + } payload, err := v.nsp.GetBitmap().PrepareExternalUnmarshal( header[:], int(size), @@ -1284,8 +1426,77 @@ func (v *Vector) readNullsWithReader(r io.Reader) error { } if _, err = io.ReadFull(r, payload); err != nil { v.nsp.Reset() + return err } - return err + return v.nsp.GetBitmap().Validate() +} + +// GroupingMarshalBinarySize returns the optional grouping bitmap wire size. +func (v *Vector) GroupingMarshalBinarySize() int { + if v == nil { + return 0 + } + return v.gsp.MarshalSize() +} + +// MarshalGroupingTo writes the optional grouping bitmap without changing the +// stable Vector wire format. +func (v *Vector) MarshalGroupingTo(w io.Writer) error { + if v == nil || w == nil { + return io.ErrClosedPipe + } + return v.gsp.MarshalTo(w) +} + +// UnmarshalGroupingFromReader restores a grouping bitmap whose size is framed +// by the caller. +func (v *Vector) UnmarshalGroupingFromReader( + r io.Reader, + size int, + mp *mpool.MPool, +) error { + if v == nil || r == nil || size < 0 { + return moerr.NewInvalidInputNoCtx("invalid vector grouping bitmap") + } + if size == 0 { + v.gsp.Reset() + return nil + } + if size < bitmap.MarshalHeaderSize { + return moerr.NewInvalidInputNoCtx("invalid vector grouping bitmap") + } + var header [bitmap.MarshalHeaderSize]byte + if _, err := io.ReadFull(r, header[:]); err != nil { + return err + } + _, bitLength, _, err := bitmap.DecodeMarshalHeader(header[:]) + if err != nil || bitLength > int64(math.MaxInt) { + return moerr.NewInvalidInputNoCtx("invalid vector grouping bitmap") + } + if v.allocationAccount == nil { + data := make([]byte, size) + copy(data, header[:]) + if _, err = io.ReadFull(r, data[len(header):]); err != nil { + return err + } + if err = validateVectorNullBitmap(data, true); err != nil { + return err + } + v.gsp.Reset() + return v.gsp.Read(data) + } + if err = v.ensureGroupingCapacity(int(bitLength), mp); err != nil { + return err + } + payload, err := v.gsp.GetBitmap().PrepareExternalUnmarshal(header[:], size) + if err != nil { + return err + } + if _, err = io.ReadFull(r, payload); err != nil { + v.gsp.Reset() + return err + } + return v.gsp.GetBitmap().Validate() } func (v *Vector) ToConst() { @@ -1303,11 +1514,24 @@ func (v *Vector) PreExtend(rows int, mp *mpool.MPool) error { } // PreExtendBitmap ensures allocation-accounted null and grouping storage can -// represent rows without allocating vector data. Legacy vectors are unchanged. +// represent rows without allocating vector data. Unaccounted vectors are +// unchanged. func (v *Vector) PreExtendBitmap(rows int, mp *mpool.MPool) error { return v.ensureBitmapCapacity(rows, mp) } +// PreExtendNulls ensures allocation-accounted null storage can represent rows. +// Unaccounted vectors are unchanged. +func (v *Vector) PreExtendNulls(rows int, mp *mpool.MPool) error { + return v.ensureNullCapacity(rows, mp) +} + +// PreExtendGrouping ensures allocation-accounted grouping storage can +// represent rows. Unaccounted vectors are unchanged. +func (v *Vector) PreExtendGrouping(rows int, mp *mpool.MPool) error { + return v.ensureGroupingCapacity(rows, mp) +} + // PreExtendArea use to expand the mpool and area of vector // extraAreaSize: the size of area to be extended // mp: mpool @@ -1359,7 +1583,7 @@ func (v *Vector) DupOffHeap(mp *mpool.MPool) (*Vector, error) { } // DupOffHeapWithAllocation copies a vector into an explicitly selected -// destination account. Passing nil creates a legacy unaccounted destination. +// destination account. Passing nil creates an unaccounted destination. func (v *Vector) DupOffHeapWithAllocation( mp *mpool.MPool, selection *AllocationAccountSelection, @@ -1392,7 +1616,7 @@ func (v *Vector) dup( w.Free(mp) return nil, mpool.ErrAllocationAccountInvalid } - if err := w.ensureBitmapCapacity( + if err := w.ensureGroupingCapacity( int(groupingRows), mp, ); err != nil { @@ -1418,17 +1642,25 @@ func (v *Vector) dup( } dataLen *= v.length } - bitmapRows := max( - v.GetNulls().GetBitmap().Len(), - v.GetGrouping().GetBitmap().Len(), - ) - if bitmapRows < 0 || bitmapRows > int64(math.MaxInt) { - w.Free(mp) - return nil, mpool.ErrAllocationAccountInvalid + if nullRows := v.GetNulls().GetBitmap().Len(); !v.GetNulls().EmptyByFlag() { + if nullRows < 0 || nullRows > int64(math.MaxInt) { + w.Free(mp) + return nil, mpool.ErrAllocationAccountInvalid + } + if err := w.ensureNullCapacity(int(nullRows), mp); err != nil { + w.Free(mp) + return nil, err + } } - if err := w.ensureBitmapCapacity(int(bitmapRows), mp); err != nil { - w.Free(mp) - return nil, err + if groupingRows := v.GetGrouping().GetBitmap().Len(); !v.GetGrouping().EmptyByFlag() { + if groupingRows < 0 || groupingRows > int64(math.MaxInt) { + w.Free(mp) + return nil, mpool.ErrAllocationAccountInvalid + } + if err := w.ensureGroupingCapacity(int(groupingRows), mp); err != nil { + w.Free(mp) + return nil, err + } } w.length = v.length w.GetNulls().InitWith(v.GetNulls()) @@ -1492,7 +1724,13 @@ func (v *Vector) cloneToFlatCompact( if v.length == 0 { return w, nil } - if err := extend(w, v.length, mp); err != nil { + if err := extendWithBitmaps( + w, + v.length, + mp, + !v.nsp.EmptyByFlag(), + !v.gsp.EmptyByFlag(), + ); err != nil { w.Free(mp) return nil, err } @@ -1855,6 +2093,22 @@ func (v *Vector) ShuffleWithBuf(sels []int64, mp *mpool.MPool, buf *[]byte) (err // Copy simply does v[vi] = w[wi] func (v *Vector) Copy(w *Vector, vi, wi int64, mp *mpool.MPool) error { + sourceGrouping := w.GetGrouping().Contains(uint64(wi)) + if sourceGrouping { + if err := v.ensureGroupingCapacity(int(vi)+1, mp); err != nil { + return err + } + v.GetGrouping().Set(uint64(vi)) + } else { + v.GetGrouping().Unset(uint64(vi)) + } + sourceNull := w.IsConstNull() || + (!w.IsConst() && w.GetNulls().Contains(uint64(wi))) + if sourceNull { + if err := v.ensureNullCapacity(int(vi)+1, mp); err != nil { + return err + } + } if w.class == CONSTANT { if w.IsConstNull() { if !v.typ.IsFixedLen() { @@ -1902,6 +2156,38 @@ func (v *Vector) Copy(w *Vector, vi, wi int64, mp *mpool.MPool) error { // GetUnionAllFunction: A more sensible function for copying vector, // which avoids having to do type conversions and type judgements every time you append. func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) error { + union := getUnionAllFunction(typ, mp) + return func(v, w *Vector) error { + oldLength := v.length + if w.gsp.Any() { + if err := v.ensureGroupingCapacity(oldLength+w.length, mp); err != nil { + return err + } + } + if err := union(v, w); err != nil { + return err + } + if w.gsp.Any() { + unionVectorBitmap(&v.gsp, &w.gsp, oldLength, w.length) + } + return nil + } +} + +func unionVectorBitmap( + destination *nulls.Nulls, + source *nulls.Nulls, + offset int, + length int, +) { + for row := 0; row < length; row++ { + if source.Contains(uint64(row)) { + destination.Set(uint64(offset + row)) + } + } +} + +func getUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) error { // a more simple and quickly union nsp but not good. unionNsp := func(dst *nulls.Nulls, more *nulls.Nulls, oldLength int, moreLength int) { u64offset := uint64(oldLength) @@ -1938,7 +2224,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -1967,7 +2253,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -1996,7 +2282,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2025,7 +2311,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2054,7 +2340,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2083,7 +2369,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2112,7 +2398,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2141,7 +2427,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2170,7 +2456,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2199,7 +2485,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2228,7 +2514,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2257,7 +2543,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2286,7 +2572,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2315,7 +2601,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2344,7 +2630,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2373,7 +2659,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2402,7 +2688,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2431,7 +2717,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2460,7 +2746,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2489,7 +2775,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2518,7 +2804,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2547,7 +2833,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2576,7 +2862,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2605,7 +2891,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2636,7 +2922,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if sz := len(v.area) + len(w.area); sz > cap(v.area) { @@ -2649,13 +2935,13 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err var err error vs := toSliceOfLengthNoTypeCheck[types.Varlena](v, v.length+w.length) + if w.gsp.Any() { + unionNsp(&v.gsp, &w.gsp, v.length, w.length) + } bm := w.nsp.GetBitmap() if bm != nil && !bm.EmptyByFlag() { for i := range ws { - if w.gsp.Contains(uint64(i)) { - nulls.Add(&v.gsp, uint64(v.length)) - } if bm.Contains(uint64(i)) { nulls.Add(&v.nsp, uint64(v.length)) } else { @@ -2692,7 +2978,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2713,7 +2999,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err // GetConstSetFunction: A more sensible function for const vector set, // which avoids having to do type conversions and type judgements every time you append. -func GetConstSetFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector, sel int64, length int) error { +func getConstSetFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector, sel int64, length int) error { switch typ.Oid { case types.T_bool: return func(v, w *Vector, sel int64, length int) error { @@ -3008,6 +3294,29 @@ func GetConstSetFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector, sel } } +func GetConstSetFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector, sel int64, length int) error { + set := getConstSetFunction(typ, mp) + return func(v, w *Vector, sel int64, length int) error { + if v == nil || w == nil || sel < 0 || sel >= int64(w.Length()) || length < 0 { + return moerr.NewInvalidInputNoCtx("invalid const vector selection") + } + grouping := w.gsp.Contains(uint64(sel)) + if grouping { + if err := v.ensureGroupingCapacity(length, mp); err != nil { + return err + } + } + if err := set(v, w, sel, length); err != nil { + return err + } + v.gsp.Reset() + if grouping && length > 0 { + v.gsp.AddRange(0, uint64(length)) + } + return nil + } +} + // fillSlice broadcasts val across s[start:end] using exponential copy doubling: // write one element, then double the filled region with copy() — O(log n) memmoves // instead of n scalar element stores. Used on the hot const-broadcast path. @@ -3059,7 +3368,10 @@ func (v *Vector) UnionNull(mp *mpool.MPool) error { // It is simply append. the purpose of retention is ease of use func (v *Vector) UnionOne(w *Vector, sel int64, mp *mpool.MPool) error { - if err := extend(v, 1, mp); err != nil { + needGrouping := nulls.Contains(&w.gsp, uint64(sel)) + needNulls := w.IsConstNull() || + (!w.IsConst() && nulls.Contains(&w.nsp, uint64(sel))) + if err := extendWithBitmaps(v, 1, mp, needNulls, needGrouping); err != nil { return err } @@ -3078,7 +3390,6 @@ func (v *Vector) UnionOne(w *Vector, sel int64, mp *mpool.MPool) error { nulls.Add(&v.nsp, uint64(oldLen)) return nil } - if v.GetType().IsVarlen() { var vs, ws []types.Varlena ToSliceNoTypeCheck(v, &vs) @@ -3112,13 +3423,32 @@ func (v *Vector) UnionOne(w *Vector, sel int64, mp *mpool.MPool) error { return nil } +func appendSelectedGrouping[T int32 | int64]( + dst *Vector, + src *Vector, + oldLength int, + sels []T, +) { + if src.gsp.EmptyByFlag() { + return + } + for i, sel := range sels { + if src.gsp.Contains(uint64(sel)) { + nulls.Add(&dst.gsp, uint64(oldLength+i)) + } + } +} + // It is simply append. the purpose of retention is ease of use func (v *Vector) UnionMulti(w *Vector, sel int64, cnt int, mp *mpool.MPool) error { if cnt == 0 { return nil } - if err := extend(v, cnt, mp); err != nil { + needGrouping := nulls.Contains(&w.gsp, uint64(sel)) + needNulls := w.IsConstNull() || + (!w.IsConst() && nulls.Contains(&w.nsp, uint64(sel))) + if err := extendWithBitmaps(v, cnt, mp, needNulls, needGrouping); err != nil { return err } @@ -3137,7 +3467,6 @@ func (v *Vector) UnionMulti(w *Vector, sel int64, cnt int, mp *mpool.MPool) erro nulls.AddRange(&v.nsp, uint64(oldLen), uint64(oldLen+cnt)) return nil } - if v.GetType().IsVarlen() { var err error var va types.Varlena @@ -3159,6 +3488,37 @@ func (v *Vector) UnionMulti(w *Vector, sel int64, cnt int, mp *mpool.MPool) erro return nil } +func appendBatchGrouping( + dst *Vector, + src *Vector, + oldLength int, + offset int64, + cnt int, + flags []uint8, +) { + if src.gsp.EmptyByFlag() { + return + } + output := oldLength + if flags == nil { + for i := range cnt { + if src.gsp.Contains(uint64(offset) + uint64(i)) { + nulls.Add(&dst.gsp, uint64(output+i)) + } + } + return + } + for i, selected := range flags { + if selected == 0 { + continue + } + if src.gsp.Contains(uint64(offset) + uint64(i)) { + nulls.Add(&dst.gsp, uint64(output)) + } + output++ + } +} + func (v *Vector) Union(w *Vector, sels []int64, mp *mpool.MPool) error { return unionT[int64](v, w, sels, mp) } @@ -3171,7 +3531,13 @@ func unionT[T int32 | int64](v, w *Vector, sels []T, mp *mpool.MPool) error { return nil } - if err := extend(v, len(sels), mp); err != nil { + if err := extendWithBitmaps( + v, + len(sels), + mp, + w.IsConstNull() || !w.nsp.EmptyByFlag(), + w.IsGrouping() || !w.gsp.EmptyByFlag(), + ); err != nil { return err } @@ -3203,6 +3569,7 @@ func unionT[T int32 | int64](v, w *Vector, sels []T, mp *mpool.MPool) error { return nil } + appendSelectedGrouping(v, w, oldLen, sels) if v.GetType().IsVarlen() { var err error @@ -3230,9 +3597,6 @@ func unionT[T int32 | int64](v, w *Vector, sels []T, mp *mpool.MPool) error { } if !w.GetNulls().EmptyByFlag() { for i, sel := range sels { - if w.gsp.Contains(uint64(sel)) { - nulls.Add(&v.gsp, uint64(oldLen+i)) - } if w.nsp.Contains(uint64(sel)) { nulls.Add(&v.nsp, uint64(oldLen+i)) continue @@ -3255,9 +3619,6 @@ func unionT[T int32 | int64](v, w *Vector, sels []T, mp *mpool.MPool) error { tlen := v.GetType().TypeSize() if !w.nsp.EmptyByFlag() { for i, sel := range sels { - if w.gsp.Contains(uint64(sel)) { - nulls.Add(&v.gsp, uint64(oldLen+i)) - } if w.nsp.Contains(uint64(sel)) { nulls.Add(&v.nsp, uint64(oldLen+i)) continue @@ -3313,7 +3674,13 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp return nil } - if err := extend(v, addCnt, mp); err != nil { + if err := extendWithBitmaps( + v, + addCnt, + mp, + w.IsConstNull() || !w.nsp.EmptyByFlag(), + w.IsGrouping() || !w.gsp.EmptyByFlag(), + ); err != nil { return err } @@ -3345,6 +3712,7 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp return nil } + appendBatchGrouping(v, w, v.length, offset, cnt, flags) if v.GetType().IsVarlen() { var err error @@ -3398,21 +3766,6 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp } } } - // propagate grouping bits (value is still real for these rows). - // Bound to [0,cnt): Foreach walks every set bit in the underlying - // bitmap, but w may carry stale bits at index >= w.length (SetLength - // shrinks length without clearing nsp/gsp, and vectors are reused). - // The per-row path only consults [0,cnt) via Contains, so we must skip - // stale bits here too — otherwise they pollute v.gsp / index past vCol. - if !w.gsp.EmptyByFlag() { - base, ucnt := uint64(oldLen), uint64(cnt) - w.gsp.Foreach(func(i uint64) bool { - if i < ucnt { - nulls.Add(&v.gsp, base+i) - } - return true - }) - } // propagate null bits and clear those (never-read) headers so a copied // big-header offset can't linger as a dangling reference into v.area. // Same [0,cnt) bound as gsp above: a stale nsp bit at i >= cnt would @@ -3472,9 +3825,6 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp if !w.nsp.EmptyByFlag() { if flags == nil { for i := 0; i < cnt; i++ { - if w.gsp.Contains(uint64(offset) + uint64(i)) { - nulls.Add(&v.gsp, uint64(v.length)) - } if w.nsp.Contains(uint64(offset) + uint64(i)) { nulls.Add(&v.nsp, uint64(v.length)) } else { @@ -3490,9 +3840,6 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp if flags[i] == 0 { continue } - if w.gsp.Contains(uint64(offset) + uint64(i)) { - nulls.Add(&v.gsp, uint64(v.length)) - } if w.nsp.Contains(uint64(offset) + uint64(i)) { nulls.Add(&v.nsp, uint64(v.length)) } else { @@ -3507,9 +3854,6 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp } else { if flags == nil { for i := 0; i < cnt; i++ { - if w.gsp.Contains(uint64(offset) + uint64(i)) { - nulls.Add(&v.gsp, uint64(v.length)) - } err = BuildVarlenaFromVarlena(v, &vCol[v.length], &wCol[int(offset)+i], &w.area, mp) if err != nil { return err @@ -3521,9 +3865,6 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp if flags[i] == 0 { continue } - if w.gsp.Contains(uint64(offset) + uint64(i)) { - nulls.Add(&v.gsp, uint64(v.length)) - } err = BuildVarlenaFromVarlena(v, &vCol[v.length], &wCol[int(offset)+i], &w.area, mp) if err != nil { return err @@ -3537,9 +3878,6 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp if !w.nsp.EmptyByFlag() { if flags == nil { for i := 0; i < cnt; i++ { - if w.gsp.Contains(uint64(offset) + uint64(i)) { - nulls.Add(&v.gsp, uint64(v.length)) - } if w.nsp.Contains(uint64(offset) + uint64(i)) { nulls.Add(&v.nsp, uint64(v.length)) } else { @@ -3552,9 +3890,6 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp if flags[i] == 0 { continue } - if w.gsp.Contains(uint64(offset) + uint64(i)) { - nulls.Add(&v.gsp, uint64(v.length)) - } if w.nsp.Contains(uint64(offset) + uint64(i)) { nulls.Add(&v.nsp, uint64(v.length)) } else { @@ -3582,9 +3917,6 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp if flags[i] == 0 { continue } - if w.gsp.Contains(uint64(offset) + uint64(i)) { - nulls.Add(&v.gsp, uint64(v.length)) - } copy(v.data[v.length*tlen:(v.length+1)*tlen], w.data[(int(offset)+i)*tlen:(int(offset)+i+1)*tlen]) v.length++ } @@ -4241,7 +4573,7 @@ func appendOneFixed[T any](vec *Vector, val T, isNull bool, mp *mpool.MPool) err return moerr.NewInternalErrorNoCtx("append to const vector") } - if err := extend(vec, 1, mp); err != nil { + if err := extendWithBitmaps(vec, 1, mp, isNull, false); err != nil { return err } length := vec.length @@ -4303,7 +4635,7 @@ func appendOneArray[T types.ArrayElement](vec *Vector, val []T, isNull bool, mp } func appendMultiFixed[T any](vec *Vector, val T, isNull bool, cnt int, mp *mpool.MPool) error { - if err := extend(vec, cnt, mp); err != nil { + if err := extendWithBitmaps(vec, cnt, mp, isNull, false); err != nil { return err } length := vec.length @@ -4322,7 +4654,7 @@ func appendMultiFixed[T any](vec *Vector, val T, isNull bool, cnt int, mp *mpool func appendMultiBytes(vec *Vector, val []byte, isNull bool, cnt int, mp *mpool.MPool) error { var err error var va types.Varlena - if err = extend(vec, cnt, mp); err != nil { + if err = extendWithBitmaps(vec, cnt, mp, isNull, false); err != nil { return err } length := vec.length @@ -4344,7 +4676,13 @@ func appendMultiBytes(vec *Vector, val []byte, isNull bool, cnt int, mp *mpool.M } func appendList[T any](vec *Vector, vals []T, isNulls []bool, mp *mpool.MPool) error { - if err := extend(vec, len(vals), mp); err != nil { + if err := extendWithBitmaps( + vec, + len(vals), + mp, + slices.Contains(isNulls, true), + false, + ); err != nil { return err } length := vec.length @@ -4362,7 +4700,13 @@ func appendList[T any](vec *Vector, vals []T, isNulls []bool, mp *mpool.MPool) e func appendBytesList(vec *Vector, vals [][]byte, isNulls []bool, mp *mpool.MPool) error { var err error - if err = extend(vec, len(vals), mp); err != nil { + if err = extendWithBitmaps( + vec, + len(vals), + mp, + slices.Contains(isNulls, true), + false, + ); err != nil { return err } length := vec.length @@ -4384,7 +4728,13 @@ func appendBytesList(vec *Vector, vals [][]byte, isNulls []bool, mp *mpool.MPool func appendStringList(vec *Vector, vals []string, isNulls []bool, mp *mpool.MPool) error { var err error - if err = extend(vec, len(vals), mp); err != nil { + if err = extendWithBitmaps( + vec, + len(vals), + mp, + slices.Contains(isNulls, true), + false, + ); err != nil { return err } length := vec.length @@ -4408,7 +4758,13 @@ func appendStringList(vec *Vector, vals []string, isNulls []bool, mp *mpool.MPoo func appendArrayList[T types.ArrayElement](vec *Vector, vals [][]T, isNulls []bool, mp *mpool.MPool) error { var err error - if err = extend(vec, len(vals), mp); err != nil { + if err = extendWithBitmaps( + vec, + len(vals), + mp, + slices.Contains(isNulls, true), + false, + ); err != nil { return err } length := vec.length @@ -4545,7 +4901,7 @@ func (s *bitmapRemapScratch) release(mp *mpool.MPool) { // storage before publishing either, so rejection cannot leave null and // grouping ownership half-mutated. func (v *Vector) remapShuffleBitmaps(sels []int64, mp *mpool.MPool) error { - if v.allocationAccount == nil || !v.allocationAccount.accountBitmaps { + if v.allocationAccount == nil { nulls.Filter(&v.gsp, sels, false) nulls.Filter(&v.nsp, sels, false) return nil @@ -4562,8 +4918,15 @@ func (v *Vector) remapShuffleBitmaps(sels []int64, mp *mpool.MPool) error { targets[1].destination.EmptyByFlag() { return nil } - if err := v.ensureBitmapCapacity(len(sels), mp); err != nil { - return err + if !targets[0].destination.EmptyByFlag() { + if err := v.ensureGroupingCapacity(len(sels), mp); err != nil { + return err + } + } + if !targets[1].destination.EmptyByFlag() { + if err := v.ensureNullCapacity(len(sels), mp); err != nil { + return err + } } var scratch [2]bitmapRemapScratch @@ -4654,26 +5017,57 @@ func vecToString[T types.FixedSizeT](v *Vector) string { // The returned object is NOT allowed to be modified ( // TODO: Nulls are deep copied. func (v *Vector) Window(start, end int) (*Vector, error) { - if v.IsConstNull() { - return NewConstNull(v.typ, end-start, nil), nil - } else if v.IsConst() { - vec := NewVec(v.typ) - vec.class = v.class - vec.data = v.data - vec.area = v.area - vec.length = end - start - vec.cantFreeArea = true - vec.cantFreeData = true - vec.sorted = v.sorted - return vec, nil + return v.window(start, end, nil, nil) +} + +// WindowWithAllocation returns a borrowed data window whose range bitmaps are +// physical allocations in selection. Accounted pressure paths must use this +// form so shrinking an operation cannot create invisible Go-heap owners. +func (v *Vector) WindowWithAllocation( + start int, + end int, + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (*Vector, error) { + if mp == nil || selection == nil { + return nil, mpool.ErrAllocationAccountInvalid + } + return v.window(start, end, mp, selection) +} + +func (v *Vector) window( + start int, + end int, + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (*Vector, error) { + if start < 0 || end < start || end > v.Length() { + return nil, moerr.NewInvalidInputNoCtx("invalid vector window") } w := NewVec(v.typ) - if start == end { - return w, nil + if selection != nil { + w.offHeap = true + if err := w.SetAllocationAccount(selection); err != nil { + return nil, err + } } - nulls.Range(&v.nsp, uint64(start), uint64(end), uint64(start), &w.nsp) - w.data = v.data[start*v.typ.TypeSize() : end*v.typ.TypeSize()] + w.class = v.class w.length = end - start + w.sorted = v.sorted + if err := v.copyWindowBitmaps(w, start, end, mp); err != nil { + w.Free(mp) + return nil, err + } + if v.IsConst() { + w.data = v.data + w.area = v.area + w.cantFreeArea = true + w.cantFreeData = true + return w, nil + } + if start != end { + w.data = v.data[start*v.typ.TypeSize() : end*v.typ.TypeSize()] + } if v.typ.IsVarlen() { w.area = v.area } @@ -4682,6 +5076,25 @@ func (v *Vector) Window(start, end int) (*Vector, error) { return w, nil } +func (v *Vector) copyWindowBitmaps(w *Vector, start, end int, mp *mpool.MPool) error { + length := end - start + hasNull := v.nsp.GetBitmap().CountRange(uint64(start), uint64(end)) > 0 + hasGrouping := v.gsp.GetBitmap().CountRange(uint64(start), uint64(end)) > 0 + if hasNull { + if err := w.PreExtendNulls(length, mp); err != nil { + return err + } + nulls.Range(&v.nsp, uint64(start), uint64(end), uint64(start), &w.nsp) + } + if hasGrouping { + if err := w.PreExtendGrouping(length, mp); err != nil { + return err + } + nulls.Range(&v.gsp, uint64(start), uint64(end), uint64(start), &w.gsp) + } + return nil +} + // CloneWindow Deep copies the content from start to end into another vector. Afterwise it's safe to destroy the original one. func (v *Vector) CloneWindow(start, end int, mp *mpool.MPool) (*Vector, error) { return v.CloneWindowWithAllocation( @@ -4731,6 +5144,9 @@ func (v *Vector) CloneWindowTo(w *Vector, start, end int, mp *mpool.MPool) error if start == end { return nil } + if err := v.copyWindowBitmaps(w, start, end, mp); err != nil { + return err + } if v.IsConstNull() { w.class = CONSTANT w.length = end - start @@ -4761,10 +5177,6 @@ func (v *Vector) CloneWindowTo(w *Vector, start, end int, mp *mpool.MPool) error return nil } } - if err := w.PreExtendBitmap(end-start, mp); err != nil { - return err - } - nulls.Range(&v.nsp, uint64(start), uint64(end), uint64(start), &w.nsp) length := (end - start) * v.typ.TypeSize() if mp == nil { if w.allocationAccount != nil { diff --git a/pkg/container/vector/vector_test.go b/pkg/container/vector/vector_test.go index cd034561dc54d..cd8ccd3c8f73a 100644 --- a/pkg/container/vector/vector_test.go +++ b/pkg/container/vector/vector_test.go @@ -1484,6 +1484,21 @@ func TestShuffle(t *testing.T) { func TestCopy(t *testing.T) { mp := mpool.MustNewZero() + { // fixed grouping provenance + dst := NewVec(types.T_int32.ToType()) + src := NewVec(types.T_int32.ToType()) + require.NoError(t, AppendFixedList(dst, []int32{0, 0}, nil, mp)) + require.NoError(t, AppendFixedList(src, []int32{1, 2}, nil, mp)) + src.GetGrouping().Add(0) + dst.GetGrouping().Add(1) + require.NoError(t, dst.Copy(src, 0, 0, mp)) + require.NoError(t, dst.Copy(src, 1, 1, mp)) + require.True(t, dst.GetGrouping().Contains(0)) + require.False(t, dst.GetGrouping().Contains(1)) + dst.Free(mp) + src.Free(mp) + require.Equal(t, int64(0), mp.CurrNB()) + } { // fixed v := NewVec(types.T_int8.ToType()) AppendFixedList(v, []int8{0, 0, 1, 0}, nil, mp) diff --git a/pkg/sql/colexec/aggexec/maxby_test.go b/pkg/sql/colexec/aggexec/maxby_test.go index 5876cbf268288..dfdeac042403f 100644 --- a/pkg/sql/colexec/aggexec/maxby_test.go +++ b/pkg/sql/colexec/aggexec/maxby_test.go @@ -91,6 +91,8 @@ func TestCompactMaxByStateVectorPreservesAllocationOwner(t *testing.T) { mpool.AllocationOwner(1), mpool.AllocationSite(1), mpool.AllocationSite(2), + mpool.AllocationSite(3), + mpool.AllocationSite(4), ) require.NoError(t, err) vec := vector.NewOffHeapVecWithType(types.T_varchar.ToType()) diff --git a/pkg/sql/colexec/allocation_state.go b/pkg/sql/colexec/allocation_state.go new file mode 100644 index 0000000000000..f1a2a9c7a78fb --- /dev/null +++ b/pkg/sql/colexec/allocation_state.go @@ -0,0 +1,65 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package colexec + +import ( + "math" + + "github.com/matrixorigin/matrixone/pkg/common/bitmap" + "github.com/matrixorigin/matrixone/pkg/common/mpool" +) + +// NewAccountedBitmap creates a bitmap whose complete backing capacity belongs +// to the statement allocation account. The caller owns the returned bitmap +// until FreeAccountedBitmap or an explicit ownership transfer. +func NewAccountedBitmap( + rows int64, + mp *mpool.MPool, + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, + site mpool.AllocationSite, +) (*bitmap.Bitmap, error) { + if rows < 0 || rows > math.MaxInt64-63 || mp == nil || account == nil { + return nil, mpool.ErrAllocationAccountInvalid + } + words := (rows + 63) / 64 + if words > int64(math.MaxInt) { + return nil, mpool.ErrAllocationAccountInvalid + } + storage, err := mpool.MakeSliceAccounted[uint64]( + int(words), + mp, + account, + owner, + site, + ) + if err != nil { + return nil, err + } + value := new(bitmap.Bitmap) + value.InstallExternalStorage(storage) + value.InitWithSize(rows) + return value, nil +} + +func FreeAccountedBitmap(value *bitmap.Bitmap, mp *mpool.MPool) { + if value == nil || !value.HasExternalStorage() { + return + } + storage := value.ReleaseExternalStorage() + if cap(storage) > 0 { + mpool.FreeSlice(mp, storage) + } +} diff --git a/pkg/sql/colexec/dedupjoin/allocation_test_helpers_test.go b/pkg/sql/colexec/dedupjoin/allocation_test_helpers_test.go new file mode 100644 index 0000000000000..ef44e8bb95c2e --- /dev/null +++ b/pkg/sql/colexec/dedupjoin/allocation_test_helpers_test.go @@ -0,0 +1,38 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package dedupjoin + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/stretchr/testify/require" +) + +type testAllocationOwner interface { + SetAllocationAccount(*mpool.AllocationAccount) error +} + +func installTestAllocation(t testing.TB, owners ...testAllocationOwner) *mpool.AllocationAccount { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.Open(1 << 60) + require.NoError(t, err) + for _, owner := range owners { + require.NoError(t, owner.SetAllocationAccount(account)) + } + return account +} diff --git a/pkg/sql/colexec/dedupjoin/expression_memory_test.go b/pkg/sql/colexec/dedupjoin/expression_memory_test.go deleted file mode 100644 index 535a32856f68f..0000000000000 --- a/pkg/sql/colexec/dedupjoin/expression_memory_test.go +++ /dev/null @@ -1,108 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package dedupjoin - -import ( - "testing" - - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/container/batch" - "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/matrixorigin/matrixone/pkg/container/vector" - "github.com/matrixorigin/matrixone/pkg/pb/plan" - "github.com/matrixorigin/matrixone/pkg/sql/colexec" - "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" - "github.com/matrixorigin/matrixone/pkg/testutil" - "github.com/matrixorigin/matrixone/pkg/vm/process" - "github.com/stretchr/testify/require" -) - -func TestDedupJoinResetReleasesProbeExpressionLease(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Lit{Lit: &plan.Literal{ - Value: &plan.Literal_I32Val{I32Val: 1}, - }}, - } - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - lease, err := hashbuild.NewExpressionMemoryLease( - generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - - arg := &DedupJoin{} - arg.ctr.evecs = []evalVector{{executor: executors[0]}} - arg.ctr.vecs = make([]*vector.Vector, len(executors)) - arg.ctr.probeExpressionLease = lease - input := batch.NewWithSize(0) - input.SetRowCount(4) - require.NoError(t, arg.ctr.evalJoinConditionBudgeted(input, proc)) - require.Positive(t, generation.Used()) - - arg.Reset(proc, false, nil) - require.Zero(t, generation.Used()) - require.Nil(t, arg.ctr.evecs) - require.Nil(t, arg.ctr.vecs) - require.Nil(t, arg.ctr.probeExpressionLease) -} - -func TestDedupJoinResetReleasesAccountedProbeExpressions(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - const capBytes = uint64(1 << 20) - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - registry, err := mpool.NewAllocationAccountRegistry(1, 16) - require.NoError(t, err) - account, err := registry.OpenWithController(capBytes, generation) - require.NoError(t, err) - expr := &plan.Expr{Typ: plan.Type{Id: int32(types.T_int32)}, Expr: &plan.Expr_Lit{Lit: &plan.Literal{Value: &plan.Literal_I32Val{I32Val: 1}}}} - executors, err := hashbuild.NewAllocationAccountedExpressionExecutorsForAccount( - proc, []*plan.Expr{expr}, account, hashbuild.HashBuildAllocationOwner) - require.NoError(t, err) - arg := &DedupJoin{allocationAccount: account} - arg.ctr.evecs = []evalVector{{executor: executors[0]}} - arg.ctr.vecs = make([]*vector.Vector, len(executors)) - arg.ctr.probeExpressionsAccounted = true - input := batch.NewWithSize(0) - input.SetRowCount(4) - require.NoError(t, arg.ctr.evalJoinConditionBudgeted(input, proc)) - require.Positive(t, account.Snapshot().Used) - - arg.Reset(proc, false, nil) - require.Zero(t, account.Snapshot().Used) - require.Zero(t, generation.Used()) - require.False(t, arg.ctr.probeExpressionsAccounted) - require.Nil(t, arg.ctr.evecs) - terminal, _, err := registry.CompleteTerminal(account) - require.NoError(t, err) - require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) -} - -func TestDedupJoinAllocationActivationRequiresBothKeySides(t *testing.T) { - col := &plan.Expr{Typ: plan.Type{Id: int32(types.T_int64)}, Expr: &plan.Expr_Col{Col: &plan.ColRef{}}} - arg := &DedupJoin{Conditions: [][]*plan.Expr{{col}, {col}}} - require.True(t, arg.AllocationAccountEnabled()) - require.False(t, arg.AllocationAccountActivationBlocked()) - arg.Conditions[0] = []*plan.Expr{nil} - require.False(t, arg.AllocationAccountEnabled()) - require.True(t, arg.AllocationAccountActivationBlocked()) -} diff --git a/pkg/sql/colexec/dedupjoin/join.go b/pkg/sql/colexec/dedupjoin/join.go index d347b161aa5c1..6c2b428ca3a01 100644 --- a/pkg/sql/colexec/dedupjoin/join.go +++ b/pkg/sql/colexec/dedupjoin/join.go @@ -19,9 +19,9 @@ import ( "strings" "github.com/matrixorigin/matrixone/pkg/catalog" - "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/hashmap" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" @@ -138,6 +138,9 @@ func (dedupJoin *DedupJoin) OpType() vm.OpType { return vm.DedupJoin } func (dedupJoin *DedupJoin) Prepare(proc *process.Process) (err error) { + if dedupJoin.allocationAccount == nil { + return mpool.ErrAllocationAccountInvalid + } if dedupJoin.OpAnalyzer == nil { dedupJoin.OpAnalyzer = process.NewAnalyzer(dedupJoin.GetIdx(), dedupJoin.IsFirst, dedupJoin.IsLast, "dedup join") } else { @@ -148,13 +151,19 @@ func (dedupJoin *DedupJoin) Prepare(proc *process.Process) (err error) { newUpdateExecs := len(dedupJoin.ctr.exprExecs) == 0 && len(dedupJoin.UpdateColExprList) > 0 var evalExecs, updateExecs []colexec.ExpressionExecutor if newEvalVectors { - evalExecs, err = colexec.NewExpressionExecutorsFromPlanExpressions(proc, dedupJoin.Conditions[0]) + evalExecs, err = hashbuild.NewExpressionExecutors( + proc, + dedupJoin.Conditions[0], + ) if err != nil { return err } } if newUpdateExecs { - updateExecs, err = colexec.NewExpressionExecutorsFromPlanExpressions(proc, dedupJoin.UpdateColExprList) + updateExecs, err = hashbuild.NewExpressionExecutors( + proc, + dedupJoin.UpdateColExprList, + ) if err != nil { for _, exec := range evalExecs { exec.Free() @@ -249,23 +258,30 @@ func (dedupJoin *DedupJoin) Call(proc *process.Process) (vm.CallResult, error) { ctr.cleanBuf(proc) // Clear previous bucket state before advancing. ctr.cleanBucketState(proc) + var allocationErr error ok, bktErr := ctr.spillEngine.AdvanceToNextBucket(proc, analyzer, func(jm *message.JoinMap, res spillutil.BucketResult) { if res == spillutil.BucketReady { ctr.mp = jm ctr.batches = jm.GetBatches() ctr.batchRowCount = jm.GetRowCount() - ctr.matched = &bitmap.Bitmap{} - if dedupJoin.OnDuplicateAction != plan.Node_UPDATE { - ctr.matched.InitWithSize(ctr.batchRowCount) - } else { - ctr.matched.InitWithSize(int64(jm.GetGroupCount())) + rows := ctr.batchRowCount + if dedupJoin.OnDuplicateAction == plan.Node_UPDATE { + rows = int64(jm.GetGroupCount()) } + ctr.matched, allocationErr = colexec.NewAccountedBitmap( + rows, proc.Mp(), dedupJoin.allocationAccount, + hashbuild.HashBuildAllocationOwner, + dedupJoinAllocationSiteMatched, + ) } }) if bktErr != nil { return result, hashbuild.TerminalBudgetError(proc.Ctx, bktErr) } + if allocationErr != nil { + return result, allocationErr + } if ok && ctr.mp != nil { // BucketReady: init capture buffers for REPLACE spill path. if ctr.batchRowCount > 0 && len(dedupJoin.OldColCapturePlaceholderIdxList) > 0 { @@ -312,45 +328,13 @@ func (dedupJoin *DedupJoin) build(analyzer process.Analyzer, proc *process.Proce if takeErr != nil { return takeErr } - var probeExpressionLease *hashbuild.ExpressionMemoryLease - var leaseErr error - if dedupJoin.allocationAccount != nil && - hashbuild.AllocationAccountedExpressionSetSupported(dedupJoin.Conditions[0]) { - ctr.cleanEvalVectors() - var probeExecutors []colexec.ExpressionExecutor - probeExecutors, leaseErr = - hashbuild.NewAllocationAccountedExpressionExecutorsForAccount( - proc, - dedupJoin.Conditions[0], - dedupJoin.allocationAccount, - hashbuild.HashBuildAllocationOwner, - ) - if leaseErr == nil { - ctr.evecs = make([]evalVector, len(probeExecutors)) - ctr.vecs = make([]*vector.Vector, len(probeExecutors)) - for i := range probeExecutors { - ctr.evecs[i].executor = probeExecutors[i] - } - ctr.probeExpressionsAccounted = true - } - } else { - probeExecutors := make([]colexec.ExpressionExecutor, len(ctr.evecs)) - for i := range ctr.evecs { - probeExecutors[i] = ctr.evecs[i].executor - } - probeExpressionLease, leaseErr = hashbuild.NewExpressionMemoryLease( - budget, dedupJoin.Conditions[0], probeExecutors, false) - } - if leaseErr != nil { + if dedupJoin.allocationAccount == nil { _ = payload.Close() ctr.mp.Free() ctr.mp = nil - ctr.cleanEvalVectors() - ctr.releaseProbeExpressionLease() - return leaseErr + return mpool.ErrAllocationAccountInvalid } - ctr.probeExpressionLease = probeExpressionLease - engine, engineErr := spillutil.NewSpillEngineForAccount(spillutil.SpillEngineConfig{ + engine, engineErr := spillutil.NewSpillEngine(spillutil.SpillEngineConfig{ BuildKeyExprs: dedupJoin.Conditions[1], ProbeKeyExprs: dedupJoin.Conditions[0], SpillThreshold: ctr.spillThreshold, @@ -366,21 +350,16 @@ func (dedupJoin *DedupJoin) build(analyzer process.Analyzer, proc *process.Proce DedupDeleteMarkerColIdx: dedupJoin.DedupDeleteMarkerColIdx, DedupDeleteKeepColIdxList: dedupJoin.DedupDeleteKeepColIdxList, Budget: budget, - ProbeExpressionLease: probeExpressionLease, }, dedupJoin.allocationAccount, hashbuild.HashBuildAllocationOwner) if engineErr != nil { _ = payload.Close() ctr.mp.Free() ctr.mp = nil ctr.cleanEvalVectors() - ctr.releaseProbeExpressionLease() return engineErr } - if len(payload.Files) > 0 { - engine.InitFromSpilledFiles(payload.Files) - } else { - engine.InitFromSpilledMap(payload.LegacyFds) - } + engine.InitFromSpilledFiles(payload.Files) + ctr.spillEngine = engine if err := engine.ScatterProbeTable(proc, func() (*batch.Batch, error) { input, err := vm.ChildrenCall(dedupJoin.GetChildren(0), proc, analyzer) @@ -388,7 +367,7 @@ func (dedupJoin *DedupJoin) build(analyzer process.Analyzer, proc *process.Proce }, analyzer, func(bat *batch.Batch) ([]*vector.Vector, error) { - if err := ctr.evalJoinConditionBudgeted(bat, proc); err != nil { + if err := ctr.evalJoinCondition(bat, proc); err != nil { return nil, err } return ctr.vecs, nil @@ -397,10 +376,10 @@ func (dedupJoin *DedupJoin) build(analyzer process.Analyzer, proc *process.Proce ctr.mp.Free() ctr.mp = nil engine.Cleanup(proc) + ctr.spillEngine = nil return err } ctr.mp.Free() - ctr.spillEngine = engine ctr.mp = nil return } @@ -411,11 +390,19 @@ func (dedupJoin *DedupJoin) build(analyzer process.Analyzer, proc *process.Proce ctr.batches = ctr.mp.GetBatches() ctr.batchRowCount = ctr.mp.GetRowCount() if ctr.batchRowCount > 0 { - ctr.matched = &bitmap.Bitmap{} - if dedupJoin.OnDuplicateAction != plan.Node_UPDATE { - ctr.matched.InitWithSize(ctr.batchRowCount) - } else { - ctr.matched.InitWithSize(int64(ctr.mp.GetGroupCount())) + rows := ctr.batchRowCount + if dedupJoin.OnDuplicateAction == plan.Node_UPDATE { + rows = int64(ctr.mp.GetGroupCount()) + } + ctr.matched, err = colexec.NewAccountedBitmap( + rows, + proc.Mp(), + dedupJoin.allocationAccount, + hashbuild.HashBuildAllocationOwner, + dedupJoinAllocationSiteMatched, + ) + if err != nil { + return err } } if ctr.batchRowCount > 0 && len(dedupJoin.OldColCapturePlaceholderIdxList) > 0 { @@ -442,7 +429,13 @@ func (ctr *container) initCaptureBuffers(ap *DedupJoin, proc *process.Process) e ctr.capturedVecs = make([]*vector.Vector, n) for i, probePos := range ap.OldColCaptureProbeIdxList { typ := ap.LeftTypes[probePos] - vec := vector.NewOffHeapVecWithType(typ) + vec, err := vector.NewOffHeapVecWithTypeAndAllocation( + typ, + ap.stateAllocation, + ) + if err != nil { + return err + } if err := vector.AppendMultiFixed(vec, 0, true, int(ctr.batchRowCount), proc.Mp()); err != nil { vec.Free(proc.Mp()) ctr.capturedVecs[i] = nil @@ -450,8 +443,18 @@ func (ctr *container) initCaptureBuffers(ap *DedupJoin, proc *process.Process) e } ctr.capturedVecs[i] = vec } - ctr.captured = &bitmap.Bitmap{} - ctr.captured.InitWithSize(ctr.batchRowCount) + var err error + ctr.captured, err = colexec.NewAccountedBitmap( + ctr.batchRowCount, + proc.Mp(), + ap.allocationAccount, + hashbuild.HashBuildAllocationOwner, + dedupJoinAllocationSiteCaptured, + ) + if err != nil { + ctr.cleanCaptured(proc) + return err + } ctr.captureResultIdx = make([]int32, len(ap.Result)) for j := range ctr.captureResultIdx { ctr.captureResultIdx[j] = -1 @@ -499,6 +502,7 @@ func (ctr *container) finalize(ap *DedupJoin, proc *process.Process) error { // point Reset/Free still owns and releases these vectors. ctr.captured = nil ctr.capturedVecs = nil + ctr.matched = nil // Publication, not acknowledgement, is the worker's single status // for this round. Mark it before waiting so concurrent cancellation // cannot make Reset enqueue a duplicate abort status. @@ -644,7 +648,18 @@ func (ctr *container) finalize(ap *DedupJoin, proc *process.Process) error { return nil } ctr.matched.Negate() - sels := make([]int32, 0, count) + sels, err := mpool.MakeSliceAccounted[int32]( + count, + proc.Mp(), + ap.allocationAccount, + hashbuild.HashBuildAllocationOwner, + dedupJoinAllocationSiteFinalizeSelections, + ) + if err != nil { + return err + } + sels = sels[:0] + defer mpool.FreeSlice(proc.Mp(), sels) itr := ctr.matched.Iterator() for itr.HasNext() { r := itr.Next() @@ -823,7 +838,7 @@ func (ctr *container) withRestoredJoinBat1Vectors(updateCols []int32, fn func() func (ctr *container) probe(bat *batch.Batch, ap *DedupJoin, proc *process.Process, analyzer process.Analyzer, result *vm.CallResult) error { ap.resetRBat() - err := ctr.evalJoinConditionBudgeted(bat, proc) + err := ctr.evalJoinCondition(bat, proc) if err != nil { return err } @@ -847,7 +862,10 @@ func (ctr *container) probe(bat *batch.Batch, ap *DedupJoin, proc *process.Proce if n > hashmap.UnitLimit { n = hashmap.UnitLimit } - vals, zvals := itr.Find(i, n, ctr.vecs) + vals, zvals, err := itr.Find(i, n, ctr.vecs) + if err != nil { + return err + } for k := 0; k < n; k++ { if zvals[k] == 0 || vals[k] == 0 { continue @@ -1004,16 +1022,6 @@ func (ctr *container) evalJoinCondition(bat *batch.Batch, proc *process.Process) return nil } -func (ctr *container) evalJoinConditionBudgeted(bat *batch.Batch, proc *process.Process) error { - if ctr.probeExpressionLease == nil { - return ctr.evalJoinCondition(bat, proc) - } - return ctr.probeExpressionLease.Eval(proc, []*batch.Batch{bat}, bat.RowCount(), func(i int, vec *vector.Vector) error { - ctr.vecs[i] = vec - ctr.evecs[i].vec = vec - return nil - }) -} func unionSelsByBatch(dst *vector.Vector, batches []*batch.Batch, colPos int32, sels []int32, proc *process.Process) error { if len(sels) <= 16 { for _, sel := range sels { diff --git a/pkg/sql/colexec/dedupjoin/join_finalize_optimize_test.go b/pkg/sql/colexec/dedupjoin/join_finalize_optimize_test.go index d14a20d3f0afc..8cdd73e366e8b 100644 --- a/pkg/sql/colexec/dedupjoin/join_finalize_optimize_test.go +++ b/pkg/sql/colexec/dedupjoin/join_finalize_optimize_test.go @@ -40,6 +40,7 @@ func runFinalizeFixture( buildBat, probeBat *batch.Batch, ) []*batch.Batch { t.Helper() + installTestAllocation(t, dedupArg, buildArg) buildArg.Children = nil buildArg.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{buildBat})) diff --git a/pkg/sql/colexec/dedupjoin/join_test.go b/pkg/sql/colexec/dedupjoin/join_test.go index 24a230c1cdb17..9d9ea4992968b 100644 --- a/pkg/sql/colexec/dedupjoin/join_test.go +++ b/pkg/sql/colexec/dedupjoin/join_test.go @@ -18,6 +18,7 @@ import ( "bytes" "context" "errors" + "io" "os" "sync" "testing" @@ -58,6 +59,30 @@ type joinTestCase struct { barg *hashbuild.HashBuild } +func newDedupTestSpillEngine( + t *testing.T, + cfg spillutil.SpillEngineConfig, +) *spillutil.SpillEngine { + t.Helper() + if cfg.Budget == nil { + budget := process.MustNewHashBuildBudget(1<<60, 1<<60) + var err error + cfg.Budget, err = budget.OpenGeneration(1) + require.NoError(t, err) + } + registry, err := mpool.NewAllocationAccountRegistry(1, 1<<20) + require.NoError(t, err) + account, err := registry.OpenWithController(1<<60, cfg.Budget) + require.NoError(t, err) + engine, err := spillutil.NewSpillEngine( + cfg, + account, + hashbuild.HashBuildAllocationOwner, + ) + require.NoError(t, err) + return engine +} + func TestDedupFinalizeCleansConsumedBuffer(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) baseline := proc.Mp().CurrNB() @@ -71,7 +96,10 @@ func TestDedupFinalizeCleansConsumedBuffer(t *testing.T) { arg.ctr.state = Finalize arg.ctr.buf = []*batch.Batch{bat} arg.ctr.lastPos = 1 - arg.ctr.spillEngine = spillutil.NewSpillEngine(spillutil.SpillEngineConfig{}) + arg.ctr.spillEngine = newDedupTestSpillEngine( + t, + spillutil.SpillEngineConfig{}, + ) res, err := arg.Call(proc) require.NoError(t, err) @@ -112,14 +140,32 @@ func writeDedupSpillBatch(t *testing.T, proc *process.Process, name string, valu require.NoError(t, err) fd, err := spillfs.CreateAndRemoveFile(proc.Ctx, name) require.NoError(t, err) - w := spillutil.BucketWriter{Name: name, Fd: fd} bat := batch.NewWithSize(1) bat.Vecs[0] = testutil.MakeInt32Vector([]int32{value}, nil, proc.Mp()) bat.SetRowCount(1) - var buf bytes.Buffer - require.NoError(t, spillutil.FlushBucketBatch(proc, bat, &w, &buf, nil)) + var payload bytes.Buffer + require.NoError(t, bat.MarshalBinaryWithGroupingTo(&payload)) + rows, size, magic := int64(1), int64(payload.Len()), uint64(spillutil.SpillMagic) + for _, part := range [][]byte{ + types.EncodeInt64(&rows), + types.EncodeInt64(&size), + payload.Bytes(), + types.EncodeUint64(&magic), + } { + _, err = fd.Write(part) + require.NoError(t, err) + } + _, err = fd.Seek(0, io.SeekStart) + require.NoError(t, err) bat.Clean(proc.Mp()) - return w.HandOffFd() + return fd +} + +func newDedupSpillFile(t *testing.T, fd *os.File, rows int64) *message.SpillFile { + t.Helper() + info, err := fd.Stat() + require.NoError(t, err) + return message.NewSpillFile(fd, rows, uint64(info.Size()), nil) } func TestDedupSpillAdvancesAfterOutput(t *testing.T) { @@ -127,15 +173,23 @@ func TestDedupSpillAdvancesAfterOutput(t *testing.T) { baseline := proc.Mp().CurrNB() typ := types.T_int32.ToType() conditions := [][]*plan.Expr{{newExpr(0, typ)}, {newExpr(0, typ)}} - engine := spillutil.NewSpillEngine(spillutil.SpillEngineConfig{ + engine := newDedupTestSpillEngine(t, spillutil.SpillEngineConfig{ BuildKeyExprs: conditions[1], NeedBatches: true, NeedsBuildForEmptyProbe: true, IsDedup: true, }) - engine.InitFromSpilledMap([]*os.File{ - writeDedupSpillBatch(t, proc, "dedup_bucket_1", 1), - writeDedupSpillBatch(t, proc, "dedup_bucket_2", 2), + engine.InitFromSpilledFiles([]*message.SpillFile{ + newDedupSpillFile( + t, + writeDedupSpillBatch(t, proc, "dedup_bucket_1", 1), + 1, + ), + newDedupSpillFile( + t, + writeDedupSpillBatch(t, proc, "dedup_bucket_2", 2), + 1, + ), }) arg := &DedupJoin{ @@ -144,6 +198,7 @@ func TestDedupSpillAdvancesAfterOutput(t *testing.T) { Result: []colexec.ResultPos{{Rel: 1, Pos: 0}}, OnDuplicateAction: plan.Node_FAIL, } + installTestAllocation(t, arg) require.NoError(t, arg.Prepare(proc)) arg.ctr.state = Finalize arg.ctr.spillEngine = engine @@ -314,6 +369,7 @@ func TestDedupPrepareFailureCanRetry(t *testing.T) { Conditions: [][]*plan.Expr{{valid}, {valid}}, UpdateColExprList: []*plan.Expr{valid, invalid}, } + installTestAllocation(t, arg) require.Error(t, arg.Prepare(proc)) require.Nil(t, arg.ctr.vecs) @@ -523,7 +579,7 @@ func newTestCase(t *testing.T, flgs []bool, ts []types.Type, rp []int32, cs [][] // }, //}) tag++ - return joinTestCase{ + tc := joinTestCase{ types: ts, flgs: flgs, proc: proc, @@ -555,6 +611,8 @@ func newTestCase(t *testing.T, flgs []bool, ts []types.Type, rp []int32, cs [][] JoinMapRefCnt: 1, }, } + installTestAllocation(t, tc.arg, tc.barg) + return tc } func resetChildren(arg *DedupJoin, m *mpool.MPool) { @@ -653,6 +711,7 @@ func TestDedupJoinCapture(t *testing.T) { JoinMapTag: curTag, JoinMapRefCnt: 1, } + installTestAllocation(t, dedupArg, buildArg) // Set up children buildOp := colexec.NewMockOperator().WithBatchs([]*batch.Batch{buildBat}) @@ -755,6 +814,7 @@ func TestDedupJoinCapturePartialMatch(t *testing.T) { JoinMapTag: curTag, JoinMapRefCnt: 1, } + installTestAllocation(t, dedupArg, buildArg) buildOp := colexec.NewMockOperator().WithBatchs([]*batch.Batch{buildBat}) buildArg.Children = nil @@ -849,6 +909,7 @@ func TestDedupJoinCaptureReset(t *testing.T) { JoinMapTag: curTag, JoinMapRefCnt: 1, } + installTestAllocation(t, dedupArg, buildArg) // --- First run --- buildBat1 := makeInt32Batch(proc.Mp(), [][]int32{{10, 20}, {0, 0}}, [][]uint64{nil, {0, 1}}) @@ -1632,6 +1693,7 @@ func TestDedupFinalizeParallelMergePreservesDataAcrossReset(t *testing.T) { IsMerger: false, Mailbox: mailbox, } + installTestAllocation(t, arg, workerArg) cleaned := false t.Cleanup(func() { if !cleaned { diff --git a/pkg/sql/colexec/dedupjoin/key_contract_test.go b/pkg/sql/colexec/dedupjoin/key_contract_test.go index 018ee5f4e6f47..199f54edfbfff 100644 --- a/pkg/sql/colexec/dedupjoin/key_contract_test.go +++ b/pkg/sql/colexec/dedupjoin/key_contract_test.go @@ -189,6 +189,7 @@ func runDedupJoinDoubleSignedZeroContract( if mode.shuffle { buildArg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: joinMapTag + 8000} } + installTestAllocation(t, dedupArg, buildArg) buildArg.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{buildBatch})) spillBefore := promtestutil.ToFloat64( diff --git a/pkg/sql/colexec/dedupjoin/types.go b/pkg/sql/colexec/dedupjoin/types.go index f1310fbe6cf10..e481c827741b2 100644 --- a/pkg/sql/colexec/dedupjoin/types.go +++ b/pkg/sql/colexec/dedupjoin/types.go @@ -44,6 +44,16 @@ const ( End ) +const ( + dedupJoinAllocationSiteMatched mpool.AllocationSite = iota + 82 + dedupJoinAllocationSiteCaptured + dedupJoinAllocationSiteCaptureData + dedupJoinAllocationSiteCaptureArea + dedupJoinAllocationSiteCaptureNulls + dedupJoinAllocationSiteCaptureGrouping + dedupJoinAllocationSiteFinalizeSelections +) + // WorkerJoinMsg carries per-worker state from non-merger workers to the // merger worker at finalize time. Regular DEDUP JOIN only populates matched; // the REPLACE INTO merged main-table scan path (OldColCapture) additionally @@ -200,6 +210,8 @@ func freeCapturedVecs(vecs []*vector.Vector, proc *process.Process) { func freeWorkerJoinMsg(msg *WorkerJoinMsg, proc *process.Process) { if msg != nil { + colexec.FreeAccountedBitmap(msg.matched, proc.Mp()) + colexec.FreeAccountedBitmap(msg.captured, proc.Mp()) freeCapturedVecs(msg.capturedVecs, proc) } } @@ -258,11 +270,6 @@ type container struct { // Spill support for large build sides. spillEngine *spillutil.SpillEngine spillThreshold int64 - // Non-nil only for spilled joins, where probe expressions are part of the - // shared HashBuild/spill working set. Resident probe expressions remain - // under normal process/mpool accounting; this is not a general query budget. - probeExpressionLease *hashbuild.ExpressionMemoryLease - probeExpressionsAccounted bool } type DedupJoin struct { @@ -301,26 +308,11 @@ type DedupJoin struct { OldColCapturePlaceholderIdxList []int32 OldColCaptureProbeIdxList []int32 allocationAccount *mpool.AllocationAccount + stateAllocation *vector.AllocationAccountSelection vm.OperatorBase } -func (dedupJoin *DedupJoin) AllocationAccountEnabled() bool { - return dedupJoin != nil && dedupJoin.allocationAccountExpressionOwnerClosed() -} - -func (dedupJoin *DedupJoin) AllocationAccountActivationBlocked() bool { - return dedupJoin != nil && !dedupJoin.allocationAccountExpressionOwnerClosed() -} - -func (dedupJoin *DedupJoin) allocationAccountExpressionOwnerClosed() bool { - if dedupJoin == nil || len(dedupJoin.Conditions) != 2 { - return false - } - return hashbuild.AllocationAccountedExpressionSetSupported(dedupJoin.Conditions[0]) && - hashbuild.AllocationAccountedExpressionSetSupported(dedupJoin.Conditions[1]) -} - func (dedupJoin *DedupJoin) SetAllocationAccount( account *mpool.AllocationAccount, ) error { @@ -331,7 +323,22 @@ func (dedupJoin *DedupJoin) SetAllocationAccount( dedupJoin.allocationAccount != account { return mpool.ErrAllocationAccountMismatch } + if dedupJoin.allocationAccount == account { + return nil + } + selection, err := vector.NewAllocationAccountSelection( + account, + hashbuild.HashBuildAllocationOwner, + dedupJoinAllocationSiteCaptureData, + dedupJoinAllocationSiteCaptureArea, + dedupJoinAllocationSiteCaptureNulls, + dedupJoinAllocationSiteCaptureGrouping, + ) + if err != nil { + return err + } dedupJoin.allocationAccount = account + dedupJoin.stateAllocation = selection return nil } @@ -345,10 +352,13 @@ func (dedupJoin *DedupJoin) ClearAllocationAccount( return mpool.ErrAllocationAccountMismatch } if dedupJoin.ctr.mp != nil || dedupJoin.ctr.spillEngine != nil || - dedupJoin.ctr.probeExpressionsAccounted { + len(dedupJoin.ctr.evecs) != 0 || len(dedupJoin.ctr.exprExecs) != 0 || + dedupJoin.ctr.matched != nil || dedupJoin.ctr.captured != nil || + len(dedupJoin.ctr.capturedVecs) != 0 { return mpool.ErrAllocationAccountInvariant } dedupJoin.allocationAccount = nil + dedupJoin.stateAllocation = nil return nil } @@ -422,21 +432,15 @@ func (dedupJoin *DedupJoin) Reset(proc *process.Process, pipelineFailed bool, er ctr.cleanBuf(proc) ctr.cleanBucketState(proc) - ctr.resetExprExecutor() + ctr.cleanExprExecutor() if ctr.spillEngine != nil { ctr.spillEngine.Cleanup(proc) ctr.spillEngine = nil } - if ctr.probeExpressionLease != nil || ctr.probeExpressionsAccounted { - ctr.cleanEvalVectors() - ctr.releaseProbeExpressionLease() - } else { - ctr.resetEvalVectors() - } + ctr.cleanEvalVectors() ctr.roundStatusPublished = false ctr.state = Build ctr.lastPos = 0 - dedupJoin.allocationAccount = nil } func (dedupJoin *DedupJoin) Free(proc *process.Process, pipelineFailed bool, err error) { @@ -455,25 +459,19 @@ func (dedupJoin *DedupJoin) Free(proc *process.Process, pipelineFailed bool, err ctr.spillEngine = nil } ctr.cleanEvalVectors() - ctr.releaseProbeExpressionLease() - dedupJoin.allocationAccount = nil } func (dedupJoin *DedupJoin) ExecProjection(proc *process.Process, input *batch.Batch) (*batch.Batch, error) { return input, nil } -func (ctr *container) resetExprExecutor() { - for i := range ctr.exprExecs { - ctr.exprExecs[i].ResetForNextQuery() - } -} - func (ctr *container) cleanExprExecutor() { for i := range ctr.exprExecs { - ctr.exprExecs[i].Free() - ctr.exprExecs[i] = nil + if ctr.exprExecs[i] != nil { + ctr.exprExecs[i].Free() + } } + ctr.exprExecs = nil } func (ctr *container) cleanBuf(proc *process.Process) { @@ -492,6 +490,7 @@ func (ctr *container) cleanCaptured(proc *process.Process) { } } ctr.capturedVecs = nil + colexec.FreeAccountedBitmap(ctr.captured, proc.Mp()) ctr.captured = nil ctr.captureResultIdx = nil } @@ -523,10 +522,12 @@ func (ctr *container) cleanBucketState(proc *process.Process) { ctr.cleanHashMap() ctr.batches = nil ctr.batchRowCount = 0 + colexec.FreeAccountedBitmap(ctr.matched, proc.Mp()) ctr.matched = nil } func (ctr *container) cleanHashMap() { + hashmap.IteratorClearOwner(ctr.cachedItr) ctr.cachedItr = nil if ctr.mp != nil { ctr.mp.Free() @@ -543,20 +544,4 @@ func (ctr *container) cleanEvalVectors() { } ctr.evecs = nil ctr.vecs = nil - ctr.probeExpressionsAccounted = false -} - -func (ctr *container) resetEvalVectors() { - for i := range ctr.evecs { - if ctr.evecs[i].executor != nil { - ctr.evecs[i].executor.ResetForNextQuery() - } - } -} - -func (ctr *container) releaseProbeExpressionLease() { - if ctr.probeExpressionLease != nil { - ctr.probeExpressionLease.Release() - ctr.probeExpressionLease = nil - } } diff --git a/pkg/sql/colexec/evalExpression.go b/pkg/sql/colexec/evalExpression.go index 3415c205bbbf4..01d9329d00f73 100644 --- a/pkg/sql/colexec/evalExpression.go +++ b/pkg/sql/colexec/evalExpression.go @@ -92,39 +92,9 @@ type ExpressionExecutor interface { } func NewExpressionExecutorsFromPlanExpressions(proc *process.Process, planExprs []*plan.Expr) (executors []ExpressionExecutor, err error) { - return newExpressionExecutorsFromPlanExpressions(proc, planExprs, nil) -} - -// NewExpressionExecutorsFromPlanExpressionsWithAllocation constructs -// allocation-accounted expression trees for callers that have already proved -// the complete expression allocation-site ledger closed. -func NewExpressionExecutorsFromPlanExpressionsWithAllocation( - proc *process.Process, - planExprs []*plan.Expr, - allocation *ExpressionAllocationAccount, -) (executors []ExpressionExecutor, err error) { - if err = allocation.validate(); err != nil { - return nil, err - } - return newExpressionExecutorsFromPlanExpressions( - proc, - planExprs, - allocation, - ) -} - -func newExpressionExecutorsFromPlanExpressions( - proc *process.Process, - planExprs []*plan.Expr, - allocation *ExpressionAllocationAccount, -) (executors []ExpressionExecutor, err error) { executors = make([]ExpressionExecutor, len(planExprs)) for i := range executors { - executors[i], err = newExpressionExecutor( - proc, - planExprs[i], - allocation, - ) + executors[i], err = NewExpressionExecutor(proc, planExprs[i]) if err != nil { for j := 0; j < i; j++ { executors[j].Free() @@ -136,39 +106,10 @@ func newExpressionExecutorsFromPlanExpressions( } func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (ExpressionExecutor, error) { - return newExpressionExecutor(proc, planExpr, nil) -} - -// NewExpressionExecutorWithAllocation is the single-root counterpart of -// NewExpressionExecutorsFromPlanExpressionsWithAllocation. -func NewExpressionExecutorWithAllocation( - proc *process.Process, - planExpr *plan.Expr, - allocation *ExpressionAllocationAccount, -) (ExpressionExecutor, error) { - if err := allocation.validate(); err != nil { - return nil, err - } - return newExpressionExecutor(proc, planExpr, allocation) -} - -func newExpressionExecutor( - proc *process.Process, - planExpr *plan.Expr, - allocation *ExpressionAllocationAccount, -) (ExpressionExecutor, error) { - if planExpr == nil { - return nil, moerr.NewInvalidInput(proc.Ctx, "nil expression") - } switch t := planExpr.Expr.(type) { case *plan.Expr_Lit: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) - vec, err := generateConstExpressionExecutor( - proc, - typ, - t.Lit, - allocation, - ) + vec, err := generateConstExpressionExecutor(proc, typ, t.Lit) if err != nil { return nil, err } @@ -176,25 +117,17 @@ func newExpressionExecutor( case *plan.Expr_T: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) - var selection *vector.AllocationAccountSelection - if allocation != nil { - selection = allocation.constant - } - vec, err := newExpressionConstNull(typ, 1, selection, proc.Mp()) - if err != nil { - return nil, err - } + vec := vector.NewConstNull(typ, 1, proc.Mp()) return NewFixedVectorExpressionExecutor(proc.Mp(), false, vec), nil case *plan.Expr_Col: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) ce := NewColumnExpressionExecutor() *ce = ColumnExpressionExecutor{ - mp: proc.Mp(), - relIndex: int(t.Col.RelPos), - colIndex: int(t.Col.ColPos), - typ: typ, - allocation: allocation, + mp: proc.Mp(), + relIndex: int(t.Col.RelPos), + colIndex: int(t.Col.ColPos), + typ: typ, } // [issue#19574] // if < 0, it's special for agg or others. @@ -205,42 +138,24 @@ func newExpressionExecutor( case *plan.Expr_P: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) - executor := NewParamExpressionExecutor(proc.Mp(), int(t.P.Pos), typ) - executor.allocation = allocation - return executor, nil + return NewParamExpressionExecutor(proc.Mp(), int(t.P.Pos), typ), nil case *plan.Expr_V: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) ve := NewVarExpressionExecutor() *ve = VarExpressionExecutor{ - mp: proc.Mp(), - name: t.V.Name, - system: t.V.System, - global: t.V.Global, - typ: typ, - allocation: allocation, + mp: proc.Mp(), + name: t.V.Name, + system: t.V.System, + global: t.V.Global, + typ: typ, } return ve, nil case *plan.Expr_Vec: - var vec *vector.Vector - var err error - if allocation == nil { - vec = vector.NewVec(types.T_any.ToType()) - err = vec.UnmarshalBinary(t.Vec.Data) - } else { - vec, err = newExpressionVector( - types.T_any.ToType(), - allocation.constant, - ) - if err == nil { - err = vec.UnmarshalBinaryWithCopy(t.Vec.Data, proc.Mp()) - } - } + vec := vector.NewVec(types.T_any.ToType()) + err := vec.UnmarshalBinary(t.Vec.Data) if err != nil { - if vec != nil { - vec.Free(proc.Mp()) - } return nil, err } return NewFixedVectorExpressionExecutor(proc.Mp(), true, vec), nil @@ -249,16 +164,9 @@ func newExpressionExecutor( executor := NewListExpressionExecutor() resultVecTyp := t.List.List[0].GetTyp() typ := types.New(types.T(resultVecTyp.Id), resultVecTyp.Width, resultVecTyp.Scale) - if err := executor.init(proc, typ, len(t.List.List), allocation); err != nil { - executor.Free() - return nil, err - } + executor.Init(proc, typ, len(t.List.List)) for i := range executor.parameterExecutor { - subExecutor, paramErr := newExpressionExecutor( - proc, - t.List.List[i], - allocation, - ) + subExecutor, paramErr := NewExpressionExecutor(proc, t.List.List[i]) if paramErr != nil { executor.Free() return nil, paramErr @@ -288,17 +196,13 @@ func newExpressionExecutor( } typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) - if err = executor.init(proc, len(t.F.Args), typ, allocation); err != nil { + if err = executor.Init(proc, len(t.F.Args), typ); err != nil { executor.Free() return nil, err } for i := range executor.parameterExecutor { - subExecutor, paramErr := newExpressionExecutor( - proc, - t.F.Args[i], - allocation, - ) + subExecutor, paramErr := NewExpressionExecutor(proc, t.F.Args[i]) if paramErr != nil { executor.Free() return nil, paramErr @@ -326,8 +230,7 @@ type FixedVectorExpressionExecutor struct { } type FunctionExpressionExecutor struct { - m *mpool.MPool - allocation *ExpressionAllocationAccount + m *mpool.MPool // resultType is the declared function return type. Some built-ins refine // result metadata (for example temporal scale or decimal width/scale) at // runtime, so reusable result vectors must start each evaluation from this @@ -359,10 +262,9 @@ type FunctionExpressionExecutor struct { } type ColumnExpressionExecutor struct { - mp *mpool.MPool - relIndex int - colIndex int - allocation *ExpressionAllocationAccount + mp *mpool.MPool + relIndex int + colIndex int // result type. typ types.Type @@ -381,9 +283,8 @@ func (expr *ColumnExpressionExecutor) GetColIndex() int { } type ParamExpressionExecutor struct { - mp *mpool.MPool - allocation *ExpressionAllocationAccount - null *vector.Vector + mp *mpool.MPool + null *vector.Vector // maskedNull is separate from null/vec because it is not a resolved // parameter value and must never participate in the folded-value cache. maskedNull *vector.Vector @@ -397,20 +298,7 @@ type ParamExpressionExecutor struct { func (expr *ParamExpressionExecutor) Eval(proc *process.Process, batches []*batch.Batch, selectList []bool) (*vector.Vector, error) { if noRowsSelected(selectList, expressionRowCount(batches)) { if expr.maskedNull == nil { - var selection *vector.AllocationAccountSelection - if expr.allocation != nil { - selection = expr.allocation.result - } - var err error - expr.maskedNull, err = newExpressionConstNull( - expr.typ, - 1, - selection, - proc.GetMPool(), - ) - if err != nil { - return nil, err - } + expr.maskedNull = vector.NewConstNull(expr.typ, 1, proc.GetMPool()) } return expr.maskedNull, nil } @@ -431,35 +319,13 @@ func (expr *ParamExpressionExecutor) Eval(proc *process.Process, batches []*batc if val == nil { if expr.null == nil { - var selection *vector.AllocationAccountSelection - if expr.allocation != nil { - selection = expr.allocation.result - } - expr.null, err = newExpressionConstNull( - expr.typ, - 1, - selection, - proc.GetMPool(), - ) - if err != nil { - return nil, err - } + expr.null = vector.NewConstNull(expr.typ, 1, proc.GetMPool()) } return expr.null, nil } if expr.vec == nil { - var selection *vector.AllocationAccountSelection - if expr.allocation != nil { - selection = expr.allocation.result - } - expr.vec, err = newExpressionConstBytes( - expr.typ, - val, - 1, - proc.Mp(), - selection, - ) + expr.vec, err = vector.NewConstBytes(expr.typ, val, 1, proc.Mp()) } else { err = vector.SetConstBytes(expr.vec, val, 1, proc.GetMPool()) } @@ -507,9 +373,8 @@ func (expr *ParamExpressionExecutor) IsColumnExpr() bool { } type VarExpressionExecutor struct { - mp *mpool.MPool - allocation *ExpressionAllocationAccount - null *vector.Vector + mp *mpool.MPool + null *vector.Vector // maskedNull lets a skipped variable avoid the resolver without changing // the value cache used by a later selected evaluation. maskedNull *vector.Vector @@ -524,20 +389,7 @@ type VarExpressionExecutor struct { func (expr *VarExpressionExecutor) Eval(proc *process.Process, batches []*batch.Batch, selectList []bool) (*vector.Vector, error) { if noRowsSelected(selectList, expressionRowCount(batches)) { if expr.maskedNull == nil { - var selection *vector.AllocationAccountSelection - if expr.allocation != nil { - selection = expr.allocation.result - } - var err error - expr.maskedNull, err = newExpressionConstNull( - expr.typ, - 1, - selection, - proc.GetMPool(), - ) - if err != nil { - return nil, err - } + expr.maskedNull = vector.NewConstNull(expr.typ, 1, proc.GetMPool()) } return expr.maskedNull, nil } @@ -559,16 +411,7 @@ func (expr *VarExpressionExecutor) Eval(proc *process.Process, batches []*batch. if val == nil { if expr.null == nil { - var selection *vector.AllocationAccountSelection - if expr.allocation != nil { - selection = expr.allocation.result - } - expr.null, err = util.GenVectorByVarValueWithAllocation( - proc, - expr.typ, - nil, - selection, - ) + expr.null, err = util.GenVectorByVarValue(proc, expr.typ, nil) } if err == nil { expr.null.SetIsBin(isBin) @@ -577,16 +420,7 @@ func (expr *VarExpressionExecutor) Eval(proc *process.Process, batches []*batch. } if expr.vec == nil { - var selection *vector.AllocationAccountSelection - if expr.allocation != nil { - selection = expr.allocation.result - } - expr.vec, err = util.GenVectorByVarValueWithAllocation( - proc, - expr.typ, - val, - selection, - ) + expr.vec, err = util.GenVectorByVarValue(proc, expr.typ, val) } else { switch v := val.(type) { case []byte: @@ -640,8 +474,7 @@ func (expr *VarExpressionExecutor) IsColumnExpr() bool { } type ListExpressionExecutor struct { - mp *mpool.MPool - allocation *ExpressionAllocationAccount + mp *mpool.MPool typ types.Type resultVector *vector.Vector @@ -651,24 +484,11 @@ type ListExpressionExecutor struct { func (expr *ListExpressionExecutor) Eval(proc *process.Process, batches []*batch.Batch, selectList []bool) (*vector.Vector, error) { if expr.resultVector == nil { - var selection *vector.AllocationAccountSelection - if expr.allocation != nil { - selection = expr.allocation.result - } - var err error - expr.resultVector, err = newExpressionVector(expr.typ, selection) - if err != nil { - return nil, err - } + expr.resultVector = vector.NewOffHeapVecWithType(expr.typ) } else { expr.resultVector.CleanOnlyData() } - if err := expr.resultVector.PreExtend( - len(expr.parameterExecutor), - proc.Mp(), - ); err != nil { - return nil, err - } + expr.resultVector.PreExtend(len(expr.parameterExecutor), proc.Mp()) for i := range expr.parameterExecutor { vec, err := expr.parameterExecutor[i].Eval(proc, batches, selectList) if err != nil { @@ -711,30 +531,12 @@ func (expr *ListExpressionExecutor) IsColumnExpr() bool { } func (expr *ListExpressionExecutor) Init(proc *process.Process, typ types.Type, parameterNum int) { - if err := expr.init(proc, typ, parameterNum, nil); err != nil { - panic(err) - } -} - -func (expr *ListExpressionExecutor) init( - proc *process.Process, - typ types.Type, - parameterNum int, - allocation *ExpressionAllocationAccount, -) error { m := proc.Mp() expr.typ = typ expr.mp = m - expr.allocation = allocation expr.parameterExecutor = make([]ExpressionExecutor, parameterNum) - var selection *vector.AllocationAccountSelection - if allocation != nil { - selection = allocation.result - } - var err error - expr.resultVector, err = newExpressionVector(typ, selection) - return err + expr.resultVector = vector.NewOffHeapVecWithType(typ) } func (expr *ListExpressionExecutor) SetParameter(index int, executor ExpressionExecutor) { @@ -751,33 +553,14 @@ func (expr *FunctionExpressionExecutor) Init( proc *process.Process, parameterNum int, retType types.Type) (err error) { - return expr.init(proc, parameterNum, retType, nil) -} - -func (expr *FunctionExpressionExecutor) init( - proc *process.Process, - parameterNum int, - retType types.Type, - allocation *ExpressionAllocationAccount, -) (err error) { m := proc.Mp() expr.m = m - expr.allocation = allocation expr.resultType = retType expr.parameterResults = make([]*vector.Vector, parameterNum) expr.parameterExecutor = make([]ExpressionExecutor, parameterNum) - if allocation == nil { - expr.resultVector = vector.NewFunctionResultWrapper(retType, m) - return nil - } - expr.resultVector, err = vector.NewFunctionResultWrapperWithFunctionAllocation( - retType, - m, - allocation.result, - allocation.function, - ) + expr.resultVector = vector.NewFunctionResultWrapper(retType, m) return err } @@ -805,26 +588,8 @@ func (expr *FunctionExpressionExecutor) EvalIff(proc *process.Process, batches [ } rowCount := expressionRowCount(batches) if len(expr.selectList1) < rowCount { - expr.selectList1, err = ensureExpressionSlice( - expr.selectList1, - rowCount, - expr.m, - expr.allocation, - ExpressionAllocationSiteSelection, - ) - if err != nil { - return err - } - expr.selectList2, err = ensureExpressionSlice( - expr.selectList2, - rowCount, - expr.m, - expr.allocation, - ExpressionAllocationSiteSelection, - ) - if err != nil { - return err - } + expr.selectList1 = make([]bool, rowCount) + expr.selectList2 = make([]bool, rowCount) } trueBranch := expr.selectList1[:rowCount] @@ -859,17 +624,14 @@ func (expr *FunctionExpressionExecutor) EvalIff(proc *process.Process, batches [ return err } } else { - expr.parameterResults[1], err = expr.iffNullResult(0, rowCount) - if err != nil { - return err - } + expr.parameterResults[1] = expr.iffNullResult(0, rowCount) } if hasSelectedRows(falseBranch) { expr.parameterResults[2], err = expr.parameterExecutor[2].Eval(proc, batches, falseBranch) return err } - expr.parameterResults[2], err = expr.iffNullResult(1, rowCount) - return err + expr.parameterResults[2] = expr.iffNullResult(1, rowCount) + return nil } func hasSelectedRows(selectList []bool) bool { @@ -881,60 +643,26 @@ func hasSelectedRows(selectList []bool) bool { return false } -func (expr *FunctionExpressionExecutor) iffNullResult( - index int, - length int, -) (*vector.Vector, error) { +func (expr *FunctionExpressionExecutor) iffNullResult(index, length int) *vector.Vector { typ := expr.resultType result := expr.iffNullResults[index] if result == nil || *result.GetType() != typ { if result != nil { result.Free(expr.m) } - var selection *vector.AllocationAccountSelection - if expr.allocation != nil { - selection = expr.allocation.result - } - var err error - result, err = newExpressionConstNull( - typ, - length, - selection, - expr.m, - ) - if err != nil { - return nil, err - } + result = vector.NewConstNull(typ, length, expr.m) expr.iffNullResults[index] = result } else { result.SetLength(length) } - return result, nil + return result } func (expr *FunctionExpressionExecutor) EvalCase(proc *process.Process, batches []*batch.Batch, selectList []bool) (err error) { rowCount := expressionRowCount(batches) if len(expr.selectList1) < rowCount { - expr.selectList1, err = ensureExpressionSlice( - expr.selectList1, - rowCount, - expr.m, - expr.allocation, - ExpressionAllocationSiteSelection, - ) - if err != nil { - return err - } - expr.selectList2, err = ensureExpressionSlice( - expr.selectList2, - rowCount, - expr.m, - expr.allocation, - ExpressionAllocationSiteSelection, - ) - if err != nil { - return err - } + expr.selectList1 = make([]bool, rowCount) + expr.selectList2 = make([]bool, rowCount) } remaining := expr.selectList1[:rowCount] selectedBranch := expr.selectList2[:rowCount] @@ -974,16 +702,7 @@ func (expr *FunctionExpressionExecutor) EvalCase(proc *process.Process, batches func (expr *FunctionExpressionExecutor) EvalCoalesce(proc *process.Process, batches []*batch.Batch, selectList []bool) (err error) { rowCount := expressionRowCount(batches) if len(expr.selectList1) < rowCount { - expr.selectList1, err = ensureExpressionSlice( - expr.selectList1, - rowCount, - expr.m, - expr.allocation, - ExpressionAllocationSiteSelection, - ) - if err != nil { - return err - } + expr.selectList1 = make([]bool, rowCount) } remaining := expr.selectList1[:rowCount] if selectList != nil { @@ -1041,19 +760,6 @@ func (expr *FunctionExpressionExecutor) evalSelectedRows( rowCount int, selectList []bool, ) (*vector.Vector, error) { - var err error - if expr.allocation != nil { - expr.selectedRows, err = ensureExpressionSlice( - expr.selectedRows, - rowCount, - expr.m, - expr.allocation, - ExpressionAllocationSiteSelectedRows, - ) - if err != nil { - return nil, err - } - } expr.selectedRows = expr.selectedRows[:0] for row := 0; row < rowCount; row++ { if selectList[row] { @@ -1080,17 +786,7 @@ func (expr *FunctionExpressionExecutor) evalSelectedRows( if rowAligned && !parameter.IsConst() { selected := expr.selectedParameterVectors[i] if selected == nil { - var selection *vector.AllocationAccountSelection - if expr.allocation != nil { - selection = expr.allocation.scratch - } - selected, err = newExpressionVector( - *parameter.GetType(), - selection, - ) - if err != nil { - return nil, err - } + selected = vector.NewOffHeapVecWithType(*parameter.GetType()) expr.selectedParameterVectors[i] = selected } else { selected.Reset(*parameter.GetType()) @@ -1110,23 +806,7 @@ func (expr *FunctionExpressionExecutor) evalSelectedRows( return nil, err } if expr.selectedResult == nil { - if expr.allocation == nil { - expr.selectedResult = vector.NewFunctionResultWrapper( - expr.resultType, - expr.m, - ) - } else { - expr.selectedResult, err = - vector.NewFunctionResultWrapperWithFunctionAllocation( - expr.resultType, - expr.m, - expr.allocation.scratch, - expr.allocation.function, - ) - if err != nil { - return nil, err - } - } + expr.selectedResult = vector.NewFunctionResultWrapper(expr.resultType, expr.m) } expr.resetResultType(expr.selectedResult) if err := expr.selectedResult.PreExtendAndReset(selectedCount); err != nil { @@ -1146,19 +826,7 @@ func (expr *FunctionExpressionExecutor) evalSelectedRows( result.SetIsBin(runtimeIsBin) result.ResetWithSameType() if expr.selectedNullResult == nil { - var selection *vector.AllocationAccountSelection - if expr.allocation != nil { - selection = expr.allocation.scratch - } - expr.selectedNullResult, err = newExpressionConstNull( - runtimeType, - 1, - selection, - expr.m, - ) - if err != nil { - return nil, err - } + expr.selectedNullResult = vector.NewConstNull(runtimeType, 1, expr.m) } else { expr.selectedNullResult.SetType(runtimeType) expr.selectedNullResult.SetLength(1) @@ -1240,16 +908,7 @@ func (expr *FunctionExpressionExecutor) Eval(proc *process.Process, batches []*b return nil, err } if selectList != nil && len(expr.selectList.SelectList) < rowCount { - expr.selectList.SelectList, err = ensureExpressionSlice( - expr.selectList.SelectList, - rowCount, - expr.m, - expr.allocation, - ExpressionAllocationSiteSelection, - ) - if err != nil { - return nil, err - } + expr.selectList.SelectList = make([]bool, rowCount) } if selectList == nil { expr.selectList.AnyNull = false @@ -1284,9 +943,6 @@ func (expr *FunctionExpressionExecutor) EvalWithoutResultReusing(proc *process.P return nil, err } if expr.folded.canFold { - if vec.AllocationAccountSelection() != nil { - return vec.DupOffHeap(proc.Mp()) - } return vec.Dup(proc.Mp()) } expr.resultVector.SetResultVector(nil) @@ -1315,18 +971,6 @@ func (expr *FunctionExpressionExecutor) Free() { parameter.Free(expr.m) } } - freeExpressionSlice(expr.selectList1, expr.m, expr.allocation) - freeExpressionSlice(expr.selectList2, expr.m, expr.allocation) - freeExpressionSlice( - expr.selectList.SelectList, - expr.m, - expr.allocation, - ) - freeExpressionSlice(expr.selectedRows, expr.m, expr.allocation) - expr.selectList1 = nil - expr.selectList2 = nil - expr.selectList.SelectList = nil - expr.selectedRows = nil for _, p := range expr.parameterExecutor { if p != nil { @@ -1364,39 +1008,19 @@ func (expr *ColumnExpressionExecutor) Eval(_ *process.Process, batches []*batch. vec := batches[relIndex].Vecs[expr.colIndex] if vec.IsConstNull() { - var err error - vec, err = expr.getConstNullVec(expr.typ, vec.Length()) - if err != nil { - return nil, err - } + vec = expr.getConstNullVec(expr.typ, vec.Length()) } return vec, nil } -func (expr *ColumnExpressionExecutor) getConstNullVec( - typ types.Type, - length int, -) (*vector.Vector, error) { +func (expr *ColumnExpressionExecutor) getConstNullVec(typ types.Type, length int) *vector.Vector { if expr.nullVecCache != nil { expr.nullVecCache.SetType(typ) expr.nullVecCache.SetLength(length) } else { - var selection *vector.AllocationAccountSelection - if expr.allocation != nil { - selection = expr.allocation.result - } - var err error - expr.nullVecCache, err = newExpressionConstNull( - typ, - length, - selection, - expr.mp, - ) - if err != nil { - return nil, err - } + expr.nullVecCache = vector.NewConstNull(typ, length, expr.mp) } - return expr.nullVecCache, nil + return expr.nullVecCache } func (expr *ColumnExpressionExecutor) EvalWithoutResultReusing(proc *process.Process, batches []*batch.Batch, _ []bool) (*vector.Vector, error) { @@ -1434,9 +1058,6 @@ func (expr *FixedVectorExpressionExecutor) EvalWithoutResultReusing(proc *proces if err != nil { return nil, err } - if vec.AllocationAccountSelection() != nil { - return vec.DupOffHeap(proc.Mp()) - } return vec.Dup(proc.Mp()) } @@ -1456,135 +1077,111 @@ func (expr *FixedVectorExpressionExecutor) IsColumnExpr() bool { return false } -func generateConstExpressionExecutor( - proc *process.Process, - typ types.Type, - con *plan.Literal, - allocation *ExpressionAllocationAccount, -) (vec *vector.Vector, err error) { - var selection *vector.AllocationAccountSelection - if allocation != nil { - selection = allocation.constant - } +func generateConstExpressionExecutor(proc *process.Process, typ types.Type, con *plan.Literal) (vec *vector.Vector, err error) { if con.GetIsnull() { - vec, err = newExpressionConstNull(typ, 1, selection, proc.Mp()) + vec = vector.NewConstNull(typ, 1, proc.Mp()) } else { switch val := con.GetValue().(type) { case *plan.Literal_Bval: - vec, err = newExpressionConstFixed(constBType, val.Bval, 1, proc.Mp(), selection) + vec, err = vector.NewConstFixed(constBType, val.Bval, 1, proc.Mp()) case *plan.Literal_I8Val: - vec, err = newExpressionConstFixed(constI8Type, int8(val.I8Val), 1, proc.Mp(), selection) + vec, err = vector.NewConstFixed(constI8Type, int8(val.I8Val), 1, proc.Mp()) case *plan.Literal_I16Val: - vec, err = newExpressionConstFixed(constI16Type, int16(val.I16Val), 1, proc.Mp(), selection) + vec, err = vector.NewConstFixed(constI16Type, int16(val.I16Val), 1, proc.Mp()) case *plan.Literal_I32Val: - vec, err = newExpressionConstFixed(constI32Type, val.I32Val, 1, proc.Mp(), selection) + vec, err = vector.NewConstFixed(constI32Type, val.I32Val, 1, proc.Mp()) case *plan.Literal_I64Val: - vec, err = newExpressionConstFixed(constI64Type, val.I64Val, 1, proc.Mp(), selection) + vec, err = vector.NewConstFixed(constI64Type, val.I64Val, 1, proc.Mp()) case *plan.Literal_U8Val: - vec, err = newExpressionConstFixed(constU8Type, uint8(val.U8Val), 1, proc.Mp(), selection) + vec, err = vector.NewConstFixed(constU8Type, uint8(val.U8Val), 1, proc.Mp()) case *plan.Literal_U16Val: - vec, err = newExpressionConstFixed(constU16Type, uint16(val.U16Val), 1, proc.Mp(), selection) + vec, err = vector.NewConstFixed(constU16Type, uint16(val.U16Val), 1, proc.Mp()) case *plan.Literal_U32Val: - vec, err = newExpressionConstFixed(constU32Type, val.U32Val, 1, proc.Mp(), selection) + vec, err = vector.NewConstFixed(constU32Type, val.U32Val, 1, proc.Mp()) case *plan.Literal_U64Val: if typ.Oid == types.T_bit { - vec, err = newExpressionConstFixed(typ, val.U64Val, 1, proc.Mp(), selection) + vec, err = vector.NewConstFixed(typ, val.U64Val, 1, proc.Mp()) } else { - vec, err = newExpressionConstFixed(constU64Type, val.U64Val, 1, proc.Mp(), selection) + vec, err = vector.NewConstFixed(constU64Type, val.U64Val, 1, proc.Mp()) } case *plan.Literal_Fval: - vec, err = newExpressionConstFixed(constFType, val.Fval, 1, proc.Mp(), selection) + vec, err = vector.NewConstFixed(constFType, val.Fval, 1, proc.Mp()) case *plan.Literal_Dval: - vec, err = newExpressionConstFixed(constDType, val.Dval, 1, proc.Mp(), selection) + vec, err = vector.NewConstFixed(constDType, val.Dval, 1, proc.Mp()) case *plan.Literal_Dateval: - vec, err = newExpressionConstFixed(constDateType, types.Date(val.Dateval), 1, proc.Mp(), selection) + vec, err = vector.NewConstFixed(constDateType, types.Date(val.Dateval), 1, proc.Mp()) case *plan.Literal_Timeval: - vec, err = newExpressionConstFixed(typ, types.Time(val.Timeval), 1, proc.Mp(), selection) + vec, err = vector.NewConstFixed(typ, types.Time(val.Timeval), 1, proc.Mp()) case *plan.Literal_Datetimeval: - vec, err = newExpressionConstFixed(typ, types.Datetime(val.Datetimeval), 1, proc.Mp(), selection) + vec, err = vector.NewConstFixed(typ, types.Datetime(val.Datetimeval), 1, proc.Mp()) case *plan.Literal_Decimal64Val: cd64 := val.Decimal64Val d64 := types.Decimal64(cd64.A) - vec, err = newExpressionConstFixed(typ, d64, 1, proc.Mp(), selection) + vec, err = vector.NewConstFixed(typ, d64, 1, proc.Mp()) case *plan.Literal_Decimal128Val: cd128 := val.Decimal128Val d128 := types.Decimal128{B0_63: uint64(cd128.A), B64_127: uint64(cd128.B)} - vec, err = newExpressionConstFixed(typ, d128, 1, proc.Mp(), selection) + vec, err = vector.NewConstFixed(typ, d128, 1, proc.Mp()) case *plan.Literal_Timestampval: scale := typ.Scale if scale < 0 || scale > 6 { return nil, moerr.NewErrTooBigPrecision(proc.Ctx, int64(scale), "TIMESTAMP", 6) } - vec, err = newExpressionConstFixed( - constTimestampTypes[scale], - types.Timestamp(val.Timestampval), - 1, - proc.Mp(), - selection, - ) + vec, err = vector.NewConstFixed(constTimestampTypes[scale], types.Timestamp(val.Timestampval), 1, proc.Mp()) case *plan.Literal_Sval: sval := val.Sval // Distinguish binary with non-binary string. if typ.Oid == types.T_binary || typ.Oid == types.T_varbinary || typ.Oid == types.T_blob { - vec, err = newExpressionConstBytes(constBinType, []byte(sval), 1, proc.Mp(), selection) + vec, err = vector.NewConstBytes(constBinType, []byte(sval), 1, proc.Mp()) } else if typ.Oid == types.T_geometry { - vec, err = newExpressionConstBytes(typ, []byte(sval), 1, proc.Mp(), selection) + vec, err = vector.NewConstBytes(typ, []byte(sval), 1, proc.Mp()) } else if typ.Oid == types.T_array_float32 { array, err1 := types.StringToArray[float32](sval) if err1 != nil { return nil, err1 } - vec, err = newExpressionConstArray(typ, array, 1, proc.Mp(), selection) + vec, err = vector.NewConstArray(typ, array, 1, proc.Mp()) } else if typ.Oid == types.T_array_float64 { array, err1 := types.StringToArray[float64](sval) if err1 != nil { return nil, err1 } - vec, err = newExpressionConstArray(typ, array, 1, proc.Mp(), selection) + vec, err = vector.NewConstArray(typ, array, 1, proc.Mp()) } else if typ.Oid == types.T_datalink { _, _, err1 := datalink.ParseDatalink(sval, proc) if err1 != nil { return nil, err1 } - vec, err = newExpressionConstBytes(constBinType, []byte(sval), 1, proc.Mp(), selection) + vec, err = vector.NewConstBytes(constBinType, []byte(sval), 1, proc.Mp()) } else { - vec, err = newExpressionConstBytes(constSType, []byte(sval), 1, proc.Mp(), selection) + vec, err = vector.NewConstBytes(constSType, []byte(sval), 1, proc.Mp()) } case *plan.Literal_Defaultval: defaultVal := val.Defaultval - vec, err = newExpressionConstFixed(constBType, defaultVal, 1, proc.Mp(), selection) + vec, err = vector.NewConstFixed(constBType, defaultVal, 1, proc.Mp()) case *plan.Literal_EnumVal: - vec, err = newExpressionConstFixed(constEnumType, types.Enum(val.EnumVal), 1, proc.Mp(), selection) + vec, err = vector.NewConstFixed(constEnumType, types.Enum(val.EnumVal), 1, proc.Mp()) case *plan.Literal_VecVal: switch typ.Oid { case types.T_array_float32: - vec, err = newExpressionConstArray(typ, types.BytesToArray[float32]([]byte(val.VecVal)), 1, proc.Mp(), selection) + vec, err = vector.NewConstArray(typ, types.BytesToArray[float32]([]byte(val.VecVal)), 1, proc.Mp()) case types.T_array_float64: - vec, err = newExpressionConstArray(typ, types.BytesToArray[float64]([]byte(val.VecVal)), 1, proc.Mp(), selection) + vec, err = vector.NewConstArray(typ, types.BytesToArray[float64]([]byte(val.VecVal)), 1, proc.Mp()) case types.T_array_bf16: - vec, err = newExpressionConstArray(typ, types.BytesToArray[types.BF16]([]byte(val.VecVal)), 1, proc.Mp(), selection) + vec, err = vector.NewConstArray(typ, types.BytesToArray[types.BF16]([]byte(val.VecVal)), 1, proc.Mp()) case types.T_array_float16: - vec, err = newExpressionConstArray(typ, types.BytesToArray[types.Float16]([]byte(val.VecVal)), 1, proc.Mp(), selection) + vec, err = vector.NewConstArray(typ, types.BytesToArray[types.Float16]([]byte(val.VecVal)), 1, proc.Mp()) case types.T_array_int8: - vec, err = newExpressionConstArray(typ, types.BytesToArray[int8]([]byte(val.VecVal)), 1, proc.Mp(), selection) + vec, err = vector.NewConstArray(typ, types.BytesToArray[int8]([]byte(val.VecVal)), 1, proc.Mp()) case types.T_array_uint8: - vec, err = newExpressionConstArray(typ, types.BytesToArray[uint8]([]byte(val.VecVal)), 1, proc.Mp(), selection) + vec, err = vector.NewConstArray(typ, types.BytesToArray[uint8]([]byte(val.VecVal)), 1, proc.Mp()) } default: return nil, moerr.NewNYI(proc.Ctx, fmt.Sprintf("const expression %v", con.GetValue())) } + vec.SetIsBin(con.IsBin) } - if err != nil { - return nil, err - } - if vec == nil { - return nil, moerr.NewNYI( - proc.Ctx, - fmt.Sprintf("const expression %v", con.GetValue()), - ) - } - vec.SetIsBin(con.IsBin) - return vec, nil + return vec, err } func GenerateConstListExpressionExecutor(proc *process.Process, exprs []*plan.Expr) (*vector.Vector, error) { diff --git a/pkg/sql/colexec/evalExpressionReset.go b/pkg/sql/colexec/evalExpressionReset.go index e17f8c29c940a..da236e888e8c8 100644 --- a/pkg/sql/colexec/evalExpressionReset.go +++ b/pkg/sql/colexec/evalExpressionReset.go @@ -209,27 +209,13 @@ func (expr *FunctionExpressionExecutor) tryFoldFlowControl( return false, nil } -func (expr *FunctionExpressionExecutor) fillSkippedFlowControlParameters() ( - func(), - error, -) { +func (expr *FunctionExpressionExecutor) fillSkippedFlowControlParameters() func() { // The registered kernels still receive their complete argument list. Supply // typed NULLs for branches that lazy folding deliberately did not evaluate; // the selected conditions make those placeholders unobservable. var boolNull *vector.Vector var resultNull *vector.Vector temporaryIndexes := make([]int, 0, len(expr.parameterResults)) - cleanup := func() { - for _, i := range temporaryIndexes { - expr.parameterResults[i] = nil - } - if boolNull != nil { - boolNull.Free(expr.m) - } - if resultNull != nil { - resultNull.Free(expr.m) - } - } parameterCount := len(expr.parameterResults) for i := range expr.parameterResults { if expr.parameterResults[i] != nil { @@ -241,45 +227,28 @@ func (expr *FunctionExpressionExecutor) fillSkippedFlowControlParameters() ( } if isCondition { if boolNull == nil { - var selection *vector.AllocationAccountSelection - if expr.allocation != nil { - selection = expr.allocation.result - } - var err error - boolNull, err = newExpressionConstNull( - types.T_bool.ToType(), - 1, - selection, - expr.m, - ) - if err != nil { - return nil, err - } + boolNull = vector.NewConstNull(types.T_bool.ToType(), 1, expr.m) } expr.parameterResults[i] = boolNull } else { if resultNull == nil { - var selection *vector.AllocationAccountSelection - if expr.allocation != nil { - selection = expr.allocation.result - } - var err error - resultNull, err = newExpressionConstNull( - expr.resultType, - 1, - selection, - expr.m, - ) - if err != nil { - cleanup() - return nil, err - } + resultNull = vector.NewConstNull(expr.resultType, 1, expr.m) } expr.parameterResults[i] = resultNull } temporaryIndexes = append(temporaryIndexes, i) } - return cleanup, nil + return func() { + for _, i := range temporaryIndexes { + expr.parameterResults[i] = nil + } + if boolNull != nil { + boolNull.Free(expr.m) + } + if resultNull != nil { + resultNull.Free(expr.m) + } + } } func (expr *FunctionExpressionExecutor) finishFolding(proc *process.Process, execLen int) error { @@ -312,10 +281,7 @@ func (expr *FunctionExpressionExecutor) doFold(proc *process.Process, atRuntime if err != nil || !folded { return err } - cleanup, err := expr.fillSkippedFlowControlParameters() - if err != nil { - return err - } + cleanup := expr.fillSkippedFlowControlParameters() defer cleanup() return expr.finishFolding(proc, 1) } diff --git a/pkg/sql/colexec/eval_expression_allocation.go b/pkg/sql/colexec/eval_expression_allocation.go deleted file mode 100644 index ca31d0830658f..0000000000000 --- a/pkg/sql/colexec/eval_expression_allocation.go +++ /dev/null @@ -1,262 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package colexec - -import ( - "math" - - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/matrixorigin/matrixone/pkg/container/vector" -) - -// Expression allocation sites are stable diagnostics within the owner chosen -// by the caller. Legacy expression constructors still do not create or select -// an account; an activated owner must opt in explicitly. -const ( - ExpressionAllocationSiteConstantData mpool.AllocationSite = iota + 1 - ExpressionAllocationSiteConstantArea - ExpressionAllocationSiteResultData - ExpressionAllocationSiteResultArea - ExpressionAllocationSiteScratchData - ExpressionAllocationSiteScratchArea - ExpressionAllocationSiteSelection - ExpressionAllocationSiteSelectedRows - ExpressionAllocationSiteConstantNulls - ExpressionAllocationSiteConstantGrouping - ExpressionAllocationSiteResultNulls - ExpressionAllocationSiteResultGrouping - ExpressionAllocationSiteScratchNulls - ExpressionAllocationSiteScratchGrouping - ExpressionAllocationSiteParameterConversion - ExpressionAllocationSiteFunctionScratch -) - -// ExpressionAllocationAccount is the immutable allocation provenance shared -// by one expression tree. Vector selections remain separate so diagnostics -// distinguish constants, results, and selected-row scratch. -type ExpressionAllocationAccount struct { - account *mpool.AllocationAccount - owner mpool.AllocationOwner - - constant *vector.AllocationAccountSelection - result *vector.AllocationAccountSelection - scratch *vector.AllocationAccountSelection - function *vector.FunctionAllocation -} - -func NewExpressionAllocationAccount( - account *mpool.AllocationAccount, - owner mpool.AllocationOwner, -) (*ExpressionAllocationAccount, error) { - constant, err := vector.NewAllocationAccountSelectionWithBitmaps( - account, - owner, - ExpressionAllocationSiteConstantData, - ExpressionAllocationSiteConstantArea, - ExpressionAllocationSiteConstantNulls, - ExpressionAllocationSiteConstantGrouping, - ) - if err != nil { - return nil, err - } - result, err := vector.NewAllocationAccountSelectionWithBitmaps( - account, - owner, - ExpressionAllocationSiteResultData, - ExpressionAllocationSiteResultArea, - ExpressionAllocationSiteResultNulls, - ExpressionAllocationSiteResultGrouping, - ) - if err != nil { - return nil, err - } - scratch, err := vector.NewAllocationAccountSelectionWithBitmaps( - account, - owner, - ExpressionAllocationSiteScratchData, - ExpressionAllocationSiteScratchArea, - ExpressionAllocationSiteScratchNulls, - ExpressionAllocationSiteScratchGrouping, - ) - if err != nil { - return nil, err - } - function, err := vector.NewFunctionAllocation( - account, - owner, - ExpressionAllocationSiteParameterConversion, - ExpressionAllocationSiteFunctionScratch, - ) - if err != nil { - return nil, err - } - return &ExpressionAllocationAccount{ - account: account, - owner: owner, - constant: constant, - result: result, - scratch: scratch, - function: function, - }, nil -} - -func (a *ExpressionAllocationAccount) validate() error { - if a == nil || a.account == nil || a.account.Handle() == 0 || - a.owner < mpool.AllocationOwnerMin || - a.owner > mpool.AllocationOwnerMax || - a.constant == nil || a.result == nil || a.scratch == nil || - a.function == nil { - return mpool.ErrAllocationAccountInvalid - } - return nil -} - -func newExpressionVector( - typ types.Type, - selection *vector.AllocationAccountSelection, -) (*vector.Vector, error) { - if selection == nil { - return vector.NewOffHeapVecWithType(typ), nil - } - return vector.NewOffHeapVecWithTypeAndAllocation(typ, selection) -} - -func newExpressionConstNull( - typ types.Type, - length int, - selection *vector.AllocationAccountSelection, - mp *mpool.MPool, -) (*vector.Vector, error) { - if selection == nil { - return vector.NewConstNull(typ, length, mp), nil - } - return vector.NewConstNullWithAllocation(typ, length, selection) -} - -func newExpressionConstFixed[T any]( - typ types.Type, - value T, - length int, - mp *mpool.MPool, - selection *vector.AllocationAccountSelection, -) (*vector.Vector, error) { - if selection == nil { - return vector.NewConstFixed(typ, value, length, mp) - } - return vector.NewConstFixedWithAllocation( - typ, - value, - length, - mp, - selection, - ) -} - -func newExpressionConstBytes( - typ types.Type, - value []byte, - length int, - mp *mpool.MPool, - selection *vector.AllocationAccountSelection, -) (*vector.Vector, error) { - if selection == nil { - return vector.NewConstBytes(typ, value, length, mp) - } - return vector.NewConstBytesWithAllocation( - typ, - value, - length, - mp, - selection, - ) -} - -func newExpressionConstArray[T types.ArrayElement]( - typ types.Type, - value []T, - length int, - mp *mpool.MPool, - selection *vector.AllocationAccountSelection, -) (*vector.Vector, error) { - if selection == nil { - return vector.NewConstArray(typ, value, length, mp) - } - return vector.NewConstArrayWithAllocation( - typ, - value, - length, - mp, - selection, - ) -} - -func ensureExpressionSlice[T any]( - values []T, - length int, - mp *mpool.MPool, - allocation *ExpressionAllocationAccount, - site mpool.AllocationSite, -) ([]T, error) { - if length < 0 { - return nil, mpool.ErrAllocationAccountInvalid - } - if length <= cap(values) { - return values[:length], nil - } - if allocation == nil { - return make([]T, length), nil - } - if err := allocation.validate(); err != nil { - return nil, err - } - - newCapacity := cap(values) - if newCapacity == 0 { - newCapacity = 1 - } - for newCapacity < length { - if newCapacity > math.MaxInt/2 { - newCapacity = length - break - } - newCapacity *= 2 - } - next, err := mpool.MakeSliceAccounted[T]( - newCapacity, - mp, - allocation.account, - allocation.owner, - site, - ) - if err != nil { - return nil, err - } - copy(next, values) - if cap(values) > 0 { - mpool.FreeSlice(mp, values) - } - return next[:length], nil -} - -func freeExpressionSlice[T any]( - values []T, - mp *mpool.MPool, - allocation *ExpressionAllocationAccount, -) { - if allocation != nil && cap(values) > 0 { - mpool.FreeSlice(mp, values) - } -} diff --git a/pkg/sql/colexec/eval_expression_allocation_test.go b/pkg/sql/colexec/eval_expression_allocation_test.go deleted file mode 100644 index af0c5bf983595..0000000000000 --- a/pkg/sql/colexec/eval_expression_allocation_test.go +++ /dev/null @@ -1,545 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package colexec - -import ( - "testing" - - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/container/batch" - "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/matrixorigin/matrixone/pkg/container/vector" - "github.com/matrixorigin/matrixone/pkg/pb/plan" - "github.com/matrixorigin/matrixone/pkg/sql/plan/function" - "github.com/matrixorigin/matrixone/pkg/testutil" - "github.com/matrixorigin/matrixone/pkg/vm/process" - "github.com/stretchr/testify/require" -) - -type testExpressionAllocationAccount struct { - registry *mpool.AllocationAccountRegistry - account *mpool.AllocationAccount - allocation *ExpressionAllocationAccount -} - -func newTestExpressionAllocationAccount( - t testing.TB, - limit uint64, - metadataSlots uint64, -) testExpressionAllocationAccount { - t.Helper() - registry, err := mpool.NewAllocationAccountRegistry(1, metadataSlots) - require.NoError(t, err) - account, err := registry.Open(limit) - require.NoError(t, err) - allocation, err := NewExpressionAllocationAccount(account, 1) - require.NoError(t, err) - return testExpressionAllocationAccount{ - registry: registry, - account: account, - allocation: allocation, - } -} - -func finalizeTestExpressionAllocationAccount( - t testing.TB, - state testExpressionAllocationAccount, -) { - t.Helper() - snapshot := state.account.Seal() - require.Zero(t, snapshot.Used) - require.Zero(t, state.registry.LiveAllocationMetadata()) - _, err := state.registry.Finalize(state.account) - require.NoError(t, err) -} - -func expressionAllocationColumn(pos int32, typ types.Type) *plan.Expr { - return &plan.Expr{ - Typ: plan.Type{ - Id: int32(typ.Oid), - Width: typ.Width, - Scale: typ.Scale, - }, - Expr: &plan.Expr_Col{ - Col: &plan.ColRef{RelPos: 0, ColPos: pos}, - }, - } -} - -func expressionAllocationString(value string) *plan.Expr { - return &plan.Expr{ - Typ: plan.Type{ - Id: int32(types.T_varchar), - NotNullable: true, - }, - Expr: &plan.Expr_Lit{Lit: &plan.Literal{ - Value: &plan.Literal_Sval{Sval: value}, - }}, - } -} - -func expressionAllocationFunction( - t testing.TB, - proc *process.Process, - name string, - args ...*plan.Expr, -) *plan.Expr { - t.Helper() - argTypes := make([]types.Type, len(args)) - for i := range args { - argTypes[i] = types.New( - types.T(args[i].Typ.Id), - args[i].Typ.Width, - args[i].Typ.Scale, - ) - } - fn, err := function.GetFunctionByName(proc.Ctx, name, argTypes) - require.NoError(t, err) - retType := fn.GetReturnType() - return &plan.Expr{ - Typ: plan.Type{ - Id: int32(retType.Oid), - Width: retType.Width, - Scale: retType.Scale, - }, - Expr: &plan.Expr_F{F: &plan.Function{ - Func: &plan.ObjectRef{ - Obj: fn.GetEncodedOverloadID(), - ObjName: name, - }, - Args: args, - }}, - } -} - -func expressionAllocationCast( - t testing.TB, - proc *process.Process, - source *plan.Expr, - targetType types.Type, -) *plan.Expr { - t.Helper() - target := &plan.Expr{ - Typ: plan.Type{ - Id: int32(targetType.Oid), - Width: targetType.Width, - Scale: targetType.Scale, - NotNullable: true, - }, - Expr: &plan.Expr_T{T: &plan.TargetType{}}, - } - return expressionAllocationFunction(t, proc, "cast", source, target) -} - -func TestExpressionAllocationAccountNestedSelectedLifecycle(t *testing.T) { - proc := testutil.NewProcessWithMPool( - t, - "", - mpool.MustNew("expression-allocation-lifecycle"), - ) - defer proc.Free() - state := newTestExpressionAllocationAccount(t, 64<<20, 128) - - input := testutil.NewBatchWithVectors([]*vector.Vector{ - testutil.NewVector( - 4, - types.T_bool.ToType(), - proc.Mp(), - false, - []bool{true, false, true, false}, - ), - testutil.NewVector( - 4, - types.T_varchar.ToType(), - proc.Mp(), - false, - []string{"a", "b", "c", "d"}, - ), - testutil.NewVector( - 4, - types.T_int64.ToType(), - proc.Mp(), - false, - []int64{10, 20, 30, 40}, - ), - }, nil) - defer input.Clean(proc.Mp()) - - thenExpr := expressionAllocationFunction( - t, - proc, - "concat", - expressionAllocationCast( - t, - proc, - expressionAllocationColumn(2, types.T_int64.ToType()), - types.T_varchar.ToType(), - ), - expressionAllocationString("-then-payload-longer-than-inline"), - ) - elseExpr := expressionAllocationFunction( - t, - proc, - "concat", - expressionAllocationColumn(1, types.T_varchar.ToType()), - expressionAllocationString("-else-payload-longer-than-inline"), - ) - caseExpr := expressionAllocationFunction( - t, - proc, - "case", - expressionAllocationColumn(0, types.T_bool.ToType()), - thenExpr, - elseExpr, - ) - executor, err := NewExpressionExecutorWithAllocation( - proc, - caseExpr, - state.allocation, - ) - require.NoError(t, err) - require.Positive(t, state.account.Snapshot().Used) - - result, err := executor.Eval( - proc, - []*batch.Batch{input}, - []bool{true, false, true, false}, - ) - require.NoError(t, err) - require.NotNil(t, result.AllocationAccountSelection()) - require.Equal(t, "10-then-payload-longer-than-inline", result.GetStringAt(0)) - require.True(t, result.IsNull(1)) - require.Equal(t, "30-then-payload-longer-than-inline", result.GetStringAt(2)) - require.True(t, result.IsNull(3)) - - root := executor.(*FunctionExpressionExecutor) - require.GreaterOrEqual(t, cap(root.selectList1), input.RowCount()) - require.GreaterOrEqual(t, cap(root.selectList2), input.RowCount()) - require.GreaterOrEqual(t, cap(root.selectedRows), input.RowCount()) - require.NotNil(t, root.selectedResult) - require.NotNil(t, root.selectedResult.GetResultVector()) - require.NotNil( - t, - root.selectedResult.GetResultVector().AllocationAccountSelection(), - ) - usedAfterPartial := state.account.Snapshot().Used - - executor.ResetForNextQuery() - result, err = executor.Eval( - proc, - []*batch.Batch{input}, - []bool{true, false, true, false}, - ) - require.NoError(t, err) - require.Equal(t, usedAfterPartial, state.account.Snapshot().Used) - - transferred, err := executor.EvalWithoutResultReusing( - proc, - []*batch.Batch{input}, - nil, - ) - require.NoError(t, err) - require.NotNil(t, transferred.AllocationAccountSelection()) - executor.Free() - require.Positive(t, state.account.Snapshot().Used) - transferred.Free(proc.Mp()) - require.Zero(t, state.account.Snapshot().Used) - finalizeTestExpressionAllocationAccount(t, state) -} - -func TestExpressionAllocationAccountConstantKinds(t *testing.T) { - proc := testutil.NewProcessWithMPool( - t, - "", - mpool.MustNew("expression-allocation-constants"), - ) - defer proc.Free() - state := newTestExpressionAllocationAccount(t, 1<<20, 8) - - fixed := &plan.Expr{ - Typ: plan.Type{ - Id: int32(types.T_int64), - NotNullable: true, - }, - Expr: &plan.Expr_Lit{Lit: &plan.Literal{ - Value: &plan.Literal_I64Val{I64Val: 42}, - }}, - } - null := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int64)}, - Expr: &plan.Expr_Lit{Lit: &plan.Literal{Isnull: true}}, - } - executors, err := - NewExpressionExecutorsFromPlanExpressionsWithAllocation( - proc, - []*plan.Expr{fixed, null}, - state.allocation, - ) - require.NoError(t, err) - require.Positive(t, state.account.Snapshot().Used) - for _, executor := range executors { - fixedExecutor := executor.(*FixedVectorExpressionExecutor) - require.NotNil( - t, - fixedExecutor.resultVector.AllocationAccountSelection(), - ) - _, err = executor.Eval( - proc, - []*batch.Batch{batch.EmptyForConstFoldBatch}, - nil, - ) - require.NoError(t, err) - } - for _, executor := range executors { - executor.Free() - } - require.Zero(t, state.account.Snapshot().Used) - finalizeTestExpressionAllocationAccount(t, state) -} - -func TestExpressionAllocationAccountDecodedVectorTransfer(t *testing.T) { - proc := testutil.NewProcessWithMPool( - t, - "", - mpool.MustNew("expression-allocation-decoded-vector"), - ) - defer proc.Free() - state := newTestExpressionAllocationAccount(t, 1<<20, 8) - - source := testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()) - data, err := source.MarshalBinary() - require.NoError(t, err) - source.Free(proc.Mp()) - - executor, err := NewExpressionExecutorWithAllocation( - proc, - &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Vec{Vec: &plan.LiteralVec{ - Len: 3, - Data: data, - }}, - }, - state.allocation, - ) - require.NoError(t, err) - fixed := executor.(*FixedVectorExpressionExecutor) - require.NotNil( - t, - fixed.resultVector.AllocationAccountSelection(), - ) - require.False(t, fixed.resultVector.NeedDup()) - require.Positive(t, state.account.Snapshot().Used) - - transferred, err := executor.EvalWithoutResultReusing( - proc, - []*batch.Batch{batch.EmptyForConstFoldBatch}, - nil, - ) - require.NoError(t, err) - executor.Free() - require.Positive(t, state.account.Snapshot().Used) - transferred.Free(proc.Mp()) - require.Zero(t, state.account.Snapshot().Used) - finalizeTestExpressionAllocationAccount(t, state) -} - -func TestExpressionAllocationAccountFoldedTransfer(t *testing.T) { - proc := testutil.NewProcessWithMPool( - t, - "", - mpool.MustNew("expression-allocation-folded-transfer"), - ) - defer proc.Free() - state := newTestExpressionAllocationAccount(t, 1<<20, 16) - - expr := expressionAllocationFunction( - t, - proc, - "concat", - expressionAllocationString("folded-payload-longer-than-inline"), - expressionAllocationString("-suffix"), - ) - executor, err := NewExpressionExecutorWithAllocation( - proc, - expr, - state.allocation, - ) - require.NoError(t, err) - - transferred, err := executor.EvalWithoutResultReusing( - proc, - []*batch.Batch{batch.EmptyForConstFoldBatch}, - nil, - ) - require.NoError(t, err) - require.Equal( - t, - "folded-payload-longer-than-inline-suffix", - transferred.GetStringAt(0), - ) - require.NotNil(t, transferred.AllocationAccountSelection()) - executor.Free() - require.Positive(t, state.account.Snapshot().Used) - transferred.Free(proc.Mp()) - require.Zero(t, state.account.Snapshot().Used) - finalizeTestExpressionAllocationAccount(t, state) -} - -func TestExpressionAllocationAccountConstructionRollback(t *testing.T) { - proc := testutil.NewProcessWithMPool( - t, - "", - mpool.MustNew("expression-allocation-construction"), - ) - defer proc.Free() - state := newTestExpressionAllocationAccount(t, 1<<20, 1) - - expr := expressionAllocationFunction( - t, - proc, - "concat", - expressionAllocationString("left"), - expressionAllocationString("right"), - ) - _, err := NewExpressionExecutorWithAllocation( - proc, - expr, - state.allocation, - ) - require.ErrorIs(t, err, mpool.ErrAllocationMetadataSlots) - require.Zero(t, state.account.Snapshot().Used) - require.Zero(t, state.registry.LiveAllocationMetadata()) - finalizeTestExpressionAllocationAccount(t, state) -} - -func TestExpressionAllocationAccountConstNullPreservesBinaryFlag(t *testing.T) { - proc := testutil.NewProcessWithMPool( - t, - "", - mpool.MustNew("expression-allocation-const-null"), - ) - defer proc.Free() - state := newTestExpressionAllocationAccount(t, 1<<20, 4) - - expr := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_varbinary)}, - Expr: &plan.Expr_Lit{Lit: &plan.Literal{ - Isnull: true, - IsBin: true, - }}, - } - executor, err := NewExpressionExecutorWithAllocation( - proc, - expr, - state.allocation, - ) - require.NoError(t, err) - result, err := executor.Eval( - proc, - []*batch.Batch{batch.EmptyForConstFoldBatch}, - nil, - ) - require.NoError(t, err) - require.True(t, result.IsConstNull()) - require.True(t, result.GetIsBin()) - - executor.Free() - finalizeTestExpressionAllocationAccount(t, state) -} - -func TestExpressionAllocationAccountScratchFailureCleanup(t *testing.T) { - proc := testutil.NewProcessWithMPool( - t, - "", - mpool.MustNew("expression-allocation-scratch-failure"), - ) - defer proc.Free() - state := newTestExpressionAllocationAccount(t, 8, 8) - - input := testutil.NewBatchWithVectors([]*vector.Vector{ - testutil.NewVector( - 8, - types.T_bool.ToType(), - proc.Mp(), - false, - []bool{true, false, true, false, true, false, true, false}, - ), - testutil.NewVector( - 8, - types.T_int64.ToType(), - proc.Mp(), - false, - []int64{1, 2, 3, 4, 5, 6, 7, 8}, - ), - }, nil) - defer input.Clean(proc.Mp()) - - expr := expressionAllocationFunction( - t, - proc, - "case", - expressionAllocationColumn(0, types.T_bool.ToType()), - expressionAllocationColumn(1, types.T_int64.ToType()), - expressionAllocationColumn(1, types.T_int64.ToType()), - ) - executor, err := NewExpressionExecutorWithAllocation( - proc, - expr, - state.allocation, - ) - require.NoError(t, err) - _, err = executor.Eval(proc, []*batch.Batch{input}, nil) - require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) - require.Equal(t, uint64(8), state.account.Snapshot().Used) - - executor.Free() - require.Zero(t, state.account.Snapshot().Used) - require.Zero(t, state.registry.LiveAllocationMetadata()) - finalizeTestExpressionAllocationAccount(t, state) -} - -func TestExpressionAllocationAccountZeroLengthScratchGrowth(t *testing.T) { - mp := mpool.MustNew("expression-allocation-zero-length-scratch") - defer mpool.DeleteMPool(mp) - state := newTestExpressionAllocationAccount(t, 1<<20, 4) - - values, err := ensureExpressionSlice( - []int64(nil), - 4, - mp, - state.allocation, - ExpressionAllocationSiteSelectedRows, - ) - require.NoError(t, err) - require.Equal(t, uint64(32), state.account.Snapshot().Used) - - values = values[:0] - values, err = ensureExpressionSlice( - values, - 8, - mp, - state.allocation, - ExpressionAllocationSiteSelectedRows, - ) - require.NoError(t, err) - require.Len(t, values, 8) - require.Equal(t, uint64(64), state.account.Snapshot().Used) - - values = values[:0] - freeExpressionSlice(values, mp, state.allocation) - require.Zero(t, state.account.Snapshot().Used) - finalizeTestExpressionAllocationAccount(t, state) -} diff --git a/pkg/sql/colexec/fuzzyfilter/filter.go b/pkg/sql/colexec/fuzzyfilter/filter.go index f51b12c88d64e..aab0610860356 100644 --- a/pkg/sql/colexec/fuzzyfilter/filter.go +++ b/pkg/sql/colexec/fuzzyfilter/filter.go @@ -19,6 +19,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/bloomfilter" "github.com/matrixorigin/matrixone/pkg/common/hashmap/keycodec" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" @@ -148,9 +149,19 @@ func (fuzzyFilter *FuzzyFilter) Prepare(proc *process.Process) (err error) { ) != keycodec.ExactRuntimeFilterUnsupported } if ctr.runtimeFilterUsable { + if fuzzyFilter.allocationAccount == nil || + fuzzyFilter.runtimeFilterAllocation == nil { + return mpool.ErrAllocationAccountInvalid + } if ctr.pass2RuntimeFilter == nil { - ctr.pass2RuntimeFilter = vector.NewOffHeapVecWithType( - plan.MakeTypeByPlan2Type(fuzzyFilter.PkTyp)) + ctr.pass2RuntimeFilter, err = + vector.NewOffHeapVecWithTypeAndAllocation( + plan.MakeTypeByPlan2Type(fuzzyFilter.PkTyp), + fuzzyFilter.runtimeFilterAllocation, + ) + if err != nil { + return err + } } } else if ctr.pass2RuntimeFilter != nil { // FuzzyFilter must still execute its uniqueness check, but an @@ -410,16 +421,13 @@ func (fuzzyFilter *FuzzyFilter) handleRuntimeFilter(proc *process.Process) error // Reset bitmap before sort to avoid corruption. ctr.pass2RuntimeFilter.GetNulls().Reset() ctr.pass2RuntimeFilter.InplaceSort() - budget, err := proc.GetHashBuildBudget() - if err != nil { - if fuzzyFilter.fallbackRuntimeFilter(proc, err) { - return nil - } - fuzzyFilter.abandonRuntimeFilter(proc) - return err - } data, release, err := runtimefilter.MarshalExactFilterVector( - ctr.pass2RuntimeFilter, budget) + ctr.pass2RuntimeFilter, + proc.Mp(), + fuzzyFilter.allocationAccount, + fuzzyFilterAllocationOwner, + fuzzyFilterAllocationSiteRuntimeFilterPayload, + ) if err != nil { if fuzzyFilter.fallbackRuntimeFilter(proc, err) { return nil @@ -536,11 +544,6 @@ func (fuzzyFilter *FuzzyFilter) generate() error { ctr := &fuzzyFilter.ctr rbat := batch.NewWithSize(1) rbat.SetVector(0, vector.NewVec(plan.MakeTypeByPlan2Type(fuzzyFilter.PkTyp))) - // Runtime-filter retention is optional and can grow to the configured IN - // cardinality. Keep it off-heap so the process pool can reject growth - // recoverably; appendPassToRuntimeFilter then abandons it and sends PASS. - ctr.pass2RuntimeFilter = vector.NewOffHeapVecWithType( - plan.MakeTypeByPlan2Type(fuzzyFilter.PkTyp)) ctr.rbat = rbat return nil } diff --git a/pkg/sql/colexec/fuzzyfilter/filter_test.go b/pkg/sql/colexec/fuzzyfilter/filter_test.go index 57a42546aa39d..c84ffd84b5d9a 100644 --- a/pkg/sql/colexec/fuzzyfilter/filter_test.go +++ b/pkg/sql/colexec/fuzzyfilter/filter_test.go @@ -160,7 +160,7 @@ func TestRuntimeFilterContract(t *testing.T) { spec := newRuntimeFilterSpec(101, probeType, payloadType) arg, proc := newRuntimeFilterTest(t, spec, payloadType) - require.NoError(t, arg.Prepare(proc)) + prepareFuzzyFilter(t, arg, proc) require.False(t, arg.ctr.runtimeFilterUsable) require.Nil(t, arg.ctr.pass2RuntimeFilter) @@ -187,7 +187,7 @@ func TestRuntimeFilterContract(t *testing.T) { spec := newRuntimeFilterSpec(102, typ, typ) arg, proc := newRuntimeFilterTest(t, spec, typ) - require.NoError(t, arg.Prepare(proc)) + prepareFuzzyFilter(t, arg, proc) require.True(t, arg.ctr.runtimeFilterUsable) require.NotNil(t, arg.ctr.pass2RuntimeFilter) require.Zero(t, arg.ctr.pass2RuntimeFilter.Length()) @@ -206,7 +206,7 @@ func TestRuntimeFilterContract(t *testing.T) { spec := newRuntimeFilterSpec(103, typ, typ) arg, proc := newRuntimeFilterTest(t, spec, typ) - require.NoError(t, arg.Prepare(proc)) + prepareFuzzyFilter(t, arg, proc) require.True(t, arg.ctr.runtimeFilterUsable) payload := testutil.MakeInt64Vector([]int64{7}, nil, proc.Mp()) require.NoError(t, arg.appendPassToRuntimeFilter(payload, proc)) @@ -225,7 +225,7 @@ func TestRuntimeFilterContract(t *testing.T) { spec := newRuntimeFilterSpec(104, typ, typ) arg, proc := newRuntimeFilterTest(t, spec, typ) - require.NoError(t, arg.Prepare(proc)) + prepareFuzzyFilter(t, arg, proc) payload := testutil.MakeInt32Vector([]int32{7, 3}, nil, proc.Mp()) require.NoError(t, arg.appendPassToRuntimeFilter(payload, proc)) @@ -275,7 +275,7 @@ func TestFuzzyRuntimeFilterCopyFailureFailsOpen(t *testing.T) { arg := newArgument(typ) arg.N = 1 arg.RuntimeFilterSpec = spec - require.NoError(t, arg.Prepare(proc)) + prepareFuzzyFilter(t, arg, proc) sourceMP := mpool.MustNewZero() payload := vector.NewVec(typ) @@ -338,7 +338,7 @@ func TestFuzzyRuntimeFilterClosureFailureFailsOpen(t *testing.T) { arg := newArgument(typ) arg.N = 1 arg.RuntimeFilterSpec = spec - require.NoError(t, arg.Prepare(proc)) + prepareFuzzyFilter(t, arg, proc) require.NoError(t, vector.AppendFixed( arg.ctr.pass2RuntimeFilter, float64(0), false, limited)) for arg.ctr.pass2RuntimeFilter.Length() < @@ -381,17 +381,26 @@ func TestFuzzyRuntimeFilterBudgetErrorPolicy(t *testing.T) { typ := types.T_int32.ToType() spec := newRuntimeFilterSpec(109, typ, typ) arg, proc := newRuntimeFilterTest(t, spec, typ) - require.NoError(t, arg.Prepare(proc)) + budget := process.MustNewHashBuildBudget(1<<20, 1<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.OpenWithController( + generation.Cap(), generation) + require.NoError(t, err) + require.NoError(t, arg.SetAllocationAccount(account)) + prepareFuzzyFilter(t, arg, proc) require.NoError(t, vector.AppendFixed( arg.ctr.pass2RuntimeFilter, int32(1), false, proc.Mp())) - generation, err := proc.GetHashBuildBudget() - require.NoError(t, err) - var held *process.HashBuildReservation + var filler []byte if test.closed { generation.Close() } else { - held, err = generation.Reserve(generation.Cap()) + remaining := account.Snapshot().Limit - account.Snapshot().Used + filler, err = proc.Mp().AllocAccounted( + int(remaining), account, 63, 255) require.NoError(t, err) } @@ -410,7 +419,8 @@ func TestFuzzyRuntimeFilterBudgetErrorPolicy(t *testing.T) { require.True(t, arg.ctr.runtimeFilterDone) require.Equal(t, int64(1), stats["FuzzyFilterRuntimeFilterBudgetFallbacks"]) - require.True(t, held.Release()) + proc.Mp().Free(filler) + generation.Close() } require.False(t, arg.ctr.runtimeFilterUsable) require.Nil(t, arg.ctr.pass2RuntimeFilter) @@ -447,7 +457,7 @@ func TestFuzzyCallErrorUnblocksRuntimeFilterBeforeReset(t *testing.T) { require.NoError(t, buildChild.Prepare(proc)) require.NoError(t, probeChild.Prepare(proc)) - require.NoError(t, arg.Prepare(proc)) + prepareFuzzyFilter(t, arg, proc) _, err := vm.Exec(arg, proc) require.ErrorIs(t, err, buildErr) @@ -473,7 +483,7 @@ func TestFuzzyCallErrorUnblocksRuntimeFilterBeforeReset(t *testing.T) { require.True(t, arg.ctr.runtimeFilterDone) proc.GetMessageBoard().Reset() - require.NoError(t, arg.Prepare(proc)) + prepareFuzzyFilter(t, arg, proc) require.False(t, arg.ctr.runtimeFilterDone, "Prepare must open the terminal gate for the next generation") arg.finalizeBuildFailure(proc) @@ -581,6 +591,22 @@ func newRuntimeFilterTest( return arg, proc } +func prepareFuzzyFilter( + t *testing.T, + arg *FuzzyFilter, + proc *process.Process, +) { + t.Helper() + if arg.allocationAccount == nil { + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.Open(1 << 60) + require.NoError(t, err) + require.NoError(t, arg.SetAllocationAccount(account)) + } + require.NoError(t, arg.Prepare(proc)) +} + func receiveRuntimeFilter( t *testing.T, proc *process.Process, diff --git a/pkg/sql/colexec/fuzzyfilter/types.go b/pkg/sql/colexec/fuzzyfilter/types.go index 40f18329461ad..48627bcfb8908 100644 --- a/pkg/sql/colexec/fuzzyfilter/types.go +++ b/pkg/sql/colexec/fuzzyfilter/types.go @@ -17,6 +17,7 @@ package fuzzyfilter import ( "github.com/matrixorigin/matrixone/pkg/common/bloomfilter" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/reuse" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" @@ -37,6 +38,16 @@ const ( End ) +const fuzzyFilterAllocationOwner mpool.AllocationOwner = 1 + +const ( + fuzzyFilterAllocationSiteRuntimeFilterData mpool.AllocationSite = iota + 1 + fuzzyFilterAllocationSiteRuntimeFilterArea + fuzzyFilterAllocationSiteRuntimeFilterNulls + fuzzyFilterAllocationSiteRuntimeFilterGrouping + fuzzyFilterAllocationSiteRuntimeFilterPayload +) + type container struct { state int @@ -55,7 +66,9 @@ type container struct { } type FuzzyFilter struct { - ctr container + ctr container + allocationAccount *mpool.AllocationAccount + runtimeFilterAllocation *vector.AllocationAccountSelection // Estimates of the number of data items obtained from statistical information N float64 @@ -69,6 +82,51 @@ type FuzzyFilter struct { vm.OperatorBase } +func (fuzzyFilter *FuzzyFilter) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if account == nil { + return mpool.ErrAllocationAccountInvalid + } + if fuzzyFilter.allocationAccount != nil { + if fuzzyFilter.allocationAccount == account { + return nil + } + return mpool.ErrAllocationAccountMismatch + } + selection, err := vector.NewAllocationAccountSelection( + account, + fuzzyFilterAllocationOwner, + fuzzyFilterAllocationSiteRuntimeFilterData, + fuzzyFilterAllocationSiteRuntimeFilterArea, + fuzzyFilterAllocationSiteRuntimeFilterNulls, + fuzzyFilterAllocationSiteRuntimeFilterGrouping, + ) + if err != nil { + return err + } + fuzzyFilter.allocationAccount = account + fuzzyFilter.runtimeFilterAllocation = selection + return nil +} + +func (fuzzyFilter *FuzzyFilter) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if fuzzyFilter.allocationAccount == nil { + return nil + } + if fuzzyFilter.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if fuzzyFilter.ctr.pass2RuntimeFilter != nil { + return mpool.ErrAllocationAccountInvariant + } + fuzzyFilter.allocationAccount = nil + fuzzyFilter.runtimeFilterAllocation = nil + return nil +} + func (fuzzyFilter *FuzzyFilter) GetOperatorBase() *vm.OperatorBase { return &fuzzyFilter.OperatorBase } @@ -127,7 +185,8 @@ func (fuzzyFilter *FuzzyFilter) Reset(proc *process.Process, pipelineFailed bool ctr.runtimeFilterUsable = false ctr.collisionCnt = 0 if ctr.pass2RuntimeFilter != nil { - ctr.pass2RuntimeFilter.CleanOnlyData() + ctr.pass2RuntimeFilter.Free(proc.Mp()) + ctr.pass2RuntimeFilter = nil } if ctr.rbat != nil { ctr.rbat.CleanOnlyData() diff --git a/pkg/sql/colexec/group/helper.go b/pkg/sql/colexec/group/helper.go index 4e5f408a927aa..d6b2d81b40759 100644 --- a/pkg/sql/colexec/group/helper.go +++ b/pkg/sql/colexec/group/helper.go @@ -391,7 +391,9 @@ func (ctr *container) spillDataToDisk(proc *process.Process, opAnalyzer process. } } gbBatch.SetRowCount(int(cnt)) - gbBatch.MarshalBinaryWithBuffer(buf, false) + if _, err := gbBatch.MarshalBinaryWithBuffer(buf, false); err != nil { + return 0, 0, err + } // write marker var magic uint64 = 0x12345678DEADBEEF diff --git a/pkg/sql/colexec/hashbuild/allocation_test_helpers_test.go b/pkg/sql/colexec/hashbuild/allocation_test_helpers_test.go new file mode 100644 index 0000000000000..5e7545f95564d --- /dev/null +++ b/pkg/sql/colexec/hashbuild/allocation_test_helpers_test.go @@ -0,0 +1,71 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package hashbuild + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/vm/process" + "github.com/stretchr/testify/require" +) + +func newTestAllocationAccount(t testing.TB) *mpool.AllocationAccount { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.Open(1 << 60) + require.NoError(t, err) + return account +} + +func installTestHashBuildBudget( + t testing.TB, + op *HashBuild, + generation *process.HashBuildBudgetGeneration, +) { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.OpenWithController( + generation.Cap(), generation) + require.NoError(t, err) + replaceTestHashBuildAllocation(t, op, account) + op.ctr.hashmapBuilder.setBudget(generation) +} + +func newTestHashmapBuilder(t testing.TB) *HashmapBuilder { + t.Helper() + builder := &HashmapBuilder{} + require.NoError(t, builder.SetAllocationAccount(newTestAllocationAccount(t))) + return builder +} + +func installTestHashBuildAllocation(t testing.TB, op *HashBuild) { + t.Helper() + require.NoError(t, op.SetAllocationAccount(newTestAllocationAccount(t))) +} + +func replaceTestHashBuildAllocation( + t testing.TB, + op *HashBuild, + account *mpool.AllocationAccount, +) { + t.Helper() + if current := op.ctr.hashmapBuilder.mapAllocationAccount; current != nil { + require.NoError(t, op.ClearAllocationAccount(current)) + } + require.NoError(t, op.SetAllocationAccount(account)) +} diff --git a/pkg/sql/colexec/hashbuild/budget.go b/pkg/sql/colexec/hashbuild/budget.go index 9d61b37ca40a8..31d9fd1ebc05e 100644 --- a/pkg/sql/colexec/hashbuild/budget.go +++ b/pkg/sql/colexec/hashbuild/budget.go @@ -15,169 +15,35 @@ package hashbuild import ( - "math" - "sync" - "github.com/matrixorigin/matrixone/pkg/common/hashmap" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/hashtable" "github.com/matrixorigin/matrixone/pkg/container/vector" - "github.com/matrixorigin/matrixone/pkg/sql/colexec" "github.com/matrixorigin/matrixone/pkg/sql/colexec/runtimefilter" "github.com/matrixorigin/matrixone/pkg/vm/message" "github.com/matrixorigin/matrixone/pkg/vm/process" ) -type hashMapResizeReservation struct { - owner *hashMapReservationOwner - token *process.HashBuildReservation -} - -func (r *hashMapResizeReservation) Commit(plan hashtable.ResizePlan) { - r.owner.commit(r.token, plan.ReuseCurrentBlocks) - r.token = nil -} - -func (r *hashMapResizeReservation) Rollback() { - if r.token != nil { - r.token.Release() - r.token = nil - } -} - -// hashMapReservationOwner follows the physical hash table across producer to -// JoinMap ownership transfer. Full-table replacement swaps the retained token; -// segmented growth keeps the existing tokens and adds one for the appended -// blocks. Resize callbacks retain this owner so consumer growth never stores -// reservations back into a reused producer. -type hashMapReservationOwner struct { - mu sync.Mutex - tokens []*process.HashBuildReservation -} - -func (o *hashMapReservationOwner) commit(token *process.HashBuildReservation, reuseCurrent bool) { - o.mu.Lock() - if reuseCurrent { - o.tokens = append(o.tokens, token) - o.mu.Unlock() - return - } - old := o.tokens - o.tokens = []*process.HashBuildReservation{token} - o.mu.Unlock() - for _, reservation := range old { - reservation.Release() - } -} - -func (o *hashMapReservationOwner) release() { - if o == nil { - return - } - o.mu.Lock() - tokens := o.tokens - o.tokens = nil - o.mu.Unlock() - for _, token := range tokens { - token.Release() - } -} - -func (hb *HashmapBuilder) setBudget(budget *process.HashBuildBudgetGeneration) { +// setBudget retains the statement generation for non-memory resource ledgers +// such as spill files and disk. Physical memory admission is exclusively +// driven by the allocation account installed through SetAllocationAccount. +func (hb *HashmapBuilder) setBudget( + budget *process.HashBuildBudgetGeneration, +) { hb.budget = budget } // SetBudget is the exported boundary used by spill and integration tests. -func (hb *HashmapBuilder) SetBudget(budget *process.HashBuildBudgetGeneration) { hb.setBudget(budget) } - -func (hb *HashmapBuilder) reserveInitialMap(size int64) error { - if hb.budget == nil || size <= 0 { - return nil - } - reservation, err := hb.budget.Reserve(uint64(size)) - if err != nil { - return err - } - hb.mapReservation = &hashMapReservationOwner{tokens: []*process.HashBuildReservation{reservation}} - return nil -} - -func resizeAdmission(budget *process.HashBuildBudgetGeneration, owner *hashMapReservationOwner, plan hashtable.ResizePlan) (hashtable.ResizeReservation, error) { - if budget == nil || plan.AdditionalBytes == 0 { - return nil, nil - } - token, err := budget.Reserve(plan.AdditionalBytes) - if err != nil { - return nil, err - } - return &hashMapResizeReservation{owner: owner, token: token}, nil -} - -// NewBudgetedEmptyJoinMap creates an initially empty JoinMap whose complete -// physical hash-table lifetime is charged to budget. The initial allocation is -// admitted before touching the mpool, every later resize uses the same -// generation, and JoinMap.Free releases all retained reservations. -// -// This is used by consumers that must grow a hash table from probe-side keys -// (for example RightDedupJoin after an empty build partition). Such maps cannot -// use the regular HashmapBuilder ownership transfer because there is no build -// batch to publish. -func NewBudgetedEmptyJoinMap( - keyWidth int, +func (hb *HashmapBuilder) SetBudget( budget *process.HashBuildBudgetGeneration, - mp *mpool.MPool, -) (*message.JoinMap, error) { - if budget == nil || mp == nil { - return nil, process.ErrHashBuildBudgetInvalid - } - - initialBytes := hashtable.Int64HashMapInitialAllocationBytes() - if keyWidth > 8 { - initialBytes = hashtable.StringHashMapInitialAllocationBytes() - } - initial, err := budget.Reserve(initialBytes) - if err != nil { - return nil, err - } - owner := &hashMapReservationOwner{ - tokens: []*process.HashBuildReservation{initial}, - } - - var ( - intHashMap *hashmap.IntHashMap - strHashMap *hashmap.StrHashMap - ) - if keyWidth <= 8 { - intHashMap, err = hashmap.NewIntHashMap(false, mp) - if err == nil { - intHashMap.SetResizeAdmission(func(plan hashtable.ResizePlan) (hashtable.ResizeReservation, error) { - return resizeAdmission(budget, owner, plan) - }) - } - } else { - strHashMap, err = hashmap.NewStrHashMap(false, mp) - if err == nil { - strHashMap.SetResizeAdmission(func(plan hashtable.ResizePlan) (hashtable.ResizeReservation, error) { - return resizeAdmission(budget, owner, plan) - }) - } - } - if err != nil { - owner.release() - return nil, err - } - - jm := message.NewJoinMap(message.GroupSels{}, intHashMap, strHashMap, nil, nil, mp) - jm.SetMemoryRelease(owner.release) - jm.IncRef(1) - return jm, nil +) { + hb.setBudget(budget) } -// NewAccountedEmptyJoinMap creates the consumer-grown empty-map variant under -// the statement allocation generation. Physical cell/descriptor Free is its -// only memory release owner; the account's controller already charges the -// shared query/CN policy, so no legacy reservation is stacked on it. +// NewAccountedEmptyJoinMap creates a consumer-grown map under the statement +// allocation generation. The map and its string-key iterator scratch carry +// the same immutable provenance as producer-built maps. func NewAccountedEmptyJoinMap( keyWidth int, account *mpool.AllocationAccount, @@ -200,9 +66,26 @@ func NewAccountedEmptyJoinMap( strHashMap *hashmap.StrHashMap ) if keyWidth <= 8 { - intHashMap, err = hashmap.NewIntHashMapWithAllocation(false, mp, selection) + intHashMap, err = hashmap.NewIntHashMapWithAllocation( + false, + mp, + selection, + ) } else { - strHashMap, err = hashmap.NewStrHashMapWithAllocation(false, mp, selection) + iteratorAllocation, allocationErr := hashmap.NewIteratorAllocation( + account, + HashBuildAllocationOwner, + HashBuildAllocationSiteHashIterator, + ) + if allocationErr != nil { + return nil, allocationErr + } + strHashMap, err = hashmap.NewStrHashMapWithAllocations( + false, + mp, + selection, + iteratorAllocation, + ) } if err != nil { return nil, err @@ -220,491 +103,35 @@ func NewAccountedEmptyJoinMap( return jm, nil } -func (hb *HashmapBuilder) attachIntHashMapAdmission(m *hashmap.IntHashMap) error { - owner := hb.mapReservation - budget := hb.budget - m.SetResizeAdmission(func(plan hashtable.ResizePlan) (hashtable.ResizeReservation, error) { - return resizeAdmission(budget, owner, plan) - }) - return nil -} - -func (hb *HashmapBuilder) attachStrHashMapAdmission(m *hashmap.StrHashMap) error { - owner := hb.mapReservation - budget := hb.budget - m.SetResizeAdmission(func(plan hashtable.ResizePlan) (hashtable.ResizeReservation, error) { - return resizeAdmission(budget, owner, plan) - }) - return nil -} - -func batchesAllocated(batches []*batch.Batch) uint64 { - var total uint64 - for _, bat := range batches { - if bat != nil { - total += uint64(bat.Allocated()) - } - } - return total -} - -type batchCopyAllocationSnapshot struct { - length int - tail *batch.Batch - tailAllocated uint64 -} - -func snapshotBatchCopyAllocation(batches []*batch.Batch) (batchCopyAllocationSnapshot, error) { - snapshot := batchCopyAllocationSnapshot{length: len(batches)} - if snapshot.length == 0 { - return snapshot, nil - } - snapshot.tail = batches[snapshot.length-1] - if snapshot.tail == nil { - return batchCopyAllocationSnapshot{}, process.ErrHashBuildBudgetInvalid - } - allocated := snapshot.tail.Allocated() - if allocated < 0 { - return batchCopyAllocationSnapshot{}, process.ErrHashBuildBudgetInvalid - } - snapshot.tailAllocated = uint64(allocated) - return snapshot, nil -} - -// batchCopyAllocatedDelta relies on CopyIntoBatches' append-only contract: it -// may grow the old partial tail and append destination batches. A full-size -// source can swap one new batch with that partial tail, so inspect the old tail -// plus the appended suffix by identity instead of rescanning every retained -// batch. Across a build this keeps retained-copy accounting linear in the -// number of destination batches rather than quadratic. -func batchCopyAllocatedDelta( - batches []*batch.Batch, - snapshot batchCopyAllocationSnapshot, -) (uint64, error) { - if snapshot.length < 0 || len(batches) < snapshot.length { - return 0, process.ErrHashBuildBudgetInvalid - } - start := 0 - seenTail := snapshot.length == 0 - if snapshot.length > 0 { - if snapshot.tail == nil { - return 0, process.ErrHashBuildBudgetInvalid - } - start = snapshot.length - 1 - } - var delta uint64 - for i := start; i < len(batches); i++ { - bat := batches[i] - if bat == nil { - return 0, process.ErrHashBuildBudgetInvalid - } - allocated := bat.Allocated() - if allocated < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - value := uint64(allocated) - if bat == snapshot.tail { - if seenTail || value < snapshot.tailAllocated { - return 0, process.ErrHashBuildBudgetInvalid - } - seenTail = true - value -= snapshot.tailAllocated - } - if delta > math.MaxUint64-value { - return 0, process.ErrHashBuildBudgetInvalid - } - delta += value - } - if !seenTail { - return 0, process.ErrHashBuildBudgetInvalid - } - return delta, nil -} - -func (hb *HashmapBuilder) copyBuildBatch(src *batch.Batch, proc *process.Process) error { - if hb.batchAllocation != nil { - return hb.Batches.CopyIntoBatchesWithAllocation( - src, - proc, - hb.batchAllocation, - ) - } - if hb.budget == nil { - return hb.Batches.CopyIntoBatches(src, proc) - } - projected, err := hb.projectedBatchCopyBytes(src) - if err != nil { - return err - } - reservation, err := hb.budget.Reserve(projected) - if err != nil { - return err - } - snapshot, err := snapshotBatchCopyAllocation(hb.Batches.Buf) - if err != nil { - reservation.Release() - return err - } - if err = hb.Batches.CopyIntoBatches(src, proc); err != nil { - reservation.Release() - hb.releaseBatchReservations() - return err - } - actual, err := batchCopyAllocatedDelta(hb.Batches.Buf, snapshot) - if err != nil { - hb.Batches.Clean(proc.Mp()) - reservation.Release() - hb.releaseBatchReservations() - return err - } - metadata, ok := retainedMetadataAllowance(src) - if !ok || actual > math.MaxUint64-metadata { - hb.Batches.Clean(proc.Mp()) - reservation.Release() - hb.releaseBatchReservations() - return process.ErrHashBuildBudgetInvalid - } - actual += metadata - if actual > projected { - // This indicates an incomplete pre-allocation bound. Fail closed after - // cleaning; never legitimize the excess with post-allocation admission. - hb.Batches.Clean(proc.Mp()) - reservation.Release() - hb.releaseBatchReservations() - return process.ErrHashBuildBudgetInvalid - } - if _, err = reservation.ReconcileDown(actual); err != nil { - hb.Batches.Clean(proc.Mp()) - reservation.Release() - hb.releaseBatchReservations() - return err +func (hb *HashmapBuilder) copyBuildBatch( + src *batch.Batch, + proc *process.Process, +) error { + if hb.batchAllocation == nil { + return mpool.ErrAllocationAccountInvalid } - hb.batchReservations = append(hb.batchReservations, reservation) - return nil + return hb.Batches.CopyIntoBatchesWithAllocation( + src, + proc, + hb.batchAllocation, + ) } -// CopyBuildBatch is an exported compatibility wrapper. -func (hb *HashmapBuilder) CopyBuildBatch(src *batch.Batch, proc *process.Process) error { +// CopyBuildBatch is the exported boundary used by spill and integration tests. +func (hb *HashmapBuilder) CopyBuildBatch( + src *batch.Batch, + proc *process.Process, +) error { return hb.copyBuildBatch(src, proc) } -func retainedMetadataAllowance(src *batch.Batch) (uint64, bool) { - if src == nil { - return 0, false - } - rows := uint64(src.RowCount()) - columns := uint64(len(src.Vecs)) - if columns > (math.MaxUint64-16)/8 { - return 0, false - } - perRow := uint64(16) + columns*8 - if rows > 0 && perRow > math.MaxUint64/rows { - return 0, false - } - return rows * perRow, true -} - -// projectedPartialTailReplacementBytes follows UnionBatch's allocation order. -// The existing tail reservation covers old capacities. At each grow, admission -// needs the complete replacement capacity plus deltas retained by earlier grows. -func projectedPartialTailReplacementBytes( - tail, src *batch.Batch, - appendRows int, -) (peak, retained uint64, err error) { - if tail == nil || src == nil || appendRows < 0 || len(tail.Vecs) != len(src.Vecs) { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - for i, srcVec := range src.Vecs { - dstVec := tail.Vecs[i] - if dstVec == nil || srcVec == nil || dstVec.Length() > math.MaxInt-appendRows { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - - typeSize := srcVec.GetType().TypeSize() - requiredRows := dstVec.Length() + appendRows - if typeSize < 0 || (typeSize > 0 && requiredRows > math.MaxInt/typeSize) { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - oldDataCap := cap(dstVec.GetData()) - if requiredData := requiredRows * typeSize; requiredData > oldDataCap { - newCap, ok := mpool.GrowCapacity(int64(oldDataCap), int64(requiredData)) - if !ok || retained > math.MaxUint64-uint64(newCap) { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - if candidate := retained + uint64(newCap); candidate > peak { - peak = candidate - } - retained += uint64(newCap) - uint64(oldDataCap) - } - - areaBytes, areaErr := - unionBatchAreaBytes(srcVec, 0, appendRows) - if areaErr != nil || areaBytes > math.MaxInt-len(dstVec.GetArea()) { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - oldAreaCap := cap(dstVec.GetArea()) - requiredArea := len(dstVec.GetArea()) + areaBytes - if requiredArea > oldAreaCap { - newCap, ok := mpool.GrowCapacity(int64(oldAreaCap), int64(requiredArea)) - if !ok || retained > math.MaxUint64-uint64(newCap) { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - if candidate := retained + uint64(newCap); candidate > peak { - peak = candidate - } - retained += uint64(newCap) - uint64(oldAreaCap) - } - } - return peak, retained, nil -} - -// projectedNewDestinationBytes follows CopyIntoBatches for destinations that -// start empty. Each destination vector is pre-extended to its final row count, -// and each varlen area is then grown once by UnionBatch. -func projectedNewDestinationBytes(src *batch.Batch, start, rows int) (uint64, error) { - if src == nil || start < 0 || rows < 0 || start > src.RowCount() || rows > src.RowCount()-start { - return 0, process.ErrHashBuildBudgetInvalid - } - end := start + rows - var total uint64 - add := func(value uint64) error { - if total > math.MaxUint64-value { - return process.ErrHashBuildBudgetInvalid - } - total += value - return nil - } - for offset := start; offset < end; { - segmentRows := end - offset - if segmentRows > colexec.DefaultBatchSize { - segmentRows = colexec.DefaultBatchSize - } - for _, vec := range src.Vecs { - if vec == nil { - return 0, process.ErrHashBuildBudgetInvalid - } - typeSize := vec.GetType().TypeSize() - if typeSize < 0 || (typeSize > 0 && segmentRows > math.MaxInt/typeSize) { - return 0, process.ErrHashBuildBudgetInvalid - } - dataCap, ok := mpool.GrowCapacity(0, int64(segmentRows*typeSize)) - if !ok || dataCap < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - if err := add(uint64(dataCap)); err != nil { - return 0, err - } - if !vec.GetType().IsVarlen() { - continue - } - - areaBytes, areaErr := - unionBatchAreaBytes(vec, offset, segmentRows) - if areaErr != nil { - return 0, process.ErrHashBuildBudgetInvalid - } - areaCap, ok := mpool.GrowCapacity(0, int64(areaBytes)) - if !ok || areaCap < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - if err := add(uint64(areaCap)); err != nil { - return 0, err - } - } - offset += segmentRows - } - return total, nil -} - -func (hb *HashmapBuilder) projectedBatchCopyBytes(src *batch.Batch) (uint64, error) { - if src == nil || src.RowCount() < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - rows := uint64(src.RowCount()) - last := len(hb.Batches.Buf) - 1 - hasPartialTail := rows != uint64(colexec.DefaultBatchSize) && - last >= 0 && hb.Batches.Buf[last] != nil && - hb.Batches.Buf[last].RowCount() != colexec.DefaultBatchSize - appendRows := 0 - if hasPartialTail { - // CopyIntoBatches appends into the partial tail. Derive each replacement - // from the destination's old capacity and the actual old+append target. - // A flat 1.25x multiplier is not a bound: GrowCapacity can take repeated - // 1.25x steps before reaching the required size. - tail := hb.Batches.Buf[last] - if tail.RowCount() < 0 || tail.RowCount() >= colexec.DefaultBatchSize { - return 0, process.ErrHashBuildBudgetInvalid - } - appendRows = colexec.DefaultBatchSize - tail.RowCount() - if appendRows > src.RowCount() { - appendRows = src.RowCount() - } - replacementPeak, retainedDelta, err := projectedPartialTailReplacementBytes(tail, src, appendRows) - if err != nil { - return 0, err - } - projected := replacementPeak - if appendRows < src.RowCount() { - // After the tail grow finishes, its retained delta stays live while - // CopyIntoBatches materializes the remaining source rows. - remaining, err := projectedNewDestinationBytes( - src, appendRows, src.RowCount()-appendRows, - ) - if err != nil { - return 0, err - } - if retainedDelta > math.MaxUint64-remaining { - return 0, process.ErrHashBuildBudgetInvalid - } - if retained := retainedDelta + remaining; retained > projected { - projected = retained - } - } - return projectedBatchCopyWithMetadata(src, projected) - } - projected, err := projectedNewDestinationBytes(src, 0, src.RowCount()) - if err != nil { - return 0, err - } - return projectedBatchCopyWithMetadata(src, projected) -} - -func projectedBatchCopyWithMetadata(src *batch.Batch, projected uint64) (uint64, error) { - // Vector null bitmaps and batch/vector slice metadata live on the Go heap - // and are therefore not included in Batch.Allocated. Charge a deliberately - // conservative per-row allowance that also scales with the column count. - // The source remains caller-owned, any retained tail already has its own - // reservation, and CopyIntoBatches reconciles this reservation to the actual - // retained delta below. - metadata, ok := retainedMetadataAllowance(src) - if !ok { - return 0, process.ErrHashBuildBudgetInvalid - } - const batchAllocationSlack = uint64(64 << 10) - if projected > math.MaxUint64-metadata || - projected+metadata > math.MaxUint64-batchAllocationSlack { - return 0, process.ErrHashBuildBudgetInvalid - } - return projected + metadata + batchAllocationSlack, nil -} - func (hb *HashmapBuilder) cleanBatches(proc *process.Process) { hb.Batches.Clean(proc.Mp()) - hb.releaseBatchReservations() -} - -func (hb *HashmapBuilder) buildAuxBytes( - needUniqueVec bool, - needAllocateSels ...bool, -) (uint64, error) { - uniqueBytes, err := hb.uniqueJoinKeyBytes() - if err != nil { - return 0, err - } - return hb.buildAuxBytesWithUniqueProjection( - needUniqueVec, uniqueBytes, needAllocateSels...) -} - -func (hb *HashmapBuilder) uniqueJoinKeyBytes() (uint64, error) { - var total uint64 - for _, vec := range hb.UniqueJoinKeys { - if vec == nil { - continue - } - allocated := vec.Allocated() - if allocated < 0 || - total > math.MaxUint64-uint64(allocated) { - return 0, process.ErrHashBuildBudgetInvalid - } - total += uint64(allocated) - } - return total, nil -} - -func (hb *HashmapBuilder) buildAuxBytesWithUniqueProjection( - needUniqueVec bool, - uniqueBytes uint64, - needAllocateSels ...bool, -) (uint64, error) { - // Covers mandatory hashmap/sels scratch plus the selected runtime-filter - // key vectors' actual persistent capacities. Before their first append, a - // bounded source-relative estimate admits the optional owner; every grow is - // then preflighted against its exact mpool capacity. Retained - // build batches are already charged by batchReservations, expression results - // have their own reservations, and runtime-filter serialization is admitted - // separately. Charging multiple whole-batch copies here double-counts those - // owners and can reject a build before any auxiliary allocation occurs. - bytes := batchesAllocated(hb.Batches.Buf) - if needUniqueVec && hb.uniqueKeyAllocation == nil { - growthSlack := bytes / 4 - if bytes%4 != 0 { - growthSlack++ - } - if uniqueBytes > growthSlack { - growthSlack = uniqueBytes - } - if bytes > math.MaxUint64-growthSlack { - return 0, process.ErrHashBuildBudgetInvalid - } - bytes += growthSlack - } - rowCount := hb.InputBatchRowCount - if hb.hashMapRowCountSet { - rowCount = hb.hashMapRowCount - } - rows := uint64(rowCount) - perRowBytes := uint64(64) - if len(needAllocateSels) > 0 && needAllocateSels[0] && - hb.batchAllocation != nil { - // GroupSels' physical slices are charged by batchAllocation. - perRowBytes -= 16 - } - const iteratorScratch = uint64(640 << 10) - if rows > math.MaxUint64/perRowBytes || - bytes > math.MaxUint64-rows*perRowBytes || - bytes+rows*perRowBytes > math.MaxUint64-iteratorScratch { - return 0, process.ErrHashBuildBudgetInvalid - } - bytes += rows*perRowBytes + iteratorScratch - return bytes, nil } -func (hb *HashmapBuilder) reserveBuildAux( - needUniqueVec bool, - needAllocateSels ...bool, -) error { - if hb.budget == nil { - return nil - } - if hb.batchAllocation != nil { - // Exact mode charges every data-scaled auxiliary owner at its physical - // allocation boundary. A second estimator reservation would double count - // the same memory and reintroduce false admission failures. - return nil - } - if hb.auxReservation != nil { - // BuildHashmap can be retried on the same retained batches with a - // different optional-runtime-filter decision. Reconcile the existing - // owner instead of silently retaining the previous projection (or, - // worse, collecting optional keys under a mandatory-only charge). - return hb.resizeBuildAuxReservation(needUniqueVec) - } - bytes, err := hb.buildAuxBytes(needUniqueVec, needAllocateSels...) - if err != nil { - return err - } - token, err := hb.budget.Reserve(bytes) - if err != nil { - return err - } - hb.auxReservation = token - return nil -} - -// abandonOptionalRuntimeFilterKeys removes only the exact-filter owner from an -// in-progress mandatory map build. No map or input batch is replayed. The -// persistent auxiliary reservation is reconciled to the same projection used -// by a build which never requested UniqueJoinKeys. +// abandonOptionalRuntimeFilterKeys removes only the optional exact-filter +// owner from an in-progress mandatory map build. Physical vector frees are the +// single source of truth for releasing the account charge. func (hb *HashmapBuilder) abandonOptionalRuntimeFilterKeys( proc *process.Process, ) error { @@ -716,9 +143,8 @@ func (hb *HashmapBuilder) abandonOptionalRuntimeFilterKeys( } // fallbackOptionalRuntimeFilterCollection converts only a proven optional -// cause into in-place key abandonment. Fatal causes are returned unchanged, -// leaving the fallback bit untouched and builder ownership with terminal -// cleanup. +// allocation failure into in-place key abandonment. Fatal causes keep builder +// ownership with terminal cleanup. func (hb *HashmapBuilder) fallbackOptionalRuntimeFilterCollection( proc *process.Process, cause error, @@ -727,19 +153,18 @@ func (hb *HashmapBuilder) fallbackOptionalRuntimeFilterCollection( runtimefilter.OptionalFallbackNone { return cause } - if err := hb.abandonOptionalRuntimeFilterKeys(proc); err != nil { - return err - } - return nil + return hb.abandonOptionalRuntimeFilterKeys(proc) } -// releaseOptionalRuntimeFilterKeys drops terminal producer-only state without -// marking collection fallback. The JoinMap retains only the mandatory -// auxiliary projection, so its transferred budget owner must be reconciled -// before publication. +// releaseOptionalRuntimeFilterKeys drops terminal producer-only vectors. The +// backing MPool allocations release their allocation-account charges exactly +// once when each vector is freed. func (hb *HashmapBuilder) releaseOptionalRuntimeFilterKeys( proc *process.Process, ) error { + if proc == nil { + return process.ErrHashBuildBudgetInvalid + } for i := range hb.UniqueJoinKeys { if hb.UniqueJoinKeys[i] != nil { hb.UniqueJoinKeys[i].Free(proc.Mp()) @@ -747,298 +172,32 @@ func (hb *HashmapBuilder) releaseOptionalRuntimeFilterKeys( } hb.UniqueJoinKeys = nil hb.uniqueSels = nil - if hb.auxReservation == nil { - return nil - } - required, err := hb.buildAuxBytes(false) - if err != nil { - return err - } - if required > hb.auxReservation.Size() { - return process.ErrHashBuildBudgetInvalid - } - _, err = hb.auxReservation.ReconcileDown(required) - return err -} - -func (hb *HashmapBuilder) resizeBuildAuxReservation( - needUniqueVec bool, -) error { - if hb.budget == nil { - return nil - } - if hb.auxReservation == nil { - return process.ErrHashBuildBudgetInvalid - } - target, err := hb.buildAuxBytes(needUniqueVec) - if err != nil { - return err - } - current := hb.auxReservation.Size() - switch { - case current < target: - return hb.auxReservation.Grow(target - current) - case current > target: - _, err = hb.auxReservation.ReconcileDown(target) - return err - default: - return nil - } + return nil } -// prepareCanonicalRuntimeFilterCollection first resizes the mandatory -// auxiliary owner for a Dedup input after its in-place canonical rewrite. It -// then attempts the optional UniqueJoinKeys delta. Failure of only that delta -// disables the runtime filter without failing the canonical map build. +// prepareCanonicalRuntimeFilterCollection restarts optional key collection +// after a destructive Dedup rewrite. The rewritten mandatory map is already +// charged at its physical allocation sites; optional growth may fail open at +// the vector allocation boundary during the rebuild. func (hb *HashmapBuilder) prepareCanonicalRuntimeFilterCollection( requested bool, ) (bool, error) { - if err := hb.resizeBuildAuxReservation(false); err != nil { - return false, err - } if !requested { return false, nil } - if err := hb.resizeBuildAuxReservation(true); err != nil { - if runtimefilter.ClassifyOptionalFallback(err) != - runtimefilter.OptionalFallbackBudgetAdmission { - return false, err - } - hb.runtimeFilterCollectionFallback = true - return false, nil - } hb.runtimeFilterCollectionFallback = false return true, nil } -func uniqueAppendAreaBytes(src *vector.Vector, start, rows int, sels []int64) (int, error) { - if src == nil || !src.GetType().IsVarlen() { - return 0, nil - } - if start < 0 || rows < 0 || (sels == nil && (start > src.Length() || rows > src.Length()-start)) { - return 0, process.ErrHashBuildBudgetInvalid - } - values, _ := vector.MustVarlenaRawData(src) - areaBytes := 0 - for i := 0; i < rows; i++ { - idx := start + i - if sels != nil { - if i >= len(sels) || sels[i] < 0 || sels[i] >= int64(src.Length()) { - return 0, process.ErrHashBuildBudgetInvalid - } - idx = int(sels[i]) - } - if src.IsConst() { - idx = 0 - } - if idx < 0 || idx >= len(values) || - (!src.GetNulls().EmptyByFlag() && src.GetNulls().Contains(uint64(idx))) || - values[idx].IsSmall() { - continue - } - _, valueLen := values[idx].OffsetLen() - valueBytes := int(valueLen) - if areaBytes > math.MaxInt-valueBytes { - return 0, process.ErrHashBuildBudgetInvalid - } - areaBytes += valueBytes - } - return areaBytes, nil -} - -// unionBatchAreaBytes mirrors Vector.UnionBatch's flags=nil varlen paths. -// In particular, its whole-vector fast path copies the complete source area, -// including bytes no longer referenced after SetLength/Shrink. Budget -// admission must therefore use len(area), not only the live row references. -func unionBatchAreaBytes( - src *vector.Vector, - start, rows int, -) (int, error) { - if src == nil || !src.GetType().IsVarlen() { - return 0, nil - } - if start < 0 || rows < 0 || - start > src.Length() || rows > src.Length()-start { - return 0, process.ErrHashBuildBudgetInvalid - } - if rows == 0 { - return 0, nil - } - if src.IsConst() { - // UnionBatch materializes the constant payload once and broadcasts its - // varlena header to the appended logical rows. - return uniqueAppendAreaBytes(src, 0, 1, nil) - } - if start == 0 && rows == src.Length() { - return len(src.GetArea()), nil - } - return uniqueAppendAreaBytes(src, start, rows, nil) -} - -func (hb *HashmapBuilder) reserveUniqueAppendOverlap(dst *vector.Vector, rows, areaBytes int) (*process.HashBuildReservation, error) { - if hb.budget == nil || hb.uniqueKeyAllocation != nil { - return nil, nil - } - if dst == nil || rows < 0 || areaBytes < 0 { - return nil, process.ErrHashBuildBudgetInvalid - } - typeSize := dst.GetType().TypeSize() - if typeSize < 0 || dst.Length() > math.MaxInt-rows || - (typeSize > 0 && dst.Length()+rows > math.MaxInt/typeSize) { - return nil, process.ErrHashBuildBudgetInvalid - } - requiredData := (dst.Length() + rows) * typeSize - dataCapacity, ok := mpool.GrowCapacity( - int64(cap(dst.GetData())), int64(requiredData)) - if !ok || dataCapacity < 0 { - return nil, process.ErrHashBuildBudgetInvalid - } - var overlap uint64 - if requiredData > cap(dst.GetData()) { - overlap = uint64(cap(dst.GetData())) - } - if len(dst.GetArea()) > math.MaxInt-areaBytes { - return nil, process.ErrHashBuildBudgetInvalid - } - requiredArea := len(dst.GetArea()) + areaBytes - areaCapacity, ok := mpool.GrowCapacity( - int64(cap(dst.GetArea())), int64(requiredArea)) - if !ok || areaCapacity < 0 { - return nil, process.ErrHashBuildBudgetInvalid - } - if requiredArea > cap(dst.GetArea()) { - if overlap > math.MaxUint64-uint64(cap(dst.GetArea())) { - return nil, process.ErrHashBuildBudgetInvalid - } - overlap += uint64(cap(dst.GetArea())) - } - if requiredData <= cap(dst.GetData()) && - requiredArea <= cap(dst.GetArea()) { - // The persistent capacities were admitted by their preceding grow. - // Avoid rescanning every retained batch for each UnitLimit append which - // stays within those capacities; only allocator growth changes either - // the retained owner or the temporary replacement overlap. - return nil, nil - } - - currentUnique, err := hb.uniqueJoinKeyBytes() - if err != nil { - return nil, err - } - oldCapacity := uint64(cap(dst.GetData())) + - uint64(cap(dst.GetArea())) - newCapacity := uint64(dataCapacity) + uint64(areaCapacity) - if currentUnique < oldCapacity || - currentUnique-oldCapacity > math.MaxUint64-newCapacity { - return nil, process.ErrHashBuildBudgetInvalid - } - projectedUnique := currentUnique - oldCapacity + newCapacity - target, err := hb.buildAuxBytesWithUniqueProjection( - true, projectedUnique) - if err != nil { - return nil, err - } - if hb.auxReservation == nil { - return nil, process.ErrHashBuildBudgetInvalid - } - if current := hb.auxReservation.Size(); current < target { - if err = hb.auxReservation.Grow(target - current); err != nil { - return nil, err - } - } - if overlap == 0 { - return nil, nil - } - return hb.budget.Reserve(overlap) -} - func (hb *HashmapBuilder) marshalRuntimeFilterVector( vec *vector.Vector, mp *mpool.MPool, ) ([]byte, func(), error) { - if vec == nil || vec.GetNulls().Any() { - return nil, nil, process.ErrHashBuildBudgetInvalid - } - if hb.mapAllocationAccount == nil { - return runtimefilter.MarshalExactFilterVector(vec, hb.budget) - } - if mp == nil { - return nil, nil, mpool.ErrAllocationAccountInvalid - } - size, err := vec.MarshalBinarySize() - if err != nil { - return nil, nil, err - } - buf, err := mpool.NewAccountedBuffer( + return runtimefilter.MarshalExactFilterVector( + vec, mp, hb.mapAllocationAccount, HashBuildAllocationOwner, HashBuildAllocationSiteRuntimeFilterPayload, ) - if err != nil { - return nil, nil, err - } - if err = buf.EnsureCapacity(size); err != nil { - buf.Free() - if mpool.IsRetryableAllocationCapacity(err) { - err = runtimefilter.MarkOptionalAllocationError(err) - } - return nil, nil, err - } - if err = vec.MarshalBinaryTo(buf); err != nil { - buf.Free() - return nil, nil, err - } - if buf.Len() != size { - buf.Free() - return nil, nil, process.ErrHashBuildBudgetInvalid - } - return buf.Bytes(), buf.Free, nil -} - -func (hb *HashmapBuilder) releaseBatchReservations() { - for _, reservation := range hb.batchReservations { - reservation.Release() - } - hb.batchReservations = nil -} - -func (hb *HashmapBuilder) releaseReservations() { - hb.releaseMapReservation() - hb.releaseBatchReservations() - if hb.auxReservation != nil { - hb.auxReservation.Release() - hb.auxReservation = nil - } -} - -func (hb *HashmapBuilder) releaseMapReservation() { - if hb.mapReservation != nil { - hb.mapReservation.release() - hb.mapReservation = nil - } -} - -func (hb *HashmapBuilder) detachReservations() func() { - mapOwner := hb.mapReservation - hb.mapReservation = nil - reservations := make([]*process.HashBuildReservation, 0, 1+len(hb.batchReservations)) - for _, reservation := range hb.batchReservations { - if token := reservation.Transfer(); token != nil { - reservations = append(reservations, token) - } - } - hb.batchReservations = nil - if hb.auxReservation != nil { - if token := hb.auxReservation.Transfer(); token != nil { - reservations = append(reservations, token) - } - hb.auxReservation = nil - } - return func() { - mapOwner.release() - for _, reservation := range reservations { - reservation.Release() - } - } } diff --git a/pkg/sql/colexec/hashbuild/build.go b/pkg/sql/colexec/hashbuild/build.go index 9a44c76c9f380..e36a01df1e4cf 100644 --- a/pkg/sql/colexec/hashbuild/build.go +++ b/pkg/sql/colexec/hashbuild/build.go @@ -76,6 +76,9 @@ func (hashBuild *HashBuild) Prepare(proc *process.Process) (err error) { return TerminalBudgetError(proc.Ctx, err) } hashBuild.ctr.hashmapBuilder.setBudget(budget) + if hashBuild.ctr.hashmapBuilder.mapAllocationAccount == nil { + return mpool.ErrAllocationAccountInvalid + } if hashBuild.IsShuffle && hashBuild.RuntimeFilterSpec == nil { return moerr.NewInternalError(proc.Ctx, "shuffle hash build must have runtime filter") } @@ -125,71 +128,11 @@ func (hashBuild *HashBuild) Call(proc *process.Process) (vm.CallResult, error) { ctr.state = SendJoinMap case SendJoinMap: - ctr.terminalMu.Lock() - if hashBuild.JoinMapTag <= 0 { - ctr.terminalMu.Unlock() - err := moerr.NewInternalError(proc.Ctx, "wrong joinmap message tag!") - hashBuild.finalizeBuildFailure(proc, err) - return result, err - } - if atomic.LoadUint32(&ctr.terminalPublished) != 0 { - ctr.terminalMu.Unlock() - return result, moerr.NewQueryInterrupted(proc.Ctx) - } - - var jm *message.JoinMap - spillMode := len(ctr.spilledFds) > 0 - var spillPayloadErr error - - if ctr.hashmapBuilder.InputBatchRowCount > 0 { - if spillMode { - // In spill mode: send empty JoinMap with spill fds, no batches - jm = message.NewJoinMap(message.GroupSels{}, nil, nil, nil, nil, proc.Mp()) - } else { - // Normal mode: send hashmap and batches - jm = ctr.hashmapBuilder.GetJoinMap(proc.Mp()) - jm.SetPushedRuntimeFilterIn(ctr.runtimeFilterIn) - } - jm.SetRowCount(int64(ctr.hashmapBuilder.InputBatchRowCount)) - jm.SetHasNullKey(ctr.hashmapBuilder.HasNullKey) - jm.IncRef(hashBuild.JoinMapRefCnt) - if spillMode { - payload := message.SpillBuildPayload{LegacyFds: ctr.spilledFds} - if ctr.spillBundle != nil { - payload = message.SpillBuildPayload{ - Files: ctr.spillBundle.accountedFiles(), - BudgetRef: ctr.hashmapBuilder.budget, - } - } - spillPayloadErr = jm.SetSpillBuildPayload(payload) - if spillPayloadErr == nil { - ctr.spilledFds = nil // ownership transferred - ctr.spillBundle = nil - } - } - } - - if spillPayloadErr != nil { - jm.FreeMemory() - ctr.terminalMu.Unlock() - err := moerr.NewInternalError(proc.Ctx, spillPayloadErr.Error()) + if err := hashBuild.sendJoinMap(proc); err != nil { hashBuild.finalizeBuildFailure(proc, err) return result, err } - - if !hashBuild.publishJoinMap(proc, jm) { - // Reset/Free may have won the terminal gate concurrently during - // cancellation. Keep the producer side successful only if this - // publication won; consumers must never see two terminal values. - if jm != nil { - jm.FreeMemory() - } - ctr.terminalMu.Unlock() - return result, moerr.NewQueryInterrupted(proc.Ctx) - } - ctr.state = SendSucceed - ctr.terminalMu.Unlock() case SendSucceed: result.Batch = nil @@ -199,6 +142,79 @@ func (hashBuild *HashBuild) Call(proc *process.Process) (vm.CallResult, error) { } } +// sendJoinMap serializes terminal publication with Reset and Free. The defer +// is part of the lifecycle contract: allocation, spill-payload, and message +// hooks may panic, and cleanup must never deadlock trying to reacquire this +// mutex while recovering the active statement. +func (hashBuild *HashBuild) sendJoinMap(proc *process.Process) error { + ctr := &hashBuild.ctr + ctr.terminalMu.Lock() + defer ctr.terminalMu.Unlock() + + if hashBuild.JoinMapTag <= 0 { + return moerr.NewInternalError(proc.Ctx, "wrong joinmap message tag!") + } + if hashBuild.JoinMapRefCnt <= 0 { + return moerr.NewInternalErrorf( + proc.Ctx, + "invalid join map reference count: %d", + hashBuild.JoinMapRefCnt, + ) + } + if atomic.LoadUint32(&ctr.terminalPublished) != 0 { + return moerr.NewQueryInterrupted(proc.Ctx) + } + + var jm *message.JoinMap + joinMapOwned := false + defer func() { + if joinMapOwned && jm != nil { + jm.FreeMemory() + } + }() + spillMode := len(ctr.spilledFds) > 0 + if ctr.hashmapBuilder.InputBatchRowCount > 0 { + if spillMode { + jm = message.NewJoinMap( + message.GroupSels{}, nil, nil, nil, nil, proc.Mp(), + ) + } else { + jm = ctr.hashmapBuilder.GetJoinMap(proc.Mp()) + if jm == nil { + return process.ErrHashBuildBudgetInvalid + } + joinMapOwned = true + jm.SetPushedRuntimeFilterIn(ctr.runtimeFilterIn) + } + if spillMode { + joinMapOwned = true + } + jm.SetRowCount(int64(ctr.hashmapBuilder.InputBatchRowCount)) + jm.SetHasNullKey(ctr.hashmapBuilder.HasNullKey) + jm.IncRef(hashBuild.JoinMapRefCnt) + if spillMode { + if ctr.spillBundle == nil || ctr.hashmapBuilder.budget == nil { + return process.ErrHashBuildBudgetInvalid + } + payload := message.SpillBuildPayload{ + Files: ctr.spillBundle.accountedFiles(), + BudgetRef: ctr.hashmapBuilder.budget, + } + if err := jm.SetSpillBuildPayload(payload); err != nil { + return moerr.NewInternalError(proc.Ctx, err.Error()) + } + ctr.spilledFds = nil + ctr.spillBundle = nil + } + } + + if !hashBuild.publishJoinMap(proc, jm) { + return moerr.NewQueryInterrupted(proc.Ctx) + } + joinMapOwned = false + return nil +} + // finalizeBuildFailure publishes every producer-side dependency before Call // returns. Consumers may already be blocked in ReceiveJoinMap/RuntimeFilter; // deferring publication until Reset could deadlock a pipeline scheduler that @@ -239,7 +255,6 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz ctr.hashmapBuilder.FreeTemporaryVectors(proc) ctr.hashmapBuilder.FreeExecutors() ctr.dropSpillScratchBuffers() - ctr.releaseSpillScratchReservation() }() startSpill := func() error { @@ -260,12 +275,6 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz hashBuild.JoinMapRefCnt, ) } - if ctr.spillBatchAllocation != nil { - // Exact mode converts the one-unit forward-progress token into the - // expression and scatter allocations that follow. Keeping both live - // would stack the same capacity and reject the recovery path itself. - ctr.releaseSpillScratchReservation() - } execs, err := ctr.initSpillExprExecs(proc, hashBuild.Conditions) if err != nil { return err @@ -326,49 +335,6 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz // particular, a rejected retained-copy admission below must not add the // same upstream batch a second time when it is spilled directly. ctr.hashmapBuilder.InputBatchRowCount += result.Batch.RowCount() - if hashBuild.IsShuffle { - // First prove that the current upstream batch can always be spilled - // directly. This uses its actual materialization semantics and never - // projects a hypothetical retained batch. - var directProofErr error - if !spillMode || ctr.spillBatchAllocation == nil { - directProofErr = ctr.ensureDirectSpillScratchReservation( - result.Batch, - analyzer, - ) - } - if directProofErr != nil { - // Existing retained batches were admitted with a future-drain - // proof. Drain them under that lease, then retry the direct proof - // after their source reservations have been released. - if spillMode || - !IsRetryableMemoryCapacity(directProofErr) || - len(ctr.hashmapBuilder.Batches.Buf) == 0 { - return directProofErr - } - if err := startSpill(); err != nil { - return err - } - if ctr.spillBatchAllocation == nil { - if err := ctr.ensureDirectSpillScratchReservation(result.Batch, analyzer); err != nil { - return err - } - } - } - if !spillMode { - // A batch may become retained only after its future spill scratch - // is admitted. If that proof does not fit, do not copy it: switch - // to the already-proven direct-spill path. - if err := ctr.ensureRetainedSpillScratchReservation(result.Batch, analyzer); err != nil { - if !IsRetryableMemoryCapacity(err) { - return err - } - if err := startSpill(); err != nil { - return err - } - } - } - } // If in spill mode, spill this batch directly to open files. if spillMode { err := ctr.spillBatchWithPressure(proc, result.Batch, spillFiles, ctr.spillExprExecs, analyzer, false) @@ -760,6 +726,14 @@ func (hashBuild *HashBuild) handleRuntimeFilter( } ctr.hashmapBuilder.uniqueKeySlots = nil }() + // A spilled build has no resident unique-key vector. Treating that absence + // as an empty build would publish DROP and incorrectly discard every probe + // row, so spill always disables this optional optimization. + if len(ctr.spilledFds) > 0 { + runtimeFilter.Typ = message.RuntimeFilter_PASS + hashBuild.sendRuntimeFilter(runtimeFilter, spec, proc) + return nil + } // send the unique join keys (doc_id membership pushdown) when requested if spec.UseMembershipFilter { @@ -794,6 +768,13 @@ func (hashBuild *HashBuild) handleRuntimeFilter( return nil } rowCount := keyVec.Length() + if keyVec.GetGrouping().GetBitmap().CountRange( + 0, uint64(keyVec.Length()), + ) > 0 { + runtimeFilter.Typ = message.RuntimeFilter_PASS + hashBuild.sendRuntimeFilter(runtimeFilter, spec, proc) + return nil + } // Always send the unique join keys; the consumer (ivfflat / fulltext // search) decides whether to use them as an exact pk IN filter or to @@ -877,15 +858,7 @@ func (hashBuild *HashBuild) handleRuntimeFilter( if err := runtimefilter.CloseFloatSignedZero( keyVec, proc.Mp(), - func() (func(), error) { - overlap, err := ctr.hashmapBuilder.reserveUniqueAppendOverlap(keyVec, 1, 0) - if err != nil || overlap == nil { - return nil, err - } - return func() { - overlap.Release() - }, nil - }, + nil, ); err != nil { if hashBuild.fallbackOptionalRuntimeFilter(err, &runtimeFilter, spec, proc) { return nil @@ -999,12 +972,11 @@ func (hashBuild *HashBuild) handleSerializedRuntimeFilter( } // materializeSerializedRuntimeFilter evaluates one proven serial/serial_full -// contract under the same query-wide HashBuild budget as the map and unique -// component vectors. It reuses the production component encoders, but +// contract under the same physical allocation account as the map and unique +// component vectors. It reuses the production component encoders and // precomputes a tight output-area bound from the actual unique values. The -// generic expression estimator must not be used here: a serial result is typed -// VARCHAR(max), which would reserve 64 KiB per tiny integer tuple and turn a -// useful index filter into PASS. +// account observes the actual vector growth rather than an estimated duplicate +// reservation. func (hashBuild *HashBuild) materializeSerializedRuntimeFilter( proc *process.Process, spec *plan.RuntimeFilterSpec, @@ -1039,27 +1011,20 @@ func (hashBuild *HashBuild) materializeSerializedRuntimeFilter( if err != nil { return nil, nil, 0, false, err } - peak, err := serializedRuntimeFilterAllocationPeak( - rowCount, areaBound, maxRowBound) - if err != nil { - return nil, nil, 0, false, err - } - - var reservation *process.HashBuildReservation - if budget := hashBuild.ctr.hashmapBuilder.budget; budget != nil { - reservation, err = budget.Reserve(peak) - if err != nil { - return nil, nil, 0, false, err - } - defer reservation.Release() - } payloadType, ok := planExprType( runtimefilter.BuildKeyExpr(spec)) - if !ok || areaBound > uint64(math.MaxInt) { + if !ok || areaBound > uint64(math.MaxInt) || + hashBuild.ctr.hashmapBuilder.uniqueKeyAllocation == nil { return nil, nil, 0, false, nil } - payload := vector.NewOffHeapVecWithType(payloadType) + payload, err := vector.NewOffHeapVecWithTypeAndAllocation( + payloadType, + hashBuild.ctr.hashmapBuilder.uniqueKeyAllocation, + ) + if err != nil { + return nil, nil, 0, false, err + } defer payload.Free(proc.Mp()) if err = payload.PreExtendWithArea( rowCount, int(areaBound), proc.Mp(), @@ -1190,53 +1155,6 @@ func serializedRuntimeFilterBounds( return areaBytes, maxRowBytes, nil } -func serializedRuntimeFilterAllocationPeak( - rowCount int, - areaBytes uint64, - maxRowBytes uint64, -) (uint64, error) { - if rowCount < 0 || - uint64(rowCount) > math.MaxUint64/types.VarlenaSize || - areaBytes > math.MaxInt64 { - return 0, process.ErrHashBuildBudgetInvalid - } - packerRequest := maxRowBytes - if packerRequest == 0 { - packerRequest = 1 - } - packerCapacity, ok := types.PackerAllocationSize(packerRequest) - if !ok { - return 0, process.ErrHashBuildBudgetInvalid - } - dataBytes := uint64(rowCount) * types.VarlenaSize - if dataBytes > math.MaxInt64 { - return 0, process.ErrHashBuildBudgetInvalid - } - dataCapacity, ok := mpool.GrowCapacity(0, int64(dataBytes)) - if !ok || dataCapacity < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - areaCapacity, ok := mpool.GrowCapacity(0, int64(areaBytes)) - if !ok || areaCapacity < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - // The output vector is pre-extended, so it has no allocate-copy-free - // growth overlap. Account the packer's actual size class rather than its - // requested slice: rounding can approach another full request. - peak := uint64(dataCapacity) - for _, part := range []uint64{ - uint64(areaCapacity), - packerCapacity, - (uint64(rowCount) + 7) / 8, - } { - if peak > math.MaxUint64-part { - return 0, process.ErrHashBuildBudgetInvalid - } - peak += part - } - return peak, nil -} - // Runtime filters are optional probe-side optimizations. Fail open only for a // query/CN admission rejection or an allocation error marked at an exact // optional payload/vector boundary. Cancellation, contract violations, and diff --git a/pkg/sql/colexec/hashbuild/build_test.go b/pkg/sql/colexec/hashbuild/build_test.go index 748d811dc4149..7646953de74ec 100644 --- a/pkg/sql/colexec/hashbuild/build_test.go +++ b/pkg/sql/colexec/hashbuild/build_test.go @@ -18,7 +18,9 @@ import ( "bytes" "context" "errors" + "fmt" "math" + "os" "strings" "sync" "testing" @@ -280,23 +282,6 @@ func TestShuffleWithoutMapRejectsMissingRuntimeFilter(t *testing.T) { tc.arg.Free(tc.proc, true, nil) } -func TestHashBuildFreeWithoutResetReleasesOwnedMemory(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) - require.NoError(t, tc.arg.Prepare(tc.proc)) - budget, err := tc.proc.GetHashBuildBudget() - require.NoError(t, err) - input := newBatch(tc.types, tc.proc, 100) - require.NoError(t, tc.arg.ctr.hashmapBuilder.copyBuildBatch(input, tc.proc)) - tc.arg.ctr.hashmapBuilder.InputBatchRowCount = input.RowCount() - input.Clean(tc.proc.Mp()) - require.NoError(t, tc.arg.ctr.hashmapBuilder.BuildHashmap(false, false, false, tc.proc)) - require.Greater(t, budget.Used(), uint64(0)) - - buildErr := errors.New("injected build failure") - tc.arg.Free(tc.proc, true, buildErr) - require.Zero(t, budget.Used()) -} - func BenchmarkBuild(b *testing.B) { for i := 0; i < b.N; i++ { tcs := []buildTestCase{ @@ -349,25 +334,27 @@ func newTestCase(t testing.TB, flgs []bool, ts []types.Type, cs []*plan.Expr) bu proc.Reg.MergeReceivers[0] = &process.WaitRegister{ Ch2: make(chan process.PipelineSignal, 10), } + arg := &HashBuild{ + JoinMapTag: 1, + JoinMapRefCnt: 1, + Conditions: cs, + NeedHashMap: true, + OperatorBase: vm.OperatorBase{ + OperatorInfo: vm.OperatorInfo{ + Idx: 0, + IsFirst: false, + IsLast: false, + }, + }, + } + installTestHashBuildAllocation(t, arg) return buildTestCase{ types: ts, flgs: flgs, proc: proc, cancel: cancel, - arg: &HashBuild{ - JoinMapTag: 1, - JoinMapRefCnt: 1, - Conditions: cs, - NeedHashMap: true, - OperatorBase: vm.OperatorBase{ - OperatorInfo: vm.OperatorInfo{ - Idx: 0, - IsFirst: false, - IsLast: false, - }, - }, - }, - marg: &merge.Merge{}, + arg: arg, + marg: &merge.Merge{}, } } @@ -378,6 +365,7 @@ func TestHashBuildPrepareDropsPriorGenerationSpillFileService(t *testing.T) { require.NoError(t, err) arg := &HashBuild{NeedHashMap: false} + installTestHashBuildAllocation(t, arg) arg.ctr.spillFS = prior require.NoError(t, arg.Prepare(proc)) require.Nil(t, arg.ctr.spillFS, "a reused operator must not retain the prior Process service") @@ -518,6 +506,7 @@ func TestHashBuildWithRuntimeFilter(t *testing.T) { }, }, } + installTestHashBuildAllocation(t, arg) err := arg.Prepare(proc) require.NoError(t, err) @@ -636,7 +625,7 @@ func TestHashBuildOptionalRuntimeFilterCollectionFallsBackToJoinMap( aggregate := process.MustNewHashBuildBudget(capBytes, capBytes) generation, err := aggregate.OpenGeneration(1) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) + installTestHashBuildBudget(t, tc.arg, generation) providerCalls := 0 forcedCollectionReject := false @@ -728,7 +717,7 @@ func TestHashBuildClosedMapBudgetDoesNotRecordCollectionFallback( aggregate := process.MustNewHashBuildBudget(capBytes, capBytes) generation, err := aggregate.OpenGeneration(1) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) + installTestHashBuildBudget(t, tc.arg, generation) providerCalls := 0 forcedClosed := false @@ -801,7 +790,7 @@ func TestHashmapBuilderUniqueGrowthFailureAbandonsOptionalKeysInPlace( aggregate := process.MustNewHashBuildBudget(capBytes, capBytes) generation, err := aggregate.OpenGeneration(1) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) + installTestHashBuildBudget(t, tc.arg, generation) const uniqueGrowthRows = hashmap.UnitLimit * 2 input := newBatch( @@ -1008,7 +997,7 @@ func TestShuffleDedupAdmissionAfterRewriteDoesNotSpillPartialInput( aggregate := process.MustNewHashBuildBudget(capBytes, capBytes) generation, err := aggregate.OpenGeneration(1) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) + installTestHashBuildBudget(t, tc.arg, generation) // Ingress happens before BuildHashmap initializes this phase. Mark the // retained source safe so the provider rejects only after Dedup crosses // its explicit in-place rewrite boundary. @@ -1227,9 +1216,7 @@ func TestHashBuildFloatRuntimeFilterAllocationFailureFallsBackToPass(t *testing. budget := process.MustNewHashBuildBudget(64<<20, 64<<20) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - arg.ctr.hashmapBuilder.setBudget(generation) - require.NoError(t, arg.ctr.hashmapBuilder.reserveBuildAux(true)) - usedWithUniqueKeys := generation.Used() + installTestHashBuildBudget(t, arg, generation) var filler []byte defer func() { @@ -1252,7 +1239,7 @@ func TestHashBuildFloatRuntimeFilterAllocationFailureFallsBackToPass(t *testing. require.False(t, arg.ctr.runtimeFilterIn) require.Nil(t, arg.ctr.hashmapBuilder.UniqueJoinKeys) require.Zero(t, generation.RejectCount()) - require.Less(t, generation.Used(), usedWithUniqueKeys) + require.Zero(t, generation.Used()) extra := arg.OpAnalyzer.GetOpStats().ExtraStats require.Equal(t, int64(1), extra["HashBuildRuntimeFilterAllocationFallbacks"]) @@ -1536,7 +1523,7 @@ func TestRuntimeFilterExplicitDecimalContractProducesIn(t *testing.T) { budget := process.MustNewHashBuildBudget(1<<20, 1<<20) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) + installTestHashBuildBudget(t, tc.arg, generation) require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) require.True(t, tc.arg.ctr.runtimeFilterDone) @@ -1584,7 +1571,7 @@ func TestDirectRuntimeFilterUsesDeclaredHashSlot(t *testing.T) { budget := process.MustNewHashBuildBudget(1<<20, 1<<20) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) + installTestHashBuildBudget(t, tc.arg, generation) require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) receiver := message.NewMessageReceiver( @@ -1684,7 +1671,7 @@ func TestHashBuildSerializedRuntimeFilterAllocationFailureFallsBackToPass(t *tes budget := process.MustNewHashBuildBudget(64<<20, 64<<20) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - arg.ctr.hashmapBuilder.setBudget(generation) + installTestHashBuildBudget(t, arg, generation) var filler []byte defer func() { @@ -1778,7 +1765,7 @@ func TestSerializedRuntimeFilterUsesTightBudgetAndProducesIn(t *testing.T) { budget := process.MustNewHashBuildBudget(512<<10, 512<<10) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) + installTestHashBuildBudget(t, tc.arg, generation) require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) require.True(t, tc.arg.ctr.runtimeFilterDone) @@ -1827,76 +1814,6 @@ func TestSerializedRuntimeFilterUsesTightBudgetAndProducesIn(t *testing.T) { require.Zero(t, tc.proc.Mp().CurrNB()) } -func TestSerializedRuntimeFilterBudgetAccountsPackerSizeClass(t *testing.T) { - componentType := types.T_varchar.ToType() - tc := newTestCase( - t, - []bool{false}, - []types.Type{componentType}, - []*plan.Expr{newExpr(0, componentType)}, - ) - spec := makeSerializedRuntimeFilterSpec( - t, tc.proc, 107, 2, []types.Type{componentType}, false) - tc.arg.RuntimeFilterSpec = spec - tc.arg.ctr.hashmapBuilder.InputBatchRowCount = 1 - value := strings.Repeat("x", 128<<10) - keys := []*vector.Vector{ - testutil.MakeVarcharVector([]string{value}, nil, tc.proc.Mp()), - } - tc.arg.ctr.hashmapBuilder.UniqueJoinKeys = keys - - service := tc.proc.GetService() - rt := moruntime.ServiceRuntime(service) - original, hadOriginal := rt.GetGlobalVariables(moruntime.MOProtocolVersion) - rt.SetGlobalVariables( - moruntime.MOProtocolVersion, defines.MORPCVersion8) - t.Cleanup(func() { - if hadOriginal { - rt.SetGlobalVariables(moruntime.MOProtocolVersion, original) - } else { - rt.SetGlobalVariables( - moruntime.MOProtocolVersion, defines.MORPCLatestVersion) - } - }) - - areaBytes, maxRowBytes, err := - serializedRuntimeFilterBounds( - tc.proc, keys, []int{0}, 1, false) - require.NoError(t, err) - require.Greater(t, maxRowBytes, uint64(128<<10)) - packerBytes, ok := types.PackerAllocationSize(maxRowBytes) - require.True(t, ok) - require.Equal(t, uint64(256<<10), packerBytes) - peak, err := serializedRuntimeFilterAllocationPeak( - 1, areaBytes, maxRowBytes) - require.NoError(t, err) - - // One byte below the true peak must fail open before constructing the - // packer. Counting only its requested slice would incorrectly admit it. - budget := process.MustNewHashBuildBudget(peak-1, peak-1) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - - require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) - receiver := message.NewMessageReceiver( - []int32{spec.Tag}, - message.AddrBroadCastOnCurrentCN(), - tc.proc.GetMessageBoard(), - ) - msgs, _, err := receiver.ReceiveMessage(false, tc.proc.Ctx) - require.NoError(t, err) - require.Len(t, msgs, 1) - require.Equal(t, int32(message.RuntimeFilter_PASS), - msgs[0].(message.RuntimeFilterMessage).Typ) - require.Zero(t, generation.Used()) - require.Zero(t, generation.Peak()) - - generation.Close() - tc.proc.Free() - require.Zero(t, tc.proc.Mp().CurrNB()) -} - func TestSerializedRuntimeFilterBoundsObserveCancellation(t *testing.T) { tc := newTestCase(t, nil, nil, nil) vec := testutil.MakeInt32Vector([]int32{1}, nil, tc.proc.Mp()) @@ -2072,7 +1989,7 @@ func TestRuntimeFilterMarshalBudgetAdmissionFallsBackToPass(t *testing.T) { budget := process.MustNewHashBuildBudget(1, 1) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) + installTestHashBuildBudget(t, tc.arg, generation) require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) require.True(t, tc.arg.ctr.runtimeFilterDone) @@ -2118,7 +2035,7 @@ func TestRuntimeFilterMarshalUsesSinglePayloadBudget(t *testing.T) { budget := process.MustNewHashBuildBudget(projected, projected) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) + installTestHashBuildBudget(t, tc.arg, generation) data, release, err := tc.arg.ctr.hashmapBuilder.marshalRuntimeFilterVector(vec, tc.proc.Mp()) require.NoError(t, err) @@ -2150,7 +2067,7 @@ func TestRuntimeFilterMarshalSinglePayloadCoversVarlenaPeak(t *testing.T) { budget := process.MustNewHashBuildBudget(projected, projected) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) + installTestHashBuildBudget(t, tc.arg, generation) data, release, err := tc.arg.ctr.hashmapBuilder.marshalRuntimeFilterVector(vec, tc.proc.Mp()) require.NoError(t, err) @@ -2184,7 +2101,7 @@ func TestRuntimeFilterMarshalAccountedPayloadMessageLifecycle(t *testing.T) { account, err := registry.OpenWithController(limit, generation) require.NoError(t, err) tc.arg.NeedHashMap = true - require.NoError(t, tc.arg.SetAllocationAccount(account)) + replaceTestHashBuildAllocation(t, tc.arg, account) tc.arg.ctr.hashmapBuilder.setBudget(generation) data, release, err := tc.arg.ctr.hashmapBuilder.marshalRuntimeFilterVector( @@ -2196,7 +2113,6 @@ func TestRuntimeFilterMarshalAccountedPayloadMessageLifecycle(t *testing.T) { require.NotNil(t, release) snapshot := account.Snapshot() require.Positive(t, snapshot.Used) - require.Equal(t, snapshot.Used, generation.Snapshot().AllocationUsed) require.Equal(t, snapshot.Used, generation.Used()) spec := &plan.RuntimeFilterSpec{Tag: 103} @@ -2241,7 +2157,7 @@ func TestRuntimeFilterMarshalAccountedOneByteShortFallsBackToPass(t *testing.T) account, err := registry.OpenWithController(limit, generation) require.NoError(t, err) tc.arg.NeedHashMap = true - require.NoError(t, tc.arg.SetAllocationAccount(account)) + replaceTestHashBuildAllocation(t, tc.arg, account) tc.arg.ctr.hashmapBuilder.setBudget(generation) tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{ Tag: 104, @@ -2283,37 +2199,69 @@ func TestRuntimeFilterMarshalAccountedOneByteShortFallsBackToPass(t *testing.T) require.Zero(t, tc.proc.Mp().CurrNB()) } -func TestRuntimeFilterMarshalClosedBudgetRemainsFatal(t *testing.T) { +func TestRuntimeFilterWithGroupingKeyFallsBackToPass(t *testing.T) { tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) - spec := rawRuntimeFilterSpec(102, 100, types.T_int32.ToType()) - tc.arg.RuntimeFilterSpec = spec - tc.arg.OpAnalyzer = process.NewAnalyzer(0, false, false, "hash build") + tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{ + Tag: 105, + UpperLimit: 100, + UseMembershipFilter: true, + } tc.arg.ctr.hashmapBuilder.InputBatchRowCount = 1 tc.arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{ - testutil.MakeInt32Vector([]int32{1}, nil, tc.proc.Mp()), + vector.NewRollupConst(types.T_int32.ToType(), 1, tc.proc.Mp()), } - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) - generation, err := budget.OpenGeneration(1) + require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) + receiver := message.NewMessageReceiver( + []int32{tc.arg.RuntimeFilterSpec.Tag}, + message.AddrBroadCastOnCurrentCN(), + tc.proc.GetMessageBoard(), + ) + msgs, done, err := receiver.ReceiveMessage(false, tc.proc.Ctx) require.NoError(t, err) - generation.Close() - tc.arg.ctr.hashmapBuilder.setBudget(generation) - - err = tc.arg.handleRuntimeFilter(tc.proc) - require.ErrorIs(t, err, process.ErrHashBuildBudgetClosed) + require.False(t, done) + require.Len(t, msgs, 1) + runtimeFilter, ok := msgs[0].(message.RuntimeFilterMessage) + require.True(t, ok) + require.Equal(t, int32(message.RuntimeFilter_PASS), runtimeFilter.Typ) + require.Empty(t, runtimeFilter.Data) require.Nil(t, tc.arg.ctr.hashmapBuilder.UniqueJoinKeys) - require.False(t, tc.arg.ctr.runtimeFilterDone) - require.Zero(t, tc.arg.OpAnalyzer.GetOpStats().ExtraStats["HashBuildRuntimeFilterBudgetFallbacks"]) - require.Zero(t, tc.arg.OpAnalyzer.GetOpStats().ExtraStats["HashBuildRuntimeFilterAllocationFallbacks"]) + tc.arg.Free(tc.proc, false, nil) + tc.proc.Free() + require.Zero(t, tc.proc.Mp().CurrNB()) +} + +func TestSpilledBuildRuntimeFilterPassesInsteadOfDropping(t *testing.T) { + tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, + []*plan.Expr{newExpr(0, types.T_int32.ToType())}) + tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{ + Tag: 106, + UpperLimit: 100, + Expr: newExpr(0, types.T_int32.ToType()), + } + tc.arg.ctr.hashmapBuilder.InputBatchRowCount = 1 + file, err := os.CreateTemp(t.TempDir(), "hashbuild-spilled-runtime-filter") + require.NoError(t, err) + tc.arg.ctr.spilledFds = []*os.File{file} + + require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) receiver := message.NewMessageReceiver( - []int32{spec.Tag}, message.AddrBroadCastOnCurrentCN(), tc.proc.GetMessageBoard()) - msgs, done, receiveErr := receiver.ReceiveMessage(false, tc.proc.Ctx) - require.NoError(t, receiveErr) + []int32{tc.arg.RuntimeFilterSpec.Tag}, + message.AddrBroadCastOnCurrentCN(), + tc.proc.GetMessageBoard(), + ) + msgs, done, err := receiver.ReceiveMessage(false, tc.proc.Ctx) + require.NoError(t, err) require.False(t, done) - require.Empty(t, msgs) + require.Len(t, msgs, 1) + runtimeFilter, ok := msgs[0].(message.RuntimeFilterMessage) + require.True(t, ok) + require.Equal(t, int32(message.RuntimeFilter_PASS), runtimeFilter.Typ) + require.Empty(t, runtimeFilter.Data) + tc.arg.Free(tc.proc, false, nil) tc.proc.Free() require.Zero(t, tc.proc.Mp().CurrNB()) } @@ -2429,6 +2377,7 @@ func TestHashBuildRuntimeFilterWithNulls(t *testing.T) { }, }, } + installTestHashBuildAllocation(t, arg) err := arg.Prepare(proc) require.NoError(t, err) @@ -2481,6 +2430,7 @@ func TestHashBuildRuntimeFilterWithNullsHashOnPK(t *testing.T) { }, }, } + installTestHashBuildAllocation(t, arg) err := arg.Prepare(proc) require.NoError(t, err) @@ -2505,55 +2455,6 @@ func TestHashBuildRuntimeFilterWithNullsHashOnPK(t *testing.T) { proc.Free() } -func TestHashBuildIsShuffle(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) - budget, budgetErr := tc.proc.GetHashBuildBudget() - require.NoError(t, budgetErr) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 1 - tc.arg.TrackNullKeys = true - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: 2} - tc.arg.SetChildren([]vm.Operator{tc.marg}) - for cycle := 0; cycle < 2; cycle++ { - if cycle > 0 { - tc.marg.Reset(tc.proc, false, nil) - tc.proc.GetMessageBoard().Reset() - } - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - build := batch.NewWithSize(1) - var buildNulls []uint64 - if cycle == 0 { - buildNulls = []uint64{1} - } - build.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 0, 2}, buildNulls, tc.proc.Mp()) - build.SetRowCount(3) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(build, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(batch.EmptyBatch, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - _, err := vm.Exec(tc.arg, tc.proc) - require.NoError(t, err) - result, receiveErr := message.ReceiveJoinMapResult(tc.arg.JoinMapTag, true, tc.arg.ShuffleIdx, tc.proc.GetMessageBoard(), tc.proc.Ctx) - require.NoError(t, receiveErr) - require.True(t, result.IsSuccess(), "cycle %d must publish a spilled JoinMap", cycle) - jm := result.JoinMap() - require.NotNil(t, jm) - require.True(t, jm.IsSpilled()) - spillPayload, err := jm.TakeSpillBuildPayload() - require.NoError(t, err) - require.Len(t, spillPayload.Files, spillNumBuckets) - require.Same(t, budget, spillPayload.BudgetRef) - require.NoError(t, spillPayload.Close()) - require.Zero(t, budget.Used()) - require.Zero(t, budget.SpillDiskUsed()) - require.Zero(t, budget.SpillFDUsed()) - tc.arg.Reset(tc.proc, false, nil) - } - tc.arg.Free(tc.proc, false, nil) - tc.proc.Free() -} - func TestBroadcastHashBuildParallelConsumersStayResident(t *testing.T) { tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) tc.arg.IsShuffle = false @@ -2635,275 +2536,50 @@ func TestHashBuildRejectsSharedSpillPayload(t *testing.T) { tc.proc.Free() } -func TestShuffleHashBuildDoesNotPreflightFutureSpill(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_varchar.ToType()}, []*plan.Expr{newExpr(0, types.T_varchar.ToType())}) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 1 << 30 - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: tc.arg.JoinMapTag + 3500} - tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - const capBytes = uint64(8 << 20) - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - - payload := make([]byte, 1<<20) - for i := range payload { - payload[i] = 'x' - } - build := batch.NewWithSize(1) - build.Vecs[0], err = vector.NewConstBytes(types.T_varchar.ToType(), payload, 1, tc.proc.Mp()) - require.NoError(t, err) - build.SetRowCount(1) - - directNeed, err := spillBudgetBytes(build) - require.NoError(t, err) - require.Less(t, directNeed, capBytes) - - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(build, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - _, err = vm.Exec(tc.arg, tc.proc) - require.NoError(t, err) - - result, err := message.ReceiveJoinMapResult(tc.arg.JoinMapTag, true, tc.arg.ShuffleIdx, tc.proc.GetMessageBoard(), tc.proc.Ctx) - require.NoError(t, err) - require.True(t, result.IsSuccess()) - jm := result.JoinMap() - require.NotNil(t, jm) - require.False(t, jm.IsSpilled(), - "a resident build must not pay for or be redirected by hypothetical future spill scratch") - require.Equal(t, int64(1), jm.GetRowCount()) - require.Zero(t, tc.arg.OpAnalyzer.GetOpStats().ExtraStats["HashBuildSpillStarts"]) - jm.Free() - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - - tc.arg.Reset(tc.proc, false, nil) - tc.arg.Free(tc.proc, false, nil) - tc.marg.Reset(tc.proc, false, nil) - generation.Close() - tc.proc.Free() - require.Zero(t, tc.proc.Mp().CurrNB()) -} - -func TestShuffleHashBuildSpillsBeforeRetainingThresholdCrossingBatch(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_varchar.ToType()}, []*plan.Expr{newExpr(0, types.T_varchar.ToType())}) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: tc.arg.JoinMapTag + 3502} - tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - makeBuildBatch := func() *batch.Batch { - values := make([]string, colexec.DefaultBatchSize) - for i := range values { - values[i] = strings.Repeat("x", 256) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeVarcharVector(values, nil, tc.proc.Mp()) - bat.SetRowCount(len(values)) - return bat - } - first := makeBuildBatch() - second := makeBuildBatch() - inputSize := int64(first.Size()) - tc.arg.SpillThreshold = inputSize + 1 - tc.arg.ctr.setSpillThreshold(tc.arg.SpillThreshold) - - copyPeak, err := tc.arg.ctr.hashmapBuilder.projectedBatchCopyBytes(first) - require.NoError(t, err) - retainedScratch, err := spillScratchBudgetBytes(first, true) - require.NoError(t, err) - directScratch, err := spillBudgetBytes(second) - require.NoError(t, err) - - // Calibrate the retained reservation and the second copy's pre-allocation - // peak. The final cap deliberately admits both copies (the old path reaches - // its post-copy threshold) but rejects scratch while both sources are live; - // the pre-copy path needs only one retained source plus scratch, then the - // direct scratch after that source is released. - calibrationBudget := process.MustNewHashBuildBudget(1<<30, 1<<30) - calibration, err := calibrationBudget.OpenGeneration(1) - require.NoError(t, err) - var calibrationBuilder HashmapBuilder - calibrationBuilder.setBudget(calibration) - require.NoError(t, calibrationBuilder.copyBuildBatch(first, tc.proc)) - actualFirst := calibration.Used() - secondCopyPeak, err := calibrationBuilder.projectedBatchCopyBytes(second) - require.NoError(t, err) - calibrationBuilder.cleanBatches(tc.proc) - require.Zero(t, calibration.Used()) - calibration.Close() - - coalesceSlack := uint64(spillNumBuckets * spillWriteCoalesceSize) - capBytes := max( - copyPeak, - actualFirst+secondCopyPeak, - actualFirst+retainedScratch, - directScratch, - ) + coalesceSlack - oldPostCopyPeak := 2*actualFirst + retainedScratch - require.Less(t, capBytes, oldPostCopyPeak, - "fixture must reject lazy scratch only after retaining the crossing batch") - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(first, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(second, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - _, buildErr := vm.Exec(tc.arg, tc.proc) - require.NoError(t, buildErr) - - result, err := message.ReceiveJoinMapResult( - tc.arg.JoinMapTag, true, tc.arg.ShuffleIdx, - tc.proc.GetMessageBoard(), tc.proc.Ctx) - require.NoError(t, err) - require.True(t, result.IsSuccess()) - jm := result.JoinMap() - require.NotNil(t, jm) - require.True(t, jm.IsSpilled()) - require.Equal(t, int64(2*colexec.DefaultBatchSize), jm.GetRowCount()) - payload, err := jm.TakeSpillBuildPayload() - require.NoError(t, err) - require.NoError(t, payload.Close()) +func TestHashBuildRejectsNonPositiveJoinMapRefCountBeforeTransfer(t *testing.T) { + for _, refCount := range []int32{0, -1} { + t.Run(fmt.Sprintf("ref-%d", refCount), func(t *testing.T) { + tc := newTestCase( + t, + []bool{false}, + []types.Type{types.T_int32.ToType()}, + []*plan.Expr{newExpr(0, types.T_int32.ToType())}, + ) + tc.arg.JoinMapRefCnt = refCount + tc.arg.SpillThreshold = math.MaxInt64 + tc.arg.SetChildren([]vm.Operator{tc.marg}) + require.NoError(t, tc.marg.Prepare(tc.proc)) + require.NoError(t, tc.arg.Prepare(tc.proc)) + account := tc.arg.ctr.hashmapBuilder.mapAllocationAccount - extra := tc.arg.OpAnalyzer.GetOpStats().ExtraStats - require.Equal(t, int64(1), extra["HashBuildSpillStarts"]) - require.Zero(t, extra["HashBuildSpillScratchReserveRejects"]) - require.Zero(t, extra["QueryHashBudgetRejects"], - "pre-copy thresholding must not consume recovery headroom first") - require.Empty(t, tc.arg.ctr.hashmapBuilder.Batches.Buf) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) + build := batch.NewWithSize(1) + build.Vecs[0] = testutil.MakeInt32Vector( + []int32{1, 2, 3}, nil, tc.proc.Mp(), + ) + build.SetRowCount(3) + tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(build, nil, tc.proc.Mp()) + tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - tc.arg.Reset(tc.proc, false, nil) - tc.marg.Reset(tc.proc, false, nil) - tc.arg.Free(tc.proc, false, nil) - first.Clean(tc.proc.Mp()) - second.Clean(tc.proc.Mp()) - generation.Close() - tc.proc.Free() - require.Zero(t, tc.proc.Mp().CurrNB()) -} + _, buildErr := vm.Exec(tc.arg, tc.proc) + require.ErrorContains(t, buildErr, "invalid join map reference count") + result, err := message.ReceiveJoinMapResult( + tc.arg.JoinMapTag, + false, + 0, + tc.proc.GetMessageBoard(), + tc.proc.Ctx, + ) + require.NoError(t, err) + require.True(t, result.IsBuildError()) + require.Nil(t, result.JoinMap()) -func TestShuffleHashBuildLazySpillAdmissionFailsClosed(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_varchar.ToType()}, []*plan.Expr{newExpr(0, types.T_varchar.ToType())}) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 1 - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{ - Tag: tc.arg.JoinMapTag + 3501, + tc.arg.Reset(tc.proc, true, buildErr) + tc.marg.Reset(tc.proc, true, buildErr) + require.Zero(t, account.Snapshot().Used) + tc.arg.Free(tc.proc, true, buildErr) + tc.proc.Free() + }) } - tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - payload := bytes.Repeat([]byte{'x'}, 1<<20) - build := batch.NewWithSize(1) - var err error - build.Vecs[0], err = vector.NewConstBytes( - types.T_varchar.ToType(), payload, 1, tc.proc.Mp()) - require.NoError(t, err) - build.SetRowCount(1) - - // Leave enough budget for the retained copy itself, but not for the - // additional scratch needed after the threshold requests spill. The lazy - // path must fail before scratch allocation instead of requiring every - // resident batch to pre-admit a hypothetical future spill. - copyPeak, err := tc.arg.ctr.hashmapBuilder.projectedBatchCopyBytes(build) - require.NoError(t, err) - budget := process.MustNewHashBuildBudget(copyPeak, copyPeak) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(build, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - _, buildErr := vm.Exec(tc.arg, tc.proc) - require.Error(t, buildErr) - require.True(t, moerr.IsMoErrCode(buildErr, moerr.ErrOOM), - "terminal spill admission must use the resource-exhausted wire code") - require.Contains(t, buildErr.Error(), "hash build memory budget exceeded") - extra := tc.arg.OpAnalyzer.GetOpStats().ExtraStats - require.Equal(t, int64(1), extra["HashBuildSpillStarts"]) - require.Equal(t, int64(1), extra["HashBuildSpillScratchReserveRejects"]) - require.Equal(t, int64(1), extra["QueryHashBudgetRejects"]) - - result, err := message.ReceiveJoinMapResult( - tc.arg.JoinMapTag, true, tc.arg.ShuffleIdx, - tc.proc.GetMessageBoard(), tc.proc.Ctx) - require.NoError(t, err) - require.True(t, result.IsBuildError()) - require.Equal(t, buildErr.Error(), result.BuildError().Error()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - - tc.arg.Reset(tc.proc, true, buildErr) - tc.marg.Reset(tc.proc, true, buildErr) - tc.arg.Free(tc.proc, true, buildErr) - build.Clean(tc.proc.Mp()) - require.Zero(t, generation.Used()) - generation.Close() - tc.proc.Free() - require.Zero(t, tc.proc.Mp().CurrNB()) -} - -func TestShuffleHashBuildSpillsExpressionKey(t *testing.T) { - bindProc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - col := newExpr(0, types.T_int32.ToType()) - modulo, err := plan2.BindFuncExprImplByPlanExpr( - bindProc.Ctx, - "%", - []*plan.Expr{col, plan2.MakePlan2Int32ConstExprWithType(2)}, - ) - require.NoError(t, err) - - tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{modulo}) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 1 - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: tc.arg.JoinMapTag + 4000} - tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - build := batch.NewWithSize(1) - build.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3, 4}, nil, tc.proc.Mp()) - build.SetRowCount(4) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(build, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - - _, err = vm.Exec(tc.arg, tc.proc) - require.NoError(t, err) - result, err := message.ReceiveJoinMapResult(tc.arg.JoinMapTag, true, tc.arg.ShuffleIdx, tc.proc.GetMessageBoard(), tc.proc.Ctx) - require.NoError(t, err) - require.True(t, result.IsSuccess()) - jm := result.JoinMap() - require.NotNil(t, jm) - require.True(t, jm.IsSpilled()) - require.Equal(t, int64(4), jm.GetRowCount()) - spillPayload, err := jm.TakeSpillBuildPayload() - require.NoError(t, err) - require.NoError(t, spillPayload.Close()) - budget, err := tc.proc.GetHashBuildBudget() - require.NoError(t, err) - require.Zero(t, budget.Used()) - require.Zero(t, budget.SpillDiskUsed()) - require.Zero(t, budget.SpillFDUsed()) - - tc.arg.Reset(tc.proc, false, nil) - tc.arg.Free(tc.proc, false, nil) - tc.proc.Free() - bindProc.Free() } func TestShuffleHashBuildAccountedSpillLifecycle(t *testing.T) { @@ -2928,7 +2604,7 @@ func TestShuffleHashBuildAccountedSpillLifecycle(t *testing.T) { require.NoError(t, err) account, err := registry.OpenWithController(limit, generation) require.NoError(t, err) - require.NoError(t, tc.arg.SetAllocationAccount(account)) + replaceTestHashBuildAllocation(t, tc.arg, account) require.NoError(t, tc.marg.Prepare(tc.proc)) require.NoError(t, tc.arg.Prepare(tc.proc)) tc.arg.ctr.hashmapBuilder.setBudget(generation) @@ -2977,138 +2653,20 @@ func TestShuffleHashBuildAccountedSpillLifecycle(t *testing.T) { require.Zero(t, tc.proc.Mp().CurrNB()) } -func TestShuffleHashBuildResizeRejectReleasesPartialMapAndSpills(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 1 << 30 - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: tc.arg.JoinMapTag + 3000} - tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - const capBytes = uint64(64 << 20) - aggregate := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := aggregate.OpenGeneration(1) - require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - - // With one full ingress batch, the first four admissions are emergency - // scratch, retained copy, build auxiliary memory, and the initial map. - // Reject exactly the fifth admission: the first resize after insertion. - providerCalls := 0 - forcedResizeReject := false - aggregate.SetAggregateCapProvider(func() (uint64, error) { - providerCalls++ - if providerCalls == 5 { - forcedResizeReject = true - return generation.Used(), nil - } - return capBytes, nil - }) - - bat := newBatch(tc.types, tc.proc, 8192) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(bat, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - _, err = vm.Exec(tc.arg, tc.proc) - require.NoError(t, err) - require.True(t, forcedResizeReject) - require.Zero(t, generation.Used(), "partial map, retained batches, and emergency scratch must be released before publication") - - result, err := message.ReceiveJoinMapResult(tc.arg.JoinMapTag, true, tc.arg.ShuffleIdx, tc.proc.GetMessageBoard(), tc.proc.Ctx) - require.NoError(t, err) - require.True(t, result.IsSuccess()) - jm := result.JoinMap() - require.NotNil(t, jm) - require.True(t, jm.IsSpilled()) - require.Equal(t, int64(8192), jm.GetRowCount()) - spillPayload, err := jm.TakeSpillBuildPayload() - require.NoError(t, err) - require.NoError(t, spillPayload.Close()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - - tc.arg.Reset(tc.proc, false, nil) - tc.arg.Free(tc.proc, false, nil) - tc.marg.Reset(tc.proc, false, nil) - tc.proc.Free() - require.Zero(t, tc.proc.Mp().CurrNB()) -} - -func TestShuffleHashBuildClosedMapBudgetDoesNotSpill(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 1 << 30 - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: tc.arg.JoinMapTag + 3001} - tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - const capBytes = uint64(64 << 20) - aggregate := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := aggregate.OpenGeneration(1) - require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - - // With one full ingress batch, the fourth admission is the initial hashmap. - // A lifecycle failure there is fatal and must never enter spill recovery. - providerCalls := 0 - forcedClosed := &process.HashBuildBudgetError{ - Kind: process.HashBuildBudgetErrorClosed, - } - aggregate.SetAggregateCapProvider(func() (uint64, error) { - providerCalls++ - if providerCalls == 4 { - return 0, forcedClosed - } - return capBytes, nil - }) - - bat := newBatch(tc.types, tc.proc, 8192) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(bat, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - _, buildErr := vm.Exec(tc.arg, tc.proc) - require.Same(t, forcedClosed, buildErr) - require.ErrorIs(t, buildErr, process.ErrHashBuildBudgetClosed) - require.NotErrorIs(t, buildErr, process.ErrHashBuildBudgetAdmission) - require.Equal(t, 4, providerCalls) - require.Zero(t, tc.arg.OpAnalyzer.GetOpStats().ExtraStats["HashBuildSpillStarts"]) - require.Empty(t, tc.arg.ctr.spilledFds) - require.Nil(t, tc.arg.ctr.spillBundle) - - result, err := message.ReceiveJoinMapResult( - tc.arg.JoinMapTag, - true, - tc.arg.ShuffleIdx, - tc.proc.GetMessageBoard(), - tc.proc.Ctx, - ) - require.NoError(t, err) - require.True(t, result.IsBuildError()) - require.False(t, result.IsSuccess()) - require.Nil(t, result.JoinMap()) - - tc.arg.Reset(tc.proc, true, buildErr) - tc.arg.Free(tc.proc, true, buildErr) - tc.marg.Reset(tc.proc, true, buildErr) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - generation.Close() - tc.proc.Free() - require.Zero(t, tc.proc.Mp().CurrNB()) -} - func TestObserveHashBuildBudgetUsesGenerationSnapshot(t *testing.T) { budget := process.MustNewHashBuildBudget(1024, 1024) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - reservation, err := generation.Reserve(128) + registry, err := mpool.NewAllocationAccountRegistry(1, 4) require.NoError(t, err) - _, err = generation.Reserve(1024) + account, err := registry.OpenWithController(2048, generation) + require.NoError(t, err) + mp := mpool.MustNewZero() + allocation, err := mp.AllocAccounted(128, account, HashBuildAllocationOwner, HashBuildAllocationSiteHashCell) + require.NoError(t, err) + _, err = mp.AllocAccounted(1024, account, HashBuildAllocationOwner, HashBuildAllocationSiteHashCell) require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - reservation.Release() + mp.Free(allocation) analyzer := process.NewAnalyzer(0, false, false, "hash build") observeHashBuildBudget(analyzer, generation) @@ -3122,115 +2680,7 @@ func TestObserveHashBuildBudgetUsesGenerationSnapshot(t *testing.T) { observeHashBuildBudget(analyzer, generation) require.Equal(t, int64(1), extra["QueryHashBudgetRejects"]) generation.Close() -} - -func TestShuffleHashBuildSpillFailureReleasesEmergencyResources(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 1 - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: tc.arg.JoinMapTag + 4000} - tc.arg.SetChildren([]vm.Operator{tc.marg}) - runtimeFilterReceiver := message.NewMessageReceiver( - []int32{tc.arg.RuntimeFilterSpec.Tag}, - message.AddrBroadCastOnCurrentCN(), - tc.proc.GetMessageBoard(), - ) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - aggregate := process.MustNewHashBuildBudget(64<<20, 64<<20) - generation, err := aggregate.OpenGenerationWithSpillCaps(1, 64<<20, 1, 32) - require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - - bat := newBatch(tc.types, tc.proc, 8192) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(bat, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - _, buildErr := vm.Exec(tc.arg, tc.proc) - require.True(t, moerr.IsMoErrCode(buildErr, moerr.ErrOOM)) - require.Contains(t, buildErr.Error(), "hash build spill disk budget exceeded") - require.Contains(t, buildErr.Error(), "requested=") - require.Contains(t, buildErr.Error(), "processLimitationSpillSize") - require.NotErrorIs(t, buildErr, process.ErrHashBuildBudgetAdmission) - require.Nil(t, tc.arg.ctr.spillScratchReservation) - require.Zero(t, cap(tc.arg.ctr.spillHashValues)) - require.Zero(t, cap(tc.arg.ctr.spillBucketRowIds)) - require.Zero(t, cap(tc.arg.ctr.spillKeyVecs)) - require.Zero(t, tc.arg.ctr.spillWriteBuf.Cap()) - - tc.arg.Reset(tc.proc, true, buildErr) - tc.arg.Reset(tc.proc, true, buildErr) - runtimeFilters, done, receiveErr := runtimeFilterReceiver.ReceiveMessage( - false, tc.proc.Ctx) - require.NoError(t, receiveErr) - require.False(t, done) - require.Len(t, runtimeFilters, 1, - "repeated Reset must publish one terminal value per generation") - tc.arg.Free(tc.proc, true, buildErr) - tc.arg.Free(tc.proc, true, buildErr) - tc.marg.Reset(tc.proc, true, buildErr) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - tc.proc.Free() - require.Zero(t, tc.proc.Mp().CurrNB()) -} - -func TestShuffleHashBuildDrainsRetainedBatchBeforeGrowingScratch(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 1 << 30 - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: tc.arg.JoinMapTag + 5000} - tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - const capBytes = uint64(64 << 20) - aggregate := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := aggregate.OpenGeneration(1) - require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - providerCalls := 0 - forcedGrowReject := false - aggregate.SetAggregateCapProvider(func() (uint64, error) { - providerCalls++ - // First ingress reserves direct scratch, grows it for future retained - // drain, then reserves its copy. Reject the fourth admission: growing - // direct scratch for the larger second ingress while that copy is live. - if providerCalls == 4 { - forcedGrowReject = true - return generation.Used(), nil - } - return capBytes, nil - }) - - first := newBatch(tc.types, tc.proc, 8192) - second := newBatch(tc.types, tc.proc, 65536) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(first, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(second, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - _, err = vm.Exec(tc.arg, tc.proc) + account.Seal() + _, err = registry.Finalize(account) require.NoError(t, err) - require.True(t, forcedGrowReject) - require.Zero(t, generation.Used()) - - result, err := message.ReceiveJoinMapResult(tc.arg.JoinMapTag, true, tc.arg.ShuffleIdx, tc.proc.GetMessageBoard(), tc.proc.Ctx) - require.NoError(t, err) - require.True(t, result.IsSuccess()) - jm := result.JoinMap() - require.True(t, jm.IsSpilled()) - require.Equal(t, int64(73728), jm.GetRowCount()) - spillPayload, err := jm.TakeSpillBuildPayload() - require.NoError(t, err) - require.NoError(t, spillPayload.Close()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - - tc.arg.Reset(tc.proc, false, nil) - tc.arg.Free(tc.proc, false, nil) - tc.marg.Reset(tc.proc, false, nil) - tc.proc.Free() - require.Zero(t, tc.proc.Mp().CurrNB()) } diff --git a/pkg/sql/colexec/hashbuild/dedup_memory.go b/pkg/sql/colexec/hashbuild/dedup_memory.go index 67634b5d88d5f..7d285d31793ee 100644 --- a/pkg/sql/colexec/hashbuild/dedup_memory.go +++ b/pkg/sql/colexec/hashbuild/dedup_memory.go @@ -21,8 +21,6 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/mpool" ) -// makeDedupSlice keeps the legacy allocation unchanged and moves only the -// activated, data-scaled owner to exact off-heap accounting. func makeDedupSlice[T any]( hb *HashmapBuilder, n int, @@ -33,7 +31,7 @@ func makeDedupSlice[T any]( return nil, mpool.ErrAllocationAccountInvalid } if hb.mapAllocationAccount == nil { - return make([]T, n), nil + return nil, mpool.ErrAllocationAccountInvalid } return mpool.MakeSliceAccounted[T]( n, @@ -45,7 +43,7 @@ func makeDedupSlice[T any]( } func freeDedupSlice[T any](hb *HashmapBuilder, values []T, mp *mpool.MPool) { - if hb.mapAllocationAccount != nil && cap(values) > 0 { + if cap(values) > 0 { mpool.FreeSlice(mp, values) } } @@ -60,8 +58,7 @@ func (hb *HashmapBuilder) newDedupBitmap( } bm := &bitmap.Bitmap{} if hb.mapAllocationAccount == nil { - bm.InitWithSize(int64(rows)) - return bm, nil + return nil, mpool.ErrAllocationAccountInvalid } words := (rows + 63) / 64 storage, err := mpool.MakeSliceAccounted[uint64]( diff --git a/pkg/sql/colexec/hashbuild/expression_memory.go b/pkg/sql/colexec/hashbuild/expression_memory.go index 283b08c85c478..6e82e6e0b9e68 100644 --- a/pkg/sql/colexec/hashbuild/expression_memory.go +++ b/pkg/sql/colexec/hashbuild/expression_memory.go @@ -15,605 +15,25 @@ package hashbuild import ( - "math" - - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/container/batch" - "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" - "github.com/matrixorigin/matrixone/pkg/sql/plan/function" "github.com/matrixorigin/matrixone/pkg/vm/process" ) -type expressionMemoryLeaseSlot struct { - tokens []*process.HashBuildReservation - admittedPeak uint64 - variableSize bool -} - -// ExpressionMemoryLease couples one stable expression-executor set to the -// HashBuild reservation that covers its retained vector capacity. It is not -// safe for concurrent use. -// -// Run retains the largest admitted expression bound independently for every -// root executor. Re-evaluating a root for an equal or smaller batch therefore -// reuses its existing high-water reservation. Growth admits only the part of -// that root's allocate-copy-free peak not covered by the old reservation, then -// reconciles the transient charge into the new high water. Unrelated roots are -// never double-charged. -// -// The owner must Free every executor and duplicate vector covered by the lease -// before calling Release. A lease cannot cross statement budget generations. -type ExpressionMemoryLease struct { - budget *process.HashBuildBudgetGeneration - exprs []*plan.Expr - executors []colexec.ExpressionExecutor - duplicate bool - slots []expressionMemoryLeaseSlot - released bool -} - -// NewAllocationAccountedExpressionExecutors constructs only expression trees -// whose complete retained and call-scoped allocation ledger is closed. The -// exact MPool leases are the sole capacity charge; unsupported function -// families continue through NewBudgetedExpressionExecutors until their own -// scratch owner is migrated. -func NewAllocationAccountedExpressionExecutors( +// NewExpressionExecutors constructs expression trees used by HashBuild and +// join operators. Expression temporaries are not retained HashBuild storage; +// only their explicit copies into retained destinations enter the account. +func NewExpressionExecutors( proc *process.Process, exprs []*plan.Expr, - allocation *colexec.ExpressionAllocationAccount, ) ([]colexec.ExpressionExecutor, error) { - if allocation == nil || !expressionSetAllocationClosed(exprs) { - return nil, process.ErrHashBuildBudgetInvalid - } - return colexec.NewExpressionExecutorsFromPlanExpressionsWithAllocation( - proc, - exprs, - allocation, - ) -} - -// NewAllocationAccountedExpressionExecutorsForAccount is the join-consumer -// entry point. It derives the same owner-scoped expression provenance used by -// HashBuild, so probe and rebuild expressions share one exact generation. -func NewAllocationAccountedExpressionExecutorsForAccount( - proc *process.Process, - exprs []*plan.Expr, - account *mpool.AllocationAccount, - owner mpool.AllocationOwner, -) ([]colexec.ExpressionExecutor, error) { - allocation, err := colexec.NewExpressionAllocationAccount(account, owner) - if err != nil { - return nil, err - } - return NewAllocationAccountedExpressionExecutors( - proc, - exprs, - allocation, - ) -} - -func expressionSetAllocationClosed(exprs []*plan.Expr) bool { if len(exprs) == 0 { - return false - } - for _, expr := range exprs { - if !expressionAllocationClosed(expr) { - return false - } - } - return true -} - -// AllocationAccountedExpressionSetSupported reports whether every execution -// path in the expression set has a closed physical-allocation ledger. Spill -// rebuild uses this gate before selecting the shared exact account. -func AllocationAccountedExpressionSetSupported(exprs []*plan.Expr) bool { - return expressionSetAllocationClosed(exprs) -} - -func expressionAllocationClosed(expr *plan.Expr) bool { - if expr == nil { - return false - } - switch node := expr.Expr.(type) { - case *plan.Expr_Col, *plan.Expr_Lit, *plan.Expr_T, - *plan.Expr_P, *plan.Expr_V, *plan.Expr_Vec, *plan.Expr_Fold: - return true - case *plan.Expr_F: - if node.F == nil || node.F.Func == nil { - return false - } - // Keep this as an implementation audit list, not a semantic function - // list. CONCAT writes directly into admitted result storage, CASE owns - // its row selections through ExpressionAllocationAccount, and varchar - // equality has no row-scaled scratch. Integer string casts use a - // stack-backed formatter; inserted casts of literals are plan-bounded. - functionID, _ := function.DecodeOverloadID(node.F.Func.Obj) - switch functionID { - case function.CONCAT, function.CASE: - case function.EQUAL: - if !closedHashBuildEqual(node.F.Args) { - return false - } - case function.CAST: - if !closedHashBuildCast(expr, node.F.Args) { - return false - } - default: - return false - } - for _, arg := range node.F.Args { - if !expressionAllocationClosed(arg) { - return false - } - } - return true - default: - return false - } -} - -func closedHashBuildEqual(args []*plan.Expr) bool { - if len(args) != 2 || args[0] == nil || args[1] == nil { - return false - } - for _, arg := range args { - oid := types.T(arg.Typ.Id) - if oid != types.T_char && oid != types.T_varchar { - return false - } - } - return true -} - -func closedHashBuildCast(result *plan.Expr, args []*plan.Expr) bool { - if result == nil || len(args) == 0 || args[0] == nil { - return false - } - source := types.T(args[0].Typ.Id) - target := types.T(result.Typ.Id) - if !source.ToType().IsIntOrUint() { - if (source == types.T_char || source == types.T_varchar) && - (target == types.T_char || target == types.T_varchar) { - _, literal := args[0].Expr.(*plan.Expr_Lit) - return literal - } - return false - } - return target == types.T_char || target == types.T_varchar -} - -// NewBudgetedExpressionExecutors admits the mpool-backed capacity owned by -// constant children before constructing them. The returned lease adopts those -// reservations, so construction and later evaluation have one continuous -// budget lifetime. -func NewBudgetedExpressionExecutors( - proc *process.Process, - budget *process.HashBuildBudgetGeneration, - exprs []*plan.Expr, - duplicate bool, -) ([]colexec.ExpressionExecutor, *ExpressionMemoryLease, error) { - if budget == nil { - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, exprs) - if err != nil { - return nil, nil, err - } - lease, err := NewExpressionMemoryLease(nil, exprs, executors, duplicate) - if err != nil { - for _, executor := range executors { - executor.Free() - } - return nil, nil, err - } - return executors, lease, err - } - - executors := make([]colexec.ExpressionExecutor, len(exprs)) - lease := &ExpressionMemoryLease{ - budget: budget, - exprs: exprs, - executors: executors, - duplicate: duplicate, - slots: make([]expressionMemoryLeaseSlot, len(exprs)), - } - cleanup := func() { - for _, executor := range executors { - if executor != nil { - executor.Free() - } - } - lease.Release() - } - - for i, expr := range exprs { - initial, err := expressionInitialOwnedBytes(expr) - if err != nil { - cleanup() - return nil, nil, err - } - if initial > 0 { - token, err := budget.Reserve(initial) - if err != nil { - cleanup() - return nil, nil, err - } - lease.slots[i].tokens = append(lease.slots[i].tokens, token) - lease.slots[i].admittedPeak = initial - } - - executor, err := colexec.NewExpressionExecutor(proc, expr) - if err != nil { - cleanup() - return nil, nil, err - } - executors[i] = executor - - retained, ok := colexec.ExpressionExecutorRetainedBytes(executor) - if !ok || retained > initial { - cleanup() - return nil, nil, process.ErrHashBuildBudgetInvalid - } - slot := &lease.slots[i] - slot.variableSize = expressionExecutorMayGrowWithinBound(expr) - if len(slot.tokens) > 0 { - token := slot.tokens[0] - if retained == 0 { - token.Release() - slot.tokens = nil - } else if _, err = token.ReconcileDown(retained); err != nil { - cleanup() - return nil, nil, err - } - } - slot.admittedPeak = retained - } - return executors, lease, nil -} - -func NewExpressionMemoryLease( - budget *process.HashBuildBudgetGeneration, - exprs []*plan.Expr, - executors []colexec.ExpressionExecutor, - duplicate bool, -) (*ExpressionMemoryLease, error) { - if len(exprs) != len(executors) { return nil, process.ErrHashBuildBudgetInvalid } - lease := &ExpressionMemoryLease{ - budget: budget, - exprs: exprs, - executors: executors, - duplicate: duplicate, - slots: make([]expressionMemoryLeaseSlot, len(executors)), - } - if budget == nil { - return lease, nil - } - - for i, executor := range executors { - retained, ok := colexec.ExpressionExecutorRetainedBytes(executor) - if !ok { - lease.Release() - return nil, process.ErrHashBuildBudgetInvalid - } - lease.slots[i].admittedPeak = retained - lease.slots[i].variableSize = expressionExecutorMayGrowWithinBound(exprs[i]) - if retained == 0 { - continue - } - token, err := budget.Reserve(retained) - if err != nil { - lease.Release() - return nil, err - } - lease.slots[i].tokens = append(lease.slots[i].tokens, token) - } - return lease, nil -} - -func expressionInitialOwnedBytes(expr *plan.Expr) (uint64, error) { - if expr == nil { - return 0, process.ErrHashBuildBudgetInvalid - } - switch typed := expr.Expr.(type) { - case *plan.Expr_Lit: - if typed.Lit == nil || typed.Lit.GetIsnull() { - return 0, nil - } - return literalInitialOwnedBytes(types.T(expr.Typ.Id), typed.Lit) - case *plan.Expr_List: - return expressionListInitialOwnedBytes(typed.List.GetList()) - case *plan.Expr_F: - return expressionListInitialOwnedBytes(typed.F.GetArgs()) - default: - return 0, nil - } -} - -func expressionListInitialOwnedBytes(exprs []*plan.Expr) (uint64, error) { - var total uint64 for _, expr := range exprs { - size, err := expressionInitialOwnedBytes(expr) - if err != nil || total > math.MaxUint64-size { - return 0, process.ErrHashBuildBudgetInvalid - } - total += size - } - return total, nil -} - -func expressionExecutorMayGrowWithinBound(expr *plan.Expr) bool { - if expr == nil { - return true - } - switch typed := expr.Expr.(type) { - case *plan.Expr_Col, *plan.Expr_Lit, *plan.Expr_T, *plan.Expr_Vec: - return false - case *plan.Expr_F: - if types.T(expr.Typ.Id).FixedLength() < 0 { - return true - } - for _, arg := range typed.F.GetArgs() { - if expressionExecutorMayGrowWithinBound(arg) { - return true - } - } - return false - case *plan.Expr_List: - if types.T(expr.Typ.Id).FixedLength() < 0 { - return true - } - for _, item := range typed.List.GetList() { - if expressionExecutorMayGrowWithinBound(item) { - return true - } - } - return false - case *plan.Expr_P, *plan.Expr_V: - return types.T(expr.Typ.Id).FixedLength() < 0 - default: - return true - } -} - -func literalInitialOwnedBytes(oid types.T, literal *plan.Literal) (uint64, error) { - var dataBytes uint64 - var payloadBytes uint64 - switch value := literal.GetValue().(type) { - case *plan.Literal_Bval, *plan.Literal_I8Val, *plan.Literal_U8Val, *plan.Literal_Defaultval: - dataBytes = 1 - case *plan.Literal_I16Val, *plan.Literal_U16Val, *plan.Literal_EnumVal: - dataBytes = 2 - case *plan.Literal_I32Val, *plan.Literal_U32Val, *plan.Literal_Fval, *plan.Literal_Dateval: - dataBytes = 4 - case *plan.Literal_I64Val, *plan.Literal_U64Val, *plan.Literal_Dval, - *plan.Literal_Timeval, *plan.Literal_Datetimeval, - *plan.Literal_Decimal64Val, *plan.Literal_Timestampval: - dataBytes = 8 - case *plan.Literal_Decimal128Val: - dataBytes = 16 - case *plan.Literal_Sval: - dataBytes = types.VarlenaSize - switch oid { - case types.T_array_float32: - if uint64(len(value.Sval)) > math.MaxUint64/4 { - return 0, process.ErrHashBuildBudgetInvalid - } - // The textual representation has at least one byte per element. - // Reserve a parsing-allocation-free upper bound; construction later - // reconciles it to the actual binary payload. - payloadBytes = uint64(len(value.Sval)) * 4 - case types.T_array_float64: - if uint64(len(value.Sval)) > math.MaxUint64/8 { - return 0, process.ErrHashBuildBudgetInvalid - } - payloadBytes = uint64(len(value.Sval)) * 8 - default: - payloadBytes = uint64(len(value.Sval)) - } - case *plan.Literal_VecVal: - dataBytes = types.VarlenaSize - payloadBytes = uint64(len(value.VecVal)) - default: - // Unsupported literal kinds are rejected by the expression factory - // before they can own an mpool-backed vector. - return 0, nil - } - - dataCapacity, err := initialAllocationCapacity(dataBytes) - if err != nil { - return 0, err - } - if payloadBytes <= types.VarlenaInlineSize { - return dataCapacity, nil - } - areaCapacity, err := initialAllocationCapacity(payloadBytes) - if err != nil || dataCapacity > math.MaxUint64-areaCapacity { - return 0, process.ErrHashBuildBudgetInvalid - } - return dataCapacity + areaCapacity, nil -} - -func initialAllocationCapacity(required uint64) (uint64, error) { - if required == 0 { - return 0, nil - } - if required > math.MaxInt64 { - return 0, process.ErrHashBuildBudgetInvalid - } - capacity, ok := mpool.GrowCapacity(0, int64(required)) - if !ok || capacity < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - return uint64(capacity), nil -} - -// Run admits and evaluates each root in index order. Growth keeps the root's -// old retained charge live, reserves only the uncovered allocate-copy-free -// overlap, and reconciles that overlap into the new high-water charge. -func (l *ExpressionMemoryLease) Run( - proc *process.Process, - rows int, - fn func(index int) error, -) (err error) { - if fn == nil { - return process.ErrHashBuildBudgetInvalid - } - if l == nil { - return process.ErrHashBuildBudgetInvalid - } - if l.released { - return process.ErrHashBuildReservationInactive - } - if rows < 0 { - return process.ErrHashBuildBudgetInvalid - } - - for i, expr := range l.exprs { - if l.budget == nil { - if err := fn(i); err != nil { - return err - } - continue - } - - peak, peakErr := expressionVectorPeak(proc, expr, rows, l.duplicate) - if peakErr != nil { - return peakErr - } - slot := &l.slots[i] - if peak <= slot.admittedPeak && !slot.variableSize { - if err := fn(i); err != nil { - return err - } - continue - } - - retained, ok := colexec.ExpressionExecutorRetainedBytes(l.executors[i]) - if !ok || retained > slot.admittedPeak { - return process.ErrHashBuildBudgetInvalid - } - var transient uint64 - if retained > math.MaxUint64-peak { - return process.ErrHashBuildBudgetInvalid - } - physicalPeak := retained + peak - if physicalPeak > slot.admittedPeak { - transient = physicalPeak - slot.admittedPeak - } - if transient == 0 { - if err := fn(i); err != nil { - return err - } - continue - } - - candidate, reserveErr := l.budget.Reserve(transient) - if reserveErr != nil { - return reserveErr - } - evalErr := fn(i) - if l.released { - candidate.Release() - return evalErr - } - if peak > slot.admittedPeak { - growth := peak - slot.admittedPeak - if _, reconcileErr := candidate.ReconcileDown(growth); reconcileErr != nil { - slot.tokens = append(slot.tokens, candidate) - slot.admittedPeak += transient - if evalErr != nil { - return evalErr - } - return reconcileErr - } - slot.tokens = append(slot.tokens, candidate) - slot.admittedPeak = peak - } else { - candidate.Release() - } - if evalErr != nil { - return evalErr - } - } - return nil -} - -// Eval evaluates the executors owned by the lease in the same per-root order -// used for admission. consume receives each successfully evaluated vector -// before the next root is admitted. -func (l *ExpressionMemoryLease) Eval( - proc *process.Process, - bats []*batch.Batch, - rows int, - consume func(index int, vec *vector.Vector) error, -) error { - if l == nil || consume == nil { - return process.ErrHashBuildBudgetInvalid - } - return l.Run(proc, rows, func(index int) error { - vec, err := l.executors[index].Eval(proc, bats, nil) - if err != nil { - return err - } - return consume(index, vec) - }) -} - -func (l *ExpressionMemoryLease) Reserved() uint64 { - if l == nil || l.released { - return 0 - } - var total uint64 - for i := range l.slots { - for _, token := range l.slots[i].tokens { - size := token.Size() - if total > math.MaxUint64-size { - return math.MaxUint64 - } - total += size - } - } - return total -} - -func (l *ExpressionMemoryLease) Len() int { - if l == nil || l.released { - return 0 - } - return len(l.executors) -} - -// Retained returns the current mpool-backed capacity physically owned by the -// executor set. Reserved may be larger: the documented delta is the retained -// high-water admission bound kept available for safe executor reuse. -func (l *ExpressionMemoryLease) Retained() (uint64, bool) { - if l == nil { - return 0, true - } - if l.released { - return 0, false - } - return colexec.ExpressionExecutorsRetainedBytes(l.executors) -} - -func (l *ExpressionMemoryLease) Release() { - if l == nil || l.released { - return - } - l.released = true - for i := range l.slots { - for _, token := range l.slots[i].tokens { - token.Release() + if expr == nil { + return nil, process.ErrHashBuildBudgetInvalid } - l.slots[i].tokens = nil - l.slots[i].admittedPeak = 0 } - l.slots = nil - l.exprs = nil - l.executors = nil - l.budget = nil + return colexec.NewExpressionExecutorsFromPlanExpressions(proc, exprs) } diff --git a/pkg/sql/colexec/hashbuild/expression_memory_test.go b/pkg/sql/colexec/hashbuild/expression_memory_test.go deleted file mode 100644 index 13e45f54a2ae8..0000000000000 --- a/pkg/sql/colexec/hashbuild/expression_memory_test.go +++ /dev/null @@ -1,1313 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package hashbuild - -import ( - "context" - "errors" - "strings" - "testing" - - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/container/batch" - "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/matrixorigin/matrixone/pkg/container/vector" - "github.com/matrixorigin/matrixone/pkg/pb/plan" - "github.com/matrixorigin/matrixone/pkg/sql/colexec" - plan2 "github.com/matrixorigin/matrixone/pkg/sql/plan" - "github.com/matrixorigin/matrixone/pkg/testutil" - "github.com/matrixorigin/matrixone/pkg/vm/process" - "github.com/stretchr/testify/require" -) - -func makeExpressionLeaseTestExpr(t *testing.T, proc *process.Process) *plan.Expr { - t.Helper() - col := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - expr, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "%", - []*plan.Expr{col, plan2.MakePlan2Int32ConstExprWithType(2)}, - ) - require.NoError(t, err) - return expr -} - -func makeExpressionLeaseTestBatch(proc *process.Process, rows int) *batch.Batch { - values := make([]int32, rows) - for i := range values { - values[i] = int32(i) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector(values, nil, proc.Mp()) - bat.SetRowCount(rows) - return bat -} - -func makeIssue26454ConcatKey(t testing.TB, proc *process.Process) *plan.Expr { - t.Helper() - cast := func(colPos int32) *plan.Expr { - col := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: colPos}}, - } - targetType := plan.Type{ - Id: int32(types.T_varchar), - Width: types.MaxVarcharLen, - } - expr, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "cast", - []*plan.Expr{ - col, - { - Typ: targetType, - Expr: &plan.Expr_T{T: &plan.TargetType{}}, - }, - }, - ) - require.NoError(t, err) - return expr - } - expr, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "concat", - []*plan.Expr{ - cast(0), - plan2.MakePlan2StringConstExprWithType("-"), - cast(1), - }, - ) - require.NoError(t, err) - return expr -} - -func makeIssue26454CaseKey(t testing.TB, proc *process.Process) *plan.Expr { - t.Helper() - column := &plan.Expr{ - Typ: plan.Type{ - Id: int32(types.T_varchar), - Width: types.MaxVarcharLen, - }, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - condition, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "=", - []*plan.Expr{ - column, - plan2.MakePlan2StringConstExprWithType("ATM_CON"), - }, - ) - require.NoError(t, err) - expr, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "case", - []*plan.Expr{ - condition, - plan2.MakePlan2StringConstExprWithType("CON_CONTRACT_HEADERS"), - plan2.MakePlan2StringConstExprWithType("CON_CONTRACT_DOC"), - }, - ) - require.NoError(t, err) - return expr -} - -func TestAllocationAccountedExpressionIssue26454AndOneByteShort(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeIssue26454ConcatKey(t, proc) - require.True(t, expressionSetAllocationClosed([]*plan.Expr{expr})) - require.False(t, expressionSetAllocationClosed( - []*plan.Expr{makeExpressionLeaseTestExpr(t, proc)}, - )) - input := batch.NewWithSize(2) - input.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp()) - input.Vecs[1] = testutil.MakeInt32Vector([]int32{3, 4}, nil, proc.Mp()) - input.SetRowCount(2) - defer input.Clean(proc.Mp()) - - run := func(limit uint64, verify bool) (uint64, error) { - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - registry, err := mpool.NewAllocationAccountRegistry(1, 64) - require.NoError(t, err) - account, err := registry.OpenWithController(limit, generation) - require.NoError(t, err) - allocation, err := colexec.NewExpressionAllocationAccount( - account, - HashBuildAllocationOwner, - ) - require.NoError(t, err) - executors, runErr := NewAllocationAccountedExpressionExecutors( - proc, - []*plan.Expr{expr}, - allocation, - ) - if runErr == nil { - var result *vector.Vector - result, runErr = executors[0].Eval( - proc, - []*batch.Batch{input}, - nil, - ) - if runErr == nil && verify { - require.Equal(t, []string{"1-3", "2-4"}, - vector.InefficientMustStrCol(result)) - } - } - peak := account.Snapshot().Peak - freeExpressionLeaseTestExecutors(executors) - require.Zero(t, account.Snapshot().Used) - require.Zero(t, generation.Used()) - _, _, terminalErr := registry.CompleteTerminal(account) - require.NoError(t, terminalErr) - return peak, runErr - } - - peak, err := run(1<<20, true) - require.NoError(t, err) - require.Positive(t, peak) - _, err = run(peak-1, false) - require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) - require.Contains(t, err.Error(), "allocation owner=1 site=") -} - -func TestAllocationAccountedExpressionIssue26454CaseKey(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeIssue26454CaseKey(t, proc) - require.True(t, expressionSetAllocationClosed([]*plan.Expr{expr})) - input := batch.NewWithSize(1) - input.Vecs[0] = testutil.MakeVarcharVector( - []string{"ATM_CON", "OTHER"}, - nil, - proc.Mp(), - ) - input.SetRowCount(2) - defer input.Clean(proc.Mp()) - - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - registry, err := mpool.NewAllocationAccountRegistry(1, 64) - require.NoError(t, err) - account, err := registry.OpenWithController(1<<20, generation) - require.NoError(t, err) - allocation, err := colexec.NewExpressionAllocationAccount( - account, - HashBuildAllocationOwner, - ) - require.NoError(t, err) - executors, err := NewAllocationAccountedExpressionExecutors( - proc, - []*plan.Expr{expr}, - allocation, - ) - require.NoError(t, err) - result, err := executors[0].Eval(proc, []*batch.Batch{input}, nil) - require.NoError(t, err) - require.Equal(t, - []string{"CON_CONTRACT_HEADERS", "CON_CONTRACT_DOC"}, - vector.InefficientMustStrCol(result), - ) - require.Positive(t, account.Snapshot().Used) - freeExpressionLeaseTestExecutors(executors) - require.Zero(t, account.Snapshot().Used) - require.Zero(t, generation.Used()) - _, _, err = registry.CompleteTerminal(account) - require.NoError(t, err) -} - -func TestAllocationAccountedExpressionRealValueOverCapRollsBack(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - column := &plan.Expr{ - Typ: plan.Type{ - Id: int32(types.T_varchar), - Width: types.MaxVarcharLen, - }, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - expr, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "concat", - []*plan.Expr{column, plan2.MakePlan2StringConstExprWithType("-suffix")}, - ) - require.NoError(t, err) - require.True(t, expressionSetAllocationClosed([]*plan.Expr{expr})) - - const capBytes = uint64(4 << 10) - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - registry, err := mpool.NewAllocationAccountRegistry(1, 64) - require.NoError(t, err) - account, err := registry.OpenWithController(capBytes, generation) - require.NoError(t, err) - allocation, err := colexec.NewExpressionAllocationAccount( - account, - HashBuildAllocationOwner, - ) - require.NoError(t, err) - executors, err := NewAllocationAccountedExpressionExecutors( - proc, - []*plan.Expr{expr}, - allocation, - ) - require.NoError(t, err) - - eval := func(value string) (*vector.Vector, error) { - input := batch.NewWithSize(1) - input.Vecs[0] = testutil.MakeVarcharVector([]string{value}, nil, proc.Mp()) - input.SetRowCount(1) - defer input.Clean(proc.Mp()) - return executors[0].Eval(proc, []*batch.Batch{input}, nil) - } - _, err = eval(strings.Repeat("x", 8<<10)) - require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) - require.Contains(t, err.Error(), "allocation owner=1 site=") - result, err := eval("ok") - require.NoError(t, err) - require.Equal(t, []string{"ok-suffix"}, vector.InefficientMustStrCol(result)) - - freeExpressionLeaseTestExecutors(executors) - require.Zero(t, account.Snapshot().Used) - require.Zero(t, generation.Used()) - _, _, err = registry.CompleteTerminal(account) - require.NoError(t, err) -} - -func TestHashmapBuilderFallsBackOnlyForUnclosedExpressionScratch(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - for _, tc := range []struct { - name string - expr *plan.Expr - accounted bool - }{ - {name: "closed concat cast", expr: makeIssue26454ConcatKey(t, proc), accounted: true}, - {name: "closed case equality", expr: makeIssue26454CaseKey(t, proc), accounted: true}, - {name: "unclosed modulo", expr: makeExpressionLeaseTestExpr(t, proc)}, - } { - t.Run(tc.name, func(t *testing.T) { - budget := process.MustNewHashBuildBudget(16<<20, 16<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - registry, err := mpool.NewAllocationAccountRegistry(1, 64) - require.NoError(t, err) - account, err := registry.OpenWithController(16<<20, generation) - require.NoError(t, err) - var op HashBuild - op.NeedHashMap = true - require.NoError(t, op.SetAllocationAccount(account)) - hb := &op.ctr.hashmapBuilder - hb.setBudget(generation) - require.NoError(t, hb.Prepare( - []*plan.Expr{tc.expr}, - -1, - -1, - nil, - proc, - )) - if tc.accounted { - require.Nil(t, hb.expressionLease) - require.Equal(t, generation.Used(), generation.Snapshot().AllocationUsed) - } else { - require.NotNil(t, hb.expressionLease) - } - hb.FreeExecutors() - require.Zero(t, account.Snapshot().Used) - require.Zero(t, generation.Used()) - require.NoError(t, op.ClearAllocationAccount(account)) - _, _, err = registry.CompleteTerminal(account) - require.NoError(t, err) - }) - } -} - -func TestHashBuildAllocationActivationRequiresClosedExpressionOwner(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - closed := &HashBuild{ - NeedHashMap: true, - Conditions: []*plan.Expr{makeIssue26454ConcatKey(t, proc)}, - } - require.True(t, closed.AllocationAccountEnabled()) - - unclosed := &HashBuild{ - NeedHashMap: true, - Conditions: []*plan.Expr{makeExpressionLeaseTestExpr(t, proc)}, - } - require.False(t, unclosed.AllocationAccountEnabled()) - - withoutMap := &HashBuild{ - Conditions: []*plan.Expr{makeIssue26454ConcatKey(t, proc)}, - } - require.False(t, withoutMap.AllocationAccountEnabled()) -} - -func BenchmarkIssue26454ExpressionAccounting(b *testing.B) { - const capBytes = uint64(8 << 30) - proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeIssue26454ConcatKey(b, proc) - input := testutil.NewBatch( - []types.Type{types.T_int32.ToType(), types.T_int32.ToType()}, - true, - colexec.DefaultBatchSize, - proc.Mp(), - ) - defer input.Clean(proc.Mp()) - - b.Run("legacy", func(b *testing.B) { - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - if err != nil { - b.Fatal(err) - } - executors, lease, err := NewBudgetedExpressionExecutors( - proc, - generation, - []*plan.Expr{expr}, - false, - ) - if err != nil { - b.Fatal(err) - } - b.ReportAllocs() - b.ResetTimer() - for range b.N { - if err = lease.Eval( - proc, - []*batch.Batch{input}, - input.RowCount(), - func(_ int, _ *vector.Vector) error { return nil }, - ); err != nil { - b.Fatal(err) - } - } - b.StopTimer() - freeExpressionLeaseTestExecutors(executors) - lease.Release() - if generation.Used() != 0 { - b.Fatalf("generation used = %d", generation.Used()) - } - }) - - b.Run("accounted", func(b *testing.B) { - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - if err != nil { - b.Fatal(err) - } - registry, err := mpool.NewAllocationAccountRegistry(1, 64) - if err != nil { - b.Fatal(err) - } - account, err := registry.OpenWithController(capBytes, generation) - if err != nil { - b.Fatal(err) - } - allocation, err := colexec.NewExpressionAllocationAccount( - account, - HashBuildAllocationOwner, - ) - if err != nil { - b.Fatal(err) - } - executors, err := NewAllocationAccountedExpressionExecutors( - proc, - []*plan.Expr{expr}, - allocation, - ) - if err != nil { - b.Fatal(err) - } - b.ReportAllocs() - b.ResetTimer() - for range b.N { - if _, err = executors[0].Eval( - proc, - []*batch.Batch{input}, - nil, - ); err != nil { - b.Fatal(err) - } - } - b.StopTimer() - freeExpressionLeaseTestExecutors(executors) - if account.Snapshot().Used != 0 || generation.Used() != 0 { - b.Fatalf( - "live account=%d generation=%d", - account.Snapshot().Used, - generation.Used(), - ) - } - if _, _, err = registry.CompleteTerminal(account); err != nil { - b.Fatal(err) - } - }) -} - -func BenchmarkIssue26454CaseExpressionAccounting(b *testing.B) { - const capBytes = uint64(64 << 20) - proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeIssue26454CaseKey(b, proc) - values := make([]string, colexec.DefaultBatchSize) - for i := range values { - if i%2 == 0 { - values[i] = "ATM_CON" - } else { - values[i] = "OTHER" - } - } - input := batch.NewWithSize(1) - input.Vecs[0] = testutil.MakeVarcharVector(values, nil, proc.Mp()) - input.SetRowCount(len(values)) - defer input.Clean(proc.Mp()) - - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - if err != nil { - b.Fatal(err) - } - registry, err := mpool.NewAllocationAccountRegistry(1, 64) - if err != nil { - b.Fatal(err) - } - account, err := registry.OpenWithController(capBytes, generation) - if err != nil { - b.Fatal(err) - } - allocation, err := colexec.NewExpressionAllocationAccount( - account, - HashBuildAllocationOwner, - ) - if err != nil { - b.Fatal(err) - } - executors, err := NewAllocationAccountedExpressionExecutors( - proc, - []*plan.Expr{expr}, - allocation, - ) - if err != nil { - b.Fatal(err) - } - b.ReportAllocs() - b.ResetTimer() - for range b.N { - if _, err = executors[0].Eval( - proc, - []*batch.Batch{input}, - nil, - ); err != nil { - b.Fatal(err) - } - } - b.StopTimer() - freeExpressionLeaseTestExecutors(executors) - if account.Snapshot().Used != 0 || generation.Used() != 0 { - b.Fatalf( - "live account=%d generation=%d", - account.Snapshot().Used, - generation.Used(), - ) - } - if _, _, err = registry.CompleteTerminal(account); err != nil { - b.Fatal(err) - } -} - -func makeMaxArrayLeaseTestVector[T types.ArrayElement]( - t *testing.T, - proc *process.Process, - oid types.T, -) *vector.Vector { - t.Helper() - typ := types.New(oid, types.MaxArrayDimension, 0) - vec := vector.NewVec(typ) - values := make([]T, types.MaxArrayDimension) - require.NoError(t, vector.AppendArrayList( - vec, - [][]T{values, values}, - nil, - proc.Mp(), - )) - return vec -} - -func evalExpressionLeaseTestExecutors( - proc *process.Process, - executors []colexec.ExpressionExecutor, - bat *batch.Batch, -) error { - for _, executor := range executors { - if _, err := executor.Eval(proc, []*batch.Batch{bat}, nil); err != nil { - return err - } - } - return nil -} - -func freeExpressionLeaseTestExecutors(executors []colexec.ExpressionExecutor) { - for _, executor := range executors { - executor.Free() - } -} - -func TestExpressionMemoryLeaseReusesRetainedHighWater(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeExpressionLeaseTestExpr(t, proc) - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - - initialRetained, ok := colexec.ExpressionExecutorsRetainedBytes(executors) - require.True(t, ok) - require.Positive(t, initialRetained) - largePeak, err := expressionVectorPeak(proc, expr, colexec.DefaultBatchSize, false) - require.NoError(t, err) - budgetCap := initialRetained + largePeak - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - lease, err := NewExpressionMemoryLease(generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - require.Equal(t, initialRetained, generation.Used()) - - large := makeExpressionLeaseTestBatch(proc, colexec.DefaultBatchSize) - defer large.Clean(proc.Mp()) - require.NoError(t, lease.Run(proc, large.RowCount(), func(_ int) error { - return evalExpressionLeaseTestExecutors(proc, executors, large) - })) - require.Equal(t, largePeak, generation.Used()) - reservesAfterLarge := generation.ReserveCount() - - retainedAfterLarge, ok := colexec.ExpressionExecutorsRetainedBytes(executors) - require.True(t, ok) - require.Greater(t, retainedAfterLarge, initialRetained) - leaseRetained, ok := lease.Retained() - require.True(t, ok) - require.Equal(t, retainedAfterLarge, leaseRetained) - require.GreaterOrEqual(t, lease.Reserved(), leaseRetained) - - small := makeExpressionLeaseTestBatch(proc, 1) - defer small.Clean(proc.Mp()) - require.NoError(t, lease.Run(proc, small.RowCount(), func(_ int) error { - return evalExpressionLeaseTestExecutors(proc, executors, small) - })) - require.Equal(t, reservesAfterLarge, generation.ReserveCount()) - require.Equal(t, largePeak, generation.Used()) - - for _, executor := range executors { - executor.ResetForNextQuery() - } - retainedAfterReset, ok := colexec.ExpressionExecutorsRetainedBytes(executors) - require.True(t, ok) - require.Equal(t, retainedAfterLarge, retainedAfterReset) - require.NoError(t, lease.Run(proc, large.RowCount(), func(_ int) error { - return evalExpressionLeaseTestExecutors(proc, executors, large) - })) - require.Equal(t, reservesAfterLarge, generation.ReserveCount()) - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) -} - -func TestExpressionTypePeakUsesArrayElementWidth(t *testing.T) { - for _, tc := range []struct { - oid types.T - elementWidth uint64 - }{ - {oid: types.T_array_float64, elementWidth: 8}, - {oid: types.T_array_float32, elementWidth: 4}, - {oid: types.T_array_bf16, elementWidth: 2}, - {oid: types.T_array_float16, elementWidth: 2}, - {oid: types.T_array_int8, elementWidth: 1}, - {oid: types.T_array_uint8, elementWidth: 1}, - } { - t.Run(tc.oid.String(), func(t *testing.T) { - peak, err := expressionTypePeak(plan.Type{ - Id: int32(tc.oid), - Width: types.MaxArrayDimension, - }, 1) - require.NoError(t, err) - require.Equal( - t, - uint64(types.MaxArrayDimension)*tc.elementWidth+32+(64<<10), - peak, - ) - }) - } -} - -func TestExpressionMemoryLeaseCoversMaxNarrowArrayPayload(t *testing.T) { - for _, tc := range []struct { - oid types.T - makeInput func(*testing.T, *process.Process) *vector.Vector - }{ - { - oid: types.T_array_bf16, - makeInput: func(t *testing.T, proc *process.Process) *vector.Vector { - return makeMaxArrayLeaseTestVector[types.BF16]( - t, proc, types.T_array_bf16) - }, - }, - { - oid: types.T_array_float16, - makeInput: func(t *testing.T, proc *process.Process) *vector.Vector { - return makeMaxArrayLeaseTestVector[types.Float16]( - t, proc, types.T_array_float16) - }, - }, - } { - t.Run(tc.oid.String(), func(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - arrayType := plan.Type{ - Id: int32(tc.oid), - Width: types.MaxArrayDimension, - } - condition := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_bool)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - left := &plan.Expr{ - Typ: arrayType, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 1}}, - } - right := &plan.Expr{ - Typ: arrayType, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 2}}, - } - expr, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "iff", - []*plan.Expr{condition, left, right}, - ) - require.NoError(t, err) - - peak, err := expressionVectorPeak(proc, expr, 2, false) - require.NoError(t, err) - budget := process.MustNewHashBuildBudget(2*peak, 2*peak) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - executors, lease, err := NewBudgetedExpressionExecutors( - proc, - generation, - []*plan.Expr{expr}, - false, - ) - require.NoError(t, err) - - input := batch.NewWithSize(3) - input.Vecs[0] = testutil.MakeBoolVector( - []bool{true, false}, nil, proc.Mp()) - input.Vecs[1] = tc.makeInput(t, proc) - input.Vecs[2] = tc.makeInput(t, proc) - input.SetRowCount(2) - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{input}, - input.RowCount(), - func(_ int, _ *vector.Vector) error { return nil }, - )) - retained, ok := lease.Retained() - require.True(t, ok) - require.LessOrEqual(t, retained, lease.Reserved(), - "retained max-width array payload must remain within admission") - - input.Clean(proc.Mp()) - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) - require.Zero(t, proc.Mp().CurrNB()) - }) - } -} - -func TestExpressionMemoryLeaseGrowthRequiresReplacementPeak(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeExpressionLeaseTestExpr(t, proc) - smallPeak, err := expressionVectorPeak(proc, expr, 1, false) - require.NoError(t, err) - largePeak, err := expressionVectorPeak(proc, expr, colexec.DefaultBatchSize, false) - require.NoError(t, err) - - t.Run("reject before evaluation", func(t *testing.T) { - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - retained, ok := colexec.ExpressionExecutorsRetainedBytes(executors) - require.True(t, ok) - require.LessOrEqual(t, retained, smallPeak) - budgetCap := retained + largePeak - 1 - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - lease, err := NewExpressionMemoryLease(generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - - require.NoError(t, lease.Run(proc, 1, func(_ int) error { return nil })) - require.Equal(t, smallPeak, generation.Used()) - evaluated := false - err = lease.Run(proc, colexec.DefaultBatchSize, func(_ int) error { - evaluated = true - return nil - }) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.False(t, evaluated) - require.Equal(t, smallPeak, generation.Used()) - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) - }) - - t.Run("commit exact replacement peak", func(t *testing.T) { - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - retained, ok := colexec.ExpressionExecutorsRetainedBytes(executors) - require.True(t, ok) - require.LessOrEqual(t, retained, smallPeak) - budgetCap := retained + largePeak - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(2) - require.NoError(t, err) - lease, err := NewExpressionMemoryLease(generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - - require.NoError(t, lease.Run(proc, 1, func(_ int) error { return nil })) - require.NoError(t, lease.Run(proc, colexec.DefaultBatchSize, func(_ int) error { return nil })) - require.Equal(t, largePeak, generation.Used()) - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) - }) -} - -func TestExpressionMemoryLeaseGrowsRootsIndependently(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeExpressionLeaseTestExpr(t, proc) - exprs := []*plan.Expr{expr, expr} - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, exprs) - require.NoError(t, err) - - smallPeak, err := expressionVectorPeak(proc, expr, 1, false) - require.NoError(t, err) - largePeak, err := expressionVectorPeak(proc, expr, colexec.DefaultBatchSize, false) - require.NoError(t, err) - require.Greater(t, largePeak, smallPeak) - for _, executor := range executors { - retained, ok := colexec.ExpressionExecutorRetainedBytes(executor) - require.True(t, ok) - require.LessOrEqual(t, retained, smallPeak) - } - - // Sequential root growth peaks at old(root 2) + new(root 1) + - // new(root 2). An aggregate replacement would incorrectly also charge - // old(root 1) and reject this exact-cap admission. - budgetCap := smallPeak + 2*largePeak - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - lease, err := NewExpressionMemoryLease(generation, exprs, executors, false) - require.NoError(t, err) - - small := makeExpressionLeaseTestBatch(proc, 1) - defer small.Clean(proc.Mp()) - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{small}, - small.RowCount(), - func(_ int, _ *vector.Vector) error { return nil }, - )) - require.Equal(t, 2*smallPeak, generation.Used()) - large := makeExpressionLeaseTestBatch(proc, colexec.DefaultBatchSize) - defer large.Clean(proc.Mp()) - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{large}, - large.RowCount(), - func(_ int, _ *vector.Vector) error { return nil }, - )) - require.Equal(t, 2*largePeak, generation.Used()) - require.LessOrEqual(t, generation.Peak(), budgetCap) - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) -} - -func TestExpressionMemoryLeaseCoversVariableWidthReuseOverlap(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - col := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_varchar), Width: types.MaxVarcharLen}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - expr, err := plan2.BindFuncExprImplByPlanExpr(proc.Ctx, "lower", []*plan.Expr{col}) - require.NoError(t, err) - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - peak, err := expressionVectorPeak(proc, expr, 1, false) - require.NoError(t, err) - - budget := process.MustNewHashBuildBudget(2*peak, 2*peak) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - lease, err := NewExpressionMemoryLease( - generation, - []*plan.Expr{expr}, - executors, - false, - ) - require.NoError(t, err) - - eval := func(value string) { - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeVarcharVector([]string{value}, nil, proc.Mp()) - bat.SetRowCount(1) - defer bat.Clean(proc.Mp()) - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{bat}, - 1, - func(_ int, _ *vector.Vector) error { return nil }, - )) - } - eval("a") - require.Equal(t, peak, generation.Used()) - retained, ok := lease.Retained() - require.True(t, ok) - require.Positive(t, retained) - reserves := generation.ReserveCount() - - eval(strings.Repeat("b", 64<<10)) - require.Equal(t, reserves+1, generation.ReserveCount(), - "same-row variable-width growth needs a transient overlap reservation") - require.Equal(t, peak, generation.Used(), - "transient overlap must not inflate the retained high-water charge") - require.GreaterOrEqual(t, generation.Peak(), peak+retained) - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) -} - -func TestExpressionMemoryLeaseCoversFlowControlSelectedScratch(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - condition := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_bool)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - makeLower := func(colPos int32) *plan.Expr { - col := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_varchar), Width: types.MaxVarcharLen}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: colPos}}, - } - expr, err := plan2.BindFuncExprImplByPlanExpr(proc.Ctx, "lower", []*plan.Expr{col}) - require.NoError(t, err) - return expr - } - expr, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "iff", - []*plan.Expr{condition, makeLower(1), makeLower(2)}, - ) - require.NoError(t, err) - peak, err := expressionVectorPeak(proc, expr, 4, false) - require.NoError(t, err) - rootOutput, err := expressionTypePeak(expr.Typ, 4) - require.NoError(t, err) - expectedPeak := rootOutput - for _, branch := range expr.GetF().Args[1:] { - branchOutput, branchErr := expressionTypePeak(branch.Typ, 4) - require.NoError(t, branchErr) - selectedParameter, parameterErr := expressionTypePeak(branch.GetF().Args[0].Typ, 4) - require.NoError(t, parameterErr) - expectedPeak += branchOutput + branchOutput + selectedParameter - } - require.Equal(t, expectedPeak, peak, - "flow-control branches need ordinary output, selected result, and selected parameter capacity") - budget := process.MustNewHashBuildBudget(2*peak, 2*peak) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - executors, lease, err := NewBudgetedExpressionExecutors( - proc, - generation, - []*plan.Expr{expr}, - false, - ) - require.NoError(t, err) - - eval := func(width int) { - bat := batch.NewWithSize(3) - bat.Vecs[0] = testutil.MakeBoolVector( - []bool{true, false, true, false}, - nil, - proc.Mp(), - ) - left := []string{ - strings.Repeat("A", width), - strings.Repeat("B", width), - strings.Repeat("C", width), - strings.Repeat("D", width), - } - right := []string{ - strings.Repeat("E", width), - strings.Repeat("F", width), - strings.Repeat("G", width), - strings.Repeat("H", width), - } - bat.Vecs[1] = testutil.MakeVarcharVector(left, nil, proc.Mp()) - bat.Vecs[2] = testutil.MakeVarcharVector(right, nil, proc.Mp()) - bat.SetRowCount(4) - defer bat.Clean(proc.Mp()) - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{bat}, - bat.RowCount(), - func(_ int, _ *vector.Vector) error { return nil }, - )) - } - eval(8) - retained, ok := lease.Retained() - require.True(t, ok) - require.LessOrEqual(t, retained, lease.Reserved(), - "selected result and parameter scratch must be covered by the admitted peak") - eval(4 << 10) - retained, ok = lease.Retained() - require.True(t, ok) - require.LessOrEqual(t, retained, lease.Reserved()) - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) -} - -func TestBudgetedExpressionConstructionAdmitsBeforeMpoolAllocation(t *testing.T) { - mp := mpool.MustNewZero() - proc := testutil.NewProcessWithMPool(t, "", mp) - defer proc.Free() - expr := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_varchar), Width: types.MaxVarcharLen}, - Expr: &plan.Expr_Lit{Lit: &plan.Literal{ - Value: &plan.Literal_Sval{Sval: strings.Repeat("x", 64<<10)}, - }}, - } - initial, err := expressionInitialOwnedBytes(expr) - require.NoError(t, err) - require.Positive(t, initial) - - budget := process.MustNewHashBuildBudget(initial-1, initial-1) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - epoch := mp.StartResourcePeakEpoch() - executors, lease, err := NewBudgetedExpressionExecutors( - proc, - generation, - []*plan.Expr{expr}, - false, - ) - peak, ok := mp.EndResourcePeakEpoch(epoch) - require.True(t, ok) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Nil(t, executors) - require.Nil(t, lease) - require.Zero(t, peak, "budget rejection must happen before constructing the literal vector") - require.Zero(t, generation.Used()) - require.Zero(t, mp.CurrNB()) -} - -func TestExpressionMemoryLeaseRetainsFailedEvaluationBound(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeExpressionLeaseTestExpr(t, proc) - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - retained, ok := colexec.ExpressionExecutorsRetainedBytes(executors) - require.True(t, ok) - peak, err := expressionVectorPeak(proc, expr, 32, false) - require.NoError(t, err) - budgetCap := retained + peak - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - lease, err := NewExpressionMemoryLease(generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - - wantErr := errors.New("expression evaluation failed") - require.ErrorIs(t, lease.Run(proc, 32, func(_ int) error { - return wantErr - }), wantErr) - require.Equal(t, peak, generation.Used(), - "a failed evaluator may retain partially grown buffers") - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) -} - -func TestExpressionMemoryLeaseDoesNotShrinkAdoptedCapacity(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeExpressionLeaseTestExpr(t, proc) - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - - large := makeExpressionLeaseTestBatch(proc, colexec.DefaultBatchSize*8) - defer large.Clean(proc.Mp()) - require.NoError(t, evalExpressionLeaseTestExecutors(proc, executors, large)) - retained, ok := colexec.ExpressionExecutorsRetainedBytes(executors) - require.True(t, ok) - require.Positive(t, retained) - - smallPeak, err := expressionVectorPeak(proc, expr, 1, false) - require.NoError(t, err) - require.Less(t, smallPeak, retained) - budget := process.MustNewHashBuildBudget(retained, retained) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - lease, err := NewExpressionMemoryLease( - generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - reservesAfterAdoption := generation.ReserveCount() - - small := makeExpressionLeaseTestBatch(proc, 1) - defer small.Clean(proc.Mp()) - require.NoError(t, lease.Run(proc, small.RowCount(), func(_ int) error { - return evalExpressionLeaseTestExecutors(proc, executors, small) - })) - require.Equal(t, reservesAfterAdoption, generation.ReserveCount()) - require.Equal(t, retained, generation.Used()) - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) -} - -func TestExpressionMemoryLeaseRejectsUnknownExecutorOwnership(t *testing.T) { - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - _, err = NewExpressionMemoryLease( - generation, - []*plan.Expr{{Typ: plan.Type{Id: int32(types.T_int32)}}}, - []colexec.ExpressionExecutor{unknownExpressionLeaseExecutor{}}, - false, - ) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - require.Zero(t, generation.Used()) -} - -func TestExpressionMemoryLeaseRejectsInvalidCalls(t *testing.T) { - _, err := NewExpressionMemoryLease( - nil, - []*plan.Expr{{Typ: plan.Type{Id: int32(types.T_int32)}}}, - nil, - false, - ) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - - var nilLease *ExpressionMemoryLease - require.ErrorIs(t, - nilLease.Run(nil, 0, func(int) error { return nil }), - process.ErrHashBuildBudgetInvalid, - ) - require.ErrorIs(t, - nilLease.Run(nil, 0, nil), - process.ErrHashBuildBudgetInvalid, - ) - require.ErrorIs(t, - nilLease.Eval(nil, nil, 0, nil), - process.ErrHashBuildBudgetInvalid, - ) - require.Zero(t, nilLease.Reserved()) - require.Zero(t, nilLease.Len()) - retained, ok := nilLease.Retained() - require.True(t, ok) - require.Zero(t, retained) - nilLease.Release() - - emptyLease := &ExpressionMemoryLease{} - require.ErrorIs(t, - emptyLease.Run(nil, -1, func(int) error { return nil }), - process.ErrHashBuildBudgetInvalid, - ) -} - -func TestExpressionMemoryAccountingHelperBoundaries(t *testing.T) { - size, err := expressionInitialOwnedBytes(nil) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - require.Zero(t, size) - - size, err = literalInitialOwnedBytes(types.T_int32, &plan.Literal{}) - require.NoError(t, err) - require.Zero(t, size) - - require.True(t, expressionExecutorMayGrowWithinBound(nil)) - require.True(t, expressionExecutorMayGrowWithinBound(&plan.Expr{})) - - size, err = initialAllocationCapacity(0) - require.NoError(t, err) - require.Zero(t, size) -} - -func TestExpressionMemoryLeaseEvalPropagatesExecutorError(t *testing.T) { - expected := errors.New("expression evaluation failed") - executor := failingExpressionLeaseExecutor{err: expected} - lease, err := NewExpressionMemoryLease( - nil, - []*plan.Expr{{Typ: plan.Type{Id: int32(types.T_int32)}}}, - []colexec.ExpressionExecutor{executor}, - false, - ) - require.NoError(t, err) - defer lease.Release() - - consumed := false - err = lease.Eval(nil, nil, 0, func(int, *vector.Vector) error { - consumed = true - return nil - }) - require.ErrorIs(t, err, expected) - require.False(t, consumed) -} - -func TestExpressionMemoryLeaseReleaseIsTerminal(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeExpressionLeaseTestExpr(t, proc) - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - lease, err := NewExpressionMemoryLease( - generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - lease.Release() - require.Zero(t, lease.Reserved()) - _, ok := lease.Retained() - require.False(t, ok) - require.Zero(t, generation.Used()) - - called := false - err = lease.Run(proc, 1, func(_ int) error { - called = true - return nil - }) - require.ErrorIs(t, err, process.ErrHashBuildReservationInactive) - require.False(t, called) - require.Zero(t, generation.Used()) - - executors, err = colexec.NewExpressionExecutorsFromPlanExpressions( - proc, []*plan.Expr{expr}) - require.NoError(t, err) - lease, err = NewExpressionMemoryLease( - generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - require.NoError(t, lease.Run(proc, 1, func(_ int) error { - freeExpressionLeaseTestExecutors(executors) - lease.Release() - return nil - })) - require.Zero(t, generation.Used(), - "release during evaluation must also discard the pending replacement") -} - -func TestExpressionMemoryLeaseCancellationAndRepeatedGenerations(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeExpressionLeaseTestExpr(t, proc) - budget := process.MustNewHashBuildBudget(8<<20, 8<<20) - - runGeneration := func(id uint64, callbackErr error) { - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions( - proc, []*plan.Expr{expr}) - require.NoError(t, err) - generation, err := budget.OpenGeneration(id) - require.NoError(t, err) - lease, err := NewExpressionMemoryLease( - generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - input := makeExpressionLeaseTestBatch(proc, 128) - - err = lease.Run(proc, input.RowCount(), func(_ int) error { - if evalErr := evalExpressionLeaseTestExecutors( - proc, executors, input); evalErr != nil { - return evalErr - } - return callbackErr - }) - if callbackErr != nil { - require.ErrorIs(t, err, callbackErr) - } else { - require.NoError(t, err) - } - require.Positive(t, generation.Used()) - - input.Clean(proc.Mp()) - freeExpressionLeaseTestExecutors(executors) - lease.Release() - generation.Close() - require.Zero(t, generation.Used()) - require.Zero(t, proc.Mp().CurrNB()) - } - - runGeneration(1, context.Canceled) - runGeneration(2, nil) -} - -type unknownExpressionLeaseExecutor struct{} - -func (unknownExpressionLeaseExecutor) Eval(*process.Process, []*batch.Batch, []bool) (*vector.Vector, error) { - return nil, nil -} -func (unknownExpressionLeaseExecutor) EvalWithoutResultReusing(*process.Process, []*batch.Batch, []bool) (*vector.Vector, error) { - return nil, nil -} -func (unknownExpressionLeaseExecutor) ResetForNextQuery() {} -func (unknownExpressionLeaseExecutor) Free() {} -func (unknownExpressionLeaseExecutor) IsColumnExpr() bool { return false } -func (unknownExpressionLeaseExecutor) TypeName() string { return "unknown" } - -type failingExpressionLeaseExecutor struct { - unknownExpressionLeaseExecutor - err error -} - -func (f failingExpressionLeaseExecutor) Eval( - *process.Process, - []*batch.Batch, - []bool, -) (*vector.Vector, error) { - return nil, f.err -} diff --git a/pkg/sql/colexec/hashbuild/expression_test_helpers_test.go b/pkg/sql/colexec/hashbuild/expression_test_helpers_test.go new file mode 100644 index 0000000000000..88bec2b987dd3 --- /dev/null +++ b/pkg/sql/colexec/hashbuild/expression_test_helpers_test.go @@ -0,0 +1,94 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package hashbuild + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/pb/plan" + plan2 "github.com/matrixorigin/matrixone/pkg/sql/plan" + "github.com/matrixorigin/matrixone/pkg/vm/process" + "github.com/stretchr/testify/require" +) + +func makeIssue26454ConcatKey(t testing.TB, proc *process.Process) *plan.Expr { + t.Helper() + cast := func(colPos int32) *plan.Expr { + col := &plan.Expr{ + Typ: plan.Type{Id: int32(types.T_int32)}, + Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: colPos}}, + } + targetType := plan.Type{ + Id: int32(types.T_varchar), + Width: types.MaxVarcharLen, + } + expr, err := plan2.BindFuncExprImplByPlanExpr( + proc.Ctx, + "cast", + []*plan.Expr{ + col, + { + Typ: targetType, + Expr: &plan.Expr_T{T: &plan.TargetType{}}, + }, + }, + ) + require.NoError(t, err) + return expr + } + expr, err := plan2.BindFuncExprImplByPlanExpr( + proc.Ctx, + "concat", + []*plan.Expr{ + cast(0), + plan2.MakePlan2StringConstExprWithType("-"), + cast(1), + }, + ) + require.NoError(t, err) + return expr +} + +func makeIssue26454CaseKey(t testing.TB, proc *process.Process) *plan.Expr { + t.Helper() + column := &plan.Expr{ + Typ: plan.Type{ + Id: int32(types.T_varchar), + Width: types.MaxVarcharLen, + }, + Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, + } + condition, err := plan2.BindFuncExprImplByPlanExpr( + proc.Ctx, + "=", + []*plan.Expr{ + column, + plan2.MakePlan2StringConstExprWithType("ATM_CON"), + }, + ) + require.NoError(t, err) + expr, err := plan2.BindFuncExprImplByPlanExpr( + proc.Ctx, + "case", + []*plan.Expr{ + condition, + plan2.MakePlan2StringConstExprWithType("CON_CONTRACT_HEADERS"), + plan2.MakePlan2StringConstExprWithType("CON_CONTRACT_DOC"), + }, + ) + require.NoError(t, err) + return expr +} diff --git a/pkg/sql/colexec/hashbuild/hashmap.go b/pkg/sql/colexec/hashbuild/hashmap.go index f79c0d5eaeba0..8d8bc7afffb75 100644 --- a/pkg/sql/colexec/hashbuild/hashmap.go +++ b/pkg/sql/colexec/hashbuild/hashmap.go @@ -31,7 +31,6 @@ import ( "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" "github.com/matrixorigin/matrixone/pkg/sql/colexec/runtimefilter" - "github.com/matrixorigin/matrixone/pkg/sql/plan/function" "github.com/matrixorigin/matrixone/pkg/vm/message" "github.com/matrixorigin/matrixone/pkg/vm/process" ) @@ -74,11 +73,7 @@ type HashmapBuilder struct { dedupDeleteKeepColIdxList []int32 DelRows *bitmap.Bitmap budget *process.HashBuildBudgetGeneration - mapReservation *hashMapReservationOwner - batchReservations []*process.HashBuildReservation - auxReservation *process.HashBuildReservation keyExprs []*plan.Expr - expressionLease *ExpressionMemoryLease // Exact runtime-filter keys are an optional owner inside the mandatory // JoinMap build. The fallback bit is observed by HashBuild for diagnostics. // @@ -91,9 +86,9 @@ type HashmapBuilder struct { retainedBatchRecoverySafe bool mapAllocationAccount *mpool.AllocationAccount mapAllocation *hashtable.AllocationAccountSelection + iteratorAllocation *hashmap.IteratorAllocation batchAllocation *vector.AllocationAccountSelection uniqueKeyAllocation *vector.AllocationAccountSelection - expressionAllocation *colexec.ExpressionAllocationAccount } func (hb *HashmapBuilder) GetSize() int64 { @@ -132,16 +127,12 @@ func (hb *HashmapBuilder) GetJoinMap(mp *mpool.MPool) *message.JoinMap { hb.DelRows = nil hb.Batches.Reset() // Iterators are producer scratch and are not part of JoinMap ownership. - // Drop budgeted cached backing before transferring the encompassing aux - // reservation to a consumer that may free it immediately after publication. hb.detachAndPruneCachedIterators() hb.freeIgnoreRows(mp) hb.uniqueSels = nil hb.curVecs = nil - release := hb.detachReservations() jm.SetMemoryRelease(func() { releaseDedupBitmap(jmDelRows, mp) - release() }) return jm } @@ -163,7 +154,20 @@ func (hb *HashmapBuilder) observeNullKeys(keyVecs []*vector.Vector) { return } for _, vec := range keyVecs { - if vec.HasNull() { + if vec == nil { + continue + } + rows := uint64(vec.Length()) + if vec.IsConstNull() { + if vec.GetGrouping().GetBitmap().CountRange(0, rows) < vec.Length() { + hb.HasNullKey = true + return + } + continue + } + if vec.GetNulls().GetBitmap().AnySetNotIn( + vec.GetGrouping().GetBitmap(), 0, rows, + ) { hb.HasNullKey = true return } @@ -189,33 +193,16 @@ func (hb *HashmapBuilder) Prepare( } keyWidth += width } - var ( - executors []colexec.ExpressionExecutor - expressionLease *ExpressionMemoryLease - err error + executors, err := NewExpressionExecutors( + proc, + keyCols, ) - if hb.expressionAllocation != nil && - expressionSetAllocationClosed(keyCols) { - executors, err = NewAllocationAccountedExpressionExecutors( - proc, - keyCols, - hb.expressionAllocation, - ) - } else { - executors, expressionLease, err = NewBudgetedExpressionExecutors( - proc, - hb.budget, - keyCols, - needDupVec, - ) - } if err != nil { return err } hb.needDupVec = needDupVec hb.executors = executors hb.keyExprs = keyCols - hb.expressionLease = expressionLease hb.keyWidth = keyWidth hb.InputBatchRowCount = 0 hb.hashMapRowCount = 0 @@ -259,14 +246,9 @@ func (hb *HashmapBuilder) Reset(proc *process.Process, hashTableHasNotSent bool) hb.UniqueJoinKeys = nil hb.uniqueKeySlots = nil // Function executors retain result-vector capacity across ResetForNextQuery. - // Free them before releasing expression reservations; Prepare recreates the - // executor set for the next generation. + // Destroy them here; immutable allocation selections remain installed until + // the statement lifecycle calls ClearAllocationAccount. hb.FreeExecutors() - hb.mapAllocationAccount = nil - hb.mapAllocation = nil - hb.batchAllocation = nil - hb.uniqueKeyAllocation = nil - hb.expressionAllocation = nil } func (hb *HashmapBuilder) Free(proc *process.Process) { @@ -290,11 +272,8 @@ func (hb *HashmapBuilder) Free(proc *process.Process) { } hb.UniqueJoinKeys = nil hb.uniqueKeySlots = nil - hb.mapAllocationAccount = nil - hb.mapAllocation = nil - hb.batchAllocation = nil - hb.uniqueKeyAllocation = nil - hb.expressionAllocation = nil + hb.runtimeFilterCollectionFallback = false + hb.retainedBatchRecoverySafe = false } func (hb *HashmapBuilder) FreeExecutors() { @@ -305,7 +284,6 @@ func (hb *HashmapBuilder) FreeExecutors() { } hb.executors = nil hb.keyExprs = nil - hb.releaseExpressionLease() } func (hb *HashmapBuilder) FreeTemporaryVectors(proc *process.Process) { @@ -332,7 +310,6 @@ func (hb *HashmapBuilder) FreeHashMapAndBatches(proc *process.Process) { hb.Batches.Clean(proc.Mp()) hb.freeIgnoreRows(proc.Mp()) hb.freeDelRows(proc.Mp()) - hb.releaseReservations() } // evalBatch evaluates join key expressions for one batch, storing results in hb.curVecs. @@ -365,13 +342,9 @@ func (hb *HashmapBuilder) evalBatch(batchIdx int, proc *process.Process) error { return nil } var err error - if hb.expressionLease != nil { - err = hb.expressionLease.Run(proc, bat.RowCount(), evalOne) - } else { - for idx := range hb.executors { - if err = evalOne(idx); err != nil { - break - } + for idx := range hb.executors { + if err = evalOne(idx); err != nil { + break } } if err != nil { @@ -383,220 +356,42 @@ func (hb *HashmapBuilder) evalBatch(batchIdx int, proc *process.Process) error { func (hb *HashmapBuilder) abortExpressionEval(proc *process.Process) { // Eval may allocate cached child/result vectors before returning an error. - // Destroy the complete executor tree before releasing its retained lease. + // Destroy the complete executor tree so every exact allocation is released. hb.FreeTemporaryVectors(proc) hb.FreeExecutors() } -// expressionVectorPeak is an execution-before-allocation upper bound based on -// the SQL result type. Varlena widths use the declared maximum (or the engine -// maximum when absent), so input-dependent expanding functions are rejected -// by admission before Eval instead of allocating first. -func expressionVectorPeak(proc *process.Process, expr *plan.Expr, rows int, duplicate bool) (uint64, error) { - if expr == nil || rows < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - total, root, err := expressionTreePeak(proc, expr, uint64(rows)) - if err != nil { - return 0, err - } - if duplicate { - if total > math.MaxUint64-root { - return 0, process.ErrHashBuildBudgetInvalid - } - total += root - } - return total, nil -} - -// ExpressionVectorPeak exposes the same execution-before-allocation bound used -// by HashmapBuilder to spill/re-spill callers. Expression evaluators cache -// intermediate and result vectors, so callers must keep the returned amount -// reserved until the corresponding executor tree is freed or evaluated again -// under a replacement reservation. -func ExpressionVectorPeak(proc *process.Process, expr *plan.Expr, rows int, duplicate bool) (uint64, error) { - return expressionVectorPeak(proc, expr, rows, duplicate) -} - -func expressionTreePeak(proc *process.Process, expr *plan.Expr, rows uint64) (total uint64, output uint64, err error) { - return expressionTreePeakWithSelection(proc, expr, rows, false) -} - -func expressionTreePeakWithSelection( - proc *process.Process, - expr *plan.Expr, - rows uint64, - mayReceivePartialSelection bool, -) (total uint64, output uint64, err error) { - if expr == nil { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - switch node := expr.Expr.(type) { - case *plan.Expr_Col: - return 0, 0, nil - case *plan.Expr_F: - if node.F == nil { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - var fid int32 = -1 - if node.F.Func != nil { - fid, _ = function.DecodeOverloadID(node.F.Func.Obj) - } - for i, arg := range node.F.Args { - childMayReceivePartialSelection := mayReceivePartialSelection - switch fid { - case function.IFF: - // IFF evaluates only its value branches through generated - // selection masks. Its condition inherits the caller mask. - childMayReceivePartialSelection = mayReceivePartialSelection || i > 0 - case function.CASE, function.COALESCE: - childMayReceivePartialSelection = true - } - child, _, childErr := expressionTreePeakWithSelection( - proc, - arg, - rows, - childMayReceivePartialSelection, - ) - if childErr != nil || total > math.MaxUint64-child { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - total += child - } - case *plan.Expr_P: - if node.P == nil || proc == nil || proc.GetPrepareParams() == nil { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - paramPeak, paramErr := expressionParamPeak(proc, node.P.Pos) - if paramErr != nil { - return 0, 0, paramErr - } - typePeak, typeErr := expressionTypePeak(expr.Typ, 1) - if typeErr != nil { - return 0, 0, typeErr - } - if paramPeak > typePeak { - output = paramPeak - } else { - output = typePeak - } - return output, output, nil - case *plan.Expr_Lit, *plan.Expr_V, *plan.Expr_Raw, *plan.Expr_Vec, *plan.Expr_Fold, *plan.Expr_T: - // These executors may materialize a vector but have no child expression - // tree. Expr_T is the target-type argument used by CAST/bit_cast and is - // evaluated as a fixed vector. Charge their declared output below. - default: - // Window, subquery, correlated, list and max nodes do not - // expose a bounded vector-evaluator tree here. - return 0, 0, process.ErrHashBuildBudgetInvalid - } - output, err = expressionTypePeak(expr.Typ, rows) - if err != nil || total > math.MaxUint64-output { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - total += output - - if _, isFunction := expr.Expr.(*plan.Expr_F); mayReceivePartialSelection && isFunction { - // A partially selected function retains both its ordinary full-row - // result and a selected-result scratch vector. Row-aligned column and - // non-folded function parameters are also copied into retained selected - // parameter vectors before the function executes. - if total > math.MaxUint64-output { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - total += output - for _, arg := range nodeFunctionArgs(expr) { - switch arg.Expr.(type) { - case *plan.Expr_Col, *plan.Expr_F: - selectedParameter, selectedErr := expressionTypePeak(arg.Typ, rows) - if selectedErr != nil || total > math.MaxUint64-selectedParameter { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - total += selectedParameter - } - } - } - return total, output, nil -} - -func nodeFunctionArgs(expr *plan.Expr) []*plan.Expr { - if node, ok := expr.Expr.(*plan.Expr_F); ok && node.F != nil { - return node.F.Args - } - return nil -} - -// expressionParamPeak returns an upper bound for the allocations made by a -// non-null ParamExpressionExecutor. Params are materialized as one-element -// const vectors, whose data is one varlena header and whose area is allocated -// only for payloads that do not fit in that header. -func expressionParamPeak(proc *process.Process, pos int32) (uint64, error) { - val, err := proc.GetPrepareParamsAt(int(pos)) - if err != nil { - return 0, err - } - if val == nil { - return 0, nil - } - - headerCap, ok := mpool.GrowCapacity(0, int64(types.VarlenaSize)) - if !ok { - return 0, process.ErrHashBuildBudgetInvalid - } - peak := uint64(headerCap) - if len(val) <= types.VarlenaInlineSize { - return peak, nil - } - - areaCap, ok := mpool.GrowCapacity(0, int64(len(val))) - if !ok || areaCap < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - if uint64(areaCap) > math.MaxUint64-peak { - return 0, process.ErrHashBuildBudgetInvalid - } - return peak + uint64(areaCap), nil -} - -func expressionTypePeak(typ plan.Type, rows uint64) (uint64, error) { - oid := types.T(typ.Id) - width := int64(oid.FixedLength()) - if width < 0 { - width = int64(typ.Width) - hardMax := int64(types.MaxVarcharLen) - if oid.IsArrayRelate() { - elementWidth := int64(oid.ToType().GetArrayElementSize()) - width *= elementWidth - hardMax = int64(types.MaxArrayDimension) * elementWidth - } else { - switch oid { - case types.T_blob, types.T_text, types.T_json, types.T_datalink, - types.T_geometry, types.T_geometry32: - hardMax = int64(types.MaxBlobLen) +// hasGroupingKey reports whether a direct build key column contains any +// GROUPING sentinel in any retained batch. Such maps use the string encoder's +// explicit key domains; an IntHashMap cannot represent a sentinel outside the +// complete uint64 raw-value domain without collisions. +func (hb *HashmapBuilder) hasGroupingKey() bool { + for _, executor := range hb.executors { + if !executor.IsColumnExpr() { + continue + } + column, ok := executor.(*colexec.ColumnExpressionExecutor) + if !ok { + continue + } + // HashBuild evaluates every key against one retained build batch. + // Some join planners preserve the original build relation index in a + // direct column expression, but ColumnExpressionExecutor deliberately + // resolves that expression against the only input batch. Inspect the + // same physical column here; filtering on RelPos would miss GROUPING + // sentinels for DedupJoin and RightDedupJoin. + position := column.GetColIndex() + for _, bat := range hb.Batches.Buf { + if bat != nil && position >= 0 && position < len(bat.Vecs) && + bat.Vecs[position] != nil && + bat.Vecs[position].GetGrouping().GetBitmap().CountRange( + 0, uint64(bat.Vecs[position].Length()), + ) > 0 { + return true } } - if width > hardMax { - // Never clamp a declared bound downward. Array width is declared - // in elements, while every other varlena width is in bytes. - hardMax = width - } - width = hardMax - } - if width < 1 { - width = 1 - } - perRow := uint64(width) + 32 - if rows > (math.MaxUint64-(64<<10))/perRow { - return 0, process.ErrHashBuildBudgetInvalid - } - return rows*perRow + (64 << 10), nil -} - -func (hb *HashmapBuilder) releaseExpressionLease() { - if hb.expressionLease != nil { - hb.expressionLease.Release() - hb.expressionLease = nil } + return false } func (hb *HashmapBuilder) BuildHashmap(hashOnPK bool, needAllocateSels bool, needUniqueVec bool, proc *process.Process) (retErr error) { @@ -645,28 +440,11 @@ func (hb *HashmapBuilder) buildHashmap( if hb.InputBatchRowCount == 0 { return nil } - if err := hb.reserveBuildAux(needUniqueVec, needAllocateSels); err != nil { - if !needUniqueVec { - return err - } - if runtimefilter.ClassifyOptionalFallback(err) != - runtimefilter.OptionalFallbackBudgetAdmission { - return err - } - // The extra auxiliary charge exists only for optional exact-filter key - // retention. Retry the admission in place without that owner before - // allocating or mutating the mandatory map. - needUniqueVec = false - if err = hb.reserveBuildAux(false, needAllocateSels); err != nil { - return err - } - // Linearize the fallback only after mandatory admission succeeds. A - // failed retry is a fatal build, not a successful optional downgrade. - hb.runtimeFilterCollectionFallback = true - } dedupBuildKeepLast = dedupBuildKeepLast && hb.IsDedup && hb.OnDuplicateAction == plan.Node_FAIL defer func() { if retErr != nil { + hashmap.IteratorClearOwner(hb.cachedIntIterator) + hashmap.IteratorClearOwner(hb.cachedStrIterator) hb.cachedIntIterator = nil hb.cachedStrIterator = nil } @@ -682,28 +460,23 @@ func (hb *HashmapBuilder) buildHashmap( var err error var itr hashmap.Iterator - if hb.keyWidth <= 8 { - if hb.mapAllocation == nil { - if err = hb.reserveInitialMap(int64(hashtable.Int64HashMapInitialAllocationBytes())); err != nil { - return err - } - hb.IntHashMap, err = hashmap.NewIntHashMap(false, proc.Mp()) - if err == nil { - err = hb.attachIntHashMapAdmission(hb.IntHashMap) - } - } else { - hb.IntHashMap, err = hashmap.NewIntHashMapWithAllocation( - false, - proc.Mp(), - hb.mapAllocation, - ) - } + hasGroupingKey := hb.hasGroupingKey() + useIntHashMap := hb.keyWidth <= 8 && !hasGroupingKey + if hb.mapAllocation == nil || hb.mapAllocationAccount == nil || + hb.iteratorAllocation == nil || hb.batchAllocation == nil { + return mpool.ErrAllocationAccountInvalid + } + if useIntHashMap { + hb.IntHashMap, err = hashmap.NewIntHashMapWithAllocation( + false, + proc.Mp(), + hb.mapAllocation, + ) if err != nil { if hb.IntHashMap != nil { hb.IntHashMap.Free() hb.IntHashMap = nil } - hb.releaseMapReservation() return err } if hb.cachedIntIterator != nil { @@ -714,29 +487,24 @@ func (hb *HashmapBuilder) buildHashmap( hb.cachedIntIterator = itr } } else { - if hb.mapAllocation == nil { - if err = hb.reserveInitialMap(int64(hashtable.StringHashMapInitialAllocationBytes())); err != nil { - return err - } - hb.StrHashMap, err = hashmap.NewStrHashMap(false, proc.Mp()) - if err == nil { - err = hb.attachStrHashMapAdmission(hb.StrHashMap) - } - } else { - hb.StrHashMap, err = hashmap.NewStrHashMapWithAllocation( - false, - proc.Mp(), - hb.mapAllocation, - ) - } + hb.StrHashMap, err = hashmap.NewStrHashMapWithAllocations( + false, + proc.Mp(), + hb.mapAllocation, + hb.iteratorAllocation, + ) if err != nil { if hb.StrHashMap != nil { hb.StrHashMap.Free() hb.StrHashMap = nil } - hb.releaseMapReservation() return err } + if hasGroupingKey { + if err = hb.StrHashMap.SetGroupingAware(); err != nil { + return err + } + } if hb.cachedStrIterator != nil { hashmap.IteratorChangeOwner(hb.cachedStrIterator, hb.StrHashMap) itr = hb.cachedStrIterator @@ -748,7 +516,7 @@ func (hb *HashmapBuilder) buildHashmap( if hashOnPK || hb.IsDedup { // if hash on primary key, prealloc hashmap size to the count of batch - if hb.keyWidth <= 8 { + if useIntHashMap { err = hb.IntHashMap.PreAlloc(uint64(hb.InputBatchRowCount)) if err != nil { return err @@ -762,18 +530,13 @@ func (hb *HashmapBuilder) buildHashmap( } if needAllocateSels { - var err error - if hb.batchAllocation == nil { - err = hb.Sels.Init(hb.InputBatchRowCount, proc.Mp()) - } else { - err = hb.Sels.InitWithAllocation( - hb.InputBatchRowCount, - proc.Mp(), - hb.mapAllocationAccount, - HashBuildAllocationOwner, - HashBuildAllocationSiteGroupSels, - ) - } + err = hb.Sels.InitWithAllocation( + hb.InputBatchRowCount, + proc.Mp(), + hb.mapAllocationAccount, + HashBuildAllocationOwner, + HashBuildAllocationSiteGroupSels, + ) if err != nil { return err } @@ -876,7 +639,7 @@ buildUnits: // if not hash on primary key, estimate the hashmap size after 8192 rows //preAlloc to improve performance and reduce memory reAlloc if !hashOnPK && !hb.IsDedup && hb.InputBatchRowCount > hashmap.HashMapSizeThreshHold && i == hashmap.HashMapSizeEstimate { - if hb.keyWidth <= 8 { + if useIntHashMap { groupCount := hb.IntHashMap.GroupCount() rate := float64(groupCount) / float64(i) hashmapCount := uint64(float64(hb.InputBatchRowCount) * rate) @@ -921,7 +684,14 @@ buildUnits: zvals = ignoreBuildZvals[:n] clear(ignoreCandidateOwnsKey[:n]) ignoreCandidateOldKey[0] = hb.Batches.Buf[vecIdx1].Vecs[hb.delColIdx] - oldVals, oldZvals := itr.Find(vecIdx2, n, ignoreCandidateOldKey) + oldVals, oldZvals, findErr := itr.Find( + vecIdx2, + n, + ignoreCandidateOldKey, + ) + if findErr != nil { + return findErr + } for k := 0; k < n; k++ { ignoreCandidateOwnsKey[k] = zvals[k] != 0 && oldZvals[k] != 0 && vals[k] != 0 && oldVals[k] == vals[k] } @@ -1006,15 +776,14 @@ buildUnits: if needUniqueVec { if len(hb.UniqueJoinKeys) == 0 { + if hb.uniqueKeyAllocation == nil { + return mpool.ErrAllocationAccountInvalid + } hb.UniqueJoinKeys = make([]*vector.Vector, len(hb.executors)) for j, vec := range hb.curVecs { if !hb.collectUniqueKeySlot(j) { continue } - if hb.uniqueKeyAllocation == nil { - hb.UniqueJoinKeys[j] = vector.NewOffHeapVecWithType(*vec.GetType()) - continue - } hb.UniqueJoinKeys[j], err = vector.NewOffHeapVecWithTypeAndAllocation( *vec.GetType(), hb.uniqueKeyAllocation, @@ -1041,32 +810,12 @@ buildUnits: if !hb.collectUniqueKeySlot(j) { continue } - areaBytes, reserveErr := - unionBatchAreaBytes(vec, vecIdx2, n) - if reserveErr != nil { - // Range and overflow failures contradict the collection - // oracle; they are never optional allocation failures. - return reserveErr - } - overlap, reserveErr := hb.reserveUniqueAppendOverlap(hb.UniqueJoinKeys[j], n, areaBytes) - if reserveErr != nil { - if fatalErr := - hb.fallbackOptionalRuntimeFilterCollection( - proc, reserveErr); fatalErr != nil { - return fatalErr - } - needUniqueVec = false - break - } err = hb.UniqueJoinKeys[j].UnionBatch(vec, int64(vecIdx2), n, nil, proc.Mp()) - if overlap != nil { - overlap.Release() - } if err != nil { - // With the range and capacity oracle above satisfied, - // UnionBatch error returns are only mpool growth failures. - allocationErr := - runtimefilter.MarkOptionalAllocationError(err) + allocationErr := err + if mpool.IsRetryableAllocationCapacity(err) { + allocationErr = runtimefilter.MarkOptionalAllocationError(err) + } if fatalErr := hb.fallbackOptionalRuntimeFilterCollection( proc, allocationErr); fatalErr != nil { @@ -1093,31 +842,12 @@ buildUnits: if !hb.collectUniqueKeySlot(j) { continue } - areaBytes, reserveErr := uniqueAppendAreaBytes(vec, 0, len(newSels), newSels) - if reserveErr != nil { - // Selector/range/overflow failures are collection - // contract errors and remain fatal. - return reserveErr - } - overlap, reserveErr := hb.reserveUniqueAppendOverlap(hb.UniqueJoinKeys[j], len(newSels), areaBytes) - if reserveErr != nil { - if fatalErr := - hb.fallbackOptionalRuntimeFilterCollection( - proc, reserveErr); fatalErr != nil { - return fatalErr - } - needUniqueVec = false - break - } err = hb.UniqueJoinKeys[j].Union(vec, newSels, proc.Mp()) - if overlap != nil { - overlap.Release() - } if err != nil { - // With generated selectors and the capacity oracle above - // satisfied, Union error returns are mpool growth failures. - allocationErr := - runtimefilter.MarkOptionalAllocationError(err) + allocationErr := err + if mpool.IsRetryableAllocationCapacity(err) { + allocationErr = runtimefilter.MarkOptionalAllocationError(err) + } if fatalErr := hb.fallbackOptionalRuntimeFilterCollection( proc, allocationErr); fatalErr != nil { @@ -1241,7 +971,14 @@ buildUnits: if err = hb.evalBatch(vecIdx1, proc); err != nil { return err } - newVals, newZvals := itr.Find(vecIdx2, n, hb.curVecs) + newVals, newZvals, findErr := itr.Find( + vecIdx2, + n, + hb.curVecs, + ) + if findErr != nil { + return findErr + } for k := 0; k < n; k++ { buildGroups[k] = 0 if newZvals[k] != 0 { @@ -1250,7 +987,10 @@ buildUnits: } } tmpVecs[0] = hb.Batches.Buf[vecIdx1].Vecs[hb.delColIdx] - vals, zvals := itr.Find(vecIdx2, n, tmpVecs) + vals, zvals, findErr := itr.Find(vecIdx2, n, tmpVecs) + if findErr != nil { + return findErr + } for k, v := range vals[:n] { if zvals[k] == 0 || v == 0 { @@ -1293,7 +1033,6 @@ func (hb *HashmapBuilder) resetHashStateForRebuild(proc *process.Process) { hb.StrHashMap.Free() hb.StrHashMap = nil } - hb.releaseMapReservation() hb.Sels.Free(proc.Mp()) for i := range hb.UniqueJoinKeys { if hb.UniqueJoinKeys[i] != nil { @@ -1311,16 +1050,11 @@ func (hb *HashmapBuilder) resetHashStateForRebuild(proc *process.Process) { } // FreeHashMapOnly discards a partial hash build while preserving the copied -// build batches and their reservations. It is the supported transition from a -// failed BuildHashmap attempt to either a less memory-intensive rebuild or -// bounded spill recovery. +// build batches for bounded spill recovery. It is the only supported +// transition from a failed BuildHashmap attempt to re-spill. func (hb *HashmapBuilder) FreeHashMapOnly(proc *process.Process) { hb.resetHashStateForRebuild(proc) hb.freeDelRows(proc.Mp()) - if hb.auxReservation != nil { - hb.auxReservation.Release() - hb.auxReservation = nil - } } func (hb *HashmapBuilder) keepDiscardedRowsForDelete(proc *process.Process) error { @@ -1409,23 +1143,25 @@ func (hb *HashmapBuilder) makeDeleteOnlyBatch(rows []int32, proc *process.Proces } bat := batch.NewOffHeapWithSize(len(hb.Batches.Buf[0].Vecs)) - if hb.mapAllocationAccount != nil { - selection, err := vector.NewAllocationAccountSelectionWithBitmaps( - hb.mapAllocationAccount, - HashBuildAllocationOwner, - HashBuildAllocationSiteDedupDeleteOnlyData, - HashBuildAllocationSiteDedupDeleteOnlyArea, - HashBuildAllocationSiteDedupDeleteOnlyNulls, - HashBuildAllocationSiteDedupDeleteOnlyGrouping, - ) - if err != nil { - bat.Clean(proc.Mp()) - return nil, err - } - if err = bat.SetAllocationAccount(selection); err != nil { - bat.Clean(proc.Mp()) - return nil, err - } + if hb.mapAllocationAccount == nil { + bat.Clean(proc.Mp()) + return nil, mpool.ErrAllocationAccountInvalid + } + selection, err := vector.NewAllocationAccountSelection( + hb.mapAllocationAccount, + HashBuildAllocationOwner, + HashBuildAllocationSiteDedupDeleteOnlyData, + HashBuildAllocationSiteDedupDeleteOnlyArea, + HashBuildAllocationSiteDedupDeleteOnlyNulls, + HashBuildAllocationSiteDedupDeleteOnlyGrouping, + ) + if err != nil { + bat.Clean(proc.Mp()) + return nil, err + } + if err = bat.SetAllocationAccount(selection); err != nil { + bat.Clean(proc.Mp()) + return nil, err } bat.Attrs = hb.Batches.Buf[0].Attrs for colIdx, vec := range hb.Batches.Buf[0].Vecs { @@ -1487,16 +1223,10 @@ func (hb *HashmapBuilder) detachAndPruneCachedIterators() { } if hb.cachedStrIterator != nil { if hashmap.StrIteratorCapacity(hb.cachedStrIterator) > hashmap.MaxStrIteratorCapacity { + hashmap.IteratorClearOwner(hb.cachedStrIterator) hb.cachedStrIterator = nil return } hashmap.IteratorClearOwner(hb.cachedStrIterator) } - if hb.budget != nil { - // Budgeted builds charge iterator scratch only for the execution that - // allocated it. Do not retain Go backing arrays in the pooled operator - // after that reservation is released or transferred. - hb.cachedIntIterator = nil - hb.cachedStrIterator = nil - } } diff --git a/pkg/sql/colexec/hashbuild/hashmap_test.go b/pkg/sql/colexec/hashbuild/hashmap_test.go index e120b56006e13..82952e6eb8f3f 100644 --- a/pkg/sql/colexec/hashbuild/hashmap_test.go +++ b/pkg/sql/colexec/hashbuild/hashmap_test.go @@ -15,10 +15,7 @@ package hashbuild import ( - "context" - "errors" "fmt" - "math" "reflect" "strconv" "strings" @@ -33,318 +30,184 @@ import ( "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" - "github.com/matrixorigin/matrixone/pkg/sql/colexec/runtimefilter" "github.com/matrixorigin/matrixone/pkg/testutil" "github.com/matrixorigin/matrixone/pkg/vm/process" "github.com/stretchr/testify/require" ) -func TestBuildHashMap(t *testing.T) { - var hb HashmapBuilder - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - err := hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc) - require.NoError(t, err) - - inputBatch := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, int(100000), proc.Mp()) - err = hb.Batches.CopyIntoBatches(inputBatch, proc) - hb.InputBatchRowCount = inputBatch.RowCount() - inputBatch.Clean(proc.Mp()) - require.NoError(t, err) - - err = hb.BuildHashmap(false, true, true, proc) - require.NoError(t, err) - require.Less(t, int64(0), hb.GetSize()) - require.Less(t, uint64(0), hb.GetGroupCount()) - hb.Reset(proc, true) - hb.Free(proc) - require.Equal(t, int64(0), proc.Mp().CurrNB()) -} - -func TestBuildHashmapOptionalAuxClosedBudgetRemainsFatal(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - generation.Close() - - hb := HashmapBuilder{InputBatchRowCount: 1} - hb.setBudget(generation) - err = hb.BuildHashmap(false, false, true, proc) - require.Error(t, err) - var budgetErr *process.HashBuildBudgetError - require.ErrorAs(t, err, &budgetErr) - require.Equal(t, process.HashBuildBudgetErrorClosed, budgetErr.Kind) - fallback, _ := hb.runtimeFilterFallbackState() - require.False(t, fallback) - require.Nil(t, hb.UniqueJoinKeys) - require.Zero(t, generation.Used()) - hb.Free(proc) -} - -func TestBuildHashmapMandatoryAuxRetryFailureDoesNotRecordFallback(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - budget := process.MustNewHashBuildBudget(1, 1) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - - hb := HashmapBuilder{InputBatchRowCount: 1} - hb.setBudget(generation) - err = hb.BuildHashmap(false, false, true, proc) - require.Error(t, err) - var budgetErr *process.HashBuildBudgetError - require.ErrorAs(t, err, &budgetErr) - require.Equal(t, process.HashBuildBudgetErrorAdmission, budgetErr.Kind) - fallback, _ := hb.runtimeFilterFallbackState() - require.False(t, fallback, - "fatal mandatory retry must not be counted as an optional fallback") - require.Equal(t, uint64(2), generation.RejectCount()) - require.Zero(t, generation.Used()) - hb.Free(proc) -} - -func TestPrepareCanonicalRuntimeFilterCollectionClosedBudgetRemainsFatal( - t *testing.T, -) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - input := testutil.NewBatch( - []types.Type{types.T_int32.ToType()}, true, 16, proc.Mp()) - defer input.Clean(proc.Mp()) - budget := process.MustNewHashBuildBudget(64<<20, 64<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - hb := HashmapBuilder{ - Batches: colexec.Batches{Buf: []*batch.Batch{input}}, - InputBatchRowCount: input.RowCount(), +func TestObserveNullKeysUsesColumnLevelGroupingSentinel(t *testing.T) { + mp := mpool.MustNewZero() + newVec := func(nullRows, groupingRows []uint64) *vector.Vector { + vec := vector.NewVec(types.T_int32.ToType()) + require.NoError(t, vector.AppendFixedList( + vec, + []int32{1, 2}, + nil, + mp, + )) + for _, row := range nullRows { + vec.GetNulls().Add(row) + } + for _, row := range groupingRows { + vec.GetGrouping().Add(row) + } + return vec } - hb.setBudget(generation) - require.NoError(t, hb.reserveBuildAux(false)) - generation.Close() - - collect, err := hb.prepareCanonicalRuntimeFilterCollection(true) - require.Error(t, err) - require.False(t, collect) - var budgetErr *process.HashBuildBudgetError - require.ErrorAs(t, err, &budgetErr) - require.Equal(t, process.HashBuildBudgetErrorClosed, budgetErr.Kind) - fallback, _ := hb.runtimeFilterFallbackState() - require.False(t, fallback) - hb.Batches.Buf = nil - hb.releaseReservations() - require.Zero(t, generation.Used()) -} -func TestOptionalRuntimeFilterCollectionCleanupFailureRemainsFatal( - t *testing.T, -) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - key := testutil.MakeInt32Vector([]int32{1}, nil, proc.Mp()) - budget := process.MustNewHashBuildBudget(64<<20, 64<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - - hb := HashmapBuilder{ - InputBatchRowCount: 1, - UniqueJoinKeys: []*vector.Vector{key}, + tests := []struct { + name string + nullRows []uint64 + groupingRows []uint64 + want bool + }{ + { + name: "grouping sentinel masks null in same row", + nullRows: []uint64{0}, + groupingRows: []uint64{0}, + want: false, + }, + { + name: "full grouping is sentinel", + nullRows: []uint64{0, 1}, + groupingRows: []uint64{0, 1}, + want: false, + }, + { + name: "partial grouping without null", + groupingRows: []uint64{0}, + want: false, + }, + { + name: "null outside partial grouping is retained", + nullRows: []uint64{1}, + groupingRows: []uint64{0}, + want: true, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + vec := newVec(test.nullRows, test.groupingRows) + defer vec.Free(mp) + builder := HashmapBuilder{TrackNullKeys: true} + builder.observeNullKeys([]*vector.Vector{vec}) + require.Equal(t, test.want, builder.HasNullKey) + }) } - hb.setBudget(generation) - require.NoError(t, hb.reserveBuildAux(true)) - require.True(t, hb.auxReservation.Release()) - - err = hb.fallbackOptionalRuntimeFilterCollection( - proc, - runtimefilter.MarkOptionalAllocationError( - errors.New("mpool allocation failed")), - ) - require.ErrorIs(t, err, process.ErrHashBuildReservationInactive) - fallback, _ := hb.runtimeFilterFallbackState() - require.False(t, fallback) - require.Nil(t, hb.UniqueJoinKeys) - require.Zero(t, generation.Used()) - hb.releaseReservations() } -func TestBuildHashmapUniqueUnionAllocationFailureFallsBack(t *testing.T) { +func TestBuildHashmapPreservesRowwiseGroupingAcrossCopiedBatchMerge(t *testing.T) { for _, test := range []struct { - name string - hashOnPK bool + name string + groupFirst bool + columns int }{ - {name: "union"}, - {name: "union-batch", hashOnPK: true}, + {name: "grouping then ordinary", groupFirst: true, columns: 1}, + {name: "ordinary then grouping", groupFirst: false, columns: 1}, + {name: "multi-column grouping pattern", groupFirst: true, columns: 2}, } { t.Run(test.name, func(t *testing.T) { - testBuildHashmapUniqueUnionAllocationFailureFallsBack( - t, test.hashOnPK) - }) - } -} - -func testBuildHashmapUniqueUnionAllocationFailureFallsBack( - t *testing.T, - hashOnPK bool, -) { - mp, err := mpool.NewMPool(t.Name(), 8<<20, mpool.NoFixed) - require.NoError(t, err) - proc := testutil.NewProcessWithMPool(t, "", mp) - - var hb HashmapBuilder - require.NoError(t, hb.Prepare( - []*plan.Expr{newExpr(0, types.T_int32.ToType())}, - -1, -1, nil, proc)) - input := testutil.NewBatch( - []types.Type{types.T_int32.ToType()}, true, 16, mp) - hb.Batches.Buf = []*batch.Batch{input} - hb.InputBatchRowCount = input.RowCount() + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + builder := newTestHashmapBuilder(t) + defer builder.Free(proc) - budget := process.MustNewHashBuildBudget(64<<20, 64<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - hb.setBudget(generation) + exprs := make([]*plan.Expr, test.columns) + for column := range exprs { + exprs[column] = newExpr(int32(column), types.T_int32.ToType()) + } + require.NoError(t, builder.Prepare(exprs, -1, -1, nil, proc)) + + appendInput := func(grouping bool, groupingColumn int) { + input := batch.NewWithSize(test.columns) + for column := 0; column < test.columns; column++ { + if grouping && column == groupingColumn { + input.Vecs[column] = vector.NewRollupConst( + types.T_int32.ToType(), 1, proc.Mp(), + ) + } else { + input.Vecs[column] = vector.NewVec(types.T_int32.ToType()) + require.NoError(t, vector.AppendFixed( + input.Vecs[column], int32(column), false, proc.Mp(), + )) + } + } + input.SetRowCount(1) + require.NoError(t, builder.CopyBuildBatch(input, proc)) + builder.InputBatchRowCount++ + input.Clean(proc.Mp()) + } - var filler []byte - defer func() { - hb.Free(proc) - require.Zero(t, generation.Used()) - if filler != nil { - mp.Free(filler) - } - generation.Close() - proc.Free() - require.Zero(t, mp.CurrNB()) - }() + appendInput(test.groupFirst, 0) + if test.columns == 2 { + // A different grouping column in the later row proves that the + // per-column bit pattern survives tail coalescing. + appendInput(true, 1) + } else { + appendInput(!test.groupFirst, 0) + } - // Calibrate the deterministic mandatory map footprint, then leave exactly - // that much headroom. The second build can recreate its required map, while - // the first optional-key Union allocation must fail at the mpool boundary. - retainedBytes := mp.CurrNB() - require.NoError(t, hb.BuildHashmap(hashOnPK, false, false, proc)) - mapBytes := mp.CurrNB() - retainedBytes - require.Greater(t, mapBytes, int64(0)) - hb.FreeHashMapOnly(proc) - require.Equal(t, retainedBytes, mp.CurrNB()) - - fillerBytes := mp.Cap() - mp.CurrNB() - mapBytes - require.Greater(t, fillerBytes, int64(0)) - filler, err = mp.Alloc(int(fillerBytes), true) - require.NoError(t, err) + require.Len(t, builder.Batches.Buf, 1) + require.Equal(t, 2, builder.Batches.Buf[0].RowCount()) + require.False(t, builder.Batches.Buf[0].Vecs[0].IsGrouping()) + require.True(t, builder.Batches.Buf[0].Vecs[0].HasGrouping()) - require.NoError(t, hb.BuildHashmap(hashOnPK, false, true, proc)) - fallback, rebuildSafe := hb.runtimeFilterFallbackState() - require.True(t, fallback) - require.True(t, rebuildSafe) - require.Nil(t, hb.UniqueJoinKeys) - require.Greater(t, hb.GetGroupCount(), uint64(0)) - require.Zero(t, generation.RejectCount(), - "mpool failure must not be misclassified as budget admission") + require.NoError(t, builder.BuildHashmap(false, false, false, proc)) + require.Nil(t, builder.IntHashMap) + require.NotNil(t, builder.StrHashMap) + require.Equal(t, uint64(2), builder.StrHashMap.GroupCount()) + }) + } } -func TestBuildHashMapBudgetRejectsResizeAndReleasesOnReset(t *testing.T) { - const budgetCap = uint64(1 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) +func TestBuildHashmapDetectsGroupingForOriginalBuildRelation(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) - - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 10_000, proc.Mp()) - require.NoError(t, hb.copyBuildBatch(input, proc)) - hb.InputBatchRowCount = input.RowCount() - input.Clean(proc.Mp()) - - err = hb.BuildHashmap(false, false, false, proc) - require.Error(t, err) - require.True(t, errors.Is(err, process.ErrHashBuildBudgetAdmission)) - require.Greater(t, generation.Used(), uint64(0)) - - hb.Reset(proc, true) - require.Zero(t, generation.Used()) -} - -func TestBuildHashMapCancellationReleasesRetainedBudgetOnReset(t *testing.T) { - const budgetCap = uint64(16 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) + builder := newTestHashmapBuilder(t) + defer builder.Free(proc) - var hb HashmapBuilder - hb.setBudget(generation) - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - ctx, cancel := context.WithCancelCause(proc.Ctx) - process.ReplacePipelineCtx(proc, ctx, cancel) - require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) + expr := newExpr(0, types.T_int32.ToType()) + expr.GetCol().RelPos = 1 + require.NoError(t, builder.Prepare([]*plan.Expr{expr}, -1, -1, nil, proc)) - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 10_000, proc.Mp()) - require.NoError(t, hb.copyBuildBatch(input, proc)) - hb.InputBatchRowCount = input.RowCount() + input := batch.NewWithSize(1) + input.Vecs[0] = vector.NewVec(types.T_int32.ToType()) + require.NoError(t, vector.AppendFixedList( + input.Vecs[0], []int32{0, 0}, nil, proc.Mp(), + )) + input.Vecs[0].GetGrouping().Add(1) + input.SetRowCount(2) + require.NoError(t, builder.CopyBuildBatch(input, proc)) + builder.InputBatchRowCount = input.RowCount() input.Clean(proc.Mp()) - require.Positive(t, generation.Used(), "retained build input must own budget before cancellation") - proc.Cancel(context.Canceled) - err = hb.BuildHashmap(false, false, false, proc) - require.ErrorIs(t, err, context.Canceled) - require.Nil(t, hb.IntHashMap) - require.Nil(t, hb.StrHashMap) - - hb.Reset(proc, true) - hb.Free(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - require.Zero(t, proc.Mp().CurrNB()) + require.NoError(t, builder.BuildHashmap(false, false, false, proc)) + require.Nil(t, builder.IntHashMap) + require.NotNil(t, builder.StrHashMap) + require.Equal(t, uint64(2), builder.StrHashMap.GroupCount()) } -func TestPublishedJoinMapResizeKeepsReservationWithConsumer(t *testing.T) { - const budgetCap = uint64(16 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) +func TestBuildHashMap(t *testing.T) { + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 100, proc.Mp()) - require.NoError(t, hb.copyBuildBatch(input, proc)) - hb.InputBatchRowCount = input.RowCount() - input.Clean(proc.Mp()) - require.NoError(t, hb.BuildHashmap(false, false, false, proc)) + err := hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc) + require.NoError(t, err) - jm := hb.GetJoinMap(proc.Mp()) - require.NotNil(t, jm) - jm.IncRef(1) - usedBeforeResize := generation.Used() - secondGeneration, err := budget.OpenGeneration(2) + inputBatch := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, int(100000), proc.Mp()) + err = hb.CopyBuildBatch(inputBatch, proc) + hb.InputBatchRowCount = inputBatch.RowCount() + inputBatch.Clean(proc.Mp()) require.NoError(t, err) - hb.setBudget(secondGeneration) - require.NoError(t, jm.PreAlloc(100_000)) - require.Greater(t, generation.Used(), usedBeforeResize) - require.Zero(t, secondGeneration.Used(), "published map must retain its original generation") - jm.Free() - require.Zero(t, generation.Used()) - hb.Reset(proc, false) + err = hb.BuildHashmap(false, true, true, proc) + require.NoError(t, err) + require.Less(t, int64(0), hb.GetSize()) + require.Less(t, uint64(0), hb.GetGroupCount()) + hb.Reset(proc, true) + hb.Free(proc) + require.Equal(t, int64(0), proc.Mp().CurrNB()) } -func TestHashmapBuilderAccountedJoinMapDoesNotStackLegacyReservations(t *testing.T) { +func TestHashmapBuilderPhysicalAllocationsChargeOnce(t *testing.T) { const budgetCap = uint64(16 << 20) budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) require.NoError(t, err) @@ -376,7 +239,6 @@ func TestHashmapBuilderAccountedJoinMapDoesNotStackLegacyReservations(t *testing proc.Mp(), ) require.NoError(t, hb.copyBuildBatch(input, proc)) - require.Empty(t, hb.batchReservations) require.NotEmpty(t, hb.Batches.Buf) for _, copied := range hb.Batches.Buf { require.Same(t, hb.batchAllocation, copied.AllocationAccountSelection()) @@ -385,13 +247,8 @@ func TestHashmapBuilderAccountedJoinMapDoesNotStackLegacyReservations(t *testing input.Clean(proc.Mp()) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) - require.Nil(t, hb.mapReservation) require.Positive(t, account.Snapshot().Used) - require.Equal( - t, - account.Snapshot().Used, - generation.Snapshot().AllocationUsed, - ) + require.Equal(t, account.Snapshot().Used, generation.Used()) jm := hb.GetJoinMap(proc.Mp()) require.NotNil(t, jm) @@ -450,7 +307,6 @@ func TestHashmapBuilderAccountedBatchCopyOneByteShortRollsBack(t *testing.T) { hb.cleanBatches(proc) } require.Empty(t, hb.Batches.Buf) - require.Empty(t, hb.batchReservations) require.Zero(t, account.Snapshot().Used) require.Zero(t, generation.Used()) require.NoError(t, op.ClearAllocationAccount(account)) @@ -527,75 +383,6 @@ func TestAccountedJoinMapTransfersBatchesAndGroupSelsToLastConsumer(t *testing.T require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) } -func TestHashMapReservationOwnerRetainsSegmentedGrowthTokens(t *testing.T) { - budget, err := process.NewHashBuildBudget(1<<20, 1<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - initial, err := generation.Reserve(100) - require.NoError(t, err) - owner := &hashMapReservationOwner{tokens: []*process.HashBuildReservation{initial}} - - incremental, err := generation.Reserve(50) - require.NoError(t, err) - (&hashMapResizeReservation{owner: owner, token: incremental}).Commit( - hashtable.ResizePlan{ReuseCurrentBlocks: true}, - ) - require.Equal(t, uint64(150), generation.Used()) - - replacement, err := generation.Reserve(200) - require.NoError(t, err) - (&hashMapResizeReservation{owner: owner, token: replacement}).Commit(hashtable.ResizePlan{}) - require.Equal(t, uint64(200), generation.Used()) - - owner.release() - require.Zero(t, generation.Used()) -} - -func TestBudgetedEmptyJoinMapRejectsUnadmittedAllocationAndResize(t *testing.T) { - for _, tc := range []struct { - name string - keyWidth int - initialBytes uint64 - }{ - {name: "int", keyWidth: 4, initialBytes: hashtable.Int64HashMapInitialAllocationBytes()}, - {name: "string", keyWidth: 128, initialBytes: hashtable.StringHashMapInitialAllocationBytes()}, - } { - t.Run(tc.name, func(t *testing.T) { - mp := mpool.MustNewZero() - - tooSmall := process.MustNewHashBuildBudget(tc.initialBytes-1, tc.initialBytes-1) - tooSmallGeneration, err := tooSmall.OpenGeneration(1) - require.NoError(t, err) - jm, err := NewBudgetedEmptyJoinMap(tc.keyWidth, tooSmallGeneration, mp) - require.Nil(t, jm) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Zero(t, tooSmallGeneration.Used()) - require.Zero(t, mp.CurrNB()) - - exact := process.MustNewHashBuildBudget(tc.initialBytes, tc.initialBytes) - generation, err := exact.OpenGeneration(2) - require.NoError(t, err) - jm, err = NewBudgetedEmptyJoinMap(tc.keyWidth, generation, mp) - require.NoError(t, err) - require.Equal(t, tc.initialBytes, generation.Used()) - require.Equal(t, int64(tc.initialBytes), mp.CurrNB()) - - err = jm.PreAlloc(10_000) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Equal(t, tc.initialBytes, generation.Used(), - "rejected growth must roll back its temporary reservation") - - jm.Free() - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - require.Zero(t, mp.CurrNB()) - }) - } -} - func TestAccountedEmptyJoinMapUsesPhysicalAllocationAsSoleCharge(t *testing.T) { for _, tc := range []struct { name string @@ -622,673 +409,73 @@ func TestAccountedEmptyJoinMapUsesPhysicalAllocationAsSoleCharge(t *testing.T) { expectedInitial := tc.initialBytes + descriptorBytes require.Equal(t, expectedInitial, account.Snapshot().Used) require.Equal(t, expectedInitial, generation.Used()) - require.Equal(t, expectedInitial, generation.Snapshot().AllocationUsed) require.NoError(t, jm.PreAlloc(10_000)) require.Equal(t, account.Snapshot().Used, generation.Used()) - require.Equal( - t, - account.Snapshot().Used, - generation.Snapshot().AllocationUsed, - ) + require.Equal(t, account.Snapshot().Used, generation.Used()) jm.Free() require.Zero(t, account.Snapshot().Used) require.Zero(t, generation.Used()) terminal, first, err := registry.CompleteTerminal(account) require.NoError(t, err) - require.True(t, first) - require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) - }) - } -} - -func TestAccountedEmptyJoinMapInitialFailureRollsBackController(t *testing.T) { - initial := hashtable.Int64HashMapInitialAllocationBytes() + - hashtable.HashMapBlockDescriptorBytes() - budget := process.MustNewHashBuildBudget(initial, initial) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - registry, err := mpool.NewAllocationAccountRegistry(1, 2) - require.NoError(t, err) - account, err := registry.OpenWithController(initial-1, generation) - require.NoError(t, err) - mp := mpool.MustNewZero() - - jm, err := NewAccountedEmptyJoinMap(4, account, mp) - require.Nil(t, jm) - require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) - require.Zero(t, account.Snapshot().Used) - require.Zero(t, generation.Used()) - require.Zero(t, registry.LiveAllocationMetadata()) - require.Zero(t, mp.CurrNB()) - _, _, err = registry.CompleteTerminal(account) - require.NoError(t, err) -} - -func TestAccountedJoinMapLateFreeKeepsOriginalGeneration(t *testing.T) { - const capBytes = uint64(64 << 20) - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - firstGeneration, err := budget.OpenGeneration(1) - require.NoError(t, err) - secondGeneration, err := budget.OpenGeneration(2) - require.NoError(t, err) - registry, err := mpool.NewAllocationAccountRegistry(2, 16) - require.NoError(t, err) - firstAccount, err := registry.OpenWithController(capBytes, firstGeneration) - require.NoError(t, err) - mp := mpool.MustNewZero() - jm, err := NewAccountedEmptyJoinMap(4, firstAccount, mp) - require.NoError(t, err) - firstUsed := firstGeneration.Used() - require.Positive(t, firstUsed) - - secondAccount, err := registry.OpenWithController(capBytes, secondGeneration) - require.NoError(t, err) - require.Zero(t, secondGeneration.Used()) - jm.Free() - require.Zero(t, firstGeneration.Used()) - require.Zero(t, firstAccount.Snapshot().Used) - require.Zero(t, secondGeneration.Used()) - - _, _, err = registry.CompleteTerminal(firstAccount) - require.NoError(t, err) - _, _, err = registry.CompleteTerminal(secondAccount) - require.NoError(t, err) -} - -func TestCopyBuildBatchBudgetsSmallIngressAfterFullBatches(t *testing.T) { - const budgetCap = uint64(32 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - for _, rows := range []int{colexec.DefaultBatchSize, colexec.DefaultBatchSize} { - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, rows, proc.Mp()) - require.NoError(t, hb.copyBuildBatch(input, proc)) - input.Clean(proc.Mp()) - } - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 100, proc.Mp()) - projected, err := hb.projectedBatchCopyBytes(input) - require.NoError(t, err) - destination, err := projectedNewDestinationBytes(input, 0, input.RowCount()) - require.NoError(t, err) - metadata, ok := retainedMetadataAllowance(input) - require.True(t, ok) - require.Equal(t, destination+metadata+uint64(64<<10), projected, - "a small ingress must not be projected as a full 8192-row allocation") - require.NoError(t, hb.copyBuildBatch(input, proc)) - input.Clean(proc.Mp()) - - require.Len(t, hb.Batches.Buf, 3) - require.Equal(t, 100, hb.Batches.Buf[2].RowCount()) - hb.FreeHashMapAndBatches(proc) - require.Zero(t, generation.Used()) -} - -func TestCopyBuildBatchBudgetsPartialTailGrowth(t *testing.T) { - const budgetCap = uint64(32 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - for range 1000 { - // Deep spill partitions contain many tiny records. They coalesce into - // one physical batch whose vector capacity grows geometrically. - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 7, proc.Mp()) - require.NoError(t, hb.copyBuildBatch(input, proc)) - input.Clean(proc.Mp()) - } - require.Len(t, hb.Batches.Buf, 1) - require.Equal(t, 7000, hb.Batches.Buf[0].RowCount()) - hb.FreeHashMapAndBatches(proc) - require.Zero(t, generation.Used()) -} - -func TestCopyBuildBatchBudgetsWideVarcharPartialTailReplacement(t *testing.T) { - const budgetCap = uint64(128 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - values := make([]string, colexec.DefaultBatchSize/2) - for i := range values { - values[i] = strings.Repeat("x", 1024) - } - input := batch.NewWithSize(1) - input.Vecs[0] = testutil.MakeVarcharVector(values, nil, proc.Mp()) - input.SetRowCount(len(values)) - defer input.Clean(proc.Mp()) - - var hb HashmapBuilder - hb.setBudget(generation) - defer hb.FreeHashMapAndBatches(proc) - require.NoError(t, hb.copyBuildBatch(input, proc)) - require.NoError(t, hb.copyBuildBatch(input, proc)) - require.Len(t, hb.Batches.Buf, 1) - require.Equal(t, colexec.DefaultBatchSize, hb.Batches.Buf[0].RowCount()) -} - -func TestProjectedPartialTailReplacementMatchesUnionBatch(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - tail := batch.NewWithSize(2) - tail.Vecs[0] = testutil.MakeInt32Vector([]int32{1}, nil, proc.Mp()) - tail.Vecs[1] = testutil.MakeVarcharVector([]string{strings.Repeat("a", 1024)}, nil, proc.Mp()) - tail.SetRowCount(1) - defer tail.Clean(proc.Mp()) - - src := batch.NewWithSize(2) - src.Vecs[0] = testutil.MakeInt32Vector([]int32{2, 3}, nil, proc.Mp()) - constVec, err := vector.NewConstBytes( - types.T_varchar.ToType(), - []byte(strings.Repeat("b", 1024)), - 2, - proc.Mp(), - ) - require.NoError(t, err) - src.Vecs[1] = constVec - src.SetRowCount(2) - defer src.Clean(proc.Mp()) - - peak, retained, err := projectedPartialTailReplacementBytes(tail, src, src.RowCount()) - require.NoError(t, err) - require.GreaterOrEqual(t, peak, retained) - before := tail.Allocated() - for i := range tail.Vecs { - require.NoError(t, tail.Vecs[i].UnionBatch(src.Vecs[i], 0, src.RowCount(), nil, proc.Mp())) - } - tail.AddRowCount(src.RowCount()) - require.Equal(t, uint64(tail.Allocated()-before), retained) - - inline := batch.NewWithSize(1) - inline.Vecs[0] = testutil.MakeVarcharVector([]string{"small"}, nil, proc.Mp()) - inline.SetRowCount(1) - defer inline.Clean(proc.Mp()) - preallocated, err := proc.NewBatchFromSrc(inline, colexec.DefaultBatchSize) - require.NoError(t, err) - defer preallocated.Clean(proc.Mp()) - require.NoError(t, preallocated.Vecs[0].UnionBatch(inline.Vecs[0], 0, 1, nil, proc.Mp())) - preallocated.AddRowCount(1) - peak, retained, err = projectedPartialTailReplacementBytes(preallocated, inline, 1) - require.NoError(t, err) - require.Zero(t, peak) - require.Zero(t, retained) -} - -func TestCopyBuildBatchBudgetsPartialTailWithRemainder(t *testing.T) { - const budgetCap = uint64(16 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - var hb HashmapBuilder - hb.setBudget(generation) - defer hb.FreeHashMapAndBatches(proc) - for _, rows := range []int{ - colexec.DefaultBatchSize, - colexec.DefaultBatchSize, - colexec.DefaultBatchSize - 1, - 2, - } { - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, rows, proc.Mp()) - require.NoError(t, hb.copyBuildBatch(input, proc)) - input.Clean(proc.Mp()) - } - require.Len(t, hb.Batches.Buf, 4) - require.Equal(t, 1, hb.Batches.Buf[3].RowCount()) -} - -func TestBatchCopyAllocatedDeltaTracksOnlyChangedSuffix(t *testing.T) { - tests := []struct { - name string - retainedRows int - ingressRows int - }{ - {name: "empty/multiple-destinations", ingressRows: 2*colexec.DefaultBatchSize + 7}, - {name: "full-tail/append", retainedRows: colexec.DefaultBatchSize, ingressRows: colexec.DefaultBatchSize}, - {name: "partial-tail/grow-and-append", retainedRows: 7000, ingressRows: 2000}, - {name: "partial-tail/full-ingress-swap", retainedRows: 7, ingressRows: colexec.DefaultBatchSize}, - } - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - var batches colexec.Batches - defer batches.Clean(proc.Mp()) - if tc.retainedRows > 0 { - retained := testutil.NewBatch( - []types.Type{types.T_int32.ToType()}, true, tc.retainedRows, proc.Mp()) - defer retained.Clean(proc.Mp()) - require.NoError(t, batches.CopyIntoBatches(retained, proc)) - } - before := batchesAllocated(batches.Buf) - snapshot, err := snapshotBatchCopyAllocation(batches.Buf) - require.NoError(t, err) - - ingress := testutil.NewBatch( - []types.Type{types.T_int32.ToType()}, true, tc.ingressRows, proc.Mp()) - defer ingress.Clean(proc.Mp()) - require.NoError(t, batches.CopyIntoBatches(ingress, proc)) - - after := batchesAllocated(batches.Buf) - require.GreaterOrEqual(t, after, before) - delta, err := batchCopyAllocatedDelta(batches.Buf, snapshot) - require.NoError(t, err) - require.Equal(t, after-before, delta) - }) - } -} - -func TestBatchCopyAllocatedDeltaRejectsLostTail(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - retained := testutil.NewBatch( - []types.Type{types.T_int32.ToType()}, true, 7, proc.Mp()) - defer retained.Clean(proc.Mp()) - snapshot, err := snapshotBatchCopyAllocation([]*batch.Batch{retained}) - require.NoError(t, err) - - replacement := testutil.NewBatch( - []types.Type{types.T_int32.ToType()}, true, 7, proc.Mp()) - defer replacement.Clean(proc.Mp()) - _, err = batchCopyAllocatedDelta([]*batch.Batch{replacement}, snapshot) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) -} - -func TestProjectedPartialTailReplacementRejectsInvalidInputs(t *testing.T) { - _, _, err := projectedPartialTailReplacementBytes(nil, nil, -1) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - _, err = (&HashmapBuilder{}).projectedBatchCopyBytes(nil) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - - tail := batch.NewOffHeapWithSize(1) - src := batch.NewOffHeapWithSize(1) - _, _, err = projectedPartialTailReplacementBytes(tail, src, 1) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - src.Vecs[0] = testutil.MakeVarcharVector([]string{strings.Repeat("x", 32)}, nil, proc.Mp()) - src.SetRowCount(1) - defer src.Clean(proc.Mp()) - tail.Vecs[0] = vector.NewOffHeapVecWithType(types.T_varchar.ToType()) - defer tail.Clean(proc.Mp()) - _, _, err = projectedPartialTailReplacementBytes(tail, src, 2) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - - invalidTail := batch.NewOffHeapWithSize(1) - invalidTail.SetRowCount(-1) - hb := HashmapBuilder{} - hb.Batches.Buf = []*batch.Batch{invalidTail} - _, err = hb.projectedBatchCopyBytes(src) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - - mismatchedTail := batch.NewOffHeapWithSize(0) - hb.Batches.Buf = []*batch.Batch{mismatchedTail} - _, err = hb.projectedBatchCopyBytes(src) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) -} - -func TestCopyBuildBatchUsesProjectedDestinationCapacity(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - values := make([]string, 4096) - for i := range values { - values[i] = strings.Repeat("x", 1024) - } - input := batch.NewWithSize(1) - input.Vecs[0] = testutil.MakeVarcharVector(values, nil, proc.Mp()) - input.SetRowCount(len(values)) - defer input.Clean(proc.Mp()) - - var hb HashmapBuilder - projected, err := hb.projectedBatchCopyBytes(input) - require.NoError(t, err) - destination, err := projectedNewDestinationBytes(input, 0, input.RowCount()) - require.NoError(t, err) - metadata, ok := retainedMetadataAllowance(input) - require.True(t, ok) - const wantSlack = uint64(64 << 10) - require.Equal(t, destination+metadata+wantSlack, projected) - - budget := process.MustNewHashBuildBudget(projected, projected) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - hb.setBudget(generation) - require.NoError(t, hb.copyBuildBatch(input, proc)) - require.Equal(t, projected, generation.Peak()) - hb.FreeHashMapAndBatches(proc) - require.Zero(t, generation.Used()) -} - -func TestCopyBuildBatchSplitsLargeIngressWithinProjection(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - const rows = 50_000 - input := testutil.NewBatch([]types.Type{types.T_uuid.ToType()}, true, rows, proc.Mp()) - part, err := vector.NewConstFixed[int32](types.T_int32.ToType(), 1, rows, proc.Mp()) - require.NoError(t, err) - input.Vecs = append(input.Vecs, part) - defer input.Clean(proc.Mp()) - - budget := process.MustNewHashBuildBudget(1<<30, 1<<30) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) - require.NoError(t, hb.copyBuildBatch(input, proc)) - require.NoError(t, hb.copyBuildBatch(input, proc)) - require.Equal(t, 2*rows, hb.Batches.RowCount()) - require.Len(t, hb.Batches.Buf, 13) - hb.FreeHashMapAndBatches(proc) - require.Zero(t, generation.Used()) -} - -func TestCopyBuildBatchSplitsLargeConstVarcharIngressWithinProjection(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - const rows = 50_000 - input := batch.NewWithSize(1) - value := make([]byte, 1<<20) - vec, err := vector.NewConstBytes(types.T_varchar.ToType(), value, rows, proc.Mp()) - require.NoError(t, err) - input.Vecs[0] = vec - input.SetRowCount(rows) - defer input.Clean(proc.Mp()) - - budget := process.MustNewHashBuildBudget(1<<30, 1<<30) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) - require.NoError(t, hb.copyBuildBatch(input, proc)) - require.NoError(t, hb.copyBuildBatch(input, proc)) - require.Equal(t, 2*rows, hb.Batches.RowCount()) - require.Len(t, hb.Batches.Buf, 13) - hb.FreeHashMapAndBatches(proc) - require.Zero(t, generation.Used()) -} - -func TestCopyBuildBatchManyExactSegmentsAvoidsFalseAdmission(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - const rows = 64 * colexec.DefaultBatchSize - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, rows, proc.Mp()) - defer input.Clean(proc.Mp()) - - destination, err := projectedNewDestinationBytes(input, 0, input.RowCount()) - require.NoError(t, err) - metadata, ok := retainedMetadataAllowance(input) - require.True(t, ok) - budgetCap := 2*(destination+metadata) + uint64(512<<10) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) - require.NoError(t, hb.copyBuildBatch(input, proc)) - require.NoError(t, hb.copyBuildBatch(input, proc)) - require.LessOrEqual(t, generation.Used(), budgetCap) - hb.FreeHashMapAndBatches(proc) - require.Zero(t, generation.Used()) -} - -func TestCopyBuildBatchSharedVarlenaRejectsBeforeAllocation(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - const rows = 50_000 - value := make([]byte, 1<<10) - constVec, err := vector.NewConstBytes(types.T_varchar.ToType(), value, rows, proc.Mp()) - require.NoError(t, err) - defer constVec.Free(proc.Mp()) - - flat := vector.NewOffHeapVecWithType(types.T_varchar.ToType()) - require.NoError(t, flat.UnionBatch(constVec, 0, rows, nil, proc.Mp())) - require.False(t, flat.IsConst()) - require.Equal(t, len(value), len(flat.GetArea())) - input := batch.NewWithSize(1) - input.Vecs[0] = flat - input.SetRowCount(rows) - defer input.Clean(proc.Mp()) - - const budgetCap = uint64(10 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) - projected, err := hb.projectedBatchCopyBytes(input) - require.NoError(t, err) - require.GreaterOrEqual(t, projected, uint64(rows*len(value))) - err = hb.copyBuildBatch(input, proc) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Empty(t, hb.Batches.Buf) - require.Zero(t, generation.Used()) -} - -func TestCopyBuildBatchWholeSharedVarlenaAvoidsFalseAdmission(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - const rows = colexec.DefaultBatchSize - value := make([]byte, 1<<10) - constVec, err := vector.NewConstBytes(types.T_varchar.ToType(), value, rows, proc.Mp()) - require.NoError(t, err) - defer constVec.Free(proc.Mp()) - - flat := vector.NewOffHeapVecWithType(types.T_varchar.ToType()) - require.NoError(t, flat.UnionBatch(constVec, 0, rows, nil, proc.Mp())) - require.False(t, flat.IsConst()) - input := batch.NewWithSize(1) - input.Vecs[0] = flat - input.SetRowCount(rows) - defer input.Clean(proc.Mp()) - - const budgetCap = uint64(2 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) - require.NoError(t, hb.copyBuildBatch(input, proc)) - require.LessOrEqual(t, generation.Used(), budgetCap) - hb.FreeHashMapAndBatches(proc) - require.Zero(t, generation.Used()) -} - -func TestReserveBuildAuxChargesOneRetainedCopy(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, colexec.DefaultBatchSize, proc.Mp()) - defer input.Clean(proc.Mp()) - - var hb HashmapBuilder - hb.Batches.Buf = []*batch.Batch{input} - hb.InputBatchRowCount = input.RowCount() - retained := batchesAllocated(hb.Batches.Buf) - const iteratorScratch = uint64(640 << 10) - want := retained + (retained+3)/4 + uint64(input.RowCount())*64 + iteratorScratch - - budget := process.MustNewHashBuildBudget(want, want) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - hb.setBudget(generation) - require.NoError(t, hb.reserveBuildAux(true, false)) - require.Equal(t, want, generation.Used()) - hb.releaseReservations() - require.Zero(t, generation.Used()) - // The batch belongs to the test rather than batchReservations. - hb.Batches.Buf = nil + require.True(t, first) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) + }) + } } -func TestReserveBuildAuxExactDoesNotDuplicatePhysicalOrHeadroomOwners(t *testing.T) { - const rows = 10_000 - const want = uint64(1) - budget := process.MustNewHashBuildBudget(want, want) +func TestAccountedEmptyJoinMapInitialFailureRollsBackController(t *testing.T) { + initial := hashtable.Int64HashMapInitialAllocationBytes() + + hashtable.HashMapBlockDescriptorBytes() + budget := process.MustNewHashBuildBudget(initial, initial) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - registry, err := mpool.NewAllocationAccountRegistry(1, 8) + registry, err := mpool.NewAllocationAccountRegistry(1, 2) require.NoError(t, err) - account, err := registry.OpenWithController(want, generation) + account, err := registry.OpenWithController(initial-1, generation) require.NoError(t, err) - var op HashBuild - op.NeedHashMap = true - require.NoError(t, op.SetAllocationAccount(account)) - hb := &op.ctr.hashmapBuilder - hb.InputBatchRowCount = rows - hb.setBudget(generation) + mp := mpool.MustNewZero() - require.NoError(t, hb.reserveBuildAux(false, true)) - require.Zero(t, generation.Used()) - require.Zero(t, generation.Snapshot().AllocationUsed) - hb.releaseReservations() + jm, err := NewAccountedEmptyJoinMap(4, account, mp) + require.Nil(t, jm) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Zero(t, account.Snapshot().Used) require.Zero(t, generation.Used()) - require.NoError(t, op.ClearAllocationAccount(account)) + require.Zero(t, registry.LiveAllocationMetadata()) + require.Zero(t, mp.CurrNB()) _, _, err = registry.CompleteTerminal(account) require.NoError(t, err) } -func TestReserveUniqueAppendOverlapChargesReplacedCapacity(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - values := make([]string, 4_096) - for i := range values { - values[i] = strings.Repeat("x", 1_024) - } - dst := testutil.MakeVarcharVector(values, nil, proc.Mp()) - defer dst.Free(proc.Mp()) - extraArea := cap(dst.GetArea()) - len(dst.GetArea()) + 1 - src := testutil.MakeVarcharVector([]string{strings.Repeat("y", extraArea)}, nil, proc.Mp()) - defer src.Free(proc.Mp()) - - want := uint64(cap(dst.GetData()) + cap(dst.GetArea())) - const budgetCap = uint64(64 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - hb := HashmapBuilder{ - budget: generation, - UniqueJoinKeys: []*vector.Vector{dst}, - } - require.NoError(t, hb.reserveBuildAux(true)) - areaBytes, err := uniqueAppendAreaBytes(src, 0, 1, nil) - require.NoError(t, err) - token, err := hb.reserveUniqueAppendOverlap(dst, 1, areaBytes) - require.NoError(t, err) - require.NotNil(t, token) - require.Equal(t, want, token.Size()) - token.Release() - require.Equal(t, hb.auxReservation.Size(), generation.Used()) - hb.releaseReservations() - require.Zero(t, generation.Used()) - - largeValue := strings.Repeat("z", 100) - selected := testutil.MakeVarcharVector([]string{"a", largeValue}, nil, proc.Mp()) - defer selected.Free(proc.Mp()) - selectedArea, err := uniqueAppendAreaBytes(selected, 0, 1, []int64{0}) +func TestAccountedJoinMapLateFreeKeepsOriginalGeneration(t *testing.T) { + const capBytes = uint64(64 << 20) + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + firstGeneration, err := budget.OpenGeneration(1) require.NoError(t, err) - require.Zero(t, selectedArea) - selectedArea, err = uniqueAppendAreaBytes(selected, 0, 1, []int64{1}) + secondGeneration, err := budget.OpenGeneration(2) require.NoError(t, err) - require.Equal(t, len(largeValue), selectedArea) - _, err = uniqueAppendAreaBytes(selected, -1, 1, nil) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - _, err = uniqueAppendAreaBytes(selected, 0, 1, []int64{int64(selected.Length())}) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - _, err = uniqueAppendAreaBytes(selected, 0, 2, []int64{0}) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - - fixed := testutil.MakeInt32Vector([]int32{1}, nil, proc.Mp()) - defer fixed.Free(proc.Mp()) - selectedArea, err = uniqueAppendAreaBytes(fixed, math.MaxInt, math.MaxInt, nil) + registry, err := mpool.NewAllocationAccountRegistry(2, 16) require.NoError(t, err) - require.Zero(t, selectedArea) - - constValue, err := vector.NewConstBytes(types.T_varchar.ToType(), []byte(largeValue), 2, proc.Mp()) + firstAccount, err := registry.OpenWithController(capBytes, firstGeneration) require.NoError(t, err) - defer constValue.Free(proc.Mp()) - selectedArea, err = uniqueAppendAreaBytes(constValue, 0, 2, nil) + mp := mpool.MustNewZero() + jm, err := NewAccountedEmptyJoinMap(4, firstAccount, mp) require.NoError(t, err) - require.Equal(t, 2*len(largeValue), selectedArea) + firstUsed := firstGeneration.Used() + require.Positive(t, firstUsed) - noBudget := HashmapBuilder{} - token, err = noBudget.reserveUniqueAppendOverlap(dst, 1, 1) + secondAccount, err := registry.OpenWithController(capBytes, secondGeneration) require.NoError(t, err) - require.Nil(t, token) - token, err = hb.reserveUniqueAppendOverlap(nil, 1, 1) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - token, err = hb.reserveUniqueAppendOverlap(dst, -1, 1) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - token, err = hb.reserveUniqueAppendOverlap(dst, 1, -1) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) -} - -func TestUniqueAppendBudgetIncludesDeadAreaCopiedByUnionBatch(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() + require.Zero(t, secondGeneration.Used()) + jm.Free() + require.Zero(t, firstGeneration.Used()) + require.Zero(t, firstAccount.Snapshot().Used) + require.Zero(t, secondGeneration.Used()) - src := testutil.MakeVarcharVector( - []string{"inline", strings.Repeat("d", 128<<10)}, nil, proc.Mp()) - defer src.Free(proc.Mp()) - // SetLength leaves the second value's area allocation behind. The sole live - // row is inline, but UnionBatch's whole-vector fast path copies all of area. - src.SetLength(1) - liveArea, err := uniqueAppendAreaBytes(src, 0, 1, nil) - require.NoError(t, err) - require.Zero(t, liveArea) - unionArea, err := unionBatchAreaBytes(src, 0, 1) + _, _, err = registry.CompleteTerminal(firstAccount) require.NoError(t, err) - require.Equal(t, len(src.GetArea()), unionArea) - require.Greater(t, unionArea, 0) - - dst := vector.NewOffHeapVecWithType(types.T_varchar.ToType()) - defer dst.Free(proc.Mp()) - const mandatoryAux = uint64(640 << 10) - budget := process.MustNewHashBuildBudget(mandatoryAux, mandatoryAux) - generation, err := budget.OpenGeneration(1) + _, _, err = registry.CompleteTerminal(secondAccount) require.NoError(t, err) - hb := HashmapBuilder{ - budget: generation, - UniqueJoinKeys: []*vector.Vector{dst}, - } - require.NoError(t, hb.reserveBuildAux(true)) - - _, err = hb.reserveUniqueAppendOverlap(dst, 1, unionArea) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Zero(t, dst.Length()) - require.Zero(t, dst.Allocated(), - "admission must fail before UnionBatch allocates copied dead area") - - hb.releaseReservations() - require.Zero(t, generation.Used()) - generation.Close() } func TestAccountedRuntimeFilterUniqueKeysDegradeWithoutFailingHashBuild(t *testing.T) { @@ -1345,138 +532,7 @@ func TestAccountedRuntimeFilterUniqueKeysDegradeWithoutFailingHashBuild(t *testi require.LessOrEqual(t, constrained.Peak, baseline.Peak) } -func TestCleanCopiedBatchReleasesCoalescedIngressReservations(t *testing.T) { - const budgetCap = uint64(4 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - for range 2 { - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, colexec.DefaultBatchSize/2, proc.Mp()) - require.NoError(t, hb.copyBuildBatch(input, proc)) - input.Clean(proc.Mp()) - } - require.Len(t, hb.Batches.Buf, 1, "small ingress batches should coalesce") - require.Len(t, hb.batchReservations, 2, "reservations follow ingress, not physical batches") - require.Greater(t, generation.Used(), uint64(0)) - require.NoError(t, hb.CleanCopiedBatchAt(0, proc)) - require.Empty(t, hb.Batches.Buf) - require.Empty(t, hb.batchReservations) - require.Zero(t, generation.Used()) -} - -func TestDrainCopiedBatchesReleasesBeforeSubsequentAdmission(t *testing.T) { - const budgetCap = uint64(4 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - for range 2 { - input := testutil.NewBatch( - []types.Type{types.T_int32.ToType()}, - true, - colexec.DefaultBatchSize/2, - proc.Mp(), - ) - require.NoError(t, hb.copyBuildBatch(input, proc)) - input.Clean(proc.Mp()) - } - require.Len(t, hb.Batches.Buf, 1, "small ingress batches should coalesce") - require.Len(t, hb.batchReservations, 2, "reservations follow ingress, not physical batches") - - visits := 0 - require.NoError(t, hb.DrainCopiedBatches(proc, func(bat *batch.Batch) error { - visits++ - require.NotNil(t, bat) - require.Positive(t, generation.Used(), "physical batch must remain charged while visited") - return nil - })) - require.Equal(t, 1, visits) - require.Empty(t, hb.Batches.Buf) - require.Empty(t, hb.batchReservations) - require.Zero(t, generation.Used(), "the final physical batch must release every coalesced ingress charge") - - // Model the expression/scatter/read reservation that follows a re-spill - // drain. It can consume the complete cap only after stale batch ownership - // has been removed from the ledger. - next, err := generation.Reserve(budgetCap) - require.NoError(t, err) - require.True(t, next.Release()) - require.Zero(t, generation.Used()) -} - -func TestDrainCopiedBatchesVisitFailureRetainsOwnership(t *testing.T) { - const budgetCap = uint64(4 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - for _, rows := range []int{ - colexec.DefaultBatchSize / 2, - colexec.DefaultBatchSize / 2, - 1024, - } { - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, rows, proc.Mp()) - require.NoError(t, hb.copyBuildBatch(input, proc)) - input.Clean(proc.Mp()) - } - require.Len(t, hb.Batches.Buf, 2) - require.Len(t, hb.batchReservations, 3) - - wantErr := errors.New("visit failed") - visits := 0 - require.ErrorIs(t, hb.DrainCopiedBatches(proc, func(*batch.Batch) error { - visits++ - if visits == 1 { - return nil - } - return wantErr - }), wantErr) - require.Equal(t, 2, visits) - require.Len(t, hb.Batches.Buf, 1, "the failed current batch remains owned after prior batches drain") - require.Len(t, hb.batchReservations, 3, - "coalesced ingress reservations stay conservative until terminal cleanup") - require.Positive(t, generation.Used()) - - hb.FreeHashMapAndBatches(proc) - require.Zero(t, generation.Used()) -} - -func TestSpillExpressionHashKeyUsesBoundedAdmission(t *testing.T) { - var ctr container - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget, err := process.NewHashBuildBudget(1<<20, 1<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - ctr.hashmapBuilder.setBudget(generation) - expr := makeExpressionLeaseTestExpr(t, proc) - _, err = ctr.initSpillExprExecs(proc, []*plan.Expr{expr}) - require.NoError(t, err) - require.NoError(t, ctr.spillExprLease.Run(proc, 8192, func(_ int) error { return nil })) - require.Positive(t, ctr.spillExprLease.Reserved()) - require.Equal(t, ctr.spillExprLease.Reserved(), generation.Used()) - ctr.freeSpillExprExecs() - require.Zero(t, generation.Used()) -} - -func TestSpillExpressionUsesExactAccountForClosedKey(t *testing.T) { +func TestSpillExpressionTemporaryIsOutsideRetainedAccount(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() budget := process.MustNewHashBuildBudget(16<<20, 16<<20) @@ -1494,8 +550,6 @@ func TestSpillExpressionUsesExactAccountForClosedKey(t *testing.T) { expr := makeIssue26454ConcatKey(t, proc) executors, err := ctr.initSpillExprExecs(proc, []*plan.Expr{expr}) require.NoError(t, err) - require.True(t, ctr.spillExprAccounted) - require.Nil(t, ctr.spillExprLease) input := batch.NewWithSize(2) input.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp()) input.Vecs[1] = testutil.MakeInt32Vector([]int32{3, 4}, nil, proc.Mp()) @@ -1504,8 +558,8 @@ func TestSpillExpressionUsesExactAccountForClosedKey(t *testing.T) { result, err := executors[0].Eval(proc, []*batch.Batch{input}, nil) require.NoError(t, err) require.Equal(t, []string{"1-3", "2-4"}, vector.InefficientMustStrCol(result)) - require.Equal(t, generation.Used(), generation.Snapshot().AllocationUsed) - require.Positive(t, account.Snapshot().Used) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) ctr.freeSpillExprExecs() require.Zero(t, account.Snapshot().Used) @@ -1556,10 +610,6 @@ func TestIssue26454ExpressionKeyBuildUsesActualCapacity(t *testing.T) { }, } { t.Run(tc.name, func(t *testing.T) { - legacyPeak, err := expressionVectorPeak(proc, tc.expr, 10_000, false) - require.NoError(t, err) - require.Greater(t, legacyPeak, capBytes) - budget := process.MustNewHashBuildBudget(capBytes, capBytes) generation, err := budget.OpenGeneration(1) require.NoError(t, err) @@ -1573,7 +623,6 @@ func TestIssue26454ExpressionKeyBuildUsesActualCapacity(t *testing.T) { hb := &op.ctr.hashmapBuilder hb.setBudget(generation) require.NoError(t, hb.Prepare([]*plan.Expr{tc.expr}, -1, -1, nil, proc)) - require.Nil(t, hb.expressionLease) input := tc.input() require.NoError(t, hb.copyBuildBatch(input, proc)) hb.InputBatchRowCount = input.RowCount() @@ -1596,75 +645,6 @@ func TestIssue26454ExpressionKeyBuildUsesActualCapacity(t *testing.T) { } } -func TestExpressionHashKeyReservesDeclaredPeakBeforeEval(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget, err := process.NewHashBuildBudget(96<<10, 96<<10) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) - require.NoError(t, hb.Prepare([]*plan.Expr{{ - Typ: plan.Type{Id: int32(types.T_varchar), Width: types.MaxVarcharLen}, - Expr: &plan.Expr_F{F: &plan.Function{}}, - }}, -1, -1, nil, proc)) - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 1, proc.Mp()) - require.NoError(t, hb.copyBuildBatch(input, proc)) - hb.InputBatchRowCount = input.RowCount() - input.Clean(proc.Mp()) - err = hb.BuildHashmap(false, true, false, proc) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - hb.Free(proc) - require.Zero(t, generation.Used()) -} - -func TestExpressionHashKeyAcceptsCastTargetType(t *testing.T) { - proc := testutil.NewProcess(t) - defer proc.Free() - expr := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_F{F: &plan.Function{Args: []*plan.Expr{ - newExpr(0, types.T_int64.ToType()), - { - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_T{T: &plan.TargetType{}}, - }, - }}}, - } - - peak, err := expressionVectorPeak(proc, expr, 1024, false) - require.NoError(t, err) - require.Equal(t, uint64(204800), peak, "charge the target-type and cast result vectors") -} - -func TestPreparedParamExpressionPeakUsesConstCardinality(t *testing.T) { - proc := testutil.NewProcess(t) - defer proc.Free() - params := vector.NewVec(types.T_text.ToType()) - defer params.Free(proc.Mp()) - proc.SetPrepareParams(params) - require.NoError(t, vector.AppendBytes(params, []byte("prepared"), false, proc.Mp())) - require.NoError(t, vector.AppendBytes(params, nil, true, proc.Mp())) - - paramExpr := func(pos int32) *plan.Expr { - return &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_text), Width: types.MaxVarcharLen}, - Expr: &plan.Expr_P{P: &plan.ParamRef{Pos: pos}}, - } - } - - peakOne, err := expressionVectorPeak(proc, paramExpr(0), 1, false) - require.NoError(t, err) - peakBatch, err := expressionVectorPeak(proc, paramExpr(0), colexec.DefaultBatchSize, false) - require.NoError(t, err) - require.Equal(t, peakOne, peakBatch, "const parameter admission must not scale with input rows") - peakNull, err := expressionVectorPeak(proc, paramExpr(1), colexec.DefaultBatchSize, false) - require.NoError(t, err) - require.Equal(t, peakOne, peakNull, "null parameter keeps the declared one-row type bound") -} - func TestPreparedParamExpressionExecutorRemainsConst(t *testing.T) { for _, tc := range []struct { name string @@ -1700,93 +680,14 @@ func TestPreparedParamExpressionExecutorRemainsConst(t *testing.T) { } } -func TestPreparedParamExpressionPeakNestedFunctionCardinality(t *testing.T) { - proc := testutil.NewProcess(t) - defer proc.Free() - params := vector.NewVec(types.T_text.ToType()) - defer params.Free(proc.Mp()) - require.NoError(t, vector.AppendBytes(params, []byte("prepared"), false, proc.Mp())) - proc.SetPrepareParams(params) - - param := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_text), Width: types.MaxVarcharLen}, - Expr: &plan.Expr_P{P: &plan.ParamRef{Pos: 0}}, - } - cast := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int64)}, - Expr: &plan.Expr_F{F: &plan.Function{Args: []*plan.Expr{ - param, - {Typ: plan.Type{Id: int32(types.T_int64)}, Expr: &plan.Expr_T{T: &plan.TargetType{}}}, - }}}, - } - modulo := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int64)}, - Expr: &plan.Expr_F{F: &plan.Function{Args: []*plan.Expr{ - {Typ: plan.Type{Id: int32(types.T_int64)}, Expr: &plan.Expr_Col{Col: &plan.ColRef{RelPos: 0, ColPos: 0}}}, - cast, - }}}, - } - - paramTotal, paramOutput, err := expressionTreePeak(proc, param, colexec.DefaultBatchSize) - require.NoError(t, err) - paramOne, _, err := expressionTreePeak(proc, param, 1) - require.NoError(t, err) - require.Equal(t, paramOne, paramTotal) - _, rootOutput, err := expressionTreePeak(proc, modulo, colexec.DefaultBatchSize) - require.NoError(t, err) - rootTypePeak, err := expressionTypePeak(modulo.Typ, colexec.DefaultBatchSize) - require.NoError(t, err) - require.Equal(t, rootTypePeak, rootOutput, "function output remains sized for input rows") - require.Greater(t, paramOutput, uint64(0)) -} - -func TestPreparedParamExpressionPeakRejectsInvalidPosition(t *testing.T) { - proc := testutil.NewProcess(t) - defer proc.Free() - params := vector.NewVec(types.T_text.ToType()) - defer params.Free(proc.Mp()) - proc.SetPrepareParams(params) - require.NoError(t, vector.AppendBytes(params, []byte("prepared"), false, proc.Mp())) - - expr := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_text)}, - Expr: &plan.Expr_P{P: &plan.ParamRef{Pos: -1}}, - } - _, err := expressionVectorPeak(proc, expr, colexec.DefaultBatchSize, false) - require.Error(t, err) -} - -func TestPreparedParamExpressionPeakAccountsLargePayload(t *testing.T) { - proc := testutil.NewProcess(t) - defer proc.Free() - params := vector.NewVec(types.T_text.ToType()) - defer params.Free(proc.Mp()) - payload := make([]byte, types.MaxBlobLen+1) - require.NoError(t, vector.AppendBytes(params, payload, false, proc.Mp())) - proc.SetPrepareParams(params) - - expr := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_varchar), Width: types.MaxVarcharLen}, - Expr: &plan.Expr_P{P: &plan.ParamRef{Pos: 0}}, - } - peak, err := expressionVectorPeak(proc, expr, colexec.DefaultBatchSize, false) - require.NoError(t, err) - header, ok := mpool.GrowCapacity(0, int64(types.VarlenaSize)) - require.True(t, ok) - area, ok := mpool.GrowCapacity(0, int64(len(payload))) - require.True(t, ok) - require.GreaterOrEqual(t, peak, uint64(header)+uint64(area)) - require.Greater(t, peak, uint64(types.MaxBlobLen)) -} - func TestGetJoinMapTransfersGroupSels(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) input := makeIntKeyValueBatch(proc, []int32{1, 1}, []int32{10, 20}) - require.NoError(t, hb.Batches.CopyIntoBatches(input, proc)) + require.NoError(t, hb.CopyBuildBatch(input, proc)) hb.InputBatchRowCount = input.RowCount() input.Clean(proc.Mp()) @@ -1810,7 +711,7 @@ func TestGetJoinMapTransfersGroupSels(t *testing.T) { } func TestDedupUpdateBuildGroupsNullKeysSeparately(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) hb.IsDedup = true hb.OnDuplicateAction = plan.Node_UPDATE @@ -1834,7 +735,7 @@ func TestDedupUpdateBuildGroupsNullKeysSeparately(t *testing.T) { bat := batch.New([]string{"id"}) bat.SetVector(0, keyVec) bat.SetRowCount(rows) - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) hb.InputBatchRowCount = bat.RowCount() bat.Clean(proc.Mp()) @@ -1854,7 +755,7 @@ func TestDedupUpdateBuildGroupsNullKeysSeparately(t *testing.T) { func TestHashMapAllocAndFree(t *testing.T) { mp := mpool.MustNewZero() - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) var err error hb.IntHashMap, err = hashmap.NewIntHashMap(false, mp) require.NoError(t, err) @@ -1879,7 +780,7 @@ func TestHashMapAllocAndFree(t *testing.T) { } func TestIteratorReuseAcrossBuilds(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) @@ -1887,7 +788,7 @@ func TestIteratorReuseAcrossBuilds(t *testing.T) { defer b.Clean(proc.Mp()) hb.InputBatchRowCount = b.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(b, proc)) + require.NoError(t, hb.CopyBuildBatch(b, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedIntIterator) itr1 := hb.cachedIntIterator @@ -1898,14 +799,15 @@ func TestIteratorReuseAcrossBuilds(t *testing.T) { require.Same(t, itr1, hb.cachedIntIterator) // Next build should reuse the same iterator instance. + require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) hb.InputBatchRowCount = b.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(b, proc)) + require.NoError(t, hb.CopyBuildBatch(b, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.Same(t, itr1, hb.cachedIntIterator) } func TestStrIteratorCapacityPrune(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_varchar.ToType())}, -1, -1, nil, proc)) @@ -1917,7 +819,7 @@ func TestStrIteratorCapacityPrune(t *testing.T) { bat.SetVector(0, vec) bat.SetRowCount(1) hb.InputBatchRowCount = bat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedStrIterator) require.Greater(t, hashmap.StrIteratorCapacity(hb.cachedStrIterator), hashmap.MaxStrIteratorCapacity) @@ -1927,7 +829,7 @@ func TestStrIteratorCapacityPrune(t *testing.T) { } func TestStrIteratorBelowThresholdIsKept(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_varchar.ToType())}, -1, -1, nil, proc)) @@ -1940,7 +842,7 @@ func TestStrIteratorBelowThresholdIsKept(t *testing.T) { bat.SetVector(0, vec) bat.SetRowCount(vec.Length()) hb.InputBatchRowCount = bat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedStrIterator) require.Less(t, hashmap.StrIteratorCapacity(hb.cachedStrIterator), hashmap.MaxStrIteratorCapacity) @@ -1950,7 +852,7 @@ func TestStrIteratorBelowThresholdIsKept(t *testing.T) { } func TestResetWithHashTableSentKeepsCache(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) @@ -1958,7 +860,7 @@ func TestResetWithHashTableSentKeepsCache(t *testing.T) { b := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 8, proc.Mp()) defer b.Clean(proc.Mp()) hb.InputBatchRowCount = b.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(b, proc)) + require.NoError(t, hb.CopyBuildBatch(b, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedIntIterator) @@ -1968,7 +870,7 @@ func TestResetWithHashTableSentKeepsCache(t *testing.T) { } func TestAlternateIntStrBuildsReuseIndependently(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) // First int build @@ -1976,7 +878,7 @@ func TestAlternateIntStrBuildsReuseIndependently(t *testing.T) { bInt := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 4, proc.Mp()) defer bInt.Clean(proc.Mp()) hb.InputBatchRowCount = bInt.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bInt, proc)) + require.NoError(t, hb.CopyBuildBatch(bInt, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedIntIterator) require.Nil(t, hb.cachedStrIterator) @@ -1993,7 +895,7 @@ func TestAlternateIntStrBuildsReuseIndependently(t *testing.T) { bat.SetVector(0, vec) bat.SetRowCount(1) hb.InputBatchRowCount = bat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedStrIterator) @@ -2007,7 +909,7 @@ func TestAlternateIntStrBuildsReuseIndependently(t *testing.T) { bInt2 := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 4, proc.Mp()) defer bInt2.Clean(proc.Mp()) hb.InputBatchRowCount = bInt2.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bInt2, proc)) + require.NoError(t, hb.CopyBuildBatch(bInt2, proc)) require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedIntIterator) @@ -2017,7 +919,7 @@ func TestAlternateIntStrBuildsReuseIndependently(t *testing.T) { } func TestBuildHashmapWithZeroInputKeepsCachesUntouched(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) @@ -2031,7 +933,7 @@ func TestBuildHashmapWithZeroInputKeepsCachesUntouched(t *testing.T) { } func TestDedupBuildDuplicateKeyStillFailsByDefault(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) hb.IsDedup = true hb.OnDuplicateAction = plan.Node_FAIL @@ -2045,7 +947,7 @@ func TestDedupBuildDuplicateKeyStillFailsByDefault(t *testing.T) { require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) bat := makeIntKeyValueBatch(proc, []int32{1, 1}, []int32{10, 20}) - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) hb.InputBatchRowCount = bat.RowCount() bat.Clean(proc.Mp()) @@ -2055,7 +957,7 @@ func TestDedupBuildDuplicateKeyStillFailsByDefault(t *testing.T) { } func TestDedupBuildKeepLastForReplace(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) hb.IsDedup = true hb.DedupBuildKeepLast = true @@ -2070,7 +972,7 @@ func TestDedupBuildKeepLastForReplace(t *testing.T) { require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) bat := makeIntKeyValueBatch(proc, []int32{1, 1, 2}, []int32{10, 20, 30}) - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) hb.InputBatchRowCount = bat.RowCount() bat.Clean(proc.Mp()) @@ -2088,7 +990,7 @@ func TestDedupBuildKeepLastForReplace(t *testing.T) { } func TestDedupBuildKeepLastPreservesDeleteOnlyRows(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) hb.IsDedup = true hb.DedupBuildKeepLast = true @@ -2109,7 +1011,7 @@ func TestDedupBuildKeepLastPreservesDeleteOnlyRows(t *testing.T) { []int32{100, 0, 0}, []uint64{1, 2}, ) - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) hb.InputBatchRowCount = bat.RowCount() bat.Clean(proc.Mp()) @@ -2173,7 +1075,7 @@ func TestAccountedDedupScratchAndDeleteBitmapFollowJoinMapLifetime(t *testing.T) require.NotNil(t, hb.DelRows) require.True(t, hb.DelRows.HasExternalStorage()) require.True(t, hb.DelRows.Contains(2)) - require.Equal(t, generation.Used(), generation.Snapshot().AllocationUsed) + require.Equal(t, account.Snapshot().Used, generation.Used()) require.Positive(t, account.Snapshot().Used) jm := hb.GetJoinMap(proc.Mp()) @@ -2212,8 +1114,7 @@ func TestAccountedDedupBitmapExactBoundaryRollsBack(t *testing.T) { require.NoError(t, err) account, err := registry.OpenWithController(tc.cap, generation) require.NoError(t, err) - var hb HashmapBuilder - hb.mapAllocationAccount = account + hb := &HashmapBuilder{mapAllocationAccount: account} bm, err := hb.newDedupBitmap( 64, proc.Mp(), @@ -2245,7 +1146,7 @@ func TestAccountedDedupBitmapExactBoundaryRollsBack(t *testing.T) { // row's old PK equals the surviving row's new key, otherwise the dedup-join // probe side raises a false DuplicateEntry for the existing row REPLACE removes. func TestDedupBuildKeepLastMarksConflictBucketForDiscardedFanout(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) hb.IsDedup = true hb.DedupBuildKeepLast = true @@ -2270,7 +1171,7 @@ func TestDedupBuildKeepLastMarksConflictBucketForDiscardedFanout(t *testing.T) { []int32{100, 200, 300}, nil, ) - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) hb.InputBatchRowCount = bat.RowCount() bat.Clean(proc.Mp()) @@ -2312,7 +1213,7 @@ func TestDedupBuildIgnoreOnlyMarksCandidateOwnOldKey(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) hb.IsDedup = true hb.OnDuplicateAction = plan.Node_IGNORE @@ -2324,7 +1225,7 @@ func TestDedupBuildIgnoreOnlyMarksCandidateOwnOldKey(t *testing.T) { require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, 1, -1, nil, proc)) bat := makeIntKeyValueBatch(proc, tt.newKeys, tt.oldKeys) - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) hb.InputBatchRowCount = bat.RowCount() bat.Clean(proc.Mp()) @@ -2338,7 +1239,7 @@ func TestDedupBuildIgnoreOnlyMarksCandidateOwnOldKey(t *testing.T) { func TestDedupBuildIgnorePrefersOriginalKeyOwner(t *testing.T) { for _, oldKeys := range [][]int32{{1, 2}, {2, 1}} { t.Run(strings.Join([]string{strconv.Itoa(int(oldKeys[0])), strconv.Itoa(int(oldKeys[1]))}, "_"), func(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) hb.IsDedup = true hb.OnDuplicateAction = plan.Node_IGNORE @@ -2350,7 +1251,7 @@ func TestDedupBuildIgnorePrefersOriginalKeyOwner(t *testing.T) { require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, 1, -1, nil, proc)) bat := makeIntKeyValueBatch(proc, []int32{2, 2}, oldKeys) - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) hb.InputBatchRowCount = bat.RowCount() bat.Clean(proc.Mp()) @@ -2365,7 +1266,7 @@ func TestDedupBuildIgnorePrefersOriginalKeyOwner(t *testing.T) { } func TestDedupBuildIgnoreRebuildsAfterOwnerReplacement(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) hb.IsDedup = true hb.OnDuplicateAction = plan.Node_IGNORE @@ -2377,7 +1278,7 @@ func TestDedupBuildIgnoreRebuildsAfterOwnerReplacement(t *testing.T) { require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, 1, -1, nil, proc)) bat := makeIntKeyValueBatch(proc, []int32{2, 1, 2}, []int32{1, 3, 2}) - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) hb.InputBatchRowCount = bat.RowCount() bat.Clean(proc.Mp()) @@ -2393,7 +1294,7 @@ func TestDedupBuildIgnoreRebuildsAfterOwnerReplacement(t *testing.T) { } func TestBuildHashmapErrorDoesNotLeakIterators(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) mp := mpool.MustNewZero() proc := testutil.NewProcessWithMPool(t, "", mp) @@ -2418,13 +1319,13 @@ func TestBuildHashmapErrorDoesNotLeakIterators(t *testing.T) { intBat := batch.New([]string{"col"}) intBat.SetVector(0, intVec) intBat.SetRowCount(1) - require.NoError(t, hb.Batches.CopyIntoBatches(intBat, proc)) + require.NoError(t, hb.CopyBuildBatch(intBat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedIntIterator) } func TestBuildHashmapReuseUniqueSelsBuffer(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) @@ -2433,7 +1334,7 @@ func TestBuildHashmapReuseUniqueSelsBuffer(t *testing.T) { // First build: should allocate uniqueSels hb.InputBatchRowCount = bat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, true, proc)) require.NotNil(t, hb.uniqueSels) require.Greater(t, cap(hb.uniqueSels), 0) @@ -2448,7 +1349,7 @@ func TestBuildHashmapReuseUniqueSelsBuffer(t *testing.T) { hb.InputBatchRowCount = bat.RowCount() hb.Batches.Reset() hb.Batches.Buf = nil - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, true, proc)) require.NotNil(t, hb.uniqueSels) require.Greater(t, len(hb.uniqueSels), 0) @@ -2456,7 +1357,7 @@ func TestBuildHashmapReuseUniqueSelsBuffer(t *testing.T) { } func TestBuildHashmapDoesNotCreateUniqueSelsWhenNotNeeded(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) @@ -2464,13 +1365,13 @@ func TestBuildHashmapDoesNotCreateUniqueSelsWhenNotNeeded(t *testing.T) { defer bat.Clean(proc.Mp()) hb.InputBatchRowCount = bat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.Nil(t, hb.uniqueSels, "should not allocate uniqueSels when needUniqueVec is false") } func TestCachedStrIteratorOwnerClearedBeforeReuse(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) // Build once to create cached str iterator. @@ -2478,7 +1379,7 @@ func TestCachedStrIteratorOwnerClearedBeforeReuse(t *testing.T) { bat := makeStrBatch(t, 4, proc) defer bat.Clean(proc.Mp()) hb.InputBatchRowCount = bat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedStrIterator) @@ -2494,7 +1395,7 @@ func TestCachedStrIteratorOwnerClearedBeforeReuse(t *testing.T) { hb.InputBatchRowCount = bat.RowCount() hb.Batches.Reset() hb.Batches.Buf = nil - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) rv := reflect.ValueOf(hb.cachedStrIterator).Elem() @@ -2504,7 +1405,7 @@ func TestCachedStrIteratorOwnerClearedBeforeReuse(t *testing.T) { } func TestSwitchKeyTypeCreatesCorrectIterator(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) // Build int first. @@ -2512,7 +1413,7 @@ func TestSwitchKeyTypeCreatesCorrectIterator(t *testing.T) { intBat := makeIntBatch(t, 2, proc) defer intBat.Clean(proc.Mp()) hb.InputBatchRowCount = intBat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(intBat, proc)) + require.NoError(t, hb.CopyBuildBatch(intBat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedIntIterator) @@ -2525,13 +1426,13 @@ func TestSwitchKeyTypeCreatesCorrectIterator(t *testing.T) { hb.InputBatchRowCount = strBat.RowCount() hb.Batches.Reset() hb.Batches.Buf = nil - require.NoError(t, hb.Batches.CopyIntoBatches(strBat, proc)) + require.NoError(t, hb.CopyBuildBatch(strBat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedStrIterator) } func TestCachedIteratorOwnerClearedBeforeReuse(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) // Build once to create cached int iterator and bind to map A. @@ -2539,7 +1440,7 @@ func TestCachedIteratorOwnerClearedBeforeReuse(t *testing.T) { bat := makeIntBatch(t, 4, proc) defer bat.Clean(proc.Mp()) hb.InputBatchRowCount = bat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedIntIterator) @@ -2555,7 +1456,7 @@ func TestCachedIteratorOwnerClearedBeforeReuse(t *testing.T) { hb.InputBatchRowCount = bat.RowCount() hb.Batches.Reset() hb.Batches.Buf = nil - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) // Owner should now be non-nil and point to the new map (i.e., not staleMap). @@ -2566,7 +1467,7 @@ func TestCachedIteratorOwnerClearedBeforeReuse(t *testing.T) { } func TestFreeThenBuildRepopulatesCache(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) // First build to populate cache. @@ -2576,7 +1477,7 @@ func TestFreeThenBuildRepopulatesCache(t *testing.T) { intBat.SetVector(0, intVec) intBat.SetRowCount(2) hb.InputBatchRowCount = intBat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(intBat, proc)) + require.NoError(t, hb.CopyBuildBatch(intBat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedIntIterator) @@ -2588,7 +1489,7 @@ func TestFreeThenBuildRepopulatesCache(t *testing.T) { // Build again after Free should succeed and repopulate cache. require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) hb.InputBatchRowCount = intBat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(intBat, proc)) + require.NoError(t, hb.CopyBuildBatch(intBat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedIntIterator) } @@ -2610,7 +1511,7 @@ func (f failingExecutor) ResetForNextQuery() {} // Benchmarks: cached vs new iterator paths for int/str. func BenchmarkBuildHashmapCachedInt(b *testing.B) { proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) - hb := &HashmapBuilder{} + hb := newTestHashmapBuilder(b) require.NoError(b, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) data := makeIntBatch(b, 1024, proc) defer data.Clean(proc.Mp()) @@ -2618,7 +1519,7 @@ func BenchmarkBuildHashmapCachedInt(b *testing.B) { b.ResetTimer() for i := 0; i < b.N; i++ { hb.InputBatchRowCount = data.RowCount() - require.NoError(b, hb.Batches.CopyIntoBatches(data, proc)) + require.NoError(b, hb.CopyBuildBatch(data, proc)) require.NoError(b, hb.BuildHashmap(false, false, false, proc)) hb.Reset(proc, true) } @@ -2626,7 +1527,7 @@ func BenchmarkBuildHashmapCachedInt(b *testing.B) { func BenchmarkBuildHashmapCachedStr(b *testing.B) { proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) - hb := &HashmapBuilder{} + hb := newTestHashmapBuilder(b) require.NoError(b, hb.Prepare([]*plan.Expr{newExpr(0, types.T_varchar.ToType())}, -1, -1, nil, proc)) data := makeStrBatch(b, 1024, proc) defer data.Clean(proc.Mp()) @@ -2634,7 +1535,7 @@ func BenchmarkBuildHashmapCachedStr(b *testing.B) { b.ResetTimer() for i := 0; i < b.N; i++ { hb.InputBatchRowCount = data.RowCount() - require.NoError(b, hb.Batches.CopyIntoBatches(data, proc)) + require.NoError(b, hb.CopyBuildBatch(data, proc)) require.NoError(b, hb.BuildHashmap(false, false, false, proc)) hb.Reset(proc, true) } @@ -2698,10 +1599,10 @@ func BenchmarkBuildHashmapColdInt(b *testing.B) { defer data.Clean(proc.Mp()) b.ResetTimer() for i := 0; i < b.N; i++ { - hb := &HashmapBuilder{} + hb := newTestHashmapBuilder(b) require.NoError(b, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) hb.InputBatchRowCount = data.RowCount() - require.NoError(b, hb.Batches.CopyIntoBatches(data, proc)) + require.NoError(b, hb.CopyBuildBatch(data, proc)) require.NoError(b, hb.BuildHashmap(false, false, false, proc)) hb.Free(proc) } @@ -2713,10 +1614,10 @@ func BenchmarkBuildHashmapColdStr(b *testing.B) { defer data.Clean(proc.Mp()) b.ResetTimer() for i := 0; i < b.N; i++ { - hb := &HashmapBuilder{} + hb := newTestHashmapBuilder(b) require.NoError(b, hb.Prepare([]*plan.Expr{newExpr(0, types.T_varchar.ToType())}, -1, -1, nil, proc)) hb.InputBatchRowCount = data.RowCount() - require.NoError(b, hb.Batches.CopyIntoBatches(data, proc)) + require.NoError(b, hb.CopyBuildBatch(data, proc)) require.NoError(b, hb.BuildHashmap(false, false, false, proc)) hb.Free(proc) } @@ -2724,71 +1625,51 @@ func BenchmarkBuildHashmapColdStr(b *testing.B) { func BenchmarkCopyBuildBatchAccounting(b *testing.B) { const capBytes = uint64(256 << 20) - for _, accounted := range []bool{false, true} { - name := "legacy" - if accounted { - name = "accounted" - } - b.Run(name, func(b *testing.B) { - proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) - defer proc.Free() - input := testutil.NewBatch( - []types.Type{types.T_int32.ToType(), types.T_varchar.ToType()}, - true, - colexec.DefaultBatchSize, - proc.Mp(), - ) - defer input.Clean(proc.Mp()) - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - if err != nil { - b.Fatal(err) - } - var hb HashmapBuilder - hb.setBudget(generation) - var ( - registry *mpool.AllocationAccountRegistry - account *mpool.AllocationAccount - op HashBuild - ) - if accounted { - registry, err = mpool.NewAllocationAccountRegistry(1, 64) - if err != nil { - b.Fatal(err) - } - account, err = registry.OpenWithController(capBytes, generation) - if err != nil { - b.Fatal(err) - } - op.NeedHashMap = true - if err = op.SetAllocationAccount(account); err != nil { - b.Fatal(err) - } - hb.batchAllocation = op.ctr.hashmapBuilder.batchAllocation - } + proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) + defer proc.Free() + input := testutil.NewBatch( + []types.Type{types.T_int32.ToType(), types.T_varchar.ToType()}, + true, + colexec.DefaultBatchSize, + proc.Mp(), + ) + defer input.Clean(proc.Mp()) + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err := budget.OpenGeneration(1) + if err != nil { + b.Fatal(err) + } + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + if err != nil { + b.Fatal(err) + } + account, err := registry.OpenWithController(capBytes, generation) + if err != nil { + b.Fatal(err) + } + hb := &HashmapBuilder{} + if err = hb.SetAllocationAccount(account); err != nil { + b.Fatal(err) + } + hb.setBudget(generation) - b.ReportAllocs() - b.ResetTimer() - for range b.N { - if err = hb.copyBuildBatch(input, proc); err != nil { - b.Fatal(err) - } - hb.cleanBatches(proc) - } - b.StopTimer() - if generation.Used() != 0 { - b.Fatalf("generation used = %d", generation.Used()) - } - if accounted { - op.ctr.hashmapBuilder.batchAllocation = nil - if err = op.ClearAllocationAccount(account); err != nil { - b.Fatal(err) - } - if _, _, err = registry.CompleteTerminal(account); err != nil { - b.Fatal(err) - } - } - }) + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if err = hb.copyBuildBatch(input, proc); err != nil { + b.Fatal(err) + } + hb.cleanBatches(proc) + } + b.StopTimer() + if generation.Used() != 0 { + b.Fatalf("generation used = %d", generation.Used()) + } + if err = hb.ClearAllocationAccount(account); err != nil { + b.Fatal(err) + } + if _, _, err = registry.CompleteTerminal(account); err != nil { + b.Fatal(err) } } @@ -2805,95 +1686,79 @@ func BenchmarkResidentHashBuildAccounting(b *testing.B) { if stringKey { kind = "varchar" } - for _, accounted := range []bool{false, true} { - mode := "legacy" - if accounted { - mode = "accounted" + b.Run(fmt.Sprintf("%s/rows-%d", kind, rows), func(b *testing.B) { + proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) + defer proc.Free() + var input *batch.Batch + var keyType types.Type + if stringKey { + input = makeStrBatch(b, rows, proc) + keyType = types.T_varchar.ToType() + } else { + input = makeIntBatch(b, rows, proc) + keyType = types.T_int32.ToType() + } + defer input.Clean(proc.Mp()) + + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err := budget.OpenGeneration(1) + if err != nil { + b.Fatal(err) + } + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + if err != nil { + b.Fatal(err) + } + account, err := registry.OpenWithController(capBytes, generation) + if err != nil { + b.Fatal(err) } - b.Run(fmt.Sprintf("%s/%s/rows-%d", mode, kind, rows), func(b *testing.B) { - proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) - defer proc.Free() - var input *batch.Batch - var keyType types.Type - if stringKey { - input = makeStrBatch(b, rows, proc) - keyType = types.T_varchar.ToType() - } else { - input = makeIntBatch(b, rows, proc) - keyType = types.T_int32.ToType() - } - defer input.Clean(proc.Mp()) - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - if err != nil { + b.ReportAllocs() + b.SetBytes(int64(input.Size())) + b.ResetTimer() + for range b.N { + hb := &HashmapBuilder{} + hb.SetBudget(generation) + if err = hb.SetAllocationAccount(account); err != nil { b.Fatal(err) } - var ( - registry *mpool.AllocationAccountRegistry - account *mpool.AllocationAccount - ) - if accounted { - registry, err = mpool.NewAllocationAccountRegistry(1, 4_096) - if err != nil { - b.Fatal(err) - } - account, err = registry.OpenWithController(capBytes, generation) - if err != nil { - b.Fatal(err) - } - } - - b.ReportAllocs() - b.SetBytes(int64(input.Size())) - b.ResetTimer() - for range b.N { - hb := &HashmapBuilder{} - hb.SetBudget(generation) - if accounted { - if err = hb.SetAllocationAccount(account); err != nil { - b.Fatal(err) - } - } - if err = hb.Prepare( - []*plan.Expr{newExpr(0, keyType)}, - -1, - -1, - nil, - proc, - ); err != nil { - b.Fatal(err) - } - hb.InputBatchRowCount = input.RowCount() - if err = hb.CopyBuildBatch(input, proc); err != nil { - b.Fatal(err) - } - if err = hb.BuildHashmap(false, false, false, proc); err != nil { - b.Fatal(err) - } - hb.Free(proc) + if err = hb.Prepare( + []*plan.Expr{newExpr(0, keyType)}, + -1, + -1, + nil, + proc, + ); err != nil { + b.Fatal(err) } - b.StopTimer() - if generation.Used() != 0 { - b.Fatalf("generation used = %d", generation.Used()) + hb.InputBatchRowCount = input.RowCount() + if err = hb.CopyBuildBatch(input, proc); err != nil { + b.Fatal(err) } - if accounted { - if account.Snapshot().Used != 0 { - b.Fatalf("account used = %d", account.Snapshot().Used) - } - if _, _, err = registry.CompleteTerminal(account); err != nil { - b.Fatal(err) - } + if err = hb.BuildHashmap(false, false, false, proc); err != nil { + b.Fatal(err) } - generation.Close() - }) - } + hb.Free(proc) + } + b.StopTimer() + if generation.Used() != 0 { + b.Fatalf("generation used = %d", generation.Used()) + } + if account.Snapshot().Used != 0 { + b.Fatalf("account used = %d", account.Snapshot().Used) + } + if _, _, err = registry.CompleteTerminal(account); err != nil { + b.Fatal(err) + } + generation.Close() + }) } } } func TestExtractRestoreCachedIterators(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) mp := mpool.MustNewZero() intMap, err := hashmap.NewIntHashMap(false, mp) @@ -2920,7 +1785,7 @@ func TestExtractRestoreCachedIterators(t *testing.T) { } func TestStrIteratorLargeStringTriggersPrune(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_varchar.ToType())}, -1, -1, nil, proc)) @@ -2933,7 +1798,7 @@ func TestStrIteratorLargeStringTriggersPrune(t *testing.T) { bat.SetRowCount(1) hb.InputBatchRowCount = bat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedStrIterator) @@ -2948,7 +1813,7 @@ func TestStrIteratorLargeStringTriggersPrune(t *testing.T) { // curVecs or UniqueJoinKeys contained nil pointers. func TestResetWithNilPointers(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) // Test case 1: curVecs with nil pointers and needDupVec = true hb.needDupVec = true @@ -2972,7 +1837,7 @@ func TestResetWithNilPointers(t *testing.T) { // TestResetWithMixedNilAndValidPointers tests Reset() with a mix of nil and valid vectors func TestResetWithMixedNilAndValidPointers(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) // Create some valid vectors vec1 := testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()) @@ -2995,7 +1860,7 @@ func TestResetWithMixedNilAndValidPointers(t *testing.T) { // TestFreeWithNilPointers tests that Free() handles nil pointers gracefully func TestFreeWithNilPointers(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) // Test case: UniqueJoinKeys with nil pointers hb.UniqueJoinKeys = make([]*vector.Vector, 3) @@ -3012,7 +1877,7 @@ func TestFreeWithNilPointers(t *testing.T) { // TestFreeWithMixedNilAndValidPointers tests Free() with a mix of nil and valid vectors func TestFreeWithMixedNilAndValidPointers(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) // Create some valid vectors vec1 := testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()) diff --git a/pkg/sql/colexec/hashbuild/pressure.go b/pkg/sql/colexec/hashbuild/pressure.go index ad1bab84acbc7..aa89f4f271746 100644 --- a/pkg/sql/colexec/hashbuild/pressure.go +++ b/pkg/sql/colexec/hashbuild/pressure.go @@ -28,7 +28,7 @@ import ( // control flow. Sealed and invariant failures are lifecycle bugs and must // remain terminal. HashBuildBudgetError now exposes disjoint lifecycle and // capacity identities; this classifier is the one control-flow boundary for -// physical-account and legacy-budget failures. +// physical-allocation and spill-resource failures. type MemoryPressureReason uint8 const ( @@ -58,12 +58,15 @@ func MemoryPressureReasonOf(err error) MemoryPressureReason { switch budgetErr.Kind { case process.HashBuildBudgetErrorAdmission: switch budgetErr.Component { + case process.HashBuildBudgetComponentMemory: + return MemoryPressureCapacity case process.HashBuildBudgetComponentSpillDisk: return MemoryPressureSpillDiskLimit case process.HashBuildBudgetComponentSpillFD: return MemoryPressureSpillFDLimit + default: + return MemoryPressureInvalid } - return MemoryPressureCapacity case process.HashBuildBudgetErrorClosed: return MemoryPressureSealed case process.HashBuildBudgetErrorInvalid, @@ -88,13 +91,11 @@ func MemoryPressureReasonOf(err error) MemoryPressureReason { return MemoryPressureInvariant } - // A few compatibility call sites return the bare sentinel. Check closed - // first; structured errors above never rely on the legacy Is alias. + // Resource-ledger helpers may still return a bare lifecycle sentinel. if errors.Is(err, process.ErrHashBuildBudgetClosed) { return MemoryPressureSealed } - if errors.Is(err, process.ErrHashBuildBudgetAdmission) || - errors.Is(err, process.ErrHashBuildBudgetRejected) { + if errors.Is(err, process.ErrHashBuildBudgetAdmission) { return MemoryPressureCapacity } if errors.Is(err, process.ErrHashBuildBudgetInvalid) || diff --git a/pkg/sql/colexec/hashbuild/pressure_test.go b/pkg/sql/colexec/hashbuild/pressure_test.go index 15127329c31a3..1ee55c7f2b2de 100644 --- a/pkg/sql/colexec/hashbuild/pressure_test.go +++ b/pkg/sql/colexec/hashbuild/pressure_test.go @@ -29,7 +29,7 @@ func TestMemoryPressureReasonSeparatesCapacityFromLifecycle(t *testing.T) { reason MemoryPressureReason }{ {nil, MemoryPressureNone}, - {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorAdmission}, MemoryPressureCapacity}, + {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorAdmission}, MemoryPressureInvalid}, {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorAdmission, Component: process.HashBuildBudgetComponentMemory}, MemoryPressureCapacity}, {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorAdmission, Component: process.HashBuildBudgetComponentSpillDisk}, MemoryPressureSpillDiskLimit}, {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorAdmission, Component: process.HashBuildBudgetComponentSpillFD}, MemoryPressureSpillFDLimit}, diff --git a/pkg/sql/colexec/hashbuild/spill.go b/pkg/sql/colexec/hashbuild/spill.go index f7a2fb0dd538a..550fa1abc6a39 100644 --- a/pkg/sql/colexec/hashbuild/spill.go +++ b/pkg/sql/colexec/hashbuild/spill.go @@ -15,7 +15,6 @@ package hashbuild import ( - "bytes" "fmt" "io" "math" @@ -25,12 +24,10 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" - "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/fileservice" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" - "github.com/matrixorigin/matrixone/pkg/sql/plan/function" "github.com/matrixorigin/matrixone/pkg/vm/process" ) @@ -44,784 +41,24 @@ const ( spillWriteCoalesceSize = 64 << 10 ) -type spillMaterializationMode uint8 - -const ( - // spillDirectMaterialization models UnionInt32 on the current upstream - // batch. A const varlen source copies its out-of-line payload once and - // broadcasts the resulting descriptor. - spillDirectMaterialization spillMaterializationMode = iota - // spillRetainedMaterialization models the compact non-const batch produced - // by CopyIntoBatches. A later UnionInt32 treats every retained row as an - // independent value, even when the ingress vector was const. - spillRetainedMaterialization -) - -func spillCheckedAdd(total, value uint64) (uint64, error) { - if total > math.MaxUint64-value { - return 0, process.ErrHashBuildBudgetInvalid - } - return total + value, nil -} - -func spillCheckedMul(left, right uint64) (uint64, error) { - if left != 0 && right > math.MaxUint64/left { - return 0, process.ErrHashBuildBudgetInvalid - } - return left * right, nil -} - -func spillCapacityReplacementOverlap(rows, keys, hashCap, rowIDCap, keyCap int) (uint64, error) { - var overlap uint64 - add := func(required, current int, width uint64) error { - if required < 0 || current < 0 { - return process.ErrHashBuildBudgetInvalid - } - if required <= current { - return nil - } - old, err := spillCheckedMul(uint64(current), width) - if err != nil { - return err - } - overlap, err = spillCheckedAdd(overlap, old) - return err - } - if err := add(keys, keyCap, 8); err != nil { - return 0, err - } - if err := add(rows, hashCap, 8); err != nil { - return 0, err - } - if err := add(rows, rowIDCap, 4); err != nil { - return 0, err - } - return overlap, nil -} - -// spillMaterializedBytes models the batch that spillBatchBounded creates with -// UnionInt32. It follows vector materialization semantics instead of retained -// capacity or stale logical length: fixed-width descriptors are per output -// row, null payload is skipped, and direct const varlen payload is copied once. -func spillMaterializedBytesFor( - bat *batch.Batch, - targetRows uint64, - mode spillMaterializationMode, -) (uint64, error) { - if bat == nil || bat.RowCount() <= 0 || targetRows == 0 { - return 0, nil - } - liveRows := uint64(bat.RowCount()) - var materialized uint64 - for _, vec := range bat.Vecs { - if vec == nil { - return 0, process.ErrHashBuildBudgetInvalid - } - typeSize := vec.GetType().TypeSize() - if typeSize < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - descriptors, err := spillCheckedMul(targetRows, uint64(typeSize)) - if err != nil { - return 0, err - } - if materialized, err = spillCheckedAdd(materialized, descriptors); err != nil { - return 0, err - } - if !vec.GetType().IsVarlen() || vec.IsConstNull() { - continue - } - - values, _ := vector.MustVarlenaRawData(vec) - valueRows := liveRows - if vec.IsConst() { - valueRows = 1 - } - if valueRows == 0 || valueRows > uint64(len(values)) { - return 0, process.ErrHashBuildBudgetInvalid - } - var livePayload uint64 - hasNull := !vec.GetNulls().EmptyByFlag() - for row := uint64(0); row < valueRows; row++ { - if hasNull && vec.GetNulls().Contains(row) { - continue - } - if values[row].IsSmall() { - continue - } - _, length := values[row].OffsetLen() - if livePayload, err = spillCheckedAdd(livePayload, uint64(length)); err != nil { - return 0, err - } - } - - projectedPayload := livePayload - if !(mode == spillDirectMaterialization && vec.IsConst()) { - // A retained CopyIntoBatches destination is non-const. Repeating - // the complete live sample is a conservative bound for any compact - // target batch assembled from ingress batches whose individual - // high-water estimates were admitted before copying. - roundedRows, err := spillCheckedAdd(targetRows, valueRows-1) - if err != nil { - return 0, err - } - repeats := roundedRows / valueRows - if projectedPayload, err = spillCheckedMul(livePayload, repeats); err != nil { - return 0, err - } - } - - if materialized, err = spillCheckedAdd(materialized, projectedPayload); err != nil { - return 0, err - } - } - return materialized, nil -} - -func spillMaterializedBytes(bat *batch.Batch) (uint64, error) { - if bat == nil || bat.RowCount() <= 0 { - return 0, nil - } - return spillMaterializedBytesFor( - bat, - uint64(bat.RowCount()), - spillDirectMaterialization, - ) -} - -func spillMarshalSlack(columns uint64) (uint64, error) { - const ( - fixedSlack = uint64(64 << 10) - perColumnSlack = uint64(128) - ) - if columns > (math.MaxUint64-fixedSlack)/perColumnSlack { - return 0, process.ErrHashBuildBudgetInvalid - } - return fixedSlack + columns*perColumnSlack, nil -} - -func spillMaterializationSlack(columns uint64) (uint64, error) { - const ( - fixedSlack = uint64(64 << 10) - perColumnSlack = uint64(16 << 10) - ) - if columns > (math.MaxUint64-fixedSlack)/perColumnSlack { - return 0, process.ErrHashBuildBudgetInvalid - } - return fixedSlack + columns*perColumnSlack, nil -} - -// spillPeakBudgetFor accounts each simultaneously live component explicitly. -// inputBytes is zero for a retained batch whose source reservation is already -// owned by HashBuild. -func spillPeakBudgetFor(rows, inputBytes, selectedBytes, columns uint64) (uint64, error) { - rowScratch, err := spillCheckedMul(rows, 12) // hashes + one row-id array - if err != nil { - return 0, err - } - total, err := spillCheckedAdd(rowScratch, inputBytes) - if err != nil { - return 0, err - } - if total, err = spillCheckedAdd(total, selectedBytes); err != nil { - return 0, err - } - // MarshalBinary creates one serialized payload. The selected estimate - // already includes its fixed-width data and varlen area, so charge that - // payload once plus bounded framing/allocation slack. - marshalSlack, err := spillMarshalSlack(columns) - if err != nil { - return 0, err - } - marshalBytes, err := spillCheckedAdd(selectedBytes, marshalSlack) - if err != nil { - return 0, err - } - if total, err = spillCheckedAdd(total, marshalBytes); err != nil { - return 0, err - } - if total > uint64(^uint(0)>>1) { - return 0, process.ErrHashBuildBudgetInvalid - } - return total, nil -} - -// spillBudgetBytes admits only the actual direct-spill path for the current -// input. It never projects a hypothetical retained batch. -func spillBudgetBytes(bat *batch.Batch) (uint64, error) { - if bat == nil || bat.RowCount() <= 0 { - return 0, nil - } - rows := uint64(bat.RowCount()) - selected, err := spillMaterializedBytesFor( - bat, - rows, - spillDirectMaterialization, - ) - if err != nil { - return 0, err - } - materializationSlack, err := spillMaterializationSlack(uint64(len(bat.Vecs))) - if err != nil { - return 0, err - } - if selected, err = spillCheckedAdd(selected, materializationSlack); err != nil { - return 0, err - } - return spillPeakBudgetFor(rows, uint64(bat.Allocated()), selected, uint64(len(bat.Vecs))) -} - -// spillScratchBudgetBytes returns the incremental spill charge. A copied -// build batch remains covered by HashmapBuilder.batchReservations while it is -// drained, so charging its source footprint again would double count it. An -// upstream batch has no HashBuild-owned reservation and keeps the full charge. -func spillScratchBudgetBytes(bat *batch.Batch, sourceAlreadyCharged bool) (uint64, error) { - need, err := spillBudgetBytes(bat) - if err != nil || !sourceAlreadyCharged || bat == nil || bat.RowCount() <= 0 { - return need, err - } - // copyBuildBatch reconciles its retained reservation against Allocated - // (plus metadata), so only that proven charge may be subtracted here. - source := uint64(bat.Allocated()) - if source > need { - return 0, process.ErrHashBuildBudgetInvalid - } - return need - source, nil -} - -// spillRetainedBudgetBytes is the future-drain proof required before -// CopyIntoBatches may retain a small input. The destination loses constness, -// so its selected payload follows retained rather than direct semantics. -func spillRetainedBudgetBytes(bat *batch.Batch) (uint64, error) { - if bat == nil || bat.RowCount() <= 0 { - return 0, nil - } - rows := uint64(bat.RowCount()) - targetRows := rows - if rows < uint64(colexec.DefaultBatchSize) { - targetRows = uint64(colexec.DefaultBatchSize) - } - selected, err := spillMaterializedBytesFor( - bat, - targetRows, - spillRetainedMaterialization, - ) - if err != nil { - return 0, err - } - metadata, ok := retainedMetadataAllowance(bat) - if !ok || metadata > math.MaxUint64/targetRows { - return 0, process.ErrHashBuildBudgetInvalid - } - projectedMetadata := metadata - if targetRows > rows { - projectedMetadata, err = spillCheckedMul(metadata, targetRows) - if err != nil { - return 0, err - } - projectedMetadata, err = spillCheckedAdd(projectedMetadata, rows-1) - if err != nil { - return 0, err - } - projectedMetadata /= rows - } - if selected, err = spillCheckedAdd(selected, projectedMetadata); err != nil { - return 0, err - } - materializationSlack, err := spillMaterializationSlack(uint64(len(bat.Vecs))) - if err != nil { - return 0, err - } - if selected, err = spillCheckedAdd(selected, materializationSlack); err != nil { - return 0, err - } - // The retained source itself is covered by batchReservations. - return spillPeakBudgetFor( - targetRows, - 0, - selected, - uint64(len(bat.Vecs)), - ) -} - -func (ctr *container) ensureSpillScratchReservationBytes( - need uint64, - analyzer process.Analyzer, -) error { - if ctr.hashmapBuilder.budget == nil || need == 0 { - return nil - } - var err error - if ctr.spillScratchReservation == nil { - ctr.spillScratchReservation, err = - ctr.hashmapBuilder.budget.Reserve(need) - if err == nil { - analyzer.GetOpStats().SetMaxExtraStat( - "HashBuildEmergencyScratchBytes", - hashBuildStatInt64(need), - ) - ctr.spillScratchEmergency = true - ctr.spillScratchBase = need - } - return err - } - if ctr.spillScratchBase >= need { - ctr.spillScratchEmergency = true - return nil - } - grow := need - ctr.spillScratchBase - if err := ctr.spillScratchReservation.Grow(grow); err != nil { - analyzer.GetOpStats().AddExtraStat( - "HashBuildEmergencyScratchGrowRejects", - 1, - ) - return err - } - analyzer.GetOpStats().AddExtraStat( - "HashBuildEmergencyScratchGrowCount", - 1, - ) - analyzer.GetOpStats().AddExtraStat( - "HashBuildEmergencyScratchGrowBytes", - hashBuildStatInt64(grow), - ) - ctr.spillScratchBase = need - ctr.spillScratchEmergency = true - return nil -} - -func (ctr *container) growSpillScratchTransient( - required uint64, - analyzer process.Analyzer, -) (uint64, bool, error) { - if ctr.hashmapBuilder.budget == nil || ctr.spillScratchReservation == nil || - required <= ctr.spillScratchBase { - return 0, false, nil - } - oldSize := ctr.spillScratchReservation.Size() - if err := ctr.spillScratchReservation.Grow(required - ctr.spillScratchBase); err != nil { - return 0, false, err - } - analyzer.GetOpStats().SetMaxExtraStat( - "HashBuildSpillScratchPeakBytes", - hashBuildStatInt64(ctr.spillScratchReservation.Size()), - ) - return oldSize, true, nil -} - -func (ctr *container) restoreSpillScratchTransient(oldSize uint64, grew bool) error { - if !grew { - return nil - } - _, err := ctr.spillScratchReservation.ReconcileDown(oldSize) - return err -} - -func (ctr *container) ensureDirectSpillScratchReservation(bat *batch.Batch, analyzer process.Analyzer) error { - if ctr.spillBatchAllocation != nil { - // A borrowed upstream batch is already live and cannot be reclaimed by - // reserving another logical token. Exact scatter/expression allocations - // admit their physical capacities and adapt the unpublished input. The - // retained-copy path below still keeps a one-unit future-progress token - // because choosing to retain is under HashBuild's control. - return nil - } - var ( - need uint64 - err error - ) - need, err = spillBudgetBytes(bat) - if err != nil { - return err - } - return ctr.ensureSpillScratchReservationBytes(need, analyzer) -} - -func (ctr *container) ensureRetainedSpillScratchReservation(bat *batch.Batch, analyzer process.Analyzer) error { - var ( - need uint64 - err error - ) - if ctr.spillBatchAllocation != nil { - need, err = spillMinimumUnitBudgetBytes(bat, ctr.spillConditions) - } else { - need, err = spillRetainedBudgetBytes(bat) - } - if err != nil { - return err - } - return ctr.ensureSpillScratchReservationBytes(need, analyzer) -} - -// spillMinimumUnitBudgetBytes keeps only the headroom for one physical spill -// unit. It derives capacities from actual input values and the closed -// expression family; it neither scales the whole batch nor applies a safety -// multiplier. The token is converted into exact allocations on spill entry. -func spillMinimumUnitBudgetBytes( - bat *batch.Batch, - exprs []*plan.Expr, -) (uint64, error) { - if bat == nil || bat.RowCount() <= 0 { - return 0, nil - } - selected, wire, err := spillMinimumSelectedAndWireBytes(bat) - if err != nil { - return 0, err - } - if wire > math.MaxUint64-24 { - return 0, process.ErrHashBuildBudgetInvalid - } - marshal, err := initialAllocationCapacity(wire + 24) - if err != nil { - return 0, err - } - expression, err := spillMinimumExpressionBytes(exprs, bat) - if err != nil { - return 0, err - } - total := uint64(12) // one hash plus one row id - for _, value := range []uint64{selected, marshal, expression} { - if total > math.MaxUint64-value { - return 0, process.ErrHashBuildBudgetInvalid - } - total += value - } - return total, nil -} - -func spillMinimumSelectedAndWireBytes( - bat *batch.Batch, -) (selected uint64, wire uint64, err error) { - // Batch framing plus Attr/ExtraBuf length prefixes. - wire = 8 + 4 + 4 + 4 + 4 + 4 + uint64(len(bat.ExtraBuf)) - for _, attr := range bat.Attrs { - if wire > math.MaxUint64-4-uint64(len(attr)) { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - wire += 4 + uint64(len(attr)) - } - for _, vec := range bat.Vecs { - if vec == nil || vec.GetType().TypeSize() < 0 { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - data, capErr := initialAllocationCapacity( - uint64(vec.GetType().TypeSize()), - ) - if capErr != nil { - return 0, 0, capErr - } - areaPayload, payloadErr := maxVectorValueBytes(vec) - if payloadErr != nil { - return 0, 0, payloadErr - } - var area uint64 - if areaPayload > types.VarlenaInlineSize { - area, capErr = initialAllocationCapacity(areaPayload) - if capErr != nil { - return 0, 0, capErr - } - } - // Accounted vectors install one null and one grouping word before - // extending their first row. - physical := data + area + 16 - if selected > math.MaxUint64-physical { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - selected += physical - // Physical capacities upper-bound one-row logical data, area, and null - // payload; only the fixed wire framing is added separately. - const vectorFraming = uint64(4 + 1 + types.TSize + 4 + 4 + 4 + 4 + 1) - if wire > math.MaxUint64-vectorFraming-physical { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - wire += vectorFraming + physical - } - return selected, wire, nil -} - -func maxVectorValueBytes(vec *vector.Vector) (uint64, error) { - if vec == nil || !vec.GetType().IsVarlen() || vec.IsConstNull() { - return 0, nil - } - values, _ := vector.MustVarlenaRawData(vec) - rows := vec.Length() - if vec.IsConst() && rows > 0 { - rows = 1 - } - if rows > len(values) { - return 0, process.ErrHashBuildBudgetInvalid - } - var maximum uint64 - for row := 0; row < rows; row++ { - if vec.GetNulls().Contains(uint64(row)) { - continue - } - var length uint64 - if values[row].IsSmall() { - length = uint64(len(values[row].GetByteSlice(nil))) - } else { - _, valueLen := values[row].OffsetLen() - length = uint64(valueLen) - } - if length > maximum { - maximum = length - } - } - return maximum, nil -} - -func spillMinimumExpressionBytes( - exprs []*plan.Expr, - bat *batch.Batch, -) (uint64, error) { - if len(exprs) == 0 { - return 0, nil - } - if !AllocationAccountedExpressionSetSupported(exprs) { - return 0, process.ErrHashBuildBudgetInvalid - } - var total uint64 - for _, expr := range exprs { - bytes, err := spillMinimumExpressionTreeBytes(expr, bat) - if err != nil || total > math.MaxUint64-bytes { - return 0, process.ErrHashBuildBudgetInvalid - } - total += bytes - } - return total, nil -} - -func spillMinimumExpressionTreeBytes( - expr *plan.Expr, - bat *batch.Batch, -) (uint64, error) { - if expr == nil { - return 0, process.ErrHashBuildBudgetInvalid - } - switch node := expr.Expr.(type) { - case *plan.Expr_Col: - return 0, nil - case *plan.Expr_Lit: - return expressionInitialOwnedBytes(expr) - case *plan.Expr_F: - if node.F == nil || node.F.Func == nil { - return 0, process.ErrHashBuildBudgetInvalid - } - var total uint64 - for _, arg := range node.F.Args { - child, err := spillMinimumExpressionTreeBytes(arg, bat) - if err != nil || total > math.MaxUint64-child { - return 0, process.ErrHashBuildBudgetInvalid - } - total += child - } - result, err := spillMinimumExpressionResultBytes(expr, node.F.Args, bat) - if err != nil || total > math.MaxUint64-result { - return 0, process.ErrHashBuildBudgetInvalid - } - total += result - functionID, _ := function.DecodeOverloadID(node.F.Func.Obj) - if functionID == function.CASE { - if total > math.MaxUint64-8 { - return 0, process.ErrHashBuildBudgetInvalid - } - total += 8 - } - return total, nil - default: - return expressionInitialOwnedBytes(expr) - } -} - -func spillMinimumExpressionResultBytes( - expr *plan.Expr, - args []*plan.Expr, - bat *batch.Batch, -) (uint64, error) { - oid := types.T(expr.Typ.Id) - typ := oid.ToType() - if typ.TypeSize() < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - data, err := initialAllocationCapacity(uint64(typ.TypeSize())) - if err != nil { - return 0, err - } - result := data + 16 - if !typ.IsVarlen() { - return result, nil - } - payload, err := spillExpressionPayloadBytes(expr, args, bat) - if err != nil { - return 0, err - } - if payload > types.VarlenaInlineSize { - area, err := initialAllocationCapacity(payload) - if err != nil || result > math.MaxUint64-area { - return 0, process.ErrHashBuildBudgetInvalid - } - result += area - } - return result, nil -} - -func spillExpressionPayloadBytes( - expr *plan.Expr, - args []*plan.Expr, - bat *batch.Batch, -) (uint64, error) { - if bat == nil || bat.RowCount() < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - var maximum uint64 - for row := 0; row < bat.RowCount(); row++ { - value, err := spillExpressionPayloadBytesAt(expr, args, bat, row) - if err != nil { - return 0, err - } - if value > maximum { - maximum = value - } - } - return maximum, nil -} - -func spillExpressionPayloadBytesAt( - expr *plan.Expr, - args []*plan.Expr, - bat *batch.Batch, - row int, -) (uint64, error) { - node, ok := expr.Expr.(*plan.Expr_F) - if !ok || node.F == nil || node.F.Func == nil { - return 0, nil - } - functionID, _ := function.DecodeOverloadID(node.F.Func.Obj) - switch functionID { - case function.CONCAT: - var total uint64 - for _, arg := range args { - value, err := spillExpressionArgPayloadBytesAt(arg, bat, row) - if err != nil || total > math.MaxUint64-value { - return 0, process.ErrHashBuildBudgetInvalid - } - total += value - } - return total, nil - case function.CASE: - // CASE's exact selected branch is evaluated later. The largest varlen - // branch value at this same row is a safe one-row bound without combining - // maxima taken from different rows. - var maximum uint64 - for _, arg := range args { - if arg == nil || !types.T(arg.Typ.Id).ToType().IsVarlen() { - continue - } - value, err := spillExpressionArgPayloadBytesAt(arg, bat, row) - if err != nil { - return 0, err - } - if value > maximum { - maximum = value - } - } - return maximum, nil - case function.CAST: - if len(args) == 0 || args[0] == nil { - return 0, process.ErrHashBuildBudgetInvalid - } - if types.T(args[0].Typ.Id).ToType().IsIntOrUint() { - return 20, nil - } - return spillExpressionArgPayloadBytesAt(args[0], bat, row) - default: - return 0, nil - } -} - -func spillExpressionArgPayloadBytesAt( - expr *plan.Expr, - bat *batch.Batch, - row int, -) (uint64, error) { - if expr == nil || bat == nil || row < 0 || row >= bat.RowCount() { - return 0, process.ErrHashBuildBudgetInvalid - } - switch node := expr.Expr.(type) { - case *plan.Expr_Col: - if node.Col == nil || node.Col.ColPos < 0 || - int(node.Col.ColPos) >= len(bat.Vecs) { - return 0, process.ErrHashBuildBudgetInvalid - } - return vectorValueBytesAt(bat.Vecs[node.Col.ColPos], row) - case *plan.Expr_Lit: - if node.Lit == nil || node.Lit.GetIsnull() { - return 0, nil - } - return uint64(len(node.Lit.GetSval())), nil - case *plan.Expr_F: - return spillExpressionPayloadBytesAt(expr, node.F.GetArgs(), bat, row) - default: - return 0, nil - } -} - -func vectorValueBytesAt(vec *vector.Vector, row int) (uint64, error) { - if vec == nil || row < 0 || row >= vec.Length() || - !vec.GetType().IsVarlen() || vec.IsConstNull() { - if vec == nil || row < 0 || row >= vec.Length() { - return 0, process.ErrHashBuildBudgetInvalid - } - return 0, nil - } - index := row - if vec.IsConst() { - index = 0 - } - if vec.GetNulls().Contains(uint64(index)) { - return 0, nil - } - values, _ := vector.MustVarlenaRawData(vec) - if index >= len(values) { - return 0, process.ErrHashBuildBudgetInvalid +func (ctr *container) dropSpillScratchBuffers() { + if cap(ctr.spillHashValues) > 0 { + mpool.FreeSlice(ctr.spillAllocationMP, ctr.spillHashValues) } - if values[index].IsSmall() { - return uint64(len(values[index].GetByteSlice(nil))), nil + if cap(ctr.spillBucketRowIds) > 0 { + mpool.FreeSlice(ctr.spillAllocationMP, ctr.spillBucketRowIds) } - _, length := values[index].OffsetLen() - return uint64(length), nil -} - -func (ctr *container) releaseSpillScratchReservation() { - if ctr.spillScratchReservation != nil { - ctr.spillScratchReservation.Release() - ctr.spillScratchReservation = nil + if ctr.spillAccountedWrite != nil { + ctr.spillAccountedWrite.Free() + ctr.spillAccountedWrite = nil } - ctr.spillScratchEmergency = false - ctr.spillScratchBase = 0 -} - -func (ctr *container) dropSpillScratchBuffers() { - if ctr.spillBatchAllocation != nil { - if cap(ctr.spillHashValues) > 0 { - mpool.FreeSlice(ctr.spillAllocationMP, ctr.spillHashValues) - } - if cap(ctr.spillBucketRowIds) > 0 { - mpool.FreeSlice(ctr.spillAllocationMP, ctr.spillBucketRowIds) - } - if ctr.spillAccountedWrite != nil { - ctr.spillAccountedWrite.Free() - ctr.spillAccountedWrite = nil - } - for i := range ctr.spillAccountedBuckets { - if ctr.spillAccountedBuckets[i] != nil { - ctr.spillAccountedBuckets[i].Free() - ctr.spillAccountedBuckets[i] = nil - } + for i := range ctr.spillAccountedBuckets { + if ctr.spillAccountedBuckets[i] != nil { + ctr.spillAccountedBuckets[i].Free() + ctr.spillAccountedBuckets[i] = nil } } - for bucket := range ctr.spillBucketWriteBufs { - ctr.spillBucketWriteBufs[bucket] = bytes.Buffer{} + for bucket := range ctr.spillBucketWriteRows { ctr.spillBucketWriteRows[bucket] = 0 } ctr.spillHashValues = nil @@ -833,7 +70,6 @@ func (ctr *container) dropSpillScratchBuffers() { ctr.spillBucketOffsets[i] = 0 } ctr.spillKeyVecs = nil - ctr.spillWriteBuf = bytes.Buffer{} ctr.spillAllocationMP = nil ctr.spillCoalesceDisabled = false } @@ -879,61 +115,6 @@ func growHashBuildSpillSlice[T any]( return next[:length], nil } -func spillMarshalGrowBytes(bat *batch.Batch) (uint64, error) { - base := uint64(bat.Allocated()) - if size := uint64(bat.Size()); size > base { - base = size - } - columns := uint64(len(bat.Vecs)) - if columns > (math.MaxUint64-24)/128 { - return 0, process.ErrHashBuildBudgetInvalid - } - return spillCheckedAdd(base, columns*128+24) -} - -func marshalSpillRecord(bat *batch.Batch, buf *bytes.Buffer) (int64, error) { - if bat == nil || bat.RowCount() == 0 { - return 0, nil - } - - cnt := int64(bat.RowCount()) - buf.Reset() - grow, err := spillMarshalGrowBytes(bat) - if err != nil { - return 0, err - } - if grow > uint64(math.MaxInt) { - return 0, process.ErrHashBuildBudgetInvalid - } - if uint64(buf.Cap()) < grow { - // Drop a smaller retained buffer before allocating the final capacity; - // otherwise bytes.Buffer's geometric growth recreates the multiplier - // that admission intentionally removed. - *buf = *bytes.NewBuffer(make([]byte, 0, int(grow))) - } - buf.Write(types.EncodeInt64(&cnt)) - // Reserve space for batchSize (filled in after marshalling) - batchSizePos := buf.Len() - var zero int64 - buf.Write(types.EncodeInt64(&zero)) - - // Write batch data directly to spillWriteBuf. The bounded partition path - // reserves this buffer's conservative upper bound before entering here. - batchStartPos := buf.Len() - if _, err := bat.MarshalBinaryWithBuffer(buf, false); err != nil { - return 0, err - } - batchSize := int64(buf.Len() - batchStartPos) - - // Write batchSize at reserved position - batchSizeBytes := types.EncodeInt64(&batchSize) - copy(buf.Bytes()[batchSizePos:batchSizePos+len(batchSizeBytes)], batchSizeBytes) - - magic := uint64(spillMagic) - buf.Write(types.EncodeUint64(&magic)) - return cnt, nil -} - func marshalSpillRecordAccounted( bat *batch.Batch, buf *mpool.AccountedBuffer, @@ -943,7 +124,7 @@ func marshalSpillRecordAccounted( } cnt := int64(bat.RowCount()) buf.Reset() - batchSize, err := bat.MarshalBinarySize() + batchSize, err := bat.MarshalBinaryWithGroupingSize() if err != nil || batchSize > math.MaxInt-24 { if err != nil { return 0, err @@ -953,25 +134,22 @@ func marshalSpillRecordAccounted( if err := buf.EnsureCapacity(batchSize + 24); err != nil { return 0, err } - if _, err := buf.Write(types.EncodeInt64(&cnt)); err != nil { + if err := buf.WriteInt64(cnt); err != nil { return 0, err } batchSizePos := buf.Len() - var zero int64 - if _, err := buf.Write(types.EncodeInt64(&zero)); err != nil { + if err := buf.WriteInt64(0); err != nil { return 0, err } batchStart := buf.Len() - if err := bat.MarshalBinaryTo(buf); err != nil { + if err := bat.MarshalBinaryWithGroupingTo(buf); err != nil { return 0, err } serializedSize := int64(buf.Len() - batchStart) - copy( - buf.Bytes()[batchSizePos:batchSizePos+8], - types.EncodeInt64(&serializedSize), - ) - magic := uint64(spillMagic) - if _, err := buf.Write(types.EncodeUint64(&magic)); err != nil { + if err := buf.SetInt64(batchSizePos, serializedSize); err != nil { + return 0, err + } + if err := buf.WriteUint64(uint64(spillMagic)); err != nil { return 0, err } return cnt, nil @@ -995,15 +173,12 @@ func (ctr *container) writeSpillPayload( return err } - var err error - if ctr.hashmapBuilder.budget != nil { - if ctr.spillBundle == nil { - return process.ErrHashBuildBudgetInvalid - } - _, _, err = ctr.spillBundle.growDisk(file, ctr.hashmapBuilder.budget, uint64(len(payload))) - if err != nil { - return err - } + if ctr.hashmapBuilder.budget == nil || ctr.spillBundle == nil { + return process.ErrHashBuildBudgetInvalid + } + _, _, err := ctr.spillBundle.growDisk(file, ctr.hashmapBuilder.budget, uint64(len(payload))) + if err != nil { + return err } if err := checkHashBuildCanceled(proc); err != nil { return err @@ -1015,12 +190,10 @@ func (ctr *container) writeSpillPayload( if written != len(payload) { return io.ErrShortWrite } - if ctr.hashmapBuilder.budget != nil { - // The exact payload length was admitted. Record logical ownership only - // after the full write; partial writes retain the conservative charge - // until the enclosing bundle closes the file. - ctr.spillBundle.recordDiskWrite(file, rows, uint64(written)) - } + // The exact payload length was admitted. Record logical ownership only + // after the full write; partial writes retain the conservative charge + // until the enclosing bundle closes the file. + ctr.spillBundle.recordDiskWrite(file, rows, uint64(written)) if analyzer != nil { analyzer.Spill(int64(written)) analyzer.SpillRows(rows) @@ -1029,20 +202,6 @@ func (ctr *container) writeSpillPayload( return nil } -func (ctr *container) flushBucketBuffer(proc *process.Process, bat *batch.Batch, file *os.File, analyzer process.Analyzer) (int64, error) { - if bat == nil || bat.RowCount() == 0 { - return 0, nil - } - cnt, err := marshalSpillRecord(bat, &ctr.spillWriteBuf) - if err != nil { - return 0, err - } - if err := ctr.writeSpillPayload(proc, file, ctr.spillWriteBuf.Bytes(), cnt, analyzer); err != nil { - return 0, err - } - return cnt, nil -} - func (ctr *container) getSpillFS(proc *process.Process) (fileservice.MutableFileService, error) { if ctr.spillFS != nil { return ctr.spillFS, nil @@ -1071,12 +230,12 @@ func (ctr *container) ensureSpillFile(proc *process.Process, files []*os.File, b return nil, err } name := fmt.Sprintf("join_%s_%d_build", ctr.spillUUID, bucket) - var fdToken *process.HashBuildSpillFDReservation - if ctr.hashmapBuilder.budget != nil { - fdToken, err = ctr.hashmapBuilder.budget.ReserveSpillFD(1) - if err != nil { - return nil, err - } + if ctr.hashmapBuilder.budget == nil { + return nil, process.ErrHashBuildBudgetInvalid + } + fdToken, err := ctr.hashmapBuilder.budget.ReserveSpillFD(1) + if err != nil { + return nil, err } f, err := spillfs.CreateAndRemoveFile(proc.Ctx, name) if err != nil { @@ -1086,12 +245,10 @@ func (ctr *container) ensureSpillFile(proc *process.Process, files []*os.File, b return nil, err } files[bucket] = f - if fdToken != nil { - if ctr.spillBundle == nil { - ctr.spillBundle = &spillFileBundle{} - } - ctr.spillBundle.addFD(f, bucket, fdToken) + if ctr.spillBundle == nil { + ctr.spillBundle = &spillFileBundle{} } + ctr.spillBundle.addFD(f, bucket, fdToken) return f, nil } @@ -1108,136 +265,54 @@ func (ctr *container) spillBatchBounded(proc *process.Process, bat *batch.Batch, if err := checkHashBuildCanceled(proc); err != nil { return err } - exact := ctr.spillBatchAllocation != nil - var ( - need uint64 - err error - ) - if exact { - if ctr.hashmapBuilder.mapAllocationAccount == nil { - return mpool.ErrAllocationAccountInvalid - } - if ctr.spillAllocationMP != nil && ctr.spillAllocationMP != proc.Mp() { - return mpool.ErrAllocationAccountInvalid - } - ctr.spillAllocationMP = proc.Mp() - // A pre-spill token is headroom, not physical ownership. Release it - // immediately before the exact scratch allocations consume that space. - ctr.releaseSpillScratchReservation() - if !sourceAlreadyCharged { - // The upstream batch is borrowed, already physically live, and cannot - // be made smaller by rejecting a new logical token. Record it as - // observation only; every new HashBuild-owned byte below is admitted - // by the exact account and the process MPool remains the global guard. - externalBytes := bat.Allocated() - if size := bat.Size(); size > externalBytes { - externalBytes = size - } - analyzer.GetOpStats().SetMaxExtraStat( - "HashBuildSpillBorrowedSourceBytes", - int64(externalBytes), - ) - } - } else { - need, err = spillScratchBudgetBytes(bat, sourceAlreadyCharged) - if err != nil { - return err - } + if ctr.spillBatchAllocation == nil || + ctr.hashmapBuilder.mapAllocationAccount == nil { + return mpool.ErrAllocationAccountInvalid } - // Scratch belongs to the execution generation, not to one batch. Build - // normally pre-admits the emergency lease before calling us; direct callers - // (including recovery/error paths and unit tests) establish the same lease - // here. Keep it live while capacities are retained and release it from - // Reset/Free/build cleanup exactly once. - if ctr.hashmapBuilder.budget != nil && !exact { - if ctr.spillScratchReservation == nil { - ctr.spillScratchReservation, err = ctr.hashmapBuilder.budget.Reserve(need) - if err != nil { - analyzer.GetOpStats().AddExtraStat("HashBuildSpillScratchReserveRejects", 1) - return err - } - ctr.spillScratchBase = need - analyzer.GetOpStats().SetMaxExtraStat( - "HashBuildSpillScratchPeakBytes", - hashBuildStatInt64(ctr.spillScratchReservation.Size()), - ) - } else if need > ctr.spillScratchBase { - grow := need - ctr.spillScratchBase - if err := ctr.spillScratchReservation.Grow(grow); err != nil { - analyzer.GetOpStats().AddExtraStat("HashBuildSpillScratchGrowRejects", 1) - return err - } - analyzer.GetOpStats().AddExtraStat("HashBuildSpillScratchGrowCount", 1) - analyzer.GetOpStats().AddExtraStat("HashBuildSpillScratchGrowBytes", hashBuildStatInt64(grow)) - ctr.spillScratchBase = need - analyzer.GetOpStats().SetMaxExtraStat( - "HashBuildSpillScratchPeakBytes", - hashBuildStatInt64(ctr.spillScratchReservation.Size()), - ) + if ctr.spillAllocationMP != nil && ctr.spillAllocationMP != proc.Mp() { + return mpool.ErrAllocationAccountInvalid + } + ctr.spillAllocationMP = proc.Mp() + if !sourceAlreadyCharged { + externalBytes := bat.Allocated() + if size := bat.Size(); size > externalBytes { + externalBytes = size } + analyzer.GetOpStats().SetMaxExtraStat( + "HashBuildSpillBorrowedSourceBytes", + int64(externalBytes), + ) } rows := bat.RowCount() - var oldScratchSize uint64 - var grewScratch bool - if !exact { - replacementOverlap, overlapErr := spillCapacityReplacementOverlap( - rows, - len(executors), - cap(ctr.spillHashValues), - cap(ctr.spillBucketRowIds), - cap(ctr.spillKeyVecs), - ) - if overlapErr != nil { - return overlapErr - } - replacementPeak, addErr := spillCheckedAdd(need, replacementOverlap) - if addErr != nil { - return addErr - } - oldScratchSize, grewScratch, err = ctr.growSpillScratchTransient( - replacementPeak, - analyzer, - ) - if err != nil { - return err - } + if !keycodec.ValidVectors(bat.Vecs, rows) { + return process.ErrHashBuildBudgetInvalid } + var err error if cap(ctr.spillKeyVecs) < len(executors) { ctr.spillKeyVecs = make([]*vector.Vector, len(executors)) } - if exact { - ctr.spillHashValues, err = growHashBuildSpillSlice( - ctr.spillHashValues, - rows, - proc.Mp(), - ctr.hashmapBuilder.mapAllocationAccount, - HashBuildSpillAllocationSiteHashValues, - ) - } else if cap(ctr.spillHashValues) < rows { - ctr.spillHashValues = make([]uint64, rows) - } + ctr.spillHashValues, err = growHashBuildSpillSlice( + ctr.spillHashValues, + rows, + proc.Mp(), + ctr.hashmapBuilder.mapAllocationAccount, + HashBuildSpillAllocationSiteHashValues, + ) if err != nil { return err } - if exact { - ctr.spillBucketRowIds, err = growHashBuildSpillSlice( - ctr.spillBucketRowIds, - rows, - proc.Mp(), - ctr.hashmapBuilder.mapAllocationAccount, - HashBuildSpillAllocationSiteRowIDs, - ) - } else if cap(ctr.spillBucketRowIds) < rows { - ctr.spillBucketRowIds = make([]int32, rows) - } + ctr.spillBucketRowIds, err = growHashBuildSpillSlice( + ctr.spillBucketRowIds, + rows, + proc.Mp(), + ctr.hashmapBuilder.mapAllocationAccount, + HashBuildSpillAllocationSiteRowIDs, + ) if err != nil { return err } - if err := ctr.restoreSpillScratchTransient(oldScratchSize, grewScratch); err != nil { - return err - } keyVecs := ctr.spillKeyVecs[:len(executors)] var selected *batch.Batch defer func() { @@ -1255,27 +330,20 @@ func (ctr *container) spillBatchBounded(proc *process.Process, bat *batch.Batch, } return evalErr } - if ctr.spillExprLease != nil { - if ctr.spillExprLease.Len() != len(executors) { - return process.ErrHashBuildBudgetInvalid - } - err = ctr.spillExprLease.Run(proc, bat.RowCount(), evalOne) - } else { - for i := range executors { - if err = evalOne(i); err != nil { - break - } + for i := range executors { + if err = evalOne(i); err != nil { + break } } if err != nil { - // Eval may leave newly allocated child/result vectors cached in the - // executor tree. Destroy that tree while both the previous and - // candidate reservations are still charged. - if !exact { - ctr.freeSpillExprExecs() - } + // Eval may leave child/result allocations cached. Destroy the tree so a + // pressure retry starts from the exact post-rollback account state. + ctr.freeSpillExprExecs() return err } + if !keycodec.ValidVectors(keyVecs, rows) { + return process.ErrHashBuildBudgetInvalid + } if err := checkHashBuildCanceled(proc); err != nil { return err } @@ -1320,29 +388,22 @@ func (ctr *container) spillBatchBounded(proc *process.Process, bat *batch.Batch, } if selected == nil { selected = batch.NewOffHeapWithSize(len(bat.Vecs)) - if exact { - if err := selected.SetAllocationAccount( - ctr.spillBatchAllocation, - ); err != nil { - return err - } + if err := selected.SetAllocationAccount( + ctr.spillBatchAllocation, + ); err != nil { + return err } - selected.Attrs = bat.Attrs for i, vec := range bat.Vecs { if vec == nil { return process.ErrHashBuildBudgetInvalid } - if exact { - selected.Vecs[i], err = - vector.NewOffHeapVecWithTypeAndAllocation( - *vec.GetType(), - ctr.spillBatchAllocation, - ) - if err != nil { - return err - } - } else { - selected.Vecs[i] = vector.NewOffHeapVecWithType(*vec.GetType()) + selected.Vecs[i], err = + vector.NewOffHeapVecWithTypeAndAllocation( + *vec.GetType(), + ctr.spillBatchAllocation, + ) + if err != nil { + return err } } } @@ -1377,7 +438,6 @@ func (ctr *container) spillBatchBounded(proc *process.Process, bat *batch.Batch, file, int(bucket), selected, - need, analyzer, ) } @@ -1387,7 +447,7 @@ func (ctr *container) spillBatchBounded(proc *process.Process, bat *batch.Batch, cursor = attemptEnd break } - if !exact || !IsRetryableMemoryCapacity(spillErr) { + if !IsRetryableMemoryCapacity(spillErr) { return spillErr } if err := checkHashBuildCanceled(proc); err != nil { @@ -1533,7 +593,9 @@ func (ctr *container) spillBatchWithPressure( current := bat if start != 0 || end != rows { var err error - current, err = bat.Window(start, end) + current, err = bat.WindowWithAllocation( + start, end, proc.Mp(), ctr.spillBatchAllocation, + ) if err != nil { return err } @@ -1626,77 +688,11 @@ func (ctr *container) appendSpillRecord( file *os.File, bucket int, bat *batch.Batch, - scratchNeed uint64, analyzer process.Analyzer, ) error { if bucket < 0 || bucket >= spillNumBuckets { return process.ErrHashBuildBudgetInvalid } - if ctr.spillBatchAllocation != nil { - return ctr.appendAccountedSpillRecord( - proc, - file, - bucket, - bat, - analyzer, - ) - } - grow, err := spillMarshalGrowBytes(bat) - if err != nil { - return err - } - var oldScratchSize uint64 - var grewScratch bool - if old := uint64(ctr.spillWriteBuf.Cap()); ctr.hashmapBuilder.budget != nil && old > 0 && old < grow { - peak, addErr := spillCheckedAdd(scratchNeed, old) - if addErr != nil { - return addErr - } - oldScratchSize, grewScratch, err = ctr.growSpillScratchTransient(peak, analyzer) - if err != nil { - return err - } - } - cnt, err := marshalSpillRecord(bat, &ctr.spillWriteBuf) - if restoreErr := ctr.restoreSpillScratchTransient(oldScratchSize, grewScratch); restoreErr != nil { - return restoreErr - } - if err != nil { - return err - } - payload := ctr.spillWriteBuf.Bytes() - buf := &ctr.spillBucketWriteBufs[bucket] - if buf.Len() > 0 && buf.Len()+len(payload) > spillWriteCoalesceSize { - if err := ctr.flushPendingSpillBucket(proc, file, bucket, analyzer); err != nil { - return err - } - } - if len(payload) > spillWriteCoalesceSize { - return ctr.writeSpillPayload(proc, file, payload, cnt, analyzer) - } - if buf.Len() == 0 { - if !ctr.ensureSpillCoalesceCapacity(buf, analyzer) { - return ctr.writeSpillPayload(proc, file, payload, cnt, analyzer) - } - if buf.Cap() < spillWriteCoalesceSize { - *buf = *bytes.NewBuffer(make([]byte, 0, spillWriteCoalesceSize)) - } - } - _, _ = buf.Write(payload) - ctr.spillBucketWriteRows[bucket] += cnt - if buf.Len() >= spillWriteCoalesceSize { - return ctr.flushPendingSpillBucket(proc, file, bucket, analyzer) - } - return nil -} - -func (ctr *container) appendAccountedSpillRecord( - proc *process.Process, - file *os.File, - bucket int, - bat *batch.Batch, - analyzer process.Analyzer, -) error { if ctr.spillAllocationMP != proc.Mp() || ctr.hashmapBuilder.mapAllocationAccount == nil { return mpool.ErrAllocationAccountInvalid @@ -1777,27 +773,6 @@ func (ctr *container) appendAccountedSpillRecord( return nil } -func (ctr *container) ensureSpillCoalesceCapacity(buf *bytes.Buffer, analyzer process.Analyzer) bool { - if buf == nil || buf.Cap() >= spillWriteCoalesceSize { - return true - } - if ctr.hashmapBuilder.budget == nil || ctr.spillScratchReservation == nil { - return ctr.hashmapBuilder.budget == nil - } - additional := uint64(spillWriteCoalesceSize - buf.Cap()) - if err := ctr.spillScratchReservation.Grow(additional); err != nil { - analyzer.GetOpStats().AddExtraStat("HashBuildCoalesceGrowRejects", 1) - return false - } - analyzer.GetOpStats().AddExtraStat("HashBuildCoalesceGrowCount", 1) - analyzer.GetOpStats().AddExtraStat("HashBuildCoalesceGrowBytes", hashBuildStatInt64(additional)) - analyzer.GetOpStats().SetMaxExtraStat( - "HashBuildSpillScratchPeakBytes", - hashBuildStatInt64(ctr.spillScratchReservation.Size()), - ) - return true -} - func (ctr *container) flushPendingSpillBucket( proc *process.Process, file *os.File, @@ -1808,28 +783,14 @@ func (ctr *container) flushPendingSpillBucket( return process.ErrHashBuildBudgetInvalid } rows := ctr.spillBucketWriteRows[bucket] - var payload []byte - if ctr.spillBatchAllocation != nil { - buffer := ctr.spillAccountedBuckets[bucket] - if buffer == nil || buffer.Len() == 0 { - return nil - } - payload = buffer.Bytes() - } else { - buf := &ctr.spillBucketWriteBufs[bucket] - if buf.Len() == 0 { - return nil - } - payload = buf.Bytes() + buffer := ctr.spillAccountedBuckets[bucket] + if buffer == nil || buffer.Len() == 0 { + return nil } - err := ctr.writeSpillPayload(proc, file, payload, rows, analyzer) + err := ctr.writeSpillPayload(proc, file, buffer.Bytes(), rows, analyzer) // Clear even on a failed/partial write. A caller's enclosing failure path // owns cleanup, and retrying the same bytes could duplicate records. - if ctr.spillBatchAllocation != nil { - ctr.spillAccountedBuckets[bucket].Reset() - } else { - ctr.spillBucketWriteBufs[bucket].Reset() - } + buffer.Reset() ctr.spillBucketWriteRows[bucket] = 0 return err } @@ -1841,16 +802,14 @@ func (ctr *container) flushPendingSpillBucket( func (ctr *container) flushSpillBuffers(proc *process.Process, files []*os.File, analyzer process.Analyzer) error { var firstErr error for bucket := 0; bucket < spillNumBuckets; bucket++ { - pending := ctr.spillBucketWriteBufs[bucket].Len() - if ctr.spillBatchAllocation != nil && - ctr.spillAccountedBuckets[bucket] != nil { + pending := 0 + if ctr.spillAccountedBuckets[bucket] != nil { pending = ctr.spillAccountedBuckets[bucket].Len() } if pending == 0 { continue } if firstErr != nil { - ctr.spillBucketWriteBufs[bucket].Reset() if ctr.spillAccountedBuckets[bucket] != nil { ctr.spillAccountedBuckets[bucket].Reset() } @@ -1859,7 +818,6 @@ func (ctr *container) flushSpillBuffers(proc *process.Process, files []*os.File, } if err := checkHashBuildCanceled(proc); err != nil { firstErr = err - ctr.spillBucketWriteBufs[bucket].Reset() if ctr.spillAccountedBuckets[bucket] != nil { ctr.spillAccountedBuckets[bucket].Reset() } @@ -1872,7 +830,6 @@ func (ctr *container) flushSpillBuffers(proc *process.Process, files []*os.File, } if file == nil { firstErr = process.ErrHashBuildBudgetInvalid - ctr.spillBucketWriteBufs[bucket].Reset() if ctr.spillAccountedBuckets[bucket] != nil { ctr.spillAccountedBuckets[bucket].Reset() } @@ -1894,47 +851,16 @@ func (ctr *container) initSpillExprExecs(proc *process.Process, conditions []*pl return nil, &process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorInvalid, Message: "nil shuffle spill key"} } } - wantAccounted := ctr.hashmapBuilder.expressionAllocation != nil && - expressionSetAllocationClosed(conditions) - if len(ctr.spillExprExecs) != len(conditions) || - ctr.spillExprAccounted != wantAccounted { - var ( - execs []colexec.ExpressionExecutor - lease *ExpressionMemoryLease - err error + if len(ctr.spillExprExecs) != len(conditions) { + execs, err := NewExpressionExecutors( + proc, + conditions, ) - if wantAccounted { - execs, err = NewAllocationAccountedExpressionExecutors( - proc, - conditions, - ctr.hashmapBuilder.expressionAllocation, - ) - } else { - execs, lease, err = NewBudgetedExpressionExecutors( - proc, - ctr.hashmapBuilder.budget, - conditions, - false, - ) - } if err != nil { return nil, err } ctr.freeSpillExprExecs() ctr.spillExprExecs = execs - ctr.spillExprLease = lease - ctr.spillExprAccounted = wantAccounted - } else if !ctr.spillExprAccounted && ctr.spillExprLease == nil { - lease, err := NewExpressionMemoryLease( - ctr.hashmapBuilder.budget, - conditions, - ctr.spillExprExecs, - false, - ) - if err != nil { - return nil, err - } - ctr.spillExprLease = lease } return ctr.spillExprExecs, nil } @@ -1947,11 +873,6 @@ func (ctr *container) freeSpillExprExecs() { } } ctr.spillExprExecs = nil - ctr.spillExprAccounted = false - if ctr.spillExprLease != nil { - ctr.spillExprLease.Release() - ctr.spillExprLease = nil - } } func (ctr *container) memUsed() int64 { diff --git a/pkg/sql/colexec/hashbuild/spill_test.go b/pkg/sql/colexec/hashbuild/spill_test.go index 4678d02ab669a..407587b8e77d4 100644 --- a/pkg/sql/colexec/hashbuild/spill_test.go +++ b/pkg/sql/colexec/hashbuild/spill_test.go @@ -16,10 +16,9 @@ package hashbuild import ( "bufio" - "bytes" "context" + "errors" "io" - "math" "os" "strings" "testing" @@ -30,1761 +29,253 @@ import ( "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" - plan2 "github.com/matrixorigin/matrixone/pkg/sql/plan" "github.com/matrixorigin/matrixone/pkg/testutil" "github.com/matrixorigin/matrixone/pkg/vm/process" "github.com/stretchr/testify/require" ) -func TestComputeXXHashBuild(t *testing.T) { - mp := mpool.MustNewZero() - - t.Run("empty", func(t *testing.T) { - computeXXHash(nil, nil) - }) - - t.Run("single_column", func(t *testing.T) { - vec := testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, mp) - hashValues := make([]uint64, 3) - computeXXHash([]*vector.Vector{vec}, hashValues) - require.NotEqual(t, uint64(0), hashValues[0]) - require.NotEqual(t, hashValues[0], hashValues[1]) - }) - - t.Run("multiple_columns", func(t *testing.T) { - vec1 := testutil.MakeInt32Vector([]int32{1, 2}, nil, mp) - vec2 := testutil.MakeVarcharVector([]string{"a", "b"}, nil, mp) - hashValues := make([]uint64, 2) - computeXXHash([]*vector.Vector{vec1, vec2}, hashValues) - require.NotEqual(t, hashValues[0], hashValues[1]) - }) - - t.Run("const_vector", func(t *testing.T) { - vec := testutil.MakeInt32Vector([]int32{5}, nil, mp) - vec.SetClass(vector.CONSTANT) - hashValues := make([]uint64, 3) - computeXXHash([]*vector.Vector{vec}, hashValues) - require.Equal(t, hashValues[0], hashValues[1]) - }) -} - -func TestFlushBucketBufferBuild(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - - file, err := spillfs.CreateFile(context.Background(), "test_build_flush") - require.NoError(t, err) - defer func() { - file.Close() - spillfs.RemoveFile(context.Background(), "test_build_flush") - }() - - analyzer := process.NewAnalyzer(0, false, false, "test") - ctr := &container{spillUUID: t.Name()} - - t.Run("empty_buffer", func(t *testing.T) { - var buf *batch.Batch - cnt, err := ctr.flushBucketBuffer(proc, buf, file, analyzer) - require.NoError(t, err) - require.Equal(t, int64(0), cnt) - }) - - t.Run("with_data", func(t *testing.T) { - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()) - bat.SetRowCount(3) - - cnt, err := ctr.flushBucketBuffer(proc, bat, file, analyzer) - require.NoError(t, err) - require.Equal(t, int64(3), cnt) - }) -} - -func TestShouldSpillBatches(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - t.Run("not_shuffle", func(t *testing.T) { - hb := &HashBuild{ - IsShuffle: false, - NeedHashMap: true, - } - hb.ctr.setSpillThreshold(1) - bat := batch.NewWithSize(0) - bat.SetRowCount(1) - hb.ctr.hashmapBuilder.Batches.Buf = []*batch.Batch{bat} - require.False(t, hb.shouldSpillBatches()) - }) - - t.Run("no_hashmap", func(t *testing.T) { - hb := &HashBuild{ - IsShuffle: true, - } - hb.ctr.setSpillThreshold(1) - bat := batch.NewWithSize(0) - bat.SetRowCount(1) - hb.ctr.hashmapBuilder.Batches.Buf = []*batch.Batch{bat} - require.False(t, hb.shouldSpillBatches()) - }) - - t.Run("below_threshold", func(t *testing.T) { - hb := &HashBuild{ - IsShuffle: true, - SpillThreshold: 1024 * 1024, // 1MB - NeedHashMap: true, - } - hb.ctr.setSpillThreshold(1024 * 1024) - hb.ctr.hashmapBuilder.Batches.Buf = []*batch.Batch{ - {Vecs: []*vector.Vector{testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp())}}, - } - require.False(t, hb.shouldSpillBatches()) - }) - - t.Run("above_threshold", func(t *testing.T) { - hb := &HashBuild{ - IsShuffle: true, - SpillThreshold: 1, // 1 byte - NeedHashMap: true, - } - hb.ctr.setSpillThreshold(1) - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3, 4, 5}, nil, proc.Mp()) - bat.SetRowCount(5) - hb.ctr.hashmapBuilder.Batches.Buf = []*batch.Batch{bat} - hb.ctr.hashmapBuilder.InputBatchRowCount = bat.RowCount() - require.True(t, hb.shouldSpillBatches()) - }) -} - -func TestShouldSpillBeforeRetain(t *testing.T) { - t.Run("byte threshold predicts crossing batch", func(t *testing.T) { - hb := &HashBuild{IsShuffle: true, NeedHashMap: true} - hb.ctr.setSpillThreshold(100_001) - hb.ctr.hashmapBuilder.Batches.MemSize = 60_000 - hb.ctr.hashmapBuilder.InputBatchRowCount = 2 - - require.False(t, hb.shouldSpillBeforeRetain(40_001), - "the byte convention spills only after the threshold") - require.True(t, hb.shouldSpillBeforeRetain(40_002), - "the crossing batch must be routed directly before it consumes headroom") - }) - - t.Run("row threshold already includes ingress batch", func(t *testing.T) { - hb := &HashBuild{IsShuffle: true, NeedHashMap: true} - hb.ctr.setSpillThreshold(10) - hb.ctr.hashmapBuilder.InputBatchRowCount = 9 - require.False(t, hb.shouldSpillBeforeRetain(1)) - hb.ctr.hashmapBuilder.InputBatchRowCount = 10 - require.True(t, hb.shouldSpillBeforeRetain(1)) - }) - - t.Run("ineligible topology stays resident", func(t *testing.T) { - hb := &HashBuild{IsShuffle: false, NeedHashMap: true} - hb.ctr.setSpillThreshold(1) - hb.ctr.hashmapBuilder.InputBatchRowCount = 1 - require.False(t, hb.shouldSpillBeforeRetain(math.MaxInt64)) - }) - - t.Run("size overflow fails toward spill", func(t *testing.T) { - hb := &HashBuild{IsShuffle: true, NeedHashMap: true} - hb.ctr.setSpillThreshold(100_001) - hb.ctr.hashmapBuilder.Batches.MemSize = math.MaxInt64 - 1 - require.True(t, hb.shouldSpillBeforeRetain(2)) - }) -} - -func TestMemUsedIncludesPartialTailAfterFullBatches(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - full := batch.NewWithSize(0) - full.SetRowCount(colexec.DefaultBatchSize) - partial := batch.NewWithSize(1) - partial.Vecs[0] = testutil.MakeVarcharVector([]string{"partial-tail"}, nil, proc.Mp()) - partial.SetRowCount(1) - defer partial.Clean(proc.Mp()) - - ctr := container{} - ctr.hashmapBuilder.Batches.Buf = []*batch.Batch{full, partial} - ctr.hashmapBuilder.Batches.MemSize = 60_000 - require.Equal(t, int64(60_000+partial.Size()), ctr.memUsed()) -} - -func TestHashDistributionBuild(t *testing.T) { - mp := mpool.MustNewZero() - vec := testutil.MakeInt32Vector([]int32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, - 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30}, nil, mp) - - hashValues := make([]uint64, 30) - computeXXHash([]*vector.Vector{vec}, hashValues) - - bucketCounts := make([]int, spillNumBuckets) - for _, hash := range hashValues { - bucketId := hash & (spillNumBuckets - 1) - bucketCounts[bucketId]++ - } - - // At least some buckets should have values - nonEmptyBuckets := 0 - for _, count := range bucketCounts { - if count > 0 { - nonEmptyBuckets++ - } - } - require.Greater(t, nonEmptyBuckets, 1) -} - -func TestLargeBufferFlushBuild(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - - analyzer := process.NewAnalyzer(0, false, false, "test") - file, err := spillfs.CreateFile(context.Background(), "test_large_build") - require.NoError(t, err) - defer func() { - file.Close() - spillfs.RemoveFile(context.Background(), "test_large_build") - }() - - // Create large batch - size := spillBufferSize + 100 - values := make([]int32, size) - for i := range values { - values[i] = int32(i) - } - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector(values, nil, proc.Mp()) - bat.SetRowCount(size) - - ctr := &container{spillUUID: t.Name()} - cnt, err := ctr.flushBucketBuffer(proc, bat, file, analyzer) - require.NoError(t, err) - require.Equal(t, int64(size), cnt) -} - -func TestMultipleDataTypesBuild(t *testing.T) { - mp := mpool.MustNewZero() - - tests := []struct { - name string - vec *vector.Vector - }{ - {"int8", testutil.MakeInt8Vector([]int8{1, 2, 3}, nil, mp)}, - {"int16", testutil.MakeInt16Vector([]int16{100, 200, 300}, nil, mp)}, - {"int64", testutil.MakeInt64Vector([]int64{1000, 2000, 3000}, nil, mp)}, - {"uint32", testutil.MakeUint32Vector([]uint32{10, 20, 30}, nil, mp)}, - {"float32", testutil.MakeFloat32Vector([]float32{1.1, 2.2, 3.3}, nil, mp)}, - {"float64", testutil.MakeFloat64Vector([]float64{10.1, 20.2, 30.3}, nil, mp)}, - {"varchar", testutil.MakeVarcharVector([]string{"abc", "def", "ghi"}, nil, mp)}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - hashValues := make([]uint64, 3) - computeXXHash([]*vector.Vector{tt.vec}, hashValues) - require.NotEqual(t, uint64(0), hashValues[0]) - require.NotEqual(t, hashValues[0], hashValues[1]) - }) - } -} - -func TestNullValuesBuild(t *testing.T) { - mp := mpool.MustNewZero() - vec := testutil.MakeInt32Vector([]int32{1, 2, 3}, []uint64{1}, mp) - hashValues := make([]uint64, 3) - computeXXHash([]*vector.Vector{vec}, hashValues) - require.NotEqual(t, uint64(0), hashValues[0]) - require.NotEqual(t, uint64(0), hashValues[2]) -} - -func TestFileWriteErrorBuild(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - analyzer := process.NewAnalyzer(0, false, false, "test") - spillfs, _ := proc.GetSpillFileService() - file, _ := spillfs.CreateFile(context.Background(), "test_error_build") - file.Close() - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1}, nil, proc.Mp()) - bat.SetRowCount(1) - - ctr := &container{spillUUID: t.Name()} - _, err := ctr.flushBucketBuffer(proc, bat, file, analyzer) - require.Error(t, err) - - spillfs.RemoveFile(context.Background(), "test_error_build") -} - -func TestWriteSpillPayloadCancellationStopsBeforePhysicalWrite(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - ctx, cancel := context.WithCancelCause(proc.Ctx) - process.ReplacePipelineCtx(proc, ctx, cancel) - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - file, err := spillfs.CreateFile(context.Background(), t.Name()) - require.NoError(t, err) - defer func() { - require.NoError(t, file.Close()) - require.NoError(t, spillfs.RemoveFile(context.Background(), t.Name())) - }() - - proc.Cancel(context.Canceled) - analyzer := process.NewAnalyzer(0, false, false, "test") - err = (&container{}).writeSpillPayload(proc, file, []byte("stale spill payload"), 1, analyzer) - require.ErrorIs(t, err, context.Canceled) - - info, err := file.Stat() - require.NoError(t, err) - require.Zero(t, info.Size(), "canceled spill must not start physical I/O") - require.Zero(t, analyzer.GetOpStats().SpillSize) - require.Zero(t, analyzer.GetOpStats().SpillRows) -} - -func TestAppendBatchToSpillFilesPartitioning(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - files := make([]*os.File, spillNumBuckets) - defer func() { - for _, file := range files { - if file != nil { - file.Close() - } - } - }() - - // Create batch with known values - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3, 4, 5, 6, 7, 8}, nil, proc.Mp()) - bat.SetRowCount(8) - - conditions := []*plan.Expr{ - { - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{ - Col: &plan.ColRef{ColPos: 0}, - }, - }, - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - ctr := &container{spillUUID: t.Name()} - _, err := ctr.initSpillExprExecs(proc, conditions) - require.NoError(t, err) - err = ctr.spillBatchBounded(proc, bat, files, ctr.spillExprExecs, analyzer, false) - require.NoError(t, err) - -} - -func TestEmptyBatchSpill(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - files := make([]*os.File, spillNumBuckets) - defer func() { - for _, file := range files { - if file != nil { - file.Close() - } - } - }() - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{}, nil, proc.Mp()) - bat.SetRowCount(0) - - conditions := []*plan.Expr{ - { - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{ - Col: &plan.ColRef{ColPos: 0}, - }, - }, - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - ctr := &container{spillUUID: t.Name()} - _, err := ctr.initSpillExprExecs(proc, conditions) - require.NoError(t, err) - err = ctr.spillBatchBounded(proc, bat, files, ctr.spillExprExecs, analyzer, false) - require.NoError(t, err) -} - -func TestAppendBuildBatchMultipleFlushes(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - files := make([]*os.File, spillNumBuckets) - defer func() { - for _, file := range files { - if file != nil { - file.Close() - } - } - }() - - // Create large batch to trigger buffer flushes - size := spillBufferSize * 2 - values := make([]int32, size) - for i := range values { - values[i] = int32(i) - } - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector(values, nil, proc.Mp()) - bat.SetRowCount(size) - - conditions := []*plan.Expr{ - { - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{ - Col: &plan.ColRef{ColPos: 0}, - }, - }, - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - ctr := &container{spillUUID: t.Name()} - - _, err := ctr.initSpillExprExecs(proc, conditions) - require.NoError(t, err) - err = ctr.spillBatchBounded(proc, bat, files, ctr.spillExprExecs, analyzer, false) - require.NoError(t, err) - -} - -func TestAppendBuildBatchWithNulls(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - files := make([]*os.File, spillNumBuckets) - defer func() { - for _, file := range files { - if file != nil { - file.Close() - } - } - }() - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3, 4}, []uint64{1}, proc.Mp()) // null at index 1 - bat.SetRowCount(4) - - conditions := []*plan.Expr{ - { - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{ - Col: &plan.ColRef{ColPos: 0}, - }, - }, - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - ctr := &container{spillUUID: t.Name()} - - _, err := ctr.initSpillExprExecs(proc, conditions) - require.NoError(t, err) - err = ctr.spillBatchBounded(proc, bat, files, ctr.spillExprExecs, analyzer, false) - require.NoError(t, err) - -} - -func TestAppendBuildBatchMultiColumn(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - files := make([]*os.File, spillNumBuckets) - defer func() { - for _, file := range files { - if file != nil { - file.Close() - } - } - }() - - bat := batch.NewWithSize(2) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()) - bat.Vecs[1] = testutil.MakeVarcharVector([]string{"a", "b", "c"}, nil, proc.Mp()) - bat.SetRowCount(3) - - conditions := []*plan.Expr{ - { - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{ - Col: &plan.ColRef{ColPos: 0}, - }, - }, - { - Typ: plan.Type{Id: int32(types.T_varchar)}, - Expr: &plan.Expr_Col{ - Col: &plan.ColRef{ColPos: 1}, - }, - }, - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - ctr := &container{spillUUID: t.Name()} - - _, err := ctr.initSpillExprExecs(proc, conditions) - require.NoError(t, err) - err = ctr.spillBatchBounded(proc, bat, files, ctr.spillExprExecs, analyzer, false) - require.NoError(t, err) - -} - -func TestShouldSpillBatchesRowThreshold(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - hb := &HashBuild{ - IsShuffle: true, - SpillThreshold: 10, // Small row threshold - NeedHashMap: true, - } - hb.ctr.setSpillThreshold(10) - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()) - bat.SetRowCount(3) - hb.ctr.hashmapBuilder.Batches.Buf = []*batch.Batch{bat} - hb.ctr.hashmapBuilder.InputBatchRowCount = bat.RowCount() - - require.False(t, hb.shouldSpillBatches()) - - // Add more batches to exceed threshold - for i := 0; i < 10; i++ { - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{int32(i)}, nil, proc.Mp()) - bat.SetRowCount(1) - hb.ctr.hashmapBuilder.Batches.Buf = append(hb.ctr.hashmapBuilder.Batches.Buf, bat) - hb.ctr.hashmapBuilder.InputBatchRowCount += bat.RowCount() - } - - require.True(t, hb.shouldSpillBatches()) -} - -func TestShouldSpillBatchesMemThreshold(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - hb := &HashBuild{ - IsShuffle: true, - SpillThreshold: 1024 * 1024, // 1MB - NeedHashMap: true, - } - hb.ctr.setSpillThreshold(1024 * 1024) - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp()) - bat.SetRowCount(2) - hb.ctr.hashmapBuilder.Batches.Buf = []*batch.Batch{bat} - - require.False(t, hb.shouldSpillBatches()) -} - -func TestHashWithConstVector(t *testing.T) { - mp := mpool.MustNewZero() - - vec := testutil.MakeInt32Vector([]int32{42}, nil, mp) - vec.SetClass(vector.CONSTANT) - - hashValues := make([]uint64, 10) - computeXXHash([]*vector.Vector{vec}, hashValues) - - // All values should be the same for const vector - for i := 1; i < len(hashValues); i++ { - require.Equal(t, hashValues[0], hashValues[i]) - } -} - -func TestHashMultiColumnCombinations(t *testing.T) { - mp := mpool.MustNewZero() - - vec1 := testutil.MakeInt32Vector([]int32{1, 1, 2}, nil, mp) - vec2 := testutil.MakeVarcharVector([]string{"a", "b", "a"}, nil, mp) - - hashValues := make([]uint64, 3) - computeXXHash([]*vector.Vector{vec1, vec2}, hashValues) - - // Different combinations should produce different hashes - require.NotEqual(t, hashValues[0], hashValues[1]) - require.NotEqual(t, hashValues[0], hashValues[2]) -} - -func TestAppendBuildBatchSingleBucket(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - files := make([]*os.File, spillNumBuckets) - defer func() { - for _, file := range files { - if file != nil { - file.Close() - } - } - }() - - // Single value should go to one bucket - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1}, nil, proc.Mp()) - bat.SetRowCount(1) - - conditions := []*plan.Expr{ - { - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{ - Col: &plan.ColRef{ColPos: 0}, - }, - }, - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - ctr := &container{spillUUID: t.Name()} - - _, err := ctr.initSpillExprExecs(proc, conditions) - require.NoError(t, err) - err = ctr.spillBatchBounded(proc, bat, files, ctr.spillExprExecs, analyzer, false) - require.NoError(t, err) - - fileCount := 0 - for _, file := range files { - if file != nil { - fileCount++ - } - } - require.Equal(t, 1, fileCount) -} - -func TestSpillScratchReuse(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - files := make([]*os.File, spillNumBuckets) - defer func() { - for _, file := range files { - if file != nil { - file.Close() - } - } - }() - - conditions := []*plan.Expr{ - { - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{ - Col: &plan.ColRef{ColPos: 0}, - }, - }, - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - ctr := &container{spillUUID: t.Name()} - - _, err := ctr.initSpillExprExecs(proc, conditions) - require.NoError(t, err) - - // First batch - bat1 := batch.NewWithSize(1) - bat1.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp()) - bat1.SetRowCount(2) - - err = ctr.spillBatchBounded(proc, bat1, files, ctr.spillExprExecs, analyzer, false) - require.NoError(t, err) - hashCapacity := cap(ctr.spillHashValues) - rowIDCapacity := cap(ctr.spillBucketRowIds) - require.Positive(t, hashCapacity) - require.Positive(t, rowIDCapacity) - - // An equal-size batch reuses the retained hash and row-id scratch. - bat2 := batch.NewWithSize(1) - bat2.Vecs[0] = testutil.MakeInt32Vector([]int32{3, 4}, nil, proc.Mp()) - bat2.SetRowCount(2) - - err = ctr.spillBatchBounded(proc, bat2, files, ctr.spillExprExecs, analyzer, false) - require.NoError(t, err) - require.Equal(t, hashCapacity, cap(ctr.spillHashValues)) - require.Equal(t, rowIDCapacity, cap(ctr.spillBucketRowIds)) -} - -func TestSpillExpressionLeaseRetainsLargeBatchHighWater(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget := process.MustNewHashBuildBudget(256<<20, 256<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - - files := make([]*os.File, spillNumBuckets) - defer func() { - for _, file := range files { - if file != nil { - _ = file.Close() - } - } - }() - expr := makeExpressionLeaseTestExpr(t, proc) - ctr := &container{spillUUID: t.Name()} - ctr.hashmapBuilder.setBudget(generation) - executors, err := ctr.initSpillExprExecs(proc, []*plan.Expr{expr}) - require.NoError(t, err) - require.NotNil(t, ctr.spillExprLease) - defer ctr.freeSpillExprExecs() - defer ctr.dropSpillScratchBuffers() - defer ctr.releaseSpillScratchReservation() - - analyzer := process.NewAnalyzer(0, false, false, "test") - large := makeExpressionLeaseTestBatch(proc, colexec.DefaultBatchSize) - defer large.Clean(proc.Mp()) - require.NoError(t, ctr.spillBatchBounded(proc, large, files, executors, analyzer, false)) - largeReserved := ctr.spillExprLease.Reserved() - require.Positive(t, largeReserved) - - small := makeExpressionLeaseTestBatch(proc, 1) - defer small.Clean(proc.Mp()) - require.NoError(t, ctr.spillBatchBounded(proc, small, files, executors, analyzer, false)) - require.Equal(t, largeReserved, ctr.spillExprLease.Reserved(), - "a small spill batch must not release retained executor headroom") - retained, ok := ctr.spillExprLease.Retained() - require.True(t, ok) - require.LessOrEqual(t, retained, ctr.spillExprLease.Reserved()) -} - -func TestSpillWriteCoalescesAcrossBatches(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget := process.MustNewHashBuildBudget(8<<20, 8<<20) - generation, err := budget.OpenGeneration(8 << 20) - require.NoError(t, err) - defer generation.Close() - files := make([]*os.File, spillNumBuckets) - defer func() { - for _, file := range files { - if file != nil { - file.Close() - } - } - }() - conditions := []*plan.Expr{{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - }} - ctr := &container{spillUUID: t.Name()} - ctr.hashmapBuilder.setBudget(generation) - _, err = ctr.initSpillExprExecs(proc, conditions) - require.NoError(t, err) - analyzer := process.NewAnalyzer(0, false, false, "test") - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 1, 1}, nil, proc.Mp()) - bat.SetRowCount(3) - defer bat.Clean(proc.Mp()) - for i := 0; i < 2; i++ { - require.NoError(t, ctr.spillBatchBounded(proc, bat, files, ctr.spillExprExecs, analyzer, false)) - } - var pending int - for i := range ctr.spillBucketWriteBufs { - pending += ctr.spillBucketWriteBufs[i].Len() - } - require.Positive(t, pending) - var file *os.File - for _, f := range files { - if f != nil { - file = f - break - } - } - require.NotNil(t, file) - stat, err := file.Stat() - require.NoError(t, err) - require.Zero(t, stat.Size(), "records stay pending until the handoff flush") - require.NoError(t, ctr.flushSpillBuffers(proc, files, analyzer)) - stat, err = file.Stat() - require.NoError(t, err) - require.Positive(t, stat.Size()) - for i := range ctr.spillBucketWriteBufs { - require.Zero(t, ctr.spillBucketWriteBufs[i].Len()) - } - _, err = file.Seek(0, io.SeekStart) - require.NoError(t, err) - reader := bufio.NewReader(file) - var totalRows int64 - for { - var header [16]byte - _, err = io.ReadFull(reader, header[:]) - if err == io.EOF { - break - } - require.NoError(t, err) - cnt := types.DecodeInt64(header[:8]) - payload := types.DecodeInt64(header[8:]) - require.GreaterOrEqual(t, cnt, int64(0)) - require.GreaterOrEqual(t, payload, int64(0)) - _, err = io.CopyN(io.Discard, reader, payload) - require.NoError(t, err) - var magic [8]byte - _, err = io.ReadFull(reader, magic[:]) - require.NoError(t, err) - require.Equal(t, uint64(spillMagic), types.DecodeUint64(magic[:])) - totalRows += cnt - } - require.Equal(t, int64(6), totalRows) - scratchPeak := analyzer.GetOpStats().ExtraStats["HashBuildSpillScratchPeakBytes"] - require.GreaterOrEqual(t, scratchPeak, hashBuildStatInt64(ctr.spillScratchReservation.Size())) - require.Greater(t, scratchPeak, hashBuildStatInt64(ctr.spillScratchBase), - "scratch peak must include retained coalesce buffers above the base lease") - for _, f := range files { - if f != nil { - _ = f.Close() - } - } - if ctr.spillBundle != nil { - ctr.spillBundle.release() - ctr.spillBundle = nil - } - ctr.dropSpillScratchBuffers() - ctr.releaseSpillScratchReservation() - require.Zero(t, generation.Used()) -} - -func TestSpillScratchBudgetDoesNotDoubleChargeRetainedSource(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := batch.NewWithSize(1) - values := make([]int32, colexec.DefaultBatchSize) - for i := range values { - values[i] = int32(i) - } - bat.Vecs[0] = testutil.MakeInt32Vector(values, nil, proc.Mp()) - bat.SetRowCount(len(values)) - defer bat.Clean(proc.Mp()) - - fullNeed, err := spillScratchBudgetBytes(bat, false) - require.NoError(t, err) - retainedNeed, err := spillScratchBudgetBytes(bat, true) - require.NoError(t, err) - source := uint64(bat.Allocated()) - require.Equal(t, source, fullNeed-retainedNeed) - require.Positive(t, retainedNeed) - - // The retained source has its own batch reservation. Only the incremental - // scratch must be admitted when the real spill path starts. - proofCap := source + fullNeed - 1 - proofBudget := process.MustNewHashBuildBudget(proofCap, proofCap) - proofGeneration, err := proofBudget.OpenGeneration(1) - require.NoError(t, err) - proofSource, err := proofGeneration.Reserve(source) - require.NoError(t, err) - _, err = proofGeneration.Reserve(fullNeed) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - proofScratch, err := proofGeneration.Reserve(retainedNeed) - require.NoError(t, err) - proofScratch.Release() - proofSource.Release() - require.Zero(t, proofGeneration.Used()) -} - -func TestSpillScratchLazyGrowSucceeds(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := batch.NewWithSize(1) - values := make([]int32, colexec.DefaultBatchSize) - for i := range values { - values[i] = int32(i) - } - bat.Vecs[0] = testutil.MakeInt32Vector(values, nil, proc.Mp()) - bat.SetRowCount(len(values)) - defer bat.Clean(proc.Mp()) - - need, err := spillScratchBudgetBytes(bat, true) - require.NoError(t, err) - require.Greater(t, need, uint64(1)) - source := uint64(bat.Allocated()) - const slack = uint64(2 << 20) - capBytes := source + need + slack - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - retainedToken, err := generation.Reserve(source) - require.NoError(t, err) - defer retainedToken.Release() - scratchToken, err := generation.Reserve(need - 1) - require.NoError(t, err) - - files := make([]*os.File, spillNumBuckets) - defer func() { - for _, file := range files { - if file != nil { - _ = file.Close() - } - } - }() - conditions := []*plan.Expr{{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - }} - ctr := &container{ - spillUUID: t.Name(), - spillScratchReservation: scratchToken, - spillScratchBase: need - 1, - } - ctr.hashmapBuilder.setBudget(generation) - _, err = ctr.initSpillExprExecs(proc, conditions) - require.NoError(t, err) - defer ctr.freeSpillExprExecs() - defer ctr.dropSpillScratchBuffers() - defer ctr.releaseSpillScratchReservation() - - analyzer := process.NewAnalyzer(0, false, false, "test") - require.NoError(t, ctr.spillBatchBounded(proc, bat, files, ctr.spillExprExecs, analyzer, true)) - require.Equal(t, need, ctr.spillScratchBase) - require.GreaterOrEqual(t, scratchToken.Size(), need) - require.Equal(t, int64(1), analyzer.GetOpStats().ExtraStats["HashBuildSpillScratchGrowCount"]) - require.Equal(t, int64(1), analyzer.GetOpStats().ExtraStats["HashBuildSpillScratchGrowBytes"]) - require.NoError(t, ctr.flushSpillBuffers(proc, files, analyzer)) -} - -func TestSpillScratchLazyGrowRejectPreservesRetainedSource(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3, 4}, nil, proc.Mp()) - bat.SetRowCount(4) - defer bat.Clean(proc.Mp()) - originalAllocated := bat.Allocated() - - need, err := spillScratchBudgetBytes(bat, true) - require.NoError(t, err) - require.Greater(t, need, uint64(1)) - source := uint64(bat.Allocated()) - capBytes := source + need - 1 - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - sourceToken, err := generation.Reserve(source) - require.NoError(t, err) - scratchToken, err := generation.Reserve(need - 1) - require.NoError(t, err) - - files := make([]*os.File, spillNumBuckets) - ctr := &container{ - spillUUID: t.Name(), - spillScratchReservation: scratchToken, - spillScratchBase: need - 1, - } - ctr.hashmapBuilder.setBudget(generation) - analyzer := process.NewAnalyzer(0, false, false, "lazy spill reject") - err = ctr.spillBatchBounded(proc, bat, files, nil, analyzer, true) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Equal(t, int64(1), analyzer.GetOpStats().ExtraStats["HashBuildSpillScratchGrowRejects"]) - require.Equal(t, need-1, scratchToken.Size()) - require.Equal(t, capBytes, generation.Used()) - require.Equal(t, 4, bat.RowCount()) - require.Equal(t, originalAllocated, bat.Allocated()) - require.Nil(t, ctr.spillHashValues) - require.Nil(t, ctr.spillBucketRowIds) - for _, file := range files { - require.Nil(t, file) - } - - ctr.releaseSpillScratchReservation() - sourceToken.Release() - require.Zero(t, generation.Used()) -} - -func TestFlushSpillBuffersCancellationDiscardsPendingWrites(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - ctx, cancel := context.WithCancelCause(proc.Ctx) - process.ReplacePipelineCtx(proc, ctx, cancel) - - ctr := &container{} - for _, bucket := range []int{0, spillNumBuckets - 1} { - _, err := ctr.spillBucketWriteBufs[bucket].Write([]byte("pending")) - require.NoError(t, err) - ctr.spillBucketWriteRows[bucket] = 1 - } - proc.Cancel(context.Canceled) - - err := ctr.flushSpillBuffers(proc, nil, process.NewAnalyzer(0, false, false, "test")) - require.ErrorIs(t, err, context.Canceled) - for bucket := 0; bucket < spillNumBuckets; bucket++ { - require.Zero(t, ctr.spillBucketWriteBufs[bucket].Len()) - require.Zero(t, ctr.spillBucketWriteRows[bucket]) - } -} - -func TestSpillMaterializedBytesDoesNotScaleShuffledConstVector(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - const sourceRows = 32 * 1024 - bat := batch.NewWithSize(2) - values := make([]int32, sourceRows) - for i := range values { - values[i] = int32(i) - } - bat.Vecs[0] = testutil.MakeInt32Vector(values, nil, proc.Mp()) - var err error - bat.Vecs[1], err = vector.NewConstBytes( - types.T_varchar.ToType(), - []byte("test create big fulltext index"), - sourceRows, - proc.Mp(), - ) - require.NoError(t, err) - bat.SetRowCount(sourceRows) - defer bat.Clean(proc.Mp()) - - // Batch.Shuffle intentionally leaves a const vector untouched while it - // changes the batch cardinality. This is the shape produced by the failed - // generate_series + const-varchar BVT query. - require.NoError(t, bat.Shuffle([]int64{0}, proc.Mp())) - require.Equal(t, 1, bat.RowCount()) - require.Equal(t, sourceRows, bat.Vecs[1].Length()) - - legacySource := uint64(bat.Allocated()) - if size := uint64(bat.Size()); size > legacySource { - legacySource = size - } - legacyScaled := legacySource * uint64(colexec.DefaultBatchSize) - legacyMetadata, ok := retainedMetadataAllowance(bat) - require.True(t, ok) - legacyScaled += legacyMetadata * uint64(colexec.DefaultBatchSize) - legacyNeed, err := spillPeakBudgetFor(uint64(colexec.DefaultBatchSize), 0, legacyScaled, uint64(len(bat.Vecs))) - require.NoError(t, err) - require.Greater(t, legacyNeed, uint64(10<<30), - "the old logical-size extrapolation must reproduce the false 10 GiB rejection") - - materialized, err := spillMaterializedBytes(bat) - require.NoError(t, err) - wantMaterialized := uint64(types.T_int32.ToType().TypeSize()) + - uint64(types.T_varchar.ToType().TypeSize()) + - uint64(len("test create big fulltext index")) - require.Equal(t, wantMaterialized, materialized, - "lazy admission must use the batch's live rows and const payload") - - directNeed, err := spillBudgetBytes(bat) - require.NoError(t, err) - require.Less(t, directNeed, uint64(16<<20), - "lazy scratch admission must not scale stale logical length") -} - -func TestSpillMaterializedBytesDoesNotScaleRetainedVectorCapacity(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - const sourceRows = 32 * 1024 - bat := batch.NewWithSize(2) - values := make([]int32, sourceRows) - strings := make([]string, sourceRows) - for i := range values { - values[i] = int32(i) - strings[i] = "test create big fulltext index" - } - bat.Vecs[0] = testutil.MakeInt32Vector(values, nil, proc.Mp()) - bat.Vecs[1] = testutil.MakeVarcharVector(strings, nil, proc.Mp()) - bat.SetRowCount(sourceRows) - defer bat.Clean(proc.Mp()) - - // Reused shuffle and table-function batches keep their allocation while - // publishing a tiny final batch. Only the first row is live, but Allocated - // still describes the original 32K-row capacity. - bat.Vecs[0].SetLength(1) - bat.Vecs[1].SetLength(1) - bat.SetRowCount(1) - require.Greater(t, bat.Allocated(), 1<<20) - - materialized, err := spillMaterializedBytes(bat) - require.NoError(t, err) - wantMaterialized := uint64(types.T_int32.ToType().TypeSize()) + - uint64(types.T_varchar.ToType().TypeSize()) + - uint64(len(strings[0])) - require.Equal(t, wantMaterialized, materialized) - - directNeed, err := spillBudgetBytes(bat) - require.NoError(t, err) - require.Less(t, directNeed, uint64(16<<20), - "source capacity is charged once, never extrapolated per live row") -} - -func TestSpillProjectedSourceSkipsStaleNullVarlenaPayload(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeVarcharVector([]string{"x"}, []uint64{0}, proc.Mp()) - bat.SetRowCount(1) - defer bat.Clean(proc.Mp()) - - // A null append does not overwrite a reused varlen slot. Plant the stale - // non-inline header that such a slot can retain; UnionInt32 skips the null - // value, so this dead payload must not be projected into the spill batch. - values, _ := vector.MustVarlenaRawData(bat.Vecs[0]) - const staleLen = uint32(1 << 20) - values[0].SetOffsetLen(0, staleLen) - - source, err := spillMaterializedBytes(bat) - require.NoError(t, err) - require.Equal(t, uint64(bat.Vecs[0].GetType().TypeSize()), source) - - need, err := spillBudgetBytes(bat) - require.NoError(t, err) - require.Less(t, need, uint64(16<<20)) +type spillTestHarness struct { + op *HashBuild + proc *process.Process + generation *process.HashBuildBudgetGeneration + registry *mpool.AllocationAccountRegistry + account *mpool.AllocationAccount + files []*os.File } -func TestSpillMaterializedBytesBoundaryInputs(t *testing.T) { - source, err := spillMaterializedBytes(nil) - require.NoError(t, err) - require.Zero(t, source) - - invalid := batch.NewWithSize(1) - invalid.SetRowCount(1) - _, err = spillMaterializedBytes(invalid) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - _, err = spillBudgetBytes(invalid) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - +func newSpillTestHarness(t *testing.T, limit uint64) *spillTestHarness { + t.Helper() proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - short := batch.NewWithSize(1) - short.Vecs[0] = testutil.MakeVarcharVector([]string{"x"}, nil, proc.Mp()) - short.SetRowCount(2) - defer short.Clean(proc.Mp()) - _, err = spillMaterializedBytes(short) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - - constNull := batch.NewWithSize(1) - constNull.Vecs[0] = vector.NewConstNull(types.T_varchar.ToType(), 1, proc.Mp()) - constNull.SetRowCount(1) - defer constNull.Clean(proc.Mp()) - - source, err = spillMaterializedBytes(constNull) - require.NoError(t, err) - require.Equal(t, uint64(types.T_varchar.ToType().TypeSize()), source) -} - -func TestSpillBudgetArithmeticFailsClosed(t *testing.T) { - value, err := spillCheckedAdd(math.MaxUint64-1, 1) - require.NoError(t, err) - require.Equal(t, uint64(math.MaxUint64), value) - _, err = spillCheckedAdd(math.MaxUint64, 1) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - - value, err = spillCheckedMul(math.MaxUint64, 1) - require.NoError(t, err) - require.Equal(t, uint64(math.MaxUint64), value) - _, err = spillCheckedMul(math.MaxUint64, 2) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - - _, err = spillPeakBudgetFor(math.MaxUint64, 0, 0, 0) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - _, err = spillPeakBudgetFor(0, math.MaxUint64, 1, 0) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - _, err = spillPeakBudgetFor(0, 0, math.MaxUint64, 0) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) -} - -func TestSpillCapacityReplacementOverlapChargesOldArrays(t *testing.T) { - got, err := spillCapacityReplacementOverlap(16, 4, 8, 8, 2) - require.NoError(t, err) - require.Equal(t, uint64(8*8+8*4+2*8), got) - - got, err = spillCapacityReplacementOverlap(8, 2, 8, 8, 2) - require.NoError(t, err) - require.Zero(t, got) - - _, err = spillCapacityReplacementOverlap(-1, 0, 0, 0, 0) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - _, err = spillCapacityReplacementOverlap(math.MaxInt, 0, math.MaxInt-1, 0, 0) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) -} - -func TestSpillReplacementPeakReusesHighWaterLease(t *testing.T) { - budget := process.MustNewHashBuildBudget(120, 120) + budget := process.MustNewHashBuildBudget(limit, limit) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - token, err := generation.Reserve(100) + registry, err := mpool.NewAllocationAccountRegistry(1, 256) require.NoError(t, err) - ctr := container{ - hashmapBuilder: HashmapBuilder{budget: generation}, - spillScratchReservation: token, - spillScratchBase: 100, - } - analyzer := process.NewAnalyzer(0, false, false, "replacement peak") - - oldSize, grew, err := ctr.growSpillScratchTransient(90, analyzer) - require.NoError(t, err) - require.False(t, grew) - require.Zero(t, oldSize) - require.Equal(t, uint64(100), generation.Used()) - - oldSize, grew, err = ctr.growSpillScratchTransient(110, analyzer) + account, err := registry.OpenWithController(limit, generation) require.NoError(t, err) - require.True(t, grew) - require.Equal(t, uint64(100), oldSize) - require.Equal(t, uint64(110), generation.Used()) - require.Equal(t, int64(110), analyzer.GetOpStats().ExtraStats["HashBuildSpillScratchPeakBytes"]) - require.NoError(t, ctr.restoreSpillScratchTransient(oldSize, grew)) - require.Equal(t, uint64(100), generation.Used()) - require.NoError(t, ctr.restoreSpillScratchTransient(0, false)) - - _, grew, err = ctr.growSpillScratchTransient(121, analyzer) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.False(t, grew) - require.Equal(t, uint64(100), generation.Used()) - - token.Release() - require.Zero(t, generation.Used()) -} - -func TestSpillPeakChargesSerializedPayloadOnce(t *testing.T) { - const ( - rows = uint64(8192) - inputBytes = uint64(3 << 20) - selectedBytes = uint64(5 << 20) - ) - got, err := spillPeakBudgetFor(rows, inputBytes, selectedBytes, 0) + op := &HashBuild{NeedHashMap: true} + require.NoError(t, op.SetAllocationAccount(account)) + op.ctr.hashmapBuilder.setBudget(generation) + op.ctr.spillUUID = t.Name() + return &spillTestHarness{ + op: op, + proc: proc, + generation: generation, + registry: registry, + account: account, + files: make([]*os.File, spillNumBuckets), + } +} + +func (h *spillTestHarness) close(t *testing.T) { + t.Helper() + h.op.ctr.dropSpillScratchBuffers() + h.op.ctr.freeSpillExprExecs() + for _, file := range h.files { + if file != nil { + require.NoError(t, file.Close()) + } + } + if h.op.ctr.spillBundle != nil { + h.op.ctr.spillBundle.release() + h.op.ctr.spillBundle = nil + } + require.Zero(t, h.account.Snapshot().Used) + require.Zero(t, h.generation.Used()) + require.NoError(t, h.op.ClearAllocationAccount(h.account)) + terminal, first, err := h.registry.CompleteTerminal(h.account) require.NoError(t, err) - want := rows*12 + inputBytes + selectedBytes + selectedBytes + 64*1024 - require.Equal(t, want, got) + require.True(t, first) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) + h.proc.Free() } -func TestMarshalSpillRecordPreallocatesSinglePayload(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - bat := batch.NewWithSize(1) - var err error - bat.Vecs[0], err = vector.NewConstBytes( - types.T_varchar.ToType(), make([]byte, 4<<20), 64, proc.Mp(), - ) - require.NoError(t, err) - bat.SetRowCount(64) - defer bat.Clean(proc.Mp()) - - buf := bytes.NewBuffer(make([]byte, 0, 1<<20)) - _, err = marshalSpillRecord(bat, buf) - require.NoError(t, err) - base := uint64(bat.Allocated()) - if size := uint64(bat.Size()); size > base { - base = size +func spillFileRows(t *testing.T, files []*os.File) int64 { + t.Helper() + var total int64 + for _, file := range files { + if file == nil { + continue + } + _, err := file.Seek(0, io.SeekStart) + require.NoError(t, err) + reader := bufio.NewReader(file) + for { + var header [16]byte + _, err = io.ReadFull(reader, header[:]) + if err == io.EOF { + break + } + require.NoError(t, err) + rows := types.DecodeInt64(header[:8]) + payload := types.DecodeInt64(header[8:]) + require.GreaterOrEqual(t, rows, int64(0)) + require.GreaterOrEqual(t, payload, int64(0)) + _, err = io.CopyN(io.Discard, reader, payload) + require.NoError(t, err) + var magic [8]byte + _, err = io.ReadFull(reader, magic[:]) + require.NoError(t, err) + require.Equal(t, uint64(spillMagic), types.DecodeUint64(magic[:])) + total += rows + } } - require.Equal(t, base+128+24, uint64(buf.Cap())) - - small := batch.NewWithSize(1) - small.Vecs[0], err = vector.NewConstBytes( - types.T_varchar.ToType(), make([]byte, 1024), 1, proc.Mp(), - ) - require.NoError(t, err) - small.SetRowCount(1) - defer small.Clean(proc.Mp()) - _, err = marshalSpillRecord(small, buf) - require.NoError(t, err, "a retained large serialization buffer must be reusable for a smaller batch") + return total } -func TestSpillLazyReservationBoundaryInputs(t *testing.T) { - budget := process.MustNewHashBuildBudget(1, 1) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - ctr := &container{} - ctr.hashmapBuilder.setBudget(generation) - analyzer := process.NewAnalyzer(0, false, false, "spill reservation boundary") - - require.NoError(t, ctr.spillBatchBounded(nil, nil, nil, nil, analyzer, false)) - require.Nil(t, ctr.spillScratchReservation) - - invalid := batch.NewWithSize(1) - invalid.SetRowCount(1) - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - err = ctr.spillBatchBounded( - proc, invalid, make([]*os.File, spillNumBuckets), nil, analyzer, false) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - require.Nil(t, ctr.spillScratchReservation) - require.Zero(t, generation.Used()) +func TestComputeXXHashBuild(t *testing.T) { + mp := mpool.MustNewZero() + first := testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, mp) + second := testutil.MakeVarcharVector([]string{"a", "b", "c"}, nil, mp) + defer first.Free(mp) + defer second.Free(mp) + hashes := make([]uint64, 3) + computeXXHash([]*vector.Vector{first, second}, hashes) + require.NotEqual(t, hashes[0], hashes[1]) + + constant := testutil.MakeInt32Vector([]int32{5}, nil, mp) + defer constant.Free(mp) + constant.SetClass(vector.CONSTANT) + computeXXHash([]*vector.Vector{constant}, hashes) + require.Equal(t, hashes[0], hashes[1]) + require.Equal(t, hashes[1], hashes[2]) } -func TestSpillMaterializedBytesFollowsConstUnionSemantics(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - payload := make([]byte, 1<<20) - for i := range payload { - payload[i] = 'x' - } - const rows = 64 - source := batch.NewWithSize(1) - var err error - source.Vecs[0], err = vector.NewConstBytes(types.T_varchar.ToType(), payload, rows, proc.Mp()) - require.NoError(t, err) - source.SetRowCount(rows) - defer source.Clean(proc.Mp()) - - directBytes, err := spillMaterializedBytes(source) - require.NoError(t, err) - require.Equal(t, - uint64(rows*types.T_varchar.ToType().TypeSize()+len(payload)), - directBytes, - "direct UnionInt32 copies one payload and broadcasts its descriptor") - - selected := batch.NewWithSize(1) - selected.Vecs[0] = vector.NewVec(types.T_varchar.ToType()) - defer selected.Clean(proc.Mp()) - sels := make([]int32, rows) - for i := range sels { - sels[i] = int32(i) +func TestShouldSpillBatches(t *testing.T) { + bat := batch.NewWithSize(0) + bat.SetRowCount(2) + op := &HashBuild{IsShuffle: true, NeedHashMap: true} + op.ctr.setSpillThreshold(1) + op.ctr.hashmapBuilder.Batches.Buf = []*batch.Batch{bat} + op.ctr.hashmapBuilder.InputBatchRowCount = bat.RowCount() + require.True(t, op.shouldSpillBatches()) + op.IsShuffle = false + require.False(t, op.shouldSpillBatches()) + op.IsShuffle = true + op.NeedHashMap = false + require.False(t, op.shouldSpillBatches()) +} + +func TestAccountedSpillAdaptsAndPreservesRows(t *testing.T) { + h := newSpillTestHarness(t, 80<<10) + defer h.close(t) + values := make([]int64, colexec.DefaultBatchSize) + for i := range values { + values[i] = int64(i) } - require.NoError(t, selected.Vecs[0].PreExtend(rows, proc.Mp())) - require.NoError(t, selected.Vecs[0].UnionInt32(source.Vecs[0], sels, proc.Mp())) - selected.SetRowCount(rows) - require.GreaterOrEqual(t, directBytes, uint64(selected.Allocated())) - - var retained colexec.Batches - defer retained.Clean(proc.Mp()) - require.NoError(t, retained.CopyIntoBatches(source, proc)) - require.Len(t, retained.Buf, 1) - require.False(t, retained.Buf[0].Vecs[0].IsConst()) - retainedBytes, err := spillMaterializedBytes(retained.Buf[0]) + input := batch.NewWithSize(1) + input.Vecs[0] = testutil.MakeInt64Vector(values, nil, h.proc.Mp()) + input.SetRowCount(len(values)) + defer input.Clean(h.proc.Mp()) + executors, err := h.op.ctr.initSpillExprExecs( + h.proc, + []*plan.Expr{newExpr(0, types.T_int64.ToType())}, + ) require.NoError(t, err) - require.Equal(t, - uint64(rows)*(uint64(types.T_varchar.ToType().TypeSize())+uint64(len(payload))), - retainedBytes, - "the actual retained batch is non-const, so selection copies each value") - require.Greater(t, retainedBytes, directBytes*32) + analyzer := process.NewAnalyzer(0, false, false, "test") + require.NoError(t, h.op.ctr.spillBatchWithPressure( + h.proc, input, h.files, executors, analyzer, false, + )) + require.Positive(t, analyzer.GetOpStats().ExtraStats["HashBuildSpillInputReductions"]) + require.NoError(t, h.op.ctr.flushSpillBuffers(h.proc, h.files, analyzer)) + require.Equal(t, int64(len(values)), spillFileRows(t, h.files)) } -func TestSpillMaterializedEstimateCoversRetainedConstCopy(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - payload := make([]byte, 1<<10) - for i := range payload { - payload[i] = 'x' - } - const ( - inputRows = 4 - totalRows = inputRows * 2 +func TestAccountedSpillCoalescesWithoutDuplicateOwnership(t *testing.T) { + h := newSpillTestHarness(t, 8<<20) + defer h.close(t) + input := batch.NewWithSize(1) + input.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 1, 1}, nil, h.proc.Mp()) + input.SetRowCount(3) + defer input.Clean(h.proc.Mp()) + executors, err := h.op.ctr.initSpillExprExecs( + h.proc, + []*plan.Expr{newExpr(0, types.T_int32.ToType())}, ) - source := batch.NewWithSize(1) - var err error - source.Vecs[0], err = vector.NewConstBytes(types.T_varchar.ToType(), payload, inputRows, proc.Mp()) - require.NoError(t, err) - source.SetRowCount(inputRows) - defer source.Clean(proc.Mp()) - - var retained colexec.Batches - defer retained.Clean(proc.Mp()) - require.NoError(t, retained.CopyIntoBatches(source, proc)) - require.NoError(t, retained.CopyIntoBatches(source, proc)) - require.Len(t, retained.Buf, 1) - require.Equal(t, totalRows, retained.Buf[0].RowCount()) - require.False(t, retained.Buf[0].Vecs[0].IsConst(), - "CopyIntoBatches materializes const ingress as retained row values") - - estimated, err := spillMaterializedBytes(retained.Buf[0]) require.NoError(t, err) - selected := batch.NewWithSize(1) - selected.Vecs[0] = vector.NewVec(types.T_varchar.ToType()) - defer selected.Clean(proc.Mp()) - sels := make([]int32, totalRows) - for i := range sels { - sels[i] = int32(i) + analyzer := process.NewAnalyzer(0, false, false, "test") + for range 2 { + require.NoError(t, h.op.ctr.spillBatchWithPressure( + h.proc, input, h.files, executors, analyzer, false, + )) } - require.NoError(t, selected.Vecs[0].PreExtend(totalRows, proc.Mp())) - require.NoError(t, selected.Vecs[0].UnionInt32(retained.Buf[0].Vecs[0], sels, proc.Mp())) - selected.SetRowCount(totalRows) - require.GreaterOrEqual(t, estimated, uint64(selected.Allocated())) -} - -func TestSpillMaterializedEstimateFollowsFullBatchCloneToSemantics(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - makeConstBatch := func(payloadBytes int) *batch.Batch { - payload := make([]byte, payloadBytes) - for i := range payload { - payload[i] = 'x' + var pending int + for _, buffer := range h.op.ctr.spillAccountedBuckets { + if buffer != nil { + pending += buffer.Len() } - source := batch.NewWithSize(1) - var err error - source.Vecs[0], err = vector.NewConstBytes( - types.T_varchar.ToType(), - payload, - colexec.DefaultBatchSize, - proc.Mp(), - ) - require.NoError(t, err) - source.SetRowCount(colexec.DefaultBatchSize) - return source - } - - large := makeConstBatch(1 << 20) - defer large.Clean(proc.Mp()) - directNeed, err := spillBudgetBytes(large) - require.NoError(t, err) - require.Less(t, directNeed, uint64(16<<20)) - - var retainedLarge colexec.Batches - defer retainedLarge.Clean(proc.Mp()) - require.NoError(t, retainedLarge.CopyIntoBatches(large, proc)) - require.Len(t, retainedLarge.Buf, 1) - require.False(t, retainedLarge.Buf[0].Vecs[0].IsConst(), - "Batch.Dup delegates to Batch.CloneTo/UnionBatch and does not call Vector.Dup") - require.Equal(t, 1<<20, len(retainedLarge.Buf[0].Vecs[0].GetArea())) - retainedNeed, err := spillMaterializedBytes(retainedLarge.Buf[0]) - require.NoError(t, err) - wantRetained := uint64(colexec.DefaultBatchSize) * - (uint64(1<<20) + uint64(types.T_varchar.ToType().TypeSize())) - require.Equal(t, wantRetained, retainedNeed, - "the actual non-const retained batch materializes one MiB plus one descriptor per row") - - // Materialize a smaller exact-full-batch payload end-to-end without - // allocating the MiB case's 8 GiB selected area. - small := makeConstBatch(4 << 10) - defer small.Clean(proc.Mp()) - var retainedSmall colexec.Batches - defer retainedSmall.Clean(proc.Mp()) - require.NoError(t, retainedSmall.CopyIntoBatches(small, proc)) - require.False(t, retainedSmall.Buf[0].Vecs[0].IsConst()) - - estimated, err := spillMaterializedBytes(retainedSmall.Buf[0]) - require.NoError(t, err) - selected := batch.NewWithSize(1) - selected.Vecs[0] = vector.NewVec(types.T_varchar.ToType()) - defer selected.Clean(proc.Mp()) - sels := make([]int32, colexec.DefaultBatchSize) - for i := range sels { - sels[i] = int32(i) } - require.NoError(t, selected.Vecs[0].PreExtend(colexec.DefaultBatchSize, proc.Mp())) - require.NoError(t, selected.Vecs[0].UnionInt32(retainedSmall.Buf[0].Vecs[0], sels, proc.Mp())) - selected.SetRowCount(colexec.DefaultBatchSize) - require.Equal(t, colexec.DefaultBatchSize*(4<<10), len(selected.Vecs[0].GetArea())) - require.GreaterOrEqual(t, estimated, uint64(selected.Allocated())) + require.Positive(t, pending) + require.NoError(t, h.op.ctr.flushSpillBuffers(h.proc, h.files, analyzer)) + require.Equal(t, int64(6), spillFileRows(t, h.files)) + require.Equal(t, h.account.Snapshot().Used, h.generation.Used()) } -func TestSpillBatchLazyReservationFailsClosed(t *testing.T) { +func TestSpillWithoutAllocationAccountFailsClosed(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3, 4}, nil, proc.Mp()) - bat.SetRowCount(4) + bat := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 1, proc.Mp()) defer bat.Clean(proc.Mp()) - - need, err := spillBudgetBytes(bat) - require.NoError(t, err) - require.Positive(t, need) - budget := process.MustNewHashBuildBudget(need-1, need-1) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - - ctr := &container{} - ctr.hashmapBuilder.setBudget(generation) - files := make([]*os.File, spillNumBuckets) - err = ctr.spillBatchBounded( - proc, bat, files, nil, - process.NewAnalyzer(0, false, false, "direct spill reject"), false) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Nil(t, ctr.spillScratchReservation) - require.Nil(t, ctr.spillHashValues) - require.Nil(t, ctr.spillBucketRowIds) - require.Equal(t, 4, bat.RowCount()) - for _, file := range files { - require.Nil(t, file) - } - require.Zero(t, generation.Used()) -} - -func TestEnsureSpillFile(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - ctr := &container{spillUUID: "test_ensure"} - files := make([]*os.File, spillNumBuckets) - - // First call creates a file. - f, err := ctr.ensureSpillFile(proc, files, 3) - require.NoError(t, err) - require.NotNil(t, f) - require.Equal(t, f, files[3]) - defer f.Close() - - // Second call returns cached file. - f2, err := ctr.ensureSpillFile(proc, files, 3) - require.NoError(t, err) - require.Same(t, f, f2, "should return the same file object") - - // Different bucket creates a different file. - f3, err := ctr.ensureSpillFile(proc, files, 7) - require.NoError(t, err) - require.NotNil(t, f3) - require.NotEqual(t, f.Fd(), f3.Fd()) - defer f3.Close() - - // Untouched buckets remain nil. - require.Nil(t, files[0]) - require.Nil(t, files[1]) -} - -func TestCleanupSpillFiles(t *testing.T) { - // Create temp files to simulate spill fds. - var fds []*os.File - for i := 0; i < 3; i++ { - f, err := os.CreateTemp("", "test_cleanup_*") - require.NoError(t, err) - defer os.Remove(f.Name()) - fds = append(fds, f) - } - // Include a nil entry. - fds = append(fds, nil) - - hb := &HashBuild{ctr: container{spilledFds: fds}} - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - hb.cleanupSpillFiles(proc) - require.Nil(t, hb.ctr.spilledFds) - - // Verify all files are closed (writing should fail). - for _, f := range fds[:3] { - _, err := f.Write([]byte("x")) - require.Error(t, err, "file should be closed") - } + err := (&container{}).spillBatchBounded( + proc, + bat, + make([]*os.File, spillNumBuckets), + nil, + process.NewAnalyzer(0, false, false, "test"), + false, + ) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) } -func TestAccountedInitialSpillConvertsHeadroomToPhysicalOwnership(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - const limit = uint64(8 << 20) - budget := process.MustNewHashBuildBudget(limit, limit) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - registry, err := mpool.NewAllocationAccountRegistry(1, 128) - require.NoError(t, err) - account, err := registry.OpenWithController(limit, generation) - require.NoError(t, err) - var op HashBuild - op.NeedHashMap = true - require.NoError(t, op.SetAllocationAccount(account)) - ctr := &op.ctr - ctr.hashmapBuilder.setBudget(generation) - ctr.spillUUID = "accounted-initial-spill" - exprs := []*plan.Expr{newExpr(0, types.T_int64.ToType())} - executors, err := ctr.initSpillExprExecs(proc, exprs) - require.NoError(t, err) - require.True(t, ctr.spillExprAccounted) - input := testutil.NewBatch( - []types.Type{types.T_int64.ToType(), types.T_varchar.ToType()}, - true, - 1_024, - proc.Mp(), +func TestSpillMinimumUnitPressureIsControlled(t *testing.T) { + h := newSpillTestHarness(t, 1<<10) + defer h.close(t) + input := batch.NewWithSize(1) + input.Vecs[0] = testutil.MakeVarcharVector( + []string{strings.Repeat("x", 64<<10)}, nil, h.proc.Mp(), + ) + input.SetRowCount(1) + defer input.Clean(h.proc.Mp()) + executors, err := h.op.ctr.initSpillExprExecs( + h.proc, + []*plan.Expr{newExpr(0, types.T_varchar.ToType())}, ) - defer input.Clean(proc.Mp()) - headroom, err := spillBudgetBytes(input) - require.NoError(t, err) - ctr.spillScratchReservation, err = generation.Reserve(headroom) require.NoError(t, err) - ctr.spillScratchBase = headroom - ctr.spillScratchEmergency = true - files := make([]*os.File, spillNumBuckets) - analyzer := process.NewAnalyzer(0, false, false, "test") - require.NoError(t, ctr.spillBatchBounded( - proc, + err = h.op.ctr.spillBatchWithPressure( + h.proc, input, - files, + h.files, executors, - analyzer, + process.NewAnalyzer(0, false, false, "test"), false, - )) - require.Nil(t, ctr.spillScratchReservation) - snapshot := generation.Snapshot() - require.Equal(t, snapshot.AllocationUsed, snapshot.Used, - "external input ownership is transient and headroom is not stacked") - require.Positive(t, snapshot.AllocationUsed) - require.NotNil(t, ctr.spillAccountedWrite) - require.NoError(t, ctr.flushSpillBuffers(proc, files, analyzer)) - - ctr.dropSpillScratchBuffers() - ctr.freeSpillExprExecs() - for _, file := range files { - if file != nil { - _ = file.Close() - } - } - if ctr.spillBundle != nil { - ctr.spillBundle.release() - ctr.spillBundle = nil - } - require.Zero(t, account.Snapshot().Used) - require.Zero(t, generation.Used()) - require.NoError(t, op.ClearAllocationAccount(account)) - _, _, err = registry.CompleteTerminal(account) - require.NoError(t, err) -} - -func TestAccountedMinimumSpillHeadroomIsOneUnitNotWholeBatch(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - makeBatch := func(rows int) *batch.Batch { - values := make([]string, rows) - for i := range values { - values[i] = strings.Repeat("x", 4<<10) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeVarcharVector(values, nil, proc.Mp()) - bat.SetRowCount(rows) - return bat - } - one := makeBatch(1) - many := makeBatch(128) - defer one.Clean(proc.Mp()) - defer many.Clean(proc.Mp()) - exprs := []*plan.Expr{newExpr(0, types.T_varchar.ToType())} - oneUnit, err := spillMinimumUnitBudgetBytes(one, exprs) - require.NoError(t, err) - manyUnits, err := spillMinimumUnitBudgetBytes(many, exprs) - require.NoError(t, err) - require.Equal(t, oneUnit, manyUnits) - require.Positive(t, oneUnit) - legacyWholeBatch, err := spillBudgetBytes(many) - require.NoError(t, err) - require.Greater(t, legacyWholeBatch, manyUnits) -} - -func TestAccountedConcatSpillHeadroomUsesOnePhysicalRow(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - column := func(pos int32) *plan.Expr { - return &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_varchar), Width: types.MaxVarcharLen}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: pos}}, - } - } - expr, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "concat", - []*plan.Expr{column(0), column(1)}, - ) - require.NoError(t, err) - bat := batch.NewWithSize(2) - bat.Vecs[0] = testutil.MakeVarcharVector( - []string{strings.Repeat("a", 4<<10), "a"}, nil, proc.Mp()) - bat.Vecs[1] = testutil.MakeVarcharVector( - []string{"b", strings.Repeat("b", 4<<10)}, nil, proc.Mp()) - bat.SetRowCount(2) - defer bat.Clean(proc.Mp()) - - payload, err := spillExpressionPayloadBytes( - expr, - expr.GetF().GetArgs(), - bat, ) - require.NoError(t, err) - require.Equal(t, uint64((4<<10)+1), payload) - headroom, err := spillMinimumUnitBudgetBytes(bat, []*plan.Expr{expr}) - require.NoError(t, err) - require.Positive(t, headroom) + var minimum *MinimumAllocationPressureError + require.True(t, errors.As(err, &minimum), "unexpected error: %v", err) } -func TestAccountedInitialSpillReducesUnpublishedInputAndPreservesRows(t *testing.T) { +func TestWriteSpillPayloadCancellationStopsBeforeIO(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - const limit = uint64(80 << 10) - budget := process.MustNewHashBuildBudget(limit, limit) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - registry, err := mpool.NewAllocationAccountRegistry(1, 128) - require.NoError(t, err) - account, err := registry.OpenWithController(limit, generation) + ctx, cancel := context.WithCancelCause(proc.Ctx) + process.ReplacePipelineCtx(proc, ctx, cancel) + spillfs, err := proc.GetSpillFileService() require.NoError(t, err) - var op HashBuild - op.NeedHashMap = true - require.NoError(t, op.SetAllocationAccount(account)) - ctr := &op.ctr - ctr.hashmapBuilder.setBudget(generation) - ctr.spillUUID = "accounted-adaptive-spill" - exprs := []*plan.Expr{newExpr(0, types.T_int64.ToType())} - executors, err := ctr.initSpillExprExecs(proc, exprs) + file, err := spillfs.CreateFile(context.Background(), t.Name()) require.NoError(t, err) - values := make([]int64, colexec.DefaultBatchSize) - for i := range values { - values[i] = int64(i) - } - input := batch.NewWithSize(1) - input.Vecs[0] = testutil.MakeInt64Vector(values, nil, proc.Mp()) - input.SetRowCount(len(values)) - defer input.Clean(proc.Mp()) - files := make([]*os.File, spillNumBuckets) - analyzer := process.NewAnalyzer(0, false, false, "test") - require.NoError(t, ctr.spillBatchWithPressure( + defer func() { + require.NoError(t, file.Close()) + require.NoError(t, spillfs.RemoveFile(context.Background(), t.Name())) + }() + proc.Cancel(context.Canceled) + err = (&container{}).writeSpillPayload( proc, - input, - files, - executors, - analyzer, - false, - )) - require.Positive(t, - analyzer.GetOpStats().ExtraStats["HashBuildSpillInputReductions"]) - require.NoError(t, ctr.flushSpillBuffers(proc, files, analyzer)) - - var totalRows int64 - for _, file := range files { - if file == nil { - continue - } - _, err = file.Seek(0, io.SeekStart) - require.NoError(t, err) - reader := bufio.NewReader(file) - for { - var header [16]byte - _, err = io.ReadFull(reader, header[:]) - if err == io.EOF { - break - } - require.NoError(t, err) - rows := types.DecodeInt64(header[:8]) - payload := types.DecodeInt64(header[8:]) - require.NoError(t, func() error { - _, copyErr := io.CopyN(io.Discard, reader, payload+8) - return copyErr - }()) - totalRows += rows - } - } - require.Equal(t, int64(len(values)), totalRows) - - ctr.dropSpillScratchBuffers() - ctr.freeSpillExprExecs() - for _, file := range files { - if file != nil { - _ = file.Close() - } - } - if ctr.spillBundle != nil { - ctr.spillBundle.release() - ctr.spillBundle = nil - } - require.Zero(t, account.Snapshot().Used) - require.Zero(t, generation.Used()) - require.NoError(t, op.ClearAllocationAccount(account)) - _, _, err = registry.CompleteTerminal(account) - require.NoError(t, err) + file, + []byte("stale"), + 1, + process.NewAnalyzer(0, false, false, "test"), + ) + require.ErrorIs(t, err, context.Canceled) + info, statErr := file.Stat() + require.NoError(t, statErr) + require.Zero(t, info.Size()) } diff --git a/pkg/sql/colexec/hashbuild/types.go b/pkg/sql/colexec/hashbuild/types.go index 8be3e85b30ce9..37724acee6c6e 100644 --- a/pkg/sql/colexec/hashbuild/types.go +++ b/pkg/sql/colexec/hashbuild/types.go @@ -15,11 +15,11 @@ package hashbuild import ( - "bytes" "os" "sync" "sync/atomic" + "github.com/matrixorigin/matrixone/pkg/common/hashmap" "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/reuse" @@ -69,6 +69,7 @@ const ( HashBuildAllocationSiteBatchNulls HashBuildAllocationSiteBatchGrouping HashBuildAllocationSiteGroupSels + HashBuildAllocationSiteHashIterator ) // Runtime-filter keys and their published wire payload have lifetimes that @@ -119,14 +120,9 @@ type container struct { // input batch. spillBucketOffsets identifies each bucket's sub-slice; // keeping one array avoids the 32 independent append/growth paths used by // the old scatter implementation. - spillBucketRowIds []int32 - spillBucketCounts [spillNumBuckets]int32 - spillBucketOffsets [spillNumBuckets + 1]int32 - spillWriteBuf bytes.Buffer - // spillBucketWriteBufs coalesce serialized records across source batches. - // Each buffer is bounded by spillWriteCoalesceSize (plus bytes.Buffer's - // bounded growth slack), so fanout does not imply fanout-sized vectors. - spillBucketWriteBufs [spillNumBuckets]bytes.Buffer + spillBucketRowIds []int32 + spillBucketCounts [spillNumBuckets]int32 + spillBucketOffsets [spillNumBuckets + 1]int32 spillBucketWriteRows [spillNumBuckets]int64 spillKeyVecs []*vector.Vector spillBatchAllocation *vector.AllocationAccountSelection @@ -134,27 +130,9 @@ type container struct { spillAccountedWrite *mpool.AccountedBuffer spillAccountedBuckets [spillNumBuckets]*mpool.AccountedBuffer spillCoalesceDisabled bool - // spillScratchReservation is a query/CN-charged emergency lease retained - // while Shuffle build batches accumulate. It prevents retained copies from - // consuming the scratch required to recover from hard-budget rejection. - spillScratchReservation *process.HashBuildReservation - // spillScratchEmergency marks a lease pre-admitted by - // ensureDirectSpillScratchReservation or - // ensureRetainedSpillScratchReservation. An uncharged upstream batch may - // not grow beyond this lease. A retained batch may grow it because its - // source memory remains charged separately while the batch is drained. - spillScratchEmergency bool - // spillScratchBase is the retained scratch floor. Coalesce-buffer growth is - // charged on top and must never be mistaken for this floor. - spillScratchBase uint64 - // cached expression executors for spill (reused across batches) spillExprExecs []colexec.ExpressionExecutor - spillExprLease *ExpressionMemoryLease spillConditions []*plan.Expr - // spillExprAccounted distinguishes an exact executor set from a legacy set - // whose retained lease has not yet been installed. - spillExprAccounted bool } // spillFileBundle is deliberately owned by hashbuild. Build converts each @@ -230,7 +208,7 @@ func (b *spillFileBundle) growDisk(file *os.File, budget *process.HashBuildBudge b.mu.Lock() defer b.mu.Unlock() if b.released { - return 0, false, process.ErrHashBuildReservationInactive + return 0, false, process.ErrHashBuildSpillReservationInactive } if b.entries == nil { b.entries = make(map[*os.File]*spillFileEntry) @@ -328,28 +306,13 @@ func (hashBuild *HashBuild) GetOperatorBase() *vm.OperatorBase { return &hashBuild.OperatorBase } -func (hashBuild *HashBuild) AllocationAccountEnabled() bool { - // Activate one complete physical owner closure. An expression family whose - // call-scoped allocation ledger is not closed must keep the whole HashBuild - // on the legacy path; mixing an exact map/batch owner with an estimator-gated - // expression would reintroduce the false-rejection mechanism that this - // activation removes. - return hashBuild != nil && hashBuild.NeedHashMap && - expressionSetAllocationClosed(hashBuild.Conditions) -} - -func (hashBuild *HashBuild) AllocationAccountActivationBlocked() bool { - return hashBuild != nil && hashBuild.NeedHashMap && - !expressionSetAllocationClosed(hashBuild.Conditions) -} - // SetAllocationAccount selects immutable provenance for the hash-table owner // before Prepare. Compile invokes it once for each execution attempt; Reset // clears the selection only after producer or JoinMap ownership has moved on. func (hashBuild *HashBuild) SetAllocationAccount( account *mpool.AllocationAccount, ) error { - selection, err := vector.NewAllocationAccountSelectionWithBitmaps( + selection, err := vector.NewAllocationAccountSelection( account, HashBuildAllocationOwner, HashBuildSpillAllocationSiteSelectedData, @@ -367,7 +330,7 @@ func (hashBuild *HashBuild) SetAllocationAccount( return nil } -// SetAllocationAccount activates the physical allocation provenance shared by +// SetAllocationAccount installs the physical allocation provenance shared by // the producer HashBuild and SpillEngine rebuild builders. A builder is always // single-generation and clears the selection only after all owned resources // have either been freed or transferred to a JoinMap. @@ -390,7 +353,15 @@ func (hb *HashmapBuilder) SetAllocationAccount( if err != nil { return err } - batchSelection, err := vector.NewAllocationAccountSelectionWithBitmaps( + iteratorAllocation, err := hashmap.NewIteratorAllocation( + account, + HashBuildAllocationOwner, + HashBuildAllocationSiteHashIterator, + ) + if err != nil { + return err + } + batchSelection, err := vector.NewAllocationAccountSelection( account, HashBuildAllocationOwner, HashBuildAllocationSiteBatchData, @@ -401,7 +372,7 @@ func (hb *HashmapBuilder) SetAllocationAccount( if err != nil { return err } - uniqueKeySelection, err := vector.NewAllocationAccountSelectionWithBitmaps( + uniqueKeySelection, err := vector.NewAllocationAccountSelection( account, HashBuildAllocationOwner, HashBuildAllocationSiteUniqueKeyData, @@ -412,18 +383,11 @@ func (hb *HashmapBuilder) SetAllocationAccount( if err != nil { return err } - expressionAllocation, err := colexec.NewExpressionAllocationAccount( - account, - HashBuildAllocationOwner, - ) - if err != nil { - return err - } builder.mapAllocationAccount = account builder.mapAllocation = selection + builder.iteratorAllocation = iteratorAllocation builder.batchAllocation = batchSelection builder.uniqueKeyAllocation = uniqueKeySelection - builder.expressionAllocation = expressionAllocation return nil } @@ -471,9 +435,9 @@ func (hb *HashmapBuilder) ClearAllocationAccount( } builder.mapAllocationAccount = nil builder.mapAllocation = nil + builder.iteratorAllocation = nil builder.batchAllocation = nil builder.uniqueKeyAllocation = nil - builder.expressionAllocation = nil return nil } @@ -527,14 +491,13 @@ func (hashBuild *HashBuild) Reset(proc *process.Process, pipelineFailed bool, er hashBuild.publishBuildError(proc, err) } else { // Preserve the established nil JoinMap convention for a true empty - // build and for legacy cleanup paths that completed without a map. + // build and for cleanup paths that completed without a map. hashBuild.publishJoinMap(proc, nil) } } hashBuild.ctr.hashmapBuilder.Reset(proc, !mapSucceed) hashBuild.ctr.dropSpillScratchBuffers() - hashBuild.ctr.releaseSpillScratchReservation() // Only clean up build files when the join map was NOT successfully sent. // When mapSucceed=true, hashjoin owns the files and deletes them after reading. if !mapSucceed { @@ -577,7 +540,6 @@ func (hashBuild *HashBuild) Free(proc *process.Process, pipelineFailed bool, err hashBuild.ctr.hashmapBuilder.Free(proc) hashBuild.ctr.freeSpillExprExecs() hashBuild.ctr.dropSpillScratchBuffers() - hashBuild.ctr.releaseSpillScratchReservation() } func (hashBuild *HashBuild) logDiagnostics(proc *process.Process, pipelineFailed bool, err error) { @@ -618,13 +580,16 @@ func (hashBuild *HashBuild) publishJoinMap(proc *process.Process, jm *message.Jo if !atomic.CompareAndSwapUint32(&hashBuild.ctr.terminalPublished, 0, 1) { return false } - message.SendJoinMapResult( + if !message.SendJoinMapResult( message.NewJoinMapResult(jm), hashBuild.JoinMapTag, hashBuild.IsShuffle, hashBuild.ShuffleIdx, proc.GetMessageBoard(), - ) + ) { + atomic.StoreUint32(&hashBuild.ctr.terminalPublished, 0) + return false + } return true } @@ -632,13 +597,16 @@ func (hashBuild *HashBuild) publishBuildError(proc *process.Process, err error) if !atomic.CompareAndSwapUint32(&hashBuild.ctr.terminalPublished, 0, 1) { return false } - message.FinalizeJoinMapBuildError( + if !message.FinalizeJoinMapBuildError( proc.GetMessageBoard(), hashBuild.JoinMapTag, hashBuild.IsShuffle, hashBuild.ShuffleIdx, err, - ) + ) { + atomic.StoreUint32(&hashBuild.ctr.terminalPublished, 0) + return false + } return true } @@ -658,9 +626,8 @@ func (hashBuild *HashBuild) cleanupSpillFiles(proc *process.Process) { } } -// CleanCopiedBatchAt is the lifecycle hook used by bounded initial spill. -// HashBuild keeps this wrapper on the operator side so batch reservation -// ownership remains private to the hashbuild package. +// CleanCopiedBatchAt releases one retained build batch after it has been +// durably transferred to spill storage. func (hb *HashmapBuilder) CleanCopiedBatchAt(idx int, proc *process.Process) error { if idx < 0 || idx >= len(hb.Batches.Buf) { return process.ErrHashBuildBudgetInvalid @@ -676,22 +643,12 @@ func (hb *HashmapBuilder) CleanCopiedBatchAt(idx int, proc *process.Process) err hb.Batches.MemSize += int64(bat.Size()) } } - // CopyIntoBatches can coalesce several ingress batches into one physical - // batch (and can reorder a full batch around a partial tail), so an ingress - // reservation cannot be matched safely to Batches.Buf[idx]. Keep the - // conservative charges until the last physical batch has been dropped. - if len(hb.Batches.Buf) == 0 { - hb.releaseBatchReservations() - } return nil } // DrainCopiedBatches visits and then releases every retained physical build // batch. A failed visit leaves the current and remaining batches owned by the -// builder so its normal cleanup path can release them. CopyIntoBatches can -// coalesce several ingress batches into one physical batch, so the associated -// reservations are released together only after the final physical batch is -// destroyed. +// builder so its normal cleanup path can release them. func (hb *HashmapBuilder) DrainCopiedBatches( proc *process.Process, visit func(*batch.Batch) error, @@ -719,7 +676,6 @@ func (hb *HashmapBuilder) DrainCopiedBatches( } hb.Batches.Buf = nil hb.Batches.MemSize = 0 - hb.releaseBatchReservations() return nil } diff --git a/pkg/sql/colexec/hashjoin/allocation_test_helpers_test.go b/pkg/sql/colexec/hashjoin/allocation_test_helpers_test.go new file mode 100644 index 0000000000000..82d16b6684660 --- /dev/null +++ b/pkg/sql/colexec/hashjoin/allocation_test_helpers_test.go @@ -0,0 +1,38 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package hashjoin + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/stretchr/testify/require" +) + +type testAllocationOwner interface { + SetAllocationAccount(*mpool.AllocationAccount) error +} + +func installTestAllocation(t testing.TB, owners ...testAllocationOwner) *mpool.AllocationAccount { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.Open(1 << 60) + require.NoError(t, err) + for _, owner := range owners { + require.NoError(t, owner.SetAllocationAccount(account)) + } + return account +} diff --git a/pkg/sql/colexec/hashjoin/bitmap_mailbox.go b/pkg/sql/colexec/hashjoin/bitmap_mailbox.go new file mode 100644 index 0000000000000..519207b70754f --- /dev/null +++ b/pkg/sql/colexec/hashjoin/bitmap_mailbox.go @@ -0,0 +1,98 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package hashjoin + +import ( + "context" + "sync" + + "github.com/matrixorigin/matrixone/pkg/common/bitmap" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/sql/colexec" +) + +// BitmapMailbox is the single ownership boundary for the parallel right-join +// bitmap exchange. A successful Send transfers ownership to the mailbox. Once +// sealed, late senders retain ownership and normal operator cleanup frees it. +type BitmapMailbox struct { + mu sync.Mutex + sealed bool + ch chan *bitmap.Bitmap +} + +func NewBitmapMailbox(workers int) *BitmapMailbox { + if workers < 1 { + workers = 1 + } + return &BitmapMailbox{ch: make(chan *bitmap.Bitmap, workers)} +} + +func (m *BitmapMailbox) Send(value *bitmap.Bitmap) bool { + if m == nil { + return false + } + m.mu.Lock() + defer m.mu.Unlock() + if m.sealed { + return false + } + // Every non-merger publishes at most once and the mailbox capacity equals + // the worker count, so publication cannot block while holding mu. + m.ch <- value + return true +} + +func (m *BitmapMailbox) Receive( + ctx context.Context, +) (*bitmap.Bitmap, bool) { + if m == nil || ctx == nil { + return nil, false + } + select { + case <-ctx.Done(): + return nil, false + case value := <-m.ch: + return value, true + } +} + +// SealAndDrain makes cancellation order-independent. It owns and frees every +// value already transferred into the mailbox; concurrent or later Send calls +// fail and leave ownership with their sender. +func (m *BitmapMailbox) SealAndDrain(mp *mpool.MPool) { + if m == nil { + return + } + m.mu.Lock() + m.sealed = true + for { + select { + case value := <-m.ch: + colexec.FreeAccountedBitmap(value, mp) + default: + m.mu.Unlock() + return + } + } +} + +func (m *BitmapMailbox) Terminal() bool { + if m == nil { + return false + } + m.mu.Lock() + defer m.mu.Unlock() + return m.sealed && len(m.ch) == 0 +} diff --git a/pkg/sql/colexec/hashjoin/bitmap_mailbox_test.go b/pkg/sql/colexec/hashjoin/bitmap_mailbox_test.go new file mode 100644 index 0000000000000..763bf52d0f1b9 --- /dev/null +++ b/pkg/sql/colexec/hashjoin/bitmap_mailbox_test.go @@ -0,0 +1,105 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package hashjoin + +import ( + "sync" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/bitmap" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" + "github.com/stretchr/testify/require" +) + +func newMailboxAccountedBitmap( + t *testing.T, + mp *mpool.MPool, + account *mpool.AllocationAccount, +) *bitmap.Bitmap { + t.Helper() + value, err := colexec.NewAccountedBitmap( + 1024, + mp, + account, + hashbuild.HashBuildAllocationOwner, + hashJoinAllocationSiteMatchedRows, + ) + require.NoError(t, err) + return value +} + +func TestBitmapMailboxSealOwnsQueuedAndRejectsLateTransfer(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + mp := mpool.MustNewZero() + mailbox := NewBitmapMailbox(2) + + queued := newMailboxAccountedBitmap(t, mp, account) + require.True(t, mailbox.Send(queued)) + mailbox.SealAndDrain(mp) + require.Zero(t, account.Snapshot().Used) + require.Empty(t, mailbox.ch) + + late := newMailboxAccountedBitmap(t, mp, account) + require.False(t, mailbox.Send(late)) + require.NotZero(t, account.Snapshot().Used) + colexec.FreeAccountedBitmap(late, mp) + require.Zero(t, account.Snapshot().Used) + + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + require.Zero(t, mp.CurrNB()) +} + +func TestBitmapMailboxConcurrentSealPreservesSingleOwner(t *testing.T) { + const workers = 16 + registry, err := mpool.NewAllocationAccountRegistry(1, workers+1) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + mp := mpool.MustNewZero() + mailbox := NewBitmapMailbox(workers) + values := make([]*bitmap.Bitmap, workers) + for i := range values { + values[i] = newMailboxAccountedBitmap(t, mp, account) + } + + start := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(workers) + for _, value := range values { + go func(value *bitmap.Bitmap) { + defer wg.Done() + <-start + if !mailbox.Send(value) { + colexec.FreeAccountedBitmap(value, mp) + } + }(value) + } + close(start) + mailbox.SealAndDrain(mp) + wg.Wait() + mailbox.SealAndDrain(mp) + + require.Zero(t, account.Snapshot().Used) + require.Empty(t, mailbox.ch) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + require.Zero(t, mp.CurrNB()) +} diff --git a/pkg/sql/colexec/hashjoin/expression_memory_test.go b/pkg/sql/colexec/hashjoin/expression_memory_test.go deleted file mode 100644 index 2baa2fd96f5cd..0000000000000 --- a/pkg/sql/colexec/hashjoin/expression_memory_test.go +++ /dev/null @@ -1,108 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package hashjoin - -import ( - "testing" - - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/container/batch" - "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/matrixorigin/matrixone/pkg/container/vector" - "github.com/matrixorigin/matrixone/pkg/pb/plan" - "github.com/matrixorigin/matrixone/pkg/sql/colexec" - "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" - "github.com/matrixorigin/matrixone/pkg/testutil" - "github.com/matrixorigin/matrixone/pkg/vm/process" - "github.com/stretchr/testify/require" -) - -func TestHashJoinResetReleasesProbeExpressionLease(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Lit{Lit: &plan.Literal{ - Value: &plan.Literal_I32Val{I32Val: 1}, - }}, - } - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - lease, err := hashbuild.NewExpressionMemoryLease( - generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - - arg := &HashJoin{} - arg.ctr.eqCondExecs = executors - arg.ctr.eqCondVecs = make([]*vector.Vector, len(executors)) - arg.ctr.probeExpressionLease = lease - input := batch.NewWithSize(0) - input.SetRowCount(4) - require.NoError(t, arg.ctr.evalJoinConditionBudgeted(input, proc)) - require.Positive(t, generation.Used()) - - arg.Reset(proc, false, nil) - require.Zero(t, generation.Used()) - require.Nil(t, arg.ctr.eqCondExecs) - require.Nil(t, arg.ctr.eqCondVecs) - require.Nil(t, arg.ctr.probeExpressionLease) -} - -func TestHashJoinResetReleasesAccountedProbeExpressions(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - const capBytes = uint64(1 << 20) - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - registry, err := mpool.NewAllocationAccountRegistry(1, 16) - require.NoError(t, err) - account, err := registry.OpenWithController(capBytes, generation) - require.NoError(t, err) - expr := &plan.Expr{Typ: plan.Type{Id: int32(types.T_int32)}, Expr: &plan.Expr_Lit{Lit: &plan.Literal{Value: &plan.Literal_I32Val{I32Val: 1}}}} - executors, err := hashbuild.NewAllocationAccountedExpressionExecutorsForAccount( - proc, []*plan.Expr{expr}, account, hashbuild.HashBuildAllocationOwner) - require.NoError(t, err) - arg := &HashJoin{allocationAccount: account} - arg.ctr.eqCondExecs = executors - arg.ctr.eqCondVecs = make([]*vector.Vector, len(executors)) - arg.ctr.probeExpressionsAccounted = true - input := batch.NewWithSize(0) - input.SetRowCount(4) - require.NoError(t, arg.ctr.evalJoinConditionBudgeted(input, proc)) - require.Positive(t, account.Snapshot().Used) - - arg.Reset(proc, false, nil) - require.Zero(t, account.Snapshot().Used) - require.Zero(t, generation.Used()) - require.False(t, arg.ctr.probeExpressionsAccounted) - require.Nil(t, arg.ctr.eqCondExecs) - terminal, _, err := registry.CompleteTerminal(account) - require.NoError(t, err) - require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) -} - -func TestHashJoinAllocationActivationRequiresBothKeySides(t *testing.T) { - col := &plan.Expr{Typ: plan.Type{Id: int32(types.T_int64)}, Expr: &plan.Expr_Col{Col: &plan.ColRef{}}} - arg := &HashJoin{EqConds: [][]*plan.Expr{{col}, {col}}} - require.True(t, arg.AllocationAccountEnabled()) - require.False(t, arg.AllocationAccountActivationBlocked()) - arg.EqConds[1] = []*plan.Expr{nil} - require.False(t, arg.AllocationAccountEnabled()) - require.True(t, arg.AllocationAccountActivationBlocked()) -} diff --git a/pkg/sql/colexec/hashjoin/join.go b/pkg/sql/colexec/hashjoin/join.go index d9f117d1b313f..b1f617b1175b8 100644 --- a/pkg/sql/colexec/hashjoin/join.go +++ b/pkg/sql/colexec/hashjoin/join.go @@ -17,9 +17,9 @@ package hashjoin import ( "bytes" - "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/hashmap" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/nulls" "github.com/matrixorigin/matrixone/pkg/container/types" @@ -76,6 +76,9 @@ func (hashJoin *HashJoin) Prepare(proc *process.Process) (err error) { return err } } + if hashJoin.allocationAccount == nil { + return mpool.ErrAllocationAccountInvalid + } if hashJoin.OpAnalyzer == nil { hashJoin.OpAnalyzer = process.NewAnalyzer(hashJoin.GetIdx(), hashJoin.IsFirst, hashJoin.IsLast, opName) @@ -91,20 +94,28 @@ func (hashJoin *HashJoin) Prepare(proc *process.Process) (err error) { } if len(ctr.eqCondVecs) == 0 { - eqCondExecs, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, hashJoin.EqConds[0]) + eqCondExecs, err := hashbuild.NewExpressionExecutors( + proc, + hashJoin.EqConds[0], + ) if err != nil { return err } var nonEqCondExec colexec.ExpressionExecutor if hashJoin.NonEqCond != nil { - nonEqCondExec, err = colexec.NewExpressionExecutor(proc, hashJoin.NonEqCond) + var nonEqExecs []colexec.ExpressionExecutor + nonEqExecs, err = hashbuild.NewExpressionExecutors( + proc, + []*plan.Expr{hashJoin.NonEqCond}, + ) if err != nil { for _, exec := range eqCondExecs { exec.Free() } return err } + nonEqCondExec = nonEqExecs[0] } ctr.eqCondVecs = make([]*vector.Vector, len(hashJoin.EqConds[0])) @@ -282,7 +293,7 @@ func (hashJoin *HashJoin) Call(proc *process.Process) (vm.CallResult, error) { // For spilled join, clean up current bucket and move to next if (ctr.spillEngine != nil) && (ctr.spillEngine.HasMoreBuckets() || ctr.spillEngine.IsProbing()) { - ctr.rightRowsMatched = nil + ctr.freeRightRowsMatched(proc) ctr.cleanHashMap() ctr.state = Probe } @@ -337,39 +348,13 @@ func (hashJoin *HashJoin) build(analyzer process.Analyzer, proc *process.Process if takeErr != nil { return takeErr } - var probeExpressionLease *hashbuild.ExpressionMemoryLease - var leaseErr error - if hashJoin.allocationAccount != nil && - hashbuild.AllocationAccountedExpressionSetSupported(hashJoin.EqConds[0]) { - ctr.cleanEqCondExecutors() - ctr.eqCondExecs, leaseErr = - hashbuild.NewAllocationAccountedExpressionExecutorsForAccount( - proc, - hashJoin.EqConds[0], - hashJoin.allocationAccount, - hashbuild.HashBuildAllocationOwner, - ) - if leaseErr == nil { - ctr.eqCondVecs = make( - []*vector.Vector, - len(hashJoin.EqConds[0]), - ) - ctr.probeExpressionsAccounted = true - } - } else { - probeExpressionLease, leaseErr = hashbuild.NewExpressionMemoryLease( - budget, hashJoin.EqConds[0], ctr.eqCondExecs, false) - } - if leaseErr != nil { + if hashJoin.allocationAccount == nil { _ = payload.Close() ctr.mp.Free() ctr.mp = nil - ctr.cleanEqCondExecutors() - ctr.releaseProbeExpressionLease() - return leaseErr + return mpool.ErrAllocationAccountInvalid } - ctr.probeExpressionLease = probeExpressionLease - engine, engineErr := spillutil.NewSpillEngineForAccount(spillutil.SpillEngineConfig{ + engine, engineErr := spillutil.NewSpillEngine(spillutil.SpillEngineConfig{ BuildKeyExprs: hashJoin.EqConds[1], ProbeKeyExprs: hashJoin.EqConds[0], SpillThreshold: ctr.spillThreshold, @@ -379,21 +364,16 @@ func (hashJoin *HashJoin) build(analyzer process.Analyzer, proc *process.Process NeedAllocateSels: !hashJoin.HashOnPK, NeedBatches: hashJoin.NeedBuildBatches(), Budget: budget, - ProbeExpressionLease: probeExpressionLease, }, hashJoin.allocationAccount, hashbuild.HashBuildAllocationOwner) if engineErr != nil { _ = payload.Close() ctr.mp.Free() ctr.mp = nil ctr.cleanEqCondExecutors() - ctr.releaseProbeExpressionLease() return engineErr } - if len(payload.Files) > 0 { - engine.InitFromSpilledFiles(payload.Files) - } else { - engine.InitFromSpilledMap(payload.LegacyFds) - } + engine.InitFromSpilledFiles(payload.Files) + ctr.spillEngine = engine if err := engine.ScatterProbeTable(proc, func() (*batch.Batch, error) { input, err := vm.ChildrenCall(hashJoin.GetChildren(0), proc, analyzer) @@ -401,7 +381,7 @@ func (hashJoin *HashJoin) build(analyzer process.Analyzer, proc *process.Process }, analyzer, func(bat *batch.Batch) ([]*vector.Vector, error) { - if err := ctr.evalJoinConditionBudgeted(bat, proc); err != nil { + if err := ctr.evalJoinCondition(bat, proc); err != nil { return nil, err } return ctr.eqCondVecs, nil @@ -410,10 +390,10 @@ func (hashJoin *HashJoin) build(analyzer process.Analyzer, proc *process.Process ctr.mp.Free() ctr.mp = nil engine.Cleanup(proc) + ctr.spillEngine = nil return err } ctr.mp.Free() - ctr.spillEngine = engine ctr.mp = nil return nil } @@ -428,8 +408,16 @@ func (hashJoin *HashJoin) build(analyzer process.Analyzer, proc *process.Process if hashJoin.EmitUnmatchedBuild() { if ctr.rightRowCnt > 0 { - ctr.rightRowsMatched = &bitmap.Bitmap{} - ctr.rightRowsMatched.InitWithSize(ctr.rightRowCnt) + ctr.rightRowsMatched, err = colexec.NewAccountedBitmap( + ctr.rightRowCnt, + proc.Mp(), + hashJoin.allocationAccount, + hashbuild.HashBuildAllocationOwner, + hashJoinAllocationSiteMatchedRows, + ) + if err != nil { + return err + } } } @@ -470,6 +458,7 @@ func (hashJoin *HashJoin) getSpilledInputBatch(proc *process.Process, analyzer p // Load next bucket via engine convenience method. if ctr.mp == nil { + var allocationErr error ok, err := engine.AdvanceToNextBucket(proc, analyzer, func(jm *message.JoinMap, res spillutil.BucketResult) { if res == spillutil.BucketReady { @@ -478,8 +467,14 @@ func (hashJoin *HashJoin) getSpilledInputBatch(proc *process.Process, analyzer p ctr.rightRowCnt = jm.GetRowCount() ctr.probeHashOnPK = hashJoin.HashOnPK || ctr.mp.HashOnUnique() if hashJoin.EmitUnmatchedBuild() && ctr.rightRowCnt > 0 { - ctr.rightRowsMatched = &bitmap.Bitmap{} - ctr.rightRowsMatched.InitWithSize(ctr.rightRowCnt) + ctr.rightRowsMatched, allocationErr = + colexec.NewAccountedBitmap( + ctr.rightRowCnt, + proc.Mp(), + hashJoin.allocationAccount, + hashbuild.HashBuildAllocationOwner, + hashJoinAllocationSiteMatchedRows, + ) ctr.rightMatchedIter = nil } } @@ -487,9 +482,13 @@ func (hashJoin *HashJoin) getSpilledInputBatch(proc *process.Process, analyzer p if err != nil { return result, err } + if allocationErr != nil { + return result, allocationErr + } if !ok { return result, nil } + hashmap.IteratorClearOwner(ctr.itr) ctr.itr = nil ctr.probeState = psNextBatch ctr.lastIdx = 0 @@ -500,7 +499,7 @@ func (hashJoin *HashJoin) getSpilledInputBatch(proc *process.Process, analyzer p } func (ctr *container) probe(hashJoin *HashJoin, proc *process.Process, result *vm.CallResult) error { - err := ctr.evalJoinConditionBudgeted(ctr.leftBat, proc) + err := ctr.evalJoinCondition(ctr.leftBat, proc) if err != nil { return err } @@ -525,7 +524,15 @@ func (ctr *container) probe(hashJoin *HashJoin, proc *process.Process, result *v case psNextBatch: if ctr.lastIdx < leftRowCnt { hashBatch := min(leftRowCnt-ctr.lastIdx, hashmap.UnitLimit) - ctr.vs, ctr.zvs = ctr.itr.Find(ctr.lastIdx, hashBatch, ctr.eqCondVecs) + var err error + ctr.vs, ctr.zvs, err = ctr.itr.Find( + ctr.lastIdx, + hashBatch, + ctr.eqCondVecs, + ) + if err != nil { + return err + } ctr.vsIdx = 0 ctr.probeState = psBatchRow } else { @@ -816,7 +823,7 @@ func (ctr *container) appendMarkForEmptyBuildBucket(marker *vector.Vector, proc return vector.SetConstNull(marker, rowCnt, proc.Mp()) } - if err := ctr.evalJoinConditionBudgeted(ctr.leftBat, proc); err != nil { + if err := ctr.evalJoinCondition(ctr.leftBat, proc); err != nil { return err } if err := vector.AppendMultiFixed(marker, false, false, rowCnt, proc.Mp()); err != nil { @@ -844,34 +851,34 @@ func (ctr *container) syncBitmap(hashJoin *HashJoin, proc *process.Process) erro if hashJoin.NumCPU > 1 { if !hashJoin.IsMerger { - hashJoin.Channel <- ctr.rightRowsMatched + if hashJoin.Mailbox.Send(ctr.rightRowsMatched) { + ctr.rightRowsMatched = nil + } return nil } else { matchedCnt := ctr.rightRowsMatched.Count() for cnt := 1; cnt < int(hashJoin.NumCPU); cnt++ { - v := colexec.ReceiveBitmapFromChannel(proc.Ctx, hashJoin.Channel) - if v == nil { + v, received := hashJoin.Mailbox.Receive(proc.Ctx) + if !received || v == nil { // A worker was torn down before syncing (its Reset sends - // nil) or the context was canceled. The merge is aborted, - // but keep draining this generation's remaining messages - // so no stale bitmap is left behind in the shared - // channel, then bail out without initializing the + // nil) or the context was canceled. Sealing transfers all + // already-published values to cleanup and makes late + // publishers retain their own value. Bail out without initializing the // iterator — Call routes to End and nothing is finalized. - for cnt++; cnt < int(hashJoin.NumCPU); cnt++ { - colexec.ReceiveBitmapFromChannel(proc.Ctx, hashJoin.Channel) - } + hashJoin.Mailbox.SealAndDrain(proc.Mp()) return nil } matchedCnt += v.Count() ctr.rightRowsMatched.Or(v) + colexec.FreeAccountedBitmap(v, proc.Mp()) } if ctr.probeSingle && matchedCnt > ctr.rightRowsMatched.Count() { return moerr.NewErrSubqueryNo1Row(proc.Ctx) } - close(hashJoin.Channel) + hashJoin.Mailbox.SealAndDrain(proc.Mp()) } } @@ -1013,16 +1020,6 @@ func (ctr *container) evalJoinCondition(bat *batch.Batch, proc *process.Process) return nil } -func (ctr *container) evalJoinConditionBudgeted(bat *batch.Batch, proc *process.Process) error { - if ctr.probeExpressionLease == nil { - return ctr.evalJoinCondition(bat, proc) - } - return ctr.probeExpressionLease.Eval(proc, []*batch.Batch{bat}, bat.RowCount(), func(i int, vec *vector.Vector) error { - ctr.eqCondVecs[i] = vec - return nil - }) -} - func (hashJoin *HashJoin) resetResultBat() { ctr := &hashJoin.ctr if ctr.resBat != nil { diff --git a/pkg/sql/colexec/hashjoin/join_test.go b/pkg/sql/colexec/hashjoin/join_test.go index e449f0f4a6ff3..9d87669172458 100644 --- a/pkg/sql/colexec/hashjoin/join_test.go +++ b/pkg/sql/colexec/hashjoin/join_test.go @@ -61,6 +61,7 @@ func TestHashJoinPrepareFailureCanRetry(t *testing.T) { EqConds: [][]*plan.Expr{{valid}, {valid}}, NonEqCond: invalid, } + installTestAllocation(t, arg) require.Error(t, arg.Prepare(proc)) require.Nil(t, arg.ctr.eqCondVecs) @@ -786,9 +787,9 @@ func TestHashJoinSingleRejectsDuplicateMatchesAcrossWorkers(t *testing.T) { JoinType: plan.Node_SINGLE, NumCPU: 2, IsMerger: true, - Channel: make(chan *bitmap.Bitmap, 1), + Mailbox: NewBitmapMailbox(2), } - hashJoin.Channel <- remoteMatches + require.True(t, hashJoin.Mailbox.Send(remoteMatches)) ctr := container{rightRowsMatched: localMatches, probeSingle: true} err := ctr.syncBitmap(hashJoin, proc) @@ -820,14 +821,14 @@ func TestHashJoinMergerSyncBitmapAborted(t *testing.T) { IsRightJoin: true, NumCPU: 3, IsMerger: true, - Channel: make(chan *bitmap.Bitmap, 3), + Mailbox: NewBitmapMailbox(3), ResultCols: []colexec.ResultPos{colexec.NewResultPos(1, 0)}, RightTypes: []types.Type{types.T_int32.ToType()}, } // Worker A was torn down before syncing (its Reset sends nil); worker B // synced normally and its bitmap lands after the abort marker. - hashJoin.Channel <- nil - hashJoin.Channel <- staleMatches + require.True(t, hashJoin.Mailbox.Send(nil)) + require.True(t, hashJoin.Mailbox.Send(staleMatches)) hashJoin.ctr.state = SyncBitmap hashJoin.ctr.rightRowsMatched = matched hashJoin.ctr.rightBats = []*batch.Batch{rightBat} @@ -836,13 +837,13 @@ func TestHashJoinMergerSyncBitmapAborted(t *testing.T) { require.NoError(t, err) require.Nil(t, result.Batch) require.Equal(t, vm.ExecStop, result.Status) - // Worker B's bitmap must not be left behind in the shared channel. - require.Empty(t, hashJoin.Channel) + // Worker B's bitmap must not be left behind in the shared mailbox. + require.Empty(t, hashJoin.Mailbox.ch) // The merger already synced this generation, so Reset must not push the // nil abort marker either. hashJoin.Reset(proc, false, nil) - require.Empty(t, hashJoin.Channel) + require.Empty(t, hashJoin.Mailbox.ch) // Next generation over the same operator and channel: a clean sync must // only observe this generation's bitmaps. @@ -856,8 +857,9 @@ func TestHashJoinMergerSyncBitmapAborted(t *testing.T) { workerMatches2.InitWithSize(4) workerMatches2.Add(2) - hashJoin.Channel <- workerMatches1 - hashJoin.Channel <- workerMatches2 + hashJoin.Mailbox = NewBitmapMailbox(3) + require.True(t, hashJoin.Mailbox.Send(workerMatches1)) + require.True(t, hashJoin.Mailbox.Send(workerMatches2)) hashJoin.ctr.state = SyncBitmap hashJoin.ctr.rightRowsMatched = matched2 hashJoin.ctr.rightBats = []*batch.Batch{rightBat} @@ -896,11 +898,11 @@ func TestHashJoinMergerFinalizeEmitsUnmatchedBuildRows(t *testing.T) { IsRightJoin: true, NumCPU: 2, IsMerger: true, - Channel: make(chan *bitmap.Bitmap, 2), + Mailbox: NewBitmapMailbox(2), ResultCols: []colexec.ResultPos{colexec.NewResultPos(1, 0)}, RightTypes: []types.Type{types.T_int32.ToType()}, } - hashJoin.Channel <- remoteMatches + require.True(t, hashJoin.Mailbox.Send(remoteMatches)) hashJoin.ctr.state = SyncBitmap hashJoin.ctr.rightRowsMatched = matched hashJoin.ctr.rightBats = []*batch.Batch{rightBat} @@ -1049,7 +1051,7 @@ func newTestCaseWithMPool( resultBatch.Vecs[i] = bat.Vecs[rp[i].Pos] } tag++ - return joinTestCase{ + tc := joinTestCase{ types: ts, flgs: flgs, proc: proc, @@ -1088,6 +1090,8 @@ func newTestCaseWithMPool( }, resultBatch: resultBatch, } + installTestAllocation(t, tc.arg, tc.barg) + return tc } func resetChildren(arg *HashJoin, m *mpool.MPool) { diff --git a/pkg/sql/colexec/hashjoin/key_contract_test.go b/pkg/sql/colexec/hashjoin/key_contract_test.go index 9f6bde1f22f51..ce84365955879 100644 --- a/pkg/sql/colexec/hashjoin/key_contract_test.go +++ b/pkg/sql/colexec/hashjoin/key_contract_test.go @@ -351,6 +351,7 @@ func runHashJoinKeyContract( if mode.shuffle { buildArg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: joinMapTag + 7000} } + installTestAllocation(t, arg, buildArg) var build, probe *batch.Batch defer func() { arg.Free(proc, false, nil) diff --git a/pkg/sql/colexec/hashjoin/mark_spill_test.go b/pkg/sql/colexec/hashjoin/mark_spill_test.go index 6549f6b56eb77..e91fcd2e39fac 100644 --- a/pkg/sql/colexec/hashjoin/mark_spill_test.go +++ b/pkg/sql/colexec/hashjoin/mark_spill_test.go @@ -15,7 +15,6 @@ package hashjoin import ( - "os" "testing" "github.com/matrixorigin/matrixone/pkg/container/batch" @@ -220,6 +219,7 @@ func TestHashMarkJoinEmptySpillBucketTruthTable(t *testing.T) { // spill scatter phase, and an empty global build makes even NULL probes FALSE. func TestHashMarkJoinSpilledEmptyBuild(t *testing.T) { tc := newMarkSpillTestCase(t) + generation, registry, account := installHashJoinTestAllocation(t, tc.arg) probe := batch.NewWithSize(1) probe.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 0}, []uint64{1}, tc.proc.Mp()) probe.SetRowCount(2) @@ -229,14 +229,14 @@ func TestHashMarkJoinSpilledEmptyBuild(t *testing.T) { jm.SetRowCount(0) jm.IncRef(1) require.NoError(t, jm.SetSpillBuildPayload(message.SpillBuildPayload{ - LegacyFds: make([]*os.File, spillutil.SpillNumBuckets), + Files: make([]*message.SpillFile, spillutil.SpillNumBuckets), + BudgetRef: generation, })) message.SendMessage(message.JoinMapMsg{ - JoinMapPtr: jm, + Result: message.NewJoinMapResult(jm), IsShuffle: true, ShuffleIdx: tc.arg.ShuffleIdx, Tag: tc.arg.JoinMapTag, - Spilled: true, }, tc.proc.GetMessageBoard()) require.NoError(t, tc.arg.Prepare(tc.proc)) @@ -245,4 +245,7 @@ func TestHashMarkJoinSpilledEmptyBuild(t *testing.T) { 1: {value: false}, }, collectMarkResults(t, &tc)) finishMarkSpillTest(t, &tc) + require.Zero(t, account.Snapshot().Used) + _, _, err := registry.CompleteTerminal(account) + require.NoError(t, err) } diff --git a/pkg/sql/colexec/hashjoin/spill_diskv2_test.go b/pkg/sql/colexec/hashjoin/spill_diskv2_test.go index 64d28bb960aad..b9ce3e4eb8fa1 100644 --- a/pkg/sql/colexec/hashjoin/spill_diskv2_test.go +++ b/pkg/sql/colexec/hashjoin/spill_diskv2_test.go @@ -17,15 +17,17 @@ package hashjoin import ( "bytes" "context" - "os" + "io" "testing" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/defines" "github.com/matrixorigin/matrixone/pkg/fileservice" "github.com/matrixorigin/matrixone/pkg/sql/colexec/spillutil" "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/matrixorigin/matrixone/pkg/vm/message" "github.com/matrixorigin/matrixone/pkg/vm/process" "github.com/stretchr/testify/require" ) @@ -74,23 +76,36 @@ func TestHashJoinSpillDiskV2(t *testing.T) { buildBat.Vecs[0] = testutil.MakeInt32Vector(buildVals, nil, proc.Mp()) buildBat.SetRowCount(100) - // spill the build bucket via FlushBucketBatch + // Write one production-format spill record to the DISK-V2 file. buildFile, err := spillfs.CreateAndRemoveFile(context.Background(), "diskv2_build") require.NoError(t, err) - var buf bytes.Buffer - bw := spillutil.BucketWriter{Name: "diskv2_build", Fd: buildFile} - err = spillutil.FlushBucketBatch(proc, buildBat, &bw, &buf, nil) + var payload bytes.Buffer + err = buildBat.MarshalBinaryWithGroupingTo(&payload) + require.NoError(t, err) + rows, size, magic := int64(buildBat.RowCount()), int64(payload.Len()), uint64(spillutil.SpillMagic) + for _, part := range [][]byte{ + types.EncodeInt64(&rows), + types.EncodeInt64(&size), + payload.Bytes(), + types.EncodeUint64(&magic), + } { + _, err = buildFile.Write(part) + require.NoError(t, err) + } + _, err = buildFile.Seek(0, io.SeekStart) + require.NoError(t, err) + info, err := buildFile.Stat() require.NoError(t, err) - buildFd := bw.HandOffFd() - require.NotNil(t, buildFd) // rebuild via SpillEngine - engine := spillutil.NewSpillEngine(spillutil.SpillEngineConfig{ + engine := newAccountedTestSpillEngine(t, spillutil.SpillEngineConfig{ BuildKeyExprs: makeKeyExpr(), NeedsBuildForEmptyProbe: true, NeedBatches: true, }) - engine.InitFromSpilledMap([]*os.File{buildFd}) + engine.InitFromSpilledFiles([]*message.SpillFile{ + message.NewSpillFile(buildFile, 100, uint64(info.Size()), nil), + }) analyzer := process.NewAnalyzer(0, false, false, "test") jm, res, err := engine.RebuildHashmap(proc, analyzer) diff --git a/pkg/sql/colexec/hashjoin/spill_integration_test.go b/pkg/sql/colexec/hashjoin/spill_integration_test.go index e3261f099ba36..9b152e09db2c7 100644 --- a/pkg/sql/colexec/hashjoin/spill_integration_test.go +++ b/pkg/sql/colexec/hashjoin/spill_integration_test.go @@ -23,6 +23,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" "github.com/matrixorigin/matrixone/pkg/sql/colexec/spillutil" "github.com/matrixorigin/matrixone/pkg/testutil" metricv2 "github.com/matrixorigin/matrixone/pkg/util/metric/v2" @@ -39,13 +40,56 @@ func makeKeyExpr() []*plan.Expr { }} } +func newAccountedTestSpillEngine( + t *testing.T, + cfg spillutil.SpillEngineConfig, +) *spillutil.SpillEngine { + t.Helper() + if cfg.Budget == nil { + budget := process.MustNewHashBuildBudget(1<<60, 1<<60) + var err error + cfg.Budget, err = budget.OpenGeneration(1) + require.NoError(t, err) + } + registry, err := mpool.NewAllocationAccountRegistry(1, 1<<20) + require.NoError(t, err) + account, err := registry.OpenWithController(1<<60, cfg.Budget) + require.NoError(t, err) + engine, err := spillutil.NewSpillEngine( + cfg, + account, + hashbuild.HashBuildAllocationOwner, + ) + require.NoError(t, err) + return engine +} + +func installHashJoinTestAllocation( + t *testing.T, + join *HashJoin, +) (*process.HashBuildBudgetGeneration, *mpool.AllocationAccountRegistry, *mpool.AllocationAccount) { + t.Helper() + budget := process.MustNewHashBuildBudget(64<<20, 64<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 1<<20) + require.NoError(t, err) + account, err := registry.OpenWithController(64<<20, generation) + require.NoError(t, err) + if join.allocationAccount != nil { + require.NoError(t, join.ClearAllocationAccount(join.allocationAccount)) + } + require.NoError(t, join.SetAllocationAccount(account)) + return generation, registry, account +} + // TestGetSpilledInputBatchNoBuckets verifies that getSpilledInputBatch // returns nil when the engine has no buckets. func TestGetSpilledInputBatchNoBuckets(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - engine := spillutil.NewSpillEngine(spillutil.SpillEngineConfig{}) + engine := newAccountedTestSpillEngine(t, spillutil.SpillEngineConfig{}) hashJoin := &HashJoin{ctr: container{spillEngine: engine}} result, err := hashJoin.getSpilledInputBatch(proc, process.NewAnalyzer(0, false, false, "test")) require.NoError(t, err) @@ -168,6 +212,101 @@ func TestShuffleJoinFiniteBudgetInitialSpillAndReSpill(t *testing.T) { require.Zero(t, tc.proc.Mp().CurrNB()) } +func TestShuffleJoinSpillUsesCanonicalGroupingPartitionKey(t *testing.T) { + for _, test := range []struct { + name string + typ types.Type + probe func(*process.Process) *vector.Vector + build func(*process.Process) *vector.Vector + }{ + { + name: "varchar", + typ: types.T_varchar.ToType(), + probe: func(proc *process.Process) *vector.Vector { + return testutil.MakeVarcharVector([]string{"probe"}, nil, proc.Mp()) + }, + build: func(proc *process.Process) *vector.Vector { + return testutil.MakeVarcharVector([]string{"build"}, nil, proc.Mp()) + }, + }, + { + name: "int32", + typ: types.T_int32.ToType(), + probe: func(proc *process.Process) *vector.Vector { + return testutil.MakeInt32Vector([]int32{222}, nil, proc.Mp()) + }, + build: func(proc *process.Process) *vector.Vector { + return testutil.MakeInt32Vector([]int32{111}, nil, proc.Mp()) + }, + }, + } { + t.Run(test.name, func(t *testing.T) { + keyExpr := []*plan.Expr{{ + Typ: plan.Type{Id: int32(test.typ.Oid), Width: test.typ.Width}, + Expr: &plan.Expr_Col{Col: &plan.ColRef{ + ColPos: 0, + }}, + }} + tc := newTestCase( + t, + []bool{false}, + []types.Type{test.typ}, + []colexec.ResultPos{colexec.NewResultPos(0, 0)}, + [][]*plan.Expr{keyExpr, keyExpr}, + ) + tc.arg.NonEqCond = nil + tc.arg.IsShuffle = true + tc.arg.ShuffleIdx = 0 + tc.arg.SpillThreshold = 1 + tc.barg.IsShuffle = true + tc.barg.ShuffleIdx = 0 + tc.barg.SpillThreshold = 1 + tc.barg.NeedBatches = false + tc.barg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{ + Tag: tc.arg.JoinMapTag + 1_500, + } + + probe := batch.NewWithSize(1) + probe.Vecs[0] = test.probe(tc.proc) + probe.SetRowCount(1) + build := batch.NewWithSize(1) + build.Vecs[0] = test.build(tc.proc) + build.SetRowCount(1) + probe.Vecs[0].GetGrouping().Add(0) + build.Vecs[0].GetGrouping().Add(0) + resetChildrenWithBatch(tc.arg, probe) + resetHashBuildChildrenWithBatch(tc.barg, build) + + require.NoError(t, tc.arg.Prepare(tc.proc)) + require.NoError(t, tc.barg.Prepare(tc.proc)) + _, err := vm.Exec(tc.barg, tc.proc) + require.NoError(t, err) + + rows := 0 + for { + result, err := vm.Exec(tc.arg, tc.proc) + require.NoError(t, err) + if result.Batch != nil { + rows += result.Batch.RowCount() + } + if result.Status == vm.ExecStop { + break + } + } + require.Equal(t, 1, rows) + require.Positive(t, + tc.barg.OpAnalyzer.GetOpStats().ExtraStats["HashBuildSpillStarts"]) + + tc.arg.Reset(tc.proc, false, nil) + tc.barg.Reset(tc.proc, false, nil) + tc.arg.Free(tc.proc, false, nil) + tc.barg.Free(tc.proc, false, nil) + tc.proc.Free() + require.Zero(t, tc.proc.Mp().CurrNB()) + }) + } +} + func TestShuffleJoinHardBudgetRejectTransitionsToSpill(t *testing.T) { tc := newTestCase( t, @@ -179,10 +318,7 @@ func TestShuffleJoinHardBudgetRejectTransitionsToSpill(t *testing.T) { // This cap admits one bounded scatter pass and per-bucket rebuild, but not // the complete 8K-row retained build/map. The very high soft threshold // proves that spill is entered from hard admission rejection, not policy. - // Lazy spill scratch no longer consumes resident headroom before spill. - // Keep this cap below the resident hashmap peak while leaving enough room - // for the bounded scatter pass after the rejected map is released. - tc.proc.Base.Lim.Size = 1536 << 10 + tc.proc.Base.Lim.Size = 200 << 10 tc.proc.Base.Lim.SpillSize = 64 << 20 const rows = 8192 @@ -207,8 +343,18 @@ func TestShuffleJoinHardBudgetRejectTransitionsToSpill(t *testing.T) { buildInput := colexec.NewMockOperator().WithBatchs([]*batch.Batch{build1, build2}) tc.barg.Children = nil tc.barg.AppendChild(buildInput) + oldAccount := tc.arg.allocationAccount + require.NoError(t, tc.arg.ClearAllocationAccount(oldAccount)) + require.NoError(t, tc.barg.ClearAllocationAccount(oldAccount)) + budget, err := tc.proc.GetHashBuildBudget() + require.NoError(t, err) + registry, err := budget.AllocationAccountRegistry() + require.NoError(t, err) + account, err := registry.OpenWithController(budget.Snapshot().Cap, budget) + require.NoError(t, err) + require.NoError(t, tc.arg.SetAllocationAccount(account)) + require.NoError(t, tc.barg.SetAllocationAccount(account)) - rejectBefore := promtestutil.ToFloat64(metricv2.HashBuildBudgetEventCounter.WithLabelValues("memory", "reject", "query")) spillBefore := promtestutil.ToFloat64(metricv2.HashBuildSpillDepthCounter.WithLabelValues("spill", "1")) require.NoError(t, tc.arg.Prepare(tc.proc)) require.NoError(t, tc.barg.Prepare(tc.proc)) @@ -228,13 +374,12 @@ func TestShuffleJoinHardBudgetRejectTransitionsToSpill(t *testing.T) { } } require.ElementsMatch(t, values, resultValues) - require.Greater(t, promtestutil.ToFloat64(metricv2.HashBuildBudgetEventCounter.WithLabelValues("memory", "reject", "query")), rejectBefore) + require.Positive(t, tc.barg.OpAnalyzer.GetOpStats().ExtraStats["HashBuildSpillStarts"]) + require.LessOrEqual(t, account.Snapshot().Peak, account.Snapshot().Limit) require.Greater(t, promtestutil.ToFloat64(metricv2.HashBuildSpillDepthCounter.WithLabelValues("spill", "1")), spillBefore) tc.arg.Free(tc.proc, false, nil) tc.barg.Free(tc.proc, false, nil) - budget, err := tc.proc.GetHashBuildBudget() - require.NoError(t, err) require.Zero(t, budget.Used()) require.Zero(t, budget.SpillDiskUsed()) require.Zero(t, budget.SpillFDUsed()) diff --git a/pkg/sql/colexec/hashjoin/types.go b/pkg/sql/colexec/hashjoin/types.go index 5dca67aef0183..70098319ec525 100644 --- a/pkg/sql/colexec/hashjoin/types.go +++ b/pkg/sql/colexec/hashjoin/types.go @@ -24,7 +24,6 @@ import ( "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" - "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" "github.com/matrixorigin/matrixone/pkg/sql/colexec/spillutil" "github.com/matrixorigin/matrixone/pkg/vm" "github.com/matrixorigin/matrixone/pkg/vm/message" @@ -49,6 +48,8 @@ const ( psBatchRow ) +const hashJoinAllocationSiteMatchedRows mpool.AllocationSite = 80 + type container struct { state int itr hashmap.Iterator @@ -105,14 +106,9 @@ type container struct { maxAllocSize int64 // spill support - spillEngine *spillutil.SpillEngine - spillThreshold int64 - // Non-nil only for spilled joins, where probe expressions are part of the - // shared HashBuild/spill working set. Resident probe expressions remain - // under normal process/mpool accounting; this is not a general query budget. - probeExpressionLease *hashbuild.ExpressionMemoryLease - probeExpressionsAccounted bool - probeBucketActive bool // true while reading probe batches from a bucket + spillEngine *spillutil.SpillEngine + spillThreshold int64 + probeBucketActive bool // true while reading probe batches from a bucket } type HashJoin struct { @@ -127,7 +123,7 @@ type HashJoin struct { NonEqCond *plan.Expr EqConds [][]*plan.Expr - Channel chan *bitmap.Bitmap + Mailbox *BitmapMailbox NumCPU uint64 HashOnPK bool @@ -145,22 +141,6 @@ type HashJoin struct { vm.OperatorBase } -func (hashJoin *HashJoin) AllocationAccountEnabled() bool { - return hashJoin != nil && hashJoin.allocationAccountExpressionOwnerClosed() -} - -func (hashJoin *HashJoin) AllocationAccountActivationBlocked() bool { - return hashJoin != nil && !hashJoin.allocationAccountExpressionOwnerClosed() -} - -func (hashJoin *HashJoin) allocationAccountExpressionOwnerClosed() bool { - if hashJoin == nil || len(hashJoin.EqConds) != 2 { - return false - } - return hashbuild.AllocationAccountedExpressionSetSupported(hashJoin.EqConds[0]) && - hashbuild.AllocationAccountedExpressionSetSupported(hashJoin.EqConds[1]) -} - func (hashJoin *HashJoin) SetAllocationAccount( account *mpool.AllocationAccount, ) error { @@ -171,6 +151,9 @@ func (hashJoin *HashJoin) SetAllocationAccount( hashJoin.allocationAccount != account { return mpool.ErrAllocationAccountMismatch } + if hashJoin.allocationAccount == account { + return nil + } hashJoin.allocationAccount = account return nil } @@ -185,7 +168,12 @@ func (hashJoin *HashJoin) ClearAllocationAccount( return mpool.ErrAllocationAccountMismatch } if hashJoin.ctr.mp != nil || hashJoin.ctr.spillEngine != nil || - hashJoin.ctr.probeExpressionsAccounted { + len(hashJoin.ctr.eqCondExecs) != 0 || + hashJoin.ctr.nonEqCondExec != nil || + hashJoin.ctr.rightRowsMatched != nil { + return mpool.ErrAllocationAccountInvariant + } + if hashJoin.NumCPU > 1 && !hashJoin.Mailbox.Terminal() { return mpool.ErrAllocationAccountInvariant } hashJoin.allocationAccount = nil @@ -241,22 +229,19 @@ func (hashJoin *HashJoin) ExecProjection(proc *process.Process, input *batch.Bat func (hashJoin *HashJoin) Reset(proc *process.Process, pipelineFailed bool, err error) { ctr := &hashJoin.ctr + hashmap.IteratorClearOwner(ctr.itr) ctr.itr = nil if !ctr.bitmapSynced && hashJoin.NumCPU > 1 && !hashJoin.IsMerger { - hashJoin.Channel <- nil + hashJoin.Mailbox.Send(nil) } - // SpillEngine borrows the probe executor lease. End that borrow before the - // join frees the executors and releases their reservation. - ctr.cleanBucketBatches(proc) - if ctr.probeExpressionLease != nil || ctr.probeExpressionsAccounted { - ctr.cleanEqCondExecutors() - ctr.releaseProbeExpressionLease() - } else { - ctr.resetEqCondExecutors() + if hashJoin.NumCPU > 1 && hashJoin.IsMerger { + hashJoin.Mailbox.SealAndDrain(proc.Mp()) } + ctr.cleanBucketBatches(proc) + ctr.cleanEqCondExecutors() ctr.cleanHashMap() - ctr.resetNonEqCondExecutor() - ctr.rightRowsMatched = nil + ctr.cleanNonEqCondExecutor() + ctr.freeRightRowsMatched(proc) ctr.rightMatchedIter = nil ctr.skipProbe = false ctr.bitmapSynced = false @@ -266,8 +251,6 @@ func (hashJoin *HashJoin) Reset(proc *process.Process, pipelineFailed bool, err ctr.state = Build ctr.probeState = psNextBatch ctr.lastIdx = 0 - hashJoin.allocationAccount = nil - if hashJoin.OpAnalyzer != nil { hashJoin.OpAnalyzer.Alloc(ctr.maxAllocSize) } @@ -280,16 +263,8 @@ func (hashJoin *HashJoin) Free(proc *process.Process, pipelineFailed bool, err e ctr.cleanBatch(proc) ctr.cleanBucketBatches(proc) ctr.cleanEqCondExecutors() - ctr.releaseProbeExpressionLease() ctr.cleanHashMap() ctr.cleanNonEqCondExecutor() - hashJoin.allocationAccount = nil -} - -func (ctr *container) resetNonEqCondExecutor() { - if ctr.nonEqCondExec != nil { - ctr.nonEqCondExec.ResetForNextQuery() - } } func (ctr *container) cleanNonEqCondExecutor() { @@ -311,9 +286,12 @@ func (ctr *container) cleanBatch(proc *process.Process) { ctr.joinBats[i] = nil } } - if ctr.rightRowsMatched != nil { - ctr.rightRowsMatched = nil - } + ctr.freeRightRowsMatched(proc) +} + +func (ctr *container) freeRightRowsMatched(proc *process.Process) { + colexec.FreeAccountedBitmap(ctr.rightRowsMatched, proc.Mp()) + ctr.rightRowsMatched = nil } func (ctr *container) cleanBucketBatches(proc *process.Process) { @@ -325,6 +303,8 @@ func (ctr *container) cleanBucketBatches(proc *process.Process) { } func (ctr *container) cleanHashMap() { + hashmap.IteratorClearOwner(ctr.itr) + ctr.itr = nil if ctr.mp != nil { ctr.mp.Free() ctr.mp = nil @@ -339,22 +319,6 @@ func (ctr *container) cleanEqCondExecutors() { } ctr.eqCondExecs = nil ctr.eqCondVecs = nil - ctr.probeExpressionsAccounted = false -} - -func (ctr *container) resetEqCondExecutors() { - for i := range ctr.eqCondExecs { - if ctr.eqCondExecs[i] != nil { - ctr.eqCondExecs[i].ResetForNextQuery() - } - } -} - -func (ctr *container) releaseProbeExpressionLease() { - if ctr.probeExpressionLease != nil { - ctr.probeExpressionLease.Release() - ctr.probeExpressionLease = nil - } } func (hashJoin *HashJoin) IsInner() bool { diff --git a/pkg/sql/colexec/indexbuild/build.go b/pkg/sql/colexec/indexbuild/build.go index 948f9e49e8221..bb2c0f7647c04 100644 --- a/pkg/sql/colexec/indexbuild/build.go +++ b/pkg/sql/colexec/indexbuild/build.go @@ -18,6 +18,7 @@ import ( "bytes" "github.com/matrixorigin/matrixone/pkg/common/hashmap/keycodec" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" @@ -66,9 +67,38 @@ func (indexBuild *IndexBuild) Prepare(proc *process.Process) (err error) { ctr.runtimeFilterUsable = runtimefilter.ExactKeyEncoding( spec, declaredType) != keycodec.ExactRuntimeFilterUnsupported } + if ctr.runtimeFilterUsable && + (indexBuild.allocationAccount == nil || + indexBuild.runtimeFilterAllocation == nil) { + return mpool.ErrAllocationAccountInvalid + } return nil } +func (indexBuild *IndexBuild) newRuntimeFilterBatch( + typ types.Type, +) (*batch.Batch, error) { + if indexBuild.runtimeFilterAllocation == nil { + return nil, mpool.ErrAllocationAccountInvalid + } + vec, err := vector.NewOffHeapVecWithTypeAndAllocation( + typ, + indexBuild.runtimeFilterAllocation, + ) + if err != nil { + return nil, err + } + buf := batch.NewOffHeapWithSize(1) + if err = buf.SetAllocationAccount( + indexBuild.runtimeFilterAllocation, + ); err != nil { + vec.Free(nil) + return nil, err + } + buf.SetVector(0, vec) + return buf, nil +} + func (indexBuild *IndexBuild) Call(proc *process.Process) (vm.CallResult, error) { analyzer := indexBuild.OpAnalyzer @@ -195,9 +225,11 @@ func (ctr *container) collectBuildBatches(indexBuild *IndexBuild, proc *process. // expanded into an equally large retained vector. if !inputVec.IsConstNull() && inputVec.Length() > 0 { if ctr.buf == nil { - ctr.buf = batch.NewOffHeapWithSize(1) - ctr.buf.Vecs[0] = vector.NewOffHeapVecWithType( + ctr.buf, err = indexBuild.newRuntimeFilterBatch( *inputVec.GetType()) + if err != nil { + return err + } } if err = ctr.buf.UnionOne(result.Batch, 0, proc.Mp()); err != nil { err = runtimefilter.MarkOptionalAllocationError(err) @@ -214,9 +246,11 @@ func (ctr *container) collectBuildBatches(indexBuild *IndexBuild, proc *process. // cardinality-bounded copy. Off-heap growth is tracked by the // process pool and can fail open to PASS instead of ending in // an unrecoverable Go-heap OOM. - ctr.buf = batch.NewOffHeapWithSize(1) - ctr.buf.Vecs[0] = vector.NewOffHeapVecWithType( + ctr.buf, err = indexBuild.newRuntimeFilterBatch( *inputVec.GetType()) + if err != nil { + return err + } } ctr.buf, err = ctr.buf.AppendWithCopy(proc.Ctx, proc.Mp(), result.Batch) if err != nil { @@ -286,7 +320,11 @@ func (ctr *container) handleRuntimeFilter(ap *IndexBuild, proc *process.Process) // Batch.Dup preserves a first-batch constant vector. Materialize only // its one distinct value: expanding every repeated row would waste // memory and AppendFixed cannot add signed-zero closure to a const vec. - flat := vector.NewOffHeapVecWithType(*vec.GetType()) + flat, err := vector.NewOffHeapVecWithTypeAndAllocation( + *vec.GetType(), ap.runtimeFilterAllocation) + if err != nil { + return err + } if !vec.IsConstNull() && vec.Length() > 0 { if err := flat.UnionOne(vec, 0, proc.Mp()); err != nil { flat.Free(proc.Mp()) @@ -335,15 +373,13 @@ func (ctr *container) handleRuntimeFilter(ap *IndexBuild, proc *process.Process) // NULLs are irrelevant for IN-filter: clear bitmap before sort. vec.GetNulls().Reset() vec.InplaceSort() - budget, err := proc.GetHashBuildBudget() - if err != nil { - if ctr.fallbackRuntimeFilter(ap, proc, err) { - return nil - } - ctr.abandonRuntimeFilter(proc) - return err - } - data, release, err := runtimefilter.MarshalExactFilterVector(vec, budget) + data, release, err := runtimefilter.MarshalExactFilterVector( + vec, + proc.Mp(), + ap.allocationAccount, + indexBuildAllocationOwner, + indexBuildAllocationSiteRuntimeFilterPayload, + ) if err != nil { if ctr.fallbackRuntimeFilter(ap, proc, err) { return nil diff --git a/pkg/sql/colexec/indexbuild/build_test.go b/pkg/sql/colexec/indexbuild/build_test.go index 122fb648ca0e5..da49eb2107f32 100644 --- a/pkg/sql/colexec/indexbuild/build_test.go +++ b/pkg/sql/colexec/indexbuild/build_test.go @@ -67,6 +67,22 @@ func indexBuildTestProcess(t *testing.T) *process.Process { return proc } +func prepareIndexBuild( + t *testing.T, + arg *IndexBuild, + proc *process.Process, +) { + t.Helper() + if arg.allocationAccount == nil { + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.Open(1 << 60) + require.NoError(t, err) + require.NoError(t, arg.SetAllocationAccount(account)) + } + require.NoError(t, arg.Prepare(proc)) +} + func indexBuildBatch(vec *vector.Vector, rows int) *batch.Batch { bat := batch.NewWithSize(1) bat.Vecs[0] = vec @@ -116,7 +132,7 @@ func executeIndexBuild( arg.AppendChild(child) require.NoError(t, child.Prepare(proc)) - require.NoError(t, arg.Prepare(proc)) + prepareIndexBuild(t, arg, proc) result, err := vm.Exec(arg, proc) require.NoError(t, err) require.Equal(t, vm.ExecStop, result.Status) @@ -149,7 +165,7 @@ func TestIndexBuildExactRuntimeFilterContract(t *testing.T) { child := colexec.NewMockOperator() arg := NewArgument() arg.RuntimeFilterSpec = spec - require.NoError(t, arg.Prepare(proc)) + prepareIndexBuild(t, arg, proc) arg.ctr.buf = indexBuildBatch(nil, 1) require.NotPanics(t, func() { require.NoError(t, arg.ctr.handleRuntimeFilter(arg, proc)) @@ -203,7 +219,7 @@ func TestIndexBuildExactRuntimeFilterContract(t *testing.T) { arg.RuntimeFilterSpec = spec arg.AppendChild(child) require.NoError(t, child.Prepare(proc)) - require.NoError(t, arg.Prepare(proc)) + prepareIndexBuild(t, arg, proc) result, err := vm.Exec(arg, proc) require.NoError(t, err) require.Equal(t, vm.ExecStop, result.Status) @@ -225,7 +241,7 @@ func TestIndexBuildExactRuntimeFilterContract(t *testing.T) { arg.RuntimeFilterSpec = spec arg.AppendChild(child) require.NoError(t, child.Prepare(proc)) - require.NoError(t, arg.Prepare(proc)) + prepareIndexBuild(t, arg, proc) result, err := vm.Exec(arg, proc) require.NoError(t, err) require.Equal(t, vm.ExecStop, result.Status) @@ -294,7 +310,7 @@ func TestIndexBuildFloatRuntimeFilterClosesConstSignedZero(t *testing.T) { arg.RuntimeFilterSpec = spec arg.AppendChild(child) require.NoError(t, child.Prepare(proc)) - require.NoError(t, arg.Prepare(proc)) + prepareIndexBuild(t, arg, proc) result, err := vm.Exec(arg, proc) require.NoError(t, err) require.Equal(t, vm.ExecStop, result.Status) @@ -353,7 +369,7 @@ func TestIndexBuildRuntimeFilterCopyFailureFailsOpen(t *testing.T) { arg.RuntimeFilterSpec = spec arg.AppendChild(child) require.NoError(t, child.Prepare(proc)) - require.NoError(t, arg.Prepare(proc)) + prepareIndexBuild(t, arg, proc) filler, err := limited.Alloc( int(limited.Cap()-limited.CurrNB()), true) require.NoError(t, err) @@ -413,7 +429,7 @@ func TestIndexBuildRuntimeFilterClosureFailureFailsOpen(t *testing.T) { proc.SetMessageBoard(message.NewMessageBoard()) arg := NewArgument() arg.RuntimeFilterSpec = spec - require.NoError(t, arg.Prepare(proc)) + prepareIndexBuild(t, arg, proc) arg.ctr.buf = batch.NewOffHeapWithSize(1) arg.ctr.buf.Vecs[0] = vector.NewOffHeapVecWithType(typ) @@ -460,19 +476,28 @@ func TestIndexBuildRuntimeFilterBudgetErrorPolicy(t *testing.T) { spec := indexBuildRawSpec(111, 16, typ) arg := NewArgument() arg.RuntimeFilterSpec = spec - require.NoError(t, arg.Prepare(proc)) + budget := process.MustNewHashBuildBudget(1<<20, 1<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.OpenWithController( + generation.Cap(), generation) + require.NoError(t, err) + require.NoError(t, arg.SetAllocationAccount(account)) + prepareIndexBuild(t, arg, proc) arg.ctr.buf = indexBuildBatch( testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()), 3, ) - generation, err := proc.GetHashBuildBudget() - require.NoError(t, err) - var held *process.HashBuildReservation + var filler []byte if test.closed { generation.Close() } else { - held, err = generation.Reserve(generation.Cap()) + remaining := account.Snapshot().Limit - account.Snapshot().Used + filler, err = proc.Mp().AllocAccounted( + int(remaining), account, 63, 255) require.NoError(t, err) } @@ -490,7 +515,8 @@ func TestIndexBuildRuntimeFilterBudgetErrorPolicy(t *testing.T) { require.True(t, arg.ctr.runtimeFilterDone) require.Equal(t, int64(1), stats["IndexBuildRuntimeFilterBudgetFallbacks"]) - require.True(t, held.Release()) + proc.Mp().Free(filler) + generation.Close() } require.False(t, arg.ctr.runtimeFilterUsable) require.Nil(t, arg.ctr.buf) @@ -545,7 +571,7 @@ func TestIndexBuildCallErrorUnblocksRuntimeFilterBeforeReset(t *testing.T) { ) require.NoError(t, child.Prepare(proc)) - require.NoError(t, arg.Prepare(proc)) + prepareIndexBuild(t, arg, proc) _, err := vm.Exec(arg, proc) require.ErrorIs(t, err, buildErr) @@ -571,7 +597,7 @@ func TestIndexBuildCallErrorUnblocksRuntimeFilterBeforeReset(t *testing.T) { require.True(t, arg.ctr.runtimeFilterDone) proc.GetMessageBoard().Reset() - require.NoError(t, arg.Prepare(proc)) + prepareIndexBuild(t, arg, proc) require.False(t, arg.ctr.runtimeFilterDone, "Prepare must open the terminal gate for the next generation") arg.finalizeBuildFailure(proc) diff --git a/pkg/sql/colexec/indexbuild/types.go b/pkg/sql/colexec/indexbuild/types.go index 2995f54323d5c..7eae66d34a174 100644 --- a/pkg/sql/colexec/indexbuild/types.go +++ b/pkg/sql/colexec/indexbuild/types.go @@ -15,8 +15,10 @@ package indexbuild import ( + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/reuse" "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/vm" "github.com/matrixorigin/matrixone/pkg/vm/message" @@ -31,6 +33,16 @@ const ( End ) +const indexBuildAllocationOwner mpool.AllocationOwner = 1 + +const ( + indexBuildAllocationSiteRuntimeFilterData mpool.AllocationSite = iota + 1 + indexBuildAllocationSiteRuntimeFilterArea + indexBuildAllocationSiteRuntimeFilterNulls + indexBuildAllocationSiteRuntimeFilterGrouping + indexBuildAllocationSiteRuntimeFilterPayload +) + type container struct { state int buf *batch.Batch @@ -39,11 +51,58 @@ type container struct { } type IndexBuild struct { - ctr container - RuntimeFilterSpec *plan.RuntimeFilterSpec + ctr container + RuntimeFilterSpec *plan.RuntimeFilterSpec + allocationAccount *mpool.AllocationAccount + runtimeFilterAllocation *vector.AllocationAccountSelection vm.OperatorBase } +func (indexBuild *IndexBuild) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if account == nil { + return mpool.ErrAllocationAccountInvalid + } + if indexBuild.allocationAccount != nil { + if indexBuild.allocationAccount == account { + return nil + } + return mpool.ErrAllocationAccountMismatch + } + selection, err := vector.NewAllocationAccountSelection( + account, + indexBuildAllocationOwner, + indexBuildAllocationSiteRuntimeFilterData, + indexBuildAllocationSiteRuntimeFilterArea, + indexBuildAllocationSiteRuntimeFilterNulls, + indexBuildAllocationSiteRuntimeFilterGrouping, + ) + if err != nil { + return err + } + indexBuild.allocationAccount = account + indexBuild.runtimeFilterAllocation = selection + return nil +} + +func (indexBuild *IndexBuild) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if indexBuild.allocationAccount == nil { + return nil + } + if indexBuild.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if indexBuild.ctr.buf != nil { + return mpool.ErrAllocationAccountInvariant + } + indexBuild.allocationAccount = nil + indexBuild.runtimeFilterAllocation = nil + return nil +} + func (indexBuild *IndexBuild) GetOperatorBase() *vm.OperatorBase { return &indexBuild.OperatorBase } @@ -92,7 +151,8 @@ func (indexBuild *IndexBuild) Reset(proc *process.Process, pipelineFailed bool, indexBuild.ctr.state = ReceiveBatch indexBuild.ctr.runtimeFilterUsable = false if indexBuild.ctr.buf != nil { - indexBuild.ctr.buf.CleanOnlyData() + indexBuild.ctr.buf.Clean(proc.Mp()) + indexBuild.ctr.buf = nil } } diff --git a/pkg/sql/colexec/intersect/intersect.go b/pkg/sql/colexec/intersect/intersect.go index 74ca37fdd4b0e..4767df195d840 100644 --- a/pkg/sql/colexec/intersect/intersect.go +++ b/pkg/sql/colexec/intersect/intersect.go @@ -172,7 +172,10 @@ func (intersect *Intersect) probeHashTable(proc *process.Process, analyzer proce copy(needInsert, resetsNeedInsert) insertcnt := 0 - vs, zs := itr.Find(i, n, input.Batch.Vecs) + vs, zs, err := itr.Find(i, n, input.Batch.Vecs) + if err != nil { + return false, err + } for j, v := range vs { diff --git a/pkg/sql/colexec/intersectall/intersectall.go b/pkg/sql/colexec/intersectall/intersectall.go index 2d283d4412e64..5c78cfcb7d808 100644 --- a/pkg/sql/colexec/intersectall/intersectall.go +++ b/pkg/sql/colexec/intersectall/intersectall.go @@ -194,7 +194,10 @@ func (intersectAll *IntersectAll) probe(proc *process.Process, analyzer process. copy(ctr.inserted[:n], ctr.resetInserted[:n]) cnt = 0 - vs, _ := itr.Find(i, n, input.Batch.Vecs) + vs, _, err := itr.Find(i, n, input.Batch.Vecs) + if err != nil { + return false, err + } for j, v := range vs { // not found diff --git a/pkg/sql/colexec/join_util.go b/pkg/sql/colexec/join_util.go index 7f951f7cb6b16..d7c42214ff6f6 100644 --- a/pkg/sql/colexec/join_util.go +++ b/pkg/sql/colexec/join_util.go @@ -60,7 +60,6 @@ func (bs *Batches) Reset() { // copy from input batch into batches // the batches structure hold data in fix size 8192 rows, and continue to append from next batch -// if error return , the batches will clean itself func (bs *Batches) CopyIntoBatches(src *batch.Batch, proc *process.Process) (err error) { return bs.CopyIntoBatchesWithAllocation(src, proc, nil) } @@ -69,17 +68,69 @@ func (bs *Batches) CopyIntoBatches(src *batch.Batch, proc *process.Process) (err // destination. The Go descriptors are bounded by one Batch per 8,192 rows and // one Vector pointer per input column; physical data, area, null, and grouping // buffers are allocation-accounted and remain owned by the copied batches. +// +// The append is transactional. In particular, an allocation rejection while +// copying a later input must not destroy batches retained from earlier inputs: +// HashBuild needs those batches intact to recover by spilling them. A partial +// tail is copied into the private staging set before it is extended, so an +// error cannot leave the published tail partially mutated either. func (bs *Batches) CopyIntoBatchesWithAllocation( src *batch.Batch, proc *process.Process, selection *vector.AllocationAccountSelection, ) (err error) { + if len(bs.Buf) > 0 && + bs.Buf[len(bs.Buf)-1].AllocationAccountSelection() != selection { + return mpool.ErrAllocationAccountMismatch + } + + var staged Batches defer func() { if err != nil { - bs.Clean(proc.Mp()) + staged.Clean(proc.Mp()) } }() + replaceTail := len(bs.Buf) > 0 && + bs.Buf[len(bs.Buf)-1].RowCount() != DefaultBatchSize + if replaceTail { + if err = staged.copyIntoBatches( + bs.Buf[len(bs.Buf)-1], + proc, + selection, + ); err != nil { + return err + } + } + if err = staged.copyIntoBatches(src, proc, selection); err != nil { + return err + } + + if replaceTail { + oldTail := bs.Buf[len(bs.Buf)-1] + bs.Buf = bs.Buf[:len(bs.Buf)-1] + bs.Buf = append(bs.Buf, staged.Buf...) + bs.MemSize += staged.MemSize + staged.Buf = nil + staged.MemSize = 0 + oldTail.Clean(proc.Mp()) + return nil + } + if bs.Buf == nil { + bs.Buf = make([]*batch.Batch, 0, max(16, len(staged.Buf))) + } + bs.Buf = append(bs.Buf, staged.Buf...) + bs.MemSize += staged.MemSize + staged.Buf = nil + staged.MemSize = 0 + return nil +} + +func (bs *Batches) copyIntoBatches( + src *batch.Batch, + proc *process.Process, + selection *vector.AllocationAccountSelection, +) (err error) { if bs.Buf == nil { bs.Buf = make([]*batch.Batch, 0, 16) } diff --git a/pkg/sql/colexec/join_util_test.go b/pkg/sql/colexec/join_util_test.go index a2b96845b29ee..94a91ca78ebff 100644 --- a/pkg/sql/colexec/join_util_test.go +++ b/pkg/sql/colexec/join_util_test.go @@ -94,7 +94,7 @@ func TestBatchesShrinkPreservesAllocationAndRollback(t *testing.T) { require.NoError(t, err) account, err := registry.Open(limit) require.NoError(t, err) - selection, err := vector.NewAllocationAccountSelectionWithBitmaps( + selection, err := vector.NewAllocationAccountSelection( account, 1, 1, 2, 3, 4) require.NoError(t, err) var batches Batches diff --git a/pkg/sql/colexec/loopjoin/join.go b/pkg/sql/colexec/loopjoin/join.go index e131626e5cc4a..499e1f4418727 100644 --- a/pkg/sql/colexec/loopjoin/join.go +++ b/pkg/sql/colexec/loopjoin/join.go @@ -17,13 +17,14 @@ package loopjoin import ( "bytes" - "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" "github.com/matrixorigin/matrixone/pkg/util/resource" "github.com/matrixorigin/matrixone/pkg/vm" "github.com/matrixorigin/matrixone/pkg/vm/message" @@ -58,6 +59,9 @@ func (loopJoin *LoopJoin) OpType() vm.OpType { func (loopJoin *LoopJoin) Prepare(proc *process.Process) error { var err error + if loopJoin.allocationAccount == nil { + return mpool.ErrAllocationAccountInvalid + } if loopJoin.OpAnalyzer == nil { loopJoin.OpAnalyzer = process.NewAnalyzer(loopJoin.GetIdx(), loopJoin.IsFirst, loopJoin.IsLast, opName) } else { @@ -65,10 +69,15 @@ func (loopJoin *LoopJoin) Prepare(proc *process.Process) error { } if loopJoin.NonEqCond != nil && loopJoin.ctr.expr == nil { - loopJoin.ctr.expr, err = colexec.NewExpressionExecutor(proc, loopJoin.NonEqCond) + var execs []colexec.ExpressionExecutor + execs, err = hashbuild.NewExpressionExecutors( + proc, + []*plan.Expr{loopJoin.NonEqCond}, + ) if err != nil { return err } + loopJoin.ctr.expr = execs[0] } return err } @@ -90,7 +99,9 @@ func (loopJoin *LoopJoin) Call(proc *process.Process) (vm.CallResult, error) { ctr.state = End } else { if loopJoin.JoinType == plan.Node_OUTER && ctr.mp != nil { - ctr.initRightMatchedBitmap() + if err = ctr.initRightMatchedBitmap(loopJoin, proc); err != nil { + return result, err + } } ctr.state = Probe } @@ -451,16 +462,43 @@ func (loopJoin *LoopJoin) resetResultBat() { } // initRightMatchedBitmap allocates the per-build-row matched bitmap. -func (ctr *container) initRightMatchedBitmap() { +func (ctr *container) initRightMatchedBitmap( + ap *LoopJoin, + proc *process.Process, +) error { bats := ctr.mp.GetBatches() - ctr.rightBatchOffset = make([]uint64, len(bats)) + var err error + ctr.rightBatchOffset, err = mpool.MakeSliceAccounted[uint64]( + len(bats), + proc.Mp(), + ap.allocationAccount, + hashbuild.HashBuildAllocationOwner, + loopJoinAllocationSiteBatchOffsets, + ) + if err != nil { + return err + } var total uint64 for i, b := range bats { ctr.rightBatchOffset[i] = total total += uint64(b.RowCount()) } - ctr.rightRowsMatched = &bitmap.Bitmap{} - ctr.rightRowsMatched.InitWithSize(int64(total)) + if total > uint64(^uint64(0)>>1) { + ctr.cleanRightMatchState(proc) + return mpool.ErrAllocationAccountInvalid + } + ctr.rightRowsMatched, err = colexec.NewAccountedBitmap( + int64(total), + proc.Mp(), + ap.allocationAccount, + hashbuild.HashBuildAllocationOwner, + loopJoinAllocationSiteMatched, + ) + if err != nil { + ctr.cleanRightMatchState(proc) + return err + } + return nil } // finalize emits one batch worth of unmatched build rows with NULL probe diff --git a/pkg/sql/colexec/loopjoin/join_test.go b/pkg/sql/colexec/loopjoin/join_test.go index 78146989376a4..93664e6b10cf3 100644 --- a/pkg/sql/colexec/loopjoin/join_test.go +++ b/pkg/sql/colexec/loopjoin/join_test.go @@ -52,6 +52,25 @@ type joinTestCase struct { resultBatch *batch.Batch } +type loopJoinTestAllocationOwner interface { + SetAllocationAccount(*mpool.AllocationAccount) error +} + +func installLoopJoinTestAllocation( + t testing.TB, + owners ...loopJoinTestAllocationOwner, +) *mpool.AllocationAccount { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.Open(1 << 60) + require.NoError(t, err) + for _, owner := range owners { + require.NoError(t, owner.SetAllocationAccount(account)) + } + return account +} + var ( tag int32 ) @@ -69,6 +88,36 @@ func TestString(t *testing.T) { } } +func TestResetRebuildsExpressionForNextAllocationGeneration(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + join := &LoopJoin{ + NonEqCond: &plan.Expr{ + Typ: plan.Type{Id: int32(types.T_bool)}, + Expr: &plan.Expr_Lit{Lit: &plan.Literal{ + Value: &plan.Literal_Bval{Bval: true}, + }}, + }, + } + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + + for range 2 { + account, openErr := registry.Open(1 << 20) + require.NoError(t, openErr) + require.NoError(t, join.SetAllocationAccount(account)) + require.NoError(t, join.Prepare(proc)) + require.NotNil(t, join.ctr.expr) + + join.Reset(proc, false, nil) + require.Nil(t, join.ctr.expr) + require.NoError(t, join.ClearAllocationAccount(account)) + terminal, _, terminalErr := registry.CompleteTerminal(account) + require.NoError(t, terminalErr) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) + } +} + func TestJoin(t *testing.T) { for _, tc := range makeTestCases(t) { @@ -296,6 +345,7 @@ func TestLoopJoinFinalizeResetsAfterPreviousEmptyProbe(t *testing.T) { }, }, } + installLoopJoinTestAllocation(t, join, build) resetChildrenWithBatch(join, makeInt32LoopJoinBatch(proc.Mp(), []int32{7})) resetHashBuildChildrenWithBatch(build, batch.EmptyBatch) @@ -398,6 +448,7 @@ func TestMarkJoinEmitsOneRowPerProbeRowAcrossBuildBatches(t *testing.T) { }, }, } + installLoopJoinTestAllocation(t, join, build) build.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{ makeInt32LoopJoinBatch(proc.Mp(), []int32{1}), makeInt32LoopJoinBatch(proc.Mp(), []int32{1}), @@ -487,6 +538,7 @@ func TestMarkJoinResumesAfterDefaultBatchSize(t *testing.T) { }, }, } + installLoopJoinTestAllocation(t, join, build) build.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{ makeInt32LoopJoinBatch(proc.Mp(), []int32{-1}), })) @@ -602,7 +654,7 @@ func newTestCase(t *testing.T, flgs []bool, ts []types.Type, rp []colexec.Result resultBatch.Vecs[i] = bat.Vecs[rp[i].Pos] } tag++ - return joinTestCase{ + testCase := joinTestCase{ types: ts, flgs: flgs, proc: proc, @@ -634,6 +686,8 @@ func newTestCase(t *testing.T, flgs []bool, ts []types.Type, rp []colexec.Result }, resultBatch: resultBatch, } + installLoopJoinTestAllocation(t, testCase.arg, testCase.barg) + return testCase } func resetChildren(arg *LoopJoin, m *mpool.MPool) { diff --git a/pkg/sql/colexec/loopjoin/types.go b/pkg/sql/colexec/loopjoin/types.go index d18eba6acd383..fa0f66b84446b 100644 --- a/pkg/sql/colexec/loopjoin/types.go +++ b/pkg/sql/colexec/loopjoin/types.go @@ -37,6 +37,11 @@ const ( End ) +const ( + loopJoinAllocationSiteMatched mpool.AllocationSite = iota + 92 + loopJoinAllocationSiteBatchOffsets +) + type container struct { state int probeIdx int @@ -60,18 +65,54 @@ type container struct { } type LoopJoin struct { - ctr container - LeftTypes []types.Type - RightTypes []types.Type - NonEqCond *plan.Expr - ResultCols []colexec.ResultPos - JoinMapTag int32 - JoinType plan.Node_JoinType - MarkPos int + ctr container + LeftTypes []types.Type + RightTypes []types.Type + NonEqCond *plan.Expr + ResultCols []colexec.ResultPos + JoinMapTag int32 + JoinType plan.Node_JoinType + MarkPos int + allocationAccount *mpool.AllocationAccount vm.OperatorBase } +func (loopJoin *LoopJoin) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if account == nil || account.Handle() == 0 { + return mpool.ErrAllocationAccountInvalid + } + if loopJoin.allocationAccount != nil && + loopJoin.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if loopJoin.allocationAccount == account { + return nil + } + loopJoin.allocationAccount = account + return nil +} + +func (loopJoin *LoopJoin) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if loopJoin.allocationAccount == nil { + return nil + } + if loopJoin.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + ctr := &loopJoin.ctr + if ctr.mp != nil || ctr.expr != nil || ctr.rightRowsMatched != nil || + len(ctr.rightBatchOffset) != 0 { + return mpool.ErrAllocationAccountInvariant + } + loopJoin.allocationAccount = nil + return nil +} + func (loopJoin *LoopJoin) GetOperatorBase() *vm.OperatorBase { return &loopJoin.OperatorBase } @@ -106,12 +147,14 @@ func (loopJoin *LoopJoin) Release() { func (loopJoin *LoopJoin) Reset(proc *process.Process, pipelineFailed bool, err error) { ctr := &loopJoin.ctr - ctr.resetNonEqCondExecutor() + // The executor owns allocations from this execution generation. Prepared + // statements must rebuild it after the next account is installed instead + // of carrying generation-bound storage across Reset. + ctr.cleanNonEqCondExecutor() ctr.cleanHashMap() ctr.state = Build ctr.inBat = nil - ctr.rightRowsMatched = nil - ctr.rightBatchOffset = nil + ctr.cleanRightMatchState(proc) ctr.rightMatchedIter = nil ctr.rightMatchedBat = 0 } @@ -121,9 +164,21 @@ func (loopJoin *LoopJoin) Free(proc *process.Process, pipelineFailed bool, err e ctr.cleanBatch(proc.Mp()) ctr.cleanNonEqCondExecutor() + ctr.cleanRightMatchState(proc) } +func (ctr *container) cleanRightMatchState(proc *process.Process) { + colexec.FreeAccountedBitmap(ctr.rightRowsMatched, proc.Mp()) + ctr.rightRowsMatched = nil + if cap(ctr.rightBatchOffset) > 0 { + mpool.FreeSlice(proc.Mp(), ctr.rightBatchOffset) + } + ctr.rightBatchOffset = nil + ctr.rightMatchedIter = nil + ctr.rightMatchedBat = 0 +} + func (loopJoin *LoopJoin) ExecProjection(proc *process.Process, input *batch.Batch) (*batch.Batch, error) { return input, nil } @@ -139,12 +194,6 @@ func (ctr *container) cleanBatch(mp *mpool.MPool) { } } -func (ctr *container) resetNonEqCondExecutor() { - if ctr.expr != nil { - ctr.expr.ResetForNextQuery() - } -} - func (ctr *container) cleanNonEqCondExecutor() { if ctr.expr != nil { ctr.expr.Free() diff --git a/pkg/sql/colexec/product/product.go b/pkg/sql/colexec/product/product.go index 4826a46fd005d..e6edd1bc8893b 100644 --- a/pkg/sql/colexec/product/product.go +++ b/pkg/sql/colexec/product/product.go @@ -17,6 +17,7 @@ package product import ( "bytes" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/sql/colexec" @@ -38,6 +39,9 @@ func (product *Product) OpType() vm.OpType { } func (product *Product) Prepare(proc *process.Process) error { + if product.allocationAccount == nil || product.resultAllocation == nil { + return mpool.ErrAllocationAccountInvalid + } if product.OpAnalyzer == nil { product.OpAnalyzer = process.NewAnalyzer(product.GetIdx(), product.IsFirst, product.IsLast, "cross join") } else { @@ -94,20 +98,30 @@ func (product *Product) Call(proc *process.Process) (vm.CallResult, error) { continue } } - if ctr.bat == nil { + if ctr.mp == nil { ctr.inBat = nil continue } if ctr.rbat == nil { + buildBat := ctr.firstBuildBatch() + if buildBat == nil { + ctr.inBat = nil + continue + } ctr.rbat = batch.NewOffHeapWithSize(len(product.Result)) for i, rp := range product.Result { if rp.Rel == 0 { ctr.rbat.Vecs[i] = vector.NewOffHeapVecWithType(*ctr.inBat.Vecs[rp.Pos].GetType()) } else { - ctr.rbat.Vecs[i] = vector.NewOffHeapVecWithType(*ctr.bat.Vecs[rp.Pos].GetType()) + ctr.rbat.Vecs[i] = vector.NewOffHeapVecWithType(*buildBat.Vecs[rp.Pos].GetType()) } } + if err := ctr.rbat.SetAllocationAccount(product.resultAllocation); err != nil { + ctr.rbat.Clean(proc.Mp()) + ctr.rbat = nil + return result, err + } } else { ctr.rbat.CleanOnlyData() } @@ -137,48 +151,63 @@ func (product *Product) build(proc *process.Process, analyzer process.Analyzer) if mp == nil { return nil } - batches := mp.GetBatches() - - //maybe optimize this in the future - for i := range batches { - ctr.bat, err = ctr.bat.AppendWithCopy(proc.Ctx, proc.Mp(), batches[i]) - if err != nil { - return err - } - } - mp.Free() + ctr.mp = mp return nil } func (ctr *container) probe(ap *Product, proc *process.Process, result *vm.CallResult) error { count := ctr.inBat.RowCount() - count2 := ctr.bat.RowCount() - var i, j int - for j = ctr.probeIdx; j < count2; j++ { - for i = 0; i < count; i++ { - for k, rp := range ap.Result { - if rp.Rel == 0 { - if err := ctr.rbat.Vecs[k].UnionOne(ctr.inBat.Vecs[rp.Pos], int64(i), proc.Mp()); err != nil { - return err - } - } else { - if err := ctr.rbat.Vecs[k].UnionOne(ctr.bat.Vecs[rp.Pos], int64(j), proc.Mp()); err != nil { - return err + batches := ctr.mp.GetBatches() + for ctr.buildBatIdx < len(batches) { + buildBat := batches[ctr.buildBatIdx] + if buildBat == nil || buildBat.RowCount() == 0 { + ctr.buildBatIdx++ + ctr.buildRowIdx = 0 + continue + } + for row := ctr.buildRowIdx; row < buildBat.RowCount(); row++ { + for probeRow := 0; probeRow < count; probeRow++ { + for k, rp := range ap.Result { + if rp.Rel == 0 { + if err := ctr.rbat.Vecs[k].UnionOne(ctr.inBat.Vecs[rp.Pos], int64(probeRow), proc.Mp()); err != nil { + return err + } + } else { + if err := ctr.rbat.Vecs[k].UnionOne(buildBat.Vecs[rp.Pos], int64(row), proc.Mp()); err != nil { + return err + } } } } + ctr.rbat.AddRowCount(count) + ctr.buildRowIdx = row + 1 + if ctr.rbat.RowCount() >= colexec.DefaultBatchSize { + if ctr.buildRowIdx == buildBat.RowCount() { + ctr.buildBatIdx++ + ctr.buildRowIdx = 0 + } + result.Batch = ctr.rbat + return nil + } } - if ctr.rbat.Vecs[0].Length() >= colexec.DefaultBatchSize { - result.Batch = ctr.rbat - ctr.rbat.SetRowCount(ctr.rbat.Vecs[0].Length()) - ctr.probeIdx = j + 1 - return nil - } + ctr.buildBatIdx++ + ctr.buildRowIdx = 0 } - // ctr.rbat.AddRowCount(count * count2) - ctr.probeIdx = 0 - ctr.rbat.SetRowCount(ctr.rbat.Vecs[0].Length()) + ctr.buildBatIdx = 0 + ctr.buildRowIdx = 0 result.Batch = ctr.rbat ctr.inBat = nil return nil } + +func (ctr *container) firstBuildBatch() *batch.Batch { + if ctr.mp == nil { + return nil + } + for _, bat := range ctr.mp.GetBatches() { + if bat != nil && bat.RowCount() > 0 { + return bat + } + } + return nil +} diff --git a/pkg/sql/colexec/product/product_test.go b/pkg/sql/colexec/product/product_test.go index 77aedf8011c95..9f3c13a81f933 100644 --- a/pkg/sql/colexec/product/product_test.go +++ b/pkg/sql/colexec/product/product_test.go @@ -73,6 +73,12 @@ func TestPrepare(t *testing.T) { } } +func TestPrepareRequiresAllocationAccount(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + require.ErrorIs(t, (&Product{}).Prepare(proc), mpool.ErrAllocationAccountInvalid) +} + func TestProduct(t *testing.T) { for _, tc := range makeTestCases(t) { @@ -98,6 +104,7 @@ func TestProduct(t *testing.T) { tc.arg.Reset(tc.proc, false, nil) tc.barg.Reset(tc.proc, false, nil) + require.Zero(t, tc.arg.allocationAccount.Snapshot().Used) resetChildren(tc.arg, tc.proc.Mp()) resetHashBuildChildren(tc.barg, tc.proc.Mp()) @@ -122,6 +129,7 @@ func TestProduct(t *testing.T) { tc.arg.Reset(tc.proc, false, nil) tc.barg.Reset(tc.proc, false, nil) + require.Zero(t, tc.arg.allocationAccount.Snapshot().Used) tc.arg.Free(tc.proc, false, nil) tc.barg.Free(tc.proc, false, nil) @@ -130,6 +138,51 @@ func TestProduct(t *testing.T) { } } +func TestProductConsumesMultipleBuildBatchesWithoutCopy(t *testing.T) { + tc := newTestCase( + t, + []bool{false}, + []types.Type{types.T_int32.ToType()}, + []colexec.ResultPos{ + colexec.NewResultPos(0, 0), + colexec.NewResultPos(1, 0), + }, + ) + probe := colexec.MakeMockBatchs(tc.proc.Mp()) + build1 := colexec.MakeMockBatchs(tc.proc.Mp()) + build2 := colexec.MakeMockBatchs(tc.proc.Mp()) + tc.arg.Children = nil + tc.arg.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{probe})) + tc.barg.Children = nil + tc.barg.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{build1, build2})) + + require.NoError(t, tc.arg.Prepare(tc.proc)) + require.NoError(t, tc.barg.Prepare(tc.proc)) + _, err := vm.Exec(tc.barg, tc.proc) + require.NoError(t, err) + wantRows := probe.RowCount() * (build1.RowCount() + build2.RowCount()) + rows := 0 + for { + result, err := vm.Exec(tc.arg, tc.proc) + require.NoError(t, err) + if result.Batch != nil { + rows += result.Batch.RowCount() + } + if result.Status == vm.ExecStop { + break + } + } + require.Equal(t, wantRows, rows) + + tc.arg.Reset(tc.proc, false, nil) + tc.barg.Reset(tc.proc, false, nil) + require.Zero(t, tc.arg.allocationAccount.Snapshot().Used) + tc.arg.Free(tc.proc, false, nil) + tc.barg.Free(tc.proc, false, nil) + tc.proc.Free() + require.Zero(t, tc.proc.Mp().CurrNB()) +} + /* func BenchmarkProduct(b *testing.B) { for i := 0; i < b.N; i++ { @@ -170,7 +223,7 @@ func newTestCase(t *testing.T, flgs []bool, ts []types.Type, rp []colexec.Result resultBatch.Vecs[i] = vector.NewVec(*bat.Vecs[rp[i].Pos].GetType()) } tag++ - return productTestCase{ + tc := productTestCase{ types: ts, flgs: flgs, proc: proc, @@ -200,6 +253,13 @@ func newTestCase(t *testing.T, flgs []bool, ts []types.Type, rp []colexec.Result }, resultBatch: resultBatch, } + registry, err := mpool.NewAllocationAccountRegistry(1, 1<<20) + require.NoError(t, err) + account, err := registry.Open(1 << 60) + require.NoError(t, err) + require.NoError(t, tc.arg.SetAllocationAccount(account)) + require.NoError(t, tc.barg.SetAllocationAccount(account)) + return tc } func resetChildren(arg *Product, m *mpool.MPool) { bat := colexec.MakeMockBatchs(m) diff --git a/pkg/sql/colexec/product/types.go b/pkg/sql/colexec/product/types.go index 733b388af223b..e548aba211337 100644 --- a/pkg/sql/colexec/product/types.go +++ b/pkg/sql/colexec/product/types.go @@ -18,8 +18,11 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/reuse" "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" "github.com/matrixorigin/matrixone/pkg/vm" + "github.com/matrixorigin/matrixone/pkg/vm/message" "github.com/matrixorigin/matrixone/pkg/vm/process" ) @@ -32,19 +35,30 @@ const ( ) type container struct { - state int - probeIdx int - bat *batch.Batch - rbat *batch.Batch - inBat *batch.Batch + state int + buildBatIdx int + buildRowIdx int + rbat *batch.Batch + inBat *batch.Batch + mp *message.JoinMap } +const ( + productAllocationSiteResultData mpool.AllocationSite = iota + 94 + productAllocationSiteResultArea + productAllocationSiteResultNulls + productAllocationSiteResultGrouping +) + type Product struct { ctr container Result []colexec.ResultPos IsShuffle bool JoinMapTag int32 + allocationAccount *mpool.AllocationAccount + resultAllocation *vector.AllocationAccountSelection + vm.OperatorBase } @@ -52,6 +66,52 @@ func (product *Product) GetOperatorBase() *vm.OperatorBase { return &product.OperatorBase } +func (product *Product) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if account == nil || account.Handle() == 0 { + return mpool.ErrAllocationAccountInvalid + } + if product.allocationAccount != nil && + product.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if product.allocationAccount == account { + return nil + } + selection, err := vector.NewAllocationAccountSelection( + account, + hashbuild.HashBuildAllocationOwner, + productAllocationSiteResultData, + productAllocationSiteResultArea, + productAllocationSiteResultNulls, + productAllocationSiteResultGrouping, + ) + if err != nil { + return err + } + product.allocationAccount = account + product.resultAllocation = selection + return nil +} + +func (product *Product) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if product.allocationAccount == nil { + return nil + } + if product.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if product.ctr.mp != nil || product.ctr.rbat != nil { + return mpool.ErrAllocationAccountInvariant + } + product.allocationAccount = nil + product.resultAllocation = nil + return nil +} + func init() { reuse.CreatePool[Product]( func() *Product { @@ -80,15 +140,8 @@ func (product *Product) Release() { } func (product *Product) Reset(proc *process.Process, pipelineFailed bool, err error) { - if product.ctr.bat != nil { - product.ctr.bat.CleanOnlyData() - } - if product.ctr.rbat != nil { - product.ctr.rbat.CleanOnlyData() - } - product.ctr.inBat = nil + product.ctr.cleanBatch(proc.Mp()) product.ctr.state = Build - product.ctr.probeIdx = 0 } func (product *Product) Free(proc *process.Process, pipelineFailed bool, err error) { @@ -100,13 +153,15 @@ func (product *Product) ExecProjection(proc *process.Process, input *batch.Batch } func (ctr *container) cleanBatch(mp *mpool.MPool) { - if ctr.bat != nil { - ctr.bat.Clean(mp) - ctr.bat = nil - } if ctr.rbat != nil { ctr.rbat.Clean(mp) ctr.rbat = nil } + if ctr.mp != nil { + ctr.mp.Free() + ctr.mp = nil + } ctr.inBat = nil + ctr.buildBatIdx = 0 + ctr.buildRowIdx = 0 } diff --git a/pkg/sql/colexec/productl2/joinmap_account_test.go b/pkg/sql/colexec/productl2/joinmap_account_test.go new file mode 100644 index 0000000000000..478f732ab9482 --- /dev/null +++ b/pkg/sql/colexec/productl2/joinmap_account_test.go @@ -0,0 +1,97 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package productl2 + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" + "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/matrixorigin/matrixone/pkg/vectorindex/metric" + "github.com/matrixorigin/matrixone/pkg/vm" + "github.com/matrixorigin/matrixone/pkg/vm/message" + "github.com/stretchr/testify/require" +) + +func TestProductL2ReleasesProducerAccountedJoinMap(t *testing.T) { + mp := mpool.MustNewZero() + proc := testutil.NewProcessWithMPool(t, "", mp) + proc.SetMessageBoard(message.NewMessageBoard()) + registry, err := mpool.NewAllocationAccountRegistry(1, 1<<12) + require.NoError(t, err) + account, err := registry.Open(1 << 30) + require.NoError(t, err) + + arrayType := types.T_array_float32.ToType() + arrayType.Width = 2 + build := batch.NewWithSize(2) + build.Vecs[0] = vector.NewVec(arrayType) + require.NoError(t, vector.AppendArrayList( + build.Vecs[0], [][]float32{{0, 0}, {10, 10}}, nil, mp, + )) + build.Vecs[1] = vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixedList( + build.Vecs[1], []int64{10, 20}, nil, mp, + )) + build.SetRowCount(2) + + probe := batch.NewWithSize(1) + probe.Vecs[0] = vector.NewVec(arrayType) + require.NoError(t, vector.AppendArrayList( + probe.Vecs[0], [][]float32{{1, 1}}, nil, mp, + )) + probe.SetRowCount(1) + + const tag = int32(7001) + producer := &hashbuild.HashBuild{ + NeedBatches: true, + JoinMapTag: tag, + JoinMapRefCnt: 1, + } + producer.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{build})) + consumer := &Productl2{ + Result: []colexec.ResultPos{colexec.NewResultPos(1, 1)}, + OnExpr: onExprWithProbeCol(0), + JoinMapTag: tag, + VectorOpType: metric.OpType_L2Distance, + } + consumer.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{probe})) + require.NoError(t, producer.SetAllocationAccount(account)) + require.NoError(t, producer.Prepare(proc)) + require.NoError(t, consumer.Prepare(proc)) + + _, err = vm.Exec(producer, proc) + require.NoError(t, err) + result, err := vm.Exec(consumer, proc) + require.NoError(t, err) + require.NotNil(t, result.Batch) + require.Equal(t, []int64{10}, vector.MustFixedColNoTypeCheck[int64](result.Batch.Vecs[0])) + + consumer.Reset(proc, false, nil) + producer.Reset(proc, false, nil) + require.Zero(t, account.Snapshot().Used) + require.NoError(t, producer.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + consumer.Free(proc, false, nil) + producer.Free(proc, false, nil) + proc.Free() + require.Zero(t, mp.CurrNB()) +} diff --git a/pkg/sql/colexec/productl2/product_l2.go b/pkg/sql/colexec/productl2/product_l2.go index 468d0bb76cb2b..e94a0a4648d46 100644 --- a/pkg/sql/colexec/productl2/product_l2.go +++ b/pkg/sql/colexec/productl2/product_l2.go @@ -186,8 +186,16 @@ func (productl2 *Productl2) build(proc *process.Process, analyzer process.Analyz return nil } batches := mp.GetBatches() - //maybe optimize this in the future + // ProductL2 index/scratch is outside the first HashBuild accounting domain. + // Create an explicit unaccounted destination instead of letting a nil + // AppendWithCopy clone inherit the producer's allocation selection. for i := range batches { + if ctr.bat == nil { + ctr.bat = batch.NewOffHeapWithSize(len(batches[i].Vecs)) + for j, source := range batches[i].Vecs { + ctr.bat.Vecs[j] = vector.NewOffHeapVecWithType(*source.GetType()) + } + } ctr.bat, err = ctr.bat.AppendWithCopy(proc.Ctx, proc.Mp(), batches[i]) if err != nil { return err diff --git a/pkg/sql/colexec/receiver_operator.go b/pkg/sql/colexec/receiver_operator.go deleted file mode 100644 index 4f9cab8451f2b..0000000000000 --- a/pkg/sql/colexec/receiver_operator.go +++ /dev/null @@ -1,32 +0,0 @@ -// Copyright 2023 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package colexec - -import ( - "context" - "github.com/matrixorigin/matrixone/pkg/common/bitmap" -) - -func ReceiveBitmapFromChannel(usr context.Context, ch chan *bitmap.Bitmap) *bitmap.Bitmap { - select { - case <-usr.Done(): - return nil - case bm, ok := <-ch: - if !ok { - return nil - } - return bm - } -} diff --git a/pkg/sql/colexec/rightdedupjoin/allocation_test_helpers_test.go b/pkg/sql/colexec/rightdedupjoin/allocation_test_helpers_test.go new file mode 100644 index 0000000000000..dbe509f9c25a7 --- /dev/null +++ b/pkg/sql/colexec/rightdedupjoin/allocation_test_helpers_test.go @@ -0,0 +1,38 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package rightdedupjoin + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/stretchr/testify/require" +) + +type testAllocationOwner interface { + SetAllocationAccount(*mpool.AllocationAccount) error +} + +func installTestAllocation(t testing.TB, owners ...testAllocationOwner) *mpool.AllocationAccount { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.Open(1 << 60) + require.NoError(t, err) + for _, owner := range owners { + require.NoError(t, owner.SetAllocationAccount(account)) + } + return account +} diff --git a/pkg/sql/colexec/rightdedupjoin/expression_memory_test.go b/pkg/sql/colexec/rightdedupjoin/expression_memory_test.go deleted file mode 100644 index a7e9f00be9e4a..0000000000000 --- a/pkg/sql/colexec/rightdedupjoin/expression_memory_test.go +++ /dev/null @@ -1,108 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package rightdedupjoin - -import ( - "testing" - - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/container/batch" - "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/matrixorigin/matrixone/pkg/container/vector" - "github.com/matrixorigin/matrixone/pkg/pb/plan" - "github.com/matrixorigin/matrixone/pkg/sql/colexec" - "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" - "github.com/matrixorigin/matrixone/pkg/testutil" - "github.com/matrixorigin/matrixone/pkg/vm/process" - "github.com/stretchr/testify/require" -) - -func TestRightDedupJoinResetReleasesProbeExpressionLease(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Lit{Lit: &plan.Literal{ - Value: &plan.Literal_I32Val{I32Val: 1}, - }}, - } - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - lease, err := hashbuild.NewExpressionMemoryLease( - generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - - arg := &RightDedupJoin{} - arg.ctr.evecs = []evalVector{{executor: executors[0]}} - arg.ctr.vecs = make([]*vector.Vector, len(executors)) - arg.ctr.probeExpressionLease = lease - input := batch.NewWithSize(0) - input.SetRowCount(4) - require.NoError(t, arg.ctr.evalJoinConditionBudgeted(input, proc)) - require.Positive(t, generation.Used()) - - arg.Reset(proc, false, nil) - require.Zero(t, generation.Used()) - require.Nil(t, arg.ctr.evecs) - require.Nil(t, arg.ctr.vecs) - require.Nil(t, arg.ctr.probeExpressionLease) -} - -func TestRightDedupJoinResetReleasesAccountedProbeExpressions(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - const capBytes = uint64(1 << 20) - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - registry, err := mpool.NewAllocationAccountRegistry(1, 16) - require.NoError(t, err) - account, err := registry.OpenWithController(capBytes, generation) - require.NoError(t, err) - expr := &plan.Expr{Typ: plan.Type{Id: int32(types.T_int32)}, Expr: &plan.Expr_Lit{Lit: &plan.Literal{Value: &plan.Literal_I32Val{I32Val: 1}}}} - executors, err := hashbuild.NewAllocationAccountedExpressionExecutorsForAccount( - proc, []*plan.Expr{expr}, account, hashbuild.HashBuildAllocationOwner) - require.NoError(t, err) - arg := &RightDedupJoin{allocationAccount: account} - arg.ctr.evecs = []evalVector{{executor: executors[0]}} - arg.ctr.vecs = make([]*vector.Vector, len(executors)) - arg.ctr.probeExpressionsAccounted = true - input := batch.NewWithSize(0) - input.SetRowCount(4) - require.NoError(t, arg.ctr.evalJoinConditionBudgeted(input, proc)) - require.Positive(t, account.Snapshot().Used) - - arg.Reset(proc, false, nil) - require.Zero(t, account.Snapshot().Used) - require.Zero(t, generation.Used()) - require.False(t, arg.ctr.probeExpressionsAccounted) - require.Nil(t, arg.ctr.evecs) - terminal, _, err := registry.CompleteTerminal(account) - require.NoError(t, err) - require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) -} - -func TestRightDedupJoinAllocationActivationRequiresBothKeySides(t *testing.T) { - col := &plan.Expr{Typ: plan.Type{Id: int32(types.T_int64)}, Expr: &plan.Expr_Col{Col: &plan.ColRef{}}} - arg := &RightDedupJoin{Conditions: [][]*plan.Expr{{col}, {col}}} - require.True(t, arg.AllocationAccountEnabled()) - require.False(t, arg.AllocationAccountActivationBlocked()) - arg.Conditions[1] = []*plan.Expr{nil} - require.False(t, arg.AllocationAccountEnabled()) - require.True(t, arg.AllocationAccountActivationBlocked()) -} diff --git a/pkg/sql/colexec/rightdedupjoin/join.go b/pkg/sql/colexec/rightdedupjoin/join.go index 0521a15ec8f54..55e157199c32b 100644 --- a/pkg/sql/colexec/rightdedupjoin/join.go +++ b/pkg/sql/colexec/rightdedupjoin/join.go @@ -19,9 +19,9 @@ import ( "strings" "github.com/matrixorigin/matrixone/pkg/catalog" - "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/hashmap" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" @@ -47,6 +47,9 @@ func (rightDedupJoin *RightDedupJoin) OpType() vm.OpType { } func (rightDedupJoin *RightDedupJoin) Prepare(proc *process.Process) (err error) { + if rightDedupJoin.allocationAccount == nil { + return mpool.ErrAllocationAccountInvalid + } if rightDedupJoin.OpAnalyzer == nil { rightDedupJoin.OpAnalyzer = process.NewAnalyzer(rightDedupJoin.GetIdx(), rightDedupJoin.IsFirst, rightDedupJoin.IsLast, "dedup join") } else { @@ -58,13 +61,19 @@ func (rightDedupJoin *RightDedupJoin) Prepare(proc *process.Process) (err error) newUpdateExecs := len(rightDedupJoin.ctr.exprExecs) == 0 && len(rightDedupJoin.UpdateColExprList) > 0 var evalExecs, updateExecs []colexec.ExpressionExecutor if newEvalVectors { - evalExecs, err = colexec.NewExpressionExecutorsFromPlanExpressions(proc, rightDedupJoin.Conditions[0]) + evalExecs, err = hashbuild.NewExpressionExecutors( + proc, + rightDedupJoin.Conditions[0], + ) if err != nil { return err } } if newUpdateExecs { - updateExecs, err = colexec.NewExpressionExecutorsFromPlanExpressions(proc, rightDedupJoin.UpdateColExprList) + updateExecs, err = hashbuild.NewExpressionExecutors( + proc, + rightDedupJoin.UpdateColExprList, + ) if err != nil { for _, exec := range evalExecs { exec.Free() @@ -147,7 +156,7 @@ func (rightDedupJoin *RightDedupJoin) Call(proc *process.Process) (vm.CallResult if ctr.spillEngine != nil { // Clear previous bucket state before advancing. ctr.cleanHashMap() - ctr.matched = nil + ctr.cleanBitmap(proc) ctr.groupCount = 0 ctr.buildGroupCount = 0 var initErr error @@ -159,8 +168,13 @@ func (rightDedupJoin *RightDedupJoin) Call(proc *process.Process) (vm.CallResult ctr.groupCount = jm.GetGroupCount() ctr.buildGroupCount = ctr.groupCount if !proc.GetTxnOperator().Txn().IsPessimistic() && ctr.buildGroupCount > 0 { - ctr.matched = &bitmap.Bitmap{} - ctr.matched.InitWithSize(int64(ctr.buildGroupCount)) + ctr.matched, initErr = colexec.NewAccountedBitmap( + int64(ctr.buildGroupCount), + proc.Mp(), + rightDedupJoin.allocationAccount, + hashbuild.HashBuildAllocationOwner, + rightDedupJoinAllocationSiteMatched, + ) } case spillutil.BucketEmptyBuild: ctr.mp, initErr = rightDedupJoin.newEmptyJoinMap(proc) @@ -212,66 +226,29 @@ func (rightDedupJoin *RightDedupJoin) build(analyzer process.Analyzer, proc *pro if takeErr != nil { return takeErr } - var probeExpressionLease *hashbuild.ExpressionMemoryLease - var leaseErr error - if rightDedupJoin.allocationAccount != nil && - hashbuild.AllocationAccountedExpressionSetSupported(rightDedupJoin.Conditions[0]) { - ctr.cleanEvalVectors() - var probeExecutors []colexec.ExpressionExecutor - probeExecutors, leaseErr = - hashbuild.NewAllocationAccountedExpressionExecutorsForAccount( - proc, - rightDedupJoin.Conditions[0], - rightDedupJoin.allocationAccount, - hashbuild.HashBuildAllocationOwner, - ) - if leaseErr == nil { - ctr.evecs = make([]evalVector, len(probeExecutors)) - ctr.vecs = make([]*vector.Vector, len(probeExecutors)) - for i := range probeExecutors { - ctr.evecs[i].executor = probeExecutors[i] - } - ctr.probeExpressionsAccounted = true - } - } else { - probeExecutors := make([]colexec.ExpressionExecutor, len(ctr.evecs)) - for i := range ctr.evecs { - probeExecutors[i] = ctr.evecs[i].executor - } - probeExpressionLease, leaseErr = hashbuild.NewExpressionMemoryLease( - budget, rightDedupJoin.Conditions[0], probeExecutors, false) - } - if leaseErr != nil { + if rightDedupJoin.allocationAccount == nil { _ = payload.Close() ctr.mp.Free() ctr.mp = nil - ctr.cleanEvalVectors() - ctr.releaseProbeExpressionLease() - return leaseErr + return mpool.ErrAllocationAccountInvalid } - ctr.probeExpressionLease = probeExpressionLease - engine, engineErr := spillutil.NewSpillEngineForAccount(spillutil.SpillEngineConfig{ + engine, engineErr := spillutil.NewSpillEngine(spillutil.SpillEngineConfig{ BuildKeyExprs: rightDedupJoin.Conditions[1], ProbeKeyExprs: rightDedupJoin.Conditions[0], SpillThreshold: ctr.spillThreshold, NeedsProbeForEmptyBuild: true, MergeProbeBatches: true, Budget: budget, - ProbeExpressionLease: probeExpressionLease, }, rightDedupJoin.allocationAccount, hashbuild.HashBuildAllocationOwner) if engineErr != nil { _ = payload.Close() ctr.mp.Free() ctr.mp = nil ctr.cleanEvalVectors() - ctr.releaseProbeExpressionLease() return engineErr } - if len(payload.Files) > 0 { - engine.InitFromSpilledFiles(payload.Files) - } else { - engine.InitFromSpilledMap(payload.LegacyFds) - } + engine.InitFromSpilledFiles(payload.Files) + ctr.spillEngine = engine if err := engine.ScatterProbeTable(proc, func() (*batch.Batch, error) { input, err := vm.ChildrenCall(rightDedupJoin.GetChildren(0), proc, analyzer) @@ -279,7 +256,7 @@ func (rightDedupJoin *RightDedupJoin) build(analyzer process.Analyzer, proc *pro }, analyzer, func(bat *batch.Batch) ([]*vector.Vector, error) { - if err := ctr.evalJoinConditionBudgeted(bat, proc); err != nil { + if err := ctr.evalJoinCondition(bat, proc); err != nil { return nil, err } return ctr.vecs, nil @@ -288,18 +265,27 @@ func (rightDedupJoin *RightDedupJoin) build(analyzer process.Analyzer, proc *pro ctr.mp.Free() ctr.mp = nil engine.Cleanup(proc) + ctr.spillEngine = nil return err } ctr.mp.Free() - ctr.spillEngine = engine ctr.mp = nil return } ctr.groupCount = ctr.mp.GetGroupCount() ctr.buildGroupCount = ctr.groupCount - if !proc.GetTxnOperator().Txn().IsPessimistic() { - ctr.matched = &bitmap.Bitmap{} - ctr.matched.InitWithSize(int64(ctr.buildGroupCount)) + if !proc.GetTxnOperator().Txn().IsPessimistic() && + ctr.buildGroupCount > 0 { + ctr.matched, err = colexec.NewAccountedBitmap( + int64(ctr.buildGroupCount), + proc.Mp(), + rightDedupJoin.allocationAccount, + hashbuild.HashBuildAllocationOwner, + rightDedupJoinAllocationSiteMatched, + ) + if err != nil { + return err + } } } @@ -317,22 +303,18 @@ func (rightDedupJoin *RightDedupJoin) newEmptyJoinMap(proc *process.Process) (*m keyWidth += width } - budget, err := proc.GetHashBuildBudget() - if err != nil { - return nil, err - } - if rightDedupJoin.allocationAccount != nil { - return hashbuild.NewAccountedEmptyJoinMap( - keyWidth, - rightDedupJoin.allocationAccount, - proc.Mp(), - ) + if rightDedupJoin.allocationAccount == nil { + return nil, mpool.ErrAllocationAccountInvalid } - return hashbuild.NewBudgetedEmptyJoinMap(keyWidth, budget, proc.Mp()) + return hashbuild.NewAccountedEmptyJoinMap( + keyWidth, + rightDedupJoin.allocationAccount, + proc.Mp(), + ) } func (ctr *container) probe(bat *batch.Batch, ap *RightDedupJoin, proc *process.Process, analyzer process.Analyzer, result *vm.CallResult) error { - err := ctr.evalJoinConditionBudgeted(bat, proc) + err := ctr.evalJoinCondition(bat, proc) if err != nil { return err } @@ -457,14 +439,3 @@ func (ctr *container) evalJoinCondition(bat *batch.Batch, proc *process.Process) } return nil } - -func (ctr *container) evalJoinConditionBudgeted(bat *batch.Batch, proc *process.Process) error { - if ctr.probeExpressionLease == nil { - return ctr.evalJoinCondition(bat, proc) - } - return ctr.probeExpressionLease.Eval(proc, []*batch.Batch{bat}, bat.RowCount(), func(i int, vec *vector.Vector) error { - ctr.vecs[i] = vec - ctr.evecs[i].vec = vec - return nil - }) -} diff --git a/pkg/sql/colexec/rightdedupjoin/join_test.go b/pkg/sql/colexec/rightdedupjoin/join_test.go index 904c953db2718..31adeb1c2d067 100644 --- a/pkg/sql/colexec/rightdedupjoin/join_test.go +++ b/pkg/sql/colexec/rightdedupjoin/join_test.go @@ -18,7 +18,6 @@ import ( "bytes" "context" "fmt" - "os" "testing" "github.com/golang/mock/gomock" @@ -104,6 +103,7 @@ func runRightDedupCase(t *testing.T, buildVals, probeVals []int32, pessimistic, JoinMapTag: curTag, } arg.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{probeBat})) + installTestAllocation(t, arg, buildArg) require.NoError(t, buildArg.Prepare(proc)) require.NoError(t, arg.Prepare(proc)) @@ -154,6 +154,13 @@ func TestRightDedupDuplicateTracking(t *testing.T) { func runRightDedupSpilledEmptyBuild(t *testing.T, pessimistic, duplicateAcrossBatches bool) { proc, ctrl := newRightDedupTestProcess(t, pessimistic) defer ctrl.Finish() + budget := process.MustNewHashBuildBudget(64<<20, 64<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 1<<20) + require.NoError(t, err) + account, err := registry.OpenWithController(64<<20, generation) + require.NoError(t, err) typ := types.T_int32.ToType() tag++ curTag := tag @@ -171,6 +178,7 @@ func runRightDedupSpilledEmptyBuild(t *testing.T, pessimistic, duplicateAcrossBa JoinMapTag: curTag, SpillThreshold: 1, } + require.NoError(t, arg.SetAllocationAccount(account)) probeValues := [][]int32{{1}, {2}} if duplicateAcrossBatches { @@ -188,14 +196,14 @@ func runRightDedupSpilledEmptyBuild(t *testing.T, pessimistic, duplicateAcrossBa jm := message.NewJoinMap(message.GroupSels{}, nil, nil, nil, nil, proc.Mp()) jm.IncRef(1) require.NoError(t, jm.SetSpillBuildPayload(message.SpillBuildPayload{ - LegacyFds: make([]*os.File, spillutil.SpillNumBuckets), + Files: make([]*message.SpillFile, spillutil.SpillNumBuckets), + BudgetRef: generation, })) message.SendMessage(message.JoinMapMsg{ - JoinMapPtr: jm, + Result: message.NewJoinMapResult(jm), IsShuffle: true, ShuffleIdx: 0, Tag: curTag, - Spilled: true, }, proc.GetMessageBoard()) require.NoError(t, arg.Prepare(proc)) @@ -220,6 +228,9 @@ func runRightDedupSpilledEmptyBuild(t *testing.T, pessimistic, duplicateAcrossBa } arg.Free(proc, false, nil) + require.Zero(t, account.Snapshot().Used) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) proc.Free() require.Equal(t, int64(0), proc.Mp().CurrNB()) } @@ -244,6 +255,7 @@ func TestRightDedupResetAndPrepareRetry(t *testing.T) { Conditions: [][]*plan.Expr{{valid}, {valid}}, UpdateColExprList: []*plan.Expr{valid, invalid}, } + installTestAllocation(t, arg) require.Error(t, arg.Prepare(proc)) require.Nil(t, arg.ctr.vecs) @@ -268,6 +280,7 @@ func TestRightDedupEmptyMapUsesEvaluatedKeyType(t *testing.T) { LeftTypes: []types.Type{types.T_int32.ToType()}, Conditions: [][]*plan.Expr{{newExpr(0, varcharTyp)}, {newExpr(0, varcharTyp)}}, } + installTestAllocation(t, arg) jm, err := arg.newEmptyJoinMap(proc) require.NoError(t, err) require.NoError(t, jm.PreAlloc(2)) @@ -311,6 +324,11 @@ func TestRightDedupEmptyBuildProbeMapHonorsHashBuildBudget(t *testing.T) { DedupColTypes: []plan.Type{{Id: int32(types.T_int32)}}, JoinMapTag: tag, } + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.OpenWithController(initialBytes, budget) + require.NoError(t, err) + require.NoError(t, arg.SetAllocationAccount(account)) arg.AppendChild(probeSource) var callErr error t.Cleanup(func() { @@ -537,7 +555,7 @@ func newTestCase(t *testing.T, flgs []bool, ts []types.Type, rp []int32, cs [][] // }, //}) tag++ - return joinTestCase{ + tc := joinTestCase{ types: ts, flgs: flgs, proc: proc, @@ -569,6 +587,8 @@ func newTestCase(t *testing.T, flgs []bool, ts []types.Type, rp []int32, cs [][] JoinMapRefCnt: 1, }, } + installTestAllocation(t, tc.arg, tc.barg) + return tc } func resetChildren(arg *RightDedupJoin, m *mpool.MPool) { diff --git a/pkg/sql/colexec/rightdedupjoin/key_contract_test.go b/pkg/sql/colexec/rightdedupjoin/key_contract_test.go index ecdc21fe2ed84..3ca8e1531e632 100644 --- a/pkg/sql/colexec/rightdedupjoin/key_contract_test.go +++ b/pkg/sql/colexec/rightdedupjoin/key_contract_test.go @@ -160,6 +160,7 @@ func runRightDedupJoinDoubleSignedZeroContract( if mode.shuffle { buildArg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: joinMapTag + 9000} } + installTestAllocation(t, rightDedupArg, buildArg) buildArg.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{buildBatch})) spillBefore := promtestutil.ToFloat64( diff --git a/pkg/sql/colexec/rightdedupjoin/types.go b/pkg/sql/colexec/rightdedupjoin/types.go index d7caa01c3041b..818141ddcdd0a 100644 --- a/pkg/sql/colexec/rightdedupjoin/types.go +++ b/pkg/sql/colexec/rightdedupjoin/types.go @@ -24,7 +24,6 @@ import ( "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" - "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" "github.com/matrixorigin/matrixone/pkg/sql/colexec/spillutil" "github.com/matrixorigin/matrixone/pkg/vm" "github.com/matrixorigin/matrixone/pkg/vm/message" @@ -40,6 +39,8 @@ const ( End ) +const rightDedupJoinAllocationSiteMatched mpool.AllocationSite = 90 + type evalVector struct { executor colexec.ExpressionExecutor vec *vector.Vector @@ -65,12 +66,7 @@ type container struct { spillEngine *spillutil.SpillEngine spillThreshold int64 - // Non-nil only for spilled joins, where probe expressions are part of the - // shared HashBuild/spill working set. Resident probe expressions remain - // under normal process/mpool accounting; this is not a general query budget. - probeExpressionLease *hashbuild.ExpressionMemoryLease - probeExpressionsAccounted bool - resultBatch *batch.Batch + resultBatch *batch.Batch } type RightDedupJoin struct { @@ -97,24 +93,6 @@ type RightDedupJoin struct { vm.OperatorBase } -func (rightDedupJoin *RightDedupJoin) AllocationAccountEnabled() bool { - return rightDedupJoin != nil && - rightDedupJoin.allocationAccountExpressionOwnerClosed() -} - -func (rightDedupJoin *RightDedupJoin) AllocationAccountActivationBlocked() bool { - return rightDedupJoin != nil && - !rightDedupJoin.allocationAccountExpressionOwnerClosed() -} - -func (rightDedupJoin *RightDedupJoin) allocationAccountExpressionOwnerClosed() bool { - if rightDedupJoin == nil || len(rightDedupJoin.Conditions) != 2 { - return false - } - return hashbuild.AllocationAccountedExpressionSetSupported(rightDedupJoin.Conditions[0]) && - hashbuild.AllocationAccountedExpressionSetSupported(rightDedupJoin.Conditions[1]) -} - func (rightDedupJoin *RightDedupJoin) SetAllocationAccount( account *mpool.AllocationAccount, ) error { @@ -125,6 +103,9 @@ func (rightDedupJoin *RightDedupJoin) SetAllocationAccount( rightDedupJoin.allocationAccount != account { return mpool.ErrAllocationAccountMismatch } + if rightDedupJoin.allocationAccount == account { + return nil + } rightDedupJoin.allocationAccount = account return nil } @@ -140,7 +121,9 @@ func (rightDedupJoin *RightDedupJoin) ClearAllocationAccount( } if rightDedupJoin.ctr.mp != nil || rightDedupJoin.ctr.spillEngine != nil || - rightDedupJoin.ctr.probeExpressionsAccounted { + len(rightDedupJoin.ctr.evecs) != 0 || + len(rightDedupJoin.ctr.exprExecs) != 0 || + rightDedupJoin.ctr.matched != nil { return mpool.ErrAllocationAccountInvariant } rightDedupJoin.allocationAccount = nil @@ -184,31 +167,26 @@ func (rightDedupJoin *RightDedupJoin) Reset(proc *process.Process, pipelineFaile rightDedupJoin.OpAnalyzer.Alloc(ctr.maxAllocSize) } ctr.maxAllocSize = 0 + hashmap.IteratorClearOwner(ctr.itr) ctr.itr = nil ctr.groupCount = 0 ctr.buildGroupCount = 0 - ctr.cleanBitmap() + ctr.cleanBitmap(proc) ctr.cleanHashMap() ctr.resetResultBatch() - ctr.resetExprExecutor() + ctr.cleanExprExecutor() if ctr.spillEngine != nil { ctr.spillEngine.Cleanup(proc) ctr.spillEngine = nil } - if ctr.probeExpressionLease != nil || ctr.probeExpressionsAccounted { - ctr.cleanEvalVectors() - ctr.releaseProbeExpressionLease() - } else { - ctr.resetEvalVectors() - } + ctr.cleanEvalVectors() ctr.state = Build - rightDedupJoin.allocationAccount = nil } func (rightDedupJoin *RightDedupJoin) Free(proc *process.Process, pipelineFailed bool, err error) { ctr := &rightDedupJoin.ctr - ctr.cleanBitmap() + ctr.cleanBitmap(proc) ctr.cleanHashMap() ctr.cleanResultBatch(proc) ctr.cleanExprExecutor() @@ -217,28 +195,23 @@ func (rightDedupJoin *RightDedupJoin) Free(proc *process.Process, pipelineFailed ctr.spillEngine = nil } ctr.cleanEvalVectors() - ctr.releaseProbeExpressionLease() - rightDedupJoin.allocationAccount = nil } func (rightDedupJoin *RightDedupJoin) ExecProjection(proc *process.Process, input *batch.Batch) (*batch.Batch, error) { return input, nil } -func (ctr *container) resetExprExecutor() { - for i := range ctr.exprExecs { - ctr.exprExecs[i].ResetForNextQuery() - } -} - func (ctr *container) cleanExprExecutor() { for i := range ctr.exprExecs { - ctr.exprExecs[i].Free() - ctr.exprExecs[i] = nil + if ctr.exprExecs[i] != nil { + ctr.exprExecs[i].Free() + } } + ctr.exprExecs = nil } func (ctr *container) cleanHashMap() { + hashmap.IteratorClearOwner(ctr.itr) ctr.itr = nil if ctr.mp != nil { ctr.mp.Free() @@ -246,7 +219,8 @@ func (ctr *container) cleanHashMap() { } } -func (ctr *container) cleanBitmap() { +func (ctr *container) cleanBitmap(proc *process.Process) { + colexec.FreeAccountedBitmap(ctr.matched, proc.Mp()) ctr.matched = nil } @@ -277,20 +251,4 @@ func (ctr *container) cleanEvalVectors() { } ctr.evecs = nil ctr.vecs = nil - ctr.probeExpressionsAccounted = false -} - -func (ctr *container) resetEvalVectors() { - for i := range ctr.evecs { - if ctr.evecs[i].executor != nil { - ctr.evecs[i].executor.ResetForNextQuery() - } - } -} - -func (ctr *container) releaseProbeExpressionLease() { - if ctr.probeExpressionLease != nil { - ctr.probeExpressionLease.Release() - ctr.probeExpressionLease = nil - } } diff --git a/pkg/sql/colexec/runtimefilter/contract.go b/pkg/sql/colexec/runtimefilter/contract.go index fbfd8e2423dc6..7b87a4ca2d2f5 100644 --- a/pkg/sql/colexec/runtimefilter/contract.go +++ b/pkg/sql/colexec/runtimefilter/contract.go @@ -18,7 +18,6 @@ package runtimefilter import ( - "bytes" "context" "errors" "math" @@ -84,9 +83,7 @@ func ClassifyOptionalFallback(err error) OptionalFallbackKind { // traversal order chosen by errors.As below. if errors.Is(err, process.ErrHashBuildBudgetClosed) || errors.Is(err, process.ErrHashBuildBudgetInvalid) || - errors.Is(err, process.ErrHashBuildCeilingMissing) || - errors.Is(err, process.ErrHashBuildReservationInactive) || - errors.Is(err, process.ErrHashBuildReservationUpward) { + errors.Is(err, process.ErrHashBuildCeilingMissing) { return OptionalFallbackNone } @@ -375,84 +372,41 @@ func CloseFloatSignedZero( vector.AppendFixed(vec, value, false, mp)) } -// MarshalExactFilterVector serializes an exact-filter vector under the -// statement/CN hash-build budget. Runtime-filter payloads live on the Go heap, -// outside mpool accounting, so every producer must retain this reservation -// until the MessageBoard destroys the message. -// -// Exact IN payloads have already discarded NULL. Requiring an empty null -// bitmap makes the wire size exact before allocation and avoids a second, -// unbudgeted roaring-bitmap serialization. +// MarshalExactFilterVector serializes an exact-filter vector into physical +// MPool storage owned by the statement allocation account. The returned +// release closure transfers that storage lifetime to the MessageBoard. func MarshalExactFilterVector( vec *vector.Vector, - budget *process.HashBuildBudgetGeneration, + mp *mpool.MPool, + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, + site mpool.AllocationSite, ) ([]byte, func(), error) { - if vec == nil || budget == nil || vec.GetNulls().Any() { - return nil, nil, process.ErrHashBuildBudgetInvalid - } - - length := vec.Length() - typeSize := vec.GetType().TypeSize() - if length < 0 || uint64(length) > math.MaxUint32 || typeSize < 0 { - return nil, nil, process.ErrHashBuildBudgetInvalid - } - dataBytes := uint64(typeSize) - if !vec.IsConst() { - if typeSize > 0 && uint64(length) > math.MaxUint64/uint64(typeSize) { - return nil, nil, process.ErrHashBuildBudgetInvalid - } - dataBytes *= uint64(length) - } else if vec.IsConstNull() { - dataBytes = 0 - } - areaBytes := uint64(len(vec.GetArea())) - if dataBytes > math.MaxUint32 || areaBytes > math.MaxUint32 || - dataBytes > uint64(len(vec.GetData())) { - return nil, nil, process.ErrHashBuildBudgetInvalid - } - - // class + encoded type + length/data/area/null lengths + sorted flag. - headerBytes := uint64(1 + len(types.EncodeType(vec.GetType())) + 4*4 + 1) - if dataBytes > math.MaxUint64-headerBytes || - areaBytes > math.MaxUint64-headerBytes-dataBytes { + if vec == nil || mp == nil || account == nil || vec.GetNulls().Any() { return nil, nil, process.ErrHashBuildBudgetInvalid } - wireBytes := headerBytes + dataBytes + areaBytes - if wireBytes > uint64(math.MaxInt) { - return nil, nil, process.ErrHashBuildBudgetInvalid - } - - // bytes.Buffer's visible capacity may be rounded above Grow's request. - // Reserve a bounded allocator overlap, verify it after allocation, then - // reconcile to the capacity retained by the message. - const allocationSlack = uint64(64 << 10) - if wireBytes > math.MaxUint64-allocationSlack { - return nil, nil, process.ErrHashBuildBudgetInvalid + plan, err := vec.PrepareMarshalBinary() + if err != nil { + return nil, nil, err } - projected := wireBytes + allocationSlack - token, err := budget.Reserve(projected) + buf, err := mpool.NewAccountedBuffer(mp, account, owner, site) if err != nil { return nil, nil, err } - - var buf bytes.Buffer - buf.Grow(int(wireBytes)) - if uint64(buf.Cap()) > projected { - token.Release() - return nil, nil, process.ErrHashBuildBudgetInvalid + if err = buf.EnsureCapacity(plan.Size()); err != nil { + buf.Free() + if mpool.IsRetryableAllocationCapacity(err) { + err = MarkOptionalAllocationError(err) + } + return nil, nil, err } - if err = vec.MarshalBinaryWithBuffer(&buf); err != nil { - token.Release() + if err = plan.MarshalTo(buf); err != nil { + buf.Free() return nil, nil, err } - data := buf.Bytes() - if uint64(len(data)) != wireBytes || uint64(cap(data)) > projected { - token.Release() + if buf.Len() != plan.Size() { + buf.Free() return nil, nil, process.ErrHashBuildBudgetInvalid } - if _, err = token.ReconcileDown(uint64(cap(data))); err != nil { - token.Release() - return nil, nil, err - } - return data, func() { token.Release() }, nil + return buf.Bytes(), buf.Free, nil } diff --git a/pkg/sql/colexec/runtimefilter/contract_test.go b/pkg/sql/colexec/runtimefilter/contract_test.go index b0dc73008731a..ae822b573fc3f 100644 --- a/pkg/sql/colexec/runtimefilter/contract_test.go +++ b/pkg/sql/colexec/runtimefilter/contract_test.go @@ -63,8 +63,6 @@ func TestClassifyOptionalFallbackFatalFirst(t *testing.T) { {name: "marked raw closed", err: marked(process.ErrHashBuildBudgetClosed), want: OptionalFallbackNone}, {name: "marked raw invalid", err: marked(process.ErrHashBuildBudgetInvalid), want: OptionalFallbackNone}, {name: "marked raw ceiling", err: marked(process.ErrHashBuildCeilingMissing), want: OptionalFallbackNone}, - {name: "marked inactive reservation", err: marked(process.ErrHashBuildReservationInactive), want: OptionalFallbackNone}, - {name: "marked upward reconciliation", err: marked(process.ErrHashBuildReservationUpward), want: OptionalFallbackNone}, } for _, test := range tests { @@ -274,7 +272,11 @@ func TestMarshalExactFilterVectorUsesWireSizedBudget(t *testing.T) { aggregate := process.MustNewHashBuildBudget(1<<20, 1<<20) budget, err := aggregate.OpenGeneration(1) require.NoError(t, err) - data, release, err := MarshalExactFilterVector(vec, budget) + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.OpenWithController(budget.Cap(), budget) + require.NoError(t, err) + data, release, err := MarshalExactFilterVector(vec, mp, account, 1, 1) require.NoError(t, err) require.Len(t, data, 34+vec.Length()) // The retained charge is the actual bytes.Buffer capacity, not a @@ -297,7 +299,11 @@ func TestMarshalExactFilterVectorAdmissionFailsBeforeAllocation(t *testing.T) { aggregate := process.MustNewHashBuildBudget(1, 1) budget, err := aggregate.OpenGeneration(1) require.NoError(t, err) - data, release, err := MarshalExactFilterVector(vec, budget) + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.OpenWithController(budget.Cap(), budget) + require.NoError(t, err) + data, release, err := MarshalExactFilterVector(vec, mp, account, 1, 1) require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) require.Nil(t, data) require.Nil(t, release) diff --git a/pkg/sql/colexec/sample/sample.go b/pkg/sql/colexec/sample/sample.go index a46ef5294c69c..6bb509b8eac65 100644 --- a/pkg/sql/colexec/sample/sample.go +++ b/pkg/sql/colexec/sample/sample.go @@ -214,42 +214,204 @@ func (ctr *container) evaluateSampleAndGroupByColumns(proc *process.Process, bat } func (ctr *container) hashAndSample(bat *batch.Batch, proc *process.Process) (err error) { - var iterator hashmap.Iterator - var groupList []uint64 count := bat.RowCount() + if !hasGroupingRows(ctr.groupVectors) { + return ctr.hashNormalRows(bat, proc, 0, count) + } + if err = ctr.enableGroupingDomain(proc); err != nil { + return err + } + groupingIterator := ctr.groupingHashMap.NewIterator() + var normalIterator hashmap.Iterator + + for offset := 0; offset < count; { + grouping := rowHasGrouping(ctr.groupVectors, offset) + end := offset + 1 + for end < count && rowHasGrouping(ctr.groupVectors, end) == grouping { + end++ + } + if grouping { + err = ctr.hashRows( + bat, + groupingIterator, + &ctr.groupingGroupIDs, + offset, + end-offset, + true, + ) + } else { + if normalIterator == nil { + normalIterator, err = ctr.normalIterator(proc) + if err != nil { + return err + } + } + err = ctr.hashRows( + bat, + normalIterator, + &ctr.normalGroupIDs, + offset, + end-offset, + true, + ) + } + if err != nil { + return err + } + offset = end + } + return nil +} + +func hasGroupingRows(vecs []*vector.Vector) bool { + for _, vec := range vecs { + if vec != nil && vec.HasGrouping() { + return true + } + } + return false +} + +func rowHasGrouping(vecs []*vector.Vector, row int) bool { + for _, vec := range vecs { + if vec != nil && vec.GetGrouping().Contains(uint64(row)) { + return true + } + } + return false +} + +func (ctr *container) normalIterator(proc *process.Process) (hashmap.Iterator, error) { + var err error if ctr.useIntHashMap { if ctr.intHashMap == nil { ctr.intHashMap, err = hashmap.NewIntHashMap(ctr.groupVectorsNullable, proc.Mp()) if err != nil { - return err + return nil, err } } - iterator = ctr.intHashMap.NewIterator() + return ctr.intHashMap.NewIterator(), nil } else { if ctr.strHashMap == nil { ctr.strHashMap, err = hashmap.NewStrHashMap(ctr.groupVectorsNullable, proc.Mp()) if err != nil { - return err + return nil, err } } - iterator = ctr.strHashMap.NewIterator() + return ctr.strHashMap.NewIterator(), nil } +} - for i := 0; i < count; i += hashmap.UnitLimit { - n := count - i +func (ctr *container) normalGroupCount() uint64 { + if ctr.useIntHashMap && ctr.intHashMap != nil { + return ctr.intHashMap.GroupCount() + } + if !ctr.useIntHashMap && ctr.strHashMap != nil { + return ctr.strHashMap.GroupCount() + } + return 0 +} + +func (ctr *container) enableGroupingDomain(proc *process.Process) error { + if ctr.groupingHashMap != nil { + return nil + } + groupingMap, err := hashmap.NewStrHashMap( + ctr.groupVectorsNullable, + proc.Mp(), + ) + if err != nil { + return err + } + if err = groupingMap.SetGroupingAware(); err != nil { + groupingMap.Free() + return err + } + + normalGroups := ctr.normalGroupCount() + ctr.normalGroupIDs = make([]uint64, normalGroups+1) + for i := uint64(1); i <= normalGroups; i++ { + ctr.normalGroupIDs[i] = i + } + ctr.groupingGroupIDs = []uint64{0} + ctr.nextGroupID = normalGroups + ctr.groupingHashMap = groupingMap + return nil +} + +func (ctr *container) globalGroupIDs( + local []uint64, + translation *[]uint64, +) []uint64 { + ids := *translation + for i, localID := range local { + if localID == 0 { + continue + } + for uint64(len(ids)) <= localID { + ctr.nextGroupID++ + ids = append(ids, ctr.nextGroupID) + } + local[i] = ids[localID] + } + *translation = ids + return local +} + +func (ctr *container) hashNormalRows( + bat *batch.Batch, + proc *process.Process, + offset int, + count int, +) error { + iterator, err := ctr.normalIterator(proc) + if err != nil { + return err + } + return ctr.hashRows( + bat, + iterator, + &ctr.normalGroupIDs, + offset, + count, + ctr.groupingHashMap != nil, + ) +} + +func (ctr *container) hashRows( + bat *batch.Batch, + iterator hashmap.Iterator, + translation *[]uint64, + offset int, + count int, + translate bool, +) error { + end := offset + count + for offset < end { + n := end - offset if n > hashmap.UnitLimit { n = hashmap.UnitLimit } - groupList, _, err = iterator.Insert(i, n, ctr.groupVectors) + groupList, _, err := iterator.Insert(offset, n, ctr.groupVectors) if err != nil { return err } - err = ctr.samplePool.BatchSample(i, n, groupList, ctr.sampleVectors, ctr.groupVectors, bat) - if err != nil { + if translate { + groupList = ctr.globalGroupIDs(groupList[:n], translation) + } + if err = ctr.samplePool.BatchSample( + offset, + n, + groupList, + ctr.sampleVectors, + ctr.groupVectors, + bat, + ); err != nil { return err } + offset += n } - return + return nil } diff --git a/pkg/sql/colexec/sample/sample_test.go b/pkg/sql/colexec/sample/sample_test.go index b51a96929c6f1..b589b593ce5c3 100644 --- a/pkg/sql/colexec/sample/sample_test.go +++ b/pkg/sql/colexec/sample/sample_test.go @@ -20,6 +20,7 @@ import ( "runtime/debug" "testing" + "github.com/matrixorigin/matrixone/pkg/common/hashmap" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" @@ -263,6 +264,340 @@ func TestSamplePool(t *testing.T) { require.Equal(t, int64(0), proc.Mp().CurrNB()) } +func TestSampleSeparatesGroupingKeyDomain(t *testing.T) { + for _, test := range []struct { + name string + pool func(*process.Process) *sPool + }{ + { + name: "row", + pool: func(proc *process.Process) *sPool { + return newSamplePoolByRows(proc, 1, 1, false) + }, + }, + { + name: "percent", + pool: func(proc *process.Process) *sPool { + return newSamplePoolByPercent(proc, 100, 1) + }, + }, + { + name: "merge", + pool: func(proc *process.Process) *sPool { + return newSamplePoolByRowsForMerge(proc, 1, 1, false) + }, + }, + } { + t.Run(test.name, func(t *testing.T) { + proc := testutil.NewProcess(t) + defer proc.Free() + ctr := &container{ + isGroupBy: true, + useIntHashMap: true, + samplePool: test.pool(proc), + } + defer freeSampleHashContainer(ctr) + + ordinary := makeSampleGroupingBatch( + t, + proc, + []int64{1, 2}, + []int64{10, 20}, + false, + test.name == "merge", + ) + ctr.groupVectors = ordinary.Vecs[:1] + ctr.sampleVectors = ordinary.Vecs[1:2] + require.NoError(t, ctr.hashAndSample(ordinary, proc)) + ordinary.Clean(proc.Mp()) + + rollup := makeSampleGroupingBatch( + t, + proc, + []int64{0}, + []int64{30}, + true, + test.name == "merge", + ) + ctr.groupVectors = rollup.Vecs[:1] + ctr.sampleVectors = rollup.Vecs[1:2] + require.NoError(t, ctr.hashAndSample(rollup, proc)) + rollup.Clean(proc.Mp()) + + result, err := ctr.samplePool.Result(true) + require.NoError(t, err) + require.Equal(t, 3, result.RowCount()) + require.Equal(t, 1, result.Vecs[0].GetGrouping().Count()) + result.Clean(proc.Mp()) + }) + } +} + +func TestSampleSeparatesSQLNullFromGrouping(t *testing.T) { + proc := testutil.NewProcess(t) + defer proc.Free() + ctr := &container{ + isGroupBy: true, + useIntHashMap: true, + groupVectorsNullable: true, + samplePool: newSamplePoolByRows(proc, 1, 1, false), + } + defer freeSampleHashContainer(ctr) + + nullBatch := batch.NewWithSize(2) + nullBatch.Vecs[0] = vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixed( + nullBatch.Vecs[0], + int64(0), + true, + proc.Mp(), + )) + nullBatch.Vecs[1] = vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixed( + nullBatch.Vecs[1], + int64(10), + false, + proc.Mp(), + )) + nullBatch.SetRowCount(1) + ctr.groupVectors = nullBatch.Vecs[:1] + ctr.sampleVectors = nullBatch.Vecs[1:] + require.NoError(t, ctr.hashAndSample(nullBatch, proc)) + nullBatch.Clean(proc.Mp()) + + rollup := makeSampleGroupingBatch( + t, + proc, + []int64{0}, + []int64{20}, + true, + false, + ) + ctr.groupVectors = rollup.Vecs[:1] + ctr.sampleVectors = rollup.Vecs[1:] + require.NoError(t, ctr.hashAndSample(rollup, proc)) + rollup.Clean(proc.Mp()) + + result, err := ctr.samplePool.Result(true) + require.NoError(t, err) + require.Equal(t, 2, result.RowCount()) + require.Equal(t, 2, result.Vecs[0].GetNulls().Count()) + require.Equal(t, 1, result.Vecs[0].GetGrouping().Count()) + result.Clean(proc.Mp()) +} + +func TestSampleSeparatesPartialRollupKeys(t *testing.T) { + proc := testutil.NewProcess(t) + defer proc.Free() + ctr := &container{ + isGroupBy: true, + useIntHashMap: true, + samplePool: newSamplePoolByRows(proc, 1, 1, false), + } + defer freeSampleHashContainer(ctr) + + input := batch.NewWithSize(3) + for column := 0; column < 2; column++ { + input.Vecs[column] = vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixedList( + input.Vecs[column], + []int64{0, 0}, + nil, + proc.Mp(), + )) + } + input.Vecs[0].GetGrouping().Add(0) + input.Vecs[1].GetGrouping().Add(1) + input.Vecs[2] = vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixedList( + input.Vecs[2], + []int64{10, 20}, + nil, + proc.Mp(), + )) + input.SetRowCount(2) + ctr.groupVectors = input.Vecs[:2] + ctr.sampleVectors = input.Vecs[2:] + require.NoError(t, ctr.hashAndSample(input, proc)) + input.Clean(proc.Mp()) + + result, err := ctr.samplePool.Result(true) + require.NoError(t, err) + require.Equal(t, 2, result.RowCount()) + require.Equal(t, 1, result.Vecs[0].GetGrouping().Count()) + require.Equal(t, 1, result.Vecs[1].GetGrouping().Count()) + result.Clean(proc.Mp()) +} + +func TestSampleAlternatingGroupingReusesIterators(t *testing.T) { + proc := testutil.NewProcess(t) + defer proc.Free() + input := makeSampleGroupingBatch( + t, + proc, + makeSequence(hashmap.UnitLimit), + makeSequence(hashmap.UnitLimit), + false, + false, + ) + defer input.Clean(proc.Mp()) + for row := 0; row < input.RowCount(); row += 2 { + input.Vecs[0].GetGrouping().Add(uint64(row)) + } + ctr := &container{ + isGroupBy: true, + useIntHashMap: true, + samplePool: newSamplePoolByRows(proc, 1, 1, false), + groupVectors: input.Vecs[:1], + sampleVectors: input.Vecs[1:], + } + defer freeSampleHashContainer(ctr) + require.NoError(t, ctr.hashAndSample(input, proc)) + + var runErr error + allocations := testing.AllocsPerRun(20, func() { + runErr = ctr.hashAndSample(input, proc) + }) + require.NoError(t, runErr) + require.Less(t, allocations, float64(32)) +} + +func BenchmarkSampleGroupedHashFastPath(b *testing.B) { + proc := testutil.NewProcess(b) + defer proc.Free() + input := makeSampleGroupingBatch( + b, + proc, + makeSequence(256), + makeSequence(256), + false, + false, + ) + defer input.Clean(proc.Mp()) + ctr := &container{ + isGroupBy: true, + useIntHashMap: true, + samplePool: newSamplePoolByRows(proc, 1, 1, false), + groupVectors: input.Vecs[:1], + sampleVectors: input.Vecs[1:], + } + defer freeSampleHashContainer(ctr) + require.NoError(b, ctr.hashAndSample(input, proc)) + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if err := ctr.hashAndSample(input, proc); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkSampleAlternatingGrouping(b *testing.B) { + proc := testutil.NewProcess(b) + defer proc.Free() + input := makeSampleGroupingBatch( + b, + proc, + makeSequence(hashmap.UnitLimit), + makeSequence(hashmap.UnitLimit), + false, + false, + ) + defer input.Clean(proc.Mp()) + for row := 0; row < input.RowCount(); row += 2 { + input.Vecs[0].GetGrouping().Add(uint64(row)) + } + ctr := &container{ + isGroupBy: true, + useIntHashMap: true, + samplePool: newSamplePoolByRows(proc, 1, 1, false), + groupVectors: input.Vecs[:1], + sampleVectors: input.Vecs[1:], + } + defer freeSampleHashContainer(ctr) + require.NoError(b, ctr.hashAndSample(input, proc)) + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if err := ctr.hashAndSample(input, proc); err != nil { + b.Fatal(err) + } + } +} + +func makeSequence(count int) []int64 { + values := make([]int64, count) + for i := range values { + values[i] = int64(i) + } + return values +} + +func makeSampleGroupingBatch( + tb testing.TB, + proc *process.Process, + groups []int64, + samples []int64, + rollup bool, + merge bool, +) *batch.Batch { + tb.Helper() + columns := 2 + if merge { + columns++ + } + bat := batch.NewWithSize(columns) + if rollup { + bat.Vecs[0] = vector.NewRollupConst( + types.T_int64.ToType(), + len(groups), + proc.Mp(), + ) + } else { + bat.Vecs[0] = vector.NewVec(types.T_int64.ToType()) + require.NoError(tb, vector.AppendFixedList( + bat.Vecs[0], + groups, + nil, + proc.Mp(), + )) + } + bat.Vecs[1] = vector.NewVec(types.T_int64.ToType()) + require.NoError(tb, vector.AppendFixedList( + bat.Vecs[1], + samples, + nil, + proc.Mp(), + )) + if merge { + var err error + bat.Vecs[2], err = vector.NewConstFixed( + types.T_int64.ToType(), + int64(len(groups)), + len(groups), + proc.Mp(), + ) + require.NoError(tb, err) + } + bat.SetRowCount(len(groups)) + return bat +} + +func freeSampleHashContainer(ctr *container) { + if ctr.intHashMap != nil { + ctr.intHashMap.Free() + } + if ctr.strHashMap != nil { + ctr.strHashMap.Free() + } + if ctr.groupingHashMap != nil { + ctr.groupingHashMap.Free() + } + ctr.samplePool.Free() +} + func genSampleBatch(proc *process.Process, rows [][]int64) (*batch.Batch, error) { b := batch.NewWithSize(len(rows[0])) diff --git a/pkg/sql/colexec/sample/types.go b/pkg/sql/colexec/sample/types.go index e10009533afee..95c3b4c0944bd 100644 --- a/pkg/sql/colexec/sample/types.go +++ b/pkg/sql/colexec/sample/types.go @@ -83,8 +83,16 @@ type container struct { buf *batch.Batch // hash map related. - intHashMap *hashmap.IntHashMap - strHashMap *hashmap.StrHashMap + intHashMap *hashmap.IntHashMap + strHashMap *hashmap.StrHashMap + groupingHashMap *hashmap.StrHashMap + + // A grouping-aware key has a domain that cannot be represented by the + // normal IntHashMap and cannot be installed into a populated StrHashMap. + // These tables translate each map's local IDs into one sample-pool domain. + normalGroupIDs []uint64 + groupingGroupIDs []uint64 + nextGroupID uint64 } func init() { @@ -209,6 +217,9 @@ func (sample *Sample) Free(proc *process.Process, pipelineFailed bool, err error if sample.ctr.strHashMap != nil { sample.ctr.strHashMap.Free() } + if sample.ctr.groupingHashMap != nil { + sample.ctr.groupingHashMap.Free() + } for _, executor := range sample.ctr.sampleExecutors { if executor != nil { executor.Free() diff --git a/pkg/sql/colexec/shuffle/shufflepool_test.go b/pkg/sql/colexec/shuffle/shufflepool_test.go index 60db6768d0448..a8d2f0acb1200 100644 --- a/pkg/sql/colexec/shuffle/shufflepool_test.go +++ b/pkg/sql/colexec/shuffle/shufflepool_test.go @@ -190,12 +190,12 @@ func TestShufflePoolReservesReadyCreditForProvenanceChange(t *testing.T) { require.NoError(t, err) account, err := registry.Open(1 << 20) require.NoError(t, err) - selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) require.NoError(t, err) sp := NewShufflePool(1, 1, true) - legacy := testutil.NewBatch([]types.Type{types.T_int64.ToType()}, false, 2, mp) - done, err := writeBatchToBucketForTest(sp, legacy, proc, 0) + unaccounted := testutil.NewBatch([]types.Type{types.T_int64.ToType()}, false, 2, mp) + done, err := writeBatchToBucketForTest(sp, unaccounted, proc, 0) require.NoError(t, err) require.True(t, done) require.Zero(t, sp.readyCount) @@ -224,7 +224,7 @@ func TestShufflePoolReservesReadyCreditForProvenanceChange(t *testing.T) { require.Same(t, selection, tail.AllocationAccountSelection()) sp.discardBatch(tail, mp) - legacy.Clean(mp) + unaccounted.Clean(mp) accounted.Clean(mp) sp.abort(mp) require.Zero(t, account.Seal().Used) diff --git a/pkg/sql/colexec/spillutil/allocation_account.go b/pkg/sql/colexec/spillutil/allocation_account.go index 451c51bb9ae82..d860a197ca62a 100644 --- a/pkg/sql/colexec/spillutil/allocation_account.go +++ b/pkg/sql/colexec/spillutil/allocation_account.go @@ -21,11 +21,9 @@ import ( "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" - "github.com/matrixorigin/matrixone/pkg/sql/colexec" ) -// Spill allocation sites use a range disjoint from colexec expression sites -// when both subsystems share one logical owner. +// Spill allocation sites occupy a dedicated range within the HashBuild owner. const ( SpillAllocationSiteDecodedData mpool.AllocationSite = iota + 32 SpillAllocationSiteDecodedArea @@ -47,16 +45,15 @@ type SpillAllocationAccount struct { account *mpool.AllocationAccount owner mpool.AllocationOwner - decoded *vector.AllocationAccountSelection - selected *vector.AllocationAccountSelection - expression *colexec.ExpressionAllocationAccount + decoded *vector.AllocationAccountSelection + selected *vector.AllocationAccountSelection } func NewSpillAllocationAccount( account *mpool.AllocationAccount, owner mpool.AllocationOwner, ) (*SpillAllocationAccount, error) { - decoded, err := vector.NewAllocationAccountSelectionWithBitmaps( + decoded, err := vector.NewAllocationAccountSelection( account, owner, SpillAllocationSiteDecodedData, @@ -67,7 +64,7 @@ func NewSpillAllocationAccount( if err != nil { return nil, err } - selected, err := vector.NewAllocationAccountSelectionWithBitmaps( + selected, err := vector.NewAllocationAccountSelection( account, owner, SpillAllocationSiteSelectedData, @@ -78,16 +75,11 @@ func NewSpillAllocationAccount( if err != nil { return nil, err } - expression, err := colexec.NewExpressionAllocationAccount(account, owner) - if err != nil { - return nil, err - } return &SpillAllocationAccount{ - account: account, - owner: owner, - decoded: decoded, - selected: selected, - expression: expression, + account: account, + owner: owner, + decoded: decoded, + selected: selected, }, nil } @@ -95,7 +87,7 @@ func (a *SpillAllocationAccount) validate() error { if a == nil || a.account == nil || a.account.Handle() == 0 || a.owner < mpool.AllocationOwnerMin || a.owner > mpool.AllocationOwnerMax || - a.decoded == nil || a.selected == nil || a.expression == nil { + a.decoded == nil || a.selected == nil { return mpool.ErrAllocationAccountInvalid } return nil @@ -105,12 +97,13 @@ func newSpillBatch( size int, selection *vector.AllocationAccountSelection, ) (*batch.Batch, error) { + if selection == nil { + return nil, mpool.ErrAllocationAccountInvalid + } bat := batch.NewOffHeapWithSize(size) - if selection != nil { - if err := bat.SetAllocationAccount(selection); err != nil { - bat.Clean(nil) - return nil, err - } + if err := bat.SetAllocationAccount(selection); err != nil { + bat.Clean(nil) + return nil, err } return bat, nil } @@ -120,7 +113,7 @@ func newSpillVector( selection *vector.AllocationAccountSelection, ) (*vector.Vector, error) { if selection == nil { - return vector.NewOffHeapVecWithType(typ), nil + return nil, mpool.ErrAllocationAccountInvalid } return vector.NewOffHeapVecWithTypeAndAllocation(typ, selection) } @@ -138,9 +131,6 @@ func growSpillSlice[T any]( if length <= cap(values) { return values[:length], nil } - if allocation == nil { - return make([]T, length), nil - } if err := allocation.validate(); err != nil { return nil, err } @@ -175,9 +165,8 @@ func growSpillSlice[T any]( func freeSpillSlice[T any]( values []T, mp *mpool.MPool, - allocation *SpillAllocationAccount, ) { - if allocation != nil && cap(values) > 0 { + if cap(values) > 0 { mpool.FreeSlice(mp, values) } } diff --git a/pkg/sql/colexec/spillutil/allocation_account_test.go b/pkg/sql/colexec/spillutil/allocation_account_test.go index 8ededc0917478..37c9fd95eafc5 100644 --- a/pkg/sql/colexec/spillutil/allocation_account_test.go +++ b/pkg/sql/colexec/spillutil/allocation_account_test.go @@ -15,7 +15,6 @@ package spillutil import ( - "bytes" "os" "path/filepath" "strings" @@ -35,6 +34,20 @@ type testSpillAllocationAccount struct { registry *mpool.AllocationAccountRegistry account *mpool.AllocationAccount allocation *SpillAllocationAccount + generation *process.HashBuildBudgetGeneration +} + +func TestNewSpillEngineRequiresBudgetGeneration(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + _, err = NewSpillEngine( + SpillEngineConfig{}, + account, + hashbuild.HashBuildAllocationOwner, + ) + require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) } func newTestSpillAllocationAccount( @@ -45,7 +58,10 @@ func newTestSpillAllocationAccount( t.Helper() registry, err := mpool.NewAllocationAccountRegistry(1, metadataSlots) require.NoError(t, err) - account, err := registry.Open(limit) + budget := process.MustNewHashBuildBudget(limit, limit) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) require.NoError(t, err) allocation, err := NewSpillAllocationAccount(account, 2) require.NoError(t, err) @@ -53,6 +69,7 @@ func newTestSpillAllocationAccount( registry: registry, account: account, allocation: allocation, + generation: generation, } } @@ -74,9 +91,7 @@ func writeSpillAllocationTestFile( truncate int, ) *os.File { t.Helper() - var encoded bytes.Buffer - require.NoError(t, marshalSpillRecord(bat, &encoded)) - payload := encoded.Bytes() + payload := marshalTestSpillRecord(bat) if truncate > 0 { payload = payload[:len(payload)-truncate] } @@ -92,15 +107,12 @@ func writeSpillAllocationTestRecords( batches ...*batch.Batch, ) *os.File { t.Helper() - var payload bytes.Buffer + payload := make([]byte, 0) for _, bat := range batches { - var encoded bytes.Buffer - require.NoError(t, marshalSpillRecord(bat, &encoded)) - _, err := payload.Write(encoded.Bytes()) - require.NoError(t, err) + payload = append(payload, marshalTestSpillRecord(bat)...) } path := filepath.Join(t.TempDir(), "spill-records.bin") - require.NoError(t, os.WriteFile(path, payload.Bytes(), 0o600)) + require.NoError(t, os.WriteFile(path, payload, 0o600)) file, err := os.Open(path) require.NoError(t, err) return file @@ -235,23 +247,24 @@ func TestSpillAllocationAccountDecodedReuseRetriesFromCleanRecord(t *testing.T) require.NoError(t, err) allocation, err := NewSpillAllocationAccount(account, 2) require.NoError(t, err) - reader = BucketReader{allocation: allocation} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForFd(writeSpillAllocationTestRecords(t, first, second)) + reader = BucketReader{ + fd: writeSpillAllocationTestRecords(t, first, second), + allocation: allocation, + } reuse = batch.NewOffHeapWithSize(0) _, err = reader.ReadBatch(proc, reuse) require.NoError(t, err) - rejects := generation.RejectCount() + rejects := generation.Snapshot().RejectCount _, err = reader.ReadBatch(proc, reuse) require.NoError(t, err) - require.Equal(t, rejects, generation.RejectCount(), + require.Equal(t, rejects, generation.Snapshot().RejectCount, "the local account rejects the overlap before the shared controller") require.Equal(t, uint64(1), reader.cleanRetries, "replacement overlap must exercise the clean-record retry") require.Equal(t, - generation.Snapshot().AllocationUsed+uint64(64<<10), + account.Snapshot().Used, generation.Used(), - "decoded payloads have no duplicate hard reservation", + "decoded payloads are charged only by their physical allocations", ) reuse.Clean(proc.Mp()) reader.Close() @@ -269,8 +282,8 @@ func TestSpillAllocationAccountScatterScratchLifecycle(t *testing.T) { ) defer proc.Free() state := newTestSpillAllocationAccount(t, 1<<20, 64) - engine, err := NewSpillEngineWithAllocation( - SpillEngineConfig{}, + engine, err := newSpillEngine( + SpillEngineConfig{Budget: state.generation}, state.allocation, ) require.NoError(t, err) @@ -285,7 +298,7 @@ func TestSpillAllocationAccountScatterScratchLifecycle(t *testing.T) { ), }, nil) defer source.Clean(proc.Mp()) - writers := MakeBucketWriters("spill_allocation_scatter") + writers := engine.makeBucketWriters("spill_allocation_scatter") defer func() { for i := range writers { writers[i].Close() @@ -341,7 +354,7 @@ func TestSpillAllocationAccountScatterDoesNotReadmitBorrowedSource(t *testing.T) require.NoError(t, err) allocation, err := NewSpillAllocationAccount(account, 2) require.NoError(t, err) - engine, err := NewSpillEngineWithAllocation( + engine, err := newSpillEngine( SpillEngineConfig{Budget: generation}, allocation, ) @@ -356,7 +369,7 @@ func TestSpillAllocationAccountScatterDoesNotReadmitBorrowedSource(t *testing.T) ), }, nil) defer source.Clean(proc.Mp()) - writers := MakeBucketWriters("spill_allocation_scatter_source") + writers := engine.makeBucketWriters("spill_allocation_scatter_source") defer func() { for i := range writers { writers[i].Close() @@ -372,8 +385,7 @@ func TestSpillAllocationAccountScatterDoesNotReadmitBorrowedSource(t *testing.T) process.NewAnalyzer(0, false, false, "test"), )) snapshot := generation.Snapshot() - require.Nil(t, engine.scatterScratchReservation) - require.Equal(t, snapshot.AllocationUsed, snapshot.Used, + require.Equal(t, account.Snapshot().Used, snapshot.Used, "borrowed input is already live; only new private spill bytes are admitted") require.Greater(t, snapshot.PeakUsed, snapshot.Used) @@ -410,9 +422,7 @@ func TestSpillAllocationAccountMarshalBufferLifecycle(t *testing.T) { ) require.NoError(t, err) require.NoError(t, marshalSpillRecordTo(source, accounted)) - var legacy bytes.Buffer - require.NoError(t, marshalSpillRecord(source, &legacy)) - require.Equal(t, legacy.Bytes(), accounted.Bytes()) + require.Equal(t, marshalTestSpillRecord(source), accounted.Bytes()) used := state.account.Snapshot().Used require.Positive(t, used) @@ -438,8 +448,8 @@ func TestSpillAllocationAccountCoalesceAdmissionFallback(t *testing.T) { // optional optimization, so its admission failure must fall back to one // direct write instead of failing the scatter. state := newTestSpillAllocationAccount(t, 1<<20, 1) - engine, err := NewSpillEngineWithAllocation( - SpillEngineConfig{}, + engine, err := newSpillEngine( + SpillEngineConfig{Budget: state.generation}, state.allocation, ) require.NoError(t, err) @@ -453,7 +463,7 @@ func TestSpillAllocationAccountCoalesceAdmissionFallback(t *testing.T) { ), }, nil) defer source.Clean(proc.Mp()) - writers := MakeBucketWriters("spill_allocation_coalesce_fallback") + writers := engine.makeBucketWriters("spill_allocation_coalesce_fallback") defer func() { for i := range writers { writers[i].Close() @@ -485,8 +495,8 @@ func TestSpillAllocationAccountScatterFailureCleanup(t *testing.T) { defer proc.Free() const rows = 8 state := newTestSpillAllocationAccount(t, rows*(8+4), 8) - engine, err := NewSpillEngineWithAllocation( - SpillEngineConfig{}, + engine, err := newSpillEngine( + SpillEngineConfig{Budget: state.generation}, state.allocation, ) require.NoError(t, err) @@ -501,7 +511,7 @@ func TestSpillAllocationAccountScatterFailureCleanup(t *testing.T) { ), }, nil) defer source.Clean(proc.Mp()) - writers := MakeBucketWriters("spill_allocation_scatter_failure") + writers := engine.makeBucketWriters("spill_allocation_scatter_failure") defer func() { for i := range writers { writers[i].Close() @@ -537,8 +547,8 @@ func TestSpillAllocationAccountScatterReducesUnpublishedInput(t *testing.T) { ) defer proc.Free() state := newTestSpillAllocationAccount(t, 80<<10, 128) - engine, err := NewSpillEngineWithAllocation( - SpillEngineConfig{}, + engine, err := newSpillEngine( + SpillEngineConfig{Budget: state.generation}, state.allocation, ) require.NoError(t, err) @@ -550,7 +560,7 @@ func TestSpillAllocationAccountScatterReducesUnpublishedInput(t *testing.T) { testutil.MakeInt64Vector(values, nil, proc.Mp()), }, nil) defer source.Clean(proc.Mp()) - writers := MakeBucketWriters("spill_allocation_scatter_reduce") + writers := engine.makeBucketWriters("spill_allocation_scatter_reduce") defer func() { for i := range writers { writers[i].Close() @@ -589,8 +599,8 @@ func TestSpillAllocationAccountExpressionPressureReducesBeforePublication(t *tes ) defer proc.Free() state := newTestSpillAllocationAccount(t, 2<<20, 4_096) - engine, err := NewSpillEngineWithAllocation( - SpillEngineConfig{}, + engine, err := newSpillEngine( + SpillEngineConfig{Budget: state.generation}, state.allocation, ) require.NoError(t, err) @@ -602,7 +612,7 @@ func TestSpillAllocationAccountExpressionPressureReducesBeforePublication(t *tes testutil.MakeInt64Vector(values, nil, proc.Mp()), }, nil) defer source.Clean(proc.Mp()) - writers := MakeBucketWriters("spill_allocation_expression_reduce") + writers := engine.makeBucketWriters("spill_allocation_expression_reduce") defer func() { for i := range writers { writers[i].Close() @@ -645,92 +655,75 @@ func TestSpillAllocationAccountExpressionPressureReducesBeforePublication(t *tes // retains the write syscall without turning repeated measurements into a disk // capacity test. func BenchmarkSpillScatterAccounting(b *testing.B) { - for _, accounted := range []bool{false, true} { - mode := "legacy" - if accounted { - mode = "accounted" + proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) + defer proc.Free() + values := make([]int64, 4_096) + for i := range values { + values[i] = int64(i) + } + source := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.MakeInt64Vector(values, nil, proc.Mp()), + }, nil) + defer source.Clean(proc.Mp()) + + state := newTestSpillAllocationAccount(b, 64<<20, 4_096) + engine, err := newSpillEngine( + SpillEngineConfig{Budget: state.generation}, + state.allocation, + ) + if err != nil { + b.Fatal(err) + } + writers := engine.makeBucketWriters("benchmark-discard") + for i := range writers { + writers[i].Fd, err = os.OpenFile(os.DevNull, os.O_WRONLY, 0) + if err != nil { + b.Fatal(err) } - b.Run(mode, func(b *testing.B) { - proc := testutil.NewProcessWithMPool( - b, - "", - mpool.MustNewZero(), - ) - defer proc.Free() - values := make([]int64, 4_096) - for i := range values { - values[i] = int64(i) - } - source := testutil.NewBatchWithVectors([]*vector.Vector{ - testutil.MakeInt64Vector(values, nil, proc.Mp()), - }, nil) - defer source.Clean(proc.Mp()) - - var ( - engine *SpillEngine - state testSpillAllocationAccount - err error - ) - if accounted { - state = newTestSpillAllocationAccount(b, 64<<20, 4_096) - engine, err = NewSpillEngineWithAllocation( - SpillEngineConfig{}, - state.allocation, - ) - if err != nil { - b.Fatal(err) - } - } else { - engine = NewSpillEngine(SpillEngineConfig{}) - } - writers := make([]BucketWriter, SpillNumBuckets) - for i := range writers { - writers[i].Name = "benchmark-discard" - writers[i].Fd, err = os.OpenFile(os.DevNull, os.O_WRONLY, 0) - if err != nil { - b.Fatal(err) - } - } - defer func() { - for i := range writers { - writers[i].Close() - } - }() - analyzer := process.NewAnalyzer(0, false, false, "benchmark") - - b.ReportAllocs() - b.SetBytes(int64(source.Size())) - b.ResetTimer() - for range b.N { - if err = engine.scatterBatchWithPressure( - proc, - source, - source.Vecs, - writers, - 0, - false, - analyzer, - ); err != nil { - b.Fatal(err) - } - if err = engine.flushScatterBuffers(proc, writers, analyzer); err != nil { - b.Fatal(err) - } - for i := range writers { - writers[i].Rows = 0 - writers[i].Bytes = 0 - } - } - b.StopTimer() - engine.Cleanup(proc) - if accounted { - if state.account.Snapshot().Used != 0 { - b.Fatalf("account used = %d", state.account.Snapshot().Used) - } - finalizeTestSpillAllocationAccount(b, state) + writers[i].diskReservation, err = state.generation.ReserveSpillDisk(0) + if err != nil { + b.Fatal(err) + } + } + defer func() { + for i := range writers { + writers[i].Close() + } + }() + analyzer := process.NewAnalyzer(0, false, false, "benchmark") + + b.ReportAllocs() + b.SetBytes(int64(source.Size())) + b.ResetTimer() + for range b.N { + if err = engine.scatterBatchWithPressure( + proc, + source, + source.Vecs, + writers, + 0, + false, + analyzer, + ); err != nil { + b.Fatal(err) + } + if err = engine.flushScatterBuffers(proc, writers, analyzer); err != nil { + b.Fatal(err) + } + for i := range writers { + if _, err = writers[i].diskReservation.ReconcileDown(0); err != nil { + b.Fatal(err) } - }) + writers[i].Rows = 0 + writers[i].Bytes = 0 + } + } + b.StopTimer() + engine.Cleanup(proc) + if state.account.Snapshot().Used != 0 { + b.Fatalf("account used = %d", state.account.Snapshot().Used) } + finalizeTestSpillAllocationAccount(b, state) } func TestSpillAllocationAccountRebuildAndRecursiveSpillLifecycle(t *testing.T) { @@ -753,7 +746,7 @@ func TestSpillAllocationAccountRebuildAndRecursiveSpillLifecycle(t *testing.T) { require.NoError(t, err) account, err := registry.OpenWithController(limit, generation) require.NoError(t, err) - engine, err := NewSpillEngineForAccount( + engine, err := NewSpillEngine( SpillEngineConfig{ BuildKeyExprs: makeTestKeyExpr(), Budget: generation, @@ -772,7 +765,7 @@ func TestSpillAllocationAccountRebuildAndRecursiveSpillLifecycle(t *testing.T) { source := makeInt32Batch(proc, values) fd := writeBuildFile(proc, "accounted_recursive_build", source) source.Clean(proc.Mp()) - engine.InitFromSpilledMap([]*os.File{fd}) + initTestSpillFiles(engine, []*os.File{fd}, int64(len(values))) analyzer := process.NewAnalyzer(0, false, false, "test") respills := 0 diff --git a/pkg/sql/colexec/spillutil/exact_test_helpers_test.go b/pkg/sql/colexec/spillutil/exact_test_helpers_test.go new file mode 100644 index 0000000000000..9a875dd28bdea --- /dev/null +++ b/pkg/sql/colexec/spillutil/exact_test_helpers_test.go @@ -0,0 +1,91 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package spillutil + +import ( + "bytes" + "os" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" + "github.com/matrixorigin/matrixone/pkg/vm/message" + "github.com/matrixorigin/matrixone/pkg/vm/process" + "github.com/stretchr/testify/require" +) + +func newExactTestSpillEngine( + t testing.TB, + cfg SpillEngineConfig, +) *SpillEngine { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, 1<<20) + require.NoError(t, err) + if cfg.Budget == nil { + budget := process.MustNewHashBuildBudget(1<<60, 1<<60) + cfg.Budget, err = budget.OpenGeneration(1) + require.NoError(t, err) + } + account, err := registry.OpenWithController(1<<60, cfg.Budget) + require.NoError(t, err) + engine, err := NewSpillEngine( + cfg, + account, + hashbuild.HashBuildAllocationOwner, + ) + require.NoError(t, err) + return engine +} + +func initTestSpillFiles(engine *SpillEngine, fds []*os.File, rows ...int64) { + if len(rows) != len(fds) { + panic("spill test file/row metadata mismatch") + } + files := make([]*message.SpillFile, len(fds)) + for i, fd := range fds { + if fd != nil { + files[i] = newTestSpillFile(fd, rows[i]) + } + } + engine.InitFromSpilledFiles(files) +} + +func newTestSpillFile(fd *os.File, rows int64) *message.SpillFile { + info, err := fd.Stat() + if err != nil { + panic(err) + } + return message.NewSpillFile(fd, rows, uint64(info.Size()), nil) +} + +type testSpillRecordBuffer struct { + bytes.Buffer +} + +func (b *testSpillRecordBuffer) EnsureCapacity(required int) error { + if b.Cap() < required { + b.Buffer = *bytes.NewBuffer(make([]byte, 0, required)) + } + return nil +} + +func marshalTestSpillRecord(bat *batch.Batch) []byte { + var encoded testSpillRecordBuffer + if err := marshalSpillRecordTo(bat, &encoded); err != nil { + panic(err) + } + return bytes.Clone(encoded.Bytes()) +} diff --git a/pkg/sql/colexec/spillutil/join_spill.go b/pkg/sql/colexec/spillutil/join_spill.go index 3d5532729360f..e13613d10e692 100644 --- a/pkg/sql/colexec/spillutil/join_spill.go +++ b/pkg/sql/colexec/spillutil/join_spill.go @@ -16,8 +16,7 @@ package spillutil import ( - "bufio" - "bytes" + "encoding/binary" "errors" "fmt" "io" @@ -53,7 +52,6 @@ const ( // spill records without retaining the pre-admission unmarshal estimate // for a large record until the reader closes. The additive bound makes the // long-lived charge independent of the largest serialized payload seen. - decodedBatchLeaseSlack = 1 << 20 ) // SpillBucket holds file descriptors for one spilled bucket. @@ -77,22 +75,24 @@ func checkSpillCanceled(proc *process.Process) error { } } -// BucketReader reads serialized batch records from an fd. +// BucketReader decodes one move-only spill file. A pending header replaces +// bufio.Peek, so all data-scaled decode storage is owned by accounted vectors +// instead of an untracked Go-heap buffer. type BucketReader struct { - fd *os.File - reader *bufio.Reader - buf [16]byte - budget *process.HashBuildBudgetGeneration - reservation *process.HashBuildReservation - batchToken *process.HashBuildReservation - batchCharge uint64 - spillFile *message.SpillFile - mergeRecords bool - allocation *SpillAllocationAccount - cleanRetries uint64 + fd *os.File + header [16]byte + headerPending bool + spillFile *message.SpillFile + mergeRecords bool + allocation *SpillAllocationAccount + cleanRetries uint64 + schema []types.Type } -func (r *BucketReader) ReadBatch(proc *process.Process, reuseBat *batch.Batch) (*batch.Batch, error) { +func (r *BucketReader) ReadBatch( + proc *process.Process, + reuseBat *batch.Batch, +) (*batch.Batch, error) { if err := checkSpillCanceled(proc); err != nil { return nil, err } @@ -105,640 +105,215 @@ func (r *BucketReader) ReadBatch(proc *process.Process, reuseBat *batch.Batch) ( "spill batch reader requires a reuse batch", ) } - if r.allocation != nil { - if err := reuseBat.SetAllocationAccount( - r.allocation.decoded, - ); err != nil { - return nil, err - } + if r.allocation == nil { + return nil, mpool.ErrAllocationAccountInvalid } - if r.reader == nil { - r.reader = bufio.NewReaderSize(r.fd, 4*1024*1024) + if err := reuseBat.SetAllocationAccount(r.allocation.decoded); err != nil { + return nil, err } - _, token, charge, err := r.readBatchRecord(proc, reuseBat, r.batchToken, r.batchCharge, true) - if err != nil { - r.releaseReadBatch(proc, reuseBat, token) + if _, err := r.readBatchRecord(proc, reuseBat); err != nil { + reuseBat.Clean(proc.Mp()) return nil, err } - oldToken := r.batchToken - r.batchToken = token - r.batchCharge = charge - if oldToken != nil && oldToken != token { - oldToken.Release() + if err := r.validateSchema(proc, reuseBat); err != nil { + reuseBat.Clean(proc.Mp()) + return nil, err } if !r.mergeRecords { if err := checkSpillCanceled(proc); err != nil { - r.releaseReadBatch(proc, reuseBat, nil) + reuseBat.Clean(proc.Mp()) return nil, err } return reuseBat, nil } - // Merge adjacent records up to the bounded historical batch payload. - // This preserves dedup/outer-join behaviour across small source batches - // without retaining one selected batch per bucket during scatter. + for reuseBat.RowCount() < colexec.DefaultBatchSize { - if err := checkSpillCanceled(proc); err != nil { - return nil, r.mergeReadError(proc, reuseBat, nil, nil, err) + nextRows, err := r.peekRecordRows(proc) + if err == io.EOF { + break } - header, err := r.reader.Peek(16) if err != nil { - if err == io.EOF && len(header) == 0 { - break - } - if err == io.EOF { - err = io.ErrUnexpectedEOF - } - return nil, r.mergeReadError(proc, reuseBat, nil, nil, err) - } - if err := checkSpillCanceled(proc); err != nil { - return nil, r.mergeReadError(proc, reuseBat, nil, nil, err) + return nil, r.mergeReadError(proc, reuseBat, nil, err) } - nextRows := types.DecodeInt64(header[:8]) - nextBatchSize := types.DecodeInt64(header[8:16]) - if nextRows < 0 || nextBatchSize < 0 { - return nil, r.mergeReadError( - proc, - reuseBat, - nil, - nil, - moerr.NewInternalError(proc.Ctx, "negative spill batch header"), - ) - } - // A source record is an indivisible ownership and budget unit. Leave it - // for the next ReadBatch rather than consuming it and growing this batch - // beyond the advertised merge bound. if nextRows > int64(colexec.DefaultBatchSize-reuseBat.RowCount()) { break } - var selection *vector.AllocationAccountSelection - if r.allocation != nil { - selection = r.allocation.decoded - } - next, err := newSpillBatch(0, selection) - if err != nil { - return nil, r.mergeReadError( - proc, - reuseBat, - nil, - nil, - err, - ) - } - _, nextToken, _, err := r.readBatchRecord(proc, next, nil, 0, false) + next, err := newSpillBatch(0, r.allocation.decoded) if err != nil { - return nil, r.mergeReadError(proc, reuseBat, next, nextToken, err) + return nil, r.mergeReadError(proc, reuseBat, nil, err) } - if err := checkSpillCanceled(proc); err != nil { - return nil, r.mergeReadError(proc, reuseBat, next, nextToken, err) - } - var mergeToken *process.HashBuildReservation - if r.budget != nil && r.allocation == nil { - // Keep the current destination (O) and the source record (N) live - // while admitting the final destination (D). UnionBatch may retain - // rounded capacities larger than O+N, so reserving O+N here is not a - // safe admission bound. - predicted, ok := predictMergedRetainedBytes(reuseBat, next) - if !ok { - return nil, r.mergeReadError(proc, reuseBat, next, nextToken, process.ErrHashBuildBudgetInvalid) - } - mergeToken, err = r.budget.Reserve(predicted) - if err != nil { - return nil, r.mergeReadError(proc, reuseBat, next, nextToken, err) - } + if _, err := r.readBatchRecord(proc, next); err != nil { + return nil, r.mergeReadError(proc, reuseBat, next, err) } - if len(reuseBat.Vecs) != len(next.Vecs) { - return nil, r.mergeReadError(proc, reuseBat, next, nextToken, process.ErrHashBuildBudgetInvalid, mergeToken) + if err := r.validateSchema(proc, next); err != nil { + return nil, r.mergeReadError(proc, reuseBat, next, err) } for i := range next.Vecs { - if err := reuseBat.Vecs[i].UnionBatch(next.Vecs[i], 0, next.RowCount(), nil, proc.Mp()); err != nil { - return nil, r.mergeReadError(proc, reuseBat, next, nextToken, err, mergeToken) + if err := reuseBat.Vecs[i].UnionBatch( + next.Vecs[i], + 0, + next.RowCount(), + nil, + proc.Mp(), + ); err != nil { + return nil, r.mergeReadError(proc, reuseBat, next, err) } } - reuseBat.SetRowCount(reuseBat.RowCount() + next.RowCount()) + reuseBat.AddRowCount(next.RowCount()) next.Clean(proc.Mp()) - if mergeToken != nil { - actual, ok := batchRetainedBytes(reuseBat) - if !ok { - return nil, r.mergeReadError(proc, reuseBat, nil, nextToken, process.ErrHashBuildBudgetInvalid, mergeToken) - } - if err := reconcileReadReservation(mergeToken, actual); err != nil { - return nil, r.mergeReadError(proc, reuseBat, nil, nextToken, err, mergeToken) - } - if r.batchToken != nil { - r.batchToken.Release() - r.batchToken = nil - r.batchCharge = 0 - } - if nextToken != nil { - nextToken.Release() - } - r.batchToken = mergeToken - r.batchCharge = actual - } } if err := checkSpillCanceled(proc); err != nil { - return nil, r.mergeReadError(proc, reuseBat, nil, nil, err) + return nil, r.mergeReadError(proc, reuseBat, nil, err) } return reuseBat, nil } -// mergeReadError unwinds all ownership acquired while appending a source -// record. The destination may have been partially mutated by UnionBatch, so it -// is cleaned as well. Reservations are exactly-once tokens; releasing an -// already released token is harmless and keeps every error path symmetric. -func (r *BucketReader) mergeReadError(proc *process.Process, dst, src *batch.Batch, srcToken *process.HashBuildReservation, err error, extra ...*process.HashBuildReservation) error { - if src != nil { - src.Clean(proc.Mp()) - } - if srcToken != nil { - srcToken.Release() +func (r *BucketReader) validateSchema( + proc *process.Process, + bat *batch.Batch, +) error { + if bat == nil { + return process.ErrHashBuildBudgetInvalid } - for _, token := range extra { - if token != nil { - token.Release() + if r.schema == nil { + r.schema = make([]types.Type, len(bat.Vecs)) + for i, vec := range bat.Vecs { + if vec == nil { + return moerr.NewInternalError(proc.Ctx, "nil vector in spill batch") + } + r.schema[i] = *vec.GetType() } - } - if dst != nil { - dst.Clean(proc.Mp()) - } - if r.batchToken != nil { - r.batchToken.Release() - r.batchToken = nil - r.batchCharge = 0 - } - return err -} - -func addUint64(a, b uint64) (uint64, bool) { - if a > math.MaxUint64-b { - return 0, false - } - return a + b, true -} - -func batchRetainedBytes(bat *batch.Batch) (uint64, bool) { - if bat == nil || bat.RowCount() < 0 { - return 0, false - } - actual := uint64(bat.Allocated()) - metadata, ok := batchRetainedMetadataBytes(uint64(bat.RowCount()), uint64(len(bat.Vecs))) - if !ok { - return 0, false - } - return addUint64(actual, metadata) -} - -func batchRetainedMetadataBytes(rows, cols uint64) (uint64, bool) { - if cols > (math.MaxUint64-16)/8 { - return 0, false - } - metadata := uint64(16) + cols*8 - if rows > 0 && metadata > math.MaxUint64/rows { - return 0, false - } - return rows * metadata, true -} - -// reconcileReadReservation shrinks a conservative read reservation to the -// retained batch size. ReconcileDown already validates that actual does not -// exceed the reservation, so callers do not need a separate Size call (and a -// second acquisition of the shared hash-build budget mutex). -func reconcileReadReservation(token *process.HashBuildReservation, actual uint64) error { - if token == nil { return nil } - if _, err := token.ReconcileDown(actual); err != nil { - if errors.Is(err, process.ErrHashBuildReservationUpward) { - return process.ErrHashBuildBudgetInvalid - } - return err - } - return nil -} - -// predictMergedRetainedBytes computes the retained upper bound after the exact -// full-record UnionBatch append used by ReadBatch. It mirrors Vector.extend's -// data-cap growth and UnionBatch's varlen fast path (which appends a complete -// non-const source area in one operation). No destination mutation is performed. -func predictMergedRetainedBytes(dst, src *batch.Batch) (uint64, bool) { - if dst == nil || src == nil || dst.RowCount() < 0 || src.RowCount() < 0 || len(dst.Vecs) != len(src.Vecs) { - return 0, false - } - oldRows, ok := intToUint64(dst.RowCount()) - if !ok { - return 0, false + if len(bat.Vecs) != len(r.schema) { + return moerr.NewInternalError(proc.Ctx, "spill batch schema changed") } - srcRows, ok := intToUint64(src.RowCount()) - if !ok { - return 0, false - } - mergedRows, ok := addUint64(oldRows, srcRows) - if !ok || mergedRows > uint64(maxIntValue()) { - return 0, false - } - - var allocated uint64 - for i := range dst.Vecs { - dv, sv := dst.Vecs[i], src.Vecs[i] - if dv == nil || sv == nil || *dv.GetType() != *sv.GetType() || dv.Length() != dst.RowCount() || sv.Length() != src.RowCount() { - return 0, false - } - typeSize := dv.GetType().TypeSize() - if typeSize < 0 { - return 0, false - } - dataRequired, ok := mulUint64(mergedRows, uint64(typeSize)) - if !ok || dataRequired > uint64(math.MaxInt64) { - return 0, false - } - dataCap, ok := predictedCapacity(cap(dv.GetData()), dataRequired) - if !ok { - return 0, false + for i, vec := range bat.Vecs { + if vec == nil || !r.schema[i].Eq(*vec.GetType()) { + return moerr.NewInternalError(proc.Ctx, "spill batch schema changed") } - if allocated, ok = addUint64(allocated, dataCap); !ok { - return 0, false - } - - if !dv.GetType().IsVarlen() { - continue - } - areaAdd, ok := mergedVarlenAreaAdd(sv, srcRows) - if !ok { - return 0, false - } - areaRequired, ok := addUint64(uint64(len(dv.GetArea())), areaAdd) - if !ok || areaRequired > uint64(math.MaxInt64) { - return 0, false - } - areaCap, ok := predictedCapacity(cap(dv.GetArea()), areaRequired) - if !ok { - return 0, false - } - if allocated, ok = addUint64(allocated, areaCap); !ok { - return 0, false - } - } - - cols := uint64(len(dst.Vecs)) - if cols > (math.MaxUint64-16)/8 { - return 0, false - } - metadata := 16 + cols*8 - rowMetadata, ok := mulUint64(mergedRows, metadata) - if !ok { - return 0, false } - return addUint64(allocated, rowMetadata) + return nil } -func mergedVarlenAreaAdd(src *vector.Vector, rows uint64) (uint64, bool) { - if src == nil || !src.GetType().IsVarlen() { - return 0, false - } - if rows == 0 { - return 0, true +func (r *BucketReader) peekRecordRows(proc *process.Process) (int64, error) { + if err := checkSpillCanceled(proc); err != nil { + return 0, err } - if src.IsConst() { - if src.IsConstNull() { - return 0, true - } - // UnionBatch materializes one const value and broadcasts its header. An - // inline value needs no area; a non-inline value appends exactly once. - if len(src.GetData()) < src.GetType().TypeSize() { - return 0, false - } - values := vector.MustFixedColNoTypeCheck[types.Varlena](src) - if len(values) != 1 { - return 0, false - } - value := &values[0] - if value.IsSmall() { - return 0, true - } - off, length := value.OffsetLen() - end, ok := addUint64(uint64(off), uint64(length)) - if !ok || end > uint64(len(src.GetArea())) { - return 0, false + if !r.headerPending { + if _, err := io.ReadFull(r.fd, r.header[:]); err != nil { + return 0, err } - return uint64(length), true - } - - // The full-record fast path copies the complete source area once, including - // stale bytes. Header validation remains UnionBatch's responsibility; avoid - // adding another per-row scan on the spill rebuild hot path. - return uint64(len(src.GetArea())), true -} - -func predictedCapacity(oldCap int, required uint64) (uint64, bool) { - if oldCap < 0 || uint64(oldCap) > uint64(math.MaxInt64) || required > uint64(math.MaxInt64) { - return 0, false - } - if required <= uint64(oldCap) { - return uint64(oldCap), true - } - cap, ok := mpool.GrowCapacity(int64(oldCap), int64(required)) - if !ok || cap < 0 { - return 0, false - } - return uint64(cap), true -} - -func intToUint64(v int) (uint64, bool) { - if v < 0 { - return 0, false - } - return uint64(v), true -} - -func mulUint64(a, b uint64) (uint64, bool) { - if a != 0 && b > math.MaxUint64/a { - return 0, false - } - return a * b, true -} - -func batchPayloadWithAllocationSlack(payload, columns uint64) (uint64, bool) { - const perVectorAllocationSlack = uint64(64 << 10) - if columns >= math.MaxUint64/perVectorAllocationSlack { - return 0, false - } - allocationSlack := (columns + 1) * perVectorAllocationSlack - if payload > math.MaxUint64-allocationSlack { - return 0, false - } - return payload + allocationSlack, true -} - -func decodedBatchProjectedBytes(payload uint64, rows int64, columns int32) (uint64, bool) { - if rows < 0 || columns < 0 { - return 0, false - } - projected, ok := batchPayloadWithAllocationSlack(payload, uint64(columns)) - if !ok { - return 0, false - } - metadata, ok := batchRetainedMetadataBytes(uint64(rows), uint64(columns)) - if !ok { - return 0, false - } - return addUint64(projected, metadata) -} - -func decodedBatchReusePeakBytes(retained, projected, payload uint64) (uint64, bool) { - // For large buffers mpool.Grow follows Go's 1.25x growth policy. The old - // allocation remains live until the replacement is allocated and copied. - // Small-buffer doubling is bounded by the per-vector slack already included - // in projected. - growthSlack := payload / 4 - if payload%4 != 0 { - growthSlack++ - } - newAllocation, ok := addUint64(projected, growthSlack) - if !ok { - return 0, false - } - return addUint64(retained, newAllocation) -} - -func maxIntValue() int { - return int(^uint(0) >> 1) -} - -func (r *BucketReader) releaseReadBatch(proc *process.Process, bat *batch.Batch, token *process.HashBuildReservation) { - if bat != nil { - bat.Clean(proc.Mp()) + r.headerPending = true } - if token != nil { - token.Release() - } - if r.batchToken != nil { - r.batchToken.Release() - r.batchToken = nil - r.batchCharge = 0 + rows := types.DecodeInt64(r.header[:8]) + batchSize := types.DecodeInt64(r.header[8:]) + if rows < 0 || batchSize < 0 { + return 0, moerr.NewInternalError( + proc.Ctx, + "negative spill batch header", + ) } + return rows, nil } func (r *BucketReader) readBatchRecord( proc *process.Process, reuseBat *batch.Batch, - token *process.HashBuildReservation, - charge uint64, - retainLease bool, -) (*batch.Batch, *process.HashBuildReservation, uint64, error) { +) (*batch.Batch, error) { if err := checkSpillCanceled(proc); err != nil { - return nil, token, charge, err + return nil, err } - if _, err := io.ReadFull(r.reader, r.buf[:]); err != nil { - if err == io.EOF { - return nil, token, charge, io.EOF + if !r.headerPending { + if _, err := io.ReadFull(r.fd, r.header[:]); err != nil { + return nil, err } - return nil, token, charge, err } - cnt := types.DecodeInt64(r.buf[:8]) - batchSize := types.DecodeInt64(r.buf[8:16]) + r.headerPending = false + cnt := types.DecodeInt64(r.header[:8]) + batchSize := types.DecodeInt64(r.header[8:]) if cnt < 0 || batchSize < 0 { - return nil, token, charge, moerr.NewInternalError(proc.Ctx, "negative spill batch header") - } - if err := checkSpillCanceled(proc); err != nil { - return nil, token, charge, err - } - if r.budget != nil && r.allocation == nil { - payload := uint64(batchSize) - if payload > uint64(maxIntValue())-(64<<10) { - return nil, token, charge, process.ErrHashBuildBudgetInvalid - } - // The batch payload starts with row count and vector count. Peek only the - // fixed header so allocator rounding can be bounded per decoded vector - // before UnmarshalFromReader performs any allocation. - header, err := r.reader.Peek(12) - if err != nil { - return nil, token, charge, err - } - rows := types.DecodeInt64(header[:8]) - columns := types.DecodeInt32(header[8:12]) - if rows != cnt { - return nil, token, charge, moerr.NewInternalError(proc.Ctx, "row count mismatch") - } - projected, ok := decodedBatchProjectedBytes(payload, rows, columns) - if !ok { - return nil, token, charge, process.ErrHashBuildBudgetInvalid - } - // The serialized payload already includes every vector's data, area, null - // bitmap, and headers. Reserve one decoded payload plus bounded allocator - // slack, then reconcile to the retained capacities reported by the batch. - // Multiplying the complete payload rejects large spill records before - // UnmarshalFromReader can establish their actual retained footprint. - if token == nil { - // A caller-provided reuse batch has no budget ownership on the first - // read. Drop it before admitting the decoded payload. - reuseBat.Clean(proc.Mp()) - var err error - token, err = r.budget.Reserve(projected) - if err != nil { - return nil, nil, 0, err - } - charge = projected - } else { - // Reusing vectors can briefly keep their old allocation alive while - // mpool.Grow allocates the replacement. Admit one complete decoded - // payload above the retained lease before unmarshal. If that transient - // peak does not fit, release the old batch and decode from a clean - // batch so a valid single payload is not rejected. - retained, retainedOK := batchRetainedBytes(reuseBat) - peak, peakOK := decodedBatchReusePeakBytes(retained, projected, payload) - var growErr error - if retainedOK && peakOK && peak > charge { - growErr = token.Grow(peak - charge) - } - if growErr != nil && - !hashbuild.IsRetryableMemoryCapacity(growErr) { - return nil, token, charge, growErr - } - if !retainedOK || !peakOK || growErr != nil { - reuseBat.Clean(proc.Mp()) - token.Release() - token = nil - var err error - token, err = r.budget.Reserve(projected) - if err != nil { - return nil, nil, 0, err - } - charge = projected - } else if peak > charge { - charge = peak - } - } + return nil, moerr.NewInternalError( + proc.Ctx, + "negative spill batch header", + ) } - - var payloadOffset int64 = -1 - if r.allocation != nil { - physical, seekErr := r.fd.Seek(0, io.SeekCurrent) - if seekErr != nil { - return nil, token, charge, seekErr - } - payloadOffset = physical - int64(r.reader.Buffered()) - if payloadOffset < 0 { - return nil, token, charge, process.ErrHashBuildBudgetInvalid - } + payloadOffset, err := r.fd.Seek(0, io.SeekCurrent) + if err != nil { + return nil, err } decode := func() (io.LimitedReader, error) { reuseBat.CleanOnlyData() if err := checkSpillCanceled(proc); err != nil { return io.LimitedReader{}, err } - limited := io.LimitedReader{R: r.reader, N: batchSize} - err := reuseBat.UnmarshalFromReader(&limited, proc.Mp()) - return limited, err + limited := io.LimitedReader{R: r.fd, N: batchSize} + return limited, reuseBat.UnmarshalFromReaderWithGrouping(&limited, proc.Mp()) } - limitReader, decodeErr := decode() - if decodeErr != nil && r.allocation != nil && - mpool.IsRetryableAllocationCapacity(decodeErr) { - // Reuse growth owns O while allocating N. If that exact overlap does - // not fit, rewind the seekable spill record, release O, and retry the - // same minimum payload from a clean batch. No multiplier predicts N. + limited, decodeErr := decode() + if decodeErr != nil && mpool.IsRetryableAllocationCapacity(decodeErr) { + // Reuse owns the old capacity while a replacement is allocated. Release + // it, rewind the unpublished record, and retry against actual allocation. reuseBat.Clean(proc.Mp()) if err := reuseBat.SetAllocationAccount(r.allocation.decoded); err != nil { - return nil, token, charge, err + return nil, err } if _, err := r.fd.Seek(payloadOffset, io.SeekStart); err != nil { - return nil, token, charge, err + return nil, err } - r.reader.Reset(r.fd) r.cleanRetries++ - limitReader, decodeErr = decode() + limited, decodeErr = decode() } if decodeErr != nil { - return nil, token, charge, decodeErr + return nil, decodeErr } - if err := checkSpillCanceled(proc); err != nil { - return nil, token, charge, err - } - - // Verify the batch unmarshal consumed exactly batchSize bytes. - if limitReader.N > 0 { - return nil, token, charge, moerr.NewInternalErrorf(proc.Ctx, "batch unmarshal did not consume all bytes: %d remaining", limitReader.N) + if limited.N != 0 { + return nil, moerr.NewInternalErrorf( + proc.Ctx, + "batch unmarshal did not consume all bytes: %d remaining", + limited.N, + ) } - - // Read magic (8 bytes) - if _, err := io.ReadFull(r.reader, r.buf[:8]); err != nil { - return nil, token, charge, err + if _, err := io.ReadFull(r.fd, r.header[:8]); err != nil { + return nil, err } - if types.DecodeUint64(r.buf[:8]) != SpillMagic { - return nil, token, charge, moerr.NewInternalError(proc.Ctx, "corrupted spill file") + if types.DecodeUint64(r.header[:8]) != SpillMagic { + return nil, moerr.NewInternalError(proc.Ctx, "corrupted spill file") } - if reuseBat.RowCount() != int(cnt) { - return nil, token, charge, moerr.NewInternalError(proc.Ctx, "row count mismatch") - } - if token != nil { - actual, ok := batchRetainedBytes(reuseBat) - if !ok || actual > charge { - return nil, token, charge, process.ErrHashBuildBudgetInvalid - } - target := actual - if retainLease { - withSlack, ok := addUint64(actual, decodedBatchLeaseSlack) - if !ok { - return nil, token, charge, process.ErrHashBuildBudgetInvalid - } - if withSlack < charge { - target = withSlack - } else { - target = charge - } - } - if target < charge { - if err := reconcileReadReservation(token, target); err != nil { - return nil, token, charge, err - } - charge = target - } + return nil, moerr.NewInternalError(proc.Ctx, "row count mismatch") } - return reuseBat, token, charge, nil + return reuseBat, checkSpillCanceled(proc) } -func (r *BucketReader) ResetForFd(fd *os.File) { - r.closeCurrentFile() - if fd == nil { - return +func (r *BucketReader) mergeReadError( + proc *process.Process, + dst *batch.Batch, + src *batch.Batch, + err error, +) error { + if src != nil { + src.Clean(proc.Mp()) } - r.fd = fd - if r.reader == nil { - r.reader = bufio.NewReaderSize(fd, 4*1024*1024) - } else { - r.reader.Reset(fd) + if dst != nil { + dst.Clean(proc.Mp()) } + return err } -func (r *BucketReader) ResetForSpillFile(file *message.SpillFile) { +func (r *BucketReader) ResetForSpillFile(file *message.SpillFile) error { r.closeCurrentFile() if file == nil { - return - } - r.spillFile = file - r.fd = file.File() - if r.reader == nil { - r.reader = bufio.NewReaderSize(r.fd, 4*1024*1024) - } else { - r.reader.Reset(r.fd) - } -} - -// EnsureBuffer admits the reader's fixed backing allocation before creating -// the bufio.Reader. Rebuild and probe readers are shared one at a time, so a -// single reservation bounds their peak rather than charging one per bucket. -func (r *BucketReader) EnsureBuffer(budget *process.HashBuildBudgetGeneration) error { - r.budget = budget - if r.reader != nil || budget == nil { return nil } - const size = uint64(64 << 10) - token, err := budget.Reserve(size) - if err != nil { + if err := file.Validate(); err != nil { + _ = file.Close() return err } - r.reservation = token - r.reader = bufio.NewReaderSize(nil, int(size)) + r.spillFile = file + r.fd = file.File() + r.headerPending = false + r.schema = nil return nil } @@ -750,25 +325,15 @@ func (r *BucketReader) closeCurrentFile() { r.fd = nil } if r.fd != nil { - r.fd.Close() + _ = r.fd.Close() r.fd = nil } + r.headerPending = false + r.schema = nil } func (r *BucketReader) Close() { r.closeCurrentFile() - if r.batchToken != nil { - r.batchToken.Release() - r.batchToken = nil - r.batchCharge = 0 - } - if r.reservation != nil { - r.reservation.Release() - r.reservation = nil - } - // A retained bufio.Reader capacity must remain charged. Drop it when the - // reservation is released; ResetForFd/ResetForSpillFile keep both alive. - r.reader = nil } // BucketWriter writes serialized batch records to an fd. @@ -828,23 +393,6 @@ func (w *BucketWriter) Close() { } } -func (w *BucketWriter) HandOffFd() *os.File { - if w.Fd == nil { - return nil - } - // A raw descriptor cannot carry accounting ownership. Budgeted writers - // must use handOffSpillFile; retain ownership here so Close can unwind it. - if w.fdReservation != nil || w.diskReservation != nil { - return nil - } - if _, err := w.Fd.Seek(0, io.SeekStart); err != nil { - return nil - } - fd := w.Fd - w.Fd = nil - return fd -} - func (w *BucketWriter) handOffSpillFile() (*message.SpillFile, error) { if w.Fd == nil { return nil, nil @@ -879,32 +427,6 @@ func MakeBucketWriters(prefix string) []BucketWriter { return writers } -// FlushBucketBatch writes one framed batch to w. It remains the low-level -// fixture/compatibility boundary for callers that already own a spill writer; -// production scatter uses SpillEngine so its retained buffers and pressure -// retries stay under the statement allocation owner. -func FlushBucketBatch( - proc *process.Process, - bat *batch.Batch, - w *BucketWriter, - bucketBuf *bytes.Buffer, - analyzer process.Analyzer, -) error { - if bat == nil || bat.RowCount() == 0 { - return nil - } - if err := marshalSpillRecord(bat, bucketBuf); err != nil { - return err - } - return writeBucketPayload( - proc, - bucketBuf.Bytes(), - int64(bat.RowCount()), - w, - analyzer, - ) -} - type spillRecordBuffer interface { io.Writer Bytes() []byte @@ -913,43 +435,6 @@ type spillRecordBuffer interface { Reset() } -type legacySpillRecordBuffer struct { - buffer *bytes.Buffer -} - -func (b legacySpillRecordBuffer) Write(value []byte) (int, error) { - return b.buffer.Write(value) -} - -func (b legacySpillRecordBuffer) Bytes() []byte { - return b.buffer.Bytes() -} - -func (b legacySpillRecordBuffer) EnsureCapacity(required int) error { - if b.buffer.Cap() < required { - *b.buffer = *bytes.NewBuffer(make([]byte, 0, required)) - } - return nil -} - -func (b legacySpillRecordBuffer) Len() int { - return b.buffer.Len() -} - -func (b legacySpillRecordBuffer) Reset() { - b.buffer.Reset() -} - -func marshalSpillRecord(bat *batch.Batch, buf *bytes.Buffer) error { - if buf == nil { - return process.ErrHashBuildBudgetInvalid - } - return marshalSpillRecordTo( - bat, - legacySpillRecordBuffer{buffer: buf}, - ) -} - func marshalSpillRecordTo( bat *batch.Batch, buf spillRecordBuffer, @@ -959,8 +444,8 @@ func marshalSpillRecordTo( } cnt := int64(bat.RowCount()) buf.Reset() - batchSize, err := bat.MarshalBinarySize() - if err != nil || batchSize > maxIntValue()-24 { + batchSize, err := bat.MarshalBinaryWithGroupingSize() + if err != nil || batchSize > math.MaxInt-24 { if err != nil { return err } @@ -969,25 +454,47 @@ func marshalSpillRecordTo( if err := buf.EnsureCapacity(batchSize + 24); err != nil { return err } - if err := writeSpillRecordBytes(buf, types.EncodeInt64(&cnt)); err != nil { + if err := writeSpillRecordInt64(buf, cnt); err != nil { return err } batchSizePos := buf.Len() - var zero int64 - if err := writeSpillRecordBytes(buf, types.EncodeInt64(&zero)); err != nil { + if err := writeSpillRecordInt64(buf, 0); err != nil { return err } batchStart := buf.Len() - if err := bat.MarshalBinaryTo(buf); err != nil { + if err := bat.MarshalBinaryWithGroupingTo(buf); err != nil { return err } serializedSize := int64(buf.Len() - batchStart) - copy( - buf.Bytes()[batchSizePos:batchSizePos+8], - types.EncodeInt64(&serializedSize), - ) - magic := uint64(SpillMagic) - return writeSpillRecordBytes(buf, types.EncodeUint64(&magic)) + if setter, ok := buf.(interface{ SetInt64(int, int64) error }); ok { + if err := setter.SetInt64(batchSizePos, serializedSize); err != nil { + return err + } + } else { + binary.NativeEndian.PutUint64( + buf.Bytes()[batchSizePos:batchSizePos+8], + uint64(serializedSize), + ) + } + return writeSpillRecordUint64(buf, uint64(SpillMagic)) +} + +func writeSpillRecordInt64(w io.Writer, value int64) error { + if typed, ok := w.(interface{ WriteInt64(int64) error }); ok { + return typed.WriteInt64(value) + } + var data [8]byte + binary.NativeEndian.PutUint64(data[:], uint64(value)) + return writeSpillRecordBytes(w, data[:]) +} + +func writeSpillRecordUint64(w io.Writer, value uint64) error { + if typed, ok := w.(interface{ WriteUint64(uint64) error }); ok { + return typed.WriteUint64(value) + } + var data [8]byte + binary.NativeEndian.PutUint64(data[:], value) + return writeSpillRecordBytes(w, data[:]) } func writeSpillRecordBytes(w io.Writer, value []byte) error { @@ -1007,6 +514,9 @@ func writeBucketPayload(proc *process.Process, payload []byte, rows int64, w *Bu if w == nil || len(payload) == 0 { return process.ErrHashBuildBudgetInvalid } + if w.Budget == nil { + return process.ErrHashBuildBudgetInvalid + } if err := checkSpillCanceled(proc); err != nil { return err } @@ -1023,32 +533,26 @@ func writeBucketPayload(proc *process.Process, payload []byte, rows int64, w *Bu } _, _ = w.diskReservation.ReconcileDown(oldDiskSize) } - if w.Budget != nil { - // Keep one growable disk token per file so bookkeeping remains bounded - // even when the input arrives as millions of tiny batches. - if w.diskReservation == nil { - diskToken, err := w.Budget.ReserveSpillDisk(uint64(len(payload))) - if err != nil { - return err - } - w.diskReservation = diskToken - newDiskToken = true - } else { - oldDiskSize = w.diskReservation.Size() - if err := w.diskReservation.Grow(uint64(len(payload))); err != nil { - return err - } + // Keep one growable disk token per file so bookkeeping remains bounded + // even when the input arrives as millions of tiny batches. + if w.diskReservation == nil { + diskToken, err := w.Budget.ReserveSpillDisk(uint64(len(payload))) + if err != nil { + return err + } + w.diskReservation = diskToken + newDiskToken = true + } else { + oldDiskSize = w.diskReservation.Size() + if err := w.diskReservation.Grow(uint64(len(payload))); err != nil { + return err } } if !w.Created() { - var fdToken *process.HashBuildSpillFDReservation - var err error - if w.Budget != nil { - fdToken, err = w.Budget.ReserveSpillFD(1) - if err != nil { - rollbackDisk() - return err - } + fdToken, err := w.Budget.ReserveSpillFD(1) + if err != nil { + rollbackDisk() + return err } fs, err := w.getSpillFileService(proc) if err != nil { @@ -1088,11 +592,6 @@ func writeBucketPayload(proc *process.Process, payload []byte, rows int64, w *Bu return nil } -// hashCombine merges a new hash value into a running hash state (Boost-style). -func hashCombine(h, val uint64) uint64 { - return keycodec.HashCombine(h, val) -} - // ComputeXXHash evaluates key vectors and computes XXHash64 values using // column-at-a-time processing for better cache locality. seed initialises every // hash slot so different spill depths produce different bucket distributions. @@ -1103,252 +602,37 @@ func ComputeXXHash(keyVecs []*vector.Vector, hashValues []uint64, seed uint64) { // classifyRows computes bucket counts, prefix offsets, and one contiguous row // id array in two linear passes. This replaces the historical bucket-by-bucket // scan of hashValues (which revisited every row once for each bucket). -func classifyRows(hashValues []uint64, bucketCount int, shift uint64, rowIDs []int32, counts []int32, offsets []int32) error { - if bucketCount <= 0 || bucketCount > SpillNumBuckets || - bucketCount&(bucketCount-1) != 0 || shift >= 64 || - len(rowIDs) < len(hashValues) || len(counts) < bucketCount || - len(offsets) < bucketCount+1 { - return process.ErrHashBuildBudgetInvalid - } - for i := 0; i < bucketCount; i++ { - counts[i] = 0 - } - mask := uint64(bucketCount - 1) - for _, hash := range hashValues { - counts[int((hash>>shift)&mask)]++ - } - offsets[0] = 0 - for i := 0; i < bucketCount; i++ { - offsets[i+1] = offsets[i] + counts[i] - } - var writePos [SpillNumBuckets]int32 - copy(writePos[:bucketCount], offsets[:bucketCount]) - for row, hash := range hashValues { - bucket := int((hash >> shift) & mask) - pos := writePos[bucket] - rowIDs[pos] = int32(row) - writePos[bucket] = pos + 1 - } - return nil -} - -func scatterTransientBudgetBytes(bat *batch.Batch, sourceAlreadyCharged bool) (uint64, error) { - if bat == nil || bat.RowCount() < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - allocated := uint64(bat.Allocated()) - if size := uint64(bat.Size()); size > allocated { - allocated = size - } - return scatterTransientBudgetFor(allocated, uint64(len(bat.Vecs)), sourceAlreadyCharged) -} - -func scatterTransientBudgetFor(allocated, columns uint64, sourceAlreadyCharged bool) (uint64, error) { - oneMaterializedBatch, ok := batchPayloadWithAllocationSlack(allocated, columns) - if !ok { - return 0, process.ErrHashBuildBudgetInvalid - } - // The selected batch and serialized payload are distinct live objects, so - // each receives one source-sized estimate plus bounded allocator/framing - // slack. The row/hash arrays are accounted separately from their capacities. - need, ok := addUint64(oneMaterializedBatch, oneMaterializedBatch) - if !ok { - return 0, process.ErrHashBuildBudgetInvalid - } - if !sourceAlreadyCharged { - if need, ok = addUint64(need, allocated); !ok { - return 0, process.ErrHashBuildBudgetInvalid - } - } - return need, nil -} - -// reserveRebuildScatterScratch protects the one-batch repartition workspace -// before the rebuild retains another decoded batch. The reservation is only an -// accounting lease: no scatter buffers are allocated until re-spill actually -// starts. Keeping this floor lets a copy admission fail early enough that the -// already-retained batches can still be repartitioned under the same hard cap. -func (e *SpillEngine) reserveRebuildScatterScratch( - builder *hashbuild.HashmapBuilder, - bat *batch.Batch, - analyzer process.Analyzer, -) error { - if e.cfg.Budget == nil { - return nil - } - if builder == nil || bat == nil || bat.RowCount() < 0 { - return process.ErrHashBuildBudgetInvalid - } - - allocated := uint64(bat.Allocated()) - if size := uint64(bat.Size()); size > allocated { - allocated = size - } - rows := bat.RowCount() - columns := len(bat.Vecs) - - // CopyIntoBatches may complete a partial physical tail with this record. - // Bound that resulting batch before the copy; reserving only for either - // input independently is not enough when two small records coalesce. - batches := builder.Batches.Buf - if bat.RowCount() != colexec.DefaultBatchSize && len(batches) > 0 { - tail := batches[len(batches)-1] - if tail == nil { - return process.ErrHashBuildBudgetInvalid - } - if tail.RowCount() != colexec.DefaultBatchSize { - merged, ok := predictMergedRetainedBytes(tail, bat) - if !ok { - return process.ErrHashBuildBudgetInvalid - } - if merged > allocated { - allocated = merged - } - if tail.RowCount() > math.MaxInt-bat.RowCount() { - return process.ErrHashBuildBudgetInvalid - } - rows += tail.RowCount() - if len(tail.Vecs) > columns { - columns = len(tail.Vecs) - } - } - } - - retained, ok := e.scatterRetainedBytes() - if !ok { - return process.ErrHashBuildBudgetInvalid - } - growth, ok := e.scatterCapacityGrowthBytes(rows, len(e.cfg.BuildKeyExprs)) - if !ok { - return process.ErrHashBuildBudgetInvalid - } - transient, err := scatterTransientBudgetFor(allocated, uint64(columns), true) - if err != nil { - return err - } - need, ok := addUint64(retained, growth) - if !ok { - return process.ErrHashBuildBudgetInvalid - } - if need, ok = addUint64(need, transient); !ok { - return process.ErrHashBuildBudgetInvalid - } - - return e.reserveRebuildScratchFloor(need, analyzer) -} - -func (e *SpillEngine) reserveRebuildScratchFloor(need uint64, analyzer process.Analyzer) error { - if need == 0 || e.cfg.Budget == nil { - return nil - } - var err error - if e.scatterScratchReservation == nil { - e.scatterScratchReservation, err = e.cfg.Budget.Reserve(need) - if err != nil { - if analyzer != nil { - analyzer.GetOpStats().AddExtraStat("JoinSpillRebuildScratchReserveRejects", 1) - } - return err - } - if analyzer != nil { - analyzer.GetOpStats().AddExtraStat("JoinSpillRebuildScratchReserveCount", 1) - } - } else if current := e.scatterScratchReservation.Size(); need > current { - grow := need - current - if err = e.scatterScratchReservation.Grow(grow); err != nil { - if analyzer != nil { - analyzer.GetOpStats().AddExtraStat("JoinSpillRebuildScratchGrowRejects", 1) - } - return err - } - if analyzer != nil { - analyzer.GetOpStats().AddExtraStat("JoinSpillRebuildScratchGrowCount", 1) - analyzer.GetOpStats().AddExtraStat("JoinSpillRebuildScratchGrowBytes", spillStatInt64(grow)) - } - } - if need > e.scatterScratchFloor { - e.scatterScratchFloor = need - } - if analyzer != nil { - analyzer.GetOpStats().SetMaxExtraStat( - "JoinSpillRebuildScratchFloorBytes", - spillStatInt64(e.scatterScratchReservation.Size()), - ) - } - return nil -} - -func spillStatInt64(v uint64) int64 { - if v > math.MaxInt64 { - return math.MaxInt64 - } - return int64(v) -} - -func (e *SpillEngine) scatterRetainedBytes() (uint64, bool) { - actual := uint64(0) - add := func(v uint64) bool { - var ok bool - actual, ok = addUint64(actual, v) - return ok - } - mul := func(v, n uint64) (uint64, bool) { - if n != 0 && v > math.MaxUint64/n { - return 0, false - } - return v * n, true - } - hashBytes, hashOK := mul(uint64(cap(e.scatterHashValues)), 8) - rowIDBytes, rowIDOK := mul(uint64(cap(e.scatterBucketRowIds)), 4) - keyBytes, keyOK := mul(uint64(cap(e.keyVecs)), 8) - countBytes, countOK := mul(uint64(len(e.scatterBucketCounts)), 4) - offsetBytes, offsetOK := mul(uint64(len(e.scatterBucketOffsets)), 4) - if !hashOK || !rowIDOK || !keyOK || !countOK || !offsetOK || - !add(hashBytes) || !add(rowIDBytes) || !add(keyBytes) || !add(countBytes) || !add(offsetBytes) || - !add(uint64(e.scatterWriteBuf.Cap())) { - return 0, false - } - for i := range e.scatterWriteBuffers { - if !add(uint64(e.scatterWriteBuffers[i].Cap())) { - return 0, false - } +func classifyRows(hashValues []uint64, bucketCount int, shift uint64, rowIDs []int32, counts []int32, offsets []int32) error { + if bucketCount <= 0 || bucketCount > SpillNumBuckets || + bucketCount&(bucketCount-1) != 0 || shift >= 64 || + len(rowIDs) < len(hashValues) || len(counts) < bucketCount || + len(offsets) < bucketCount+1 { + return process.ErrHashBuildBudgetInvalid } - return actual, true -} - -func (e *SpillEngine) scatterCapacityGrowthBytes(rows, keys int) (uint64, bool) { - if rows < 0 || keys < 0 { - return 0, false + for i := 0; i < bucketCount; i++ { + counts[i] = 0 } - var growth uint64 - addGrowth := func(required, current uint64) bool { - if required <= current { - return true - } - var ok bool - // make allocates the complete replacement before assignment drops the - // old slice. retained already includes current, so admit all of required. - growth, ok = addUint64(growth, required) - return ok + mask := uint64(bucketCount - 1) + for _, hash := range hashValues { + counts[int((hash>>shift)&mask)]++ } - rowCount := uint64(rows) - keyCount := uint64(keys) - if rowCount > math.MaxUint64/8 || keyCount > math.MaxUint64/8 { - return 0, false + offsets[0] = 0 + for i := 0; i < bucketCount; i++ { + offsets[i+1] = offsets[i] + counts[i] } - if !addGrowth(rowCount*8, uint64(cap(e.scatterHashValues))*8) || - !addGrowth(rowCount*4, uint64(cap(e.scatterBucketRowIds))*4) || - !addGrowth(keyCount*8, uint64(cap(e.keyVecs))*8) { - return 0, false + var writePos [SpillNumBuckets]int32 + copy(writePos[:bucketCount], offsets[:bucketCount]) + for row, hash := range hashValues { + bucket := int((hash >> shift) & mask) + pos := writePos[bucket] + rowIDs[pos] = int32(row) + writePos[bucket] = pos + 1 } - return growth, true + return nil } -// scatterBatchBounded writes one bucket at a time. The historical path kept -// SpillNumBuckets selected batches alive for the full input stream; that made -// a repartition pass itself exceed the hash-build budget. This implementation -// keeps one selected batch and one row-id slice, flushing it before advancing -// to the next bucket. +// scatterBatchBounded writes one bucket at a time. It keeps one selected batch +// and one row-id slice, flushing them before advancing to the next bucket. func (e *SpillEngine) scatterBatchBounded( proc *process.Process, bat *batch.Batch, @@ -1367,56 +651,25 @@ func (e *SpillEngine) scatterBatchBounded( if len(writers) == 0 || len(writers) > SpillNumBuckets { return process.ErrHashBuildBudgetInvalid } + if e.allocation == nil { + return mpool.ErrAllocationAccountInvalid + } rows := bat.RowCount() + if !keycodec.ValidVectors(bat.Vecs, rows) || + !keycodec.ValidVectors(keyVecs, rows) { + return process.ErrHashBuildBudgetInvalid + } var selected *batch.Batch defer func() { if selected != nil { selected.Clean(proc.Mp()) selected = nil } - reconcileErr := e.reconcileScatterScratch() - if reconcileErr != nil && retErr == nil { - retErr = reconcileErr - } - if retErr != nil && - !(e.allocation != nil && hashbuild.IsRetryableMemoryCapacity(retErr)) { + if retErr != nil && !hashbuild.IsRetryableMemoryCapacity(retErr) { e.discardScatterBuffers() } }() - if e.cfg.Budget != nil && e.allocation == nil { - // Start with retained capacities already owned by this token, add only - // row/hash capacity growth, then add each per-batch transient once. - retained, ok := e.scatterRetainedBytes() - if !ok { - return process.ErrHashBuildBudgetInvalid - } - growth, ok := e.scatterCapacityGrowthBytes(rows, len(keyVecs)) - if !ok { - return process.ErrHashBuildBudgetInvalid - } - transient, err := scatterTransientBudgetBytes(bat, sourceAlreadyCharged) - if err != nil { - return err - } - need, ok := addUint64(retained, growth) - if !ok { - return process.ErrHashBuildBudgetInvalid - } - if need, ok = addUint64(need, transient); !ok { - return process.ErrHashBuildBudgetInvalid - } - if e.scatterScratchReservation == nil { - e.scatterScratchReservation, err = e.cfg.Budget.Reserve(need) - } else if current := e.scatterScratchReservation.Size(); need > current { - // Grow the retained scratch token to the complete batch peak. Its - // current hash/row-id/coalesce capacities are components of need, - // not an additional allocation to charge a second time. - err = e.scatterScratchReservation.Grow(need - current) - } - if err != nil { - return err - } - } else if e.cfg.Budget != nil && !sourceAlreadyCharged { + if !sourceAlreadyCharged { // The child batch is borrowed and already physically live. Rejecting a // new logical token cannot reclaim it, so observe it while exact-account // admission governs every new scatter allocation. @@ -1430,12 +683,10 @@ func (e *SpillEngine) scatterBatchBounded( ) } - if e.allocation != nil { - if e.allocationMP != nil && e.allocationMP != proc.Mp() { - return mpool.ErrAllocationAccountInvalid - } - e.allocationMP = proc.Mp() + if e.allocationMP != nil && e.allocationMP != proc.Mp() { + return mpool.ErrAllocationAccountInvalid } + e.allocationMP = proc.Mp() var err error e.scatterHashValues, err = growSpillSlice( e.scatterHashValues, @@ -1485,18 +736,17 @@ func (e *SpillEngine) scatterBatchBounded( continue } if selected == nil { - var selection *vector.AllocationAccountSelection - if e.allocation != nil { - selection = e.allocation.selected - } - selected, err = newSpillBatch(len(bat.Vecs), selection) + selected, err = newSpillBatch( + len(bat.Vecs), + e.allocation.selected, + ) if err != nil { return err } for j, vec := range bat.Vecs { selected.Vecs[j], err = newSpillVector( *vec.GetType(), - selection, + e.allocation.selected, ) if err != nil { return err @@ -1535,8 +785,7 @@ func (e *SpillEngine) scatterBatchBounded( cursor = attemptEnd break } - if e.allocation == nil || - !hashbuild.IsRetryableMemoryCapacity(scatterErr) { + if !hashbuild.IsRetryableMemoryCapacity(scatterErr) { return scatterErr } if err := checkSpillCanceled(proc); err != nil { @@ -1620,11 +869,11 @@ func (e *SpillEngine) reclaimOptionalScatterBuffers( } func (e *SpillEngine) releaseScatterComputeScratch() { - if e.allocation == nil || e.allocationMP == nil { + if e.allocationMP == nil { return } - freeSpillSlice(e.scatterHashValues, e.allocationMP, e.allocation) - freeSpillSlice(e.scatterBucketRowIds, e.allocationMP, e.allocation) + freeSpillSlice(e.scatterHashValues, e.allocationMP) + freeSpillSlice(e.scatterBucketRowIds, e.allocationMP) e.scatterHashValues = nil e.scatterBucketRowIds = nil } @@ -1638,7 +887,7 @@ func (e *SpillEngine) scatterBatchWithPressure( sourceAlreadyCharged bool, analyzer process.Analyzer, ) error { - if e.allocation == nil || bat == nil || bat.RowCount() == 0 { + if bat == nil || bat.RowCount() == 0 { return e.scatterBatchBounded( proc, bat, @@ -1666,13 +915,17 @@ func (e *SpillEngine) scatterBatchWithPressure( currentKeys := keyVecs if start != 0 || end != rows { var err error - current, err = bat.Window(start, end) + current, err = bat.WindowWithAllocation( + start, end, proc.Mp(), e.allocation.selected, + ) if err != nil { return err } currentKeys = make([]*vector.Vector, len(keyVecs)) for i, key := range keyVecs { - currentKeys[i], err = key.Window(start, end) + currentKeys[i], err = key.WindowWithAllocation( + start, end, proc.Mp(), e.allocation.selected, + ) if err != nil { for j := 0; j < i; j++ { currentKeys[j].Free(proc.Mp()) @@ -1787,21 +1040,6 @@ func (e *SpillEngine) scatterEvaluatedBatchWithPressure( if eval == nil { return process.ErrHashBuildBudgetInvalid } - if e.allocation == nil { - keyVecs, err := eval(bat) - if err != nil { - return err - } - return e.scatterBatchWithPressure( - proc, - bat, - keyVecs, - writers, - partitionLevel, - sourceAlreadyCharged, - analyzer, - ) - } rows := bat.RowCount() chunk := rows @@ -1819,7 +1057,9 @@ func (e *SpillEngine) scatterEvaluatedBatchWithPressure( current := bat if start != 0 || end != rows { var err error - current, err = bat.Window(start, end) + current, err = bat.WindowWithAllocation( + start, end, proc.Mp(), e.allocation.selected, + ) if err != nil { return err } @@ -1924,44 +1164,14 @@ func (e *SpillEngine) appendScatterRecord(proc *process.Process, bat *batch.Batc if bucket < 0 || bucket >= SpillNumBuckets || writer == nil { return process.ErrHashBuildBudgetInvalid } - cnt := int64(bat.RowCount()) - if e.allocation != nil { - return e.appendAccountedScatterRecord( - proc, - bat, - writer, - bucket, - cnt, - analyzer, - ) - } - if err := marshalSpillRecord(bat, &e.scatterWriteBuf); err != nil { - return err - } - payload := e.scatterWriteBuf.Bytes() - buf := &e.scatterWriteBuffers[bucket] - if buf.Len() > 0 && buf.Len()+len(payload) > spillWriteCoalesceSize { - if err := e.flushPendingScatterBucket(proc, writer, bucket, analyzer); err != nil { - return err - } - } - if len(payload) > spillWriteCoalesceSize { - return writeBucketPayload(proc, payload, cnt, writer, analyzer) - } - if buf.Len() == 0 { - if !e.ensureScatterCoalesceCapacity(buf) { - return writeBucketPayload(proc, payload, cnt, writer, analyzer) - } - if buf.Cap() < spillWriteCoalesceSize { - *buf = *bytes.NewBuffer(make([]byte, 0, spillWriteCoalesceSize)) - } - } - _, _ = buf.Write(payload) - e.scatterWriteRows[bucket] += cnt - if buf.Len() >= spillWriteCoalesceSize { - return e.flushPendingScatterBucket(proc, writer, bucket, analyzer) - } - return nil + return e.appendAccountedScatterRecord( + proc, + bat, + writer, + bucket, + int64(bat.RowCount()), + analyzer, + ) } func (e *SpillEngine) appendAccountedScatterRecord( @@ -2047,45 +1257,21 @@ func (e *SpillEngine) appendAccountedScatterRecord( return nil } -func (e *SpillEngine) ensureScatterCoalesceCapacity(buf *bytes.Buffer) bool { - if buf == nil || buf.Cap() >= spillWriteCoalesceSize { - return true - } - if e.cfg.Budget == nil || e.scatterScratchReservation == nil { - return e.cfg.Budget == nil - } - additional := uint64(spillWriteCoalesceSize - buf.Cap()) - if err := e.scatterScratchReservation.Grow(additional); err != nil { - return false - } - return true -} - func (e *SpillEngine) flushPendingScatterBucket(proc *process.Process, writer *BucketWriter, bucket int, analyzer process.Analyzer) error { if bucket < 0 || bucket >= SpillNumBuckets || writer == nil { return process.ErrHashBuildBudgetInvalid } - if e.allocation != nil { - buf := e.scatterAccountedWriteBuffers[bucket] - if buf == nil || buf.Len() == 0 { - return nil - } - err := writeBucketPayload( - proc, - buf.Bytes(), - e.scatterWriteRows[bucket], - writer, - analyzer, - ) - buf.Reset() - e.scatterWriteRows[bucket] = 0 - return err - } - buf := &e.scatterWriteBuffers[bucket] - if buf.Len() == 0 { + buf := e.scatterAccountedWriteBuffers[bucket] + if buf == nil || buf.Len() == 0 { return nil } - err := writeBucketPayload(proc, buf.Bytes(), e.scatterWriteRows[bucket], writer, analyzer) + err := writeBucketPayload( + proc, + buf.Bytes(), + e.scatterWriteRows[bucket], + writer, + analyzer, + ) buf.Reset() e.scatterWriteRows[bucket] = 0 return err @@ -2097,9 +1283,8 @@ func (e *SpillEngine) flushPendingScatterBucket(proc *process.Process, writer *B func (e *SpillEngine) flushScatterBuffers(proc *process.Process, writers []BucketWriter, analyzer process.Analyzer) error { var firstErr error for bucket := 0; bucket < SpillNumBuckets; bucket++ { - pending := e.scatterWriteBuffers[bucket].Len() - if e.allocation != nil && - e.scatterAccountedWriteBuffers[bucket] != nil { + pending := 0 + if e.scatterAccountedWriteBuffers[bucket] != nil { pending = e.scatterAccountedWriteBuffers[bucket].Len() } if pending == 0 { @@ -2117,8 +1302,7 @@ func (e *SpillEngine) flushScatterBuffers(proc *process.Process, writers []Bucke } func (e *SpillEngine) discardScatterBuffers() { - for bucket := range e.scatterWriteBuffers { - e.scatterWriteBuffers[bucket].Reset() + for bucket := range e.scatterAccountedWriteBuffers { if e.scatterAccountedWriteBuffers[bucket] != nil { e.scatterAccountedWriteBuffers[bucket].Reset() } @@ -2134,17 +1318,14 @@ func (e *SpillEngine) releaseScatterScratch() { freeSpillSlice( e.scatterHashValues, e.allocationMP, - e.allocation, ) freeSpillSlice( e.scatterBucketRowIds, e.allocationMP, - e.allocation, ) e.scatterHashValues = nil e.scatterBucketRowIds = nil e.keyVecs = nil - e.scatterWriteBuf = bytes.Buffer{} if e.scatterAccountedWriteBuf != nil { e.scatterAccountedWriteBuf.Free() e.scatterAccountedWriteBuf = nil @@ -2155,8 +1336,7 @@ func (e *SpillEngine) releaseScatterScratch() { for i := range e.scatterBucketOffsets { e.scatterBucketOffsets[i] = 0 } - for i := range e.scatterWriteBuffers { - e.scatterWriteBuffers[i] = bytes.Buffer{} + for i := range e.scatterAccountedWriteBuffers { if e.scatterAccountedWriteBuffers[i] != nil { e.scatterAccountedWriteBuffers[i].Free() e.scatterAccountedWriteBuffers[i] = nil @@ -2165,37 +1345,6 @@ func (e *SpillEngine) releaseScatterScratch() { } e.allocationMP = nil e.scatterCoalesceDisabled = false - if e.scatterScratchReservation != nil { - e.scatterScratchReservation.Release() - e.scatterScratchReservation = nil - } - e.scatterScratchFloor = 0 -} - -// reconcileScatterScratch leaves only the capacities retained by the engine -// charged after a batch completes. Source and selected vectors are transient; -// reusable marshal and coalesce buffers remain charged only for the phase. -func (e *SpillEngine) reconcileScatterScratch() error { - if e.scatterScratchReservation == nil { - return nil - } - actual, ok := e.scatterRetainedBytes() - if !ok { - return process.ErrHashBuildBudgetInvalid - } - if actual < e.scatterScratchFloor { - actual = e.scatterScratchFloor - } - reserved := e.scatterScratchReservation.Size() - if actual > reserved { - return process.ErrHashBuildBudgetInvalid - } - if actual < reserved { - if _, err := e.scatterScratchReservation.ReconcileDown(actual); err != nil { - return err - } - } - return nil } // ReusableBufferPool maintains a persistent pool of spill buffers, preserving @@ -2256,12 +1405,8 @@ type SpillEngineConfig struct { // Budget is the statement generation shared with HashBuild. Rebuild and // re-spill must charge this exact generation; creating a fresh generation // would bypass aggregate admission and make ownership impossible to audit. - Budget *process.HashBuildBudgetGeneration - // ProbeExpressionLease is owned by the consuming join operator and borrowed - // by SpillEngine while it scatters or re-scatters probe batches. The join - // must free its probe executors before releasing this lease. - ProbeExpressionLease *hashbuild.ExpressionMemoryLease - MaxQueue int + Budget *process.HashBuildBudgetGeneration + MaxQueue int } // BucketResult encodes the outcome of a RebuildHashmap call. @@ -2284,85 +1429,66 @@ type SpillEngine struct { allocationMP *mpool.MPool // Current bucket state - buildReader BucketReader - probeReader BucketReader - buildReadBatch *batch.Batch - probeReadBatch *batch.Batch + buildReader BucketReader + probeReader BucketReader + buildReadBatch *batch.Batch + probeReadBatch *batch.Batch + probeExpected int64 + probeDecoded int64 + probeExpectedSet bool // Reusable scatter state buildPool ReusableBufferPool probePool ReusableBufferPool // Cached key executors for re-spill - keyExecs []colexec.ExpressionExecutor - keyVecs []*vector.Vector - buildExprLease *hashbuild.ExpressionMemoryLease + keyExecs []colexec.ExpressionExecutor + keyVecs []*vector.Vector // Reusable scatter buffers to avoid per-batch allocations. scatterHashValues []uint64 scatterBucketRowIds []int32 scatterBucketCounts [SpillNumBuckets]int32 scatterBucketOffsets [SpillNumBuckets + 1]int32 - scatterWriteBuf bytes.Buffer - scatterWriteBuffers [SpillNumBuckets]bytes.Buffer scatterAccountedWriteBuf *mpool.AccountedBuffer scatterAccountedWriteBuffers [SpillNumBuckets]*mpool.AccountedBuffer scatterCoalesceDisabled bool scatterWriteRows [SpillNumBuckets]int64 - // The lease follows the reusable scratch capacities for the engine - // lifetime. It is released only by Cleanup, after all backing arrays have - // been dropped. - scatterScratchReservation *process.HashBuildReservation - // scatterScratchFloor is pre-admitted only while rebuilding an already - // spilled bucket. It keeps one bounded repartition workspace available if the - // next retained-copy admission or threshold decision requires re-spill. It - // is a conservative bound, not a measurement of later physical allocations. - scatterScratchFloor uint64 - - // probeKeyEval is the unbudgeted fallback for probe re-scatter. Production - // spilled joins evaluate the probe executors owned by ProbeExpressionLease. + // probeKeyEval evaluates the consuming join's allocation-accounted probe + // executors during re-scatter. probeKeyEval func(*batch.Batch) ([]*vector.Vector, error) } -// NewSpillEngine creates an engine from configuration. Call InitFromSpilledMap next. -func NewSpillEngine(cfg SpillEngineConfig) *SpillEngine { - return newSpillEngine(cfg, nil) -} - -// NewSpillEngineWithAllocation constructs the allocation-accounted spill path. -// NewSpillEngine remains available to callers outside a statement account. -func NewSpillEngineWithAllocation( - cfg SpillEngineConfig, - allocation *SpillAllocationAccount, -) (*SpillEngine, error) { - if err := allocation.validate(); err != nil { - return nil, err - } - return newSpillEngine(cfg, allocation), nil -} - -// NewSpillEngineForAccount activates exact spill ownership when an execution -// attempt supplied an account and preserves the legacy constructor for direct -// operator tests and callers outside the statement lifecycle. -func NewSpillEngineForAccount( +// NewSpillEngine binds every spill allocation to one execution generation. +// A spill engine cannot exist outside that account. +func NewSpillEngine( cfg SpillEngineConfig, account *mpool.AllocationAccount, owner mpool.AllocationOwner, ) (*SpillEngine, error) { if account == nil { - return NewSpillEngine(cfg), nil + return nil, mpool.ErrAllocationAccountInvalid } allocation, err := NewSpillAllocationAccount(account, owner) if err != nil { return nil, err } - return NewSpillEngineWithAllocation(cfg, allocation) + return newSpillEngine(cfg, allocation) } func newSpillEngine( cfg SpillEngineConfig, allocation *SpillAllocationAccount, -) *SpillEngine { +) (*SpillEngine, error) { + if err := allocation.validate(); err != nil { + return nil, err + } + if cfg.Budget == nil { + return nil, process.ErrHashBuildBudgetInvalid + } + if cfg.Budget.Closed() { + return nil, process.ErrHashBuildBudgetClosed + } if cfg.MaxQueue <= 0 { cfg.MaxQueue = SpillNumBuckets * SpillNumBuckets } @@ -2372,13 +1498,14 @@ func newSpillEngine( } engine.buildReader.allocation = allocation engine.probeReader.allocation = allocation - return engine + return engine, nil } func (e *SpillEngine) makeBucketWriters(prefix string) []BucketWriter { writers := MakeBucketWriters(prefix) for i := range writers { writers[i].spillFS = &e.spillFS + writers[i].Budget = e.cfg.Budget } return writers } @@ -2390,51 +1517,62 @@ func TakeSpillBuildPayload( proc *process.Process, jm *message.JoinMap, ) (message.SpillBuildPayload, *process.HashBuildBudgetGeneration, error) { + if jm == nil { + return message.SpillBuildPayload{}, nil, moerr.NewInternalError( + proc.Ctx, + message.ErrSpillBuildPayloadEmpty.Error(), + ) + } + expectedRows := jm.GetRowCount() payload, err := jm.TakeSpillBuildPayload() if err != nil { return message.SpillBuildPayload{}, nil, moerr.NewInternalError(proc.Ctx, err.Error()) } - - var budget *process.HashBuildBudgetGeneration - if len(payload.Files) > 0 { - var ok bool - budget, ok = payload.BudgetRef.(*process.HashBuildBudgetGeneration) - if !ok || budget == nil { - _ = payload.Close() - return message.SpillBuildPayload{}, nil, moerr.NewInternalError( - proc.Ctx, - "spilled join map is missing its producer budget generation", - ) + actualRows := int64(0) + validRows := expectedRows >= 0 + for _, file := range payload.Files { + if file == nil { + continue } - } else { - budget, err = proc.GetHashBuildBudget() - if err != nil { + rows := file.Rows() + if rows <= 0 || actualRows > math.MaxInt64-rows { + validRows = false + break + } + actualRows += rows + } + if !validRows || actualRows != expectedRows { + _ = payload.Close() + return message.SpillBuildPayload{}, nil, spillRowCountMismatch( + proc, + "build payload", + expectedRows, + actualRows, + ) + } + for _, file := range payload.Files { + if file == nil { + continue + } + if err := file.Validate(); err != nil { _ = payload.Close() return message.SpillBuildPayload{}, nil, err } } - return payload, budget, nil -} -// InitFromSpilledMap creates SpillBucket entries from build FDs. -// Empty (nil) FDs become placeholder buckets for outer-join semantics. -func (e *SpillEngine) InitFromSpilledMap(buildFds []*os.File) { - e.buckets = make([]SpillBucket, 0, len(buildFds)) - for _, fd := range buildFds { - var file *message.SpillFile - if fd != nil { - file = message.NewSpillFile(fd, 0, 0, nil) - } - e.buckets = append(e.buckets, SpillBucket{ - BuildFd: file, - Depth: 1, - }) + budget, ok := payload.BudgetRef.(*process.HashBuildBudgetGeneration) + if !ok || budget == nil { + _ = payload.Close() + return message.SpillBuildPayload{}, nil, moerr.NewInternalError( + proc.Ctx, + "spilled join map is missing its producer budget generation", + ) } + return payload, budget, nil } -// InitFromSpilledFiles is the ownership-preserving counterpart of the legacy -// descriptor initializer. Each SpillFile remains the sole owner of its fd and -// reservations while it moves through the bucket queue. +// InitFromSpilledFiles transfers the sole ownership of each build spill file +// and its resource reservations into the bucket queue. func (e *SpillEngine) InitFromSpilledFiles(files []*message.SpillFile) { e.buckets = make([]SpillBucket, 0, len(files)) for _, file := range files { @@ -2465,11 +1603,8 @@ func (e *SpillEngine) ScatterProbeTable( // The build payload defines the partition fanout. Using the production // maximum unconditionally would hash probe rows into writers that have no // corresponding build bucket; those files are never handed off and their - // rows would be silently discarded for legacy or reduced-fanout payloads. + // rows would be silently discarded for reduced-fanout payloads. writers := e.makeBucketWriters("probe")[:bucketCount] - for i := range writers { - writers[i].Budget = e.cfg.Budget - } // Disable writers for empty-build buckets unless outer join requires probe output. if !e.cfg.NeedsProbeForEmptyBuild { @@ -2566,12 +1701,8 @@ func (e *SpillEngine) NextProbeBatch(proc *process.Process) (*batch.Batch, error return nil, nil } if e.probeReadBatch == nil { - var selection *vector.AllocationAccountSelection - if e.allocation != nil { - selection = e.allocation.decoded - } var err error - e.probeReadBatch, err = newSpillBatch(0, selection) + e.probeReadBatch, err = newSpillBatch(0, e.allocation.decoded) if err != nil { return nil, err } @@ -2579,6 +1710,14 @@ func (e *SpillEngine) NextProbeBatch(proc *process.Process) (*batch.Batch, error e.probeReader.mergeRecords = e.cfg.MergeProbeBatches || e.cfg.IsDedup bat, err := e.probeReader.ReadBatch(proc, e.probeReadBatch) if err == io.EOF { + if e.probeExpectedSet && e.probeDecoded != e.probeExpected { + return nil, spillRowCountMismatch( + proc, + "probe", + e.probeExpected, + e.probeDecoded, + ) + } return nil, nil } if err != nil { @@ -2587,30 +1726,63 @@ func (e *SpillEngine) NextProbeBatch(proc *process.Process) (*batch.Batch, error // Cancellation can race the reader's final record-boundary check. Do not // hand a freshly decoded batch to the join probe loop after that point. if err := checkSpillCanceled(proc); err != nil { - e.probeReader.releaseReadBatch(proc, e.probeReadBatch, nil) + e.probeReadBatch.Clean(proc.Mp()) return nil, err } + rows := int64(bat.RowCount()) + if rows < 0 || e.probeDecoded > math.MaxInt64-rows { + e.probeReadBatch.Clean(proc.Mp()) + return nil, spillRowCountMismatch( + proc, + "probe", + e.probeExpected, + math.MaxInt64, + ) + } + decoded := e.probeDecoded + rows + if e.probeExpectedSet && decoded > e.probeExpected { + e.probeReadBatch.Clean(proc.Mp()) + return nil, spillRowCountMismatch( + proc, + "probe", + e.probeExpected, + decoded, + ) + } + e.probeDecoded = decoded return bat, nil } -// builderMemSize computes total memory used by a HashmapBuilder during the rebuild -// loop. MemSize covers completed fixed-size batches; include the one permitted -// partial tail as well. The full scan is only a fallback for directly assembled -// state where MemSize has not been maintained. -func builderMemSize(builder *hashbuild.HashmapBuilder) int64 { - sz := builder.GetSize() + builder.Batches.MemSize - batches := builder.Batches.Buf - if builder.Batches.MemSize == 0 { - for _, b := range builder.Batches.Buf { - sz += int64(b.Size()) - } - } else if len(batches) > 0 { - tail := batches[len(batches)-1] - if tail != nil && tail.RowCount() != colexec.DefaultBatchSize { - sz += int64(tail.Size()) - } +func (e *SpillEngine) startProbe(file *message.SpillFile, expected int64) error { + if err := e.probeReader.ResetForSpillFile(file); err != nil { + return err } - return sz + e.probeExpected = expected + e.probeDecoded = 0 + e.probeExpectedSet = true + return nil +} + +func spillRowCountMismatch( + proc *process.Process, + side string, + expected int64, + actual int64, +) error { + return moerr.NewInternalErrorf( + proc.Ctx, + "corrupted spill %s row count: expected=%d actual=%d", + side, + expected, + actual, + ) +} + +// builderMemSize computes total memory used by a HashmapBuilder during the +// rebuild loop. GetSize covers hashmap structures and Batches.MemSize covers +// the raw accumulated batches maintained by the builder's copy API. +func builderMemSize(builder *hashbuild.HashmapBuilder) int64 { + return builder.GetSize() + builder.Batches.MemSize } func shouldReSpillBeforeRetain( @@ -2647,6 +1819,14 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana return nil, BucketQueueEmpty, nil } bucket := e.buckets[0] + if (bucket.BuildFd == nil && bucket.BuildRows != 0) || + (bucket.ProbeFd == nil && bucket.ProbeRows != 0) || + bucket.BuildRows < 0 || bucket.ProbeRows < 0 { + return nil, BucketSkip, moerr.NewInternalError( + proc.Ctx, + "corrupted spill bucket file/row metadata", + ) + } // A build-only bucket cannot contribute to joins that never emit unmatched // build rows. Close and pop it before allocating a reader, copying batches, @@ -2665,11 +1845,10 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana e.buckets[0].ProbeFd = nil // transferred to reader below; prevent Cleanup double-close e.buckets = e.buckets[1:] if e.cfg.NeedsProbeForEmptyBuild && bucket.ProbeFd != nil { - if err := e.probeReader.EnsureBuffer(e.cfg.Budget); err != nil { - bucket.ProbeFd.Close() + if err := e.startProbe(bucket.ProbeFd, bucket.ProbeRows); err != nil { + bucket.ProbeFd = nil return nil, BucketSkip, err } - e.probeReader.ResetForSpillFile(bucket.ProbeFd) bucket.ProbeFd = nil return nil, BucketEmptyBuild, nil } @@ -2680,12 +1859,14 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana } builder := &hashbuild.HashmapBuilder{} + // The rebuild builder is stack-owned until GetJoinMap detaches its durable + // state. Keep a panic-safe terminal guard: the outer pipeline recover cannot + // otherwise reach a local builder abandoned by expression, hash, or spill + // hooks. + defer builder.Free(proc) builder.SetBudget(e.cfg.Budget) - if e.allocation != nil { - if err := builder.SetAllocationAccount(e.allocation.account); err != nil { - builder.Free(proc) - return nil, BucketSkip, err - } + if err := builder.SetAllocationAccount(e.allocation.account); err != nil { + return nil, BucketSkip, err } builder.IsDedup = e.cfg.IsDedup builder.OnDuplicateAction = e.cfg.OnDuplicateAction @@ -2697,21 +1878,13 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana return nil, BucketSkip, err } - if err := e.buildReader.EnsureBuffer(e.cfg.Budget); err != nil { - builder.Free(proc) - bucket.BuildFd.Close() - bucket.BuildFd = nil + if err := e.buildReader.ResetForSpillFile(bucket.BuildFd); err != nil { return nil, BucketSkip, err } - e.buildReader.ResetForSpillFile(bucket.BuildFd) e.buckets[0].BuildFd = nil // prevent Cleanup double-close on error defer e.buildReader.closeCurrentFile() if e.buildReadBatch == nil { - var selection *vector.AllocationAccountSelection - if e.allocation != nil { - selection = e.allocation.decoded - } - readBatch, err := newSpillBatch(0, selection) + readBatch, err := newSpillBatch(0, e.allocation.decoded) if err != nil { builder.Free(proc) return nil, BucketSkip, err @@ -2756,17 +1929,6 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana return nil, BucketSkip, err } if bucket.Depth < SpillMaxPass { - if err := e.reserveRebuildScatterScratch(builder, bat, analyzer); err != nil { - // Scratch is contingency headroom, not a prerequisite for a - // bucket that may still rebuild within the cap. Admission misses - // are observable but best-effort; lifecycle/accounting failures - // remain terminal and are returned unchanged. - if !isBudgetAdmission(err) { - builder.FreeHashMapAndBatches(proc) - builder.Free(proc) - return nil, BucketSkip, err - } - } if shouldReSpillBeforeRetain(builder, bat, e.cfg.SpillThreshold) { if analyzer != nil { analyzer.GetOpStats().AddExtraStat("JoinSpillRebuildPreCopyReSpillAttempts", 1) @@ -2808,6 +1970,16 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana return nil, BucketReSpilled, nil } } + if int64(builder.InputBatchRowCount) != bucket.BuildRows { + builder.FreeHashMapAndBatches(proc) + builder.Free(proc) + return nil, BucketSkip, spillRowCountMismatch( + proc, + "build", + bucket.BuildRows, + int64(builder.InputBatchRowCount), + ) + } if err := checkSpillCanceled(proc); err != nil { builder.FreeHashMapAndBatches(proc) @@ -2851,7 +2023,14 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana } } - jm := builder.GetJoinMap(proc.Mp()) + var jm *message.JoinMap + joinMapTransferred := false + defer func() { + if jm != nil && !joinMapTransferred { + jm.FreeMemory() + } + }() + jm = builder.GetJoinMap(proc.Mp()) if jm == nil { // GetJoinMap transfers nothing when the decoded build contains no // rows. Release executors and every residual builder allocation before @@ -2860,11 +2039,10 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana e.buckets[0].ProbeFd = nil // transferred to reader below; prevent Cleanup double-close e.buckets = e.buckets[1:] if e.cfg.NeedsProbeForEmptyBuild && bucket.ProbeFd != nil { - if err := e.probeReader.EnsureBuffer(e.cfg.Budget); err != nil { - bucket.ProbeFd.Close() + if err := e.startProbe(bucket.ProbeFd, bucket.ProbeRows); err != nil { + bucket.ProbeFd = nil return nil, BucketSkip, err } - e.probeReader.ResetForSpillFile(bucket.ProbeFd) bucket.ProbeFd = nil return nil, BucketEmptyBuild, nil } @@ -2881,14 +2059,13 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana // Pop the head bucket and open probe reader. e.buckets = e.buckets[1:] if bucket.ProbeFd != nil { - if err := e.probeReader.EnsureBuffer(e.cfg.Budget); err != nil { - bucket.ProbeFd.Close() - jm.Free() + if err := e.startProbe(bucket.ProbeFd, bucket.ProbeRows); err != nil { + bucket.ProbeFd = nil return nil, BucketSkip, err } - e.probeReader.ResetForSpillFile(bucket.ProbeFd) bucket.ProbeFd = nil } + joinMapTransferred = true return jm, BucketReady, nil } @@ -2901,13 +2078,7 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal // executors, so the two equivalent retained working sets never overlap. builder.FreeExecutors() buildWriters := e.makeBucketWriters("build_sub") - for i := range buildWriters { - buildWriters[i].Budget = e.cfg.Budget - } probeWriters := e.makeBucketWriters("probe_sub") - for i := range probeWriters { - probeWriters[i].Budget = e.cfg.Budget - } partitionLevel := uint64(bucket.Depth) probeFdConsumed := false @@ -2941,39 +2112,13 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal // Cache key executors. if len(e.keyExecs) != len(e.cfg.BuildKeyExprs) { - var execs []colexec.ExpressionExecutor - var lease *hashbuild.ExpressionMemoryLease - var err error - if e.allocation == nil || - !hashbuild.AllocationAccountedExpressionSetSupported( - e.cfg.BuildKeyExprs, - ) { - execs, lease, err = - hashbuild.NewBudgetedExpressionExecutors( - proc, - e.cfg.Budget, - e.cfg.BuildKeyExprs, - false, - ) - } else { - execs, err = - hashbuild.NewAllocationAccountedExpressionExecutors( - proc, - e.cfg.BuildKeyExprs, - e.allocation.expression, - ) - if err == nil { - lease, err = hashbuild.NewExpressionMemoryLease( - nil, - e.cfg.BuildKeyExprs, - execs, - false, - ) - } - if err != nil { - for _, exec := range execs { - exec.Free() - } + execs, err := hashbuild.NewExpressionExecutors( + proc, + e.cfg.BuildKeyExprs, + ) + if err != nil { + for _, exec := range execs { + exec.Free() } } if err != nil { @@ -2981,7 +2126,6 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal } e.freeKeyExecs() e.keyExecs = execs - e.buildExprLease = lease } // evalAndScatter builds key vectors using the given executors and scatters. @@ -3011,24 +2155,19 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal for i := range keyVecs { keyVecs[i] = nil } - err := e.buildExprLease.Run(proc, current.RowCount(), func(i int) error { - vec, evalErr := execs[i].Eval(proc, []*batch.Batch{current}, nil) - if evalErr != nil { - return evalErr + for i := range execs { + vec, err := execs[i].Eval(proc, []*batch.Batch{current}, nil) + if err != nil { + // Exact capacity pressure keeps the executor tree as the + // rollback checkpoint: admitted child/result capacities may + // make a smaller immutable window fit. Every other failure is + // terminal and can destroy the private tree immediately. + if !hashbuild.IsRetryableMemoryCapacity(err) { + e.freeKeyExecs() + } + return nil, err } keyVecs[i] = vec - return nil - }) - if err != nil { - // Exact capacity pressure keeps the executor tree as the - // rollback checkpoint: admitted child/result capacities may - // make a smaller immutable window fit. Every other failure is - // terminal and can destroy the private tree immediately. - if e.allocation == nil || - !hashbuild.IsRetryableMemoryCapacity(err) { - e.freeKeyExecs() - } - return nil, err } return keyVecs, nil }, @@ -3074,29 +2213,31 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal return nil, err } } + if buildRows != bucket.BuildRows { + return nil, spillRowCountMismatch( + proc, + "build", + bucket.BuildRows, + buildRows, + ) + } if err := e.flushScatterBuffers(proc, buildWriters, analyzer); err != nil { return nil, err } if e.probeReadBatch == nil { - var selection *vector.AllocationAccountSelection - if e.allocation != nil { - selection = e.allocation.decoded - } - readBatch, err := newSpillBatch(0, selection) + readBatch, err := newSpillBatch(0, e.allocation.decoded) if err != nil { return nil, err } e.probeReadBatch = readBatch } - // Scatter the probe file through the same admitted 64 KiB reader buffer - // used for the build pass. + var probeRows int64 if bucket.ProbeFd != nil { - if err := reader.EnsureBuffer(e.cfg.Budget); err != nil { + if err := reader.ResetForSpillFile(bucket.ProbeFd); err != nil { return nil, err } - reader.ResetForSpillFile(bucket.ProbeFd) probeFdConsumed = true // Disable probe writers for empty sub-build buckets (unless outer join). if !e.cfg.NeedsProbeForEmptyBuild { @@ -3117,10 +2258,19 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal if err != nil { return nil, err } + probeRows += int64(bat.RowCount()) if err := scatterProbe(proc, e, bat, probeWriters, partitionLevel, analyzer); err != nil { return nil, err } } + if probeRows != bucket.ProbeRows { + return nil, spillRowCountMismatch( + proc, + "probe", + bucket.ProbeRows, + probeRows, + ) + } if err := e.flushScatterBuffers(proc, probeWriters, analyzer); err != nil { return nil, err } @@ -3277,9 +2427,7 @@ func (e *SpillEngine) AdvanceToNextBucket( return true, nil } -// scatterProbe evaluates probe-side keys (EqConds[0]) for probe re-scatter. -// It uses the borrowed probe lease, not build-side keyExecs; probeKeyEval is -// retained only as the unbudgeted fallback. +// scatterProbe evaluates the consuming join's probe-side keys for re-scatter. func scatterProbe(proc *process.Process, e *SpillEngine, bat *batch.Batch, writers []BucketWriter, seed uint64, analyzer process.Analyzer) error { return e.scatterEvaluatedBatchWithPressure( proc, @@ -3297,37 +2445,12 @@ func scatterProbe(proc *process.Process, e *SpillEngine, bat *batch.Batch, write func (e *SpillEngine) evalProbeKeys( proc *process.Process, bat *batch.Batch, - fallback func(*batch.Batch) ([]*vector.Vector, error), + eval func(*batch.Batch) ([]*vector.Vector, error), ) ([]*vector.Vector, error) { - if e.cfg.ProbeExpressionLease == nil { - if fallback == nil { - return nil, process.ErrHashBuildBudgetInvalid - } - return fallback(bat) - } - if e.cfg.ProbeExpressionLease.Len() != len(e.cfg.ProbeKeyExprs) { + if eval == nil { return nil, process.ErrHashBuildBudgetInvalid } - if cap(e.keyVecs) < len(e.cfg.ProbeKeyExprs) { - e.keyVecs = make([]*vector.Vector, len(e.cfg.ProbeKeyExprs)) - } - keyVecs := e.keyVecs[:len(e.cfg.ProbeKeyExprs)] - err := e.cfg.ProbeExpressionLease.Eval( - proc, - []*batch.Batch{bat}, - bat.RowCount(), - func(index int, vec *vector.Vector) error { - keyVecs[index] = vec - return nil - }, - ) - if err != nil { - for i := range keyVecs { - keyVecs[i] = nil - } - return nil, err - } - return keyVecs, nil + return eval(bat) } func (e *SpillEngine) freeKeyExecs() { @@ -3337,10 +2460,6 @@ func (e *SpillEngine) freeKeyExecs() { } } e.keyExecs = nil - if e.buildExprLease != nil { - e.buildExprLease.Release() - e.buildExprLease = nil - } } func isBudgetAdmission(err error) bool { diff --git a/pkg/sql/colexec/spillutil/join_spill_test.go b/pkg/sql/colexec/spillutil/join_spill_test.go index 0281c17a441d6..b3bd33ef23e90 100644 --- a/pkg/sql/colexec/spillutil/join_spill_test.go +++ b/pkg/sql/colexec/spillutil/join_spill_test.go @@ -15,4577 +15,577 @@ package spillutil import ( - "bufio" "bytes" "context" - "errors" "fmt" "io" - "math" "os" - "runtime" - "strings" "testing" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" - "github.com/matrixorigin/matrixone/pkg/defines" - "github.com/matrixorigin/matrixone/pkg/fileservice" "github.com/matrixorigin/matrixone/pkg/pb/plan" - "github.com/matrixorigin/matrixone/pkg/sql/colexec" "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" - plan2 "github.com/matrixorigin/matrixone/pkg/sql/plan" "github.com/matrixorigin/matrixone/pkg/testutil" "github.com/matrixorigin/matrixone/pkg/vm/message" "github.com/matrixorigin/matrixone/pkg/vm/process" "github.com/stretchr/testify/require" ) -type boundaryCancelReader struct { - reader *bytes.Reader - boundary int64 - read int64 - cancel func() - canceled bool +func makeTestKeyExpr() []*plan.Expr { + return []*plan.Expr{{ + Typ: plan.Type{Id: int32(types.T_int32), Width: 32}, + Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, + }} } -func (r *boundaryCancelReader) Read(p []byte) (int, error) { - if !r.canceled && r.read >= r.boundary { - r.canceled = true - r.cancel() - } - if !r.canceled { - remaining := r.boundary - r.read - if int64(len(p)) > remaining { - p = p[:remaining] - } +func makeTestEvalKeysFn() func(*batch.Batch) ([]*vector.Vector, error) { + return func(bat *batch.Batch) ([]*vector.Vector, error) { + return bat.Vecs[:1], nil } - n, err := r.reader.Read(p) - r.read += int64(n) - return n, err } -// scatterImpl retains the former buffered implementation solely as a -// compatibility oracle. It is intentionally test-only: production owns one -// selected batch at a time in SpillEngine.scatterBatchBounded. -func scatterImpl( +func makeInt32Batch(proc *process.Process, values []int32) *batch.Batch { + bat := batch.NewWithSize(1) + bat.Vecs[0] = testutil.MakeInt32Vector(values, nil, proc.Mp()) + bat.SetRowCount(len(values)) + return bat +} + +func writeBuildFile( proc *process.Process, + name string, bat *batch.Batch, - keyVecs []*vector.Vector, - writers []BucketWriter, - buffers []*batch.Batch, - seed uint64, - bucketBuf *bytes.Buffer, - analyzer process.Analyzer, - reuseHashValues *[]uint64, - reuseBucketRowIds *[][]int32, -) error { - rowCount := bat.RowCount() - if rowCount == 0 { - return nil - } +) *os.File { + return writeBuildRecords(proc, name, bat) +} - var hashValues []uint64 - if reuseHashValues != nil && cap(*reuseHashValues) >= rowCount { - hashValues = (*reuseHashValues)[:rowCount] - } else { - hashValues = make([]uint64, rowCount) - if reuseHashValues != nil { - *reuseHashValues = hashValues - } +func writeBuildRecords( + proc *process.Process, + name string, + batches ...*batch.Batch, +) *os.File { + spillfs, err := proc.GetSpillFileService() + if err != nil { + panic(err) } - ComputeXXHash(keyVecs, hashValues, seed) - - if len(writers) == 0 || len(writers) > SpillNumBuckets || - len(writers)&(len(writers)-1) != 0 { - return process.ErrHashBuildBudgetInvalid + file, err := spillfs.CreateAndRemoveFile(context.Background(), name) + if err != nil { + panic(err) } - var bucketRowIds [][]int32 - if reuseBucketRowIds != nil { - bucketRowIds = *reuseBucketRowIds - if cap(bucketRowIds) < len(writers) { - bucketRowIds = make([][]int32, len(writers)) - *reuseBucketRowIds = bucketRowIds - } else { - bucketRowIds = bucketRowIds[:len(writers)] + for _, bat := range batches { + if _, err := file.Write(marshalTestSpillRecord(bat)); err != nil { + panic(err) } - } else { - bucketRowIds = make([][]int32, len(writers)) - } - var rowIDs []int32 - if len(bucketRowIds) > 0 && cap(bucketRowIds[0]) >= rowCount { - rowIDs = bucketRowIds[0][:rowCount] - } else { - rowIDs = make([]int32, rowCount) } - var counts [SpillNumBuckets]int32 - var offsets [SpillNumBuckets + 1]int32 - if err := classifyRows( - hashValues, - len(writers), - 0, - rowIDs, - counts[:], - offsets[:], - ); err != nil { - return err - } - for i := range bucketRowIds { - bucketRowIds[i] = rowIDs[offsets[i]:offsets[i+1]] + if _, err := file.Seek(0, io.SeekStart); err != nil { + panic(err) } + return file +} - for bucketID, sels := range bucketRowIds { - if len(sels) == 0 || writers[bucketID].Name == "" { - continue - } - buf := buffers[bucketID] - if buf == nil { - buf = batch.NewOffHeapWithSize(len(bat.Vecs)) - for i, vec := range bat.Vecs { - buf.Vecs[i] = vector.NewOffHeapVecWithType(*vec.GetType()) - if err := buf.Vecs[i].PreExtend(8192, proc.Mp()); err != nil { - return err - } - } - buffers[bucketID] = buf - } - for i, vec := range bat.Vecs { - if err := buf.Vecs[i].UnionInt32(vec, sels, proc.Mp()); err != nil { - return err - } - } - buf.SetRowCount(buf.RowCount() + len(sels)) - if buf.RowCount() >= 8192 { - if err := FlushBucketBatch( - proc, - buf, - &writers[bucketID], - bucketBuf, - analyzer, - ); err != nil { - return err - } - buf.CleanOnlyData() - } - } - return nil +func makeCorruptBatchFile(t *testing.T) *os.File { + t.Helper() + file, err := os.CreateTemp(t.TempDir(), "corrupt-spill") + require.NoError(t, err) + rowCount, batchSize := int64(1), int64(1) + var encoded bytes.Buffer + encoded.Write(types.EncodeInt64(&rowCount)) + encoded.Write(types.EncodeInt64(&batchSize)) + encoded.WriteByte(0xff) + _, err = file.Write(encoded.Bytes()) + require.NoError(t, err) + _, err = file.Seek(0, io.SeekStart) + require.NoError(t, err) + return file } func TestTakeSpillBuildPayloadRejectsWrongBudgetRef(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - _, _, err := TakeSpillBuildPayload(proc, nil) require.ErrorContains(t, err, message.ErrSpillBuildPayloadEmpty.Error()) fd, err := os.CreateTemp(t.TempDir(), "wrong-budget-ref") require.NoError(t, err) - t.Cleanup(func() { _ = fd.Close() }) + _, err = fd.Write([]byte{1}) + require.NoError(t, err) + _, err = fd.Seek(0, io.SeekStart) + require.NoError(t, err) releases := 0 file := message.NewSpillFile(fd, 1, 1, func() { releases++ }) - jm := message.NewJoinMap(message.GroupSels{}, nil, nil, nil, nil, proc.Mp()) + jm := message.NewJoinMap( + message.GroupSels{}, nil, nil, nil, nil, proc.Mp(), + ) + jm.SetRowCount(1) jm.IncRef(1) - jmFreed := false - t.Cleanup(func() { - if !jmFreed { - jm.Free() - } - }) require.NoError(t, jm.SetSpillBuildPayload(message.SpillBuildPayload{ Files: []*message.SpillFile{file}, BudgetRef: struct{}{}, })) - _, _, err = TakeSpillBuildPayload(proc, jm) require.ErrorContains(t, err, "missing its producer budget generation") require.Equal(t, 1, releases) - _, err = fd.Stat() - require.Error(t, err) - jm.Free() - jmFreed = true - require.Equal(t, 1, releases) } -func TestTakeSpillBuildPayloadLegacyResolvesConsumerBudget(t *testing.T) { +func TestTakeSpillBuildPayloadRejectsGlobalRowMismatch(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - - fd, err := os.CreateTemp(t.TempDir(), "legacy-build-payload") + budget := process.MustNewHashBuildBudget(1<<20, 1<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + fd, err := os.CreateTemp(t.TempDir(), "payload-row-mismatch") require.NoError(t, err) - t.Cleanup(func() { _ = fd.Close() }) - jm := message.NewJoinMap(message.GroupSels{}, nil, nil, nil, nil, proc.Mp()) + releases := 0 + jm := message.NewJoinMap( + message.GroupSels{}, nil, nil, nil, nil, proc.Mp(), + ) + jm.SetRowCount(2) jm.IncRef(1) - jmFreed := false - t.Cleanup(func() { - if !jmFreed { - jm.Free() - } - }) require.NoError(t, jm.SetSpillBuildPayload(message.SpillBuildPayload{ - LegacyFds: []*os.File{fd}, + Files: []*message.SpillFile{ + message.NewSpillFile(fd, 1, 0, func() { releases++ }), + }, + BudgetRef: generation, })) - - wantBudget, err := proc.GetHashBuildBudget() - require.NoError(t, err) - payload, budget, err := TakeSpillBuildPayload(proc, jm) - require.NoError(t, err) - require.Same(t, fd, payload.LegacyFds[0]) - require.Same(t, wantBudget, budget) - t.Cleanup(func() { _ = payload.Close() }) - require.NoError(t, payload.Close()) - _, err = fd.Stat() - require.Error(t, err) - + _, _, err = TakeSpillBuildPayload(proc, jm) + require.ErrorContains(t, err, "row count") + require.Equal(t, 1, releases) jm.Free() - jmFreed = true -} - -func TestComputeXXHash(t *testing.T) { - mp := mpool.MustNewZero() - ComputeXXHash(nil, nil, 0) - - vec := testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, mp) - defer vec.Free(mp) - hashValues := make([]uint64, 3) - ComputeXXHash([]*vector.Vector{vec}, hashValues, 0) - require.NotEqual(t, uint64(0), hashValues[0]) - - constVec, err := vector.NewConstFixed(types.T_int32.ToType(), int32(7), 3, mp) - require.NoError(t, err) - defer constVec.Free(mp) - ComputeXXHash([]*vector.Vector{constVec}, hashValues, 1) - require.Equal(t, hashValues[0], hashValues[2]) - - shortVec := testutil.MakeInt32Vector([]int32{9}, nil, mp) - defer shortVec.Free(mp) - ComputeXXHash([]*vector.Vector{shortVec}, hashValues, 2) - - constNull := vector.NewConstNull(types.T_int32.ToType(), 3, mp) - defer constNull.Free(mp) - nullHashes := make([]uint64, 3) - ComputeXXHash([]*vector.Vector{constNull}, nullHashes, 7) - require.Equal(t, hashCombine(uint64(7), uint64(0)), nullHashes[0]) - require.Equal(t, nullHashes[0], nullHashes[2]) } func TestClassifyRowsConservesRows(t *testing.T) { hashes := make([]uint64, 257) for i := range hashes { - // Include skew, an empty bucket, and rows that differ only at the - // re-spill bit offset. hashes[i] = uint64(i%7) | (uint64(i&3) << 5) } rowIDs := make([]int32, len(hashes)) counts := make([]int32, SpillNumBuckets) offsets := make([]int32, SpillNumBuckets+1) - require.NoError(t, classifyRows(hashes, SpillNumBuckets, 0, rowIDs, counts, offsets)) - require.Equal(t, int32(len(hashes)), offsets[SpillNumBuckets]) - seen := make([]bool, len(hashes)) - for bucket := 0; bucket < SpillNumBuckets; bucket++ { - for _, rowID := range rowIDs[offsets[bucket]:offsets[bucket+1]] { + for _, shift := range []uint64{0, 5} { + require.NoError(t, classifyRows( + hashes, + SpillNumBuckets, + shift, + rowIDs, + counts, + offsets, + )) + require.Equal(t, int32(len(hashes)), offsets[SpillNumBuckets]) + seen := make([]bool, len(hashes)) + for _, rowID := range rowIDs { row := int(rowID) require.GreaterOrEqual(t, row, 0) require.Less(t, row, len(hashes)) require.False(t, seen[row]) seen[row] = true - require.Equal(t, bucket, int(hashes[row]&(SpillNumBuckets-1))) } } - for _, ok := range seen { - require.True(t, ok) - } - - // Re-spill consumes the next five hash bits without changing the row - // conservation invariant. - require.NoError(t, classifyRows(hashes, SpillNumBuckets, 5, rowIDs, counts, offsets)) - require.Equal(t, int32(len(hashes)), offsets[SpillNumBuckets]) -} - -func TestClassifyRowsRejectsNonProductionFanout(t *testing.T) { - const fanout = SpillNumBuckets * 2 - err := classifyRows( + require.ErrorIs(t, classifyRows( []uint64{0}, - fanout, + SpillNumBuckets*2, 0, make([]int32, 1), - make([]int32, fanout), - make([]int32, fanout+1), - ) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) -} - -func legacyClassifyRows(hashes []uint64, rowIDs []int32) { - pos := 0 - for bucket := uint64(0); bucket < SpillNumBuckets; bucket++ { - for row, hash := range hashes { - if hash&(SpillNumBuckets-1) == bucket { - rowIDs[pos] = int32(row) - pos++ - } - } - } -} - -func BenchmarkClassifyRows(b *testing.B) { - hashes := make([]uint64, 8192) - for i := range hashes { - hashes[i] = uint64(i*2654435761) ^ uint64(i>>3) - } - rowIDs := make([]int32, len(hashes)) - counts := make([]int32, SpillNumBuckets) - offsets := make([]int32, SpillNumBuckets+1) - b.Run("counts_prefix_rowids", func(b *testing.B) { - b.ReportAllocs() - b.ResetTimer() - for i := 0; i < b.N; i++ { - if err := classifyRows(hashes, SpillNumBuckets, 0, rowIDs, counts, offsets); err != nil { - b.Fatal(err) - } - } - }) - b.Run("legacy_bucket_scan", func(b *testing.B) { - b.ReportAllocs() - b.ResetTimer() - for i := 0; i < b.N; i++ { - legacyClassifyRows(hashes, rowIDs) - } - }) + make([]int32, SpillNumBuckets*2), + make([]int32, SpillNumBuckets*2+1), + ), process.ErrHashBuildBudgetInvalid) } -func TestBucketWriterAccountedHandOffSeekFailureRetainsOwnership(t *testing.T) { +func TestAccountedBucketReaderRoundTripAndCorruption(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - budget, err := process.NewHashBuildBudget(1<<20, 1<<20) - require.NoError(t, err) + budget := process.MustNewHashBuildBudget(8<<20, 8<<20) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - fdToken, err := generation.ReserveSpillFD(1) - require.NoError(t, err) - fd, err := os.CreateTemp(t.TempDir(), "closed-spill") + source := makeInt32Batch(proc, []int32{1, 2, 3}) + defer source.Clean(proc.Mp()) + registry, err := mpool.NewAllocationAccountRegistry(1, 1<<20) require.NoError(t, err) - require.NoError(t, fd.Close()) - w := BucketWriter{Fd: fd, fdReservation: fdToken} - file, err := w.handOffSpillFile() - require.Error(t, err) - require.Nil(t, file) - require.Same(t, fd, w.Fd, "failed rewind must retain file ownership") - require.Equal(t, uint64(1), generation.SpillFDUsed()) - w.Close() - require.Zero(t, generation.SpillFDUsed()) -} - -func TestFlushBucketBatchAndReadRoundtrip(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() + account, err := registry.OpenWithController(8<<20, generation) require.NoError(t, err) - - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_rt") + allocation, err := NewSpillAllocationAccount( + account, + hashbuild.HashBuildAllocationOwner, + ) require.NoError(t, err) - defer f.Close() - - var buf bytes.Buffer - w := BucketWriter{Name: "test_rt", Fd: f} - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{10, 20, 30}, nil, proc.Mp()) - bat.SetRowCount(3) - require.NoError(t, FlushBucketBatch(proc, nil, &w, &buf, nil)) - err = FlushBucketBatch(proc, bat, &w, &buf, process.NewAnalyzer(0, false, false, "test")) + reader := &BucketReader{ + fd: writeBuildFile(proc, t.Name(), source), + allocation: allocation, + } + decoded, err := newSpillBatch(0, reader.allocation.decoded) require.NoError(t, err) - - fd := w.HandOffFd() - reader := BucketReader{fd: fd} - reuseBat := batch.NewOffHeapWithSize(0) - got, err := reader.ReadBatch(proc, reuseBat) + got, err := reader.ReadBatch(proc, decoded) require.NoError(t, err) require.Equal(t, 3, got.RowCount()) + require.Equal(t, []int32{1, 2, 3}, vector.MustFixedColNoTypeCheck[int32](got.Vecs[0])) + got.Clean(proc.Mp()) + _, err = reader.ReadBatch(proc, decoded) + require.ErrorIs(t, err, io.EOF) reader.Close() -} - -func TestBucketReaderAccountedLifecycle(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget, err := process.NewHashBuildBudget(8<<20, 8<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - bat := makeInt32Batch(proc, []int32{1, 2, 3, 4}) - defer bat.Clean(proc.Mp()) - var buf bytes.Buffer - w := BucketWriter{Name: "test_accounted_reader", Budget: generation} - require.NoError(t, FlushBucketBatch(proc, bat, &w, &buf, nil)) - file, err := w.handOffSpillFile() - require.NoError(t, err) - require.NotNil(t, file) - require.Positive(t, generation.SpillDiskUsed()) - require.Equal(t, uint64(1), generation.SpillFDUsed()) - - reader := BucketReader{} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForSpillFile(file) - reuseBat := batch.NewOffHeapWithSize(0) - got, err := reader.ReadBatch(proc, reuseBat) + require.Zero(t, account.Snapshot().Used) + _, _, err = registry.CompleteTerminal(account) require.NoError(t, err) - require.Equal(t, 4, got.RowCount()) - require.Positive(t, generation.Used()) - - reuseBat.Clean(proc.Mp()) - reader.Close() require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) -} -func TestReconcileReadReservation(t *testing.T) { - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) - generation, err := budget.OpenGeneration(1) + corruptState := newTestSpillAllocationAccount(t, 8<<20, 16) + corrupt := &BucketReader{ + fd: makeCorruptBatchFile(t), + allocation: corruptState.allocation, + } + bad, err := newSpillBatch(0, corruptState.allocation.decoded) require.NoError(t, err) - - t.Run("shrink", func(t *testing.T) { - token, err := generation.Reserve(1024) - require.NoError(t, err) - require.NoError(t, reconcileReadReservation(token, 256)) - require.Equal(t, uint64(256), generation.Used()) - require.True(t, token.Release()) - require.Zero(t, generation.Used()) - }) - - t.Run("underestimated-retained-bytes", func(t *testing.T) { - token, err := generation.Reserve(256) - require.NoError(t, err) - require.ErrorIs(t, reconcileReadReservation(token, 257), process.ErrHashBuildBudgetInvalid) - // Failed upward reconciliation keeps the original token live so both - // reader cleanup paths can release the complete reservation exactly once. - require.Equal(t, uint64(256), generation.Used()) - require.True(t, token.Release()) - require.Zero(t, generation.Used()) - }) + _, err = corrupt.ReadBatch(proc, bad) + require.Error(t, err) + corrupt.Close() + finalizeTestSpillAllocationAccount(t, corruptState) } -func TestPredictMergedRetainedBytesMatchesUnionBatch(t *testing.T) { +func TestBucketReaderRejectsSchemaChangeBeforeMerge(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - tests := []struct { - name string - dst func() (*batch.Batch, *batch.Batch) - }{ - { - name: "fixed-and-varlen-multi-column", - dst: func() (*batch.Batch, *batch.Batch) { - dst := batch.NewWithSize(2) - dst.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp()) - dst.Vecs[1] = testutil.MakeVarcharVector([]string{"left", "side"}, nil, proc.Mp()) - dst.SetRowCount(2) - src := batch.NewWithSize(2) - src.Vecs[0] = testutil.MakeInt32Vector([]int32{3, 4, 5}, nil, proc.Mp()) - src.Vecs[1] = testutil.MakeVarcharVector([]string{"right", "hand", "rows"}, nil, proc.Mp()) - src.SetRowCount(3) - return dst, src - }, - }, - { - name: "const-fixed", - dst: func() (*batch.Batch, *batch.Batch) { - dst := makeInt32Batch(proc, []int32{1, 2}) - src := batch.NewWithSize(1) - var err error - src.Vecs[0], err = vector.NewConstFixed(types.T_int32.ToType(), int32(9), 3, proc.Mp()) - require.NoError(t, err) - src.SetRowCount(3) - return dst, src - }, - }, - { - name: "const-inline-varlen", - dst: func() (*batch.Batch, *batch.Batch) { - dst := batch.NewWithSize(1) - dst.Vecs[0] = testutil.MakeVarcharVector([]string{"left", "side"}, nil, proc.Mp()) - dst.SetRowCount(2) - src := batch.NewWithSize(1) - var err error - src.Vecs[0], err = vector.NewConstBytes(types.T_varchar.ToType(), []byte("inline"), 3, proc.Mp()) - require.NoError(t, err) - src.SetRowCount(3) - return dst, src - }, - }, - { - name: "const-non-inline", - dst: func() (*batch.Batch, *batch.Batch) { - dst := batch.NewWithSize(1) - dst.Vecs[0] = testutil.MakeVarcharVector([]string{"left", "side"}, nil, proc.Mp()) - dst.SetRowCount(2) - src := batch.NewWithSize(1) - var err error - src.Vecs[0], err = vector.NewConstBytes(types.T_varchar.ToType(), []byte("a sufficiently long constant value"), 3, proc.Mp()) - require.NoError(t, err) - src.SetRowCount(3) - return dst, src - }, - }, - { - name: "const-null", - dst: func() (*batch.Batch, *batch.Batch) { - dst := batch.NewWithSize(1) - dst.Vecs[0] = testutil.MakeVarcharVector([]string{"left", "side"}, nil, proc.Mp()) - dst.SetRowCount(2) - src := batch.NewWithSize(1) - src.Vecs[0] = vector.NewConstNull(types.T_varchar.ToType(), 3, proc.Mp()) - src.SetRowCount(3) - return dst, src - }, - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - dst, src := tc.dst() - defer dst.Clean(proc.Mp()) - defer src.Clean(proc.Mp()) - predicted, ok := predictMergedRetainedBytes(dst, src) - require.True(t, ok) - require.NoError(t, dst.UnionWindow(src, 0, src.RowCount(), proc.Mp())) - actual, ok := batchRetainedBytes(dst) - require.True(t, ok) - require.LessOrEqual(t, actual, predicted) - require.Equal(t, src.RowCount()+2, dst.RowCount()) - }) + textBatch := batch.NewWithSize(1) + textBatch.Vecs[0] = vector.NewVec(types.T_varchar.ToType()) + require.NoError(t, vector.AppendBytes( + textBatch.Vecs[0], []byte("x"), false, proc.Mp(), + )) + textBatch.SetRowCount(1) + defer textBatch.Clean(proc.Mp()) + intBatch := makeInt32Batch(proc, []int32{1}) + defer intBatch.Clean(proc.Mp()) + + state := newTestSpillAllocationAccount(t, 8<<20, 16) + reader := &BucketReader{ + fd: writeBuildRecords( + proc, + t.Name(), + textBatch, + intBatch, + ), + mergeRecords: true, + allocation: state.allocation, } + reuse, err := newSpillBatch(0, state.allocation.decoded) + require.NoError(t, err) + _, err = reader.ReadBatch(proc, reuse) + require.ErrorContains(t, err, "spill batch schema changed") + reader.Close() + require.Zero(t, state.account.Snapshot().Used) + finalizeTestSpillAllocationAccount(t, state) } -func TestPredictMergedRetainedBytesAdmissionBudget(t *testing.T) { +func TestRebuildHashmapBasic(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - - dstVals := make([]int32, 100) - srcVals := make([]int32, 100) - for i := range dstVals { - dstVals[i] = int32(i) - srcVals[i] = int32(i + len(dstVals)) - } - dst := makeInt32Batch(proc, dstVals) - src := makeInt32Batch(proc, srcVals) - defer dst.Clean(proc.Mp()) - defer src.Clean(proc.Mp()) - old, ok := batchRetainedBytes(dst) - require.True(t, ok) - next, ok := batchRetainedBytes(src) - require.True(t, ok) - predicted, ok := predictMergedRetainedBytes(dst, src) - require.True(t, ok) - require.Greater(t, predicted, old+next, "rounded destination growth must be admitted independently") - - reserveAll := func(t *testing.T, cap uint64) { - budget := process.MustNewHashBuildBudget(cap, cap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - o, err := generation.Reserve(old) - require.NoError(t, err) - n, err := generation.Reserve(next) - require.NoError(t, err) - d, err := generation.Reserve(predicted) - require.NoError(t, err) - o.Release() - n.Release() - d.Release() - require.Zero(t, generation.Used()) + values := make([]int32, 100) + for i := range values { + values[i] = int32(i) } - reserveAll(t, old+next+predicted) - - budget := process.MustNewHashBuildBudget(old+next+predicted-1, old+next+predicted-1) - generation, err := budget.OpenGeneration(2) - require.NoError(t, err) - o, err := generation.Reserve(old) + build := makeInt32Batch(proc, values) + defer build.Clean(proc.Mp()) + engine := newExactTestSpillEngine(t, SpillEngineConfig{ + BuildKeyExprs: makeTestKeyExpr(), + NeedsBuildForEmptyProbe: true, + }) + initTestSpillFiles(engine, []*os.File{ + writeBuildFile(proc, t.Name(), build), + }, int64(len(values))) + jm, result, err := engine.RebuildHashmap( + proc, + process.NewAnalyzer(0, false, false, "test"), + ) require.NoError(t, err) - n, err := generation.Reserve(next) + require.Equal(t, BucketReady, result) + require.Equal(t, int64(100), jm.GetRowCount()) + jm.Free() + _, result, err = engine.RebuildHashmap( + proc, + process.NewAnalyzer(0, false, false, "test"), + ) require.NoError(t, err) - _, err = generation.Reserve(predicted) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - // Admission happens before UnionBatch, so the destination remains intact. - require.Equal(t, len(dstVals), dst.RowCount()) - require.Equal(t, dstVals[0], vector.GetFixedAtNoTypeCheck[int32](dst.Vecs[0], 0)) - o.Release() - n.Release() - require.Zero(t, generation.Used()) + require.Equal(t, BucketQueueEmpty, result) + engine.Cleanup(proc) } -func TestBucketWriterAggregatesDiskAccountingPerFile(t *testing.T) { +func TestReSpillConservesBuildAndProbeRows(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - budget, err := process.NewHashBuildBudget(8<<20, 8<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) + values := make([]int32, 5_000) + for i := range values { + values[i] = int32(i) + } + build := makeInt32Batch(proc, values) + probe := makeInt32Batch(proc, values) + defer build.Clean(proc.Mp()) + defer probe.Clean(proc.Mp()) + engine := newExactTestSpillEngine(t, SpillEngineConfig{ + BuildKeyExprs: makeTestKeyExpr(), + SpillThreshold: 500, + NeedsBuildForEmptyProbe: true, + NeedsProbeForEmptyBuild: true, + }) + initTestSpillFiles(engine, []*os.File{ + writeBuildFile(proc, t.Name()+"-build", build), + }, int64(len(values))) + engine.buckets[0].ProbeFd = newTestSpillFile( + writeBuildFile(proc, t.Name()+"-probe", probe), + int64(len(values)), + ) + engine.buckets[0].ProbeRows = int64(len(values)) + engine.probeKeyEval = makeTestEvalKeysFn() + jm, result, err := engine.RebuildHashmap( + proc, + process.NewAnalyzer(0, false, false, "test"), + ) require.NoError(t, err) - bat := makeInt32Batch(proc, []int32{1, 2, 3}) - defer bat.Clean(proc.Mp()) + require.Nil(t, jm) + require.Equal(t, BucketReSpilled, result) + var buildRows, probeRows int64 + for _, bucket := range engine.buckets { + buildRows += bucket.BuildRows + probeRows += bucket.ProbeRows + } + require.Equal(t, int64(len(values)), buildRows) + require.Equal(t, int64(len(values)), probeRows) + engine.Cleanup(proc) +} - w := BucketWriter{Name: "aggregate_disk_token", Budget: generation} - var buf bytes.Buffer - require.NoError(t, FlushBucketBatch(proc, bat, &w, &buf, nil)) - first := w.diskReservation - require.NotNil(t, first) - firstSize := first.Size() - require.NoError(t, FlushBucketBatch(proc, bat, &w, &buf, nil)) - require.Same(t, first, w.diskReservation) - require.Greater(t, w.diskReservation.Size(), firstSize) - require.Equal(t, w.diskReservation.Size(), generation.SpillDiskUsed()) - w.Close() - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) +func TestSpillRejectsCompleteRecordTruncation(t *testing.T) { + for _, test := range []struct { + name string + threshold int64 + }{ + {name: "rebuild", threshold: 1 << 30}, + {name: "re-spill", threshold: 1}, + } { + for _, metadataRows := range []int64{0, 6} { + t.Run(fmt.Sprintf("%s/metadata-%d", test.name, metadataRows), func(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + first := makeInt32Batch(proc, []int32{1, 2, 3}) + defer first.Clean(proc.Mp()) + engine := newExactTestSpillEngine(t, SpillEngineConfig{ + BuildKeyExprs: makeTestKeyExpr(), + SpillThreshold: test.threshold, + NeedsBuildForEmptyProbe: true, + }) + engine.InitFromSpilledFiles([]*message.SpillFile{ + newTestSpillFile( + writeBuildFile(proc, fmt.Sprintf("truncate-%s-%d", test.name, metadataRows), first), + metadataRows, + ), + }) + _, _, err := engine.RebuildHashmap( + proc, + process.NewAnalyzer(0, false, false, "test"), + ) + require.ErrorContains(t, err, "row count") + engine.Cleanup(proc) + }) + } + } } -func TestBucketReaderEOF(t *testing.T) { +func TestSpillRejectsPhysicalTruncationBeforeFirstRecord(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() + baseline := proc.Mp().CurrNB() + first := makeInt32Batch(proc, []int32{1}) + second := makeInt32Batch(proc, []int32{2}) + defer first.Clean(proc.Mp()) + defer second.Clean(proc.Mp()) - spillfs, err := proc.GetSpillFileService() + file := writeBuildRecords(proc, t.Name(), first, second) + info, err := file.Stat() require.NoError(t, err) - - var buf bytes.Buffer - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_eof") + require.NoError(t, file.Truncate(int64(len(marshalTestSpillRecord(first))))) + _, err = file.Seek(0, io.SeekStart) require.NoError(t, err) - w := BucketWriter{Name: "test_eof", Fd: f} - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp()) - bat.SetRowCount(2) - FlushBucketBatch(proc, bat, &w, &buf, nil) - fd := w.HandOffFd() - reader := BucketReader{} - reader.ResetForFd(fd) - reuseBat := batch.NewOffHeapWithSize(0) - got, err := reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - require.Equal(t, 2, got.RowCount()) - _, err = reader.ReadBatch(proc, reuseBat) - require.Equal(t, io.EOF, err) - reader.Close() + engine := newExactTestSpillEngine(t, SpillEngineConfig{ + BuildKeyExprs: makeTestKeyExpr(), + NeedsBuildForEmptyProbe: true, + }) + engine.InitFromSpilledFiles([]*message.SpillFile{ + message.NewSpillFile(file, 2, uint64(info.Size()), nil), + }) + jm, _, err := engine.RebuildHashmap( + proc, + process.NewAnalyzer(0, false, false, "test"), + ) + require.Nil(t, jm) + require.ErrorContains(t, err, "corrupted spill file size") + engine.Cleanup(proc) + require.Zero(t, engine.allocation.account.Snapshot().Used) + require.Equal(t, baseline, proc.Mp().CurrNB()) +} + +func TestProbeRejectsCompleteRecordTruncation(t *testing.T) { + for _, test := range []struct { + name string + values []int32 + metadataRows int64 + rebuildError bool + firstError bool + }{ + {name: "zero metadata", values: []int32{1}, metadataRows: 0, rebuildError: true}, + {name: "row excess", values: []int32{1, 2}, metadataRows: 1, firstError: true}, + {name: "complete record truncation", values: []int32{1}, metadataRows: 2}, + } { + t.Run(test.name, func(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + build := makeInt32Batch(proc, []int32{1}) + probe := makeInt32Batch(proc, test.values) + defer build.Clean(proc.Mp()) + defer probe.Clean(proc.Mp()) + engine := newExactTestSpillEngine(t, SpillEngineConfig{ + BuildKeyExprs: makeTestKeyExpr(), + NeedsBuildForEmptyProbe: true, + }) + engine.InitFromSpilledFiles([]*message.SpillFile{ + newTestSpillFile(writeBuildFile(proc, "probe-build", build), 1), + }) + engine.buckets[0].ProbeFd = newTestSpillFile( + writeBuildFile(proc, "probe-data", probe), + test.metadataRows, + ) + engine.buckets[0].ProbeRows = test.metadataRows + jm, result, err := engine.RebuildHashmap( + proc, + process.NewAnalyzer(0, false, false, "test"), + ) + if test.rebuildError { + require.Nil(t, jm) + require.ErrorContains(t, err, "row count") + engine.Cleanup(proc) + return + } + require.NoError(t, err) + require.Equal(t, BucketReady, result) + jm.Free() + got, err := engine.NextProbeBatch(proc) + if test.firstError { + require.Nil(t, got) + require.ErrorContains(t, err, "row count") + } else { + require.NoError(t, err) + require.Equal(t, 1, got.RowCount()) + _, err = engine.NextProbeBatch(proc) + require.ErrorContains(t, err, "row count") + } + engine.Cleanup(proc) + }) + } } -func TestBucketReaderCorruptedMagic(t *testing.T) { +func TestReSpillRejectsProbeRowMetadata(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - budget := process.MustNewHashBuildBudget(8<<20, 8<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_corrupt") - require.NoError(t, err) - - // Write a valid batch via FlushBucketBatch, then corrupt the magic. - var buf bytes.Buffer - w := BucketWriter{Name: "test_corrupt", Fd: f} - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1}, nil, proc.Mp()) - bat.SetRowCount(1) - err = FlushBucketBatch(proc, bat, &w, &buf, nil) - require.NoError(t, err) - - // Overwrite last 8 bytes (magic) with zeros. - f.Seek(-8, io.SeekEnd) - var zeroMagic uint64 - f.Write(types.EncodeUint64(&zeroMagic)) - f.Seek(0, io.SeekStart) - - reader := BucketReader{} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForFd(f) - reuseBat := batch.NewOffHeapWithSize(0) - _, err = reader.ReadBatch(proc, reuseBat) - require.Error(t, err) - require.Contains(t, err.Error(), "corrupted") - require.Equal(t, uint64(64<<10), generation.Used(), "failed read must release its decoded-batch lease") - reader.Close() - require.Zero(t, generation.Used()) - f.Close() -} - -func TestBucketReaderTruncatedMagic(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := makeInt32Batch(proc, []int32{1}) - fd := writeBuildFile(proc, "test_truncated_magic", bat) - bat.Clean(proc.Mp()) - info, err := fd.Stat() - require.NoError(t, err) - require.NoError(t, fd.Truncate(info.Size()-4)) - _, err = fd.Seek(0, io.SeekStart) - require.NoError(t, err) - - reader := BucketReader{} - reader.ResetForFd(fd) - reuseBat := batch.NewOffHeapWithSize(0) - _, err = reader.ReadBatch(proc, reuseBat) - require.ErrorIs(t, err, io.ErrUnexpectedEOF) - reader.Close() -} - -func TestBucketWriterHandOffFd(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - spillfs, _ := proc.GetSpillFileService() - f, _ := spillfs.CreateAndRemoveFile(context.Background(), "test_handoff") - w := BucketWriter{Fd: f} - fd := w.HandOffFd() - require.NotNil(t, fd) - require.Nil(t, w.Fd) - require.False(t, w.Created()) - fd.Close() - - budget, err := process.NewHashBuildBudget(8<<20, 8<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(2) - require.NoError(t, err) - accounted := BucketWriter{Name: "test_accounted_raw_handoff", Budget: generation} - bat := makeInt32Batch(proc, []int32{1}) - var buf bytes.Buffer - require.NoError(t, FlushBucketBatch(proc, bat, &accounted, &buf, nil)) - bat.Clean(proc.Mp()) - require.Nil(t, accounted.HandOffFd(), "raw handoff must not orphan accounting tokens") - accounted.Close() - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) -} - -func TestMakeBucketWriters(t *testing.T) { - writers := MakeBucketWriters("test") - require.Equal(t, SpillNumBuckets, len(writers)) - for i := range writers { - require.NotEmpty(t, writers[i].Name) - require.Nil(t, writers[i].Fd) - } -} - -type countingMutableFileService struct { - fileservice.MutableFileService - ensureCalls int - closeCalls int -} - -func (s *countingMutableFileService) EnsureDir(ctx context.Context, path string) error { - s.ensureCalls++ - return s.MutableFileService.EnsureDir(ctx, path) -} - -func (s *countingMutableFileService) Close(ctx context.Context) { - s.closeCalls++ - s.MutableFileService.Close(ctx) -} - -func TestSpillEngineSharesFileServiceAcrossWriters(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - local, err := fileservice.Get[fileservice.MutableFileService]( - proc.GetFileService(), - defines.LocalFileServiceName, - ) - require.NoError(t, err) - countingLocal := &countingMutableFileService{MutableFileService: local} - services, err := fileservice.NewFileServices("", countingLocal) - require.NoError(t, err) - proc.SetFileService(services) - - engine := NewSpillEngine(SpillEngineConfig{}) - first := engine.makeBucketWriters("cached_first") - second := engine.makeBucketWriters("cached_second") - require.Same(t, first[0].spillFS, second[0].spillFS) - require.Same(t, &engine.spillFS, first[0].spillFS) - require.Zero(t, countingLocal.ensureCalls, "resolution remains lazy") - - require.NoError(t, writeBucketPayload(proc, []byte("first"), 1, &first[0], nil)) - require.Equal(t, 1, countingLocal.ensureCalls) - cached := engine.spillFS.fs - require.NotNil(t, cached) - - require.NoError(t, writeBucketPayload(proc, []byte("second"), 1, &second[1], nil)) - require.Equal(t, 1, countingLocal.ensureCalls, "all engine writers reuse one resolved service") - require.Equal(t, cached, engine.spillFS.fs) - - first[0].Close() - second[1].Close() - - // The service is borrowed. Cleanup releases engine-owned files and memory, - // but must neither close nor invalidate the process-owned service. - borrowed := &countingMutableFileService{MutableFileService: cached} - engine.spillFS.fs = borrowed - engine.Cleanup(proc) - require.Zero(t, borrowed.closeCalls) - file, err := borrowed.CreateAndRemoveFile(proc.Ctx, "after_engine_cleanup") - require.NoError(t, err) - require.NoError(t, file.Close()) - - // A writer constructed outside SpillEngine has no shared cache and keeps - // the historical process lookup fallback. - direct := BucketWriter{Name: "direct_writer_fallback"} - require.NoError(t, writeBucketPayload(proc, []byte("direct"), 1, &direct, nil)) - require.Equal(t, 2, countingLocal.ensureCalls) - direct.Close() -} - -func TestScatterProbeTableRejectsRecursiveMarker(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - marker := batch.NewWithSize(0) - marker.SetRowCount(1) - marker.SetLast() - engine := NewSpillEngine(SpillEngineConfig{}) - engine.InitFromSpilledMap([]*os.File{nil}) - called := false - err := engine.ScatterProbeTable( - proc, - func() (*batch.Batch, error) { - if called { - return nil, nil - } - called = true - return marker, nil - }, - nil, - func(*batch.Batch) ([]*vector.Vector, error) { - t.Fatal("recursive marker must not be evaluated as data") - return nil, nil - }, - ) - require.Error(t, err) - require.Contains(t, err.Error(), "recursive input") - engine.Cleanup(proc) -} - -func TestScatterProbeTableErrors(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - wantErr := errors.New("scatter probe failure") - - engine := NewSpillEngine(SpillEngineConfig{}) - engine.InitFromSpilledMap([]*os.File{nil}) - err := engine.ScatterProbeTable(proc, - func() (*batch.Batch, error) { return nil, wantErr }, nil, - func(*batch.Batch) ([]*vector.Vector, error) { return nil, nil }) - require.ErrorIs(t, err, wantErr) - engine.Cleanup(proc) - - bat := makeInt32Batch(proc, []int32{1}) - engine = NewSpillEngine(SpillEngineConfig{}) - engine.InitFromSpilledMap([]*os.File{nil}) - err = engine.ScatterProbeTable(proc, - func() (*batch.Batch, error) { return bat, nil }, nil, - func(*batch.Batch) ([]*vector.Vector, error) { return nil, wantErr }) - require.ErrorIs(t, err, wantErr) - engine.Cleanup(proc) - bat.Clean(proc.Mp()) -} - -func TestReusableBufferPool(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - pool := ReusableBufferPool{} - bufs := pool.Acquire(SpillNumBuckets) - require.Equal(t, SpillNumBuckets, len(bufs)) - for i := range bufs { - require.Nil(t, bufs[i]) - } - pool.Release(proc) -} - -func TestBucketReaderCancellationStopsBeforeMergingNextRecord(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - ctx, cancel := context.WithCancelCause(proc.Ctx) - process.ReplacePipelineCtx(proc, ctx, cancel) - - first := makeInt32Batch(proc, []int32{1, 2}) - second := makeInt32Batch(proc, []int32{3, 4}) - var encoded bytes.Buffer - require.NoError(t, marshalSpillRecord(first, &encoded)) - firstRecord := bytes.Clone(encoded.Bytes()) - require.NoError(t, marshalSpillRecord(second, &encoded)) - stream := append(firstRecord, encoded.Bytes()...) - first.Clean(proc.Mp()) - second.Clean(proc.Mp()) - - source := &boundaryCancelReader{ - reader: bytes.NewReader(stream), - boundary: int64(len(firstRecord)), - cancel: func() { proc.Cancel(context.Canceled) }, - } - fd, err := os.CreateTemp(t.TempDir(), "bucket-reader-cancel") - require.NoError(t, err) - reader := BucketReader{ - fd: fd, - reader: bufio.NewReaderSize(source, 16), - mergeRecords: true, - } - reuseBat := batch.NewOffHeapWithSize(0) - - got, err := reader.ReadBatch(proc, reuseBat) - require.ErrorIs(t, err, context.Canceled) - require.Nil(t, got) - require.True(t, source.canceled) - require.Equal(t, int64(len(firstRecord)+16), source.read, - "reader may inspect the next header but must not decode its payload after cancellation") - require.Zero(t, reuseBat.RowCount()) - - reader.Close() - reuseBat.Clean(proc.Mp()) - require.Zero(t, proc.Mp().CurrNB()) -} - -func TestBucketReaderEmptyFile(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - r := BucketReader{} - r.ResetForFd(nil) - reuseBat := batch.NewOffHeapWithSize(0) - _, err := r.ReadBatch(proc, reuseBat) - require.Equal(t, io.EOF, err) -} - -func TestLazySpillFileCreation(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - writers := MakeBucketWriters("test_lazy") - for i := range writers { - require.Nil(t, writers[i].Fd, "all writers should start with nil Fd") - } - - // Write a batch that will populate some buckets - var buf bytes.Buffer - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3, 4, 5}, nil, proc.Mp()) - bat.SetRowCount(5) - buffers := make([]*batch.Batch, len(writers)) - err := scatterImpl(proc, bat, bat.Vecs[:1], writers, buffers, 0, &buf, nil, nil, nil) - require.NoError(t, err) - - // Flush remaining buffers — files are created lazily on first write - for i, b := range buffers { - if b != nil && b.RowCount() > 0 { - err := FlushBucketBatch(proc, b, &writers[i], &buf, nil) - require.NoError(t, err) - require.True(t, writers[i].Created(), "writer should have created file on first flush") - } - } - - // Clean up - for i := range writers { - writers[i].Close() - } -} - -func TestReaderRowCountMismatch(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_mismatch") - require.NoError(t, err) - - // Write valid batch then corrupt the row count in the header. - var buf bytes.Buffer - w := BucketWriter{Name: "test", Fd: f} - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp()) - bat.SetRowCount(2) - err = FlushBucketBatch(proc, bat, &w, &buf, nil) - require.NoError(t, err) - - // Overwrite the count (first 8 bytes) with a wrong value. - f.Seek(0, io.SeekStart) - wrongCnt := int64(999) - f.Write(types.EncodeInt64(&wrongCnt)) - f.Seek(0, io.SeekStart) - - reader := BucketReader{} - reader.ResetForFd(f) - reuseBat := batch.NewOffHeapWithSize(0) - _, err = reader.ReadBatch(proc, reuseBat) - require.Error(t, err) - require.Contains(t, err.Error(), "mismatch") - reader.Close() - f.Close() -} - -func TestScatterBatchDistribution(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - require.NoError(t, scatterImpl(proc, batch.NewWithSize(0), nil, nil, nil, 0, nil, nil, nil, nil)) - - writers := MakeBucketWriters("test_dist") - buffers := make([]*batch.Batch, len(writers)) - var buf bytes.Buffer - - nRows := SpillNumBuckets * 100 - vals := make([]int32, nRows) - for i := range vals { - vals[i] = int32(i) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector(vals, nil, proc.Mp()) - bat.SetRowCount(nRows) - - err := scatterImpl(proc, bat, bat.Vecs[:1], writers, buffers, 0, &buf, nil, nil, nil) - require.NoError(t, err) - - // Most buckets should have data with enough rows. - nonEmpty := 0 - for _, b := range buffers { - if b != nil && b.RowCount() > 0 { - nonEmpty++ - } - } - require.Greater(t, nonEmpty, SpillNumBuckets/2, "at least half the buckets should have data") - - // Total rows should be preserved. - totalRows := 0 - for i, b := range buffers { - if b != nil { - FlushBucketBatch(proc, b, &writers[i], &buf, nil) - totalRows += b.RowCount() - } - } - require.Equal(t, nRows, totalRows) -} - -func TestBucketReaderPartialRead(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_partial") - require.NoError(t, err) - - // Write incomplete data (only a count, no batch body). - cnt := int64(5) - f.Write(types.EncodeInt64(&cnt)) - f.Seek(0, io.SeekStart) - - reader := BucketReader{} - reader.ResetForFd(f) - reuseBat := batch.NewOffHeapWithSize(0) - _, err = reader.ReadBatch(proc, reuseBat) - require.Error(t, err) - reader.Close() - f.Close() -} - -func TestBucketReaderDoubleClose(t *testing.T) { - r := BucketReader{} - r.Close() - r.Close() // should not panic -} - -func TestComputeXXHashWithNulls(t *testing.T) { - mp := mpool.MustNewZero() - vec := testutil.MakeInt32Vector([]int32{1, 2, 3}, []uint64{1}, mp) // null at index 1 - hashValues := make([]uint64, 3) - ComputeXXHash([]*vector.Vector{vec}, hashValues, 0) - require.NotEqual(t, uint64(0), hashValues[0]) - require.NotEqual(t, uint64(0), hashValues[2]) -} - -func TestComputeXXHashMultipleColumns(t *testing.T) { - mp := mpool.MustNewZero() - vec1 := testutil.MakeInt32Vector([]int32{1, 1, 2}, nil, mp) - vec2 := testutil.MakeVarcharVector([]string{"a", "b", "a"}, nil, mp) - hashValues := make([]uint64, 3) - ComputeXXHash([]*vector.Vector{vec1, vec2}, hashValues, 0) - // Same (col1, col2) pairs should hash differently. - require.NotEqual(t, hashValues[0], hashValues[1]) - require.NotEqual(t, hashValues[0], hashValues[2]) -} - -func TestHandOffFdSeeksToStart(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_seek") - require.NoError(t, err) - - var buf bytes.Buffer - w := BucketWriter{Name: "test_seek", Fd: f} - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{10, 20, 30}, nil, proc.Mp()) - bat.SetRowCount(3) - err = FlushBucketBatch(proc, bat, &w, &buf, nil) - require.NoError(t, err) - - // Position should be past data. - pos, _ := w.Fd.Seek(0, io.SeekCurrent) - require.Greater(t, pos, int64(0)) - - fd := w.HandOffFd() - require.NotNil(t, fd) - pos, _ = fd.Seek(0, io.SeekCurrent) - require.Equal(t, int64(0), pos, "HandOffFd must seek to start") - fd.Close() -} - -func TestComputeXXHashMultipleTypes(t *testing.T) { - mp := mpool.MustNewZero() - tests := []struct { - name string - vec *vector.Vector - }{ - {"int8", testutil.MakeInt8Vector([]int8{1, 2, 3}, nil, mp)}, - {"int16", testutil.MakeInt16Vector([]int16{100, 200, 300}, nil, mp)}, - {"int64", testutil.MakeInt64Vector([]int64{1000, 2000, 3000}, nil, mp)}, - {"uint32", testutil.MakeUint32Vector([]uint32{10, 20, 30}, nil, mp)}, - {"float32", testutil.MakeFloat32Vector([]float32{1.1, 2.2, 3.3}, nil, mp)}, - {"float64", testutil.MakeFloat64Vector([]float64{10.1, 20.2, 30.3}, nil, mp)}, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - hashValues := make([]uint64, 3) - ComputeXXHash([]*vector.Vector{tt.vec}, hashValues, 0) - require.NotEqual(t, uint64(0), hashValues[0]) - }) - } -} - -func TestScatterBatchLargeData(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - writers := MakeBucketWriters("test_large") - buffers := make([]*batch.Batch, len(writers)) - var buf bytes.Buffer - - // Large enough to trigger internal flush (>8192 rows). - size := 10000 - vals := make([]int32, size) - for i := range vals { - vals[i] = int32(i) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector(vals, nil, proc.Mp()) - bat.SetRowCount(size) - - err := scatterImpl(proc, bat, bat.Vecs[:1], writers, buffers, 0, &buf, nil, nil, nil) - require.NoError(t, err) - - // Verify total rows preserved. - totalRows := 0 - for i, b := range buffers { - if b != nil { - FlushBucketBatch(proc, b, &writers[i], &buf, nil) - totalRows += b.RowCount() - } - } - require.Equal(t, size, totalRows) -} - -func TestResetForFdReusesReader(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - - var buf bytes.Buffer - f1, _ := spillfs.CreateAndRemoveFile(context.Background(), "test_reuse_1") - w1 := BucketWriter{Name: "test1", Fd: f1} - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()) - bat.SetRowCount(3) - FlushBucketBatch(proc, bat, &w1, &buf, nil) - fd1 := w1.HandOffFd() - - f2, _ := spillfs.CreateAndRemoveFile(context.Background(), "test_reuse_2") - w2 := BucketWriter{Name: "test2", Fd: f2} - FlushBucketBatch(proc, bat, &w2, &buf, nil) - fd2 := w2.HandOffFd() - - r := BucketReader{} - r.ResetForFd(fd1) - require.NotNil(t, r.fd) - - // Second ResetForFd reuses internal state. - r.ResetForFd(fd2) - require.NotNil(t, r.fd) - - r.Close() - fd1.Close() -} - -func TestFlushBucketBatchMultipleCalls(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_multi_flush") - require.NoError(t, err) - - var buf bytes.Buffer - w := BucketWriter{Name: "test_multi_flush", Fd: f} - totalRows := 0 - for i := 0; i < 10; i++ { - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{int32(i)}, nil, proc.Mp()) - bat.SetRowCount(1) - err := FlushBucketBatch(proc, bat, &w, &buf, nil) - require.NoError(t, err) - totalRows++ - } - - fd := w.HandOffFd() - reader := BucketReader{} - reader.ResetForFd(fd) - reuseBat := batch.NewOffHeapWithSize(0) - readRows := 0 - for { - got, err := reader.ReadBatch(proc, reuseBat) - if err == io.EOF { - break - } - require.NoError(t, err) - readRows += got.RowCount() - } - require.Equal(t, totalRows, readRows) - reader.Close() -} - -func TestHashDistribution(t *testing.T) { - mp := mpool.MustNewZero() - vec := testutil.MakeInt32Vector([]int32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20}, nil, mp) - hashValues := make([]uint64, 20) - ComputeXXHash([]*vector.Vector{vec}, hashValues, 0) - - bucketCounts := make([]int, SpillNumBuckets) - for _, h := range hashValues { - bucketCounts[h&(SpillNumBuckets-1)]++ - } - nonEmpty := 0 - for _, c := range bucketCounts { - if c > 0 { - nonEmpty++ - } - } - require.Greater(t, nonEmpty, 1, "hashes must distribute across multiple buckets") -} - -func TestSpillFileCleanup(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - - file, err := spillfs.CreateFile(context.Background(), "test_cleanup") - require.NoError(t, err) - - var buf bytes.Buffer - w := BucketWriter{Name: "test_cleanup", Fd: file} - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1}, nil, proc.Mp()) - bat.SetRowCount(1) - err = FlushBucketBatch(proc, bat, &w, &buf, nil) - require.NoError(t, err) - file.Close() - - // File should still exist (it was CreateFile, not CreateAndRemoveFile). - f2, err := spillfs.OpenFile(context.Background(), "test_cleanup") - require.NoError(t, err) - f2.Close() - - spillfs.RemoveFile(context.Background(), "test_cleanup") -} - -func TestFileWriteError(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - file, err := spillfs.CreateFile(context.Background(), "test_error") - require.NoError(t, err) - file.Close() // close before write - - var buf bytes.Buffer - w := BucketWriter{Name: "test_error", Fd: file} - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1}, nil, proc.Mp()) - bat.SetRowCount(1) - err = FlushBucketBatch(proc, bat, &w, &buf, nil) - require.Error(t, err) - - spillfs.RemoveFile(context.Background(), "test_error") -} - -func TestScatterBatchWithNulls(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - writers := MakeBucketWriters("test_null") - buffers := make([]*batch.Batch, len(writers)) - var buf bytes.Buffer - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3, 4}, []uint64{1}, proc.Mp()) // null at index 1 - bat.SetRowCount(4) - - err := scatterImpl(proc, bat, bat.Vecs[:1], writers, buffers, 0, &buf, nil, nil, nil) - require.NoError(t, err) - // Should not panic with nulls. -} - -func TestReaderBatchReuse(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget := process.MustNewHashBuildBudget(8<<20, 8<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - var buf bytes.Buffer - f, _ := spillfs.CreateAndRemoveFile(context.Background(), "test_reuse_read") - w := BucketWriter{Name: "test", Fd: f} - - // Write batches with different sizes. - for _, size := range []int{5, 2} { - vals := make([]int32, size) - for i := range vals { - vals[i] = int32(i) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector(vals, nil, proc.Mp()) - bat.SetRowCount(size) - FlushBucketBatch(proc, bat, &w, &buf, nil) - } - - fd := w.HandOffFd() - reader := BucketReader{} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForFd(fd) - reuseBat := batch.NewOffHeapWithSize(0) - before := generation.Snapshot() - - // Read with the same reuseBat. The second record admits the one-time - // old-plus-new transition and then keeps the bounded high-water lease. - got, err := reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - require.Equal(t, 5, got.RowCount()) - afterFirst := generation.Snapshot() - require.Equal(t, before.ReserveCount+1, afterFirst.ReserveCount) - require.Equal(t, before.ReconcileCount, afterFirst.ReconcileCount) - - got, err = reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - require.Equal(t, 2, got.RowCount()) - afterSecond := generation.Snapshot() - require.Equal(t, afterFirst.ReserveCount+1, afterSecond.ReserveCount) - require.Equal(t, afterFirst.ReconcileCount, afterSecond.ReconcileCount) - require.Equal(t, afterFirst.ReleaseCount, afterSecond.ReleaseCount) - - reuseBat.Clean(proc.Mp()) - reader.Close() - require.Zero(t, generation.Used()) -} - -func TestReaderBatchLeaseGrowsForLargerRecord(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget := process.MustNewHashBuildBudget(8<<20, 8<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_read_lease_grow") - require.NoError(t, err) - var buf bytes.Buffer - w := BucketWriter{Name: "test_read_lease_grow", Fd: f} - for _, size := range []int{2, 1_000} { - vals := make([]int32, size) - bat := makeInt32Batch(proc, vals) - require.NoError(t, FlushBucketBatch(proc, bat, &w, &buf, nil)) - bat.Clean(proc.Mp()) - } - - reader := BucketReader{} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForFd(w.HandOffFd()) - reuseBat := batch.NewOffHeapWithSize(0) - before := generation.Snapshot() - - got, err := reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - require.Equal(t, 2, got.RowCount()) - afterFirst := generation.Snapshot() - require.Equal(t, before.ReserveCount+1, afterFirst.ReserveCount) - - got, err = reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - require.Equal(t, 1_000, got.RowCount()) - afterSecond := generation.Snapshot() - require.Equal(t, afterFirst.ReserveCount+1, afterSecond.ReserveCount, "larger record should grow the existing lease once") - require.Greater(t, afterSecond.Used, afterFirst.Used) - require.Equal(t, afterFirst.ReconcileCount, afterSecond.ReconcileCount) - - reuseBat.Clean(proc.Mp()) - reader.Close() - require.Zero(t, generation.Used()) -} - -func TestDecodedBatchReusePeakCoversMpoolGrowth(t *testing.T) { - const oldCapacity = int64(4 << 20) - required := oldCapacity + 1 - newCapacity, ok := mpool.GrowCapacity(oldCapacity, required) - require.True(t, ok) - projected, ok := decodedBatchProjectedBytes(uint64(required), 1, 1) - require.True(t, ok) - peak, ok := decodedBatchReusePeakBytes(uint64(oldCapacity), projected, uint64(required)) - require.True(t, ok) - require.GreaterOrEqual(t, peak, uint64(oldCapacity+newCapacity)) -} - -func TestReaderBatchReuseFallsBackBeforeTransientGrowth(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_read_reuse_transient") - require.NoError(t, err) - var buf bytes.Buffer - w := BucketWriter{Name: "test_read_reuse_transient", Fd: f} - for _, width := range []int{1_024, 1_025} { - values := make([]string, 4_096) - for i := range values { - values[i] = strings.Repeat("x", width) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeVarcharVector(values, nil, proc.Mp()) - bat.SetRowCount(len(values)) - require.NoError(t, FlushBucketBatch(proc, bat, &w, &buf, nil)) - bat.Clean(proc.Mp()) - } - - fd := w.HandOffFd() - var projected [2]uint64 - for i := range projected { - var header [16]byte - _, err = io.ReadFull(fd, header[:]) - require.NoError(t, err) - rows := types.DecodeInt64(header[:8]) - payload := types.DecodeInt64(header[8:]) - require.Positive(t, rows) - require.Positive(t, payload) - var ok bool - projected[i], ok = decodedBatchProjectedBytes(uint64(payload), rows, 1) - require.True(t, ok) - _, err = fd.Seek(payload+8, io.SeekCurrent) - require.NoError(t, err) - } - require.Greater(t, projected[1], projected[0]) - _, err = fd.Seek(0, io.SeekStart) - require.NoError(t, err) - - // The cap fits the retained record plus one logical new payload, but not the - // allocator's 1.25x replacement capacity. The reader must release the old - // lease and decode the second record fresh. - cap := uint64(64<<10) + projected[0] + projected[1] - budget := process.MustNewHashBuildBudget(cap, cap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - reader := BucketReader{} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForFd(fd) - reuseBat := batch.NewOffHeapWithSize(0) - baseline := uint64(proc.Mp().CurrNB()) - epoch := proc.Mp().StartResourcePeakEpoch() - require.NotNil(t, epoch) - - got, err := reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - require.Equal(t, 4_096, got.RowCount()) - got, err = reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - require.Equal(t, 4_096, got.RowCount()) - - peak, ok := proc.Mp().EndResourcePeakEpoch(epoch) - require.True(t, ok) - require.LessOrEqual(t, peak, baseline+projected[1]) - require.LessOrEqual(t, generation.Peak(), cap) - - reuseBat.Clean(proc.Mp()) - reader.Close() - require.Zero(t, generation.Used()) -} - -func TestReaderBatchLeaseUsesSinglePayloadEstimate(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget := process.MustNewHashBuildBudget(64<<20, 64<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_read_lease_trim") - require.NoError(t, err) - values := make([]string, 4_096) - for i := range values { - values[i] = strings.Repeat("x", 1_024) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeVarcharVector(values, nil, proc.Mp()) - bat.SetRowCount(len(values)) - var buf bytes.Buffer - w := BucketWriter{Name: "test_read_lease_trim", Fd: f} - require.NoError(t, FlushBucketBatch(proc, bat, &w, &buf, nil)) - bat.Clean(proc.Mp()) - - reader := BucketReader{} - require.NoError(t, reader.EnsureBuffer(generation)) - fd := w.HandOffFd() - var header [16]byte - _, err = io.ReadFull(fd, header[:]) - require.NoError(t, err) - payload := types.DecodeInt64(header[8:]) - _, err = fd.Seek(0, io.SeekStart) - require.NoError(t, err) - reader.ResetForFd(fd) - reuseBat := batch.NewOffHeapWithSize(0) - before := generation.Snapshot() - got, err := reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - actual, ok := batchRetainedBytes(got) - require.True(t, ok) - after := generation.Snapshot() - require.Equal(t, before.ReconcileCount, after.ReconcileCount) - require.GreaterOrEqual(t, after.Used, before.Used+actual) - projected, ok := decodedBatchProjectedBytes(uint64(payload), int64(got.RowCount()), int32(len(got.Vecs))) - require.True(t, ok) - require.Equal(t, before.Used+projected, after.Used) - - reuseBat.Clean(proc.Mp()) - reader.Close() - require.Zero(t, generation.Used()) -} - -func TestMarshalSpillRecordPreallocatesExactPayload(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - bat := batch.NewWithSize(1) - var err error - bat.Vecs[0], err = vector.NewConstBytes( - types.T_varchar.ToType(), make([]byte, 4<<20), 64, proc.Mp(), - ) - require.NoError(t, err) - bat.SetRowCount(64) - defer bat.Clean(proc.Mp()) - - buf := bytes.NewBuffer(make([]byte, 0, 1<<20)) - require.NoError(t, marshalSpillRecord(bat, buf)) - size, err := bat.MarshalBinarySize() - require.NoError(t, err) - require.Equal(t, size+24, buf.Cap()) - - small := batch.NewWithSize(1) - small.Vecs[0], err = vector.NewConstBytes( - types.T_varchar.ToType(), make([]byte, 1024), 1, proc.Mp(), - ) - require.NoError(t, err) - small.SetRowCount(1) - defer small.Clean(proc.Mp()) - require.NoError(t, marshalSpillRecord(small, buf), - "a retained large serialization buffer must be reusable for a smaller batch") -} - -func TestReaderBatchLeaseGrowRejectionReleasesToken(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_read_lease_reject") - require.NoError(t, err) - var buf bytes.Buffer - w := BucketWriter{Name: "test_read_lease_reject", Fd: f} - for _, size := range []int{2, 1_000} { - bat := makeInt32Batch(proc, make([]int32, size)) - require.NoError(t, FlushBucketBatch(proc, bat, &w, &buf, nil)) - bat.Clean(proc.Mp()) - } - fd := w.HandOffFd() - - var header [16]byte - _, err = io.ReadFull(fd, header[:]) - require.NoError(t, err) - firstPayload := types.DecodeInt64(header[8:]) - _, err = fd.Seek(firstPayload+8, io.SeekCurrent) - require.NoError(t, err) - _, err = io.ReadFull(fd, header[:]) - require.NoError(t, err) - secondPayload := types.DecodeInt64(header[8:]) - require.Greater(t, secondPayload, firstPayload) - _, err = fd.Seek(0, io.SeekStart) - require.NoError(t, err) - - firstProjected, ok := decodedBatchProjectedBytes(uint64(firstPayload), 2, 1) - require.True(t, ok) - secondProjected, ok := decodedBatchProjectedBytes(uint64(secondPayload), 1_000, 1) - require.True(t, ok) - cap := uint64(64<<10) + firstProjected + (secondProjected-firstProjected)/2 - budget := process.MustNewHashBuildBudget(cap, cap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - reader := BucketReader{} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForFd(fd) - reuseBat := batch.NewOffHeapWithSize(0) - - got, err := reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - require.Equal(t, 2, got.RowCount()) - _, err = reader.ReadBatch(proc, reuseBat) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Equal(t, uint64(64<<10), generation.Used(), "grow rejection must release the existing decoded-batch lease") - - reader.Close() - require.Zero(t, generation.Used()) -} - -func TestReaderBatchClosedLeaseDoesNotRetryAdmission(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_read_closed_lease") - require.NoError(t, err) - var buf bytes.Buffer - w := BucketWriter{Name: "test_read_closed_lease", Fd: f} - for _, size := range []int{2, 1_000} { - bat := makeInt32Batch(proc, make([]int32, size)) - require.NoError(t, FlushBucketBatch(proc, bat, &w, &buf, nil)) - bat.Clean(proc.Mp()) - } - - budget := process.MustNewHashBuildBudget(8<<20, 8<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - reader := BucketReader{} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForFd(w.HandOffFd()) - reuseBat := batch.NewOffHeapWithSize(0) - - got, err := reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - require.Equal(t, 2, got.RowCount()) - before := generation.Snapshot() - generation.Close() - - _, err = reader.ReadBatch(proc, reuseBat) - require.ErrorIs(t, err, process.ErrHashBuildBudgetClosed) - require.NotErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - after := generation.Snapshot() - require.Equal(t, before.RejectCount+1, after.RejectCount, - "a closed lease must not fall through to a second Reserve attempt") - require.Nil(t, reader.batchToken) - require.Zero(t, reader.batchCharge) - require.Zero(t, reuseBat.RowCount()) - - reader.Close() - require.Zero(t, generation.Used()) -} - -func TestScatterWithMultiColumn(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - writers := MakeBucketWriters("test_multi_col") - buffers := make([]*batch.Batch, len(writers)) - var buf bytes.Buffer - - bat := batch.NewWithSize(2) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3, 4, 5}, nil, proc.Mp()) - bat.Vecs[1] = testutil.MakeVarcharVector([]string{"a", "b", "c", "d", "e"}, nil, proc.Mp()) - bat.SetRowCount(5) - - err := scatterImpl(proc, bat, bat.Vecs[:1], writers, buffers, 0, &buf, nil, nil, nil) - require.NoError(t, err) - - // All 5 rows must be distributed across buffers. - totalRows := 0 - hasTwoCols := false - for _, b := range buffers { - if b != nil && b.RowCount() > 0 { - totalRows += b.RowCount() - if len(b.Vecs) == 2 { - hasTwoCols = true - } - } - } - require.Equal(t, 5, totalRows, "all rows must be accounted for in buffers") - require.True(t, hasTwoCols, "buffer batches must preserve column count") -} - -func TestScatterLargeVarchar(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - writers := MakeBucketWriters("test_large_varchar") - buffers := make([]*batch.Batch, len(writers)) - var buf bytes.Buffer - - size := 100 - vals := make([]string, size) - for i := range vals { - vals[i] = fmt.Sprintf("large_string_value_%d_with_padding", i) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeVarcharVector(vals, nil, proc.Mp()) - bat.SetRowCount(size) - - err := scatterImpl(proc, bat, bat.Vecs[:1], writers, buffers, 0, &buf, nil, nil, nil) - require.NoError(t, err) - - totalRows := 0 - for i, b := range buffers { - if b != nil && b.RowCount() > 0 { - FlushBucketBatch(proc, b, &writers[i], &buf, nil) - totalRows += b.RowCount() - } - } - require.Equal(t, size, totalRows) -} - -func TestBucketBufferReuse(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_buf_reuse") - require.NoError(t, err) - - var buf bytes.Buffer - w := BucketWriter{Name: "test", Fd: f} - - // Reuse same writer across multiple batches. - for range 2 { - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp()) - bat.SetRowCount(2) - err := FlushBucketBatch(proc, bat, &w, &buf, nil) - require.NoError(t, err) - } - - fd := w.HandOffFd() - reader := BucketReader{} - reader.ResetForFd(fd) - reuseBat := batch.NewOffHeapWithSize(0) - totalRows := 0 - for { - got, err := reader.ReadBatch(proc, reuseBat) - if err == io.EOF { - break - } - require.NoError(t, err) - totalRows += got.RowCount() - } - require.Equal(t, 4, totalRows) // 2 batches × 2 rows - reader.Close() -} - -func TestReusableBufferPoolWithData(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - pool := ReusableBufferPool{} - bufs := pool.Acquire(3) - - // Populate buffers with data. - for i := range bufs { - bufs[i] = batch.NewWithSize(1) - bufs[i].Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()) - bufs[i].SetRowCount(3) - } - - // Release should clean everything. - pool.Release(proc) -} - -func TestSpillFileFormatMultipleBatches(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_format") - require.NoError(t, err) - - var buf bytes.Buffer - w := BucketWriter{Name: "test", Fd: f} - for i := 0; i < 3; i++ { - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{int32(i * 10), int32(i*10 + 1)}, nil, proc.Mp()) - bat.SetRowCount(2) - err := FlushBucketBatch(proc, bat, &w, &buf, nil) - require.NoError(t, err) - } - - fd := w.HandOffFd() - reader := BucketReader{} - reader.ResetForFd(fd) - reuseBat := batch.NewOffHeapWithSize(0) - batchCount := 0 - totalRows := 0 - for { - got, err := reader.ReadBatch(proc, reuseBat) - if err == io.EOF { - break - } - require.NoError(t, err) - batchCount++ - totalRows += got.RowCount() - } - require.Equal(t, 3, batchCount) - require.Equal(t, 6, totalRows) - reader.Close() -} - -func TestBucketReaderMergesAdjacentAccountedRecords(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget, err := process.NewHashBuildBudget(16<<20, 16<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var buf bytes.Buffer - writer := BucketWriter{Name: "merge_records", Budget: generation} - for _, values := range [][]int32{{1, 2}, {3, 4, 5}} { - bat := makeInt32Batch(proc, values) - require.NoError(t, FlushBucketBatch(proc, bat, &writer, &buf, nil)) - bat.Clean(proc.Mp()) - } - file, err := writer.handOffSpillFile() - require.NoError(t, err) - - reader := BucketReader{mergeRecords: true} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForSpillFile(file) - reuse := batch.NewOffHeapWithSize(0) - got, err := reader.ReadBatch(proc, reuse) - require.NoError(t, err) - require.Equal(t, 5, got.RowCount()) - require.Positive(t, generation.Used()) - _, err = reader.ReadBatch(proc, reuse) - require.ErrorIs(t, err, io.EOF) - reuse.Clean(proc.Mp()) - reader.Close() - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) -} - -func TestBucketReaderMergeRejectsTruncatedTrailingHeader(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget, err := process.NewHashBuildBudget(16<<20, 16<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - bat := makeInt32Batch(proc, []int32{1, 2, 3}) - var payload bytes.Buffer - require.NoError(t, marshalSpillRecord(bat, &payload)) - bat.Clean(proc.Mp()) - - fd, err := os.CreateTemp(t.TempDir(), "truncated-spill") - require.NoError(t, err) - t.Cleanup(func() { _ = fd.Close() }) - _, err = fd.Write(payload.Bytes()) - require.NoError(t, err) - // A clean file boundary has zero bytes left. Any non-empty fragment of the - // next 16-byte frame header is corruption and must not be accepted as EOF. - _, err = fd.Write(types.EncodeInt64(new(int64))) - require.NoError(t, err) - _, err = fd.Seek(0, io.SeekStart) - require.NoError(t, err) - - reader := BucketReader{mergeRecords: true} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForFd(fd) - reuse := batch.NewOffHeapWithSize(0) - got, err := reader.ReadBatch(proc, reuse) - require.Nil(t, got) - require.ErrorIs(t, err, io.ErrUnexpectedEOF) - require.Zero(t, reuse.RowCount()) - require.Nil(t, reader.batchToken) - require.Zero(t, reader.batchCharge) - - reuse.Clean(proc.Mp()) - reader.Close() - require.Zero(t, generation.Used()) -} - -func TestBucketReaderMergeRecordsRespectsBatchBoundary(t *testing.T) { - tests := []struct { - name string - recordRows []int - wantRows []int - }{ - { - name: "two medium records stay separate", - recordRows: []int{5000, 5000}, - wantRows: []int{5000, 5000}, - }, - { - name: "records exactly fill the boundary", - recordRows: []int{8191, 1}, - wantRows: []int{8192}, - }, - { - name: "record crossing the boundary stays separate", - recordRows: []int{8191, 2}, - wantRows: []int{8191, 2}, - }, - { - name: "one oversized source record remains indivisible", - recordRows: []int{9000}, - wantRows: []int{9000}, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget, err := process.NewHashBuildBudget(16<<20, 16<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var buf bytes.Buffer - writer := BucketWriter{Name: "merge_boundary", Budget: generation} - for record, rows := range tt.recordRows { - values := make([]int32, rows) - for row := range values { - values[row] = int32(record*10000 + row) - } - bat := makeInt32Batch(proc, values) - require.NoError(t, FlushBucketBatch(proc, bat, &writer, &buf, nil)) - bat.Clean(proc.Mp()) - } - file, err := writer.handOffSpillFile() - require.NoError(t, err) - - reader := BucketReader{mergeRecords: true} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForSpillFile(file) - reuse := batch.NewOffHeapWithSize(0) - var gotRows []int - for { - got, err := reader.ReadBatch(proc, reuse) - if err == io.EOF { - break - } - require.NoError(t, err) - gotRows = append(gotRows, got.RowCount()) - } - require.Equal(t, tt.wantRows, gotRows) - - reuse.Clean(proc.Mp()) - reader.Close() - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - }) - } -} - -func TestBucketReaderMergeErrorReleasesAllOwnership(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget, err := process.NewHashBuildBudget(1<<20, 1<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - readerToken, err := generation.Reserve(1) - require.NoError(t, err) - sourceToken, err := generation.Reserve(1) - require.NoError(t, err) - extraToken, err := generation.Reserve(1) - require.NoError(t, err) - reader := BucketReader{batchToken: readerToken, batchCharge: 1} - dst := makeInt32Batch(proc, []int32{1}) - src := makeInt32Batch(proc, []int32{2}) - want := errors.New("merge failed") - require.ErrorIs(t, reader.mergeReadError(proc, dst, src, sourceToken, want, extraToken, nil), want) - require.Nil(t, reader.batchToken) - require.Zero(t, reader.batchCharge) - require.Zero(t, generation.Used()) - require.True(t, readerToken.Released()) - require.True(t, sourceToken.Released()) - require.True(t, extraToken.Released()) - - require.ErrorIs(t, reader.mergeReadError(proc, nil, nil, nil, want, nil), want) -} - -func TestSpillEngineInitFromOwnedFilesAndErrorClassification(t *testing.T) { - first, err := os.CreateTemp(t.TempDir(), "owned-build") - require.NoError(t, err) - owned := message.NewSpillFile(first, 7, 11, nil) - engine := NewSpillEngine(SpillEngineConfig{}) - engine.InitFromSpilledFiles([]*message.SpillFile{owned, nil}) - require.Len(t, engine.buckets, 2) - require.Same(t, owned, engine.buckets[0].BuildFd) - require.Equal(t, int64(7), engine.buckets[0].BuildRows) - require.Equal(t, 1, engine.buckets[0].Depth) - require.Nil(t, engine.buckets[1].BuildFd) - require.Zero(t, engine.buckets[1].BuildRows) - - require.False(t, isBudgetAdmission(nil)) - require.False(t, isBudgetAdmission(io.EOF)) - require.True(t, isBudgetAdmission(process.ErrHashBuildBudgetAdmission)) - require.False(t, isBudgetAdmission(process.ErrHashBuildBudgetClosed)) - require.False(t, isBudgetAdmission(&process.HashBuildBudgetError{ - Kind: process.HashBuildBudgetErrorAdmission, - Component: process.HashBuildBudgetComponentSpillDisk, - })) - require.False(t, isBudgetAdmission(&process.HashBuildBudgetError{ - Kind: process.HashBuildBudgetErrorAdmission, - Component: process.HashBuildBudgetComponentSpillFD, - })) - require.Equal(t, - hashbuild.MemoryPressureMinimumUnit, - hashbuild.MemoryPressureReasonOf(noProgressError(nil, 3))) - require.NoError(t, owned.Close()) -} - -func TestSpillSizeHelpersRejectInvalidAndOverflowInputs(t *testing.T) { - require.ErrorIs(t, writeBucketPayload(nil, nil, 0, nil, nil), process.ErrHashBuildBudgetInvalid) - require.NoError(t, marshalSpillRecord(nil, &bytes.Buffer{})) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - var header bytes.Buffer - negativeRows := int64(-1) - zero := int64(0) - header.Write(types.EncodeInt64(&negativeRows)) - header.Write(types.EncodeInt64(&zero)) - reader := BucketReader{reader: bufio.NewReader(&header)} - _, _, _, err := reader.readBatchRecord(proc, batch.NewOffHeapWithSize(0), nil, 0, false) - require.Error(t, err) - - makeHeader := func(batchSize int64) *bufio.Reader { - var data bytes.Buffer - rows := int64(0) - data.Write(types.EncodeInt64(&rows)) - data.Write(types.EncodeInt64(&batchSize)) - data.Write(make([]byte, 12)) - return bufio.NewReader(&data) - } - budgetForHeader, err := process.NewHashBuildBudget(1, 1) - require.NoError(t, err) - headerGeneration, err := budgetForHeader.OpenGeneration(1) - require.NoError(t, err) - reader = BucketReader{reader: makeHeader(math.MaxInt64), budget: headerGeneration} - _, _, _, err = reader.readBatchRecord(proc, batch.NewOffHeapWithSize(0), nil, 0, false) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - - var truncatedPayload bytes.Buffer - truncatedPayload.Write(types.EncodeInt64(&zero)) - one := int64(1) - truncatedPayload.Write(types.EncodeInt64(&one)) - truncatedPayload.Write(make([]byte, 4)) - reader = BucketReader{reader: bufio.NewReader(&truncatedPayload), budget: headerGeneration} - _, _, _, err = reader.readBatchRecord(proc, batch.NewOffHeapWithSize(0), nil, 0, false) - require.Error(t, err) - - var mismatchedRows bytes.Buffer - twelve := int64(12) - mismatchedRows.Write(types.EncodeInt64(&one)) - mismatchedRows.Write(types.EncodeInt64(&twelve)) - mismatchedRows.Write(types.EncodeInt64(&zero)) - mismatchedRows.Write(make([]byte, 4)) - reader = BucketReader{reader: bufio.NewReader(&mismatchedRows), budget: headerGeneration} - _, _, _, err = reader.readBatchRecord(proc, batch.NewOffHeapWithSize(0), nil, 0, false) - require.Error(t, err) - require.Contains(t, err.Error(), "row count mismatch") - - reader = BucketReader{reader: makeHeader(1), budget: headerGeneration} - _, _, _, err = reader.readBatchRecord(proc, batch.NewOffHeapWithSize(0), nil, 0, false) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - headerToken, err := headerGeneration.Reserve(1) - require.NoError(t, err) - reader = BucketReader{reader: makeHeader(1), budget: headerGeneration} - _, returnedToken, _, err := reader.readBatchRecord(proc, batch.NewOffHeapWithSize(0), headerToken, 1, false) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Nil(t, returnedToken) - require.True(t, headerToken.Released()) - - _, ok := addUint64(math.MaxUint64, 1) - require.False(t, ok) - _, ok = mulUint64(math.MaxUint64, 2) - require.False(t, ok) - _, ok = decodedBatchReusePeakBytes(math.MaxUint64, 1, 1) - require.False(t, ok) - _, ok = decodedBatchReusePeakBytes(0, math.MaxUint64, 1) - require.False(t, ok) - _, ok = decodedBatchProjectedBytes(0, -1, 0) - require.False(t, ok) - _, ok = decodedBatchProjectedBytes(0, 0, -1) - require.False(t, ok) - _, ok = decodedBatchProjectedBytes(math.MaxUint64, 0, 0) - require.False(t, ok) - _, ok = decodedBatchProjectedBytes(0, math.MaxInt64, math.MaxInt32) - require.False(t, ok) - _, ok = batchRetainedMetadataBytes(1, math.MaxUint64) - require.False(t, ok) - _, ok = batchRetainedMetadataBytes(math.MaxUint64, 1) - require.False(t, ok) - _, ok = batchPayloadWithAllocationSlack(0, math.MaxUint64) - require.False(t, ok) - _, ok = batchPayloadWithAllocationSlack(math.MaxUint64, 0) - require.False(t, ok) - _, ok = intToUint64(-1) - require.False(t, ok) - _, ok = predictedCapacity(-1, 1) - require.False(t, ok) - _, ok = predictedCapacity(1, math.MaxUint64) - require.False(t, ok) - - _, ok = batchRetainedBytes(nil) - require.False(t, ok) - invalidRows := batch.NewOffHeapWithSize(0) - invalidRows.SetRowCount(-1) - _, ok = batchRetainedBytes(invalidRows) - require.False(t, ok) - _, err = scatterTransientBudgetBytes(nil, false) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - _, ok = (&SpillEngine{}).scatterCapacityGrowthBytes(-1, 0) - require.False(t, ok) - _, ok = (&SpillEngine{}).scatterCapacityGrowthBytes(math.MaxInt, 0) - require.False(t, ok) - - valid := batch.NewOffHeapWithSize(0) - valid.SetRowCount(0) - _, ok = predictMergedRetainedBytes(nil, valid) - require.False(t, ok) - invalidRows = batch.NewOffHeapWithSize(0) - invalidRows.SetRowCount(-1) - _, ok = predictMergedRetainedBytes(invalidRows, valid) - require.False(t, ok) - mismatched := batch.NewOffHeapWithSize(1) - mismatched.SetRowCount(0) - _, ok = predictMergedRetainedBytes(valid, mismatched) - require.False(t, ok) - nilVectorDst := batch.NewOffHeapWithSize(1) - nilVectorDst.SetRowCount(0) - nilVectorSrc := batch.NewOffHeapWithSize(1) - nilVectorSrc.SetRowCount(0) - _, ok = predictMergedRetainedBytes(nilVectorDst, nilVectorSrc) - require.False(t, ok) - hugeRowsDst := batch.NewOffHeapWithSize(0) - hugeRowsDst.SetRowCount(maxIntValue()) - hugeRowsSrc := batch.NewOffHeapWithSize(0) - hugeRowsSrc.SetRowCount(1) - _, ok = predictMergedRetainedBytes(hugeRowsDst, hugeRowsSrc) - require.False(t, ok) - - mp := mpool.MustNewZero() - fixed := testutil.MakeInt32Vector([]int32{1}, nil, mp) - defer fixed.Free(mp) - _, ok = mergedVarlenAreaAdd(nil, 1) - require.False(t, ok) - _, ok = mergedVarlenAreaAdd(fixed, 1) - require.False(t, ok) - constNull := vector.NewConstNull(types.T_varchar.ToType(), 1, mp) - defer constNull.Free(mp) - bytes, ok := mergedVarlenAreaAdd(constNull, 0) - require.True(t, ok) - require.Zero(t, bytes) - bytes, ok = mergedVarlenAreaAdd(constNull, 1) - require.True(t, ok) - require.Zero(t, bytes) - - require.NoError(t, reconcileReadReservation(nil, 0)) - budget, err := process.NewHashBuildBudget(10, 10) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - token, err := generation.Reserve(1) - require.NoError(t, err) - require.ErrorIs(t, reconcileReadReservation(token, 2), process.ErrHashBuildBudgetInvalid) - require.True(t, token.Release()) - require.ErrorIs(t, reconcileReadReservation(token, 0), process.ErrHashBuildReservationInactive) -} - -func TestScatterSkipsDisabledWriters(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - writers := MakeBucketWriters("test_skip") - for i := 0; i < len(writers); i += 2 { - writers[i].Name = "" // disable even buckets - } - - buffers := make([]*batch.Batch, len(writers)) - var buf bytes.Buffer - nRows := SpillNumBuckets * 100 - vals := make([]int32, nRows) - for i := range vals { - vals[i] = int32(i) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector(vals, nil, proc.Mp()) - bat.SetRowCount(nRows) - - err := scatterImpl(proc, bat, bat.Vecs[:1], writers, buffers, 0, &buf, nil, nil, nil) - require.NoError(t, err) - - // Disabled buckets must have no buffer and no file created. - for i := 0; i < len(writers); i += 2 { - require.Nil(t, buffers[i], "disabled bucket %d must have no buffer", i) - require.False(t, writers[i].Created(), "disabled bucket %d must not have file", i) - } - - // Enabled buckets should have received data. - hashValues := make([]uint64, nRows) - ComputeXXHash(bat.Vecs[:1], hashValues, 0) - var expectedOddRows int - for _, h := range hashValues { - if h&uint64(SpillNumBuckets-1)&1 == 1 { - expectedOddRows++ - } - } - require.Greater(t, expectedOddRows, 0) - - var oddRows int - for i := 1; i < len(writers); i += 2 { - if buffers[i] != nil { - oddRows += buffers[i].RowCount() - } - } - require.Equal(t, expectedOddRows, oddRows, "all odd-bucket rows must be in buffers") -} - -// --- SpillEngine tests --- - -func makeTestKeyExpr() []*plan.Expr { - return []*plan.Expr{{ - Typ: plan.Type{Id: int32(types.T_int32), Width: 32}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - }} -} - -func makeTestEvalKeysFn() func(*batch.Batch) ([]*vector.Vector, error) { - return func(bat *batch.Batch) ([]*vector.Vector, error) { - return bat.Vecs[:1], nil - } -} - -func TestScatterProbeAdmitsExpressionBeforeEvaluation(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - col := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - modulo, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "%", - []*plan.Expr{col, plan2.MakePlan2Int32ConstExprWithType(2)}, - ) - require.NoError(t, err) - execs, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{modulo}) - require.NoError(t, err) - budget, err := process.NewHashBuildBudget(8<<20, 8<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - probeLease, err := hashbuild.NewExpressionMemoryLease( - generation, []*plan.Expr{modulo}, execs, false) - require.NoError(t, err) - engine := NewSpillEngine(SpillEngineConfig{ - ProbeKeyExprs: []*plan.Expr{modulo}, - Budget: generation, - ProbeExpressionLease: probeLease, - }) - engine.InitFromSpilledMap(make([]*os.File, SpillNumBuckets)) - input := makeInt32Batch(proc, []int32{1, 2, 3, 4}) - defer input.Clean(proc.Mp()) - childrenCalls := 0 - fallbackCalled := false - err = engine.ScatterProbeTable( - proc, - func() (*batch.Batch, error) { - childrenCalls++ - if childrenCalls == 1 { - return input, nil - } - return nil, nil - }, - process.NewAnalyzer(0, false, false, "test"), - func(*batch.Batch) ([]*vector.Vector, error) { - fallbackCalled = true - return nil, errors.New("budgeted probe must evaluate its leased executors") - }, - ) - require.NoError(t, err) - require.Equal(t, 2, childrenCalls) - require.False(t, fallbackCalled) - require.Positive(t, probeLease.Reserved()) - require.Positive(t, generation.Used()) - engine.Cleanup(proc) - require.Positive(t, generation.Used(), "SpillEngine only borrows the probe lease") - for _, exec := range execs { - exec.Free() - } - probeLease.Release() - require.Zero(t, generation.Used()) -} - -func TestScatterProbeExpressionAdmissionRejectsBeforeEval(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - col := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - modulo, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "%", - []*plan.Expr{col, plan2.MakePlan2Int32ConstExprWithType(2)}, - ) - require.NoError(t, err) - execs, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{modulo}) - require.NoError(t, err) - retained, ok := colexec.ExpressionExecutorsRetainedBytes(execs) - require.True(t, ok) - peak, err := hashbuild.ExpressionVectorPeak(proc, modulo, 4, false) - require.NoError(t, err) - budgetCap := retained + peak - 1 - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - probeLease, err := hashbuild.NewExpressionMemoryLease( - generation, []*plan.Expr{modulo}, execs, false) - require.NoError(t, err) - engine := NewSpillEngine(SpillEngineConfig{ - ProbeKeyExprs: []*plan.Expr{modulo}, - Budget: generation, - ProbeExpressionLease: probeLease, - }) - engine.InitFromSpilledMap(make([]*os.File, SpillNumBuckets)) - input := makeInt32Batch(proc, []int32{1, 2, 3, 4}) - defer input.Clean(proc.Mp()) - childrenCalls := 0 - evalCalled := false - err = engine.ScatterProbeTable( - proc, - func() (*batch.Batch, error) { - childrenCalls++ - if childrenCalls == 1 { - return input, nil - } - return nil, nil - }, - process.NewAnalyzer(0, false, false, "test"), - func(*batch.Batch) ([]*vector.Vector, error) { - evalCalled = true - return nil, nil - }, - ) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Equal(t, 1, childrenCalls) - require.False(t, evalCalled) - engine.Cleanup(proc) - for _, exec := range execs { - exec.Free() - } - probeLease.Release() - require.Zero(t, generation.Used()) -} - -func makeInt32Batch(proc *process.Process, vals []int32) *batch.Batch { - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector(vals, nil, proc.Mp()) - bat.SetRowCount(len(vals)) - return bat -} - -func makeInt32PayloadBatch(t *testing.T, proc *process.Process, start, rows, payloadBytes int) *batch.Batch { - t.Helper() - bat := batch.NewWithSize(2) - vals := make([]int32, rows) - for i := range vals { - vals[i] = int32(start + i) - } - bat.Vecs[0] = testutil.MakeInt32Vector(vals, nil, proc.Mp()) - bat.Vecs[1] = vector.NewVec(types.T_varchar.ToType()) - payload := bytes.Repeat([]byte{'x'}, payloadBytes) - for i := 0; i < rows; i++ { - require.NoError(t, vector.AppendBytes(bat.Vecs[1], payload, false, proc.Mp())) - } - bat.SetRowCount(rows) - return bat -} - -func makeDedupKeepLastSpillBatch(proc *process.Process) *batch.Batch { - bat := batch.NewWithSize(3) - bat.Vecs[0] = testutil.MakeInt32Vector( - []int32{1, 1, 2}, nil, proc.Mp()) - bat.Vecs[1] = testutil.MakeInt32Vector( - []int32{10, 20, 30}, nil, proc.Mp()) - bat.Vecs[2] = testutil.MakeInt32Vector( - []int32{100, 0, 0}, []uint64{1, 2}, proc.Mp()) - bat.SetRowCount(3) - return bat -} - -func runtimeStackHasFunctionSuffix(suffix string) bool { - var callers [32]uintptr - n := runtime.Callers(2, callers[:]) - frames := runtime.CallersFrames(callers[:n]) - for { - frame, more := frames.Next() - if strings.HasSuffix(frame.Function, suffix) { - return true - } - if !more { - return false - } - } -} - -func writeBuildFile(proc *process.Process, name string, bat *batch.Batch) *os.File { - return writeBuildRecords(proc, name, bat) -} - -func writeBuildRecords(proc *process.Process, name string, batches ...*batch.Batch) *os.File { - spillfs, _ := proc.GetSpillFileService() - f, _ := spillfs.CreateAndRemoveFile(context.Background(), name) - var buf bytes.Buffer - w := BucketWriter{Name: name, Fd: f} - for _, bat := range batches { - FlushBucketBatch(proc, bat, &w, &buf, nil) - } - return w.HandOffFd() -} - -func makeCorruptBatchFile(t *testing.T) *os.File { - f, err := os.CreateTemp(t.TempDir(), "corrupt-spill") - require.NoError(t, err) - rowCount, batchSize := int64(1), int64(1) - var buf bytes.Buffer - buf.Write(types.EncodeInt64(&rowCount)) - buf.Write(types.EncodeInt64(&batchSize)) - buf.WriteByte(0xff) - _, err = f.Write(buf.Bytes()) - require.NoError(t, err) - _, err = f.Seek(0, io.SeekStart) - require.NoError(t, err) - return f -} - -func TestInitFromSpilledMapMixed(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := makeInt32Batch(proc, []int32{1, 2, 3}) - fd1 := writeBuildFile(proc, "test_mixed_1", bat) - fd2 := writeBuildFile(proc, "test_mixed_2", bat) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - }) - engine.InitFromSpilledMap([]*os.File{fd1, nil, fd2}) - require.Equal(t, 3, len(engine.buckets)) - require.NotNil(t, engine.buckets[0].BuildFd) - require.Nil(t, engine.buckets[1].BuildFd) - require.NotNil(t, engine.buckets[2].BuildFd) - engine.Cleanup(proc) -} - -func TestRebuildHashmapBasic(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - vals := make([]int32, 100) - for i := range vals { - vals[i] = int32(i) - } - bat := makeInt32Batch(proc, vals) - fd := writeBuildFile(proc, "test_rebuild", bat) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - }) - engine.InitFromSpilledMap([]*os.File{fd}) - - analyzer := process.NewAnalyzer(0, false, false, "test") - jm, res, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.Equal(t, BucketReady, res) - require.NotNil(t, jm) - require.Equal(t, int64(100), jm.GetRowCount()) - - jm2, res2, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.Equal(t, BucketQueueEmpty, res2) - require.Nil(t, jm2) - - jm.Free() - engine.Cleanup(proc) -} - -func TestRebuildHashmapCancellationKeepsFileOwnedUntilCleanup(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - ctx, cancel := context.WithCancelCause(proc.Ctx) - process.ReplacePipelineCtx(proc, ctx, cancel) - budget, err := process.NewHashBuildBudget(16<<20, 16<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - build := makeInt32Batch(proc, []int32{1, 2, 3, 4}) - var serialized bytes.Buffer - writer := BucketWriter{Name: "rebuild_cancel", Budget: generation} - defer writer.Close() - require.NoError(t, FlushBucketBatch(proc, build, &writer, &serialized, nil)) - build.Clean(proc.Mp()) - file, err := writer.handOffSpillFile() - require.NoError(t, err) - require.Positive(t, generation.SpillDiskUsed()) - require.Positive(t, generation.SpillFDUsed()) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - Budget: generation, - }) - engine.InitFromSpilledFiles([]*message.SpillFile{file}) - defer engine.Cleanup(proc) - - proc.Cancel(context.Canceled) - jm, res, err := engine.RebuildHashmap(proc, process.NewAnalyzer(0, false, false, "test")) - require.ErrorIs(t, err, context.Canceled) - require.Nil(t, jm) - require.Equal(t, BucketSkip, res) - require.True(t, engine.HasMoreBuckets(), "cancellation must leave the queued file with the engine cleanup owner") - require.Positive(t, generation.SpillDiskUsed()) - require.Positive(t, generation.SpillFDUsed()) - - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - require.Zero(t, proc.Mp().CurrNB()) -} - -func TestRebuildHashmapRespectsNeedFlags(t *testing.T) { - tests := []struct { - name string - needAllocateSels bool - needBatches bool - }{ - {name: "sels only", needAllocateSels: true}, - {name: "batches only", needBatches: true}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := makeInt32Batch(proc, []int32{1, 1}) - fd := writeBuildFile(proc, "test_rebuild_flags", bat) - bat.Clean(proc.Mp()) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - NeedAllocateSels: tt.needAllocateSels, - NeedBatches: tt.needBatches, - }) - engine.InitFromSpilledMap([]*os.File{fd}) - - jm, res, err := engine.RebuildHashmap(proc, process.NewAnalyzer(0, false, false, "test")) - require.NoError(t, err) - require.Equal(t, BucketReady, res) - require.NotNil(t, jm) - - if tt.needAllocateSels { - require.Equal(t, []int32{0, 1}, jm.GetSels(0)) - } else { - require.Nil(t, jm.GetSels(0)) - } - if tt.needBatches { - require.Len(t, jm.GetBatches(), 1) - require.Equal(t, 2, jm.GetBatches()[0].RowCount()) - } else { - require.Empty(t, jm.GetBatches()) - } - - jm.Free() - engine.Cleanup(proc) - require.Equal(t, int64(0), proc.Mp().CurrNB()) - }) - } -} - -func TestRebuildHashmapWithoutBatchesDropsBatchBudgetBeforeProbe(t *testing.T) { - run := func(t *testing.T, needBatches bool) uint64 { - t.Helper() - const budgetCap = uint64(64 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - values := make([]int32, colexec.DefaultBatchSize/2) - for i := range values { - values[i] = int32(i) - } - buildBat := makeInt32Batch(proc, values) - buildFd := writeBuildFile(proc, "test_rebuild_batch_budget", buildBat) - buildBat.Clean(proc.Mp()) - probeBat := makeInt32Batch(proc, []int32{1}) - probeFd := writeBuildFile(proc, "test_rebuild_batch_budget_probe", probeBat) - probeBat.Clean(proc.Mp()) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedBatches: needBatches, - Budget: generation, - }) - engine.InitFromSpilledMap([]*os.File{buildFd}) - engine.buckets[0].ProbeFd = message.NewSpillFile(probeFd, 1, 0, nil) - - jm, res, err := engine.RebuildHashmap(proc, process.NewAnalyzer(0, false, false, "test")) - require.NoError(t, err) - require.Equal(t, BucketReady, res) - require.NotNil(t, jm) - if needBatches { - require.NotEmpty(t, jm.GetBatches()) - } else { - require.Empty(t, jm.GetBatches()) - } - used := generation.Used() - require.Positive(t, used, "hash map and reader ownership remain live during probe") - - jm.Free() - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - return used - } - - withoutBatches := run(t, false) - withBatches := run(t, true) - require.Greater(t, withBatches, withoutBatches, - "NeedBatches=false must not transfer destroyed batch reservations into the JoinMap") -} - -func TestRebuildHashmapEmptyBuild(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - }) - engine.InitFromSpilledMap([]*os.File{nil}) - - analyzer := process.NewAnalyzer(0, false, false, "test") - jm, res, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.Equal(t, BucketSkip, res) - require.Nil(t, jm) - require.False(t, engine.HasMoreBuckets()) - engine.Cleanup(proc) -} - -func TestRebuildHashmapEmptyBuildOuterJoin(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - var buf bytes.Buffer - spillfs, _ := proc.GetSpillFileService() - f, _ := spillfs.CreateAndRemoveFile(context.Background(), "test_outer_probe") - bat := makeInt32Batch(proc, []int32{1, 2}) - w := BucketWriter{Name: "test_outer_probe", Fd: f} - FlushBucketBatch(proc, bat, &w, &buf, nil) - probeFd := w.HandOffFd() - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsProbeForEmptyBuild: true, - }) - engine.InitFromSpilledMap([]*os.File{nil}) - engine.buckets[0].ProbeFd = message.NewSpillFile(probeFd, 0, 0, nil) - - analyzer := process.NewAnalyzer(0, false, false, "test") - jm, res, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.Equal(t, BucketEmptyBuild, res) - require.Nil(t, jm) - require.True(t, engine.IsProbing()) - require.False(t, engine.HasMoreBuckets()) - - got, err := engine.NextProbeBatch(proc) - require.NoError(t, err) - require.NotNil(t, got) - require.Equal(t, 2, got.RowCount()) - - engine.Cleanup(proc) -} - -func TestRebuildHashmapEmptyFile(t *testing.T) { - for _, keepProbe := range []bool{false, true} { - t.Run(fmt.Sprintf("keep_probe_%t", keepProbe), func(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - baseline := proc.Mp().CurrNB() - budget, err := process.NewHashBuildBudget(16<<20, 16<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - buildFd, err := spillfs.CreateAndRemoveFile(proc.Ctx, "test_empty_build_file") - require.NoError(t, err) - probeBat := makeInt32Batch(proc, []int32{1}) - probeFd := writeBuildFile(proc, "test_empty_build_probe", probeBat) - probeBat.Clean(proc.Mp()) - - engine := NewSpillEngine(SpillEngineConfig{ - // A literal executor owns an mpool vector as soon as Prepare - // succeeds. This makes the empty-build branch's builder.Free - // observable instead of relying on a zero-allocation column - // executor. - BuildKeyExprs: []*plan.Expr{ - plan2.MakePlan2Int32ConstExprWithType(1), - }, - NeedsProbeForEmptyBuild: keepProbe, - Budget: generation, - }) - engine.InitFromSpilledMap([]*os.File{buildFd}) - engine.buckets[0].ProbeFd = message.NewSpillFile(probeFd, 0, 0, nil) - - jm, res, err := engine.RebuildHashmap(proc, process.NewAnalyzer(0, false, false, "test")) - require.NoError(t, err) - require.Nil(t, jm) - if keepProbe { - require.Equal(t, BucketEmptyBuild, res) - require.True(t, engine.IsProbing()) - } else { - require.Equal(t, BucketSkip, res) - } - engine.Cleanup(proc) - require.Equal(t, baseline, proc.Mp().CurrNB()) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - }) - } -} - -func TestScatterProbeTable(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := makeInt32Batch(proc, []int32{10, 20, 30}) - fd := writeBuildFile(proc, "test_sp_build", bat) - bat.Clean(proc.Mp()) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - }) - engine.InitFromSpilledMap([]*os.File{fd}) - - // Use many rows to ensure distribution across buckets. - vals := make([]int32, 1000) - for i := range vals { - vals[i] = int32(i) - } - batches := []*batch.Batch{makeInt32Batch(proc, vals)} - defer batches[0].Clean(proc.Mp()) - idx := 0 - children := func() (*batch.Batch, error) { - if idx >= len(batches) { - return nil, nil - } - b := batches[idx] - idx++ - return b, nil - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - err := engine.ScatterProbeTable(proc, children, analyzer, makeTestEvalKeysFn()) - require.NoError(t, err) - require.NotNil(t, engine.probeKeyEval) - require.Len(t, engine.buckets, 1) - require.NotNil(t, engine.buckets[0].ProbeFd) - require.Equal(t, int64(len(vals)), engine.buckets[0].ProbeRows, - "probe partitioning must conserve every row at the build payload's fanout") - - engine.Cleanup(proc) -} - -func TestScatterProbeTableRejectsInvalidBuildFanoutBeforeInput(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - for _, bucketCount := range []int{0, 3, SpillNumBuckets + 1} { - t.Run(fmt.Sprintf("buckets_%d", bucketCount), func(t *testing.T) { - engine := NewSpillEngine(SpillEngineConfig{}) - engine.InitFromSpilledMap(make([]*os.File, bucketCount)) - inputCalled := false - err := engine.ScatterProbeTable( - proc, - func() (*batch.Batch, error) { - inputCalled = true - return nil, nil - }, - process.NewAnalyzer(0, false, false, "test"), - makeTestEvalKeysFn(), - ) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - require.False(t, inputCalled, "invalid fanout must fail before consuming probe input") - engine.Cleanup(proc) - }) - } -} - -func TestScatterProbeCancellationReleasesPhysicalAndMemoryBudget(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - ctx, cancel := context.WithCancelCause(proc.Ctx) - process.ReplacePipelineCtx(proc, ctx, cancel) - budget, err := process.NewHashBuildBudget(64<<20, 64<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - // Identical keys force one large selected payload through a real spill - // writer before the second upstream call cancels the pipeline. - const rows = 8192 - keys := make([]int32, rows) - payload := make([]string, rows) - for i := range keys { - keys[i] = 7 - payload[i] = strings.Repeat("x", 128) - } - input := batch.NewWithSize(2) - input.Vecs[0] = testutil.MakeInt32Vector(keys, nil, proc.Mp()) - input.Vecs[1] = testutil.MakeVarcharVector(payload, nil, proc.Mp()) - input.SetRowCount(rows) - defer func() { - if input != nil { - input.Clean(proc.Mp()) - } - }() - - engine := NewSpillEngine(SpillEngineConfig{ - ProbeKeyExprs: makeTestKeyExpr(), - NeedsProbeForEmptyBuild: true, - Budget: generation, - }) - engine.InitFromSpilledMap(make([]*os.File, SpillNumBuckets)) - analyzer := process.NewAnalyzer(0, false, false, "test") - childrenCalls := 0 - var peakDisk, peakFD uint64 - err = engine.ScatterProbeTable( - proc, - func() (*batch.Batch, error) { - childrenCalls++ - if childrenCalls == 1 { - return input, nil - } - peakDisk = generation.SpillDiskUsed() - peakFD = generation.SpillFDUsed() - proc.Cancel(context.Canceled) - return input, nil - }, - analyzer, - makeTestEvalKeysFn(), - ) - require.ErrorIs(t, err, context.Canceled) - require.Equal(t, 2, childrenCalls) - require.Positive(t, peakDisk, "first batch must reach a physical spill file") - require.Positive(t, peakFD, "first batch must own an admitted spill descriptor") - - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - input.Clean(proc.Mp()) - input = nil - require.Zero(t, proc.Mp().CurrNB()) -} - -func TestSpillEntryPointsRejectPreCanceledProcessWithoutOwnershipTransfer(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - t.Cleanup(proc.Free) - ctx, cancel := context.WithCancelCause(proc.Ctx) - process.ReplacePipelineCtx(proc, ctx, cancel) - proc.Cancel(context.Canceled) - - reuseBat := batch.NewOffHeapWithSize(0) - reuseCleaned := false - cleanReuse := func() { - if !reuseCleaned { - reuseBat.Clean(proc.Mp()) - reuseCleaned = true - } - } - t.Cleanup(cleanReuse) - reader := BucketReader{} - got, err := reader.ReadBatch(proc, reuseBat) - require.Nil(t, got) - require.ErrorIs(t, err, context.Canceled) - require.Zero(t, reuseBat.RowCount()) - - writer := BucketWriter{Name: "must_not_be_created"} - t.Cleanup(writer.Close) - err = writeBucketPayload( - proc, - []byte{1}, - 1, - &writer, - process.NewAnalyzer(0, false, false, "test"), - ) - require.ErrorIs(t, err, context.Canceled) - require.False(t, writer.Created()) - require.Nil(t, writer.diskReservation) - require.Nil(t, writer.fdReservation) - - input := makeInt32Batch(proc, []int32{1}) - inputCleaned := false - cleanInput := func() { - if !inputCleaned { - input.Clean(proc.Mp()) - inputCleaned = true - } - } - t.Cleanup(cleanInput) - engine := NewSpillEngine(SpillEngineConfig{}) - engineCleaned := false - cleanEngine := func() { - if !engineCleaned { - engine.Cleanup(proc) - engineCleaned = true - } - } - t.Cleanup(cleanEngine) - scatterWriters := []BucketWriter{{Name: "must_not_be_created"}} - t.Cleanup(scatterWriters[0].Close) - err = engine.scatterBatchWithPressure( - proc, - input, - []*vector.Vector{input.Vecs[0]}, - scatterWriters, - 0, - false, - process.NewAnalyzer(0, false, false, "test"), - ) - require.ErrorIs(t, err, context.Canceled) - require.False(t, scatterWriters[0].Created()) - require.Nil(t, scatterWriters[0].diskReservation) - require.Nil(t, scatterWriters[0].fdReservation) - - engine.InitFromSpilledMap([]*os.File{nil}) - childrenCalled := false - err = engine.ScatterProbeTable( - proc, - func() (*batch.Batch, error) { - childrenCalled = true - return input, nil - }, - process.NewAnalyzer(0, false, false, "test"), - makeTestEvalKeysFn(), - ) - require.ErrorIs(t, err, context.Canceled) - require.False(t, childrenCalled) - - probeFd, err := os.CreateTemp(t.TempDir(), "pre-canceled-probe") - require.NoError(t, err) - t.Cleanup(func() { _ = probeFd.Close() }) - probeReleases := 0 - probeFile := message.NewSpillFile(probeFd, 1, 1, func() { probeReleases++ }) - engine.probeReader.ResetForSpillFile(probeFile) - got, err = engine.NextProbeBatch(proc) - require.Nil(t, got) - require.ErrorIs(t, err, context.Canceled) - require.Zero(t, probeReleases) - _, err = probeFd.Stat() - require.NoError(t, err) - - reSpillFd, err := os.CreateTemp(t.TempDir(), "pre-canceled-respill") - require.NoError(t, err) - t.Cleanup(func() { _ = reSpillFd.Close() }) - reSpillReleases := 0 - reSpillFile := message.NewSpillFile(reSpillFd, 1, 1, func() { reSpillReleases++ }) - t.Cleanup(func() { _ = reSpillFile.Close() }) - subBuckets, err := engine.reSpillBucket( - proc, - process.NewAnalyzer(0, false, false, "test"), - SpillBucket{BuildFd: reSpillFile}, - nil, - nil, - nil, - ) - require.Nil(t, subBuckets) - require.ErrorIs(t, err, context.Canceled) - require.Zero(t, reSpillReleases) - _, err = reSpillFd.Stat() - require.NoError(t, err) - require.NoError(t, reSpillFile.Close()) - require.Equal(t, 1, reSpillReleases) - _, err = reSpillFd.Stat() - require.Error(t, err) - - cleanEngine() - require.Equal(t, 1, probeReleases) - _, err = probeFd.Stat() - require.Error(t, err) - cleanInput() - cleanReuse() - require.Zero(t, proc.Mp().CurrNB()) -} - -func TestScatterProbeTableSkipEmptyBuild(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := makeInt32Batch(proc, []int32{1, 2, 3}) - fd1 := writeBuildFile(proc, "test_skip_build_1", bat) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - }) - engine.InitFromSpilledMap([]*os.File{fd1, nil}) - - // Use many rows to ensure distribution. - vals := make([]int32, 1000) - for i := range vals { - vals[i] = int32(i) - } - batches := []*batch.Batch{makeInt32Batch(proc, vals)} - idx := 0 - children := func() (*batch.Batch, error) { - if idx >= len(batches) { - return nil, nil - } - b := batches[idx] - idx++ - return b, nil - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - err := engine.ScatterProbeTable(proc, children, analyzer, makeTestEvalKeysFn()) - require.NoError(t, err) - - // Bucket 1 (nil build, not outer join) should have no probe data. - require.Nil(t, engine.buckets[1].ProbeFd) - - engine.Cleanup(proc) -} - -func TestScatterProbeTableWithEmptyBatches(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := makeInt32Batch(proc, []int32{1, 2}) - fd := writeBuildFile(proc, "test_empty_bat", bat) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - }) - engine.InitFromSpilledMap([]*os.File{fd}) - - batches := []*batch.Batch{ - batch.NewWithSize(0), - makeInt32Batch(proc, []int32{5, 6}), - } - idx := 0 - children := func() (*batch.Batch, error) { - if idx >= len(batches) { - return nil, nil - } - b := batches[idx] - idx++ - return b, nil - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - err := engine.ScatterProbeTable(proc, children, analyzer, makeTestEvalKeysFn()) - require.NoError(t, err) - - engine.Cleanup(proc) -} - -func TestNextProbeBatch(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - // Write build file. - bat := makeInt32Batch(proc, []int32{1, 2, 3}) - fd := writeBuildFile(proc, "test_npb_build", bat) - - // Write probe file manually (ensures bucket 0 has probe data). - probeFd := writeBuildFile(proc, "test_npb_probe", makeInt32Batch(proc, []int32{5, 6, 7})) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - }) - engine.InitFromSpilledMap([]*os.File{fd}) - engine.buckets[0].ProbeFd = message.NewSpillFile(probeFd, 0, 0, nil) - - analyzer := process.NewAnalyzer(0, false, false, "test") - jm, res, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.Equal(t, BucketReady, res) - require.True(t, engine.IsProbing()) - - got, err := engine.NextProbeBatch(proc) - require.NoError(t, err) - require.NotNil(t, got) - require.Equal(t, 3, got.RowCount()) - - got2, err2 := engine.NextProbeBatch(proc) - require.NoError(t, err2) - require.Nil(t, got2) - - engine.FinishBucket() - require.False(t, engine.IsProbing()) - got3, err3 := engine.NextProbeBatch(proc) - require.NoError(t, err3) - require.Nil(t, got3) - - jm.Free() - engine.Cleanup(proc) -} - -func TestCorruptSpillBatchErrors(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - engine := NewSpillEngine(SpillEngineConfig{}) - engine.probeReader.ResetForFd(makeCorruptBatchFile(t)) - _, err := engine.NextProbeBatch(proc) - require.Error(t, err) - engine.Cleanup(proc) - - engine = NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - }) - engine.InitFromSpilledMap([]*os.File{makeCorruptBatchFile(t)}) - jm, res, err := engine.RebuildHashmap(proc, process.NewAnalyzer(0, false, false, "test")) - require.Error(t, err) - require.Nil(t, jm) - require.Equal(t, BucketSkip, res) - engine.Cleanup(proc) -} - -func TestRebuildSkipsBuildOnlyBucketWhenJoinCannotUseIt(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - fd := makeCorruptBatchFile(t) - releases := 0 - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - }) - engine.InitFromSpilledFiles([]*message.SpillFile{ - message.NewSpillFile(fd, 1, 17, func() { releases++ }), - }) - - jm, res, err := engine.RebuildHashmap(proc, process.NewAnalyzer(0, false, false, "test")) - require.NoError(t, err, "an irrelevant build-only file must not be read") - require.Nil(t, jm) - require.Equal(t, BucketSkip, res) - require.False(t, engine.HasMoreBuckets()) - require.Equal(t, 1, releases) - _, err = fd.Stat() - require.Error(t, err) - - engine.Cleanup(proc) - require.Equal(t, 1, releases) -} - -func TestAdvanceToNextBucket(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - vals := make([]int32, 50) - for i := range vals { - vals[i] = int32(i) - } - bat := makeInt32Batch(proc, vals) - fd := writeBuildFile(proc, "test_advance_build", bat) - probeFd := writeBuildFile(proc, "test_advance_probe", makeInt32Batch(proc, []int32{5, 6})) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - }) - engine.InitFromSpilledMap([]*os.File{fd}) - engine.buckets[0].ProbeFd = message.NewSpillFile(probeFd, 0, 0, nil) - - var capturedJM *message.JoinMap - var capturedRes BucketResult - analyzer := process.NewAnalyzer(0, false, false, "test") - - ok, err := engine.AdvanceToNextBucket(proc, analyzer, func(jm *message.JoinMap, res BucketResult) { - capturedJM = jm - capturedRes = res - }) - require.NoError(t, err) - require.True(t, ok) - require.NotNil(t, capturedJM) - require.Equal(t, BucketReady, capturedRes) - require.True(t, engine.IsProbing()) - require.False(t, engine.HasMoreBuckets()) - - capturedJM.Free() - engine.FinishBucket() - ok, err = engine.AdvanceToNextBucket(proc, analyzer, func(*message.JoinMap, BucketResult) { - t.Fatal("queue exhaustion must not invoke callback") - }) - require.NoError(t, err) - require.False(t, ok) - engine.Cleanup(proc) -} - -func TestReSpillBucket(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - vals := make([]int32, 5000) - for i := range vals { - vals[i] = int32(i) - } - bat := makeInt32Batch(proc, vals) - fd := writeBuildFile(proc, "test_respill_build", bat) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - SpillThreshold: 100, - }) - engine.InitFromSpilledMap([]*os.File{fd}) - - analyzer := process.NewAnalyzer(0, false, false, "test") - jm, res, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.Equal(t, BucketReSpilled, res) - require.Nil(t, jm) - - // Drain all remaining buckets. - for engine.HasMoreBuckets() { - jm2, _, err2 := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err2) - if jm2 != nil { - jm2.Free() - } - } - - engine.Cleanup(proc) -} - -func TestRebuildHashmapKeepsScratchHeadroomForCopyAdmissionReSpill(t *testing.T) { - const ( - budgetCap = uint64(12 << 20) - recordRows = colexec.DefaultBatchSize - recordCount = 3 - payloadSize = 256 - ) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - records := make([]*batch.Batch, recordCount) - for i := range records { - records[i] = makeInt32PayloadBatch(t, proc, i*recordRows, recordRows, payloadSize) - } - buildFd := writeBuildRecords(proc, "rebuild_scratch_headroom", records...) - for i := range records { - records[i].Clean(proc.Mp()) - } - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - SpillThreshold: 1 << 30, - Budget: generation, - }) - engine.InitFromSpilledMap([]*os.File{buildFd}) - analyzer := process.NewAnalyzer(0, false, false, "test") - - jm, result, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.Nil(t, jm) - require.Equal(t, BucketReSpilled, result) - var childRows int64 - for _, child := range engine.buckets { - childRows += child.BuildRows - } - require.Equal(t, int64(recordRows*recordCount), childRows) - extra := analyzer.GetOpStats().ExtraStats - require.Positive(t, extra["JoinSpillRebuildScratchFloorBytes"]) - require.Equal(t, int64(1), extra["JoinSpillRebuildCopyAdmissionReSpillAttempts"]) - require.Zero(t, extra["JoinSpillRebuildPreCopyReSpillAttempts"]) - - // The rejected record is reader-owned pending state: retained, pending, and - // unread rows must form an exact partition of the original stream. - seen := make([]uint8, recordRows*recordCount) - reuse := batch.NewOffHeapWithSize(0) - reader := BucketReader{} - for i := range engine.buckets { - file := engine.buckets[i].BuildFd - engine.buckets[i].BuildFd = nil - reader.ResetForSpillFile(file) - for { - bat, readErr := reader.ReadBatch(proc, reuse) - if readErr == io.EOF { - break - } - require.NoError(t, readErr) - for _, key := range vector.MustFixedColNoTypeCheck[int32](bat.Vecs[0]) { - require.GreaterOrEqual(t, key, int32(0)) - require.Less(t, key, int32(len(seen))) - seen[key]++ - } - } - reader.closeCurrentFile() - } - reader.Close() - reuse.Clean(proc.Mp()) - for key, count := range seen { - require.Equalf(t, uint8(1), count, "key %d must be emitted exactly once", key) - } - - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - require.Zero(t, proc.Mp().CurrNB()) - generation.Close() - proc.Free() -} - -func TestRebuildScratchAdmissionIsBestEffortForResidentBucket(t *testing.T) { - const budgetCap = uint64(64 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - build := makeInt32PayloadBatch(t, proc, 0, 1024, 1024) - buildFd := writeBuildFile(proc, "rebuild_scratch_best_effort", build) - build.Clean(proc.Mp()) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - NeedBatches: true, - SpillThreshold: 1 << 30, - Budget: generation, - }) - engine.InitFromSpilledMap([]*os.File{buildFd}) - floorRejected := false - budget.SetAggregateCapProvider(func() (uint64, error) { - if !floorRejected && runtimeStackHasFunctionSuffix( - "spillutil.(*SpillEngine).reserveRebuildScatterScratch", - ) { - floorRejected = true - return generation.Used(), nil - } - return budgetCap, nil - }) - analyzer := process.NewAnalyzer(0, false, false, "test") - - jm, result, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.True(t, floorRejected) - require.Equal(t, BucketReady, result) - require.NotNil(t, jm) - require.Equal(t, int64(1024), jm.GetRowCount()) - require.Equal(t, int64(1), - analyzer.GetOpStats().ExtraStats["JoinSpillRebuildScratchReserveRejects"]) - jm.Free() - - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - require.Zero(t, proc.Mp().CurrNB()) - generation.Close() - proc.Free() -} - -func TestRebuildScratchLifecycleFailureIsNotRecoveredAsAdmission(t *testing.T) { - const budgetCap = uint64(64 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - build := makeInt32Batch(proc, []int32{1, 2, 3}) - buildFd := writeBuildFile(proc, "rebuild_scratch_closed", build) - build.Clean(proc.Mp()) - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - SpillThreshold: 1 << 30, - Budget: generation, - }) - engine.InitFromSpilledMap([]*os.File{buildFd}) - closedErr := &process.HashBuildBudgetError{ - Kind: process.HashBuildBudgetErrorClosed, - Message: "forced closed rebuild scratch budget", - } - budget.SetAggregateCapProvider(func() (uint64, error) { - if runtimeStackHasFunctionSuffix( - "spillutil.(*SpillEngine).reserveRebuildScatterScratch", - ) { - return 0, closedErr - } - return budgetCap, nil - }) - analyzer := process.NewAnalyzer(0, false, false, "test") - - jm, result, err := engine.RebuildHashmap(proc, analyzer) - require.Same(t, closedErr, err) - require.ErrorIs(t, err, process.ErrHashBuildBudgetClosed) - require.NotErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Nil(t, jm) - require.Equal(t, BucketSkip, result) - require.Zero(t, analyzer.GetOpStats().ExtraStats["JoinSpillRebuildCopyAdmissionReSpillAttempts"]) - require.Zero(t, analyzer.GetOpStats().ExtraStats["JoinSpillRebuildPreCopyReSpillAttempts"]) - - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - require.Zero(t, proc.Mp().CurrNB()) - generation.Close() - proc.Free() -} - -func TestRebuildHashmapRejectsReSpillAfterDedupRewrite(t *testing.T) { - const budgetCap = uint64(64 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - build := makeDedupKeepLastSpillBatch(proc) - buildFd := writeBuildFile(proc, "dedup_unsafe_respill", build) - build.Clean(proc.Mp()) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - NeedBatches: true, - IsDedup: true, - OnDuplicateAction: plan.Node_FAIL, - DedupBuildKeepLast: true, - DedupColName: "id", - DedupColTypes: []plan.Type{{Id: int32(types.T_int32)}}, - DelColIdx: -1, - DedupDeleteMarkerColIdx: 2, - DedupDeleteKeepColIdxList: []int32{2}, - SpillThreshold: 1 << 30, - Budget: generation, - }) - engine.InitFromSpilledMap([]*os.File{buildFd}) - - forcedUnsafeReject := false - budget.SetAggregateCapProvider(func() (uint64, error) { - // keepDiscardedRowsForDelete has already compacted the retained input - // when it asks copyBuildBatch to admit the delete-only rows. Reject that - // exact transition without depending on a fragile global call ordinal. - if runtimeStackHasFunctionSuffix( - "hashbuild.(*HashmapBuilder).keepDiscardedRowsForDelete", - ) { - forcedUnsafeReject = true - return max(uint64(1), generation.Used()), nil - } - return budgetCap, nil - }) - - analyzer := process.NewAnalyzer(0, false, false, "test") - jm, result, err := engine.RebuildHashmap(proc, analyzer) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.True(t, forcedUnsafeReject) - require.Nil(t, jm) - require.Equal(t, BucketSkip, result) - require.Len(t, engine.buckets, 1, - "unsafe recovery must not replace the parent with child buckets") - require.Equal(t, 1, engine.buckets[0].Depth) - require.Nil(t, engine.buckets[0].BuildFd, - "the consumed parent file must not be republished as a JoinMap or child") - - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - require.Zero(t, proc.Mp().CurrNB()) - generation.Close() - proc.Free() -} - -func TestRebuildHashmapReSpillsAdmissionBeforeDedupRewrite(t *testing.T) { - const budgetCap = uint64(64 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - build := makeDedupKeepLastSpillBatch(proc) - buildFd := writeBuildFile(proc, "dedup_safe_respill", build) - build.Clean(proc.Mp()) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - NeedBatches: true, - IsDedup: true, - OnDuplicateAction: plan.Node_FAIL, - DedupBuildKeepLast: true, - DedupColName: "id", - DedupColTypes: []plan.Type{{Id: int32(types.T_int32)}}, - DelColIdx: -1, - DedupDeleteMarkerColIdx: 2, - DedupDeleteKeepColIdxList: []int32{2}, - SpillThreshold: 1 << 30, - Budget: generation, - }) - engine.InitFromSpilledMap([]*os.File{buildFd}) - - forcedSafeReject := false - budget.SetAggregateCapProvider(func() (uint64, error) { - // The first budget request made from buildHashmap is reserveBuildAux, - // before any Dedup batch rewrite. Reject once; re-spill itself does not - // call buildHashmap and therefore retains the normal cap. - if !forcedSafeReject && - runtimeStackHasFunctionSuffix( - "hashbuild.(*HashmapBuilder).buildHashmap", - ) && - !runtimeStackHasFunctionSuffix( - "hashbuild.(*HashmapBuilder).keepDiscardedRowsForDelete", - ) { - forcedSafeReject = true - return max(uint64(1), generation.Used()), nil - } - return budgetCap, nil - }) - - analyzer := process.NewAnalyzer(0, false, false, "test") - jm, result, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.True(t, forcedSafeReject) - require.Nil(t, jm) - require.Equal(t, BucketReSpilled, result) - require.NotEmpty(t, engine.buckets) - var childRows int64 - for _, child := range engine.buckets { - require.Equal(t, 2, child.Depth) - childRows += child.BuildRows - } - require.Equal(t, int64(3), childRows, - "safe recovery must conserve the original retained rows") - require.Equal(t, int64(1), - analyzer.GetOpStats().ExtraStats["JoinSpillRebuildMapAdmissionReSpillAttempts"]) - - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - require.Zero(t, proc.Mp().CurrNB()) - generation.Close() - proc.Free() -} - -func TestRebuildHashmapClosedBudgetDoesNotReSpill(t *testing.T) { - const budgetCap = uint64(64 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - build := makeInt32Batch(proc, []int32{1, 2, 3}) - buildFd := writeBuildFile(proc, "closed_budget_no_respill", build) - build.Clean(proc.Mp()) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - SpillThreshold: 1 << 30, - Budget: generation, - }) - engine.InitFromSpilledMap([]*os.File{buildFd}) - - closedErr := &process.HashBuildBudgetError{ - Kind: process.HashBuildBudgetErrorClosed, - Message: "forced closed hash-build budget", - } - forcedClosed := false - budget.SetAggregateCapProvider(func() (uint64, error) { - if !forcedClosed && runtimeStackHasFunctionSuffix( - "hashbuild.(*HashmapBuilder).buildHashmap", - ) { - forcedClosed = true - return 0, closedErr - } - return budgetCap, nil - }) - - jm, result, err := engine.RebuildHashmap( - proc, process.NewAnalyzer(0, false, false, "test")) - require.True(t, forcedClosed) - require.Same(t, closedErr, err, - "a lifecycle failure must be returned unchanged") - require.ErrorIs(t, err, process.ErrHashBuildBudgetClosed) - require.NotErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Nil(t, jm) - require.Equal(t, BucketSkip, result) - require.Len(t, engine.buckets, 1, - "a lifecycle failure must not replace the parent with child buckets") - require.Equal(t, 1, engine.buckets[0].Depth) - require.Nil(t, engine.buckets[0].BuildFd) - - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - require.Zero(t, proc.Mp().CurrNB()) - generation.Close() - proc.Free() -} - -func TestReSpillReleasesBuilderExecutorsBeforeReplacementAdmission(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - col := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - modulo, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "%", - []*plan.Expr{col, plan2.MakePlan2Int32ConstExprWithType(2)}, - ) - require.NoError(t, err) - exprs := []*plan.Expr{modulo} - - probeExecs, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, exprs) - require.NoError(t, err) - retained, ok := colexec.ExpressionExecutorsRetainedBytes(probeExecs) - require.True(t, ok) - require.Positive(t, retained) - for _, executor := range probeExecs { - executor.Free() - } - - // The cap intentionally fits exactly one executor set. reSpillBucket must - // release the failed builder's equivalent set before constructing its own. - budget := process.MustNewHashBuildBudget(retained, retained) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - builder := &hashbuild.HashmapBuilder{} - builder.SetBudget(generation) - require.NoError(t, builder.Prepare(exprs, -1, -1, nil, proc)) - require.Equal(t, retained, generation.Used()) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: exprs, - Budget: generation, - }) - subBuckets, err := engine.reSpillBucket( - proc, - process.NewAnalyzer(0, false, false, "test"), - SpillBucket{}, - builder, - &BucketReader{}, - nil, - ) - require.NoError(t, err) - require.Empty(t, subBuckets) - require.Equal(t, retained, generation.Used()) - require.NotNil(t, engine.buildExprLease) - - builder.Free(proc) - engine.Cleanup(proc) - require.Zero(t, generation.Used()) -} - -func TestReSpillBucketReleasesDrainedBatchBudget(t *testing.T) { - const budgetCap = uint64(64 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - builder := &hashbuild.HashmapBuilder{} - builder.SetBudget(generation) - require.NoError(t, builder.Prepare(makeTestKeyExpr(), -1, -1, nil, proc)) - defer builder.Free(proc) - - values := make([]int32, colexec.DefaultBatchSize/2) - for i := range values { - values[i] = int32(i) - } - input := makeInt32Batch(proc, values) - require.NoError(t, builder.CopyBuildBatch(input, proc)) - builder.InputBatchRowCount = input.RowCount() - input.Clean(proc.Mp()) - batchCharge := generation.Used() - require.Positive(t, batchCharge) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - Budget: generation, - }) - subBuckets, err := engine.reSpillBucket( - proc, - process.NewAnalyzer(0, false, false, "test"), - SpillBucket{Depth: 1, BuildRows: int64(len(values))}, - builder, - &BucketReader{}, - nil, - ) - require.NoError(t, err) - for i := range subBuckets { - if subBuckets[i].BuildFd != nil { - require.NoError(t, subBuckets[i].BuildFd.Close()) - } - if subBuckets[i].ProbeFd != nil { - require.NoError(t, subBuckets[i].ProbeFd.Close()) - } - } - engine.Cleanup(proc) - - require.Empty(t, builder.Batches.Buf) - require.Less(t, generation.Used(), batchCharge, - "re-spill must not retain the destroyed build-batch reservation") - require.Zero(t, generation.Used()) - - // Model the scratch/read admission that follows the drain. The full cap is - // available only when re-spill released the stale batch ownership itself, - // instead of relying on its caller to free the builder later. - next, err := generation.Reserve(budgetCap) - require.NoError(t, err) - require.True(t, next.Release()) -} - -func TestReSpillConservesBuildAndProbeRows(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - vals := make([]int32, 5000) - for i := range vals { - vals[i] = int32(i) - } - build := makeInt32Batch(proc, vals) - probe := makeInt32Batch(proc, vals) - buildFd := writeBuildFile(proc, "test_conserve_build", build) - probeFd := writeBuildFile(proc, "test_conserve_probe", probe) - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - SpillThreshold: 500, - NeedsBuildForEmptyProbe: true, - NeedsProbeForEmptyBuild: true, - }) - engine.InitFromSpilledMap([]*os.File{buildFd}) - engine.buckets[0].ProbeFd = message.NewSpillFile(probeFd, int64(len(vals)), 0, nil) - engine.probeKeyEval = makeTestEvalKeysFn() - - jm, res, err := engine.RebuildHashmap(proc, process.NewAnalyzer(0, false, false, "test")) - require.NoError(t, err) - require.Equal(t, BucketReSpilled, res) - require.Nil(t, jm) - var buildRows, probeRows, largest int64 - for _, child := range engine.buckets { - buildRows += child.BuildRows - probeRows += child.ProbeRows - if child.BuildRows > largest { - largest = child.BuildRows - } - } - require.Equal(t, int64(len(vals)), buildRows) - require.Equal(t, int64(len(vals)), probeRows) - require.Less(t, largest, int64(len(vals))) - - for engine.HasMoreBuckets() { - jm, _, err = engine.RebuildHashmap(proc, process.NewAnalyzer(0, false, false, "test")) - require.NoError(t, err) - if jm != nil { - jm.Free() - } - } - engine.Cleanup(proc) -} - -func TestReSpillDepthLimit(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget, err := process.NewHashBuildBudget(64<<10, 64<<10) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - - bat := makeInt32Batch(proc, []int32{1, 2, 3, 4, 5}) - fd := writeBuildFile(proc, "test_depth_build", bat) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - SpillThreshold: 1, - Budget: generation, - }) - engine.InitFromSpilledMap([]*os.File{fd}) - engine.buckets[0].Depth = SpillMaxPass - - analyzer := process.NewAnalyzer(0, false, false, "test") - jm, res, err := engine.RebuildHashmap(proc, analyzer) - require.Error(t, err, "depth limit must not force an over-budget hashmap build") - require.Equal(t, BucketSkip, res) - require.Nil(t, jm) - engine.Cleanup(proc) -} - -func TestReSpillWithProbe(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - vals := make([]int32, 5000) - for i := range vals { - vals[i] = int32(i) - } - bat := makeInt32Batch(proc, vals) - fd := writeBuildFile(proc, "test_rsp_probe_build", bat) - - // Manually write probe file. - probeFd := writeBuildFile(proc, "test_rsp_probe", makeInt32Batch(proc, []int32{100, 200, 300})) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - SpillThreshold: 100, - }) - engine.InitFromSpilledMap([]*os.File{fd}) - engine.buckets[0].ProbeFd = message.NewSpillFile(probeFd, 0, 0, nil) - - // Set probeKeyEval so scatterProbe works during re-spill. - engine.probeKeyEval = makeTestEvalKeysFn() - - analyzer := process.NewAnalyzer(0, false, false, "test") - jm, res, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.Equal(t, BucketReSpilled, res) - require.Nil(t, jm) - - for engine.HasMoreBuckets() { - jm2, _, err := engine.RebuildHashmap(proc, analyzer) - if jm2 != nil { - jm2.Free() - } - if err != nil { - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - break - } - } - - engine.Cleanup(proc) -} - -func TestAdvanceToNextBucketReSpilled(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - vals := make([]int32, 5000) - for i := range vals { - vals[i] = int32(i) + values := make([]int32, 100) + for i := range values { + values[i] = int32(i) } - bat := makeInt32Batch(proc, vals) - fd := writeBuildFile(proc, "test_adv_re_build", bat) - - engine := NewSpillEngine(SpillEngineConfig{ + build := makeInt32Batch(proc, values) + probe := makeInt32Batch(proc, []int32{1}) + defer build.Clean(proc.Mp()) + defer probe.Clean(proc.Mp()) + engine := newExactTestSpillEngine(t, SpillEngineConfig{ BuildKeyExprs: makeTestKeyExpr(), + SpillThreshold: 1, NeedsBuildForEmptyProbe: true, - SpillThreshold: 100, - }) - engine.InitFromSpilledMap([]*os.File{fd}) - - analyzer := process.NewAnalyzer(0, false, false, "test") - - callbackCalled := false - ok, err := engine.AdvanceToNextBucket(proc, analyzer, func(jm *message.JoinMap, _ BucketResult) { - callbackCalled = true - if jm != nil { - jm.Free() - } - }) - require.NoError(t, err) - require.True(t, ok) - require.False(t, callbackCalled, "re-spill is consumed before the callback") - - for engine.HasMoreBuckets() { - ok, err := engine.AdvanceToNextBucket(proc, analyzer, func(jm *message.JoinMap, _ BucketResult) { - if jm != nil { - jm.Free() - } - }) - require.NoError(t, err) - if ok { - engine.FinishBucket() - } - } - // Test passes if the loop terminates without errors. - engine.Cleanup(proc) -} - -func TestBuilderMemSize(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - builder := &hashbuild.HashmapBuilder{} - err := builder.Prepare(makeTestKeyExpr(), -1, -1, nil, proc) - require.NoError(t, err) - - sz := builderMemSize(builder) - require.Equal(t, int64(0), sz) - - bat := makeInt32Batch(proc, []int32{1, 2, 3, 4, 5}) - err = builder.Batches.CopyIntoBatches(bat, proc) - require.NoError(t, err) - builder.InputBatchRowCount += bat.RowCount() - - sz2 := builderMemSize(builder) - require.Greater(t, sz2, int64(0), "size should grow after adding batches") - - builder.FreeHashMapAndBatches(proc) - builder.Free(proc) -} - -func TestBuilderMemSizeIncludesCompletedBatchesAndPartialTail(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - builder := &hashbuild.HashmapBuilder{} - fullValues := make([]int32, colexec.DefaultBatchSize) - full := makeInt32Batch(proc, fullValues) - partial := makeInt32Batch(proc, []int32{1, 2, 3}) - require.NoError(t, builder.Batches.CopyIntoBatches(full, proc)) - require.NoError(t, builder.Batches.CopyIntoBatches(partial, proc)) - require.Len(t, builder.Batches.Buf, 2) - require.Equal(t, colexec.DefaultBatchSize, builder.Batches.Buf[0].RowCount()) - require.Equal(t, 3, builder.Batches.Buf[1].RowCount()) - - want := builder.Batches.MemSize + int64(builder.Batches.Buf[1].Size()) - require.Equal(t, want, builderMemSize(builder)) - - full.Clean(proc.Mp()) - partial.Clean(proc.Mp()) - builder.FreeHashMapAndBatches(proc) - builder.Free(proc) -} - -func TestShouldReSpillBeforeRetainUsesPredictedBytesAndRows(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - builder := &hashbuild.HashmapBuilder{} - retained := makeInt32Batch(proc, []int32{1, 2, 3}) - next := makeInt32Batch(proc, []int32{4, 5}) - require.NoError(t, builder.Batches.CopyIntoBatches(retained, proc)) - builder.InputBatchRowCount = retained.RowCount() - builder.Batches.MemSize = 200_000 - - predictedBytes := builderMemSize(builder) + int64(next.Size()) - require.False(t, shouldReSpillBeforeRetain(builder, next, predictedBytes)) - require.True(t, shouldReSpillBeforeRetain(builder, next, predictedBytes-1)) - require.False(t, shouldReSpillBeforeRetain(builder, next, 6)) - require.True(t, shouldReSpillBeforeRetain(builder, next, 5)) - require.False(t, shouldReSpillBeforeRetain(builder, next, 0)) - - retained.Clean(proc.Mp()) - next.Clean(proc.Mp()) - builder.FreeHashMapAndBatches(proc) - builder.Free(proc) -} - -func TestRebuildScratchFloorBoundsCoalescedPhysicalBatch(t *testing.T) { - const budgetCap = uint64(64 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - builder := &hashbuild.HashmapBuilder{} - builder.SetBudget(generation) - require.NoError(t, builder.Prepare(makeTestKeyExpr(), -1, -1, nil, proc)) - firstRows := make([]int32, colexec.DefaultBatchSize/2) - first := makeInt32Batch(proc, firstRows) - second := makeInt32Batch(proc, firstRows) - require.NoError(t, builder.CopyBuildBatch(first, proc)) - builder.InputBatchRowCount = first.RowCount() - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - Budget: generation, - }) - analyzer := process.NewAnalyzer(0, false, false, "test") - require.NoError(t, engine.reserveRebuildScatterScratch(builder, second, analyzer)) - require.Positive(t, engine.scatterScratchFloor) - require.Equal(t, engine.scatterScratchFloor, engine.scatterScratchReservation.Size()) - require.NoError(t, builder.CopyBuildBatch(second, proc)) - builder.InputBatchRowCount += second.RowCount() - require.Len(t, builder.Batches.Buf, 1, "two half records must coalesce") - physical := builder.Batches.Buf[0] - transient, err := scatterTransientBudgetBytes(physical, true) - require.NoError(t, err) - growth, ok := engine.scatterCapacityGrowthBytes(physical.RowCount(), 1) - require.True(t, ok) - require.GreaterOrEqual(t, engine.scatterScratchFloor, transient+growth) - - require.NoError(t, engine.reconcileScatterScratch()) - require.Equal(t, engine.scatterScratchFloor, engine.scatterScratchReservation.Size(), - "repartition headroom must survive per-batch reconciliation") - extra := analyzer.GetOpStats().ExtraStats - require.Equal(t, int64(1), extra["JoinSpillRebuildScratchReserveCount"]) - require.Equal(t, spillStatInt64(engine.scatterScratchFloor), extra["JoinSpillRebuildScratchFloorBytes"]) - - first.Clean(proc.Mp()) - second.Clean(proc.Mp()) - builder.FreeHashMapAndBatches(proc) - builder.Free(proc) - engine.Cleanup(proc) - require.Zero(t, generation.Used()) -} - -func TestFinishBucket(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), + NeedsProbeForEmptyBuild: true, }) - engine.InitFromSpilledMap([]*os.File{nil}) - - fd, err := os.CreateTemp(t.TempDir(), "probe") - require.NoError(t, err) - engine.probeReader.ResetForFd(fd) - require.True(t, engine.IsProbing()) - - engine.FinishBucket() - require.False(t, engine.IsProbing()) - - engine.FinishBucket() - require.False(t, engine.IsProbing()) - - engine.Cleanup(proc) -} - -func TestCleanupSpillEngine(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := makeInt32Batch(proc, []int32{1, 2, 3}) - fd1 := writeBuildFile(proc, "test_cl_build", bat) - - var buf bytes.Buffer - spillfs, _ := proc.GetSpillFileService() - probeFile, _ := spillfs.CreateAndRemoveFile(context.Background(), "test_cl_probe") - pw := BucketWriter{Name: "test_cl_probe", Fd: probeFile} - FlushBucketBatch(proc, bat, &pw, &buf, nil) - fd2 := pw.HandOffFd() - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), + engine.InitFromSpilledFiles([]*message.SpillFile{ + newTestSpillFile(writeBuildFile(proc, t.Name()+"-build", build), int64(len(values))), }) - engine.InitFromSpilledMap([]*os.File{fd1}) - engine.probeReader.ResetForFd(fd2) - - engine.buildReadBatch = batch.NewOffHeapWithSize(0) - engine.probeReadBatch = batch.NewOffHeapWithSize(0) - - engine.keyExecs = make([]colexec.ExpressionExecutor, 1) - exec, _ := colexec.NewExpressionExecutor(proc, makeTestKeyExpr()[0]) - engine.keyExecs[0] = exec - + engine.buckets[0].ProbeFd = newTestSpillFile( + writeBuildFile(proc, t.Name()+"-probe", probe), + 0, + ) + engine.probeKeyEval = makeTestEvalKeysFn() + _, _, err := engine.RebuildHashmap( + proc, + process.NewAnalyzer(0, false, false, "test"), + ) + require.ErrorContains(t, err, "row count") engine.Cleanup(proc) - - require.False(t, engine.IsProbing()) - require.Nil(t, engine.buckets) - require.Nil(t, engine.buildReadBatch) - require.Nil(t, engine.probeReadBatch) - require.Nil(t, engine.keyExecs) - - b := make([]byte, 1) - _, err := fd1.Read(b) - require.Error(t, err, "fd1 should be closed") - _, err = fd2.Read(b) - require.Error(t, err, "fd2 should be closed") -} - -func TestScatterProbeFunctionUsesStoredEval(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - }) - - evalCalled := false - engine.probeKeyEval = func(bat *batch.Batch) ([]*vector.Vector, error) { - evalCalled = true - return []*vector.Vector{bat.Vecs[0]}, nil - } - - writers := MakeBucketWriters("test_scatter_func") - bat := makeInt32Batch(proc, []int32{5, 15, 25}) - - err := scatterProbe(proc, engine, bat, writers, 1, nil) - require.NoError(t, err) - require.True(t, evalCalled, "probeKeyEval must be used for scatterProbe") - - wantErr := errors.New("probe key evaluation failed") - engine.probeKeyEval = func(*batch.Batch) ([]*vector.Vector, error) { return nil, wantErr } - require.ErrorIs(t, scatterProbe(proc, engine, bat, writers, 1, nil), wantErr) - - for i := range writers { - writers[i].Close() - } } -func TestScatterPeakDoesNotDoubleChargeReservedSource(t *testing.T) { +func TestRebuildRejectsRowsWithoutFile(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - values := make([]int32, 8192) - for i := range values { - values[i] = int32(i) - } - bat := makeInt32Batch(proc, values) - defer bat.Clean(proc.Mp()) - - charged, err := scatterTransientBudgetBytes(bat, true) - require.NoError(t, err) - uncharged, err := scatterTransientBudgetBytes(bat, false) - require.NoError(t, err) - source := uint64(bat.Allocated()) - if size := uint64(bat.Size()); size > source { - source = size - } - require.Equal(t, source, uncharged-charged) - - emptyEngine := NewSpillEngine(SpillEngineConfig{}) - retained, ok := emptyEngine.scatterRetainedBytes() - require.True(t, ok) - growth, ok := emptyEngine.scatterCapacityGrowthBytes(bat.RowCount(), 1) - require.True(t, ok) - marshalSize, err := bat.MarshalBinarySize() - require.NoError(t, err) - capacity := source + retained + growth + charged + uint64(marshalSize+24) - budget, err := process.NewHashBuildBudget(capacity, capacity) - require.NoError(t, err) - generation, err := budget.OpenGeneration(capacity) - require.NoError(t, err) - defer generation.Close() - sourceReservation, err := generation.Reserve(source) - require.NoError(t, err) - defer sourceReservation.Release() - - engine := NewSpillEngine(SpillEngineConfig{Budget: generation}) - writers := MakeBucketWriters("test_scatter_charged_source") - defer func() { - for i := range writers { - writers[i].Close() - } + for _, bucket := range []SpillBucket{ + {BuildRows: 1}, + {ProbeRows: 1}, + } { + engine := newExactTestSpillEngine(t, SpillEngineConfig{ + BuildKeyExprs: makeTestKeyExpr(), + }) + engine.buckets = []SpillBucket{bucket} + _, _, err := engine.RebuildHashmap( + proc, + process.NewAnalyzer(0, false, false, "test"), + ) + require.ErrorContains(t, err, "file/row metadata") engine.Cleanup(proc) - }() - analyzer := process.NewAnalyzer(0, false, false, "test") - require.NoError(t, engine.scatterBatchBounded( - proc, bat, []*vector.Vector{bat.Vecs[0]}, writers, 0, true, analyzer, - )) -} - -func TestScatterCapacityGrowthChargesCompleteReplacement(t *testing.T) { - engine := NewSpillEngine(SpillEngineConfig{}) - engine.scatterHashValues = make([]uint64, 8) - engine.scatterBucketRowIds = make([]int32, 8) - engine.keyVecs = make([]*vector.Vector, 1) - growth, ok := engine.scatterCapacityGrowthBytes(9, 2) - require.True(t, ok) - require.Equal(t, uint64(9*8+9*4+2*8), growth) -} - -func TestScatterScratchLifecycle(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget, err := process.NewHashBuildBudget(64<<20, 64<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(64 << 20) - require.NoError(t, err) - defer generation.Close() - - engine := NewSpillEngine(SpillEngineConfig{Budget: generation}) - writers := MakeBucketWriters("test_scatter_scratch") - defer func() { - for i := range writers { - writers[i].Close() - } - }() - values := make([]int32, 8192) - for i := range values { - values[i] = int32(i) - } - bat := makeInt32Batch(proc, values) - defer bat.Clean(proc.Mp()) - analyzer := process.NewAnalyzer(0, false, false, "test") - keys := []*vector.Vector{bat.Vecs[0]} - require.NoError(t, engine.scatterBatchBounded(proc, bat, keys, writers, 0, false, analyzer)) - require.NotNil(t, engine.scatterScratchReservation) - firstHashCap := cap(engine.scatterHashValues) - firstRowIDCap := cap(engine.scatterBucketRowIds) - require.Equal(t, len(values), firstHashCap) - require.Equal(t, len(values), firstRowIDCap) - firstHash := &engine.scatterHashValues[0] - firstRowID := &engine.scatterBucketRowIds[0] - retained := generation.Used() - firstReserveCount := generation.ReserveCount() - require.NoError(t, engine.scatterBatchBounded(proc, bat, keys, writers, 0, false, analyzer)) - require.Greater(t, generation.ReserveCount(), firstReserveCount, "each batch peak must be admitted above retained scratch") - require.Equal(t, retained, generation.Used(), "batch peak must reconcile to retained scratch") - require.Equal(t, firstHashCap, cap(engine.scatterHashValues)) - require.Equal(t, firstRowIDCap, cap(engine.scatterBucketRowIds)) - require.Equal(t, firstHash, &engine.scatterHashValues[0]) - require.Equal(t, firstRowID, &engine.scatterBucketRowIds[0]) - - for i := range writers { - writers[i].Close() - } - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - for i := range engine.scatterWriteBuffers { - require.Zero(t, engine.scatterWriteBuffers[i].Cap()) } - // Cleanup is an idempotent terminal release point. - engine.Cleanup(proc) - require.Zero(t, generation.Used()) } -func TestScatterScratchRejectsPeakAboveRetainedBudget(t *testing.T) { +func TestReSpillOmitsUnusedBatchMetadata(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - budget, err := process.NewHashBuildBudget(1<<20, 1<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1 << 20) - require.NoError(t, err) - defer generation.Close() - engine := NewSpillEngine(SpillEngineConfig{Budget: generation}) - writers := MakeBucketWriters("test_scatter_peak_reject") - defer func() { - for i := range writers { - writers[i].Close() - } - }() - values := make([]int32, 8192) + values := make([]int32, 100) for i := range values { values[i] = int32(i) } - bat := makeInt32Batch(proc, values) - defer bat.Clean(proc.Mp()) - keys := []*vector.Vector{bat.Vecs[0]} - analyzer := process.NewAnalyzer(0, false, false, "test") - require.NoError(t, engine.scatterBatchBounded(proc, bat, keys, writers, 0, false, analyzer)) - rowCap := cap(engine.scatterBucketRowIds) - used := generation.Used() - require.Positive(t, used) - largerValues := make([]int32, len(values)*2) - for i := range largerValues { - largerValues[i] = int32(i) - } - larger := makeInt32Batch(proc, largerValues) - defer larger.Clean(proc.Mp()) - err = engine.scatterBatchBounded( - proc, larger, []*vector.Vector{larger.Vecs[0]}, writers, 0, false, analyzer, - ) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Equal(t, rowCap, cap(engine.scatterBucketRowIds), "rejection must precede new scratch allocation") - engine.Cleanup(proc) - require.Zero(t, generation.Used()) -} - -func TestScatterPhaseReleasesScratchKeepsSpillOwnership(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget, err := process.NewHashBuildBudget(64<<20, 64<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(64 << 20) - require.NoError(t, err) - defer generation.Close() - buildBat := makeInt32Batch(proc, []int32{1}) - buildFd := writeBuildFile(proc, "test_phase_build", buildBat) - defer buildBat.Clean(proc.Mp()) - engine := NewSpillEngine(SpillEngineConfig{ - ProbeKeyExprs: makeTestKeyExpr(), - NeedsProbeForEmptyBuild: true, - Budget: generation, + build := makeInt32Batch(proc, values) + build.Attrs = []string{"key"} + defer build.Clean(proc.Mp()) + engine := newExactTestSpillEngine(t, SpillEngineConfig{ + BuildKeyExprs: makeTestKeyExpr(), + NeedBatches: true, + SpillThreshold: 1, + NeedsBuildForEmptyProbe: true, }) - buildFds := make([]*os.File, SpillNumBuckets) - buildFds[0] = buildFd - engine.InitFromSpilledMap(buildFds) - probeBat := makeInt32Batch(proc, []int32{2, 2, 2, 2}) - defer probeBat.Clean(proc.Mp()) - childrenDone := false - analyzer := process.NewAnalyzer(0, false, false, "test") - err = engine.ScatterProbeTable(proc, func() (*batch.Batch, error) { - if childrenDone { - return nil, nil - } - childrenDone = true - return probeBat, nil - }, analyzer, func(bat *batch.Batch) ([]*vector.Vector, error) { - return []*vector.Vector{bat.Vecs[0]}, nil + engine.InitFromSpilledFiles([]*message.SpillFile{ + newTestSpillFile(writeBuildFile(proc, t.Name(), build), int64(len(values))), }) - require.NoError(t, err) - require.Zero(t, generation.Used(), "scatter memory scratch must end with the phase") - require.Nil(t, engine.scatterScratchReservation) - for i := range engine.scatterWriteBuffers { - require.Zero(t, engine.scatterWriteBuffers[i].Cap()) - } - require.Positive(t, generation.SpillDiskUsed(), "handed-off probe file keeps disk accounting") - require.Positive(t, generation.SpillFDUsed(), "handed-off probe file keeps FD accounting") - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) -} - -func TestScatterCoalescesAcrossBatchesUntilFlush(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - engine := NewSpillEngine(SpillEngineConfig{}) - writers := MakeBucketWriters("test_scatter_coalesce") - defer func() { - for i := range writers { - writers[i].Close() - } - }() - bat := makeInt32Batch(proc, []int32{1, 1, 1}) - defer bat.Clean(proc.Mp()) - keys := []*vector.Vector{bat.Vecs[0]} - analyzer := process.NewAnalyzer(0, false, false, "test") - require.NoError(t, engine.scatterBatchBounded(proc, bat, keys, writers, 0, false, analyzer)) - require.NoError(t, engine.scatterBatchBounded(proc, bat, keys, writers, 0, false, analyzer)) - var pending int - for i := range engine.scatterWriteBuffers { - pending += engine.scatterWriteBuffers[i].Len() - } - require.Positive(t, pending) - for i := range writers { - require.Zero(t, writers[i].Rows) - } - require.NoError(t, engine.flushScatterBuffers(proc, writers, analyzer)) - for i := range engine.scatterWriteBuffers { - require.Zero(t, engine.scatterWriteBuffers[i].Len()) - } - var rows int64 - for i := range writers { - rows += writers[i].Rows - } - require.Equal(t, int64(6), rows) -} - -func TestScatterCoalescedRecordRoundTrip(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - engine := NewSpillEngine(SpillEngineConfig{}) - writers := MakeBucketWriters("test_scatter_coalesce_roundtrip") - defer func() { - for i := range writers { - writers[i].Close() - } - }() - bat := makeInt32Batch(proc, []int32{7, 7, 7}) - defer bat.Clean(proc.Mp()) - keys := []*vector.Vector{bat.Vecs[0]} + engine.buckets[0].Depth = SpillMaxPass - 1 analyzer := process.NewAnalyzer(0, false, false, "test") - for i := 0; i < 3; i++ { - require.NoError(t, engine.scatterBatchBounded(proc, bat, keys, writers, 0, false, analyzer)) - } - require.NoError(t, engine.flushScatterBuffers(proc, writers, analyzer)) - var target *BucketWriter - for i := range writers { - if writers[i].Rows > 0 { - target = &writers[i] - break - } - } - require.NotNil(t, target) - require.Equal(t, int64(9), target.Rows) - _, err := target.Fd.Seek(0, io.SeekStart) + _, result, err := engine.RebuildHashmap(proc, analyzer) require.NoError(t, err) - reader := BucketReader{fd: target.Fd} - reuse := batch.NewOffHeapWithSize(0) - got, err := reader.ReadBatch(proc, reuse) - require.NoError(t, err) - require.Equal(t, 3, got.RowCount()) - got, err = reader.ReadBatch(proc, reuse) - require.NoError(t, err) - require.Equal(t, 3, got.RowCount()) - got, err = reader.ReadBatch(proc, reuse) - require.NoError(t, err) - require.Equal(t, 3, got.RowCount()) - _, err = reader.ReadBatch(proc, reuse) - require.ErrorIs(t, err, io.EOF) - reader.Close() -} - -func TestScatterCoalesceFlushErrorClearsPending(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - engine := NewSpillEngine(SpillEngineConfig{}) - writers := MakeBucketWriters("test_scatter_coalesce_error") - defer func() { - for i := range writers { - writers[i].Close() - } - }() - bat := makeInt32Batch(proc, []int32{11, 11, 11}) - defer bat.Clean(proc.Mp()) - keys := []*vector.Vector{bat.Vecs[0]} - analyzer := process.NewAnalyzer(0, false, false, "test") - require.NoError(t, engine.scatterBatchBounded(proc, bat, keys, writers, 0, false, analyzer)) - require.NoError(t, engine.flushScatterBuffers(proc, writers, analyzer)) - var target *BucketWriter - for i := range writers { - if writers[i].Rows > 0 { - target = &writers[i] - break - } - } - require.NotNil(t, target) - require.NoError(t, engine.scatterBatchBounded(proc, bat, keys, writers, 0, false, analyzer)) - require.Positive(t, engine.scatterWriteBuffers[targetIndex(writers, target)].Len()) - require.NoError(t, target.Fd.Close()) - require.Error(t, engine.flushScatterBuffers(proc, writers, analyzer)) - for i := range engine.scatterWriteBuffers { - require.Zero(t, engine.scatterWriteBuffers[i].Len()) - } -} + require.Equal(t, BucketReSpilled, result) -func targetIndex(writers []BucketWriter, target *BucketWriter) int { - for i := range writers { - if &writers[i] == target { - return i + found := false + for len(engine.buckets) > 0 { + jm, next, err := engine.RebuildHashmap(proc, analyzer) + require.NoError(t, err) + if next != BucketReady { + continue } - } - return -1 -} - -func TestScatterProbeTableOuterJoinKeepsProbe(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := makeInt32Batch(proc, []int32{10, 20, 30}) - fd1 := writeBuildFile(proc, "test_outer_build", bat) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsProbeForEmptyBuild: true, - }) - engine.InitFromSpilledMap([]*os.File{fd1, nil}) - - vals := make([]int32, 1000) - for i := range vals { - vals[i] = int32(i) - } - batches := []*batch.Batch{makeInt32Batch(proc, vals)} - idx := 0 - children := func() (*batch.Batch, error) { - if idx >= len(batches) { - return nil, nil + found = true + for _, bat := range jm.GetBatches() { + require.Empty(t, bat.Attrs) } - b := batches[idx] - idx++ - return b, nil + jm.Free() } - - analyzer := process.NewAnalyzer(0, false, false, "test") - err := engine.ScatterProbeTable(proc, children, analyzer, makeTestEvalKeysFn()) - require.NoError(t, err) - - require.NotNil(t, engine.buckets[0].ProbeFd) - require.NotNil(t, engine.buckets[1].ProbeFd, "outer join must keep probe for empty build") - + require.True(t, found) engine.Cleanup(proc) } func TestCleanupDoubleSafe(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - - engine := NewSpillEngine(SpillEngineConfig{ + engine := newExactTestSpillEngine(t, SpillEngineConfig{ BuildKeyExprs: makeTestKeyExpr(), }) - engine.InitFromSpilledMap([]*os.File{nil, nil, nil}) - + initTestSpillFiles(engine, []*os.File{nil, nil, nil}, 0, 0, 0) engine.Cleanup(proc) engine.Cleanup(proc) - require.Nil(t, engine.buckets) } - -// TestRebuildHashmapPrepareError covers the builder.Free(proc) path -// when HashmapBuilder.Prepare fails (e.g., with an invalid key expression). -func TestRebuildHashmapPrepareError(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := makeInt32Batch(proc, []int32{1, 2, 3}) - fd := writeBuildFile(proc, "test_prep_err", bat) - - // Use an expression that will fail in Prepare (nil Expr field). - badExpr := []*plan.Expr{{}} - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: badExpr, - NeedsBuildForEmptyProbe: true, - }) - engine.InitFromSpilledMap([]*os.File{fd}) - - callbackCalled := false - ok, err := engine.AdvanceToNextBucket(proc, process.NewAnalyzer(0, false, false, "test"), - func(*message.JoinMap, BucketResult) { callbackCalled = true }) - require.Error(t, err) - require.False(t, ok) - require.False(t, callbackCalled) - - engine.Cleanup(proc) -} diff --git a/pkg/sql/compile/allocation_account_lifecycle.go b/pkg/sql/compile/allocation_account_lifecycle.go index 2de7d883f4986..4b34b1452977c 100644 --- a/pkg/sql/compile/allocation_account_lifecycle.go +++ b/pkg/sql/compile/allocation_account_lifecycle.go @@ -25,20 +25,10 @@ import ( ) type executionAllocationAccountOwner interface { - AllocationAccountEnabled() bool SetAllocationAccount(*mpool.AllocationAccount) error ClearAllocationAccount(*mpool.AllocationAccount) error } -// executionAllocationAccountBlocker marks an operator whose physical owner is -// known but whose allocation-site closure is not yet complete. Automatic -// activation is statement-atomic: one blocker keeps every participating -// operator on the legacy path instead of creating a mixed exact/estimated -// generation. -type executionAllocationAccountBlocker interface { - AllocationAccountActivationBlocked() bool -} - // statementAllocationAttempt owns one local execution generation. The // MessageBoard pointer is captured at open so prepared/retry Reset cannot make // terminal cleanup drain a newer board. @@ -47,7 +37,11 @@ type statementAllocationAttempt struct { account *mpool.AllocationAccount board *message.MessageBoard exporter func(mpool.AllocationAccountTerminalSnapshot) + + ownersMu sync.Mutex owners []executionAllocationAccountOwner + ownerSet map[executionAllocationAccountOwner]struct{} + closing bool once sync.Once snapshot mpool.AllocationAccountTerminalSnapshot @@ -65,16 +59,24 @@ func (c *Compile) beginAllocationAccountAttempt() ( c.allocationTerminalExporter == nil { return nil, mpool.ErrAllocationAccountInvariant } - var controller mpool.AllocationCapacityController + owners := c.allocationAccountOwners var err error - if c.allocationControllerProvider != nil { - controller, err = c.allocationControllerProvider() + if owners == nil { + owners, err = collectAllocationAccountOwners(c.scopes) if err != nil { return nil, err } - if controller == nil { - return nil, mpool.ErrAllocationAccountInvariant - } + } + c.allocationAccountOwners = nil + if c.allocationControllerProvider == nil { + return nil, mpool.ErrAllocationAccountInvariant + } + controller, err := c.allocationControllerProvider() + if err != nil { + return nil, err + } + if controller == nil { + return nil, mpool.ErrAllocationAccountInvariant } account, err := c.allocationAccountRegistry.OpenWithController( c.allocationAccountLimit, @@ -83,7 +85,7 @@ func (c *Compile) beginAllocationAccountAttempt() ( if err != nil { return nil, err } - owners, err := configureAllocationAccountOwners(c.scopes, account) + owners, err = configureAllocationAccountOwners(owners, account) if err != nil { snapshot, first, finalizeErr := c.allocationAccountRegistry. CompleteTerminalWithError(account, err) @@ -101,24 +103,68 @@ func (c *Compile) beginAllocationAccountAttempt() ( board: c.MessageBoard, exporter: c.allocationTerminalExporter, owners: owners, + ownerSet: make(map[executionAllocationAccountOwner]struct{}, len(owners)), + } + for _, owner := range owners { + attempt.ownerSet[owner] = struct{}{} } c.allocationAttempt = attempt return attempt, nil } +// attachRuntimeOwners binds operators cloned after runOnce starts to the same +// attempt. Parallel scan/load workers are execution-local and do not exist +// when the template scopes are collected. +func (a *statementAllocationAttempt) attachRuntimeOwners(scopes []*Scope) error { + if a == nil || a.account == nil { + return mpool.ErrAllocationAccountInvariant + } + owners, err := collectAllocationAccountOwners(scopes) + if err != nil || len(owners) == 0 { + return err + } + + a.ownersMu.Lock() + defer a.ownersMu.Unlock() + if a.closing { + return mpool.ErrAllocationAccountInvariant + } + newOwners := make([]executionAllocationAccountOwner, 0, len(owners)) + for _, owner := range owners { + if _, exists := a.ownerSet[owner]; !exists { + newOwners = append(newOwners, owner) + } + } + configured, err := configureAllocationAccountOwners(newOwners, a.account) + if err != nil { + return err + } + for _, owner := range configured { + a.ownerSet[owner] = struct{}{} + } + a.owners = append(a.owners, configured...) + return nil +} + +func (c *Compile) attachRuntimeAllocationOwners(scopes []*Scope) error { + if c == nil { + return mpool.ErrAllocationAccountInvariant + } + if c.allocationAttempt == nil { + owners, err := collectAllocationAccountOwners(scopes) + if err != nil || len(owners) == 0 { + return err + } + return mpool.ErrAllocationAccountInvariant + } + return c.allocationAttempt.attachRuntimeOwners(scopes) +} + func configureAllocationAccountOwners( - scopes []*Scope, + owners []executionAllocationAccountOwner, account *mpool.AllocationAccount, ) ([]executionAllocationAccountOwner, error) { - var configured []executionAllocationAccountOwner - isConfigured := func(candidate executionAllocationAccountOwner) bool { - for _, owner := range configured { - if owner == candidate { - return true - } - } - return false - } + configured := make([]executionAllocationAccountOwner, 0, len(owners)) rollback := func(cause error) error { for i := len(configured) - 1; i >= 0; i-- { cause = errors.Join( @@ -128,98 +174,75 @@ func configureAllocationAccountOwners( } return cause } - var configure func(*Scope) error - configure = func(scope *Scope) error { - if scope == nil { - return nil - } - if err := vm.HandleAllOp( - scope.RootOp, - func(_ vm.Operator, op vm.Operator) error { - if blocker, ok := op.(executionAllocationAccountBlocker); ok && - blocker.AllocationAccountActivationBlocked() { - return mpool.ErrAllocationAccountInvariant - } - if owner, ok := op.(executionAllocationAccountOwner); ok && - owner.AllocationAccountEnabled() { - if isConfigured(owner) { - return nil - } - if err := owner.SetAllocationAccount(account); err != nil { - return err - } - configured = append(configured, owner) - } - return nil - }, - ); err != nil { - return err - } - for _, preScope := range scope.PreScopes { - if err := configure(preScope); err != nil { - return err - } - } - return nil - } - for _, scope := range scopes { - if err := configure(scope); err != nil { + for _, owner := range owners { + if err := owner.SetAllocationAccount(account); err != nil { return nil, rollback(err) } + configured = append(configured, owner) } return configured, nil } -func hasAllocationAccountOwner(scopes []*Scope) bool { - found, blocked := false, false - var inspect func(*Scope) - inspect = func(scope *Scope) { +func collectAllocationAccountOwners( + scopes []*Scope, +) ([]executionAllocationAccountOwner, error) { + owners := make([]executionAllocationAccountOwner, 0) + seen := make(map[executionAllocationAccountOwner]struct{}) + var inspect func(*Scope) error + inspect = func(scope *Scope) error { if scope == nil { - return + return nil } - _ = vm.HandleAllOp(scope.RootOp, func(_ vm.Operator, op vm.Operator) error { - if blocker, ok := op.(executionAllocationAccountBlocker); ok && - blocker.AllocationAccountActivationBlocked() { - blocked = true - } - if owner, ok := op.(executionAllocationAccountOwner); ok && - owner.AllocationAccountEnabled() { - found = true + if err := vm.HandleAllOp(scope.RootOp, func(_ vm.Operator, op vm.Operator) error { + if owner, ok := op.(executionAllocationAccountOwner); ok { + if _, exists := seen[owner]; !exists { + seen[owner] = struct{}{} + owners = append(owners, owner) + } } return nil - }) + }); err != nil { + return err + } for _, preScope := range scope.PreScopes { - inspect(preScope) + if err := inspect(preScope); err != nil { + return err + } } + return nil } for _, scope := range scopes { - inspect(scope) + if err := inspect(scope); err != nil { + return nil, err + } } - return found && !blocked + return owners, nil } -// ensureAllocationAccountLifecycle activates accounting only when the physical -// plan contains a complete migrated owner. Legacy plans never open a registry -// slot or initialize the HashBuild budget. +// ensureAllocationAccountLifecycle installs one account whenever the physical +// plan contains a HashBuild/join allocation owner. Implementing the owner +// contract is the boundary: there is no per-owner activation switch. func (c *Compile) ensureAllocationAccountLifecycle( exporter func(mpool.AllocationAccountTerminalSnapshot), ) error { if c == nil { return nil } - if !hasAllocationAccountOwner(c.scopes) { - if c.allocationLifecycleAutomatic { + owners, err := collectAllocationAccountOwners(c.scopes) + if err != nil { + return err + } + c.allocationAccountOwners = owners + if len(owners) == 0 { + c.allocationAccountOwners = nil + if c.allocationControllerProvider != nil { c.allocationAccountRegistry = nil c.allocationAccountLimit = 0 c.allocationControllerProvider = nil c.allocationTerminalExporter = nil - c.allocationLifecycleAutomatic = false } return nil } - if c.allocationAccountRegistry != nil && !c.allocationLifecycleAutomatic { - return nil - } if exporter == nil { return mpool.ErrAllocationAccountInvariant } @@ -238,18 +261,15 @@ func (c *Compile) ensureAllocationAccountLifecycle( if limit == 0 { return mpool.ErrAllocationAccountInvariant } - c.ConfigureAllocationAccountLifecycleWithController( - registry, - limit, - func() (mpool.AllocationCapacityController, error) { - if budget.Closed() { - return nil, process.ErrHashBuildBudgetClosed - } - return budget, nil - }, - exporter, - ) - c.allocationLifecycleAutomatic = true + c.allocationAccountRegistry = registry + c.allocationAccountLimit = limit + c.allocationControllerProvider = func() (mpool.AllocationCapacityController, error) { + if budget.Closed() { + return nil, process.ErrHashBuildBudgetClosed + } + return budget, nil + } + c.allocationTerminalExporter = exporter return nil } @@ -265,13 +285,18 @@ func (a *statementAllocationAttempt) finish() ( // before this point. Draining the board first releases queued JoinMap // and spill payload ownership through their normal Destroy methods. a.board.CloseAndDrain() - for i := len(a.owners) - 1; i >= 0; i-- { + a.ownersMu.Lock() + a.closing = true + owners := a.owners + a.owners = nil + a.ownerSet = nil + a.ownersMu.Unlock() + for i := len(owners) - 1; i >= 0; i-- { a.err = errors.Join( a.err, - a.owners[i].ClearAllocationAccount(a.account), + owners[i].ClearAllocationAccount(a.account), ) } - a.owners = nil var first bool var terminalErr error a.snapshot, first, terminalErr = a.registry.CompleteTerminalWithError( @@ -300,11 +325,8 @@ func (c *Compile) copyAllocationAccountLifecycleTo(dst *Compile) { if c == nil || dst == nil { return } - dst.ConfigureAllocationAccountLifecycle( - c.allocationAccountRegistry, - c.allocationAccountLimit, - c.allocationTerminalExporter, - ) + dst.allocationAccountRegistry = c.allocationAccountRegistry + dst.allocationAccountLimit = c.allocationAccountLimit + dst.allocationTerminalExporter = c.allocationTerminalExporter dst.allocationControllerProvider = c.allocationControllerProvider - dst.allocationLifecycleAutomatic = c.allocationLifecycleAutomatic } diff --git a/pkg/sql/compile/allocation_account_lifecycle_test.go b/pkg/sql/compile/allocation_account_lifecycle_test.go index a478bbafb6c87..2d39ebfa133a4 100644 --- a/pkg/sql/compile/allocation_account_lifecycle_test.go +++ b/pkg/sql/compile/allocation_account_lifecycle_test.go @@ -29,8 +29,10 @@ import ( "github.com/matrixorigin/matrixone/pkg/pb/txn" "github.com/matrixorigin/matrixone/pkg/sql/colexec" "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/product" "github.com/matrixorigin/matrixone/pkg/testutil" "github.com/matrixorigin/matrixone/pkg/vm" + "github.com/matrixorigin/matrixone/pkg/vm/engine" "github.com/matrixorigin/matrixone/pkg/vm/message" "github.com/matrixorigin/matrixone/pkg/vm/process" "github.com/stretchr/testify/require" @@ -38,7 +40,8 @@ import ( type allocationLifecycleErrorOperator struct { *colexec.MockOperator - err error + err error + account *mpool.AllocationAccount } type allocationLifecycleOwnerOperator struct { @@ -46,7 +49,6 @@ type allocationLifecycleOwnerOperator struct { account *mpool.AllocationAccount failSet bool failClear bool - blocked bool clears int released bool releaseSawLiveAccount bool @@ -57,14 +59,6 @@ func (op *allocationLifecycleOwnerOperator) Release() { op.releaseSawLiveAccount = op.account != nil } -func (op *allocationLifecycleOwnerOperator) AllocationAccountEnabled() bool { - return true -} - -func (op *allocationLifecycleOwnerOperator) AllocationAccountActivationBlocked() bool { - return op.blocked -} - func (op *allocationLifecycleOwnerOperator) SetAllocationAccount( account *mpool.AllocationAccount, ) error { @@ -101,9 +95,36 @@ func (op *allocationLifecycleErrorOperator) Call( return vm.CancelResult, op.err } +func (op *allocationLifecycleErrorOperator) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if op.account != nil && op.account != account { + return mpool.ErrAllocationAccountMismatch + } + op.account = account + return nil +} + +func (op *allocationLifecycleErrorOperator) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if op.account != account { + return mpool.ErrAllocationAccountMismatch + } + op.account = nil + return nil +} + +type allocationLifecycleTestController struct{} + +func (*allocationLifecycleTestController) AcquireAllocationCapacity(uint64) error { + return nil +} + +func (*allocationLifecycleTestController) ReleaseAllocationCapacity(uint64) {} + func newRunLifecycleCompile( t *testing.T, - exporter func(mpool.AllocationAccountTerminalSnapshot), ) (*Compile, *mpool.AllocationAccountRegistry) { t.Helper() proc := testutil.NewProcess(t) @@ -130,9 +151,10 @@ func newRunLifecycleCompile( ) c.pn = &plan.Plan{Plan: &plan.Plan_Query{Query: &plan.Query{}}} c.anal = newAnalyzeModule() - registry, err := mpool.NewAllocationAccountRegistry(2, 2) + budget, err := proc.GetHashBuildBudget() + require.NoError(t, err) + registry, err := budget.AllocationAccountRegistry() require.NoError(t, err) - c.ConfigureAllocationAccountLifecycle(registry, 1<<20, exporter) return c, registry } @@ -143,10 +165,13 @@ func newTestAllocationLifecycleCompile( ) *Compile { t.Helper() return &Compile{ - proc: testutil.NewProcess(t), - MessageBoard: message.NewMessageBoard(), - allocationAccountRegistry: registry, - allocationAccountLimit: 1 << 20, + proc: testutil.NewProcess(t), + MessageBoard: message.NewMessageBoard(), + allocationAccountRegistry: registry, + allocationAccountLimit: 1 << 20, + allocationControllerProvider: func() (mpool.AllocationCapacityController, error) { + return &allocationLifecycleTestController{}, nil + }, allocationTerminalExporter: exporter, } } @@ -344,31 +369,106 @@ func TestStatementAllocationAttemptOwnerConfigurationRollsBack(t *testing.T) { require.NoError(t, err) } -func TestAllocationAccountActivationIsStatementAtomic(t *testing.T) { - eligible := &allocationLifecycleOwnerOperator{ +func TestAllocationAccountConfiguresEveryStatementOwner(t *testing.T) { + first := &allocationLifecycleOwnerOperator{ MockOperator: colexec.NewMockOperator(), } - blocker := &allocationLifecycleOwnerOperator{ + second := &allocationLifecycleOwnerOperator{ MockOperator: colexec.NewMockOperator(), - blocked: true, } - scopes := []*Scope{{RootOp: eligible}, {RootOp: blocker}} - require.False(t, hasAllocationAccountOwner(scopes), - "one unclosed physical owner must keep the whole statement legacy") + scopes := []*Scope{{RootOp: first}, {RootOp: second}} + owners, err := collectAllocationAccountOwners(scopes) + require.NoError(t, err) + require.Len(t, owners, 2) - registry, err := mpool.NewAllocationAccountRegistry(1, 1) + registry, err := mpool.NewAllocationAccountRegistry(1, 8) require.NoError(t, err) account, err := registry.Open(1 << 20) require.NoError(t, err) - configured, err := configureAllocationAccountOwners(scopes, account) - require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) - require.Nil(t, configured) - require.Nil(t, eligible.account) - require.Equal(t, 1, eligible.clears) + configured, err := configureAllocationAccountOwners(owners, account) + require.NoError(t, err) + require.Len(t, configured, 2) + require.Same(t, account, first.account) + require.Same(t, account, second.account) + for i := len(configured) - 1; i >= 0; i-- { + require.NoError(t, configured[i].ClearAllocationAccount(account)) + } _, _, err = registry.CompleteTerminal(account) require.NoError(t, err) } +func TestAllocationAccountCollectsProductConsumerAndHashBuild(t *testing.T) { + consumer := product.NewArgument() + producer := hashbuild.NewArgument() + scopes := []*Scope{{ + RootOp: consumer, + PreScopes: []*Scope{ + {RootOp: producer}, + }, + }} + owners, err := collectAllocationAccountOwners(scopes) + require.NoError(t, err) + require.Len(t, owners, 2) + + registry, err := mpool.NewAllocationAccountRegistry(1, 32) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + configured, err := configureAllocationAccountOwners(owners, account) + require.NoError(t, err) + require.Len(t, configured, 2) + for i := len(configured) - 1; i >= 0; i-- { + require.NoError(t, configured[i].ClearAllocationAccount(account)) + } + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + consumer.Release() + producer.Release() +} + +func TestParallelRuntimeClonesJoinAllocationAttempt(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + var exported []mpool.AllocationAccountTerminalSnapshot + c := newTestAllocationLifecycleCompile(t, registry, func( + snapshot mpool.AllocationAccountTerminalSnapshot, + ) { + exported = append(exported, snapshot) + }) + template := hashbuild.NewArgument() + template.NeedHashMap = false + source := &Scope{ + RootOp: template, + Proc: c.proc, + NodeInfo: engine.Node{ + Mcpu: 2, + }, + } + c.scopes = []*Scope{source} + attempt, err := c.beginAllocationAccountAttempt() + require.NoError(t, err) + require.NotNil(t, attempt) + + parallel, workers := newParallelScope(source) + require.Len(t, workers, 2) + for _, worker := range workers { + hb := worker.RootOp.(*hashbuild.HashBuild) + require.ErrorIs(t, hb.Prepare(worker.Proc), mpool.ErrAllocationAccountInvalid) + } + require.NoError(t, c.attachRuntimeAllocationOwners(workers)) + for _, worker := range workers { + hb := worker.RootOp.(*hashbuild.HashBuild) + require.NoError(t, hb.Prepare(worker.Proc)) + } + + require.NoError(t, c.finishAllocationAccountAttempt()) + require.Len(t, exported, 1) + require.Equal(t, mpool.AllocationAccountTerminalValid, exported[0].State) + require.Zero(t, exported[0].Used) + parallel.release() + template.Release() +} + func TestStatementAllocationAttemptOwnerTeardownFailureExportsFailure(t *testing.T) { registry, err := mpool.NewAllocationAccountRegistry(1, 1) require.NoError(t, err) @@ -446,7 +546,7 @@ func TestCompileAutomaticallyActivatesCompleteHashTableOwner(t *testing.T) { mpool.AllocationAccountTerminalSnapshot, ) { })) - require.True(t, c.allocationLifecycleAutomatic) + require.NotNil(t, c.allocationControllerProvider) require.NotNil(t, c.allocationAccountRegistry) attempt, err := c.beginAllocationAccountAttempt() require.NoError(t, err) @@ -459,71 +559,56 @@ func TestCompileAutomaticallyActivatesCompleteHashTableOwner(t *testing.T) { } func TestCompileRunFinalizesAllocationAttemptOnCancellation(t *testing.T) { - var snapshots []mpool.AllocationAccountTerminalSnapshot - c, registry := newRunLifecycleCompile(t, func( - snapshot mpool.AllocationAccountTerminalSnapshot, - ) { - snapshots = append(snapshots, snapshot) - }) + c, registry := newRunLifecycleCompile(t) // The canceled outer context is observed after runOnce, after the // allocation generation has opened. canceled, cancel := context.WithCancel(context.Background()) cancel() c.proc.ReplaceTopCtx(canceled) - c.scopes = []*Scope{newScope(magicType(255))} + scope := newScope(magicType(255)) + owner := &allocationLifecycleOwnerOperator{MockOperator: colexec.NewMockOperator()} + scope.RootOp = owner + c.scopes = []*Scope{scope} _, err := c.Run(0) require.ErrorIs(t, err, context.Canceled) - require.Len(t, snapshots, 1) - require.Equal(t, mpool.AllocationAccountTerminalValid, snapshots[0].State) - require.Zero(t, snapshots[0].Used) - _, ok := registry.Resolve(snapshots[0].Handle) - require.False(t, ok) + require.Nil(t, owner.account) + require.Zero(t, registry.LiveAllocationMetadata()) c.Release() } func TestCompileRunFinalizesAllocationAttemptOnExecutionError(t *testing.T) { - var snapshots []mpool.AllocationAccountTerminalSnapshot - c, registry := newRunLifecycleCompile(t, func( - snapshot mpool.AllocationAccountTerminalSnapshot, - ) { - snapshots = append(snapshots, snapshot) - }) + c, registry := newRunLifecycleCompile(t) executionErr := moerr.NewInternalErrorNoCtx("allocation lifecycle test") scope := newScope(Normal) scope.Proc = c.proc.NewNoContextChildProc(0) - scope.RootOp = &allocationLifecycleErrorOperator{ + owner := &allocationLifecycleErrorOperator{ MockOperator: colexec.NewMockOperator(), err: executionErr, } + scope.RootOp = owner c.scopes = []*Scope{scope} _, err := c.Run(0) require.ErrorIs(t, err, executionErr) - require.Len(t, snapshots, 1) - require.Equal(t, mpool.AllocationAccountTerminalValid, snapshots[0].State) - _, ok := registry.Resolve(snapshots[0].Handle) - require.False(t, ok) + require.Nil(t, owner.account) + require.Zero(t, registry.LiveAllocationMetadata()) c.Release() } func TestCompileRunFinalizesAllocationAttemptOnPanic(t *testing.T) { - var snapshots []mpool.AllocationAccountTerminalSnapshot - c, registry := newRunLifecycleCompile(t, func( - snapshot mpool.AllocationAccountTerminalSnapshot, - ) { - snapshots = append(snapshots, snapshot) - }) - c.scopes = []*Scope{newScope(magicType(255))} + c, registry := newRunLifecycleCompile(t) + scope := newScope(magicType(255)) + owner := &allocationLifecycleOwnerOperator{MockOperator: colexec.NewMockOperator()} + scope.RootOp = owner + c.scopes = []*Scope{scope} // Force the panic after beginAllocationAccountAttempt and before runOnce. c.lockMeta = nil require.Panics(t, func() { _, _ = c.Run(0) }) - require.Len(t, snapshots, 1) - require.Equal(t, mpool.AllocationAccountTerminalValid, snapshots[0].State) - _, ok := registry.Resolve(snapshots[0].Handle) - require.False(t, ok) + require.Nil(t, owner.account) + require.Zero(t, registry.LiveAllocationMetadata()) c.Release() } diff --git a/pkg/sql/compile/compile.go b/pkg/sql/compile/compile.go index a1397dd9af400..aa349b82213cc 100644 --- a/pkg/sql/compile/compile.go +++ b/pkg/sql/compile/compile.go @@ -327,7 +327,7 @@ func (c *Compile) clear() { c.allocationAccountLimit = 0 c.allocationControllerProvider = nil c.allocationTerminalExporter = nil - c.allocationLifecycleAutomatic = false + c.allocationAccountOwners = nil c.allocationAttempt = nil c.isPrepare = false c.hasMergeOp = false @@ -628,11 +628,7 @@ func (c *Compile) prePipelineInitializer() (err error) { func newMaterializedSpillBudget(proc *process.Process) materialized.SpillBudget { return materialized.SpillBudget{ ReserveMemory: func(size uint64) (materialized.Reservation, error) { - budget, err := proc.GetHashBuildBudget() - if err != nil { - return nil, err - } - return budget.Reserve(size) + return proc.GetCTEMemoryBudget().Reserve(proc.Ctx, size) }, ReserveDisk: func(size uint64) (materialized.GrowingReservation, error) { budget, err := proc.GetHashBuildBudget() @@ -7000,35 +6996,6 @@ func (c *Compile) SetResourceAttemptOwnerEligible() { c.resourceAttemptOwnerEligible = true } -// ConfigureAllocationAccountLifecycle installs the generation provider used -// by allocation-accounted owners. A nil registry keeps production on the -// legacy path and opens no generation. -func (c *Compile) ConfigureAllocationAccountLifecycle( - registry *mpool.AllocationAccountRegistry, - limit uint64, - exporter func(mpool.AllocationAccountTerminalSnapshot), -) { - c.ConfigureAllocationAccountLifecycleWithController( - registry, - limit, - nil, - exporter, - ) -} - -func (c *Compile) ConfigureAllocationAccountLifecycleWithController( - registry *mpool.AllocationAccountRegistry, - limit uint64, - controllerProvider func() (mpool.AllocationCapacityController, error), - exporter func(mpool.AllocationAccountTerminalSnapshot), -) { - c.allocationAccountRegistry = registry - c.allocationAccountLimit = limit - c.allocationControllerProvider = controllerProvider - c.allocationTerminalExporter = exporter - c.allocationLifecycleAutomatic = false -} - func (c *Compile) SetBuildPlanFunc(buildPlanFunc func(ctx context.Context) (*plan2.Plan, error)) { c.buildPlanFunc = buildPlanFunc } diff --git a/pkg/sql/compile/compile_test.go b/pkg/sql/compile/compile_test.go index 7cc7877d0b5b6..1387b977ed0b3 100644 --- a/pkg/sql/compile/compile_test.go +++ b/pkg/sql/compile/compile_test.go @@ -128,16 +128,6 @@ func TestCompileRunPreservesBinaryPrepareParamAcrossRetries(t *testing.T) { } c := NewCompile("test", "test", "select ?", "", "", newStubEngine(), proc, stmts[0], false, nil, time.Now()) - registry, err := mpool.NewAllocationAccountRegistry(4, 4) - require.NoError(t, err) - var terminalSnapshots []mpool.AllocationAccountTerminalSnapshot - c.ConfigureAllocationAccountLifecycle( - registry, - 1<<20, - func(snapshot mpool.AllocationAccountTerminalSnapshot) { - terminalSnapshots = append(terminalSnapshots, snapshot) - }, - ) require.NoError(t, c.Compile(ctx, pn, fill)) _, err = c.Run(0) require.NoError(t, err) @@ -146,19 +136,6 @@ func TestCompileRunPreservesBinaryPrepareParamAcrossRetries(t *testing.T) { require.Zero(t, params.Length()) require.Nil(t, params.GetData()) require.Nil(t, params.GetArea()) - require.Len(t, terminalSnapshots, 3) - seenHandles := make(map[mpool.AllocationAccountHandle]struct{}, 3) - for _, snapshot := range terminalSnapshots { - require.Equal(t, mpool.AllocationAccountTerminalValid, snapshot.State) - require.Zero(t, snapshot.Used) - require.True(t, snapshot.Sealed) - _, duplicate := seenHandles[snapshot.Handle] - require.False(t, duplicate) - seenHandles[snapshot.Handle] = struct{}{} - _, ok := registry.Resolve(snapshot.Handle) - require.False(t, ok) - } - c.Release() proc.Free() proc.GetSessionInfo().Buf.Free() diff --git a/pkg/sql/compile/operator.go b/pkg/sql/compile/operator.go index 165db85aee2e0..e6c82daed1adb 100644 --- a/pkg/sql/compile/operator.go +++ b/pkg/sql/compile/operator.go @@ -27,7 +27,6 @@ import ( "github.com/gogo/protobuf/proto" "github.com/google/uuid" "github.com/matrixorigin/matrixone/pkg/catalog" - "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/common/sqlquote" "github.com/matrixorigin/matrixone/pkg/container/batch" @@ -130,14 +129,14 @@ func mergeReceiverChannelBufferSize(s *Scope) int { type operatorDupContext struct { shufflePools map[*shuffle.Shuffle]*shuffle.ShufflePool - hashJoinChannels map[*hashjoin.HashJoin]chan *bitmap.Bitmap + hashJoinMailboxes map[*hashjoin.HashJoin]*hashjoin.BitmapMailbox dedupJoinMailboxes map[*dedupjoin.DedupJoin]*dedupjoin.WorkerJoinMailbox } func newOperatorDupContext() *operatorDupContext { return &operatorDupContext{ shufflePools: make(map[*shuffle.Shuffle]*shuffle.ShufflePool), - hashJoinChannels: make(map[*hashjoin.HashJoin]chan *bitmap.Bitmap), + hashJoinMailboxes: make(map[*hashjoin.HashJoin]*hashjoin.BitmapMailbox), dedupJoinMailboxes: make(map[*dedupjoin.DedupJoin]*dedupjoin.WorkerJoinMailbox), } } @@ -230,12 +229,12 @@ func dupOperatorWithContext(sourceOp vm.Operator, index int, maxParallel int, du op.CanSkipProbe = t.CanSkipProbe op.IsShuffle = t.IsShuffle if !t.IsShuffle { - channel := dupCtx.hashJoinChannels[t] - if channel == nil { - channel = make(chan *bitmap.Bitmap, maxParallel) - dupCtx.hashJoinChannels[t] = channel + mailbox := dupCtx.hashJoinMailboxes[t] + if mailbox == nil { + mailbox = hashjoin.NewBitmapMailbox(maxParallel) + dupCtx.hashJoinMailboxes[t] = mailbox } - op.Channel = channel + op.Mailbox = mailbox op.NumCPU = uint64(maxParallel) op.IsMerger = (index == 0) } diff --git a/pkg/sql/compile/operator_test.go b/pkg/sql/compile/operator_test.go index e2e19ef55291d..ba15db4ea7119 100644 --- a/pkg/sql/compile/operator_test.go +++ b/pkg/sql/compile/operator_test.go @@ -19,7 +19,6 @@ import ( "math" "testing" - "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" @@ -599,21 +598,21 @@ func TestDupOperatorDedupJoinSharesMailboxOnlyWithinGeneration(t *testing.T) { require.NotSame(t, dup1.Mailbox, nextGeneration.Mailbox) } -func TestDupOperatorHashJoinSharesChannelOnlyWithinGeneration(t *testing.T) { +func TestDupOperatorHashJoinSharesMailboxOnlyWithinGeneration(t *testing.T) { op := hashjoin.NewArgument() - staleChannel := make(chan *bitmap.Bitmap, 2) - close(staleChannel) - op.Channel = staleChannel + staleMailbox := hashjoin.NewBitmapMailbox(2) + staleMailbox.SealAndDrain(mpool.MustNewZero()) + op.Mailbox = staleMailbox dupCtx := newOperatorDupContext() dup1 := dupOperatorWithContext(op, 0, 2, dupCtx).(*hashjoin.HashJoin) dup2 := dupOperatorWithContext(op, 1, 2, dupCtx).(*hashjoin.HashJoin) - require.Equal(t, staleChannel, op.Channel, "duplicating must not mutate the reusable template") - require.NotEqual(t, staleChannel, dup1.Channel, "a stale closed template channel must not enter a new execution") - require.Equal(t, dup1.Channel, dup2.Channel) + require.Same(t, staleMailbox, op.Mailbox, "duplicating must not mutate the reusable template") + require.NotSame(t, staleMailbox, dup1.Mailbox, "a stale template mailbox must not enter a new execution") + require.Same(t, dup1.Mailbox, dup2.Mailbox) nextGeneration := dupOperatorWithContext(op, 0, 2, newOperatorDupContext()).(*hashjoin.HashJoin) - require.NotEqual(t, dup1.Channel, nextGeneration.Channel) + require.NotSame(t, dup1.Mailbox, nextGeneration.Mailbox) } func TestDupOperatorAssignsSharedShuffleConsumerIndex(t *testing.T) { diff --git a/pkg/sql/compile/scope.go b/pkg/sql/compile/scope.go index 37d329b63c001..3b11a1d798dcf 100644 --- a/pkg/sql/compile/scope.go +++ b/pkg/sql/compile/scope.go @@ -751,6 +751,10 @@ func buildLoadParallelRun(s *Scope, c *Compile) (*Scope, error) { return nil, err } } + if err := c.attachRuntimeAllocationOwners(ss); err != nil { + s.discardParallelGeneration(ms) + return nil, err + } return ms, nil } @@ -802,6 +806,10 @@ func buildScanParallelRun(s *Scope, c *Compile) (*Scope, error) { RecvMsgList: recvMsgList, } } + if err := c.attachRuntimeAllocationOwners(ss); err != nil { + s.discardParallelGeneration(ms) + return nil, err + } return ms, nil } diff --git a/pkg/sql/compile/types.go b/pkg/sql/compile/types.go index cd7a13db799ec..05c75669e54a4 100644 --- a/pkg/sql/compile/types.go +++ b/pkg/sql/compile/types.go @@ -348,7 +348,7 @@ type Compile struct { allocationAccountLimit uint64 allocationControllerProvider func() (mpool.AllocationCapacityController, error) allocationTerminalExporter func(mpool.AllocationAccountTerminalSnapshot) - allocationLifecycleAutomatic bool + allocationAccountOwners []executionAllocationAccountOwner allocationAttempt *statementAllocationAttempt hasMergeOp bool diff --git a/pkg/sql/plan/function/baseTemplate.go b/pkg/sql/plan/function/baseTemplate.go index c56c4ece2fd38..ad6f163f554cf 100644 --- a/pkg/sql/plan/function/baseTemplate.go +++ b/pkg/sql/plan/function/baseTemplate.go @@ -58,14 +58,8 @@ func generalFunctionTemplateFactor[T1 templateTp1, T2 templateTr1]( return func(parameters []*vector.Vector, result vector.FunctionResultWrapper, _ *process.Process, length int, selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[T2](result) - p1, err := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) - if err != nil { - return err - } - p2, err := vector.OptGetParamFromWrapper[T1](rs, 1, parameters[1]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) + p2 := vector.OptGetParamFromWrapper[T1](rs, 1, parameters[1]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[T2](rsVec) @@ -161,14 +155,8 @@ func generalFunctionTemplateFactor[T1 templateTp1, T2 templateTr1]( return func(parameters []*vector.Vector, result vector.FunctionResultWrapper, _ *process.Process, length int, selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[T2](result) - p1, err := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) - if err != nil { - return err - } - p2, err := vector.OptGetParamFromWrapper[T1](rs, 1, parameters[1]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) + p2 := vector.OptGetParamFromWrapper[T1](rs, 1, parameters[1]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[T2](rsVec) @@ -395,14 +383,8 @@ func decimalBatchArith[TIn templateDec, TOut templateDecOut](parameters []*vecto arithFn func(v1, v2 []TIn, rs []TOut, scale1, scale2 int32, rsnull *nulls.Nulls) error, selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[TOut](result) - p1, err := vector.OptGetParamFromWrapper[TIn](rs, 0, parameters[0]) - if err != nil { - return err - } - p2, err := vector.OptGetParamFromWrapper[TIn](rs, 1, parameters[1]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[TIn](rs, 0, parameters[0]) + p2 := vector.OptGetParamFromWrapper[TIn](rs, 1, parameters[1]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[TOut](rsVec) scale1 := p1.GetType().Scale @@ -459,7 +441,7 @@ func decimalBatchArith[TIn templateDec, TOut templateDecOut](parameters []*vecto } else { v2 = p2.UnSafeGetAllValue() } - err = arithFn(v1, v2, rss, scale1, scale2, rsNull) + err := arithFn(v1, v2, rss, scale1, scale2, rsNull) if err != nil { if moerr.IsMoErrCode(err, moerr.ErrInvalidInput) { return moerr.NewOutOfRange(proc.Ctx, "DECIMAL", err.Error()) @@ -485,14 +467,8 @@ func opBinaryFixedFixedToFixed[ resultFn func(v1 T1, v2 T2) Tr, selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[Tr](result) - p1, err := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) - if err != nil { - return err - } - p2, err := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) + p2 := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[Tr](rsVec) @@ -614,14 +590,8 @@ func opBinaryFixedFixedToFixedWithErrorCheck[ resultFn func(v1 T1, v2 T2) (Tr, error), selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[Tr](result) - p1, err := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) - if err != nil { - return err - } - p2, err := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) + p2 := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[Tr](rsVec) @@ -770,14 +740,8 @@ func opBinaryFixedFixedToFixedWithNullOnError[ resultFn func(v1 T1, v2 T2) (Tr, error), selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[Tr](result) - p1, err := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) - if err != nil { - return err - } - p2, err := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) + p2 := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[Tr](rsVec) @@ -933,10 +897,7 @@ func opBinaryStrFixedToFixedWithErrorCheck[ result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[Tr](result) p1 := vector.OptGetBytesParamFromWrapper(rs, 0, parameters[0]) - p2, err := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) - if err != nil { - return err - } + p2 := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[Tr](rsVec) @@ -1086,10 +1047,7 @@ func opBinaryStrFixedToStrWithErrorCheck[ result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[types.Varlena](result) p1 := vector.OptGetBytesParamFromWrapper(rs, 0, parameters[0]) - p2, err := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) - if err != nil { - return err - } + p2 := vector.OptGetParamFromWrapper[T2](rs, 1, parameters[1]) rsVec := rs.GetResultVector() c1, c2 := parameters[0].IsConst(), parameters[1].IsConst() @@ -1252,10 +1210,7 @@ func opBinaryFixedStrToFixedWithErrorCheck[ resultFn func(v1 T1, v2 string) (Tr, error), selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[Tr](result) - p1, err := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[T1](rs, 0, parameters[0]) p2 := vector.OptGetBytesParamFromWrapper(rs, 1, parameters[1]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[Tr](rsVec) @@ -1404,14 +1359,8 @@ func specialTemplateForModFunction[ modFn func(v1, v2 T) T, selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[T](result) - p1, err := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) - if err != nil { - return err - } - p2, err := vector.OptGetParamFromWrapper[T](rs, 1, parameters[1]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) + p2 := vector.OptGetParamFromWrapper[T](rs, 1, parameters[1]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[T](rsVec) @@ -1682,14 +1631,8 @@ func specialTemplateForDivFunction[ divFn func(v1, v2 T) (T2, error), selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[T2](result) - p1, err := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) - if err != nil { - return err - } - p2, err := vector.OptGetParamFromWrapper[T](rs, 1, parameters[1]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) + p2 := vector.OptGetParamFromWrapper[T](rs, 1, parameters[1]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[T2](rsVec) @@ -2643,10 +2586,7 @@ func opUnaryFixedToFixed[ resultFn func(v T) Tr, selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(1) rs := vector.MustFunctionResult[Tr](result) - p1, err := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[Tr](rsVec) @@ -3057,10 +2997,7 @@ func opUnaryFixedToStr[ resultFn func(v T) string, selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(1) rs := vector.MustFunctionResult[types.Varlena](result) - p1, err := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) rsVec := rs.GetResultVector() c1 := parameters[0].IsConst() @@ -3137,10 +3074,7 @@ func opUnaryFixedToStrWithNullOnError[ resultFn func(v T) (string, error), selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(1) rs := vector.MustFunctionResult[types.Varlena](result) - p1, err := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) var constValue []byte constNull := false @@ -3203,10 +3137,7 @@ func opUnaryFixedToStrWithErrorCheck[ resultFn func(v T) (string, error), selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(1) rs := vector.MustFunctionResult[types.Varlena](result) - p1, err := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) rsVec := rs.GetResultVector() c1 := parameters[0].IsConst() @@ -3633,10 +3564,7 @@ func opUnaryFixedToFixedWithErrorCheck[ resultFn func(v T) (Tr, error), selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(1) rs := vector.MustFunctionResult[Tr](result) - p1, err := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[Tr](rsVec) @@ -3676,6 +3604,7 @@ func opUnaryFixedToFixedWithErrorCheck[ } // basic case. + var err error if p1.WithAnyNullValue() || rsAnyNull { nulls.Or(rsNull, parameters[0].GetNulls(), rsNull) rowCount := uint64(length) @@ -3709,10 +3638,7 @@ func opUnaryFixedToFixedWithNullOnError[ resultFn func(v T) (Tr, error), selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(1) rs := vector.MustFunctionResult[Tr](result) - p1, err := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[Tr](rsVec) diff --git a/pkg/sql/plan/function/func_binary.go b/pkg/sql/plan/function/func_binary.go index f22eaf04c2e22..4fbd6e60ae8f0 100644 --- a/pkg/sql/plan/function/func_binary.go +++ b/pkg/sql/plan/function/func_binary.go @@ -25,7 +25,6 @@ import ( "encoding/hex" "encoding/json" "fmt" - "io" "math" "math/big" "math/bits" @@ -34,15 +33,12 @@ import ( "strconv" "strings" "time" - "unicode/utf8" "github.com/matrixorigin/matrixone/pkg/util/fault" "go.uber.org/zap" "github.com/matrixorigin/matrixone/pkg/clusterservice" "github.com/matrixorigin/matrixone/pkg/common/moerr" - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/common/util" "github.com/matrixorigin/matrixone/pkg/container/nulls" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" @@ -1822,14 +1818,8 @@ func DateAdd(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc * // Use custom implementation to handle maximum overflow (return NULL) result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[types.Date](result) - p1, err := vector.OptGetParamFromWrapper[types.Date](rs, 0, ivecs[0]) - if err != nil { - return err - } - p2, err := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[types.Date](rs, 0, ivecs[0]) + p2 := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[types.Date](rsVec) rsNull := rsVec.GetNulls() @@ -1873,14 +1863,8 @@ func DatetimeAdd(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pr // Use custom implementation to handle maximum overflow (return NULL) result.UseOptFunctionParamFrame(2) - p1, err := vector.OptGetParamFromWrapper[types.Datetime](rs, 0, ivecs[0]) - if err != nil { - return err - } - p2, err := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[types.Datetime](rs, 0, ivecs[0]) + p2 := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[types.Datetime](rsVec) rsNull := rsVec.GetNulls() @@ -2001,14 +1985,8 @@ func TimestampAdd(ivecs []*vector.Vector, result vector.FunctionResultWrapper, p rs.TempSetType(types.New(types.T_timestamp, 0, scale)) result.UseOptFunctionParamFrame(2) - p1, err := vector.OptGetParamFromWrapper[types.Timestamp](rs, 0, ivecs[0]) - if err != nil { - return err - } - p2, err := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[types.Timestamp](rs, 0, ivecs[0]) + p2 := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[types.Timestamp](rsVec) rsNull := rsVec.GetNulls() @@ -2101,12 +2079,7 @@ func TimestampAddDate(ivecs []*vector.Vector, result vector.FunctionResultWrappe if resultType == types.T_date { // Result wrapper is DATE, but we need to return DATETIME // Convert to DATETIME type - if err := vec.SetTypeAndFixData( - types.New(types.T_datetime, 0, scale), - proc.GetMPool(), - ); err != nil { - return err - } + vec.SetTypeAndFixData(types.New(types.T_datetime, 0, scale), proc.GetMPool()) rss := vector.MustFixedColNoTypeCheck[types.Datetime](vec) rsNull := vec.GetNulls() @@ -2190,12 +2163,7 @@ func TimestampAddDate(ivecs []*vector.Vector, result vector.FunctionResultWrappe } else { // Result wrapper is DATETIME (backward compatibility) // Use SetType to change vector type to DATE - if err := vec.SetTypeAndFixData( - types.New(types.T_date, 0, 0), - proc.GetMPool(), - ); err != nil { - return err - } + vec.SetTypeAndFixData(types.New(types.T_date, 0, 0), proc.GetMPool()) rss := vector.MustFixedColNoTypeCheck[types.Date](vec) rsNull := vec.GetNulls() @@ -2253,12 +2221,7 @@ func TimestampAddDate(ivecs []*vector.Vector, result vector.FunctionResultWrappe scale := maxScale if resultType == types.T_date { // Result wrapper is DATE, but we need to return DATETIME - if err := vec.SetTypeAndFixData( - types.New(types.T_datetime, 0, scale), - proc.GetMPool(), - ); err != nil { - return err - } + vec.SetTypeAndFixData(types.New(types.T_datetime, 0, scale), proc.GetMPool()) rss := vector.MustFixedColNoTypeCheck[types.Datetime](vec) rsNull := vec.GetNulls() @@ -2355,12 +2318,7 @@ func TimestampAddDate(ivecs []*vector.Vector, result vector.FunctionResultWrappe } } else { // Result wrapper is DATETIME, but all units are date units, so return DATE - if err := vec.SetTypeAndFixData( - types.New(types.T_date, 0, 0), - proc.GetMPool(), - ); err != nil { - return err - } + vec.SetTypeAndFixData(types.New(types.T_date, 0, 0), proc.GetMPool()) rss := vector.MustFixedColNoTypeCheck[types.Date](vec) rsNull := vec.GetNulls() @@ -3687,7 +3645,7 @@ func DateFormat(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro } //format := "%b %D %M" -> []func{func1,func2, func3} - var legacy bytes.Buffer + var buf bytes.Buffer for i := uint64(0); i < uint64(length); i++ { d, null1 := dates.GetValue(i) if null1 || null2 { @@ -3695,19 +3653,15 @@ func DateFormat(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro return err } } else { + buf.Reset() var isNull bool - if isNull, err = appendFormattedBytesForResult( - result, - rs, - &legacy, - func(buf formatBuffer) (bool, error) { - return dateFmtOperator(proc.Ctx, d, string(fmt), buf) - }, - ); err != nil { + if isNull, err = dateFmtOperator(proc.Ctx, d, string(fmt), &buf); err != nil { return err } if isNull { err = rs.AppendBytes(nil, true) + } else { + err = rs.AppendBytes(buf.Bytes(), false) } if err != nil { return err @@ -3717,117 +3671,11 @@ func DateFormat(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro return nil } -type formatBuffer interface { - io.Writer - io.StringWriter - io.ByteWriter - WriteRune(rune) (int, error) - Grow(int) -} - -type countingFormatBuffer struct { - written int - err error -} - -func (w *countingFormatBuffer) add(size int) (int, error) { - if w.err != nil { - return 0, w.err - } - if size < 0 || size > math.MaxInt-w.written { - w.err = io.ErrShortBuffer - return 0, w.err - } - w.written += size - return size, nil -} - -func (w *countingFormatBuffer) Write(value []byte) (int, error) { - return w.add(len(value)) -} - -func (w *countingFormatBuffer) WriteString(value string) (int, error) { - return w.add(len(value)) -} - -func (w *countingFormatBuffer) WriteByte(byte) error { - _, err := w.add(1) - return err -} - -func (w *countingFormatBuffer) WriteRune(value rune) (int, error) { - size := utf8.RuneLen(value) - if size < 0 { - size = utf8.RuneLen(utf8.RuneError) - } - return w.add(size) -} - -func (w *countingFormatBuffer) Grow(int) {} - -func appendFormattedBytes( - rs *vector.FunctionResult[types.Varlena], - build func(formatBuffer) (bool, error), -) (bool, error) { - var counter countingFormatBuffer - isNull, err := build(&counter) - if err != nil || isNull { - return isNull, err - } - if counter.err != nil { - return false, counter.err - } - var output fixedSliceWriter - err = rs.AppendBytesWithBuilder(counter.written, func(dst []byte) (int, error) { - output.Reset(dst) - secondNull, buildErr := build(&output) - if buildErr != nil { - return 0, buildErr - } - if output.Err() != nil { - return 0, output.Err() - } - if secondNull { - return 0, moerr.NewInternalErrorNoCtx( - "format result changed between sizing and encoding", - ) - } - if output.Written() != counter.written { - return 0, moerr.NewInternalErrorNoCtx( - "format result size changed between sizing and encoding", - ) - } - return output.Written(), nil - }) - return false, err -} - -// appendFormattedBytesForResult preserves the legacy one-pass buffer path and -// uses exact two-pass publication only when dormant allocation accounting is -// selected. This avoids imposing duplicate formatting work on production -// callers before the expression owner is activated. -func appendFormattedBytesForResult( - result vector.FunctionResultWrapper, - rs *vector.FunctionResult[types.Varlena], - legacy *bytes.Buffer, - build func(formatBuffer) (bool, error), -) (bool, error) { - if result.HasFunctionScratch() { - return appendFormattedBytes(rs, build) - } - legacy.Reset() - isNull, err := build(legacy) - if err != nil || isNull { - return isNull, err - } - return false, rs.AppendBytes(legacy.Bytes(), false) -} - -type DateFormatFunc func(ctx context.Context, datetime types.Datetime, format string, buf formatBuffer) (isNull bool, err error) +type DateFormatFunc func(ctx context.Context, datetime types.Datetime, format string, buf *bytes.Buffer) (isNull bool, err error) // DATE_FORMAT datetime // handle '%d/%m/%Y' -> 22/04/2021 -func date_format_combine_pattern1(_ context.Context, t types.Datetime, format string, buf formatBuffer) (bool, error) { +func date_format_combine_pattern1(_ context.Context, t types.Datetime, format string, buf *bytes.Buffer) (bool, error) { month := int(t.Month()) day := int(t.Day()) year := int(t.Year()) @@ -3852,7 +3700,7 @@ func date_format_combine_pattern1(_ context.Context, t types.Datetime, format st } // handle '%Y%m%d' -> 20210422 -func date_format_combine_pattern2(_ context.Context, t types.Datetime, format string, buf formatBuffer) (bool, error) { +func date_format_combine_pattern2(_ context.Context, t types.Datetime, format string, buf *bytes.Buffer) (bool, error) { year := t.Year() month := int(t.Month()) day := int(t.Day()) @@ -3877,7 +3725,7 @@ func date_format_combine_pattern2(_ context.Context, t types.Datetime, format st } // handle '%Y' -> 2021 -func date_format_combine_pattern3(_ context.Context, t types.Datetime, format string, buf formatBuffer) (bool, error) { +func date_format_combine_pattern3(_ context.Context, t types.Datetime, format string, buf *bytes.Buffer) (bool, error) { year := t.Year() // Year conversion buf.WriteByte(byte('0' + (year / 1000 % 10))) @@ -3888,7 +3736,7 @@ func date_format_combine_pattern3(_ context.Context, t types.Datetime, format st } // %Y-%m-%d 2021-04-22 -func date_format_combine_pattern4(_ context.Context, t types.Datetime, format string, buf formatBuffer) (bool, error) { +func date_format_combine_pattern4(_ context.Context, t types.Datetime, format string, buf *bytes.Buffer) (bool, error) { year := t.Year() month := int(t.Month()) day := int(t.Day()) @@ -3916,7 +3764,7 @@ func date_format_combine_pattern4(_ context.Context, t types.Datetime, format st // handle '%Y-%m-%d %H:%i:%s' -> 2004-04-03 13:11:10 // handle ' %Y-%m-%d %T' -> 2004-04-03 13:11:10 -func date_format_combine_pattern5(_ context.Context, t types.Datetime, format string, buf formatBuffer) (bool, error) { +func date_format_combine_pattern5(_ context.Context, t types.Datetime, format string, buf *bytes.Buffer) (bool, error) { year := int(t.Year()) month := int(t.Month()) day := int(t.Day()) @@ -3966,7 +3814,7 @@ func date_format_combine_pattern5(_ context.Context, t types.Datetime, format st } // handle '%Y/%m/%d' -> 2010/01/07 -func date_format_combine_pattern6(_ context.Context, t types.Datetime, format string, buf formatBuffer) (bool, error) { +func date_format_combine_pattern6(_ context.Context, t types.Datetime, format string, buf *bytes.Buffer) (bool, error) { year := t.Year() month := int(t.Month()) day := int(t.Day()) @@ -3994,7 +3842,7 @@ func date_format_combine_pattern6(_ context.Context, t types.Datetime, format st // handle '%Y/%m/%d %H:%i:%s' -> 2010/01/07 23:12:34 // handle '%Y/%m/%d %T' -> 2010/01/07 23:12:34 -func date_format_combine_pattern7(_ context.Context, t types.Datetime, format string, buf formatBuffer) (bool, error) { +func date_format_combine_pattern7(_ context.Context, t types.Datetime, format string, buf *bytes.Buffer) (bool, error) { year := int(t.Year()) month := int(t.Month()) day := int(t.Day()) @@ -4044,7 +3892,7 @@ func date_format_combine_pattern7(_ context.Context, t types.Datetime, format st } // datetimeFormat: format the datetime value according to the format string. -func datetimeFormat(ctx context.Context, datetime types.Datetime, format string, buf formatBuffer) (bool, error) { +func datetimeFormat(ctx context.Context, datetime types.Datetime, format string, buf *bytes.Buffer) (bool, error) { inPatternMatch := false for _, b := range format { if inPatternMatch { @@ -4116,7 +3964,7 @@ var ( ) // makeDateFormat: Get the format string corresponding to the date according to a single format character -func makeDateFormat(_ context.Context, t types.Datetime, b rune, buf formatBuffer) (bool, error) { +func makeDateFormat(_ context.Context, t types.Datetime, b rune, buf *bytes.Buffer) (bool, error) { switch b { case 'b': m := t.Month() @@ -4282,7 +4130,7 @@ func TimeFormat(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro fmt, null2 := formats.GetStrValue(0) emptyFormat := len(fmt) == 0 - var legacy bytes.Buffer + var buf bytes.Buffer for i := uint64(0); i < uint64(length); i++ { t, null1 := times.GetValue(i) if null1 || null2 || emptyFormat { @@ -4290,15 +4138,11 @@ func TimeFormat(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro return err } } else { - _, err = appendFormattedBytesForResult( - result, - rs, - &legacy, - func(buf formatBuffer) (bool, error) { - return false, timeFormat(proc.Ctx, t, string(fmt), buf) - }, - ) - if err != nil { + buf.Reset() + if err = timeFormat(proc.Ctx, t, string(fmt), &buf); err != nil { + return err + } + if err = rs.AppendBytes(buf.Bytes(), false); err != nil { return err } } @@ -4308,7 +4152,7 @@ func TimeFormat(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro // timeFormat: Get the format string corresponding to the time according to format specifiers // Only supports time-related format specifiers: %H, %h, %I, %i, %k, %l, %S, %s, %f, %p, %r, %T -func timeFormat(ctx context.Context, t types.Time, format string, buf formatBuffer) error { +func timeFormat(ctx context.Context, t types.Time, format string, buf *bytes.Buffer) error { hour, minute, sec, msec, isNeg := t.ClockFormat() if isNeg && len(format) > 0 { buf.WriteByte('-') @@ -4335,7 +4179,7 @@ func timeFormat(ctx context.Context, t types.Time, format string, buf formatBuff // makeTimeFormat: Get the format string corresponding to the time according to a single format character // Only supports time-related format specifiers -func makeTimeFormat(ctx context.Context, hour uint64, minute, sec uint8, msec uint64, b rune, buf formatBuffer) error { +func makeTimeFormat(ctx context.Context, hour uint64, minute, sec uint8, msec uint64, b rune, buf *bytes.Buffer) error { switch b { case 'f': fmt.Fprintf(buf, "%06d", msec) @@ -4405,7 +4249,7 @@ func FormatIntByWidth(num, n int) string { return builder.String() } -func FormatInt2BufByWidth(num, n int, buf formatBuffer) { +func FormatInt2BufByWidth(num, n int, buf *bytes.Buffer) { numStr := strconv.Itoa(num) if len(numStr) >= n { buf.WriteString(numStr) @@ -4582,14 +4426,8 @@ func DateSub(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc * result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[types.Date](result) - p1, err := vector.OptGetParamFromWrapper[types.Date](rs, 0, ivecs[0]) - if err != nil { - return err - } - p2, err := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[types.Date](rs, 0, ivecs[0]) + p2 := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[types.Date](rsVec) rsNull := rsVec.GetNulls() @@ -4737,14 +4575,8 @@ func DatetimeSub(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pr // Use custom implementation to handle maximum overflow (return NULL) result.UseOptFunctionParamFrame(2) - p1, err := vector.OptGetParamFromWrapper[types.Datetime](rs, 0, ivecs[0]) - if err != nil { - return err - } - p2, err := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[types.Datetime](rs, 0, ivecs[0]) + p2 := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[types.Datetime](rsVec) rsNull := rsVec.GetNulls() @@ -4866,14 +4698,8 @@ func TimestampSub(ivecs []*vector.Vector, result vector.FunctionResultWrapper, p // Use custom implementation to handle maximum overflow (return NULL) result.UseOptFunctionParamFrame(2) - p1, err := vector.OptGetParamFromWrapper[types.Timestamp](rs, 0, ivecs[0]) - if err != nil { - return err - } - p2, err := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[types.Timestamp](rs, 0, ivecs[0]) + p2 := vector.OptGetParamFromWrapper[int64](rs, 1, ivecs[1]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[types.Timestamp](rsVec) rsNull := rsVec.GetNulls() @@ -4979,26 +4805,21 @@ func fieldCheck(overloads []overload, inputs []types.Type) checkResult { } func FieldNumber[T number](ivecs []*vector.Vector, result vector.FunctionResultWrapper, _ *process.Process, length int, selectList *FunctionSelectList) (err error) { - result.UseOptFunctionParamFrame(len(ivecs)) rs := vector.MustFunctionResult[uint64](result) - first, err := vector.OptGetParamFromWrapper[T](rs, 0, ivecs[0]) - if err != nil { - return err + + fs := make([]vector.FunctionParameterWrapper[T], len(ivecs)) + for i := range ivecs { + fs[i] = vector.GenerateFunctionFixedTypeParameter[T](ivecs[i]) } - nums := vector.MustFixedColNoTypeCheck[uint64](rs.GetResultVector()) - clear(nums[:length]) + nums := make([]uint64, length) for j := 1; j < len(ivecs); j++ { - candidate, err := vector.OptGetParamFromWrapper[T](rs, j, ivecs[j]) - if err != nil { - return err - } for i := uint64(0); i < uint64(length); i++ { - v1, null1 := first.GetValue(i) - v2, null2 := candidate.GetValue(i) + v1, null1 := fs[0].GetValue(i) + v2, null2 := fs[j].GetValue(i) if (nums[i] != 0) || (null1 || null2) { continue @@ -5010,23 +4831,32 @@ func FieldNumber[T number](ivecs []*vector.Vector, result vector.FunctionResultW } } + + for i := uint64(0); i < uint64(length); i++ { + if err := rs.Append(nums[i], false); err != nil { + return err + } + } + return nil } func FieldString(ivecs []*vector.Vector, result vector.FunctionResultWrapper, _ *process.Process, length int, selectList *FunctionSelectList) (err error) { - result.UseOptFunctionParamFrame(len(ivecs)) rs := vector.MustFunctionResult[uint64](result) - first := vector.OptGetBytesParamFromWrapper(rs, 0, ivecs[0]) - nums := vector.MustFixedColNoTypeCheck[uint64](rs.GetResultVector()) - clear(nums[:length]) + + fs := make([]vector.FunctionParameterWrapper[types.Varlena], len(ivecs)) + for i := range ivecs { + fs[i] = vector.GenerateFunctionStrParameter(ivecs[i]) + } + + nums := make([]uint64, length) for j := 1; j < len(ivecs); j++ { - candidate := vector.OptGetBytesParamFromWrapper(rs, j, ivecs[j]) for i := uint64(0); i < uint64(length); i++ { - v1, null1 := first.GetStrValue(i) - v2, null2 := candidate.GetStrValue(i) + v1, null1 := fs[0].GetStrValue(i) + v2, null2 := fs[j].GetStrValue(i) if (nums[i] != 0) || (null1 || null2) { continue @@ -5038,6 +4868,13 @@ func FieldString(ivecs []*vector.Vector, result vector.FunctionResultWrapper, _ } } + + for i := uint64(0); i < uint64(length); i++ { + if err := rs.Append(nums[i], false); err != nil { + return err + } + } + return nil } @@ -5348,46 +5185,21 @@ func MakeSet(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc * continue } - resultSize := 0 - partCount := 0 + // Build the result string by checking each bit position + var parts []string for j := 0; j < len(strParams); j++ { // Check if bit j is set (0-based, so bit 0 corresponds to str1, bit 1 to str2, etc.) if (bitsUint>>uint(j))&1 == 1 { str, null := strParams[j].GetStrValue(i) if !null { - if len(str) > math.MaxInt-resultSize { - return moerr.NewInvalidInputNoCtx("MAKE_SET result is too large") - } - resultSize += len(str) - partCount++ + parts = append(parts, functionUtil.QuickBytesToStr(str)) } } } - if partCount > 1 { - if partCount-1 > math.MaxInt-resultSize { - return moerr.NewInvalidInputNoCtx("MAKE_SET result is too large") - } - resultSize += partCount - 1 - } - if err := rs.AppendBytesWithFill(resultSize, func(dst []byte) { - written := 0 - parts := 0 - for j := 0; j < len(strParams); j++ { - if (bitsUint>>uint(j))&1 == 0 { - continue - } - str, null := strParams[j].GetStrValue(i) - if null { - continue - } - if parts > 0 { - dst[written] = ',' - written++ - } - written += copy(dst[written:], str) - parts++ - } - }); err != nil { + + // Join with comma separator + resultStr := strings.Join(parts, ",") + if err := rs.AppendBytes([]byte(resultStr), false); err != nil { return err } } @@ -5667,38 +5479,18 @@ func ExportSet(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc } } - resultSize := 0 + // Build the result string + var parts []string for j := int64(0); j < numberOfBits; j++ { - partSize := len(off) if (bitsUint>>uint(j))&1 == 1 { - partSize = len(on) - } - if partSize > math.MaxInt-resultSize { - return moerr.NewInvalidInputNoCtx("EXPORT_SET result is too large") - } - resultSize += partSize - } - separatorBytes := functionUtil.QuickStrToBytes(separator) - if numberOfBits > 1 { - separatorSize := uint64(numberOfBits-1) * uint64(len(separatorBytes)) - if separatorSize > uint64(math.MaxInt-resultSize) { - return moerr.NewInvalidInputNoCtx("EXPORT_SET result is too large") + parts = append(parts, functionUtil.QuickBytesToStr(on)) + } else { + parts = append(parts, functionUtil.QuickBytesToStr(off)) } - resultSize += int(separatorSize) } - if err := rs.AppendBytesWithFill(resultSize, func(dst []byte) { - written := 0 - for j := int64(0); j < numberOfBits; j++ { - if j > 0 { - written += copy(dst[written:], separatorBytes) - } - part := off - if (bitsUint>>uint(j))&1 == 1 { - part = on - } - written += copy(dst[written:], part) - } - }); err != nil { + + resultStr := strings.Join(parts, separator) + if err := rs.AppendBytes([]byte(resultStr), false); err != nil { return err } } @@ -6056,7 +5848,7 @@ func FromUnixTimeInt64Format(ivecs []*vector.Vector, result vector.FunctionResul formatMask, null1 := vector.GenerateFunctionStrParameter(ivecs[1]).GetStrValue(0) f := string(formatMask) - var legacy bytes.Buffer + var buf bytes.Buffer for i := uint64(0); i < uint64(length); i++ { v, null := vs.GetValue(i) @@ -6065,10 +5857,12 @@ func FromUnixTimeInt64Format(ivecs []*vector.Vector, result vector.FunctionResul return err } } else { + buf.Reset() r := types.DatetimeFromUnix(proc.GetSessionInfo().TimeZone, v) - if _, err = appendFormattedBytesForResult(result, rs, &legacy, func(buf formatBuffer) (bool, error) { - return datetimeFormat(proc.Ctx, r, f, buf) - }); err != nil { + if _, err = datetimeFormat(proc.Ctx, r, f, &buf); err != nil { + return err + } + if err = rs.AppendBytes(buf.Bytes(), false); err != nil { return err } } @@ -6086,7 +5880,7 @@ func FromUnixTimeUint64Format(ivecs []*vector.Vector, result vector.FunctionResu formatMask, null1 := vector.GenerateFunctionStrParameter(ivecs[1]).GetStrValue(0) f := string(formatMask) - var legacy bytes.Buffer + var buf bytes.Buffer for i := uint64(0); i < uint64(length); i++ { v, null := vs.GetValue(i) @@ -6095,10 +5889,12 @@ func FromUnixTimeUint64Format(ivecs []*vector.Vector, result vector.FunctionResu return err } } else { + buf.Reset() r := types.DatetimeFromUnix(proc.GetSessionInfo().TimeZone, int64(v)) - if _, err = appendFormattedBytesForResult(result, rs, &legacy, func(buf formatBuffer) (bool, error) { - return datetimeFormat(proc.Ctx, r, f, buf) - }); err != nil { + if _, err = datetimeFormat(proc.Ctx, r, f, &buf); err != nil { + return err + } + if err = rs.AppendBytes(buf.Bytes(), false); err != nil { return err } } @@ -6116,7 +5912,7 @@ func FromUnixTimeFloat64Format(ivecs []*vector.Vector, result vector.FunctionRes formatMask, null1 := vector.GenerateFunctionStrParameter(ivecs[1]).GetStrValue(0) f := string(formatMask) - var legacy bytes.Buffer + var buf bytes.Buffer for i := uint64(0); i < uint64(length); i++ { v, null := vs.GetValue(i) @@ -6125,11 +5921,13 @@ func FromUnixTimeFloat64Format(ivecs []*vector.Vector, result vector.FunctionRes return err } } else { + buf.Reset() x, y := splitDecimalToIntAndFrac(v) r := types.DatetimeFromUnixWithNsec(proc.GetSessionInfo().TimeZone, x, y) - if _, err = appendFormattedBytesForResult(result, rs, &legacy, func(buf formatBuffer) (bool, error) { - return datetimeFormat(proc.Ctx, r, f, buf) - }); err != nil { + if _, err = datetimeFormat(proc.Ctx, r, f, &buf); err != nil { + return err + } + if err = rs.AppendBytes(buf.Bytes(), false); err != nil { return err } } @@ -6148,7 +5946,7 @@ func FromUnixTimeDecimal256Format(ivecs []*vector.Vector, result vector.Function formatMask, null1 := vector.GenerateFunctionStrParameter(ivecs[1]).GetStrValue(0) f := string(formatMask) - var legacy bytes.Buffer + var buf bytes.Buffer for i := uint64(0); i < uint64(length); i++ { v, null := vs.GetValue(i) sec, nsec, ok, convErr := decimal256UnixTimeParts(v, scale) @@ -6161,10 +5959,12 @@ func FromUnixTimeDecimal256Format(ivecs []*vector.Vector, result vector.Function return err } } else { + buf.Reset() r := types.DatetimeFromUnixWithNsec(proc.GetSessionInfo().TimeZone, sec, nsec) - if _, err = appendFormattedBytesForResult(result, rs, &legacy, func(buf formatBuffer) (bool, error) { - return datetimeFormat(proc.Ctx, r, f, buf) - }); err != nil { + if _, err = datetimeFormat(proc.Ctx, r, f, &buf); err != nil { + return err + } + if err = rs.AppendBytes(buf.Bytes(), false); err != nil { return err } } @@ -8758,8 +8558,6 @@ func batchArrayDistanceSync[T types.RealNumbers]( length int, m metric.MetricType, proc *process.Process, - dist []float32, - result vector.FunctionResultWrapper, ) ([]float32, bool, error) { c0, c1 := ivecs[0].IsConst(), ivecs[1].IsConst() if c0 == c1 { @@ -8783,9 +8581,15 @@ func batchArrayDistanceSync[T types.RealNumbers]( if len(queryBytes) == 0 { return nil, false, nil } - query := types.BytesToArray[T](queryBytes) + x := [][]T{types.BytesToArray[T](queryBytes)} col := ivecs[colIdx] + y := make([][]T, length) + for i := range y { + y[i] = types.BytesToArray[T](col.GetBytesAt(i)) + } + + dist := make([]float32, length) // proc is non-nil under SQL execution; the nil branch keeps unit // tests (which don't synthesize a process) compiling and lets // EffectiveGpuMode fall back to the build-tag default. @@ -8794,39 +8598,7 @@ func batchArrayDistanceSync[T types.RealNumbers]( resolver = proc.GetResolveVariableFunc() } gpuMode := gpumode.EffectiveGpuMode(resolver) - scratchSize, usesGPU, err := metric.PairwiseDistanceOneToManyScratchSize( - query, - length, - m, - metric.GPUThresholdSQL, - gpuMode, - ) - if err != nil { - return nil, false, err - } - var scratch []byte - if usesGPU && result != nil { - var selected bool - scratch, selected, err = result.ResizeFunctionScratch(scratchSize) - if err != nil { - return nil, false, err - } - if !selected { - scratch = nil - } - } - handle, err := metric.PairwiseDistanceLaunchOneToManyWithScratch( - query, - length, - func(row int) []T { - return types.BytesToArray[T](col.GetBytesAt(row)) - }, - m, - dist, - metric.GPUThresholdSQL, - gpuMode, - scratch, - ) + handle, err := metric.PairwiseDistanceLaunch(x, y, m, dist, metric.GPUThresholdSQL, gpuMode) if err != nil { return nil, false, err } @@ -8837,60 +8609,15 @@ func batchArrayDistanceSync[T types.RealNumbers]( return dist, true, nil } -func tryBatchArrayDistance[T types.RealNumbers]( - ivecs []*vector.Vector, - result vector.FunctionResultWrapper, - proc *process.Process, - length int, - m metric.MetricType, - cosineSimilarity bool, -) (bool, error) { - rs := vector.MustFunctionResult[float64](result) - output := vector.MustFixedColNoTypeCheck[float64](rs.GetResultVector()) - if len(output) < length { - return false, moerr.NewInternalErrorNoCtx( - "array distance result is smaller than the input batch", - ) - } - - // The float64 result already owns 8*length admitted bytes. Pairwise distance - // needs 4*length temporary bytes, so use the upper half of that same backing - // store and convert forward after Wait. Each float64 write can only overwrite - // float32 values that were already read. - outputBytes := util.UnsafeSliceCast[float32](output[:length]) - distScratch := outputBytes[length:] - dist, ok, err := batchArrayDistanceSync[T]( - ivecs, - length, - m, - proc, - distScratch, - result, - ) - if err != nil || !ok { - return ok, err - } - for idx, value := range dist { - if cosineSimilarity { - output[idx] = 1 - float64(value) - } else { - output[idx] = float64(value) - } - } - return true, nil -} - func InnerProductArray[T types.RealNumbers](ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - if ok, err := tryBatchArrayDistance[T]( - ivecs, - result, - proc, - length, - metric.Metric_InnerProduct, - false, - ); err != nil { + if dist, ok, err := batchArrayDistanceSync[T](ivecs, length, metric.Metric_InnerProduct, proc); err != nil { return err } else if ok { + rs := vector.MustFunctionResult[float64](result) + rss := vector.MustFixedColNoTypeCheck[float64](rs.GetResultVector()) + for i, d := range dist { + rss[i] = float64(d) + } return nil } return opBinaryBytesBytesToFixedWithErrorCheck[float64](ivecs, result, proc, length, func(v1, v2 []byte) (out float64, err error) { @@ -8902,16 +8629,14 @@ func InnerProductArray[T types.RealNumbers](ivecs []*vector.Vector, result vecto func CosineSimilarityArray[T types.RealNumbers](ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { // Use Metric_CosineDistance and convert: similarity = 1 - distance. - if ok, err := tryBatchArrayDistance[T]( - ivecs, - result, - proc, - length, - metric.Metric_CosineDistance, - true, - ); err != nil { + if dist, ok, err := batchArrayDistanceSync[T](ivecs, length, metric.Metric_CosineDistance, proc); err != nil { return err } else if ok { + rs := vector.MustFunctionResult[float64](result) + rss := vector.MustFixedColNoTypeCheck[float64](rs.GetResultVector()) + for i, d := range dist { + rss[i] = 1.0 - float64(d) + } return nil } return opBinaryBytesBytesToFixedWithErrorCheck[float64](ivecs, result, proc, length, func(v1, v2 []byte) (out float64, err error) { @@ -8922,16 +8647,14 @@ func CosineSimilarityArray[T types.RealNumbers](ivecs []*vector.Vector, result v } func L2DistanceArray[T types.RealNumbers](ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - if ok, err := tryBatchArrayDistance[T]( - ivecs, - result, - proc, - length, - metric.Metric_L2Distance, - false, - ); err != nil { + if dist, ok, err := batchArrayDistanceSync[T](ivecs, length, metric.Metric_L2Distance, proc); err != nil { return err } else if ok { + rs := vector.MustFunctionResult[float64](result) + rss := vector.MustFixedColNoTypeCheck[float64](rs.GetResultVector()) + for i, d := range dist { + rss[i] = float64(d) + } return nil } return opBinaryBytesBytesToFixedWithErrorCheck[float64](ivecs, result, proc, length, func(v1, v2 []byte) (out float64, err error) { @@ -12403,16 +12126,14 @@ func sameGeometryPoint(a, b geometryPoint2D) bool { } func L2DistanceSqArray[T types.RealNumbers](ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - if ok, err := tryBatchArrayDistance[T]( - ivecs, - result, - proc, - length, - metric.Metric_L2sqDistance, - false, - ); err != nil { + if dist, ok, err := batchArrayDistanceSync[T](ivecs, length, metric.Metric_L2sqDistance, proc); err != nil { return err } else if ok { + rs := vector.MustFunctionResult[float64](result) + rss := vector.MustFixedColNoTypeCheck[float64](rs.GetResultVector()) + for i, d := range dist { + rss[i] = float64(d) + } return nil } return opBinaryBytesBytesToFixedWithErrorCheck[float64](ivecs, result, proc, length, func(v1, v2 []byte) (out float64, err error) { @@ -12423,16 +12144,14 @@ func L2DistanceSqArray[T types.RealNumbers](ivecs []*vector.Vector, result vecto } func CosineDistanceArray[T types.RealNumbers](ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - if ok, err := tryBatchArrayDistance[T]( - ivecs, - result, - proc, - length, - metric.Metric_CosineDistance, - false, - ); err != nil { + if dist, ok, err := batchArrayDistanceSync[T](ivecs, length, metric.Metric_CosineDistance, proc); err != nil { return err } else if ok { + rs := vector.MustFunctionResult[float64](result) + rss := vector.MustFixedColNoTypeCheck[float64](rs.GetResultVector()) + for i, d := range dist { + rss[i] = float64(d) + } return nil } return opBinaryBytesBytesToFixedWithErrorCheck[float64](ivecs, result, proc, length, func(v1, v2 []byte) (out float64, err error) { @@ -12482,28 +12201,6 @@ func arrayDistanceNarrow[T types.ArrayElement]( func arrayDistanceViaF32[T types.ArrayElement]( ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList, kernel func(v1, v2 []float32) (float64, error)) error { - if result.HasFunctionScratch() { - return opBinaryBytesBytesToFixedWithErrorCheck[float64](ivecs, result, proc, length, func(v1, v2 []byte) (float64, error) { - left := types.BytesToArray[T](v1) - right := types.BytesToArray[T](v2) - if len(right) > math.MaxInt-len(left) || len(left)+len(right) > math.MaxInt/4 { - return 0, mpool.ErrAllocationAccountInvalid - } - scratch, selected, err := result.ResizeFunctionScratch((len(left) + len(right)) * 4) - if err != nil { - return 0, err - } - if !selected { - return 0, mpool.ErrAllocationAccountInvalid - } - values := util.UnsafeSliceCast[float32](scratch) - leftValues := values[:len(left)] - rightValues := values[len(left):] - arrayToFloat32Into(leftValues, left) - arrayToFloat32Into(rightValues, right) - return kernel(leftValues, rightValues) - }, selectList) - } return opBinaryBytesBytesToFixedWithErrorCheck[float64](ivecs, result, proc, length, func(v1, v2 []byte) (float64, error) { f1 := types.ToFloat32Array[T](types.BytesToArray[T](v1)) f2 := types.ToFloat32Array[T](types.BytesToArray[T](v2)) @@ -12511,35 +12208,6 @@ func arrayDistanceViaF32[T types.ArrayElement]( }, selectList) } -func arrayToFloat32Into[T types.ArrayElement](dst []float32, src []T) { - switch values := any(src).(type) { - case []float32: - copy(dst, values) - case []float64: - for idx, value := range values { - dst[idx] = float32(value) - } - case []types.BF16: - for idx, value := range values { - dst[idx] = value.ToFloat32() - } - case []types.Float16: - for idx, value := range values { - dst[idx] = value.ToFloat32() - } - case []int8: - for idx, value := range values { - dst[idx] = float32(value) - } - case []uint8: - for idx, value := range values { - dst[idx] = float32(value) - } - default: - panic(moerr.NewInternalErrorNoCtx("unsupported array element type")) - } -} - func L2DistanceArrayViaF32[T types.ArrayElement](ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { return arrayDistanceNarrow[T](ivecs, result, proc, length, selectList, metric.Metric_L2Distance, true) } @@ -12580,110 +12248,107 @@ func generateAESKey(key []byte, keyLen int) ([]byte, error) { return out, nil } -func aesPaddedSize(plaintextSize int) (int, error) { - padding := aes.BlockSize - plaintextSize%aes.BlockSize - if plaintextSize > int(^uint(0)>>1)-padding { - return 0, moerr.NewInvalidInputNoCtx("plaintext is too large") +// pkcs7Padding adds PKCS7 padding to the data +func pkcs7Padding(data []byte, blockSize int) []byte { + padding := blockSize - len(data)%blockSize + padtext := make([]byte, padding) + for i := range padtext { + padtext[i] = byte(padding) } - return plaintextSize + padding, nil + return append(data, padtext...) } -func encryptAESPKCS7Into( - block cipher.Block, - plaintext []byte, - iv []byte, - useCBC bool, - ciphertext []byte, -) { - fullBytes := len(plaintext) - len(plaintext)%aes.BlockSize - var finalBlock [aes.BlockSize]byte - copy(finalBlock[:], plaintext[fullBytes:]) - padding := byte(aes.BlockSize - len(plaintext)%aes.BlockSize) - for idx := len(plaintext) % aes.BlockSize; idx < aes.BlockSize; idx++ { - finalBlock[idx] = padding +// pkcs7Unpadding removes PKCS7 padding from the data +func pkcs7Unpadding(data []byte) ([]byte, error) { + if len(data) == 0 { + return nil, moerr.NewInvalidInputNoCtx("invalid padding") } - - if useCBC { - mode := cipher.NewCBCEncrypter(block, iv[:aes.BlockSize]) - if fullBytes > 0 { - mode.CryptBlocks(ciphertext[:fullBytes], plaintext[:fullBytes]) - } - mode.CryptBlocks(ciphertext[fullBytes:], finalBlock[:]) - return + padding := int(data[len(data)-1]) + if padding > len(data) || padding == 0 { + return nil, moerr.NewInvalidInputNoCtx("invalid padding") } - for offset := 0; offset < fullBytes; offset += aes.BlockSize { - block.Encrypt( - ciphertext[offset:offset+aes.BlockSize], - plaintext[offset:offset+aes.BlockSize], - ) - } - block.Encrypt(ciphertext[fullBytes:], finalBlock[:]) -} - -func decryptAESPKCS7LastBlock( - block cipher.Block, - ciphertext []byte, - iv []byte, - useCBC bool, -) ([aes.BlockSize]byte, int, error) { - var finalBlock [aes.BlockSize]byte - if len(ciphertext) == 0 || len(ciphertext)%aes.BlockSize != 0 { - return finalBlock, 0, moerr.NewInvalidInputNoCtx( - "invalid ciphertext length", - ) - } - lastOffset := len(ciphertext) - aes.BlockSize - block.Decrypt(finalBlock[:], ciphertext[lastOffset:]) - if useCBC { - previous := iv[:aes.BlockSize] - if lastOffset > 0 { - previous = ciphertext[lastOffset-aes.BlockSize : lastOffset] - } - for idx := range finalBlock { - finalBlock[idx] ^= previous[idx] - } - } - - padding := int(finalBlock[aes.BlockSize-1]) - if padding == 0 || padding > aes.BlockSize { - return finalBlock, 0, moerr.NewInvalidInputNoCtx("invalid padding") - } - for idx := aes.BlockSize - padding; idx < aes.BlockSize; idx++ { - if finalBlock[idx] != byte(padding) { - return finalBlock, 0, moerr.NewInvalidInputNoCtx( - "invalid padding", - ) - } - } - return finalBlock, padding, nil -} - -func decryptAESPKCS7Into( - block cipher.Block, - ciphertext []byte, - iv []byte, - useCBC bool, - finalBlock [aes.BlockSize]byte, - padding int, - plaintext []byte, -) { - fullBytes := len(ciphertext) - aes.BlockSize - if useCBC { - if fullBytes > 0 { - cipher.NewCBCDecrypter( - block, - iv[:aes.BlockSize], - ).CryptBlocks(plaintext[:fullBytes], ciphertext[:fullBytes]) - } - } else { - for offset := 0; offset < fullBytes; offset += aes.BlockSize { - block.Decrypt( - plaintext[offset:offset+aes.BlockSize], - ciphertext[offset:offset+aes.BlockSize], - ) + // Verify padding + for i := len(data) - padding; i < len(data); i++ { + if data[i] != byte(padding) { + return nil, moerr.NewInvalidInputNoCtx("invalid padding") } } - copy(plaintext[fullBytes:], finalBlock[:aes.BlockSize-padding]) + return data[:len(data)-padding], nil +} + +// encryptECB encrypts data using AES-128-ECB mode +func encryptECB(plaintext, key []byte) ([]byte, error) { + block, err := aes.NewCipher(key) + if err != nil { + return nil, err + } + + // Add PKCS7 padding + padded := pkcs7Padding(plaintext, aes.BlockSize) + + // Encrypt each block independently (ECB mode) + ciphertext := make([]byte, len(padded)) + for i := 0; i < len(padded); i += aes.BlockSize { + block.Encrypt(ciphertext[i:i+aes.BlockSize], padded[i:i+aes.BlockSize]) + } + + return ciphertext, nil +} + +// decryptECB decrypts data using AES-128-ECB mode +func decryptECB(ciphertext, key []byte) ([]byte, error) { + block, err := aes.NewCipher(key) + if err != nil { + return nil, err + } + + // Check that ciphertext length is a multiple of block size + if len(ciphertext)%aes.BlockSize != 0 { + return nil, moerr.NewInvalidInputNoCtx("invalid ciphertext length") + } + + // Decrypt each block independently (ECB mode) + plaintext := make([]byte, len(ciphertext)) + for i := 0; i < len(ciphertext); i += aes.BlockSize { + block.Decrypt(plaintext[i:i+aes.BlockSize], ciphertext[i:i+aes.BlockSize]) + } + + // Remove PKCS7 padding + return pkcs7Unpadding(plaintext) +} + +// encryptCBC encrypts data using AES-CBC mode +func encryptCBC(plaintext, key, iv []byte) ([]byte, error) { + block, err := aes.NewCipher(key) + if err != nil { + return nil, err + } + if len(iv) < aes.BlockSize { + return nil, moerr.NewInvalidInputNoCtx("invalid iv length") + } + padded := pkcs7Padding(plaintext, aes.BlockSize) + ciphertext := make([]byte, len(padded)) + mode := cipher.NewCBCEncrypter(block, iv[:aes.BlockSize]) + mode.CryptBlocks(ciphertext, padded) + return ciphertext, nil +} + +// decryptCBC decrypts data using AES-CBC mode +func decryptCBC(ciphertext, key, iv []byte) ([]byte, error) { + block, err := aes.NewCipher(key) + if err != nil { + return nil, err + } + if len(iv) < aes.BlockSize { + return nil, moerr.NewInvalidInputNoCtx("invalid iv length") + } + if len(ciphertext)%aes.BlockSize != 0 { + return nil, moerr.NewInvalidInputNoCtx("invalid ciphertext length") + } + plaintext := make([]byte, len(ciphertext)) + mode := cipher.NewCBCDecrypter(block, iv[:aes.BlockSize]) + mode.CryptBlocks(plaintext, ciphertext) + return pkcs7Unpadding(plaintext) } type aesModeInfo struct { @@ -12762,9 +12427,14 @@ func AESEncrypt(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro continue } - block, blockErr := aes.NewCipher(aesKey) - ciphertextSize, sizeErr := aesPaddedSize(len(str)) - if blockErr != nil || sizeErr != nil { + var ciphertext []byte + var encErr error + if modeInfo.useCBC { + ciphertext, encErr = encryptCBC(str, aesKey, iv) + } else { + ciphertext, encErr = encryptECB(str, aesKey) + } + if encErr != nil { // On error, return NULL (MySQL behavior) if err := rs.AppendBytes(nil, true); err != nil { return err @@ -12772,18 +12442,7 @@ func AESEncrypt(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro continue } - if err := rs.AppendBytesWithFill( - ciphertextSize, - func(ciphertext []byte) { - encryptAESPKCS7Into( - block, - str, - iv, - modeInfo.useCBC, - ciphertext, - ) - }, - ); err != nil { + if err := rs.AppendBytes(ciphertext, false); err != nil { return err } } @@ -12841,42 +12500,22 @@ func AESDecrypt(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro continue } - block, blockErr := aes.NewCipher(aesKey) - if blockErr != nil { - // On error, return NULL (MySQL behavior) - if err := rs.AppendBytes(nil, true); err != nil { - return err - } - continue + var plaintext []byte + var decErr error + if modeInfo.useCBC { + plaintext, decErr = decryptCBC(crypt, aesKey, iv) + } else { + plaintext, decErr = decryptECB(crypt, aesKey) } - - finalBlock, padding, decErr := decryptAESPKCS7LastBlock( - block, - crypt, - iv, - modeInfo.useCBC, - ) if decErr != nil { + // On error, return NULL (MySQL behavior) if err := rs.AppendBytes(nil, true); err != nil { return err } continue } - plaintextSize := len(crypt) - padding - if err := rs.AppendBytesWithFill( - plaintextSize, - func(plaintext []byte) { - decryptAESPKCS7Into( - block, - crypt, - iv, - modeInfo.useCBC, - finalBlock, - padding, - plaintext, - ) - }, - ); err != nil { + + if err := rs.AppendBytes(plaintext, false); err != nil { return err } } diff --git a/pkg/sql/plan/function/func_binary_aes_test.go b/pkg/sql/plan/function/func_binary_aes_test.go index 4f25c712a6e0d..c30f901349bde 100644 --- a/pkg/sql/plan/function/func_binary_aes_test.go +++ b/pkg/sql/plan/function/func_binary_aes_test.go @@ -15,13 +15,10 @@ package function import ( - "crypto/aes" "fmt" - "strings" "testing" "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/testutil" "github.com/matrixorigin/matrixone/pkg/vm/process" "github.com/stretchr/testify/require" @@ -45,12 +42,8 @@ func TestAESEncryptDecryptECB(t *testing.T) { aesKey, err := generateAESKey([]byte(key), 16) require.NoError(t, err) - block, err := aes.NewCipher(aesKey) + ciphertext, err := encryptECB([]byte(plain), aesKey) require.NoError(t, err) - ciphertextSize, err := aesPaddedSize(len(plain)) - require.NoError(t, err) - ciphertext := make([]byte, ciphertextSize) - encryptAESPKCS7Into(block, []byte(plain), nil, false, ciphertext) encryptCase := NewFunctionTestCase(proc, []FunctionTestInput{ @@ -83,18 +76,8 @@ func TestAESEncryptDecryptCBC(t *testing.T) { aesKey, err := generateAESKey([]byte(key), 32) require.NoError(t, err) - block, err := aes.NewCipher(aesKey) - require.NoError(t, err) - ciphertextSize, err := aesPaddedSize(len(plain)) + ciphertext, err := encryptCBC([]byte(plain), aesKey, []byte(iv)) require.NoError(t, err) - ciphertext := make([]byte, ciphertextSize) - encryptAESPKCS7Into( - block, - []byte(plain), - []byte(iv), - true, - ciphertext, - ) encryptCase := NewFunctionTestCase(proc, []FunctionTestInput{ @@ -121,105 +104,6 @@ func TestAESEncryptDecryptCBC(t *testing.T) { require.True(t, ok, fmt.Sprintf("decrypt cbc failed: %s", info)) } -func TestAESEncryptDecryptBlockBoundaries(t *testing.T) { - plaintexts := []string{ - "", - strings.Repeat("a", aes.BlockSize-1), - strings.Repeat("b", aes.BlockSize), - strings.Repeat("c", aes.BlockSize+1), - strings.Repeat("d", 2*aes.BlockSize), - } - for _, tc := range []struct { - name string - mode string - key string - iv string - }{ - {name: "ecb", mode: "aes-128-ecb", key: "boundary-key"}, - {name: "cbc", mode: "aes-256-cbc", key: "boundary-key", iv: "0123456789abcdef"}, - } { - t.Run(tc.name, func(t *testing.T) { - proc := newAESProcess(t, tc.mode) - mp := proc.Mp() - plainVec := newVectorByType( - mp, - types.T_varchar.ToType(), - plaintexts, - nil, - ) - defer plainVec.Free(mp) - keys := make([]string, len(plaintexts)) - for idx := range keys { - keys[idx] = tc.key - } - keyVec := newVectorByType( - mp, - types.T_varchar.ToType(), - keys, - nil, - ) - defer keyVec.Free(mp) - params := []*vector.Vector{plainVec, keyVec} - if tc.iv != "" { - ivs := make([]string, len(plaintexts)) - for idx := range ivs { - ivs[idx] = tc.iv - } - ivVec := newVectorByType( - mp, - types.T_varchar.ToType(), - ivs, - nil, - ) - defer ivVec.Free(mp) - params = append(params, ivVec) - } - - encrypted := vector.NewFunctionResultWrapper( - types.T_blob.ToType(), - mp, - ) - defer encrypted.Free() - require.NoError(t, encrypted.PreExtendAndReset(len(plaintexts))) - require.NoError(t, AESEncrypt( - params, - encrypted, - proc, - len(plaintexts), - nil, - )) - - decryptParams := []*vector.Vector{ - encrypted.GetResultVector(), - keyVec, - } - if len(params) == 3 { - decryptParams = append(decryptParams, params[2]) - } - decrypted := vector.NewFunctionResultWrapper( - types.T_varchar.ToType(), - mp, - ) - defer decrypted.Free() - require.NoError(t, decrypted.PreExtendAndReset(len(plaintexts))) - require.NoError(t, AESDecrypt( - decryptParams, - decrypted, - proc, - len(plaintexts), - nil, - )) - for idx, plaintext := range plaintexts { - require.Equal( - t, - []byte(plaintext), - decrypted.GetResultVector().GetBytesAt(idx), - ) - } - }) - } -} - func TestAESEncryptCBCMissingIV(t *testing.T) { proc := newAESProcess(t, "aes-256-cbc") plain := "missing iv" diff --git a/pkg/sql/plan/function/func_binary_array_distance_gpu_test.go b/pkg/sql/plan/function/func_binary_array_distance_gpu_test.go index c019e8619f124..0b9aee421d619 100644 --- a/pkg/sql/plan/function/func_binary_array_distance_gpu_test.go +++ b/pkg/sql/plan/function/func_binary_array_distance_gpu_test.go @@ -60,8 +60,8 @@ func TestBatchArrayDistanceSync_GPU_L2sq(t *testing.T) { constVec := makeConstArrayVec[float32](t, mp, query, N) colVec := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - gpuDist, ok, err := testBatchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, N, metric.Metric_L2sqDistance) + gpuDist, ok, err := batchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, N, metric.Metric_L2sqDistance, nil) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(gpuDist)) @@ -102,8 +102,8 @@ func TestBatchArrayDistanceSync_GPU_InnerProduct(t *testing.T) { constVec := makeConstArrayVec[float32](t, mp, query, N) colVec := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - gpuDist, ok, err := testBatchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, N, metric.Metric_InnerProduct) + gpuDist, ok, err := batchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, N, metric.Metric_InnerProduct, nil) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(gpuDist)) @@ -143,8 +143,8 @@ func TestBatchArrayDistanceSync_GPU_CosineDistance(t *testing.T) { constVec := makeConstArrayVec[float32](t, mp, query, N) colVec := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - gpuDist, ok, err := testBatchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, N, metric.Metric_CosineDistance) + gpuDist, ok, err := batchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, N, metric.Metric_CosineDistance, nil) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(gpuDist)) @@ -184,8 +184,8 @@ func TestBatchArrayDistanceSync_GPU_L2Distance(t *testing.T) { constVec := makeConstArrayVec[float32](t, mp, query, N) colVec := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - gpuDist, ok, err := testBatchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, N, metric.Metric_L2Distance) + gpuDist, ok, err := batchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, N, metric.Metric_L2Distance, nil) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(gpuDist)) diff --git a/pkg/sql/plan/function/func_binary_array_distance_test.go b/pkg/sql/plan/function/func_binary_array_distance_test.go index 0810d0bec2b54..8d8b1f6283e7a 100644 --- a/pkg/sql/plan/function/func_binary_array_distance_test.go +++ b/pkg/sql/plan/function/func_binary_array_distance_test.go @@ -26,7 +26,7 @@ import ( ) // makeConstArrayVec creates a constant vector holding a single array value repeated length times. -func makeConstArrayVec[T types.RealNumbers](t testing.TB, mp *mpool.MPool, arr []T, length int) *vector.Vector { +func makeConstArrayVec[T types.RealNumbers](t *testing.T, mp *mpool.MPool, arr []T, length int) *vector.Vector { t.Helper() b := types.ArrayToBytes[T](arr) v, err := vector.NewConstBytes(types.T_array_float32.ToType(), b, length, mp) @@ -35,7 +35,7 @@ func makeConstArrayVec[T types.RealNumbers](t testing.TB, mp *mpool.MPool, arr [ } // makeConstArrayVec64 is the float64 variant. -func makeConstArrayVec64(t testing.TB, mp *mpool.MPool, arr []float64, length int) *vector.Vector { +func makeConstArrayVec64(t *testing.T, mp *mpool.MPool, arr []float64, length int) *vector.Vector { t.Helper() b := types.ArrayToBytes[float64](arr) v, err := vector.NewConstBytes(types.T_array_float64.ToType(), b, length, mp) @@ -44,7 +44,7 @@ func makeConstArrayVec64(t testing.TB, mp *mpool.MPool, arr []float64, length in } // makeColArrayVec creates a column vector holding one array per row. -func makeColArrayVec[T types.RealNumbers](t testing.TB, mp *mpool.MPool, typ types.Type, rows [][]T) *vector.Vector { +func makeColArrayVec[T types.RealNumbers](t *testing.T, mp *mpool.MPool, typ types.Type, rows [][]T) *vector.Vector { t.Helper() v := vector.NewVec(typ) for _, row := range rows { @@ -66,58 +66,6 @@ func approxEqF32(a, b float32) bool { return diff/avg < 1e-4 } -func testBatchArrayDistanceSync[T types.RealNumbers]( - ivecs []*vector.Vector, - length int, - m metric.MetricType, -) ([]float32, bool, error) { - return batchArrayDistanceSync[T]( - ivecs, - length, - m, - nil, - make([]float32, length), - nil, - ) -} - -func BenchmarkBatchArrayDistanceSync8192(b *testing.B) { - mp := mpool.MustNewZero() - defer mpool.DeleteMPool(mp) - - const dimension = 128 - query := make([]float32, dimension) - rows := make([][]float32, 8192) - for row := range rows { - rows[row] = make([]float32, dimension) - for col := range rows[row] { - rows[row][col] = float32((row + col) % 17) - } - } - constVec := makeConstArrayVec[float32](b, mp, query, len(rows)) - defer constVec.Free(mp) - colVec := makeColArrayVec[float32](b, mp, types.T_array_float32.ToType(), rows) - defer colVec.Free(mp) - inputs := []*vector.Vector{constVec, colVec} - distScratch := make([]float32, len(rows)) - - b.ReportAllocs() - b.ResetTimer() - for i := 0; i < b.N; i++ { - dist, ok, err := batchArrayDistanceSync[float32]( - inputs, - len(rows), - metric.Metric_L2sqDistance, - nil, - distScratch, - nil, - ) - require.NoError(b, err) - require.True(b, ok) - require.Len(b, dist, len(rows)) - } -} - // TestBatchArrayDistanceSync_L2Sq verifies batchArrayDistanceSync with Metric_L2sqDistance // on a small const-vs-column input (always CPU path). func TestBatchArrayDistanceSync_L2Sq(t *testing.T) { @@ -133,8 +81,8 @@ func TestBatchArrayDistanceSync_L2Sq(t *testing.T) { constVec := makeConstArrayVec[float32](t, mp, query, N) colVec := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - dist, ok, err := testBatchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, N, metric.Metric_L2sqDistance) + dist, ok, err := batchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, N, metric.Metric_L2sqDistance, nil) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(dist)) @@ -143,43 +91,6 @@ func TestBatchArrayDistanceSync_L2Sq(t *testing.T) { } } -func TestBatchArrayDistanceResultScratchAlias(t *testing.T) { - mp := mpool.MustNewZero() - defer mpool.DeleteMPool(mp) - - query := []float32{1, 0, 0} - rows := [][]float32{ - {1, 0, 0}, - {0, 1, 0}, - {0, 0, 1}, - } - constVec := makeConstArrayVec[float32](t, mp, query, len(rows)) - defer constVec.Free(mp) - colVec := makeColArrayVec[float32]( - t, - mp, - types.T_array_float32.ToType(), - rows, - ) - defer colVec.Free(mp) - - result := vector.NewFunctionResultWrapper(types.T_float64.ToType(), mp) - defer result.Free() - require.NoError(t, result.PreExtendAndReset(len(rows))) - require.NoError(t, L2DistanceSqArray[float32]( - []*vector.Vector{constVec, colVec}, - result, - nil, - len(rows), - nil, - )) - require.Equal( - t, - []float64{0, 2, 2}, - vector.MustFixedColNoTypeCheck[float64](result.GetResultVector()), - ) -} - // TestBatchArrayDistanceSync_L2 verifies batchArrayDistanceSync with Metric_L2Distance. func TestBatchArrayDistanceSync_L2(t *testing.T) { mp := mpool.MustNewZero() @@ -194,8 +105,8 @@ func TestBatchArrayDistanceSync_L2(t *testing.T) { constVec := makeConstArrayVec[float32](t, mp, query, N) colVec := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - dist, ok, err := testBatchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, N, metric.Metric_L2Distance) + dist, ok, err := batchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, N, metric.Metric_L2Distance, nil) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(dist)) @@ -219,8 +130,8 @@ func TestBatchArrayDistanceSync_InnerProduct(t *testing.T) { constVec := makeConstArrayVec[float32](t, mp, query, N) colVec := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - dist, ok, err := testBatchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, N, metric.Metric_InnerProduct) + dist, ok, err := batchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, N, metric.Metric_InnerProduct, nil) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(dist)) @@ -244,8 +155,8 @@ func TestBatchArrayDistanceSync_CosineDistance(t *testing.T) { constVec := makeConstArrayVec[float32](t, mp, query, N) colVec := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - dist, ok, err := testBatchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, N, metric.Metric_CosineDistance) + dist, ok, err := batchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, N, metric.Metric_CosineDistance, nil) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(dist)) @@ -269,8 +180,8 @@ func TestBatchArrayDistanceSync_QueryAsSecondArg(t *testing.T) { constVec := makeConstArrayVec[float32](t, mp, query, N) // Note: const is ivecs[1], column is ivecs[0] - dist, ok, err := testBatchArrayDistanceSync[float32]( - []*vector.Vector{colVec, constVec}, N, metric.Metric_L2sqDistance) + dist, ok, err := batchArrayDistanceSync[float32]( + []*vector.Vector{colVec, constVec}, N, metric.Metric_L2sqDistance, nil) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(dist)) @@ -292,8 +203,8 @@ func TestBatchArrayDistanceSync_Float64(t *testing.T) { constVec := makeConstArrayVec64(t, mp, query, N) colVec := makeColArrayVec[float64](t, mp, types.T_array_float64.ToType(), rows) - dist, ok, err := testBatchArrayDistanceSync[float64]( - []*vector.Vector{constVec, colVec}, N, metric.Metric_L2sqDistance) + dist, ok, err := batchArrayDistanceSync[float64]( + []*vector.Vector{constVec, colVec}, N, metric.Metric_L2sqDistance, nil) require.NoError(t, err) require.True(t, ok) require.Equal(t, N, len(dist)) @@ -313,8 +224,8 @@ func TestBatchArrayDistanceSync_BothConst(t *testing.T) { v1, err := vector.NewConstBytes(types.T_array_float32.ToType(), b, 4, mp) require.NoError(t, err) - _, ok, err := testBatchArrayDistanceSync[float32]( - []*vector.Vector{v0, v1}, 4, metric.Metric_L2sqDistance) + _, ok, err := batchArrayDistanceSync[float32]( + []*vector.Vector{v0, v1}, 4, metric.Metric_L2sqDistance, nil) require.NoError(t, err) require.False(t, ok, "both-const should return ok=false") } @@ -328,8 +239,8 @@ func TestBatchArrayDistanceSync_BothCol(t *testing.T) { v0 := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) v1 := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - _, ok, err := testBatchArrayDistanceSync[float32]( - []*vector.Vector{v0, v1}, 2, metric.Metric_L2sqDistance) + _, ok, err := batchArrayDistanceSync[float32]( + []*vector.Vector{v0, v1}, 2, metric.Metric_L2sqDistance, nil) require.NoError(t, err) require.False(t, ok, "col-vs-col should return ok=false") } @@ -343,8 +254,8 @@ func TestBatchArrayDistanceSync_NullConst(t *testing.T) { rows := [][]float32{{1, 0, 0}, {0, 1, 0}, {0, 0, 1}, {1, 1, 0}} colVec := makeColArrayVec[float32](t, mp, types.T_array_float32.ToType(), rows) - _, ok, err := testBatchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, 4, metric.Metric_L2sqDistance) + _, ok, err := batchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, 4, metric.Metric_L2sqDistance, nil) require.NoError(t, err) require.False(t, ok, "null const should return ok=false") } @@ -363,8 +274,8 @@ func TestBatchArrayDistanceSync_NullInColumn(t *testing.T) { require.NoError(t, vector.AppendBytes(colVec, nil, true, mp)) // null row require.NoError(t, vector.AppendBytes(colVec, types.ArrayToBytes[float32]([]float32{0, 1, 0}), false, mp)) - _, ok, err := testBatchArrayDistanceSync[float32]( - []*vector.Vector{constVec, colVec}, 3, metric.Metric_L2sqDistance) + _, ok, err := batchArrayDistanceSync[float32]( + []*vector.Vector{constVec, colVec}, 3, metric.Metric_L2sqDistance, nil) require.NoError(t, err) require.False(t, ok, "column with nulls should return ok=false") } diff --git a/pkg/sql/plan/function/func_builtin.go b/pkg/sql/plan/function/func_builtin.go index ace53043744e0..5bb92efb50ed5 100644 --- a/pkg/sql/plan/function/func_builtin.go +++ b/pkg/sql/plan/function/func_builtin.go @@ -29,7 +29,6 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/hashmap" "github.com/matrixorigin/matrixone/pkg/common/moerr" - "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/runtime" "github.com/matrixorigin/matrixone/pkg/common/util" "github.com/matrixorigin/matrixone/pkg/config" @@ -233,8 +232,8 @@ func parseLeadingInteger(s string) (int64, bool) { return v, true } -// appendCharBytes appends one MySQL CHAR() integer as big-endian bytes. MySQL -// treats CHAR(N) values as unsigned 32-bit integers and expands +// encodeCharBytes converts an int64 argument for MySQL CHAR() into big-endian +// bytes. MySQL treats CHAR(N) values as unsigned 32-bit integers and expands // values > 255 into multiple big-endian bytes: // // CHAR(256) → 0x0100 (two bytes) @@ -242,10 +241,10 @@ func parseLeadingInteger(s string) (int64, bool) { // CHAR(-1) → 0xFFFFFFFF (four bytes, via two's complement uint32) // // See MySQL docs: https://dev.mysql.com/doc/refman/8.4/en/string-functions.html#function_char -func appendCharBytes(dst []byte, v int64) []byte { +func encodeCharBytes(v int64) []byte { uv := uint32(v) if uv == 0 { - return append(dst, 0) + return []byte{0} } // Encode as big-endian 32-bit, then strip leading zero bytes. var buf [4]byte @@ -258,7 +257,7 @@ func appendCharBytes(dst []byte, v int64) []byte { for start < 3 && buf[start] == 0 { start++ } - return append(dst, buf[start:]...) + return buf[start:] } const ( @@ -799,33 +798,25 @@ func builtInConcat(parameters []*vector.Vector, result vector.FunctionResultWrap } for i := uint64(0); i < uint64(length); i++ { - total := 0 - null := false + var vs string + apv := true + for _, p := range ps { - v, isNull := p.GetStrValue(i) - if isNull { + v, null := p.GetStrValue(i) + if null { if err := rs.AppendBytes(nil, true); err != nil { return err } - null = true + apv = false break + } else { + vs += string(v) } - if len(v) > math.MaxInt-total { - return moerr.NewInvalidInputNoCtx("CONCAT result is too large") - } - total += len(v) } - if null { - continue - } - if err := rs.AppendBytesWithFill(total, func(dst []byte) { - offset := 0 - for _, p := range ps { - v, _ := p.GetStrValue(i) - offset += copy(dst[offset:], v) + if apv { + if err := rs.AppendBytes([]byte(vs), false); err != nil { + return err } - }); err != nil { - return err } } return nil @@ -1248,21 +1239,25 @@ func builtInChar(parameters []*vector.Vector, result vector.FunctionResultWrappe continue } - if len(getters) > math.MaxInt/4 { - return moerr.NewInvalidInputNoCtx("CHAR has too many arguments") - } - // MySQL skips NULL arguments and expands each remaining value to at - // most four bytes. Build directly in the admitted result backing. - if err := rs.AppendBytesWithBuilder(len(getters)*4, func(dst []byte) (int, error) { - output := dst[:0] - for _, getter := range getters { - value, null := getter(i) - if !null { - output = appendCharBytes(output, value) - } + var resultBytes []byte + + // Process all arguments + for _, getter := range getters { + v, null := getter(i) + if null { + // MySQL skips NULL arguments instead of returning NULL + continue } - return len(output), nil - }); err != nil { + // Convert argument to big-endian multi-byte sequence. + // MySQL treats CHAR(N) as unsigned 32-bit ints and expands + // values > 255 into multiple big-endian bytes (CHAR(256) → 0x0100). + // Negative values use two's complement uint32 (CHAR(-1) → 0xFFFFFFFF). + resultBytes = append(resultBytes, encodeCharBytes(v)...) + } + + // resultBytes is empty only when every argument was NULL. MySQL returns + // an empty (non-NULL) string in that case, e.g. CHAR(NULL, NULL) -> ''. + if err := rs.AppendBytes(resultBytes, false); err != nil { return err } } @@ -2090,95 +2085,8 @@ func builtInUnixTimestampVarcharToDecimal128(parameters []*vector.Vector, result return nil } -func builtInHashAccounted( - parameters []*vector.Vector, - result vector.FunctionResultWrapper, - length int, - appendState func(uint64), -) error { - var keys [hashmap.UnitLimit][]byte - var states [hashmap.UnitLimit][3]uint64 - var keySizes [hashmap.UnitLimit]int - for start := 0; start < length; start += hashmap.UnitLimit { - count := min(length-start, hashmap.UnitLimit) - total := 0 - for localRow := 0; localRow < count; localRow++ { - row := start + localRow - size := 0 - for _, parameter := range parameters { - if size == math.MaxInt { - return moerr.NewInvalidInputNoCtx("HASH input is too large") - } - size++ // one NULL marker per parameter - if !parameter.IsNull(uint64(row)) { - valueSize := len(parameter.GetRawBytesAt(row)) - if valueSize > math.MaxInt-size { - return moerr.NewInvalidInputNoCtx("HASH input is too large") - } - size += valueSize - } - } - if size < len(hashtable.StrKeyPadding) { - size = len(hashtable.StrKeyPadding) - } - if size > math.MaxInt-total { - return moerr.NewInvalidInputNoCtx("HASH input is too large") - } - keySizes[localRow] = size - total += size - } - - scratch, selected, err := result.ResizeFunctionScratch(total) - if err != nil { - return err - } - if !selected { - return mpool.ErrAllocationAccountInvalid - } - offset := 0 - for localRow := 0; localRow < count; localRow++ { - size := keySizes[localRow] - keys[localRow] = scratch[offset : offset+size] - offset += size - } - for localRow := 0; localRow < count; localRow++ { - row := start + localRow - key := keys[localRow] - written := 0 - for _, parameter := range parameters { - isNull := parameter.IsNull(uint64(row)) - if isNull { - key[written] = 1 - written++ - continue - } - key[written] = 0 - written++ - written += copy(key[written:], parameter.GetRawBytesAt(row)) - } - if written < len(hashtable.StrKeyPadding) { - copy(key[written:], hashtable.StrKeyPadding[written:]) - } - } - hashtable.BytesBatchGenHashStates(&keys[0], &states[0], count) - for localRow := 0; localRow < count; localRow++ { - appendState(states[localRow][0]) - } - } - return nil -} - // XXX I just copy this function. func builtInHash(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - if result.HasFunctionScratch() { - rs := vector.MustFunctionResult[int64](result) - return builtInHashAccounted( - parameters, - result, - length, - func(state uint64) { rs.AppendMustValue(int64(state)) }, - ) - } fillStringGroupStr := func(keys [][]byte, vec *vector.Vector, n int, start int) { if vec.IsConst() { area := vec.GetArea() @@ -2292,15 +2200,6 @@ func builtInHash(parameters []*vector.Vector, result vector.FunctionResultWrappe // builtInHashPartition mirrors builtInHash but returns uint64 so downstream modulo results are non-negative. func builtInHashPartition(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - if result.HasFunctionScratch() { - rs := vector.MustFunctionResult[uint64](result) - return builtInHashAccounted( - parameters, - result, - length, - rs.AppendMustValue, - ) - } fillStringGroupStr := func(keys [][]byte, vec *vector.Vector, n int, start int) { if vec.IsConst() { area := vec.GetArea() diff --git a/pkg/sql/plan/function/func_builtin_jq.go b/pkg/sql/plan/function/func_builtin_jq.go index 9c76af9ae3a8f..ec33dbe8f0442 100644 --- a/pkg/sql/plan/function/func_builtin_jq.go +++ b/pkg/sql/plan/function/func_builtin_jq.go @@ -18,9 +18,7 @@ import ( "bytes" "cmp" "encoding/json" - "errors" "fmt" - "io" "math" "math/big" "slices" @@ -29,11 +27,8 @@ import ( "github.com/itchyny/gojq" "github.com/matrixorigin/matrixone/pkg/common/moerr" - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/container/bytejson" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" - "github.com/matrixorigin/matrixone/pkg/sql/plan/function/functionUtil" "github.com/matrixorigin/matrixone/pkg/vm/process" "golang.org/x/exp/constraints" ) @@ -77,13 +72,6 @@ func (op *opBuiltInJq) tryJq(params []*vector.Vector, result vector.FunctionResu func (op *opBuiltInJq) tryJqImpl(params []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList, isTry bool) error { - var scratchOutput functionScratchOutput - if result.HasFunctionScratch() { - scratchOutput.result = result - op.enc.useWriter(&scratchOutput) - defer op.enc.restoreWriter() - } - p1 := vector.GenerateFunctionStrParameter(params[0]) p2 := vector.GenerateFunctionStrParameter(params[1]) rs := vector.MustFunctionResult[types.Varlena](result) @@ -107,16 +95,14 @@ func (op *opBuiltInJq) tryJqImpl(params []*vector.Vector, result vector.Function err = op.jqImpl(v1, code) } if err != nil { - if isTry && !isJqOutputError(err) { + if isTry { rs.AddNullRange(0, uint64(length)) return nil } else { return err } } - if err := rs.AppendBytes(op.enc.bytes(), false); err != nil { - return err - } + rs.AppendBytes(op.enc.bytes(), false) op.enc.done() } return nil @@ -131,26 +117,20 @@ func (op *opBuiltInJq) tryJqImpl(params []*vector.Vector, result vector.Function for i := uint64(0); i < uint64(length); i++ { v2, null2 := p2.GetStrValue(i) if null2 || selectList.Contains(i) { - if err := rs.AppendBytes(nil, true); err != nil { - return err - } + rs.AppendBytes(nil, true) } else { code, err := op.getJqCode(string(v2)) if err == nil { err = op.jqImpl(v1, code) } if err != nil { - if isTry && !isJqOutputError(err) { - if err := rs.AppendBytes(nil, true); err != nil { - return err - } + if isTry { + rs.AppendBytes(nil, true) } else { return err } } else { - if err := rs.AppendBytes(op.enc.bytes(), false); err != nil { - return err - } + rs.AppendBytes(op.enc.bytes(), false) op.enc.done() } } @@ -177,23 +157,17 @@ func (op *opBuiltInJq) tryJqImpl(params []*vector.Vector, result vector.Function for i := uint64(0); i < uint64(length); i++ { v1, null1 := p1.GetStrValue(i) if null1 || selectList.Contains(i) { - if err := rs.AppendBytes(nil, true); err != nil { - return err - } + rs.AppendBytes(nil, true) } else { err = op.jqImpl(v1, code) if err != nil { - if isTry && !isJqOutputError(err) { - if err := rs.AppendBytes(nil, true); err != nil { - return err - } + if isTry { + rs.AppendBytes(nil, true) } else { return err } } else { - if err := rs.AppendBytes(op.enc.bytes(), false); err != nil { - return err - } + rs.AppendBytes(op.enc.bytes(), false) op.enc.done() } } @@ -204,9 +178,7 @@ func (op *opBuiltInJq) tryJqImpl(params []*vector.Vector, result vector.Function v1, null1 := p1.GetStrValue(i) v2, null2 := p2.GetStrValue(i) if null1 || null2 || selectList.Contains(i) { - if err := rs.AppendBytes(nil, true); err != nil { - return err - } + rs.AppendBytes(nil, true) } else { code, err := op.getJqCode(string(v2)) if err == nil { @@ -214,18 +186,14 @@ func (op *opBuiltInJq) tryJqImpl(params []*vector.Vector, result vector.Function } if err != nil { - if isTry && !isJqOutputError(err) { - if err := rs.AppendBytes(nil, true); err != nil { - return err - } + if isTry { + rs.AppendBytes(nil, true) // continue } else { return err } } else { - if err := rs.AppendBytes(op.enc.bytes(), false); err != nil { - return err - } + rs.AppendBytes(op.enc.bytes(), false) op.enc.done() } } @@ -299,162 +267,19 @@ func (op *opBuiltInJq) getJqCode(jq string) (*gojq.Code, error) { // We removed all the terminal color related code and we write to buffer w // and do not flush until the encoding is done. type JqEncoder struct { - legacy jqLegacyOutput - w jqOutput + w *bytes.Buffer tab bool indent int depth int buf [64]byte } -type jqOutputError struct { - err error -} - -func (e *jqOutputError) Error() string { return e.err.Error() } -func (e *jqOutputError) Unwrap() error { return e.err } - -func isJqOutputError(err error) bool { - var outputErr *jqOutputError - return errors.As(err, &outputErr) -} - -type jqOutput interface { - formatBuffer - Bytes() []byte - Len() int - Reset() - Err() error -} - -type jqLegacyOutput struct { - bytes.Buffer -} - -func (*jqLegacyOutput) Err() error { return nil } - -type functionScratchOutput struct { - result vector.FunctionResultWrapper - data []byte - written int - err error -} - -func (w *functionScratchOutput) ensure(required int) error { - if w.err != nil { - return w.err - } - if required < 0 { - w.err = mpool.ErrAllocationAccountInvalid - return w.err - } - if required <= cap(w.data) { - w.data = w.data[:required] - return nil - } - capacity, ok := mpool.GrowCapacity(int64(cap(w.data)), int64(required)) - if !ok || capacity > int64(math.MaxInt) { - w.err = mpool.ErrAllocationAccountInvalid - return w.err - } - data, selected, err := w.result.ResizeFunctionScratch(int(capacity)) - if err != nil { - w.err = err - return err - } - if !selected { - w.err = mpool.ErrAllocationAccountInvalid - return w.err - } - w.data = data[:required] - return nil -} - -func (w *functionScratchOutput) Write(value []byte) (int, error) { - if len(value) > math.MaxInt-w.written { - w.err = io.ErrShortBuffer - return 0, w.err - } - if err := w.ensure(w.written + len(value)); err != nil { - return 0, err - } - copy(w.data[w.written:], value) - w.written += len(value) - return len(value), nil -} - -func (w *functionScratchOutput) WriteString(value string) (int, error) { - if len(value) > math.MaxInt-w.written { - w.err = io.ErrShortBuffer - return 0, w.err - } - if err := w.ensure(w.written + len(value)); err != nil { - return 0, err - } - copy(w.data[w.written:], value) - w.written += len(value) - return len(value), nil -} - -func (w *functionScratchOutput) WriteByte(value byte) error { - if w.written == math.MaxInt { - w.err = io.ErrShortBuffer - return w.err - } - if err := w.ensure(w.written + 1); err != nil { - return err - } - w.data[w.written] = value - w.written++ - return nil -} - -func (w *functionScratchOutput) WriteRune(value rune) (int, error) { - var encoded [utf8.UTFMax]byte - size := utf8.EncodeRune(encoded[:], value) - return w.Write(encoded[:size]) -} - -func (w *functionScratchOutput) Grow(size int) { - if size < 0 || size > math.MaxInt-w.written { - w.err = io.ErrShortBuffer - return - } - _ = w.ensure(w.written + size) -} - -func (w *functionScratchOutput) Bytes() []byte { - return w.data[:w.written] -} - -func (w *functionScratchOutput) Len() int { return w.written } - -func (w *functionScratchOutput) Reset() { - w.data = w.data[:0] - w.written = 0 - w.err = nil -} - -func (w *functionScratchOutput) Err() error { return w.err } - func (e *JqEncoder) intialize(tab bool, indent int) { - e.legacy.Reset() - e.w = &e.legacy + e.w = new(bytes.Buffer) e.tab = tab e.indent = indent } -func (e *JqEncoder) useWriter(w jqOutput) { - e.w = w - e.done() -} - -func (e *JqEncoder) restoreWriter() { - e.done() - e.w = &e.legacy - e.legacy.Reset() -} - func (e *JqEncoder) bytes() []byte { return e.w.Bytes() } @@ -463,13 +288,6 @@ func (e *JqEncoder) done() { e.depth = 0 } -func (e *JqEncoder) err() error { - if err := e.w.Err(); err != nil { - return &jqOutputError{err: err} - } - return nil -} - func (e *JqEncoder) encode(v any) error { switch v := v.(type) { case nil: @@ -499,7 +317,7 @@ func (e *JqEncoder) encode(v any) error { default: panic(fmt.Sprintf("invalid type: %[1]T (%[1]v)", v)) } - return e.err() + return nil } // ref: floatEncoder in encoding/json @@ -530,22 +348,16 @@ func (e *JqEncoder) encodeFloat64(f float64) { // ref: encodeState#string in encoding/json func (e *JqEncoder) encodeString(s string) { - e.encodeBytes(functionUtil.QuickStrToBytes(s)) -} - -// encodeBytes preserves JSON_ROW's legacy replacement behavior for invalid -// UTF-8 while avoiding a per-row []byte-to-string allocation. -func (e *JqEncoder) encodeBytes(value []byte) { e.w.WriteByte('"') start := 0 - for i := 0; i < len(value); { - if b := value[i]; b < utf8.RuneSelf { + for i := 0; i < len(s); { + if b := s[i]; b < utf8.RuneSelf { if ' ' <= b && b <= '~' && b != '"' && b != '\\' { i++ continue } if start < i { - e.w.Write(value[start:i]) + e.w.WriteString(s[start:i]) } switch b { case '"': @@ -572,20 +384,20 @@ func (e *JqEncoder) encodeBytes(value []byte) { start = i continue } - c, size := utf8.DecodeRune(value[i:]) + c, size := utf8.DecodeRuneInString(s[i:]) if c == utf8.RuneError && size == 1 { if start < i { - e.w.Write(value[start:i]) + e.w.WriteString(s[start:i]) } e.w.WriteString(`\ufffd`) - i++ + i += size start = i continue } i += size } - if start < len(value) { - e.w.Write(value[start:]) + if start < len(s) { + e.w.WriteString(s[start:]) } e.w.WriteByte('"') } @@ -664,290 +476,295 @@ func (e *JqEncoder) writeIndent() { } func (e *JqEncoder) writeIndentInternal(n int, spaces string) { - for n > 0 { - length := min(n, len(spaces)) - e.w.WriteString(spaces[:length]) - n -= length + if l := len(spaces); n <= l { + e.w.WriteString(spaces[:n]) + } else { + e.w.WriteString(spaces) + for n -= l; n > 0; n, l = n-l, l*2 { + if n < l { + l = n + } + e.w.Write(e.w.Bytes()[e.w.Len()-l:]) + } } } type opBuiltInJsonRow struct { - enc JqEncoder - columns []jsonRowColumnEncoder + enc []JqEncoder } func newOpBuiltInJsonRow() *opBuiltInJsonRow { var op opBuiltInJsonRow - op.enc.intialize(false, 0) return &op } -func (op *opBuiltInJsonRow) jsonRow(params []*vector.Vector, result vector.FunctionResultWrapper, - proc *process.Process, length int, selectList *FunctionSelectList) error { - var scratchOutput functionScratchOutput - if result.HasFunctionScratch() { - scratchOutput.result = result - op.enc.useWriter(&scratchOutput) - defer op.enc.restoreWriter() +func (op *opBuiltInJsonRow) grow(length int) { + if len(op.enc) == 0 { + op.enc = make([]JqEncoder, length) + for i := 0; i < length; i++ { + op.enc[i].intialize(false, 0) + } + } else if length > len(op.enc) { + for i := len(op.enc); i < length; i++ { + op.enc = append(op.enc, JqEncoder{}) + op.enc[i].intialize(false, 0) + } } +} +func (op *opBuiltInJsonRow) jsonRow(params []*vector.Vector, result vector.FunctionResultWrapper, + proc *process.Process, length int, selectList *FunctionSelectList) error { + op.grow(length) rs := vector.MustFunctionResult[types.Varlena](result) - if cap(op.columns) < len(params) { - op.columns = make([]jsonRowColumnEncoder, len(params)) - } else { - op.columns = op.columns[:len(params)] - } - for idx, param := range params { - column, err := prepareJSONRowColumn(param, proc) - if err != nil { - clear(op.columns) - return err - } - op.columns[idx] = column + ulen := uint64(length) + + for j := 0; j < length; j++ { + op.enc[j].w.WriteByte('[') } - defer clear(op.columns) - op.enc.done() - defer op.enc.done() - for row := uint64(0); row < uint64(length); row++ { - if selectList.Contains(row) { - if err := rs.AppendBytes(nil, true); err != nil { - return err + for i := 0; i < len(params); i++ { + // write separator first + if i > 0 { + for j := 0; j < length; j++ { + op.enc[j].w.WriteByte(',') } - continue } - op.enc.w.WriteByte('[') - for paramIdx := range op.columns { - if paramIdx > 0 { - op.enc.w.WriteByte(',') - } - if err := op.columns[paramIdx](&op.enc, row); err != nil { + + // oh the dreaded type switch + fromType := params[i].GetType() + switch fromType.Oid { + case types.T_any: // scalar null + op.encodeScalarNull(ulen) + case types.T_bool: + op.encodeBool(params[i], ulen) + case types.T_int8: + encodeInt[int8](op, params[i], ulen) + case types.T_int16: + encodeInt[int16](op, params[i], ulen) + case types.T_int32: + encodeInt[int32](op, params[i], ulen) + case types.T_int64: + encodeInt[int64](op, params[i], ulen) + case types.T_uint8: + encodeInt[uint8](op, params[i], ulen) + case types.T_uint16: + encodeInt[uint16](op, params[i], ulen) + case types.T_uint32: + encodeInt[uint32](op, params[i], ulen) + case types.T_uint64: + encodeInt[uint64](op, params[i], ulen) + case types.T_float32: + encodeFloat[float32](op, params[i], ulen) + case types.T_float64: + encodeFloat[float64](op, params[i], ulen) + case types.T_decimal64: + encodeDecimal[types.Decimal64](op, params[i], ulen) + case types.T_decimal128: + encodeDecimal[types.Decimal128](op, params[i], ulen) + case types.T_date: + encodeFixedStringer[types.Date](op, params[i], ulen) + case types.T_time: + encodeFixedStringer[types.Time](op, params[i], ulen) + case types.T_datetime: + encodeFixedStringer[types.Datetime](op, params[i], ulen) + case types.T_timestamp: + encodeFixedStringer[types.Timestamp](op, params[i], ulen) + case types.T_char, types.T_varchar, types.T_text: + encodeString(op, params[i], ulen) + case types.T_binary, types.T_varbinary, types.T_blob: + // well, in cast, we handle binary as if they are string. + // However it id deemed too dangerous to do so in json_row. + return moerr.NewInvalidInputf(proc.Ctx, "binary data not supported json_row: %v", fromType.String()) + case types.T_array_float32: + // vector of float, we will encode them as json array + encodeFloatArray[float32](op, params[i], ulen) + case types.T_array_float64: + // vector of float, we will encode them as json array + encodeFloatArray[float64](op, params[i], ulen) + case types.T_array_bf16: + encodeNarrowArray[types.BF16](op, params[i], ulen, + func(x types.BF16) float64 { return float64(x.ToFloat32()) }) + case types.T_array_float16: + encodeNarrowArray[types.Float16](op, params[i], ulen, + func(x types.Float16) float64 { return float64(x.ToFloat32()) }) + case types.T_array_int8: + encodeNarrowArray[int8](op, params[i], ulen, + func(x int8) float64 { return float64(x) }) + case types.T_array_uint8: + encodeNarrowArray[uint8](op, params[i], ulen, + func(x uint8) float64 { return float64(x) }) + case types.T_uuid: + encodeFixedStringer[types.Uuid](op, params[i], ulen) + case types.T_json: + if err := encodeJson(op, params[i], ulen); err != nil { return err } + default: + return moerr.NewInvalidInputf(proc.Ctx, "unsupported type for json_row: %v", fromType.String()) } - op.enc.w.WriteByte(']') - if err := op.enc.err(); err != nil { - return err - } - if err := rs.AppendBytes(op.enc.bytes(), false); err != nil { - return err + } + + for j := 0; j < length; j++ { + op.enc[j].w.WriteByte(']') + if selectList.Contains(uint64(j)) { + rs.AppendBytes(nil, true) + } else { + rs.AppendBytes(op.enc[j].bytes(), false) } - op.enc.done() + op.enc[j].done() } return nil } -type jsonRowColumnEncoder func(*JqEncoder, uint64) error - -func encodeJSONRowNull(e *JqEncoder, _ uint64) error { - e.w.WriteString("null") - return nil +func (op *opBuiltInJsonRow) encodeScalarNull(length uint64) { + for i := uint64(0); i < length; i++ { + op.enc[i].w.WriteString("null") + } } -func prepareJSONRowColumn( - v *vector.Vector, - proc *process.Process, -) (jsonRowColumnEncoder, error) { - switch fromType := v.GetType(); fromType.Oid { - case types.T_any: - return encodeJSONRowNull, nil - case types.T_bool: - param := vector.GenerateFunctionFixedTypeParameter[bool](v) - return func(e *JqEncoder, row uint64) error { - value, isNull := param.GetValue(row) - if isNull { - return encodeJSONRowNull(e, row) - } - if value { - e.w.WriteString("true") +func (op *opBuiltInJsonRow) encodeBool(v *vector.Vector, length uint64) { + p := vector.GenerateFunctionFixedTypeParameter[bool](v) + for i := uint64(0); i < length; i++ { + v, null := p.GetValue(i) + if null { + op.enc[i].w.WriteString("null") + } else { + if v { + op.enc[i].w.WriteString("true") } else { - e.w.WriteString("false") - } - return nil - }, nil - case types.T_int8: - return prepareJSONRowSignedColumn[int8](v), nil - case types.T_int16: - return prepareJSONRowSignedColumn[int16](v), nil - case types.T_int32: - return prepareJSONRowSignedColumn[int32](v), nil - case types.T_int64: - return prepareJSONRowSignedColumn[int64](v), nil - case types.T_uint8: - return prepareJSONRowUnsignedColumn[uint8](v), nil - case types.T_uint16: - return prepareJSONRowUnsignedColumn[uint16](v), nil - case types.T_uint32: - return prepareJSONRowUnsignedColumn[uint32](v), nil - case types.T_uint64: - return prepareJSONRowUnsignedColumn[uint64](v), nil - case types.T_float32: - return prepareJSONRowFloatColumn[float32](v), nil - case types.T_float64: - return prepareJSONRowFloatColumn[float64](v), nil - case types.T_decimal64: - return prepareJSONRowDecimalColumn[types.Decimal64](v), nil - case types.T_decimal128: - return prepareJSONRowDecimalColumn[types.Decimal128](v), nil - case types.T_date: - return prepareJSONRowStringerColumn[types.Date](v), nil - case types.T_time: - return prepareJSONRowStringerColumn[types.Time](v), nil - case types.T_datetime: - return prepareJSONRowStringerColumn[types.Datetime](v), nil - case types.T_timestamp: - return prepareJSONRowStringerColumn[types.Timestamp](v), nil - case types.T_char, types.T_varchar, types.T_text: - return prepareJSONRowStringColumn(v), nil - case types.T_array_float32: - return prepareJSONRowArrayColumn(v, - func(value float32) float64 { return float64(value) }), nil - case types.T_array_float64: - return prepareJSONRowArrayColumn(v, - func(value float64) float64 { return value }), nil - case types.T_array_bf16: - return prepareJSONRowArrayColumn(v, - func(value types.BF16) float64 { return float64(value.ToFloat32()) }), nil - case types.T_array_float16: - return prepareJSONRowArrayColumn(v, - func(value types.Float16) float64 { return float64(value.ToFloat32()) }), nil - case types.T_array_int8: - return prepareJSONRowArrayColumn(v, - func(value int8) float64 { return float64(value) }), nil - case types.T_array_uint8: - return prepareJSONRowArrayColumn(v, - func(value uint8) float64 { return float64(value) }), nil - case types.T_uuid: - return prepareJSONRowStringerColumn[types.Uuid](v), nil - case types.T_json: - param := vector.GenerateFunctionStrParameter(v) - return func(e *JqEncoder, row uint64) error { - value, isNull := param.GetStrValue(row) - if isNull { - return encodeJSONRowNull(e, row) + op.enc[i].w.WriteString("false") } - return bytejson.WriteJSONText(e.w, types.DecodeJson(value)) - }, nil - case types.T_binary, types.T_varbinary, types.T_blob: - return nil, moerr.NewInvalidInputf(proc.Ctx, - "binary data not supported json_row: %v", - fromType.String()) - default: - return nil, moerr.NewInvalidInputf(proc.Ctx, - "unsupported type for json_row: %v", - fromType.String()) + } } } -func prepareJSONRowSignedColumn[T constraints.Signed]( - v *vector.Vector, -) jsonRowColumnEncoder { - param := vector.GenerateFunctionFixedTypeParameter[T](v) - return func(e *JqEncoder, row uint64) error { - value, isNull := param.GetValue(row) - if isNull { - return encodeJSONRowNull(e, row) +func encodeInt[T constraints.Integer](op *opBuiltInJsonRow, v *vector.Vector, length uint64) { + p := vector.GenerateFunctionFixedTypeParameter[T](v) + for i := uint64(0); i < length; i++ { + v, null := p.GetValue(i) + if null { + op.enc[i].w.WriteString("null") + } else { + op.enc[i].w.Write(strconv.AppendInt(op.enc[i].buf[:0], int64(v), 10)) } - e.w.Write(strconv.AppendInt(e.buf[:0], int64(value), 10)) - return nil } } -func prepareJSONRowUnsignedColumn[T constraints.Unsigned]( - v *vector.Vector, -) jsonRowColumnEncoder { - param := vector.GenerateFunctionFixedTypeParameter[T](v) - return func(e *JqEncoder, row uint64) error { - value, isNull := param.GetValue(row) - if isNull { - return encodeJSONRowNull(e, row) +func encodeFloat[T constraints.Float](op *opBuiltInJsonRow, v *vector.Vector, length uint64) { + p := vector.GenerateFunctionFixedTypeParameter[T](v) + for i := uint64(0); i < length; i++ { + v, null := p.GetValue(i) + if null { + op.enc[i].w.WriteString("null") + } else { + op.enc[i].encodeFloat64(float64(v)) } - e.w.Write(strconv.AppendUint(e.buf[:0], uint64(value), 10)) - return nil } } -func prepareJSONRowFloatColumn[T constraints.Float]( - v *vector.Vector, -) jsonRowColumnEncoder { - param := vector.GenerateFunctionFixedTypeParameter[T](v) - return func(e *JqEncoder, row uint64) error { - value, isNull := param.GetValue(row) - if isNull { - return encodeJSONRowNull(e, row) +func encodeDecimal[T types.DecimalWithFormat](op *opBuiltInJsonRow, v *vector.Vector, length uint64) { + p := vector.GenerateFunctionFixedTypeParameter[T](v) + fromTyp := v.GetType() + for i := uint64(0); i < length; i++ { + v, null := p.GetValue(i) + if null { + op.enc[i].w.WriteString("null") + } else { + bs := []byte(v.Format(fromTyp.Scale)) + op.enc[i].w.Write(bs) } - e.encodeFloat64(float64(value)) - return nil } } -func prepareJSONRowDecimalColumn[T types.DecimalWithFormat]( - v *vector.Vector, -) jsonRowColumnEncoder { - param := vector.GenerateFunctionFixedTypeParameter[T](v) - scale := v.GetType().Scale - return func(e *JqEncoder, row uint64) error { - value, isNull := param.GetValue(row) - if isNull { - return encodeJSONRowNull(e, row) +func encodeFixedStringer[T types.FixedWithStringer](op *opBuiltInJsonRow, v *vector.Vector, length uint64) { + p := vector.GenerateFunctionFixedTypeParameter[T](v) + for i := uint64(0); i < length; i++ { + v, null := p.GetValue(i) + if null { + op.enc[i].w.WriteString("null") + } else { + op.enc[i].encodeString(v.String()) } - e.w.WriteString(value.Format(scale)) - return nil } } -func prepareJSONRowStringerColumn[T types.FixedWithStringer]( - v *vector.Vector, -) jsonRowColumnEncoder { - param := vector.GenerateFunctionFixedTypeParameter[T](v) - return func(e *JqEncoder, row uint64) error { - value, isNull := param.GetValue(row) - if isNull { - return encodeJSONRowNull(e, row) +func encodeString(op *opBuiltInJsonRow, v *vector.Vector, length uint64) { + p := vector.GenerateFunctionStrParameter(v) + for i := uint64(0); i < length; i++ { + v, null := p.GetStrValue(i) + if null { + op.enc[i].w.WriteString("null") + } else { + op.enc[i].encodeString(string(v)) } - e.encodeString(value.String()) - return nil } } -func prepareJSONRowStringColumn(v *vector.Vector) jsonRowColumnEncoder { - param := vector.GenerateFunctionStrParameter(v) - return func(e *JqEncoder, row uint64) error { - value, isNull := param.GetStrValue(row) - if isNull { - return encodeJSONRowNull(e, row) +func encodeFloatArray[T constraints.Float](op *opBuiltInJsonRow, v *vector.Vector, length uint64) { + // GenStrParam: array is varlena also. + p := vector.GenerateFunctionStrParameter(v) + for i := uint64(0); i < length; i++ { + v, null := p.GetStrValue(i) + if null { + op.enc[i].w.WriteString("null") + } else { + vv := types.BytesToArray[T](v) + op.enc[i].w.WriteByte('[') + for j, val := range vv { + if j > 0 { + op.enc[i].w.WriteByte(',') + } + ff := float64(val) + op.enc[i].encodeFloat64(ff) + } + op.enc[i].w.WriteByte(']') } - e.encodeBytes(value) - return e.err() } } -func prepareJSONRowArrayColumn[T types.ArrayElement]( - v *vector.Vector, - toFloat64 func(T) float64, -) jsonRowColumnEncoder { - param := vector.GenerateFunctionStrParameter(v) - return func(e *JqEncoder, row uint64) error { - value, isNull := param.GetStrValue(row) - if isNull { - return encodeJSONRowNull(e, row) +// encodeNarrowArray is encodeFloatArray for element types that are not +// constraints.Float: BF16/Float16 need ToFloat32(), int8/uint8 are plain +// integers. The JSON shape emitted is identical to the f32/f64 arrays. +func encodeNarrowArray[T types.ArrayElement](op *opBuiltInJsonRow, v *vector.Vector, length uint64, toF64 func(T) float64) { + p := vector.GenerateFunctionStrParameter(v) + for i := uint64(0); i < length; i++ { + v, null := p.GetStrValue(i) + if null { + op.enc[i].w.WriteString("null") + } else { + vv := types.BytesToArray[T](v) + op.enc[i].w.WriteByte('[') + for j, val := range vv { + if j > 0 { + op.enc[i].w.WriteByte(',') + } + op.enc[i].encodeFloat64(toF64(val)) + } + op.enc[i].w.WriteByte(']') } - encodeJSONRowArray(e, types.BytesToArray[T](value), toFloat64) - return nil } } -func encodeJSONRowArray[T types.ArrayElement]( - e *JqEncoder, - values []T, - toFloat64 func(T) float64, -) { - e.w.WriteByte('[') - for idx, value := range values { - if idx > 0 { - e.w.WriteByte(',') +func encodeJson(op *opBuiltInJsonRow, v *vector.Vector, length uint64) error { + // GenStrParam: json is varlena also. + p := vector.GenerateFunctionStrParameter(v) + for i := uint64(0); i < length; i++ { + v, null := p.GetStrValue(i) + if null { + op.enc[i].w.WriteString("null") + } else { + bj := types.DecodeJson(v) + val, err := bj.MarshalJSON() + // this should a valid json and we should never + // error here. Check it anyway. + if err != nil { + return err + } + // note here we already have a valid json string + // do NOT use encodeString, which will escape + // the string again. + op.enc[i].w.Write(val) } - e.encodeFloat64(toFloat64(value)) } - e.w.WriteByte(']') + return nil } diff --git a/pkg/sql/plan/function/func_builtin_json.go b/pkg/sql/plan/function/func_builtin_json.go index 7889b3556a404..8579a6e4c4f31 100644 --- a/pkg/sql/plan/function/func_builtin_json.go +++ b/pkg/sql/plan/function/func_builtin_json.go @@ -19,7 +19,7 @@ import ( "context" "encoding/binary" "encoding/json" - "math" + "fmt" "strconv" "strings" "time" @@ -86,12 +86,11 @@ func encodeJsonOrderingParam(value []byte) ([]byte, error) { } type opBuiltInJsonExtract struct { - allConst bool - npath int - pathStrs []string - paths []*bytejson.Path - pathWrappers []vector.FunctionParameterWrapper[types.Varlena] - simple bool + allConst bool + npath int + pathStrs []string + paths []*bytejson.Path + simple bool } type opBuiltInJsonContains struct{} @@ -980,7 +979,7 @@ func computeStringJsonRemove(json []byte, paths []*bytejson.Path) (bytejson.Byte return bj.Remove(paths) } -func (op *opBuiltInJsonExtract) buildPath(params []*vector.Vector, _ int, selectList *FunctionSelectList) error { +func (op *opBuiltInJsonExtract) buildPath(params []*vector.Vector, length int, selectList *FunctionSelectList) error { op.npath = len(params) - 1 if op.npath == 0 { return nil @@ -1026,20 +1025,19 @@ func (op *opBuiltInJsonExtract) buildPath(params []*vector.Vector, _ int, select return nil } } - } else if len(op.pathStrs) != op.npath { - op.pathStrs = make([]string, op.npath) - op.paths = make([]*bytejson.Path, op.npath) + } else { + op.pathStrs = make([]string, op.npath*length) + op.paths = make([]*bytejson.Path, op.npath*length) } - if len(op.pathWrappers) != op.npath { - op.pathWrappers = make([]vector.FunctionParameterWrapper[types.Varlena], op.npath) - } + // Do it! + pathWrapers := make([]vector.FunctionParameterWrapper[types.Varlena], op.npath) for i := 0; i < op.npath; i++ { - op.pathWrappers[i] = vector.GenerateFunctionStrParameter(params[i+1]) + pathWrapers[i] = vector.GenerateFunctionStrParameter(params[i+1]) } if op.allConst { - if _, err := op.buildOnePath(0); err != nil { + if err := op.buildOnePath(pathWrapers, 0, op.pathStrs, op.paths); err != nil { return err } op.simple = true @@ -1051,45 +1049,63 @@ func (op *opBuiltInJsonExtract) buildPath(params []*vector.Vector, _ int, select } return nil } else { - clear(op.pathStrs) - clear(op.paths) - op.simple = false + op.simple = true + for i := 0; i < length; i++ { + strs := op.pathStrs[i*op.npath : (i+1)*op.npath] + paths := op.paths[i*op.npath : (i+1)*op.npath] + if selectList.Contains(uint64(i)) { + for j := 0; j < op.npath; j++ { + strs[j] = "" + paths[j] = nil + } + continue + } + if err := op.buildOnePath(pathWrapers, i, strs, paths); err != nil { + return err + } + for _, p := range paths { + if p == nil { + continue + } + op.simple = op.simple && p.IsSimple() + } + } } return nil } func (op *opBuiltInJsonExtract) getPaths(i uint64) []*bytejson.Path { - if !op.allConst && len(op.paths) > op.npath { - return op.paths[i*uint64(op.npath) : (i+1)*uint64(op.npath)] + if op.allConst { + return op.paths } - return op.paths + return op.paths[i*uint64(op.npath) : (i+1)*uint64(op.npath)] } -func (op *opBuiltInJsonExtract) buildOnePath(i int) (bool, error) { +func (op *opBuiltInJsonExtract) buildOnePath(paramWrappers []vector.FunctionParameterWrapper[types.Varlena], i int, strs []string, paths []*bytejson.Path) error { skip := false - simple := true - for j := 0; j < len(op.pathWrappers); j++ { - pathBytes, pIsNull := op.pathWrappers[j].GetStrValue(uint64(i)) + for j := 0; j < len(paramWrappers); j++ { + pathBytes, pIsNull := paramWrappers[j].GetStrValue(uint64(i)) if pIsNull { skip = true break } - op.pathStrs[j] = string(pathBytes) - p, err := types.ParseStringToPath(op.pathStrs[j]) + strs[j] = string(pathBytes) + p, err := types.ParseStringToPath(strs[j]) if err != nil { - return false, err + return err } - op.paths[j] = &p - simple = simple && p.IsSimple() + paths[j] = &p } if skip { - clear(op.pathStrs) - clear(op.paths) + for j := 0; j < len(paramWrappers); j++ { + strs[j] = "" + paths[j] = nil + } } - return simple, nil + return nil } func (op *opBuiltInJsonExtract) jsonExtract(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { @@ -1135,13 +1151,6 @@ func (op *opBuiltInJsonExtract) jsonExtract(parameters []*vector.Vector, result } continue } - rowSimple := op.simple - if !op.allConst { - rowSimple, err = op.buildOnePath(int(i)) - if err != nil { - return err - } - } jsonBytes, jIsNull := jsonWrapper.GetStrValue(i) if jIsNull { if err = rs.AppendBytes(nil, true); err != nil { @@ -1157,15 +1166,7 @@ func (op *opBuiltInJsonExtract) jsonExtract(parameters []*vector.Vector, result } continue } else { - rowFn := fn - if rowSimple { - if jsonVec.GetType().Oid == types.T_json { - rowFn = computeJsonSimpleWithExists - } else { - rowFn = computeStringSimpleWithExists - } - } - out, exists, err := rowFn(jsonBytes, paths) + out, exists, err := fn(jsonBytes, paths) if err != nil { return err } @@ -1206,7 +1207,7 @@ func (op *opBuiltInJsonExtract) jsonExtractString(parameters []*vector.Vector, r return err } - if op.allConst && (!op.simple || op.npath > 1) { + if !op.simple || op.npath > 1 { return moerr.NewInvalidInput(proc.Ctx, "json_extract_string should use a path that retrives a single value") } if jsonVec.GetType().Oid == types.T_json { @@ -1222,16 +1223,6 @@ func (op *opBuiltInJsonExtract) jsonExtractString(parameters []*vector.Vector, r } continue } - rowSimple := op.simple - if !op.allConst { - rowSimple, err = op.buildOnePath(int(i)) - if err != nil { - return err - } - } - if !rowSimple || op.npath > 1 { - return moerr.NewInvalidInput(proc.Ctx, "json_extract_string should use a path that retrives a single value") - } jsonBytes, jIsNull := jsonWrapper.GetStrValue(i) if jIsNull { if err = rs.AppendBytes(nil, true); err != nil { @@ -1299,7 +1290,7 @@ func (op *opBuiltInJsonExtract) jsonExtractFloat64(parameters []*vector.Vector, if err = op.buildPath(parameters, length, selectList); err != nil { return err } - if op.allConst && (!op.simple || op.npath > 1) { + if !op.simple || op.npath > 1 { return moerr.NewInvalidInput(proc.Ctx, "json_extract_float64 should use a path that retrives a single value") } @@ -1316,16 +1307,6 @@ func (op *opBuiltInJsonExtract) jsonExtractFloat64(parameters []*vector.Vector, } continue } - rowSimple := op.simple - if !op.allConst { - rowSimple, err = op.buildOnePath(int(i)) - if err != nil { - return err - } - } - if !rowSimple || op.npath > 1 { - return moerr.NewInvalidInput(proc.Ctx, "json_extract_float64 should use a path that retrives a single value") - } jsonBytes, jIsNull := jsonWrapper.GetStrValue(i) if jIsNull { if err = rs.Append(0, true); err != nil { @@ -1757,13 +1738,6 @@ func (op *opBuiltInJsonSet) buildJsonFunction(parameters []*vector.Vector, resul jsonVec := parameters[0] jsonWrapper := vector.GenerateFunctionStrParameter(jsonVec) rs := vector.MustFunctionResult[types.Varlena](result) - valueCount := (len(parameters) - 1) / 2 - pathExprs := make([]*bytejson.Path, valueCount) - valExprs := make([]bytejson.ByteJson, valueCount) - valueEncoders := make([]bytejson.ByteJsonDataEncoder, valueCount) - defer clear(pathExprs) - defer clear(valueEncoders) - accountedScratch := result.HasFunctionScratch() if selectList.IgnoreAllRow() { for i := 0; i < length; i++ { @@ -1814,6 +1788,7 @@ rowLoop: } // build all paths + pathExprs := make([]*bytejson.Path, 0, (len(parameters)-1)/2+1) for j := 1; j < len(parameters); j += 2 { pathBytes, pIsNull := vector.GenerateFunctionStrParameter(parameters[j]).GetStrValue(uint64(i)) if pIsNull { @@ -1832,53 +1807,17 @@ rowLoop: return moerr.NewInvalidArg(proc.Ctx, jsonModifyFunctionName(jsonFuncType), "invalid path expression") } - pathExprs[j/2] = &p + pathExprs = append(pathExprs, &p) } - // Build one storage-compatible representation for all values. Exact - // execution owns this backing through FunctionResult scratch; legacy - // execution preserves a row-local Go buffer. - valueBytes := 0 + // build all values + valExprs := make([]bytejson.ByteJson, 0, (len(parameters)-1)/2+1) for j := 2; j < len(parameters); j += 2 { - encoder, err := (&opBuiltInJsonArray{}).buildValueEncoder( - proc, - parameters[j], - int(i), - ) - if err != nil { - return err - } - valueEncoders[j/2-1] = encoder - size := uint64(encoder.DataSize()) + 1 - if size > uint64(math.MaxInt-valueBytes) { - return moerr.NewInvalidArg(proc.Ctx, jsonModifyFunctionName(jsonFuncType), "JSON value is too large") - } - valueBytes += int(size) - } - var valueStorage []byte - if accountedScratch { - valueStorage, _, err = result.ResizeFunctionScratch(valueBytes) + val, err := op.buildJsonModifyValue(proc, parameters[j], int(i)) if err != nil { return err } - } else { - valueStorage = make([]byte, valueBytes) - } - offset := 0 - for idx, encoder := range valueEncoders { - valueStorage[offset] = byte(encoder.TypeCode()) - size := int(encoder.DataSize()) - written, encodeErr := encoder.EncodeDataInto( - valueStorage[offset+1 : offset+1+size], - ) - if encodeErr != nil { - return encodeErr - } - if written != size { - return moerr.NewInternalErrorNoCtx("JSON value encoder size mismatch") - } - valExprs[idx] = types.DecodeJson(valueStorage[offset : offset+1+size]) - offset += size + 1 + valExprs = append(valExprs, val) } out, err := fn(jsonBytes, pathExprs, valExprs) @@ -1911,6 +1850,14 @@ func jsonModifyFunctionName(jsonFuncType bytejson.JsonModifyType) string { } } +func (op *opBuiltInJsonSet) buildJsonModifyValue(proc *process.Process, v *vector.Vector, row int) (bytejson.ByteJson, error) { + elem, err := (&opBuiltInJsonArray{}).convertToAny(proc, v, row) + if err != nil { + return bytejson.Null, err + } + return bytejson.CreateByteJSON(elem) +} + type opBuiltInJsonArray struct{} func newOpBuiltInJsonArray() *opBuiltInJsonArray { @@ -1920,8 +1867,6 @@ func newOpBuiltInJsonArray() *opBuiltInJsonArray { func (op *opBuiltInJsonArray) jsonArray(params []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { rs := vector.MustFunctionResult[types.Varlena](result) - encoders := make([]bytejson.ByteJsonDataEncoder, len(params)) - defer clear(encoders) if selectList != nil && selectList.IgnoreAllRow() { for j := 0; j < length; j++ { @@ -1939,225 +1884,250 @@ func (op *opBuiltInJsonArray) jsonArray(params []*vector.Vector, result vector.F } continue } - for i := range params { - encoder, err := op.buildValueEncoder(proc, params[i], j) + elems := make([]any, 0, len(params)) + for i := 0; i < len(params); i++ { + elem, err := op.convertToAny(proc, params[i], j) if err != nil { return err } - encoders[i] = encoder + elems = append(elems, elem) + } + + bj, err := bytejson.CreateByteJSON(elems) + if err != nil { + return err } - encoder, err := bytejson.NewArrayDataEncoder(encoders) + dt, err := bj.Marshal() if err != nil { return err } - if err := rs.AppendByteJsonEncoded(encoder); err != nil { + if err := rs.AppendBytes(dt, false); err != nil { return err } } return nil } -func (op *opBuiltInJsonArray) buildValueEncoder( - proc *process.Process, - v *vector.Vector, - row int, -) (bytejson.ByteJsonDataEncoder, error) { - if v.IsNull(uint64(row)) { - return bytejson.NewLiteralDataEncoder(bytejson.LiteralNull), nil +func (op *opBuiltInJsonArray) convertToAny(proc *process.Process, v *vector.Vector, row int) (any, error) { + ctx := context.Background() + if proc != nil { + ctx = proc.Ctx } fromType := v.GetType() switch fromType.Oid { case types.T_bool: - literal := bytejson.LiteralFalse - if vector.GetFixedAtNoTypeCheck[bool](v, row) { - literal = bytejson.LiteralTrue + if v.IsNull(uint64(row)) { + return nil, nil } - return bytejson.NewLiteralDataEncoder(literal), nil + return vector.GetFixedAtNoTypeCheck[bool](v, row), nil case types.T_int8: - return bytejson.NewInt64DataEncoder( - int64(vector.GetFixedAtNoTypeCheck[int8](v, row)), - ), nil + if v.IsNull(uint64(row)) { + return nil, nil + } + return int64(vector.GetFixedAtNoTypeCheck[int8](v, row)), nil case types.T_int16: - return bytejson.NewInt64DataEncoder( - int64(vector.GetFixedAtNoTypeCheck[int16](v, row)), - ), nil + if v.IsNull(uint64(row)) { + return nil, nil + } + return int64(vector.GetFixedAtNoTypeCheck[int16](v, row)), nil case types.T_int32: - return bytejson.NewInt64DataEncoder( - int64(vector.GetFixedAtNoTypeCheck[int32](v, row)), - ), nil + if v.IsNull(uint64(row)) { + return nil, nil + } + return int64(vector.GetFixedAtNoTypeCheck[int32](v, row)), nil case types.T_int64: - return bytejson.NewInt64DataEncoder( - vector.GetFixedAtNoTypeCheck[int64](v, row), - ), nil + if v.IsNull(uint64(row)) { + return nil, nil + } + return vector.GetFixedAtNoTypeCheck[int64](v, row), nil case types.T_uint8: - return bytejson.NewUint64DataEncoder( - uint64(vector.GetFixedAtNoTypeCheck[uint8](v, row)), - ), nil + if v.IsNull(uint64(row)) { + return nil, nil + } + return uint64(vector.GetFixedAtNoTypeCheck[uint8](v, row)), nil case types.T_uint16: - return bytejson.NewUint64DataEncoder( - uint64(vector.GetFixedAtNoTypeCheck[uint16](v, row)), - ), nil + if v.IsNull(uint64(row)) { + return nil, nil + } + return uint64(vector.GetFixedAtNoTypeCheck[uint16](v, row)), nil case types.T_uint32: - return bytejson.NewUint64DataEncoder( - uint64(vector.GetFixedAtNoTypeCheck[uint32](v, row)), - ), nil + if v.IsNull(uint64(row)) { + return nil, nil + } + return uint64(vector.GetFixedAtNoTypeCheck[uint32](v, row)), nil case types.T_uint64: - return bytejson.NewUint64DataEncoder( - vector.GetFixedAtNoTypeCheck[uint64](v, row), - ), nil + if v.IsNull(uint64(row)) { + return nil, nil + } + return vector.GetFixedAtNoTypeCheck[uint64](v, row), nil case types.T_float32: - return bytejson.NewFloat64DataEncoder( - float64(vector.GetFixedAtNoTypeCheck[float32](v, row)), - ), nil + if v.IsNull(uint64(row)) { + return nil, nil + } + return float64(vector.GetFixedAtNoTypeCheck[float32](v, row)), nil case types.T_float64: - return bytejson.NewFloat64DataEncoder( - vector.GetFixedAtNoTypeCheck[float64](v, row), - ), nil - case types.T_char, types.T_varchar, types.T_text, types.T_geometry: - return bytejson.NewTypedStringDataEncoder( - bytejson.TpCodeString, - v.GetBytesAt(row), - ) + if v.IsNull(uint64(row)) { + return nil, nil + } + return vector.GetFixedAtNoTypeCheck[float64](v, row), nil + case types.T_char, types.T_varchar, types.T_text: + if v.IsNull(uint64(row)) { + return nil, nil + } + return string(v.GetBytesAt(row)), nil case types.T_json: + if v.IsNull(uint64(row)) { + return nil, nil + } data := v.GetBytesAt(row) if len(data) == 0 { - return bytejson.NewLiteralDataEncoder(bytejson.LiteralNull), nil + return nil, nil } - return bytejson.NewRawDataEncoder(types.DecodeJson(data)) + bj := types.DecodeJson(data) + return bj, nil case types.T_date: - return newJSONTypedStringEncoder( - bytejson.TpCodeDate, - vector.GetFixedAtNoTypeCheck[types.Date](v, row).String(), - ) + if v.IsNull(uint64(row)) { + return nil, nil + } + return newTypedByteJson(bytejson.TpCodeDate, vector.GetFixedAtNoTypeCheck[types.Date](v, row).String()), nil case types.T_time: - return newJSONTypedStringEncoder( - bytejson.TpCodeTime, - vector.GetFixedAtNoTypeCheck[types.Time](v, row).String2(fromType.Scale), - ) + if v.IsNull(uint64(row)) { + return nil, nil + } + return newTypedByteJson(bytejson.TpCodeTime, vector.GetFixedAtNoTypeCheck[types.Time](v, row).String2(fromType.Scale)), nil case types.T_datetime: - return newJSONTypedStringEncoder( - bytejson.TpCodeDatetime, - vector.GetFixedAtNoTypeCheck[types.Datetime](v, row).String2(fromType.Scale), - ) + if v.IsNull(uint64(row)) { + return nil, nil + } + return newTypedByteJson(bytejson.TpCodeDatetime, vector.GetFixedAtNoTypeCheck[types.Datetime](v, row).String2(fromType.Scale)), nil case types.T_timestamp: - return newJSONTypedStringEncoder( - bytejson.TpCodeDatetime, - vector.GetFixedAtNoTypeCheck[types.Timestamp](v, row).String2( - jsonSessionTimeZone(proc), - fromType.Scale, - ), - ) + if v.IsNull(uint64(row)) { + return nil, nil + } + return newTypedByteJson(bytejson.TpCodeDatetime, vector.GetFixedAtNoTypeCheck[types.Timestamp](v, row).String2(jsonSessionTimeZone(proc), fromType.Scale)), nil case types.T_decimal64: - return newJSONTypedStringEncoder( - bytejson.TpCodeDecimal, - string(vector.GetFixedAtNoTypeCheck[types.Decimal64](v, row).Format(fromType.Scale)), - ) + if v.IsNull(uint64(row)) { + return nil, nil + } + val := vector.GetFixedAtNoTypeCheck[types.Decimal64](v, row) + return newTypedByteJson(bytejson.TpCodeDecimal, string(val.Format(fromType.Scale))), nil case types.T_decimal128: - return newJSONTypedStringEncoder( - bytejson.TpCodeDecimal, - string(vector.GetFixedAtNoTypeCheck[types.Decimal128](v, row).Format(fromType.Scale)), - ) - case types.T_decimal256: - return newJSONTypedStringEncoder( - bytejson.TpCodeDecimal, - string(vector.GetFixedAtNoTypeCheck[types.Decimal256](v, row).Format(fromType.Scale)), - ) + if v.IsNull(uint64(row)) { + return nil, nil + } + val := vector.GetFixedAtNoTypeCheck[types.Decimal128](v, row) + return newTypedByteJson(bytejson.TpCodeDecimal, string(val.Format(fromType.Scale))), nil case types.T_binary, types.T_varbinary, types.T_blob: - return bytejson.NewOpaqueDataEncoder(v.GetBytesAt(row)) + if v.IsNull(uint64(row)) { + return nil, nil + } + return newTypedByteJson(bytejson.TpCodeOpaque, string(v.GetBytesAt(row))), nil + case types.T_decimal256: + if v.IsNull(uint64(row)) { + return nil, nil + } + val := vector.GetFixedAtNoTypeCheck[types.Decimal256](v, row) + return newTypedByteJson(bytejson.TpCodeDecimal, string(val.Format(fromType.Scale))), nil case types.T_year: - return newJSONTypedStringEncoder( - bytejson.TpCodeString, - strconv.FormatInt( - int64(vector.GetFixedAtNoTypeCheck[int16](v, row)), - 10, - ), - ) - case types.T_bit: - width := fromType.Width - if width <= 0 { - width = 1 + if v.IsNull(uint64(row)) { + return nil, nil } - if width > 64 { - ctx := context.Background() - if proc != nil && proc.Ctx != nil { - ctx = proc.Ctx - } - return nil, moerr.NewInvalidInputf(ctx, "cannot cast BIT(%d) to json", width) + val := vector.GetFixedAtNoTypeCheck[int16](v, row) + return strconv.FormatInt(int64(val), 10), nil + case types.T_bit: + if v.IsNull(uint64(row)) { + return nil, nil } - value := vector.GetFixedAtNoTypeCheck[uint64](v, row) - if width < 64 { - value &= uint64(1)<= 2 && bj[0] == '"' { + key = string(bj[1 : len(bj)-1]) + } else { + key = fmt.Sprint(v) + } + default: + if bj, err := v.MarshalJSON(); err == nil { + key = string(bj) + } else { + key = fmt.Sprint(v) + } + } + case nil: + return moerr.NewInvalidInputf(proc.Ctx, "JSON documents may not contain NULL member names") + default: + key = fmt.Sprint(v) } - keyOffsets[entryIdx] = [2]int{start, keyOutput.Len()} - value, err := arrayOp.buildValueEncoder(proc, params[i+1], j) + elem, err := arrayOp.convertToAny(proc, params[i+1], j) if err != nil { return err } - entries[entryIdx].Value = value + obj[key] = elem } - keyBytes := keyOutput.Bytes() - for idx, offsets := range keyOffsets { - entries[idx].Key = keyBytes[offsets[0]:offsets[1]] - } - encoder, err := bytejson.NewObjectDataEncoder(entries) + + bj, err := bytejson.CreateByteJSON(obj) if err != nil { return err } - if err := rs.AppendByteJsonEncoded(encoder); err != nil { + dt, err := bj.Marshal() + if err != nil { return err } - } - return nil -} - -func writeJSONObjectKey( - w formatBuffer, - proc *process.Process, - v *vector.Vector, - row int, -) error { - fromType := v.GetType() - var numeric [64]byte - write := func(value []byte) error { - _, err := w.Write(value) - return err - } - switch fromType.Oid { - case types.T_bool: - return write(strconv.AppendBool( - numeric[:0], - vector.GetFixedAtNoTypeCheck[bool](v, row), - )) - case types.T_int8: - return write(strconv.AppendInt(numeric[:0], int64(vector.GetFixedAtNoTypeCheck[int8](v, row)), 10)) - case types.T_int16: - return write(strconv.AppendInt(numeric[:0], int64(vector.GetFixedAtNoTypeCheck[int16](v, row)), 10)) - case types.T_int32: - return write(strconv.AppendInt(numeric[:0], int64(vector.GetFixedAtNoTypeCheck[int32](v, row)), 10)) - case types.T_int64: - return write(strconv.AppendInt(numeric[:0], vector.GetFixedAtNoTypeCheck[int64](v, row), 10)) - case types.T_uint8: - return write(strconv.AppendUint(numeric[:0], uint64(vector.GetFixedAtNoTypeCheck[uint8](v, row)), 10)) - case types.T_uint16: - return write(strconv.AppendUint(numeric[:0], uint64(vector.GetFixedAtNoTypeCheck[uint16](v, row)), 10)) - case types.T_uint32: - return write(strconv.AppendUint(numeric[:0], uint64(vector.GetFixedAtNoTypeCheck[uint32](v, row)), 10)) - case types.T_uint64: - return write(strconv.AppendUint(numeric[:0], vector.GetFixedAtNoTypeCheck[uint64](v, row), 10)) - case types.T_float32: - return write(strconv.AppendFloat(numeric[:0], float64(vector.GetFixedAtNoTypeCheck[float32](v, row)), 'g', -1, 64)) - case types.T_float64: - return write(strconv.AppendFloat(numeric[:0], vector.GetFixedAtNoTypeCheck[float64](v, row), 'g', -1, 64)) - case types.T_char, types.T_varchar, types.T_text, types.T_geometry: - return write(v.GetBytesAt(row)) - case types.T_json: - return bytejson.WriteJSONObjectKeyText(w, types.DecodeJson(v.GetBytesAt(row))) - case types.T_date: - _, err := w.WriteString(vector.GetFixedAtNoTypeCheck[types.Date](v, row).String()) - return err - case types.T_time: - _, err := w.WriteString(vector.GetFixedAtNoTypeCheck[types.Time](v, row).String2(fromType.Scale)) - return err - case types.T_datetime: - _, err := w.WriteString(vector.GetFixedAtNoTypeCheck[types.Datetime](v, row).String2(fromType.Scale)) - return err - case types.T_timestamp: - _, err := w.WriteString(vector.GetFixedAtNoTypeCheck[types.Timestamp](v, row).String2(jsonSessionTimeZone(proc), fromType.Scale)) - return err - case types.T_decimal64: - _, err := w.WriteString(vector.GetFixedAtNoTypeCheck[types.Decimal64](v, row).Format(fromType.Scale)) - return err - case types.T_decimal128: - _, err := w.WriteString(vector.GetFixedAtNoTypeCheck[types.Decimal128](v, row).Format(fromType.Scale)) - return err - case types.T_decimal256: - _, err := w.WriteString(vector.GetFixedAtNoTypeCheck[types.Decimal256](v, row).Format(fromType.Scale)) - return err - case types.T_binary, types.T_varbinary, types.T_blob: - return bytejson.WriteJSONBase64Text(w, v.GetBytesAt(row)) - case types.T_year: - return write(strconv.AppendInt(numeric[:0], int64(vector.GetFixedAtNoTypeCheck[int16](v, row)), 10)) - case types.T_bit: - width := fromType.Width - if width <= 0 { - width = 1 - } - if width > 64 { - ctx := context.Background() - if proc != nil && proc.Ctx != nil { - ctx = proc.Ctx - } - return moerr.NewInvalidInputf(ctx, "cannot cast BIT(%d) to json", width) - } - value := vector.GetFixedAtNoTypeCheck[uint64](v, row) - if width < 64 { - value &= uint64(1)< 0 { - if err := w.WriteByte(' '); err != nil { - return err - } - } - value := strconv.AppendFloat(numeric[:0], valueAt(idx), 'g', -1, 64) - if _, err := w.Write(value); err != nil { + if err := rs.AppendBytes(dt, false); err != nil { return err } } - return w.WriteByte(']') + return nil } type opBuiltInJsonType struct{} @@ -2572,25 +2419,19 @@ func jsonKeysRoot(ivecs []*vector.Vector, result vector.FunctionResultWrapper, p if selectList != nil && selectList.IgnoreAllRow() { for i := 0; i < length; i++ { - if err := rs.AppendMustNullForBytesResult(); err != nil { - return err - } + rs.AppendMustNullForBytesResult() } return nil } for i := uint64(0); i < uint64(length); i++ { if selectList.Contains(i) { - if err := rs.AppendMustNullForBytesResult(); err != nil { - return err - } + rs.AppendMustNullForBytesResult() continue } v, null := p1.GetStrValue(i) if null { - if err := rs.AppendMustNullForBytesResult(); err != nil { - return err - } + rs.AppendMustNullForBytesResult() continue } var bj bytejson.ByteJson @@ -2603,14 +2444,14 @@ func jsonKeysRoot(ivecs []*vector.Vector, result vector.FunctionResultWrapper, p if err != nil { return moerr.NewInvalidArg(proc.Ctx, "json_keys", "invalid JSON document") } - appended, err := appendJsonKeysArray(rs, bj) - if err != nil { - return err - } - if !appended { + keysArray, err := buildJsonKeysArray(bj) + if err != nil || keysArray.IsNull() { rs.AppendMustNullForBytesResult() continue } + if err := rs.AppendByteJson(keysArray, false); err != nil { + return err + } } return nil } @@ -2660,33 +2501,28 @@ func jsonKeysWithPath(ivecs []*vector.Vector, result vector.FunctionResultWrappe rs.AppendMustNullForBytesResult() continue } - appended, err := appendJsonKeysArray(rs, val) - if err != nil { - return err - } - if !appended { + keysArray, err := buildJsonKeysArray(val) + if err != nil || keysArray.IsNull() { rs.AppendMustNullForBytesResult() continue } + if err := rs.AppendByteJson(keysArray, false); err != nil { + return err + } } return nil } -func appendJsonKeysArray( - rs *vector.FunctionResult[types.Varlena], - bj bytejson.ByteJson, -) (bool, error) { +func buildJsonKeysArray(bj bytejson.ByteJson) (bytejson.ByteJson, error) { if bj.Type != bytejson.TpCodeObject { - return false, nil + return bytejson.Null, nil } - encoder, err := bytejson.NewObjectKeysArrayEncoder(bj) - if err != nil { - return false, err - } - if err := rs.AppendByteJsonEncoded(encoder); err != nil { - return false, err + cnt := bj.GetElemCnt() + keys := make([]any, cnt) + for i := 0; i < cnt; i++ { + keys[i] = string(bj.GetObjectKey(i)) } - return true, nil + return bytejson.CreateByteJSON(keys) } // JSON_PRETTY @@ -2695,7 +2531,6 @@ func JsonPretty(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro result.UseOptFunctionParamFrame(1) rs := vector.MustFunctionResult[types.Varlena](result) p1 := vector.OptGetBytesParamFromWrapper(rs, 0, ivecs[0]) - var legacy bytes.Buffer if selectList != nil && selectList.IgnoreAllRow() { for i := 0; i < length; i++ { @@ -2724,17 +2559,25 @@ func JsonPretty(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro if err != nil { return moerr.NewInvalidArg(proc.Ctx, "json_pretty", "invalid JSON document") } - _, err = appendFormattedBytesForResult(result, rs, &legacy, func(w formatBuffer) (bool, error) { - return false, jsonPrettyPrintTo(w, bj, 0) - }) + out, err := jsonPrettyPrint(bj, 0) if err != nil { return err } + rs.AppendMustBytesValue(out) } return nil } -func jsonPrettyPrintTo(w formatBuffer, bj bytejson.ByteJson, depth int) error { +func jsonPrettyPrint(bj bytejson.ByteJson, depth int) ([]byte, error) { + var buf bytes.Buffer + err := jsonPrettyPrintTo(&buf, bj, depth) + if err != nil { + return nil, err + } + return buf.Bytes(), nil +} + +func jsonPrettyPrintTo(w *bytes.Buffer, bj bytejson.ByteJson, depth int) error { switch bj.Type { case bytejson.TpCodeObject: return prettyPrintObject(w, bj, depth) @@ -2745,94 +2588,69 @@ func jsonPrettyPrintTo(w formatBuffer, bj bytejson.ByteJson, depth int) error { } } -func prettyPrintObject(w formatBuffer, bj bytejson.ByteJson, depth int) error { +func prettyPrintObject(w *bytes.Buffer, bj bytejson.ByteJson, depth int) error { cnt := bj.GetElemCnt() if cnt == 0 { - _, err := w.WriteString("{}") - return err - } - if _, err := w.WriteString("{\n"); err != nil { - return err + w.WriteString("{}") + return nil } + indent := strings.Repeat(" ", depth+1) + w.WriteString("{\n") for i := 0; i < cnt; i++ { key := bj.GetObjectKey(i) - if err := writePrettyIndent(w, depth+1); err != nil { - return err - } - if err := bytejson.WriteJSONString(w, key); err != nil { - return err - } - if _, err := w.WriteString(": "); err != nil { - return err - } + // Escape key the same way JSON_QUOTE would. + keyJSON, _ := json.Marshal(string(key)) + w.WriteString(indent) + w.Write(keyJSON) + w.WriteString(": ") val := bj.GetObjectVal(i) if err := jsonPrettyPrintTo(w, val, depth+1); err != nil { return err } if i < cnt-1 { - if err := w.WriteByte(','); err != nil { - return err - } - } - if err := w.WriteByte('\n'); err != nil { - return err + w.WriteString(",") } + w.WriteString("\n") } - if err := writePrettyIndent(w, depth); err != nil { - return err - } - return w.WriteByte('}') + w.WriteString(strings.Repeat(" ", depth)) + w.WriteString("}") + return nil } -func prettyPrintArray(w formatBuffer, bj bytejson.ByteJson, depth int) error { +func prettyPrintArray(w *bytes.Buffer, bj bytejson.ByteJson, depth int) error { cnt := bj.GetElemCnt() if cnt == 0 { - _, err := w.WriteString("[]") - return err - } - if _, err := w.WriteString("[\n"); err != nil { - return err + w.WriteString("[]") + return nil } + indent := strings.Repeat(" ", depth+1) + w.WriteString("[\n") for i := 0; i < cnt; i++ { - if err := writePrettyIndent(w, depth+1); err != nil { - return err - } + w.WriteString(indent) elem := bj.GetArrayElem(i) if err := jsonPrettyPrintTo(w, elem, depth+1); err != nil { return err } if i < cnt-1 { - if err := w.WriteByte(','); err != nil { - return err - } + w.WriteString(",") } - if err := w.WriteByte('\n'); err != nil { - return err - } - } - if err := writePrettyIndent(w, depth); err != nil { - return err + w.WriteString("\n") } - return w.WriteByte(']') + w.WriteString(strings.Repeat(" ", depth)) + w.WriteString("]") + return nil } -func writePrettyIndent(w formatBuffer, depth int) error { - const spaces = " " - remaining := depth * 2 - for remaining > 0 { - length := min(remaining, len(spaces)) - if _, err := w.WriteString(spaces[:length]); err != nil { - return err - } - remaining -= length +func prettyPrintScalar(w *bytes.Buffer, bj bytejson.ByteJson) error { + // Use MarshalJSON to get properly formatted/escaped scalar value. + text, err := bj.MarshalJSON() + if err != nil { + return err } + w.Write(text) return nil } -func prettyPrintScalar(w formatBuffer, bj bytejson.ByteJson) error { - return bytejson.WriteJSONText(w, bj) -} - // JSON_SCHEMA_VALID func JsonSchemaValid(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { result.UseOptFunctionParamFrame(2) diff --git a/pkg/sql/plan/function/func_builtin_json_row_test.go b/pkg/sql/plan/function/func_builtin_json_row_test.go deleted file mode 100644 index 0c1f53397095e..0000000000000 --- a/pkg/sql/plan/function/func_builtin_json_row_test.go +++ /dev/null @@ -1,179 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package function - -import ( - "strconv" - "testing" - - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/matrixorigin/matrixone/pkg/container/vector" - "github.com/matrixorigin/matrixone/pkg/testutil" - "github.com/stretchr/testify/require" -) - -func TestJSONRowStreamsRowsAndResetsAfterError(t *testing.T) { - mp := mpool.MustNewZero() - defer mpool.DeleteMPool(mp) - proc := testutil.NewProcessWithMPool(t, "", mp) - - ints := vector.NewVec(types.T_int64.ToType()) - require.NoError(t, vector.AppendFixedList( - ints, - []int64{1, 2, 3}, - []bool{false, true, false}, - mp, - )) - defer ints.Free(mp) - strings := vector.NewVec(types.T_varchar.ToType()) - require.NoError(t, vector.AppendBytesList( - strings, - [][]byte{[]byte("a"), []byte("b"), []byte("c")}, - nil, - mp, - )) - defer strings.Free(mp) - bools := vector.NewVec(types.T_bool.ToType()) - require.NoError(t, vector.AppendFixedList( - bools, - []bool{true, false, true}, - nil, - mp, - )) - defer bools.Free(mp) - uints := vector.NewVec(types.T_uint64.ToType()) - require.NoError(t, vector.AppendFixedList( - uints, - []uint64{^uint64(0), 2, 3}, - nil, - mp, - )) - defer uints.Free(mp) - - result := vector.NewFunctionResultWrapper(types.T_varchar.ToType(), mp) - defer result.Free() - op := newOpBuiltInJsonRow() - require.NoError(t, result.PreExtendAndReset(3)) - require.NoError(t, op.jsonRow( - []*vector.Vector{ints, strings, bools, uints}, - result, - proc, - 3, - &FunctionSelectList{AnyNull: true, SelectList: []bool{true, false, true}}, - )) - out := result.GetResultVector() - require.Equal( - t, - []byte(`[1,"a",true,18446744073709551615]`), - out.GetBytesAt(0), - ) - require.True(t, out.IsNull(1)) - require.Equal(t, []byte(`[3,"c",true,3]`), out.GetBytesAt(2)) - - binary := vector.NewVec(types.T_binary.ToType()) - require.NoError(t, vector.AppendBytesList(binary, [][]byte{[]byte("x")}, nil, mp)) - defer binary.Free(mp) - require.NoError(t, result.PreExtendAndReset(1)) - require.Error(t, op.jsonRow( - []*vector.Vector{binary}, - result, - proc, - 1, - nil, - )) - for _, column := range op.columns { - require.Nil(t, column) - } - require.Zero(t, op.enc.w.Len()) - - require.NoError(t, result.PreExtendAndReset(3)) - require.NoError(t, op.jsonRow( - []*vector.Vector{ints, strings}, - result, - proc, - 3, - nil, - )) - out = result.GetResultVector() - require.Equal(t, []byte(`[1,"a"]`), out.GetBytesAt(0)) - require.Equal(t, []byte(`[null,"b"]`), out.GetBytesAt(1)) - require.Equal(t, []byte(`[3,"c"]`), out.GetBytesAt(2)) -} - -func newJSONRowBenchmarkParameters( - b *testing.B, - mp *mpool.MPool, - rows int, -) []*vector.Vector { - b.Helper() - ints := vector.NewVec(types.T_int64.ToType()) - intValues := make([]int64, rows) - for i := range intValues { - intValues[i] = int64(i) - } - require.NoError(b, vector.AppendFixedList(ints, intValues, nil, mp)) - - strings := vector.NewVec(types.T_varchar.ToType()) - stringValues := make([][]byte, rows) - for i := range stringValues { - stringValues[i] = []byte("value-" + strconv.Itoa(i%100)) - } - require.NoError(b, vector.AppendBytesList(strings, stringValues, nil, mp)) - return []*vector.Vector{ints, strings} -} - -func BenchmarkJSONRowFreshOperator8192(b *testing.B) { - mp := mpool.MustNewZero() - defer mpool.DeleteMPool(mp) - proc := testutil.NewProcessWithMPool(b, "", mp) - params := newJSONRowBenchmarkParameters(b, mp, 8192) - defer params[0].Free(mp) - defer params[1].Free(mp) - result := vector.NewFunctionResultWrapper(types.T_varchar.ToType(), mp) - defer result.Free() - - b.ReportAllocs() - b.ResetTimer() - for i := 0; i < b.N; i++ { - require.NoError(b, result.PreExtendAndReset(8192)) - require.NoError(b, newOpBuiltInJsonRow().jsonRow( - params, - result, - proc, - 8192, - nil, - )) - } -} - -func BenchmarkJSONRowReusedOperator8192(b *testing.B) { - mp := mpool.MustNewZero() - defer mpool.DeleteMPool(mp) - proc := testutil.NewProcessWithMPool(b, "", mp) - params := newJSONRowBenchmarkParameters(b, mp, 8192) - defer params[0].Free(mp) - defer params[1].Free(mp) - result := vector.NewFunctionResultWrapper(types.T_varchar.ToType(), mp) - defer result.Free() - op := newOpBuiltInJsonRow() - - b.ReportAllocs() - b.ResetTimer() - for i := 0; i < b.N; i++ { - require.NoError(b, result.PreExtendAndReset(8192)) - require.NoError(b, op.jsonRow(params, result, proc, 8192, nil)) - } -} diff --git a/pkg/sql/plan/function/func_cast.go b/pkg/sql/plan/function/func_cast.go index ca8bfe3a70853..29616b6810cc4 100644 --- a/pkg/sql/plan/function/func_cast.go +++ b/pkg/sql/plan/function/func_cast.go @@ -3613,8 +3613,7 @@ func signedToStr[T constraints.Integer]( return err } } else { - var scratch [20]byte - result := strconv.AppendInt(scratch[:0], int64(v), 10) + result := []byte(strconv.FormatInt(int64(v), 10)) if toType.Oid == types.T_binary || toType.Oid == types.T_varbinary { if int32(len(result)) > toType.Width { return moerr.NewDataTruncatedNoCtx("Signed", " truncated for binary/varbinary") @@ -3663,8 +3662,7 @@ func unsignedToStr[T constraints.Unsigned]( return err } } else { - var scratch [20]byte - result := strconv.AppendUint(scratch[:0], uint64(v), 10) + result := []byte(strconv.FormatUint(uint64(v), 10)) if toType.Oid == types.T_binary || toType.Oid == types.T_varbinary { if int32(len(result)) > toType.Width { return moerr.NewDataTruncatedNoCtx("Unsigned", "truncated for binary/varbinary") @@ -7477,18 +7475,11 @@ func arrayToArray[I types.ArrayElement, O types.ArrayElement]( // upcast the source element type to []float32, then narrow to the // target element type (int8 rounds+clamps; bf16/f16 round-to-even). // This replaces moarray.Cast[I,O], which only handled float pairs. - values := types.BytesToArray[I](v) - var outputElement O - elementSize := int(unsafe.Sizeof(outputElement)) - if len(values) > math.MaxInt/elementSize { - return moerr.NewInvalidInputNoCtx("array cast result is too large") - } - if err := to.AppendBytesWithFill(len(values)*elementSize, func(dst []byte) { - output := util.UnsafeSliceCast[O](dst) - for idx, value := range values { - output[idx] = float32ToArrayElement[O](arrayElementToFloat32(value)) - } - }); err != nil { + _v := types.BytesToArray[I](v) + f32 := types.ToFloat32Array[I](_v) + out := types.FromFloat32Array[O](f32) + bytes := types.ArrayToBytes[O](out) + if err := to.AppendBytes(bytes, false); err != nil { return err } } @@ -7497,47 +7488,6 @@ func arrayToArray[I types.ArrayElement, O types.ArrayElement]( return nil } -func arrayElementToFloat32[T types.ArrayElement](value T) float32 { - switch typed := any(value).(type) { - case float32: - return typed - case float64: - return float32(typed) - case types.BF16: - return typed.ToFloat32() - case types.Float16: - return typed.ToFloat32() - case int8: - return float32(typed) - case uint8: - return float32(typed) - default: - panic(moerr.NewInternalErrorNoCtx("unsupported array element type")) - } -} - -func float32ToArrayElement[T types.ArrayElement](value float32) T { - var output any - var zero T - switch any(zero).(type) { - case float32: - output = value - case float64: - output = float64(value) - case types.BF16: - output = types.BF16FromFloat32(value) - case types.Float16: - output = types.Float16FromFloat32(value) - case int8: - output = types.Float32ToInt8(value) - case uint8: - output = types.Float32ToUint8(value) - default: - panic(moerr.NewInternalErrorNoCtx("unsupported array element type")) - } - return output.(T) -} - func uuidToStr( ctx context.Context, from vector.FunctionParameterWrapper[types.Uuid], diff --git a/pkg/sql/plan/function/func_compare.go b/pkg/sql/plan/function/func_compare.go index bc4cd2d05f3dd..b22d59c60d6bd 100644 --- a/pkg/sql/plan/function/func_compare.go +++ b/pkg/sql/plan/function/func_compare.go @@ -104,14 +104,8 @@ func opBinaryFixedFixedToFixedNullSafe[T types.FixedSizeTExceptStrType]( ) error { result.UseOptFunctionParamFrame(2) rs := vector.MustFunctionResult[bool](result) - p1, err := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) - if err != nil { - return err - } - p2, err := vector.OptGetParamFromWrapper[T](rs, 1, parameters[1]) - if err != nil { - return err - } + p1 := vector.OptGetParamFromWrapper[T](rs, 0, parameters[0]) + p2 := vector.OptGetParamFromWrapper[T](rs, 1, parameters[1]) rsVec := rs.GetResultVector() rss := vector.MustFixedColNoTypeCheck[bool](rsVec) @@ -1581,7 +1575,7 @@ func operatorOpInt64Uint64Fn( func operatorOpStrFn( parameters []*vector.Vector, result vector.FunctionResultWrapper, _ *process.Process, length int, - fn func([]byte, []byte, []byte)) error { + fn func([]byte, []byte) ([]byte, error)) error { p1 := vector.GenerateFunctionStrParameter(parameters[0]) p2 := vector.GenerateFunctionStrParameter(parameters[1]) rs := vector.MustFunctionResult[types.Varlena](result) @@ -1593,14 +1587,11 @@ func operatorOpStrFn( return err } } else { - if len(v1) != len(v2) { - return moerr.NewInternalErrorNoCtx( - "Binary operands of bitwise operators must be of equal length", - ) + rv, err := fn(v1, v2) + if err != nil { + return err } - if err := rs.AppendBytesWithFill(len(v1), func(dst []byte) { - fn(dst, v1, v2) - }); err != nil { + if err = rs.AppendBytes(rv, false); err != nil { return err } } @@ -1625,10 +1616,15 @@ func operatorOpBitAndInt64Uint64Fn(parameters []*vector.Vector, result vector.Fu } func operatorOpBitAndStrFn(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - return operatorOpStrFn(parameters, result, proc, length, func(dst, left, right []byte) { - for idx := range dst { - dst[idx] = left[idx] & right[idx] + return operatorOpStrFn(parameters, result, proc, length, func(i []byte, i2 []byte) ([]byte, error) { + if len(i) != len(i2) { + return nil, moerr.NewInternalErrorNoCtx("Binary operands of bitwise operators must be of equal length") } + rv := make([]byte, len(i)) + for j := range rv { + rv[j] = i[j] & i2[j] + } + return rv, nil }) } @@ -1649,10 +1645,15 @@ func operatorOpBitXorInt64Uint64Fn(parameters []*vector.Vector, result vector.Fu } func operatorOpBitXorStrFn(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - return operatorOpStrFn(parameters, result, proc, length, func(dst, left, right []byte) { - for idx := range dst { - dst[idx] = left[idx] ^ right[idx] + return operatorOpStrFn(parameters, result, proc, length, func(i []byte, i2 []byte) ([]byte, error) { + if len(i) != len(i2) { + return nil, moerr.NewInternalErrorNoCtx("Binary operands of bitwise operators must be of equal length") + } + rv := make([]byte, len(i)) + for j := range rv { + rv[j] = i[j] ^ i2[j] } + return rv, nil }) } @@ -1673,10 +1674,15 @@ func operatorOpBitOrInt64Uint64Fn(parameters []*vector.Vector, result vector.Fun } func operatorOpBitOrStrFn(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - return operatorOpStrFn(parameters, result, proc, length, func(dst, left, right []byte) { - for idx := range dst { - dst[idx] = left[idx] | right[idx] + return operatorOpStrFn(parameters, result, proc, length, func(i []byte, i2 []byte) ([]byte, error) { + if len(i) != len(i2) { + return nil, moerr.NewInternalErrorNoCtx("Binary operands of bitwise operators must be of equal length") + } + rv := make([]byte, len(i)) + for j := range rv { + rv[j] = i[j] | i2[j] } + return rv, nil }) } diff --git a/pkg/sql/plan/function/func_prefix.go b/pkg/sql/plan/function/func_prefix.go index 761415482b083..7c7429e75f906 100644 --- a/pkg/sql/plan/function/func_prefix.go +++ b/pkg/sql/plan/function/func_prefix.go @@ -16,8 +16,6 @@ package function import ( "bytes" - "encoding/binary" - "math" "sort" "github.com/matrixorigin/matrixone/pkg/common/mpool" @@ -86,10 +84,8 @@ func PrefixInRange(parameters []*vector.Vector, result vector.FunctionResultWrap } type implPrefixIn struct { - ready bool - vals [][]byte - scratch []byte - scratchCount int + ready bool + vals [][]byte } func newImplPrefixIn() *implPrefixIn { @@ -97,6 +93,7 @@ func newImplPrefixIn() *implPrefixIn { } func (op *implPrefixIn) init(rvec *vector.Vector, mp *mpool.MPool) error { + op.ready = true op.vals = make([][]byte, rvec.Length()) vlen := 0 @@ -126,124 +123,12 @@ func (op *implPrefixIn) init(rvec *vector.Vector, mp *mpool.MPool) error { } } op.vals = op.vals[:vlen] - op.ready = true return nil } -const prefixScratchEntrySize = 8 - -type prefixScratchEntries struct { - data []byte - count int -} - -func (e prefixScratchEntries) Len() int { - return e.count -} - -func (e prefixScratchEntries) Less(left, right int) bool { - return bytes.Compare(e.value(left), e.value(right)) < 0 -} - -func (e prefixScratchEntries) Swap(left, right int) { - leftEntry := e.data[left*prefixScratchEntrySize : (left+1)*prefixScratchEntrySize] - rightEntry := e.data[right*prefixScratchEntrySize : (right+1)*prefixScratchEntrySize] - var saved [prefixScratchEntrySize]byte - copy(saved[:], leftEntry) - copy(leftEntry, rightEntry) - copy(rightEntry, saved[:]) -} - -func (e prefixScratchEntries) value(index int) []byte { - entry := e.data[index*prefixScratchEntrySize:] - offset := binary.LittleEndian.Uint32(entry) - length := binary.LittleEndian.Uint32(entry[4:]) - return e.data[int(offset):int(offset+length)] -} - -func (op *implPrefixIn) initAccounted( - rvec *vector.Vector, - result vector.FunctionResultWrapper, -) error { - rowCount := rvec.Length() - if rowCount < 0 || rowCount > math.MaxInt/prefixScratchEntrySize { - return mpool.ErrAllocationAccountInvalid - } - total := rowCount * prefixScratchEntrySize - for row := 0; row < rowCount; row++ { - valueSize := len(rvec.GetBytesAt(row)) - if valueSize > math.MaxInt-total { - return mpool.ErrAllocationAccountInvalid - } - total += valueSize - } - if uint64(total) > math.MaxUint32 { - return mpool.ErrAllocationAccountInvalid - } - scratch, selected, err := result.ResizeFunctionScratch(total) - if err != nil { - return err - } - if !selected { - return mpool.ErrAllocationAccountInvalid - } - entries := prefixScratchEntries{data: scratch, count: rowCount} - payloadOffset := rowCount * prefixScratchEntrySize - for row := 0; row < rowCount; row++ { - value := rvec.GetBytesAt(row) - entry := scratch[row*prefixScratchEntrySize:] - binary.LittleEndian.PutUint32(entry, uint32(payloadOffset)) - binary.LittleEndian.PutUint32(entry[4:], uint32(len(value))) - payloadOffset += copy(scratch[payloadOffset:], value) - } - if !rvec.GetSorted() { - sort.Sort(entries) - } - compactCount := 0 - for row := 0; row < rowCount; row++ { - value := entries.value(row) - if compactCount != 0 && bytes.HasPrefix(value, entries.value(compactCount-1)) { - continue - } - if compactCount != row { - copy( - scratch[compactCount*prefixScratchEntrySize:], - scratch[row*prefixScratchEntrySize:(row+1)*prefixScratchEntrySize], - ) - } - compactCount++ - } - op.scratch = scratch - op.scratchCount = compactCount - op.ready = true - return nil -} - -func (op *implPrefixIn) valueCount() int { - if op.scratch != nil { - return op.scratchCount - } - return len(op.vals) -} - -func (op *implPrefixIn) valueAt(index int) []byte { - if op.scratch != nil { - return (prefixScratchEntries{ - data: op.scratch, - count: op.scratchCount, - }).value(index) - } - return op.vals[index] -} - func (op *implPrefixIn) doPrefixIn(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { if !op.ready { - var err error - if result.HasFunctionScratch() { - err = op.initAccounted(parameters[1], result) - } else { - err = op.init(parameters[1], proc.Mp()) - } + err := op.init(parameters[1], proc.Mp()) if err != nil { return err } @@ -251,7 +136,7 @@ func (op *implPrefixIn) doPrefixIn(parameters []*vector.Vector, result vector.Fu lvec := parameters[0] res := vector.MustFixedColWithTypeCheck[bool](result.GetResultVector()) - if op.valueCount() == 0 { + if len(op.vals) == 0 { for i := range length { res[i] = false } @@ -262,9 +147,9 @@ func (op *implPrefixIn) doPrefixIn(parameters []*vector.Vector, result vector.Fu lvecHasNull := lvec.HasNull() if lvec.GetSorted() && !lvecHasNull { - rval := op.valueAt(0) + rval := op.vals[0] rpos := 0 - rlen := op.valueCount() + rlen := len(op.vals) for i := range length { lval := lcol[i].GetByteSlice(larea) @@ -277,7 +162,7 @@ func (op *implPrefixIn) doPrefixIn(parameters []*vector.Vector, result vector.Fu return nil } - rval = op.valueAt(rpos) + rval = op.vals[rpos] } res[i] = bytes.HasPrefix(lval, rval) @@ -292,21 +177,21 @@ func (op *implPrefixIn) doPrefixIn(parameters []*vector.Vector, result vector.Fu rNulls.Add(i) } else { lval := lcol[i].GetByteSlice(larea) - rpos, _ := sort.Find(op.valueCount(), func(j int) int { - return types.PrefixCompare(lval, op.valueAt(j)) + rpos, _ := sort.Find(len(op.vals), func(j int) int { + return types.PrefixCompare(lval, op.vals[j]) }) - res[i] = rpos < op.valueCount() && bytes.HasPrefix(lval, op.valueAt(rpos)) + res[i] = rpos < len(op.vals) && bytes.HasPrefix(lval, op.vals[rpos]) } } } else { for i := range length { lval := lcol[i].GetByteSlice(larea) - rpos, _ := sort.Find(op.valueCount(), func(j int) int { - return types.PrefixCompare(lval, op.valueAt(j)) + rpos, _ := sort.Find(len(op.vals), func(j int) int { + return types.PrefixCompare(lval, op.vals[j]) }) - res[i] = rpos < op.valueCount() && bytes.HasPrefix(lval, op.valueAt(rpos)) + res[i] = rpos < len(op.vals) && bytes.HasPrefix(lval, op.vals[rpos]) } } } diff --git a/pkg/sql/plan/function/func_string_complex_test.go b/pkg/sql/plan/function/func_string_complex_test.go index 3712a3f96a825..88877aecd6f00 100644 --- a/pkg/sql/plan/function/func_string_complex_test.go +++ b/pkg/sql/plan/function/func_string_complex_test.go @@ -838,8 +838,8 @@ func Test_EncodeCharBytes(t *testing.T) { {math.MinInt64, []byte{0x00}}, } for _, c := range cases { - got := appendCharBytes(nil, c.input) - require.Equal(t, c.want, got, "appendCharBytes(%d)", c.input) + got := encodeCharBytes(c.input) + require.Equal(t, c.want, got, "encodeCharBytes(%d)", c.input) } } diff --git a/pkg/sql/plan/function/func_unary.go b/pkg/sql/plan/function/func_unary.go index e15486326426c..f9c59f806cdbf 100644 --- a/pkg/sql/plan/function/func_unary.go +++ b/pkg/sql/plan/function/func_unary.go @@ -50,7 +50,6 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/system" "github.com/matrixorigin/matrixone/pkg/common/util" - "github.com/matrixorigin/matrixone/pkg/container/bytejson" "github.com/matrixorigin/matrixone/pkg/container/nulls" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" @@ -217,146 +216,120 @@ func AbsArray[T types.RealNumbers](ivecs []*vector.Vector, result vector.Functio }, selectList) } +var ( + arrayF32Pool = sync.Pool{ + New: func() interface{} { + s := make([]float32, 128) + return &s + }, + } + + arrayF64Pool = sync.Pool{ + New: func() interface{} { + s := make([]float64, 128) + return &s + }, + } +) + func NormalizeL2Array[T types.ArrayElement](parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { source := vector.GenerateFunctionStrParameter(parameters[0]) rs := vector.MustFunctionResult[types.Varlena](result) rowCount := uint64(length) + var inArrayF32 []float32 + var outArrayF32Ptr *[]float32 + var outArrayF32 []float32 + + var inArrayF64 []float64 + var outArrayF64Ptr *[]float64 + var outArrayF64 []float64 + + var data []byte + var null bool + for i := uint64(0); i < rowCount; i++ { - data, null := source.GetStrValue(i) + data, null = source.GetStrValue(i) if null { - if err := rs.AppendMustNullForBytesResult(); err != nil { - return err - } + _ = rs.AppendMustNullForBytesResult() continue } switch t := parameters[0].GetType().Oid; t { case types.T_array_float32: - if err := appendNormalizedRealArray[float32](rs, data); err != nil { - return err + inArrayF32 = types.BytesToArray[float32](data) + + outArrayF32Ptr = arrayF32Pool.Get().(*[]float32) + outArrayF32 = *outArrayF32Ptr + + if cap(outArrayF32) < len(inArrayF32) { + outArrayF32 = make([]float32, len(inArrayF32)) + } else { + outArrayF32 = outArrayF32[:len(inArrayF32)] } + _ = moarray.NormalizeL2(inArrayF32, outArrayF32) + _ = rs.AppendBytes(types.ArrayToBytes[float32](outArrayF32), false) + + *outArrayF32Ptr = outArrayF32 + arrayF32Pool.Put(outArrayF32Ptr) case types.T_array_float64: - if err := appendNormalizedRealArray[float64](rs, data); err != nil { - return err + inArrayF64 = types.BytesToArray[float64](data) + + outArrayF64Ptr = arrayF64Pool.Get().(*[]float64) + outArrayF64 = *outArrayF64Ptr + + if cap(outArrayF64) < len(inArrayF64) { + outArrayF64 = make([]float64, len(inArrayF64)) + } else { + outArrayF64 = outArrayF64[:len(inArrayF64)] } + _ = moarray.NormalizeL2(inArrayF64, outArrayF64) + _ = rs.AppendBytes(types.ArrayToBytes[float64](outArrayF64), false) + + *outArrayF64Ptr = outArrayF64 + arrayF64Pool.Put(outArrayF64Ptr) case types.T_array_bf16: - if err := appendNormalizedNarrowArray( - rs, - data, - types.BF16.ToFloat32, - types.BF16FromFloat32, - ); err != nil { - return err - } + _ = appendNormalizedNarrowArray[types.BF16](rs, data) case types.T_array_float16: - if err := appendNormalizedNarrowArray( - rs, - data, - types.Float16.ToFloat32, - types.Float16FromFloat32, - ); err != nil { - return err - } + _ = appendNormalizedNarrowArray[types.Float16](rs, data) case types.T_array_int8: // A normalized vector is a unit vector, which cannot be represented in // an integer element type (components round to 0/±1 and the norm is no // longer 1), so int8/uint8 normalize_l2 widens the result to vecf32. // The overload's retType is T_array_float32 to match (see list_builtIn). - if err := appendNormalizedArrayAsFloat32( - rs, - data, - func(value int8) float32 { return float32(value) }, - ); err != nil { - return err - } + _ = appendNormalizedIntArrayAsFloat32[int8](rs, data) case types.T_array_uint8: - if err := appendNormalizedArrayAsFloat32( - rs, - data, - func(value uint8) float32 { return float32(value) }, - ); err != nil { - return err - } + _ = appendNormalizedIntArrayAsFloat32[uint8](rs, data) } + } return nil } -func appendNormalizedRealArray[T types.RealNumbers]( - rs *vector.FunctionResult[types.Varlena], - data []byte, -) error { - input := types.BytesToArray[T](data) - return rs.AppendBytesWithFill(len(data), func(dst []byte) { - _ = moarray.NormalizeL2(input, types.BytesToArray[T](dst)) - }) -} - -// appendNormalizedNarrowArray normalizes a bf16/f16 vector through float32 -// arithmetic while writing the narrowed values directly into the result. -func appendNormalizedNarrowArray[T types.ArrayElement]( - rs *vector.FunctionResult[types.Varlena], - data []byte, - toFloat32 func(T) float32, - fromFloat32 func(float32) T, -) error { - input := types.BytesToArray[T](data) - return rs.AppendBytesWithFill(len(data), func(dst []byte) { - output := types.BytesToArray[T](dst) - normalizeArrayInto(input, output, toFloat32, fromFloat32) - }) +// appendNormalizedNarrowArray normalizes a bf16/f16 vector by upcasting to +// float32, normalizing in float32, then narrowing back to T. bf16/f16 are +// floating-point so they can hold a (near-)unit vector; int8/uint8 cannot and +// use appendNormalizedIntArrayAsFloat32 instead. +func appendNormalizedNarrowArray[T types.ArrayElement](rs *vector.FunctionResult[types.Varlena], data []byte) error { + in := types.ToFloat32Array[T](types.BytesToArray[T](data)) + out := make([]float32, len(in)) + _ = moarray.NormalizeL2(in, out) + return rs.AppendBytes(types.ArrayToBytes[T](types.FromFloat32Array[T](out)), false) } -// appendNormalizedArrayAsFloat32 normalizes an integer-typed (int8/uint8) +// appendNormalizedIntArrayAsFloat32 normalizes an integer-typed (int8/uint8) // vector and writes the result as float32. A unit vector cannot be represented // in an integer element type — narrowing back would round components to 0/±1 so // the norm is no longer 1 (e.g. normalize_l2([0,1,2,3]::vecuint8) would become // [0,0,1,1], whose norm is √2). Widening the result to vecf32 keeps the unit-norm // contract; the int8/uint8 overloads declare retType T_array_float32 to match. -func appendNormalizedArrayAsFloat32[T types.ArrayElement]( - rs *vector.FunctionResult[types.Varlena], - data []byte, - toFloat32 func(T) float32, -) error { - input := types.BytesToArray[T](data) - if len(input) > int(^uint(0)>>1)/4 { - return moerr.NewInternalErrorNoCtx("normalized array result is too large") - } - return rs.AppendBytesWithFill(len(input)*4, func(dst []byte) { - output := types.BytesToArray[float32](dst) - normalizeArrayInto( - input, - output, - toFloat32, - func(value float32) float32 { return value }, - ) - }) -} - -func normalizeArrayInto[TIn, TOut types.ArrayElement]( - input []TIn, - output []TOut, - toFloat32 func(TIn) float32, - fromFloat32 func(float32) TOut, -) { - var sumSquares float64 - for _, value := range input { - converted := float64(toFloat32(value)) - sumSquares += converted * converted - } - norm := math.Sqrt(sumSquares) - if norm == 0 { - for idx, value := range input { - output[idx] = fromFloat32(toFloat32(value)) - } - return - } - for idx, value := range input { - output[idx] = fromFloat32(float32(float64(toFloat32(value)) / norm)) - } +func appendNormalizedIntArrayAsFloat32[T types.ArrayElement](rs *vector.FunctionResult[types.Varlena], data []byte) error { + in := types.ToFloat32Array[T](types.BytesToArray[T](data)) + out := make([]float32, len(in)) + _ = moarray.NormalizeL2(in, out) + return rs.AppendBytes(types.ArrayToBytes[float32](out), false) } func L1NormArray[T types.RealNumbers](ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { @@ -1011,25 +984,15 @@ func Empty(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *pr } func JsonQuote(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - source := vector.GenerateFunctionStrParameter(ivecs[0]) - rs := vector.MustFunctionResult[types.Varlena](result) - for row := uint64(0); row < uint64(length); row++ { - value, isNull := source.GetStrValue(row) - if isNull || selectList.Contains(row) { - if err := rs.AppendBytes(nil, true); err != nil { - return err - } - continue - } - encoder, err := bytejson.NewStringDataEncoder(value) + single := func(str string) ([]byte, error) { + bj, err := types.ParseStringToByteJson(strconv.Quote(str)) if err != nil { - return err - } - if err := rs.AppendByteJsonEncoded(encoder); err != nil { - return err + return nil, err } + return bj.Marshal() } - return nil + + return opUnaryStrToBytesWithErrorCheck(ivecs, result, proc, length, single, selectList) } func JsonUnquote(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { @@ -1058,11 +1021,6 @@ func JsonUnquote(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pr // Escapes single quotes by doubling them, backslashes, and control characters func QuoteString(str string) string { var result strings.Builder - writeQuotedString(&result, str) - return result.String() -} - -func writeQuotedString(result formatBuffer, str string) { result.WriteByte('\'') for _, r := range str { @@ -1095,28 +1053,15 @@ func writeQuotedString(result formatBuffer, str string) { } result.WriteByte('\'') + return result.String() } func Quote(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - source := vector.GenerateFunctionStrParameter(ivecs[0]) - rs := vector.MustFunctionResult[types.Varlena](result) - var legacy bytes.Buffer - for row := uint64(0); row < uint64(length); row++ { - value, isNull := source.GetStrValue(row) - if isNull || selectList.Contains(row) { - if err := rs.AppendBytes(nil, true); err != nil { - return err - } - continue - } - if _, err := appendFormattedBytesForResult(result, rs, &legacy, func(output formatBuffer) (bool, error) { - writeQuotedString(output, functionUtil.QuickBytesToStr(value)) - return false, nil - }); err != nil { - return err - } - } - return nil + return opUnaryBytesToBytes(ivecs, result, proc, length, func(v []byte) []byte { + str := functionUtil.QuickBytesToStr(v) + quoted := QuoteString(str) + return functionUtil.QuickStrToBytes(quoted) + }, selectList) } func StAsText(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { @@ -4727,13 +4672,13 @@ func Values(parameters []*vector.Vector, result vector.FunctionResultWrapper, pr toVec := result.GetResultVector() toVec.Reset(*toVec.GetType()) - return toVec.UnionBatch( - fromVec, - 0, - fromVec.Length(), - nil, - proc.GetMPool(), - ) + sels := make([]int64, fromVec.Length()) + for j := 0; j < len(sels); j++ { + sels[j] = int64(j) + } + + err := toVec.Union(fromVec, sels, proc.GetMPool()) + return err } func builtInNameConst(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { @@ -4984,24 +4929,11 @@ func HexFloat64(ivecs []*vector.Vector, result vector.FunctionResultWrapper, pro } func HexArray(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - source := vector.GenerateFunctionStrParameter(ivecs[0]) - rs := vector.MustFunctionResult[types.Varlena](result) - for row := uint64(0); row < uint64(length); row++ { - data, isNull := source.GetStrValue(row) - if isNull || selectList.Contains(row) { - if err := rs.AppendBytes(nil, true); err != nil { - return err - } - continue - } - encodedSize := hex.EncodedLen(len(data)) - if err := rs.AppendBytesWithFill(encodedSize, func(dst []byte) { - hex.Encode(dst, data) - }); err != nil { - return err - } - } - return nil + return opUnaryBytesToBytesWithErrorCheck(ivecs, result, proc, length, func(data []byte) ([]byte, error) { + buf := make([]byte, hex.EncodedLen(len(functionUtil.QuickBytesToStr(data)))) + hex.Encode(buf, data) + return buf, nil + }, selectList) } func hexEncodeString(xs []byte) string { @@ -5825,45 +5757,17 @@ func unhexToBytes(data []byte, null bool, rs *vector.FunctionResult[types.Varlen return rs.AppendMustNullForBytesResult() } - decodedSize := (len(data) + 1) / 2 - var decodeErr error - err := rs.AppendBytesWithBuilder(decodedSize, func(dst []byte) (int, error) { - written, err := decodeHexInto(dst, data) - decodeErr = err - return written, err - }) - if decodeErr != nil { - return rs.AppendMustNullForBytesResult() + // Add a '0' to the front, if the length is not the multiple of 2 + str := functionUtil.QuickBytesToStr(data) + if len(str)%2 != 0 { + str = "0" + str } - return err -} -func decodeHexInto(dst, src []byte) (int, error) { - written := 0 - if len(src)%2 != 0 { - value, ok := decodeHexNibble(src[0]) - if !ok { - return 0, hex.InvalidByteError(src[0]) - } - dst[0] = value - written = 1 - src = src[1:] - } - n, err := hex.Decode(dst[written:], src) - return written + n, err -} - -func decodeHexNibble(value byte) (byte, bool) { - switch { - case value >= '0' && value <= '9': - return value - '0', true - case value >= 'a' && value <= 'f': - return value - 'a' + 10, true - case value >= 'A' && value <= 'F': - return value - 'A' + 10, true - default: - return 0, false + bs, err := hex.DecodeString(str) + if err != nil { + return rs.AppendMustNullForBytesResult() } + return rs.AppendMustBytesValue(bs) } func Unhex(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { @@ -5882,24 +5786,10 @@ func Unhex(parameters []*vector.Vector, result vector.FunctionResultWrapper, pro } func Md5(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - source := vector.GenerateFunctionStrParameter(parameters[0]) - rs := vector.MustFunctionResult[types.Varlena](result) - for row := uint64(0); row < uint64(length); row++ { - data, isNull := source.GetStrValue(row) - if isNull || selectList.Contains(row) { - if err := rs.AppendBytes(nil, true); err != nil { - return err - } - continue - } + return opUnaryBytesToBytes(parameters, result, proc, length, func(data []byte) []byte { sum := md5.Sum(data) - if err := rs.AppendBytesWithFill(hex.EncodedLen(len(sum)), func(dst []byte) { - hex.Encode(dst, sum[:]) - }); err != nil { - return err - } - } - return nil + return []byte(hex.EncodeToString(sum[:])) + }, selectList) } @@ -5927,24 +5817,11 @@ func (content *crc32ExecContext) builtInCrc32(parameters []*vector.Vector, resul } func ToBase64(ivecs []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) (err error) { - source := vector.GenerateFunctionStrParameter(ivecs[0]) - rs := vector.MustFunctionResult[types.Varlena](result) - for row := uint64(0); row < uint64(length); row++ { - data, isNull := source.GetStrValue(row) - if isNull || selectList.Contains(row) { - if err = rs.AppendBytes(nil, true); err != nil { - return err - } - continue - } - encodedSize := base64.StdEncoding.EncodedLen(len(data)) - if err = rs.AppendBytesWithFill(encodedSize, func(dst []byte) { - base64.StdEncoding.Encode(dst, data) - }); err != nil { - return err - } - } - return nil + return opUnaryBytesToBytesWithErrorCheck(ivecs, result, proc, length, func(data []byte) ([]byte, error) { + buf := make([]byte, base64.StdEncoding.EncodedLen(len(functionUtil.QuickBytesToStr(data)))) + base64.StdEncoding.Encode(buf, data) + return buf, nil + }, selectList) } func FromBase64(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { @@ -5954,31 +5831,16 @@ func FromBase64(parameters []*vector.Vector, result vector.FunctionResultWrapper rowCount := uint64(length) for i := uint64(0); i < rowCount; i++ { data, null := source.GetStrValue(i) - if null || selectList.Contains(i) { - if err := rs.AppendMustNullForBytesResult(); err != nil { - return err - } - continue + if null { + return rs.AppendMustNullForBytesResult() } - var decodeErr error - err := rs.AppendBytesWithBuilder( - base64.StdEncoding.DecodedLen(len(data)), - func(dst []byte) (int, error) { - var written int - written, decodeErr = base64.StdEncoding.Decode(dst, data) - return written, decodeErr - }, - ) - if decodeErr != nil { - if err = rs.AppendMustNullForBytesResult(); err != nil { - return err - } - continue - } + buf := make([]byte, base64.StdEncoding.DecodedLen(len(functionUtil.QuickBytesToStr(data)))) + _, err := base64.StdEncoding.Decode(buf, data) if err != nil { - return err + return rs.AppendMustNullForBytesResult() } + _ = rs.AppendMustBytesValue(buf) } return nil @@ -6019,10 +5881,11 @@ func VecFromBase64[T types.ArrayElement](parameters []*vector.Vector, result vec } } + var buf []byte rowCount := uint64(length) for i := uint64(0); i < rowCount; i++ { data, null := source.GetStrValue(i) - if null || selectList.Contains(i) { + if null { if err := rs.AppendBytes(nil, true); err != nil { return err } @@ -6030,22 +5893,21 @@ func VecFromBase64[T types.ArrayElement](parameters []*vector.Vector, result vec } need := base64.StdEncoding.DecodedLen(len(data)) - if err := rs.AppendBytesWithBuilder(need, func(dst []byte) (int, error) { - written, err := base64.StdEncoding.Decode(dst, data) - if err != nil { - return 0, moerr.NewInternalErrorNoCtx( - "vec_from_base64: invalid base64 input", - ) - } - if written%elemSize != 0 { - return 0, moerr.NewInternalErrorNoCtxf( - "vec_from_base64: decoded length %d is not a multiple of %d bytes", - written, - elemSize, - ) - } - return written, nil - }); err != nil { + if cap(buf) < need { + buf = make([]byte, need) + } else { + buf = buf[:need] + } + n, err := base64.StdEncoding.Decode(buf, data) + if err != nil { + return moerr.NewInternalErrorNoCtx("vec_from_base64: invalid base64 input") + } + + if n%elemSize != 0 { + return moerr.NewInternalErrorNoCtxf("vec_from_base64: decoded length %d is not a multiple of %d bytes", n, elemSize) + } + + if err = rs.AppendBytes(buf[:n], false); err != nil { return err } } @@ -6058,8 +5920,6 @@ func VecFromBase64[T types.ArrayElement](parameters []*vector.Vector, result vec func Compress(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { source := vector.GenerateFunctionStrParameter(parameters[0]) rs := vector.MustFunctionResult[types.Varlena](result) - var writer *flate.Writer - var output fixedSliceWriter rowCount := uint64(length) for i := uint64(0); i < rowCount; i++ { @@ -6078,126 +5938,49 @@ func Compress(parameters []*vector.Vector, result vector.FunctionResultWrapper, continue } - capacity, err := flateResultCapacity(len(data)) + // Compress using zlib (flate) + var buf bytes.Buffer + writer, err := flate.NewWriter(&buf, flate.DefaultCompression) if err != nil { - return err - } - if writer == nil { - writer, err = flate.NewWriter(io.Discard, flate.DefaultCompression) - if err != nil { + if err := rs.AppendBytes(nil, true); err != nil { return err } + continue } - var compressionErr error - err = rs.AppendBytesWithBuilder(capacity, func(dst []byte) (int, error) { - output.Reset(dst[4:]) - writer.Reset(&output) - if _, compressionErr = writer.Write(data); compressionErr == nil { - compressionErr = writer.Close() - } else { - _ = writer.Close() - } - if compressionErr != nil { - return 0, compressionErr - } - binary.LittleEndian.PutUint32(dst[:4], uint32(len(data))) - return 4 + output.Written(), nil - }) - if compressionErr != nil { - if nullErr := rs.AppendBytes(nil, true); nullErr != nil { - return nullErr + + _, err = writer.Write(data) + if err != nil { + writer.Close() + if err := rs.AppendBytes(nil, true); err != nil { + return err } continue } + + err = writer.Close() if err != nil { - return err + if err := rs.AppendBytes(nil, true); err != nil { + return err + } + continue } - } - - return nil -} - -type fixedSliceWriter struct { - dst []byte - written int - err error -} -func (w *fixedSliceWriter) Reset(dst []byte) { - w.dst = dst - w.written = 0 - w.err = nil -} + compressed := buf.Bytes() -func (w *fixedSliceWriter) Write(value []byte) (int, error) { - if w.err != nil { - return 0, w.err - } - if len(value) > len(w.dst)-w.written { - w.err = io.ErrShortBuffer - return 0, w.err - } - copy(w.dst[w.written:], value) - w.written += len(value) - return len(value), nil -} + // MySQL format: 4-byte length (little-endian) + compressed data + originalLen := uint32(len(data)) + result := make([]byte, 4+len(compressed)) + binary.LittleEndian.PutUint32(result[0:4], originalLen) + copy(result[4:], compressed) -func (w *fixedSliceWriter) WriteString(value string) (int, error) { - if w.err != nil { - return 0, w.err - } - if len(value) > len(w.dst)-w.written { - w.err = io.ErrShortBuffer - return 0, w.err + if err := rs.AppendBytes(result, false); err != nil { + return err + } } - copy(w.dst[w.written:], value) - w.written += len(value) - return len(value), nil -} -func (w *fixedSliceWriter) WriteByte(value byte) error { - if w.err != nil { - return w.err - } - if w.written == len(w.dst) { - w.err = io.ErrShortBuffer - return w.err - } - w.dst[w.written] = value - w.written++ return nil } -func (w *fixedSliceWriter) WriteRune(value rune) (int, error) { - var encoded [utf8.UTFMax]byte - size := utf8.EncodeRune(encoded[:], value) - return w.Write(encoded[:size]) -} - -func (w *fixedSliceWriter) Grow(int) {} - -func (w *fixedSliceWriter) Written() int { - return w.written -} - -func (w *fixedSliceWriter) Err() error { - return w.err -} - -func flateResultCapacity(inputSize int) (int, error) { - if inputSize < 0 || uint64(inputSize) > math.MaxUint32 { - return 0, moerr.NewInvalidInputNoCtx("compress input is too large") - } - // zlib's compressBound is also an upper bound for the contained raw - // DEFLATE stream; add four bytes for MySQL's original-length prefix. - size := uint64(inputSize) - bound := size + (size >> 12) + (size >> 14) + (size >> 25) + 13 - if bound > uint64(^uint(0)>>1)-4 { - return 0, moerr.NewInvalidInputNoCtx("compress result is too large") - } - return int(bound) + 4, nil -} - // Uncompress: UNCOMPRESS(string) - Uncompresses a string compressed by COMPRESS() // Reads 4-byte length, then decompresses the rest func Uncompress(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { @@ -6234,33 +6017,30 @@ func Uncompress(parameters []*vector.Vector, result vector.FunctionResultWrapper originalLen := binary.LittleEndian.Uint32(data[0:4]) compressed := data[4:] - var decodeErr error - err := rs.AppendBytesWithBuilder(int(originalLen), func(dst []byte) (int, error) { - reader := flate.NewReader(bytes.NewReader(compressed)) - defer reader.Close() - if _, decodeErr = io.ReadFull(reader, dst); decodeErr != nil { - return 0, decodeErr - } - var extra [1]byte - n, err := reader.Read(extra[:]) - if err != io.EOF || n != 0 { - if err == nil { - err = moerr.NewInvalidInputNoCtx( - "decompressed length exceeds header", - ) - } - decodeErr = err - return 0, err + // Decompress using zlib (flate) + reader := flate.NewReader(bytes.NewReader(compressed)) + decompressed := make([]byte, originalLen) + n, err := reader.Read(decompressed) + reader.Close() + + if err != nil && err != io.EOF { + // Decompression failed, return NULL + if err := rs.AppendBytes(nil, true); err != nil { + return err } - return len(dst), nil - }) - if decodeErr != nil { + continue + } + + // Check if we got the expected length + if uint32(n) != originalLen { + // Length mismatch, return NULL if err := rs.AppendBytes(nil, true); err != nil { return err } continue } - if err != nil { + + if err := rs.AppendBytes(decompressed, false); err != nil { return err } } @@ -6425,12 +6205,9 @@ func generateSHAKey(key []byte) []byte { } func generateInitializationVector(key []byte, length int) []byte { - hasher := sha256.New() - _, _ = hasher.Write(key) - var lengthByte [1]byte - lengthByte[0] = byte(length) - _, _ = hasher.Write(lengthByte[:]) - return hasher.Sum(nil)[:aes.BlockSize] + data := append(key, byte(length)) + hash := sha256.Sum256(data) + return hash[:aes.BlockSize] } // encode function encrypts a string, returns a binary string of the same length of the original string. @@ -6445,10 +6222,10 @@ func encodeByAES(plaintext []byte, key []byte, null bool, rs *vector.FunctionRes return err } initializationVector := generateInitializationVector(key, len(plaintext)) + ciphertext := make([]byte, len(plaintext)) stream := cipher.NewCTR(block, initializationVector) - return rs.AppendBytesWithFill(len(plaintext), func(ciphertext []byte) { - stream.XORKeyStream(ciphertext, plaintext) - }) + stream.XORKeyStream(ciphertext, plaintext) + return rs.AppendMustBytesValue(ciphertext) } func Encode(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { @@ -6480,10 +6257,10 @@ func decodeByAES(ciphertext []byte, key []byte, null bool, rs *vector.FunctionRe return err } iv := generateInitializationVector(key, len(ciphertext)) + plaintext := make([]byte, len(ciphertext)) stream := cipher.NewCTR(block, iv) - return rs.AppendBytesWithFill(len(ciphertext), func(plaintext []byte) { - stream.XORKeyStream(plaintext, ciphertext) - }) + stream.XORKeyStream(plaintext, ciphertext) + return rs.AppendMustBytesValue(plaintext) } func Decode(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { @@ -6613,20 +6390,18 @@ func RandomBytes(parameters []*vector.Vector, result vector.FunctionResultWrappe continue } - var randomErr error - err := rs.AppendBytesWithBuilder(int(lenVal), func(dst []byte) (int, error) { - var written int - written, randomErr = rand.Read(dst) - return written, randomErr - }) - if randomErr != nil { + // Generate random bytes using crypto/rand + randomBytes := make([]byte, lenVal) + _, err := rand.Read(randomBytes) + if err != nil { // On error, return NULL if err := rs.AppendBytes(nil, true); err != nil { return err } continue } - if err != nil { + + if err := rs.AppendBytes(randomBytes, false); err != nil { return err } } diff --git a/pkg/sql/plan/function/func_unary_codec_scratch_test.go b/pkg/sql/plan/function/func_unary_codec_scratch_test.go deleted file mode 100644 index c9fca4abd8734..0000000000000 --- a/pkg/sql/plan/function/func_unary_codec_scratch_test.go +++ /dev/null @@ -1,161 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package function - -import ( - "strings" - "testing" - - "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/matrixorigin/matrixone/pkg/container/vector" - "github.com/matrixorigin/matrixone/pkg/testutil" - "github.com/stretchr/testify/require" -) - -func TestCompressUncompressDirectOutput(t *testing.T) { - proc := testutil.NewProcess(t) - mp := proc.Mp() - values := []string{ - "", - "a", - strings.Repeat("compressible-value-", 4096), - } - source := newVectorByType(mp, types.T_blob.ToType(), values, nil) - defer source.Free(mp) - - compressed := vector.NewFunctionResultWrapper(types.T_blob.ToType(), mp) - defer compressed.Free() - require.NoError(t, compressed.PreExtendAndReset(len(values))) - require.NoError(t, Compress( - []*vector.Vector{source}, - compressed, - proc, - len(values), - nil, - )) - - decoded := vector.NewFunctionResultWrapper(types.T_blob.ToType(), mp) - defer decoded.Free() - require.NoError(t, decoded.PreExtendAndReset(len(values))) - require.NoError(t, Uncompress( - []*vector.Vector{compressed.GetResultVector()}, - decoded, - proc, - len(values), - nil, - )) - for row, value := range values { - require.Equal(t, []byte(value), decoded.GetResultVector().GetBytesAt(row)) - } -} - -func TestCompressDirectOutputBound(t *testing.T) { - proc := testutil.NewProcess(t) - mp := proc.Mp() - for _, size := range []int{ - 0, 1, 15, 16, 17, 127, 128, 255, 256, - 4095, 4096, 4097, 16383, 16384, 16385, - 65534, 65535, 65536, 1 << 20, - } { - value := make([]byte, size) - state := uint64(size) + 1 - for idx := range value { - state = state*6364136223846793005 + 1442695040888963407 - value[idx] = byte(state >> 56) - } - source := newVectorByType( - mp, - types.T_blob.ToType(), - []string{string(value)}, - nil, - ) - compressed := vector.NewFunctionResultWrapper(types.T_blob.ToType(), mp) - require.NoError(t, compressed.PreExtendAndReset(1)) - require.NoErrorf(t, Compress( - []*vector.Vector{source}, - compressed, - proc, - 1, - nil, - ), "size %d", size) - - decoded := vector.NewFunctionResultWrapper(types.T_blob.ToType(), mp) - require.NoError(t, decoded.PreExtendAndReset(1)) - require.NoErrorf(t, Uncompress( - []*vector.Vector{compressed.GetResultVector()}, - decoded, - proc, - 1, - nil, - ), "size %d", size) - require.Equalf(t, value, decoded.GetResultVector().GetBytesAt(0), "size %d", size) - decoded.Free() - compressed.Free() - source.Free(mp) - } -} - -func TestFromBase64DirectOutputNullAndSelection(t *testing.T) { - proc := testutil.NewProcess(t) - mp := proc.Mp() - source := vector.NewVec(types.T_varchar.ToType()) - defer source.Free(mp) - require.NoError(t, vector.AppendBytes(source, []byte("YWJj"), false, mp)) - require.NoError(t, vector.AppendBytes(source, nil, true, mp)) - require.NoError(t, vector.AppendBytes(source, []byte("ZGVm"), false, mp)) - - result := vector.NewFunctionResultWrapper(types.T_blob.ToType(), mp) - defer result.Free() - require.NoError(t, result.PreExtendAndReset(3)) - require.NoError(t, FromBase64( - []*vector.Vector{source}, - result, - proc, - 3, - &FunctionSelectList{ - AnyNull: true, - SelectList: []bool{true, true, false}, - }, - )) - require.Equal(t, []byte("abc"), result.GetResultVector().GetBytesAt(0)) - require.True(t, result.GetResultVector().IsNull(1)) - require.True(t, result.GetResultVector().IsNull(2)) -} - -func TestRandomBytesDirectOutput(t *testing.T) { - proc := testutil.NewProcess(t) - mp := proc.Mp() - source := newVectorByType( - mp, - types.T_int64.ToType(), - []int64{1, 1024, 0}, - nil, - ) - defer source.Free(mp) - - result := vector.NewFunctionResultWrapper(types.T_blob.ToType(), mp) - defer result.Free() - require.NoError(t, result.PreExtendAndReset(3)) - require.NoError(t, RandomBytes( - []*vector.Vector{source}, - result, - proc, - 3, - nil, - )) - require.Len(t, result.GetResultVector().GetBytesAt(0), 1) - require.Len(t, result.GetResultVector().GetBytesAt(1), 1024) - require.True(t, result.GetResultVector().IsNull(2)) -} diff --git a/pkg/sql/plan/function/function_allocation_scratch_test.go b/pkg/sql/plan/function/function_allocation_scratch_test.go deleted file mode 100644 index 1a0678f67e1df..0000000000000 --- a/pkg/sql/plan/function/function_allocation_scratch_test.go +++ /dev/null @@ -1,763 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package function - -import ( - "math" - "strings" - "testing" - - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/container/nulls" - "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/matrixorigin/matrixone/pkg/container/vector" - "github.com/matrixorigin/matrixone/pkg/testutil" - "github.com/stretchr/testify/require" -) - -const ( - testFunctionOwner mpool.AllocationOwner = 1 - testFunctionResultData mpool.AllocationSite = 1 - testFunctionResultArea mpool.AllocationSite = 2 - testFunctionResultNulls mpool.AllocationSite = 3 - testFunctionResultGroup mpool.AllocationSite = 4 - testFunctionParam mpool.AllocationSite = 5 - testFunctionScratch mpool.AllocationSite = 6 -) - -func newAccountedFunctionResult( - t *testing.T, - typ types.Type, - mp *mpool.MPool, - limit uint64, -) (vector.FunctionResultWrapper, *mpool.AllocationAccountRegistry, *mpool.AllocationAccount) { - t.Helper() - registry, err := mpool.NewAllocationAccountRegistry(1, 32) - require.NoError(t, err) - account, err := registry.Open(limit) - require.NoError(t, err) - selection, err := vector.NewAllocationAccountSelectionWithBitmaps( - account, - testFunctionOwner, - testFunctionResultData, - testFunctionResultArea, - testFunctionResultNulls, - testFunctionResultGroup, - ) - require.NoError(t, err) - allocation, err := vector.NewFunctionAllocation( - account, - testFunctionOwner, - testFunctionParam, - testFunctionScratch, - ) - require.NoError(t, err) - result, err := vector.NewFunctionResultWrapperWithFunctionAllocation( - typ, - mp, - selection, - allocation, - ) - require.NoError(t, err) - return result, registry, account -} - -func finalizeAccountedFunctionResult( - t *testing.T, - result vector.FunctionResultWrapper, - registry *mpool.AllocationAccountRegistry, - account *mpool.AllocationAccount, -) { - t.Helper() - result.Free() - require.Zero(t, account.Snapshot().Used) - account.Seal() - _, err := registry.Finalize(account) - require.NoError(t, err) - require.Zero(t, registry.LiveAllocationMetadata()) -} - -func TestAppendFormattedBytesRollsBackChangedSecondPass(t *testing.T) { - proc := testutil.NewProcess(t) - result, registry, account := newAccountedFunctionResult( - t, - types.T_varchar.ToType(), - proc.Mp(), - 1<<20, - ) - require.NoError(t, result.PreExtendAndReset(1)) - calls := 0 - _, err := appendFormattedBytes( - vector.MustFunctionResult[types.Varlena](result), - func(w formatBuffer) (bool, error) { - calls++ - if calls == 1 { - _, err := w.WriteString("two") - return false, err - } - _, err := w.WriteString("x") - return false, err - }, - ) - require.Error(t, err) - require.Zero(t, result.GetResultVector().Length()) - finalizeAccountedFunctionResult(t, result, registry, account) -} - -func TestBuiltInHashAccountedScratchMatchesLegacy(t *testing.T) { - proc := testutil.NewProcess(t) - mp := proc.Mp() - rows := 3 - nsp := nulls.NewWithSize(rows) - nsp.Set(1) - stringsInput := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{"alpha", "unused", strings.Repeat("wide-value", 256)}, - nsp, - ) - defer stringsInput.Free(mp) - integersInput := newVectorByType( - mp, - types.T_int64.ToType(), - []int64{7, 8, 9}, - nil, - ) - defer integersInput.Free(mp) - inputs := []*vector.Vector{stringsInput, integersInput} - - legacy := vector.NewFunctionResultWrapper(types.T_int64.ToType(), mp) - require.NoError(t, legacy.PreExtendAndReset(rows)) - require.NoError(t, builtInHash(inputs, legacy, proc, rows, nil)) - want := append( - []int64(nil), - vector.MustFixedColWithTypeCheck[int64](legacy.GetResultVector())..., - ) - legacy.Free() - - accounted, registry, account := newAccountedFunctionResult( - t, - types.T_int64.ToType(), - mp, - 1<<20, - ) - require.NoError(t, accounted.PreExtendAndReset(rows)) - require.NoError(t, builtInHash(inputs, accounted, proc, rows, nil)) - require.Equal( - t, - want, - vector.MustFixedColWithTypeCheck[int64](accounted.GetResultVector()), - ) - require.Positive(t, account.Snapshot().Used) - finalizeAccountedFunctionResult(t, accounted, registry, account) -} - -func TestBuiltInHashAccountedScratchRejectsCapacity(t *testing.T) { - proc := testutil.NewProcess(t) - mp := proc.Mp() - input := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{strings.Repeat("x", 4096)}, - nil, - ) - defer input.Free(mp) - result, registry, account := newAccountedFunctionResult( - t, - types.T_int64.ToType(), - mp, - 1024, - ) - require.NoError(t, result.PreExtendAndReset(1)) - err := builtInHash([]*vector.Vector{input}, result, proc, 1, nil) - require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) - finalizeAccountedFunctionResult(t, result, registry, account) -} - -func TestPrefixInAccountedScratchMatchesLegacy(t *testing.T) { - proc := testutil.NewProcess(t) - mp := proc.Mp() - left := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{"abc", "zzz", "other"}, - nil, - ) - defer left.Free(mp) - right := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{"z", "ab", "a"}, - nil, - ) - defer right.Free(mp) - inputs := []*vector.Vector{left, right} - - legacy := vector.NewFunctionResultWrapper(types.T_bool.ToType(), mp) - require.NoError(t, legacy.PreExtendAndReset(3)) - require.NoError(t, newImplPrefixIn().doPrefixIn(inputs, legacy, proc, 3, nil)) - want := append( - []bool(nil), - vector.MustFixedColWithTypeCheck[bool](legacy.GetResultVector())..., - ) - legacy.Free() - - accounted, registry, account := newAccountedFunctionResult( - t, - types.T_bool.ToType(), - mp, - 1<<20, - ) - require.NoError(t, accounted.PreExtendAndReset(3)) - op := newImplPrefixIn() - require.NoError(t, op.doPrefixIn(inputs, accounted, proc, 3, nil)) - require.Equal( - t, - want, - vector.MustFixedColWithTypeCheck[bool](accounted.GetResultVector()), - ) - require.NotEmpty(t, op.scratch) - require.Empty(t, op.vals) - require.Positive(t, account.Snapshot().Used) - - require.NoError(t, accounted.PreExtendAndReset(3)) - require.NoError(t, op.doPrefixIn(inputs, accounted, proc, 3, nil)) - require.Equal( - t, - want, - vector.MustFixedColWithTypeCheck[bool](accounted.GetResultVector()), - ) - finalizeAccountedFunctionResult(t, accounted, registry, account) -} - -func TestPrefixInAccountedScratchRejectsCapacity(t *testing.T) { - proc := testutil.NewProcess(t) - mp := proc.Mp() - left := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{"value"}, - nil, - ) - defer left.Free(mp) - right := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{strings.Repeat("x", 4096)}, - nil, - ) - defer right.Free(mp) - result, registry, account := newAccountedFunctionResult( - t, - types.T_bool.ToType(), - mp, - 1024, - ) - require.NoError(t, result.PreExtendAndReset(1)) - err := newImplPrefixIn().doPrefixIn( - []*vector.Vector{left, right}, - result, - proc, - 1, - nil, - ) - require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) - finalizeAccountedFunctionResult(t, result, registry, account) -} - -func TestJqAccountedOutputMatchesLegacy(t *testing.T) { - proc := testutil.NewProcess(t) - mp := proc.Mp() - jsonInput := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{`{"values":[1,2,3]}`}, - nil, - ) - defer jsonInput.Free(mp) - queryInput := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{`.values | map(. * 2)`}, - nil, - ) - defer queryInput.Free(mp) - params := []*vector.Vector{jsonInput, queryInput} - - legacy := vector.NewFunctionResultWrapper(types.T_varchar.ToType(), mp) - require.NoError(t, legacy.PreExtendAndReset(1)) - require.NoError(t, newOpBuiltInJq().jq(params, legacy, proc, 1, nil)) - want := append([]byte(nil), legacy.GetResultVector().GetBytesAt(0)...) - legacy.Free() - require.Equal(t, []byte(`[2,4,6]`), want) - - accounted, registry, account := newAccountedFunctionResult( - t, - types.T_varchar.ToType(), - mp, - 1<<20, - ) - require.NoError(t, accounted.PreExtendAndReset(1)) - op := newOpBuiltInJq() - require.NoError(t, op.jq(params, accounted, proc, 1, nil)) - require.Equal(t, want, accounted.GetResultVector().GetBytesAt(0)) - require.Positive(t, account.Snapshot().Used) - finalizeAccountedFunctionResult(t, accounted, registry, account) -} - -func TestJqAccountedOutputRejectsCapacity(t *testing.T) { - proc := testutil.NewProcess(t) - mp := proc.Mp() - jsonInput := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{`"` + strings.Repeat("x", 4096) + `"`}, - nil, - ) - defer jsonInput.Free(mp) - queryInput := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{"."}, - nil, - ) - defer queryInput.Free(mp) - params := []*vector.Vector{jsonInput, queryInput} - for _, test := range []struct { - name string - run func(*opBuiltInJq, vector.FunctionResultWrapper) error - }{ - { - name: "jq", - run: func(op *opBuiltInJq, result vector.FunctionResultWrapper) error { - return op.jq(params, result, proc, 1, nil) - }, - }, - { - name: "try_jq", - run: func(op *opBuiltInJq, result vector.FunctionResultWrapper) error { - return op.tryJq(params, result, proc, 1, nil) - }, - }, - } { - t.Run(test.name, func(t *testing.T) { - result, registry, account := newAccountedFunctionResult( - t, - types.T_varchar.ToType(), - mp, - 1024, - ) - require.NoError(t, result.PreExtendAndReset(1)) - err := test.run(newOpBuiltInJq(), result) - require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) - finalizeAccountedFunctionResult(t, result, registry, account) - }) - } -} - -func TestJSONRowAccountedOutputMatchesLegacy(t *testing.T) { - proc := testutil.NewProcess(t) - mp := proc.Mp() - integers := newVectorByType( - mp, - types.T_int64.ToType(), - []int64{7, 8}, - nil, - ) - defer integers.Free(mp) - stringsInput := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{"alpha", "beta"}, - nil, - ) - defer stringsInput.Free(mp) - params := []*vector.Vector{integers, stringsInput} - - legacy := vector.NewFunctionResultWrapper(types.T_varchar.ToType(), mp) - require.NoError(t, legacy.PreExtendAndReset(2)) - require.NoError(t, newOpBuiltInJsonRow().jsonRow(params, legacy, proc, 2, nil)) - want0 := append([]byte(nil), legacy.GetResultVector().GetBytesAt(0)...) - want1 := append([]byte(nil), legacy.GetResultVector().GetBytesAt(1)...) - legacy.Free() - - accounted, registry, account := newAccountedFunctionResult( - t, - types.T_varchar.ToType(), - mp, - 1<<20, - ) - require.NoError(t, accounted.PreExtendAndReset(2)) - require.NoError(t, newOpBuiltInJsonRow().jsonRow( - params, - accounted, - proc, - 2, - nil, - )) - require.Equal(t, want0, accounted.GetResultVector().GetBytesAt(0)) - require.Equal(t, want1, accounted.GetResultVector().GetBytesAt(1)) - require.Positive(t, account.Snapshot().Used) - finalizeAccountedFunctionResult(t, accounted, registry, account) -} - -func TestJSONObjectAccountedKeyScratchMatchesLegacy(t *testing.T) { - proc := testutil.NewProcess(t) - mp := proc.Mp() - keys := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{strings.Repeat("key", 256), "second"}, - nil, - ) - defer keys.Free(mp) - values := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{"first", "value"}, - nil, - ) - defer values.Free(mp) - params := []*vector.Vector{keys, values} - - legacy := vector.NewFunctionResultWrapper(types.T_json.ToType(), mp) - require.NoError(t, legacy.PreExtendAndReset(2)) - require.NoError(t, newOpBuiltInJsonObject().jsonObject(params, legacy, proc, 2, nil)) - want0 := append([]byte(nil), legacy.GetResultVector().GetBytesAt(0)...) - want1 := append([]byte(nil), legacy.GetResultVector().GetBytesAt(1)...) - legacy.Free() - - accounted, registry, account := newAccountedFunctionResult( - t, - types.T_json.ToType(), - mp, - 1<<20, - ) - require.NoError(t, accounted.PreExtendAndReset(2)) - require.NoError(t, newOpBuiltInJsonObject().jsonObject( - params, - accounted, - proc, - 2, - nil, - )) - require.Equal(t, want0, accounted.GetResultVector().GetBytesAt(0)) - require.Equal(t, want1, accounted.GetResultVector().GetBytesAt(1)) - require.Positive(t, account.Snapshot().Used) - finalizeAccountedFunctionResult(t, accounted, registry, account) -} - -func TestJSONObjectAccountedKeyScratchRejectsCapacity(t *testing.T) { - proc := testutil.NewProcess(t) - mp := proc.Mp() - keys := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{strings.Repeat("key", 2048)}, - nil, - ) - defer keys.Free(mp) - values := newVectorByType( - mp, - types.T_int64.ToType(), - []int64{1}, - nil, - ) - defer values.Free(mp) - result, registry, account := newAccountedFunctionResult( - t, - types.T_json.ToType(), - mp, - 1024, - ) - require.NoError(t, result.PreExtendAndReset(1)) - err := newOpBuiltInJsonObject().jsonObject( - []*vector.Vector{keys, values}, - result, - proc, - 1, - nil, - ) - require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) - finalizeAccountedFunctionResult(t, result, registry, account) -} - -func TestJSONModifyAccountedValueScratchRejectsCapacity(t *testing.T) { - proc := testutil.NewProcess(t) - mp := proc.Mp() - document := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{`{"value":0}`}, - nil, - ) - defer document.Free(mp) - path := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{"$.value"}, - nil, - ) - defer path.Free(mp) - arrayType := types.T_array_float32.ToType() - array := vector.NewVec(arrayType) - values := make([]float32, 1024) - for idx := range values { - values[idx] = float32(idx) - } - require.NoError(t, vector.AppendBytes( - array, - types.ArrayToBytes(values), - false, - mp, - )) - defer array.Free(mp) - result, registry, account := newAccountedFunctionResult( - t, - types.T_json.ToType(), - mp, - 1024, - ) - require.NoError(t, result.PreExtendAndReset(1)) - err := newOpBuiltInJsonSet().buildJsonSet( - []*vector.Vector{document, path, array}, - result, - proc, - 1, - nil, - ) - require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) - finalizeAccountedFunctionResult(t, result, registry, account) -} - -func TestFixedInAccountedScratchMatchesMapSemantics(t *testing.T) { - proc := testutil.NewProcess(t) - mp := proc.Mp() - left := newVectorByType( - mp, - types.T_float64.ToType(), - []float64{math.NaN(), math.Copysign(0, -1), 1, 2}, - nil, - ) - defer left.Free(mp) - tupleNulls := nulls.NewWithSize(3) - tupleNulls.Set(2) - tuple := newVectorByType( - mp, - types.T_float64.ToType(), - []float64{math.NaN(), 0, 0}, - tupleNulls, - ) - defer tuple.Free(mp) - params := []*vector.Vector{left, tuple} - - legacy := vector.NewFunctionResultWrapper(types.T_bool.ToType(), mp) - require.NoError(t, legacy.PreExtendAndReset(4)) - require.NoError(t, newOpOperatorFixedIn[float64]().operatorIn( - params, - legacy, - proc, - 4, - nil, - )) - wantValues := append( - []bool(nil), - vector.MustFixedColWithTypeCheck[bool](legacy.GetResultVector())..., - ) - wantNulls := make([]bool, 4) - for row := range wantNulls { - wantNulls[row] = legacy.GetResultVector().IsNull(uint64(row)) - } - legacy.Free() - - accounted, registry, account := newAccountedFunctionResult( - t, - types.T_bool.ToType(), - mp, - 1<<20, - ) - require.NoError(t, accounted.PreExtendAndReset(4)) - op := newOpOperatorFixedIn[float64]() - require.NoError(t, op.operatorIn(params, accounted, proc, 4, nil)) - require.True(t, op.accounted) - require.Nil(t, op.mp) - require.Equal( - t, - wantValues, - vector.MustFixedColWithTypeCheck[bool](accounted.GetResultVector()), - ) - for row, wantNull := range wantNulls { - require.Equal(t, wantNull, accounted.GetResultVector().IsNull(uint64(row))) - } - finalizeAccountedFunctionResult(t, accounted, registry, account) -} - -func TestFixedInAccountedScratchRejectsCapacity(t *testing.T) { - proc := testutil.NewProcess(t) - mp := proc.Mp() - left := newVectorByType(mp, types.T_int64.ToType(), []int64{1}, nil) - defer left.Free(mp) - tupleValues := make([]int64, 1024) - for idx := range tupleValues { - tupleValues[idx] = int64(idx) - } - tuple := newVectorByType(mp, types.T_int64.ToType(), tupleValues, nil) - defer tuple.Free(mp) - result, registry, account := newAccountedFunctionResult( - t, - types.T_bool.ToType(), - mp, - 1024, - ) - require.NoError(t, result.PreExtendAndReset(1)) - err := newOpOperatorFixedIn[int64]().operatorIn( - []*vector.Vector{left, tuple}, - result, - proc, - 1, - nil, - ) - require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) - finalizeAccountedFunctionResult(t, result, registry, account) -} - -func TestStringInAccountedScratchMatchesMapSemantics(t *testing.T) { - proc := testutil.NewProcess(t) - mp := proc.Mp() - left := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{"alpha", "missing", "", "wide"}, - nil, - ) - defer left.Free(mp) - tupleNulls := nulls.NewWithSize(5) - tupleNulls.Set(4) - tuple := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{"wide", "alpha", "alpha", "", "unused"}, - tupleNulls, - ) - defer tuple.Free(mp) - params := []*vector.Vector{left, tuple} - - legacy := vector.NewFunctionResultWrapper(types.T_bool.ToType(), mp) - require.NoError(t, legacy.PreExtendAndReset(4)) - require.NoError(t, newOpOperatorStrIn().operatorIn( - params, - legacy, - proc, - 4, - nil, - )) - wantValues := append( - []bool(nil), - vector.MustFixedColWithTypeCheck[bool](legacy.GetResultVector())..., - ) - wantNulls := make([]bool, 4) - for row := range wantNulls { - wantNulls[row] = legacy.GetResultVector().IsNull(uint64(row)) - } - legacy.Free() - - accounted, registry, account := newAccountedFunctionResult( - t, - types.T_bool.ToType(), - mp, - 1<<20, - ) - require.NoError(t, accounted.PreExtendAndReset(4)) - op := newOpOperatorStrIn() - require.NoError(t, op.operatorIn(params, accounted, proc, 4, nil)) - require.True(t, op.accounted) - require.Nil(t, op.mp) - require.Equal( - t, - wantValues, - vector.MustFixedColWithTypeCheck[bool](accounted.GetResultVector()), - ) - for row, wantNull := range wantNulls { - require.Equal(t, wantNull, accounted.GetResultVector().IsNull(uint64(row))) - } - finalizeAccountedFunctionResult(t, accounted, registry, account) -} - -func TestStringInAccountedScratchRejectsCapacity(t *testing.T) { - proc := testutil.NewProcess(t) - mp := proc.Mp() - left := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{"value"}, - nil, - ) - defer left.Free(mp) - tuple := newVectorByType( - mp, - types.T_varchar.ToType(), - []string{strings.Repeat("x", 4096)}, - nil, - ) - defer tuple.Free(mp) - result, registry, account := newAccountedFunctionResult( - t, - types.T_bool.ToType(), - mp, - 1024, - ) - require.NoError(t, result.PreExtendAndReset(1)) - err := newOpOperatorStrIn().operatorIn( - []*vector.Vector{left, tuple}, - result, - proc, - 1, - nil, - ) - require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) - finalizeAccountedFunctionResult(t, result, registry, account) -} - -func TestNarrowCosineSimilarityAccountedScratchRejectsCapacity(t *testing.T) { - proc := testutil.NewProcess(t) - mp := proc.Mp() - arrayType := types.T_array_uint8.ToType() - left := vector.NewVec(arrayType) - right := vector.NewVec(arrayType) - values := make([]uint8, 4096) - for idx := range values { - values[idx] = uint8(idx) - } - require.NoError(t, vector.AppendBytes(left, types.ArrayToBytes(values), false, mp)) - require.NoError(t, vector.AppendBytes(right, types.ArrayToBytes(values), false, mp)) - defer left.Free(mp) - defer right.Free(mp) - result, registry, account := newAccountedFunctionResult( - t, - types.T_float64.ToType(), - mp, - 1024, - ) - require.NoError(t, result.PreExtendAndReset(1)) - err := CosineSimilarityArrayViaF32[uint8]( - []*vector.Vector{left, right}, - result, - proc, - 1, - nil, - ) - require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) - finalizeAccountedFunctionResult(t, result, registry, account) -} diff --git a/pkg/sql/plan/function/operator_in.go b/pkg/sql/plan/function/operator_in.go index 6e5b99f69fcab..7e57b1158940b 100644 --- a/pkg/sql/plan/function/operator_in.go +++ b/pkg/sql/plan/function/operator_in.go @@ -15,16 +15,6 @@ package function import ( - "bytes" - "cmp" - "encoding/binary" - "math" - "slices" - "sort" - "unsafe" - - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/common/util" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/vm/process" @@ -37,205 +27,15 @@ type TGenericOfIn interface { } type opOperatorFixedIn[T TGenericOfIn] struct { - ready bool - hasNull bool - mp map[T]bool - accounted bool - scratch []byte - count int + ready bool + hasNull bool + mp map[T]bool } type opOperatorStrIn struct { - ready bool - hasNull bool - mp map[string]bool - accounted bool - scratch []byte - count int -} - -func compareInValues[T TGenericOfIn](left, right T) int { - switch value := any(left).(type) { - case uint8: - return cmp.Compare(value, any(right).(uint8)) - case uint16: - return cmp.Compare(value, any(right).(uint16)) - case uint32: - return cmp.Compare(value, any(right).(uint32)) - case uint64: - return cmp.Compare(value, any(right).(uint64)) - case int8: - return cmp.Compare(value, any(right).(int8)) - case int16: - return cmp.Compare(value, any(right).(int16)) - case int32: - return cmp.Compare(value, any(right).(int32)) - case int64: - return cmp.Compare(value, any(right).(int64)) - case float32: - return compareInFloat64(float64(value), float64(any(right).(float32))) - case float64: - return compareInFloat64(value, any(right).(float64)) - case bool: - other := any(right).(bool) - if value == other { - return 0 - } - if !value { - return -1 - } - return 1 - case types.Uuid: - return types.CompareUuid(value, any(right).(types.Uuid)) - case types.Time: - return cmp.Compare(value, any(right).(types.Time)) - case types.Timestamp: - return cmp.Compare(value, any(right).(types.Timestamp)) - case types.Date: - return cmp.Compare(value, any(right).(types.Date)) - case types.Datetime: - return cmp.Compare(value, any(right).(types.Datetime)) - case types.Decimal64: - return value.Compare(any(right).(types.Decimal64)) - case types.Decimal128: - return value.Compare(any(right).(types.Decimal128)) - case types.Decimal256: - return value.Compare(any(right).(types.Decimal256)) - case types.MoYear: - return cmp.Compare(value, any(right).(types.MoYear)) - default: - panic("unsupported IN value type") - } -} - -func compareInFloat64(left, right float64) int { - leftNaN := math.IsNaN(left) - rightNaN := math.IsNaN(right) - switch { - case leftNaN && rightNaN: - return cmp.Compare(math.Float64bits(left), math.Float64bits(right)) - case leftNaN: - return 1 - case rightNaN: - return -1 - default: - return cmp.Compare(left, right) - } -} - -func (op *opOperatorFixedIn[T]) initAccounted( - tuple *vector.Vector, - result vector.FunctionResultWrapper, -) error { - op.hasNull = false - count := 0 - parameter := vector.GenerateFunctionFixedTypeParameter[T](tuple) - for row := uint64(0); row < uint64(tuple.Length()); row++ { - _, isNull := parameter.GetValue(row) - if isNull { - op.hasNull = true - } else { - count++ - } - } - var zero T - elementSize := int(unsafe.Sizeof(zero)) - if count > math.MaxInt/elementSize { - return mpool.ErrAllocationAccountInvalid - } - scratch, selected, err := result.ResizeFunctionScratch(count * elementSize) - if err != nil { - return err - } - if !selected { - return mpool.ErrAllocationAccountInvalid - } - values := util.UnsafeSliceCast[T](scratch)[:count] - write := 0 - for row := uint64(0); row < uint64(tuple.Length()); row++ { - value, isNull := parameter.GetValue(row) - if !isNull { - values[write] = value - write++ - } - } - slices.SortFunc(values, compareInValues[T]) - op.accounted = true - op.scratch = scratch - op.count = count - op.ready = true - return nil -} - -func (op *opOperatorFixedIn[T]) containsAccounted(value T) bool { - values := util.UnsafeSliceCast[T](op.scratch)[:op.count] - idx, found := slices.BinarySearchFunc(values, value, compareInValues[T]) - return found && values[idx] == value -} - -func (op *opOperatorStrIn) initAccounted( - tuple *vector.Vector, - result vector.FunctionResultWrapper, -) error { - op.hasNull = false - count := 0 - payloadSize := 0 - parameter := vector.GenerateFunctionStrParameter(tuple) - for row := uint64(0); row < uint64(tuple.Length()); row++ { - value, isNull := parameter.GetStrValue(row) - if isNull { - op.hasNull = true - continue - } - if len(value) > math.MaxInt-payloadSize { - return mpool.ErrAllocationAccountInvalid - } - payloadSize += len(value) - count++ - } - if count > math.MaxInt/prefixScratchEntrySize || - payloadSize > math.MaxInt-count*prefixScratchEntrySize { - return mpool.ErrAllocationAccountInvalid - } - total := count*prefixScratchEntrySize + payloadSize - if uint64(total) > math.MaxUint32 { - return mpool.ErrAllocationAccountInvalid - } - scratch, selected, err := result.ResizeFunctionScratch(total) - if err != nil { - return err - } - if !selected { - return mpool.ErrAllocationAccountInvalid - } - entries := prefixScratchEntries{data: scratch, count: count} - payloadOffset := count * prefixScratchEntrySize - write := 0 - for row := uint64(0); row < uint64(tuple.Length()); row++ { - value, isNull := parameter.GetStrValue(row) - if isNull { - continue - } - entry := scratch[write*prefixScratchEntrySize:] - binary.LittleEndian.PutUint32(entry, uint32(payloadOffset)) - binary.LittleEndian.PutUint32(entry[4:], uint32(len(value))) - payloadOffset += copy(scratch[payloadOffset:], value) - write++ - } - sort.Sort(entries) - op.accounted = true - op.scratch = scratch - op.count = count - op.ready = true - return nil -} - -func (op *opOperatorStrIn) containsAccounted(value []byte) bool { - entries := prefixScratchEntries{data: op.scratch, count: op.count} - idx := sort.Search(op.count, func(idx int) bool { - return bytes.Compare(entries.value(idx), value) >= 0 - }) - return idx < op.count && bytes.Equal(entries.value(idx), value) + ready bool + hasNull bool + mp map[string]bool } func newOpOperatorFixedIn[T TGenericOfIn]() *opOperatorFixedIn[T] { @@ -318,51 +118,9 @@ func (op *opOperatorStrIn) init(tuple *vector.Vector) { } } -func (op *opOperatorFixedIn[T]) ensureInitialized( - tuple *vector.Vector, - result vector.FunctionResultWrapper, -) error { - if op.ready { - return nil - } - if result.HasFunctionScratch() { - return op.initAccounted(tuple, result) - } - op.init(tuple) - return nil -} - -func (op *opOperatorStrIn) ensureInitialized( - tuple *vector.Vector, - result vector.FunctionResultWrapper, -) error { - if op.ready { - return nil - } - if result.HasFunctionScratch() { - return op.initAccounted(tuple, result) - } - op.init(tuple) - return nil -} - -func (op *opOperatorFixedIn[T]) contains(value T) bool { - if op.accounted { - return op.containsAccounted(value) - } - return op.mp[value] -} - -func (op *opOperatorStrIn) contains(value []byte) bool { - if op.accounted { - return op.containsAccounted(value) - } - return op.mp[string(value)] -} - func (op *opOperatorFixedIn[T]) operatorIn(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - if err := op.ensureInitialized(parameters[1], result); err != nil { - return err + if !op.ready { + op.init(parameters[1]) } p := vector.GenerateFunctionFixedTypeParameter[T](parameters[0]) @@ -374,7 +132,7 @@ func (op *opOperatorFixedIn[T]) operatorIn(parameters []*vector.Vector, result v return err } } else { - ok := op.contains(v) + _, ok := op.mp[v] if !ok && op.hasNull { if err := rs.Append(false, true); err != nil { return err @@ -390,8 +148,8 @@ func (op *opOperatorFixedIn[T]) operatorIn(parameters []*vector.Vector, result v } func (op *opOperatorFixedIn[T]) operatorNotIn(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - if err := op.ensureInitialized(parameters[1], result); err != nil { - return err + if !op.ready { + op.init(parameters[1]) } p := vector.GenerateFunctionFixedTypeParameter[T](parameters[0]) @@ -403,7 +161,7 @@ func (op *opOperatorFixedIn[T]) operatorNotIn(parameters []*vector.Vector, resul return err } } else { - ok := op.contains(v) + _, ok := op.mp[v] if !ok && op.hasNull { if err := rs.Append(false, true); err != nil { return err @@ -419,8 +177,8 @@ func (op *opOperatorFixedIn[T]) operatorNotIn(parameters []*vector.Vector, resul } func (op *opOperatorStrIn) operatorIn(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - if err := op.ensureInitialized(parameters[1], result); err != nil { - return err + if !op.ready { + op.init(parameters[1]) } p := vector.GenerateFunctionStrParameter(parameters[0]) @@ -432,7 +190,7 @@ func (op *opOperatorStrIn) operatorIn(parameters []*vector.Vector, result vector return err } } else { - ok := op.contains(v) + _, ok := op.mp[string(v)] if !ok && op.hasNull { if err := rs.Append(false, true); err != nil { return err @@ -448,8 +206,8 @@ func (op *opOperatorStrIn) operatorIn(parameters []*vector.Vector, result vector } func (op *opOperatorStrIn) operatorNotIn(parameters []*vector.Vector, result vector.FunctionResultWrapper, proc *process.Process, length int, selectList *FunctionSelectList) error { - if err := op.ensureInitialized(parameters[1], result); err != nil { - return err + if !op.ready { + op.init(parameters[1]) } p := vector.GenerateFunctionStrParameter(parameters[0]) @@ -461,7 +219,7 @@ func (op *opOperatorStrIn) operatorNotIn(parameters []*vector.Vector, result vec return err } } else { - ok := op.contains(v) + _, ok := op.mp[string(v)] if !ok && op.hasNull { if err := rs.Append(false, true); err != nil { return err diff --git a/pkg/sql/util/copy_batch_test.go b/pkg/sql/util/copy_batch_test.go index 6f36bb0a7ae9c..659ff1db4f210 100644 --- a/pkg/sql/util/copy_batch_test.go +++ b/pkg/sql/util/copy_batch_test.go @@ -98,7 +98,7 @@ func TestCopyBatchCrossesAllocationOwnershipBoundary(t *testing.T) { require.NoError(t, err) account, err := registry.Open(1 << 20) require.NoError(t, err) - selection, err := vector.NewAllocationAccountSelectionWithBitmaps( + selection, err := vector.NewAllocationAccountSelection( account, 1, 1, diff --git a/pkg/sql/util/eval_expr_util.go b/pkg/sql/util/eval_expr_util.go index 2de311ce9baf5..da58ac7b71d38 100644 --- a/pkg/sql/util/eval_expr_util.go +++ b/pkg/sql/util/eval_expr_util.go @@ -111,40 +111,12 @@ func DecodeBinaryString(s string) ([]byte, error) { } func GenVectorByVarValue(proc *process.Process, typ types.Type, val any) (*vector.Vector, error) { - return GenVectorByVarValueWithAllocation(proc, typ, val, nil) -} - -// GenVectorByVarValueWithAllocation is the dormant allocation-accounted -// variant used by expression executors. A nil selection preserves the legacy -// allocation mode. -func GenVectorByVarValueWithAllocation( - proc *process.Process, - typ types.Type, - val any, - selection *vector.AllocationAccountSelection, -) (*vector.Vector, error) { if val == nil { - if selection == nil { - return vector.NewConstNull(typ, 1, proc.Mp()), nil - } - return vector.NewConstNullWithAllocation(typ, 1, selection) + vec := vector.NewConstNull(typ, 1, proc.Mp()) + return vec, nil } else { strVal := getVal(val) - if selection == nil { - return vector.NewConstBytes( - typ, - []byte(strVal), - 1, - proc.Mp(), - ) - } - return vector.NewConstBytesWithAllocation( - typ, - []byte(strVal), - 1, - proc.Mp(), - selection, - ) + return vector.NewConstBytes(typ, []byte(strVal), 1, proc.Mp()) } } diff --git a/pkg/sql/util/eval_expr_util_test.go b/pkg/sql/util/eval_expr_util_test.go index 46f87f38c685d..679c5c1d738e9 100644 --- a/pkg/sql/util/eval_expr_util_test.go +++ b/pkg/sql/util/eval_expr_util_test.go @@ -19,9 +19,7 @@ import ( "testing" "time" - "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/sql/parsers/tree" "github.com/matrixorigin/matrixone/pkg/testutil" "github.com/matrixorigin/matrixone/pkg/vm/process" @@ -44,53 +42,6 @@ func TestHexToInt(t *testing.T) { require.Error(t, err) } -func TestGenVectorByVarValueWithAllocation(t *testing.T) { - registry, err := mpool.NewAllocationAccountRegistry(1, 4) - require.NoError(t, err) - account, err := registry.Open(1 << 20) - require.NoError(t, err) - selection, err := vector.NewAllocationAccountSelection( - account, - 1, - 1, - 2, - ) - require.NoError(t, err) - proc := testutil.NewProcessWithMPool( - t, - "", - mpool.MustNew("variable-value-allocation"), - ) - defer proc.Free() - - nullVec, err := GenVectorByVarValueWithAllocation( - proc, - types.T_varchar.ToType(), - nil, - selection, - ) - require.NoError(t, err) - require.Same(t, selection, nullVec.AllocationAccountSelection()) - - valueVec, err := GenVectorByVarValueWithAllocation( - proc, - types.T_varchar.ToType(), - "variable payload longer than the inline varlena limit", - selection, - ) - require.NoError(t, err) - require.Same(t, selection, valueVec.AllocationAccountSelection()) - require.Positive(t, account.Snapshot().Used) - - nullVec.Free(proc.Mp()) - valueVec.Free(proc.Mp()) - snapshot := account.Seal() - require.Zero(t, snapshot.Used) - require.Zero(t, registry.LiveAllocationMetadata()) - _, err = registry.Finalize(account) - require.NoError(t, err) -} - func TestSetInsertValueStringBinaryHexPadding(t *testing.T) { proc := testutil.NewProcess(t) diff --git a/pkg/util/resource/summary.go b/pkg/util/resource/summary.go index 05c69fd8541b0..74389d8619690 100644 --- a/pkg/util/resource/summary.go +++ b/pkg/util/resource/summary.go @@ -57,7 +57,7 @@ type MemoryTotals struct { CrossPoolFreeCount uint64 } -// AllocationAccountTotals is the fixed-size terminal observation of activated +// AllocationAccountTotals is the fixed-size terminal observation of accounted // allocation generations. It is diagnostic only: these bytes are a subset of // allocator memory and are never added to MemoryTotals or fed back into // admission. diff --git a/pkg/vectorindex/metric/cpu.go b/pkg/vectorindex/metric/cpu.go index 19acc1e02077e..3ca9222647712 100644 --- a/pkg/vectorindex/metric/cpu.go +++ b/pkg/vectorindex/metric/cpu.go @@ -49,56 +49,6 @@ func PairwiseDistanceLaunch[T types.ArrayElement]( return PairwiseDistanceLaunchCPU(x, y, metric, dist) } -func PairwiseDistanceLaunchOneToMany[T types.RealNumbers]( - query []T, - rowCount int, - rowAt func(int) []T, - metric MetricType, - dist []float32, - _ uint64, - _ bool, -) (PairwiseJobHandle, error) { - return PairwiseDistanceLaunchOneToManyWithScratch( - query, - rowCount, - rowAt, - metric, - dist, - 0, - false, - nil, - ) -} - -func PairwiseDistanceOneToManyScratchSize[T types.RealNumbers]( - _ []T, - _ int, - _ MetricType, - _ uint64, - _ bool, -) (int, bool, error) { - return 0, false, nil -} - -func PairwiseDistanceLaunchOneToManyWithScratch[T types.RealNumbers]( - query []T, - rowCount int, - rowAt func(int) []T, - metric MetricType, - dist []float32, - _ uint64, - _ bool, - _ []byte, -) (PairwiseJobHandle, error) { - return PairwiseDistanceLaunchOneToManyCPU( - query, - rowCount, - rowAt, - metric, - dist, - ) -} - func PairwiseDistanceWait(handle PairwiseJobHandle, metric MetricType) ([]float32, error) { return PairwiseDistanceWaitCPU(handle, metric) } diff --git a/pkg/vectorindex/metric/gpu.go b/pkg/vectorindex/metric/gpu.go index 03a4ace196f97..5c474bd0ab76e 100644 --- a/pkg/vectorindex/metric/gpu.go +++ b/pkg/vectorindex/metric/gpu.go @@ -21,7 +21,6 @@ import ( "sync" "github.com/matrixorigin/matrixone/pkg/common/malloc" - "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/common/util" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/cuvs" @@ -106,7 +105,6 @@ type gpuJob struct { cuvsJobID uint64 deallocators []malloc.Deallocator dist []float32 - scratch []byte } type gpuJobManager struct { @@ -143,20 +141,6 @@ func (m *gpuJobManager) update(jobID uint64, cuvsID uint64, d ...malloc.Dealloca } } -func (m *gpuJobManager) updateScratch( - jobID uint64, - cuvsID uint64, - scratch []byte, -) { - m.mu.Lock() - defer m.mu.Unlock() - job := m.jobs[jobID] - if job != nil { - job.cuvsJobID = cuvsID - job.scratch = scratch - } -} - func (m *gpuJobManager) pop(jobID uint64) *gpuJob { m.mu.Lock() defer m.mu.Unlock() @@ -218,141 +202,6 @@ func PairwiseDistanceLaunch[T types.ArrayElement]( return PairwiseDistanceLaunchCPU(x, y, metric, dist) } -func PairwiseDistanceLaunchOneToMany[T types.RealNumbers]( - query []T, - rowCount int, - rowAt func(int) []T, - metric MetricType, - dist []float32, - minWorkSize uint64, - gpuMode bool, -) (PairwiseJobHandle, error) { - return PairwiseDistanceLaunchOneToManyWithScratch( - query, - rowCount, - rowAt, - metric, - dist, - minWorkSize, - gpuMode, - nil, - ) -} - -func PairwiseDistanceOneToManyScratchSize[T types.RealNumbers]( - query []T, - rowCount int, - metric MetricType, - minWorkSize uint64, - gpuMode bool, -) (int, bool, error) { - if !gpuMode || rowCount <= 0 { - return 0, false, nil - } - dim := len(query) - work := uint64(rowCount) - if dim != 0 && work > ^uint64(0)/uint64(dim) { - work = ^uint64(0) - } else { - work *= uint64(dim) - } - _, supportedMetric := MetricTypeToCuvsMetric[metric] - if !supportedMetric || work < minWorkSize { - return 0, false, nil - } - if _, ok := any(query).([]float32); !ok { - return 0, false, nil - } - rows := uint64(rowCount) + 1 - if rows == 0 || dim != 0 && rows > ^uint64(0)/uint64(dim) { - return 0, false, moerr.NewInternalErrorNoCtx( - "pairwise distance input is too large", - ) - } - elements := rows * uint64(dim) - if elements > uint64(^uint(0)>>1)/4 { - return 0, false, moerr.NewInternalErrorNoCtx( - "pairwise distance input is too large", - ) - } - return int(elements * 4), true, nil -} - -// PairwiseDistanceLaunchOneToManyWithScratch uses caller-owned scratch for -// GPU input flattening when provided. The caller must retain the buffer until -// PairwiseDistanceWait returns. A nil buffer preserves the legacy C-allocator -// path. -func PairwiseDistanceLaunchOneToManyWithScratch[T types.RealNumbers]( - query []T, - rowCount int, - rowAt func(int) []T, - metric MetricType, - dist []float32, - minWorkSize uint64, - gpuMode bool, - scratch []byte, -) (PairwiseJobHandle, error) { - if !gpuMode { - return PairwiseDistanceLaunchOneToManyCPU( - query, - rowCount, - rowAt, - metric, - dist, - ) - } - if rowCount < 0 || len(dist) < rowCount { - return 0, moerr.NewInternalErrorNoCtx( - "pairwise distance output is smaller than the row count", - ) - } - if rowCount == 0 { - return PairwiseDistanceLaunchOneToManyCPU( - query, - rowCount, - rowAt, - metric, - dist, - ) - } - - dim := len(query) - work := uint64(rowCount) - if dim != 0 && work > ^uint64(0)/uint64(dim) { - work = ^uint64(0) - } else { - work *= uint64(dim) - } - cuvsMetric, supportedMetric := MetricTypeToCuvsMetric[metric] - if supportedMetric && - work >= minWorkSize { - if typedQuery, ok := any(query).([]float32); ok { - return gpuPairwiseLaunchRowsWithScratch( - 1, - rowCount, - dim, - func(_ int) []float32 { - return typedQuery - }, - func(row int) []float32 { - return any(rowAt(row)).([]float32) - }, - cuvsMetric, - dist[:rowCount], - 4, - scratch, - ) - } - } - return PairwiseDistanceLaunchOneToManyCPU( - query, - rowCount, - rowAt, - metric, - dist, - ) -} - // gpuPairwiseLaunch flattens [][]C into a C-allocator buffer (elemSize bytes per // element) and launches the async cuVS pairwise distance. C is float32 (4B) or // cuvs.Float16 (2B). Mirrors the old f32-only path, generalized over the element. @@ -363,123 +212,27 @@ func gpuPairwiseLaunch[C cuvs.VectorType]( dist []float32, elemSize int, ) (PairwiseJobHandle, error) { - return gpuPairwiseLaunchRows( - len(x), - len(y), - dim, - func(row int) []C { - return x[row] - }, - func(row int) []C { - return y[row] - }, - cuvsMetric, - dist, - elemSize, - ) -} - -func gpuPairwiseLaunchRows[C cuvs.VectorType]( - nX, nY, dim int, - xAt, yAt func(int) []C, - cuvsMetric cuvs.DistanceType, - dist []float32, - elemSize int, -) (PairwiseJobHandle, error) { - return gpuPairwiseLaunchRowsWithScratch( - nX, - nY, - dim, - xAt, - yAt, - cuvsMetric, - dist, - elemSize, - nil, - ) -} - -func gpuPairwiseLaunchRowsWithScratch[C cuvs.VectorType]( - nX, nY, dim int, - xAt, yAt func(int) []C, - cuvsMetric cuvs.DistanceType, - dist []float32, - elemSize int, - scratch []byte, -) (PairwiseJobHandle, error) { - if nX < 0 || - nY < 0 || - dim < 0 || - elemSize <= 0 || - uint64(dim) > uint64(^uint32(0)) || - uint64(dim) > ^uint64(0)/uint64(elemSize) { - return 0, moerr.NewInternalErrorNoCtx( - "pairwise distance input is too large", - ) - } - rowBytes := uint64(dim) * uint64(elemSize) - if rowBytes != 0 && - (uint64(nX) > ^uint64(0)/rowBytes || - uint64(nY) > ^uint64(0)/rowBytes) { - return 0, moerr.NewInternalErrorNoCtx( - "pairwise distance input is too large", - ) - } - if scratch != nil { - return gpuPairwiseLaunchRowsFromScratch( - nX, - nY, - dim, - xAt, - yAt, - cuvsMetric, - dist, - elemSize, - rowBytes, - scratch, - ) - } + nX, nY := len(x), len(y) allocator := malloc.NewCAllocator() // 1. Flatten Y - yBuf, yDeallocator, err := allocator.Allocate( - uint64(nY)*rowBytes, - malloc.NoClear, - ) + yBuf, yDeallocator, err := allocator.Allocate(uint64(nY*dim*elemSize), malloc.NoClear) if err != nil { return 0, err } yf := util.UnsafeSliceCast[C](yBuf) - for i := 0; i < nY; i++ { - v := yAt(i) - if len(v) != dim { - yDeallocator.Deallocate() - return 0, moerr.NewInternalErrorNoCtx( - "vector dimension not matched", - ) - } + for i, v := range y { copy(yf[i*dim:(i+1)*dim], v) } // 2. Flatten X - xBuf, xDeallocator, err := allocator.Allocate( - uint64(nX)*rowBytes, - malloc.NoClear, - ) + xBuf, xDeallocator, err := allocator.Allocate(uint64(nX*dim*elemSize), malloc.NoClear) if err != nil { yDeallocator.Deallocate() return 0, err } xf := util.UnsafeSliceCast[C](xBuf) - for i := 0; i < nX; i++ { - v := xAt(i) - if len(v) != dim { - xDeallocator.Deallocate() - yDeallocator.Deallocate() - return 0, moerr.NewInternalErrorNoCtx( - "vector dimension not matched", - ) - } + for i, v := range x { copy(xf[i*dim:(i+1)*dim], v) } @@ -508,66 +261,6 @@ func gpuPairwiseLaunchRowsWithScratch[C cuvs.VectorType]( return PairwiseJobHandle(gpuID), nil } -func gpuPairwiseLaunchRowsFromScratch[C cuvs.VectorType]( - nX, nY, dim int, - xAt, yAt func(int) []C, - cuvsMetric cuvs.DistanceType, - dist []float32, - elemSize int, - rowBytes uint64, - scratch []byte, -) (PairwiseJobHandle, error) { - yBytes := uint64(nY) * rowBytes - xBytes := uint64(nX) * rowBytes - if yBytes > uint64(len(scratch)) || xBytes > uint64(len(scratch))-yBytes { - return 0, moerr.NewInternalErrorNoCtx( - "pairwise distance scratch is smaller than the flattened input", - ) - } - yf := util.UnsafeSliceCast[C](scratch[:yBytes]) - for row := 0; row < nY; row++ { - value := yAt(row) - if len(value) != dim { - return 0, moerr.NewInternalErrorNoCtx( - "vector dimension not matched", - ) - } - copy(yf[row*dim:(row+1)*dim], value) - } - xScratch := scratch[yBytes : yBytes+xBytes] - xf := util.UnsafeSliceCast[C](xScratch) - for row := 0; row < nX; row++ { - value := xAt(row) - if len(value) != dim { - return 0, moerr.NewInternalErrorNoCtx( - "vector dimension not matched", - ) - } - copy(xf[row*dim:(row+1)*dim], value) - } - - gpuID := globalGpuJobManager.add(dist) - cuvsID, err := cuvs.PairwiseDistanceLaunch( - xf, - uint64(nX), - yf, - uint64(nY), - uint32(dim), - cuvsMetric, - dist, - ) - if err != nil { - globalGpuJobManager.pop(gpuID) - return 0, err - } - globalGpuJobManager.updateScratch( - gpuID, - cuvsID, - scratch[:yBytes+xBytes], - ) - return PairwiseJobHandle(gpuID), nil -} - // PairwiseDistanceWait waits for the completion of the asynchronous GPU distance // calculation initiated by Launch. func PairwiseDistanceWait(handle PairwiseJobHandle, metric MetricType) ([]float32, error) { diff --git a/pkg/vectorindex/metric/pairwise.go b/pkg/vectorindex/metric/pairwise.go index e0c866aeda284..a012c5f5ade2c 100644 --- a/pkg/vectorindex/metric/pairwise.go +++ b/pkg/vectorindex/metric/pairwise.go @@ -72,15 +72,18 @@ func PairwiseDistanceLaunchCPU[T types.ArrayElement]( dist = make([]float32, nX*nY) } + job := &pairWiseJob{ + dist: dist, + } + // One unified loop over any ArrayElement type — the resolver handles f32/f64 // and the narrow kernels (bf16/f16/int8/uint8) uniformly. - var jobErr error for r := 0; r < nX; r++ { xr := x[r] for c := 0; c < nY; c++ { d, err := distFn(xr, y[c]) if err != nil { - jobErr = err + job.err = err goto DONE } dist[r*nY+c] = d @@ -94,56 +97,6 @@ func PairwiseDistanceLaunchCPU[T types.ArrayElement]( } DONE: - return registerPairwiseCPUJob(dist, jobErr), nil -} - -// PairwiseDistanceLaunchOneToManyCPU computes the distances between one query -// and rowCount caller-owned rows. The caller supplies the output storage so the -// SQL expression path does not need a row-scaled [][]T descriptor slice or an -// additional result allocation. -func PairwiseDistanceLaunchOneToManyCPU[T types.RealNumbers]( - query []T, - rowCount int, - rowAt func(int) []T, - metric MetricType, - dist []float32, -) (PairwiseJobHandle, error) { - if rowCount < 0 || len(dist) < rowCount { - return 0, moerr.NewInternalErrorNoCtx( - "pairwise distance output is smaller than the row count", - ) - } - dist = dist[:rowCount] - distFn, err := ResolveDistanceFn[T, float32](metric) - if err != nil { - return 0, err - } - - var jobErr error - for row := 0; row < rowCount; row++ { - value, err := distFn(query, rowAt(row)) - if err != nil { - jobErr = err - break - } - dist[row] = value - } - if jobErr == nil && metric == Metric_L2Distance { - for idx := range dist { - dist[idx] = float32(math.Sqrt(float64(dist[idx]))) - } - } - return registerPairwiseCPUJob(dist, jobErr), nil -} - -func registerPairwiseCPUJob( - dist []float32, - err error, -) PairwiseJobHandle { - job := &pairWiseJob{ - dist: dist, - err: err, - } jobMu.Lock() id := nextID nextID++ @@ -154,7 +107,7 @@ func registerPairwiseCPUJob( jobMap[uint64(handle)] = job jobMu.Unlock() - return handle + return handle, nil } // PairwiseDistanceWaitCPU returns the results of the pairwise distance calculation diff --git a/pkg/vectorindex/metric/pairwise_test.go b/pkg/vectorindex/metric/pairwise_test.go index aba18191b4417..8aab26132f5b1 100644 --- a/pkg/vectorindex/metric/pairwise_test.go +++ b/pkg/vectorindex/metric/pairwise_test.go @@ -126,38 +126,6 @@ func TestPairwiseDistanceLaunchWaitCPU_Float64_L2(t *testing.T) { require.InDelta(t, 1.0, float64(out[3]), 1e-5) } -func TestPairwiseDistanceLaunchOneToManyCPUUsesCallerOutput(t *testing.T) { - query := []float32{1, 0} - rows := [][]float32{{1, 0}, {1, 1}, {0, 1}} - dist := make([]float32, len(rows)) - - handle, err := PairwiseDistanceLaunchOneToManyCPU( - query, - len(rows), - func(row int) []float32 { - return rows[row] - }, - Metric_L2sqDistance, - dist, - ) - require.NoError(t, err) - out, err := PairwiseDistanceWaitCPU(handle, Metric_L2sqDistance) - require.NoError(t, err) - require.Equal(t, []float32{0, 1, 2}, out) - require.Equal(t, &dist[0], &out[0]) - - _, err = PairwiseDistanceLaunchOneToManyCPU( - query, - len(rows), - func(row int) []float32 { - return rows[row] - }, - Metric_L2sqDistance, - dist[:len(rows)-1], - ) - require.Error(t, err) -} - func TestPairwiseDistanceWaitCPU_InvalidHandle(t *testing.T) { _, err := PairwiseDistanceWaitCPU(PairwiseJobHandle(0), Metric_L2sqDistance) require.Error(t, err) diff --git a/pkg/vm/engine/disttae/txn_test.go b/pkg/vm/engine/disttae/txn_test.go index 5cf3b861c72bd..a739a457e62e5 100644 --- a/pkg/vm/engine/disttae/txn_test.go +++ b/pkg/vm/engine/disttae/txn_test.go @@ -1262,6 +1262,8 @@ func TestDupVectorWithoutNullsLeavesSealedStatementOwner(t *testing.T) { mpool.AllocationOwner(1), mpool.AllocationSite(1), mpool.AllocationSite(2), + mpool.AllocationSite(3), + mpool.AllocationSite(4), ) require.NoError(t, err) source, err := vector.NewOffHeapVecWithTypeAndAllocation( diff --git a/pkg/vm/message/group_sels_test.go b/pkg/vm/message/group_sels_test.go index 91940adb45ec2..bf5fab352ddac 100644 --- a/pkg/vm/message/group_sels_test.go +++ b/pkg/vm/message/group_sels_test.go @@ -25,6 +25,25 @@ func testMp() *mpool.MPool { return mpool.MustNewZero() } +func initTestGroupSels( + t *testing.T, + sels *GroupSels, + n int, + mp *mpool.MPool, +) { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + require.NoError(t, sels.InitWithAllocation(n, mp, account, 1, 1)) + t.Cleanup(func() { + sels.Free(mp) + _, _, err := registry.CompleteTerminal(account) + require.NoError(t, err) + }) +} + func TestGroupSels_NilBeforeInit(t *testing.T) { var js GroupSels require.NoError(t, js.Finalize(3, 3, testMp())) @@ -35,7 +54,7 @@ func TestGroupSels_NilBeforeInit(t *testing.T) { func TestGroupSels_AllUnique(t *testing.T) { mp := testMp() var js GroupSels - require.NoError(t, js.Init(3, mp)) + initTestGroupSels(t, &js, 3, mp) js.Insert(0, 0) js.Insert(1, 1) js.Insert(2, 2) @@ -47,7 +66,7 @@ func TestGroupSels_AllUnique(t *testing.T) { func TestGroupSels_Normal0Based(t *testing.T) { mp := testMp() var js GroupSels - require.NoError(t, js.Init(4, mp)) + initTestGroupSels(t, &js, 4, mp) js.Insert(0, 10) js.Insert(1, 11) js.Insert(0, 12) @@ -63,7 +82,7 @@ func TestGroupSels_Normal0Based(t *testing.T) { func TestGroupSels_Dedup1Based(t *testing.T) { mp := testMp() var js GroupSels - require.NoError(t, js.Init(3, mp)) + initTestGroupSels(t, &js, 3, mp) js.Insert(1, 0) js.Insert(2, 1) js.Insert(1, 2) @@ -78,7 +97,7 @@ func TestGroupSels_Dedup1Based(t *testing.T) { func TestGroupSels_Free(t *testing.T) { mp := testMp() var js GroupSels - require.NoError(t, js.Init(2, mp)) + initTestGroupSels(t, &js, 2, mp) js.Insert(0, 5) js.Free(mp) require.NoError(t, js.Finalize(1, 1, mp)) @@ -89,7 +108,7 @@ func TestGroupSels_AllNulls(t *testing.T) { // all rows are null — Init called but Insert never called mp := testMp() var js GroupSels - require.NoError(t, js.Init(3, mp)) + initTestGroupSels(t, &js, 3, mp) // no Insert calls require.NoError(t, js.Finalize(0, 3, mp)) require.Nil(t, js.offsets) @@ -100,7 +119,7 @@ func TestGroupSels_NullsSkipped(t *testing.T) { // must NOT trigger all-unique path mp := testMp() var js GroupSels - require.NoError(t, js.Init(4, mp)) + initTestGroupSels(t, &js, 4, mp) js.Insert(0, 0) // row 0 → group 0 js.Insert(1, 1) // row 1 → group 1 // row 2 is null, skipped diff --git a/pkg/vm/message/joinMapDependency_test.go b/pkg/vm/message/joinMapDependency_test.go index 77b3965369220..50513eca3cc4c 100644 --- a/pkg/vm/message/joinMapDependency_test.go +++ b/pkg/vm/message/joinMapDependency_test.go @@ -48,6 +48,19 @@ func TestJoinMapResultDistinguishesSuccessEmptyAndBuildError(t *testing.T) { require.Equal(t, baseErr.ErrorCode(), got.ErrorCode()) } +func TestSendJoinMapResultRetainsOwnershipWhenBoardUnavailable(t *testing.T) { + jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, nil) + require.False(t, SendJoinMapResult( + NewJoinMapResult(jm), + 1, + false, + 0, + nil, + )) + require.True(t, jm.IsValid()) + jm.FreeMemory() +} + func TestRuntimeFilterMemoryReleaseIsSharedAcrossMessageCopies(t *testing.T) { var releases atomic.Int32 msg := RuntimeFilterMessage{Data: make([]byte, 128)} diff --git a/pkg/vm/message/joinMapMsg.go b/pkg/vm/message/joinMapMsg.go index ed3864b2d06e8..d7385949e22a9 100644 --- a/pkg/vm/message/joinMapMsg.go +++ b/pkg/vm/message/joinMapMsg.go @@ -50,10 +50,6 @@ func freeSlice[T any](mp *mpool.MPool, s []T) { mpool.FreeSlice(mp, s[:cap(s)]) } -func (sels *GroupSels) Init(n int, mp *mpool.MPool) error { - return sels.InitWithAllocation(n, mp, nil, 0, 0) -} - // InitWithAllocation makes the complete temporary/final row-index owner use // one immutable allocation generation. GroupSels is copied into JoinMap at // publication, so its physical slices retain this provenance until the last @@ -71,21 +67,19 @@ func (sels *GroupSels) InitWithAllocation( if sels.tmp != nil || sels.vals != nil || sels.offsets != nil { return mpool.ErrAllocationAccountInvariant } - var err error - if account == nil { - if owner != 0 || site != 0 { - return mpool.ErrAllocationAccountInvalid - } - sels.tmp, err = mpool.MakeSlice[int32](n*2, mp, false) - } else { - sels.tmp, err = mpool.MakeSliceAccounted[int32]( - n*2, - mp, - account, - owner, - site, - ) + if account == nil || account.Handle() == 0 || + owner < mpool.AllocationOwnerMin || owner > mpool.AllocationOwnerMax || + site < mpool.AllocationSiteMin { + return mpool.ErrAllocationAccountInvalid } + var err error + sels.tmp, err = mpool.MakeSliceAccounted[int32]( + n*2, + mp, + account, + owner, + site, + ) if err != nil { return err } @@ -98,7 +92,7 @@ func (sels *GroupSels) InitWithAllocation( func (sels *GroupSels) makeSlice(n int, mp *mpool.MPool) ([]int32, error) { if sels.account == nil { - return mpool.MakeSlice[int32](n, mp, false) + return nil, mpool.ErrAllocationAccountInvalid } return mpool.MakeSliceAccounted[int32]( n, @@ -224,9 +218,9 @@ type JoinMap struct { memoryReleaseOnce sync.Once // A resident JoinMap may be broadcast to multiple consumers, but a spill - // payload is move-only. Keep the complete payload behind one lock so files, - // legacy descriptors, and the producer budget generation cannot be claimed - // by different consumers. + // payload is move-only. Keep the complete payload behind one lock so files + // and the producer budget generation cannot be claimed by different + // consumers. spillMu sync.Mutex spilled atomic.Bool spillPayload SpillBuildPayload @@ -236,9 +230,7 @@ type JoinMap struct { var ( ErrSpillBuildPayloadEmpty = moerr.NewInternalErrorNoCtx("spill build payload is empty") - ErrSpillBuildPayloadMixed = moerr.NewInternalErrorNoCtx("spill build payload mixes accounted files and legacy descriptors") ErrSpillBuildBudgetRef = moerr.NewInternalErrorNoCtx("accounted spill build payload is missing its budget reference") - ErrSpillBuildLegacyBudget = moerr.NewInternalErrorNoCtx("legacy spill build payload must not carry a budget reference") ErrSpillBuildPayloadSet = moerr.NewInternalErrorNoCtx("spill build payload is already set") ErrSpillBuildPayloadTaken = moerr.NewInternalErrorNoCtx("spill build payload is already taken") ErrSpillBuildShared = moerr.NewInternalErrorNoCtx("spill build payload requires exactly one consumer") @@ -250,7 +242,6 @@ var ( // transferring the files to SpillEngine. type SpillBuildPayload struct { Files []*SpillFile - LegacyFds []*os.File BudgetRef any } @@ -271,15 +262,6 @@ func (p *SpillBuildPayload) Close() error { } } p.Files = nil - for i, fd := range p.LegacyFds { - if fd != nil { - if err := fd.Close(); err != nil && firstErr == nil { - firstErr = err - } - p.LegacyFds[i] = nil - } - } - p.LegacyFds = nil p.BudgetRef = nil return firstErr } @@ -327,6 +309,40 @@ func (f *SpillFile) Bytes() uint64 { return f.bytes } +// Validate proves that the move-only descriptor still names the complete +// physical file recorded by its producer. Writers publish Rows and Bytes only +// after complete-record writes, so a non-empty spill file must have positive +// metadata and an exact physical size before any record is decoded. +func (f *SpillFile) Validate() error { + if f == nil { + return moerr.NewInternalErrorNoCtx("nil spill file") + } + f.mu.Lock() + defer f.mu.Unlock() + if f.fd == nil { + return moerr.NewInternalErrorNoCtx("invalid spill file metadata") + } + if f.rows <= 0 { + return moerr.NewInternalErrorNoCtx("corrupted spill file row count metadata") + } + if f.bytes == 0 { + return moerr.NewInternalErrorNoCtx("corrupted spill file size metadata") + } + info, err := f.fd.Stat() + if err != nil { + return err + } + if info.Size() < 0 || uint64(info.Size()) != f.bytes { + return moerr.NewInternalErrorf( + context.Background(), + "corrupted spill file size: expected=%d actual=%d", + f.bytes, + info.Size(), + ) + } + return nil +} + func (f *SpillFile) Close() error { if f == nil { return nil @@ -465,18 +481,12 @@ func (jm *JoinMap) SetMemoryRelease(release func()) { // // On error, ownership remains with the caller. func (jm *JoinMap) SetSpillBuildPayload(payload SpillBuildPayload) error { - if jm == nil || (len(payload.Files) == 0 && len(payload.LegacyFds) == 0) { + if jm == nil || len(payload.Files) == 0 { return ErrSpillBuildPayloadEmpty } - if len(payload.Files) > 0 && len(payload.LegacyFds) > 0 { - return ErrSpillBuildPayloadMixed - } - if len(payload.Files) > 0 && payload.BudgetRef == nil { + if payload.BudgetRef == nil { return ErrSpillBuildBudgetRef } - if len(payload.LegacyFds) > 0 && payload.BudgetRef != nil { - return ErrSpillBuildLegacyBudget - } jm.spillMu.Lock() defer jm.spillMu.Unlock() if jm.spillPayloadTaken { @@ -495,8 +505,8 @@ func (jm *JoinMap) SetSpillBuildPayload(payload SpillBuildPayload) error { return nil } -// TakeSpillBuildPayload atomically transfers files, legacy descriptors, and -// the producer budget generation to the sole spill consumer. +// TakeSpillBuildPayload atomically transfers files and the producer budget +// generation to the sole spill consumer. func (jm *JoinMap) TakeSpillBuildPayload() (SpillBuildPayload, error) { if jm == nil { return SpillBuildPayload{}, ErrSpillBuildPayloadEmpty @@ -574,25 +584,10 @@ func (jm *JoinMap) PreAlloc(n uint64) error { } type JoinMapMsg struct { - JoinMapPtr *JoinMap IsShuffle bool ShuffleIdx int32 Tag int32 - Spilled bool - // Result is the terminal dependency state. The zero value is retained for - // source compatibility with older direct JoinMapMsg literals; those are - // interpreted as an explicit successful result (including nil for an empty - // build) by terminalResult. - Result JoinMapResult -} - -func (t JoinMapMsg) terminalResult() JoinMapResult { - if t.Result.Finalized() { - return t.Result - } - // Legacy messages predate the explicit result field. A nil JoinMap in a - // legacy message is the established empty-build success convention. - return NewJoinMapResult(t.JoinMapPtr) + Result JoinMapResult } func (t JoinMapMsg) Serialize() []byte { @@ -608,10 +603,7 @@ func (t JoinMapMsg) NeedBlock() bool { } func (t JoinMapMsg) Destroy() { - jm := t.JoinMapPtr - if jm == nil && t.Result.IsSuccess() { - jm = t.Result.JoinMap() - } + jm := t.Result.JoinMap() if jm != nil { jm.FreeMemory() } @@ -627,13 +619,13 @@ func (t JoinMapMsg) DebugString() string { if t.IsShuffle { buf.WriteString("shuffle index " + strconv.Itoa(int(t.ShuffleIdx)) + "\n") } - if t.JoinMapPtr != nil { - buf.WriteString("joinmap rowcnt " + strconv.Itoa(int(t.JoinMapPtr.rowCnt)) + "\n") - buf.WriteString("joinmap refcnt " + strconv.Itoa(int(t.JoinMapPtr.GetRefCount())) + "\n") + if jm := t.Result.JoinMap(); jm != nil { + buf.WriteString("joinmap rowcnt " + strconv.Itoa(int(jm.rowCnt)) + "\n") + buf.WriteString("joinmap refcnt " + strconv.Itoa(int(jm.GetRefCount())) + "\n") } else if t.Result.IsBuildError() { buf.WriteString("joinmap build error " + t.Result.BuildError().Error() + "\n") } else { - buf.WriteString("joinmapPtr is nil \n") + buf.WriteString("joinmap is nil \n") } return buf.String() } @@ -679,7 +671,7 @@ func ReceiveJoinMapResult(tag int32, isShuffle bool, shuffleIdx int32, mb *Messa continue } } - result := msg.terminalResult() + result := msg.Result if !result.Finalized() { // A malformed/zero result must not be interpreted as empty. Keep // waiting for the producer's terminal publication. @@ -701,36 +693,24 @@ func ReceiveJoinMapResult(tag int32, isShuffle bool, shuffleIdx int32, mb *Messa } // SendJoinMapResult publishes one terminal dependency value without waiting -// for any consumer acknowledgement. The caller owns exactly-once admission -// (typically an atomic generation gate in HashBuild); this function only -// performs the non-blocking MessageBoard publication. -func SendJoinMapResult(result JoinMapResult, tag int32, isShuffle bool, shuffleIdx int32, mb *MessageBoard) { - if !result.Finalized() { - return +// for consumer acknowledgement. True means the MessageBoard accepted the +// value's ownership (or was already closed and destroyed it); false leaves +// ownership with the caller. +func SendJoinMapResult(result JoinMapResult, tag int32, isShuffle bool, shuffleIdx int32, mb *MessageBoard) bool { + if !result.Finalized() || mb == nil || mb.rwMutex == nil { + return false } msg := JoinMapMsg{ - JoinMapPtr: result.JoinMap(), IsShuffle: isShuffle, ShuffleIdx: shuffleIdx, Tag: tag, Result: result, } - if jm := result.JoinMap(); jm != nil { - msg.Spilled = jm.IsSpilled() - } SendMessage(msg, mb) + return true } // FinalizeJoinMapBuildError publishes a typed BuildError terminal value. -// It is kept separate from FinalizeJoinMapMessage so legacy nil-map empty -// build compatibility cannot accidentally turn an admission failure into a -// successful empty dependency. -func FinalizeJoinMapBuildError(mb *MessageBoard, tag int32, isShuffle bool, shuffleIdx int32, err error) { - SendJoinMapResult(NewJoinMapBuildErrorResult(err), tag, isShuffle, shuffleIdx, mb) -} - -func FinalizeJoinMapMessage(mb *MessageBoard, tag int32, isShuffle bool, shuffleIdx int32, sendMapSucceed bool) { - if !sendMapSucceed { - SendJoinMapResult(NewJoinMapResult(nil), tag, isShuffle, shuffleIdx, mb) - } +func FinalizeJoinMapBuildError(mb *MessageBoard, tag int32, isShuffle bool, shuffleIdx int32, err error) bool { + return SendJoinMapResult(NewJoinMapBuildErrorResult(err), tag, isShuffle, shuffleIdx, mb) } diff --git a/pkg/vm/message/message_test.go b/pkg/vm/message/message_test.go index 011392d217b18..549bd8153b364 100644 --- a/pkg/vm/message/message_test.go +++ b/pkg/vm/message/message_test.go @@ -102,7 +102,7 @@ func TestJoinMapMsgDestroyReleasesJoinMapMemory(t *testing.T) { shm: shm, } - JoinMapMsg{JoinMapPtr: jm, Tag: 1}.Destroy() + JoinMapMsg{Result: NewJoinMapResult(jm), Tag: 1}.Destroy() require.Nil(t, jm.shm) require.False(t, jm.valid) @@ -215,164 +215,6 @@ func TestMessageBoardFinalizerDestroysQueuedMessages(t *testing.T) { }, 5*time.Second, 20*time.Millisecond) } -func TestLegacySpillBuildPayload(t *testing.T) { - t.Run("transfers_ownership", func(t *testing.T) { - f1, err := os.CreateTemp("", "test_fd_*") - require.NoError(t, err) - defer os.Remove(f1.Name()) - f2, err := os.CreateTemp("", "test_fd_*") - require.NoError(t, err) - defer os.Remove(f2.Name()) - - jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, mpool.MustNewZero()) - jm.IncRef(1) - require.NoError(t, jm.SetSpillBuildPayload(SpillBuildPayload{ - LegacyFds: []*os.File{f1, f2}, - })) - payload, err := jm.TakeSpillBuildPayload() - require.NoError(t, err) - - require.Len(t, payload.LegacyFds, 2) - require.Same(t, f1, payload.LegacyFds[0]) - require.Same(t, f2, payload.LegacyFds[1]) - - require.NoError(t, payload.Close()) - }) - - t.Run("second_call_returns_explicit_error", func(t *testing.T) { - f, err := os.CreateTemp("", "test_fd_*") - require.NoError(t, err) - defer os.Remove(f.Name()) - - jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, mpool.MustNewZero()) - jm.IncRef(1) - require.NoError(t, jm.SetSpillBuildPayload(SpillBuildPayload{ - LegacyFds: []*os.File{f}, - })) - payload, err := jm.TakeSpillBuildPayload() - require.NoError(t, err) - - _, err = jm.TakeSpillBuildPayload() - require.ErrorIs(t, err, ErrSpillBuildPayloadTaken) - require.NoError(t, payload.Close()) - }) - - t.Run("empty_payload_is_rejected", func(t *testing.T) { - jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, mpool.MustNewZero()) - jm.IncRef(1) - require.ErrorIs(t, jm.SetSpillBuildPayload(SpillBuildPayload{}), ErrSpillBuildPayloadEmpty) - }) - - t.Run("mixed_payload_is_rejected", func(t *testing.T) { - jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, mpool.MustNewZero()) - jm.IncRef(1) - require.ErrorIs(t, jm.SetSpillBuildPayload(SpillBuildPayload{ - Files: []*SpillFile{nil}, - LegacyFds: []*os.File{nil}, - BudgetRef: struct{}{}, - }), ErrSpillBuildPayloadMixed) - }) - - t.Run("legacy_payload_with_budget_is_rejected", func(t *testing.T) { - jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, mpool.MustNewZero()) - jm.IncRef(1) - require.ErrorIs(t, jm.SetSpillBuildPayload(SpillBuildPayload{ - LegacyFds: []*os.File{nil}, - BudgetRef: struct{}{}, - }), ErrSpillBuildLegacyBudget) - }) - - t.Run("free_before_set_rejects_late_publication", func(t *testing.T) { - f, err := os.CreateTemp("", "test_fd_*") - require.NoError(t, err) - defer os.Remove(f.Name()) - - jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, mpool.MustNewZero()) - jm.IncRef(1) - jm.FreeMemory() - payload := SpillBuildPayload{LegacyFds: []*os.File{f}} - require.ErrorIs(t, jm.SetSpillBuildPayload(payload), ErrSpillBuildPayloadTaken) - _, err = f.Stat() - require.NoError(t, err, "rejected late payload remains caller-owned") - require.NoError(t, payload.Close()) - }) -} - -func TestFreeMemoryClosesSpillFds(t *testing.T) { - t.Run("closes_all_fds", func(t *testing.T) { - mp := mpool.MustNewZero() - f1, err := os.CreateTemp("", "test_fd_*") - require.NoError(t, err) - defer os.Remove(f1.Name()) - f2, err := os.CreateTemp("", "test_fd_*") - require.NoError(t, err) - defer os.Remove(f2.Name()) - - jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, mp) - jm.IncRef(1) - require.NoError(t, jm.SetSpillBuildPayload(SpillBuildPayload{ - LegacyFds: []*os.File{f1, f2}, - })) - jm.FreeMemory() - - require.False(t, jm.valid) - - // Verify fds are closed - _, err = f1.Stat() - require.Error(t, err) - _, err = f2.Stat() - require.Error(t, err) - }) - - t.Run("handles_nil_in_fd_slice", func(t *testing.T) { - mp := mpool.MustNewZero() - f, err := os.CreateTemp("", "test_fd_*") - require.NoError(t, err) - defer os.Remove(f.Name()) - - jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, mp) - jm.IncRef(1) - require.NoError(t, jm.SetSpillBuildPayload(SpillBuildPayload{ - LegacyFds: []*os.File{f, nil}, - })) - jm.FreeMemory() // must not panic on nil entry - }) - - t.Run("take_then_free_does_not_double_close", func(t *testing.T) { - mp := mpool.MustNewZero() - f, err := os.CreateTemp("", "test_fd_*") - require.NoError(t, err) - defer os.Remove(f.Name()) - - jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, mp) - jm.IncRef(1) - require.NoError(t, jm.SetSpillBuildPayload(SpillBuildPayload{ - LegacyFds: []*os.File{f}, - })) - payload, err := jm.TakeSpillBuildPayload() - require.NoError(t, err) - require.Len(t, payload.LegacyFds, 1) - - // FreeMemory after TakeSpillBuildPayload should not close the fds. - jm.FreeMemory() - - // fd is still open (caller owns it) - _, err = f.Stat() - require.NoError(t, err) - require.NoError(t, payload.Close()) - }) - - t.Run("double_free_safe", func(t *testing.T) { - mp := mpool.MustNewZero() - jm := &JoinMap{ - valid: true, - mpool: mp, - } - jm.FreeMemory() - jm.FreeMemory() // must not panic - }) -} - func TestAccountedSpillFileOwnership(t *testing.T) { var releases atomic.Int32 newFile := func() (*SpillFile, string) { diff --git a/pkg/vm/process/cte_memory_budget.go b/pkg/vm/process/cte_memory_budget.go index eb3de96d2cf3e..af68a441cd81f 100644 --- a/pkg/vm/process/cte_memory_budget.go +++ b/pkg/vm/process/cte_memory_budget.go @@ -113,15 +113,15 @@ func (r *CTEMemoryReservation) Resize(ctx context.Context, bytes uint64) error { return nil } -func (r *CTEMemoryReservation) Release() { +func (r *CTEMemoryReservation) Release() bool { if r == nil || r.budget == nil { - return + return false } b := r.budget b.mu.Lock() defer b.mu.Unlock() if !r.active { - return + return false } if !b.closed { if r.bytes <= b.used { @@ -132,6 +132,7 @@ func (r *CTEMemoryReservation) Release() { } r.bytes = 0 r.active = false + return true } func (r *CTEMemoryReservation) Bytes() uint64 { diff --git a/pkg/vm/process/hashbuild_budget.go b/pkg/vm/process/hashbuild_budget.go index 4734ceffc1c15..266d71b5c43c9 100644 --- a/pkg/vm/process/hashbuild_budget.go +++ b/pkg/vm/process/hashbuild_budget.go @@ -126,18 +126,12 @@ func observeHashBuildBudget(component, event, scope string, bytes uint64) { // process rather than the SQL layer so that operators and remote execution // code can make an admission decision without importing frontend packages. var ( - ErrHashBuildBudgetAdmission = moerr.NewInternalErrorNoCtx("hash build budget admission rejected") - // ErrHashBuildBudgetRejected is kept as a more discoverable spelling of the - // admission sentinel. It is the same value, so errors.Is works with either. - ErrHashBuildBudgetRejected = ErrHashBuildBudgetAdmission - ErrHashBuildBudgetClosed = moerr.NewInternalErrorNoCtx("hash build budget is closed") - ErrHashBuildBudgetInvalid = moerr.NewInternalErrorNoCtx("invalid hash build budget") - ErrHashBuildCeilingMissing = moerr.NewInternalErrorNoCtx("hash build budget ceiling unavailable") - ErrHashBuildBudgetUnavailable = ErrHashBuildCeilingMissing - ErrHashBuildReservationInactive = moerr.NewInternalErrorNoCtx("hash build reservation is inactive") - ErrHashBuildReservationUpward = moerr.NewInternalErrorNoCtx("hash build reservation reconciliation would increase charge") - ErrHashBuildReservationReconcileUpward = ErrHashBuildReservationUpward - ErrHashBuildReservationClosed = ErrHashBuildReservationInactive + ErrHashBuildBudgetAdmission = moerr.NewInternalErrorNoCtx("hash build budget admission rejected") + ErrHashBuildBudgetClosed = moerr.NewInternalErrorNoCtx("hash build budget is closed") + ErrHashBuildBudgetInvalid = moerr.NewInternalErrorNoCtx("invalid hash build budget") + ErrHashBuildCeilingMissing = moerr.NewInternalErrorNoCtx("hash build budget ceiling unavailable") + ErrHashBuildSpillReservationInactive = moerr.NewInternalErrorNoCtx("hash build spill reservation is inactive") + ErrHashBuildSpillReservationUpward = moerr.NewInternalErrorNoCtx("hash build spill reservation reconciliation would increase charge") ) // HashBuildBudgetErrorKind identifies the class of a budget error. @@ -151,11 +145,10 @@ const ( ) // HashBuildBudgetComponent identifies the independently bounded resource that -// rejected an admission. The zero value remains the memory component for -// compatibility with older callers that construct HashBuildBudgetError -// directly. A spill-disk or spill-FD rejection must never enter the memory -// reclaim/reduce loop: reducing an in-memory batch cannot create either -// resource and may replay already-published spill records. +// rejected an admission. Zero is invalid: every admission error must name its +// physical resource. A spill-disk or spill-FD rejection must never enter the +// memory reclaim/reduce loop because reducing an in-memory batch cannot create +// either resource and may replay already-published spill records. type HashBuildBudgetComponent uint8 const ( @@ -223,7 +216,7 @@ func (e *HashBuildBudgetError) Is(target error) bool { if e == nil { return false } - if target == ErrHashBuildBudgetAdmission || target == ErrHashBuildBudgetRejected { + if target == ErrHashBuildBudgetAdmission { return e.Kind == HashBuildBudgetErrorAdmission } switch e.Kind { @@ -475,17 +468,6 @@ func MustNewHashBuildBudget(aggregateCap, queryCap uint64) *HashBuildBudget { return b } -func NewHashBuildBudgetWithSpillCaps(aggregateCap, queryCap, spillDiskCap, spillFDCap uint64) (*HashBuildBudget, error) { - b, err := NewHashBuildBudget(aggregateCap, queryCap) - if err != nil { - return nil, err - } - if err = b.SetSpillCaps(spillDiskCap, spillFDCap); err != nil { - return nil, err - } - return b, nil -} - // AggregateCap returns the configured local-CN cap. func (b *HashBuildBudget) AggregateCap() uint64 { if b == nil { @@ -496,10 +478,6 @@ func (b *HashBuildBudget) AggregateCap() uint64 { return b.aggregateCap } -// CNHashCap is an alias useful to callers that describe the aggregate as the -// CN hash cap. -func (b *HashBuildBudget) CNHashCap() uint64 { return b.AggregateCap() } - // QueryCap returns the per-generation, per-target-CN cap. func (b *HashBuildBudget) QueryCap() uint64 { if b == nil { @@ -520,15 +498,6 @@ func (b *HashBuildBudget) AggregateUsed() uint64 { return b.aggregateUsed } -// CNHashUsed is an alias for AggregateUsed. -func (b *HashBuildBudget) CNHashUsed() uint64 { return b.AggregateUsed() } - -// Current is a concise alias for AggregateUsed. -func (b *HashBuildBudget) Current() uint64 { return b.AggregateUsed() } - -// Capacity is a concise alias for AggregateCap. -func (b *HashBuildBudget) Capacity() uint64 { return b.AggregateCap() } - // Closed reports whether no new generation or reservation may be opened. func (b *HashBuildBudget) Closed() bool { if b == nil { @@ -849,7 +818,6 @@ type HashBuildBudgetGeneration struct { id uint64 cap uint64 used uint64 - allocationUsed uint64 closed bool spillDiskCap, spillDiskUsed uint64 spillFDConfiguredCap, spillFDCap, spillFDUsed uint64 @@ -862,21 +830,12 @@ var _ mpool.AllocationCapacityController = (*HashBuildBudgetGeneration)(nil) // HashBuildBudgetGenerationSnapshot is an immutable fixed-cardinality view. type HashBuildBudgetGenerationSnapshot struct { ID, Cap, Used, PeakUsed uint64 - AllocationUsed uint64 ReserveCount, RejectCount, ReconcileCount, ReleaseCount uint64 SpillDiskCap, SpillDiskUsed, SpillFDCap uint64 SpillFDUsed uint64 Closed bool } -// HashBuildGeneration is a shorter spelling retained for call sites. -type HashBuildGeneration = HashBuildBudgetGeneration - -// HashBuildQueryBudget makes the per-generation/per-target-CN scope explicit -// at call sites. It is an alias, so tokens and methods retain one ownership -// implementation. -type HashBuildQueryBudget = HashBuildBudgetGeneration - // OpenGeneration opens a per-statement execution generation. The budget's // query cap is copied by reference (and remains immutable), while used bytes // belong solely to the returned generation. @@ -987,10 +946,7 @@ func (b *HashBuildBudget) openProcessGeneration( } } - memoryCap := requestedMemoryCap - if memoryCap > b.aggregateCap { - memoryCap = b.aggregateCap - } + memoryCap := min(requestedMemoryCap, b.aggregateCap) if memoryCap == 0 { return nil, &HashBuildBudgetError{ Kind: HashBuildBudgetErrorInvalid, @@ -1003,17 +959,12 @@ func (b *HashBuildBudget) openProcessGeneration( if spillDiskCap == 0 { spillDiskCap = defaultSpillCap(memoryCap) } - if spillDiskCap > b.spillDiskCap { - spillDiskCap = b.spillDiskCap - } - configuredFDCap := configuredSpillFDCap(memoryCap) - if configuredFDCap > b.spillFDConfiguredCap { - configuredFDCap = b.spillFDConfiguredCap - } - effectiveFDCap := configuredFDCap - if effectiveFDCap > b.spillFDCap { - effectiveFDCap = b.spillFDCap - } + spillDiskCap = min(spillDiskCap, b.spillDiskCap) + configuredFDCap := min( + configuredSpillFDCap(memoryCap), + b.spillFDConfiguredCap, + ) + effectiveFDCap := min(configuredFDCap, b.spillFDCap) return &HashBuildBudgetGeneration{ budget: b, @@ -1025,25 +976,6 @@ func (b *HashBuildBudget) openProcessGeneration( }, nil } -// OpenGenerationWithLimits is a compatibility spelling for explicit spill caps. -func (b *HashBuildBudget) OpenGenerationWithLimits(id, memoryCap, spillDiskCap, spillFDCap uint64) (*HashBuildBudgetGeneration, error) { - return b.OpenGenerationWithSpillCaps(id, memoryCap, spillDiskCap, spillFDCap) -} - -func (b *HashBuildBudget) OpenGenerationWithCapAndSpill(id, memoryCap, spillDiskCap, spillFDCap uint64) (*HashBuildBudgetGeneration, error) { - return b.OpenGenerationWithSpillCaps(id, memoryCap, spillDiskCap, spillFDCap) -} - -// NewGeneration is an alias for OpenGeneration. -func (b *HashBuildBudget) NewGeneration(id uint64) (*HashBuildBudgetGeneration, error) { - return b.OpenGeneration(id) -} - -// OpenQueryBudget is the explicit per-query-CN spelling of OpenGeneration. -func (b *HashBuildBudget) OpenQueryBudget(id uint64) (*HashBuildQueryBudget, error) { - return b.OpenGeneration(id) -} - // ID returns the execution generation identity. func (g *HashBuildBudgetGeneration) ID() uint64 { if g == nil { @@ -1060,12 +992,6 @@ func (g *HashBuildBudgetGeneration) Cap() uint64 { return g.cap } -// QueryCap returns this generation's query-CN cap. -func (g *HashBuildBudgetGeneration) QueryCap() uint64 { return g.Cap() } - -// Capacity is a concise alias for Cap. -func (g *HashBuildBudgetGeneration) Capacity() uint64 { return g.Cap() } - // Used reports bytes reserved by this generation. func (g *HashBuildBudgetGeneration) Used() uint64 { if g == nil || g.budget == nil { @@ -1076,6 +1002,19 @@ func (g *HashBuildBudgetGeneration) Used() uint64 { return g.used } +// Peak reports the maximum physically owned bytes observed by this +// generation. It is observational only; admission and release remain owned by +// AllocationAccount-backed MPool allocations. +func (g *HashBuildBudgetGeneration) Peak() uint64 { + return g.Snapshot().PeakUsed +} + +// RejectCount reports physical allocation-capacity rejections observed by +// this generation. +func (g *HashBuildBudgetGeneration) RejectCount() uint64 { + return g.Snapshot().RejectCount +} + func (g *HashBuildBudgetGeneration) SpillDiskCap() uint64 { if g == nil || g.budget == nil { return 0 @@ -1117,24 +1056,12 @@ func (g *HashBuildBudgetGeneration) Snapshot() HashBuildBudgetGenerationSnapshot defer g.budget.mu.Unlock() return HashBuildBudgetGenerationSnapshot{ ID: g.id, Cap: g.cap, Used: g.used, PeakUsed: g.peakUsed, - AllocationUsed: g.allocationUsed, - ReserveCount: g.reserveCount, RejectCount: g.rejectCount, ReconcileCount: g.reconcileCount, ReleaseCount: g.releaseCount, + ReserveCount: g.reserveCount, RejectCount: g.rejectCount, ReconcileCount: g.reconcileCount, ReleaseCount: g.releaseCount, SpillDiskCap: g.spillDiskCap, SpillDiskUsed: g.spillDiskUsed, SpillFDCap: g.spillFDCap, SpillFDUsed: g.spillFDUsed, Closed: g.closed || g.budget.closed, } } -// Stats is an alias retained for observability call sites. -func (g *HashBuildBudgetGeneration) Stats() HashBuildBudgetGenerationSnapshot { return g.Snapshot() } -func (g *HashBuildBudgetGeneration) Peak() uint64 { return g.Snapshot().PeakUsed } -func (g *HashBuildBudgetGeneration) ReserveCount() uint64 { return g.Snapshot().ReserveCount } -func (g *HashBuildBudgetGeneration) RejectCount() uint64 { return g.Snapshot().RejectCount } -func (g *HashBuildBudgetGeneration) ReconcileCount() uint64 { return g.Snapshot().ReconcileCount } -func (g *HashBuildBudgetGeneration) ReleaseCount() uint64 { return g.Snapshot().ReleaseCount } - -// Current is a concise alias for Used. -func (g *HashBuildBudgetGeneration) Current() uint64 { return g.Used() } - // Closed reports whether this generation rejects new reservations. func (g *HashBuildBudgetGeneration) Closed() bool { if g == nil || g.budget == nil { @@ -1146,7 +1073,7 @@ func (g *HashBuildBudgetGeneration) Closed() bool { } // AllocationAccountRegistry returns the bounded CN-local registry shared by -// every activated HashBuild generation under this aggregate budget. The slot +// every HashBuild generation under this aggregate budget. The slot // bound follows a conservation fact rather than a per-operator multiplier: // every live allocation owns at least one byte and all accounts share the // aggregate byte cap, so live metadata cannot exceed aggregate capacity. A @@ -1188,15 +1115,14 @@ func (g *HashBuildBudgetGeneration) Close() { g.budget.mu.Unlock() } -// AcquireAllocationCapacity adapts allocation-accounted MPool ownership into -// the existing HashBuild query/CN policy during migration. It creates no -// independently releasable reservation token: the physical allocation lease -// is the sole release owner. +// AcquireAllocationCapacity applies the HashBuild query/CN policy to a physical +// MPool allocation. It creates no independently releasable reservation token: +// the physical allocation lease is the sole release owner. func (g *HashBuildBudgetGeneration) AcquireAllocationCapacity(size uint64) error { if size == 0 { return nil } - _, err := g.reserve(size, true) + err := g.acquireMemory(size) if err == nil { return nil } @@ -1221,11 +1147,10 @@ func (g *HashBuildBudgetGeneration) ReleaseAllocationCapacity(size uint64) { } b := g.budget b.mu.Lock() - if g.allocationUsed < size || g.used < size || b.aggregateUsed < size { + if g.used < size || b.aggregateUsed < size { b.mu.Unlock() panic("hash build allocation capacity release underflow") } - g.allocationUsed -= size g.used -= size b.aggregateUsed -= size g.releaseCount++ @@ -1234,19 +1159,12 @@ func (g *HashBuildBudgetGeneration) ReleaseAllocationCapacity(size uint64) { observeHashBuildBudget("memory", "release", "cn", size) } -// Reserve performs the required two-level sequence: charge CN aggregate, -// then charge query-CN. If query-CN rejects, aggregate is rolled back before -// returning, so callers never observe a partial reservation. -func (g *HashBuildBudgetGeneration) Reserve(size uint64) (*HashBuildReservation, error) { - return g.reserve(size, false) -} - -func (g *HashBuildBudgetGeneration) reserve( - size uint64, - allocationOwned bool, -) (*HashBuildReservation, error) { +// acquireMemory admits one physical MPool allocation. The allocation account +// is the only owner of the charge and releases it from MPool.Free; there is no +// parallel estimate/reservation token. +func (g *HashBuildBudgetGeneration) acquireMemory(size uint64) error { if g == nil || g.budget == nil { - return nil, &HashBuildBudgetError{Kind: HashBuildBudgetErrorInvalid, Message: "nil hash build generation"} + return &HashBuildBudgetError{Kind: HashBuildBudgetErrorInvalid, Message: "nil hash build generation"} } b := g.budget // A closed budget/generation has a deterministic lifecycle result and does @@ -1257,7 +1175,7 @@ func (g *HashBuildBudgetGeneration) reserve( observeHashBuildBudget("memory", "reject", "query", size) err := &HashBuildBudgetError{Kind: HashBuildBudgetErrorClosed, Requested: size, Used: g.used, Cap: g.cap} b.mu.Unlock() - return nil, err + return err } // The common cached-cap path decides whether a refresh is needed and updates @@ -1267,12 +1185,11 @@ func (g *HashBuildBudgetGeneration) reserve( if cached { if err != nil { b.mu.Unlock() - return nil, err + return err } - token, firstErr, aggregateRejected := g.reserveLocked( + firstErr, aggregateRejected := g.acquireMemoryLocked( size, false, - allocationOwned, ) b.mu.Unlock() if firstErr == nil && !aggregateRejected { @@ -1280,7 +1197,7 @@ func (g *HashBuildBudgetGeneration) reserve( observeHashBuildBudget("memory", "reserve", "cn", size) } if !aggregateRejected { - return token, firstErr + return firstErr } } else { b.mu.Unlock() @@ -1290,13 +1207,12 @@ func (g *HashBuildBudgetGeneration) reserve( var refreshed bool epoch, hasProvider, refreshed, err = b.refreshAggregateCap(false, 0) if err != nil { - return nil, err + return err } b.mu.Lock() - token, firstErr, aggregateRejected := g.reserveLocked( + firstErr, aggregateRejected := g.acquireMemoryLocked( size, false, - allocationOwned, ) b.mu.Unlock() if firstErr == nil && !aggregateRejected { @@ -1304,7 +1220,7 @@ func (g *HashBuildBudgetGeneration) reserve( observeHashBuildBudget("memory", "reserve", "cn", size) } if !aggregateRejected { - return token, firstErr + return firstErr } if refreshed { hasProvider = false @@ -1316,15 +1232,14 @@ func (g *HashBuildBudgetGeneration) reserve( // The epoch check turns concurrent retries into a single-flight operation. if hasProvider { if _, _, _, err = b.refreshAggregateCap(true, epoch); err != nil { - return nil, err + return err } } b.mu.Lock() - token, err, aggregateRejected := g.reserveLocked( + err, aggregateRejected := g.acquireMemoryLocked( size, true, - allocationOwned, ) b.mu.Unlock() if err == nil && !aggregateRejected { @@ -1332,24 +1247,24 @@ func (g *HashBuildBudgetGeneration) reserve( observeHashBuildBudget("memory", "reserve", "cn", size) } if aggregateRejected { - return nil, err + return err } - return token, err + return err } -// reserveLocked attempts one memory reservation. b.mu must be held. The bool -// result identifies an aggregate-cap failure so Reserve can trigger a forced +// acquireMemoryLocked attempts one physical-memory admission. b.mu must be +// held. The bool result identifies an aggregate-cap failure so the caller can +// trigger a forced // live-ceiling refresh without counting a transient failure as a rejection. -func (g *HashBuildBudgetGeneration) reserveLocked( +func (g *HashBuildBudgetGeneration) acquireMemoryLocked( size uint64, recordAggregateReject bool, - allocationOwned bool, -) (*HashBuildReservation, error, bool) { +) (error, bool) { b := g.budget if b.closed || g.closed { g.rejectCount++ observeHashBuildBudget("memory", "reject", "query", size) - return nil, &HashBuildBudgetError{Kind: HashBuildBudgetErrorClosed, Requested: size, Used: g.used, Cap: g.cap}, false + return &HashBuildBudgetError{Kind: HashBuildBudgetErrorClosed, Requested: size, Used: g.used, Cap: g.cap}, false } // Check by subtraction rather than used+size: this is safe for // math.MaxUint64 and rejects every overflow-sized request. @@ -1358,7 +1273,7 @@ func (g *HashBuildBudgetGeneration) reserveLocked( g.rejectCount++ observeHashBuildBudget("memory", "reject", "cn", size) } - return nil, newAdmissionError(size, b.aggregateUsed, b.aggregateCap), true + return newAdmissionError(size, b.aggregateUsed, b.aggregateCap), true } b.aggregateUsed += size if g.used > g.cap || size > g.cap-g.used { @@ -1366,150 +1281,13 @@ func (g *HashBuildBudgetGeneration) reserveLocked( b.aggregateUsed -= size g.rejectCount++ observeHashBuildBudget("memory", "reject", "query", size) - return nil, newAdmissionError(size, g.used, g.cap), false + return newAdmissionError(size, g.used, g.cap), false } g.used += size g.reserveCount++ if g.used > g.peakUsed { g.peakUsed = g.used } - if allocationOwned { - if size > math.MaxUint64-g.allocationUsed { - panic("hash build allocation capacity overflow") - } - g.allocationUsed += size - return nil, nil, false - } - return &HashBuildReservation{budget: b, generation: g, core: &hashBuildReservationCore{size: size}}, nil, false -} - -// TryReserve is a boolean convenience for admission-only call sites. -func (g *HashBuildBudgetGeneration) TryReserve(size uint64) bool { - t, err := g.Reserve(size) - if err != nil { - return false - } - // A TryReserve caller has no token to retain; immediately release it. Use - // Release rather than manually decrementing to preserve exactly-once state. - t.Release() - return true -} - -// Grow increases a live memory reservation atomically. It is used for the -// Shuffle emergency spill-scratch lease so retained copies cannot consume the -// memory required to recover from a later admission rejection. -func (r *HashBuildReservation) Grow(additional uint64) error { - if r == nil || r.core == nil || r.budget == nil || r.generation == nil { - return ErrHashBuildReservationInactive - } - if additional == 0 { - return nil - } - b := r.budget - b.mu.Lock() - if r.core.state.Load() != hashBuildReservationActive { - b.mu.Unlock() - return ErrHashBuildReservationInactive - } - if b.closed || r.generation.closed { - r.generation.rejectCount++ - observeHashBuildBudget("memory", "reject", "query", additional) - err := &HashBuildBudgetError{Kind: HashBuildBudgetErrorClosed, Requested: additional, Used: r.generation.used, Cap: r.generation.cap} - b.mu.Unlock() - return err - } - - _, epoch, hasProvider, cached, err := b.aggregateCapRefreshDecisionLocked(false, 0) - if cached { - if err != nil { - b.mu.Unlock() - return err - } - firstErr, aggregateRejected := r.growLocked(additional, false) - b.mu.Unlock() - if firstErr == nil { - observeHashBuildBudget("memory", "reserve", "query", additional) - observeHashBuildBudget("memory", "reserve", "cn", additional) - } - if !aggregateRejected { - return firstErr - } - } else { - b.mu.Unlock() - var refreshed bool - epoch, hasProvider, refreshed, err = b.refreshAggregateCap(false, 0) - if err != nil { - return err - } - b.mu.Lock() - firstErr, aggregateRejected := r.growLocked(additional, false) - b.mu.Unlock() - if firstErr == nil { - observeHashBuildBudget("memory", "reserve", "query", additional) - observeHashBuildBudget("memory", "reserve", "cn", additional) - } - if !aggregateRejected { - return firstErr - } - if refreshed { - hasProvider = false - } - } - - if hasProvider { - if _, _, _, err = b.refreshAggregateCap(true, epoch); err != nil { - return err - } - } - b.mu.Lock() - err, aggregateRejected := r.growLocked(additional, true) - b.mu.Unlock() - if err == nil { - observeHashBuildBudget("memory", "reserve", "query", additional) - observeHashBuildBudget("memory", "reserve", "cn", additional) - } - if aggregateRejected { - return err - } - return err -} - -// growLocked attempts one memory reservation growth. b.mu must be held. -func (r *HashBuildReservation) growLocked(additional uint64, recordAggregateReject bool) (error, bool) { - b := r.budget - g := r.generation - if r.core.state.Load() != hashBuildReservationActive { - return ErrHashBuildReservationInactive, false - } - if b.closed || g.closed { - g.rejectCount++ - observeHashBuildBudget("memory", "reject", "query", additional) - return &HashBuildBudgetError{Kind: HashBuildBudgetErrorClosed, Requested: additional, Used: g.used, Cap: g.cap}, false - } - if b.aggregateUsed > b.aggregateCap || additional > b.aggregateCap-b.aggregateUsed { - // Defer the counter/metric until the caller knows whether a forced - // refresh can make this transient failure admissible. - if recordAggregateReject { - g.rejectCount++ - observeHashBuildBudget("memory", "reject", "cn", additional) - } - return newAdmissionError(additional, b.aggregateUsed, b.aggregateCap), true - } - if g.used > g.cap || additional > g.cap-g.used { - g.rejectCount++ - observeHashBuildBudget("memory", "reject", "query", additional) - return newAdmissionError(additional, g.used, g.cap), false - } - if r.core.size > math.MaxUint64-additional { - return &HashBuildBudgetError{Kind: HashBuildBudgetErrorInvalid, Requested: additional, Message: "hash build reservation size overflow"}, false - } - b.aggregateUsed += additional - g.used += additional - r.core.size += additional - if g.used > g.peakUsed { - g.peakUsed = g.used - } - g.reserveCount++ return nil, false } @@ -1538,168 +1316,30 @@ func newComponentAdmissionError( } } -// HashBuildReservation is an exactly-once ownership token for one charge in -// both the CN aggregate and its generation. State transitions are atomic: -// active -> released or active -> transferred. A late release therefore -// always affects the original generation and can never decrement a newer one. -type HashBuildReservation struct { - budget *HashBuildBudget - generation *HashBuildBudgetGeneration - // core is shared by accidental token copies, keeping mutable charge and - // exactly-once state together under the budget mutex. - core *hashBuildReservationCore -} - -type hashBuildReservationCore struct { +type hashBuildSpillReservationCore struct { size uint64 state atomic.Uint32 } const ( - hashBuildReservationActive uint32 = iota - hashBuildReservationReleased - hashBuildReservationTransferred + hashBuildSpillReservationActive uint32 = iota + hashBuildSpillReservationReleased ) -// Size returns the reservation's current reconciled charge. -func (r *HashBuildReservation) Size() uint64 { - if r == nil || r.budget == nil || r.core == nil { - return 0 - } - r.budget.mu.Lock() - defer r.budget.mu.Unlock() - return r.core.size -} - -// GenerationID returns the generation charged by this token. -func (r *HashBuildReservation) GenerationID() uint64 { - if r == nil || r.generation == nil { - return 0 - } - return r.generation.id -} - -// Released reports whether this token has relinquished its ownership. A -// transferred token is not released, but no longer owns the charge. -func (r *HashBuildReservation) Released() bool { - if r == nil || r.core == nil { - return true - } - if r.budget == nil { - return r.core.state.Load() != hashBuildReservationActive - } - r.budget.mu.Lock() - defer r.budget.mu.Unlock() - return r.core.state.Load() != hashBuildReservationActive -} - -// Release relinquishes this token once. It returns true only for the caller -// that won the active -> released transition. -func (r *HashBuildReservation) Release() bool { - if r == nil || r.core == nil || r.budget == nil || r.generation == nil { - return false - } - r.budget.mu.Lock() - if !r.core.state.CompareAndSwap(hashBuildReservationActive, hashBuildReservationReleased) { - r.budget.mu.Unlock() - return false - } - size := r.core.size - // The subtraction is exact for a live token. Keep a defensive branch so - // corrupted state cannot underflow and turn into an apparent huge charge. - if r.generation.used >= size { - r.generation.used -= size - } else { - r.generation.used = 0 - } - if r.budget.aggregateUsed >= size { - r.budget.aggregateUsed -= size - } else { - r.budget.aggregateUsed = 0 - } - r.generation.releaseCount++ - r.budget.mu.Unlock() - observeHashBuildBudget("memory", "release", "query", size) - observeHashBuildBudget("memory", "release", "cn", size) - return true -} - -// ReconcileDown shrinks a live charge to actual bytes. It is linearized with -// reserve/release/transfer under the owning budget mutex. Upward reconciliation -// is rejected and inactive tokens never mutate counters. -func (r *HashBuildReservation) ReconcileDown(actual uint64) (bool, error) { - if r == nil || r.core == nil || r.budget == nil || r.generation == nil { - return false, ErrHashBuildReservationInactive - } - r.budget.mu.Lock() - if r.core.state.Load() != hashBuildReservationActive { - r.budget.mu.Unlock() - return false, ErrHashBuildReservationInactive - } - if actual > r.core.size { - r.budget.mu.Unlock() - return false, ErrHashBuildReservationUpward - } - delta := r.core.size - actual - if delta > 0 { - if r.generation.used < delta || r.budget.aggregateUsed < delta { - r.budget.mu.Unlock() - return false, ErrHashBuildReservationInactive - } - r.generation.used -= delta - r.budget.aggregateUsed -= delta - r.core.size = actual - } - r.generation.reconcileCount++ - r.budget.mu.Unlock() - if delta > 0 { - observeHashBuildBudget("memory", "reconcile", "query", delta) - observeHashBuildBudget("memory", "reconcile", "cn", delta) - } - return true, nil -} - -// Reconcile is a compatibility alias. -func (r *HashBuildReservation) Reconcile(actual uint64) (bool, error) { return r.ReconcileDown(actual) } - -// Transfer moves ownership to a fresh token exactly once. The original token -// becomes inert; releasing it after a successful transfer cannot decrement the -// budget. If Release wins the race, Transfer returns nil. -func (r *HashBuildReservation) Transfer() *HashBuildReservation { - if r == nil || r.core == nil || r.budget == nil { - return nil - } - r.budget.mu.Lock() - defer r.budget.mu.Unlock() - if !r.core.state.CompareAndSwap(hashBuildReservationActive, hashBuildReservationTransferred) { - return nil - } - return &HashBuildReservation{budget: r.budget, generation: r.generation, core: &hashBuildReservationCore{size: r.core.size}} -} - -// TransferOwnership is a descriptive alias for Transfer. -func (r *HashBuildReservation) TransferOwnership() *HashBuildReservation { return r.Transfer() } - -// TransferTo is another descriptive spelling for ownership transfer. -func (r *HashBuildReservation) TransferTo() *HashBuildReservation { return r.Transfer() } - // HashBuildSpillDiskReservation owns query and CN spill-disk bytes. type HashBuildSpillDiskReservation struct { budget *HashBuildBudget generation *HashBuildBudgetGeneration - core *hashBuildReservationCore + core *hashBuildSpillReservationCore } // HashBuildSpillFDReservation owns query and CN spill file descriptors. type HashBuildSpillFDReservation struct { budget *HashBuildBudget generation *HashBuildBudgetGeneration - core *hashBuildReservationCore + core *hashBuildSpillReservationCore } -type SpillDiskReservation = HashBuildSpillDiskReservation -type SpillFDReservation = HashBuildSpillFDReservation - func (r *HashBuildSpillDiskReservation) Size() uint64 { if r == nil || r.budget == nil || r.core == nil { return 0 @@ -1716,13 +1356,6 @@ func (r *HashBuildSpillFDReservation) Size() uint64 { defer r.budget.mu.Unlock() return r.core.size } -func (r *HashBuildSpillDiskReservation) Released() bool { - return r == nil || r.core == nil || r.core.state.Load() != hashBuildReservationActive -} -func (r *HashBuildSpillFDReservation) Released() bool { - return r == nil || r.core == nil || r.core.state.Load() != hashBuildReservationActive -} - func (g *HashBuildBudgetGeneration) ReserveSpillDisk(size uint64) (*HashBuildSpillDiskReservation, error) { if g == nil || g.budget == nil { return nil, &HashBuildBudgetError{Kind: HashBuildBudgetErrorInvalid} @@ -1749,11 +1382,7 @@ func (g *HashBuildBudgetGeneration) ReserveSpillDisk(size uint64) (*HashBuildSpi g.spillDiskUsed += size observeHashBuildBudget("spill_disk", "reserve", "query", size) observeHashBuildBudget("spill_disk", "reserve", "cn", size) - return &HashBuildSpillDiskReservation{budget: b, generation: g, core: &hashBuildReservationCore{size: size}}, nil -} - -func (g *HashBuildBudgetGeneration) ReserveSpillDiskBytes(size uint64) (*HashBuildSpillDiskReservation, error) { - return g.ReserveSpillDisk(size) + return &HashBuildSpillDiskReservation{budget: b, generation: g, core: &hashBuildSpillReservationCore{size: size}}, nil } // Grow increases one live per-file disk reservation without allocating a new @@ -1761,7 +1390,7 @@ func (g *HashBuildBudgetGeneration) ReserveSpillDiskBytes(size uint64) (*HashBui // rather than to the number of tiny batch records written to those files. func (r *HashBuildSpillDiskReservation) Grow(additional uint64) error { if r == nil || r.core == nil || r.budget == nil || r.generation == nil { - return ErrHashBuildReservationInactive + return ErrHashBuildSpillReservationInactive } if additional == 0 { return nil @@ -1770,8 +1399,8 @@ func (r *HashBuildSpillDiskReservation) Grow(additional uint64) error { g := r.generation b.mu.Lock() defer b.mu.Unlock() - if r.core.state.Load() != hashBuildReservationActive { - return ErrHashBuildReservationInactive + if r.core.state.Load() != hashBuildSpillReservationActive { + return ErrHashBuildSpillReservationInactive } if b.closed || g.closed { g.rejectCount++ @@ -1835,11 +1464,7 @@ func (g *HashBuildBudgetGeneration) ReserveSpillFD(size uint64) (*HashBuildSpill g.spillFDUsed += size observeHashBuildBudget("spill_fd", "reserve", "query", size) observeHashBuildBudget("spill_fd", "reserve", "cn", size) - return &HashBuildSpillFDReservation{budget: b, generation: g, core: &hashBuildReservationCore{size: size}}, nil -} - -func (g *HashBuildBudgetGeneration) ReserveSpillFileDescriptors(size uint64) (*HashBuildSpillFDReservation, error) { - return g.ReserveSpillFD(size) + return &HashBuildSpillFDReservation{budget: b, generation: g, core: &hashBuildSpillReservationCore{size: size}}, nil } func (r *HashBuildSpillDiskReservation) Release() bool { @@ -1848,19 +1473,15 @@ func (r *HashBuildSpillDiskReservation) Release() bool { } r.budget.mu.Lock() defer r.budget.mu.Unlock() - if !r.core.state.CompareAndSwap(hashBuildReservationActive, hashBuildReservationReleased) { + if !r.core.state.CompareAndSwap(hashBuildSpillReservationActive, hashBuildSpillReservationReleased) { return false } - if r.generation.spillDiskUsed >= r.core.size { - r.generation.spillDiskUsed -= r.core.size - } else { - r.generation.spillDiskUsed = 0 - } - if r.budget.spillDiskUsed >= r.core.size { - r.budget.spillDiskUsed -= r.core.size - } else { - r.budget.spillDiskUsed = 0 + if r.generation.spillDiskUsed < r.core.size || + r.budget.spillDiskUsed < r.core.size { + panic("hash build spill disk reservation release underflow") } + r.generation.spillDiskUsed -= r.core.size + r.budget.spillDiskUsed -= r.core.size observeHashBuildBudget("spill_disk", "release", "query", r.core.size) observeHashBuildBudget("spill_disk", "release", "cn", r.core.size) r.generation.releaseCount++ @@ -1873,19 +1494,15 @@ func (r *HashBuildSpillFDReservation) Release() bool { } r.budget.mu.Lock() defer r.budget.mu.Unlock() - if !r.core.state.CompareAndSwap(hashBuildReservationActive, hashBuildReservationReleased) { + if !r.core.state.CompareAndSwap(hashBuildSpillReservationActive, hashBuildSpillReservationReleased) { return false } - if r.generation.spillFDUsed >= r.core.size { - r.generation.spillFDUsed -= r.core.size - } else { - r.generation.spillFDUsed = 0 - } - if r.budget.spillFDUsed >= r.core.size { - r.budget.spillFDUsed -= r.core.size - } else { - r.budget.spillFDUsed = 0 + if r.generation.spillFDUsed < r.core.size || + r.budget.spillFDUsed < r.core.size { + panic("hash build spill fd reservation release underflow") } + r.generation.spillFDUsed -= r.core.size + r.budget.spillFDUsed -= r.core.size observeHashBuildBudget("spill_fd", "release", "query", r.core.size) observeHashBuildBudget("spill_fd", "release", "cn", r.core.size) r.generation.releaseCount++ @@ -1894,20 +1511,20 @@ func (r *HashBuildSpillFDReservation) Release() bool { func (r *HashBuildSpillDiskReservation) ReconcileDown(actual uint64) (bool, error) { if r == nil || r.core == nil || r.budget == nil || r.generation == nil { - return false, ErrHashBuildReservationInactive + return false, ErrHashBuildSpillReservationInactive } r.budget.mu.Lock() defer r.budget.mu.Unlock() - if r.core.state.Load() != hashBuildReservationActive { - return false, ErrHashBuildReservationInactive + if r.core.state.Load() != hashBuildSpillReservationActive { + return false, ErrHashBuildSpillReservationInactive } if actual > r.core.size { - return false, ErrHashBuildReservationUpward + return false, ErrHashBuildSpillReservationUpward } delta := r.core.size - actual if delta > 0 { if r.generation.spillDiskUsed < delta || r.budget.spillDiskUsed < delta { - return false, ErrHashBuildReservationInactive + return false, ErrHashBuildSpillReservationInactive } r.generation.spillDiskUsed -= delta r.budget.spillDiskUsed -= delta @@ -1918,71 +1535,6 @@ func (r *HashBuildSpillDiskReservation) ReconcileDown(actual uint64) (bool, erro r.generation.reconcileCount++ return true, nil } -func (r *HashBuildSpillFDReservation) ReconcileDown(actual uint64) (bool, error) { - if r == nil || r.core == nil || r.budget == nil || r.generation == nil { - return false, ErrHashBuildReservationInactive - } - r.budget.mu.Lock() - defer r.budget.mu.Unlock() - if r.core.state.Load() != hashBuildReservationActive { - return false, ErrHashBuildReservationInactive - } - if actual > r.core.size { - return false, ErrHashBuildReservationUpward - } - delta := r.core.size - actual - if delta > 0 { - if r.generation.spillFDUsed < delta || r.budget.spillFDUsed < delta { - return false, ErrHashBuildReservationInactive - } - r.generation.spillFDUsed -= delta - r.budget.spillFDUsed -= delta - r.core.size = actual - observeHashBuildBudget("spill_fd", "reconcile", "query", delta) - observeHashBuildBudget("spill_fd", "reconcile", "cn", delta) - } - r.generation.reconcileCount++ - return true, nil -} -func (r *HashBuildSpillDiskReservation) Reconcile(actual uint64) (bool, error) { - return r.ReconcileDown(actual) -} -func (r *HashBuildSpillFDReservation) Reconcile(actual uint64) (bool, error) { - return r.ReconcileDown(actual) -} - -func (r *HashBuildSpillDiskReservation) Transfer() *HashBuildSpillDiskReservation { - if r == nil || r.core == nil || r.budget == nil { - return nil - } - r.budget.mu.Lock() - defer r.budget.mu.Unlock() - if !r.core.state.CompareAndSwap(hashBuildReservationActive, hashBuildReservationTransferred) { - return nil - } - return &HashBuildSpillDiskReservation{budget: r.budget, generation: r.generation, core: &hashBuildReservationCore{size: r.core.size}} -} -func (r *HashBuildSpillFDReservation) Transfer() *HashBuildSpillFDReservation { - if r == nil || r.core == nil || r.budget == nil { - return nil - } - r.budget.mu.Lock() - defer r.budget.mu.Unlock() - if !r.core.state.CompareAndSwap(hashBuildReservationActive, hashBuildReservationTransferred) { - return nil - } - return &HashBuildSpillFDReservation{budget: r.budget, generation: r.generation, core: &hashBuildReservationCore{size: r.core.size}} -} -func (r *HashBuildSpillDiskReservation) TransferOwnership() *HashBuildSpillDiskReservation { - return r.Transfer() -} -func (r *HashBuildSpillFDReservation) TransferOwnership() *HashBuildSpillFDReservation { - return r.Transfer() -} -func (r *HashBuildSpillDiskReservation) TransferTo() *HashBuildSpillDiskReservation { - return r.Transfer() -} -func (r *HashBuildSpillFDReservation) TransferTo() *HashBuildSpillFDReservation { return r.Transfer() } // HashBuildCeilingInputs are the finite resource sources used by // ResolveHashBuildCeiling. A zero or math.MaxUint64 source means unavailable @@ -2066,18 +1618,6 @@ func ResolveHashBuildCeiling(in HashBuildCeilingInputs) (HashBuildCeiling, error }, nil } -// ResolveHashBuildBudget is a semantic alias used by budget initialization -// callers. -func ResolveHashBuildBudget(in HashBuildCeilingInputs) (HashBuildCeiling, error) { - return ResolveHashBuildCeiling(in) -} - -// NewHashBuildBudgetFromCeiling wires a resolved ceiling into the local-CN -// aggregate/generation budget. -func NewHashBuildBudgetFromCeiling(ceiling HashBuildCeiling) (*HashBuildBudget, error) { - return NewHashBuildBudget(ceiling.CNHashCap, ceiling.QueryCap) -} - // GetHashBuildBudget returns the statement generation shared by every child // process in this BaseProcess. Different top-level processes on the same CN // charge a shared aggregate budget. diff --git a/pkg/vm/process/hashbuild_budget_test.go b/pkg/vm/process/hashbuild_budget_test.go index 3d250e54e9f79..2a657433c9cc4 100644 --- a/pkg/vm/process/hashbuild_budget_test.go +++ b/pkg/vm/process/hashbuild_budget_test.go @@ -17,50 +17,63 @@ package process import ( "errors" "math" - "os" - "os/exec" - "runtime" - "sort" - "strconv" "sync" "sync/atomic" "testing" "time" - commonmpool "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/fileservice" ) -func TestHashBuildBudgetExactLimitAndOverflow(t *testing.T) { - b, err := NewHashBuildBudget(math.MaxUint64, math.MaxUint64) - if err != nil { - t.Fatal(err) +// testPhysicalAllocation exercises the capacity controller at the same +// acquire/free boundary used by MPool without exposing a second production +// reservation API. +type testPhysicalAllocation struct { + generation *HashBuildBudgetGeneration + size uint64 + released atomic.Bool +} + +func acquireTestPhysicalAllocation( + generation *HashBuildBudgetGeneration, + size uint64, +) (*testPhysicalAllocation, error) { + if err := generation.AcquireAllocationCapacity(size); err != nil { + return nil, err } - g, err := b.OpenGeneration(1) + return &testPhysicalAllocation{generation: generation, size: size}, nil +} + +func (a *testPhysicalAllocation) Release() bool { + if a == nil || a.generation == nil || !a.released.CompareAndSwap(false, true) { + return false + } + a.generation.ReleaseAllocationCapacity(a.size) + return true +} + +func TestHashBuildBudgetPhysicalAllocationLimit(t *testing.T) { + budget := MustNewHashBuildBudget(math.MaxUint64, math.MaxUint64) + generation, err := budget.OpenGeneration(1) if err != nil { t.Fatal(err) } - tok, err := g.Reserve(math.MaxUint64) + allocation, err := acquireTestPhysicalAllocation(generation, math.MaxUint64) if err != nil { t.Fatalf("exact limit rejected: %v", err) } - if got := b.AggregateUsed(); got != math.MaxUint64 { - t.Fatalf("aggregate used = %d, want max uint64", got) - } - if _, err = g.Reserve(1); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("limit+1 error = %v, want admission rejection", err) - } - if got := b.AggregateUsed(); got != math.MaxUint64 { - t.Fatalf("failed query reservation changed aggregate: %d", got) + if _, err = acquireTestPhysicalAllocation(generation, 1); !errors.Is(err, ErrHashBuildBudgetAdmission) { + t.Fatalf("limit+1 error = %v", err) } - if got := g.Used(); got != math.MaxUint64 { - t.Fatalf("failed query reservation changed generation: %d", got) + if budget.AggregateUsed() != math.MaxUint64 || generation.Used() != math.MaxUint64 { + t.Fatal("failed admission changed the physical allocation ledger") } - if !tok.Release() || tok.Release() { - t.Fatal("release must transition exactly once") + if !allocation.Release() || allocation.Release() { + t.Fatal("physical allocation must release exactly once") } - if b.AggregateUsed() != 0 || g.Used() != 0 { - t.Fatalf("released reservation remains: cn=%d query=%d", b.AggregateUsed(), g.Used()) + if budget.AggregateUsed() != 0 || generation.Used() != 0 { + t.Fatal("physical allocation release leaked capacity") } } @@ -81,8 +94,7 @@ func TestHashBuildBudgetAdmissionIdentifiesResource(t *testing.T) { name: "memory", want: HashBuildBudgetComponentMemory, call: func() error { - _, reserveErr := g.Reserve(11) - return reserveErr + return g.AcquireAllocationCapacity(11) }, }, { @@ -116,109 +128,45 @@ func TestHashBuildBudgetAdmissionIdentifiesResource(t *testing.T) { } } -func TestHashBuildBudgetAllocationAccountAdapter(t *testing.T) { +func TestHashBuildBudgetAllocationAccountIsSoleOwner(t *testing.T) { budget := MustNewHashBuildBudget(10, 10) generation, err := budget.OpenGeneration(1) if err != nil { t.Fatal(err) } - legacy, err := generation.Reserve(4) - if err != nil { - t.Fatal(err) - } - - registry, err := commonmpool.NewAllocationAccountRegistry(1, 2) - if err != nil { - t.Fatal(err) - } - account, err := registry.OpenWithController(10, generation) - if err != nil { - t.Fatal(err) - } - mp := commonmpool.MustNew("hash-build-allocation-account-adapter") - defer commonmpool.DeleteMPool(mp) - - noMetadataRegistry, err := commonmpool.NewAllocationAccountRegistry(1, 0) + registry, err := mpool.NewAllocationAccountRegistry(1, 2) if err != nil { t.Fatal(err) } - noMetadataAccount, err := noMetadataRegistry.OpenWithController( - 10, - generation, - ) + account, err := registry.OpenWithController(11, generation) if err != nil { t.Fatal(err) } - if _, err = mp.AllocAccounted( - 1, - noMetadataAccount, - 1, - 1, - ); !errors.Is(err, commonmpool.ErrAllocationMetadataSlots) { - t.Fatalf("metadata admission error = %v", err) - } - if generation.Used() != 4 || - generation.Snapshot().AllocationUsed != 0 { - t.Fatalf("metadata failure leaked controller capacity: %+v", - generation.Snapshot()) - } - noMetadataAccount.Seal() - if _, err = noMetadataRegistry.Finalize(noMetadataAccount); err != nil { - t.Fatal(err) - } + mp := mpool.MustNewZero() - buffer, err := mp.AllocAccounted(6, account, 1, 1) + buffer, err := mp.AllocAccounted(10, account, 1, 1) if err != nil { t.Fatal(err) } - snapshot := generation.Snapshot() - if snapshot.Used != 10 || snapshot.AllocationUsed != 6 { - t.Fatalf("unexpected generation snapshot: %+v", snapshot) + if generation.Used() != 10 || account.Snapshot().Used != 10 { + t.Fatal("physical allocation was not charged exactly once") } - if account.Snapshot().Used != 6 { - t.Fatalf("account used = %d, want 6", account.Snapshot().Used) + if _, err = mp.AllocAccounted(1, account, 1, 1); !errors.Is(err, ErrHashBuildBudgetAdmission) || + !errors.Is(err, mpool.ErrAllocationAccountCapacity) { + t.Fatalf("capacity error = %v", err) } - - if _, err = mp.AllocAccounted(1, account, 1, 1); !errors.Is( - err, - ErrHashBuildBudgetAdmission, - ) { - t.Fatalf("combined legacy/exact admission error = %v", err) - } - if !errors.Is(err, commonmpool.ErrAllocationAccountCapacity) || - commonmpool.AllocationFailureReasonOf(err) != - commonmpool.AllocationFailureCapacity { - t.Fatalf("adapter did not type policy pressure as capacity: %v", err) - } - if account.Snapshot().Used != 6 || + if generation.Used() != 10 || account.Snapshot().Used != 10 || registry.LiveAllocationMetadata() != 1 { - t.Fatal("failed adapter admission did not roll back") + t.Fatal("failed allocation changed account state") } generation.Close() - if _, err = mp.AllocAccounted(1, account, 1, 1); !errors.Is( - err, - ErrHashBuildBudgetClosed, - ) { - t.Fatalf("closed generation admission error = %v", err) - } - if !errors.Is(err, commonmpool.ErrAllocationAccountSealed) || - commonmpool.IsRetryableAllocationCapacity(err) { - t.Fatalf("closed adapter error entered capacity retry: %v", err) - } - if account.Snapshot().Used != 6 || - registry.LiveAllocationMetadata() != 1 { - t.Fatal("closed adapter admission did not roll back") + if _, err = mp.AllocAccounted(1, account, 1, 1); !errors.Is(err, mpool.ErrAllocationAccountSealed) { + t.Fatalf("closed generation error = %v", err) } - - // Close rejects new capacity but cannot invalidate a live physical lease. mp.Free(buffer) - if generation.Used() != 4 || generation.Snapshot().AllocationUsed != 0 { - t.Fatalf("allocation release did not retain only legacy charge: %+v", - generation.Snapshot()) - } - if !legacy.Release() || generation.Used() != 0 { - t.Fatal("legacy reservation did not release") + if generation.Used() != 0 || account.Snapshot().Used != 0 { + t.Fatal("MPool.Free did not release the sole charge") } account.Seal() if _, err = registry.Finalize(account); err != nil { @@ -226,7 +174,7 @@ func TestHashBuildBudgetAllocationAccountAdapter(t *testing.T) { } } -func TestHashBuildAllocationAccountRegistryUsesByteConservationBound(t *testing.T) { +func TestHashBuildAllocationAccountRegistryBounds(t *testing.T) { budget := MustNewHashBuildBudget(16<<10, 16<<10) first, err := budget.OpenGeneration(1) if err != nil { @@ -236,2024 +184,347 @@ func TestHashBuildAllocationAccountRegistryUsesByteConservationBound(t *testing. if err != nil { t.Fatal(err) } - if registry.GenerationCapacity() != hashBuildAllocationGenerationSlots { - t.Fatalf("generation slots = %d", registry.GenerationCapacity()) - } - if registry.MaxAllocationMetadata() != 16<<10 { - t.Fatalf( - "allocation slots = %d, want %d", - registry.MaxAllocationMetadata(), - uint64(16<<10), - ) + if registry.GenerationCapacity() != hashBuildAllocationGenerationSlots || + registry.MaxAllocationMetadata() != 16<<10 { + t.Fatal("allocation registry does not follow byte-conservation bounds") } second, err := budget.OpenGeneration(2) if err != nil { t.Fatal(err) } - secondRegistry, err := second.AllocationAccountRegistry() - if err != nil { - t.Fatal(err) - } - if secondRegistry != registry { + shared, err := second.AllocationAccountRegistry() + if err != nil || shared != registry { t.Fatal("one CN budget created multiple allocation registries") } -} -func TestHashBuildAllocationAccountRegistryCapsMetadataHeadroom(t *testing.T) { - budget := MustNewHashBuildBudget( - hashBuildAllocationMetadataMaxSlots+1, - hashBuildAllocationMetadataMaxSlots+1, - ) - generation, err := budget.OpenGeneration(1) + large := MustNewHashBuildBudget(hashBuildAllocationMetadataMaxSlots+1, hashBuildAllocationMetadataMaxSlots+1) + largeGeneration, err := large.OpenGeneration(1) if err != nil { t.Fatal(err) } - registry, err := generation.AllocationAccountRegistry() + largeRegistry, err := largeGeneration.AllocationAccountRegistry() if err != nil { t.Fatal(err) } - if got := registry.MaxAllocationMetadata(); got != hashBuildAllocationMetadataMaxSlots { - t.Fatalf("allocation slots = %d, want %d", got, hashBuildAllocationMetadataMaxSlots) + if largeRegistry.MaxAllocationMetadata() != hashBuildAllocationMetadataMaxSlots { + t.Fatal("allocation metadata exceeded its fixed headroom") } } func TestHashBuildBudgetQueryRejectRollsBackCN(t *testing.T) { - b := MustNewHashBuildBudget(10, 4) - g1, _ := b.OpenGeneration(1) - g2, _ := b.OpenGeneration(2) - first, err := g1.Reserve(4) + budget := MustNewHashBuildBudget(10, 4) + first, _ := budget.OpenGeneration(1) + second, _ := budget.OpenGeneration(2) + allocation, err := acquireTestPhysicalAllocation(first, 4) if err != nil { t.Fatal(err) } - if _, err = g2.Reserve(7); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("error = %v, want admission rejection", err) + if _, err = acquireTestPhysicalAllocation(second, 7); !errors.Is(err, ErrHashBuildBudgetAdmission) { + t.Fatalf("query rejection error = %v", err) } - if b.AggregateUsed() != 4 || g2.Used() != 0 { - t.Fatalf("query rejection did not roll back CN: cn=%d g2=%d", b.AggregateUsed(), g2.Used()) + if budget.AggregateUsed() != 4 || second.Used() != 0 { + t.Fatal("query rejection did not roll back the CN charge") } - first.Release() + allocation.Release() } -func TestHashBuildBudgetConcurrentReserveRelease(t *testing.T) { - const workers = 32 - b := MustNewHashBuildBudget(workers, workers) - gens := make([]*HashBuildBudgetGeneration, workers) - for i := range gens { - gens[i], _ = b.OpenGeneration(uint64(i + 1)) - } - start := make(chan struct{}) - acquired := make(chan *HashBuildReservation, workers) +func TestHashBuildBudgetConcurrentPhysicalAllocations(t *testing.T) { + const workers = 64 + budget := MustNewHashBuildBudget(workers, workers) + allocations := make(chan *testPhysicalAllocation, workers) var wg sync.WaitGroup - for i := range gens { + for i := 0; i < workers; i++ { + generation, err := budget.OpenGeneration(uint64(i + 1)) + if err != nil { + t.Fatal(err) + } wg.Add(1) - go func(g *HashBuildBudgetGeneration) { + go func() { defer wg.Done() - <-start - tok, err := g.Reserve(1) - if err == nil { - acquired <- tok + allocation, acquireErr := acquireTestPhysicalAllocation(generation, 1) + if acquireErr != nil { + t.Errorf("acquire: %v", acquireErr) + return } - }(gens[i]) + allocations <- allocation + }() } - close(start) wg.Wait() - if len(acquired) != workers { - t.Fatalf("acquired %d reservations, want %d", len(acquired), workers) - } - if b.AggregateUsed() != workers { - t.Fatalf("aggregate used = %d, want %d", b.AggregateUsed(), workers) + close(allocations) + if budget.AggregateUsed() != workers { + t.Fatalf("aggregate used = %d", budget.AggregateUsed()) } - for i := 0; i < workers; i++ { - (<-acquired).Release() + for allocation := range allocations { + allocation.Release() } - if b.AggregateUsed() != 0 { - t.Fatalf("aggregate used after release = %d", b.AggregateUsed()) + if budget.AggregateUsed() != 0 { + t.Fatal("concurrent physical allocations leaked") } } -func TestHashBuildBudgetTransferAndClose(t *testing.T) { - b := MustNewHashBuildBudget(8, 8) - g, _ := b.OpenGeneration(7) - tok, err := g.Reserve(3) - if err != nil { - t.Fatal(err) - } - moved := tok.Transfer() - if moved == nil || tok.Release() { - t.Fatal("transfer must make original token inert") - } - if b.AggregateUsed() != 3 { - t.Fatalf("transfer changed charge: %d", b.AggregateUsed()) - } - g.Close() - if _, err = g.Reserve(1); !errors.Is(err, ErrHashBuildBudgetClosed) { +func TestHashBuildBudgetGenerationIsolationAndClose(t *testing.T) { + budget := MustNewHashBuildBudget(8, 8) + oldGeneration, _ := budget.OpenGeneration(1) + oldAllocation, _ := acquireTestPhysicalAllocation(oldGeneration, 6) + oldGeneration.Close() + if _, err := acquireTestPhysicalAllocation(oldGeneration, 1); !errors.Is(err, ErrHashBuildBudgetClosed) { t.Fatalf("closed generation error = %v", err) } - if !moved.Release() || moved.Release() { - t.Fatal("transferred token release must be exactly once") - } - if b.AggregateUsed() != 0 { - t.Fatalf("live token release after close leaked: %d", b.AggregateUsed()) - } - b.Close() - if _, err = b.OpenGeneration(8); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed budget error = %v", err) - } -} - -func TestHashBuildBudgetGenerationIsolation(t *testing.T) { - b := MustNewHashBuildBudget(8, 8) - old, _ := b.OpenGeneration(1) - oldToken, _ := old.Reserve(6) - old.Close() - newGeneration, _ := b.OpenGeneration(2) - newToken, err := newGeneration.Reserve(2) - if err != nil { - t.Fatal(err) - } - oldToken.Release() - if newGeneration.Used() != 2 || b.AggregateUsed() != 2 { - t.Fatalf("old release affected new generation: new=%d aggregate=%d", newGeneration.Used(), b.AggregateUsed()) - } - newToken.Release() -} - -func TestHashBuildBudgetCapReductionFailsClosedUntilRelease(t *testing.T) { - b := MustNewHashBuildBudget(10, 10) - g1, _ := b.OpenGenerationWithCap(1, 10) - owned, err := g1.Reserve(8) - if err != nil { - t.Fatal(err) - } - if err = b.UpdateAggregateCap(6); err != nil { - t.Fatal(err) - } - g2, _ := b.OpenGenerationWithCap(2, 6) - if _, err = g2.Reserve(1); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("cap shrink did not fail closed: %v", err) - } - owned.Release() - newToken, err := g2.Reserve(6) - if err != nil { - t.Fatalf("reservation after release failed: %v", err) - } - newToken.Release() -} - -func TestHashBuildReservationReconcileCopyAlias(t *testing.T) { - b := MustNewHashBuildBudget(20, 20) - g, _ := b.OpenGeneration(1) - tok, err := g.Reserve(10) + newGeneration, _ := budget.OpenGeneration(2) + newAllocation, err := acquireTestPhysicalAllocation(newGeneration, 2) if err != nil { t.Fatal(err) } - alias := *tok - if err := tok.Grow(2); err != nil { - t.Fatalf("grow: %v", err) - } - if tok.Size() != 12 || alias.Size() != 12 || g.Used() != 12 { - t.Fatalf("alias grow diverged: token=%d alias=%d used=%d", tok.Size(), alias.Size(), g.Used()) - } - if ok, err := alias.ReconcileDown(4); !ok || err != nil { - t.Fatalf("reconcile: ok=%v err=%v", ok, err) - } - if tok.Size() != 4 || g.Used() != 4 || b.AggregateUsed() != 4 { - t.Fatalf("alias reconcile diverged: size=%d gen=%d cn=%d", tok.Size(), g.Used(), b.AggregateUsed()) + oldAllocation.Release() + if newGeneration.Used() != 2 || budget.AggregateUsed() != 2 { + t.Fatal("old generation release affected the new generation") } - if _, err := tok.ReconcileDown(5); !errors.Is(err, ErrHashBuildReservationUpward) { - t.Fatalf("upward err=%v", err) - } - if !tok.Release() || alias.Release() { - t.Fatal("copy aliases must release exactly once") + newAllocation.Release() + budget.Close() + if _, err = budget.OpenGeneration(3); !errors.Is(err, ErrHashBuildBudgetClosed) { + t.Fatalf("closed budget error = %v", err) } } -func TestHashBuildReservationGrowRejectsWithoutChangingCharge(t *testing.T) { - b := MustNewHashBuildBudget(10, 10) - g, _ := b.OpenGeneration(1) - tok, err := g.Reserve(8) +func TestHashBuildBudgetLiveCapRefresh(t *testing.T) { + budget := MustNewHashBuildBudget(10, 10) + var capValue atomic.Uint64 + capValue.Store(10) + var calls atomic.Uint64 + budget.SetAggregateCapProvider(func() (uint64, error) { + calls.Add(1) + return capValue.Load(), nil + }) + budget.capRefreshTTL = time.Hour + generation, _ := budget.OpenGeneration(1) + first, err := acquireTestPhysicalAllocation(generation, 8) if err != nil { t.Fatal(err) } - before := g.Snapshot() - if err = tok.Grow(3); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("grow rejection=%v", err) + if calls.Load() != 1 { + t.Fatalf("provider calls = %d", calls.Load()) } - after := g.Snapshot() - if tok.Size() != 8 || after.Used != before.Used || b.AggregateUsed() != 8 { - t.Fatalf("rejected grow changed charge: token=%d generation=%d aggregate=%d", tok.Size(), after.Used, b.AggregateUsed()) - } - if after.RejectCount != before.RejectCount+1 { - t.Fatalf("reject count=%d, want %d", after.RejectCount, before.RejectCount+1) - } - tok.Release() -} - -func TestHashBuildReservationGrowHonorsInactiveClosedAndLiveCap(t *testing.T) { - b := MustNewHashBuildBudget(20, 20) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGeneration(1) - cap := uint64(20) - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - b.SetAggregateCapProvider(func() (uint64, error) { return cap, nil }) - tok, err := g.Reserve(8) + capValue.Store(6) + cachedAllocation, err := acquireTestPhysicalAllocation(generation, 1) if err != nil { - t.Fatal(err) - } - cap = 8 - now = now.Add(hashBuildBudgetCapRefreshTTL + time.Nanosecond) - if err = tok.Grow(1); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("live-cap grow=%v", err) - } - if tok.Size() != 8 || g.Used() != 8 || b.AggregateUsed() != 8 { - t.Fatalf("live-cap rejection changed charge") - } - tok.Release() - if err = tok.Grow(1); !errors.Is(err, ErrHashBuildReservationInactive) { - t.Fatalf("released grow=%v", err) + t.Fatalf("cached cap should remain valid: %v", err) } - - cap = 20 - closed, err := g.Reserve(2) - if err != nil { - t.Fatal(err) + if calls.Load() != 1 { + t.Fatal("cached fast path sampled the provider") } - g.Close() - if err = closed.Grow(1); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed grow=%v", err) + first.Release() + cachedAllocation.Release() + // Invalidate the cache to model the next observation interval. + budget.mu.Lock() + budget.capCached = false + budget.mu.Unlock() + if _, err = acquireTestPhysicalAllocation(generation, 7); !errors.Is(err, ErrHashBuildBudgetAdmission) { + t.Fatalf("shrunk live cap error = %v", err) } - closed.Release() -} - -func TestHashBuildReservationGrowConcurrentTerminalTransitions(t *testing.T) { - for i := 0; i < 100; i++ { - b := MustNewHashBuildBudget(64, 64) - g, _ := b.OpenGeneration(uint64(i + 1)) - tok, err := g.Reserve(10) - if err != nil { - t.Fatal(err) - } - var wg sync.WaitGroup - wg.Add(3) - movedC := make(chan *HashBuildReservation, 1) - go func() { - defer wg.Done() - _ = tok.Grow(5) - }() - go func() { - defer wg.Done() - _, _ = tok.ReconcileDown(4) - }() - go func() { - defer wg.Done() - movedC <- tok.Transfer() - }() - wg.Wait() - close(movedC) - tok.Release() - if moved := <-movedC; moved != nil { - moved.Release() - } - if g.Used() != 0 || b.AggregateUsed() != 0 { - t.Fatalf("iteration %d leaked charge: generation=%d aggregate=%d", i, g.Used(), b.AggregateUsed()) - } + if budget.AggregateCap() != 6 { + t.Fatalf("aggregate cap = %d", budget.AggregateCap()) } } -func TestHashBuildSpillLedgersTransferReconcile(t *testing.T) { - b := MustNewHashBuildBudget(64, 64) - g, _ := b.OpenGeneration(1) - disk, err := g.ReserveSpillDisk(100) +func TestHashBuildBudgetSpillLedgers(t *testing.T) { + budget := MustNewHashBuildBudget(64, 64) + generation, _ := budget.OpenGeneration(1) + disk, err := generation.ReserveSpillDisk(100) if err != nil { t.Fatal(err) } - fd, err := g.ReserveSpillFD(2) + fd, err := generation.ReserveSpillFD(2) if err != nil { t.Fatal(err) } - if b.SpillDiskUsed() != 100 || b.SpillFDUsed() != 2 { - t.Fatalf("used disk=%d fd=%d", b.SpillDiskUsed(), b.SpillFDUsed()) + if err = disk.Grow(25); err != nil || disk.Size() != 125 { + t.Fatalf("disk grow: size=%d err=%v", disk.Size(), err) } - if err := disk.Grow(25); err != nil { - t.Fatalf("disk grow: %v", err) + if ok, reconcileErr := disk.ReconcileDown(40); !ok || reconcileErr != nil { + t.Fatalf("disk reconcile: ok=%v err=%v", ok, reconcileErr) } - if disk.Size() != 125 || b.SpillDiskUsed() != 125 { - t.Fatalf("grown disk token=%d used=%d", disk.Size(), b.SpillDiskUsed()) - } - if ok, err := disk.ReconcileDown(40); !ok || err != nil { - t.Fatalf("disk reconcile: %v %v", ok, err) - } - moved := fd.Transfer() - if moved == nil || fd.Release() { - t.Fatal("fd transfer") - } - g.Close() - if _, err := g.ReserveSpillDisk(1); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed spill reserve=%v", err) + generation.Close() + if _, err = generation.ReserveSpillDisk(1); !errors.Is(err, ErrHashBuildBudgetClosed) { + t.Fatalf("closed spill error = %v", err) } disk.Release() - moved.Release() - if b.SpillDiskUsed() != 0 || b.SpillFDUsed() != 0 { - t.Fatalf("spill leak disk=%d fd=%d", b.SpillDiskUsed(), b.SpillFDUsed()) + fd.Release() + if budget.SpillDiskUsed() != 0 || budget.SpillFDUsed() != 0 { + t.Fatal("spill reservations leaked") } } -func TestConfiguredSpillFDCapCushionsFirstShuffleRepartitionPeak(t *testing.T) { - const firstRepartitionPeak = uint64(16 * (64 + 64)) - if got := configuredSpillFDCap(192 << 20); got < firstRepartitionPeak { - t.Fatalf("configured spill fd cap=%d, want at least first 16-way repartition peak=%d", got, firstRepartitionPeak) - } -} - -func TestClampSpillFDCapBoundaries(t *testing.T) { - for _, tc := range []struct { - name string - configured, processLimit uint64 - limitKnown bool - want uint64 +func TestHashBuildBudgetSpillReleaseRejectsLedgerUnderflow(t *testing.T) { + for _, test := range []struct { + name string + reserve func(*HashBuildBudgetGeneration) (func() bool, error) + corrupt func(*HashBuildBudget, *HashBuildBudgetGeneration) }{ - {name: "unknown fails closed", configured: 2048, processLimit: 1 << 20, want: 0}, - {name: "zero configured", configured: 0, processLimit: 1024, limitKnown: true, want: 0}, - {name: "below absolute headroom", configured: 2048, processLimit: 63, limitKnown: true, want: 0}, - {name: "at absolute headroom", configured: 2048, processLimit: 64, limitKnown: true, want: 0}, - {name: "one fd above headroom", configured: 2048, processLimit: 65, limitKnown: true, want: 1}, - {name: "absolute headroom dominates", configured: 2048, processLimit: 128, limitKnown: true, want: 64}, - {name: "quarter headroom dominates", configured: 2048, processLimit: 1024, limitKnown: true, want: 768}, - {name: "explicit finite cap retained", configured: 10, processLimit: 1024, limitKnown: true, want: 10}, - {name: "unlimited retains configured", configured: 2048, processLimit: math.MaxUint64, limitKnown: true, want: 2048}, + { + name: "disk", + reserve: func(generation *HashBuildBudgetGeneration) (func() bool, error) { + reservation, err := generation.ReserveSpillDisk(2) + return reservation.Release, err + }, + corrupt: func(budget *HashBuildBudget, generation *HashBuildBudgetGeneration) { + generation.spillDiskUsed = 1 + budget.spillDiskUsed = 1 + }, + }, + { + name: "fd", + reserve: func(generation *HashBuildBudgetGeneration) (func() bool, error) { + reservation, err := generation.ReserveSpillFD(2) + return reservation.Release, err + }, + corrupt: func(budget *HashBuildBudget, generation *HashBuildBudgetGeneration) { + generation.spillFDUsed = 1 + budget.spillFDUsed = 1 + }, + }, } { - t.Run(tc.name, func(t *testing.T) { - if got := clampSpillFDCap(tc.configured, tc.processLimit, tc.limitKnown); got != tc.want { - t.Fatalf("clampSpillFDCap(%d, %d, %v)=%d, want %d", - tc.configured, tc.processLimit, tc.limitKnown, got, tc.want) + t.Run(test.name, func(t *testing.T) { + budget := MustNewHashBuildBudget(64, 64) + generation, err := budget.OpenGeneration(1) + if err != nil { + t.Fatal(err) + } + release, err := test.reserve(generation) + if err != nil { + t.Fatal(err) } + budget.mu.Lock() + test.corrupt(budget, generation) + budget.mu.Unlock() + defer func() { + if recover() == nil { + t.Fatal("corrupt spill ledger release did not panic") + } + }() + release() }) } } -func TestDefaultSpillFDCapMatchesProcessLimit(t *testing.T) { - limit, ok := processOpenFileLimit() - want := clampSpillFDCap(configuredSpillFDCap(192<<20), limit, ok) - b := MustNewHashBuildBudget(192<<20, 192<<20) - if got := b.SpillFDCap(); got != want { - t.Fatalf("spill fd cap=%d, want process-clamped cap=%d (limit=%d known=%v)", got, want, limit, ok) - } -} - -func TestHashBuildSpillFDCapUnderRLIMIT(t *testing.T) { - const ( - childEnv = "MO_HASHBUILD_RLIMIT_CHILD" - limitEnv = "MO_HASHBUILD_RLIMIT_NOFILE" - ) - if os.Getenv(childEnv) == "1" { - limit, ok := processOpenFileLimit() - if !ok { - t.Fatal("RLIMIT_NOFILE unavailable in RLIMIT child") - } - rawTarget := os.Getenv(limitEnv) - target, err := strconv.ParseUint(rawTarget, 10, 64) - if err != nil { - t.Fatalf("parse target %q: %v", rawTarget, err) - } - if limit != target { - t.Fatalf("child RLIMIT_NOFILE=%d, want %d", limit, target) - } - - configured := configuredSpillFDCap(192 << 20) - want := clampSpillFDCap(configured, limit, true) - b := MustNewHashBuildBudget(192<<20, 192<<20) - if got := b.SpillFDCap(); got != want { - t.Fatalf("child spill fd cap=%d, want %d", got, want) - } - g, err := b.OpenGeneration(1) - if err != nil { - t.Fatal(err) - } - // Simulate a budget/generation opened while the process limit was - // higher. ReserveSpillFD must sample the current RLIMIT again instead - // of trusting these stale effective caps. - b.mu.Lock() - b.spillFDCap = configured - g.spillFDCap = configured - b.mu.Unlock() - if _, err = g.ReserveSpillFD(want + 1); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("RLIMIT+headroom overflow error=%v, want admission rejection", err) - } - if got := b.SpillFDCap(); got != want { - t.Fatalf("runtime preflight refreshed spill fd cap=%d, want %d", got, want) - } - token, err := g.ReserveSpillFD(want) - if err != nil { - t.Fatalf("exact safe FD cap rejected: %v", err) - } - if !token.Release() { - t.Fatal("exact safe FD reservation did not release") - } - - if err = b.SetSpillCaps(0, 10); err != nil { - t.Fatal(err) - } - wantExplicit := clampSpillFDCap(10, limit, true) - if got := b.SpillFDCap(); got != wantExplicit { - t.Fatalf("explicit finite FD cap=%d, want process-clamped %d", got, wantExplicit) +func TestClampSpillFDCapBoundaries(t *testing.T) { + for _, test := range []struct { + configured, limit uint64 + known bool + want uint64 + }{ + {2048, 1 << 20, false, 0}, + {0, 1024, true, 0}, + {2048, 64, true, 0}, + {2048, 65, true, 1}, + {2048, 128, true, 64}, + {2048, 1024, true, 768}, + {10, 1024, true, 10}, + {2048, math.MaxUint64, true, 2048}, + } { + if got := clampSpillFDCap(test.configured, test.limit, test.known); got != test.want { + t.Fatalf("clampSpillFDCap(%d, %d, %v) = %d, want %d", + test.configured, test.limit, test.known, got, test.want) } - return - } - - switch runtime.GOOS { - case "darwin", "linux": - default: - t.Skip("RLIMIT_NOFILE subprocess is only supported on Darwin and Linux") - } - parentLimit, ok := processOpenFileLimit() - if !ok || parentLimit < hashBuildNonSpillFDHeadroom+1 { - t.Skipf("parent RLIMIT_NOFILE=%d known=%v is too small for isolated child test", parentLimit, ok) - } - target := uint64(128) - if parentLimit < target { - target = parentLimit - } - targetText := strconv.FormatUint(target, 10) - cmd := exec.Command( - "/bin/sh", "-c", - `ulimit -S -n "$MO_HASHBUILD_RLIMIT_NOFILE" && -ulimit -H -n "$MO_HASHBUILD_RLIMIT_NOFILE" && -exec "$@"`, - "sh", os.Args[0], "-test.run=^TestHashBuildSpillFDCapUnderRLIMIT$", "-test.count=1", - ) - cmd.Env = append(os.Environ(), childEnv+"=1", limitEnv+"="+targetText) - if output, err := cmd.CombinedOutput(); err != nil { - t.Fatalf("RLIMIT child failed: %v\n%s", err, output) } } -func TestHashBuildBudgetLiveCapProviderShrinksOpenGeneration(t *testing.T) { - b := MustNewHashBuildBudget(10, 10) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGenerationWithCap(1, 10) - cap := uint64(10) - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - b.SetAggregateCapProvider(func() (uint64, error) { return cap, nil }) - owned, err := g.Reserve(6) - if err != nil { - t.Fatal(err) - } - cap = 5 - now = now.Add(hashBuildBudgetCapRefreshTTL + time.Nanosecond) - if _, err = g.Reserve(1); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("open generation ignored live cap shrink: %v", err) - } - owned.Release() - token, err := g.Reserve(5) - if err != nil { - t.Fatalf("reservation at refreshed cap failed: %v", err) - } - token.Release() -} +func TestGetHashBuildBudgetInitializesAndReusesCNAggregate(t *testing.T) { + const localService = "__process_local_cn__" + hashBuildCNBudgets.Delete(localService) + t.Cleanup(func() { hashBuildCNBudgets.Delete(localService) }) -func TestHashBuildBudgetCapProviderCachesWithinTTLAndRefreshes(t *testing.T) { - b := MustNewHashBuildBudget(10, 10) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGeneration(1) - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - var calls atomic.Int32 - cap := uint64(10) - b.SetAggregateCapProvider(func() (uint64, error) { - calls.Add(1) - return cap, nil - }) - first, err := g.Reserve(4) + first := &Process{Base: &BaseProcess{Lim: Limitation{Size: 2 << 20, SpillSize: 4 << 20}}} + firstGeneration, err := first.GetHashBuildBudget() if err != nil { t.Fatal(err) } - if got := calls.Load(); got != 1 { - t.Fatalf("first reservation provider calls=%d, want 1", got) + if firstGeneration.Cap() != 2<<20 || firstGeneration.SpillDiskCap() != 4<<20 { + t.Fatalf("first generation limits: %+v", firstGeneration.Snapshot()) + } + if cached, cachedErr := first.GetHashBuildBudget(); cachedErr != nil || cached != firstGeneration { + t.Fatal("process generation was not cached") } - now = now.Add(hashBuildBudgetCapRefreshTTL - time.Nanosecond) - second, err := g.Reserve(1) + second := &Process{Base: &BaseProcess{Lim: Limitation{Size: 1 << 20}}} + secondGeneration, err := second.GetHashBuildBudget() if err != nil { t.Fatal(err) } - if got := calls.Load(); got != 1 { - t.Fatalf("TTL reservation provider calls=%d, want 1", got) - } - cap = 3 - now = now.Add(2 * time.Nanosecond) - if _, err = g.Reserve(1); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("expired shrink reservation=%v, want admission rejection", err) - } - if got := calls.Load(); got != 2 { - t.Fatalf("expired reservation provider calls=%d, want 2", got) + if secondGeneration == firstGeneration || secondGeneration.budget != firstGeneration.budget || + secondGeneration.Cap() != 1<<20 { + t.Fatal("second process did not reuse the CN aggregate") } - first.Release() - second.Release() + firstGeneration.Close() + secondGeneration.Close() + firstGeneration.budget.Close() } -func TestHashBuildBudgetCachedFastPathSkipsRefreshGate(t *testing.T) { - b := MustNewHashBuildBudget(10, 10) - b.capRefreshTTL = time.Hour - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - b.SetAggregateCapProvider(func() (uint64, error) { return 10, nil }) - if _, _, refreshed, err := b.refreshAggregateCap(false, 0); err != nil || !refreshed { - t.Fatalf("seed refresh: refreshed=%v err=%v", refreshed, err) - } - - b.refreshMu.Lock() - resultC := make(chan error, 1) - go func() { - _, _, refreshed, err := b.refreshAggregateCap(false, 0) - if err == nil && refreshed { - err = errors.New("cached refresh unexpectedly sampled provider") - } - resultC <- err - }() - - var err error - select { - case err = <-resultC: - b.refreshMu.Unlock() - case <-time.After(time.Second): - b.refreshMu.Unlock() - <-resultC - t.Fatal("cached refresh waited for refreshMu") - } +func TestResolveHashBuildCeiling(t *testing.T) { + const gib = uint64(1 << 30) + ceiling, err := ResolveHashBuildCeiling(HashBuildCeilingInputs{ + CgroupMemoryMax: 20 * gib, + HostMemTotal: 10 * gib, + GlobalMpoolCap: 30 * gib, + FileCacheHint: gib, + ProcessLimitationSize: 2 * gib, + }) if err != nil { t.Fatal(err) } -} - -func TestHashBuildBudgetUnchangedObservationSkipsRefreshGate(t *testing.T) { - inputs := HashBuildCeilingInputs{ - CgroupMemoryMax: 8 << 30, - HostMemTotal: 16 << 30, - GlobalMpoolCap: 6 << 30, - FileCacheHint: 512 << 20, + if ceiling.EffectiveCN != 10*gib || ceiling.Reserve != 4*gib || + ceiling.CNHashCap != 6*gib || ceiling.QueryCap != 2*gib { + t.Fatalf("ceiling = %+v", ceiling) } - ceiling, err := ResolveHashBuildCeiling(inputs) - if err != nil { - t.Fatal(err) + if _, err = ResolveHashBuildCeiling(HashBuildCeilingInputs{ + CgroupMemoryMax: math.MaxUint64, + }); !errors.Is(err, ErrHashBuildCeilingMissing) { + t.Fatalf("missing finite source error = %v", err) } - b := MustNewHashBuildBudget(ceiling.CNHashCap, ceiling.CNHashCap) - b.installCNCapProvider(inputs) - - b.refreshMu.Lock() - doneC := make(chan struct{}) - go func() { - b.mergeObservedCNCap(inputs, ceiling.CNHashCap) - close(doneC) - }() - select { - case <-doneC: - b.refreshMu.Unlock() - case <-time.After(time.Second): - b.refreshMu.Unlock() - <-doneC - t.Fatal("unchanged CN cap observation waited for refreshMu") + if small, smallErr := ResolveHashBuildCeiling(HashBuildCeilingInputs{ + HostMemTotal: 3 * gib, + FileCacheHint: 3 * gib, + }); smallErr != nil || small.CNHashCap != 3*gib/20 { + t.Fatalf("small-CN ceiling = %+v, err=%v", small, smallErr) } } -func TestHashBuildBudgetCNProviderKeepsProcessMemorySnapshot(t *testing.T) { - previousInputs := hashBuildProcessMemoryInputs - previousCap := commonmpool.GlobalCap() - previousHint := fileservice.GlobalMemoryCacheSizeHint.Swap(0) - t.Cleanup(func() { - hashBuildProcessMemoryInputs = previousInputs - commonmpool.InitCap(previousCap) - fileservice.GlobalMemoryCacheSizeHint.Store(previousHint) - }) - - hashBuildProcessMemoryInputs = HashBuildCeilingInputs{ - CgroupMemoryMax: 8 << 30, - HostMemTotal: 16 << 30, - } - commonmpool.InitCap(commonmpool.PB) +func TestHashBuildBudgetUsesCurrentMemoryInputs(t *testing.T) { + previous := hashBuildProcessMemoryInputs + hashBuildProcessMemoryInputs = HashBuildCeilingInputs{HostMemTotal: 8 << 30} + t.Cleanup(func() { hashBuildProcessMemoryInputs = previous }) + previousHint := fileservice.GlobalMemoryCacheSizeHint.Load() + fileservice.GlobalMemoryCacheSizeHint.Store(1 << 30) + t.Cleanup(func() { fileservice.GlobalMemoryCacheSizeHint.Store(previousHint) }) - b := MustNewHashBuildBudget(4<<30, 4<<30) - b.installCNCapProvider(HashBuildCeilingInputs{ - CgroupMemoryMax: 4 << 30, - HostMemTotal: 8 << 30, - }) - if _, err := b.sampleCNCap(); err != nil { - t.Fatal(err) - } - if b.liveCapInputs.CgroupMemoryMax != 8<<30 || b.liveCapInputs.HostMemTotal != 16<<30 { - t.Fatalf("physical snapshot changed: %+v", b.liveCapInputs) + inputs := hashBuildProcessMemoryInputs + inputs.FileCacheHint = uint64(fileservice.GlobalMemoryCacheSizeHint.Load()) + ceiling, err := ResolveHashBuildCeiling(inputs) + if err != nil || ceiling.CNHashCap == 0 { + t.Fatalf("current memory inputs did not produce a finite cap: %+v %v", ceiling, err) } } -func TestHashBuildBudgetCapProviderGrowthOnAggregateReject(t *testing.T) { - b := MustNewHashBuildBudget(10, 10) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGeneration(1) - g2, _ := b.OpenGeneration(2) - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - var calls atomic.Int32 - cap := uint64(10) - b.SetAggregateCapProvider(func() (uint64, error) { - calls.Add(1) - return cap, nil - }) - owned, err := g.Reserve(10) +func BenchmarkHashBuildBudgetAllocationAccount(b *testing.B) { + budget := MustNewHashBuildBudget(math.MaxUint64, math.MaxUint64) + generation, err := budget.OpenGeneration(1) if err != nil { - t.Fatal(err) + b.Fatal(err) } - cap = 20 - // The cached cap is still 10, so this failed aggregate admission forces a - // refresh even though the TTL has not elapsed and then succeeds at 20. - grown, err := g2.Reserve(1) + registry, err := mpool.NewAllocationAccountRegistry(1, uint64(b.N)+1) if err != nil { - t.Fatalf("growth refresh reservation=%v", err) + b.Fatal(err) } - if got := calls.Load(); got != 2 { - t.Fatalf("growth refresh provider calls=%d, want 2", got) + account, err := registry.OpenWithController(math.MaxUint64, generation) + if err != nil { + b.Fatal(err) } - grown.Release() - owned.Release() -} - -func TestHashBuildBudgetCapProviderConcurrentSingleFlight(t *testing.T) { - b := MustNewHashBuildBudget(128, 128) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGeneration(1) - var calls atomic.Int32 - started := make(chan struct{}) - release := make(chan struct{}) - b.SetAggregateCapProvider(func() (uint64, error) { - if calls.Add(1) == 1 { - close(started) - <-release + mp := mpool.MustNewZero() + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + allocation, allocErr := mp.AllocAccounted(1, account, 1, 1) + if allocErr != nil { + b.Fatal(allocErr) } - return 128, nil - }) - const workers = 16 - tokens := make(chan *HashBuildReservation, workers) - var wg sync.WaitGroup - for i := 0; i < workers; i++ { - wg.Add(1) - go func() { - defer wg.Done() - tok, err := g.Reserve(1) - if err != nil { - t.Errorf("concurrent reserve: %v", err) - return - } - tokens <- tok - }() - } - <-started - close(release) - wg.Wait() - if got := calls.Load(); got != 1 { - t.Fatalf("concurrent provider calls=%d, want 1", got) - } - for i := 0; i < workers; i++ { - (<-tokens).Release() - } -} - -func TestHashBuildBudgetCapProviderErrorCachedFailClosed(t *testing.T) { - b := MustNewHashBuildBudget(10, 10) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGeneration(1) - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - var calls atomic.Int32 - want := errors.New("cgroup unavailable") - b.SetAggregateCapProvider(func() (uint64, error) { - calls.Add(1) - return 0, want - }) - if _, err := g.Reserve(1); !errors.Is(err, want) { - t.Fatalf("provider error=%v, want %v", err, want) - } - if _, err := g.Reserve(1); !errors.Is(err, want) { - t.Fatalf("cached provider error=%v, want %v", err, want) - } - if got := calls.Load(); got != 1 { - t.Fatalf("cached error provider calls=%d, want 1", got) - } - now = now.Add(hashBuildBudgetCapRefreshTTL + time.Nanosecond) - if _, err := g.Reserve(1); !errors.Is(err, want) { - t.Fatalf("expired provider error=%v, want %v", err, want) - } - if got := calls.Load(); got != 2 { - t.Fatalf("expired error provider calls=%d, want 2", got) - } -} - -func TestHashBuildBudgetCapProviderSharedByReserveAndGrow(t *testing.T) { - b := MustNewHashBuildBudget(20, 20) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGeneration(1) - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - var calls atomic.Int32 - b.SetAggregateCapProvider(func() (uint64, error) { - calls.Add(1) - return 20, nil - }) - tok, err := g.Reserve(2) - if err != nil { - t.Fatal(err) - } - if err = tok.Grow(3); err != nil { - t.Fatal(err) - } - if got := calls.Load(); got != 1 { - t.Fatalf("Reserve+Grow provider calls=%d, want 1", got) - } - tok.Release() -} - -func TestHashBuildBudgetCapProviderZeroTTLRefreshesEveryReservation(t *testing.T) { - b := MustNewHashBuildBudget(16, 16) - g, _ := b.OpenGeneration(1) - b.capRefreshTTL = 0 - var calls atomic.Int32 - b.SetAggregateCapProvider(func() (uint64, error) { - calls.Add(1) - return 16, nil - }) - for i := 0; i < 2; i++ { - tok, err := g.Reserve(1) - if err != nil { - t.Fatal(err) - } - tok.Release() - } - if got := calls.Load(); got != 2 { - t.Fatalf("zero-TTL provider calls=%d, want 2", got) - } -} - -func TestHashBuildBudgetUpdateAndProviderReinstallReuseCache(t *testing.T) { - b := MustNewHashBuildBudget(16, 16) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGeneration(1) - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - var calls atomic.Int32 - provider := func() (uint64, error) { - calls.Add(1) - return 16, nil - } - b.SetAggregateCapProvider(provider) - first, err := g.Reserve(1) - if err != nil { - t.Fatal(err) - } - if err = b.UpdateAggregateCap(16); err != nil { - t.Fatal(err) - } - // GetHashBuildBudget re-installs an equivalent closure before updating the - // freshly resolved cap. The update seeds the cache for this query. - b.SetAggregateCapProvider(provider) - if err = b.UpdateAggregateCap(16); err != nil { - t.Fatal(err) - } - second, err := g.Reserve(1) - if err != nil { - t.Fatal(err) - } - if got := calls.Load(); got != 1 { - t.Fatalf("Update+Set provider calls=%d, want 1", got) - } - first.Release() - second.Release() -} - -func TestHashBuildBudgetProviderReplacementInvalidatesCache(t *testing.T) { - b := MustNewHashBuildBudget(16, 16) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGeneration(1) - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - var oldCalls, newCalls atomic.Int32 - b.SetAggregateCapProvider(func() (uint64, error) { - oldCalls.Add(1) - return 16, nil - }) - owned, err := g.Reserve(1) - if err != nil { - t.Fatal(err) - } - b.SetAggregateCapProvider(func() (uint64, error) { - newCalls.Add(1) - return 2, nil - }) - second, err := g.Reserve(1) - if err != nil { - t.Fatalf("replacement reserve=%v", err) - } - if oldCalls.Load() != 1 || newCalls.Load() != 1 { - t.Fatalf("provider calls old=%d new=%d, want 1,1", oldCalls.Load(), newCalls.Load()) - } - owned.Release() - second.Release() -} - -func TestHashBuildBudgetCNSourceTurnoverRetainsRestrictiveFallback(t *testing.T) { - const gib = uint64(1 << 30) - oldInputs := HashBuildCeilingInputs{HostMemTotal: 10 * gib, CgroupMemoryMax: 20 * gib} - newInputs := HashBuildCeilingInputs{HostMemTotal: 20 * gib, CgroupMemoryMax: 10 * gib} - oldCeiling, err := ResolveHashBuildCeiling(oldInputs) - if err != nil { - t.Fatal(err) - } - newCeiling, err := ResolveHashBuildCeiling(newInputs) - if err != nil { - t.Fatal(err) - } - if oldCeiling.CNHashCap != newCeiling.CNHashCap { - t.Fatalf("test setup caps old=%d new=%d, want equal", oldCeiling.CNHashCap, newCeiling.CNHashCap) - } - b := MustNewHashBuildBudget(oldCeiling.CNHashCap, oldCeiling.CNHashCap) - b.installCNCapProvider(oldInputs) - b.mergeObservedCNCap(newInputs, newCeiling.CNHashCap) - - // A zero source sample models transient read failures before the installed - // provider gets a complete view of the turnover. The merged shared snapshot - // must retain both restrictive observations and cannot reopen the cap. - b.refreshMu.Lock() - got, err := b.resolveCNCapSample(HashBuildCeilingInputs{}) - b.refreshMu.Unlock() - if err != nil { - t.Fatal(err) - } - if got > newCeiling.CNHashCap { - t.Fatalf("fallback cap=%d, exceeds restrictive cap %d", got, newCeiling.CNHashCap) - } -} - -func TestHashBuildBudgetSlowProviderTTLStartsAfterSample(t *testing.T) { - b := MustNewHashBuildBudget(20, 20) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGeneration(1) - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - var calls atomic.Int32 - b.SetAggregateCapProvider(func() (uint64, error) { - calls.Add(1) - now = now.Add(2 * hashBuildBudgetCapRefreshTTL) - return 20, nil - }) - first, err := g.Reserve(1) - if err != nil { - t.Fatal(err) - } - second, err := g.Reserve(1) - if err != nil { - t.Fatal(err) - } - if got := calls.Load(); got != 1 { - t.Fatalf("slow provider calls=%d, want freshly completed sample reused", got) - } - first.Release() - second.Release() -} - -func TestHashBuildBudgetClosedSkipsCapProvider(t *testing.T) { - b := MustNewHashBuildBudget(10, 10) - g, _ := b.OpenGeneration(1) - var calls atomic.Int32 - b.SetAggregateCapProvider(func() (uint64, error) { - calls.Add(1) - return 10, nil - }) - b.Close() - if _, err := g.Reserve(1); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed reserve=%v", err) - } - if got := calls.Load(); got != 0 { - t.Fatalf("closed provider calls=%d, want 0", got) - } -} - -func TestHashBuildBudgetCompatibilityAndObservabilitySurface(t *testing.T) { - var nilBudget *HashBuildBudget - if !nilBudget.Snapshot().Closed || - nilBudget.AggregateCap() != 0 || - nilBudget.CNHashCap() != 0 || - nilBudget.QueryCap() != 0 || - nilBudget.AggregateUsed() != 0 || - nilBudget.CNHashUsed() != 0 || - nilBudget.Current() != 0 || - nilBudget.Capacity() != 0 || - !nilBudget.Closed() || - nilBudget.SpillDiskCap() != 0 || - nilBudget.SpillDiskUsed() != 0 || - nilBudget.SpillFDCap() != 0 || - nilBudget.SpillFDUsed() != 0 { - t.Fatal("nil budget accessors must report an inert closed budget") - } - nilBudget.Close() - if err := nilBudget.SetSpillCaps(1, 1); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("nil SetSpillCaps error = %v", err) - } - - var nilGeneration *HashBuildBudgetGeneration - if nilGeneration.ID() != 0 || - nilGeneration.Cap() != 0 || - nilGeneration.QueryCap() != 0 || - nilGeneration.Capacity() != 0 || - nilGeneration.Used() != 0 || - nilGeneration.Current() != 0 || - !nilGeneration.Closed() || - nilGeneration.SpillDiskCap() != 0 || - nilGeneration.SpillDiskUsed() != 0 || - nilGeneration.SpillFDCap() != 0 || - nilGeneration.SpillFDUsed() != 0 || - !nilGeneration.Snapshot().Closed { - t.Fatal("nil generation accessors must report an inert closed generation") - } - nilGeneration.Close() - if nilGeneration.TryReserve(1) { - t.Fatal("nil generation reservation succeeded") - } - - for _, kind := range []HashBuildBudgetErrorKind{ - HashBuildBudgetErrorAdmission, - HashBuildBudgetErrorClosed, - HashBuildBudgetErrorInvalid, - HashBuildBudgetErrorCeilingMissing, - } { - err := &HashBuildBudgetError{Kind: kind} - if err.Error() == "" || err.Unwrap() == nil { - t.Fatalf("kind %d did not expose an error", kind) - } - } - var nilBudgetErr *HashBuildBudgetError - if nilBudgetErr.Error() != "" || nilBudgetErr.Unwrap() != nil || nilBudgetErr.Is(ErrHashBuildBudgetAdmission) { - t.Fatal("nil budget error must remain inert") - } - for _, tc := range []struct { - kind HashBuildBudgetErrorKind - target error - }{ - {HashBuildBudgetErrorAdmission, ErrHashBuildBudgetAdmission}, - {HashBuildBudgetErrorClosed, ErrHashBuildBudgetClosed}, - {HashBuildBudgetErrorInvalid, ErrHashBuildBudgetInvalid}, - {HashBuildBudgetErrorCeilingMissing, ErrHashBuildCeilingMissing}, - } { - if !errors.Is(&HashBuildBudgetError{Kind: tc.kind}, tc.target) { - t.Fatalf("kind %d did not match %v", tc.kind, tc.target) - } - } - if errors.Is(&HashBuildBudgetError{Kind: HashBuildBudgetErrorClosed}, - ErrHashBuildBudgetAdmission, - ) { - t.Fatal("closed budget must not match a recoverable capacity admission") - } - unknown := &HashBuildBudgetError{Kind: HashBuildBudgetErrorKind(255)} - if errors.Is(unknown, ErrHashBuildBudgetAdmission) || - !errors.Is(unknown, ErrHashBuildBudgetInvalid) { - t.Fatal("unknown error kind must remain a fatal invalid error") - } - if unknown.Is(ErrHashBuildBudgetInvalid) { - t.Fatal("unknown error kind matched a sentinel") - } - message := &HashBuildBudgetError{Message: "explicit"} - if message.Error() != "explicit" { - t.Fatalf("explicit message = %q", message.Error()) - } - - b, err := NewHashBuildBudgetWithSpillCaps(100, 80, 200, 10) - if err != nil { - t.Fatal(err) - } - if b.AggregateCap() != 100 || b.CNHashCap() != 100 || b.Capacity() != 100 || b.QueryCap() != 80 { - t.Fatal("budget cap aliases disagree") - } - if b.SpillDiskCap() != 200 || b.SpillFDCap() != 10 { - t.Fatal("explicit spill caps were not installed") - } - if err = b.SetSpillCaps(0, 0); err != nil { - t.Fatal(err) - } - snapshot := b.Snapshot() - if snapshot.AggregateCap != 100 || snapshot.AggregateUsed != 0 || snapshot.Closed { - t.Fatalf("unexpected budget snapshot: %+v", snapshot) - } - - g, err := b.OpenGenerationWithLimits(7, 50, 100, 5) - if err != nil { - t.Fatal(err) - } - if g.ID() != 7 || g.Cap() != 50 || g.QueryCap() != 50 || g.Capacity() != 50 { - t.Fatal("generation identity or cap aliases disagree") - } - if g.SpillDiskCap() != 100 || g.SpillFDCap() != 5 || g.Current() != 0 { - t.Fatal("generation spill caps or current usage are wrong") - } - if !g.TryReserve(1) { - t.Fatal("TryReserve rejected an admissible charge") - } - token, err := g.Reserve(8) - if err != nil { - t.Fatal(err) - } - if token.GenerationID() != 7 || token.Size() != 8 || token.Released() { - t.Fatal("memory reservation accessors are inconsistent") - } - if ok, reconcileErr := token.Reconcile(6); !ok || reconcileErr != nil { - t.Fatalf("compatibility reconcile failed: ok=%v err=%v", ok, reconcileErr) - } - moved := token.TransferOwnership() - if moved == nil || !token.Released() || moved.GenerationID() != 7 { - t.Fatal("memory ownership transfer failed") - } - if !moved.Release() { - t.Fatal("transferred memory reservation did not release") - } - - disk, err := g.ReserveSpillDiskBytes(12) - if err != nil { - t.Fatal(err) - } - fd, err := g.ReserveSpillFileDescriptors(2) - if err != nil { - t.Fatal(err) - } - if disk.Size() != 12 || disk.Released() || fd.Size() != 2 || fd.Released() { - t.Fatal("spill reservation accessors are inconsistent") - } - if ok, reconcileErr := disk.Reconcile(10); !ok || reconcileErr != nil { - t.Fatalf("disk reconcile failed: ok=%v err=%v", ok, reconcileErr) - } - if ok, reconcileErr := fd.Reconcile(1); !ok || reconcileErr != nil { - t.Fatalf("fd reconcile failed: ok=%v err=%v", ok, reconcileErr) - } - movedDisk := disk.TransferTo() - movedFD := fd.TransferOwnership() - if movedDisk == nil || movedFD == nil || !disk.Released() || !fd.Released() { - t.Fatal("spill ownership transfer failed") - } - if !movedDisk.Release() || !movedFD.Release() { - t.Fatal("transferred spill reservations did not release") - } - stats := g.Stats() - if stats.ID != 7 || - g.Peak() == 0 || - g.ReserveCount() == 0 || - g.ReconcileCount() == 0 || - g.ReleaseCount() == 0 || - g.RejectCount() != 0 { - t.Fatalf("unexpected generation stats: %+v", stats) - } - - other, err := b.NewGeneration(8) - if err != nil { - t.Fatal(err) - } - other.Close() - query, err := b.OpenQueryBudget(9) - if err != nil { - t.Fatal(err) - } - query.Close() - explicit, err := b.OpenGenerationWithCapAndSpill(10, 40, 80, 4) - if err != nil { - t.Fatal(err) - } - explicit.Close() - g.Close() - b.Close() - if !b.Closed() || !g.Closed() { - t.Fatal("close accessors did not observe terminal state") - } - - ceiling, err := ResolveHashBuildBudget(HashBuildCeilingInputs{HostMemTotal: 8 << 30}) - if err != nil { - t.Fatal(err) - } - fromCeiling, err := NewHashBuildBudgetFromCeiling(ceiling) - if err != nil { - t.Fatal(err) - } - fromCeiling.Close() -} - -func TestHashBuildBudgetCompatibilityUnhappyPaths(t *testing.T) { - var nilProcess *Process - if _, err := nilProcess.GetHashBuildBudget(); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("nil process error = %v", err) - } - - var nilBudget *HashBuildBudget - if _, err := nilBudget.OpenGeneration(1); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("nil budget generation error = %v", err) - } - if _, err := nilBudget.OpenGenerationWithCap(1, 1); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("nil budget explicit generation error = %v", err) - } - - b, err := NewHashBuildBudgetWithSpillCaps(100, 80, 20, 4) - if err != nil { - t.Fatal(err) - } - if _, err = b.OpenGenerationWithCap(1, 0); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("zero generation cap error = %v", err) - } - if _, err = b.OpenGenerationWithCap(1, 101); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("oversized generation cap error = %v", err) - } - - g, err := b.OpenGenerationWithSpillCaps(1, 80, 10, 2) - if err != nil { - t.Fatal(err) - } - other, err := b.OpenGenerationWithSpillCaps(2, 80, 20, 4) - if err != nil { - t.Fatal(err) - } - defaults, err := b.OpenGenerationWithSpillCaps(3, 80, 0, 0) - if err != nil { - t.Fatal(err) - } - if defaults.SpillDiskCap() != 20 || defaults.SpillFDCap() != 4 { - t.Fatal("default generation spill caps were not clamped to the CN caps") - } - defaults.Close() - - disk, err := g.ReserveSpillDisk(8) - if err != nil { - t.Fatal(err) - } - if err = disk.Grow(2); err != nil { - t.Fatal(err) - } - if err = disk.Grow(1); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("query disk admission error = %v", err) - } - if _, err = other.ReserveSpillDisk(11); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("CN disk admission error = %v", err) - } - if ok, reconcileErr := disk.ReconcileDown(11); ok || !errors.Is(reconcileErr, ErrHashBuildReservationUpward) { - t.Fatalf("upward disk reconcile: ok=%v err=%v", ok, reconcileErr) - } - - fd, err := g.ReserveSpillFD(1) - if err != nil { - t.Fatal(err) - } - if _, err = g.ReserveSpillFD(2); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("query FD admission error = %v", err) - } - otherFD, err := other.ReserveSpillFD(3) - if err != nil { - t.Fatal(err) - } - if _, err = other.ReserveSpillFD(1); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("CN FD admission error = %v", err) - } - if ok, reconcileErr := fd.ReconcileDown(2); ok || !errors.Is(reconcileErr, ErrHashBuildReservationUpward) { - t.Fatalf("upward FD reconcile: ok=%v err=%v", ok, reconcileErr) - } - - memory, err := g.Reserve(1) - if err != nil { - t.Fatal(err) - } - movedMemory := memory.TransferTo() - if movedMemory == nil || memory.Transfer() != nil || memory.Release() { - t.Fatal("inactive memory token accepted a second terminal transition") - } - if !movedMemory.Release() || movedMemory.Release() { - t.Fatal("memory release was not exactly once") - } - - g.Close() - if _, err = g.ReserveSpillDisk(1); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed disk reservation error = %v", err) - } - if _, err = g.ReserveSpillFD(1); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed FD reservation error = %v", err) - } - if err = disk.Grow(1); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed disk growth error = %v", err) - } - if !disk.Release() || disk.Release() || disk.Transfer() != nil { - t.Fatal("disk release was not exactly once") - } - if _, reconcileErr := disk.ReconcileDown(0); !errors.Is(reconcileErr, ErrHashBuildReservationInactive) { - t.Fatalf("inactive disk reconcile error = %v", reconcileErr) - } - if !fd.Release() || fd.Release() || fd.Transfer() != nil { - t.Fatal("FD release was not exactly once") - } - if _, reconcileErr := fd.ReconcileDown(0); !errors.Is(reconcileErr, ErrHashBuildReservationInactive) { - t.Fatalf("inactive FD reconcile error = %v", reconcileErr) - } - otherFD.Release() - - var nilMemory *HashBuildReservation - var nilDisk *HashBuildSpillDiskReservation - var nilFD *HashBuildSpillFDReservation - var nilGeneration *HashBuildBudgetGeneration - if nilMemory.Size() != 0 || nilMemory.GenerationID() != 0 || !nilMemory.Released() || - nilMemory.Release() || nilMemory.Transfer() != nil || nilMemory.TransferTo() != nil || - nilDisk.Size() != 0 || !nilDisk.Released() || nilDisk.Release() || - nilDisk.Transfer() != nil || nilDisk.TransferOwnership() != nil || - nilFD.Size() != 0 || !nilFD.Released() || nilFD.Release() || - nilFD.Transfer() != nil || nilFD.TransferTo() != nil { - t.Fatal("nil reservation must remain inert") - } - if _, err = nilGeneration.ReserveSpillDisk(1); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("nil disk generation error = %v", err) - } - if _, err = nilGeneration.ReserveSpillFD(1); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("nil FD generation error = %v", err) - } - if _, reconcileErr := nilMemory.Reconcile(0); !errors.Is(reconcileErr, ErrHashBuildReservationInactive) { - t.Fatalf("nil memory reconcile error = %v", reconcileErr) - } - if _, reconcileErr := nilDisk.Reconcile(0); !errors.Is(reconcileErr, ErrHashBuildReservationInactive) { - t.Fatalf("nil disk reconcile error = %v", reconcileErr) - } - if _, reconcileErr := nilFD.Reconcile(0); !errors.Is(reconcileErr, ErrHashBuildReservationInactive) { - t.Fatalf("nil FD reconcile error = %v", reconcileErr) - } - if err = nilMemory.Grow(1); !errors.Is(err, ErrHashBuildReservationInactive) { - t.Fatalf("nil memory growth error = %v", err) - } - if err = nilDisk.Grow(1); !errors.Is(err, ErrHashBuildReservationInactive) { - t.Fatalf("nil disk growth error = %v", err) - } - - b.Close() - if _, err = b.OpenGeneration(3); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed budget generation error = %v", err) - } - if _, err = b.OpenGenerationWithCap(3, 1); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed budget explicit generation error = %v", err) - } -} - -func TestHashBuildBudgetAdmissionNamesIndependentComponent(t *testing.T) { - b, err := NewHashBuildBudgetWithSpillCaps(8, 8, 1, 1) - if err != nil { - t.Fatal(err) - } - g, err := b.OpenGenerationWithSpillCaps(1, 8, 1, 1) - if err != nil { - t.Fatal(err) - } - defer g.Close() - - assertComponent := func(err error, want HashBuildBudgetComponent) { - t.Helper() - var budgetErr *HashBuildBudgetError - if !errors.As(err, &budgetErr) || budgetErr.Component != want { - t.Fatalf("admission component: err=%v got=%v want=%v", err, budgetErr, want) - } - } - - memory, err := g.Reserve(8) - if err != nil { - t.Fatal(err) - } - defer memory.Release() - _, err = g.Reserve(1) - assertComponent(err, HashBuildBudgetComponentMemory) - - disk, err := g.ReserveSpillDisk(1) - if err != nil { - t.Fatal(err) - } - defer disk.Release() - _, err = g.ReserveSpillDisk(1) - assertComponent(err, HashBuildBudgetComponentSpillDisk) - - fd, err := g.ReserveSpillFD(1) - if err != nil { - t.Fatal(err) - } - defer fd.Release() - _, err = g.ReserveSpillFD(1) - assertComponent(err, HashBuildBudgetComponentSpillFD) - - err = b.SetSpillCaps(0, 1) - if err != nil { - t.Fatal(err) - } - err = b.SetSpillCaps(1, 0) - if err != nil { - t.Fatal(err) - } -} - -func TestGetHashBuildBudgetInitializesAndReusesCNAggregate(t *testing.T) { - const localService = "__process_local_cn__" - hashBuildCNBudgets.Delete(localService) - t.Cleanup(func() { hashBuildCNBudgets.Delete(localService) }) - - first := &Process{Base: &BaseProcess{Lim: Limitation{ - Size: 2 << 20, - SpillSize: 4 << 20, - }}} - firstGeneration, err := first.GetHashBuildBudget() - if err != nil { - t.Fatal(err) - } - if firstGeneration.Cap() != 2<<20 || firstGeneration.SpillDiskCap() != 4<<20 { - t.Fatalf("unexpected first generation limits: %+v", firstGeneration.Snapshot()) - } - cached, err := first.GetHashBuildBudget() - if err != nil || cached != firstGeneration { - t.Fatalf("process-local generation was not cached: generation=%p err=%v", cached, err) - } - - second := &Process{Base: &BaseProcess{Lim: Limitation{Size: 1 << 20}}} - secondGeneration, err := second.GetHashBuildBudget() - if err != nil { - t.Fatal(err) - } - if secondGeneration == firstGeneration || - secondGeneration.budget != firstGeneration.budget || - secondGeneration.Cap() != 1<<20 { - t.Fatal("second process did not reuse the CN aggregate with its own generation") - } - - aggregate := firstGeneration.budget - defaultAggregateSpillCap := aggregate.SpillDiskCap() - raisedSpillCap := defaultAggregateSpillCap + 1<<20 - third := &Process{Base: &BaseProcess{Lim: Limitation{ - Size: 1 << 20, - SpillSize: int64(raisedSpillCap), - }}} - thirdGeneration, err := third.GetHashBuildBudget() - if err != nil { - t.Fatal(err) - } - if thirdGeneration.SpillDiskCap() != raisedSpillCap || - aggregate.SpillDiskCap() != raisedSpillCap { - t.Fatalf("explicit spill cap was not raised at the shared ledger: generation=%d aggregate=%d want=%d", - thirdGeneration.SpillDiskCap(), aggregate.SpillDiskCap(), raisedSpillCap) - } - - lower := &Process{Base: &BaseProcess{Lim: Limitation{ - Size: 1 << 20, - SpillSize: 2 << 20, - }}} - lowerGeneration, err := lower.GetHashBuildBudget() - if err != nil { - t.Fatal(err) - } - if lowerGeneration.SpillDiskCap() != 2<<20 || - aggregate.SpillDiskCap() != raisedSpillCap { - t.Fatalf("lower per-query spill cap changed the shared ceiling: generation=%d aggregate=%d want aggregate=%d", - lowerGeneration.SpillDiskCap(), aggregate.SpillDiskCap(), raisedSpillCap) - } - - firstGeneration.Close() - secondGeneration.Close() - thirdGeneration.Close() - lowerGeneration.Close() - aggregate.Close() -} - -func TestHashBuildBudgetExplicitSpillCapConcurrentRaise(t *testing.T) { - budget := MustNewHashBuildBudget(100, 100) - t.Cleanup(budget.Close) - generation, err := budget.OpenGenerationWithSpillCaps(1, 100, 800, 1) - if err != nil { - t.Fatal(err) - } - t.Cleanup(generation.Close) - reservation, err := generation.ReserveSpillDisk(700) - if err != nil { - t.Fatal(err) - } - t.Cleanup(func() { reservation.Release() }) - - caps := []uint64{801, 900, 1200, 1100} - start := make(chan struct{}) - errs := make(chan error, len(caps)) - var wg sync.WaitGroup - for _, cap := range caps { - wg.Add(1) - go func() { - defer wg.Done() - <-start - errs <- budget.raiseSpillDiskCapToExplicitLimit(cap) - }() - } - close(start) - wg.Wait() - close(errs) - for raiseErr := range errs { - if raiseErr != nil { - t.Fatal(raiseErr) - } - } - if got := budget.SpillDiskCap(); got != 1200 { - t.Fatalf("concurrent raised spill cap = %d, want 1200", got) - } - if !reservation.Release() || budget.SpillDiskUsed() != 0 { - t.Fatalf("live reservation did not release after cap growth: %+v", budget.Snapshot()) - } - - budget.Close() - if err = budget.raiseSpillDiskCapToExplicitLimit(1300); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed budget raise error = %v, want %v", err, ErrHashBuildBudgetClosed) - } - if got := budget.SpillDiskCap(); got != 1200 { - t.Fatalf("closed budget changed spill cap to %d", got) - } -} - -func TestOpenProcessGenerationClampsStaleResolvedCapAtomically(t *testing.T) { - budget := MustNewHashBuildBudget(100, 100) - if err := budget.UpdateAggregateCap(40); err != nil { - t.Fatal(err) - } - - generation, err := budget.openProcessGeneration(1, 100, 0) - if err != nil { - t.Fatal(err) - } - if generation.Cap() != 40 { - t.Fatalf("generation cap = %d, want current aggregate cap 40", - generation.Cap()) - } - if generation.SpillDiskCap() != defaultSpillCap(40) { - t.Fatalf("spill disk cap = %d, want %d", - generation.SpillDiskCap(), defaultSpillCap(40)) - } - - // Explicit public configuration remains strict. Only the process path may - // clamp a ceiling sample that became stale between resolution and opening. - if _, err = budget.OpenGenerationWithCap(2, 100); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("explicit oversized generation cap returned %v", err) - } - generation.Close() - budget.Close() -} - -func TestHashBuildBudgetDefensiveAndProviderFailurePaths(t *testing.T) { - for _, limits := range [][2]uint64{{0, 1}, {1, 0}, {1, 2}} { - if _, err := NewHashBuildBudget(limits[0], limits[1]); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("invalid limits %v returned %v", limits, err) - } - if _, err := NewHashBuildBudgetWithSpillCaps(limits[0], limits[1], 1, 1); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("invalid spill budget limits %v returned %v", limits, err) - } - } - func() { - defer func() { - if recover() == nil { - t.Fatal("MustNewHashBuildBudget did not panic for invalid limits") - } - }() - MustNewHashBuildBudget(0, 0) - }() - - var nilBudget *HashBuildBudget - if _, _, _, err := nilBudget.refreshAggregateCap(false, 0); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("nil refresh error = %v", err) - } - nilBudget.SetAggregateCapProvider(func() (uint64, error) { return 1, nil }) - if err := nilBudget.UpdateAggregateCap(1); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("nil cap update error = %v", err) - } - - b, err := NewHashBuildBudgetWithSpillCaps(math.MaxUint64, math.MaxUint64, 10, 10) - if err != nil { - t.Fatal(err) - } - g, err := b.OpenGenerationWithSpillCaps(1, math.MaxUint64, 10, 10) - if err != nil { - t.Fatal(err) - } - if _, err = b.OpenGenerationWithSpillCaps(2, 0, 0, 0); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("invalid spill generation error = %v", err) - } - disk, err := g.ReserveSpillDisk(5) - if err != nil { - t.Fatal(err) - } - if err = b.SetSpillCaps(4, 10); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("spill cap reduction error = %v", err) - } - disk.Release() - - b.capNow = nil - if err = b.UpdateAggregateCap(math.MaxUint64); err != nil { - t.Fatal(err) - } - b.SetAggregateCapProvider(func() (uint64, error) { return 0, nil }) - if _, err = g.Reserve(1); !errors.Is(err, ErrHashBuildCeilingMissing) { - t.Fatalf("zero provider ceiling error = %v", err) - } - b.capCached = true - b.capRefreshTTL = time.Hour - b.capRefreshAt = time.Now() - b.capRefreshEpoch = 2 - b.capRefreshErr = nil - called := false - b.capProvider = func() (uint64, error) { - called = true - return 1, nil - } - if _, _, refreshed, err := b.refreshAggregateCap(true, 1); err != nil || refreshed || called { - t.Fatalf("concurrent refresh was not reused: refreshed=%v called=%v err=%v", refreshed, called, err) - } - - failing, err := NewHashBuildBudget(100, 100) - if err != nil { - t.Fatal(err) - } - failingGeneration, err := failing.OpenGeneration(1) - if err != nil { - t.Fatal(err) - } - providerErr := errors.New("provider failed") - failing.SetAggregateCapProvider(func() (uint64, error) { return 0, providerErr }) - if _, err = failingGeneration.Reserve(1); !errors.Is(err, providerErr) { - t.Fatalf("reserve provider error = %v", err) - } - token := &HashBuildReservation{ - budget: failing, - generation: failingGeneration, - core: &hashBuildReservationCore{size: 1}, - } - if err = token.Grow(1); !errors.Is(err, providerErr) { - t.Fatalf("grow provider error = %v", err) - } - - forceReserve, err := NewHashBuildBudget(10, 10) - if err != nil { - t.Fatal(err) - } - forceReserve.capRefreshTTL = time.Hour - forceReserveGeneration, err := forceReserve.OpenGeneration(1) - if err != nil { - t.Fatal(err) - } - reserveCalls := 0 - forceReserve.SetAggregateCapProvider(func() (uint64, error) { - reserveCalls++ - if reserveCalls == 1 { - return 5, nil - } - return 0, providerErr - }) - seed, err := forceReserveGeneration.Reserve(1) - if err != nil { - t.Fatal(err) - } - seed.Release() - if _, err = forceReserveGeneration.Reserve(6); !errors.Is(err, providerErr) { - t.Fatalf("forced reserve refresh error = %v", err) - } - - forceGrow, err := NewHashBuildBudget(10, 10) - if err != nil { - t.Fatal(err) - } - forceGrow.capRefreshTTL = time.Hour - forceGrowGeneration, err := forceGrow.OpenGeneration(1) - if err != nil { - t.Fatal(err) - } - growCalls := 0 - forceGrow.SetAggregateCapProvider(func() (uint64, error) { - growCalls++ - if growCalls == 1 { - return 5, nil - } - return 0, providerErr - }) - growToken, err := forceGrowGeneration.Reserve(1) - if err != nil { - t.Fatal(err) - } - if err = growToken.Grow(5); !errors.Is(err, providerErr) { - t.Fatalf("forced grow refresh error = %v", err) - } - growToken.Release() - - rescueGrow, err := NewHashBuildBudget(10, 10) - if err != nil { - t.Fatal(err) - } - rescueGrow.capRefreshTTL = time.Hour - rescueGeneration, err := rescueGrow.OpenGeneration(1) - if err != nil { - t.Fatal(err) - } - rescueCalls := 0 - rescueGrow.SetAggregateCapProvider(func() (uint64, error) { - rescueCalls++ - if rescueCalls == 1 { - return 5, nil - } - return 10, nil - }) - rescueToken, err := rescueGeneration.Reserve(1) - if err != nil { - t.Fatal(err) - } - if err = rescueToken.Grow(5); err != nil || rescueToken.Size() != 6 { - t.Fatalf("forced growth rescue: size=%d err=%v", rescueToken.Size(), err) - } - rescueToken.Release() - - noProvider, err := NewHashBuildBudget(5, 5) - if err != nil { - t.Fatal(err) - } - noProviderGeneration, err := noProvider.OpenGeneration(1) - if err != nil { - t.Fatal(err) - } - noProviderToken, err := noProviderGeneration.Reserve(1) - if err != nil { - t.Fatal(err) - } - if err = noProviderToken.Grow(5); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("non-provider aggregate growth error = %v", err) - } - noProviderToken.Release() - - empty, err := NewHashBuildBudget(100, 100) - if err != nil { - t.Fatal(err) - } - if _, err = empty.resolveCNCapSample(HashBuildCeilingInputs{}); !errors.Is(err, ErrHashBuildCeilingMissing) { - t.Fatalf("empty live sample error = %v", err) - } - previousCap := commonmpool.GlobalCap() - commonmpool.InitCap(2 << 30) - previousHint := fileservice.GlobalMemoryCacheSizeHint.Swap(32 << 20) - func() { - defer commonmpool.InitCap(previousCap) - defer fileservice.GlobalMemoryCacheSizeHint.Store(previousHint) - if _, err = empty.sampleCNCap(); err != nil { - t.Fatalf("live CN sample error = %v", err) - } - }() - empty.capNow = nil - empty.installCNCapProvider(HashBuildCeilingInputs{HostMemTotal: 1 << 30}) - empty.mergeObservedCNCap(HashBuildCeilingInputs{ - CgroupMemoryMax: 512 << 20, - HostMemTotal: 768 << 20, - GlobalMpoolCap: 640 << 20, - FileCacheHint: 32 << 20, - }, 50) - if empty.AggregateCap() != 50 { - t.Fatalf("merged aggregate cap = %d", empty.AggregateCap()) - } - - closedBudget, err := NewHashBuildBudget(10, 10) - if err != nil { - t.Fatal(err) - } - closedGeneration, err := closedBudget.OpenGeneration(1) - if err != nil { - t.Fatal(err) - } - closedGeneration.closed = true - closedBudget.mu.Lock() - _, err, rejected := closedGeneration.reserveLocked(1, true, false) - closedBudget.mu.Unlock() - if rejected || !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed reserveLocked: rejected=%v err=%v", rejected, err) - } - closedToken := &HashBuildReservation{ - budget: closedBudget, - generation: closedGeneration, - core: &hashBuildReservationCore{size: 1}, - } - if err, rejected = closedToken.growLocked(1, true); rejected || !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed growLocked: rejected=%v err=%v", rejected, err) - } - closedToken.core.state.Store(hashBuildReservationReleased) - if err, rejected = closedToken.growLocked(1, true); rejected || !errors.Is(err, ErrHashBuildReservationInactive) { - t.Fatalf("inactive growLocked: rejected=%v err=%v", rejected, err) - } - budgetless := &HashBuildReservation{core: &hashBuildReservationCore{}} - if budgetless.Released() { - t.Fatal("active budgetless token reported released") - } - - overflowBudget, err := NewHashBuildBudget(math.MaxUint64, math.MaxUint64) - if err != nil { - t.Fatal(err) - } - overflowGeneration, err := overflowBudget.OpenGeneration(1) - if err != nil { - t.Fatal(err) - } - overflow := &HashBuildReservation{ - budget: overflowBudget, - generation: overflowGeneration, - core: &hashBuildReservationCore{size: math.MaxUint64}, - } - if err = overflow.Grow(0); err != nil { - t.Fatal(err) - } - if err = overflow.Grow(1); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("overflow growth error = %v", err) - } - overflowDisk := &HashBuildSpillDiskReservation{ - budget: overflowBudget, - generation: overflowGeneration, - core: &hashBuildReservationCore{size: math.MaxUint64}, - } - overflowBudget.spillDiskCap = math.MaxUint64 - overflowGeneration.spillDiskCap = math.MaxUint64 - if err = overflowDisk.Grow(0); err != nil { - t.Fatal(err) - } - if err = overflowDisk.Grow(1); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("overflow disk growth error = %v", err) - } - overflowDisk.core.state.Store(hashBuildReservationReleased) - if err = overflowDisk.Grow(1); !errors.Is(err, ErrHashBuildReservationInactive) { - t.Fatalf("inactive disk growth error = %v", err) - } - cnDiskBudget, err := NewHashBuildBudgetWithSpillCaps(100, 100, 10, 10) - if err != nil { - t.Fatal(err) - } - cnDiskFirst, err := cnDiskBudget.OpenGenerationWithSpillCaps(1, 100, 10, 10) - if err != nil { - t.Fatal(err) - } - cnDiskSecond, err := cnDiskBudget.OpenGenerationWithSpillCaps(2, 100, 10, 10) - if err != nil { - t.Fatal(err) - } - firstDisk, err := cnDiskFirst.ReserveSpillDisk(6) - if err != nil { - t.Fatal(err) - } - secondDisk, err := cnDiskSecond.ReserveSpillDisk(1) - if err != nil { - t.Fatal(err) - } - if err = secondDisk.Grow(4); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("CN disk growth error = %v", err) - } - firstDisk.Release() - secondDisk.Release() - - corruptMemory := &HashBuildReservation{ - budget: overflowBudget, - generation: overflowGeneration, - core: &hashBuildReservationCore{size: 5}, - } - if ok, reconcileErr := corruptMemory.ReconcileDown(0); ok || !errors.Is(reconcileErr, ErrHashBuildReservationInactive) { - t.Fatalf("corrupt memory reconcile: ok=%v err=%v", ok, reconcileErr) - } - if !corruptMemory.Release() || overflowBudget.AggregateUsed() != 0 || overflowGeneration.Used() != 0 { - t.Fatal("defensive memory release did not clamp corrupt counters") - } - - corruptDisk := &HashBuildSpillDiskReservation{ - budget: overflowBudget, - generation: overflowGeneration, - core: &hashBuildReservationCore{size: 5}, - } - if ok, reconcileErr := corruptDisk.ReconcileDown(0); ok || !errors.Is(reconcileErr, ErrHashBuildReservationInactive) { - t.Fatalf("corrupt disk reconcile: ok=%v err=%v", ok, reconcileErr) - } - if !corruptDisk.Release() { - t.Fatal("defensive disk release failed") - } - corruptFD := &HashBuildSpillFDReservation{ - budget: overflowBudget, - generation: overflowGeneration, - core: &hashBuildReservationCore{size: 5}, - } - if ok, reconcileErr := corruptFD.ReconcileDown(0); ok || !errors.Is(reconcileErr, ErrHashBuildReservationInactive) { - t.Fatalf("corrupt FD reconcile: ok=%v err=%v", ok, reconcileErr) - } - if !corruptFD.Release() { - t.Fatal("defensive FD release failed") - } - - largeHint, err := ResolveHashBuildCeiling(HashBuildCeilingInputs{ - HostMemTotal: 10 << 30, - FileCacheHint: 9 << 30, - }) - if err != nil || largeHint.RequestedReserve != 9<<30 { - t.Fatalf("large cache hint ceiling = %+v, err=%v", largeHint, err) - } - tiny, err := ResolveHashBuildCeiling(HashBuildCeilingInputs{HostMemTotal: 128 << 20}) - if err != nil || tiny.CNHashCap == 0 { - t.Fatalf("tiny ceiling = %+v, err=%v", tiny, err) - } - if _, err = ResolveHashBuildCeiling(HashBuildCeilingInputs{HostMemTotal: 1}); !errors.Is(err, ErrHashBuildCeilingMissing) { - t.Fatalf("zero resulting CN cap error = %v", err) - } -} - -func BenchmarkHashBuildBudgetReserveCachedProvider(b *testing.B) { - budget := MustNewHashBuildBudget(uint64(b.N)+1, uint64(b.N)+1) - budget.capRefreshTTL = hashBuildBudgetCapRefreshTTL - gen, _ := budget.OpenGeneration(1) - var calls atomic.Int64 - budget.SetAggregateCapProvider(func() (uint64, error) { - calls.Add(1) - return uint64(b.N) + 1, nil - }) - b.ReportAllocs() - b.ResetTimer() - for i := 0; i < b.N; i++ { - tok, err := gen.Reserve(1) - if err != nil { - b.Fatal(err) - } - tok.Release() - } - b.StopTimer() - b.ReportMetric(float64(calls.Load()), "provider-calls") -} - -func BenchmarkHashBuildBudgetAllocationAccount(b *testing.B) { - const capacity = uint64(1 << 60) - budget := MustNewHashBuildBudget(capacity, capacity) - generation, err := budget.OpenGeneration(1) - if err != nil { - b.Fatal(err) - } - registry, err := commonmpool.NewAllocationAccountRegistry(1, 1) - if err != nil { - b.Fatal(err) - } - account, err := registry.OpenWithController(capacity, generation) - if err != nil { - b.Fatal(err) - } - mp := commonmpool.MustNew("hash-build-allocation-account-benchmark") - defer commonmpool.DeleteMPool(mp) - - b.ReportAllocs() - b.SetBytes(64 << 10) - b.ResetTimer() - for range b.N { - buffer, allocErr := mp.AllocAccounted(64<<10, account, 1, 1) - if allocErr != nil { - b.Fatal(allocErr) - } - mp.Free(buffer) - } - b.StopTimer() - account.Seal() - if _, err = registry.Finalize(account); err != nil { - b.Fatal(err) - } -} - -// BenchmarkHashBuildAllocationAttemptLifecycle covers the production control -// plane around high-frequency statements: concurrent generation open, account -// publication, one physical owner allocation, release, terminal snapshot, and -// slot reuse. The explicit quantiles make the contention tail visible when the -// benchmark is run with -cpu=1,8. -func BenchmarkHashBuildAllocationAttemptLifecycle(b *testing.B) { - const ( - aggregateCap = uint64(1 << 50) - attemptCap = uint64(1 << 20) - ) - budget := MustNewHashBuildBudget(aggregateCap, attemptCap) - registry, err := commonmpool.NewAllocationAccountRegistry(4_096, 4_096) - if err != nil { - b.Fatal(err) - } - mp := commonmpool.MustNew("hash-build-attempt-lifecycle-benchmark") - defer commonmpool.DeleteMPool(mp) - latencies := make([]int64, b.N) - var ( - nextID atomic.Uint64 - sample atomic.Uint64 - ) - - b.ReportAllocs() - b.SetBytes(4 << 10) - b.ResetTimer() - b.RunParallel(func(pb *testing.PB) { - for pb.Next() { - started := time.Now() - generation, openErr := budget.OpenGenerationWithCap( - nextID.Add(1), - attemptCap, - ) - if openErr != nil { - b.Errorf("open generation: %v", openErr) - return - } - account, openErr := registry.OpenWithController(attemptCap, generation) - if openErr != nil { - generation.Close() - b.Errorf("open account: %v", openErr) - return - } - buffer, allocErr := mp.AllocAccounted(4<<10, account, 1, 1) - if allocErr == nil { - mp.Free(buffer) - } - _, _, terminalErr := registry.CompleteTerminal(account) - generation.Close() - if allocErr != nil || terminalErr != nil { - b.Errorf("attempt alloc=%v terminal=%v", allocErr, terminalErr) - return - } - index := sample.Add(1) - 1 - latencies[index] = time.Since(started).Nanoseconds() - } - }) - b.StopTimer() - count := int(sample.Load()) - if count != b.N { - b.Fatalf("completed attempts = %d, want %d", count, b.N) - } - sort.Slice(latencies, func(i, j int) bool { return latencies[i] < latencies[j] }) - b.ReportMetric(float64(latencies[(count-1)*50/100]), "p50-ns/op") - b.ReportMetric(float64(latencies[(count-1)*99/100]), "p99-ns/op") - if budget.AggregateUsed() != 0 || registry.LiveAllocationMetadata() != 0 { - b.Fatalf( - "terminal leak: budget=%d metadata=%d", - budget.AggregateUsed(), - registry.LiveAllocationMetadata(), - ) - } -} - -// BenchmarkHashBuildAllocationReleaseStorm isolates concurrent physical -// alloc/free against one live generation, the shape produced when broadcast -// consumers and spill buffers drain together. -func BenchmarkHashBuildAllocationReleaseStorm(b *testing.B) { - const capacity = uint64(1 << 50) - budget := MustNewHashBuildBudget(capacity, capacity) - generation, err := budget.OpenGeneration(1) - if err != nil { - b.Fatal(err) - } - registry, err := commonmpool.NewAllocationAccountRegistry(1, 65_536) - if err != nil { - b.Fatal(err) - } - account, err := registry.OpenWithController(capacity, generation) - if err != nil { - b.Fatal(err) - } - mp := commonmpool.MustNew("hash-build-release-storm-benchmark") - defer commonmpool.DeleteMPool(mp) - - b.ReportAllocs() - b.SetBytes(4 << 10) - b.ResetTimer() - b.RunParallel(func(pb *testing.PB) { - for pb.Next() { - buffer, allocErr := mp.AllocAccounted(4<<10, account, 1, 1) - if allocErr != nil { - b.Errorf("allocate: %v", allocErr) - return - } - mp.Free(buffer) - } - }) - b.StopTimer() - if account.Snapshot().Used != 0 || generation.Used() != 0 { - b.Fatalf( - "release storm leak: account=%d generation=%d", - account.Snapshot().Used, - generation.Used(), - ) - } - if _, _, err = registry.CompleteTerminal(account); err != nil { - b.Fatal(err) - } - generation.Close() -} - -func TestResolveHashBuildCeiling(t *testing.T) { - const gib = uint64(1 << 30) - got, err := ResolveHashBuildCeiling(HashBuildCeilingInputs{ - CgroupMemoryMax: 20 * gib, - HostMemTotal: 10 * gib, - GlobalMpoolCap: 30 * gib, - FileCacheHint: gib, - ProcessLimitationSize: 2 * gib, - }) - if err != nil { - t.Fatal(err) - } - if got.EffectiveCN != 10*gib || got.RequestedReserve != 4*gib || got.Reserve != 4*gib || got.CNHashCap != 6*gib || got.QueryCap != 2*gib { - t.Fatalf("unexpected ceiling: %+v", got) - } - if _, err = ResolveHashBuildCeiling(HashBuildCeilingInputs{CgroupMemoryMax: math.MaxUint64, HostMemTotal: 0, GlobalMpoolCap: 0}); !errors.Is(err, ErrHashBuildCeilingMissing) { - t.Fatalf("missing finite source error = %v", err) - } - small, err := ResolveHashBuildCeiling(HashBuildCeilingInputs{HostMemTotal: 3 * gib, FileCacheHint: 3 * gib}) - if err != nil || small.CNHashCap != 3*gib/20 { - t.Fatalf("small-CN bounded allowance = %+v, err=%v", small, err) + mp.Free(allocation) } } From edc2731f34576c5405be7d47c629c5f2a05ed390 Mon Sep 17 00:00:00 2001 From: aptend Date: Sat, 1 Aug 2026 11:45:54 +0800 Subject: [PATCH 39/61] executor: close allocation-account CI regressions --- pkg/common/hashmap/strhashmap.go | 5 +-- pkg/common/hashmap/strhashmap_test.go | 4 +-- .../batch/allocation_account_test.go | 36 +++++++++++++++++++ pkg/container/batch/batch.go | 25 +++++++++++-- .../vector/allocation_account_test.go | 9 +++-- pkg/container/vector/vector.go | 4 +-- pkg/vm/engine/tae/txn/txnimpl/anode.go | 19 +++++++++- pkg/vm/engine/tae/txn/txnimpl/table.go | 1 + pkg/vm/engine/test/workspace_test.go | 10 ++++++ 9 files changed, 98 insertions(+), 15 deletions(-) diff --git a/pkg/common/hashmap/strhashmap.go b/pkg/common/hashmap/strhashmap.go index b4ddb84f68a62..e0c14ebefc06c 100644 --- a/pkg/common/hashmap/strhashmap.go +++ b/pkg/common/hashmap/strhashmap.go @@ -155,10 +155,7 @@ func (itr *strHashmapIterator) prepareHashKeys( } values, area := vector.MustVarlenaRawData(vec) for i := 0; i < count; i++ { - value := values[start+i].ByteSlice() - if area != nil { - value = values[start+i].GetByteSlice(area) - } + value := values[start+i].GetByteSlice(area) if err := add(i, prefix+4+len(value)); err != nil { return err } diff --git a/pkg/common/hashmap/strhashmap_test.go b/pkg/common/hashmap/strhashmap_test.go index 25b74b3610d85..b671c9b4ea97f 100644 --- a/pkg/common/hashmap/strhashmap_test.go +++ b/pkg/common/hashmap/strhashmap_test.go @@ -463,8 +463,8 @@ func TestHashMapIteratorsRejectMalformedRowShapes(t *testing.T) { require.NoError(t, err) for _, vecs := range [][]*vector.Vector{ nil, - []*vector.Vector{nil}, - []*vector.Vector{short}, + {nil}, + {short}, } { _, _, err = iterator.Insert(0, 2, vecs) require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) diff --git a/pkg/container/batch/allocation_account_test.go b/pkg/container/batch/allocation_account_test.go index a2861f63390f3..723d38d2ef69c 100644 --- a/pkg/container/batch/allocation_account_test.go +++ b/pkg/container/batch/allocation_account_test.go @@ -439,6 +439,42 @@ func TestMixedBatchAllocationClonePreservesVectorProvenance(t *testing.T) { finalizeTestBatchAllocationAccount(t, state) } +func TestBatchImplicitClonesPromoteAccountedVectorsOffHeap(t *testing.T) { + state := newTestBatchAllocationAccount(t, 16) + mp := mpool.MustNewZero() + source := NewWithSize(2) + source.Attrs = []string{"accounted", "unaccounted"} + source.Vecs[0] = vector.NewOffHeapVecWithType(types.T_int64.ToType()) + require.NoError(t, source.Vecs[0].SetAllocationAccount(state.selection)) + source.Vecs[1] = vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixed(source.Vecs[0], int64(1), false, mp)) + require.NoError(t, vector.AppendFixed(source.Vecs[1], int64(2), false, mp)) + source.SetRowCount(1) + sourceUsed := state.account.Snapshot().Used + + dup, err := source.Dup(mp) + require.NoError(t, err) + require.True(t, dup.offHeap) + require.Same(t, state.selection, dup.Vecs[0].AllocationAccountSelection()) + require.Nil(t, dup.Vecs[1].AllocationAccountSelection()) + dup.Clean(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + selected, err := source.CloneSelectedColumns( + []int{0}, + []string{"accounted"}, + mp, + ) + require.NoError(t, err) + require.True(t, selected.offHeap) + require.Same(t, state.selection, selected.Vecs[0].AllocationAccountSelection()) + selected.Clean(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + func TestMixedBatchAllocationBatchSetPreservesVectorProvenance(t *testing.T) { state := newTestBatchAllocationAccount(t, 128) mp := mpool.MustNewZero() diff --git a/pkg/container/batch/batch.go b/pkg/container/batch/batch.go index 880646822327b..e37767ace8fcd 100644 --- a/pkg/container/batch/batch.go +++ b/pkg/container/batch/batch.go @@ -928,10 +928,10 @@ func (bat *Batch) CloneSelectedColumns( ) (cloned *Batch, err error) { cloned = NewWithSize(len(selectCols)) cloned.Attrs = selectAttrs - cloned.offHeap = bat.offHeap + cloned.offHeap = bat.offHeap || bat.selectedColumnsHaveAllocationAccount(selectCols) var typ types.Type for idx := range selectCols { - if bat.offHeap { + if cloned.offHeap { cloned.Vecs[idx] = vector.NewOffHeapVecWithType(typ) } else { cloned.Vecs[idx] = vector.NewVec(typ) @@ -1155,7 +1155,26 @@ func (bat *Batch) CloneTo(toBat *Batch, mp *mpool.MPool) (err error) { // Dup used to copy a Batch object, this method will create a new batch // and copy all vectors (Vecs) of the current batch to the new batch. func (bat *Batch) Dup(mp *mpool.MPool) (*Batch, error) { - return bat.Clone(mp, bat.offHeap) + return bat.Clone(mp, bat.offHeap || bat.hasAllocationAccountVector()) +} + +func (bat *Batch) hasAllocationAccountVector() bool { + for _, vec := range bat.Vecs { + if vec != nil && vec.AllocationAccountSelection() != nil { + return true + } + } + return false +} + +func (bat *Batch) selectedColumnsHaveAllocationAccount(selectCols []int) bool { + for _, sourceIdx := range selectCols { + vec := bat.Vecs[sourceIdx] + if vec != nil && vec.AllocationAccountSelection() != nil { + return true + } + } + return false } func (bat *Batch) Union(bat2 *Batch, sels []int64, m *mpool.MPool) error { diff --git a/pkg/container/vector/allocation_account_test.go b/pkg/container/vector/allocation_account_test.go index 793e8b8a7ca98..b9be24e07544c 100644 --- a/pkg/container/vector/allocation_account_test.go +++ b/pkg/container/vector/allocation_account_test.go @@ -121,8 +121,13 @@ func TestVectorAllocationAccountConfiguration(t *testing.T) { require.Panics(t, func() { vec.SetOffHeap(false) }) - _, err = vec.Dup(mp) - require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + sourceUsed := state.account.Snapshot().Used + dup, err := vec.Dup(mp) + require.NoError(t, err) + require.Same(t, state.selection, dup.AllocationAccountSelection()) + require.Greater(t, state.account.Snapshot().Used, sourceUsed) + dup.Free(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) _, err = vec.CloneToFlatCompact(mp) require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) diff --git a/pkg/container/vector/vector.go b/pkg/container/vector/vector.go index e8cde3a52d936..ed3645c369f6c 100644 --- a/pkg/container/vector/vector.go +++ b/pkg/container/vector/vector.go @@ -1570,9 +1570,7 @@ func (v *Vector) PreExtendWithArea(rows int, extraAreaSize int, mp *mpool.MPool) // Dup use to copy an identical vector func (v *Vector) Dup(mp *mpool.MPool) (*Vector, error) { if v.allocationAccount != nil { - return nil, allocationAccountInvalid( - "accounted vector duplication requires an off-heap destination", - ) + return v.dup(mp, true, true, v.allocationAccount) } return v.dup(mp, false, v.offHeap, nil) } diff --git a/pkg/vm/engine/tae/txn/txnimpl/anode.go b/pkg/vm/engine/tae/txn/txnimpl/anode.go index b41477515ab64..05990d6aa7ef8 100644 --- a/pkg/vm/engine/tae/txn/txnimpl/anode.go +++ b/pkg/vm/engine/tae/txn/txnimpl/anode.go @@ -19,6 +19,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/nulls" + cnvector "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/objectio" "github.com/matrixorigin/matrixone/pkg/vm/engine/tae/catalog" "github.com/matrixorigin/matrixone/pkg/vm/engine/tae/common" @@ -118,6 +119,7 @@ func (n *anode) Append(data *containers.Batch, offset uint32) (an uint32, err er from := uint32(n.data.Length()) an = n.PrepareAppend(data, offset) + fakePKProvided := false for _, attr := range data.Attrs { if attr == catalog.PhyAddrColumnName { continue @@ -128,12 +130,27 @@ func (n *anode) Append(data *containers.Batch, offset uint32) (an uint32, err er // } // } def := schema.ColDefs[schema.GetColIdx(attr)] + fakePKProvided = fakePKProvided || def.FakePK destVec := n.data.Vecs[def.Idx] // logutil.Infof("destVec: %s, %d, %d", destVec.String(), cnt, data.Length()) destVec.ExtendWithOffset(data.Vecs[def.Idx], int(offset), int(an)) } + if schema.HasFakePK() && !fakePKProvided { + fakePK := n.data.Vecs[schema.GetPrimaryKey().Idx] + if err = cnvector.AppendMultiFixed( + fakePK.GetDownstreamVector(), + uint64(0), + false, + int(an), + fakePK.GetAllocator(), + ); err != nil { + return + } + } + if err = n.FillPhyAddrColumn(from, an); err != nil { + return + } n.rows = uint32(n.data.Length()) - err = n.FillPhyAddrColumn(from, an) return } diff --git a/pkg/vm/engine/tae/txn/txnimpl/table.go b/pkg/vm/engine/tae/txn/txnimpl/table.go index d334cd820d779..1bfec495efafb 100644 --- a/pkg/vm/engine/tae/txn/txnimpl/table.go +++ b/pkg/vm/engine/tae/txn/txnimpl/table.go @@ -558,6 +558,7 @@ func (tbl *txnTable) TransferDeletes( } deletes.Vecs[i].CompactByBitmap(&transferd) } + tbl.tombstoneTable.tableSpace.node.rows = uint32(deletes.Length()) return } diff --git a/pkg/vm/engine/test/workspace_test.go b/pkg/vm/engine/test/workspace_test.go index 5946fd9176e7f..ec4c0f96cf157 100644 --- a/pkg/vm/engine/test/workspace_test.go +++ b/pkg/vm/engine/test/workspace_test.go @@ -1079,10 +1079,15 @@ func Test_MultiTxnInsertDelete(t *testing.T) { { require.NoError(t, testutil.WriteToRelation(ctx, txn, relation, bat2, false, true)) + localPKs := vector.MustFixedColWithTypeCheck[int64](bat2.Vecs[primaryKeyIdx]) + localPKOffset := 0 txn.GetWorkspace().(*disttae.Transaction).ForEachTableWrites( relation.GetDBID(ctx), relation.GetTableID(ctx), 1, func(entry disttae.Entry) { waitedDeletes := vector.MustFixedColWithTypeCheck[types.Rowid](entry.Bat().GetVector(0)) require.NoError(t, vector.AppendFixedList[types.Rowid](tombstoneBat.Vecs[0], waitedDeletes, nil, mp)) + require.NoError(t, vector.AppendFixedList[int64]( + tombstoneBat.Vecs[1], localPKs[localPKOffset:localPKOffset+len(waitedDeletes)], nil, mp)) + localPKOffset += len(waitedDeletes) tombstoneBat.SetRowCount(tombstoneBat.RowCount() + len(waitedDeletes)) }) @@ -1882,10 +1887,15 @@ func Test_MultiTxnRollbackStatement(t *testing.T) { require.NoError(t, txn.GetWorkspace().RollbackLastStatement(ctx)) require.NoError(t, txn.GetWorkspace().IncrStatementID(ctx, false)) + localPKs := vector.MustFixedColWithTypeCheck[int64](bat2.Vecs[primaryKeyIdx]) + localPKOffset := 0 txn.GetWorkspace().(*disttae.Transaction).ForEachTableWrites( relation.GetDBID(ctx), relation.GetTableID(ctx), 1, func(entry disttae.Entry) { waitedDeletes := vector.MustFixedColWithTypeCheck[types.Rowid](entry.Bat().GetVector(0)) require.NoError(t, vector.AppendFixedList[types.Rowid](tombstoneBat.Vecs[0], waitedDeletes, nil, mp)) + require.NoError(t, vector.AppendFixedList[int64]( + tombstoneBat.Vecs[1], localPKs[localPKOffset:localPKOffset+len(waitedDeletes)], nil, mp)) + localPKOffset += len(waitedDeletes) tombstoneBat.SetRowCount(tombstoneBat.RowCount() + len(waitedDeletes)) }) From c7ee02b6e42e47c5e88d305dfb1090f309441602 Mon Sep 17 00:00:00 2001 From: aptend Date: Sat, 1 Aug 2026 15:26:00 +0800 Subject: [PATCH 40/61] fix: account expression and join result storage --- ...ocation_accounted_memory_admission_impl.md | 28 +- .../design/evidence/26459_local_validation.md | 18 +- ...0_allocation_accounted_memory_admission.md | 25 +- pkg/container/vector/functionTool_test.go | 28 ++ pkg/container/vector/functionTools.go | 46 ++- .../dedupjoin/allocation_test_helpers_test.go | 87 +++++ pkg/sql/colexec/dedupjoin/join.go | 69 +++- pkg/sql/colexec/dedupjoin/types.go | 36 +- pkg/sql/colexec/evalExpression.go | 355 +++++++++++++----- .../colexec/evalExpressionAllocation_test.go | 106 ++++++ pkg/sql/colexec/evalExpressionMemory.go | 131 ------- pkg/sql/colexec/evalExpressionMemory_test.go | 156 -------- .../colexec/hashbuild/expression_memory.go | 29 +- pkg/sql/colexec/hashbuild/hashmap.go | 1 + pkg/sql/colexec/hashbuild/hashmap_test.go | 53 ++- pkg/sql/colexec/hashbuild/spill.go | 1 + .../hashjoin/allocation_test_helpers_test.go | 46 +++ pkg/sql/colexec/hashjoin/join.go | 24 +- pkg/sql/colexec/hashjoin/mark_spill_test.go | 2 +- pkg/sql/colexec/hashjoin/types.go | 29 +- pkg/sql/colexec/loopjoin/join.go | 51 +-- pkg/sql/colexec/loopjoin/join_test.go | 42 +++ pkg/sql/colexec/loopjoin/types.go | 28 +- .../allocation_test_helpers_test.go | 46 +++ pkg/sql/colexec/rightdedupjoin/join.go | 14 +- pkg/sql/colexec/rightdedupjoin/types.go | 53 ++- .../spillutil/allocation_account_test.go | 34 +- pkg/sql/colexec/spillutil/join_spill.go | 1 + pkg/sql/util/eval_expr_util.go | 23 +- 29 files changed, 1094 insertions(+), 468 deletions(-) create mode 100644 pkg/sql/colexec/evalExpressionAllocation_test.go delete mode 100644 pkg/sql/colexec/evalExpressionMemory.go delete mode 100644 pkg/sql/colexec/evalExpressionMemory_test.go diff --git a/docs/design/allocation_accounted_memory_admission_impl.md b/docs/design/allocation_accounted_memory_admission_impl.md index 89cd9873f685c..13e8015fabb5f 100644 --- a/docs/design/allocation_accounted_memory_admission_impl.md +++ b/docs/design/allocation_accounted_memory_admission_impl.md @@ -29,21 +29,27 @@ family: - hash table cells, descriptors, iterator keys, and selection lists; - copied build batches and retained unique keys; - JoinMap-owned batches and grouping metadata; -- join matched/capture/result state; +- HashJoin, LoopJoin, DedupJoin, and RightDedupJoin result/finalize state; +- join matched/capture state; - Product result state; - runtime-filter payloads until ownership transfer; - spill encode/decode/scatter buffers and rebuilt retained state. -ExpressionExecutor results, caches, and library-internal Go heap objects are -not HashBuild-retained storage. They remain in the existing MPool/Go runtime -domain. When an expression result is copied into a retained batch, key, or -result vector, that destination allocation is charged. This boundary avoids a -misleading partial "exact expression memory" account for regexp, JSON, JQ, and -other libraries that do not expose allocator/free hooks. - -The account is therefore not advertised as total query RSS. General transient -expression admission is a separate problem and must not be represented by an -estimated charge inside this exact, terminal-zero ledger. +HashBuild and join expression trees receive the same account before +construction. Their fixed/variable values, nested function results, selected +row buffers, list results, and other owned MPool vectors allocate directly +through it. Function-result wrappers retain the immutable selection even when +the current vector is transferred, so later reuse cannot fall back to an +unaccounted vector. Borrowed input and serialized-plan vectors keep their +source ownership. + +Library-internal Go heap objects remain outside the controlled domain. This +explicit boundary avoids pretending that regexp, JSON, JQ, and other libraries +with no allocator/free hooks are exactly charged. + +The account is therefore not advertised as total query RSS. Unobservable +library allocations must not be represented by an estimated charge inside +this exact, terminal-zero ledger. ProductL2 scratch and native CPU/GPU index storage are also outside this first controlled domain. ProductL2 still consumes an accounted JoinMap: those source diff --git a/docs/design/evidence/26459_local_validation.md b/docs/design/evidence/26459_local_validation.md index 673d75aa60b5a..60cffb99a851e 100644 --- a/docs/design/evidence/26459_local_validation.md +++ b/docs/design/evidence/26459_local_validation.md @@ -8,7 +8,8 @@ are not retained as validation dimensions. - no production allocation-account enable switch; - no HashBuild logical-size memory reservation token; -- no join/HashBuild expression account pretending to cover library Go heap; +- every join/HashBuild expression-owned MPool vector is constructed with the + attempt account, while opaque library Go heap remains an explicit boundary; - SpillEngine construction rejects a missing or closed budget generation; - runtime scan/load clones are attached to the current attempt before worker `Prepare`; @@ -24,7 +25,8 @@ link flags, and runtime paths match the MatrixOne build contract. .agents/skills/mo-dev/scripts/mo-cgo-test -count=1 -timeout=240s \ ./pkg/common/mpool ./pkg/common/bitmap ./pkg/common/hashmap/... \ ./pkg/container/vector ./pkg/container/batch ./pkg/vm/message \ - ./pkg/vm/process ./pkg/sql/colexec/hashbuild \ + ./pkg/vm/process ./pkg/sql/util ./pkg/sql/colexec \ + ./pkg/sql/colexec/hashbuild \ ./pkg/sql/colexec/hashjoin ./pkg/sql/colexec/dedupjoin \ ./pkg/sql/colexec/rightdedupjoin ./pkg/sql/colexec/loopjoin \ ./pkg/sql/colexec/product ./pkg/sql/colexec/productl2 \ @@ -35,8 +37,11 @@ Selected race coverage: ```text .agents/skills/mo-dev/scripts/mo-cgo-test -race -count=1 -timeout=300s \ - ./pkg/common/mpool ./pkg/sql/colexec/hashjoin \ - ./pkg/container/pSpool ./pkg/sql/colexec/productl2 \ + ./pkg/common/mpool ./pkg/container/vector ./pkg/container/pSpool \ + ./pkg/sql/util ./pkg/sql/colexec ./pkg/sql/colexec/hashbuild \ + ./pkg/sql/colexec/hashjoin ./pkg/sql/colexec/loopjoin \ + ./pkg/sql/colexec/dedupjoin ./pkg/sql/colexec/rightdedupjoin \ + ./pkg/sql/colexec/productl2 \ ./pkg/sql/colexec/spillutil ./pkg/sql/colexec/sample ./pkg/sql/compile ``` @@ -78,6 +83,11 @@ The local suite covers: - recursive spill row/schema/file validation; - minimum-unit and monotonic pressure termination; - optional runtime-filter degradation; +- nested/selected expression result accounting, capacity rejection, transfer, + and terminal release; +- HashJoin, LoopJoin, DedupJoin multi-batch finalize, and RightDedupJoin result + accounting, batch-local reuse, capacity rejection, and prepared + `Reset -> ClearAllocationAccount` terminal release; - grouping-aware copy, hash, equality, and ordering; - late and alternating Sample grouping domains across row, percent, and merge modes; diff --git a/docs/rfcs/00000000_allocation_accounted_memory_admission.md b/docs/rfcs/00000000_allocation_accounted_memory_admission.md index 2d5c6b341e719..250ff2faa729a 100644 --- a/docs/rfcs/00000000_allocation_accounted_memory_admission.md +++ b/docs/rfcs/00000000_allocation_accounted_memory_admission.md @@ -91,11 +91,17 @@ its join consumers: hash tables, copied build batches, JoinMap state, retained keys, join bitmaps/capture/result state, Product result state, runtime-filter payloads, and spill encode/decode/rebuild buffers. -ExpressionExecutor temporary results, caches, and library-internal Go heap -remain in the existing MPool/Go runtime domain. They cannot truthfully be put -in an exact terminal-zero account while regexp, JSON, JQ, and similar libraries -do not expose allocator/free hooks. An expression value becomes accounted when -it is physically copied into retained HashBuild/join storage. +Every MPool-backed vector owned by a HashBuild/join ExpressionExecutor tree is +inside the account: fixed and variable values, reusable function results, +nested child results, selected-row buffers, and list results. Allocation +provenance is installed when the tree is constructed, before the first owned +capacity allocation. Borrowed input and serialized-plan vectors keep their +source ownership. + +Library-internal Go heap used by regexp, JSON, JQ, and similar implementations +remains outside this controlled domain because those libraries do not expose +allocator/free hooks. The account is exact for its stated allocator-visible +domain; it is not advertised as total expression memory or total query RSS. This is a static ownership boundary, not a runtime fallback. A future general expression-memory design must either use allocator-aware implementations or a @@ -221,11 +227,12 @@ also creates a second release owner beside MPool. Sampling is observational, process-wide, and too late to authorize a specific allocation. -### Account only selected expression buffers +### Estimate or manually charge selected expression buffers -This reports false exactness while opaque library Go heap remains untracked. -The implementation therefore accounts retained copies, not partial expression -internals. +This creates another ledger and still reports false exactness. The +implementation instead propagates one immutable allocation selection through +the complete expression tree and charges every physical MPool capacity change; +opaque library Go heap remains an explicit boundary rather than an estimate. ### Keep old and new production paths behind a switch diff --git a/pkg/container/vector/functionTool_test.go b/pkg/container/vector/functionTool_test.go index 6b3ba1f7e5cf4..73c66c9703134 100644 --- a/pkg/container/vector/functionTool_test.go +++ b/pkg/container/vector/functionTool_test.go @@ -26,6 +26,34 @@ import ( "github.com/stretchr/testify/require" ) +func TestFunctionResultAllocationSurvivesVectorTransfer(t *testing.T) { + mp := mpool.MustNewZeroNoFixed() + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + wrapper, err := NewFunctionResultWrapperWithAllocation( + types.T_int64.ToType(), mp, selection, + ) + require.NoError(t, err) + require.NoError(t, wrapper.PreExtendAndReset(64)) + transferred := wrapper.GetResultVector() + require.Same(t, selection, transferred.AllocationAccountSelection()) + firstUsed := account.Snapshot().Used + require.Positive(t, firstUsed) + + wrapper.SetResultVector(nil) + require.NoError(t, wrapper.PreExtendAndReset(64)) + require.Same(t, selection, wrapper.GetResultVector().AllocationAccountSelection()) + require.Greater(t, account.Snapshot().Used, firstUsed) + + transferred.Free(mp) + wrapper.Free() + require.Zero(t, account.Snapshot().Used) +} + func TestAppendByteJsonUsesStorageCompatibleTypeCodes(t *testing.T) { mp := mpool.MustNewZeroNoFixed() wrapper := NewFunctionResultWrapper(types.T_json.ToType(), mp) diff --git a/pkg/container/vector/functionTools.go b/pkg/container/vector/functionTools.go index 9605a068a0ad4..438613bfd6089 100644 --- a/pkg/container/vector/functionTools.go +++ b/pkg/container/vector/functionTools.go @@ -571,9 +571,10 @@ func OptGetBytesParamFromWrapper(wrapper FunctionResultWrapper, idx int, src *Ve var _ FunctionResultWrapper = &FunctionResult[int64]{} type FunctionResult[T types.FixedSizeT] struct { - typ types.Type - vec *Vector - mp *mpool.MPool + typ types.Type + vec *Vector + mp *mpool.MPool + allocation *AllocationAccountSelection isVarlena bool cols []T @@ -621,7 +622,11 @@ func (fr *FunctionResult[T]) getConvenientParamList() []reusableParameterWrapper func (fr *FunctionResult[T]) PreExtendAndReset(targetSize int) error { if fr.vec == nil { - fr.vec = NewOffHeapVecWithType(fr.typ) + var err error + fr.vec, err = NewOffHeapVecWithTypeAndAllocation(fr.typ, fr.allocation) + if err != nil { + return err + } } oldLength := fr.vec.Length() @@ -631,6 +636,9 @@ func (fr *FunctionResult[T]) PreExtendAndReset(targetSize int) error { return err } } + if err := fr.vec.PreExtendNulls(targetSize, fr.mp); err != nil { + return err + } fr.vec.ResetWithSameType() if !fr.isVarlena { @@ -760,6 +768,14 @@ func (fr *FunctionResult[T]) Free() { fr.convenientParam = nil } +func (fr *FunctionResult[T]) setAllocation(selection *AllocationAccountSelection) { + fr.allocation = selection +} + +type functionResultAllocationSetter interface { + setAllocation(*AllocationAccountSelection) +} + func NewFunctionResultWrapper(typ types.Type, mp *mpool.MPool) FunctionResultWrapper { if typ.IsVarlen() { return newResultFunc[types.Varlena](typ, mp) @@ -819,3 +835,25 @@ func NewFunctionResultWrapper(typ types.Type, mp *mpool.MPool) FunctionResultWra } panic(fmt.Sprintf("unexpected type %s for function result", typ)) } + +// NewFunctionResultWrapperWithAllocation constructs a reusable result whose +// current and future vectors allocate through selection. The selection stays +// with the wrapper when EvalWithoutResultReusing transfers its current vector. +func NewFunctionResultWrapperWithAllocation( + typ types.Type, + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (FunctionResultWrapper, error) { + if selection != nil { + if err := selection.validate(); err != nil { + return nil, err + } + } + result := NewFunctionResultWrapper(typ, mp) + setter, ok := result.(functionResultAllocationSetter) + if !ok { + return nil, mpool.ErrAllocationAccountInvariant + } + setter.setAllocation(selection) + return result, nil +} diff --git a/pkg/sql/colexec/dedupjoin/allocation_test_helpers_test.go b/pkg/sql/colexec/dedupjoin/allocation_test_helpers_test.go index ef44e8bb95c2e..b0c6f5132c28e 100644 --- a/pkg/sql/colexec/dedupjoin/allocation_test_helpers_test.go +++ b/pkg/sql/colexec/dedupjoin/allocation_test_helpers_test.go @@ -17,7 +17,13 @@ package dedupjoin import ( "testing" + "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/testutil" "github.com/stretchr/testify/require" ) @@ -36,3 +42,84 @@ func installTestAllocation(t testing.TB, owners ...testAllocationOwner) *mpool.A } return account } + +func TestDedupJoinResultAndFinalizeBatchesUseAllocationAccount(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + arg := &DedupJoin{ + Result: []colexec.ResultPos{{Rel: 0, Pos: 0}}, + LeftTypes: []types.Type{types.T_int64.ToType()}, + RightTypes: []types.Type{types.T_int64.ToType()}, + } + account := installTestAllocation(t, arg) + require.NoError(t, arg.resetRBat()) + require.Same(t, arg.resultAllocation, arg.ctr.rbat.Vecs[0].AllocationAccountSelection()) + + arg.ctr.matched = &bitmap.Bitmap{} + first := batch.NewWithSize(0) + first.SetRowCount(colexec.DefaultBatchSize) + second := batch.NewWithSize(0) + second.SetRowCount(1) + arg.ctr.batches = []*batch.Batch{first, second} + require.NoError(t, arg.ctr.finalize(arg, proc)) + require.Len(t, arg.ctr.buf, 2) + for _, result := range arg.ctr.buf { + require.Same(t, arg.resultAllocation, result.Vecs[0].AllocationAccountSelection()) + } + require.Positive(t, account.Snapshot().Used) + + arg.Reset(proc, false, nil) + require.Nil(t, arg.ctr.rbat) + require.Empty(t, arg.ctr.buf) + require.Zero(t, account.Snapshot().Used) + require.NoError(t, arg.ClearAllocationAccount(account)) +} + +func TestDedupJoinResultAndFinalizeBatchesHonorAllocationCapacity(t *testing.T) { + newAccount := func(t *testing.T) *mpool.AllocationAccount { + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1) + require.NoError(t, err) + return account + } + + t.Run("probe result", func(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + account := newAccount(t) + arg := &DedupJoin{ + Result: []colexec.ResultPos{{Rel: 0, Pos: 0}}, + LeftTypes: []types.Type{types.T_int64.ToType()}, + } + require.NoError(t, arg.SetAllocationAccount(account)) + require.NoError(t, arg.resetRBat()) + err := vector.AppendFixed(arg.ctr.rbat.Vecs[0], int64(1), false, proc.Mp()) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Zero(t, account.Snapshot().Used) + arg.Reset(proc, false, nil) + require.NoError(t, arg.ClearAllocationAccount(account)) + }) + + t.Run("multi batch finalize", func(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + account := newAccount(t) + arg := &DedupJoin{ + Result: []colexec.ResultPos{{Rel: 0, Pos: 0}}, + LeftTypes: []types.Type{types.T_int64.ToType()}, + } + require.NoError(t, arg.SetAllocationAccount(account)) + arg.ctr.matched = &bitmap.Bitmap{} + first := batch.NewWithSize(0) + first.SetRowCount(colexec.DefaultBatchSize) + second := batch.NewWithSize(0) + second.SetRowCount(1) + arg.ctr.batches = []*batch.Batch{first, second} + err := arg.ctr.finalize(arg, proc) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Zero(t, account.Snapshot().Used) + arg.Reset(proc, false, nil) + require.NoError(t, arg.ClearAllocationAccount(account)) + }) +} diff --git a/pkg/sql/colexec/dedupjoin/join.go b/pkg/sql/colexec/dedupjoin/join.go index 6c2b428ca3a01..d6dbebb91babf 100644 --- a/pkg/sql/colexec/dedupjoin/join.go +++ b/pkg/sql/colexec/dedupjoin/join.go @@ -154,6 +154,7 @@ func (dedupJoin *DedupJoin) Prepare(proc *process.Process) (err error) { evalExecs, err = hashbuild.NewExpressionExecutors( proc, dedupJoin.Conditions[0], + dedupJoin.allocationAccount, ) if err != nil { return err @@ -163,6 +164,7 @@ func (dedupJoin *DedupJoin) Prepare(proc *process.Process) (err error) { updateExecs, err = hashbuild.NewExpressionExecutors( proc, dedupJoin.UpdateColExprList, + dedupJoin.allocationAccount, ) if err != nil { for _, exec := range evalExecs { @@ -473,6 +475,7 @@ func (ctr *container) initCaptureBuffers(ap *DedupJoin, proc *process.Process) e return nil } func (ctr *container) finalize(ap *DedupJoin, proc *process.Process) error { + var err error if ap.needsFinalizeMerge() { if !ap.IsMerger { if ap.Mailbox == nil { @@ -605,12 +608,18 @@ func (ctr *container) finalize(ap *DedupJoin, proc *process.Process) error { cIdx := ctr.captureResultIdx[j] if ctr.captured != nil && ctr.captured.Count() > 0 { typ := ap.RightTypes[rp.Pos] - ap.ctr.buf[i].Vecs[j] = vector.NewOffHeapVecWithType(typ) + ap.ctr.buf[i].Vecs[j], err = ap.newResultVector(typ) + if err != nil { + return err + } if err := ap.ctr.buf[i].Vecs[j].UnionBatch(ctr.capturedVecs[cIdx], capOffset, batSize, nil, proc.Mp()); err != nil { return err } } else { - ap.ctr.buf[i].Vecs[j] = vector.NewOffHeapVecWithType(ap.RightTypes[rp.Pos]) + ap.ctr.buf[i].Vecs[j], err = ap.newResultVector(ap.RightTypes[rp.Pos]) + if err != nil { + return err + } if err := vector.AppendMultiFixed(ap.ctr.buf[i].Vecs[j], 0, true, batSize, proc.Mp()); err != nil { return err } @@ -627,13 +636,19 @@ func (ctr *container) finalize(ap *DedupJoin, proc *process.Process) error { // its own valid vector; ownership transfer here would // leave later references reading a nil vector. typ := ap.RightTypes[rp.Pos] - ap.ctr.buf[i].Vecs[j] = vector.NewOffHeapVecWithType(typ) + ap.ctr.buf[i].Vecs[j], err = ap.newResultVector(typ) + if err != nil { + return err + } if err := vector.GetUnionAllFunction(typ, proc.Mp())(ap.ctr.buf[i].Vecs[j], bat.Vecs[rp.Pos]); err != nil { return err } } } else { - ap.ctr.buf[i].Vecs[j] = vector.NewOffHeapVecWithType(ap.LeftTypes[rp.Pos]) + ap.ctr.buf[i].Vecs[j], err = ap.newResultVector(ap.LeftTypes[rp.Pos]) + if err != nil { + return err + } if err := vector.AppendMultiFixed(ap.ctr.buf[i].Vecs[j], 0, true, batSize, proc.Mp()); err != nil { return err } @@ -677,12 +692,18 @@ func (ctr *container) finalize(ap *DedupJoin, proc *process.Process) error { ap.ctr.buf[i] = batch.NewOffHeapWithSize(len(ap.Result)) for j, rp := range ap.Result { if rp.Rel == 1 { - ap.ctr.buf[i].Vecs[j] = vector.NewOffHeapVecWithType(ap.RightTypes[rp.Pos]) + ap.ctr.buf[i].Vecs[j], err = ap.newResultVector(ap.RightTypes[rp.Pos]) + if err != nil { + return err + } if err := unionSelsByBatch(ap.ctr.buf[i].Vecs[j], ctr.batches, rp.Pos, newSels, proc); err != nil { return err } } else { - ap.ctr.buf[i].Vecs[j] = vector.NewOffHeapVecWithType(ap.LeftTypes[rp.Pos]) + ap.ctr.buf[i].Vecs[j], err = ap.newResultVector(ap.LeftTypes[rp.Pos]) + if err != nil { + return err + } if err := vector.AppendMultiFixed(ap.ctr.buf[i].Vecs[j], 0, true, len(newSels), proc.Mp()); err != nil { return err } @@ -708,7 +729,10 @@ func (ctr *container) finalize(ap *DedupJoin, proc *process.Process) error { ap.ctr.buf[batIdx] = batch.NewOffHeapWithSize(len(ap.Result)) for i, rp := range ap.Result { if rp.Rel == 1 { - ap.ctr.buf[batIdx].Vecs[i] = vector.NewOffHeapVecWithType(ap.RightTypes[rp.Pos]) + ap.ctr.buf[batIdx].Vecs[i], err = ap.newResultVector(ap.RightTypes[rp.Pos]) + if err != nil { + return err + } for _, sel := range sels[fillCnt : fillCnt+batSize] { idx1, idx2 := sel/colexec.DefaultBatchSize, sel%colexec.DefaultBatchSize if err := ap.ctr.buf[batIdx].Vecs[i].UnionOne(ctr.batches[idx1].Vecs[rp.Pos], int64(idx2), proc.Mp()); err != nil { @@ -716,7 +740,10 @@ func (ctr *container) finalize(ap *DedupJoin, proc *process.Process) error { } } } else { - ap.ctr.buf[batIdx].Vecs[i] = vector.NewOffHeapVecWithType(ap.LeftTypes[rp.Pos]) + ap.ctr.buf[batIdx].Vecs[i], err = ap.newResultVector(ap.LeftTypes[rp.Pos]) + if err != nil { + return err + } if err := vector.AppendMultiFixed(ap.ctr.buf[batIdx].Vecs[i], 0, true, batSize, proc.Mp()); err != nil { return err } @@ -740,9 +767,15 @@ func (ctr *container) finalize(ap *DedupJoin, proc *process.Process) error { ap.ctr.buf[batIdx] = batch.NewOffHeapWithSize(len(ap.Result)) for i, rp := range ap.Result { if rp.Rel == 1 { - ap.ctr.buf[batIdx].Vecs[i] = vector.NewOffHeapVecWithType(ap.RightTypes[rp.Pos]) + ap.ctr.buf[batIdx].Vecs[i], err = ap.newResultVector(ap.RightTypes[rp.Pos]) + if err != nil { + return err + } } else { - ap.ctr.buf[batIdx].Vecs[i] = vector.NewOffHeapVecWithType(ap.LeftTypes[rp.Pos]) + ap.ctr.buf[batIdx].Vecs[i], err = ap.newResultVector(ap.LeftTypes[rp.Pos]) + if err != nil { + return err + } } } } @@ -837,7 +870,9 @@ func (ctr *container) withRestoredJoinBat1Vectors(updateCols []int32, fn func() } func (ctr *container) probe(bat *batch.Batch, ap *DedupJoin, proc *process.Process, analyzer process.Analyzer, result *vm.CallResult) error { - ap.resetRBat() + if err := ap.resetRBat(); err != nil { + return err + } err := ctr.evalJoinCondition(bat, proc) if err != nil { return err @@ -1055,7 +1090,11 @@ func unionSelsByBatch(dst *vector.Vector, batches []*batch.Batch, colPos int32, } return nil } -func (dedupJoin *DedupJoin) resetRBat() { +func (dedupJoin *DedupJoin) newResultVector(typ types.Type) (*vector.Vector, error) { + return vector.NewOffHeapVecWithTypeAndAllocation(typ, dedupJoin.resultAllocation) +} + +func (dedupJoin *DedupJoin) resetRBat() error { ctr := &dedupJoin.ctr if ctr.rbat != nil { ctr.rbat.CleanOnlyData() @@ -1068,5 +1107,11 @@ func (dedupJoin *DedupJoin) resetRBat() { ctr.rbat.Vecs[i] = vector.NewOffHeapVecWithType(dedupJoin.RightTypes[rp.Pos]) } } + if err := ctr.rbat.SetAllocationAccount(dedupJoin.resultAllocation); err != nil { + ctr.rbat.Clean(nil) + ctr.rbat = nil + return err + } } + return nil } diff --git a/pkg/sql/colexec/dedupjoin/types.go b/pkg/sql/colexec/dedupjoin/types.go index e481c827741b2..7831404cabe6b 100644 --- a/pkg/sql/colexec/dedupjoin/types.go +++ b/pkg/sql/colexec/dedupjoin/types.go @@ -54,6 +54,13 @@ const ( dedupJoinAllocationSiteFinalizeSelections ) +const ( + dedupJoinAllocationSiteResultData mpool.AllocationSite = iota + 110 + dedupJoinAllocationSiteResultArea + dedupJoinAllocationSiteResultNulls + dedupJoinAllocationSiteResultGrouping +) + // WorkerJoinMsg carries per-worker state from non-merger workers to the // merger worker at finalize time. Regular DEDUP JOIN only populates matched; // the REPLACE INTO merged main-table scan path (OldColCapture) additionally @@ -309,6 +316,7 @@ type DedupJoin struct { OldColCaptureProbeIdxList []int32 allocationAccount *mpool.AllocationAccount stateAllocation *vector.AllocationAccountSelection + resultAllocation *vector.AllocationAccountSelection vm.OperatorBase } @@ -337,8 +345,20 @@ func (dedupJoin *DedupJoin) SetAllocationAccount( if err != nil { return err } + resultSelection, err := vector.NewAllocationAccountSelection( + account, + hashbuild.HashBuildAllocationOwner, + dedupJoinAllocationSiteResultData, + dedupJoinAllocationSiteResultArea, + dedupJoinAllocationSiteResultNulls, + dedupJoinAllocationSiteResultGrouping, + ) + if err != nil { + return err + } dedupJoin.allocationAccount = account dedupJoin.stateAllocation = selection + dedupJoin.resultAllocation = resultSelection return nil } @@ -354,11 +374,13 @@ func (dedupJoin *DedupJoin) ClearAllocationAccount( if dedupJoin.ctr.mp != nil || dedupJoin.ctr.spillEngine != nil || len(dedupJoin.ctr.evecs) != 0 || len(dedupJoin.ctr.exprExecs) != 0 || dedupJoin.ctr.matched != nil || dedupJoin.ctr.captured != nil || - len(dedupJoin.ctr.capturedVecs) != 0 { + len(dedupJoin.ctr.capturedVecs) != 0 || dedupJoin.ctr.rbat != nil || + len(dedupJoin.ctr.buf) != 0 { return mpool.ErrAllocationAccountInvariant } dedupJoin.allocationAccount = nil dedupJoin.stateAllocation = nil + dedupJoin.resultAllocation = nil return nil } @@ -430,7 +452,7 @@ func (dedupJoin *DedupJoin) Reset(proc *process.Process, pipelineFailed bool, er } ctr.maxAllocSize = 0 - ctr.cleanBuf(proc) + ctr.cleanResultBatches(proc) ctr.cleanBucketState(proc) ctr.cleanExprExecutor() if ctr.spillEngine != nil { @@ -450,7 +472,7 @@ func (dedupJoin *DedupJoin) Free(proc *process.Process, pipelineFailed bool, err // reopened prepared-pipeline generation stopped from Free. dedupJoin.Mailbox.drain(proc) } - ctr.cleanBuf(proc) + ctr.cleanResultBatches(proc) ctr.cleanBucketState(proc) ctr.cleanBatch(proc) ctr.cleanExprExecutor() @@ -483,6 +505,14 @@ func (ctr *container) cleanBuf(proc *process.Process) { ctr.buf = nil } +func (ctr *container) cleanResultBatches(proc *process.Process) { + ctr.cleanBuf(proc) + if ctr.rbat != nil { + ctr.rbat.Clean(proc.GetMPool()) + ctr.rbat = nil + } +} + func (ctr *container) cleanCaptured(proc *process.Process) { for _, v := range ctr.capturedVecs { if v != nil { diff --git a/pkg/sql/colexec/evalExpression.go b/pkg/sql/colexec/evalExpression.go index 01d9329d00f73..2d1018771b4db 100644 --- a/pkg/sql/colexec/evalExpression.go +++ b/pkg/sql/colexec/evalExpression.go @@ -92,9 +92,20 @@ type ExpressionExecutor interface { } func NewExpressionExecutorsFromPlanExpressions(proc *process.Process, planExprs []*plan.Expr) (executors []ExpressionExecutor, err error) { + return NewExpressionExecutorsFromPlanExpressionsWithAllocation(proc, planExprs, nil) +} + +// NewExpressionExecutorsFromPlanExpressionsWithAllocation builds a complete +// expression tree whose owned MPool vectors use one immutable allocation +// selection. Borrowed input vectors remain owned and charged by their source. +func NewExpressionExecutorsFromPlanExpressionsWithAllocation( + proc *process.Process, + planExprs []*plan.Expr, + selection *vector.AllocationAccountSelection, +) (executors []ExpressionExecutor, err error) { executors = make([]ExpressionExecutor, len(planExprs)) for i := range executors { - executors[i], err = NewExpressionExecutor(proc, planExprs[i]) + executors[i], err = NewExpressionExecutorWithAllocation(proc, planExprs[i], selection) if err != nil { for j := 0; j < i; j++ { executors[j].Free() @@ -106,10 +117,18 @@ func NewExpressionExecutorsFromPlanExpressions(proc *process.Process, planExprs } func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (ExpressionExecutor, error) { + return NewExpressionExecutorWithAllocation(proc, planExpr, nil) +} + +func NewExpressionExecutorWithAllocation( + proc *process.Process, + planExpr *plan.Expr, + selection *vector.AllocationAccountSelection, +) (ExpressionExecutor, error) { switch t := planExpr.Expr.(type) { case *plan.Expr_Lit: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) - vec, err := generateConstExpressionExecutor(proc, typ, t.Lit) + vec, err := generateConstExpressionExecutor(proc, typ, t.Lit, selection) if err != nil { return nil, err } @@ -117,17 +136,21 @@ func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (Expressi case *plan.Expr_T: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) - vec := vector.NewConstNull(typ, 1, proc.Mp()) + vec, err := newExpressionConstNull(typ, 1, selection) + if err != nil { + return nil, err + } return NewFixedVectorExpressionExecutor(proc.Mp(), false, vec), nil case *plan.Expr_Col: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) ce := NewColumnExpressionExecutor() *ce = ColumnExpressionExecutor{ - mp: proc.Mp(), - relIndex: int(t.Col.RelPos), - colIndex: int(t.Col.ColPos), - typ: typ, + mp: proc.Mp(), + relIndex: int(t.Col.RelPos), + colIndex: int(t.Col.ColPos), + typ: typ, + allocation: selection, } // [issue#19574] // if < 0, it's special for agg or others. @@ -138,17 +161,20 @@ func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (Expressi case *plan.Expr_P: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) - return NewParamExpressionExecutor(proc.Mp(), int(t.P.Pos), typ), nil + executor := NewParamExpressionExecutor(proc.Mp(), int(t.P.Pos), typ) + executor.allocation = selection + return executor, nil case *plan.Expr_V: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) ve := NewVarExpressionExecutor() *ve = VarExpressionExecutor{ - mp: proc.Mp(), - name: t.V.Name, - system: t.V.System, - global: t.V.Global, - typ: typ, + mp: proc.Mp(), + name: t.V.Name, + system: t.V.System, + global: t.V.Global, + typ: typ, + allocation: selection, } return ve, nil @@ -156,6 +182,7 @@ func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (Expressi vec := vector.NewVec(types.T_any.ToType()) err := vec.UnmarshalBinary(t.Vec.Data) if err != nil { + vec.Free(proc.Mp()) return nil, err } return NewFixedVectorExpressionExecutor(proc.Mp(), true, vec), nil @@ -164,9 +191,12 @@ func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (Expressi executor := NewListExpressionExecutor() resultVecTyp := t.List.List[0].GetTyp() typ := types.New(types.T(resultVecTyp.Id), resultVecTyp.Width, resultVecTyp.Scale) - executor.Init(proc, typ, len(t.List.List)) + if err := executor.init(proc, typ, len(t.List.List), selection); err != nil { + executor.Free() + return nil, err + } for i := range executor.parameterExecutor { - subExecutor, paramErr := NewExpressionExecutor(proc, t.List.List[i]) + subExecutor, paramErr := NewExpressionExecutorWithAllocation(proc, t.List.List[i], selection) if paramErr != nil { executor.Free() return nil, paramErr @@ -196,13 +226,13 @@ func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (Expressi } typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) - if err = executor.Init(proc, len(t.F.Args), typ); err != nil { + if err = executor.init(proc, len(t.F.Args), typ, selection); err != nil { executor.Free() return nil, err } for i := range executor.parameterExecutor { - subExecutor, paramErr := NewExpressionExecutor(proc, t.F.Args[i]) + subExecutor, paramErr := NewExpressionExecutorWithAllocation(proc, t.F.Args[i], selection) if paramErr != nil { executor.Free() return nil, paramErr @@ -215,6 +245,66 @@ func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (Expressi return nil, moerr.NewNYI(proc.Ctx, fmt.Sprintf("unsupported expression executor for %v now", planExpr)) } +func newExpressionOffHeapVector( + typ types.Type, + selection *vector.AllocationAccountSelection, +) (*vector.Vector, error) { + if selection == nil { + return vector.NewOffHeapVecWithType(typ), nil + } + return vector.NewOffHeapVecWithTypeAndAllocation(typ, selection) +} + +func newExpressionConstNull( + typ types.Type, + length int, + selection *vector.AllocationAccountSelection, +) (*vector.Vector, error) { + if selection == nil { + return vector.NewConstNull(typ, length, nil), nil + } + return vector.NewConstNullWithAllocation(typ, length, selection) +} + +func newExpressionConstFixed[T any]( + typ types.Type, + value T, + length int, + mp *mpool.MPool, + selection *vector.AllocationAccountSelection, +) (*vector.Vector, error) { + if selection == nil { + return vector.NewConstFixed(typ, value, length, mp) + } + return vector.NewConstFixedWithAllocation(typ, value, length, mp, selection) +} + +func newExpressionConstBytes( + typ types.Type, + value []byte, + length int, + mp *mpool.MPool, + selection *vector.AllocationAccountSelection, +) (*vector.Vector, error) { + if selection == nil { + return vector.NewConstBytes(typ, value, length, mp) + } + return vector.NewConstBytesWithAllocation(typ, value, length, mp, selection) +} + +func newExpressionConstArray[T types.ArrayElement]( + typ types.Type, + value []T, + length int, + mp *mpool.MPool, + selection *vector.AllocationAccountSelection, +) (*vector.Vector, error) { + if selection == nil { + return vector.NewConstArray(typ, value, length, mp) + } + return vector.NewConstArrayWithAllocation(typ, value, length, mp, selection) +} + // FixedVectorExpressionExecutor // the content of its vector is fixed. // e.g. @@ -230,7 +320,8 @@ type FixedVectorExpressionExecutor struct { } type FunctionExpressionExecutor struct { - m *mpool.MPool + m *mpool.MPool + allocation *vector.AllocationAccountSelection // resultType is the declared function return type. Some built-ins refine // result metadata (for example temporal scale or decimal width/scale) at // runtime, so reusable result vectors must start each evaluation from this @@ -262,9 +353,10 @@ type FunctionExpressionExecutor struct { } type ColumnExpressionExecutor struct { - mp *mpool.MPool - relIndex int - colIndex int + mp *mpool.MPool + allocation *vector.AllocationAccountSelection + relIndex int + colIndex int // result type. typ types.Type @@ -283,8 +375,9 @@ func (expr *ColumnExpressionExecutor) GetColIndex() int { } type ParamExpressionExecutor struct { - mp *mpool.MPool - null *vector.Vector + mp *mpool.MPool + allocation *vector.AllocationAccountSelection + null *vector.Vector // maskedNull is separate from null/vec because it is not a resolved // parameter value and must never participate in the folded-value cache. maskedNull *vector.Vector @@ -298,7 +391,11 @@ type ParamExpressionExecutor struct { func (expr *ParamExpressionExecutor) Eval(proc *process.Process, batches []*batch.Batch, selectList []bool) (*vector.Vector, error) { if noRowsSelected(selectList, expressionRowCount(batches)) { if expr.maskedNull == nil { - expr.maskedNull = vector.NewConstNull(expr.typ, 1, proc.GetMPool()) + var err error + expr.maskedNull, err = newExpressionConstNull(expr.typ, 1, expr.allocation) + if err != nil { + return nil, err + } } return expr.maskedNull, nil } @@ -319,13 +416,18 @@ func (expr *ParamExpressionExecutor) Eval(proc *process.Process, batches []*batc if val == nil { if expr.null == nil { - expr.null = vector.NewConstNull(expr.typ, 1, proc.GetMPool()) + expr.null, err = newExpressionConstNull(expr.typ, 1, expr.allocation) + if err != nil { + return nil, err + } } return expr.null, nil } if expr.vec == nil { - expr.vec, err = vector.NewConstBytes(expr.typ, val, 1, proc.Mp()) + expr.vec, err = newExpressionConstBytes( + expr.typ, val, 1, proc.Mp(), expr.allocation, + ) } else { err = vector.SetConstBytes(expr.vec, val, 1, proc.GetMPool()) } @@ -373,8 +475,9 @@ func (expr *ParamExpressionExecutor) IsColumnExpr() bool { } type VarExpressionExecutor struct { - mp *mpool.MPool - null *vector.Vector + mp *mpool.MPool + allocation *vector.AllocationAccountSelection + null *vector.Vector // maskedNull lets a skipped variable avoid the resolver without changing // the value cache used by a later selected evaluation. maskedNull *vector.Vector @@ -389,7 +492,11 @@ type VarExpressionExecutor struct { func (expr *VarExpressionExecutor) Eval(proc *process.Process, batches []*batch.Batch, selectList []bool) (*vector.Vector, error) { if noRowsSelected(selectList, expressionRowCount(batches)) { if expr.maskedNull == nil { - expr.maskedNull = vector.NewConstNull(expr.typ, 1, proc.GetMPool()) + var err error + expr.maskedNull, err = newExpressionConstNull(expr.typ, 1, expr.allocation) + if err != nil { + return nil, err + } } return expr.maskedNull, nil } @@ -411,7 +518,9 @@ func (expr *VarExpressionExecutor) Eval(proc *process.Process, batches []*batch. if val == nil { if expr.null == nil { - expr.null, err = util.GenVectorByVarValue(proc, expr.typ, nil) + expr.null, err = util.GenVectorByVarValueWithAllocation( + proc, expr.typ, nil, expr.allocation, + ) } if err == nil { expr.null.SetIsBin(isBin) @@ -420,7 +529,9 @@ func (expr *VarExpressionExecutor) Eval(proc *process.Process, batches []*batch. } if expr.vec == nil { - expr.vec, err = util.GenVectorByVarValue(proc, expr.typ, val) + expr.vec, err = util.GenVectorByVarValueWithAllocation( + proc, expr.typ, val, expr.allocation, + ) } else { switch v := val.(type) { case []byte: @@ -474,7 +585,8 @@ func (expr *VarExpressionExecutor) IsColumnExpr() bool { } type ListExpressionExecutor struct { - mp *mpool.MPool + mp *mpool.MPool + allocation *vector.AllocationAccountSelection typ types.Type resultVector *vector.Vector @@ -484,11 +596,17 @@ type ListExpressionExecutor struct { func (expr *ListExpressionExecutor) Eval(proc *process.Process, batches []*batch.Batch, selectList []bool) (*vector.Vector, error) { if expr.resultVector == nil { - expr.resultVector = vector.NewOffHeapVecWithType(expr.typ) + var err error + expr.resultVector, err = newExpressionOffHeapVector(expr.typ, expr.allocation) + if err != nil { + return nil, err + } } else { expr.resultVector.CleanOnlyData() } - expr.resultVector.PreExtend(len(expr.parameterExecutor), proc.Mp()) + if err := expr.resultVector.PreExtend(len(expr.parameterExecutor), proc.Mp()); err != nil { + return nil, err + } for i := range expr.parameterExecutor { vec, err := expr.parameterExecutor[i].Eval(proc, batches, selectList) if err != nil { @@ -517,7 +635,9 @@ func (expr *ListExpressionExecutor) Free() { return } for _, e := range expr.parameterExecutor { - e.Free() + if e != nil { + e.Free() + } } if expr.resultVector != nil { expr.resultVector.Free(expr.mp) @@ -531,12 +651,26 @@ func (expr *ListExpressionExecutor) IsColumnExpr() bool { } func (expr *ListExpressionExecutor) Init(proc *process.Process, typ types.Type, parameterNum int) { + if err := expr.init(proc, typ, parameterNum, nil); err != nil { + panic(err) + } +} + +func (expr *ListExpressionExecutor) init( + proc *process.Process, + typ types.Type, + parameterNum int, + selection *vector.AllocationAccountSelection, +) error { m := proc.Mp() expr.typ = typ expr.mp = m + expr.allocation = selection expr.parameterExecutor = make([]ExpressionExecutor, parameterNum) - expr.resultVector = vector.NewOffHeapVecWithType(typ) + var err error + expr.resultVector, err = newExpressionOffHeapVector(typ, selection) + return err } func (expr *ListExpressionExecutor) SetParameter(index int, executor ExpressionExecutor) { @@ -553,14 +687,24 @@ func (expr *FunctionExpressionExecutor) Init( proc *process.Process, parameterNum int, retType types.Type) (err error) { + return expr.init(proc, parameterNum, retType, nil) +} + +func (expr *FunctionExpressionExecutor) init( + proc *process.Process, + parameterNum int, + retType types.Type, + selection *vector.AllocationAccountSelection, +) (err error) { m := proc.Mp() expr.m = m + expr.allocation = selection expr.resultType = retType expr.parameterResults = make([]*vector.Vector, parameterNum) expr.parameterExecutor = make([]ExpressionExecutor, parameterNum) - expr.resultVector = vector.NewFunctionResultWrapper(retType, m) + expr.resultVector, err = vector.NewFunctionResultWrapperWithAllocation(retType, m, selection) return err } @@ -624,14 +768,17 @@ func (expr *FunctionExpressionExecutor) EvalIff(proc *process.Process, batches [ return err } } else { - expr.parameterResults[1] = expr.iffNullResult(0, rowCount) + expr.parameterResults[1], err = expr.iffNullResult(0, rowCount) + if err != nil { + return err + } } if hasSelectedRows(falseBranch) { expr.parameterResults[2], err = expr.parameterExecutor[2].Eval(proc, batches, falseBranch) return err } - expr.parameterResults[2] = expr.iffNullResult(1, rowCount) - return nil + expr.parameterResults[2], err = expr.iffNullResult(1, rowCount) + return err } func hasSelectedRows(selectList []bool) bool { @@ -643,19 +790,24 @@ func hasSelectedRows(selectList []bool) bool { return false } -func (expr *FunctionExpressionExecutor) iffNullResult(index, length int) *vector.Vector { +func (expr *FunctionExpressionExecutor) iffNullResult(index, length int) (*vector.Vector, error) { typ := expr.resultType result := expr.iffNullResults[index] if result == nil || *result.GetType() != typ { if result != nil { result.Free(expr.m) + expr.iffNullResults[index] = nil + } + var err error + result, err = newExpressionConstNull(typ, length, expr.allocation) + if err != nil { + return nil, err } - result = vector.NewConstNull(typ, length, expr.m) expr.iffNullResults[index] = result } else { result.SetLength(length) } - return result + return result, nil } func (expr *FunctionExpressionExecutor) EvalCase(proc *process.Process, batches []*batch.Batch, selectList []bool) (err error) { @@ -786,7 +938,13 @@ func (expr *FunctionExpressionExecutor) evalSelectedRows( if rowAligned && !parameter.IsConst() { selected := expr.selectedParameterVectors[i] if selected == nil { - selected = vector.NewOffHeapVecWithType(*parameter.GetType()) + var err error + selected, err = newExpressionOffHeapVector( + *parameter.GetType(), expr.allocation, + ) + if err != nil { + return nil, err + } expr.selectedParameterVectors[i] = selected } else { selected.Reset(*parameter.GetType()) @@ -806,7 +964,13 @@ func (expr *FunctionExpressionExecutor) evalSelectedRows( return nil, err } if expr.selectedResult == nil { - expr.selectedResult = vector.NewFunctionResultWrapper(expr.resultType, expr.m) + var err error + expr.selectedResult, err = vector.NewFunctionResultWrapperWithAllocation( + expr.resultType, expr.m, expr.allocation, + ) + if err != nil { + return nil, err + } } expr.resetResultType(expr.selectedResult) if err := expr.selectedResult.PreExtendAndReset(selectedCount); err != nil { @@ -826,7 +990,13 @@ func (expr *FunctionExpressionExecutor) evalSelectedRows( result.SetIsBin(runtimeIsBin) result.ResetWithSameType() if expr.selectedNullResult == nil { - expr.selectedNullResult = vector.NewConstNull(runtimeType, 1, expr.m) + var err error + expr.selectedNullResult, err = newExpressionConstNull( + runtimeType, 1, expr.allocation, + ) + if err != nil { + return nil, err + } } else { expr.selectedNullResult.SetType(runtimeType) expr.selectedNullResult.SetLength(1) @@ -1008,19 +1178,27 @@ func (expr *ColumnExpressionExecutor) Eval(_ *process.Process, batches []*batch. vec := batches[relIndex].Vecs[expr.colIndex] if vec.IsConstNull() { - vec = expr.getConstNullVec(expr.typ, vec.Length()) + var err error + vec, err = expr.getConstNullVec(expr.typ, vec.Length()) + if err != nil { + return nil, err + } } return vec, nil } -func (expr *ColumnExpressionExecutor) getConstNullVec(typ types.Type, length int) *vector.Vector { +func (expr *ColumnExpressionExecutor) getConstNullVec(typ types.Type, length int) (*vector.Vector, error) { if expr.nullVecCache != nil { expr.nullVecCache.SetType(typ) expr.nullVecCache.SetLength(length) } else { - expr.nullVecCache = vector.NewConstNull(typ, length, expr.mp) + var err error + expr.nullVecCache, err = newExpressionConstNull(typ, length, expr.allocation) + if err != nil { + return nil, err + } } - return expr.nullVecCache + return expr.nullVecCache, nil } func (expr *ColumnExpressionExecutor) EvalWithoutResultReusing(proc *process.Process, batches []*batch.Batch, _ []bool) (*vector.Vector, error) { @@ -1077,109 +1255,116 @@ func (expr *FixedVectorExpressionExecutor) IsColumnExpr() bool { return false } -func generateConstExpressionExecutor(proc *process.Process, typ types.Type, con *plan.Literal) (vec *vector.Vector, err error) { +func generateConstExpressionExecutor( + proc *process.Process, + typ types.Type, + con *plan.Literal, + selection *vector.AllocationAccountSelection, +) (vec *vector.Vector, err error) { if con.GetIsnull() { - vec = vector.NewConstNull(typ, 1, proc.Mp()) + vec, err = newExpressionConstNull(typ, 1, selection) } else { switch val := con.GetValue().(type) { case *plan.Literal_Bval: - vec, err = vector.NewConstFixed(constBType, val.Bval, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constBType, val.Bval, 1, proc.Mp(), selection) case *plan.Literal_I8Val: - vec, err = vector.NewConstFixed(constI8Type, int8(val.I8Val), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constI8Type, int8(val.I8Val), 1, proc.Mp(), selection) case *plan.Literal_I16Val: - vec, err = vector.NewConstFixed(constI16Type, int16(val.I16Val), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constI16Type, int16(val.I16Val), 1, proc.Mp(), selection) case *plan.Literal_I32Val: - vec, err = vector.NewConstFixed(constI32Type, val.I32Val, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constI32Type, val.I32Val, 1, proc.Mp(), selection) case *plan.Literal_I64Val: - vec, err = vector.NewConstFixed(constI64Type, val.I64Val, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constI64Type, val.I64Val, 1, proc.Mp(), selection) case *plan.Literal_U8Val: - vec, err = vector.NewConstFixed(constU8Type, uint8(val.U8Val), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constU8Type, uint8(val.U8Val), 1, proc.Mp(), selection) case *plan.Literal_U16Val: - vec, err = vector.NewConstFixed(constU16Type, uint16(val.U16Val), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constU16Type, uint16(val.U16Val), 1, proc.Mp(), selection) case *plan.Literal_U32Val: - vec, err = vector.NewConstFixed(constU32Type, val.U32Val, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constU32Type, val.U32Val, 1, proc.Mp(), selection) case *plan.Literal_U64Val: if typ.Oid == types.T_bit { - vec, err = vector.NewConstFixed(typ, val.U64Val, 1, proc.Mp()) + vec, err = newExpressionConstFixed(typ, val.U64Val, 1, proc.Mp(), selection) } else { - vec, err = vector.NewConstFixed(constU64Type, val.U64Val, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constU64Type, val.U64Val, 1, proc.Mp(), selection) } case *plan.Literal_Fval: - vec, err = vector.NewConstFixed(constFType, val.Fval, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constFType, val.Fval, 1, proc.Mp(), selection) case *plan.Literal_Dval: - vec, err = vector.NewConstFixed(constDType, val.Dval, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constDType, val.Dval, 1, proc.Mp(), selection) case *plan.Literal_Dateval: - vec, err = vector.NewConstFixed(constDateType, types.Date(val.Dateval), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constDateType, types.Date(val.Dateval), 1, proc.Mp(), selection) case *plan.Literal_Timeval: - vec, err = vector.NewConstFixed(typ, types.Time(val.Timeval), 1, proc.Mp()) + vec, err = newExpressionConstFixed(typ, types.Time(val.Timeval), 1, proc.Mp(), selection) case *plan.Literal_Datetimeval: - vec, err = vector.NewConstFixed(typ, types.Datetime(val.Datetimeval), 1, proc.Mp()) + vec, err = newExpressionConstFixed(typ, types.Datetime(val.Datetimeval), 1, proc.Mp(), selection) case *plan.Literal_Decimal64Val: cd64 := val.Decimal64Val d64 := types.Decimal64(cd64.A) - vec, err = vector.NewConstFixed(typ, d64, 1, proc.Mp()) + vec, err = newExpressionConstFixed(typ, d64, 1, proc.Mp(), selection) case *plan.Literal_Decimal128Val: cd128 := val.Decimal128Val d128 := types.Decimal128{B0_63: uint64(cd128.A), B64_127: uint64(cd128.B)} - vec, err = vector.NewConstFixed(typ, d128, 1, proc.Mp()) + vec, err = newExpressionConstFixed(typ, d128, 1, proc.Mp(), selection) case *plan.Literal_Timestampval: scale := typ.Scale if scale < 0 || scale > 6 { return nil, moerr.NewErrTooBigPrecision(proc.Ctx, int64(scale), "TIMESTAMP", 6) } - vec, err = vector.NewConstFixed(constTimestampTypes[scale], types.Timestamp(val.Timestampval), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constTimestampTypes[scale], types.Timestamp(val.Timestampval), 1, proc.Mp(), selection) case *plan.Literal_Sval: sval := val.Sval // Distinguish binary with non-binary string. if typ.Oid == types.T_binary || typ.Oid == types.T_varbinary || typ.Oid == types.T_blob { - vec, err = vector.NewConstBytes(constBinType, []byte(sval), 1, proc.Mp()) + vec, err = newExpressionConstBytes(constBinType, []byte(sval), 1, proc.Mp(), selection) } else if typ.Oid == types.T_geometry { - vec, err = vector.NewConstBytes(typ, []byte(sval), 1, proc.Mp()) + vec, err = newExpressionConstBytes(typ, []byte(sval), 1, proc.Mp(), selection) } else if typ.Oid == types.T_array_float32 { array, err1 := types.StringToArray[float32](sval) if err1 != nil { return nil, err1 } - vec, err = vector.NewConstArray(typ, array, 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, array, 1, proc.Mp(), selection) } else if typ.Oid == types.T_array_float64 { array, err1 := types.StringToArray[float64](sval) if err1 != nil { return nil, err1 } - vec, err = vector.NewConstArray(typ, array, 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, array, 1, proc.Mp(), selection) } else if typ.Oid == types.T_datalink { _, _, err1 := datalink.ParseDatalink(sval, proc) if err1 != nil { return nil, err1 } - vec, err = vector.NewConstBytes(constBinType, []byte(sval), 1, proc.Mp()) + vec, err = newExpressionConstBytes(constBinType, []byte(sval), 1, proc.Mp(), selection) } else { - vec, err = vector.NewConstBytes(constSType, []byte(sval), 1, proc.Mp()) + vec, err = newExpressionConstBytes(constSType, []byte(sval), 1, proc.Mp(), selection) } case *plan.Literal_Defaultval: defaultVal := val.Defaultval - vec, err = vector.NewConstFixed(constBType, defaultVal, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constBType, defaultVal, 1, proc.Mp(), selection) case *plan.Literal_EnumVal: - vec, err = vector.NewConstFixed(constEnumType, types.Enum(val.EnumVal), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constEnumType, types.Enum(val.EnumVal), 1, proc.Mp(), selection) case *plan.Literal_VecVal: switch typ.Oid { case types.T_array_float32: - vec, err = vector.NewConstArray(typ, types.BytesToArray[float32]([]byte(val.VecVal)), 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, types.BytesToArray[float32]([]byte(val.VecVal)), 1, proc.Mp(), selection) case types.T_array_float64: - vec, err = vector.NewConstArray(typ, types.BytesToArray[float64]([]byte(val.VecVal)), 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, types.BytesToArray[float64]([]byte(val.VecVal)), 1, proc.Mp(), selection) case types.T_array_bf16: - vec, err = vector.NewConstArray(typ, types.BytesToArray[types.BF16]([]byte(val.VecVal)), 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, types.BytesToArray[types.BF16]([]byte(val.VecVal)), 1, proc.Mp(), selection) case types.T_array_float16: - vec, err = vector.NewConstArray(typ, types.BytesToArray[types.Float16]([]byte(val.VecVal)), 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, types.BytesToArray[types.Float16]([]byte(val.VecVal)), 1, proc.Mp(), selection) case types.T_array_int8: - vec, err = vector.NewConstArray(typ, types.BytesToArray[int8]([]byte(val.VecVal)), 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, types.BytesToArray[int8]([]byte(val.VecVal)), 1, proc.Mp(), selection) case types.T_array_uint8: - vec, err = vector.NewConstArray(typ, types.BytesToArray[uint8]([]byte(val.VecVal)), 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, types.BytesToArray[uint8]([]byte(val.VecVal)), 1, proc.Mp(), selection) } default: return nil, moerr.NewNYI(proc.Ctx, fmt.Sprintf("const expression %v", con.GetValue())) } - vec.SetIsBin(con.IsBin) + if err == nil { + vec.SetIsBin(con.IsBin) + } } return vec, err } diff --git a/pkg/sql/colexec/evalExpressionAllocation_test.go b/pkg/sql/colexec/evalExpressionAllocation_test.go new file mode 100644 index 0000000000000..32dc2e9a299f0 --- /dev/null +++ b/pkg/sql/colexec/evalExpressionAllocation_test.go @@ -0,0 +1,106 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package colexec + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/pb/plan" + "github.com/matrixorigin/matrixone/pkg/sql/plan/function" + "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/stretchr/testify/require" +) + +func TestAccountedExpressionTreeCoversNestedAndSelectedResults(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + typ := types.T_varchar.ToType() + column := &plan.Expr{ + Typ: plan.Type{Id: int32(typ.Oid), Width: typ.Width}, + Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, + } + bindFunction := func(name string, args ...*plan.Expr) *plan.Expr { + argTypes := make([]types.Type, len(args)) + for i := range args { + argTypes[i] = types.New( + types.T(args[i].Typ.Id), args[i].Typ.Width, args[i].Typ.Scale, + ) + } + fn, bindErr := function.GetFunctionByName(proc.Ctx, name, argTypes) + require.NoError(t, bindErr) + retType := fn.GetReturnType() + return &plan.Expr{ + Typ: plan.Type{ + Id: int32(retType.Oid), Width: retType.Width, Scale: retType.Scale, + }, + Expr: &plan.Expr_F{F: &plan.Function{ + Func: &plan.ObjectRef{Obj: fn.GetEncodedOverloadID(), ObjName: name}, + Args: args, + }}, + } + } + literal := &plan.Expr{ + Typ: plan.Type{Id: int32(typ.Oid), Width: typ.Width}, + Expr: &plan.Expr_Lit{Lit: &plan.Literal{ + Value: &plan.Literal_Sval{Sval: "-"}, + }}, + } + expression := bindFunction("concat", bindFunction("lower", column), column, literal) + executor, err := NewExpressionExecutorWithAllocation(proc, expression, selection) + require.NoError(t, err) + root := executor.(*FunctionExpressionExecutor) + nested := root.parameterExecutor[0].(*FunctionExpressionExecutor) + + input := batch.NewWithSize(1) + input.Vecs[0] = testutil.MakeVarcharVector( + []string{"AA", "BB", "CC", "DD"}, nil, proc.Mp(), + ) + input.SetRowCount(4) + defer input.Clean(proc.Mp()) + result, err := executor.Eval( + proc, []*batch.Batch{input}, []bool{true, false, true, false}, + ) + require.NoError(t, err) + require.Equal(t, []string{"aaAA-", "", "ccCC-", ""}, vector.InefficientMustStrCol(result)) + + assertFunctionStorage := func(function *FunctionExpressionExecutor) { + t.Helper() + require.Same(t, selection, function.resultVector.GetResultVector().AllocationAccountSelection()) + require.Same(t, selection, function.selectedResult.GetResultVector().AllocationAccountSelection()) + for _, selected := range function.selectedParameterVectors { + if selected != nil { + require.Same(t, selection, selected.AllocationAccountSelection()) + } + } + } + assertFunctionStorage(root) + assertFunctionStorage(nested) + fixed := root.parameterExecutor[2].(*FixedVectorExpressionExecutor) + require.Same(t, selection, fixed.resultVector.AllocationAccountSelection()) + require.Positive(t, account.Snapshot().Used) + + executor.Free() + require.Zero(t, account.Snapshot().Used) +} diff --git a/pkg/sql/colexec/evalExpressionMemory.go b/pkg/sql/colexec/evalExpressionMemory.go deleted file mode 100644 index b263deb729bf2..0000000000000 --- a/pkg/sql/colexec/evalExpressionMemory.go +++ /dev/null @@ -1,131 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package colexec - -import ( - "math" - - "github.com/matrixorigin/matrixone/pkg/container/vector" -) - -// ExpressionExecutorRetainedBytes returns the mpool-backed vector capacity -// owned by an executor tree. Borrowed evaluation results are deliberately not -// counted: every owned vector is reached through exactly one executor field. -// -// Go-managed executor metadata is not included here. HashBuild's expression -// admission bound retains a separate per-vector allowance for that metadata. -func ExpressionExecutorRetainedBytes(executor ExpressionExecutor) (uint64, bool) { - switch expr := executor.(type) { - case nil: - return 0, true - case *ColumnExpressionExecutor: - return expressionVectorRetainedBytes(expr.nullVecCache) - case *FixedVectorExpressionExecutor: - return fixedExpressionVectorRetainedBytes(expr.resultVector) - case *ParamExpressionExecutor: - return sumExpressionVectorRetainedBytes(expr.null, expr.maskedNull, expr.vec) - case *VarExpressionExecutor: - return sumExpressionVectorRetainedBytes(expr.null, expr.maskedNull, expr.vec) - case *ListExpressionExecutor: - total, ok := expressionVectorRetainedBytes(expr.resultVector) - if !ok { - return 0, false - } - return addExpressionExecutorRetainedBytes(total, expr.parameterExecutor) - case *FunctionExpressionExecutor: - var result *vector.Vector - if expr.resultVector != nil { - result = expr.resultVector.GetResultVector() - } - var selectedResult *vector.Vector - if expr.selectedResult != nil { - selectedResult = expr.selectedResult.GetResultVector() - } - total, ok := sumExpressionVectorRetainedBytes( - result, - selectedResult, - expr.selectedNullResult, - expr.iffNullResults[0], - expr.iffNullResults[1], - ) - if !ok { - return 0, false - } - for _, parameter := range expr.selectedParameterVectors { - bytes, valid := expressionVectorRetainedBytes(parameter) - if !valid || total > math.MaxUint64-bytes { - return 0, false - } - total += bytes - } - return addExpressionExecutorRetainedBytes(total, expr.parameterExecutor) - default: - // External test or extension executors do not expose an ownership graph. - // Treat their retained capacity as unknown instead of silently claiming - // that they own no memory. - return 0, false - } -} - -func ExpressionExecutorsRetainedBytes(executors []ExpressionExecutor) (uint64, bool) { - return addExpressionExecutorRetainedBytes(0, executors) -} - -func addExpressionExecutorRetainedBytes(total uint64, executors []ExpressionExecutor) (uint64, bool) { - for _, executor := range executors { - bytes, ok := ExpressionExecutorRetainedBytes(executor) - if !ok || total > math.MaxUint64-bytes { - return 0, false - } - total += bytes - } - return total, true -} - -func sumExpressionVectorRetainedBytes(vectors ...*vector.Vector) (uint64, bool) { - var total uint64 - for _, vec := range vectors { - bytes, ok := expressionVectorRetainedBytes(vec) - if !ok || total > math.MaxUint64-bytes { - return 0, false - } - total += bytes - } - return total, true -} - -func expressionVectorRetainedBytes(vec *vector.Vector) (uint64, bool) { - if vec == nil { - return 0, true - } - if vec.NeedDup() { - // Mmap/no-copy vectors point at bytes owned by the serialized plan or - // another vector. They are stable fixed-expression inputs, not mpool - // capacity owned by this executor. - return 0, false - } - allocated := vec.Allocated() - if allocated < 0 { - return 0, false - } - return uint64(allocated), true -} - -func fixedExpressionVectorRetainedBytes(vec *vector.Vector) (uint64, bool) { - if vec == nil || vec.NeedDup() { - return 0, true - } - return expressionVectorRetainedBytes(vec) -} diff --git a/pkg/sql/colexec/evalExpressionMemory_test.go b/pkg/sql/colexec/evalExpressionMemory_test.go deleted file mode 100644 index 300fa5439157a..0000000000000 --- a/pkg/sql/colexec/evalExpressionMemory_test.go +++ /dev/null @@ -1,156 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package colexec - -import ( - "testing" - - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/container/batch" - "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/matrixorigin/matrixone/pkg/container/vector" - "github.com/matrixorigin/matrixone/pkg/pb/plan" - "github.com/matrixorigin/matrixone/pkg/testutil" - "github.com/matrixorigin/matrixone/pkg/vm/process" - "github.com/stretchr/testify/require" -) - -func TestExpressionExecutorRetainedBytesExcludesBorrowedPlanVector(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - source := testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()) - data, err := source.MarshalBinary() - require.NoError(t, err) - source.Free(proc.Mp()) - require.Zero(t, proc.Mp().CurrNB()) - - executor, err := NewExpressionExecutor(proc, &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Vec{Vec: &plan.LiteralVec{ - Len: 3, - Data: data, - }}, - }) - require.NoError(t, err) - fixed, ok := executor.(*FixedVectorExpressionExecutor) - require.True(t, ok) - require.True(t, fixed.resultVector.NeedDup()) - require.Positive(t, fixed.resultVector.Allocated()) - - retained, known := ExpressionExecutorRetainedBytes(executor) - require.True(t, known) - require.Zero(t, retained) - executor.Free() - require.Zero(t, proc.Mp().CurrNB()) -} - -func TestExpressionExecutorRetainedBytesCoversOwnedExecutorTree(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - newVec := func(values ...int32) *vector.Vector { - return testutil.MakeInt32Vector(values, nil, proc.Mp()) - } - - fixedVec := newVec(1) - listResult := newVec(1, 2) - paramNull := newVec(1) - paramMaskedNull := newVec(2) - paramVec := newVec(3) - varNull := newVec(4) - varMaskedNull := newVec(5) - varVec := newVec(6) - selectedParameter := newVec(7, 8) - selectedNull := vector.NewConstNull(types.T_int32.ToType(), 2, proc.Mp()) - iffNull0 := vector.NewConstNull(types.T_int32.ToType(), 2, proc.Mp()) - iffNull1 := vector.NewConstNull(types.T_int32.ToType(), 2, proc.Mp()) - ownedVectors := []*vector.Vector{ - fixedVec, - listResult, - paramNull, - paramMaskedNull, - paramVec, - varNull, - varMaskedNull, - varVec, - selectedParameter, - selectedNull, - iffNull0, - iffNull1, - } - - result := vector.NewFunctionResultWrapper(types.T_int32.ToType(), proc.Mp()) - require.NoError(t, result.PreExtendAndReset(4)) - selectedResult := vector.NewFunctionResultWrapper(types.T_int32.ToType(), proc.Mp()) - require.NoError(t, selectedResult.PreExtendAndReset(2)) - - fixed := &FixedVectorExpressionExecutor{resultVector: fixedVec} - list := &ListExpressionExecutor{ - resultVector: listResult, - parameterExecutor: []ExpressionExecutor{fixed}, - } - param := &ParamExpressionExecutor{ - null: paramNull, - maskedNull: paramMaskedNull, - vec: paramVec, - } - variable := &VarExpressionExecutor{ - null: varNull, - maskedNull: varMaskedNull, - vec: varVec, - } - function := &FunctionExpressionExecutor{ - resultVector: result, - selectedResult: selectedResult, - selectedNullResult: selectedNull, - selectedParameterVectors: []*vector.Vector{selectedParameter}, - parameterExecutor: []ExpressionExecutor{param}, - iffNullResults: [2]*vector.Vector{iffNull0, iffNull1}, - } - column := &ColumnExpressionExecutor{ - nullVecCache: vector.NewConstNull(types.T_int32.ToType(), 2, proc.Mp()), - } - - executors := []ExpressionExecutor{list, variable, function, column, nil} - retained, known := ExpressionExecutorsRetainedBytes(executors) - require.True(t, known) - require.Positive(t, retained) - for _, executor := range executors { - _, known = ExpressionExecutorRetainedBytes(executor) - require.True(t, known) - } - _, known = ExpressionExecutorRetainedBytes(retainedBytesUnknownExecutor{}) - require.False(t, known) - - result.Free() - selectedResult.Free() - column.nullVecCache.Free(proc.Mp()) - for _, vec := range ownedVectors { - vec.Free(proc.Mp()) - } - require.Zero(t, proc.Mp().CurrNB()) -} - -type retainedBytesUnknownExecutor struct{} - -func (retainedBytesUnknownExecutor) Eval(*process.Process, []*batch.Batch, []bool) (*vector.Vector, error) { - return nil, nil -} -func (retainedBytesUnknownExecutor) EvalWithoutResultReusing(*process.Process, []*batch.Batch, []bool) (*vector.Vector, error) { - return nil, nil -} -func (retainedBytesUnknownExecutor) ResetForNextQuery() {} -func (retainedBytesUnknownExecutor) Free() {} -func (retainedBytesUnknownExecutor) IsColumnExpr() bool { return false } -func (retainedBytesUnknownExecutor) TypeName() string { return "unknown" } diff --git a/pkg/sql/colexec/hashbuild/expression_memory.go b/pkg/sql/colexec/hashbuild/expression_memory.go index 6e82e6e0b9e68..4b94912414c25 100644 --- a/pkg/sql/colexec/hashbuild/expression_memory.go +++ b/pkg/sql/colexec/hashbuild/expression_memory.go @@ -15,17 +15,27 @@ package hashbuild import ( + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" "github.com/matrixorigin/matrixone/pkg/vm/process" ) +const ( + hashBuildAllocationSiteExpressionData mpool.AllocationSite = iota + 98 + hashBuildAllocationSiteExpressionArea + hashBuildAllocationSiteExpressionNulls + hashBuildAllocationSiteExpressionGrouping +) + // NewExpressionExecutors constructs expression trees used by HashBuild and -// join operators. Expression temporaries are not retained HashBuild storage; -// only their explicit copies into retained destinations enter the account. +// join operators. Every MPool vector owned by the tree, including nested +// function results and reusable selection buffers, shares the query account. func NewExpressionExecutors( proc *process.Process, exprs []*plan.Expr, + account *mpool.AllocationAccount, ) ([]colexec.ExpressionExecutor, error) { if len(exprs) == 0 { return nil, process.ErrHashBuildBudgetInvalid @@ -35,5 +45,18 @@ func NewExpressionExecutors( return nil, process.ErrHashBuildBudgetInvalid } } - return colexec.NewExpressionExecutorsFromPlanExpressions(proc, exprs) + selection, err := vector.NewAllocationAccountSelection( + account, + HashBuildAllocationOwner, + hashBuildAllocationSiteExpressionData, + hashBuildAllocationSiteExpressionArea, + hashBuildAllocationSiteExpressionNulls, + hashBuildAllocationSiteExpressionGrouping, + ) + if err != nil { + return nil, err + } + return colexec.NewExpressionExecutorsFromPlanExpressionsWithAllocation( + proc, exprs, selection, + ) } diff --git a/pkg/sql/colexec/hashbuild/hashmap.go b/pkg/sql/colexec/hashbuild/hashmap.go index 8d8bc7afffb75..ea941bb798612 100644 --- a/pkg/sql/colexec/hashbuild/hashmap.go +++ b/pkg/sql/colexec/hashbuild/hashmap.go @@ -196,6 +196,7 @@ func (hb *HashmapBuilder) Prepare( executors, err := NewExpressionExecutors( proc, keyCols, + hb.mapAllocationAccount, ) if err != nil { return err diff --git a/pkg/sql/colexec/hashbuild/hashmap_test.go b/pkg/sql/colexec/hashbuild/hashmap_test.go index 82952e6eb8f3f..900123d85d10a 100644 --- a/pkg/sql/colexec/hashbuild/hashmap_test.go +++ b/pkg/sql/colexec/hashbuild/hashmap_test.go @@ -532,7 +532,7 @@ func TestAccountedRuntimeFilterUniqueKeysDegradeWithoutFailingHashBuild(t *testi require.LessOrEqual(t, constrained.Peak, baseline.Peak) } -func TestSpillExpressionTemporaryIsOutsideRetainedAccount(t *testing.T) { +func TestSpillExpressionStorageUsesRetainedAccount(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() budget := process.MustNewHashBuildBudget(16<<20, 16<<20) @@ -550,6 +550,7 @@ func TestSpillExpressionTemporaryIsOutsideRetainedAccount(t *testing.T) { expr := makeIssue26454ConcatKey(t, proc) executors, err := ctr.initSpillExprExecs(proc, []*plan.Expr{expr}) require.NoError(t, err) + constructorUsed := account.Snapshot().Used input := batch.NewWithSize(2) input.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp()) input.Vecs[1] = testutil.MakeInt32Vector([]int32{3, 4}, nil, proc.Mp()) @@ -558,8 +559,8 @@ func TestSpillExpressionTemporaryIsOutsideRetainedAccount(t *testing.T) { result, err := executors[0].Eval(proc, []*batch.Batch{input}, nil) require.NoError(t, err) require.Equal(t, []string{"1-3", "2-4"}, vector.InefficientMustStrCol(result)) - require.Zero(t, account.Snapshot().Used) - require.Zero(t, generation.Used()) + require.Greater(t, account.Snapshot().Used, constructorUsed) + require.Equal(t, account.Snapshot().Used, generation.Used()) ctr.freeSpillExprExecs() require.Zero(t, account.Snapshot().Used) @@ -569,6 +570,52 @@ func TestSpillExpressionTemporaryIsOutsideRetainedAccount(t *testing.T) { require.NoError(t, err) } +func TestSpillExpressionStorageHonorsAccountCapacity(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + run := func(limit uint64) (uint64, error) { + budget := process.MustNewHashBuildBudget(16<<20, 16<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + var op HashBuild + op.NeedHashMap = true + require.NoError(t, op.SetAllocationAccount(account)) + op.ctr.hashmapBuilder.setBudget(generation) + executors, evalErr := op.ctr.initSpillExprExecs( + proc, + []*plan.Expr{makeIssue26454ConcatKey(t, proc)}, + ) + if evalErr == nil { + input := testutil.NewBatch( + []types.Type{types.T_int32.ToType(), types.T_int32.ToType()}, + true, + 10_000, + proc.Mp(), + ) + _, evalErr = executors[0].Eval(proc, []*batch.Batch{input}, nil) + input.Clean(proc.Mp()) + } + peak := account.Snapshot().Peak + op.ctr.freeSpillExprExecs() + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.NoError(t, op.ClearAllocationAccount(account)) + _, _, terminalErr := registry.CompleteTerminal(account) + require.NoError(t, terminalErr) + return peak, evalErr + } + + peak, err := run(16 << 20) + require.NoError(t, err) + require.Greater(t, peak, uint64(1)) + _, err = run(peak - 1) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) +} + func TestIssue26454ExpressionKeyBuildUsesActualCapacity(t *testing.T) { const capBytes = uint64(16 << 20) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) diff --git a/pkg/sql/colexec/hashbuild/spill.go b/pkg/sql/colexec/hashbuild/spill.go index 550fa1abc6a39..ef13add04a2b5 100644 --- a/pkg/sql/colexec/hashbuild/spill.go +++ b/pkg/sql/colexec/hashbuild/spill.go @@ -855,6 +855,7 @@ func (ctr *container) initSpillExprExecs(proc *process.Process, conditions []*pl execs, err := NewExpressionExecutors( proc, conditions, + ctr.hashmapBuilder.mapAllocationAccount, ) if err != nil { return nil, err diff --git a/pkg/sql/colexec/hashjoin/allocation_test_helpers_test.go b/pkg/sql/colexec/hashjoin/allocation_test_helpers_test.go index 82d16b6684660..5c00c5dca350d 100644 --- a/pkg/sql/colexec/hashjoin/allocation_test_helpers_test.go +++ b/pkg/sql/colexec/hashjoin/allocation_test_helpers_test.go @@ -18,6 +18,10 @@ import ( "testing" "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/testutil" "github.com/stretchr/testify/require" ) @@ -36,3 +40,45 @@ func installTestAllocation(t testing.TB, owners ...testAllocationOwner) *mpool.A } return account } + +func TestHashJoinResultBatchUsesAllocationAccount(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + arg := &HashJoin{ + ResultCols: []colexec.ResultPos{{Rel: 0, Pos: 0}}, + LeftTypes: []types.Type{types.T_int64.ToType()}, + } + account := installTestAllocation(t, arg) + require.NoError(t, arg.resetResultBat()) + require.Same(t, arg.resultAllocation, arg.ctr.resBat.Vecs[0].AllocationAccountSelection()) + require.NoError(t, vector.AppendFixed(arg.ctr.resBat.Vecs[0], int64(1), false, proc.Mp())) + used := account.Snapshot().Used + require.Positive(t, used) + require.NoError(t, arg.resetResultBat()) + require.Equal(t, used, account.Snapshot().Used) + + arg.Reset(proc, false, nil) + require.Nil(t, arg.ctr.resBat) + require.Zero(t, account.Snapshot().Used) + require.NoError(t, arg.ClearAllocationAccount(account)) +} + +func TestHashJoinResultBatchHonorsAllocationCapacity(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1) + require.NoError(t, err) + arg := &HashJoin{ + ResultCols: []colexec.ResultPos{{Rel: 0, Pos: 0}}, + LeftTypes: []types.Type{types.T_int64.ToType()}, + } + require.NoError(t, arg.SetAllocationAccount(account)) + require.NoError(t, arg.resetResultBat()) + err = vector.AppendFixed(arg.ctr.resBat.Vecs[0], int64(1), false, proc.Mp()) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Zero(t, account.Snapshot().Used) + arg.Reset(proc, false, nil) + require.NoError(t, arg.ClearAllocationAccount(account)) +} diff --git a/pkg/sql/colexec/hashjoin/join.go b/pkg/sql/colexec/hashjoin/join.go index b1f617b1175b8..84d6bba355821 100644 --- a/pkg/sql/colexec/hashjoin/join.go +++ b/pkg/sql/colexec/hashjoin/join.go @@ -97,6 +97,7 @@ func (hashJoin *HashJoin) Prepare(proc *process.Process) (err error) { eqCondExecs, err := hashbuild.NewExpressionExecutors( proc, hashJoin.EqConds[0], + hashJoin.allocationAccount, ) if err != nil { return err @@ -108,6 +109,7 @@ func (hashJoin *HashJoin) Prepare(proc *process.Process) (err error) { nonEqExecs, err = hashbuild.NewExpressionExecutors( proc, []*plan.Expr{hashJoin.NonEqCond}, + hashJoin.allocationAccount, ) if err != nil { for _, exec := range eqCondExecs { @@ -213,7 +215,9 @@ func (hashJoin *HashJoin) Call(proc *process.Process) (vm.CallResult, error) { ctr.lastIdx = 0 } - hashJoin.resetResultBat() + if err = hashJoin.resetResultBat(); err != nil { + return result, err + } for i, rp := range hashJoin.ResultCols { if rp.Rel == 0 { ctr.resBat.Vecs[i].SetSorted(ctr.leftBat.Vecs[rp.Pos].GetSorted()) @@ -831,12 +835,18 @@ func (ctr *container) appendMarkForEmptyBuildBucket(marker *vector.Vector, proc } for _, vec := range ctr.eqCondVecs { if vec.IsConstNull() { + if err := marker.PreExtendNulls(rowCnt, proc.Mp()); err != nil { + return err + } marker.GetNulls().AddRange(0, uint64(rowCnt)) return nil } if !vec.GetNulls().Any() { continue } + if err := marker.PreExtendNulls(rowCnt, proc.Mp()); err != nil { + return err + } nulls.Or(marker.GetNulls(), vec.GetNulls(), marker.GetNulls()) } return nil @@ -892,7 +902,9 @@ func (ctr *container) syncBitmap(hashJoin *HashJoin, proc *process.Process) erro } func (ctr *container) finalize(hashJoin *HashJoin, proc *process.Process, result *vm.CallResult) error { - hashJoin.resetResultBat() + if err := hashJoin.resetResultBat(); err != nil { + return err + } rowCnt := 0 for ; rowCnt < colexec.DefaultBatchSize && ctr.rightMatchedIter.HasNext(); rowCnt++ { @@ -1020,7 +1032,7 @@ func (ctr *container) evalJoinCondition(bat *batch.Batch, proc *process.Process) return nil } -func (hashJoin *HashJoin) resetResultBat() { +func (hashJoin *HashJoin) resetResultBat() error { ctr := &hashJoin.ctr if ctr.resBat != nil { ctr.resBat.CleanOnlyData() @@ -1041,5 +1053,11 @@ func (hashJoin *HashJoin) resetResultBat() { ctr.resBat.Vecs[i] = vector.NewOffHeapVecWithType(types.T_bool.ToType()) } } + if err := ctr.resBat.SetAllocationAccount(hashJoin.resultAllocation); err != nil { + ctr.resBat.Clean(nil) + ctr.resBat = nil + return err + } } + return nil } diff --git a/pkg/sql/colexec/hashjoin/mark_spill_test.go b/pkg/sql/colexec/hashjoin/mark_spill_test.go index e91fcd2e39fac..72927a3642172 100644 --- a/pkg/sql/colexec/hashjoin/mark_spill_test.go +++ b/pkg/sql/colexec/hashjoin/mark_spill_test.go @@ -190,7 +190,7 @@ func TestHashMarkJoinEmptySpillBucketTruthTable(t *testing.T) { tc.arg.ctr.leftBat = probe tc.arg.ctr.globalBuildRowCnt = tt.globalBuildRowCnt tc.arg.ctr.buildHasNullKey = tt.buildHasNullKey - tc.arg.resetResultBat() + require.NoError(t, tc.arg.resetResultBat()) var result vm.CallResult require.NoError(t, tc.arg.ctr.emptyProbe(tc.arg, tc.proc, &result)) diff --git a/pkg/sql/colexec/hashjoin/types.go b/pkg/sql/colexec/hashjoin/types.go index 70098319ec525..77adc90291530 100644 --- a/pkg/sql/colexec/hashjoin/types.go +++ b/pkg/sql/colexec/hashjoin/types.go @@ -24,6 +24,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" "github.com/matrixorigin/matrixone/pkg/sql/colexec/spillutil" "github.com/matrixorigin/matrixone/pkg/vm" "github.com/matrixorigin/matrixone/pkg/vm/message" @@ -50,6 +51,13 @@ const ( const hashJoinAllocationSiteMatchedRows mpool.AllocationSite = 80 +const ( + hashJoinAllocationSiteResultData mpool.AllocationSite = iota + 102 + hashJoinAllocationSiteResultArea + hashJoinAllocationSiteResultNulls + hashJoinAllocationSiteResultGrouping +) + type container struct { state int itr hashmap.Iterator @@ -137,6 +145,7 @@ type HashJoin struct { JoinMapTag int32 SpillThreshold int64 allocationAccount *mpool.AllocationAccount + resultAllocation *vector.AllocationAccountSelection vm.OperatorBase } @@ -154,7 +163,19 @@ func (hashJoin *HashJoin) SetAllocationAccount( if hashJoin.allocationAccount == account { return nil } + selection, err := vector.NewAllocationAccountSelection( + account, + hashbuild.HashBuildAllocationOwner, + hashJoinAllocationSiteResultData, + hashJoinAllocationSiteResultArea, + hashJoinAllocationSiteResultNulls, + hashJoinAllocationSiteResultGrouping, + ) + if err != nil { + return err + } hashJoin.allocationAccount = account + hashJoin.resultAllocation = selection return nil } @@ -170,13 +191,15 @@ func (hashJoin *HashJoin) ClearAllocationAccount( if hashJoin.ctr.mp != nil || hashJoin.ctr.spillEngine != nil || len(hashJoin.ctr.eqCondExecs) != 0 || hashJoin.ctr.nonEqCondExec != nil || - hashJoin.ctr.rightRowsMatched != nil { + hashJoin.ctr.rightRowsMatched != nil || + hashJoin.ctr.resBat != nil { return mpool.ErrAllocationAccountInvariant } if hashJoin.NumCPU > 1 && !hashJoin.Mailbox.Terminal() { return mpool.ErrAllocationAccountInvariant } hashJoin.allocationAccount = nil + hashJoin.resultAllocation = nil return nil } @@ -241,6 +264,10 @@ func (hashJoin *HashJoin) Reset(proc *process.Process, pipelineFailed bool, err ctr.cleanEqCondExecutors() ctr.cleanHashMap() ctr.cleanNonEqCondExecutor() + if ctr.resBat != nil { + ctr.resBat.Clean(proc.GetMPool()) + ctr.resBat = nil + } ctr.freeRightRowsMatched(proc) ctr.rightMatchedIter = nil ctr.skipProbe = false diff --git a/pkg/sql/colexec/loopjoin/join.go b/pkg/sql/colexec/loopjoin/join.go index 499e1f4418727..b1e7a0fece2b7 100644 --- a/pkg/sql/colexec/loopjoin/join.go +++ b/pkg/sql/colexec/loopjoin/join.go @@ -73,6 +73,7 @@ func (loopJoin *LoopJoin) Prepare(proc *process.Process) error { execs, err = hashbuild.NewExpressionExecutors( proc, []*plan.Expr{loopJoin.NonEqCond}, + loopJoin.allocationAccount, ) if err != nil { return err @@ -130,7 +131,9 @@ func (loopJoin *LoopJoin) Call(proc *process.Process) (vm.CallResult, error) { ctr.batIdx = 0 } - loopJoin.resetResultBat() + if err = loopJoin.resetResultBat(); err != nil { + return result, err + } for i, rp := range loopJoin.ResultCols { if rp.Rel == 0 { ctr.resBat.Vecs[i].SetSorted(ctr.inBat.Vecs[rp.Pos].GetSorted()) @@ -435,7 +438,7 @@ func (ctr *container) probe(ap *LoopJoin, proc *process.Process, result *vm.Call return nil } -func (loopJoin *LoopJoin) resetResultBat() { +func (loopJoin *LoopJoin) resetResultBat() error { ctr := &loopJoin.ctr if ctr.resBat != nil { ctr.resBat.CleanOnlyData() @@ -444,21 +447,37 @@ func (loopJoin *LoopJoin) resetResultBat() { ctr.resBat.Vecs[i].SetLength(0) } } else { - ctr.resBat = batch.NewWithSize(len(loopJoin.ResultCols)) + ctr.resBat = batch.NewOffHeapWithSize(len(loopJoin.ResultCols)) for i, rp := range loopJoin.ResultCols { switch rp.Rel { case 0: - ctr.resBat.Vecs[i] = vector.NewVec(*ctr.inBat.Vecs[rp.Pos].GetType()) + var leftType types.Type + if ctr.inBat != nil && int(rp.Pos) < len(ctr.inBat.Vecs) { + leftType = *ctr.inBat.Vecs[rp.Pos].GetType() + } else if int(rp.Pos) < len(loopJoin.LeftTypes) { + leftType = loopJoin.LeftTypes[rp.Pos] + } else { + ctr.resBat.Clean(nil) + ctr.resBat = nil + return process.ErrHashBuildBudgetInvalid + } + ctr.resBat.Vecs[i] = vector.NewOffHeapVecWithType(leftType) case 1: - ctr.resBat.Vecs[i] = vector.NewVec(loopJoin.RightTypes[rp.Pos]) + ctr.resBat.Vecs[i] = vector.NewOffHeapVecWithType(loopJoin.RightTypes[rp.Pos]) case -1: - ctr.resBat.Vecs[i] = vector.NewVec(types.T_bool.ToType()) + ctr.resBat.Vecs[i] = vector.NewOffHeapVecWithType(types.T_bool.ToType()) } } + if err := ctr.resBat.SetAllocationAccount(loopJoin.resultAllocation); err != nil { + ctr.resBat.Clean(nil) + ctr.resBat = nil + return err + } } + return nil } // initRightMatchedBitmap allocates the per-build-row matched bitmap. @@ -505,24 +524,8 @@ func (ctr *container) initRightMatchedBitmap( // columns. Iterator is monotonic, so rightMatchedBat only advances. func (ctr *container) finalize(ap *LoopJoin, proc *process.Process, result *vm.CallResult) error { bats := ctr.mp.GetBatches() - if ctr.resBat == nil { - ctr.resBat = batch.NewWithSize(len(ap.ResultCols)) - for i, rp := range ap.ResultCols { - switch rp.Rel { - case 0: - ctr.resBat.Vecs[i] = vector.NewVec(ap.LeftTypes[rp.Pos]) - case 1: - ctr.resBat.Vecs[i] = vector.NewVec(ap.RightTypes[rp.Pos]) - default: - ctr.resBat.Vecs[i] = vector.NewVec(types.T_bool.ToType()) - } - } - } else { - ctr.resBat.CleanOnlyData() - for i := range ctr.resBat.Vecs { - ctr.resBat.Vecs[i].SetClass(vector.FLAT) - ctr.resBat.Vecs[i].SetLength(0) - } + if err := ap.resetResultBat(); err != nil { + return err } rowCnt := 0 diff --git a/pkg/sql/colexec/loopjoin/join_test.go b/pkg/sql/colexec/loopjoin/join_test.go index 93664e6b10cf3..fece6216d6c01 100644 --- a/pkg/sql/colexec/loopjoin/join_test.go +++ b/pkg/sql/colexec/loopjoin/join_test.go @@ -71,6 +71,48 @@ func installLoopJoinTestAllocation( return account } +func TestLoopJoinResultBatchUsesAllocationAccount(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + arg := &LoopJoin{ + ResultCols: []colexec.ResultPos{{Rel: 0, Pos: 0}}, + LeftTypes: []types.Type{types.T_int64.ToType()}, + } + account := installLoopJoinTestAllocation(t, arg) + require.NoError(t, arg.resetResultBat()) + require.Same(t, arg.resultAllocation, arg.ctr.resBat.Vecs[0].AllocationAccountSelection()) + require.NoError(t, vector.AppendFixed(arg.ctr.resBat.Vecs[0], int64(1), false, proc.Mp())) + used := account.Snapshot().Used + require.Positive(t, used) + require.NoError(t, arg.resetResultBat()) + require.Equal(t, used, account.Snapshot().Used) + + arg.Reset(proc, false, nil) + require.Nil(t, arg.ctr.resBat) + require.Zero(t, account.Snapshot().Used) + require.NoError(t, arg.ClearAllocationAccount(account)) +} + +func TestLoopJoinResultBatchHonorsAllocationCapacity(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1) + require.NoError(t, err) + arg := &LoopJoin{ + ResultCols: []colexec.ResultPos{{Rel: 0, Pos: 0}}, + LeftTypes: []types.Type{types.T_int64.ToType()}, + } + require.NoError(t, arg.SetAllocationAccount(account)) + require.NoError(t, arg.resetResultBat()) + err = vector.AppendFixed(arg.ctr.resBat.Vecs[0], int64(1), false, proc.Mp()) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Zero(t, account.Snapshot().Used) + arg.Reset(proc, false, nil) + require.NoError(t, arg.ClearAllocationAccount(account)) +} + var ( tag int32 ) diff --git a/pkg/sql/colexec/loopjoin/types.go b/pkg/sql/colexec/loopjoin/types.go index fa0f66b84446b..791e5915272b0 100644 --- a/pkg/sql/colexec/loopjoin/types.go +++ b/pkg/sql/colexec/loopjoin/types.go @@ -23,6 +23,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" "github.com/matrixorigin/matrixone/pkg/vm" "github.com/matrixorigin/matrixone/pkg/vm/message" "github.com/matrixorigin/matrixone/pkg/vm/process" @@ -42,6 +43,13 @@ const ( loopJoinAllocationSiteBatchOffsets ) +const ( + loopJoinAllocationSiteResultData mpool.AllocationSite = iota + 106 + loopJoinAllocationSiteResultArea + loopJoinAllocationSiteResultNulls + loopJoinAllocationSiteResultGrouping +) + type container struct { state int probeIdx int @@ -74,6 +82,7 @@ type LoopJoin struct { JoinType plan.Node_JoinType MarkPos int allocationAccount *mpool.AllocationAccount + resultAllocation *vector.AllocationAccountSelection vm.OperatorBase } @@ -91,7 +100,19 @@ func (loopJoin *LoopJoin) SetAllocationAccount( if loopJoin.allocationAccount == account { return nil } + selection, err := vector.NewAllocationAccountSelection( + account, + hashbuild.HashBuildAllocationOwner, + loopJoinAllocationSiteResultData, + loopJoinAllocationSiteResultArea, + loopJoinAllocationSiteResultNulls, + loopJoinAllocationSiteResultGrouping, + ) + if err != nil { + return err + } loopJoin.allocationAccount = account + loopJoin.resultAllocation = selection return nil } @@ -106,10 +127,11 @@ func (loopJoin *LoopJoin) ClearAllocationAccount( } ctr := &loopJoin.ctr if ctr.mp != nil || ctr.expr != nil || ctr.rightRowsMatched != nil || - len(ctr.rightBatchOffset) != 0 { + len(ctr.rightBatchOffset) != 0 || ctr.resBat != nil { return mpool.ErrAllocationAccountInvariant } loopJoin.allocationAccount = nil + loopJoin.resultAllocation = nil return nil } @@ -152,6 +174,10 @@ func (loopJoin *LoopJoin) Reset(proc *process.Process, pipelineFailed bool, err // of carrying generation-bound storage across Reset. ctr.cleanNonEqCondExecutor() ctr.cleanHashMap() + if ctr.resBat != nil { + ctr.resBat.Clean(proc.GetMPool()) + ctr.resBat = nil + } ctr.state = Build ctr.inBat = nil ctr.cleanRightMatchState(proc) diff --git a/pkg/sql/colexec/rightdedupjoin/allocation_test_helpers_test.go b/pkg/sql/colexec/rightdedupjoin/allocation_test_helpers_test.go index dbe509f9c25a7..08768aaa6b6f6 100644 --- a/pkg/sql/colexec/rightdedupjoin/allocation_test_helpers_test.go +++ b/pkg/sql/colexec/rightdedupjoin/allocation_test_helpers_test.go @@ -18,6 +18,10 @@ import ( "testing" "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/testutil" "github.com/stretchr/testify/require" ) @@ -36,3 +40,45 @@ func installTestAllocation(t testing.TB, owners ...testAllocationOwner) *mpool.A } return account } + +func TestRightDedupJoinResultBatchUsesAllocationAccount(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + arg := &RightDedupJoin{ + Result: []colexec.ResultPos{{Rel: 0, Pos: 0}}, + LeftTypes: []types.Type{types.T_int64.ToType()}, + } + account := installTestAllocation(t, arg) + require.NoError(t, arg.resetResultBatch()) + require.Same(t, arg.resultAllocation, arg.ctr.resultBatch.Vecs[0].AllocationAccountSelection()) + require.NoError(t, vector.AppendFixed(arg.ctr.resultBatch.Vecs[0], int64(1), false, proc.Mp())) + used := account.Snapshot().Used + require.Positive(t, used) + require.NoError(t, arg.resetResultBatch()) + require.Equal(t, used, account.Snapshot().Used) + + arg.Reset(proc, false, nil) + require.Nil(t, arg.ctr.resultBatch) + require.Zero(t, account.Snapshot().Used) + require.NoError(t, arg.ClearAllocationAccount(account)) +} + +func TestRightDedupJoinResultBatchHonorsAllocationCapacity(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1) + require.NoError(t, err) + arg := &RightDedupJoin{ + Result: []colexec.ResultPos{{Rel: 0, Pos: 0}}, + LeftTypes: []types.Type{types.T_int64.ToType()}, + } + require.NoError(t, arg.SetAllocationAccount(account)) + require.NoError(t, arg.resetResultBatch()) + err = vector.AppendFixed(arg.ctr.resultBatch.Vecs[0], int64(1), false, proc.Mp()) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Zero(t, account.Snapshot().Used) + arg.Reset(proc, false, nil) + require.NoError(t, arg.ClearAllocationAccount(account)) +} diff --git a/pkg/sql/colexec/rightdedupjoin/join.go b/pkg/sql/colexec/rightdedupjoin/join.go index 55e157199c32b..27b9f3225cc61 100644 --- a/pkg/sql/colexec/rightdedupjoin/join.go +++ b/pkg/sql/colexec/rightdedupjoin/join.go @@ -64,6 +64,7 @@ func (rightDedupJoin *RightDedupJoin) Prepare(proc *process.Process) (err error) evalExecs, err = hashbuild.NewExpressionExecutors( proc, rightDedupJoin.Conditions[0], + rightDedupJoin.allocationAccount, ) if err != nil { return err @@ -73,6 +74,7 @@ func (rightDedupJoin *RightDedupJoin) Prepare(proc *process.Process) (err error) updateExecs, err = hashbuild.NewExpressionExecutors( proc, rightDedupJoin.UpdateColExprList, + rightDedupJoin.allocationAccount, ) if err != nil { for _, exec := range evalExecs { @@ -399,16 +401,8 @@ func (ctr *container) probe(bat *batch.Batch, ap *RightDedupJoin, proc *process. } } - ctr.resetResultBatch() - if ctr.resultBatch == nil { - ctr.resultBatch = batch.NewOffHeapWithSize(len(ap.Result)) - for i, rp := range ap.Result { - if rp.Rel == 0 { - ctr.resultBatch.Vecs[i] = vector.NewOffHeapVecWithType(ap.LeftTypes[rp.Pos]) - } else { - ctr.resultBatch.Vecs[i] = vector.NewOffHeapVecWithType(ap.RightTypes[rp.Pos]) - } - } + if err := ap.resetResultBatch(); err != nil { + return err } for i, rp := range ap.Result { diff --git a/pkg/sql/colexec/rightdedupjoin/types.go b/pkg/sql/colexec/rightdedupjoin/types.go index 818141ddcdd0a..54bbd08894387 100644 --- a/pkg/sql/colexec/rightdedupjoin/types.go +++ b/pkg/sql/colexec/rightdedupjoin/types.go @@ -24,6 +24,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" "github.com/matrixorigin/matrixone/pkg/sql/colexec/spillutil" "github.com/matrixorigin/matrixone/pkg/vm" "github.com/matrixorigin/matrixone/pkg/vm/message" @@ -41,6 +42,13 @@ const ( const rightDedupJoinAllocationSiteMatched mpool.AllocationSite = 90 +const ( + rightDedupJoinAllocationSiteResultData mpool.AllocationSite = iota + 114 + rightDedupJoinAllocationSiteResultArea + rightDedupJoinAllocationSiteResultNulls + rightDedupJoinAllocationSiteResultGrouping +) + type evalVector struct { executor colexec.ExpressionExecutor vec *vector.Vector @@ -89,6 +97,7 @@ type RightDedupJoin struct { UpdateColIdxList []int32 UpdateColExprList []*plan.Expr allocationAccount *mpool.AllocationAccount + resultAllocation *vector.AllocationAccountSelection vm.OperatorBase } @@ -106,7 +115,19 @@ func (rightDedupJoin *RightDedupJoin) SetAllocationAccount( if rightDedupJoin.allocationAccount == account { return nil } + selection, err := vector.NewAllocationAccountSelection( + account, + hashbuild.HashBuildAllocationOwner, + rightDedupJoinAllocationSiteResultData, + rightDedupJoinAllocationSiteResultArea, + rightDedupJoinAllocationSiteResultNulls, + rightDedupJoinAllocationSiteResultGrouping, + ) + if err != nil { + return err + } rightDedupJoin.allocationAccount = account + rightDedupJoin.resultAllocation = selection return nil } @@ -123,10 +144,12 @@ func (rightDedupJoin *RightDedupJoin) ClearAllocationAccount( rightDedupJoin.ctr.spillEngine != nil || len(rightDedupJoin.ctr.evecs) != 0 || len(rightDedupJoin.ctr.exprExecs) != 0 || - rightDedupJoin.ctr.matched != nil { + rightDedupJoin.ctr.matched != nil || + rightDedupJoin.ctr.resultBatch != nil { return mpool.ErrAllocationAccountInvariant } rightDedupJoin.allocationAccount = nil + rightDedupJoin.resultAllocation = nil return nil } @@ -174,7 +197,7 @@ func (rightDedupJoin *RightDedupJoin) Reset(proc *process.Process, pipelineFaile ctr.cleanBitmap(proc) ctr.cleanHashMap() - ctr.resetResultBatch() + ctr.cleanResultBatch(proc) ctr.cleanExprExecutor() if ctr.spillEngine != nil { ctr.spillEngine.Cleanup(proc) @@ -235,6 +258,32 @@ func (ctr *container) resetResultBatch() { } } +func (rightDedupJoin *RightDedupJoin) resetResultBatch() error { + ctr := &rightDedupJoin.ctr + ctr.resetResultBatch() + if ctr.resultBatch != nil { + return nil + } + ctr.resultBatch = batch.NewOffHeapWithSize(len(rightDedupJoin.Result)) + for i, rp := range rightDedupJoin.Result { + if rp.Rel == 0 { + ctr.resultBatch.Vecs[i] = vector.NewOffHeapVecWithType( + rightDedupJoin.LeftTypes[rp.Pos], + ) + } else { + ctr.resultBatch.Vecs[i] = vector.NewOffHeapVecWithType( + rightDedupJoin.RightTypes[rp.Pos], + ) + } + } + if err := ctr.resultBatch.SetAllocationAccount(rightDedupJoin.resultAllocation); err != nil { + ctr.resultBatch.Clean(nil) + ctr.resultBatch = nil + return err + } + return nil +} + func (ctr *container) cleanResultBatch(proc *process.Process) { if ctr.resultBatch != nil { ctr.resultBatch.Clean(proc.Mp()) diff --git a/pkg/sql/colexec/spillutil/allocation_account_test.go b/pkg/sql/colexec/spillutil/allocation_account_test.go index 37c9fd95eafc5..249daac7e970b 100644 --- a/pkg/sql/colexec/spillutil/allocation_account_test.go +++ b/pkg/sql/colexec/spillutil/allocation_account_test.go @@ -24,7 +24,9 @@ import ( "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" + plan2 "github.com/matrixorigin/matrixone/pkg/sql/plan" "github.com/matrixorigin/matrixone/pkg/testutil" "github.com/matrixorigin/matrixone/pkg/vm/process" "github.com/stretchr/testify/require" @@ -37,6 +39,32 @@ type testSpillAllocationAccount struct { generation *process.HashBuildBudgetGeneration } +func makeTestCastKeyExpr( + t testing.TB, + proc *process.Process, +) []*plan.Expr { + t.Helper() + expr, err := plan2.BindFuncExprImplByPlanExpr( + proc.Ctx, + "cast", + []*plan.Expr{ + { + Typ: plan.Type{Id: int32(types.T_int32)}, + Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, + }, + { + Typ: plan.Type{ + Id: int32(types.T_varchar), + Width: types.MaxVarcharLen, + }, + Expr: &plan.Expr_T{T: &plan.TargetType{}}, + }, + }, + ) + require.NoError(t, err) + return []*plan.Expr{expr} +} + func TestNewSpillEngineRequiresBudgetGeneration(t *testing.T) { registry, err := mpool.NewAllocationAccountRegistry(1, 1) require.NoError(t, err) @@ -748,7 +776,7 @@ func TestSpillAllocationAccountRebuildAndRecursiveSpillLifecycle(t *testing.T) { require.NoError(t, err) engine, err := NewSpillEngine( SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), + BuildKeyExprs: makeTestCastKeyExpr(t, proc), Budget: generation, SpillThreshold: 100, NeedsBuildForEmptyProbe: true, @@ -770,6 +798,7 @@ func TestSpillAllocationAccountRebuildAndRecursiveSpillLifecycle(t *testing.T) { respills := 0 ready := 0 + expressionStorageObserved := false for steps := 0; engine.HasMoreBuckets(); steps++ { require.Less(t, steps, 4_096, "recursive spill queue made no progress") jm, result, rebuildErr := engine.RebuildHashmap(proc, analyzer) @@ -777,6 +806,8 @@ func TestSpillAllocationAccountRebuildAndRecursiveSpillLifecycle(t *testing.T) { switch result { case BucketReSpilled: respills++ + expressionStorageObserved = expressionStorageObserved || + (len(engine.keyExecs) == 1 && account.Snapshot().Used > 0) case BucketReady: ready++ require.NotNil(t, jm) @@ -789,6 +820,7 @@ func TestSpillAllocationAccountRebuildAndRecursiveSpillLifecycle(t *testing.T) { } require.Positive(t, respills) require.Positive(t, ready) + require.True(t, expressionStorageObserved) require.Positive(t, account.Snapshot().Peak) engine.Cleanup(proc) diff --git a/pkg/sql/colexec/spillutil/join_spill.go b/pkg/sql/colexec/spillutil/join_spill.go index e13613d10e692..ef881bea66daa 100644 --- a/pkg/sql/colexec/spillutil/join_spill.go +++ b/pkg/sql/colexec/spillutil/join_spill.go @@ -2115,6 +2115,7 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal execs, err := hashbuild.NewExpressionExecutors( proc, e.cfg.BuildKeyExprs, + e.allocation.account, ) if err != nil { for _, exec := range execs { diff --git a/pkg/sql/util/eval_expr_util.go b/pkg/sql/util/eval_expr_util.go index da58ac7b71d38..56bf9d641aee3 100644 --- a/pkg/sql/util/eval_expr_util.go +++ b/pkg/sql/util/eval_expr_util.go @@ -111,13 +111,28 @@ func DecodeBinaryString(s string) ([]byte, error) { } func GenVectorByVarValue(proc *process.Process, typ types.Type, val any) (*vector.Vector, error) { + return GenVectorByVarValueWithAllocation(proc, typ, val, nil) +} + +func GenVectorByVarValueWithAllocation( + proc *process.Process, + typ types.Type, + val any, + selection *vector.AllocationAccountSelection, +) (*vector.Vector, error) { if val == nil { - vec := vector.NewConstNull(typ, 1, proc.Mp()) - return vec, nil - } else { - strVal := getVal(val) + if selection == nil { + return vector.NewConstNull(typ, 1, proc.Mp()), nil + } + return vector.NewConstNullWithAllocation(typ, 1, selection) + } + strVal := getVal(val) + if selection == nil { return vector.NewConstBytes(typ, []byte(strVal), 1, proc.Mp()) } + return vector.NewConstBytesWithAllocation( + typ, []byte(strVal), 1, proc.Mp(), selection, + ) } func AppendAnyToStringVector(proc *process.Process, val any, vec *vector.Vector) error { From d50c3404c6ed4bf9a19cb86c97299543a8507028 Mon Sep 17 00:00:00 2001 From: aptend Date: Sat, 1 Aug 2026 18:09:59 +0800 Subject: [PATCH 41/61] executor: finalize remote allocation accounts per statement --- ...ocation_accounted_memory_admission_impl.md | 41 + .../design/evidence/26459_local_validation.md | 27 + ...0_allocation_accounted_memory_admission.md | 30 + pkg/pb/pipeline/pipeline.pb.go | 1179 ++++++++++------- .../compile/allocation_account_lifecycle.go | 123 +- .../allocation_account_lifecycle_test.go | 32 + pkg/sql/compile/analyze_module.go | 108 +- pkg/sql/compile/compile.go | 5 + pkg/sql/compile/compile2.go | 12 + .../remote_allocation_statement_group.go | 466 +++++++ .../remote_allocation_statement_group_test.go | 394 ++++++ pkg/sql/compile/remoterun.go | 7 + pkg/sql/compile/remoterunClient.go | 26 +- pkg/sql/compile/remoterunServer.go | 233 +++- pkg/sql/compile/remoterun_test.go | 43 +- pkg/sql/compile/resource_accounting.go | 142 +- pkg/sql/compile/resource_accounting_test.go | 123 ++ pkg/sql/compile/types.go | 2 + proto/pipeline.proto | 6 + 19 files changed, 2410 insertions(+), 589 deletions(-) create mode 100644 pkg/sql/compile/remote_allocation_statement_group.go create mode 100644 pkg/sql/compile/remote_allocation_statement_group_test.go diff --git a/docs/design/allocation_accounted_memory_admission_impl.md b/docs/design/allocation_accounted_memory_admission_impl.md index 13e8015fabb5f..901f94611dce6 100644 --- a/docs/design/allocation_accounted_memory_admission_impl.md +++ b/docs/design/allocation_accounted_memory_admission_impl.md @@ -100,6 +100,44 @@ The sequence is: 7. seal and finalize the account; 8. export exactly one terminal snapshot. +For a coordinator-local attempt, that sequence owns one MessageBoard. Remote +PipelineMessage handlers are different: all fragments of the same statement on +one CN share the board, and accounted JoinMap/result storage may remain owned by +a producer account until a sibling consumes it. The coordinator therefore +computes the complete physical scope graph's RPC count per CN, creates a unique +physical execution ID, and carries both in ProcessInfo. Each remote handler +registers before decoding/execution, stages its account and MPool after its +operators quiesce, and detaches the shared board before Compile cleanup. A +quiesced handler returns without waiting for siblings and adds one counted +pending terminal-memory-domain signal. This avoids a B-to-C-to-B response +dependency cycle. +The final fragment drains the board, completes every staged account, samples +every staged MPool, and publishes the aggregate plus a completion marker exactly +once. The coordinator reduces pending/completed markers independent of response +order; unresolved counts preserve every suppressed fragment MPool domain and +make the attempt explicitly partial instead of silently omitting memory. + +If dispatch fails before every planned handler arrives, the first quiesced +fragment starts a bounded orphan-registration timer. Registration of the full +set cancels it. Expiry closes only that board generation; active registered +fragments retain their accounts until they quiesce, after which the final +registered fragment completes the group. If no response remains to carry a +completion marker, the already-returned pending marker preserves the missing +domain as an explicit partial result. A fragment failure aborts an incomplete +group without waiting for the timer and cancels active registered siblings. +This bounds the MessageCenter/group maps without sealing live +allocations or imposing an execution timeout on a fully registered query. +Prepared executions and retries use fresh execution IDs, so a stale group and +board cannot capture the next generation. + +ProcessInfo without both topology and execution-ID metadata remains decodable. +Remote plans outside the accounted owner domain keep their legacy lifecycle; +plans containing an allocation-account owner are rejected explicitly. The +server never guesses a group size and never silently falls back to unaccounted +HashBuild execution. In the opposite upgrade direction, query-candidate +discovery already excludes CNs whose CommitID differs from the coordinator +binary, so a new coordinator cannot dispatch this lifecycle to an old handler. + Prepared statements and retries create a new generation. Reset frees all generation-bound state; it does not carry an executor, bitmap, mailbox payload, or allocation selection into the next attempt. Runtime parallel clones are @@ -185,5 +223,8 @@ The implementation is complete only when: - every transfer has exactly one owner after success and on cancellation; - memory, disk, and FD rejection remain distinct; - prepared/retry generations terminate independently at zero; +- remote fragments sharing one MessageBoard terminate at the per-CN statement + boundary, with incomplete dispatch bounded and old board generations unable + to close a replacement; - local unit, race, build, vet, lifecycle, spill, and performance checks pass; - independent reviews report no blocker or major correctness/performance issue. diff --git a/docs/design/evidence/26459_local_validation.md b/docs/design/evidence/26459_local_validation.md index 60cffb99a851e..081d4400ee31a 100644 --- a/docs/design/evidence/26459_local_validation.md +++ b/docs/design/evidence/26459_local_validation.md @@ -69,6 +69,17 @@ The final semantic edit was followed by a clean local run on 2026-08-01. No partial or still-running session is counted as a pass. +After the distributed q7 failure exposed the shared remote-MessageBoard +boundary, the final remote-lifecycle amendment was validated separately: + +- the complete `pkg/sql/compile` suite passed in normal and race modes; +- `pkg/pb/pipeline`, `pkg/vm/message`, and `pkg/vm/process` passed; +- protobuf regeneration was clean and reproducible; +- vet passed for all four affected packages; +- an independent lifecycle/concurrency review converged with no blocker or + major finding after its pending-domain cardinality and mixed-version concerns + were resolved in code or by existing scheduling evidence. + ## Behavioral coverage The local suite covers: @@ -93,6 +104,22 @@ The local suite covers: modes; - Product cleanup and account terminal zero; - accounted JoinMap release after an unaccounted ProductL2 consumer frees it. +- remote scope-graph fragment counting across nested CN execution addresses; +- ProcessInfo topology-map and execution-ID wire round trips, with distinct + MessageBoard generations across retries; +- remote statement-group board drain after a producer finishes before a later + sibling registers, including accounted queued-message destruction, exact + account terminal zero, and aggregate MPool terminal sampling; +- incomplete remote dispatch expiry releases staged account and MPool domains, + while an unresolved pending terminal marker makes the coordinator summary + explicitly partial; this includes the case where another registered fragment + is still active when the old board generation closes; +- fragment failure aborts an incomplete group immediately, and a legacy remote + plan containing an accounted owner is rejected instead of running without an + account; +- counted pending/completed group markers resolve independent of + terminal-response order; a four-fragment lost-final-response case preserves + all three suppressed reported domains plus the directly missing domain. ## Performance evidence diff --git a/docs/rfcs/00000000_allocation_accounted_memory_admission.md b/docs/rfcs/00000000_allocation_accounted_memory_admission.md index 250ff2faa729a..58eaaf9b25244 100644 --- a/docs/rfcs/00000000_allocation_accounted_memory_admission.md +++ b/docs/rfcs/00000000_allocation_accounted_memory_admission.md @@ -154,6 +154,36 @@ live metadata. A late allocation or release mismatch is a lifecycle invariant, not capacity pressure. Prepared statements and retries use new generations; generation-bound state cannot survive Reset. +Remote pipeline RPCs for one statement on one CN share a MessageBoard, so an +individual RPC is not a terminal ownership boundary. The coordinator carries +the exact number of planned fragments per target CN plus a unique physical +execution ID. Every received fragment joins the execution-and-board-keyed +statement group, clears its reachable operators when it quiesces, and transfers +its account and MPool terminal ownership to that group. Quiesced handlers do not +wait for siblings, because nested B-to-C-to-B execution can otherwise create a +response dependency cycle. Their terminal responses carry a counted pending +memory-domain signal. The finalizer drains the board, completes all accounts, samples all +fragment MPool domains, and lets exactly one terminal response publish the +aggregate plus a completion marker; the other responses publish no duplicate +terminal facts. Pending/completed signals reduce independent of response order; +unresolved counts preserve the cardinality of every suppressed fragment MPool +domain and make the root summary explicitly partial. +A missing planned RPC is bounded by the MessageBoard receive interval: expiry +closes the old board generation and cancels active registered fragments; their +accounts remain live until those fragments actually quiesce. A fragment failure +aborts the incomplete group immediately. The timer is canceled once all planned +fragments register, so it never limits a fully dispatched statement. Prepared +executions and retries use new execution IDs, so an incomplete old group cannot +capture a replacement. + +The topology and execution-ID fields are a required capability for remote +plans containing allocation-account owners. A legacy ProcessInfo remains +decodable, and remote plans outside the accounted domain keep their legacy +lifecycle. An accounted legacy plan is rejected explicitly rather than running +with a guessed fragment count or an unaccounted fallback. Query-candidate +discovery admits only CNs with the coordinator binary's CommitID, which prevents +the inverse new-coordinator-to-old-handler execution during a rolling upgrade. + ## Pressure protocol Typed reasons keep control flow honest: diff --git a/pkg/pb/pipeline/pipeline.pb.go b/pkg/pb/pipeline/pipeline.pb.go index f06d2144764d6..9025de3bb79f6 100644 --- a/pkg/pb/pipeline/pipeline.pb.go +++ b/pkg/pb/pipeline/pipeline.pb.go @@ -5363,19 +5363,25 @@ func (m *PrepareParamInfo) GetIsBin() []bool { } type ProcessInfo struct { - Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` - Sql string `protobuf:"bytes,2,opt,name=sql,proto3" json:"sql,omitempty"` - Lim ProcessLimitation `protobuf:"bytes,3,opt,name=lim,proto3" json:"lim"` - UnixTime int64 `protobuf:"varint,4,opt,name=unix_time,json=unixTime,proto3" json:"unix_time,omitempty"` - AccountId uint32 `protobuf:"varint,5,opt,name=account_id,json=accountId,proto3" json:"account_id,omitempty"` - Snapshot txn.CNTxnSnapshot `protobuf:"bytes,6,opt,name=snapshot,proto3" json:"snapshot"` - SessionInfo SessionInfo `protobuf:"bytes,7,opt,name=session_info,json=sessionInfo,proto3" json:"session_info"` - SessionLogger SessionLoggerInfo `protobuf:"bytes,8,opt,name=session_logger,json=sessionLogger,proto3" json:"session_logger"` - PrepareParams PrepareParamInfo `protobuf:"bytes,9,opt,name=prepare_params,json=prepareParams,proto3" json:"prepare_params"` - AffectedRows int64 `protobuf:"varint,10,opt,name=affected_rows,json=affectedRows,proto3" json:"affected_rows,omitempty"` - XXX_NoUnkeyedLiteral struct{} `json:"-"` - XXX_unrecognized []byte `json:"-"` - XXX_sizecache int32 `json:"-"` + Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` + Sql string `protobuf:"bytes,2,opt,name=sql,proto3" json:"sql,omitempty"` + Lim ProcessLimitation `protobuf:"bytes,3,opt,name=lim,proto3" json:"lim"` + UnixTime int64 `protobuf:"varint,4,opt,name=unix_time,json=unixTime,proto3" json:"unix_time,omitempty"` + AccountId uint32 `protobuf:"varint,5,opt,name=account_id,json=accountId,proto3" json:"account_id,omitempty"` + Snapshot txn.CNTxnSnapshot `protobuf:"bytes,6,opt,name=snapshot,proto3" json:"snapshot"` + SessionInfo SessionInfo `protobuf:"bytes,7,opt,name=session_info,json=sessionInfo,proto3" json:"session_info"` + SessionLogger SessionLoggerInfo `protobuf:"bytes,8,opt,name=session_logger,json=sessionLogger,proto3" json:"session_logger"` + PrepareParams PrepareParamInfo `protobuf:"bytes,9,opt,name=prepare_params,json=prepareParams,proto3" json:"prepare_params"` + AffectedRows int64 `protobuf:"varint,10,opt,name=affected_rows,json=affectedRows,proto3" json:"affected_rows,omitempty"` + // Planned PipelineMessage RPCs per target CN for statement-level remote + // MessageBoard and resource-account terminal ownership. + RemoteFragmentCounts map[string]uint32 `protobuf:"bytes,11,rep,name=remote_fragment_counts,json=remoteFragmentCounts,proto3" json:"remote_fragment_counts,omitempty" protobuf_key:"bytes,1,opt,name=key,proto3" protobuf_val:"varint,2,opt,name=value,proto3"` + // Unique physical execution attempt. Unlike the SQL statement ID, this + // changes across retries and prepared-statement executions. + RemoteExecutionId []byte `protobuf:"bytes,12,opt,name=remote_execution_id,json=remoteExecutionId,proto3" json:"remote_execution_id,omitempty"` + XXX_NoUnkeyedLiteral struct{} `json:"-"` + XXX_unrecognized []byte `json:"-"` + XXX_sizecache int32 `json:"-"` } func (m *ProcessInfo) Reset() { *m = ProcessInfo{} } @@ -5481,6 +5487,20 @@ func (m *ProcessInfo) GetAffectedRows() int64 { return 0 } +func (m *ProcessInfo) GetRemoteFragmentCounts() map[string]uint32 { + if m != nil { + return m.RemoteFragmentCounts + } + return nil +} + +func (m *ProcessInfo) GetRemoteExecutionId() []byte { + if m != nil { + return m.RemoteExecutionId + } + return nil +} + type SessionInfo struct { User string `protobuf:"bytes,1,opt,name=user,proto3" json:"user,omitempty"` Host string `protobuf:"bytes,2,opt,name=host,proto3" json:"host,omitempty"` @@ -6103,6 +6123,7 @@ func init() { proto.RegisterType((*ProcessLimitation)(nil), "pipeline.ProcessLimitation") proto.RegisterType((*PrepareParamInfo)(nil), "pipeline.PrepareParamInfo") proto.RegisterType((*ProcessInfo)(nil), "pipeline.ProcessInfo") + proto.RegisterMapType((map[string]uint32)(nil), "pipeline.ProcessInfo.RemoteFragmentCountsEntry") proto.RegisterType((*SessionInfo)(nil), "pipeline.SessionInfo") proto.RegisterType((*SessionLoggerInfo)(nil), "pipeline.SessionLoggerInfo") proto.RegisterType((*Pipeline)(nil), "pipeline.Pipeline") @@ -6114,481 +6135,486 @@ func init() { func init() { proto.RegisterFile("pipeline.proto", fileDescriptor_7ac67a7adf3df9c7) } var fileDescriptor_7ac67a7adf3df9c7 = []byte{ - // 7584 bytes of a gzipped FileDescriptorProto + // 7656 bytes of a gzipped FileDescriptorProto 0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0xec, 0x7c, 0x4d, 0x6f, 0x1c, 0xc9, - 0x92, 0xd8, 0x34, 0xfb, 0x3b, 0xfa, 0x83, 0xcd, 0x24, 0x45, 0xb5, 0xa4, 0x79, 0x23, 0x4d, 0xcf, - 0x48, 0xc3, 0xa7, 0xd1, 0x50, 0x12, 0x67, 0xe6, 0xbd, 0xd9, 0xf7, 0xf6, 0xed, 0x5b, 0x8a, 0x92, - 0xde, 0xf0, 0x3d, 0x51, 0xe2, 0x16, 0x29, 0x0f, 0x30, 0x80, 0x5d, 0x28, 0x56, 0x65, 0x77, 0xd7, - 0xb0, 0xba, 0xb2, 0x54, 0x99, 0x25, 0x91, 0xba, 0xd8, 0x07, 0x9f, 0x7c, 0xf1, 0x71, 0xf7, 0xb8, - 0x80, 0x7d, 0x58, 0xdb, 0x07, 0x1f, 0x8c, 0xe7, 0x9f, 0x60, 0x2c, 0x6c, 0xc3, 0x58, 0xf8, 0xe0, - 0xa3, 0x61, 0xbc, 0x3d, 0x1a, 0x30, 0x0c, 0x03, 0x36, 0x16, 0x30, 0x0c, 0x18, 0x11, 0x99, 0x59, - 0x55, 0xdd, 0x4d, 0x69, 0x3e, 0x6c, 0xec, 0x65, 0xf7, 0xd4, 0x95, 0x11, 0x91, 0x59, 0x59, 0x91, - 0x91, 0x11, 0x91, 0x11, 0x91, 0x0d, 0xfd, 0x24, 0x4c, 0x78, 0x14, 0xc6, 0x7c, 0x3b, 0x49, 0x85, - 0x12, 0xac, 0x65, 0xdb, 0x57, 0x3f, 0x99, 0x84, 0x6a, 0x9a, 0x9d, 0x6c, 0xfb, 0x62, 0x76, 0x77, - 0x22, 0x26, 0xe2, 0x2e, 0x11, 0x9c, 0x64, 0x63, 0x6a, 0x51, 0x83, 0x9e, 0x74, 0xc7, 0xab, 0x10, - 0x09, 0xff, 0xd4, 0x3e, 0x27, 0x91, 0x17, 0x9b, 0xe7, 0x55, 0x15, 0xce, 0xb8, 0x54, 0xde, 0x2c, - 0x31, 0x80, 0xb6, 0x3a, 0x33, 0xb8, 0xd1, 0xbf, 0xab, 0x42, 0xf3, 0x80, 0x4b, 0xe9, 0x4d, 0x38, - 0x1b, 0x41, 0x55, 0x86, 0xc1, 0xb0, 0x72, 0xa3, 0xb2, 0xd5, 0xdf, 0x19, 0x6c, 0xe7, 0xd3, 0x3a, - 0x52, 0x9e, 0xca, 0xa4, 0x83, 0x48, 0xa4, 0xf1, 0x67, 0xc1, 0x70, 0x65, 0x91, 0xe6, 0x80, 0xab, - 0xa9, 0x08, 0x1c, 0x44, 0xb2, 0x01, 0x54, 0x79, 0x9a, 0x0e, 0xab, 0x37, 0x2a, 0x5b, 0x5d, 0x07, - 0x1f, 0x19, 0x83, 0x5a, 0xe0, 0x29, 0x6f, 0x58, 0x23, 0x10, 0x3d, 0xb3, 0x0f, 0xa1, 0x9f, 0xa4, - 0xc2, 0x77, 0xc3, 0x78, 0x2c, 0x5c, 0xc2, 0xd6, 0x09, 0xdb, 0x45, 0xe8, 0x7e, 0x3c, 0x16, 0x0f, - 0x91, 0x6a, 0x08, 0x4d, 0x2f, 0xf6, 0xa2, 0x73, 0xc9, 0x87, 0x0d, 0x42, 0xdb, 0x26, 0xeb, 0xc3, - 0x4a, 0x18, 0x0c, 0x9b, 0x37, 0x2a, 0x5b, 0x35, 0x67, 0x25, 0x0c, 0xf0, 0x1d, 0x59, 0x16, 0x06, - 0xc3, 0x96, 0x7e, 0x07, 0x3e, 0xb3, 0x11, 0x74, 0x63, 0xce, 0x83, 0xa7, 0x42, 0x39, 0x3c, 0x89, - 0xce, 0x87, 0xed, 0x1b, 0x95, 0xad, 0x96, 0x33, 0x07, 0x63, 0x57, 0xa1, 0x15, 0xf0, 0x93, 0x6c, - 0x72, 0x20, 0x27, 0x43, 0xb8, 0x51, 0xd9, 0x6a, 0x3b, 0x79, 0x9b, 0x1d, 0xc3, 0xe5, 0x94, 0xbf, - 0xc8, 0xb8, 0x54, 0x3c, 0x70, 0x15, 0xf7, 0xd2, 0x40, 0xbc, 0x8a, 0xdd, 0x99, 0x08, 0xf8, 0xb0, - 0x43, 0x1c, 0x78, 0xb7, 0xcc, 0xa5, 0x94, 0x7b, 0xb3, 0x63, 0x43, 0x74, 0x20, 0x02, 0xee, 0x5c, - 0xca, 0x3b, 0x97, 0xc1, 0xcc, 0x81, 0x4d, 0xcf, 0xf7, 0x79, 0xb2, 0x3c, 0x68, 0xf7, 0x3b, 0x0c, - 0xba, 0x61, 0xfb, 0x96, 0xa1, 0x3f, 0xab, 0xfd, 0xc9, 0x9f, 0x5e, 0x7f, 0x67, 0xf4, 0x1c, 0xda, - 0x7b, 0x22, 0x8e, 0xb9, 0xaf, 0x44, 0xca, 0xae, 0x43, 0xc7, 0x8e, 0xe3, 0x9a, 0x65, 0xad, 0x3b, - 0x60, 0x41, 0xfb, 0x01, 0xfb, 0x08, 0x56, 0x7d, 0x4b, 0xed, 0x86, 0x71, 0xc0, 0xcf, 0x68, 0x5d, - 0xeb, 0x4e, 0x3f, 0x07, 0xef, 0x23, 0x74, 0xf4, 0x6f, 0xaa, 0xd0, 0x3c, 0x9a, 0x66, 0xe3, 0x71, - 0xc4, 0xd9, 0x87, 0xd0, 0x33, 0x8f, 0x7b, 0x22, 0xda, 0x0f, 0xce, 0xcc, 0xb8, 0xf3, 0x40, 0x76, - 0x03, 0x3a, 0x06, 0x70, 0x7c, 0x9e, 0x70, 0x33, 0x6c, 0x19, 0x34, 0x3f, 0xce, 0x41, 0x18, 0x93, - 0xb8, 0x54, 0x9d, 0x79, 0xe0, 0x02, 0x95, 0x77, 0x46, 0x12, 0x34, 0x4f, 0xe5, 0xd1, 0xdb, 0x76, - 0xa3, 0xf0, 0x25, 0x77, 0xf8, 0x64, 0x2f, 0x56, 0x24, 0x47, 0x75, 0xa7, 0x0c, 0x62, 0x3b, 0x70, - 0x49, 0xea, 0x2e, 0x6e, 0xea, 0xc5, 0x13, 0x2e, 0xdd, 0x2c, 0x8c, 0xd5, 0x4f, 0x3e, 0x1b, 0x36, - 0x6e, 0x54, 0xb7, 0x6a, 0xce, 0xba, 0x41, 0x3a, 0x84, 0x7b, 0x4e, 0x28, 0x76, 0x0f, 0x36, 0x16, - 0xfa, 0xe8, 0x2e, 0xcd, 0x1b, 0xd5, 0xad, 0xaa, 0xc3, 0xe6, 0xba, 0xec, 0x53, 0x8f, 0x47, 0xb0, - 0x96, 0x66, 0x31, 0xee, 0xb6, 0xc7, 0x61, 0xa4, 0x78, 0x7a, 0x94, 0x70, 0x9f, 0xe4, 0xb1, 0xb3, - 0x73, 0x79, 0x9b, 0x36, 0xa4, 0xb3, 0x88, 0x76, 0x96, 0x7b, 0xb0, 0x3b, 0x39, 0xf3, 0x1e, 0x9d, - 0x25, 0x29, 0x09, 0x6d, 0x67, 0x07, 0xf4, 0x00, 0x08, 0x71, 0xca, 0x68, 0x76, 0x1b, 0xd6, 0x82, - 0xd4, 0x0b, 0x63, 0xd7, 0x8b, 0x22, 0xf7, 0x24, 0xf3, 0x4f, 0xb9, 0x92, 0x24, 0xc8, 0x2d, 0x67, - 0x95, 0x10, 0xbb, 0x51, 0xf4, 0x40, 0x83, 0x47, 0x7f, 0xb5, 0x02, 0xad, 0x87, 0xa1, 0x4c, 0x3c, - 0xe5, 0x4f, 0xd9, 0x65, 0x68, 0x8e, 0xb3, 0xd8, 0x2f, 0x64, 0xa3, 0x81, 0xcd, 0xfd, 0x80, 0xfd, - 0x3e, 0xac, 0x46, 0xc2, 0xf7, 0x22, 0x37, 0x17, 0x83, 0xe1, 0xca, 0x8d, 0xea, 0x56, 0x67, 0x67, - 0xbd, 0x10, 0xcc, 0x5c, 0xcc, 0x9c, 0x3e, 0xd1, 0x16, 0x62, 0xf7, 0x0b, 0x18, 0xa4, 0x7c, 0x26, - 0x14, 0x2f, 0x75, 0xaf, 0x52, 0x77, 0x56, 0x74, 0xff, 0x2a, 0xf5, 0x92, 0xa7, 0x28, 0xcd, 0xab, - 0x9a, 0xb6, 0xe8, 0x7e, 0xbf, 0xb4, 0x52, 0x7c, 0xe2, 0x86, 0xc1, 0x99, 0x4b, 0x2f, 0x18, 0xd6, - 0x6e, 0x54, 0xb7, 0xea, 0x05, 0xdb, 0xf9, 0x64, 0x3f, 0x38, 0x7b, 0x82, 0x18, 0xf6, 0x29, 0x6c, - 0x2e, 0x76, 0xd1, 0xa3, 0x0e, 0xeb, 0xd4, 0x67, 0x7d, 0xae, 0x8f, 0x43, 0x28, 0xf6, 0x3e, 0x74, - 0x6d, 0x27, 0x85, 0x22, 0xda, 0xd0, 0x42, 0x23, 0x4b, 0x22, 0x7a, 0x19, 0x9a, 0xa1, 0x74, 0x65, - 0x18, 0x9f, 0x92, 0x9a, 0x69, 0x39, 0x8d, 0x50, 0x1e, 0x85, 0xf1, 0x29, 0xbb, 0x02, 0xad, 0x94, - 0xfb, 0x1a, 0xd3, 0x22, 0x4c, 0x33, 0xe5, 0x3e, 0xa1, 0x2e, 0x03, 0x3e, 0xba, 0xbe, 0xe2, 0x46, - 0xd9, 0x34, 0x52, 0xee, 0xef, 0x29, 0x3e, 0x92, 0x50, 0x3f, 0xe0, 0xe9, 0x84, 0xa3, 0xbe, 0xc1, - 0x8e, 0x47, 0xbe, 0x17, 0x13, 0xdf, 0x5b, 0x4e, 0xde, 0x46, 0x6d, 0x97, 0x78, 0xa9, 0x0a, 0xbd, - 0x88, 0xb6, 0x4c, 0xcb, 0xb1, 0x4d, 0x76, 0x0d, 0xda, 0x52, 0x79, 0xa9, 0xc2, 0xaf, 0xa3, 0xad, - 0x52, 0x77, 0x5a, 0x04, 0xc0, 0xdd, 0x76, 0x19, 0x9a, 0x3c, 0x0e, 0x08, 0x55, 0xd3, 0x2b, 0xc9, - 0xe3, 0x60, 0x3f, 0x38, 0x1b, 0xfd, 0xab, 0x0a, 0xf4, 0x0e, 0xb2, 0x48, 0x85, 0xbb, 0xe9, 0x24, - 0xe3, 0xb3, 0x58, 0xa1, 0x96, 0x7c, 0x18, 0x4a, 0x65, 0xde, 0x4c, 0xcf, 0x6c, 0x0b, 0xda, 0xbf, - 0x4a, 0x45, 0x96, 0x90, 0xb4, 0xe9, 0x95, 0x2e, 0x4b, 0x5b, 0x81, 0x44, 0xc9, 0x7c, 0x96, 0x06, - 0x3c, 0x7d, 0x70, 0x4e, 0xb4, 0xd5, 0x25, 0xda, 0x32, 0x9a, 0xbd, 0x0b, 0xed, 0x23, 0x9e, 0x78, - 0xa9, 0x87, 0x22, 0x50, 0x23, 0xd5, 0x5a, 0x00, 0xf0, 0x5b, 0x89, 0x78, 0x3f, 0x30, 0x1b, 0xd6, - 0x36, 0x47, 0xff, 0xa4, 0x02, 0xed, 0xdd, 0xc9, 0x24, 0xe5, 0x13, 0x4f, 0x91, 0x9e, 0x17, 0x09, - 0xcd, 0xb7, 0xea, 0xac, 0x88, 0x84, 0x6c, 0x09, 0x7e, 0x81, 0x66, 0x10, 0x3d, 0xb3, 0xf7, 0xa0, - 0xc6, 0x2f, 0x9e, 0x10, 0xc1, 0xd9, 0x26, 0x34, 0x7c, 0x11, 0x8f, 0xc3, 0x89, 0xb1, 0x40, 0xa6, - 0xc5, 0x7e, 0x06, 0x1d, 0xfd, 0xa4, 0x65, 0xa0, 0x4e, 0xea, 0xf7, 0x8a, 0xee, 0x9e, 0xcf, 0x60, - 0x8f, 0x28, 0x50, 0x22, 0x1c, 0xf0, 0xf3, 0xe7, 0xd1, 0x3f, 0xaf, 0x42, 0x9d, 0x38, 0x83, 0x6b, - 0x83, 0x16, 0xc5, 0xe5, 0x2f, 0xbd, 0xc8, 0x2e, 0x29, 0x02, 0x1e, 0xbd, 0xf4, 0x22, 0x76, 0x03, - 0xea, 0x38, 0x05, 0x79, 0x01, 0x63, 0x35, 0x82, 0xdd, 0x82, 0x3a, 0xbe, 0x5d, 0xce, 0xcf, 0x1e, - 0xdf, 0xf1, 0xa0, 0xf6, 0xe7, 0xff, 0xf9, 0xfa, 0x3b, 0x8e, 0x46, 0xb3, 0x8f, 0xa0, 0xe6, 0x4d, - 0x26, 0x92, 0x36, 0xc2, 0xdc, 0x5e, 0xcc, 0x67, 0xea, 0x10, 0x01, 0xfb, 0x1c, 0xda, 0x7a, 0xd1, - 0x91, 0xba, 0x4e, 0xd4, 0x97, 0x4b, 0x96, 0xba, 0x2c, 0x0f, 0x4e, 0x41, 0x89, 0xcb, 0x15, 0x4a, - 0xa3, 0x59, 0x68, 0x3b, 0xb4, 0x9c, 0x02, 0x80, 0xa6, 0x34, 0x49, 0xf9, 0x6e, 0x14, 0x09, 0xff, - 0x28, 0x7c, 0xcd, 0x8d, 0xe1, 0x9d, 0x83, 0xb1, 0x5b, 0xd0, 0x3f, 0xd4, 0xf2, 0xea, 0x70, 0x99, - 0x45, 0x4a, 0x1a, 0x63, 0xbc, 0x00, 0x65, 0xdb, 0xc0, 0xe6, 0x20, 0xc7, 0xf4, 0xf9, 0xed, 0x1b, - 0xd5, 0xad, 0x9e, 0x73, 0x01, 0x86, 0x7d, 0x00, 0xbd, 0x09, 0x72, 0x3a, 0x8c, 0x27, 0xee, 0x38, - 0xf2, 0xd0, 0x4e, 0x57, 0xd1, 0x8e, 0x5b, 0xe0, 0xe3, 0xc8, 0x9b, 0xd0, 0x0e, 0x49, 0xc2, 0x28, - 0x72, 0x67, 0x7c, 0x46, 0xd6, 0xb9, 0xea, 0xb4, 0x08, 0x70, 0xc0, 0x67, 0xa3, 0x7f, 0x51, 0x83, - 0xc6, 0x7e, 0x2c, 0x79, 0xaa, 0x70, 0xff, 0x79, 0xe3, 0x31, 0xf7, 0x15, 0xd7, 0x7a, 0xaf, 0xe6, - 0xe4, 0x6d, 0x64, 0xc1, 0xb1, 0xf8, 0x2a, 0x0d, 0x15, 0x3f, 0xfa, 0xd4, 0x08, 0x58, 0x01, 0x40, - 0x4d, 0xeb, 0x05, 0x81, 0x6b, 0xa9, 0xdd, 0x54, 0xbc, 0x92, 0xb4, 0x17, 0x5b, 0xce, 0xaa, 0x17, - 0x04, 0xbb, 0x06, 0xee, 0x88, 0x57, 0x92, 0xbd, 0x0f, 0xd5, 0x94, 0x8f, 0x49, 0xdc, 0x3a, 0x3b, - 0xab, 0x7a, 0x49, 0x9f, 0x9d, 0x7c, 0xc3, 0x7d, 0xe5, 0xf0, 0xb1, 0x83, 0x38, 0xb6, 0x01, 0x75, - 0x4f, 0xa9, 0x54, 0x2f, 0x51, 0xdb, 0xd1, 0x0d, 0xb6, 0x0d, 0xeb, 0xb4, 0xe7, 0x55, 0x28, 0x62, - 0x57, 0x79, 0x27, 0x11, 0x1a, 0x6f, 0x69, 0xec, 0xd4, 0x5a, 0x8e, 0x3a, 0x46, 0xcc, 0x7e, 0x20, - 0xd1, 0xb2, 0x2d, 0xd2, 0xc7, 0xde, 0x8c, 0x4b, 0x32, 0x53, 0x6d, 0x67, 0x7d, 0xbe, 0xc7, 0x53, - 0x44, 0x21, 0x3f, 0x8b, 0x3e, 0xa8, 0x35, 0x5a, 0xb4, 0x01, 0xbb, 0x39, 0x10, 0x95, 0xca, 0x25, - 0x68, 0x84, 0xd2, 0xe5, 0x71, 0x60, 0x14, 0x59, 0x3d, 0x94, 0x8f, 0xe2, 0x80, 0x7d, 0x0c, 0x6d, - 0xfd, 0x96, 0x80, 0x8f, 0xc9, 0xcc, 0x74, 0x76, 0xfa, 0x46, 0x62, 0x11, 0xfc, 0x90, 0x8f, 0x9d, - 0x96, 0x32, 0x4f, 0xe8, 0x82, 0x28, 0xe1, 0xf2, 0x33, 0xc5, 0xd3, 0xd8, 0x8b, 0x68, 0x55, 0x5a, - 0x0e, 0x28, 0xf1, 0xc8, 0x40, 0xd8, 0xe7, 0x70, 0xd9, 0x62, 0x5d, 0xa9, 0x66, 0xca, 0xcd, 0xe2, - 0xf0, 0xcc, 0x8d, 0xbd, 0x58, 0x90, 0x2f, 0x54, 0x75, 0x36, 0x2c, 0xfa, 0x48, 0xcd, 0xd4, 0xf3, - 0x38, 0x3c, 0x7b, 0xea, 0xc5, 0x82, 0x6d, 0xc1, 0x20, 0xef, 0xa6, 0x5e, 0xd3, 0x07, 0x0f, 0x7b, - 0xa4, 0x60, 0xfa, 0x16, 0x7e, 0xfc, 0x1a, 0xbf, 0x15, 0x6d, 0x43, 0x99, 0x52, 0x8c, 0xc7, 0x92, - 0x2b, 0x57, 0x72, 0x7f, 0xd8, 0xa7, 0x6f, 0x5e, 0x2f, 0xe8, 0x9f, 0x11, 0xee, 0x88, 0xfb, 0xa3, - 0xdf, 0x56, 0xa0, 0x43, 0xfb, 0xe2, 0x79, 0x12, 0xa0, 0x0a, 0xfa, 0x00, 0x7a, 0xf3, 0x8b, 0xae, - 0xe5, 0xa6, 0xeb, 0x95, 0x57, 0x7c, 0x13, 0x1a, 0xbb, 0x3e, 0x32, 0x8f, 0x04, 0xa7, 0xe7, 0x98, - 0x16, 0xfb, 0x29, 0xac, 0x66, 0x34, 0x8c, 0xeb, 0xab, 0x33, 0x37, 0x42, 0xd5, 0xa5, 0x37, 0xba, - 0x91, 0x0a, 0xfd, 0x8e, 0x3d, 0x75, 0xe6, 0xf4, 0x32, 0xfb, 0xf8, 0x04, 0x95, 0xda, 0x3d, 0xd8, - 0x48, 0x39, 0x4a, 0x8c, 0xfb, 0x9a, 0xa7, 0xc2, 0x55, 0x7c, 0x96, 0x88, 0x94, 0x0c, 0x21, 0x72, - 0x91, 0x69, 0xdc, 0xd7, 0x3c, 0x15, 0xc7, 0x06, 0x33, 0xfa, 0x11, 0xd4, 0x77, 0xd3, 0xd4, 0x3b, - 0x27, 0xd1, 0xc2, 0x87, 0x61, 0x85, 0x0c, 0xa0, 0x6e, 0x8c, 0x7c, 0xa8, 0x1e, 0x78, 0x09, 0xbb, - 0x09, 0x2b, 0xb3, 0x84, 0x30, 0x9d, 0x9d, 0x4b, 0x25, 0xbd, 0xe0, 0x25, 0xdb, 0x07, 0xc9, 0xa3, - 0x58, 0xa5, 0xe7, 0xce, 0xca, 0x2c, 0xb9, 0xfa, 0x39, 0x34, 0x4d, 0x13, 0x1d, 0xfa, 0x53, 0x7e, - 0x4e, 0x5f, 0xdd, 0x76, 0xf0, 0x11, 0x5f, 0xf0, 0xd2, 0x8b, 0x32, 0xeb, 0xd9, 0xe9, 0xc6, 0xcf, - 0x56, 0xbe, 0xa8, 0x8c, 0xfe, 0x67, 0x0d, 0x5a, 0x0f, 0x79, 0xc4, 0xe9, 0xdb, 0x47, 0xd0, 0x2d, - 0xef, 0x0a, 0xcb, 0xb7, 0xb9, 0x9d, 0x32, 0x82, 0xae, 0x36, 0xc9, 0xd4, 0x8b, 0x9b, 0x6d, 0x37, - 0x07, 0x43, 0x5b, 0xb1, 0xaf, 0x7d, 0x18, 0xda, 0x6f, 0x3d, 0xc7, 0x36, 0x11, 0xf3, 0xd4, 0x60, - 0x6a, 0x1a, 0x63, 0x9a, 0xec, 0x5d, 0x80, 0x54, 0xbc, 0x72, 0x43, 0x6d, 0x17, 0xb5, 0x89, 0x69, - 0xa5, 0xe2, 0xd5, 0x3e, 0x5a, 0xc6, 0xbf, 0x96, 0x6d, 0xf6, 0x53, 0x18, 0x96, 0xb6, 0x19, 0x7a, - 0xd2, 0x6e, 0x18, 0xbb, 0x27, 0xe8, 0x7c, 0x99, 0x1d, 0x57, 0x8c, 0x49, 0x8e, 0xf6, 0x7e, 0xfc, - 0x80, 0x3c, 0x33, 0xa3, 0x3c, 0xda, 0x6f, 0x51, 0x1e, 0x17, 0xea, 0x22, 0xb8, 0x58, 0x17, 0x3d, - 0x00, 0x38, 0xe2, 0x93, 0x19, 0x8f, 0xd5, 0x81, 0x97, 0x0c, 0x3b, 0xb4, 0xf0, 0xa3, 0x62, 0xe1, - 0xed, 0x6a, 0x6d, 0x17, 0x44, 0x5a, 0x0a, 0x4a, 0xbd, 0xd0, 0x5d, 0xf2, 0xbd, 0xd8, 0x55, 0x69, - 0x16, 0xfb, 0x9e, 0xd2, 0x27, 0x95, 0x96, 0xd3, 0xf1, 0xbd, 0xf8, 0xd8, 0x80, 0x4a, 0x0a, 0xa3, - 0x57, 0x56, 0x18, 0xb7, 0x60, 0x35, 0x49, 0xc3, 0x99, 0x97, 0x9e, 0xbb, 0xa7, 0xfc, 0x9c, 0x16, - 0x43, 0x6f, 0xbd, 0x9e, 0x01, 0xff, 0x86, 0x9f, 0xef, 0x07, 0x67, 0x57, 0x7f, 0x01, 0xab, 0x0b, - 0x13, 0xf8, 0x5e, 0x72, 0xf7, 0x1f, 0xab, 0xd0, 0x3e, 0x4c, 0xb9, 0x51, 0xf2, 0xd7, 0xa1, 0x23, - 0xfd, 0x29, 0x9f, 0x79, 0x5a, 0x37, 0xe8, 0x11, 0x40, 0x83, 0x48, 0x2f, 0xcc, 0xa9, 0xb1, 0x95, - 0x6f, 0x51, 0x63, 0x03, 0xa8, 0x6a, 0xb7, 0x0b, 0x37, 0x13, 0x3e, 0x16, 0xba, 0xbb, 0x56, 0xd6, - 0xdd, 0x37, 0xa0, 0x3b, 0xf5, 0xa4, 0xeb, 0x65, 0x4a, 0xb8, 0xbe, 0x88, 0x48, 0xe8, 0x5a, 0x0e, - 0x4c, 0x3d, 0xb9, 0x9b, 0x29, 0xb1, 0x27, 0x22, 0xf6, 0x23, 0x00, 0x5f, 0x44, 0x46, 0x0d, 0x19, - 0x9f, 0xb3, 0xed, 0x8b, 0x48, 0xeb, 0x1e, 0x94, 0x4a, 0x2e, 0x55, 0x38, 0xf3, 0xcc, 0x92, 0xba, - 0xbe, 0xc8, 0x62, 0x45, 0xb6, 0xb6, 0xea, 0xac, 0xe5, 0x28, 0x47, 0xbc, 0xda, 0x43, 0x04, 0xbb, - 0x07, 0x7d, 0x5f, 0xcc, 0x12, 0x37, 0x41, 0xce, 0x92, 0x07, 0xd4, 0x5a, 0x3a, 0x2c, 0x74, 0x91, - 0xe2, 0xf0, 0x94, 0x6b, 0x9f, 0x6c, 0x07, 0x56, 0xfd, 0x28, 0x93, 0x8a, 0xa7, 0xee, 0x89, 0xe9, - 0xb2, 0x7c, 0xbe, 0xe8, 0x19, 0x12, 0xe3, 0xc7, 0x8d, 0xa0, 0x17, 0x4a, 0x57, 0x44, 0x81, 0xab, - 0x15, 0x94, 0x91, 0xb3, 0x4e, 0x28, 0x9f, 0x45, 0x81, 0x51, 0x91, 0x9a, 0x26, 0xe6, 0xaf, 0x2c, - 0x4d, 0xc7, 0xd2, 0x3c, 0xe5, 0xaf, 0x0c, 0xcd, 0x9b, 0x14, 0x5a, 0xf7, 0x8d, 0x0a, 0xed, 0x3f, - 0xad, 0x40, 0xf3, 0x50, 0x48, 0xf5, 0x70, 0x16, 0xd9, 0x4d, 0x51, 0xf9, 0xbe, 0x9b, 0x62, 0xe5, - 0xe2, 0x4d, 0x71, 0x81, 0x58, 0x56, 0x2f, 0x10, 0x4b, 0x34, 0x35, 0x65, 0x3a, 0x12, 0x27, 0xed, - 0xcb, 0xf6, 0x0b, 0x42, 0x12, 0xa9, 0x6b, 0xe8, 0x3f, 0xb9, 0x81, 0xd6, 0x62, 0x7a, 0xe9, 0x5b, - 0xa1, 0x34, 0x1a, 0x4c, 0x23, 0x43, 0x92, 0x4e, 0xe3, 0x5c, 0xb5, 0x42, 0x69, 0xa4, 0xf5, 0xf7, - 0xe0, 0x4a, 0xde, 0xd3, 0x7d, 0x15, 0xaa, 0xa9, 0xc8, 0x94, 0x3b, 0xa6, 0x03, 0xa1, 0x34, 0x47, - 0x8f, 0x4d, 0x3b, 0xd2, 0x57, 0x1a, 0xad, 0x8f, 0x8b, 0xe4, 0xeb, 0x8d, 0xb3, 0x28, 0x72, 0x15, - 0x3f, 0x53, 0x66, 0xf1, 0x87, 0x9a, 0x37, 0x86, 0x6f, 0x8f, 0xb3, 0x28, 0x3a, 0xe6, 0x67, 0x0a, - 0x0d, 0x4c, 0x6b, 0x6c, 0x1a, 0xa3, 0x3f, 0xae, 0x01, 0x3c, 0x11, 0xfe, 0xe9, 0xb1, 0x97, 0x4e, - 0xb8, 0xc2, 0x03, 0x8d, 0xd5, 0x81, 0x46, 0x47, 0x37, 0x95, 0xd6, 0x7c, 0x6c, 0x07, 0x36, 0xed, - 0xf7, 0xa3, 0xe4, 0xe2, 0xe1, 0x4a, 0x2b, 0x31, 0xb3, 0x05, 0x99, 0xc1, 0xea, 0x83, 0x3f, 0x69, - 0x30, 0xf6, 0x45, 0xc1, 0x5b, 0xec, 0xa3, 0xce, 0x13, 0xe2, 0xed, 0x45, 0xbe, 0x6d, 0xaf, 0xe8, - 0x7e, 0x7c, 0x9e, 0xb0, 0x7b, 0x70, 0x29, 0xe5, 0xe3, 0x94, 0xcb, 0xa9, 0xab, 0x64, 0xf9, 0x65, - 0xfa, 0x5c, 0xb3, 0x66, 0x90, 0xc7, 0x32, 0x7f, 0xd7, 0x3d, 0xb8, 0xa4, 0x39, 0xb5, 0x38, 0x3d, - 0xad, 0xf1, 0xd7, 0x34, 0xb2, 0x3c, 0xbb, 0x1f, 0x01, 0x05, 0xca, 0xb4, 0x16, 0xb7, 0x8e, 0x6e, - 0x44, 0xcc, 0x38, 0x89, 0x38, 0xfa, 0x80, 0x7b, 0x53, 0x3c, 0xd4, 0x3f, 0xe4, 0x63, 0xc3, 0xfc, - 0x02, 0xc0, 0x46, 0x50, 0x3b, 0x10, 0x01, 0x27, 0x56, 0xf7, 0x77, 0xfa, 0xdb, 0x14, 0x72, 0x43, - 0x4e, 0x52, 0x6c, 0x86, 0x70, 0xec, 0x23, 0xa0, 0xe1, 0xb4, 0xf8, 0x2d, 0xef, 0xae, 0x16, 0x22, - 0x49, 0x06, 0xef, 0xc1, 0xa5, 0x62, 0x26, 0xae, 0xa7, 0x5c, 0x35, 0xe5, 0xa4, 0x40, 0xf5, 0x06, - 0x5b, 0xcb, 0x27, 0xb5, 0xab, 0x8e, 0xa7, 0x1c, 0x95, 0xe9, 0x16, 0x34, 0xc5, 0xc9, 0x37, 0x2e, - 0x6e, 0x84, 0xce, 0xc5, 0x1b, 0xa1, 0x21, 0x4e, 0xbe, 0x71, 0xf8, 0x98, 0xfd, 0xa4, 0x6c, 0x7c, - 0x16, 0x58, 0xd3, 0x25, 0xd6, 0x6c, 0xe4, 0xf8, 0x12, 0x77, 0x46, 0x5f, 0x40, 0x03, 0x3f, 0xe7, - 0x59, 0xc2, 0xb6, 0xa1, 0xa9, 0x48, 0x3c, 0xa4, 0x71, 0x16, 0x36, 0x0a, 0x9b, 0x51, 0xc8, 0x8e, - 0x63, 0x89, 0x46, 0x0e, 0xac, 0xe6, 0x0a, 0xf8, 0x79, 0x1c, 0xbe, 0xc8, 0x38, 0xfb, 0x25, 0xac, - 0x25, 0x29, 0x37, 0x62, 0xef, 0x66, 0xa7, 0xe8, 0x02, 0x99, 0x1d, 0xbc, 0x61, 0xa4, 0x34, 0xef, - 0x71, 0x8a, 0x12, 0xda, 0x4f, 0xe6, 0xda, 0xa3, 0xaf, 0xe1, 0x72, 0x4e, 0x71, 0xc4, 0x7d, 0x11, - 0x07, 0x5e, 0x7a, 0x4e, 0xb6, 0x72, 0x61, 0x6c, 0xf9, 0x7d, 0xc6, 0x3e, 0xa2, 0xb1, 0xff, 0x7b, - 0x05, 0x3a, 0x8f, 0xb3, 0xd7, 0xaf, 0xcf, 0xf5, 0x5e, 0x62, 0x5d, 0xa8, 0x3c, 0xa5, 0x01, 0x56, - 0x9c, 0xca, 0x53, 0x74, 0xe7, 0x0e, 0x4f, 0x71, 0x5f, 0x93, 0x9c, 0xb7, 0x1d, 0xd3, 0xc2, 0xd3, - 0xda, 0xe1, 0xe9, 0xf1, 0x5b, 0x24, 0x5a, 0xa3, 0xf1, 0x98, 0xf1, 0x20, 0x0b, 0x23, 0x74, 0x36, - 0x8c, 0xf0, 0xe6, 0x6d, 0x3c, 0xff, 0xec, 0x8f, 0xf5, 0x54, 0x1e, 0xa7, 0x62, 0xa6, 0x99, 0x65, - 0x54, 0xc6, 0x05, 0x18, 0xf6, 0x2b, 0x58, 0x37, 0x51, 0x22, 0xa3, 0x15, 0x5c, 0x99, 0x70, 0x9f, - 0x44, 0xf7, 0x7b, 0x45, 0x96, 0x46, 0x7f, 0x55, 0x83, 0xd6, 0x97, 0x9e, 0x9c, 0xfe, 0x5a, 0x84, - 0x31, 0xbb, 0x07, 0xed, 0x6f, 0x44, 0x18, 0xeb, 0xa3, 0xaf, 0x0e, 0xfa, 0xae, 0xeb, 0xb1, 0x9e, - 0x8a, 0x80, 0x6f, 0x23, 0x0d, 0x1d, 0x7a, 0x5b, 0xdf, 0x98, 0x27, 0xa3, 0xe4, 0xd3, 0x70, 0x32, - 0x55, 0x2e, 0x02, 0x8d, 0x6e, 0xed, 0x84, 0xd2, 0x41, 0x18, 0x8d, 0xfa, 0x2e, 0xa0, 0xbd, 0x9b, - 0xba, 0x22, 0x76, 0x93, 0x53, 0x73, 0x3a, 0x6a, 0x21, 0xe4, 0x59, 0x7c, 0x78, 0x8a, 0x7b, 0x2f, - 0x94, 0xae, 0x09, 0xb2, 0x18, 0x4f, 0xb6, 0x74, 0xc8, 0xfc, 0x10, 0xfa, 0xe8, 0x65, 0xc8, 0xd3, - 0x30, 0x71, 0x93, 0x54, 0x9c, 0x58, 0xa6, 0xa0, 0xef, 0x71, 0x74, 0x1a, 0x26, 0x87, 0x08, 0x23, - 0xe3, 0x6e, 0x42, 0x37, 0xa8, 0xb6, 0xb5, 0x15, 0x05, 0x03, 0x42, 0xfe, 0x52, 0x7c, 0x26, 0xd2, - 0xbe, 0x76, 0x93, 0x8c, 0x76, 0x33, 0xe5, 0x11, 0x39, 0xd5, 0x57, 0xa0, 0x85, 0x9b, 0x81, 0x50, - 0x2d, 0x8d, 0xf2, 0x85, 0x46, 0xfd, 0x18, 0x20, 0xe2, 0x63, 0xe5, 0xa2, 0x94, 0xe9, 0xd3, 0xe8, - 0x42, 0x1c, 0x04, 0xb1, 0x7b, 0x88, 0x64, 0x1f, 0x43, 0x47, 0x73, 0x41, 0xd3, 0xc2, 0x12, 0x2d, - 0x10, 0x5a, 0x13, 0xdf, 0x86, 0x4e, 0x2c, 0x62, 0x97, 0xbf, 0x20, 0x6a, 0xb3, 0x6f, 0xe7, 0x06, - 0x8e, 0x45, 0xfc, 0xe8, 0x05, 0x12, 0xb3, 0xbb, 0x66, 0x0e, 0x3a, 0x20, 0xd0, 0x7d, 0x43, 0x40, - 0x80, 0x66, 0xa2, 0x8f, 0xc6, 0xf7, 0xed, 0x4c, 0x74, 0x8f, 0xde, 0x1b, 0x7a, 0xe8, 0xf9, 0xe8, - 0x2e, 0x37, 0xa0, 0x4b, 0xeb, 0x3e, 0xf3, 0x12, 0x57, 0x79, 0x13, 0xe3, 0x8d, 0x01, 0xc2, 0x0e, - 0xbc, 0xe4, 0xd8, 0x9b, 0x30, 0x07, 0xae, 0x2c, 0xc8, 0xdb, 0x09, 0x8a, 0xae, 0xe6, 0xda, 0xaa, - 0x0d, 0x28, 0x5c, 0x2c, 0x75, 0x9b, 0x73, 0x52, 0x47, 0x22, 0x8f, 0xdc, 0x1d, 0xfd, 0xd3, 0x15, - 0x68, 0x3d, 0x11, 0x22, 0xf9, 0x81, 0xa2, 0x57, 0x5e, 0xd2, 0x95, 0x37, 0x2f, 0x69, 0x75, 0x7e, - 0x49, 0x17, 0x58, 0x5f, 0xfb, 0xee, 0xac, 0xaf, 0x7f, 0x6f, 0xd6, 0x37, 0x7e, 0x00, 0xeb, 0x9b, - 0x8b, 0xac, 0x1f, 0x35, 0xa1, 0x7e, 0xc4, 0xd5, 0xb3, 0x64, 0xf4, 0x2f, 0x5b, 0xd0, 0x7e, 0xc8, - 0x83, 0x4c, 0x33, 0xac, 0xfc, 0xf9, 0x95, 0x37, 0x7f, 0xfe, 0xca, 0xfc, 0xe7, 0xa3, 0x21, 0xb2, - 0x12, 0x7d, 0x41, 0x6c, 0xac, 0x65, 0x05, 0x1a, 0x45, 0xbf, 0x90, 0x67, 0x13, 0x60, 0x9a, 0x63, - 0x53, 0x2e, 0xce, 0x6f, 0x97, 0x8d, 0xfa, 0x0f, 0x92, 0x8d, 0x05, 0xad, 0xb0, 0x14, 0x7a, 0xfa, - 0x56, 0xae, 0x2d, 0x6a, 0x84, 0xd6, 0x92, 0x46, 0x78, 0x02, 0xeb, 0x22, 0x76, 0x83, 0x2c, 0x89, - 0x42, 0x3c, 0xab, 0xb8, 0x9e, 0x3e, 0xa9, 0xb7, 0x6d, 0xbe, 0x25, 0x17, 0xbd, 0x67, 0xf1, 0x43, - 0x4b, 0xa4, 0xcf, 0xef, 0xce, 0x9a, 0x58, 0x04, 0xa1, 0x9a, 0x0a, 0x70, 0x69, 0xc8, 0xae, 0x92, - 0x47, 0xa8, 0x13, 0x47, 0x5d, 0x82, 0xee, 0x89, 0x88, 0x2c, 0xc5, 0x17, 0xb0, 0x5a, 0x50, 0x69, - 0x19, 0xe9, 0xbc, 0x41, 0x46, 0x7a, 0xb6, 0xa3, 0x16, 0x93, 0xbf, 0x0e, 0x2d, 0xf0, 0x09, 0xac, - 0xdb, 0xb0, 0x84, 0x71, 0x0e, 0x68, 0x05, 0xfb, 0x24, 0x41, 0x03, 0x13, 0x89, 0x20, 0xbf, 0x80, - 0x96, 0xe8, 0xe7, 0xb0, 0x51, 0x22, 0xc7, 0x73, 0x43, 0x59, 0x1b, 0x94, 0x65, 0x65, 0x2d, 0xef, - 0x8b, 0xcd, 0x27, 0x3a, 0x3c, 0xdb, 0x09, 0x78, 0x64, 0x5f, 0x34, 0x1c, 0xe8, 0x63, 0x4f, 0xc0, - 0x23, 0x93, 0x2d, 0x3a, 0x80, 0x0f, 0xf1, 0x74, 0x81, 0x78, 0xdf, 0x4b, 0x54, 0x96, 0x72, 0x37, - 0x89, 0x3c, 0x9f, 0x4f, 0x45, 0x14, 0xf0, 0xb4, 0x98, 0xdc, 0x1a, 0x4d, 0xee, 0xba, 0x88, 0x82, - 0x3d, 0x11, 0xed, 0x69, 0xca, 0xc3, 0x82, 0xd0, 0xce, 0x75, 0x17, 0xde, 0x5b, 0x1a, 0x0e, 0x0d, - 0x47, 0x31, 0x10, 0xa3, 0x81, 0xae, 0xcc, 0x0f, 0x84, 0x24, 0x76, 0x88, 0xfb, 0x70, 0x49, 0xaf, - 0x9d, 0x16, 0xee, 0x53, 0xce, 0x13, 0x37, 0xf2, 0xa4, 0x1a, 0xae, 0x6b, 0x23, 0x4d, 0x48, 0x12, - 0xe0, 0xdf, 0x70, 0x9e, 0x3c, 0xf1, 0xf4, 0x5b, 0x75, 0x17, 0xe3, 0xc7, 0x53, 0x9f, 0x39, 0xde, - 0x6e, 0xe8, 0xb7, 0x12, 0x95, 0x76, 0xe6, 0xb1, 0x73, 0x89, 0xc9, 0xbf, 0x0f, 0xd7, 0xe6, 0x86, - 0x98, 0x79, 0xe9, 0x69, 0xe1, 0xd8, 0x0e, 0x2f, 0x11, 0xdf, 0x2e, 0x97, 0xfa, 0x1f, 0x10, 0x81, - 0x1e, 0x61, 0xf4, 0xdf, 0xea, 0xd0, 0x27, 0x3b, 0xfc, 0xb7, 0x6a, 0xe3, 0x6f, 0xd5, 0xc6, 0xdf, - 0x00, 0xb5, 0x31, 0xfa, 0x07, 0x15, 0x68, 0x1e, 0xa6, 0x22, 0xc8, 0x7c, 0xf5, 0x03, 0x25, 0x7d, - 0x5e, 0x82, 0xaa, 0xdf, 0x26, 0x41, 0xb5, 0x25, 0x73, 0xfd, 0xcf, 0x2a, 0xd0, 0x36, 0x53, 0x78, - 0xb2, 0xf3, 0x03, 0x27, 0x51, 0x24, 0xaf, 0x2a, 0x17, 0x26, 0xaf, 0xbe, 0x75, 0x16, 0x28, 0x58, - 0x2f, 0x75, 0x16, 0x5f, 0x24, 0x45, 0x26, 0xab, 0xed, 0x74, 0x35, 0xf4, 0x59, 0x42, 0x09, 0xab, - 0x57, 0xd0, 0xa6, 0x93, 0x13, 0x69, 0x86, 0x4d, 0x68, 0xa4, 0x94, 0x61, 0x31, 0x13, 0x35, 0xad, - 0xb7, 0xef, 0xd3, 0x95, 0x1f, 0xe6, 0xfa, 0xfd, 0xdb, 0x15, 0xe8, 0xd1, 0x31, 0xf6, 0x71, 0x16, - 0xeb, 0x9d, 0x90, 0x87, 0xcf, 0x2a, 0xf3, 0xe1, 0xb3, 0x5a, 0x8a, 0xa7, 0x4d, 0xfd, 0x9a, 0xae, - 0x7e, 0xcd, 0x9e, 0x88, 0x1e, 0xf2, 0xb1, 0x43, 0x18, 0x64, 0x95, 0x97, 0x4e, 0xe4, 0x45, 0x79, - 0x3e, 0x84, 0xe3, 0x57, 0x25, 0x5e, 0xea, 0xcd, 0xa4, 0xcd, 0xf3, 0xe9, 0x16, 0x63, 0x50, 0xa3, - 0xfd, 0xa6, 0xd9, 0x42, 0xcf, 0x26, 0x22, 0x23, 0xc3, 0x78, 0x92, 0x2b, 0x8f, 0x16, 0xe5, 0x77, - 0x27, 0x11, 0x67, 0x0f, 0x81, 0xe9, 0x80, 0x6d, 0xca, 0x3d, 0x34, 0x41, 0x34, 0x0e, 0x69, 0x90, - 0xce, 0xce, 0xa6, 0x7e, 0x2d, 0xf1, 0xd2, 0x21, 0xf4, 0x21, 0x62, 0x9d, 0x41, 0xb8, 0x00, 0xb9, - 0x80, 0x99, 0xda, 0x0e, 0xe5, 0xa7, 0x8f, 0xef, 0xcc, 0x4c, 0x32, 0x4e, 0xc4, 0xcc, 0x5d, 0xb8, - 0x64, 0xb3, 0x27, 0xa8, 0x2e, 0x76, 0x70, 0x2f, 0xd0, 0x79, 0xd8, 0x7e, 0x63, 0xa5, 0xf4, 0x8d, - 0x1b, 0x50, 0x2f, 0xd7, 0x75, 0xe8, 0xc6, 0xe8, 0x26, 0x74, 0xc6, 0x61, 0xc4, 0x4d, 0x14, 0x12, - 0x99, 0x66, 0xe2, 0x91, 0x15, 0xaa, 0x6c, 0x30, 0xad, 0xd1, 0x6f, 0x2b, 0x70, 0x39, 0xf1, 0xd2, - 0x17, 0x19, 0x57, 0x14, 0x8b, 0xa4, 0x6c, 0x9b, 0x2b, 0xa7, 0x5e, 0x1a, 0xe0, 0xc6, 0xa1, 0x21, - 0xf4, 0xe8, 0xba, 0x7c, 0xa0, 0x8d, 0x10, 0x3d, 0x97, 0x5b, 0xb0, 0x5a, 0xea, 0xa1, 0xbc, 0xd4, - 0x46, 0x8b, 0x7a, 0xa9, 0x78, 0x45, 0x49, 0xd3, 0x23, 0x04, 0xe2, 0x81, 0xb2, 0xa0, 0xe3, 0x64, - 0x6d, 0x28, 0x0b, 0x6f, 0xa9, 0x1e, 0xc5, 0x01, 0xee, 0x9c, 0x38, 0x9b, 0xe9, 0x60, 0x8a, 0xae, - 0xfe, 0x68, 0xc6, 0xd9, 0x8c, 0xe2, 0x27, 0x1b, 0x50, 0x3f, 0x39, 0x57, 0xe4, 0xad, 0x23, 0x5c, - 0x37, 0x46, 0x7f, 0x51, 0x87, 0xf5, 0x7d, 0x9f, 0x9f, 0xf0, 0x74, 0xf2, 0xd0, 0x53, 0xde, 0xe3, - 0x30, 0xe2, 0xc7, 0x9e, 0x3c, 0xc5, 0x05, 0xa7, 0x39, 0x27, 0x9e, 0x9a, 0x1a, 0x2e, 0xb5, 0x10, - 0x70, 0xe8, 0xa9, 0x29, 0x9a, 0x02, 0x42, 0x8e, 0x45, 0x3a, 0x33, 0xb1, 0xad, 0xb6, 0x43, 0xdf, - 0xf8, 0x98, 0x20, 0x79, 0x6f, 0x19, 0xbe, 0xe6, 0xa6, 0x56, 0x85, 0x7a, 0x53, 0xe2, 0xf3, 0x7d, - 0xe8, 0xa6, 0xdc, 0x17, 0x69, 0x60, 0x02, 0xb6, 0x7a, 0x9e, 0x1d, 0x0d, 0xd3, 0xa1, 0xda, 0xdb, - 0x50, 0x64, 0x15, 0xe8, 0xf8, 0xee, 0x86, 0x36, 0xf1, 0xbd, 0x9a, 0x23, 0x70, 0xe5, 0xf7, 0x03, - 0xf6, 0x77, 0x61, 0x50, 0xd0, 0x52, 0x88, 0xdb, 0x1e, 0x2f, 0x76, 0x8a, 0x10, 0xcc, 0x05, 0x9f, - 0xb8, 0x7d, 0x68, 0x7b, 0xfd, 0x1d, 0xea, 0xa4, 0xc3, 0xf8, 0xc5, 0xf0, 0x1a, 0xca, 0x3e, 0x80, - 0x9e, 0x4c, 0xa2, 0x50, 0x19, 0x01, 0x90, 0xa6, 0xa2, 0xa5, 0x4b, 0x40, 0x1d, 0x89, 0x96, 0x17, - 0x2d, 0x61, 0xeb, 0x3b, 0x2d, 0x61, 0x7b, 0x79, 0x09, 0x7f, 0x0c, 0x03, 0x3f, 0xe5, 0x01, 0x8f, - 0x55, 0xe8, 0x45, 0xae, 0xf4, 0x45, 0x62, 0x4d, 0xdf, 0x6a, 0x01, 0x3f, 0x42, 0x30, 0xfb, 0x09, - 0x5c, 0xf6, 0x45, 0xac, 0x78, 0xac, 0x5c, 0xc9, 0x5f, 0x64, 0x3c, 0xf6, 0xb9, 0x1b, 0x67, 0xb3, - 0x13, 0x9e, 0x9a, 0x9c, 0xee, 0x25, 0x83, 0x3e, 0x32, 0xd8, 0xa7, 0x84, 0x64, 0xf7, 0x60, 0x43, - 0x2f, 0xcf, 0x42, 0x27, 0x9d, 0x45, 0x64, 0xb4, 0x52, 0xf3, 0x3d, 0xb6, 0x61, 0x7d, 0xea, 0x49, - 0x37, 0xe5, 0x32, 0x0c, 0x32, 0x2f, 0x32, 0x3b, 0xd4, 0xe4, 0x2e, 0xd6, 0xa6, 0x9e, 0x74, 0x0c, - 0xc6, 0x84, 0x87, 0x28, 0x7a, 0x3d, 0x47, 0xeb, 0x4e, 0x3d, 0x39, 0xa5, 0xe3, 0x73, 0xdb, 0x61, - 0xe9, 0x1c, 0xf5, 0x97, 0x9e, 0x9c, 0x5e, 0x7d, 0x00, 0x1b, 0x17, 0x2d, 0xc8, 0xb7, 0xa5, 0x35, - 0xda, 0xa5, 0xb4, 0x86, 0xa9, 0xeb, 0xfa, 0x1f, 0x2b, 0x70, 0xc9, 0xae, 0x37, 0x39, 0x7e, 0xb9, - 0x50, 0x5f, 0x27, 0x1b, 0x89, 0xce, 0x62, 0x7e, 0x96, 0x6e, 0x3b, 0xa0, 0x41, 0x74, 0x70, 0xde, - 0x82, 0x81, 0x21, 0x28, 0x84, 0x5f, 0xbf, 0xa5, 0x1f, 0xe4, 0x43, 0xd1, 0x16, 0xa0, 0x0f, 0x1c, - 0xf3, 0x14, 0x79, 0x14, 0x50, 0x45, 0x1e, 0x75, 0x21, 0x61, 0xa7, 0x0f, 0xb4, 0x38, 0x2b, 0x72, - 0xec, 0x0e, 0x30, 0xfe, 0x22, 0xf3, 0xa2, 0x50, 0x9d, 0xbb, 0xe3, 0x90, 0x47, 0x01, 0xe5, 0xd0, - 0x74, 0xa1, 0xce, 0xc0, 0x62, 0x1e, 0x23, 0x62, 0x3f, 0x90, 0xa5, 0x99, 0x98, 0xd4, 0x4c, 0xbe, - 0x01, 0xcc, 0x4c, 0x8e, 0x08, 0xbc, 0x1f, 0x5c, 0xbc, 0x57, 0x1a, 0x17, 0xef, 0x95, 0x8f, 0x60, - 0x75, 0x71, 0xcd, 0x75, 0xba, 0xa4, 0x2f, 0xe7, 0xd7, 0xfb, 0x22, 0x21, 0x6c, 0x5d, 0x28, 0x84, - 0x86, 0xe9, 0xff, 0x6b, 0x05, 0x36, 0x0c, 0xd3, 0xf7, 0x44, 0x94, 0xcd, 0xd0, 0xda, 0x26, 0x61, - 0x3c, 0x41, 0x83, 0x3c, 0x13, 0xda, 0x2d, 0x29, 0xa9, 0x3f, 0x98, 0x89, 0x5c, 0x17, 0x6f, 0xc1, - 0x20, 0xd4, 0x3d, 0x73, 0xbe, 0xd8, 0xd2, 0x3a, 0x03, 0x37, 0x5c, 0x41, 0x29, 0x94, 0xb1, 0x97, - 0xc8, 0xa9, 0x50, 0x86, 0x94, 0x94, 0xb8, 0xe6, 0xf9, 0x9a, 0x45, 0x11, 0x35, 0x79, 0x87, 0x77, - 0x80, 0xf9, 0x59, 0x9a, 0xe2, 0xfe, 0x28, 0x91, 0xeb, 0x84, 0xc4, 0xc0, 0x60, 0x0a, 0xea, 0x0f, - 0xa0, 0x39, 0x13, 0x85, 0x47, 0x30, 0xe7, 0xdc, 0x39, 0x8d, 0x99, 0x20, 0x09, 0xb9, 0x8a, 0x5e, - 0xcb, 0x8b, 0x2c, 0x4c, 0x79, 0x60, 0xed, 0xa0, 0x6d, 0x1b, 0x23, 0x39, 0x0d, 0x83, 0x80, 0xc7, - 0x26, 0x18, 0xde, 0x0a, 0xe5, 0x97, 0xd4, 0xa6, 0xca, 0x33, 0x3e, 0xf6, 0xb2, 0x48, 0xb9, 0x71, - 0x16, 0xd1, 0xae, 0x88, 0x4c, 0x3d, 0xd4, 0xaa, 0x41, 0x3c, 0xcd, 0x22, 0xdc, 0x11, 0x91, 0x59, - 0x52, 0xb2, 0x25, 0x28, 0x82, 0xee, 0x34, 0x8c, 0x15, 0xa9, 0x8a, 0x36, 0x2d, 0x29, 0x22, 0x50, - 0x08, 0xbf, 0x0c, 0x63, 0x35, 0xfa, 0xb3, 0x15, 0xd8, 0x34, 0x8c, 0x3f, 0x32, 0x0c, 0x30, 0xf6, - 0x91, 0x3c, 0x76, 0xcb, 0x2e, 0x93, 0xab, 0xa8, 0x3a, 0x60, 0x41, 0xfb, 0x34, 0xe1, 0x42, 0xba, - 0x56, 0x4c, 0x9d, 0x94, 0x95, 0xab, 0x3b, 0xc0, 0x96, 0xe4, 0x4a, 0x9a, 0x98, 0xd1, 0x60, 0x41, - 0xb0, 0x24, 0xfb, 0x0c, 0x36, 0x67, 0x5c, 0x79, 0xb4, 0x11, 0x22, 0xe1, 0x7b, 0xd4, 0x8b, 0xb6, - 0xbc, 0x66, 0xf7, 0x86, 0xc5, 0x3e, 0x31, 0x48, 0xdc, 0xf4, 0xf8, 0x8e, 0x99, 0x17, 0x87, 0x63, - 0x2e, 0x15, 0xd9, 0x79, 0xdd, 0x43, 0x3b, 0x1e, 0x03, 0x8b, 0x41, 0x4b, 0x4e, 0xd4, 0xe4, 0x31, - 0x8e, 0xf5, 0x22, 0x36, 0x88, 0xa6, 0x99, 0xf2, 0xb1, 0x59, 0xbb, 0x1e, 0xae, 0x55, 0x1c, 0xc6, - 0x13, 0x5d, 0x1c, 0xda, 0xd4, 0x3e, 0x9d, 0x05, 0x1e, 0x88, 0x80, 0x8f, 0xfe, 0xa4, 0x96, 0xcb, - 0xe8, 0xa1, 0x81, 0x1f, 0x29, 0x4f, 0x49, 0x76, 0x13, 0xfa, 0xf9, 0xe4, 0xb5, 0x8d, 0xd4, 0xbc, - 0xea, 0x59, 0xe8, 0x03, 0x04, 0xa2, 0xf8, 0xcd, 0xcf, 0x56, 0xd3, 0xae, 0xe8, 0x84, 0x63, 0x79, - 0xba, 0x9a, 0x1e, 0x87, 0xb5, 0xf4, 0x9a, 0xd4, 0x94, 0x6d, 0x5a, 0xa8, 0x26, 0xfb, 0xa4, 0x60, - 0x82, 0x74, 0x25, 0x8f, 0x74, 0xb5, 0x4d, 0x6d, 0x7e, 0x54, 0x79, 0x64, 0x10, 0xb8, 0x35, 0x0b, - 0xf2, 0x24, 0xcd, 0x62, 0x1e, 0x18, 0x93, 0xbe, 0x9a, 0xc3, 0x0f, 0x09, 0x8c, 0x13, 0xce, 0x35, - 0x53, 0x69, 0xe8, 0x86, 0x1e, 0x3a, 0x30, 0x9a, 0xa9, 0x18, 0x1a, 0x65, 0xb4, 0xa0, 0x37, 0x63, - 0x6b, 0x05, 0xb1, 0x9a, 0x53, 0x9b, 0xb1, 0x7f, 0x0a, 0xc3, 0x9c, 0x56, 0x7f, 0x5d, 0xf1, 0x82, - 0x96, 0x36, 0x3e, 0xb6, 0x0b, 0x7d, 0x66, 0xfe, 0x92, 0x4f, 0x61, 0x73, 0xb1, 0xa3, 0x79, 0x53, - 0x9b, 0xba, 0xad, 0xcf, 0x75, 0x2b, 0xbe, 0x24, 0x5f, 0x5f, 0xdf, 0xf3, 0xa7, 0xdc, 0x9d, 0x86, - 0xa6, 0x72, 0xb3, 0xea, 0xac, 0x59, 0xd4, 0x1e, 0x62, 0xbe, 0x0c, 0x95, 0xbc, 0x80, 0x7e, 0x16, - 0x4a, 0x69, 0xac, 0xe2, 0x3c, 0xfd, 0x41, 0x28, 0xe5, 0xe8, 0x1f, 0x03, 0x74, 0xad, 0xa7, 0x48, - 0xc5, 0x85, 0x77, 0xca, 0x4e, 0x77, 0x67, 0x67, 0x60, 0xbd, 0x67, 0x24, 0xd9, 0x55, 0x2a, 0xb5, - 0xf9, 0x0b, 0xed, 0x8c, 0xcf, 0xf9, 0x3b, 0x2b, 0xe4, 0x20, 0x14, 0xfe, 0xce, 0x2e, 0xac, 0x95, - 0x3c, 0x48, 0x57, 0x09, 0xe5, 0x45, 0xc6, 0x29, 0x2f, 0x55, 0x94, 0x94, 0x48, 0x9c, 0x55, 0x6c, - 0x68, 0xdf, 0xe2, 0x18, 0xa9, 0xd1, 0xd9, 0xf7, 0x45, 0x64, 0xab, 0xd9, 0x16, 0x9c, 0x7d, 0xc4, - 0x50, 0xae, 0x3c, 0xe5, 0x78, 0x76, 0x94, 0x2f, 0x22, 0xb3, 0x83, 0xda, 0x1a, 0x72, 0xf4, 0x22, - 0xca, 0x27, 0x48, 0xce, 0x74, 0x83, 0xce, 0x11, 0x34, 0x41, 0x3a, 0x53, 0x7d, 0x02, 0x1d, 0x91, - 0x86, 0x93, 0x90, 0x52, 0x5f, 0xda, 0xc1, 0x59, 0x7c, 0x09, 0x68, 0x82, 0x3d, 0x7c, 0xd5, 0x08, - 0x1a, 0xc6, 0xfc, 0x2f, 0xe7, 0xcf, 0x0d, 0x06, 0x1d, 0x22, 0xa9, 0xd2, 0xd0, 0x57, 0x38, 0x1d, - 0xbd, 0x23, 0x75, 0x61, 0x54, 0x4f, 0x83, 0x8f, 0x5e, 0x44, 0x94, 0xfd, 0xbb, 0x05, 0xab, 0x3e, - 0x99, 0x0b, 0xbd, 0xa1, 0x22, 0x1e, 0xd3, 0x9a, 0xd6, 0x9d, 0x9e, 0x06, 0xe3, 0xfc, 0x9e, 0xf0, - 0xd8, 0x14, 0x61, 0x79, 0x51, 0x84, 0x27, 0x46, 0xe1, 0x05, 0x26, 0x63, 0xde, 0xb5, 0xc0, 0x27, - 0xc2, 0x0b, 0xd8, 0xcf, 0xe0, 0x2a, 0xe2, 0x5c, 0x3e, 0x4b, 0xd4, 0x39, 0xda, 0x37, 0x9e, 0x86, - 0xbe, 0xeb, 0x49, 0xca, 0xa0, 0x9b, 0xc4, 0xf9, 0x26, 0x52, 0x3c, 0x42, 0x82, 0xa7, 0x1a, 0xbf, - 0x2b, 0xbf, 0xe6, 0xa9, 0x60, 0x5f, 0x53, 0x06, 0xf0, 0x22, 0xf7, 0xdd, 0x1e, 0xf5, 0xdf, 0x2f, - 0xd6, 0xea, 0x0d, 0x94, 0x54, 0xa1, 0x82, 0x08, 0xc7, 0x3a, 0x7d, 0xd4, 0x9f, 0xfd, 0x06, 0x98, - 0x35, 0x70, 0x24, 0xf9, 0xca, 0x93, 0xa7, 0x92, 0xa2, 0x00, 0x9d, 0x9d, 0x1f, 0xbd, 0xd5, 0x47, - 0x75, 0xac, 0x65, 0x44, 0x20, 0x02, 0x24, 0xfb, 0x23, 0xd8, 0xc8, 0x07, 0x33, 0xbe, 0x0c, 0x0d, - 0xa7, 0x83, 0x04, 0xd7, 0x97, 0x87, 0x9b, 0x73, 0x81, 0x1c, 0x3b, 0x13, 0x0d, 0xd6, 0x43, 0xfe, - 0x0a, 0x56, 0xed, 0x90, 0x9a, 0xeb, 0x72, 0x38, 0xa0, 0xd1, 0xde, 0x5b, 0x1a, 0x6d, 0xce, 0xb6, - 0xe7, 0xf6, 0x59, 0x43, 0xf1, 0x43, 0x73, 0x4b, 0x6e, 0xad, 0xcc, 0x70, 0x8d, 0x64, 0xe4, 0xc6, - 0xd2, 0x48, 0x0b, 0xc6, 0xca, 0xb1, 0x53, 0xb0, 0x70, 0x76, 0x1f, 0x2e, 0xd9, 0xc1, 0x04, 0x25, - 0x6c, 0xdd, 0x50, 0x50, 0x2e, 0x97, 0x69, 0x17, 0xcb, 0x20, 0x75, 0x32, 0x77, 0x5f, 0x38, 0x7c, - 0xcc, 0x7e, 0x01, 0xd7, 0x6c, 0x17, 0x6d, 0x85, 0xe9, 0x44, 0x9a, 0x7f, 0xd4, 0x3a, 0xd9, 0xae, - 0xa1, 0x21, 0xd1, 0x76, 0x19, 0x4f, 0xa0, 0x76, 0xfa, 0x5b, 0x30, 0xa0, 0xd2, 0x54, 0x5c, 0x56, - 0x91, 0x06, 0x61, 0xec, 0x45, 0xc3, 0x0d, 0x92, 0x9a, 0x3e, 0xc2, 0x1d, 0xf1, 0xea, 0x99, 0x86, - 0xb2, 0x63, 0xd8, 0xb4, 0x2f, 0xca, 0xd5, 0x8c, 0x44, 0x53, 0x42, 0x61, 0xc7, 0x8b, 0x18, 0x37, - 0x67, 0x70, 0x1c, 0xbb, 0x84, 0xf3, 0x66, 0xe8, 0x21, 0x5c, 0x5f, 0x58, 0xda, 0x99, 0x77, 0xe6, - 0xce, 0xf8, 0x4c, 0xa4, 0xe7, 0xc6, 0x80, 0x6c, 0x92, 0x02, 0xbb, 0x36, 0xb7, 0x88, 0x07, 0xde, - 0xd9, 0x01, 0xd1, 0x68, 0x73, 0xf2, 0x4b, 0x78, 0x77, 0x61, 0x14, 0x5d, 0xe9, 0xc9, 0x63, 0xef, - 0x24, 0xe2, 0xc1, 0xf0, 0x32, 0x7d, 0xd1, 0x95, 0xb9, 0x21, 0x8e, 0x90, 0xe2, 0x91, 0x26, 0x30, - 0x0e, 0xdd, 0x09, 0xb4, 0x29, 0x0c, 0x41, 0xda, 0x30, 0xaf, 0xba, 0xad, 0xbc, 0xbd, 0xea, 0xf6, - 0x13, 0xe8, 0x1a, 0x6f, 0xff, 0x4d, 0x65, 0xbc, 0x1d, 0x8d, 0xc7, 0x67, 0x39, 0xba, 0x03, 0x6d, - 0x72, 0xf5, 0xe9, 0x1d, 0xd7, 0xa1, 0x43, 0xd5, 0x5e, 0xee, 0x49, 0x24, 0xfc, 0x53, 0xeb, 0x9c, - 0x13, 0xe8, 0x01, 0x42, 0x46, 0x00, 0xad, 0xe7, 0x71, 0x28, 0xe2, 0xdd, 0x28, 0x1a, 0xfd, 0x65, - 0x03, 0xda, 0xe8, 0x13, 0x50, 0xdc, 0x04, 0x8f, 0x55, 0xb4, 0x70, 0x94, 0x4b, 0x9d, 0x79, 0x89, - 0xa9, 0x2b, 0xee, 0x20, 0x10, 0xa9, 0x0e, 0xbc, 0x64, 0x21, 0xd5, 0xba, 0xb2, 0x90, 0x6a, 0x7d, - 0x5f, 0xdf, 0x7d, 0xd1, 0xf5, 0x66, 0xdc, 0x16, 0xaa, 0xd2, 0x00, 0x0f, 0x34, 0x08, 0x7d, 0x15, - 0x22, 0xf1, 0x22, 0xf2, 0x6f, 0xf0, 0xf4, 0x14, 0x49, 0x93, 0x95, 0x25, 0xb9, 0xd9, 0x35, 0x88, - 0x23, 0xae, 0xf5, 0x71, 0x29, 0x58, 0x56, 0x5f, 0x0c, 0x96, 0xdd, 0x06, 0xf0, 0x45, 0x1c, 0x90, - 0x0b, 0xb5, 0x90, 0x0d, 0xd3, 0x29, 0xd1, 0x02, 0xfb, 0x1d, 0x42, 0xb3, 0x1f, 0xc1, 0x20, 0xa7, - 0x40, 0x0f, 0xc9, 0x8f, 0xf3, 0xf3, 0xa7, 0xa1, 0x72, 0xf8, 0x78, 0x2f, 0x56, 0x8b, 0x31, 0xdc, - 0xf6, 0x52, 0x0c, 0xf7, 0x0d, 0xc9, 0x73, 0xf8, 0xde, 0xd7, 0x32, 0xae, 0x40, 0x8b, 0xaa, 0x74, - 0x82, 0x2c, 0x31, 0xba, 0xba, 0x19, 0x4a, 0x8a, 0xb5, 0xbf, 0x29, 0x4e, 0xdc, 0xfd, 0xff, 0x15, - 0x27, 0xee, 0x7d, 0xb7, 0x38, 0x71, 0xff, 0xbb, 0xc5, 0x89, 0x17, 0xe2, 0xaa, 0xab, 0x8b, 0xe9, - 0x98, 0x37, 0x26, 0x3f, 0x06, 0x6f, 0x4c, 0x7e, 0x7c, 0x4b, 0xe6, 0x62, 0xed, 0xad, 0x99, 0x8b, - 0xef, 0x90, 0x3a, 0x61, 0xdf, 0x96, 0x3a, 0xb9, 0x05, 0xab, 0x2a, 0xf5, 0xfc, 0x53, 0x7d, 0x12, - 0x39, 0xe5, 0xe7, 0xd2, 0xa4, 0x6a, 0x7a, 0x04, 0xc6, 0x73, 0xc8, 0x6f, 0xf8, 0xb9, 0x1c, 0x3d, - 0x07, 0xa0, 0x23, 0x1a, 0x7d, 0xda, 0x9b, 0x64, 0xa3, 0xf2, 0xbd, 0x0b, 0x2b, 0xfe, 0x4f, 0x05, - 0xe0, 0xc8, 0x9b, 0x25, 0x3a, 0xc6, 0xc9, 0xfe, 0x10, 0x3a, 0x92, 0x5a, 0xe5, 0x0c, 0x77, 0xc9, - 0x90, 0x15, 0xa4, 0xe6, 0x51, 0xdf, 0x2e, 0x90, 0xf9, 0x33, 0x89, 0xb5, 0x1e, 0x21, 0x2f, 0x62, - 0xab, 0x5b, 0x02, 0x8a, 0x7d, 0xdd, 0x84, 0xbe, 0x21, 0x48, 0x78, 0xea, 0xf3, 0x58, 0x57, 0xc6, - 0x56, 0x9c, 0x9e, 0x86, 0x1e, 0x6a, 0x20, 0xbb, 0x9f, 0x93, 0x59, 0x93, 0xb1, 0x9c, 0xa6, 0x31, - 0x5d, 0x8c, 0xcd, 0x18, 0xed, 0xd8, 0x4f, 0xa1, 0x89, 0xb4, 0xa0, 0x86, 0xef, 0x1b, 0xbc, 0xc3, - 0x3a, 0xd0, 0x34, 0xa3, 0x0e, 0x2a, 0xac, 0x07, 0x6d, 0xba, 0xce, 0x42, 0xb8, 0x95, 0xd1, 0x1f, - 0xaf, 0x41, 0x67, 0x3f, 0x96, 0x2a, 0xcd, 0xb4, 0x08, 0x17, 0x97, 0x36, 0xea, 0x74, 0x69, 0xc3, - 0x54, 0x50, 0xea, 0xcf, 0xa0, 0x0a, 0xca, 0x4f, 0xa0, 0x69, 0xee, 0x07, 0x99, 0xc0, 0xf7, 0x85, - 0x97, 0x8b, 0x2c, 0x0d, 0xdb, 0x86, 0x56, 0x60, 0x2e, 0x2e, 0x99, 0x34, 0x7e, 0xe9, 0x36, 0x91, - 0xbd, 0xd2, 0xe4, 0xe4, 0x34, 0xec, 0x7d, 0xa8, 0x7a, 0x93, 0x89, 0x39, 0xf5, 0xae, 0x16, 0xa4, - 0xe4, 0xc4, 0x38, 0x88, 0x63, 0x77, 0xa1, 0x4d, 0xea, 0x93, 0x2a, 0x59, 0x1a, 0x8b, 0x63, 0xda, - 0x32, 0x19, 0xad, 0x51, 0x29, 0x66, 0x7e, 0x17, 0xda, 0x91, 0x10, 0x89, 0xee, 0xd0, 0x5c, 0xec, + 0x92, 0xd8, 0x34, 0xfb, 0x3b, 0xba, 0x9b, 0x6c, 0x26, 0x29, 0xaa, 0x25, 0xcd, 0x8c, 0x34, 0x3d, + 0x23, 0x0d, 0x9f, 0x46, 0x43, 0x49, 0x9c, 0x99, 0xf7, 0x66, 0xdf, 0xdb, 0xb7, 0x6f, 0x29, 0x4a, + 0x9a, 0xe1, 0x7b, 0xa2, 0xc4, 0x2d, 0x52, 0x1e, 0x60, 0x00, 0xbb, 0x50, 0xac, 0xca, 0xee, 0xae, + 0x61, 0x75, 0x65, 0xa9, 0x32, 0x4b, 0x22, 0x75, 0xb1, 0x0f, 0x3e, 0xf9, 0xe2, 0xe3, 0xee, 0x71, + 0x01, 0xfb, 0xb0, 0xb6, 0x01, 0xfb, 0x60, 0x3c, 0xff, 0x04, 0x63, 0x61, 0xfb, 0xb0, 0xf0, 0xc1, + 0x47, 0xc3, 0x78, 0x7b, 0x34, 0x60, 0x18, 0x06, 0x6c, 0x2c, 0x60, 0x18, 0x30, 0x22, 0x32, 0xb3, + 0xaa, 0xba, 0x9b, 0x94, 0x66, 0xc6, 0xc6, 0x5e, 0x76, 0x4f, 0x5d, 0x19, 0x11, 0x99, 0x95, 0x15, + 0x19, 0x19, 0x11, 0x19, 0x11, 0xd9, 0xb0, 0x9c, 0x84, 0x09, 0x8f, 0xc2, 0x98, 0x6f, 0x25, 0xa9, + 0x50, 0x82, 0xb5, 0x6c, 0xfb, 0xea, 0xa7, 0xe3, 0x50, 0x4d, 0xb2, 0xe3, 0x2d, 0x5f, 0x4c, 0xef, + 0x8e, 0xc5, 0x58, 0xdc, 0x25, 0x82, 0xe3, 0x6c, 0x44, 0x2d, 0x6a, 0xd0, 0x93, 0xee, 0x78, 0x15, + 0x22, 0xe1, 0x9f, 0xd8, 0xe7, 0x24, 0xf2, 0x62, 0xf3, 0xbc, 0xa2, 0xc2, 0x29, 0x97, 0xca, 0x9b, + 0x26, 0x06, 0xd0, 0x56, 0xa7, 0x06, 0x37, 0xfc, 0xf7, 0x55, 0x68, 0xee, 0x73, 0x29, 0xbd, 0x31, + 0x67, 0x43, 0xa8, 0xca, 0x30, 0x18, 0x54, 0x6e, 0x54, 0x36, 0x97, 0xb7, 0xfb, 0x5b, 0xf9, 0xb4, + 0x0e, 0x95, 0xa7, 0x32, 0xe9, 0x20, 0x12, 0x69, 0xfc, 0x69, 0x30, 0x58, 0x9a, 0xa7, 0xd9, 0xe7, + 0x6a, 0x22, 0x02, 0x07, 0x91, 0xac, 0x0f, 0x55, 0x9e, 0xa6, 0x83, 0xea, 0x8d, 0xca, 0x66, 0xd7, + 0xc1, 0x47, 0xc6, 0xa0, 0x16, 0x78, 0xca, 0x1b, 0xd4, 0x08, 0x44, 0xcf, 0xec, 0x23, 0x58, 0x4e, + 0x52, 0xe1, 0xbb, 0x61, 0x3c, 0x12, 0x2e, 0x61, 0xeb, 0x84, 0xed, 0x22, 0x74, 0x2f, 0x1e, 0x89, + 0x87, 0x48, 0x35, 0x80, 0xa6, 0x17, 0x7b, 0xd1, 0x99, 0xe4, 0x83, 0x06, 0xa1, 0x6d, 0x93, 0x2d, + 0xc3, 0x52, 0x18, 0x0c, 0x9a, 0x37, 0x2a, 0x9b, 0x35, 0x67, 0x29, 0x0c, 0xf0, 0x1d, 0x59, 0x16, + 0x06, 0x83, 0x96, 0x7e, 0x07, 0x3e, 0xb3, 0x21, 0x74, 0x63, 0xce, 0x83, 0xa7, 0x42, 0x39, 0x3c, + 0x89, 0xce, 0x06, 0xed, 0x1b, 0x95, 0xcd, 0x96, 0x33, 0x03, 0x63, 0x57, 0xa1, 0x15, 0xf0, 0xe3, + 0x6c, 0xbc, 0x2f, 0xc7, 0x03, 0xb8, 0x51, 0xd9, 0x6c, 0x3b, 0x79, 0x9b, 0x1d, 0xc1, 0xe5, 0x94, + 0xbf, 0xc8, 0xb8, 0x54, 0x3c, 0x70, 0x15, 0xf7, 0xd2, 0x40, 0xbc, 0x8a, 0xdd, 0xa9, 0x08, 0xf8, + 0xa0, 0x43, 0x1c, 0x78, 0xb7, 0xcc, 0xa5, 0x94, 0x7b, 0xd3, 0x23, 0x43, 0xb4, 0x2f, 0x02, 0xee, + 0x5c, 0xca, 0x3b, 0x97, 0xc1, 0xcc, 0x81, 0x0d, 0xcf, 0xf7, 0x79, 0xb2, 0x38, 0x68, 0xf7, 0x7b, + 0x0c, 0xba, 0x6e, 0xfb, 0x96, 0xa1, 0x3f, 0xaf, 0xfd, 0xc9, 0x9f, 0x5e, 0x7f, 0x67, 0xf8, 0x1c, + 0xda, 0xbb, 0x22, 0x8e, 0xb9, 0xaf, 0x44, 0xca, 0xae, 0x43, 0xc7, 0x8e, 0xe3, 0x9a, 0x65, 0xad, + 0x3b, 0x60, 0x41, 0x7b, 0x01, 0xfb, 0x18, 0x56, 0x7c, 0x4b, 0xed, 0x86, 0x71, 0xc0, 0x4f, 0x69, + 0x5d, 0xeb, 0xce, 0x72, 0x0e, 0xde, 0x43, 0xe8, 0xf0, 0xdf, 0x56, 0xa1, 0x79, 0x38, 0xc9, 0x46, + 0xa3, 0x88, 0xb3, 0x8f, 0xa0, 0x67, 0x1e, 0x77, 0x45, 0xb4, 0x17, 0x9c, 0x9a, 0x71, 0x67, 0x81, + 0xec, 0x06, 0x74, 0x0c, 0xe0, 0xe8, 0x2c, 0xe1, 0x66, 0xd8, 0x32, 0x68, 0x76, 0x9c, 0xfd, 0x30, + 0x26, 0x71, 0xa9, 0x3a, 0xb3, 0xc0, 0x39, 0x2a, 0xef, 0x94, 0x24, 0x68, 0x96, 0xca, 0xa3, 0xb7, + 0xed, 0x44, 0xe1, 0x4b, 0xee, 0xf0, 0xf1, 0x6e, 0xac, 0x48, 0x8e, 0xea, 0x4e, 0x19, 0xc4, 0xb6, + 0xe1, 0x92, 0xd4, 0x5d, 0xdc, 0xd4, 0x8b, 0xc7, 0x5c, 0xba, 0x59, 0x18, 0xab, 0x9f, 0x7e, 0x3e, + 0x68, 0xdc, 0xa8, 0x6e, 0xd6, 0x9c, 0x35, 0x83, 0x74, 0x08, 0xf7, 0x9c, 0x50, 0xec, 0x1e, 0xac, + 0xcf, 0xf5, 0xd1, 0x5d, 0x9a, 0x37, 0xaa, 0x9b, 0x55, 0x87, 0xcd, 0x74, 0xd9, 0xa3, 0x1e, 0x8f, + 0x60, 0x35, 0xcd, 0x62, 0xdc, 0x6d, 0x8f, 0xc3, 0x48, 0xf1, 0xf4, 0x30, 0xe1, 0x3e, 0xc9, 0x63, + 0x67, 0xfb, 0xf2, 0x16, 0x6d, 0x48, 0x67, 0x1e, 0xed, 0x2c, 0xf6, 0x60, 0x77, 0x72, 0xe6, 0x3d, + 0x3a, 0x4d, 0x52, 0x12, 0xda, 0xce, 0x36, 0xe8, 0x01, 0x10, 0xe2, 0x94, 0xd1, 0xec, 0x36, 0xac, + 0x06, 0xa9, 0x17, 0xc6, 0xae, 0x17, 0x45, 0xee, 0x71, 0xe6, 0x9f, 0x70, 0x25, 0x49, 0x90, 0x5b, + 0xce, 0x0a, 0x21, 0x76, 0xa2, 0xe8, 0x81, 0x06, 0x0f, 0xff, 0x6a, 0x09, 0x5a, 0x0f, 0x43, 0x99, + 0x78, 0xca, 0x9f, 0xb0, 0xcb, 0xd0, 0x1c, 0x65, 0xb1, 0x5f, 0xc8, 0x46, 0x03, 0x9b, 0x7b, 0x01, + 0xfb, 0x7d, 0x58, 0x89, 0x84, 0xef, 0x45, 0x6e, 0x2e, 0x06, 0x83, 0xa5, 0x1b, 0xd5, 0xcd, 0xce, + 0xf6, 0x5a, 0x21, 0x98, 0xb9, 0x98, 0x39, 0xcb, 0x44, 0x5b, 0x88, 0xdd, 0x2f, 0xa1, 0x9f, 0xf2, + 0xa9, 0x50, 0xbc, 0xd4, 0xbd, 0x4a, 0xdd, 0x59, 0xd1, 0xfd, 0x9b, 0xd4, 0x4b, 0x9e, 0xa2, 0x34, + 0xaf, 0x68, 0xda, 0xa2, 0xfb, 0xfd, 0xd2, 0x4a, 0xf1, 0xb1, 0x1b, 0x06, 0xa7, 0x2e, 0xbd, 0x60, + 0x50, 0xbb, 0x51, 0xdd, 0xac, 0x17, 0x6c, 0xe7, 0xe3, 0xbd, 0xe0, 0xf4, 0x09, 0x62, 0xd8, 0x67, + 0xb0, 0x31, 0xdf, 0x45, 0x8f, 0x3a, 0xa8, 0x53, 0x9f, 0xb5, 0x99, 0x3e, 0x0e, 0xa1, 0xd8, 0x07, + 0xd0, 0xb5, 0x9d, 0x14, 0x8a, 0x68, 0x43, 0x0b, 0x8d, 0x2c, 0x89, 0xe8, 0x65, 0x68, 0x86, 0xd2, + 0x95, 0x61, 0x7c, 0x42, 0x6a, 0xa6, 0xe5, 0x34, 0x42, 0x79, 0x18, 0xc6, 0x27, 0xec, 0x0a, 0xb4, + 0x52, 0xee, 0x6b, 0x4c, 0x8b, 0x30, 0xcd, 0x94, 0xfb, 0x84, 0xba, 0x0c, 0xf8, 0xe8, 0xfa, 0x8a, + 0x1b, 0x65, 0xd3, 0x48, 0xb9, 0xbf, 0xab, 0xf8, 0x50, 0x42, 0x7d, 0x9f, 0xa7, 0x63, 0x8e, 0xfa, + 0x06, 0x3b, 0x1e, 0xfa, 0x5e, 0x4c, 0x7c, 0x6f, 0x39, 0x79, 0x1b, 0xb5, 0x5d, 0xe2, 0xa5, 0x2a, + 0xf4, 0x22, 0xda, 0x32, 0x2d, 0xc7, 0x36, 0xd9, 0x35, 0x68, 0x4b, 0xe5, 0xa5, 0x0a, 0xbf, 0x8e, + 0xb6, 0x4a, 0xdd, 0x69, 0x11, 0x00, 0x77, 0xdb, 0x65, 0x68, 0xf2, 0x38, 0x20, 0x54, 0x4d, 0xaf, + 0x24, 0x8f, 0x83, 0xbd, 0xe0, 0x74, 0xf8, 0xaf, 0x2b, 0xd0, 0xdb, 0xcf, 0x22, 0x15, 0xee, 0xa4, + 0xe3, 0x8c, 0x4f, 0x63, 0x85, 0x5a, 0xf2, 0x61, 0x28, 0x95, 0x79, 0x33, 0x3d, 0xb3, 0x4d, 0x68, + 0x7f, 0x95, 0x8a, 0x2c, 0x21, 0x69, 0xd3, 0x2b, 0x5d, 0x96, 0xb6, 0x02, 0x89, 0x92, 0xf9, 0x2c, + 0x0d, 0x78, 0xfa, 0xe0, 0x8c, 0x68, 0xab, 0x0b, 0xb4, 0x65, 0x34, 0x7b, 0x17, 0xda, 0x87, 0x3c, + 0xf1, 0x52, 0x0f, 0x45, 0xa0, 0x46, 0xaa, 0xb5, 0x00, 0xe0, 0xb7, 0x12, 0xf1, 0x5e, 0x60, 0x36, + 0xac, 0x6d, 0x0e, 0xff, 0x49, 0x05, 0xda, 0x3b, 0xe3, 0x71, 0xca, 0xc7, 0x9e, 0x22, 0x3d, 0x2f, + 0x12, 0x9a, 0x6f, 0xd5, 0x59, 0x12, 0x09, 0xd9, 0x12, 0xfc, 0x02, 0xcd, 0x20, 0x7a, 0x66, 0xef, + 0x43, 0x8d, 0x9f, 0x3f, 0x21, 0x82, 0xb3, 0x0d, 0x68, 0xf8, 0x22, 0x1e, 0x85, 0x63, 0x63, 0x81, + 0x4c, 0x8b, 0xfd, 0x1c, 0x3a, 0xfa, 0x49, 0xcb, 0x40, 0x9d, 0xd4, 0xef, 0x15, 0xdd, 0x3d, 0x9f, + 0xc1, 0x2e, 0x51, 0xa0, 0x44, 0x38, 0xe0, 0xe7, 0xcf, 0xc3, 0x7f, 0x5e, 0x85, 0x3a, 0x71, 0x06, + 0xd7, 0x06, 0x2d, 0x8a, 0xcb, 0x5f, 0x7a, 0x91, 0x5d, 0x52, 0x04, 0x3c, 0x7a, 0xe9, 0x45, 0xec, + 0x06, 0xd4, 0x71, 0x0a, 0xf2, 0x1c, 0xc6, 0x6a, 0x04, 0xbb, 0x05, 0x75, 0x7c, 0xbb, 0x9c, 0x9d, + 0x3d, 0xbe, 0xe3, 0x41, 0xed, 0xcf, 0xff, 0xf3, 0xf5, 0x77, 0x1c, 0x8d, 0x66, 0x1f, 0x43, 0xcd, + 0x1b, 0x8f, 0x25, 0x6d, 0x84, 0x99, 0xbd, 0x98, 0xcf, 0xd4, 0x21, 0x02, 0xf6, 0x05, 0xb4, 0xf5, + 0xa2, 0x23, 0x75, 0x9d, 0xa8, 0x2f, 0x97, 0x2c, 0x75, 0x59, 0x1e, 0x9c, 0x82, 0x12, 0x97, 0x2b, + 0x94, 0x46, 0xb3, 0xd0, 0x76, 0x68, 0x39, 0x05, 0x00, 0x4d, 0x69, 0x92, 0xf2, 0x9d, 0x28, 0x12, + 0xfe, 0x61, 0xf8, 0x9a, 0x1b, 0xc3, 0x3b, 0x03, 0x63, 0xb7, 0x60, 0xf9, 0x40, 0xcb, 0xab, 0xc3, + 0x65, 0x16, 0x29, 0x69, 0x8c, 0xf1, 0x1c, 0x94, 0x6d, 0x01, 0x9b, 0x81, 0x1c, 0xd1, 0xe7, 0xb7, + 0x6f, 0x54, 0x37, 0x7b, 0xce, 0x39, 0x18, 0xf6, 0x21, 0xf4, 0xc6, 0xc8, 0xe9, 0x30, 0x1e, 0xbb, + 0xa3, 0xc8, 0x43, 0x3b, 0x5d, 0x45, 0x3b, 0x6e, 0x81, 0x8f, 0x23, 0x6f, 0x4c, 0x3b, 0x24, 0x09, + 0xa3, 0xc8, 0x9d, 0xf2, 0x29, 0x59, 0xe7, 0xaa, 0xd3, 0x22, 0xc0, 0x3e, 0x9f, 0x0e, 0xff, 0x45, + 0x0d, 0x1a, 0x7b, 0xb1, 0xe4, 0xa9, 0xc2, 0xfd, 0xe7, 0x8d, 0x46, 0xdc, 0x57, 0x5c, 0xeb, 0xbd, + 0x9a, 0x93, 0xb7, 0x91, 0x05, 0x47, 0xe2, 0x9b, 0x34, 0x54, 0xfc, 0xf0, 0x33, 0x23, 0x60, 0x05, + 0x00, 0x35, 0xad, 0x17, 0x04, 0xae, 0xa5, 0x76, 0x53, 0xf1, 0x4a, 0xd2, 0x5e, 0x6c, 0x39, 0x2b, + 0x5e, 0x10, 0xec, 0x18, 0xb8, 0x23, 0x5e, 0x49, 0xf6, 0x01, 0x54, 0x53, 0x3e, 0x22, 0x71, 0xeb, + 0x6c, 0xaf, 0xe8, 0x25, 0x7d, 0x76, 0xfc, 0x1d, 0xf7, 0x95, 0xc3, 0x47, 0x0e, 0xe2, 0xd8, 0x3a, + 0xd4, 0x3d, 0xa5, 0x52, 0xbd, 0x44, 0x6d, 0x47, 0x37, 0xd8, 0x16, 0xac, 0xd1, 0x9e, 0x57, 0xa1, + 0x88, 0x5d, 0xe5, 0x1d, 0x47, 0x68, 0xbc, 0xa5, 0xb1, 0x53, 0xab, 0x39, 0xea, 0x08, 0x31, 0x7b, + 0x81, 0x44, 0xcb, 0x36, 0x4f, 0x1f, 0x7b, 0x53, 0x2e, 0xc9, 0x4c, 0xb5, 0x9d, 0xb5, 0xd9, 0x1e, + 0x4f, 0x11, 0x85, 0xfc, 0x2c, 0xfa, 0xa0, 0xd6, 0x68, 0xd1, 0x06, 0xec, 0xe6, 0x40, 0x54, 0x2a, + 0x97, 0xa0, 0x11, 0x4a, 0x97, 0xc7, 0x81, 0x51, 0x64, 0xf5, 0x50, 0x3e, 0x8a, 0x03, 0xf6, 0x09, + 0xb4, 0xf5, 0x5b, 0x02, 0x3e, 0x22, 0x33, 0xd3, 0xd9, 0x5e, 0x36, 0x12, 0x8b, 0xe0, 0x87, 0x7c, + 0xe4, 0xb4, 0x94, 0x79, 0x42, 0x17, 0x44, 0x09, 0x97, 0x9f, 0x2a, 0x9e, 0xc6, 0x5e, 0x44, 0xab, + 0xd2, 0x72, 0x40, 0x89, 0x47, 0x06, 0xc2, 0xbe, 0x80, 0xcb, 0x16, 0xeb, 0x4a, 0x35, 0x55, 0x6e, + 0x16, 0x87, 0xa7, 0x6e, 0xec, 0xc5, 0x82, 0x7c, 0xa1, 0xaa, 0xb3, 0x6e, 0xd1, 0x87, 0x6a, 0xaa, + 0x9e, 0xc7, 0xe1, 0xe9, 0x53, 0x2f, 0x16, 0x6c, 0x13, 0xfa, 0x79, 0x37, 0xf5, 0x9a, 0x3e, 0x78, + 0xd0, 0x23, 0x05, 0xb3, 0x6c, 0xe1, 0x47, 0xaf, 0xf1, 0x5b, 0xd1, 0x36, 0x94, 0x29, 0xc5, 0x68, + 0x24, 0xb9, 0x72, 0x25, 0xf7, 0x07, 0xcb, 0xf4, 0xcd, 0x6b, 0x05, 0xfd, 0x33, 0xc2, 0x1d, 0x72, + 0x7f, 0xf8, 0xdb, 0x0a, 0x74, 0x68, 0x5f, 0x3c, 0x4f, 0x02, 0x54, 0x41, 0x1f, 0x42, 0x6f, 0x76, + 0xd1, 0xb5, 0xdc, 0x74, 0xbd, 0xf2, 0x8a, 0x6f, 0x40, 0x63, 0xc7, 0x47, 0xe6, 0x91, 0xe0, 0xf4, + 0x1c, 0xd3, 0x62, 0x3f, 0x83, 0x95, 0x8c, 0x86, 0x71, 0x7d, 0x75, 0xea, 0x46, 0xa8, 0xba, 0xf4, + 0x46, 0x37, 0x52, 0xa1, 0xdf, 0xb1, 0xab, 0x4e, 0x9d, 0x5e, 0x66, 0x1f, 0x9f, 0xa0, 0x52, 0xbb, + 0x07, 0xeb, 0x29, 0x47, 0x89, 0x71, 0x5f, 0xf3, 0x54, 0xb8, 0x8a, 0x4f, 0x13, 0x91, 0x92, 0x21, + 0x44, 0x2e, 0x32, 0x8d, 0xfb, 0x96, 0xa7, 0xe2, 0xc8, 0x60, 0x86, 0xef, 0x41, 0x7d, 0x27, 0x4d, + 0xbd, 0x33, 0x12, 0x2d, 0x7c, 0x18, 0x54, 0xc8, 0x00, 0xea, 0xc6, 0xd0, 0x87, 0xea, 0xbe, 0x97, + 0xb0, 0x9b, 0xb0, 0x34, 0x4d, 0x08, 0xd3, 0xd9, 0xbe, 0x54, 0xd2, 0x0b, 0x5e, 0xb2, 0xb5, 0x9f, + 0x3c, 0x8a, 0x55, 0x7a, 0xe6, 0x2c, 0x4d, 0x93, 0xab, 0x5f, 0x40, 0xd3, 0x34, 0xd1, 0xa1, 0x3f, + 0xe1, 0x67, 0xf4, 0xd5, 0x6d, 0x07, 0x1f, 0xf1, 0x05, 0x2f, 0xbd, 0x28, 0xb3, 0x9e, 0x9d, 0x6e, + 0xfc, 0x7c, 0xe9, 0xcb, 0xca, 0xf0, 0x7f, 0xd6, 0xa0, 0xf5, 0x90, 0x47, 0x9c, 0xbe, 0x7d, 0x08, + 0xdd, 0xf2, 0xae, 0xb0, 0x7c, 0x9b, 0xd9, 0x29, 0x43, 0xe8, 0x6a, 0x93, 0x4c, 0xbd, 0xb8, 0xd9, + 0x76, 0x33, 0x30, 0xb4, 0x15, 0x7b, 0xda, 0x87, 0xa1, 0xfd, 0xd6, 0x73, 0x6c, 0x13, 0x31, 0x4f, + 0x0d, 0xa6, 0xa6, 0x31, 0xa6, 0xc9, 0xde, 0x05, 0x48, 0xc5, 0x2b, 0x37, 0xd4, 0x76, 0x51, 0x9b, + 0x98, 0x56, 0x2a, 0x5e, 0xed, 0xa1, 0x65, 0xfc, 0x6b, 0xd9, 0x66, 0x3f, 0x83, 0x41, 0x69, 0x9b, + 0xa1, 0x27, 0xed, 0x86, 0xb1, 0x7b, 0x8c, 0xce, 0x97, 0xd9, 0x71, 0xc5, 0x98, 0xe4, 0x68, 0xef, + 0xc5, 0x0f, 0xc8, 0x33, 0x33, 0xca, 0xa3, 0xfd, 0x06, 0xe5, 0x71, 0xae, 0x2e, 0x82, 0xf3, 0x75, + 0xd1, 0x03, 0x80, 0x43, 0x3e, 0x9e, 0xf2, 0x58, 0xed, 0x7b, 0xc9, 0xa0, 0x43, 0x0b, 0x3f, 0x2c, + 0x16, 0xde, 0xae, 0xd6, 0x56, 0x41, 0xa4, 0xa5, 0xa0, 0xd4, 0x0b, 0xdd, 0x25, 0xdf, 0x8b, 0x5d, + 0x95, 0x66, 0xb1, 0xef, 0x29, 0x7d, 0x52, 0x69, 0x39, 0x1d, 0xdf, 0x8b, 0x8f, 0x0c, 0xa8, 0xa4, + 0x30, 0x7a, 0x65, 0x85, 0x71, 0x0b, 0x56, 0x92, 0x34, 0x9c, 0x7a, 0xe9, 0x99, 0x7b, 0xc2, 0xcf, + 0x68, 0x31, 0xf4, 0xd6, 0xeb, 0x19, 0xf0, 0x6f, 0xf8, 0xd9, 0x5e, 0x70, 0x7a, 0xf5, 0x97, 0xb0, + 0x32, 0x37, 0x81, 0x1f, 0x24, 0x77, 0xff, 0xb1, 0x0a, 0xed, 0x83, 0x94, 0x1b, 0x25, 0x7f, 0x1d, + 0x3a, 0xd2, 0x9f, 0xf0, 0xa9, 0xa7, 0x75, 0x83, 0x1e, 0x01, 0x34, 0x88, 0xf4, 0xc2, 0x8c, 0x1a, + 0x5b, 0x7a, 0x8b, 0x1a, 0xeb, 0x43, 0x55, 0xbb, 0x5d, 0xb8, 0x99, 0xf0, 0xb1, 0xd0, 0xdd, 0xb5, + 0xb2, 0xee, 0xbe, 0x01, 0xdd, 0x89, 0x27, 0x5d, 0x2f, 0x53, 0xc2, 0xf5, 0x45, 0x44, 0x42, 0xd7, + 0x72, 0x60, 0xe2, 0xc9, 0x9d, 0x4c, 0x89, 0x5d, 0x11, 0xb1, 0xf7, 0x00, 0x7c, 0x11, 0x19, 0x35, + 0x64, 0x7c, 0xce, 0xb6, 0x2f, 0x22, 0xad, 0x7b, 0x50, 0x2a, 0xb9, 0x54, 0xe1, 0xd4, 0x33, 0x4b, + 0xea, 0xfa, 0x22, 0x8b, 0x15, 0xd9, 0xda, 0xaa, 0xb3, 0x9a, 0xa3, 0x1c, 0xf1, 0x6a, 0x17, 0x11, + 0xec, 0x1e, 0x2c, 0xfb, 0x62, 0x9a, 0xb8, 0x09, 0x72, 0x96, 0x3c, 0xa0, 0xd6, 0xc2, 0x61, 0xa1, + 0x8b, 0x14, 0x07, 0x27, 0x5c, 0xfb, 0x64, 0xdb, 0xb0, 0xe2, 0x47, 0x99, 0x54, 0x3c, 0x75, 0x8f, + 0x4d, 0x97, 0xc5, 0xf3, 0x45, 0xcf, 0x90, 0x18, 0x3f, 0x6e, 0x08, 0xbd, 0x50, 0xba, 0x22, 0x0a, + 0x5c, 0xad, 0xa0, 0x8c, 0x9c, 0x75, 0x42, 0xf9, 0x2c, 0x0a, 0x8c, 0x8a, 0xd4, 0x34, 0x31, 0x7f, + 0x65, 0x69, 0x3a, 0x96, 0xe6, 0x29, 0x7f, 0x65, 0x68, 0x2e, 0x52, 0x68, 0xdd, 0x0b, 0x15, 0xda, + 0x7f, 0x5a, 0x82, 0xe6, 0x81, 0x90, 0xea, 0xe1, 0x34, 0xb2, 0x9b, 0xa2, 0xf2, 0x43, 0x37, 0xc5, + 0xd2, 0xf9, 0x9b, 0xe2, 0x1c, 0xb1, 0xac, 0x9e, 0x23, 0x96, 0x68, 0x6a, 0xca, 0x74, 0x24, 0x4e, + 0xda, 0x97, 0x5d, 0x2e, 0x08, 0x49, 0xa4, 0xae, 0xa1, 0xff, 0xe4, 0x06, 0x5a, 0x8b, 0xe9, 0xa5, + 0x6f, 0x85, 0xd2, 0x68, 0x30, 0x8d, 0x0c, 0x49, 0x3a, 0x8d, 0x73, 0xd5, 0x0a, 0xa5, 0x91, 0xd6, + 0xdf, 0x83, 0x2b, 0x79, 0x4f, 0xf7, 0x55, 0xa8, 0x26, 0x22, 0x53, 0xee, 0x88, 0x0e, 0x84, 0xd2, + 0x1c, 0x3d, 0x36, 0xec, 0x48, 0xdf, 0x68, 0xb4, 0x3e, 0x2e, 0x92, 0xaf, 0x37, 0xca, 0xa2, 0xc8, + 0x55, 0xfc, 0x54, 0x99, 0xc5, 0x1f, 0x68, 0xde, 0x18, 0xbe, 0x3d, 0xce, 0xa2, 0xe8, 0x88, 0x9f, + 0x2a, 0x34, 0x30, 0xad, 0x91, 0x69, 0x0c, 0xff, 0xb8, 0x06, 0xf0, 0x44, 0xf8, 0x27, 0x47, 0x5e, + 0x3a, 0xe6, 0x0a, 0x0f, 0x34, 0x56, 0x07, 0x1a, 0x1d, 0xdd, 0x54, 0x5a, 0xf3, 0xb1, 0x6d, 0xd8, + 0xb0, 0xdf, 0x8f, 0x92, 0x8b, 0x87, 0x2b, 0xad, 0xc4, 0xcc, 0x16, 0x64, 0x06, 0xab, 0x0f, 0xfe, + 0xa4, 0xc1, 0xd8, 0x97, 0x05, 0x6f, 0xb1, 0x8f, 0x3a, 0x4b, 0x88, 0xb7, 0xe7, 0xf9, 0xb6, 0xbd, + 0xa2, 0xfb, 0xd1, 0x59, 0xc2, 0xee, 0xc1, 0xa5, 0x94, 0x8f, 0x52, 0x2e, 0x27, 0xae, 0x92, 0xe5, + 0x97, 0xe9, 0x73, 0xcd, 0xaa, 0x41, 0x1e, 0xc9, 0xfc, 0x5d, 0xf7, 0xe0, 0x92, 0xe6, 0xd4, 0xfc, + 0xf4, 0xb4, 0xc6, 0x5f, 0xd5, 0xc8, 0xf2, 0xec, 0xde, 0x03, 0x0a, 0x94, 0x69, 0x2d, 0x6e, 0x1d, + 0xdd, 0x88, 0x98, 0x71, 0x1c, 0x71, 0xf4, 0x01, 0x77, 0x27, 0x78, 0xa8, 0x7f, 0xc8, 0x47, 0x86, + 0xf9, 0x05, 0x80, 0x0d, 0xa1, 0xb6, 0x2f, 0x02, 0x4e, 0xac, 0x5e, 0xde, 0x5e, 0xde, 0xa2, 0x90, + 0x1b, 0x72, 0x92, 0x62, 0x33, 0x84, 0x63, 0x1f, 0x03, 0x0d, 0xa7, 0xc5, 0x6f, 0x71, 0x77, 0xb5, + 0x10, 0x49, 0x32, 0x78, 0x0f, 0x2e, 0x15, 0x33, 0x71, 0x3d, 0xe5, 0xaa, 0x09, 0x27, 0x05, 0xaa, + 0x37, 0xd8, 0x6a, 0x3e, 0xa9, 0x1d, 0x75, 0x34, 0xe1, 0xa8, 0x4c, 0x37, 0xa1, 0x29, 0x8e, 0xbf, + 0x73, 0x71, 0x23, 0x74, 0xce, 0xdf, 0x08, 0x0d, 0x71, 0xfc, 0x9d, 0xc3, 0x47, 0xec, 0xa7, 0x65, + 0xe3, 0x33, 0xc7, 0x9a, 0x2e, 0xb1, 0x66, 0x3d, 0xc7, 0x97, 0xb8, 0x33, 0xfc, 0x12, 0x1a, 0xf8, + 0x39, 0xcf, 0x12, 0xb6, 0x05, 0x4d, 0x45, 0xe2, 0x21, 0x8d, 0xb3, 0xb0, 0x5e, 0xd8, 0x8c, 0x42, + 0x76, 0x1c, 0x4b, 0x34, 0x74, 0x60, 0x25, 0x57, 0xc0, 0xcf, 0xe3, 0xf0, 0x45, 0xc6, 0xd9, 0xaf, + 0x60, 0x35, 0x49, 0xb9, 0x11, 0x7b, 0x37, 0x3b, 0x41, 0x17, 0xc8, 0xec, 0xe0, 0x75, 0x23, 0xa5, + 0x79, 0x8f, 0x13, 0x94, 0xd0, 0xe5, 0x64, 0xa6, 0x3d, 0xfc, 0x16, 0x2e, 0xe7, 0x14, 0x87, 0xdc, + 0x17, 0x71, 0xe0, 0xa5, 0x67, 0x64, 0x2b, 0xe7, 0xc6, 0x96, 0x3f, 0x64, 0xec, 0x43, 0x1a, 0xfb, + 0xbf, 0x57, 0xa0, 0xf3, 0x38, 0x7b, 0xfd, 0xfa, 0x4c, 0xef, 0x25, 0xd6, 0x85, 0xca, 0x53, 0x1a, + 0x60, 0xc9, 0xa9, 0x3c, 0x45, 0x77, 0xee, 0xe0, 0x04, 0xf7, 0x35, 0xc9, 0x79, 0xdb, 0x31, 0x2d, + 0x3c, 0xad, 0x1d, 0x9c, 0x1c, 0xbd, 0x41, 0xa2, 0x35, 0x1a, 0x8f, 0x19, 0x0f, 0xb2, 0x30, 0x42, + 0x67, 0xc3, 0x08, 0x6f, 0xde, 0xc6, 0xf3, 0xcf, 0xde, 0x48, 0x4f, 0xe5, 0x71, 0x2a, 0xa6, 0x9a, + 0x59, 0x46, 0x65, 0x9c, 0x83, 0x61, 0x5f, 0xc1, 0x9a, 0x89, 0x12, 0x19, 0xad, 0xe0, 0xca, 0x84, + 0xfb, 0x24, 0xba, 0x3f, 0x28, 0xb2, 0x34, 0xfc, 0xab, 0x1a, 0xb4, 0xbe, 0xf6, 0xe4, 0xe4, 0xd7, + 0x22, 0x8c, 0xd9, 0x3d, 0x68, 0x7f, 0x27, 0xc2, 0x58, 0x1f, 0x7d, 0x75, 0xd0, 0x77, 0x4d, 0x8f, + 0xf5, 0x54, 0x04, 0x7c, 0x0b, 0x69, 0xe8, 0xd0, 0xdb, 0xfa, 0xce, 0x3c, 0x19, 0x25, 0x9f, 0x86, + 0xe3, 0x89, 0x72, 0x11, 0x68, 0x74, 0x6b, 0x27, 0x94, 0x0e, 0xc2, 0x68, 0xd4, 0x77, 0x01, 0xed, + 0xdd, 0xc4, 0x15, 0xb1, 0x9b, 0x9c, 0x98, 0xd3, 0x51, 0x0b, 0x21, 0xcf, 0xe2, 0x83, 0x13, 0xdc, + 0x7b, 0xa1, 0x74, 0x4d, 0x90, 0xc5, 0x78, 0xb2, 0xa5, 0x43, 0xe6, 0x47, 0xb0, 0x8c, 0x5e, 0x86, + 0x3c, 0x09, 0x13, 0x37, 0x49, 0xc5, 0xb1, 0x65, 0x0a, 0xfa, 0x1e, 0x87, 0x27, 0x61, 0x72, 0x80, + 0x30, 0x32, 0xee, 0x26, 0x74, 0x83, 0x6a, 0x5b, 0x5b, 0x51, 0x30, 0x20, 0xe4, 0x2f, 0xc5, 0x67, + 0x22, 0xed, 0x6b, 0x37, 0xc9, 0x68, 0x37, 0x53, 0x1e, 0x91, 0x53, 0x7d, 0x05, 0x5a, 0xb8, 0x19, + 0x08, 0xd5, 0xd2, 0x28, 0x5f, 0x68, 0xd4, 0x4f, 0x00, 0x22, 0x3e, 0x52, 0x2e, 0x4a, 0x99, 0x3e, + 0x8d, 0xce, 0xc5, 0x41, 0x10, 0xbb, 0x8b, 0x48, 0xf6, 0x09, 0x74, 0x34, 0x17, 0x34, 0x2d, 0x2c, + 0xd0, 0x02, 0xa1, 0x35, 0xf1, 0x6d, 0xe8, 0xc4, 0x22, 0x76, 0xf9, 0x0b, 0xa2, 0x36, 0xfb, 0x76, + 0x66, 0xe0, 0x58, 0xc4, 0x8f, 0x5e, 0x20, 0x31, 0xbb, 0x6b, 0xe6, 0xa0, 0x03, 0x02, 0xdd, 0x0b, + 0x02, 0x02, 0x34, 0x13, 0x7d, 0x34, 0xbe, 0x6f, 0x67, 0xa2, 0x7b, 0xf4, 0x2e, 0xe8, 0xa1, 0xe7, + 0xa3, 0xbb, 0xdc, 0x80, 0x2e, 0xad, 0xfb, 0xd4, 0x4b, 0x5c, 0xe5, 0x8d, 0x8d, 0x37, 0x06, 0x08, + 0xdb, 0xf7, 0x92, 0x23, 0x6f, 0xcc, 0x1c, 0xb8, 0x32, 0x27, 0x6f, 0xc7, 0x28, 0xba, 0x9a, 0x6b, + 0x2b, 0x36, 0xa0, 0x70, 0xbe, 0xd4, 0x6d, 0xcc, 0x48, 0x1d, 0x89, 0x3c, 0x72, 0x77, 0xf8, 0x4f, + 0x97, 0xa0, 0xf5, 0x44, 0x88, 0xe4, 0x47, 0x8a, 0x5e, 0x79, 0x49, 0x97, 0x2e, 0x5e, 0xd2, 0xea, + 0xec, 0x92, 0xce, 0xb1, 0xbe, 0xf6, 0xfd, 0x59, 0x5f, 0xff, 0xc1, 0xac, 0x6f, 0xfc, 0x08, 0xd6, + 0x37, 0xe7, 0x59, 0x3f, 0x6c, 0x42, 0xfd, 0x90, 0xab, 0x67, 0xc9, 0xf0, 0x5f, 0xb5, 0xa0, 0xfd, + 0x90, 0x07, 0x99, 0x66, 0x58, 0xf9, 0xf3, 0x2b, 0x17, 0x7f, 0xfe, 0xd2, 0xec, 0xe7, 0xa3, 0x21, + 0xb2, 0x12, 0x7d, 0x4e, 0x6c, 0xac, 0x65, 0x05, 0x1a, 0x45, 0xbf, 0x90, 0x67, 0x13, 0x60, 0x9a, + 0x61, 0x53, 0x2e, 0xce, 0x6f, 0x96, 0x8d, 0xfa, 0x8f, 0x92, 0x8d, 0x39, 0xad, 0xb0, 0x10, 0x7a, + 0x7a, 0x2b, 0xd7, 0xe6, 0x35, 0x42, 0x6b, 0x41, 0x23, 0x3c, 0x81, 0x35, 0x11, 0xbb, 0x41, 0x96, + 0x44, 0x21, 0x9e, 0x55, 0x5c, 0x4f, 0x9f, 0xd4, 0xdb, 0x36, 0xdf, 0x92, 0x8b, 0xde, 0xb3, 0xf8, + 0xa1, 0x25, 0xd2, 0xe7, 0x77, 0x67, 0x55, 0xcc, 0x83, 0x50, 0x4d, 0x05, 0xb8, 0x34, 0x64, 0x57, + 0xc9, 0x23, 0xd4, 0x89, 0xa3, 0x2e, 0x41, 0x77, 0x45, 0x44, 0x96, 0xe2, 0x4b, 0x58, 0x29, 0xa8, + 0xb4, 0x8c, 0x74, 0x2e, 0x90, 0x91, 0x9e, 0xed, 0xa8, 0xc5, 0xe4, 0xaf, 0x43, 0x0b, 0x7c, 0x0a, + 0x6b, 0x36, 0x2c, 0x61, 0x9c, 0x03, 0x5a, 0xc1, 0x65, 0x92, 0xa0, 0xbe, 0x89, 0x44, 0x90, 0x5f, + 0x40, 0x4b, 0xf4, 0x0b, 0x58, 0x2f, 0x91, 0xe3, 0xb9, 0xa1, 0xac, 0x0d, 0xca, 0xb2, 0xb2, 0x9a, + 0xf7, 0xc5, 0xe6, 0x13, 0x1d, 0x9e, 0xed, 0x04, 0x3c, 0xb2, 0x2f, 0x1a, 0xf4, 0xf5, 0xb1, 0x27, + 0xe0, 0x91, 0xc9, 0x16, 0xed, 0xc3, 0x47, 0x78, 0xba, 0x40, 0xbc, 0xef, 0x25, 0x2a, 0x4b, 0xb9, + 0x9b, 0x44, 0x9e, 0xcf, 0x27, 0x22, 0x0a, 0x78, 0x5a, 0x4c, 0x6e, 0x95, 0x26, 0x77, 0x5d, 0x44, + 0xc1, 0xae, 0x88, 0x76, 0x35, 0xe5, 0x41, 0x41, 0x68, 0xe7, 0xba, 0x03, 0xef, 0x2f, 0x0c, 0x87, + 0x86, 0xa3, 0x18, 0x88, 0xd1, 0x40, 0x57, 0x66, 0x07, 0x42, 0x12, 0x3b, 0xc4, 0x7d, 0xb8, 0xa4, + 0xd7, 0x4e, 0x0b, 0xf7, 0x09, 0xe7, 0x89, 0x1b, 0x79, 0x52, 0x0d, 0xd6, 0xb4, 0x91, 0x26, 0x24, + 0x09, 0xf0, 0x6f, 0x38, 0x4f, 0x9e, 0x78, 0xfa, 0xad, 0xba, 0x8b, 0xf1, 0xe3, 0xa9, 0xcf, 0x0c, + 0x6f, 0xd7, 0xf5, 0x5b, 0x89, 0x4a, 0x3b, 0xf3, 0xd8, 0xb9, 0xc4, 0xe4, 0xdf, 0x87, 0x6b, 0x33, + 0x43, 0x4c, 0xbd, 0xf4, 0xa4, 0x70, 0x6c, 0x07, 0x97, 0x88, 0x6f, 0x97, 0x4b, 0xfd, 0xf7, 0x89, + 0x40, 0x8f, 0x30, 0xfc, 0x6f, 0x75, 0x58, 0x26, 0x3b, 0xfc, 0xb7, 0x6a, 0xe3, 0x6f, 0xd5, 0xc6, + 0xdf, 0x00, 0xb5, 0x31, 0xfc, 0x07, 0x15, 0x68, 0x1e, 0xa4, 0x22, 0xc8, 0x7c, 0xf5, 0x23, 0x25, + 0x7d, 0x56, 0x82, 0xaa, 0x6f, 0x93, 0xa0, 0xda, 0x82, 0xb9, 0xfe, 0x67, 0x15, 0x68, 0x9b, 0x29, + 0x3c, 0xd9, 0xfe, 0x91, 0x93, 0x28, 0x92, 0x57, 0x95, 0x73, 0x93, 0x57, 0x6f, 0x9d, 0x05, 0x0a, + 0xd6, 0x4b, 0x9d, 0xc5, 0x17, 0x49, 0x91, 0xc9, 0x6a, 0x3b, 0x5d, 0x0d, 0x7d, 0x96, 0x50, 0xc2, + 0xea, 0x15, 0xb4, 0xe9, 0xe4, 0x44, 0x9a, 0x61, 0x03, 0x1a, 0x29, 0x65, 0x58, 0xcc, 0x44, 0x4d, + 0xeb, 0xcd, 0xfb, 0x74, 0xe9, 0xc7, 0xb9, 0x7e, 0xff, 0x6e, 0x09, 0x7a, 0x74, 0x8c, 0x7d, 0x9c, + 0xc5, 0x7a, 0x27, 0xe4, 0xe1, 0xb3, 0xca, 0x6c, 0xf8, 0xac, 0x96, 0xe2, 0x69, 0x53, 0xbf, 0xa6, + 0xab, 0x5f, 0xb3, 0x2b, 0xa2, 0x87, 0x7c, 0xe4, 0x10, 0x06, 0x59, 0xe5, 0xa5, 0x63, 0x79, 0x5e, + 0x9e, 0x0f, 0xe1, 0xf8, 0x55, 0x89, 0x97, 0x7a, 0x53, 0x69, 0xf3, 0x7c, 0xba, 0xc5, 0x18, 0xd4, + 0x68, 0xbf, 0x69, 0xb6, 0xd0, 0xb3, 0x89, 0xc8, 0xc8, 0x30, 0x1e, 0xe7, 0xca, 0xa3, 0x45, 0xf9, + 0xdd, 0x71, 0xc4, 0xd9, 0x43, 0x60, 0x3a, 0x60, 0x9b, 0x72, 0x0f, 0x4d, 0x10, 0x8d, 0x43, 0x1a, + 0xa4, 0xb3, 0xbd, 0xa1, 0x5f, 0x4b, 0xbc, 0x74, 0x08, 0x7d, 0x80, 0x58, 0xa7, 0x1f, 0xce, 0x41, + 0xce, 0x61, 0xa6, 0xb6, 0x43, 0xf9, 0xe9, 0xe3, 0x7b, 0x33, 0x93, 0x8c, 0x13, 0x31, 0x73, 0x07, + 0x2e, 0xd9, 0xec, 0x09, 0xaa, 0x8b, 0x6d, 0xdc, 0x0b, 0x74, 0x1e, 0xb6, 0xdf, 0x58, 0x29, 0x7d, + 0xe3, 0x3a, 0xd4, 0xcb, 0x75, 0x1d, 0xba, 0x31, 0xbc, 0x09, 0x9d, 0x51, 0x18, 0x71, 0x13, 0x85, + 0x44, 0xa6, 0x99, 0x78, 0x64, 0x85, 0x2a, 0x1b, 0x4c, 0x6b, 0xf8, 0xdb, 0x0a, 0x5c, 0x4e, 0xbc, + 0xf4, 0x45, 0xc6, 0x15, 0xc5, 0x22, 0x29, 0xdb, 0xe6, 0xca, 0x89, 0x97, 0x06, 0xb8, 0x71, 0x68, + 0x08, 0x3d, 0xba, 0x2e, 0x1f, 0x68, 0x23, 0x44, 0xcf, 0xe5, 0x16, 0xac, 0x94, 0x7a, 0x28, 0x2f, + 0xb5, 0xd1, 0xa2, 0x5e, 0x2a, 0x5e, 0x51, 0xd2, 0xf4, 0x10, 0x81, 0x78, 0xa0, 0x2c, 0xe8, 0x38, + 0x59, 0x1b, 0xca, 0xc2, 0x5b, 0xaa, 0x47, 0x71, 0x80, 0x3b, 0x27, 0xce, 0xa6, 0x3a, 0x98, 0xa2, + 0xab, 0x3f, 0x9a, 0x71, 0x36, 0xa5, 0xf8, 0xc9, 0x3a, 0xd4, 0x8f, 0xcf, 0x14, 0x79, 0xeb, 0x08, + 0xd7, 0x8d, 0xe1, 0x5f, 0xd4, 0x61, 0x6d, 0xcf, 0xe7, 0xc7, 0x3c, 0x1d, 0x3f, 0xf4, 0x94, 0xf7, + 0x38, 0x8c, 0xf8, 0x91, 0x27, 0x4f, 0x70, 0xc1, 0x69, 0xce, 0x89, 0xa7, 0x26, 0x86, 0x4b, 0x2d, + 0x04, 0x1c, 0x78, 0x6a, 0x82, 0xa6, 0x80, 0x90, 0x23, 0x91, 0x4e, 0x4d, 0x6c, 0xab, 0xed, 0xd0, + 0x37, 0x3e, 0x26, 0x48, 0xde, 0x5b, 0x86, 0xaf, 0xb9, 0xa9, 0x55, 0xa1, 0xde, 0x94, 0xf8, 0xfc, + 0x00, 0xba, 0x29, 0xf7, 0x45, 0x1a, 0x98, 0x80, 0xad, 0x9e, 0x67, 0x47, 0xc3, 0x74, 0xa8, 0xf6, + 0x36, 0x14, 0x59, 0x05, 0x3a, 0xbe, 0xbb, 0xa1, 0x4d, 0x7c, 0xaf, 0xe4, 0x08, 0x5c, 0xf9, 0xbd, + 0x80, 0xfd, 0x5d, 0xe8, 0x17, 0xb4, 0x14, 0xe2, 0xb6, 0xc7, 0x8b, 0xed, 0x22, 0x04, 0x73, 0xce, + 0x27, 0x6e, 0x1d, 0xd8, 0x5e, 0x7f, 0x87, 0x3a, 0xe9, 0x30, 0x7e, 0x31, 0xbc, 0x86, 0xb2, 0x0f, + 0xa1, 0x27, 0x93, 0x28, 0x54, 0x46, 0x00, 0xa4, 0xa9, 0x68, 0xe9, 0x12, 0x50, 0x47, 0xa2, 0xe5, + 0x79, 0x4b, 0xd8, 0xfa, 0x5e, 0x4b, 0xd8, 0x5e, 0x5c, 0xc2, 0x9f, 0x40, 0xdf, 0x4f, 0x79, 0xc0, + 0x63, 0x15, 0x7a, 0x91, 0x2b, 0x7d, 0x91, 0x58, 0xd3, 0xb7, 0x52, 0xc0, 0x0f, 0x11, 0xcc, 0x7e, + 0x0a, 0x97, 0x7d, 0x11, 0x2b, 0x1e, 0x2b, 0x57, 0xf2, 0x17, 0x19, 0x8f, 0x7d, 0xee, 0xc6, 0xd9, + 0xf4, 0x98, 0xa7, 0x26, 0xa7, 0x7b, 0xc9, 0xa0, 0x0f, 0x0d, 0xf6, 0x29, 0x21, 0xd9, 0x3d, 0x58, + 0xd7, 0xcb, 0x33, 0xd7, 0x49, 0x67, 0x11, 0x19, 0xad, 0xd4, 0x6c, 0x8f, 0x2d, 0x58, 0x9b, 0x78, + 0xd2, 0x4d, 0xb9, 0x0c, 0x83, 0xcc, 0x8b, 0xcc, 0x0e, 0x35, 0xb9, 0x8b, 0xd5, 0x89, 0x27, 0x1d, + 0x83, 0x31, 0xe1, 0x21, 0x8a, 0x5e, 0xcf, 0xd0, 0xba, 0x13, 0x4f, 0x4e, 0xe8, 0xf8, 0xdc, 0x76, + 0x58, 0x3a, 0x43, 0xfd, 0xb5, 0x27, 0x27, 0x57, 0x1f, 0xc0, 0xfa, 0x79, 0x0b, 0xf2, 0xb6, 0xb4, + 0x46, 0xbb, 0x94, 0xd6, 0x30, 0x75, 0x5d, 0xff, 0x63, 0x09, 0x2e, 0xd9, 0xf5, 0x26, 0xc7, 0x2f, + 0x17, 0xea, 0xeb, 0x64, 0x23, 0xd1, 0x59, 0xcc, 0xcf, 0xd2, 0x6d, 0x07, 0x34, 0x88, 0x0e, 0xce, + 0x9b, 0xd0, 0x37, 0x04, 0x85, 0xf0, 0xeb, 0xb7, 0x2c, 0x07, 0xf9, 0x50, 0xb4, 0x05, 0xe8, 0x03, + 0x47, 0x3c, 0x45, 0x1e, 0x05, 0x54, 0x91, 0x47, 0x5d, 0x48, 0xd8, 0xe9, 0x03, 0x2d, 0xce, 0x8a, + 0x1c, 0xbb, 0x03, 0x8c, 0xbf, 0xc8, 0xbc, 0x28, 0x54, 0x67, 0xee, 0x28, 0xe4, 0x51, 0x40, 0x39, + 0x34, 0x5d, 0xa8, 0xd3, 0xb7, 0x98, 0xc7, 0x88, 0xd8, 0x0b, 0x64, 0x69, 0x26, 0x26, 0x35, 0x93, + 0x6f, 0x00, 0x33, 0x93, 0x43, 0x02, 0xef, 0x05, 0xe7, 0xef, 0x95, 0xc6, 0xf9, 0x7b, 0xe5, 0x63, + 0x58, 0x99, 0x5f, 0x73, 0x9d, 0x2e, 0x59, 0x96, 0xb3, 0xeb, 0x7d, 0x9e, 0x10, 0xb6, 0xce, 0x15, + 0x42, 0xc3, 0xf4, 0xff, 0xb5, 0x04, 0xeb, 0x86, 0xe9, 0xbb, 0x22, 0xca, 0xa6, 0x68, 0x6d, 0x93, + 0x30, 0x1e, 0xa3, 0x41, 0x9e, 0x0a, 0xed, 0x96, 0x94, 0xd4, 0x1f, 0x4c, 0x45, 0xae, 0x8b, 0x37, + 0xa1, 0x1f, 0xea, 0x9e, 0x39, 0x5f, 0x6c, 0x69, 0x9d, 0x81, 0x1b, 0xae, 0xa0, 0x14, 0xca, 0xd8, + 0x4b, 0xe4, 0x44, 0x28, 0x43, 0x4a, 0x4a, 0x5c, 0xf3, 0x7c, 0xd5, 0xa2, 0x88, 0x9a, 0xbc, 0xc3, + 0x3b, 0xc0, 0xfc, 0x2c, 0x4d, 0x71, 0x7f, 0x94, 0xc8, 0x75, 0x42, 0xa2, 0x6f, 0x30, 0x05, 0xf5, + 0x87, 0xd0, 0x9c, 0x8a, 0xc2, 0x23, 0x98, 0x71, 0xee, 0x9c, 0xc6, 0x54, 0x90, 0x84, 0x5c, 0x45, + 0xaf, 0xe5, 0x45, 0x16, 0xa6, 0x3c, 0xb0, 0x76, 0xd0, 0xb6, 0x8d, 0x91, 0x9c, 0x84, 0x41, 0xc0, + 0x63, 0x13, 0x0c, 0x6f, 0x85, 0xf2, 0x6b, 0x6a, 0x53, 0xe5, 0x19, 0x1f, 0x79, 0x59, 0xa4, 0xdc, + 0x38, 0x8b, 0x68, 0x57, 0x44, 0xa6, 0x1e, 0x6a, 0xc5, 0x20, 0x9e, 0x66, 0x11, 0xee, 0x88, 0xc8, + 0x2c, 0x29, 0xd9, 0x12, 0x14, 0x41, 0x77, 0x12, 0xc6, 0x8a, 0x54, 0x45, 0x9b, 0x96, 0x14, 0x11, + 0x28, 0x84, 0x5f, 0x87, 0xb1, 0x1a, 0xfe, 0xd9, 0x12, 0x6c, 0x18, 0xc6, 0x1f, 0x1a, 0x06, 0x18, + 0xfb, 0x48, 0x1e, 0xbb, 0x65, 0x97, 0xc9, 0x55, 0x54, 0x1d, 0xb0, 0xa0, 0x3d, 0x9a, 0x70, 0x21, + 0x5d, 0x4b, 0xa6, 0x4e, 0xca, 0xca, 0xd5, 0x1d, 0x60, 0x0b, 0x72, 0x25, 0x4d, 0xcc, 0xa8, 0x3f, + 0x27, 0x58, 0x92, 0x7d, 0x0e, 0x1b, 0x53, 0xae, 0x3c, 0xda, 0x08, 0x91, 0xf0, 0x3d, 0xea, 0x45, + 0x5b, 0x5e, 0xb3, 0x7b, 0xdd, 0x62, 0x9f, 0x18, 0x24, 0x6e, 0x7a, 0x7c, 0xc7, 0xd4, 0x8b, 0xc3, + 0x11, 0x97, 0x8a, 0xec, 0xbc, 0xee, 0xa1, 0x1d, 0x8f, 0xbe, 0xc5, 0xa0, 0x25, 0x27, 0x6a, 0xf2, + 0x18, 0x47, 0x7a, 0x11, 0x1b, 0x44, 0xd3, 0x4c, 0xf9, 0xc8, 0xac, 0x5d, 0x0f, 0xd7, 0x2a, 0x0e, + 0xe3, 0xb1, 0x2e, 0x0e, 0x6d, 0x6a, 0x9f, 0xce, 0x02, 0xf7, 0x45, 0xc0, 0x87, 0x7f, 0x52, 0xcb, + 0x65, 0xf4, 0xc0, 0xc0, 0x0f, 0x95, 0xa7, 0x24, 0xbb, 0x09, 0xcb, 0xf9, 0xe4, 0xb5, 0x8d, 0xd4, + 0xbc, 0xea, 0x59, 0xe8, 0x03, 0x04, 0xa2, 0xf8, 0xcd, 0xce, 0x56, 0xd3, 0x2e, 0xe9, 0x84, 0x63, + 0x79, 0xba, 0x9a, 0x1e, 0x87, 0xb5, 0xf4, 0x9a, 0xd4, 0x94, 0x6d, 0x5a, 0xa8, 0x26, 0xfb, 0xb4, + 0x60, 0x82, 0x74, 0x25, 0x8f, 0x74, 0xb5, 0x4d, 0x6d, 0x76, 0x54, 0x79, 0x68, 0x10, 0xb8, 0x35, + 0x0b, 0xf2, 0x24, 0xcd, 0x62, 0x1e, 0x18, 0x93, 0xbe, 0x92, 0xc3, 0x0f, 0x08, 0x8c, 0x13, 0xce, + 0x35, 0x53, 0x69, 0xe8, 0x86, 0x1e, 0x3a, 0x30, 0x9a, 0xa9, 0x18, 0x1a, 0x65, 0xb4, 0xa0, 0x37, + 0x63, 0x6b, 0x05, 0xb1, 0x92, 0x53, 0x9b, 0xb1, 0x7f, 0x06, 0x83, 0x9c, 0x56, 0x7f, 0x5d, 0xf1, + 0x82, 0x96, 0x36, 0x3e, 0xb6, 0x0b, 0x7d, 0x66, 0xfe, 0x92, 0xcf, 0x60, 0x63, 0xbe, 0xa3, 0x79, + 0x53, 0x9b, 0xba, 0xad, 0xcd, 0x74, 0x2b, 0xbe, 0x24, 0x5f, 0x5f, 0xdf, 0xf3, 0x27, 0xdc, 0x9d, + 0x84, 0xa6, 0x72, 0xb3, 0xea, 0xac, 0x5a, 0xd4, 0x2e, 0x62, 0xbe, 0x0e, 0x95, 0x3c, 0x87, 0x7e, + 0x1a, 0x4a, 0x69, 0xac, 0xe2, 0x2c, 0xfd, 0x7e, 0x28, 0xe5, 0xf0, 0x1f, 0x03, 0x74, 0xad, 0xa7, + 0x48, 0xc5, 0x85, 0x77, 0xca, 0x4e, 0x77, 0x67, 0xbb, 0x6f, 0xbd, 0x67, 0x24, 0xd9, 0x51, 0x2a, + 0xb5, 0xf9, 0x0b, 0xed, 0x8c, 0xcf, 0xf8, 0x3b, 0x4b, 0xe4, 0x20, 0x14, 0xfe, 0xce, 0x0e, 0xac, + 0x96, 0x3c, 0x48, 0x57, 0x09, 0xe5, 0x45, 0xc6, 0x29, 0x2f, 0x55, 0x94, 0x94, 0x48, 0x9c, 0x15, + 0x6c, 0x68, 0xdf, 0xe2, 0x08, 0xa9, 0xd1, 0xd9, 0xf7, 0x45, 0x64, 0xab, 0xd9, 0xe6, 0x9c, 0x7d, + 0xc4, 0x50, 0xae, 0x3c, 0xe5, 0x78, 0x76, 0x94, 0x2f, 0x22, 0xb3, 0x83, 0xda, 0x1a, 0x72, 0xf8, + 0x22, 0xca, 0x27, 0x48, 0xce, 0x74, 0x83, 0xce, 0x11, 0x34, 0x41, 0x3a, 0x53, 0x7d, 0x0a, 0x1d, + 0x91, 0x86, 0xe3, 0x90, 0x52, 0x5f, 0xda, 0xc1, 0x99, 0x7f, 0x09, 0x68, 0x82, 0x5d, 0x7c, 0xd5, + 0x10, 0x1a, 0xc6, 0xfc, 0x2f, 0xe6, 0xcf, 0x0d, 0x06, 0x1d, 0x22, 0xa9, 0xd2, 0xd0, 0x57, 0x38, + 0x1d, 0xbd, 0x23, 0x75, 0x61, 0x54, 0x4f, 0x83, 0x0f, 0x5f, 0x44, 0x94, 0xfd, 0xbb, 0x05, 0x2b, + 0x3e, 0x99, 0x0b, 0xbd, 0xa1, 0x22, 0x1e, 0xd3, 0x9a, 0xd6, 0x9d, 0x9e, 0x06, 0xe3, 0xfc, 0x9e, + 0xf0, 0xd8, 0x14, 0x61, 0x79, 0x51, 0x84, 0x27, 0x46, 0xe1, 0x05, 0x26, 0x63, 0xde, 0xb5, 0xc0, + 0x27, 0xc2, 0x0b, 0xd8, 0xcf, 0xe1, 0x2a, 0xe2, 0x5c, 0x3e, 0x4d, 0xd4, 0x19, 0xda, 0x37, 0x9e, + 0x86, 0xbe, 0xeb, 0x49, 0xca, 0xa0, 0x9b, 0xc4, 0xf9, 0x06, 0x52, 0x3c, 0x42, 0x82, 0xa7, 0x1a, + 0xbf, 0x23, 0xbf, 0xe5, 0xa9, 0x60, 0xdf, 0x52, 0x06, 0xf0, 0x3c, 0xf7, 0xdd, 0x1e, 0xf5, 0x3f, + 0x28, 0xd6, 0xea, 0x02, 0x4a, 0xaa, 0x50, 0x41, 0x84, 0x63, 0x9d, 0x3e, 0xea, 0xcf, 0x7e, 0x03, + 0xcc, 0x1a, 0x38, 0x92, 0x7c, 0xe5, 0xc9, 0x13, 0x49, 0x51, 0x80, 0xce, 0xf6, 0x7b, 0x6f, 0xf4, + 0x51, 0x1d, 0x6b, 0x19, 0x11, 0x88, 0x00, 0xc9, 0xfe, 0x08, 0xd6, 0xf3, 0xc1, 0x8c, 0x2f, 0x43, + 0xc3, 0xe9, 0x20, 0xc1, 0xf5, 0xc5, 0xe1, 0x66, 0x5c, 0x20, 0xc7, 0xce, 0x44, 0x83, 0xf5, 0x90, + 0x5f, 0xc1, 0x8a, 0x1d, 0x52, 0x73, 0x5d, 0x0e, 0xfa, 0x34, 0xda, 0xfb, 0x0b, 0xa3, 0xcd, 0xd8, + 0xf6, 0xdc, 0x3e, 0x6b, 0x28, 0x7e, 0x68, 0x6e, 0xc9, 0xad, 0x95, 0x19, 0xac, 0x92, 0x8c, 0xdc, + 0x58, 0x18, 0x69, 0xce, 0x58, 0x39, 0x76, 0x0a, 0x16, 0xce, 0xee, 0xc3, 0x25, 0x3b, 0x98, 0xa0, + 0x84, 0xad, 0x1b, 0x0a, 0xca, 0xe5, 0x32, 0xed, 0x62, 0x19, 0xa4, 0x4e, 0xe6, 0xee, 0x09, 0x87, + 0x8f, 0xd8, 0x2f, 0xe1, 0x9a, 0xed, 0xa2, 0xad, 0x30, 0x9d, 0x48, 0xf3, 0x8f, 0x5a, 0x23, 0xdb, + 0x35, 0x30, 0x24, 0xda, 0x2e, 0xe3, 0x09, 0xd4, 0x4e, 0x7f, 0x13, 0xfa, 0x54, 0x9a, 0x8a, 0xcb, + 0x2a, 0xd2, 0x20, 0x8c, 0xbd, 0x68, 0xb0, 0x4e, 0x52, 0xb3, 0x8c, 0x70, 0x47, 0xbc, 0x7a, 0xa6, + 0xa1, 0xec, 0x08, 0x36, 0xec, 0x8b, 0x72, 0x35, 0x23, 0xd1, 0x94, 0x50, 0xd8, 0xf1, 0x3c, 0xc6, + 0xcd, 0x18, 0x1c, 0xc7, 0x2e, 0xe1, 0xac, 0x19, 0x7a, 0x08, 0xd7, 0xe7, 0x96, 0x76, 0xea, 0x9d, + 0xba, 0x53, 0x3e, 0x15, 0xe9, 0x99, 0x31, 0x20, 0x1b, 0xa4, 0xc0, 0xae, 0xcd, 0x2c, 0xe2, 0xbe, + 0x77, 0xba, 0x4f, 0x34, 0xda, 0x9c, 0xfc, 0x0a, 0xde, 0x9d, 0x1b, 0x45, 0x57, 0x7a, 0xf2, 0xd8, + 0x3b, 0x8e, 0x78, 0x30, 0xb8, 0x4c, 0x5f, 0x74, 0x65, 0x66, 0x88, 0x43, 0xa4, 0x78, 0xa4, 0x09, + 0x8c, 0x43, 0x77, 0x0c, 0x6d, 0x0a, 0x43, 0x90, 0x36, 0xcc, 0xab, 0x6e, 0x2b, 0x6f, 0xae, 0xba, + 0xfd, 0x14, 0xba, 0xc6, 0xdb, 0xbf, 0xa8, 0x8c, 0xb7, 0xa3, 0xf1, 0xf8, 0x2c, 0x87, 0x77, 0xa0, + 0x4d, 0xae, 0x3e, 0xbd, 0xe3, 0x3a, 0x74, 0xa8, 0xda, 0xcb, 0x3d, 0x8e, 0x84, 0x7f, 0x62, 0x9d, + 0x73, 0x02, 0x3d, 0x40, 0xc8, 0x10, 0xa0, 0xf5, 0x3c, 0x0e, 0x45, 0xbc, 0x13, 0x45, 0xc3, 0xbf, + 0x6c, 0x40, 0x1b, 0x7d, 0x02, 0x8a, 0x9b, 0xe0, 0xb1, 0x8a, 0x16, 0x8e, 0x72, 0xa9, 0x53, 0x2f, + 0x31, 0x75, 0xc5, 0x1d, 0x04, 0x22, 0xd5, 0xbe, 0x97, 0xcc, 0xa5, 0x5a, 0x97, 0xe6, 0x52, 0xad, + 0x1f, 0xe8, 0xbb, 0x2f, 0xba, 0xde, 0x8c, 0xdb, 0x42, 0x55, 0x1a, 0xe0, 0x81, 0x06, 0xa1, 0xaf, + 0x42, 0x24, 0x5e, 0x44, 0xfe, 0x0d, 0x9e, 0x9e, 0x22, 0x69, 0xb2, 0xb2, 0x24, 0x37, 0x3b, 0x06, + 0x71, 0xc8, 0xb5, 0x3e, 0x2e, 0x05, 0xcb, 0xea, 0xf3, 0xc1, 0xb2, 0xdb, 0x00, 0xbe, 0x88, 0x03, + 0x72, 0xa1, 0xe6, 0xb2, 0x61, 0x3a, 0x25, 0x5a, 0x60, 0xbf, 0x47, 0x68, 0xf6, 0x63, 0xe8, 0xe7, + 0x14, 0xe8, 0x21, 0xf9, 0x71, 0x7e, 0xfe, 0x34, 0x54, 0x0e, 0x1f, 0xed, 0xc6, 0x6a, 0x3e, 0x86, + 0xdb, 0x5e, 0x88, 0xe1, 0x5e, 0x90, 0x3c, 0x87, 0x1f, 0x7c, 0x2d, 0xe3, 0x0a, 0xb4, 0xa8, 0x4a, + 0x27, 0xc8, 0x12, 0xa3, 0xab, 0x9b, 0xa1, 0xa4, 0x58, 0xfb, 0x45, 0x71, 0xe2, 0xee, 0xff, 0xaf, + 0x38, 0x71, 0xef, 0xfb, 0xc5, 0x89, 0x97, 0xbf, 0x5f, 0x9c, 0x78, 0x2e, 0xae, 0xba, 0x32, 0x9f, + 0x8e, 0xb9, 0x30, 0xf9, 0xd1, 0xbf, 0x30, 0xf9, 0xf1, 0x96, 0xcc, 0xc5, 0xea, 0x1b, 0x33, 0x17, + 0xdf, 0x23, 0x75, 0xc2, 0xde, 0x96, 0x3a, 0xb9, 0x05, 0x2b, 0x2a, 0xf5, 0xfc, 0x13, 0x7d, 0x12, + 0x39, 0xe1, 0x67, 0xd2, 0xa4, 0x6a, 0x7a, 0x04, 0xc6, 0x73, 0xc8, 0x6f, 0xf8, 0x99, 0x1c, 0x3e, + 0x07, 0xa0, 0x23, 0x1a, 0x7d, 0xda, 0x45, 0xb2, 0x51, 0xf9, 0xc1, 0x85, 0x15, 0xff, 0xa7, 0x02, + 0x70, 0xe8, 0x4d, 0x13, 0x1d, 0xe3, 0x64, 0x7f, 0x08, 0x1d, 0x49, 0xad, 0x72, 0x86, 0xbb, 0x64, + 0xc8, 0x0a, 0x52, 0xf3, 0xa8, 0x6f, 0x17, 0xc8, 0xfc, 0x99, 0xc4, 0x5a, 0x8f, 0x90, 0x17, 0xb1, + 0xd5, 0x2d, 0x01, 0xc5, 0xbe, 0x6e, 0xc2, 0xb2, 0x21, 0x48, 0x78, 0xea, 0xf3, 0x58, 0x57, 0xc6, + 0x56, 0x9c, 0x9e, 0x86, 0x1e, 0x68, 0x20, 0xbb, 0x9f, 0x93, 0x59, 0x93, 0xb1, 0x98, 0xa6, 0x31, + 0x5d, 0x8c, 0xcd, 0x18, 0x6e, 0xdb, 0x4f, 0xa1, 0x89, 0xb4, 0xa0, 0x86, 0xef, 0xeb, 0xbf, 0xc3, + 0x3a, 0xd0, 0x34, 0xa3, 0xf6, 0x2b, 0xac, 0x07, 0x6d, 0xba, 0xce, 0x42, 0xb8, 0xa5, 0xe1, 0x1f, + 0xaf, 0x42, 0x67, 0x2f, 0x96, 0x2a, 0xcd, 0xb4, 0x08, 0x17, 0x97, 0x36, 0xea, 0x74, 0x69, 0xc3, + 0x54, 0x50, 0xea, 0xcf, 0xa0, 0x0a, 0xca, 0x4f, 0xa1, 0x69, 0xee, 0x07, 0x99, 0xc0, 0xf7, 0xb9, + 0x97, 0x8b, 0x2c, 0x0d, 0xdb, 0x82, 0x56, 0x60, 0x2e, 0x2e, 0x99, 0x34, 0x7e, 0xe9, 0x36, 0x91, + 0xbd, 0xd2, 0xe4, 0xe4, 0x34, 0xec, 0x03, 0xa8, 0x7a, 0xe3, 0xb1, 0x39, 0xf5, 0xae, 0x14, 0xa4, + 0xe4, 0xc4, 0x38, 0x88, 0x63, 0x77, 0xa1, 0x4d, 0xea, 0x93, 0x2a, 0x59, 0x1a, 0xf3, 0x63, 0xda, + 0x32, 0x19, 0xad, 0x51, 0x29, 0x66, 0x7e, 0x17, 0xda, 0x91, 0x10, 0x89, 0xee, 0xd0, 0x9c, 0xef, 0x60, 0x8b, 0x1b, 0x9c, 0x56, 0x64, 0xcb, 0x1c, 0x6e, 0x41, 0x03, 0xdd, 0x63, 0x91, 0x18, 0xb7, - 0xb2, 0x34, 0x0f, 0x4a, 0xf2, 0x3b, 0x75, 0x89, 0x3f, 0x6c, 0x07, 0x40, 0xcb, 0x3f, 0x8d, 0xdc, - 0x5e, 0x64, 0x47, 0x9e, 0xcf, 0xc3, 0x4d, 0x6a, 0x53, 0x7b, 0x0f, 0x60, 0xa0, 0x73, 0x37, 0xa5, - 0x9e, 0x60, 0xeb, 0xff, 0x6c, 0xcf, 0xf9, 0x74, 0xa0, 0xd3, 0x4f, 0xe7, 0xd3, 0x83, 0x1f, 0x43, - 0x33, 0xd1, 0xc9, 0x0b, 0xd2, 0x30, 0x9d, 0x9d, 0xb5, 0xa2, 0xab, 0xc9, 0x6a, 0x38, 0x96, 0x82, - 0xfd, 0x01, 0xf4, 0x75, 0x9d, 0xda, 0xd8, 0x44, 0xf1, 0x29, 0xf2, 0x35, 0x77, 0xb5, 0x64, 0x2e, - 0xc8, 0xef, 0xf4, 0xd4, 0x5c, 0xcc, 0xff, 0xe7, 0xd0, 0x2b, 0x4a, 0xfd, 0x7d, 0x2f, 0x26, 0xbd, - 0x43, 0xd1, 0x74, 0xdb, 0xbd, 0x7c, 0x5a, 0x71, 0xba, 0xbc, 0x7c, 0x76, 0xd9, 0x82, 0x86, 0xa9, - 0x9d, 0x1c, 0x50, 0xaf, 0xd2, 0xcd, 0x53, 0x5d, 0x2d, 0xe5, 0x18, 0x3c, 0xf2, 0xb2, 0x28, 0x0b, - 0x23, 0xc7, 0x6a, 0x8e, 0x97, 0x79, 0x4d, 0x98, 0xd3, 0xce, 0xcb, 0xc1, 0xd8, 0xa3, 0xf9, 0x32, - 0x35, 0x5d, 0x8e, 0xb5, 0x4e, 0x5d, 0xaf, 0x5c, 0xd0, 0x55, 0x57, 0x65, 0x39, 0xab, 0xc9, 0x42, - 0xb5, 0xdb, 0x1d, 0x68, 0x89, 0x34, 0xa0, 0xca, 0x5a, 0xca, 0xf5, 0x12, 0x3f, 0xa9, 0x3a, 0x4f, - 0x5f, 0x8a, 0x22, 0xe5, 0xd1, 0x14, 0xba, 0x81, 0x8e, 0x45, 0x92, 0x0a, 0xf2, 0x02, 0x49, 0xc5, - 0x5d, 0x5a, 0x76, 0x2c, 0x0c, 0x9e, 0x14, 0xdc, 0x87, 0xd0, 0xb4, 0x15, 0xa1, 0x9b, 0x4b, 0x94, - 0x16, 0xc5, 0x3e, 0x85, 0xd5, 0x79, 0x85, 0x26, 0x87, 0x97, 0x97, 0xa8, 0xfb, 0x73, 0xfa, 0x0b, - 0xad, 0x71, 0x3d, 0x0a, 0x67, 0xa1, 0x1a, 0x0e, 0x97, 0x0e, 0x3f, 0x1a, 0x81, 0xe7, 0x23, 0x93, - 0x22, 0xb8, 0xb2, 0x7c, 0x3e, 0x32, 0x69, 0x84, 0x21, 0x34, 0x43, 0xf9, 0x38, 0x4c, 0xa5, 0x1a, - 0x5e, 0xb5, 0xd6, 0x91, 0x9a, 0x6c, 0x13, 0x1a, 0xa1, 0x44, 0x33, 0x31, 0xbc, 0x66, 0xaf, 0xd1, - 0x91, 0xd1, 0xb8, 0x0d, 0x0d, 0x53, 0x2d, 0x7b, 0x63, 0x69, 0x47, 0x9b, 0x9a, 0x74, 0xc7, 0x50, - 0xb0, 0x1f, 0x43, 0x93, 0x4a, 0x25, 0x45, 0x32, 0x7c, 0x7f, 0x51, 0x02, 0x74, 0xbd, 0xa2, 0xd3, - 0x88, 0x74, 0xdd, 0xe2, 0xc7, 0xd0, 0xb4, 0x4e, 0xca, 0x68, 0x51, 0xaa, 0x8d, 0xb3, 0xe2, 0x58, - 0x0a, 0x76, 0x13, 0xea, 0x33, 0xd4, 0x63, 0xc3, 0x0f, 0x16, 0x77, 0xa8, 0x56, 0x6f, 0x1a, 0xcb, - 0xfe, 0x1e, 0x5c, 0x2d, 0x17, 0x1b, 0xda, 0x4a, 0x44, 0x13, 0x00, 0xbc, 0x49, 0x7d, 0xdf, 0xbf, - 0x40, 0x54, 0xe6, 0x6b, 0x16, 0x9d, 0xcb, 0xc9, 0x1b, 0x8a, 0x19, 0x3f, 0xcf, 0xd5, 0x3d, 0xee, - 0xae, 0xe1, 0x2d, 0x5b, 0xc6, 0xb8, 0x6c, 0x30, 0xac, 0x11, 0x20, 0x3b, 0xf3, 0x05, 0x74, 0xc7, - 0xd9, 0xeb, 0xd7, 0xe7, 0x36, 0x78, 0xfd, 0x11, 0xf5, 0x2b, 0x1d, 0xc1, 0x4b, 0xf5, 0x8d, 0x4e, - 0x67, 0x5c, 0x2a, 0x76, 0xbc, 0x0c, 0x4d, 0x3f, 0x76, 0xbd, 0x20, 0x48, 0x87, 0x5b, 0xba, 0xbe, - 0xd1, 0x8f, 0x77, 0x83, 0x80, 0x6e, 0x0d, 0x8b, 0x84, 0xd3, 0x15, 0x3d, 0x37, 0x0c, 0x86, 0x3f, - 0xd6, 0x86, 0xc7, 0x82, 0xf6, 0x03, 0xba, 0x56, 0x6c, 0xcf, 0xad, 0x61, 0x30, 0xbc, 0x6d, 0xae, - 0x15, 0x1b, 0xd0, 0x7e, 0x80, 0x8e, 0x27, 0x3a, 0xf9, 0x16, 0x32, 0xfc, 0x58, 0x27, 0x04, 0x66, - 0xde, 0xd9, 0xa1, 0x01, 0xe1, 0x26, 0xd5, 0xe9, 0x35, 0x52, 0x5b, 0x77, 0x16, 0x37, 0x69, 0x9e, - 0xa6, 0x74, 0xda, 0x61, 0x9e, 0xb1, 0xa4, 0x8d, 0x4d, 0xaa, 0xc8, 0x8d, 0x76, 0x86, 0x9f, 0x2c, - 0x6f, 0x6c, 0x93, 0x85, 0xc5, 0x8d, 0x6d, 0x13, 0xb2, 0x3b, 0x00, 0x5a, 0x67, 0x91, 0xc2, 0xd9, - 0x5e, 0xec, 0x93, 0x9f, 0x06, 0x1c, 0x7d, 0x19, 0x80, 0x54, 0xcd, 0x0e, 0x00, 0x85, 0xdf, 0x75, - 0x9f, 0xbb, 0x8b, 0x7d, 0x72, 0xef, 0xde, 0x69, 0xbf, 0xcc, 0x1d, 0xfd, 0xbb, 0xd0, 0xce, 0xd0, - 0x8f, 0x47, 0x4f, 0x7a, 0x78, 0x6f, 0x51, 0x98, 0xad, 0x8b, 0xef, 0xb4, 0x32, 0xf3, 0x84, 0x2f, - 0x21, 0xdb, 0x43, 0x6e, 0xc8, 0xf0, 0xfe, 0xe2, 0x4b, 0xf2, 0x73, 0x80, 0x43, 0x26, 0x4a, 0x1f, - 0x09, 0x3e, 0x87, 0x8e, 0x66, 0x9a, 0xee, 0xb4, 0xb3, 0x28, 0x23, 0x85, 0x5f, 0xe3, 0x68, 0xee, - 0xea, 0x6e, 0x37, 0xa1, 0xee, 0x25, 0x49, 0x74, 0x3e, 0xfc, 0x74, 0x51, 0xc2, 0x77, 0x11, 0xec, - 0x68, 0x2c, 0x8a, 0xd2, 0x2c, 0x8b, 0x54, 0x68, 0xeb, 0xf7, 0x3f, 0x5b, 0x14, 0xa5, 0xd2, 0x85, - 0x28, 0xa7, 0x33, 0x2b, 0xdd, 0x8e, 0xba, 0x03, 0xad, 0x44, 0x48, 0xe5, 0x06, 0xb3, 0x68, 0xf8, - 0xf9, 0x92, 0x19, 0xd1, 0x55, 0xe8, 0x4e, 0x33, 0x31, 0x65, 0xfc, 0x73, 0xd7, 0xf4, 0x7e, 0x32, - 0x7f, 0x4d, 0x8f, 0xed, 0x40, 0x77, 0x26, 0xe2, 0x89, 0x08, 0x4e, 0x34, 0xf7, 0x7f, 0x5a, 0xae, - 0x71, 0x3e, 0x40, 0x0c, 0x71, 0xbe, 0x63, 0x88, 0xb0, 0xa1, 0xcf, 0x76, 0xbf, 0xae, 0xb5, 0xd6, - 0x06, 0xec, 0xd7, 0xb5, 0xd6, 0x87, 0x83, 0x9b, 0x4e, 0x47, 0xd2, 0xad, 0x79, 0x1a, 0x62, 0xf4, - 0x39, 0x74, 0x77, 0xe9, 0x1f, 0x03, 0x42, 0x49, 0x7a, 0xf4, 0x26, 0xd4, 0xf2, 0xfc, 0x7c, 0xae, - 0xa0, 0x89, 0xe2, 0x35, 0xdf, 0x8f, 0xc7, 0xc2, 0x21, 0xf4, 0xe8, 0x5f, 0xd7, 0xa0, 0x71, 0x24, - 0xb2, 0xd4, 0xe7, 0xdf, 0x7e, 0x97, 0xe4, 0x47, 0x56, 0xca, 0xe2, 0xa2, 0x5c, 0x58, 0x0b, 0x14, - 0xa1, 0x17, 0xeb, 0x13, 0xdb, 0x45, 0xea, 0x7f, 0x03, 0xea, 0xfa, 0x68, 0xa8, 0x23, 0xca, 0xba, - 0x41, 0x3b, 0x2c, 0x93, 0x53, 0xfa, 0x5b, 0x00, 0x93, 0x23, 0xa9, 0x39, 0x60, 0x41, 0xfb, 0x01, + 0xb2, 0x34, 0x0f, 0x4a, 0xf2, 0x3b, 0x75, 0x89, 0x3f, 0x6c, 0x1b, 0x40, 0xcb, 0x3f, 0x8d, 0xdc, + 0x9e, 0x67, 0x47, 0x9e, 0xcf, 0xc3, 0x4d, 0x6a, 0x53, 0x7b, 0x0f, 0xa0, 0xaf, 0x73, 0x37, 0xa5, + 0x9e, 0x60, 0xeb, 0xff, 0x6c, 0xcf, 0xd9, 0x74, 0xa0, 0xb3, 0x9c, 0xce, 0xa6, 0x07, 0x3f, 0x81, + 0x66, 0xa2, 0x93, 0x17, 0xa4, 0x61, 0x3a, 0xdb, 0xab, 0x45, 0x57, 0x93, 0xd5, 0x70, 0x2c, 0x05, + 0xfb, 0x03, 0x58, 0xd6, 0x75, 0x6a, 0x23, 0x13, 0xc5, 0xa7, 0xc8, 0xd7, 0xcc, 0xd5, 0x92, 0x99, + 0x20, 0xbf, 0xd3, 0x53, 0x33, 0x31, 0xff, 0x5f, 0x40, 0xaf, 0x28, 0xf5, 0xf7, 0xbd, 0x98, 0xf4, + 0x0e, 0x45, 0xd3, 0x6d, 0xf7, 0xf2, 0x69, 0xc5, 0xe9, 0xf2, 0xf2, 0xd9, 0x65, 0x13, 0x1a, 0xa6, + 0x76, 0xb2, 0x4f, 0xbd, 0x4a, 0x37, 0x4f, 0x75, 0xb5, 0x94, 0x63, 0xf0, 0xc8, 0xcb, 0xa2, 0x2c, + 0x8c, 0x1c, 0xab, 0x19, 0x5e, 0xe6, 0x35, 0x61, 0x4e, 0x3b, 0x2f, 0x07, 0x63, 0x8f, 0x66, 0xcb, + 0xd4, 0x74, 0x39, 0xd6, 0x1a, 0x75, 0xbd, 0x72, 0x4e, 0x57, 0x5d, 0x95, 0xe5, 0xac, 0x24, 0x73, + 0xd5, 0x6e, 0x77, 0xa0, 0x25, 0xd2, 0x80, 0x2a, 0x6b, 0x29, 0xd7, 0x4b, 0xfc, 0xa4, 0xea, 0x3c, + 0x7d, 0x29, 0x8a, 0x94, 0x47, 0x53, 0xe8, 0x06, 0x3a, 0x16, 0x49, 0x2a, 0xc8, 0x0b, 0x24, 0x15, + 0x77, 0x69, 0xd1, 0xb1, 0x30, 0x78, 0x52, 0x70, 0x1f, 0x41, 0xd3, 0x56, 0x84, 0x6e, 0x2c, 0x50, + 0x5a, 0x14, 0xfb, 0x0c, 0x56, 0x66, 0x15, 0x9a, 0x1c, 0x5c, 0x5e, 0xa0, 0x5e, 0x9e, 0xd1, 0x5f, + 0x68, 0x8d, 0xeb, 0x51, 0x38, 0x0d, 0xd5, 0x60, 0xb0, 0x70, 0xf8, 0xd1, 0x08, 0x3c, 0x1f, 0x99, + 0x14, 0xc1, 0x95, 0xc5, 0xf3, 0x91, 0x49, 0x23, 0x0c, 0xa0, 0x19, 0xca, 0xc7, 0x61, 0x2a, 0xd5, + 0xe0, 0xaa, 0xb5, 0x8e, 0xd4, 0x64, 0x1b, 0xd0, 0x08, 0x25, 0x9a, 0x89, 0xc1, 0x35, 0x7b, 0x8d, + 0x8e, 0x8c, 0xc6, 0x6d, 0x68, 0x98, 0x6a, 0xd9, 0x1b, 0x0b, 0x3b, 0xda, 0xd4, 0xa4, 0x3b, 0x86, + 0x82, 0xfd, 0x04, 0x9a, 0x54, 0x2a, 0x29, 0x92, 0xc1, 0x07, 0xf3, 0x12, 0xa0, 0xeb, 0x15, 0x9d, + 0x46, 0xa4, 0xeb, 0x16, 0x3f, 0x81, 0xa6, 0x75, 0x52, 0x86, 0xf3, 0x52, 0x6d, 0x9c, 0x15, 0xc7, + 0x52, 0xb0, 0x9b, 0x50, 0x9f, 0xa2, 0x1e, 0x1b, 0x7c, 0x38, 0xbf, 0x43, 0xb5, 0x7a, 0xd3, 0x58, + 0xf6, 0xf7, 0xe0, 0x6a, 0xb9, 0xd8, 0xd0, 0x56, 0x22, 0x9a, 0x00, 0xe0, 0x4d, 0xea, 0xfb, 0xc1, + 0x39, 0xa2, 0x32, 0x5b, 0xb3, 0xe8, 0x5c, 0x4e, 0x2e, 0x28, 0x66, 0xfc, 0x22, 0x57, 0xf7, 0xb8, + 0xbb, 0x06, 0xb7, 0x6c, 0x19, 0xe3, 0xa2, 0xc1, 0xb0, 0x46, 0x80, 0xec, 0xcc, 0x97, 0xd0, 0x1d, + 0x65, 0xaf, 0x5f, 0x9f, 0xd9, 0xe0, 0xf5, 0xc7, 0xd4, 0xaf, 0x74, 0x04, 0x2f, 0xd5, 0x37, 0x3a, + 0x9d, 0x51, 0xa9, 0xd8, 0xf1, 0x32, 0x34, 0xfd, 0xd8, 0xf5, 0x82, 0x20, 0x1d, 0x6c, 0xea, 0xfa, + 0x46, 0x3f, 0xde, 0x09, 0x02, 0xba, 0x35, 0x2c, 0x12, 0x4e, 0x57, 0xf4, 0xdc, 0x30, 0x18, 0xfc, + 0x44, 0x1b, 0x1e, 0x0b, 0xda, 0x0b, 0xe8, 0x5a, 0xb1, 0x3d, 0xb7, 0x86, 0xc1, 0xe0, 0xb6, 0xb9, + 0x56, 0x6c, 0x40, 0x7b, 0x01, 0x3a, 0x9e, 0xe8, 0xe4, 0x5b, 0xc8, 0xe0, 0x13, 0x9d, 0x10, 0x98, + 0x7a, 0xa7, 0x07, 0x06, 0x84, 0x9b, 0x54, 0xa7, 0xd7, 0x48, 0x6d, 0xdd, 0x99, 0xdf, 0xa4, 0x79, + 0x9a, 0xd2, 0x69, 0x87, 0x79, 0xc6, 0x92, 0x36, 0x36, 0xa9, 0x22, 0x37, 0xda, 0x1e, 0x7c, 0xba, + 0xb8, 0xb1, 0x4d, 0x16, 0x16, 0x37, 0xb6, 0x4d, 0xc8, 0x6e, 0x03, 0x68, 0x9d, 0x45, 0x0a, 0x67, + 0x6b, 0xbe, 0x4f, 0x7e, 0x1a, 0x70, 0xf4, 0x65, 0x00, 0x52, 0x35, 0xdb, 0x00, 0x14, 0x7e, 0xd7, + 0x7d, 0xee, 0xce, 0xf7, 0xc9, 0xbd, 0x7b, 0xa7, 0xfd, 0x32, 0x77, 0xf4, 0xef, 0x42, 0x3b, 0x43, + 0x3f, 0x1e, 0x3d, 0xe9, 0xc1, 0xbd, 0x79, 0x61, 0xb6, 0x2e, 0xbe, 0xd3, 0xca, 0xcc, 0x13, 0xbe, + 0x84, 0x6c, 0x0f, 0xb9, 0x21, 0x83, 0xfb, 0xf3, 0x2f, 0xc9, 0xcf, 0x01, 0x0e, 0x99, 0x28, 0x7d, + 0x24, 0xf8, 0x02, 0x3a, 0x9a, 0x69, 0xba, 0xd3, 0xf6, 0xbc, 0x8c, 0x14, 0x7e, 0x8d, 0xa3, 0xb9, + 0xab, 0xbb, 0xdd, 0x84, 0xba, 0x97, 0x24, 0xd1, 0xd9, 0xe0, 0xb3, 0x79, 0x09, 0xdf, 0x41, 0xb0, + 0xa3, 0xb1, 0x28, 0x4a, 0xd3, 0x2c, 0x52, 0xa1, 0xad, 0xdf, 0xff, 0x7c, 0x5e, 0x94, 0x4a, 0x17, + 0xa2, 0x9c, 0xce, 0xb4, 0x74, 0x3b, 0xea, 0x0e, 0xb4, 0x12, 0x21, 0x95, 0x1b, 0x4c, 0xa3, 0xc1, + 0x17, 0x0b, 0x66, 0x44, 0x57, 0xa1, 0x3b, 0xcd, 0xc4, 0x94, 0xf1, 0xcf, 0x5c, 0xd3, 0xfb, 0xe9, + 0xec, 0x35, 0x3d, 0xb6, 0x0d, 0xdd, 0xa9, 0x88, 0xc7, 0x22, 0x38, 0xd6, 0xdc, 0xff, 0x59, 0xb9, + 0xc6, 0x79, 0x1f, 0x31, 0xc4, 0xf9, 0x8e, 0x21, 0xc2, 0x86, 0x3e, 0xdb, 0xfd, 0xba, 0xd6, 0x5a, + 0xed, 0xb3, 0x5f, 0xd7, 0x5a, 0x1f, 0xf5, 0x6f, 0x3a, 0x1d, 0x49, 0xb7, 0xe6, 0x69, 0x88, 0xe1, + 0x17, 0xd0, 0xdd, 0xa1, 0x7f, 0x0c, 0x08, 0x25, 0xe9, 0xd1, 0x9b, 0x50, 0xcb, 0xf3, 0xf3, 0xb9, + 0x82, 0x26, 0x8a, 0xd7, 0x7c, 0x2f, 0x1e, 0x09, 0x87, 0xd0, 0xc3, 0x7f, 0x53, 0x83, 0xc6, 0xa1, + 0xc8, 0x52, 0x9f, 0xbf, 0xfd, 0x2e, 0xc9, 0x7b, 0x56, 0xca, 0xe2, 0xa2, 0x5c, 0x58, 0x0b, 0x14, + 0xa1, 0xe7, 0xeb, 0x13, 0xdb, 0x45, 0xea, 0x7f, 0x1d, 0xea, 0xfa, 0x68, 0xa8, 0x23, 0xca, 0xba, + 0x41, 0x3b, 0x2c, 0x93, 0x13, 0xfa, 0x5b, 0x00, 0x93, 0x23, 0xa9, 0x39, 0x60, 0x41, 0x7b, 0x01, 0x85, 0x8e, 0x2c, 0x01, 0x6d, 0xe1, 0x86, 0x09, 0x0d, 0x1b, 0x20, 0x6d, 0x64, 0x5b, 0x56, 0xd0, - 0x7c, 0x43, 0x59, 0xc1, 0x7b, 0x50, 0x8b, 0x6d, 0x25, 0x7b, 0x8e, 0xa7, 0x3b, 0xd9, 0x04, 0x67, - 0xb7, 0x21, 0xbf, 0x00, 0x63, 0x5c, 0x92, 0x37, 0x5f, 0x90, 0xd9, 0x81, 0x76, 0xfe, 0x1f, 0x13, - 0xc6, 0x0b, 0xd9, 0xd8, 0x2e, 0xfe, 0x75, 0xe2, 0xd8, 0x3e, 0x39, 0x05, 0xd9, 0xdb, 0x93, 0xe3, - 0x9d, 0x1f, 0x94, 0x1c, 0x37, 0x47, 0x34, 0x5f, 0xc4, 0x52, 0x99, 0xe0, 0x58, 0x33, 0x94, 0x7b, - 0xd8, 0x64, 0xbf, 0x07, 0xbd, 0x94, 0xfb, 0x2f, 0xdd, 0x99, 0x9c, 0xe8, 0x57, 0xf4, 0xca, 0x97, - 0xf0, 0x66, 0x72, 0xf2, 0x25, 0x25, 0xee, 0xcd, 0x89, 0xa9, 0x83, 0xb4, 0x07, 0x72, 0x42, 0xa3, - 0x7e, 0x0c, 0x6b, 0x33, 0x3e, 0x3b, 0xe1, 0xa9, 0x9c, 0x86, 0x89, 0x55, 0xb5, 0x7d, 0x2a, 0x30, - 0x18, 0x14, 0x08, 0x3d, 0x97, 0xd1, 0x3f, 0xaa, 0x40, 0x0b, 0xb9, 0x88, 0xb2, 0xc4, 0x18, 0xd4, - 0x66, 0x7e, 0x92, 0x19, 0x47, 0x98, 0x9e, 0xcd, 0xff, 0x56, 0x68, 0x29, 0x31, 0xff, 0x5b, 0x41, - 0x6b, 0xa8, 0x53, 0x3e, 0xf4, 0xac, 0xef, 0x81, 0x9f, 0x53, 0x54, 0x50, 0x4b, 0x86, 0x6d, 0xb2, - 0x4b, 0xd0, 0xf0, 0x63, 0x3a, 0x0d, 0xeb, 0xd4, 0x59, 0xdd, 0x8f, 0xf1, 0x14, 0xac, 0xc1, 0x45, - 0x35, 0x74, 0xdd, 0x8f, 0xf7, 0x83, 0xb3, 0xd1, 0xbf, 0xaf, 0xc0, 0xda, 0x61, 0x2a, 0x7c, 0x2e, - 0xe5, 0x13, 0x34, 0xe4, 0x94, 0xa6, 0xc0, 0x37, 0x52, 0x54, 0x57, 0x67, 0x04, 0xe8, 0x19, 0x65, - 0x58, 0x87, 0x2a, 0xf2, 0xe3, 0x46, 0xd5, 0x69, 0x13, 0x84, 0x4e, 0x1b, 0x39, 0xba, 0x94, 0xfe, - 0xd6, 0x68, 0x8a, 0x07, 0xdf, 0x84, 0x7e, 0x91, 0x58, 0x29, 0x65, 0xea, 0x8b, 0x6b, 0xa6, 0x34, - 0xca, 0x75, 0xe8, 0x98, 0x7a, 0x0a, 0x1a, 0x46, 0x87, 0xf8, 0x41, 0x83, 0x8e, 0xcc, 0x2c, 0xb4, - 0x72, 0x20, 0xbc, 0x0e, 0xea, 0x6b, 0x75, 0x81, 0xe8, 0xd1, 0xdf, 0x87, 0xc1, 0x61, 0xca, 0x13, - 0x2f, 0xe5, 0x54, 0x5f, 0x41, 0x2c, 0xde, 0x84, 0x46, 0xc4, 0xe3, 0x89, 0x49, 0xe9, 0x57, 0x1d, - 0xd3, 0xca, 0xff, 0x72, 0x64, 0xa5, 0xf4, 0x97, 0x23, 0xc8, 0xea, 0x94, 0x7b, 0xe6, 0x9f, 0x49, - 0xe8, 0x19, 0xb7, 0x20, 0x1e, 0x19, 0xf5, 0xb9, 0xa8, 0xe5, 0xe8, 0x86, 0xb9, 0xcb, 0x76, 0x12, - 0xc6, 0x54, 0x9b, 0x46, 0x77, 0xd9, 0x1e, 0x84, 0xf1, 0xe8, 0x3f, 0x54, 0xa1, 0x63, 0xf8, 0x49, - 0x2f, 0xd7, 0x6b, 0x59, 0xc9, 0xd7, 0x72, 0x00, 0x55, 0xf9, 0x22, 0x32, 0x8b, 0x8b, 0x8f, 0xec, - 0x53, 0xa8, 0x46, 0xe1, 0xcc, 0x1c, 0x71, 0xae, 0xcd, 0x99, 0xab, 0xf9, 0x55, 0x31, 0x82, 0x87, - 0xd4, 0xa8, 0x23, 0xe9, 0x1e, 0x2c, 0x8a, 0xb8, 0xe1, 0x24, 0x9a, 0x8e, 0x33, 0xdc, 0x47, 0xc8, - 0x23, 0xcf, 0xa7, 0x32, 0x03, 0xab, 0x1c, 0x7a, 0x4e, 0xdb, 0x40, 0xf6, 0x03, 0xf6, 0x19, 0xb4, - 0xf2, 0x40, 0xa5, 0x3d, 0xd4, 0xa8, 0xb3, 0x78, 0x7b, 0xef, 0xe9, 0xf1, 0x59, 0x6c, 0x23, 0x91, - 0xe6, 0x65, 0x39, 0x25, 0xfb, 0x03, 0xe8, 0x4a, 0x2e, 0xa5, 0xbe, 0xa8, 0x38, 0x16, 0x46, 0x69, - 0x5c, 0x2a, 0x9f, 0x57, 0x08, 0x8b, 0x5f, 0x6d, 0xb7, 0x88, 0x2c, 0x40, 0xec, 0x4b, 0xe8, 0xdb, - 0xfe, 0x91, 0x98, 0x4c, 0xf2, 0x40, 0xfa, 0xb5, 0xa5, 0x11, 0x9e, 0x10, 0xba, 0x34, 0x4e, 0x4f, - 0x96, 0x11, 0xec, 0x57, 0xd0, 0x4f, 0xf4, 0x1a, 0xbb, 0xa6, 0x94, 0x47, 0x2b, 0x9f, 0xab, 0x73, - 0xde, 0xd5, 0x9c, 0x0c, 0x14, 0x57, 0x89, 0x0a, 0xb8, 0x5c, 0xbe, 0xb4, 0xab, 0x33, 0x2b, 0x73, - 0x97, 0x76, 0x47, 0xff, 0xbb, 0x0a, 0x9d, 0xd2, 0xa7, 0xd1, 0x9f, 0xc8, 0x48, 0x9e, 0xda, 0x22, - 0x1a, 0x7c, 0x46, 0xd8, 0x54, 0x48, 0x5b, 0x13, 0x42, 0xcf, 0x08, 0x4b, 0x45, 0x9e, 0x1b, 0xa7, - 0x67, 0x7c, 0xa1, 0x39, 0xad, 0xea, 0x6b, 0xd5, 0xb4, 0x72, 0x35, 0xa7, 0x5b, 0x00, 0xf7, 0x03, - 0xfa, 0xb7, 0x19, 0x4f, 0x79, 0x27, 0x9e, 0xb4, 0xd5, 0x48, 0x79, 0x1b, 0x77, 0xfd, 0x4b, 0x9e, - 0xe2, 0x5c, 0x6c, 0x2e, 0xd0, 0x34, 0x51, 0x20, 0x48, 0x51, 0xbe, 0x16, 0xb1, 0xce, 0x03, 0x76, - 0x9d, 0x16, 0x02, 0xbe, 0x16, 0x31, 0x75, 0x33, 0xcb, 0x6f, 0xf2, 0xd9, 0xb6, 0x89, 0xea, 0xf0, - 0x45, 0xc6, 0xd1, 0x4d, 0x0d, 0xe8, 0x3e, 0x43, 0xdb, 0x69, 0x52, 0x5b, 0xa7, 0xd8, 0xc9, 0x9f, - 0x7e, 0xe5, 0x85, 0x8a, 0xe4, 0x4c, 0x64, 0xca, 0x70, 0x68, 0x15, 0x11, 0x5f, 0x79, 0xa1, 0x3a, - 0xd6, 0x60, 0x76, 0xdf, 0x5c, 0x53, 0x2a, 0xd3, 0xba, 0x78, 0x18, 0xd0, 0x51, 0x30, 0xb6, 0x40, - 0x7f, 0xc4, 0xe9, 0xff, 0x56, 0x66, 0x9e, 0x4a, 0xc3, 0x33, 0x11, 0xa3, 0x59, 0x54, 0xe1, 0x4b, - 0x5e, 0xfc, 0xc3, 0x4d, 0xcb, 0x59, 0xcf, 0x91, 0x4f, 0x09, 0x47, 0x89, 0x93, 0xe7, 0xb0, 0xc5, - 0xcf, 0x92, 0x28, 0xf4, 0xc3, 0x85, 0x0b, 0x82, 0xae, 0xef, 0x49, 0xe5, 0xa6, 0x5c, 0x65, 0x69, - 0x2c, 0x29, 0xc0, 0x63, 0xaa, 0x34, 0x3e, 0xb0, 0xf4, 0xe5, 0x4b, 0x83, 0x7b, 0x9e, 0x54, 0x8e, - 0xa6, 0x7d, 0x9a, 0x45, 0x11, 0x32, 0x21, 0x4f, 0xd8, 0xe8, 0x5a, 0x8d, 0xa6, 0xd4, 0xa9, 0x9a, - 0xd1, 0x7f, 0xad, 0xc0, 0xda, 0x92, 0x58, 0xa2, 0x6b, 0x8c, 0x22, 0x69, 0xf3, 0xcb, 0x5d, 0xa7, - 0x81, 0xcd, 0xfd, 0x80, 0x10, 0x6a, 0xa6, 0x6c, 0x66, 0x19, 0x11, 0x6a, 0x86, 0x7b, 0xee, 0x12, - 0x34, 0xd4, 0x19, 0x2d, 0xb9, 0xd6, 0x2c, 0x75, 0x75, 0x86, 0x6b, 0xbd, 0x0b, 0xed, 0x48, 0x4c, - 0xdc, 0x88, 0xbf, 0xe4, 0xfa, 0xd6, 0x76, 0x7f, 0xe7, 0xc3, 0xb7, 0xec, 0x87, 0xed, 0x27, 0x62, - 0xf2, 0x04, 0x69, 0x9d, 0x56, 0x64, 0x9e, 0x46, 0xbf, 0x86, 0x96, 0x85, 0xb2, 0x36, 0xd4, 0x1f, - 0xf2, 0x93, 0x6c, 0x32, 0x78, 0x87, 0xb5, 0xa0, 0x86, 0x3d, 0x06, 0x15, 0x7c, 0xfa, 0xca, 0x4b, - 0xe3, 0xc1, 0x0a, 0xa2, 0x1f, 0xa5, 0xa9, 0x48, 0x07, 0x55, 0x7c, 0x3c, 0xf4, 0xe2, 0xd0, 0x1f, - 0xd4, 0xf0, 0xf1, 0xb1, 0xa7, 0xbc, 0x68, 0x50, 0x1f, 0xfd, 0xb6, 0x0e, 0xad, 0x43, 0xf3, 0x76, - 0xf6, 0x10, 0x7a, 0xf9, 0x9f, 0x03, 0x5d, 0x1c, 0xa3, 0x3a, 0x5c, 0x7c, 0xa0, 0x18, 0x55, 0x37, - 0x29, 0xb5, 0x16, 0xff, 0x62, 0x68, 0x65, 0xe9, 0x2f, 0x86, 0xde, 0x85, 0xea, 0x8b, 0xf4, 0x7c, - 0xbe, 0xb4, 0xf1, 0x30, 0xf2, 0x62, 0x07, 0xc1, 0xec, 0x3e, 0x74, 0x28, 0x7b, 0x24, 0xc9, 0xd1, - 0x32, 0x71, 0x9d, 0xf2, 0x1f, 0x4f, 0x11, 0xdc, 0x01, 0x24, 0x32, 0xce, 0xd8, 0x36, 0xb4, 0xfc, - 0x69, 0x18, 0x05, 0x29, 0x8f, 0x4d, 0xd9, 0x30, 0x5b, 0x9e, 0xb2, 0x93, 0xd3, 0xb0, 0x3f, 0x84, - 0x41, 0x58, 0xc4, 0xa5, 0x8a, 0x64, 0xe1, 0x9c, 0x72, 0x2b, 0x45, 0xae, 0x9c, 0xd5, 0x12, 0x39, - 0x59, 0xff, 0xe2, 0x5a, 0x73, 0xb3, 0x7c, 0xad, 0x59, 0xff, 0x39, 0x0c, 0x99, 0xe8, 0x56, 0x7e, - 0xaa, 0x45, 0x0b, 0x7d, 0xcb, 0xf8, 0x55, 0xed, 0xc5, 0x63, 0x80, 0xf5, 0x0a, 0x8c, 0x7f, 0xf5, - 0x21, 0xf4, 0xd1, 0x5f, 0x73, 0xb5, 0x9b, 0x87, 0x4a, 0x17, 0xcc, 0xbf, 0x30, 0x64, 0x72, 0xfa, - 0x10, 0x1d, 0x3d, 0x14, 0xc6, 0x9b, 0xd0, 0xb7, 0xdf, 0x62, 0x6a, 0xcb, 0x3a, 0x26, 0x99, 0x68, - 0xa0, 0xba, 0xba, 0x6c, 0x1b, 0xd6, 0xfd, 0xa9, 0x17, 0xc7, 0x3c, 0x72, 0x4f, 0xb2, 0xf1, 0xd8, - 0x5a, 0xd8, 0x2e, 0x85, 0x4d, 0xd7, 0x0c, 0xea, 0x01, 0x61, 0xc8, 0xd0, 0x8e, 0xa0, 0x17, 0x87, - 0x91, 0xce, 0x0d, 0x90, 0x37, 0xd1, 0x23, 0xca, 0x4e, 0x1c, 0x46, 0x94, 0x1c, 0x40, 0x9f, 0xe2, - 0x97, 0x30, 0xc8, 0xb2, 0x30, 0x90, 0xae, 0x12, 0xf6, 0x7f, 0x75, 0x4c, 0x84, 0xb9, 0x14, 0xb3, - 0x79, 0x9e, 0x85, 0xc1, 0xb1, 0x30, 0xff, 0xac, 0xd3, 0x23, 0x7a, 0xdb, 0x1c, 0xfd, 0x12, 0xba, - 0x65, 0xd9, 0x41, 0x59, 0xa4, 0x43, 0xf5, 0xe0, 0x1d, 0x06, 0xd0, 0x78, 0x2a, 0xd2, 0x99, 0x17, - 0x0d, 0x2a, 0xf8, 0xac, 0x2f, 0xfb, 0x0f, 0x56, 0x58, 0x17, 0x5a, 0xf6, 0x90, 0x38, 0xa8, 0x9a, - 0xb4, 0xcd, 0xcf, 0xa1, 0x65, 0xff, 0x2e, 0x88, 0xfe, 0x6a, 0x45, 0x04, 0x5c, 0x7b, 0xbd, 0xa6, - 0x86, 0x0f, 0x01, 0xe4, 0xf1, 0xda, 0x7f, 0x00, 0x5b, 0x29, 0xfe, 0x01, 0x6c, 0xf4, 0x47, 0xd0, - 0x2d, 0x4f, 0xd1, 0x06, 0x22, 0x2b, 0x45, 0x20, 0xf2, 0x82, 0x5e, 0x94, 0x5b, 0x4e, 0xc5, 0xcc, - 0x2d, 0x39, 0x66, 0x2d, 0x04, 0xe0, 0x6b, 0x46, 0xff, 0xb0, 0x02, 0x75, 0x3a, 0x3a, 0x91, 0x29, - 0xc6, 0x87, 0x62, 0x07, 0xd5, 0x9d, 0x36, 0x41, 0xfe, 0x1f, 0xee, 0x2c, 0xe5, 0x89, 0xa9, 0xda, - 0x5b, 0x13, 0x53, 0xb7, 0xff, 0xac, 0x02, 0x0d, 0xfd, 0xaf, 0x6b, 0x6c, 0x0d, 0x7a, 0xcf, 0xe3, - 0xd3, 0x58, 0xbc, 0x8a, 0x35, 0x60, 0xf0, 0x0e, 0x5b, 0x87, 0x55, 0xcb, 0x7b, 0xf3, 0xf7, 0x6e, - 0x83, 0x0a, 0x1b, 0x40, 0x97, 0x56, 0xd7, 0x42, 0x56, 0xd8, 0xbb, 0x30, 0x34, 0xd6, 0xf4, 0x21, - 0x2a, 0x63, 0xa1, 0xc2, 0xf1, 0xb9, 0xc5, 0x56, 0xd9, 0x2a, 0x74, 0x8e, 0x94, 0x48, 0x8e, 0x78, - 0x1c, 0x84, 0xf1, 0x64, 0x50, 0x63, 0x43, 0xd8, 0xb0, 0xa3, 0xea, 0x7f, 0x26, 0x7b, 0x1c, 0xc6, - 0xa1, 0x9c, 0x0e, 0xea, 0xec, 0x1a, 0x5c, 0xbe, 0x08, 0xb3, 0xeb, 0x9f, 0x0e, 0x1a, 0xb7, 0x3f, - 0x03, 0xb6, 0xfc, 0x47, 0x66, 0x38, 0xfa, 0x13, 0x3e, 0xf1, 0xfc, 0xf3, 0xbd, 0x48, 0x48, 0x14, - 0x8a, 0x1e, 0xb4, 0x8b, 0x5e, 0x95, 0xdb, 0x8f, 0xa1, 0xa1, 0xff, 0x79, 0xae, 0xf4, 0x7d, 0x1a, - 0x30, 0x78, 0x07, 0x3b, 0xa3, 0xc9, 0x09, 0xe3, 0xc9, 0x53, 0x7e, 0xa6, 0xb4, 0x22, 0x7c, 0xe2, - 0x49, 0x35, 0x58, 0x61, 0x7d, 0x00, 0xf3, 0x09, 0x8f, 0xe2, 0x60, 0x50, 0x7d, 0xb0, 0xf7, 0xe7, - 0xbf, 0x7b, 0xaf, 0xf2, 0x17, 0xbf, 0x7b, 0xaf, 0xf2, 0x5f, 0x7e, 0xf7, 0xde, 0x3b, 0x7f, 0xfa, - 0x97, 0xef, 0x55, 0xbe, 0xbe, 0x5f, 0xfa, 0x5f, 0x3d, 0x63, 0x89, 0xa8, 0x16, 0xe0, 0x6e, 0x6e, - 0x96, 0xee, 0x26, 0xa7, 0x93, 0xbb, 0xc9, 0xc9, 0x5d, 0x2b, 0xe7, 0x27, 0x0d, 0xfa, 0xbb, 0xbc, - 0x4f, 0xff, 0x6f, 0x00, 0x00, 0x00, 0xff, 0xff, 0x1b, 0xfe, 0xe5, 0x5d, 0xad, 0x4f, 0x00, 0x00, + 0xbc, 0xa0, 0xac, 0xe0, 0x7d, 0xa8, 0xc5, 0xb6, 0x92, 0x3d, 0xc7, 0xd3, 0x9d, 0x6c, 0x82, 0xb3, + 0xdb, 0x90, 0x5f, 0x80, 0x31, 0x2e, 0xc9, 0xc5, 0x17, 0x64, 0xb6, 0xa1, 0x9d, 0xff, 0xc7, 0x84, + 0xf1, 0x42, 0xd6, 0xb7, 0x8a, 0x7f, 0x9d, 0x38, 0xb2, 0x4f, 0x4e, 0x41, 0xf6, 0xe6, 0xe4, 0x78, + 0xe7, 0x47, 0x25, 0xc7, 0xcd, 0x11, 0xcd, 0x17, 0xb1, 0x54, 0x26, 0x38, 0xd6, 0x0c, 0xe5, 0x2e, + 0x36, 0xd9, 0xef, 0x41, 0x2f, 0xe5, 0xfe, 0x4b, 0x77, 0x2a, 0xc7, 0xfa, 0x15, 0xbd, 0xf2, 0x25, + 0xbc, 0xa9, 0x1c, 0x7f, 0x4d, 0x89, 0x7b, 0x73, 0x62, 0xea, 0x20, 0xed, 0xbe, 0x1c, 0xd3, 0xa8, + 0x9f, 0xc0, 0xea, 0x94, 0x4f, 0x8f, 0x79, 0x2a, 0x27, 0x61, 0x62, 0x55, 0xed, 0x32, 0x15, 0x18, + 0xf4, 0x0b, 0x84, 0x9e, 0xcb, 0xf0, 0x1f, 0x55, 0xa0, 0x85, 0x5c, 0x44, 0x59, 0x62, 0x0c, 0x6a, + 0x53, 0x3f, 0xc9, 0x8c, 0x23, 0x4c, 0xcf, 0xe6, 0x7f, 0x2b, 0xb4, 0x94, 0x98, 0xff, 0xad, 0xa0, + 0x35, 0xd4, 0x29, 0x1f, 0x7a, 0xd6, 0xf7, 0xc0, 0xcf, 0x28, 0x2a, 0xa8, 0x25, 0xc3, 0x36, 0xd9, + 0x25, 0x68, 0xf8, 0x31, 0x9d, 0x86, 0x75, 0xea, 0xac, 0xee, 0xc7, 0x78, 0x0a, 0xd6, 0xe0, 0xa2, + 0x1a, 0xba, 0xee, 0xc7, 0x7b, 0xc1, 0xe9, 0xf0, 0x3f, 0x54, 0x60, 0xf5, 0x20, 0x15, 0x3e, 0x97, + 0xf2, 0x09, 0x1a, 0x72, 0x4a, 0x53, 0xe0, 0x1b, 0x29, 0xaa, 0xab, 0x33, 0x02, 0xf4, 0x8c, 0x32, + 0xac, 0x43, 0x15, 0xf9, 0x71, 0xa3, 0xea, 0xb4, 0x09, 0x42, 0xa7, 0x8d, 0x1c, 0x5d, 0x4a, 0x7f, + 0x6b, 0x34, 0xc5, 0x83, 0x6f, 0xc2, 0x72, 0x91, 0x58, 0x29, 0x65, 0xea, 0x8b, 0x6b, 0xa6, 0x34, + 0xca, 0x75, 0xe8, 0x98, 0x7a, 0x0a, 0x1a, 0x46, 0x87, 0xf8, 0x41, 0x83, 0x0e, 0xcd, 0x2c, 0xb4, + 0x72, 0x20, 0xbc, 0x0e, 0xea, 0x6b, 0x75, 0x81, 0xe8, 0xe1, 0xdf, 0x87, 0xfe, 0x41, 0xca, 0x13, + 0x2f, 0xe5, 0x54, 0x5f, 0x41, 0x2c, 0xde, 0x80, 0x46, 0xc4, 0xe3, 0xb1, 0x49, 0xe9, 0x57, 0x1d, + 0xd3, 0xca, 0xff, 0x72, 0x64, 0xa9, 0xf4, 0x97, 0x23, 0xc8, 0xea, 0x94, 0x7b, 0xe6, 0x9f, 0x49, + 0xe8, 0x19, 0xb7, 0x20, 0x1e, 0x19, 0xf5, 0xb9, 0xa8, 0xe5, 0xe8, 0x86, 0xb9, 0xcb, 0x76, 0x1c, + 0xc6, 0x54, 0x9b, 0x46, 0x77, 0xd9, 0x1e, 0x84, 0xf1, 0xf0, 0x5f, 0xd6, 0xa1, 0x63, 0xf8, 0x49, + 0x2f, 0xd7, 0x6b, 0x59, 0xc9, 0xd7, 0xb2, 0x0f, 0x55, 0xf9, 0x22, 0x32, 0x8b, 0x8b, 0x8f, 0xec, + 0x33, 0xa8, 0x46, 0xe1, 0xd4, 0x1c, 0x71, 0xae, 0xcd, 0x98, 0xab, 0xd9, 0x55, 0x31, 0x82, 0x87, + 0xd4, 0xa8, 0x23, 0xe9, 0x1e, 0x2c, 0x8a, 0xb8, 0xe1, 0x24, 0x9a, 0x8e, 0x53, 0xdc, 0x47, 0xc8, + 0x23, 0xcf, 0xa7, 0x32, 0x03, 0xab, 0x1c, 0x7a, 0x4e, 0xdb, 0x40, 0xf6, 0x02, 0xf6, 0x39, 0xb4, + 0xf2, 0x40, 0xa5, 0x3d, 0xd4, 0xa8, 0xd3, 0x78, 0x6b, 0xf7, 0xe9, 0xd1, 0x69, 0x6c, 0x23, 0x91, + 0xe6, 0x65, 0x39, 0x25, 0xfb, 0x03, 0xe8, 0x4a, 0x2e, 0xa5, 0xbe, 0xa8, 0x38, 0x12, 0x46, 0x69, + 0x5c, 0x2a, 0x9f, 0x57, 0x08, 0x8b, 0x5f, 0x6d, 0xb7, 0x88, 0x2c, 0x40, 0xec, 0x6b, 0x58, 0xb6, + 0xfd, 0x23, 0x31, 0x1e, 0xe7, 0x81, 0xf4, 0x6b, 0x0b, 0x23, 0x3c, 0x21, 0x74, 0x69, 0x9c, 0x9e, + 0x2c, 0x23, 0xd8, 0x57, 0xb0, 0x9c, 0xe8, 0x35, 0x76, 0x4d, 0x29, 0x8f, 0x56, 0x3e, 0x57, 0x67, + 0xbc, 0xab, 0x19, 0x19, 0x28, 0xae, 0x12, 0x15, 0x70, 0xb9, 0x78, 0x69, 0x57, 0x67, 0x56, 0x66, + 0x2f, 0xed, 0x72, 0xd8, 0x30, 0x7f, 0x56, 0x31, 0x4a, 0x3d, 0xba, 0x7c, 0xa8, 0x2b, 0x38, 0x6c, + 0xcd, 0xdd, 0xdd, 0x85, 0x15, 0xc3, 0x17, 0x6e, 0xe9, 0xcb, 0xa9, 0x8f, 0x4d, 0x17, 0x2a, 0xf0, + 0x30, 0xc5, 0x16, 0xeb, 0xe9, 0x39, 0x28, 0xb6, 0x05, 0x6b, 0xe6, 0x35, 0xfc, 0x94, 0xfb, 0x99, + 0xb9, 0x77, 0x4d, 0x2a, 0xaa, 0xeb, 0xac, 0x6a, 0xd4, 0x23, 0x8b, 0xd9, 0x0b, 0xae, 0x7e, 0x05, + 0x57, 0x2e, 0x7c, 0xc5, 0xdb, 0xca, 0x07, 0x7a, 0xe5, 0x5b, 0x91, 0xff, 0xbb, 0x0a, 0x9d, 0xd2, + 0xd2, 0xd1, 0x9f, 0xe4, 0x48, 0x9e, 0xda, 0x22, 0x21, 0x7c, 0x46, 0xd8, 0x44, 0x48, 0x5b, 0xf3, + 0x42, 0xcf, 0x08, 0x4b, 0x45, 0x9e, 0xfb, 0xa7, 0x67, 0x64, 0xa8, 0x39, 0x8d, 0x9b, 0xe9, 0xd7, + 0xf4, 0x6d, 0xde, 0x02, 0xb8, 0x17, 0xd0, 0xbf, 0xe9, 0x78, 0xca, 0x3b, 0xf6, 0xa4, 0xad, 0xb6, + 0xca, 0xdb, 0xa8, 0xd5, 0x5e, 0xf2, 0x14, 0xe7, 0x62, 0x73, 0x9d, 0xa6, 0x89, 0x02, 0x4f, 0x86, + 0xe0, 0xb5, 0x88, 0x75, 0x9e, 0xb3, 0xeb, 0xb4, 0x10, 0xf0, 0xad, 0x88, 0xa9, 0x9b, 0x11, 0x6f, + 0x93, 0xaf, 0xb7, 0x4d, 0x54, 0xf7, 0x2f, 0x32, 0x8e, 0x6e, 0x78, 0x40, 0xf7, 0x35, 0xda, 0x4e, + 0x93, 0xda, 0xba, 0x84, 0x80, 0xce, 0x0b, 0xaf, 0xbc, 0x50, 0xd1, 0x3e, 0x12, 0x99, 0x32, 0x12, + 0xb0, 0x82, 0x88, 0x6f, 0xbc, 0x50, 0x1d, 0x69, 0x30, 0xbb, 0x6f, 0xae, 0x61, 0x95, 0x69, 0x5d, + 0x3c, 0xec, 0xe8, 0x28, 0x1f, 0x9b, 0xa3, 0x3f, 0xe4, 0xf4, 0x7f, 0x32, 0x53, 0x4f, 0xa5, 0xe1, + 0xa9, 0x88, 0xd1, 0xec, 0xab, 0xf0, 0x25, 0x2f, 0xfe, 0xc1, 0xa7, 0xe5, 0xac, 0xe5, 0xc8, 0xa7, + 0x84, 0xa3, 0xc4, 0xd0, 0x73, 0xd8, 0xe4, 0xa7, 0x49, 0x14, 0xfa, 0xe1, 0xdc, 0x05, 0x48, 0xd7, + 0xf7, 0xa4, 0x72, 0x53, 0xae, 0xb2, 0x34, 0x96, 0x14, 0xc0, 0x32, 0x55, 0x28, 0x1f, 0x5a, 0xfa, + 0xf2, 0xa5, 0xc8, 0x5d, 0x4f, 0x2a, 0x47, 0xd3, 0x3e, 0xcd, 0xa2, 0x08, 0x99, 0x90, 0x27, 0xa4, + 0x74, 0x2d, 0x4a, 0x53, 0xea, 0x54, 0xd4, 0xf0, 0xbf, 0x56, 0x60, 0x75, 0x61, 0xdb, 0xa1, 0xeb, + 0x8f, 0x5b, 0xce, 0xe6, 0xcf, 0xbb, 0x4e, 0x03, 0x9b, 0x7b, 0x01, 0x21, 0xd4, 0x54, 0xd9, 0xcc, + 0x39, 0x22, 0xd4, 0x14, 0x75, 0xca, 0x25, 0x68, 0xa8, 0x53, 0x5a, 0x72, 0xad, 0x39, 0xeb, 0xea, + 0x14, 0xd7, 0x7a, 0x07, 0xda, 0x91, 0x18, 0xbb, 0x11, 0x7f, 0xc9, 0xf5, 0xad, 0xf4, 0xe5, 0xed, + 0x8f, 0xde, 0xb0, 0xdf, 0xb7, 0x9e, 0x88, 0xf1, 0x13, 0xa4, 0x75, 0x5a, 0x91, 0x79, 0x1a, 0xfe, + 0x1a, 0x5a, 0x16, 0xca, 0xda, 0x50, 0x7f, 0xc8, 0x8f, 0xb3, 0x71, 0xff, 0x1d, 0xd6, 0x82, 0x1a, + 0xf6, 0xe8, 0x57, 0xf0, 0xe9, 0x1b, 0x2f, 0x8d, 0xfb, 0x4b, 0x88, 0x7e, 0x94, 0xa6, 0x22, 0xed, + 0x57, 0xf1, 0xf1, 0xc0, 0x8b, 0x43, 0xbf, 0x5f, 0xc3, 0xc7, 0xc7, 0x9e, 0xf2, 0xa2, 0x7e, 0x7d, + 0xf8, 0xdb, 0x3a, 0xb4, 0x0e, 0xcc, 0xdb, 0xd9, 0x43, 0xe8, 0xe5, 0x7f, 0x7e, 0x74, 0x7e, 0x0c, + 0xee, 0x60, 0xfe, 0x81, 0x62, 0x70, 0xdd, 0xa4, 0xd4, 0x9a, 0xff, 0x0b, 0xa5, 0xa5, 0x85, 0xbf, + 0x50, 0x7a, 0x17, 0xaa, 0x2f, 0xd2, 0xb3, 0xd9, 0xd2, 0xcd, 0x83, 0xc8, 0x8b, 0x1d, 0x04, 0xb3, + 0xfb, 0xd0, 0xa1, 0xec, 0x98, 0x24, 0x47, 0xd2, 0xc4, 0xad, 0xca, 0x7f, 0xac, 0x45, 0x70, 0x07, + 0x90, 0xc8, 0x38, 0x9b, 0x5b, 0xd0, 0xf2, 0x27, 0x61, 0x14, 0xa4, 0x3c, 0x36, 0x65, 0xd1, 0x6c, + 0x71, 0xca, 0x4e, 0x4e, 0xc3, 0xfe, 0x10, 0xfa, 0x61, 0x11, 0x77, 0x2b, 0x92, 0xa1, 0x33, 0xca, + 0xbb, 0x14, 0x99, 0x73, 0x56, 0x4a, 0xe4, 0xe4, 0xdd, 0x14, 0xd7, 0xb6, 0x9b, 0xe5, 0x6b, 0xdb, + 0xfa, 0xcf, 0x6f, 0xc8, 0x05, 0x69, 0xe5, 0xa7, 0x76, 0xf4, 0x40, 0x6e, 0x19, 0xbf, 0xb1, 0x3d, + 0x7f, 0xcc, 0xb1, 0x5e, 0x8f, 0xf1, 0x1f, 0x3f, 0x82, 0x65, 0xf4, 0x47, 0x5d, 0xed, 0xc6, 0xa2, + 0x51, 0x01, 0xf3, 0x2f, 0x13, 0x99, 0x9c, 0x3c, 0x44, 0x47, 0x16, 0x85, 0xf1, 0x26, 0x2c, 0xdb, + 0x6f, 0x31, 0xb5, 0x73, 0x1d, 0x93, 0x2c, 0x35, 0x50, 0x5d, 0x3d, 0xb7, 0x05, 0x6b, 0xfe, 0xc4, + 0x8b, 0x63, 0x1e, 0xb9, 0xc7, 0xd9, 0x68, 0x64, 0x3d, 0x88, 0x2e, 0x85, 0x85, 0x57, 0x0d, 0xea, + 0x01, 0x61, 0xc8, 0x91, 0x18, 0x42, 0x2f, 0x0e, 0x23, 0x9d, 0xfb, 0x20, 0x6f, 0xa9, 0x47, 0x94, + 0x9d, 0x38, 0x8c, 0x28, 0xf9, 0x81, 0x3e, 0xd3, 0xaf, 0xa0, 0x9f, 0x65, 0x61, 0x20, 0x5d, 0x25, + 0xec, 0xff, 0x06, 0x99, 0x08, 0x7a, 0x29, 0x26, 0xf5, 0x3c, 0x0b, 0x83, 0x23, 0x61, 0xfe, 0x39, + 0xa8, 0x47, 0xf4, 0xb6, 0x39, 0xfc, 0x15, 0x74, 0xcb, 0xb2, 0x83, 0xb2, 0x48, 0x41, 0x83, 0xfe, + 0x3b, 0x0c, 0xa0, 0xf1, 0x54, 0xa4, 0x53, 0x2f, 0xea, 0x57, 0xf0, 0x59, 0x2b, 0xf3, 0xfe, 0x12, + 0xeb, 0x42, 0xcb, 0x1e, 0x82, 0xfb, 0x55, 0x93, 0x96, 0xfa, 0x05, 0xb4, 0xec, 0xdf, 0x21, 0xd1, + 0x5f, 0xc9, 0x88, 0x80, 0x6b, 0xaf, 0xde, 0xd4, 0x28, 0x22, 0x80, 0x3c, 0x7a, 0xfb, 0x0f, 0x67, + 0x4b, 0xc5, 0x3f, 0x9c, 0x0d, 0xff, 0x08, 0xba, 0xe5, 0x29, 0xda, 0x40, 0x6b, 0xa5, 0x08, 0xb4, + 0x9e, 0xd3, 0x8b, 0x72, 0xe7, 0xa9, 0x98, 0xba, 0x25, 0xc7, 0xb3, 0x85, 0x00, 0x7c, 0xcd, 0xf0, + 0x1f, 0x56, 0xa0, 0x4e, 0x47, 0x43, 0x72, 0x35, 0xf0, 0xa1, 0xd8, 0x41, 0x75, 0xa7, 0x4d, 0x90, + 0xff, 0x87, 0x3b, 0x59, 0x79, 0xe2, 0xad, 0xf6, 0xc6, 0xc4, 0xdb, 0xed, 0x3f, 0xab, 0x40, 0x43, + 0xff, 0xab, 0x1c, 0x5b, 0x85, 0xde, 0xf3, 0xf8, 0x24, 0x16, 0xaf, 0x62, 0x0d, 0xe8, 0xbf, 0xc3, + 0xd6, 0x60, 0xc5, 0xf2, 0xde, 0xfc, 0x7d, 0x5d, 0xbf, 0xc2, 0xfa, 0xd0, 0xa5, 0xd5, 0xb5, 0x90, + 0x25, 0xf6, 0x2e, 0x0c, 0x8c, 0xb7, 0xf0, 0x10, 0x95, 0xb1, 0x50, 0xe1, 0xe8, 0xcc, 0x62, 0xab, + 0x6c, 0x05, 0x3a, 0x87, 0x4a, 0x24, 0x87, 0x3c, 0x0e, 0xc2, 0x78, 0xdc, 0xaf, 0xb1, 0x01, 0xac, + 0xdb, 0x51, 0xf5, 0x3f, 0xaf, 0x3d, 0x0e, 0xe3, 0x50, 0x4e, 0xfa, 0x75, 0x76, 0x0d, 0x2e, 0x9f, + 0x87, 0xd9, 0xf1, 0x4f, 0xfa, 0x8d, 0xdb, 0x9f, 0x03, 0x5b, 0xfc, 0xa3, 0x36, 0x1c, 0xfd, 0x09, + 0x1f, 0x7b, 0xfe, 0xd9, 0x6e, 0x24, 0x24, 0x0a, 0x45, 0x0f, 0xda, 0x45, 0xaf, 0xca, 0xed, 0xc7, + 0xd0, 0xd0, 0xff, 0xac, 0x57, 0xfa, 0x3e, 0x0d, 0xe8, 0xbf, 0x83, 0x9d, 0xd1, 0xe4, 0x84, 0xf1, + 0xf8, 0x29, 0x3f, 0x55, 0x5a, 0x11, 0x3e, 0xf1, 0xa4, 0xea, 0x2f, 0xb1, 0x65, 0x00, 0xf3, 0x09, + 0x8f, 0xe2, 0xa0, 0x5f, 0x7d, 0xb0, 0xfb, 0xe7, 0xbf, 0x7b, 0xbf, 0xf2, 0x17, 0xbf, 0x7b, 0xbf, + 0xf2, 0x5f, 0x7e, 0xf7, 0xfe, 0x3b, 0x7f, 0xfa, 0x97, 0xef, 0x57, 0xbe, 0xbd, 0x5f, 0xfa, 0xdf, + 0x40, 0x63, 0x89, 0xa8, 0xd6, 0xe1, 0x6e, 0x6e, 0x96, 0xee, 0x26, 0x27, 0xe3, 0xbb, 0xc9, 0xf1, + 0x5d, 0x2b, 0xe7, 0xc7, 0x0d, 0xfa, 0x3b, 0xc0, 0xcf, 0xfe, 0x6f, 0x00, 0x00, 0x00, 0xff, 0xff, + 0x8e, 0x30, 0xc6, 0xcd, 0x8d, 0x50, 0x00, 0x00, } func (m *Message) Marshal() (dAtA []byte, err error) { @@ -11978,6 +12004,30 @@ func (m *ProcessInfo) MarshalToSizedBuffer(dAtA []byte) (int, error) { i -= len(m.XXX_unrecognized) copy(dAtA[i:], m.XXX_unrecognized) } + if len(m.RemoteExecutionId) > 0 { + i -= len(m.RemoteExecutionId) + copy(dAtA[i:], m.RemoteExecutionId) + i = encodeVarintPipeline(dAtA, i, uint64(len(m.RemoteExecutionId))) + i-- + dAtA[i] = 0x62 + } + if len(m.RemoteFragmentCounts) > 0 { + for k := range m.RemoteFragmentCounts { + v := m.RemoteFragmentCounts[k] + baseI := i + i = encodeVarintPipeline(dAtA, i, uint64(v)) + i-- + dAtA[i] = 0x10 + i -= len(k) + copy(dAtA[i:], k) + i = encodeVarintPipeline(dAtA, i, uint64(len(k))) + i-- + dAtA[i] = 0xa + i = encodeVarintPipeline(dAtA, i, uint64(baseI-i)) + i-- + dAtA[i] = 0x5a + } + } if m.AffectedRows != 0 { i = encodeVarintPipeline(dAtA, i, uint64(m.AffectedRows)) i-- @@ -14884,6 +14934,18 @@ func (m *ProcessInfo) ProtoSize() (n int) { if m.AffectedRows != 0 { n += 1 + sovPipeline(uint64(m.AffectedRows)) } + if len(m.RemoteFragmentCounts) > 0 { + for k, v := range m.RemoteFragmentCounts { + _ = k + _ = v + mapEntrySize := 1 + len(k) + sovPipeline(uint64(len(k))) + 1 + sovPipeline(uint64(v)) + n += mapEntrySize + 1 + sovPipeline(uint64(mapEntrySize)) + } + } + l = len(m.RemoteExecutionId) + if l > 0 { + n += 1 + l + sovPipeline(uint64(l)) + } if m.XXX_unrecognized != nil { n += len(m.XXX_unrecognized) } @@ -30669,6 +30731,153 @@ func (m *ProcessInfo) Unmarshal(dAtA []byte) error { break } } + case 11: + if wireType != 2 { + return fmt.Errorf("proto: wrong wireType = %d for field RemoteFragmentCounts", wireType) + } + var msglen int + for shift := uint(0); ; shift += 7 { + if shift >= 64 { + return ErrIntOverflowPipeline + } + if iNdEx >= l { + return io.ErrUnexpectedEOF + } + b := dAtA[iNdEx] + iNdEx++ + msglen |= int(b&0x7F) << shift + if b < 0x80 { + break + } + } + if msglen < 0 { + return ErrInvalidLengthPipeline + } + postIndex := iNdEx + msglen + if postIndex < 0 { + return ErrInvalidLengthPipeline + } + if postIndex > l { + return io.ErrUnexpectedEOF + } + if m.RemoteFragmentCounts == nil { + m.RemoteFragmentCounts = make(map[string]uint32) + } + var mapkey string + var mapvalue uint32 + for iNdEx < postIndex { + entryPreIndex := iNdEx + var wire uint64 + for shift := uint(0); ; shift += 7 { + if shift >= 64 { + return ErrIntOverflowPipeline + } + if iNdEx >= l { + return io.ErrUnexpectedEOF + } + b := dAtA[iNdEx] + iNdEx++ + wire |= uint64(b&0x7F) << shift + if b < 0x80 { + break + } + } + fieldNum := int32(wire >> 3) + if fieldNum == 1 { + var stringLenmapkey uint64 + for shift := uint(0); ; shift += 7 { + if shift >= 64 { + return ErrIntOverflowPipeline + } + if iNdEx >= l { + return io.ErrUnexpectedEOF + } + b := dAtA[iNdEx] + iNdEx++ + stringLenmapkey |= uint64(b&0x7F) << shift + if b < 0x80 { + break + } + } + intStringLenmapkey := int(stringLenmapkey) + if intStringLenmapkey < 0 { + return ErrInvalidLengthPipeline + } + postStringIndexmapkey := iNdEx + intStringLenmapkey + if postStringIndexmapkey < 0 { + return ErrInvalidLengthPipeline + } + if postStringIndexmapkey > l { + return io.ErrUnexpectedEOF + } + mapkey = string(dAtA[iNdEx:postStringIndexmapkey]) + iNdEx = postStringIndexmapkey + } else if fieldNum == 2 { + for shift := uint(0); ; shift += 7 { + if shift >= 64 { + return ErrIntOverflowPipeline + } + if iNdEx >= l { + return io.ErrUnexpectedEOF + } + b := dAtA[iNdEx] + iNdEx++ + mapvalue |= uint32(b&0x7F) << shift + if b < 0x80 { + break + } + } + } else { + iNdEx = entryPreIndex + skippy, err := skipPipeline(dAtA[iNdEx:]) + if err != nil { + return err + } + if (skippy < 0) || (iNdEx+skippy) < 0 { + return ErrInvalidLengthPipeline + } + if (iNdEx + skippy) > postIndex { + return io.ErrUnexpectedEOF + } + iNdEx += skippy + } + } + m.RemoteFragmentCounts[mapkey] = mapvalue + iNdEx = postIndex + case 12: + if wireType != 2 { + return fmt.Errorf("proto: wrong wireType = %d for field RemoteExecutionId", wireType) + } + var byteLen int + for shift := uint(0); ; shift += 7 { + if shift >= 64 { + return ErrIntOverflowPipeline + } + if iNdEx >= l { + return io.ErrUnexpectedEOF + } + b := dAtA[iNdEx] + iNdEx++ + byteLen |= int(b&0x7F) << shift + if b < 0x80 { + break + } + } + if byteLen < 0 { + return ErrInvalidLengthPipeline + } + postIndex := iNdEx + byteLen + if postIndex < 0 { + return ErrInvalidLengthPipeline + } + if postIndex > l { + return io.ErrUnexpectedEOF + } + m.RemoteExecutionId = append(m.RemoteExecutionId[:0], dAtA[iNdEx:postIndex]...) + if m.RemoteExecutionId == nil { + m.RemoteExecutionId = []byte{} + } + iNdEx = postIndex default: iNdEx = preIndex skippy, err := skipPipeline(dAtA[iNdEx:]) diff --git a/pkg/sql/compile/allocation_account_lifecycle.go b/pkg/sql/compile/allocation_account_lifecycle.go index 4b34b1452977c..e3281f7ebc599 100644 --- a/pkg/sql/compile/allocation_account_lifecycle.go +++ b/pkg/sql/compile/allocation_account_lifecycle.go @@ -16,6 +16,7 @@ package compile import ( "errors" + "fmt" "sync" "github.com/matrixorigin/matrixone/pkg/common/mpool" @@ -24,6 +25,19 @@ import ( "github.com/matrixorigin/matrixone/pkg/vm/process" ) +func allocationLifecycleCall(call func() error) (err error) { + defer func() { + if recovered := recover(); recovered != nil { + err = errors.Join( + err, + mpool.ErrAllocationAccountInvariant, + fmt.Errorf("allocation lifecycle panic: %v", recovered), + ) + } + }() + return call() +} + type executionAllocationAccountOwner interface { SetAllocationAccount(*mpool.AllocationAccount) error ClearAllocationAccount(*mpool.AllocationAccount) error @@ -43,9 +57,11 @@ type statementAllocationAttempt struct { ownerSet map[executionAllocationAccountOwner]struct{} closing bool - once sync.Once - snapshot mpool.AllocationAccountTerminalSnapshot - err error + prepareOnce sync.Once + completeOnce sync.Once + snapshot mpool.AllocationAccountTerminalSnapshot + prepareErr error + completeErr error } func (c *Compile) beginAllocationAccountAttempt() ( @@ -87,13 +103,25 @@ func (c *Compile) beginAllocationAccountAttempt() ( } owners, err = configureAllocationAccountOwners(owners, account) if err != nil { - snapshot, first, finalizeErr := c.allocationAccountRegistry. - CompleteTerminalWithError(account, err) + var snapshot mpool.AllocationAccountTerminalSnapshot + var first bool + finalizeErr := allocationLifecycleCall(func() error { + var terminalErr error + snapshot, first, terminalErr = c.allocationAccountRegistry. + CompleteTerminalWithError(account, err) + return terminalErr + }) if first { - c.allocationTerminalExporter(snapshot) + finalizeErr = errors.Join( + finalizeErr, + allocationLifecycleCall(func() error { + c.allocationTerminalExporter(snapshot) + return nil + }), + ) } if finalizeErr != nil { - return nil, finalizeErr + return nil, errors.Join(err, finalizeErr) } return nil, err } @@ -169,13 +197,17 @@ func configureAllocationAccountOwners( for i := len(configured) - 1; i >= 0; i-- { cause = errors.Join( cause, - configured[i].ClearAllocationAccount(account), + allocationLifecycleCall(func() error { + return configured[i].ClearAllocationAccount(account) + }), ) } return cause } for _, owner := range owners { - if err := owner.SetAllocationAccount(account); err != nil { + if err := allocationLifecycleCall(func() error { + return owner.SetAllocationAccount(account) + }); err != nil { return nil, rollback(err) } configured = append(configured, owner) @@ -280,11 +312,28 @@ func (a *statementAllocationAttempt) finish() ( if a == nil { return mpool.AllocationAccountTerminalSnapshot{}, nil } - a.once.Do(func() { - // Scope.Run/MergeRun and remote notifier barriers must have returned - // before this point. Draining the board first releases queued JoinMap - // and spill payload ownership through their normal Destroy methods. - a.board.CloseAndDrain() + a.prepareTerminal(true) + return a.completeTerminal() +} + +// prepareTerminal closes the operator-owned part of an attempt after every +// scope producer has quiesced. A coordinator-owned board can be closed here. +// Remote fragments share one board on a CN, so their statement group closes it +// only after every expected fragment has reached this boundary. +func (a *statementAllocationAttempt) prepareTerminal(closeBoard bool) error { + if a == nil { + return nil + } + a.prepareOnce.Do(func() { + if closeBoard { + a.prepareErr = errors.Join( + a.prepareErr, + allocationLifecycleCall(func() error { + a.board.CloseAndDrain() + return nil + }), + ) + } a.ownersMu.Lock() a.closing = true owners := a.owners @@ -292,23 +341,47 @@ func (a *statementAllocationAttempt) finish() ( a.ownerSet = nil a.ownersMu.Unlock() for i := len(owners) - 1; i >= 0; i-- { - a.err = errors.Join( - a.err, - owners[i].ClearAllocationAccount(a.account), + a.prepareErr = errors.Join( + a.prepareErr, + allocationLifecycleCall(func() error { + return owners[i].ClearAllocationAccount(a.account) + }), ) } + }) + return a.prepareErr + +} + +func (a *statementAllocationAttempt) completeTerminal() ( + mpool.AllocationAccountTerminalSnapshot, + error, +) { + if a == nil { + return mpool.AllocationAccountTerminalSnapshot{}, nil + } + a.completeOnce.Do(func() { + prepareErr := a.prepareTerminal(false) var first bool - var terminalErr error - a.snapshot, first, terminalErr = a.registry.CompleteTerminalWithError( - a.account, - a.err, - ) - a.err = terminalErr + a.completeErr = allocationLifecycleCall(func() error { + var terminalErr error + a.snapshot, first, terminalErr = a.registry.CompleteTerminalWithError( + a.account, + prepareErr, + ) + return terminalErr + }) if first && a.exporter != nil { - a.exporter(a.snapshot) + a.completeErr = errors.Join( + a.completeErr, + allocationLifecycleCall(func() error { + a.exporter(a.snapshot) + return nil + }), + ) } }) - return a.snapshot, a.err + return a.snapshot, a.completeErr } func (c *Compile) finishAllocationAccountAttempt() error { diff --git a/pkg/sql/compile/allocation_account_lifecycle_test.go b/pkg/sql/compile/allocation_account_lifecycle_test.go index 2d39ebfa133a4..5b6462bf83e5c 100644 --- a/pkg/sql/compile/allocation_account_lifecycle_test.go +++ b/pkg/sql/compile/allocation_account_lifecycle_test.go @@ -49,6 +49,7 @@ type allocationLifecycleOwnerOperator struct { account *mpool.AllocationAccount failSet bool failClear bool + panicClear bool clears int released bool releaseSawLiveAccount bool @@ -75,6 +76,9 @@ func (op *allocationLifecycleOwnerOperator) SetAllocationAccount( func (op *allocationLifecycleOwnerOperator) ClearAllocationAccount( account *mpool.AllocationAccount, ) error { + if op.panicClear { + panic("test allocation owner clear panic") + } if op.account == nil { return nil } @@ -89,6 +93,34 @@ func (op *allocationLifecycleOwnerOperator) ClearAllocationAccount( return nil } +func TestStatementAllocationAttemptOwnerTeardownPanicIsTerminalFailure(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + var exported []mpool.AllocationAccountTerminalSnapshot + c := newTestAllocationLifecycleCompile(t, registry, func( + snapshot mpool.AllocationAccountTerminalSnapshot, + ) { + exported = append(exported, snapshot) + }) + owner := &allocationLifecycleOwnerOperator{ + MockOperator: colexec.NewMockOperator(), + panicClear: true, + } + c.scopes = []*Scope{{RootOp: owner}} + + _, err = c.beginAllocationAccountAttempt() + require.NoError(t, err) + err = c.finishAllocationAccountAttempt() + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + require.Len(t, exported, 1) + require.Equal( + t, + mpool.AllocationAccountTerminalInvariantFailure, + exported[0].State, + ) + require.Zero(t, registry.LiveAllocationMetadata()) +} + func (op *allocationLifecycleErrorOperator) Call( *process.Process, ) (vm.CallResult, error) { diff --git a/pkg/sql/compile/analyze_module.go b/pkg/sql/compile/analyze_module.go index f7f39436eea30..cff5353e893d9 100644 --- a/pkg/sql/compile/analyze_module.go +++ b/pkg/sql/compile/analyze_module.go @@ -17,6 +17,7 @@ package compile import ( "bytes" "fmt" + "slices" "strings" "sync" @@ -44,16 +45,18 @@ type AnalyzeModule struct { phyPlan *models.PhyPlan remotePhyPlans []models.PhyPlan // Added read-write lock - mu sync.RWMutex - retryTimes int - explainPhyBuffer *bytes.Buffer - remoteUsage resource.Usage - remoteMemory resource.MemoryTotals - remoteAllocation resource.AllocationAccountTotals - remoteQuality resource.QualityFlags - remoteMissingFragments uint64 - remoteMissingMemoryDomains uint64 - remoteReports uint64 + mu sync.RWMutex + retryTimes int + explainPhyBuffer *bytes.Buffer + remoteUsage resource.Usage + remoteMemory resource.MemoryTotals + remoteAllocation resource.AllocationAccountTotals + remoteQuality resource.QualityFlags + remoteMissingFragments uint64 + remoteMissingMemoryDomains uint64 + remoteReports uint64 + remotePendingAllocationGroups map[string]uint64 + remoteCompletedAllocationGroups map[string]struct{} } // remoteResourceSnapshot is a by-value view of the terminal resource facts @@ -61,13 +64,15 @@ type AnalyzeModule struct { // separately from the nested missing counts so each hop can account for the // remote scopes it expected to hear from exactly once. type remoteResourceSnapshot struct { - Usage resource.Usage - Memory resource.MemoryTotals - Allocation resource.AllocationAccountTotals - Quality resource.QualityFlags - MissingFragmentCount uint64 - MissingMemoryDomainCount uint64 - DirectReportCount uint64 + Usage resource.Usage + Memory resource.MemoryTotals + Allocation resource.AllocationAccountTotals + Quality resource.QualityFlags + MissingFragmentCount uint64 + MissingMemoryDomainCount uint64 + DirectReportCount uint64 + PendingAllocationGroups []remoteAllocationGroupPending + CompletedAllocationGroups []string } // Reset When Compile reused, reset AnalyzeModule to prevent resource accumulation @@ -89,6 +94,8 @@ func (anal *AnalyzeModule) Reset(isPrepare bool, isTpQuery bool) { anal.remoteMissingFragments = 0 anal.remoteMissingMemoryDomains = 0 anal.remoteReports = 0 + anal.remotePendingAllocationGroups = nil + anal.remoteCompletedAllocationGroups = nil if anal.qry != nil { for _, node := range anal.qry.Nodes { if node.AnalyzeInfo == nil { @@ -108,6 +115,8 @@ func (anal *AnalyzeModule) appendRemoteResource( allocation resource.AllocationAccountTotals, missingFragments uint64, missingMemoryDomains uint64, + pendingAllocationGroups []remoteAllocationGroupPending, + completedAllocationGroups []string, ) { if anal == nil { return @@ -126,9 +135,68 @@ func (anal *AnalyzeModule) appendRemoteResource( anal.remoteMissingMemoryDomains, quality = addCheckedRemoteCounter( anal.remoteMissingMemoryDomains, missingMemoryDomains, quality) anal.remoteReports, quality = addCheckedRemoteCounter(anal.remoteReports, 1, quality) + if len(completedAllocationGroups) > 0 && anal.remoteCompletedAllocationGroups == nil { + anal.remoteCompletedAllocationGroups = make(map[string]struct{}) + } + for _, key := range completedAllocationGroups { + if key == "" { + quality |= resource.QualityInvariantFailure + continue + } + anal.remoteCompletedAllocationGroups[key] = struct{}{} + delete(anal.remotePendingAllocationGroups, key) + } + if len(pendingAllocationGroups) > 0 && anal.remotePendingAllocationGroups == nil { + anal.remotePendingAllocationGroups = make(map[string]uint64) + } + for _, signal := range pendingAllocationGroups { + if signal.Key == "" || signal.Count == 0 { + quality |= resource.QualityInvariantFailure + continue + } + if _, completed := anal.remoteCompletedAllocationGroups[signal.Key]; !completed { + anal.remotePendingAllocationGroups[signal.Key], quality = + addCheckedRemoteCounter( + anal.remotePendingAllocationGroups[signal.Key], + signal.Count, + quality, + ) + } + } anal.remoteQuality |= quality } +func sortedRemoteAllocationGroupKeys(values map[string]struct{}) []string { + if len(values) == 0 { + return nil + } + keys := make([]string, 0, len(values)) + for key := range values { + keys = append(keys, key) + } + slices.Sort(keys) + return keys +} + +func sortedRemoteAllocationGroupPending( + values map[string]uint64, +) []remoteAllocationGroupPending { + if len(values) == 0 { + return nil + } + signals := make([]remoteAllocationGroupPending, 0, len(values)) + for key, count := range values { + signals = append(signals, remoteAllocationGroupPending{ + Key: key, + Count: count, + }) + } + slices.SortFunc(signals, func(a, b remoteAllocationGroupPending) int { + return strings.Compare(a.Key, b.Key) + }) + return signals +} + func (anal *AnalyzeModule) remoteResourceSummary() remoteResourceSnapshot { if anal == nil { return remoteResourceSnapshot{} @@ -143,6 +211,12 @@ func (anal *AnalyzeModule) remoteResourceSummary() remoteResourceSnapshot { MissingFragmentCount: anal.remoteMissingFragments, MissingMemoryDomainCount: anal.remoteMissingMemoryDomains, DirectReportCount: anal.remoteReports, + PendingAllocationGroups: sortedRemoteAllocationGroupPending( + anal.remotePendingAllocationGroups, + ), + CompletedAllocationGroups: sortedRemoteAllocationGroupKeys( + anal.remoteCompletedAllocationGroups, + ), } } diff --git a/pkg/sql/compile/compile.go b/pkg/sql/compile/compile.go index aa349b82213cc..3c8707ae45dc7 100644 --- a/pkg/sql/compile/compile.go +++ b/pkg/sql/compile/compile.go @@ -28,6 +28,7 @@ import ( "strings" "time" + "github.com/google/uuid" "github.com/parquet-go/parquet-go" "github.com/matrixorigin/matrixone/pkg/catalog" @@ -240,6 +241,8 @@ func (c *Compile) Reset(proc *process.Process, startAt time.Time, fill func(*bat c.MessageBoard = c.MessageBoard.Reset() proc.SetMessageBoard(c.MessageBoard) + c.remoteFragmentCounts = nil + c.remoteExecutionID = uuid.Nil c.counterSet.Reset() for _, f := range c.fuzzys { @@ -329,6 +332,8 @@ func (c *Compile) clear() { c.allocationTerminalExporter = nil c.allocationAccountOwners = nil c.allocationAttempt = nil + c.remoteFragmentCounts = nil + c.remoteExecutionID = uuid.Nil c.isPrepare = false c.hasMergeOp = false c.needBlock = false diff --git a/pkg/sql/compile/compile2.go b/pkg/sql/compile/compile2.go index 30bd84bdf65b8..4def9d28eb934 100644 --- a/pkg/sql/compile/compile2.go +++ b/pkg/sql/compile/compile2.go @@ -23,6 +23,7 @@ import ( "strings" "time" + "github.com/google/uuid" "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/common/mpool" commonutil "github.com/matrixorigin/matrixone/pkg/common/util" @@ -304,6 +305,17 @@ func (c *Compile) Run(_ uint64) (queryResult *util2.RunResult, err error) { // Before compile.runOnce, Reset the 'StatsInfo' execution related resources in context // running. + if runC.remoteFragmentCounts == nil { + runC.remoteFragmentCounts = collectRemoteFragmentCounts(runC.scopes, runC.addr) + } + // A retry is a new physical execution generation. Reusing the previous + // ID could attach late RPCs from the failed generation to the new + // generation's shared board and terminal-account group. + if len(runC.remoteFragmentCounts) > 0 { + runC.remoteExecutionID = newRemoteExecutionID() + } else { + runC.remoteExecutionID = uuid.Nil + } exporter := func(snapshot mpool.AllocationAccountTerminalSnapshot) { if resourceRecorder != nil { resourceRecorder.recordAllocationAccountTerminal(snapshot) diff --git a/pkg/sql/compile/remote_allocation_statement_group.go b/pkg/sql/compile/remote_allocation_statement_group.go new file mode 100644 index 0000000000000..46f43d7099b06 --- /dev/null +++ b/pkg/sql/compile/remote_allocation_statement_group.go @@ -0,0 +1,466 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package compile + +import ( + "errors" + "sync" + "time" + + "github.com/google/uuid" + "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/logutil" + "github.com/matrixorigin/matrixone/pkg/util/resource" + "github.com/matrixorigin/matrixone/pkg/vm/message" + "go.uber.org/zap" +) + +func newRemoteExecutionID() uuid.UUID { + return uuid.New() +} + +func remoteMessageBoardID( + statementID uuid.UUID, + remoteExecutionID uuid.UUID, +) uuid.UUID { + if remoteExecutionID != uuid.Nil { + return remoteExecutionID + } + return statementID +} + +func remoteAllocationStatementGroupKey( + remoteExecutionID uuid.UUID, + address string, +) string { + if remoteExecutionID == uuid.Nil || address == "" { + return "" + } + return remoteExecutionID.String() + "@" + address +} + +// A missing planned RPC means the coordinator failed before dispatch +// completed. Bound the orphan lifetime by the same five-minute interval used +// by blocking MessageBoard receives. The timer is canceled as soon as every +// expected fragment has registered, so it never limits a fully dispatched +// statement's execution time. +var remoteAllocationStatementRegistrationTimeout = 5 * time.Minute + +// collectRemoteFragmentCounts computes the number of pipeline RPCs that the +// complete physical scope graph will send to each CN. The execution address +// changes when traversal crosses a Remote scope: nested scopes targeting that +// same address execute inside the received pipeline and do not create another +// RPC, while a different target does. +func collectRemoteFragmentCounts( + scopes []*Scope, + rootAddress string, +) map[string]uint32 { + counts := make(map[string]uint32) + var visit func(*Scope, string) + visit = func(scope *Scope, executionAddress string) { + if scope == nil { + return + } + if scope.Magic == Remote && !scope.ipAddrMatch(executionAddress) { + target := scope.NodeInfo.Addr + counts[target]++ + executionAddress = target + } + for _, pre := range scope.PreScopes { + visit(pre, executionAddress) + } + } + for _, scope := range scopes { + visit(scope, rootAddress) + } + return counts +} + +func validateRemoteAllocationTopologyCapability( + scopes []*Scope, + remoteFragmentCounts map[string]uint32, +) error { + if len(remoteFragmentCounts) > 0 { + return nil + } + owners, err := collectAllocationAccountOwners(scopes) + if err != nil || len(owners) == 0 { + return err + } + return moerr.NewNotSupportedNoCtx( + "remote allocation-accounted execution requires fragment topology metadata", + ) +} + +var remoteAllocationStatementGroups = struct { + sync.Mutex + byBoard map[*message.MessageBoard]*remoteAllocationStatementGroup +}{ + byBoard: make(map[*message.MessageBoard]*remoteAllocationStatementGroup), +} + +// remoteAllocationStatementGroup is the terminal owner for all pipeline RPCs +// of one statement that execute on one CN. Those RPCs share a MessageBoard, so +// no individual fragment may close it or validate transferred allocations +// while a sibling can still consume them. +type remoteAllocationStatementGroup struct { + board *message.MessageBoard + expected uint32 + registered uint32 + finished uint32 + attempts []*statementAllocationAttempt + pools []*mpool.MPool + participants []*remoteAllocationStatementParticipant + timer *time.Timer + expired bool + finalized bool + err error +} + +type remoteAllocationStatementParticipant struct { + group *remoteAllocationStatementGroup + cancel func(error) + finished bool + + stageOnce sync.Once + finishOnce sync.Once + terminal remoteAllocationStatementTerminal + err error +} + +type remoteAllocationStatementTerminal struct { + allocation []mpool.AllocationAccountTerminalSnapshot + memory resource.MemoryTotals + quality resource.QualityFlags + complete bool +} + +func acquireRemoteAllocationStatementParticipant( + board *message.MessageBoard, + expected uint32, + cancel func(error), +) (*remoteAllocationStatementParticipant, error) { + if board == nil { + return nil, mpool.ErrAllocationAccountInvariant + } + if expected == 0 { + return nil, mpool.ErrAllocationAccountInvariant + } + + remoteAllocationStatementGroups.Lock() + defer remoteAllocationStatementGroups.Unlock() + group := remoteAllocationStatementGroups.byBoard[board] + if group == nil { + group = &remoteAllocationStatementGroup{ + board: board, + expected: expected, + } + remoteAllocationStatementGroups.byBoard[board] = group + } + if group.expected != expected || group.finalized || + group.expired || group.registered >= group.expected { + return nil, errors.Join( + mpool.ErrAllocationAccountInvariant, + moerr.NewInternalErrorNoCtx("invalid remote allocation statement group registration"), + ) + } + participant := &remoteAllocationStatementParticipant{ + group: group, + cancel: cancel, + } + group.registered++ + group.participants = append(group.participants, participant) + if group.registered == group.expected && group.timer != nil { + group.timer.Stop() + group.timer = nil + } + return participant, nil +} + +// stage clears fragment-local operators while they are still reachable and +// retains the fragment MPool for a statement-boundary snapshot. Account and +// allocator completion are deferred until every expected fragment has +// finished. +func (p *remoteAllocationStatementParticipant) stage( + attempt *statementAllocationAttempt, + pool *mpool.MPool, +) { + if p == nil || p.group == nil { + return + } + p.stageOnce.Do(func() { + if attempt != nil { + attempt.exporter = nil + } + + remoteAllocationStatementGroups.Lock() + if p.group.finalized { + p.err = mpool.ErrAllocationAccountInvariant + remoteAllocationStatementGroups.Unlock() + return + } + if attempt != nil { + p.group.attempts = append(p.group.attempts, attempt) + } + if pool != nil { + p.group.pools = append(p.group.pools, pool) + } + remoteAllocationStatementGroups.Unlock() + + // Transfer terminal ownership before clearing operators. If a cleanup + // hook panics, the handler defer can still finish this participant and + // the group retains the account and allocator domain. + if attempt != nil { + _ = attempt.prepareTerminal(false) + } + }) +} + +// finish marks one remote fragment quiescent without imposing a cross-RPC +// response barrier. Nested remote paths can re-enter the same CN, so waiting +// for sibling completion here would create a B -> C -> B dependency cycle. +// The fragment that completes the group publishes the aggregate exactly once; +// earlier responses carry no duplicate terminal totals. +func (p *remoteAllocationStatementParticipant) finish(cause error) ( + remoteAllocationStatementTerminal, + error, +) { + if p == nil || p.group == nil { + return remoteAllocationStatementTerminal{}, nil + } + p.stage(nil, nil) + p.finishOnce.Do(func() { + remoteAllocationStatementGroups.Lock() + group := p.group + if group.finalized || + (!group.expired && group.finished >= group.expected) || + (group.expired && group.finished >= group.registered) { + p.err = errors.Join(p.err, mpool.ErrAllocationAccountInvariant) + remoteAllocationStatementGroups.Unlock() + return + } + abort := cause != nil && !group.expired + if cause != nil { + group.err = errors.Join(group.err, cause) + group.expired = true + if group.timer != nil { + group.timer.Stop() + group.timer = nil + } + } + group.finished++ + p.finished = true + var cancels []func(error) + if abort { + cancels = activeRemoteAllocationStatementCancelsLocked(group) + } + if !group.expired && group.registered < group.expected && group.timer == nil { + group.timer = time.AfterFunc( + remoteAllocationStatementRegistrationTimeout, + func() { expireRemoteAllocationStatementGroup(group) }, + ) + } + complete := group.expired && group.finished == group.registered || + group.registered == group.expected && group.finished == group.expected + if !complete { + remoteAllocationStatementGroups.Unlock() + if abort { + abortErr := allocationLifecycleCall(func() error { + group.board.CloseAndDrain() + return nil + }) + abortErr = errors.Join( + abortErr, + cancelRemoteAllocationStatementParticipants(cancels, cause), + ) + if abortErr != nil { + remoteAllocationStatementGroups.Lock() + group.err = errors.Join(group.err, abortErr) + remoteAllocationStatementGroups.Unlock() + p.err = errors.Join(p.err, abortErr) + } + } + } else { + attempts, pools := takeRemoteAllocationStatementGroupLocked(group) + terminalErr := group.err + remoteAllocationStatementGroups.Unlock() + defer releaseRemoteAllocationStatementGroup(group) + p.terminal, terminalErr = completeRemoteAllocationStatementGroup( + group, + attempts, + pools, + terminalErr, + ) + p.err = errors.Join(p.err, terminalErr) + } + }) + return p.terminal, p.err +} + +func releaseRemoteAllocationStatementGroup(group *remoteAllocationStatementGroup) { + remoteAllocationStatementGroups.Lock() + if remoteAllocationStatementGroups.byBoard[group.board] == group { + delete(remoteAllocationStatementGroups.byBoard, group.board) + } + remoteAllocationStatementGroups.Unlock() +} + +func cancelRemoteAllocationStatementParticipants( + cancels []func(error), + cause error, +) error { + var err error + for _, cancel := range cancels { + if cancel != nil { + err = errors.Join(err, allocationLifecycleCall(func() error { + cancel(cause) + return nil + })) + } + } + return err +} + +func activeRemoteAllocationStatementCancelsLocked( + group *remoteAllocationStatementGroup, +) []func(error) { + cancels := make([]func(error), 0, len(group.participants)) + for _, participant := range group.participants { + if participant != nil && !participant.finished && participant.cancel != nil { + cancels = append(cancels, participant.cancel) + } + } + return cancels +} + +func takeRemoteAllocationStatementGroupLocked( + group *remoteAllocationStatementGroup, +) ([]*statementAllocationAttempt, []*mpool.MPool) { + group.finalized = true + if group.timer != nil { + group.timer.Stop() + group.timer = nil + } + attempts := group.attempts + group.attempts = nil + pools := group.pools + group.pools = nil + return attempts, pools +} + +func completeRemoteAllocationStatementGroup( + group *remoteAllocationStatementGroup, + attempts []*statementAllocationAttempt, + pools []*mpool.MPool, + terminalErr error, +) ( + remoteAllocationStatementTerminal, + error, +) { + terminalErr = errors.Join( + terminalErr, + allocationLifecycleCall(func() error { + group.board.CloseAndDrain() + return nil + }), + ) + terminal := remoteAllocationStatementTerminal{ + complete: true, + allocation: make( + []mpool.AllocationAccountTerminalSnapshot, + 0, + len(attempts), + ), + } + for _, attempt := range attempts { + snapshot, err := attempt.completeTerminal() + terminal.allocation = append(terminal.allocation, snapshot) + terminalErr = errors.Join(terminalErr, err) + } + for _, pool := range pools { + terminalErr = errors.Join( + terminalErr, + allocationLifecycleCall(func() error { + domain, quality := pool.ResourceSnapshot() + terminal.quality |= quality | + resource.MergeMemoryDomain(&terminal.memory, domain) + return nil + }), + ) + } + return terminal, terminalErr +} + +func expireRemoteAllocationStatementGroup( + group *remoteAllocationStatementGroup, +) { + remoteAllocationStatementGroups.Lock() + if group.finalized || group.registered == group.expected { + remoteAllocationStatementGroups.Unlock() + return + } + group.expired = true + timeoutErr := moerr.NewInternalErrorNoCtx( + "remote allocation statement group registration timed out", + ) + group.err = errors.Join(group.err, timeoutErr) + expected, registered, finished := group.expected, group.registered, group.finished + cancels := activeRemoteAllocationStatementCancelsLocked(group) + var attempts []*statementAllocationAttempt + var pools []*mpool.MPool + complete := group.finished == group.registered + if complete { + attempts, pools = takeRemoteAllocationStatementGroupLocked(group) + } + terminalErr := group.err + remoteAllocationStatementGroups.Unlock() + + if complete { + defer releaseRemoteAllocationStatementGroup(group) + _, terminalErr = completeRemoteAllocationStatementGroup( + group, + attempts, + pools, + terminalErr, + ) + } else { + terminalErr = errors.Join( + terminalErr, + allocationLifecycleCall(func() error { + group.board.CloseAndDrain() + return nil + }), + ) + } + cancelErr := cancelRemoteAllocationStatementParticipants(cancels, timeoutErr) + terminalErr = errors.Join(terminalErr, cancelErr) + if !complete && cancelErr != nil { + remoteAllocationStatementGroups.Lock() + if !group.finalized { + group.err = errors.Join(group.err, cancelErr) + } + remoteAllocationStatementGroups.Unlock() + } + fields := []zap.Field{ + zap.Uint32("expected-fragments", expected), + zap.Uint32("registered-fragments", registered), + zap.Uint32("finished-fragments", finished), + } + if terminalErr != nil { + fields = append(fields, zap.Error(terminalErr)) + } + logutil.Warn("expired incomplete remote allocation statement group", fields...) +} diff --git a/pkg/sql/compile/remote_allocation_statement_group_test.go b/pkg/sql/compile/remote_allocation_statement_group_test.go new file mode 100644 index 0000000000000..110a0948af125 --- /dev/null +++ b/pkg/sql/compile/remote_allocation_statement_group_test.go @@ -0,0 +1,394 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package compile + +import ( + "errors" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/vm/engine" + "github.com/matrixorigin/matrixone/pkg/vm/message" + "github.com/stretchr/testify/require" +) + +type remoteAllocationAccountedMessage struct { + mp *mpool.MPool + buffer []byte + destroyed *atomic.Int32 +} + +func remoteAllocationStatementGroupRegistered(board *message.MessageBoard) bool { + remoteAllocationStatementGroups.Lock() + defer remoteAllocationStatementGroups.Unlock() + _, registered := remoteAllocationStatementGroups.byBoard[board] + return registered +} + +func (m *remoteAllocationAccountedMessage) Serialize() []byte { return nil } + +func (m *remoteAllocationAccountedMessage) Deserialize([]byte) message.Message { + return m +} + +func (m *remoteAllocationAccountedMessage) NeedBlock() bool { return true } + +func (m *remoteAllocationAccountedMessage) GetMsgTag() int32 { return 1 } + +func (m *remoteAllocationAccountedMessage) GetReceiverAddr() message.MessageAddress { + return message.AddrBroadCastOnCurrentCN() +} + +func (m *remoteAllocationAccountedMessage) DebugString() string { + return "remote allocation-accounted message" +} + +func (m *remoteAllocationAccountedMessage) Destroy() { + if m.buffer != nil { + m.mp.Free(m.buffer) + m.buffer = nil + m.destroyed.Add(1) + } +} + +func TestCollectRemoteFragmentCountsCarriesExecutionAddress(t *testing.T) { + remoteC := &Scope{ + Magic: Remote, + NodeInfo: engine.Node{Addr: "cn-c:6001"}, + } + sameB := &Scope{ + Magic: Remote, + NodeInfo: engine.Node{Addr: "cn-b:6001"}, + PreScopes: []*Scope{remoteC}, + } + firstB := &Scope{ + Magic: Remote, + NodeInfo: engine.Node{Addr: "cn-b:6001"}, + PreScopes: []*Scope{sameB}, + } + secondB := &Scope{ + Magic: Remote, + NodeInfo: engine.Node{Addr: "cn-b:6001"}, + } + local := &Scope{ + Magic: Remote, + NodeInfo: engine.Node{Addr: "cn-a:6001"}, + } + + require.Equal(t, map[string]uint32{ + "cn-b:6001": 2, + "cn-c:6001": 1, + }, collectRemoteFragmentCounts( + []*Scope{firstB, nil, secondB, local}, + "cn-a:6001", + )) +} + +func TestRemoteExecutionIDSeparatesRetryMessageBoards(t *testing.T) { + statementID := newRemoteExecutionID() + firstAttempt := newRemoteExecutionID() + secondAttempt := newRemoteExecutionID() + require.NotEqual(t, firstAttempt, secondAttempt) + require.Equal(t, statementID, remoteMessageBoardID(statementID, [16]byte{})) + require.Equal(t, firstAttempt, remoteMessageBoardID(statementID, firstAttempt)) + require.Equal(t, secondAttempt, remoteMessageBoardID(statementID, secondAttempt)) + require.NotEqual(t, + remoteAllocationStatementGroupKey(firstAttempt, "cn-a:6001"), + remoteAllocationStatementGroupKey(secondAttempt, "cn-a:6001"), + ) +} + +func TestRemoteAllocationTopologyCapabilityIsRequiredForOwners(t *testing.T) { + owner := &allocationLifecycleOwnerOperator{ + MockOperator: colexec.NewMockOperator(), + } + scopes := []*Scope{{RootOp: owner}} + require.Error(t, validateRemoteAllocationTopologyCapability(scopes, nil)) + require.NoError(t, validateRemoteAllocationTopologyCapability( + scopes, + map[string]uint32{"cn-a:6001": 1}, + )) +} + +func TestRemoteAllocationStatementGroupDefersSharedBoardTerminal(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(2, 4) + require.NoError(t, err) + board := message.NewMessageBoard() + producer := newTestAllocationLifecycleCompile(t, registry, func( + mpool.AllocationAccountTerminalSnapshot, + ) { + t.Fatal("remote statement group must own terminal export") + }) + producer.MessageBoard = board + attempt, err := producer.beginAllocationAccountAttempt() + require.NoError(t, err) + + buffer, err := producer.proc.Mp().AllocAccounted( + 64, + attempt.account, + mpool.AllocationOwner(1), + mpool.AllocationSite(1), + ) + require.NoError(t, err) + var destroyed atomic.Int32 + message.SendMessage(&remoteAllocationAccountedMessage{ + mp: producer.proc.Mp(), + buffer: buffer, + destroyed: &destroyed, + }, board) + + first, err := acquireRemoteAllocationStatementParticipant(board, 2, nil) + require.NoError(t, err) + first.stage(attempt, producer.proc.Mp()) + terminal, err := first.finish(nil) + require.NoError(t, err) + require.Empty(t, terminal.allocation) + require.False(t, terminal.complete) + require.Zero(t, destroyed.Load()) + require.Equal(t, uint64(cap(buffer)), attempt.account.Snapshot().Used) + require.NotContains(t, board.DebugString(), "closed") + require.False(t, registry.AdmissionSuspended()) + + second, err := acquireRemoteAllocationStatementParticipant(board, 2, nil) + require.NoError(t, err) + // The second fragment has no allocation owner. It still participates in + // the statement boundary and, as the last fragment, drains the producer's + // queued ownership before completing that producer's account. + terminal, err = second.finish(nil) + require.NoError(t, err) + require.True(t, terminal.complete) + require.Len(t, terminal.allocation, 1) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.allocation[0].State) + require.Zero(t, terminal.allocation[0].Used) + require.Equal(t, uint64(cap(buffer)), terminal.allocation[0].Peak) + require.Equal(t, uint64(cap(buffer)), terminal.memory.AllocatedBytes) + require.Equal(t, uint64(cap(buffer)), terminal.memory.FreedBytes) + require.Zero(t, terminal.memory.LiveBytesAtSeal) + require.Zero(t, terminal.quality) + require.Equal(t, int32(1), destroyed.Load()) + require.Contains(t, board.DebugString(), "closed") + require.False(t, registry.AdmissionSuspended()) + require.Zero(t, registry.LiveAllocationMetadata()) + require.False(t, remoteAllocationStatementGroupRegistered(board)) +} + +func TestRemoteAllocationStatementGroupRejectsTopologyMismatch(t *testing.T) { + board := message.NewMessageBoard() + participant, err := acquireRemoteAllocationStatementParticipant(board, 2, nil) + require.NoError(t, err) + _, err = acquireRemoteAllocationStatementParticipant(board, 3, nil) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + + second, err := acquireRemoteAllocationStatementParticipant(board, 2, nil) + require.NoError(t, err) + _, err = participant.finish(nil) + require.NoError(t, err) + _, err = second.finish(nil) + require.NoError(t, err) + require.False(t, remoteAllocationStatementGroupRegistered(board)) +} + +func TestRemoteAllocationStatementGroupExpiresMissingFragment(t *testing.T) { + previousTimeout := remoteAllocationStatementRegistrationTimeout + remoteAllocationStatementRegistrationTimeout = 10 * time.Millisecond + t.Cleanup(func() { + remoteAllocationStatementRegistrationTimeout = previousTimeout + }) + + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + board := message.NewMessageBoard() + producer := newTestAllocationLifecycleCompile(t, registry, func( + mpool.AllocationAccountTerminalSnapshot, + ) { + t.Fatal("expired remote statement group must own terminal export") + }) + producer.MessageBoard = board + attempt, err := producer.beginAllocationAccountAttempt() + require.NoError(t, err) + buffer, err := producer.proc.Mp().AllocAccounted( + 64, + attempt.account, + mpool.AllocationOwner(1), + mpool.AllocationSite(1), + ) + require.NoError(t, err) + var destroyed atomic.Int32 + message.SendMessage(&remoteAllocationAccountedMessage{ + mp: producer.proc.Mp(), + buffer: buffer, + destroyed: &destroyed, + }, board) + + participant, err := acquireRemoteAllocationStatementParticipant(board, 2, nil) + require.NoError(t, err) + participant.stage(attempt, producer.proc.Mp()) + terminal, err := participant.finish(nil) + require.NoError(t, err) + require.Empty(t, terminal.allocation) + require.False(t, terminal.complete) + require.Eventually(t, func() bool { + return destroyed.Load() == 1 && + registry.LiveAllocationMetadata() == 0 && + !remoteAllocationStatementGroupRegistered(board) + }, time.Second, time.Millisecond) + require.Contains(t, board.DebugString(), "closed") + require.False(t, registry.AdmissionSuspended()) +} + +func TestRemoteAllocationStatementGroupFailureAbortsMissingFragment(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + board := message.NewMessageBoard() + producer := newTestAllocationLifecycleCompile(t, registry, func( + mpool.AllocationAccountTerminalSnapshot, + ) { + t.Fatal("remote statement group must own terminal export") + }) + producer.MessageBoard = board + attempt, err := producer.beginAllocationAccountAttempt() + require.NoError(t, err) + participant, err := acquireRemoteAllocationStatementParticipant(board, 2, nil) + require.NoError(t, err) + participant.stage(attempt, producer.proc.Mp()) + + failure := errors.New("remote fragment failed") + terminal, err := participant.finish(failure) + require.ErrorIs(t, err, failure) + require.True(t, terminal.complete) + require.Len(t, terminal.allocation, 1) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.allocation[0].State) + require.Contains(t, board.DebugString(), "closed") + require.Zero(t, registry.LiveAllocationMetadata()) + require.False(t, remoteAllocationStatementGroupRegistered(board)) +} + +func TestRemoteAllocationStatementGroupFailureCancelsActiveSibling(t *testing.T) { + board := message.NewMessageBoard() + canceled := make(chan error, 2) + first, err := acquireRemoteAllocationStatementParticipant( + board, + 3, + func(cause error) { canceled <- cause }, + ) + require.NoError(t, err) + second, err := acquireRemoteAllocationStatementParticipant( + board, + 3, + func(cause error) { canceled <- cause }, + ) + require.NoError(t, err) + + failure := errors.New("first remote fragment failed") + terminal, err := first.finish(failure) + require.NoError(t, err) + require.Empty(t, terminal.allocation) + require.False(t, terminal.complete) + // The failing participant is already quiescent; only its active sibling + // needs cancellation. + require.ErrorIs(t, <-canceled, failure) + select { + case unexpected := <-canceled: + t.Fatalf("finished participant was canceled: %v", unexpected) + default: + } + require.Contains(t, board.DebugString(), "closed") + + terminal, err = second.finish(errors.New("active sibling canceled")) + require.Error(t, err) + require.Empty(t, terminal.allocation) + require.False(t, remoteAllocationStatementGroupRegistered(board)) +} + +func TestRemoteAllocationStatementGroupExpirationWaitsForActiveFragment(t *testing.T) { + previousTimeout := remoteAllocationStatementRegistrationTimeout + remoteAllocationStatementRegistrationTimeout = 10 * time.Millisecond + t.Cleanup(func() { + remoteAllocationStatementRegistrationTimeout = previousTimeout + }) + + registry, err := mpool.NewAllocationAccountRegistry(2, 2) + require.NoError(t, err) + board := message.NewMessageBoard() + newAttempt := func() (*Compile, *statementAllocationAttempt, []byte) { + c := newTestAllocationLifecycleCompile(t, registry, func( + mpool.AllocationAccountTerminalSnapshot, + ) { + t.Fatal("remote statement group must own terminal export") + }) + c.MessageBoard = board + attempt, openErr := c.beginAllocationAccountAttempt() + require.NoError(t, openErr) + buffer, allocErr := c.proc.Mp().AllocAccounted( + 64, + attempt.account, + mpool.AllocationOwner(1), + mpool.AllocationSite(1), + ) + require.NoError(t, allocErr) + return c, attempt, buffer + } + + firstCompile, firstAttempt, firstBuffer := newAttempt() + secondCompile, secondAttempt, secondBuffer := newAttempt() + canceled := make(chan error, 2) + first, err := acquireRemoteAllocationStatementParticipant( + board, 3, func(cause error) { canceled <- cause }, + ) + require.NoError(t, err) + second, err := acquireRemoteAllocationStatementParticipant( + board, 3, func(cause error) { canceled <- cause }, + ) + require.NoError(t, err) + firstCompile.proc.Mp().Free(firstBuffer) + first.stage(firstAttempt, firstCompile.proc.Mp()) + terminal, err := first.finish(nil) + require.NoError(t, err) + require.Empty(t, terminal.allocation) + + require.Eventually(t, func() bool { + return strings.Contains(board.DebugString(), "closed") + }, time.Second, time.Millisecond) + require.Error(t, <-canceled) + select { + case unexpected := <-canceled: + t.Fatalf("finished participant was canceled: %v", unexpected) + default: + } + // Expiration closes the shared transport but cannot terminally inspect a + // registered fragment that is still executing. + _, live := registry.Resolve(secondAttempt.account.Handle()) + require.True(t, live) + require.Equal(t, uint64(cap(secondBuffer)), secondAttempt.account.Snapshot().Used) + require.False(t, secondAttempt.account.Snapshot().Sealed) + + secondCompile.proc.Mp().Free(secondBuffer) + second.stage(secondAttempt, secondCompile.proc.Mp()) + terminal, err = second.finish(errors.New("active fragment observed cancellation")) + require.Error(t, err) + require.Len(t, terminal.allocation, 2) + for _, snapshot := range terminal.allocation { + require.Equal(t, mpool.AllocationAccountTerminalValid, snapshot.State) + require.Zero(t, snapshot.Used) + } + require.Zero(t, terminal.memory.LiveBytesAtSeal) + require.Zero(t, terminal.quality) + require.Zero(t, registry.LiveAllocationMetadata()) + require.False(t, remoteAllocationStatementGroupRegistered(board)) +} diff --git a/pkg/sql/compile/remoterun.go b/pkg/sql/compile/remoterun.go index 5bfe6fd98806d..5819e20e9e999 100644 --- a/pkg/sql/compile/remoterun.go +++ b/pkg/sql/compile/remoterun.go @@ -17,6 +17,7 @@ package compile import ( "context" "fmt" + "maps" "time" "unsafe" @@ -170,11 +171,17 @@ func decodeScope(data []byte, proc *process.Process, isRemote bool, eng engine.E func encodeProcessInfo( proc *process.Process, sql string, + remoteFragmentCounts map[string]uint32, + remoteExecutionID uuid.UUID, ) ([]byte, error) { v, err := proc.BuildProcessInfo(sql) if err != nil { return nil, err } + v.RemoteFragmentCounts = maps.Clone(remoteFragmentCounts) + if remoteExecutionID != uuid.Nil { + v.RemoteExecutionId = append([]byte(nil), remoteExecutionID[:]...) + } return v.Marshal() } diff --git a/pkg/sql/compile/remoterunClient.go b/pkg/sql/compile/remoterunClient.go index 16b84fe9aab1c..41cef01d9aac6 100644 --- a/pkg/sql/compile/remoterunClient.go +++ b/pkg/sql/compile/remoterunClient.go @@ -23,6 +23,7 @@ import ( "sync" "time" + "github.com/google/uuid" "github.com/matrixorigin/matrixone/pkg/cnservice/cnclient" "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/common/morpc" @@ -69,7 +70,13 @@ func (s *Scope) remoteRun(c *Compile) (sender *messageSenderOnClient, err error) // encode structures which need to send. var scopeEncodeData, processEncodeData []byte var withoutOutput, folded bool - scopeEncodeData, withoutOutput, processEncodeData, folded, err = prepareRemoteRunSendingData(c.sql, s, c.proc) + scopeEncodeData, withoutOutput, processEncodeData, folded, err = prepareRemoteRunSendingData( + c.sql, + s, + c.proc, + c.remoteFragmentCounts, + c.remoteExecutionID, + ) if err != nil { return nil, err } @@ -180,7 +187,13 @@ func checkPipelineStandaloneExecutableAtRemote(s *Scope) bool { return true } -func prepareRemoteRunSendingData(sqlStr string, s *Scope, proc *process.Process) (scopeData []byte, withoutOutput bool, processData []byte, folded bool, err error) { +func prepareRemoteRunSendingData( + sqlStr string, + s *Scope, + proc *process.Process, + remoteFragmentCounts map[string]uint32, + remoteExecutionID uuid.UUID, +) (scopeData []byte, withoutOutput bool, processData []byte, folded bool, err error) { encodedScope, withoutOutput := getScopeForRemoteRunEncoding(s) encodedScope, folded, err = foldVarExprsInRemoteRunScope(encodedScope, proc) if err != nil { @@ -193,7 +206,12 @@ func prepareRemoteRunSendingData(sqlStr string, s *Scope, proc *process.Process) } // Encode the Process related information. - if processData, err = encodeProcessInfo(s.Proc, sqlStr); err != nil { + if processData, err = encodeProcessInfo( + s.Proc, + sqlStr, + remoteFragmentCounts, + remoteExecutionID, + ); err != nil { return nil, false, nil, false, err } @@ -808,6 +826,8 @@ func (sender *messageSenderOnClient) dealRemoteTerminal(data []byte) error { envelope.Allocation, envelope.MissingFragmentCount, envelope.MissingMemoryDomainCount, + envelope.PendingAllocationGroups, + envelope.CompletedAllocationGroups, ) } sender.terminalSeen = true diff --git a/pkg/sql/compile/remoterunServer.go b/pkg/sql/compile/remoterunServer.go index d7bfa5b87b05e..5ad79d63db0b3 100644 --- a/pkg/sql/compile/remoterunServer.go +++ b/pkg/sql/compile/remoterunServer.go @@ -19,6 +19,7 @@ import ( "encoding/json" "errors" "fmt" + "maps" "math" "sync" "time" @@ -51,6 +52,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/vm" "github.com/matrixorigin/matrixone/pkg/vm/engine" "github.com/matrixorigin/matrixone/pkg/vm/engine/disttae" + "github.com/matrixorigin/matrixone/pkg/vm/message" "github.com/matrixorigin/matrixone/pkg/vm/process" "go.uber.org/zap" ) @@ -252,26 +254,107 @@ func handlePipelineMessage(receiver *messageReceiverOnServer) (err error) { return errBuildCompile } var allocationAttempt *statementAllocationAttempt + var allocationParticipant *remoteAllocationStatementParticipant + var allocationGroupKey string var localAllocation resource.AllocationAccountTotals var localAllocationQuality resource.QualityFlags var runErr error + sharedMessageBoard := runCompile.MessageBoard + memoryPool := runCompile.proc.Mp() + statementGroupEnabled := len(runCompile.remoteFragmentCounts) > 0 + participantFinished := false + // This outer defer is the last-resort owner for a participant if the + // normal terminal path itself panics. Keep the Compile alive until after + // this guard runs: Release can make the pooled object reachable by a new + // RPC while terminal cleanup still needs its stable attempt references. defer func() { + recovered := recover() + defer runCompile.Release() + if allocationParticipant != nil && !participantFinished { + if allocationAttempt != nil { + runCompile.allocationAttempt = nil + } + allocationParticipant.stage(allocationAttempt, memoryPool) + cause := err + if recovered != nil { + cause = errors.Join( + cause, + moerr.ConvertPanicError(receiver.messageCtx, recovered), + ) + } + _, terminalErr := allocationParticipant.finish(cause) + err = errors.Join(err, terminalErr) + participantFinished = true + } + if recovered != nil { + panic(recovered) + } + }() + defer func() { + if recovered := recover(); recovered != nil { + err = errors.Join( + err, + moerr.ConvertPanicError(receiver.messageCtx, recovered), + ) + } // Capture operator and descendant facts before cleanup. The MPool // snapshot intentionally follows Compile.clear so temporary execution // allocations are released before LiveBytesAtSeal is measured. The // descendant snapshot is already reduced under AnalyzeModule's mutex; // sender quiescence remains the lifecycle contract for this boundary. - localDelta := collectScopeResourceDelta(runCompile.scopes, receiver.cnInformation.cnAddr) - descendant := runCompile.anal.remoteResourceSummary() - expectedDirect := countExpectedRemoteScopes(runCompile.scopes, receiver.cnInformation.cnAddr) - memoryPool := runCompile.proc.Mp() + var localDelta resource.Delta + var descendant remoteResourceSnapshot + var expectedDirect uint64 + err = errors.Join(err, allocationLifecycleCall(func() error { + localDelta = collectScopeResourceDelta( + runCompile.scopes, + receiver.cnInformation.cnAddr, + ) + descendant = runCompile.anal.remoteResourceSummary() + expectedDirect = countExpectedRemoteScopes( + runCompile.scopes, + receiver.cnInformation.cnAddr, + ) + return nil + })) + var terminal remoteAllocationStatementTerminal if allocationAttempt != nil { runCompile.allocationAttempt = nil - _, terminalErr := allocationAttempt.finish() - err = errors.Join(err, terminalErr) } - runCompile.clear() - localMemory, localMemoryQuality := memoryPool.ResourceSnapshot() + allocationParticipant.stage(allocationAttempt, memoryPool) + err = errors.Join(err, allocationLifecycleCall(func() error { + if statementGroupEnabled { + // The remote statement group, rather than any one fragment Compile, + // owns the shared multi-CN board. Detach it before clear resets the + // fragment-local Compile state. A fragment rejected before it joins a + // group is terminal for the statement, so it closes its unowned board. + if allocationParticipant == nil { + sharedMessageBoard.CloseAndDrain() + } + runCompile.MessageBoard = message.NewMessageBoard() + runCompile.proc.SetMessageBoard(runCompile.MessageBoard) + } + runCompile.clear() + return nil + })) + terminal, terminalErr := allocationParticipant.finish(err) + participantFinished = true + err = errors.Join(err, terminalErr) + for _, snapshot := range terminal.allocation { + localAllocationQuality |= localAllocation.AddGeneration( + snapshot.Peak, + snapshot.Used, + snapshot.State == mpool.AllocationAccountTerminalValid, + ) + } + var localMemory resource.MemoryDomainSummary + var localMemoryQuality resource.QualityFlags + if !statementGroupEnabled { + err = errors.Join(err, allocationLifecycleCall(func() error { + localMemory, localMemoryQuality = memoryPool.ResourceSnapshot() + return nil + })) + } aggregate := composeRemoteResourceAggregate( localDelta, localMemory, @@ -284,14 +367,54 @@ func handlePipelineMessage(receiver *messageReceiverOnServer) (err error) { &aggregate.Allocation, localAllocation, ) + aggregate.Delta.Quality |= terminal.quality | + resource.MergeMemoryTotals(&aggregate.Memory, terminal.memory) + if allocationParticipant != nil { + addRemoteAllocationGroupSignal( + &aggregate, + allocationGroupKey, + terminal.complete, + ) + } receiver.resourceDelta = aggregate.Delta receiver.resourceMemory = aggregate.Memory receiver.resourceAllocation = aggregate.Allocation receiver.resourceMissingFragments = aggregate.MissingFragmentCount receiver.resourceMissingMemoryDomains = aggregate.MissingMemoryDomainCount - - runCompile.Release() + receiver.resourcePendingAllocationGroups = aggregate.PendingAllocationGroups + receiver.resourceCompletedAllocationGroups = aggregate.CompletedAllocationGroups }() + if statementGroupEnabled { + expectedFragments, ok := runCompile.remoteFragmentCounts[runCompile.addr] + if !ok || expectedFragments == 0 { + return errors.Join( + mpool.ErrAllocationAccountInvariant, + moerr.NewInternalErrorNoCtx( + "remote fragment topology has no local CN entry", + ), + ) + } + // Capture the function value for this execution generation. Both the + // Compile and Process fields are reset after an early sibling returns; + // looking up proc.Cancel later could cancel an unrelated generation. + remoteCancel := runCompile.proc.Cancel + allocationGroupKey = remoteAllocationStatementGroupKey( + runCompile.remoteExecutionID, + runCompile.addr, + ) + allocationParticipant, runErr = acquireRemoteAllocationStatementParticipant( + runCompile.MessageBoard, + expectedFragments, + func(cause error) { + if remoteCancel != nil { + remoteCancel(cause) + } + }, + ) + if runErr != nil { + return runErr + } + } // decode and running the pipeline. s, runErr := decodeScope(receiver.scopeData, runCompile.proc, true, runCompile.e) @@ -315,17 +438,19 @@ func handlePipelineMessage(receiver *messageReceiverOnServer) (err error) { } runCompile.scopes = []*Scope{s} - runErr = runCompile.ensureAllocationAccountLifecycle(func( - snapshot mpool.AllocationAccountTerminalSnapshot, - ) { - localAllocationQuality |= localAllocation.AddGeneration( - snapshot.Peak, - snapshot.Used, - snapshot.State == mpool.AllocationAccountTerminalValid, + if runErr = validateRemoteAllocationTopologyCapability( + runCompile.scopes, + runCompile.remoteFragmentCounts, + ); runErr != nil { + return runErr + } + if statementGroupEnabled { + runErr = runCompile.ensureAllocationAccountLifecycle( + func(mpool.AllocationAccountTerminalSnapshot) {}, ) - }) - if runErr == nil { - allocationAttempt, runErr = runCompile.beginAllocationAccountAttempt() + if runErr == nil { + allocationAttempt, runErr = runCompile.beginAllocationAccountAttempt() + } } if runErr != nil { return runErr @@ -554,9 +679,11 @@ type processHelper struct { txnClient client.TxnClient sessionInfo process.SessionInfo //analysisNodeList []int32 - StmtId uuid.UUID - prepareParams pipeline.PrepareParamInfo - affectedRows int64 + StmtId uuid.UUID + prepareParams pipeline.PrepareParamInfo + affectedRows int64 + remoteFragmentCounts map[string]uint32 + remoteExecutionID uuid.UUID } // messageReceiverOnServer supported a series methods to write back results. @@ -585,12 +712,14 @@ type messageReceiverOnServer struct { colexecServer *colexec.Server // result. - phyPlan *models.PhyPlan - resourceDelta resource.Delta - resourceMemory resource.MemoryTotals - resourceAllocation resource.AllocationAccountTotals - resourceMissingFragments uint64 - resourceMissingMemoryDomains uint64 + phyPlan *models.PhyPlan + resourceDelta resource.Delta + resourceMemory resource.MemoryTotals + resourceAllocation resource.AllocationAccountTotals + resourceMissingFragments uint64 + resourceMissingMemoryDomains uint64 + resourcePendingAllocationGroups []remoteAllocationGroupPending + resourceCompletedAllocationGroups []string } func newMessageReceiverOnServer( @@ -745,10 +874,16 @@ func (receiver *messageReceiverOnServer) newCompile() (*Compile, error) { c := allocateNewCompile(proc) c.execType = plan2.ExecTypeAP_MULTICN c.e = cnInfo.storeEngine - c.MessageBoard = c.MessageBoard.SetMultiCN(c.GetMessageCenter(), c.proc.GetStmtProfile().GetStmtId()) + messageBoardID := remoteMessageBoardID( + c.proc.GetStmtProfile().GetStmtId(), + pHelper.remoteExecutionID, + ) + c.MessageBoard = c.MessageBoard.SetMultiCN(c.GetMessageCenter(), messageBoardID) c.proc.SetMessageBoard(c.MessageBoard) c.anal = newAnalyzeModule() c.addr = receiver.cnInformation.cnAddr + c.remoteFragmentCounts = maps.Clone(pHelper.remoteFragmentCounts) + c.remoteExecutionID = pHelper.remoteExecutionID // a method to send back. c.execType = plan2.ExecTypeAP_MULTICN @@ -888,12 +1023,14 @@ func (receiver *messageReceiverOnServer) sendEndMessage() error { func (receiver *messageReceiverOnServer) setTerminalAnalysis(message *pipeline.Message) error { envelope := remoteTerminalEnvelope{ - TerminalResourceVersion: remoteTerminalResourceVersion, - Delta: receiver.resourceDelta, - Memory: receiver.resourceMemory, - Allocation: receiver.resourceAllocation, - MissingFragmentCount: receiver.resourceMissingFragments, - MissingMemoryDomainCount: receiver.resourceMissingMemoryDomains, + TerminalResourceVersion: remoteTerminalResourceVersion, + Delta: receiver.resourceDelta, + Memory: receiver.resourceMemory, + Allocation: receiver.resourceAllocation, + MissingFragmentCount: receiver.resourceMissingFragments, + MissingMemoryDomainCount: receiver.resourceMissingMemoryDomains, + PendingAllocationGroups: receiver.resourcePendingAllocationGroups, + CompletedAllocationGroups: receiver.resourceCompletedAllocationGroups, } if receiver.phyPlan != nil { envelope.PhyPlan = *receiver.phyPlan @@ -914,12 +1051,24 @@ func generateProcessHelper(ctx context.Context, data []byte, cli client.TxnClien } result := processHelper{ - id: procInfo.Id, - lim: process.ConvertToProcessLimitation(procInfo.Lim), - unixTime: procInfo.UnixTime, - accountId: procInfo.AccountId, - txnClient: cli, - affectedRows: procInfo.AffectedRows, + id: procInfo.Id, + lim: process.ConvertToProcessLimitation(procInfo.Lim), + unixTime: procInfo.UnixTime, + accountId: procInfo.AccountId, + txnClient: cli, + affectedRows: procInfo.AffectedRows, + remoteFragmentCounts: maps.Clone(procInfo.RemoteFragmentCounts), + } + if len(procInfo.RemoteExecutionId) > 0 { + result.remoteExecutionID, err = uuid.FromBytes(procInfo.RemoteExecutionId) + if err != nil { + return processHelper{}, err + } + } + if (len(result.remoteFragmentCounts) == 0) != (result.remoteExecutionID == uuid.Nil) { + return processHelper{}, moerr.NewInternalErrorNoCtx( + "incomplete remote allocation lifecycle metadata", + ) } result.txnOperator, err = cli.NewWithSnapshot(ctx, procInfo.Snapshot) if err != nil { diff --git a/pkg/sql/compile/remoterun_test.go b/pkg/sql/compile/remoterun_test.go index 804a0748c2604..392faa5d37188 100644 --- a/pkg/sql/compile/remoterun_test.go +++ b/pkg/sql/compile/remoterun_test.go @@ -145,8 +145,41 @@ func Test_EncodeProcessInfo(t *testing.T) { SqlHelper: nil, } - _, err := encodeProcessInfo(proc, "") + remoteExecutionID := uuid.New() + data, err := encodeProcessInfo(proc, "", map[string]uint32{ + "cn-a:6001": 2, + "cn-b:6001": 1, + }, remoteExecutionID) require.Nil(t, err) + restored := new(pipeline.ProcessInfo) + require.NoError(t, restored.Unmarshal(data)) + require.Equal(t, map[string]uint32{ + "cn-a:6001": 2, + "cn-b:6001": 1, + }, restored.RemoteFragmentCounts) + restoredExecutionID, err := uuid.FromBytes(restored.RemoteExecutionId) + require.NoError(t, err) + require.Equal(t, remoteExecutionID, restoredExecutionID) +} + +func TestGenerateProcessHelperRejectsIncompleteRemoteLifecycleMetadata(t *testing.T) { + tests := []pipeline.ProcessInfo{ + {RemoteFragmentCounts: map[string]uint32{"cn-a:6001": 1}}, + {RemoteExecutionId: func() []byte { + id := uuid.New() + return id[:] + }()}, + { + RemoteFragmentCounts: map[string]uint32{"cn-a:6001": 1}, + RemoteExecutionId: []byte{1}, + }, + } + for i := range tests { + data, err := tests[i].Marshal() + require.NoError(t, err) + _, err = generateProcessHelper(context.Background(), data, nil) + require.Error(t, err) + } } func Test_refactorScope(t *testing.T) { @@ -2753,7 +2786,7 @@ func Test_prepareRemoteRunSendingData(t *testing.T) { Proc: proc, RootOp: connector.NewArgument(), } - _, withoutOut, _, _, err := prepareRemoteRunSendingData("", s1, proc) + _, withoutOut, _, _, err := prepareRemoteRunSendingData("", s1, proc, nil, uuid.Nil) require.NoError(t, err) require.False(t, withoutOut) require.NotNil(t, s1.RootOp) @@ -2767,7 +2800,7 @@ func Test_prepareRemoteRunSendingData(t *testing.T) { } s2.RootOp.AppendChild(value_scan.NewArgument()) originChild := s2.RootOp.GetOperatorBase().GetChildren(0) - _, withoutOut, _, _, err = prepareRemoteRunSendingData("", s2, proc) + _, withoutOut, _, _, err = prepareRemoteRunSendingData("", s2, proc, nil, uuid.Nil) require.NoError(t, err) require.False(t, withoutOut) require.Equal(t, 1, s2.RootOp.GetOperatorBase().NumChildren()) @@ -2780,7 +2813,7 @@ func Test_prepareRemoteRunSendingData(t *testing.T) { RootOp: value_scan.NewArgument(), } s3.RootOp.AppendChild(value_scan.NewArgument()) - _, withoutOut, _, _, err = prepareRemoteRunSendingData("", s3, proc) + _, withoutOut, _, _, err = prepareRemoteRunSendingData("", s3, proc, nil, uuid.Nil) require.NoError(t, err) require.True(t, withoutOut) } @@ -2808,7 +2841,7 @@ func TestPrepareRemoteRunSendingDataKeepsConnectorChildTableFunctionParams(t *te RootOp: conn, } - scopeData, withoutOut, _, _, err := prepareRemoteRunSendingData("", s, proc) + scopeData, withoutOut, _, _, err := prepareRemoteRunSendingData("", s, proc, nil, uuid.Nil) require.NoError(t, err) require.False(t, withoutOut) diff --git a/pkg/sql/compile/resource_accounting.go b/pkg/sql/compile/resource_accounting.go index 1d695199ae88e..b921571570135 100644 --- a/pkg/sql/compile/resource_accounting.go +++ b/pkg/sql/compile/resource_accounting.go @@ -17,6 +17,8 @@ package compile import ( "context" "math" + "slices" + "strings" "sync/atomic" "time" @@ -37,7 +39,7 @@ type executionResourceRecorder struct { pendingAllocationQuality resource.QualityFlags } -const remoteTerminalResourceVersion = 2 +const remoteTerminalResourceVersion = 3 // remoteTerminalEnvelope keeps PhyPlan fields at the top level so clients from // before resource accounting can still decode the terminal plan during a @@ -45,22 +47,31 @@ const remoteTerminalResourceVersion = 2 // appended resource facts from a legacy bare PhyPlan payload. type remoteTerminalEnvelope struct { models.PhyPlan - TerminalResourceVersion uint32 `json:"terminal_resource_version,omitempty"` - Delta resource.Delta `json:"resource_delta"` - Memory resource.MemoryTotals `json:"memory"` - Allocation resource.AllocationAccountTotals `json:"allocation_account"` - MissingFragmentCount uint64 `json:"missing_fragment_count,omitempty"` - MissingMemoryDomainCount uint64 `json:"missing_memory_domain_count,omitempty"` + TerminalResourceVersion uint32 `json:"terminal_resource_version,omitempty"` + Delta resource.Delta `json:"resource_delta"` + Memory resource.MemoryTotals `json:"memory"` + Allocation resource.AllocationAccountTotals `json:"allocation_account"` + MissingFragmentCount uint64 `json:"missing_fragment_count,omitempty"` + MissingMemoryDomainCount uint64 `json:"missing_memory_domain_count,omitempty"` + PendingAllocationGroups []remoteAllocationGroupPending `json:"pending_allocation_groups,omitempty"` + CompletedAllocationGroups []string `json:"completed_allocation_groups,omitempty"` +} + +type remoteAllocationGroupPending struct { + Key string `json:"key"` + Count uint64 `json:"count"` } // remoteResourceAggregate is the already-reduced terminal output sent by one // remote hop. Delta contains local plus descendant usage and quality. type remoteResourceAggregate struct { - Delta resource.Delta - Memory resource.MemoryTotals - Allocation resource.AllocationAccountTotals - MissingFragmentCount uint64 - MissingMemoryDomainCount uint64 + Delta resource.Delta + Memory resource.MemoryTotals + Allocation resource.AllocationAccountTotals + MissingFragmentCount uint64 + MissingMemoryDomainCount uint64 + PendingAllocationGroups []remoteAllocationGroupPending + CompletedAllocationGroups []string } func (r *executionResourceRecorder) recordAllocationAccountTerminal( @@ -135,6 +146,24 @@ func (r *executionResourceRecorder) finishAttempt( remote, countExpectedRemoteScopes(scopes, localAddress), ) + var pending uint64 + for _, signal := range remoteAggregate.PendingAllocationGroups { + pending, remoteAggregate.Delta.Quality = addCheckedRemoteCounter( + pending, + signal.Count, + remoteAggregate.Delta.Quality, + ) + } + if pending > 0 { + remoteAggregate.MissingMemoryDomainCount, remoteAggregate.Delta.Quality = + addCheckedRemoteCounter( + remoteAggregate.MissingMemoryDomainCount, + pending, + remoteAggregate.Delta.Quality, + ) + remoteAggregate.Delta.Quality |= resource.QualityPartial | + resource.QualityMissingMemoryDomain + } delta = remoteAggregate.Delta var coordinator resource.LocalRecorder @@ -250,6 +279,80 @@ func addCheckedRemoteCounter(value, add uint64, quality resource.QualityFlags) ( return value + add, quality } +func reduceRemoteAllocationGroupSignals( + pending []remoteAllocationGroupPending, + completed []string, +) ([]remoteAllocationGroupPending, []string, resource.QualityFlags) { + var quality resource.QualityFlags + completedSet := make(map[string]struct{}, len(completed)) + for _, key := range completed { + if key == "" { + quality |= resource.QualityInvariantFailure + continue + } + completedSet[key] = struct{}{} + } + pendingCounts := make(map[string]uint64, len(pending)) + for _, signal := range pending { + if signal.Key == "" || signal.Count == 0 { + quality |= resource.QualityInvariantFailure + continue + } + if _, resolved := completedSet[signal.Key]; !resolved { + pendingCounts[signal.Key], quality = addCheckedRemoteCounter( + pendingCounts[signal.Key], + signal.Count, + quality, + ) + } + } + pending = pending[:0] + for key, count := range pendingCounts { + pending = append(pending, remoteAllocationGroupPending{ + Key: key, + Count: count, + }) + } + completed = completed[:0] + for key := range completedSet { + completed = append(completed, key) + } + slices.SortFunc(pending, func(a, b remoteAllocationGroupPending) int { + return strings.Compare(a.Key, b.Key) + }) + slices.Sort(completed) + return pending, completed, quality +} + +func addRemoteAllocationGroupSignal( + aggregate *remoteResourceAggregate, + key string, + completed bool, +) { + if aggregate == nil { + return + } + if completed { + aggregate.CompletedAllocationGroups = append( + aggregate.CompletedAllocationGroups, + key, + ) + } else { + aggregate.PendingAllocationGroups = append( + aggregate.PendingAllocationGroups, + remoteAllocationGroupPending{Key: key, Count: 1}, + ) + } + var quality resource.QualityFlags + aggregate.PendingAllocationGroups, + aggregate.CompletedAllocationGroups, + quality = reduceRemoteAllocationGroupSignals( + aggregate.PendingAllocationGroups, + aggregate.CompletedAllocationGroups, + ) + aggregate.Delta.Quality |= quality +} + // composeRemoteResourceAggregate composes one hop's captured local resource // facts with an already-reduced descendant aggregate. It is pure so every // remote hop and the coordinator use the same merge algebra. @@ -271,12 +374,27 @@ func composeRemoteResourceAggregate( ) result.MissingFragmentCount = descendant.MissingFragmentCount result.MissingMemoryDomainCount = descendant.MissingMemoryDomainCount + result.PendingAllocationGroups = append( + []remoteAllocationGroupPending(nil), + descendant.PendingAllocationGroups..., + ) + result.CompletedAllocationGroups = append( + []string(nil), descendant.CompletedAllocationGroups..., + ) if descendant.MissingFragmentCount > 0 { result.Delta.Quality |= resource.QualityPartial | resource.QualityMissingFragment } if descendant.MissingMemoryDomainCount > 0 { result.Delta.Quality |= resource.QualityPartial | resource.QualityMissingMemoryDomain } + var groupQuality resource.QualityFlags + result.PendingAllocationGroups, + result.CompletedAllocationGroups, + groupQuality = reduceRemoteAllocationGroupSignals( + result.PendingAllocationGroups, + result.CompletedAllocationGroups, + ) + result.Delta.Quality |= groupQuality if descendant.DirectReportCount < expectedDirect { directMissing := expectedDirect - descendant.DirectReportCount result.MissingFragmentCount, result.Delta.Quality = addCheckedRemoteCounter( diff --git a/pkg/sql/compile/resource_accounting_test.go b/pkg/sql/compile/resource_accounting_test.go index c313049e8a99d..ff60d30fa8dee 100644 --- a/pkg/sql/compile/resource_accounting_test.go +++ b/pkg/sql/compile/resource_accounting_test.go @@ -55,6 +55,8 @@ func TestExecutionResourceRecorder(t *testing.T) { resource.AllocationAccountTotals{}, 0, 0, + nil, + nil, ) recorder.finishAttempt( 0, @@ -302,6 +304,11 @@ func TestRemoteTerminalEnvelope(t *testing.T) { MaxGenerationPeak: 17, SumGenerationPeak: 17, }, + PendingAllocationGroups: []remoteAllocationGroupPending{{ + Key: "pending@cn", + Count: 2, + }}, + CompletedAllocationGroups: []string{"completed@cn"}, } data, err := json.Marshal(envelope) require.NoError(t, err) @@ -311,6 +318,11 @@ func TestRemoteTerminalEnvelope(t *testing.T) { summary := anal.remoteResourceSummary() require.Equal(t, uint64(1), summary.DirectReportCount) require.Equal(t, uint64(11), summary.Usage.ExclusiveActiveNS) + require.Equal(t, []remoteAllocationGroupPending{{ + Key: "pending@cn", + Count: 2, + }}, summary.PendingAllocationGroups) + require.Equal(t, []string{"completed@cn"}, summary.CompletedAllocationGroups) require.Equal(t, uint64(12), summary.Usage.S3ReadBytes) require.Equal(t, uint64(15), summary.Memory.MaxDomainPeakLiveBytes) require.Equal(t, uint64(1), summary.Allocation.GenerationCount) @@ -496,6 +508,8 @@ func TestRemoteResourceCounterSaturates(t *testing.T) { resource.AllocationAccountTotals{}, 1, 1, + nil, + nil, ) snapshot := anal.remoteResourceSummary() require.Equal(t, uint64(math.MaxUint64), snapshot.MissingFragmentCount) @@ -522,12 +536,119 @@ func TestAnalyzeModuleResetClearsRemoteResourceAggregate(t *testing.T) { resource.AllocationAccountTotals{}, 2, 3, + []remoteAllocationGroupPending{{Key: "pending", Count: 1}}, + []string{"completed"}, ) anal.Reset(false, false) snapshot := anal.remoteResourceSummary() require.Equal(t, remoteResourceSnapshot{}, snapshot) } +func TestAnalyzeModuleResolvesRemoteAllocationGroupsInEitherOrder(t *testing.T) { + for _, tc := range []struct { + name string + completeFirst bool + }{ + {name: "pending-before-complete"}, + {name: "complete-before-pending", completeFirst: true}, + } { + t.Run(tc.name, func(t *testing.T) { + anal := &AnalyzeModule{} + pending := []remoteAllocationGroupPending{{ + Key: "execution@cn", + Count: 1, + }} + var firstPending, secondPending []remoteAllocationGroupPending + var firstCompleted, secondCompleted []string + if tc.completeFirst { + firstCompleted = []string{"execution@cn"} + secondPending = pending + } else { + firstPending = pending + secondCompleted = []string{"execution@cn"} + } + anal.appendRemoteResource( + resource.Delta{}, resource.MemoryTotals{}, + resource.AllocationAccountTotals{}, 0, 0, + firstPending, firstCompleted, + ) + anal.appendRemoteResource( + resource.Delta{}, resource.MemoryTotals{}, + resource.AllocationAccountTotals{}, 0, 0, + secondPending, secondCompleted, + ) + + snapshot := anal.remoteResourceSummary() + require.Empty(t, snapshot.PendingAllocationGroups) + require.Equal(t, []string{"execution@cn"}, snapshot.CompletedAllocationGroups) + require.Zero(t, snapshot.Quality&resource.QualityInvariantFailure) + }) + } +} + +func TestExecutionResourceRecorderMarksUnresolvedAllocationGroupPartial(t *testing.T) { + root := resource.NewRoot(resource.ConnExternal) + recorder := newExecutionResourceRecorder( + resource.ContextWithRoot(context.Background(), root), + true, + ) + require.NotNil(t, recorder) + anal := &AnalyzeModule{} + anal.appendRemoteResource( + resource.Delta{}, resource.MemoryTotals{}, + resource.AllocationAccountTotals{}, 0, 0, + []remoteAllocationGroupPending{{Key: "execution@cn", Count: 1}}, nil, + ) + + recorder.finishAttempt( + 0, time.Now(), 0, 0, nil, nil, anal, "local:6001", false, + ) + recorder.publish() + + summary := root.PreResponseSummary() + require.Equal(t, uint64(1), summary.MissingMemoryDomainCount) + require.NotZero(t, summary.Quality&resource.QualityPartial) + require.NotZero(t, summary.Quality&resource.QualityMissingMemoryDomain) + require.Zero(t, summary.Quality&resource.QualityMissingFragment) +} + +func TestExecutionResourceRecorderPreservesPendingGroupCardinality(t *testing.T) { + root := resource.NewRoot(resource.ConnExternal) + recorder := newExecutionResourceRecorder( + resource.ContextWithRoot(context.Background(), root), + true, + ) + require.NotNil(t, recorder) + anal := &AnalyzeModule{} + for range 3 { + anal.appendRemoteResource( + resource.Delta{}, resource.MemoryTotals{}, + resource.AllocationAccountTotals{}, 0, 0, + []remoteAllocationGroupPending{{Key: "execution@cn", Count: 1}}, + nil, + ) + } + scopes := make([]*Scope, 4) + for i := range scopes { + scopes[i] = &Scope{ + Magic: Remote, + NodeInfo: engine.Node{Addr: "remote:6001"}, + } + } + + recorder.finishAttempt( + 0, time.Now(), 0, 0, nil, scopes, anal, "local:6001", false, + ) + recorder.publish() + + summary := root.PreResponseSummary() + require.Equal(t, uint64(1), summary.MissingFragmentCount) + require.Equal(t, uint64(4), summary.MissingMemoryDomainCount) + require.NotZero(t, summary.Quality&resource.QualityPartial) + require.NotZero(t, summary.Quality&resource.QualityMissingFragment) + require.NotZero(t, summary.Quality&resource.QualityMissingMemoryDomain) +} + func TestAnalyzeModuleRemoteResourceConcurrentAccess(t *testing.T) { anal := &AnalyzeModule{} var wg sync.WaitGroup @@ -542,6 +663,8 @@ func TestAnalyzeModuleRemoteResourceConcurrentAccess(t *testing.T) { resource.AllocationAccountTotals{}, 1, 1, + []remoteAllocationGroupPending{{Key: "pending", Count: 1}}, + []string{"completed"}, ) _ = anal.remoteResourceSummary() } diff --git a/pkg/sql/compile/types.go b/pkg/sql/compile/types.go index 05c75669e54a4..c4365a7b82908 100644 --- a/pkg/sql/compile/types.go +++ b/pkg/sql/compile/types.go @@ -350,6 +350,8 @@ type Compile struct { allocationTerminalExporter func(mpool.AllocationAccountTerminalSnapshot) allocationAccountOwners []executionAllocationAccountOwner allocationAttempt *statementAllocationAttempt + remoteFragmentCounts map[string]uint32 + remoteExecutionID uuid.UUID hasMergeOp bool // ncpu set as system.GoRoutines() while NewCompile, instead of global static value. diff --git a/proto/pipeline.proto b/proto/pipeline.proto index 9df84b6d9ef93..0339ea23ac174 100644 --- a/proto/pipeline.proto +++ b/proto/pipeline.proto @@ -645,6 +645,12 @@ message ProcessInfo { SessionLoggerInfo session_logger = 8 [(gogoproto.nullable) = false]; PrepareParamInfo prepare_params = 9 [(gogoproto.nullable) = false]; int64 affected_rows = 10; + // Planned PipelineMessage RPCs per target CN for statement-level remote + // MessageBoard and resource-account terminal ownership. + map remote_fragment_counts = 11; + // Unique physical execution attempt. Unlike the SQL statement ID, this + // changes across retries and prepared-statement executions. + bytes remote_execution_id = 12; } message SessionInfo { From 3cf0777c1ee764f73420adf6801b12ddee2d0726 Mon Sep 17 00:00:00 2001 From: aptend Date: Sat, 1 Aug 2026 18:30:13 +0800 Subject: [PATCH 42/61] fix: preserve repeated vector alias decoding --- pkg/container/vector/allocation_account.go | 11 +++++++++ .../vector/allocation_account_test.go | 24 +++++++++++++++++++ pkg/container/vector/vector.go | 2 +- pkg/container/vector/vector_test.go | 10 ++++---- 4 files changed, 42 insertions(+), 5 deletions(-) diff --git a/pkg/container/vector/allocation_account.go b/pkg/container/vector/allocation_account.go index af7f6c878a9cf..6163657940ddf 100644 --- a/pkg/container/vector/allocation_account.go +++ b/pkg/container/vector/allocation_account.go @@ -227,6 +227,17 @@ func (v *Vector) hasBackingStorage() bool { v.gsp.GetBitmap().ExternalStorageCapacity() != 0 } +// hasOwnedBackingStorage reports storage that UnmarshalBinary cannot replace +// without losing an MPool-owned allocation. Data and area marked cantFree are +// borrowed aliases; ordinary bitmap backing is Go-owned and remains GC-visible +// after replacement. Accounted bitmap storage is explicit external storage. +func (v *Vector) hasOwnedBackingStorage() bool { + return cap(v.data) != 0 && !v.cantFreeData || + cap(v.area) != 0 && !v.cantFreeArea || + v.nsp.GetBitmap().ExternalStorageCapacity() != 0 || + v.gsp.GetBitmap().ExternalStorageCapacity() != 0 +} + // SetAllocationAccount selects the account used by future owned allocations. // It is intentionally explicit and is legal only before the first backing // allocation. Reset retains the selection; Free clears it. diff --git a/pkg/container/vector/allocation_account_test.go b/pkg/container/vector/allocation_account_test.go index b9be24e07544c..2ae706b9e241f 100644 --- a/pkg/container/vector/allocation_account_test.go +++ b/pkg/container/vector/allocation_account_test.go @@ -1229,3 +1229,27 @@ func TestUnmarshalBinaryRejectsOwnedDestinationWithoutLosingBacking(t *testing.T target.Free(mp) require.Zero(t, mp.CurrNB()) } + +func TestUnmarshalBinaryReplacesBorrowedAliases(t *testing.T) { + mp := mpool.MustNewZero() + first := NewVec(types.T_int64.ToType()) + second := NewVec(types.T_int64.ToType()) + require.NoError(t, AppendFixedList(first, []int64{1, 2}, nil, mp)) + require.NoError(t, AppendFixedList(second, []int64{3, 4}, []bool{true, false}, mp)) + firstData, err := first.MarshalBinary() + require.NoError(t, err) + secondData, err := second.MarshalBinary() + require.NoError(t, err) + + var target Vector + require.NoError(t, target.UnmarshalBinary(firstData)) + require.Equal(t, []int64{1, 2}, MustFixedColWithTypeCheck[int64](&target)) + require.NoError(t, target.UnmarshalBinaryTrusted(secondData)) + require.Equal(t, []int64{0, 4}, MustFixedColWithTypeCheck[int64](&target)) + require.True(t, target.GetNulls().Contains(0)) + + first.Free(mp) + second.Free(mp) + target.Free(mp) + require.Zero(t, mp.CurrNB()) +} diff --git a/pkg/container/vector/vector.go b/pkg/container/vector/vector.go index ed3645c369f6c..33bd3b913a8c3 100644 --- a/pkg/container/vector/vector.go +++ b/pkg/container/vector/vector.go @@ -1098,7 +1098,7 @@ func (v *Vector) unmarshalBinary(data []byte, validateValues bool) error { if err != nil { return err } - if v.hasBackingStorage() { + if v.hasOwnedBackingStorage() { return allocationAccountInvalid( "cannot replace owned vector storage with aliases", ) diff --git a/pkg/container/vector/vector_test.go b/pkg/container/vector/vector_test.go index cd8ccd3c8f73a..e54c763b2ad58 100644 --- a/pkg/container/vector/vector_test.go +++ b/pkg/container/vector/vector_test.go @@ -3020,15 +3020,17 @@ func TestGetAny(t *testing.T) { func BenchmarkUnmarshal(b *testing.B) { mp := mpool.MustNewZero() - vec := NewVec(types.T_int8.ToType()) - AppendAny(vec, int8(42), false, mp) - data, err := vec.MarshalBinary() + source := NewVec(types.T_int8.ToType()) + AppendAny(source, int8(42), false, mp) + data, err := source.MarshalBinary() if err != nil { b.Fatal(err) } + source.Free(mp) + var target Vector b.ResetTimer() for i := 0; i < b.N; i++ { - err := vec.UnmarshalBinary(data) + err := target.UnmarshalBinary(data) if err != nil { b.Fatal(err) } From 0abf799477973c825bc99de1228a04f7bb2ce783 Mon Sep 17 00:00:00 2001 From: aptend Date: Sat, 1 Aug 2026 20:03:58 +0800 Subject: [PATCH 43/61] executor: close statement result ownership boundaries --- .../batch/allocation_account_test.go | 42 ++++++ pkg/container/batch/batch.go | 24 +++ pkg/frontend/export.go | 55 +++++-- pkg/frontend/export_test.go | 71 +++++++++ pkg/frontend/types.go | 4 +- pkg/frontend/types_test.go | 47 ++++++ .../compile/allocation_account_lifecycle.go | 6 +- .../remote_allocation_statement_group.go | 13 +- .../remote_allocation_statement_group_test.go | 62 ++++++++ pkg/sql/compile/sql_executor.go | 63 ++++++-- .../compile/sql_executor_allocation_test.go | 138 ++++++++++++++++++ 11 files changed, 489 insertions(+), 36 deletions(-) create mode 100644 pkg/sql/compile/sql_executor_allocation_test.go diff --git a/pkg/container/batch/allocation_account_test.go b/pkg/container/batch/allocation_account_test.go index 723d38d2ef69c..ba710e90a90be 100644 --- a/pkg/container/batch/allocation_account_test.go +++ b/pkg/container/batch/allocation_account_test.go @@ -170,6 +170,48 @@ func TestBatchAllocationAccountCloneDupAndWindow(t *testing.T) { finalizeTestBatchAllocationAccount(t, state) } +func TestBatchDupWithoutAllocationAccountCrossesStatementBoundary(t *testing.T) { + state := newTestBatchAllocationAccount(t, 64) + mp := mpool.MustNewZero() + source := newBatchAllocationTestSource(t, mp, state.selection) + sourceUsed := state.account.Snapshot().Used + require.Positive(t, sourceUsed) + + cloned, err := source.DupWithoutAllocationAccount(mp) + require.NoError(t, err) + require.Nil(t, cloned.AllocationAccountSelection()) + for _, vec := range cloned.Vecs { + require.Nil(t, vec.AllocationAccountSelection()) + } + require.Equal(t, source.RowCount(), cloned.RowCount()) + require.Equal( + t, + int64(0), + vector.GetFixedAtNoTypeCheck[int64](cloned.Vecs[0], 0), + ) + require.Equal( + t, + []byte("batch allocation payload that is not inline"), + cloned.Vecs[1].GetBytesAt(0), + ) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) + require.Equal( + t, + int64(0), + vector.GetFixedAtNoTypeCheck[int64](cloned.Vecs[0], 0), + ) + require.Equal( + t, + []byte("batch allocation payload that is not inline"), + cloned.Vecs[1].GetBytesAt(0), + ) + cloned.Clean(mp) + require.Zero(t, mp.CurrNB()) +} + func TestBatchAccountedReaderAcceptsBitmapCapacityBeyondLogicalRows(t *testing.T) { state := newTestBatchAllocationAccount(t, 64) mp := mpool.MustNewZero() diff --git a/pkg/container/batch/batch.go b/pkg/container/batch/batch.go index e37767ace8fcd..d766f404f3f49 100644 --- a/pkg/container/batch/batch.go +++ b/pkg/container/batch/batch.go @@ -1158,6 +1158,30 @@ func (bat *Batch) Dup(mp *mpool.MPool) (*Batch, error) { return bat.Clone(mp, bat.offHeap || bat.hasAllocationAccountVector()) } +// CloneWithoutAllocationAccount deep-copies bat without carrying an +// allocation account into the destination. Use it only at an ownership +// boundary whose destination can outlive the source account. +func (bat *Batch) CloneWithoutAllocationAccount( + mp *mpool.MPool, + offHeap bool, +) (*Batch, error) { + attrs, attrTypes := bat.GetSchema() + cloned := NewWithSchema(offHeap, attrs, attrTypes) + cloned.Recursive = bat.Recursive + if err := bat.CloneTo(cloned, mp); err != nil { + return nil, err + } + return cloned, nil +} + +// DupWithoutAllocationAccount is the ownership-boundary counterpart of Dup. +func (bat *Batch) DupWithoutAllocationAccount(mp *mpool.MPool) (*Batch, error) { + return bat.CloneWithoutAllocationAccount( + mp, + bat.offHeap || bat.hasAllocationAccountVector(), + ) +} + func (bat *Batch) hasAllocationAccountVector() bool { for _, vec := range bat.Vecs { if vec != nil && vec.AllocationAccountSelection() != nil { diff --git a/pkg/frontend/export.go b/pkg/frontend/export.go index f2e46dfedbcdc..51216a78fa80c 100644 --- a/pkg/frontend/export.go +++ b/pkg/frontend/export.go @@ -377,6 +377,19 @@ func escapeJSONControlChars(s string) string { return builder.String() } +func sendExportBatchByte( + ctx context.Context, + byteChan chan *BatchByte, + value *BatchByte, +) bool { + select { + case byteChan <- value: + return true + case <-ctx.Done(): + return false + } +} + func constructByte(ctx context.Context, obj FeSession, bat *batch.Batch, index int32, ByteChan chan *BatchByte, ep *ExportConfig) { var ( ok bool @@ -398,15 +411,6 @@ func constructByte(ctx context.Context, obj FeSession, bat *batch.Batch, index i return } - sendByte := func(bb *BatchByte) bool { - select { - case ByteChan <- bb: - return true - case <-ctx.Done(): - return false - } - } - symbol := ep.Symbol closeby := ep.userConfig.Fields.EnclosedBy.Value flag := ep.ColumnFlag @@ -482,7 +486,7 @@ func constructByte(ctx context.Context, obj FeSession, bat *batch.Batch, index i case types.T_geometry, types.T_geometry32: text, err := planfunction.GeometryPayloadToText(vec.GetBytesAt(i)) if err != nil { - sendByte(&BatchByte{err: err}) + sendExportBatchByte(ctx, ByteChan, &BatchByte{err: err}) bat.Clean(mp) return } @@ -574,7 +578,7 @@ func constructByte(ctx context.Context, obj FeSession, bat *batch.Batch, index i } // stop early if downstream already failed - sendByte(&BatchByte{ + sendExportBatchByte(ctx, ByteChan, &BatchByte{ err: moerr.NewInternalErrorf(ctx, "constructByte : unsupported type %d", vec.GetType().Oid), }) bat.Clean(mp) @@ -589,7 +593,7 @@ func constructByte(ctx context.Context, obj FeSession, bat *batch.Batch, index i copy(result, buffer.Bytes()) buffer = nil - if !sendByte(&BatchByte{ + if !sendExportBatchByte(ctx, ByteChan, &BatchByte{ index: index, writeByte: result, err: nil, @@ -948,7 +952,10 @@ func (ec *ExportConfig) init() { func (ec *ExportConfig) Write(execCtx *ExecCtx, crs *perfcounter.CounterSet, bat *batch.Batch) error { ec.Index.Add(1) - copied, err := bat.Dup(execCtx.ses.GetMemPool()) + // CSV and JSON conversion runs asynchronously and can outlive runner.Run. + // The worker copy therefore belongs to the export pipeline, not to the + // producing statement's allocation account. + copied, err := cloneExportWorkerBatch(bat, execCtx.ses.GetMemPool()) if err != nil { return err } @@ -973,6 +980,13 @@ func (ec *ExportConfig) Write(execCtx *ExecCtx, crs *perfcounter.CounterSet, bat return nil } +func cloneExportWorkerBatch( + bat *batch.Batch, + mp *mpool.MPool, +) (*batch.Batch, error) { + return bat.DupWithoutAllocationAccount(mp) +} + // writeParquet writes a batch to the parquet writer func (ec *ExportConfig) writeParquet(execCtx *ExecCtx, bat *batch.Batch) error { defer bat.Clean(execCtx.ses.GetMemPool()) @@ -1156,8 +1170,11 @@ func constructJSONLine(ctx context.Context, obj FeSession, bat *batch.Batch, ind } val, err := vectorValueToJSON(vec, i, ss, backSes) if err != nil { - ByteChan <- &BatchByte{ + if !sendExportBatchByte(ctx, ByteChan, &BatchByte{ err: err, + }) { + bat.Clean(mp) + return } bat.Clean(mp) return @@ -1166,8 +1183,11 @@ func constructJSONLine(ctx context.Context, obj FeSession, bat *batch.Batch, ind } jsonBytes, err := json.Marshal(row) if err != nil { - ByteChan <- &BatchByte{ + if !sendExportBatchByte(ctx, ByteChan, &BatchByte{ err: moerr.NewInternalErrorf(ctx, "failed to marshal JSON: %v", err), + }) { + bat.Clean(mp) + return } bat.Clean(mp) return @@ -1181,10 +1201,13 @@ func constructJSONLine(ctx context.Context, obj FeSession, bat *batch.Batch, ind copy(result, buffer.Bytes()) buffer = nil - ByteChan <- &BatchByte{ + if !sendExportBatchByte(ctx, ByteChan, &BatchByte{ index: index, writeByte: result, err: nil, + }) { + bat.Clean(mp) + return } bat.Clean(mp) diff --git a/pkg/frontend/export_test.go b/pkg/frontend/export_test.go index f2a472fb2f8e0..ab006c0024eab 100644 --- a/pkg/frontend/export_test.go +++ b/pkg/frontend/export_test.go @@ -19,6 +19,7 @@ import ( "context" "strings" "testing" + "time" "github.com/prashantv/gostub" "github.com/smartystreets/goconvey/convey" @@ -27,9 +28,11 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/defines" "github.com/matrixorigin/matrixone/pkg/sql/parsers/tree" "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/stretchr/testify/require" ) var colName1, colName2 = "DATABASE()", "VARIABLE_VALUE" @@ -370,6 +373,74 @@ func TestConstructByteFormatsUnscaledFloat64WithFullPrecision(t *testing.T) { }) } +func TestExportWorkerBatchEndsStatementAllocationOwnership(t *testing.T) { + mp := mpool.MustNewZero() + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + + source := batch.NewWithSchema( + true, + []string{"value"}, + []types.Type{types.T_int64.ToType()}, + ) + require.NoError(t, source.SetAllocationAccount(selection)) + require.NoError(t, vector.AppendFixed(source.Vecs[0], int64(42), false, mp)) + source.SetRowCount(1) + require.Positive(t, account.Snapshot().Used) + + workerBatch, err := cloneExportWorkerBatch(source, mp) + require.NoError(t, err) + require.Nil(t, workerBatch.AllocationAccountSelection()) + require.Nil(t, workerBatch.Vecs[0].AllocationAccountSelection()) + + source.Clean(mp) + snapshot := account.Seal() + require.Zero(t, snapshot.Used) + require.Zero(t, registry.LiveAllocationMetadata()) + _, err = registry.Finalize(account) + require.NoError(t, err) + require.Equal(t, int64(42), vector.GetFixedAtNoTypeCheck[int64](workerBatch.Vecs[0], 0)) + workerBatch.Clean(mp) + require.Zero(t, mp.CurrNB()) +} + +func TestJSONExportCancellationCleansBlockedWorkerBatch(t *testing.T) { + mp := mpool.MustNewZero() + bat := batch.NewOffHeapWithSize(1) + bat.Vecs[0] = vector.NewOffHeapVecWithType(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixed(bat.Vecs[0], int64(42), false, mp)) + bat.SetRowCount(1) + require.Positive(t, mp.CurrNB()) + + mrs := &MysqlResultSet{} + column := &MysqlColumn{} + column.SetName("value") + mrs.AddColumn(column) + ep := &ExportConfig{mrs: mrs} + bytesChan := make(chan *BatchByte, 1) + bytesChan <- &BatchByte{index: 0} + ctx, cancel := context.WithCancel(context.Background()) + cancel() + ses := &backSession{feSessionImpl: feSessionImpl{pool: mp}} + + done := make(chan struct{}) + go func() { + defer close(done) + constructJSONLine(ctx, ses, bat, 1, bytesChan, ep) + }() + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("canceled JSON export worker remained blocked on a full channel") + } + require.Zero(t, mp.CurrNB()) +} + func Test_getExportFormat(t *testing.T) { convey.Convey("getExportFormat returns correct format", t, func() { // Test default format (empty string returns "csv") diff --git a/pkg/frontend/types.go b/pkg/frontend/types.go index 34db16b934b9e..308f03ba6d3cc 100644 --- a/pkg/frontend/types.go +++ b/pkg/frontend/types.go @@ -1359,7 +1359,9 @@ func (ses *feSessionImpl) GetResultBatches() []*batch.Batch { } func (ses *feSessionImpl) AppendResultBatch(bat *batch.Batch) error { - copied, err := bat.Dup(ses.pool) + // Result batches belong to the session and can remain reachable after the + // producing statement has sealed its allocation account. + copied, err := bat.DupWithoutAllocationAccount(ses.pool) if err != nil { return err } diff --git a/pkg/frontend/types_test.go b/pkg/frontend/types_test.go index 8d5c30c300386..395eb1bc89aaa 100644 --- a/pkg/frontend/types_test.go +++ b/pkg/frontend/types_test.go @@ -17,8 +17,11 @@ package frontend import ( "testing" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/config" + "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/sql/parsers/tree" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -392,6 +395,50 @@ func TestPrepareStmt_Close(t *testing.T) { ps.Close() } +func TestAppendResultBatchEndsStatementAllocationOwnership(t *testing.T) { + mp := mpool.MustNewZero() + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection( + account, + 1, + 1, + 2, + 3, + 4, + ) + require.NoError(t, err) + + source := batch.NewOffHeapWithSize(1) + source.Vecs[0] = vector.NewOffHeapVecWithType(types.T_int64.ToType()) + require.NoError(t, source.SetAllocationAccount(selection)) + require.NoError(t, vector.AppendFixed(source.Vecs[0], int64(42), false, mp)) + source.SetRowCount(1) + + ses := &feSessionImpl{pool: mp} + require.NoError(t, ses.AppendResultBatch(source)) + require.Len(t, ses.resultBatches, 1) + require.Nil(t, ses.resultBatches[0].AllocationAccountSelection()) + require.Nil(t, ses.resultBatches[0].Vecs[0].AllocationAccountSelection()) + + source.Clean(mp) + snapshot := account.Seal() + require.Zero(t, snapshot.Used) + require.Zero(t, registry.LiveAllocationMetadata()) + _, err = registry.Finalize(account) + require.NoError(t, err) + require.Equal( + t, + int64(42), + vector.GetFixedAtNoTypeCheck[int64](ses.resultBatches[0].Vecs[0], 0), + ) + + ses.ClearResultBatches() + require.Zero(t, mp.CurrNB()) +} + func BenchmarkSessionAllocator(b *testing.B) { allocator := NewSessionAllocator(&config.ParameterUnit{ SV: &config.FrontendParameters{ diff --git a/pkg/sql/compile/allocation_account_lifecycle.go b/pkg/sql/compile/allocation_account_lifecycle.go index e3281f7ebc599..3638fad8e9341 100644 --- a/pkg/sql/compile/allocation_account_lifecycle.go +++ b/pkg/sql/compile/allocation_account_lifecycle.go @@ -16,9 +16,9 @@ package compile import ( "errors" - "fmt" "sync" + "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/vm" "github.com/matrixorigin/matrixone/pkg/vm/message" @@ -31,7 +31,9 @@ func allocationLifecycleCall(call func() error) (err error) { err = errors.Join( err, mpool.ErrAllocationAccountInvariant, - fmt.Errorf("allocation lifecycle panic: %v", recovered), + moerr.NewInternalErrorNoCtxf( + "allocation lifecycle panic: %v", recovered, + ), ) } }() diff --git a/pkg/sql/compile/remote_allocation_statement_group.go b/pkg/sql/compile/remote_allocation_statement_group.go index 46f43d7099b06..64616d3080cb2 100644 --- a/pkg/sql/compile/remote_allocation_statement_group.go +++ b/pkg/sql/compile/remote_allocation_statement_group.go @@ -183,9 +183,16 @@ func acquireRemoteAllocationStatementParticipant( } group.registered++ group.participants = append(group.participants, participant) - if group.registered == group.expected && group.timer != nil { - group.timer.Stop() - group.timer = nil + if group.registered == group.expected { + if group.timer != nil { + group.timer.Stop() + group.timer = nil + } + } else if group.timer == nil { + group.timer = time.AfterFunc( + remoteAllocationStatementRegistrationTimeout, + func() { expireRemoteAllocationStatementGroup(group) }, + ) } return participant, nil } diff --git a/pkg/sql/compile/remote_allocation_statement_group_test.go b/pkg/sql/compile/remote_allocation_statement_group_test.go index 110a0948af125..370873c942d71 100644 --- a/pkg/sql/compile/remote_allocation_statement_group_test.go +++ b/pkg/sql/compile/remote_allocation_statement_group_test.go @@ -252,6 +252,68 @@ func TestRemoteAllocationStatementGroupExpiresMissingFragment(t *testing.T) { require.False(t, registry.AdmissionSuspended()) } +func TestRemoteAllocationStatementRegistrationTimerStartsBeforeFinish(t *testing.T) { + previousTimeout := remoteAllocationStatementRegistrationTimeout + remoteAllocationStatementRegistrationTimeout = 10 * time.Millisecond + t.Cleanup(func() { + remoteAllocationStatementRegistrationTimeout = previousTimeout + }) + + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + board := message.NewMessageBoard() + producer := newTestAllocationLifecycleCompile(t, registry, func( + mpool.AllocationAccountTerminalSnapshot, + ) { + t.Fatal("expired remote statement group must own terminal export") + }) + producer.MessageBoard = board + attempt, err := producer.beginAllocationAccountAttempt() + require.NoError(t, err) + buffer, err := producer.proc.Mp().AllocAccounted( + 64, + attempt.account, + mpool.AllocationOwner(1), + mpool.AllocationSite(1), + ) + require.NoError(t, err) + var destroyed atomic.Int32 + message.SendMessage(&remoteAllocationAccountedMessage{ + mp: producer.proc.Mp(), + buffer: buffer, + destroyed: &destroyed, + }, board) + + canceled := make(chan error, 1) + participant, err := acquireRemoteAllocationStatementParticipant( + board, + 2, + func(cause error) { canceled <- cause }, + ) + require.NoError(t, err) + + select { + case cause := <-canceled: + require.Error(t, cause) + case <-time.After(time.Second): + t.Fatal("registration timeout did not cancel the active fragment") + } + require.Eventually(t, func() bool { + return destroyed.Load() == 1 && strings.Contains(board.DebugString(), "closed") + }, time.Second, time.Millisecond) + + participant.stage(attempt, producer.proc.Mp()) + terminal, err := participant.finish(errors.New("active fragment observed cancellation")) + require.Error(t, err) + require.True(t, terminal.complete) + require.Len(t, terminal.allocation, 1) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.allocation[0].State) + require.Zero(t, terminal.allocation[0].Used) + require.Zero(t, terminal.memory.LiveBytesAtSeal) + require.Zero(t, registry.LiveAllocationMetadata()) + require.False(t, remoteAllocationStatementGroupRegistered(board)) +} + func TestRemoteAllocationStatementGroupFailureAbortsMissingFragment(t *testing.T) { registry, err := mpool.NewAllocationAccountRegistry(1, 1) require.NoError(t, err) diff --git a/pkg/sql/compile/sql_executor.go b/pkg/sql/compile/sql_executor.go index ac962d14f1c7d..605cb7d6b05e5 100644 --- a/pkg/sql/compile/sql_executor.go +++ b/pkg/sql/compile/sql_executor.go @@ -534,26 +534,28 @@ func (exec *txnExecutor) Exec( // the bat is valid only in current method. So we need copy data. // FIXME: add a custom streaming apply handler to consume readed data. Now // our current internal sql will never read too much data. - rows, err := bat.Clone(exec.s.mp, streaming) + // Internal executor results outlive Compile.Run. Keep the + // existing physical clone mode, but end statement allocation + // ownership before returning or publishing the batch. + rows, err := cloneInternalExecutorResultBatch( + bat, + exec.s.mp, + streaming, + ) if err != nil { return err } if streaming { stream_result := executor.NewResult(exec.s.mp) - for len(stream_chan) == cap(stream_chan) { - select { - case <-proc.Ctx.Done(): - err_chan <- moerr.NewInternalError(proc.Ctx, "context cancelled") - return moerr.NewInternalError(proc.Ctx, "context cancelled") - case <-exec.ctx.Done(): - err_chan <- exec.ctx.Err() - return exec.ctx.Err() - default: - time.Sleep(1 * time.Millisecond) - } - } stream_result.Batches = []*batch.Batch{rows} - stream_chan <- stream_result + if err := publishInternalExecutorStreamResult( + proc.Ctx, + exec.ctx, + stream_chan, + stream_result, + ); err != nil { + return err + } } else { batches = append(batches, rows) } @@ -599,6 +601,39 @@ func (exec *txnExecutor) Exec( return result, nil } +func cloneInternalExecutorResultBatch( + bat *batch.Batch, + mp *mpool.MPool, + streaming bool, +) (*batch.Batch, error) { + return bat.CloneWithoutAllocationAccount(mp, streaming) +} + +// publishInternalExecutorStreamResult transfers result ownership only after a +// successful send. Cancellation keeps ownership here and closes the result. +func publishInternalExecutorStreamResult( + procCtx context.Context, + execCtx context.Context, + streamChan chan executor.Result, + result executor.Result, +) error { + select { + case streamChan <- result: + return nil + default: + } + select { + case streamChan <- result: + return nil + case <-procCtx.Done(): + result.Close() + return moerr.NewInternalError(procCtx, "context cancelled") + case <-execCtx.Done(): + result.Close() + return execCtx.Err() + } +} + func (exec *txnExecutor) LockTable(table string) error { txnOp := exec.opts.Txn() ctx := exec.ctx diff --git a/pkg/sql/compile/sql_executor_allocation_test.go b/pkg/sql/compile/sql_executor_allocation_test.go new file mode 100644 index 0000000000000..7382860e6292f --- /dev/null +++ b/pkg/sql/compile/sql_executor_allocation_test.go @@ -0,0 +1,138 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package compile + +import ( + "context" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/util/executor" + "github.com/stretchr/testify/require" +) + +func TestInternalExecutorResultEndsStatementAllocationOwnership(t *testing.T) { + for _, streaming := range []bool{false, true} { + t.Run(map[bool]string{false: "retained", true: "streaming"}[streaming], func(t *testing.T) { + mp := mpool.MustNewZero() + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection( + account, + 1, + 1, + 2, + 3, + 4, + ) + require.NoError(t, err) + + source := batch.NewWithSchema( + true, + []string{"id", "value"}, + []types.Type{types.T_int64.ToType(), types.T_varchar.ToType()}, + ) + require.NoError(t, source.SetAllocationAccount(selection)) + require.NoError(t, vector.AppendFixed(source.Vecs[0], int64(42), false, mp)) + require.NoError(t, vector.AppendBytes(source.Vecs[1], []byte("result"), false, mp)) + source.SetRowCount(1) + + cloned, err := cloneInternalExecutorResultBatch(source, mp, streaming) + require.NoError(t, err) + require.Nil(t, cloned.AllocationAccountSelection()) + for _, vec := range cloned.Vecs { + require.Nil(t, vec.AllocationAccountSelection()) + } + + var result executor.Result + if streaming { + results := make(chan executor.Result, 1) + published := executor.NewResult(mp) + published.Batches = []*batch.Batch{cloned} + results <- published + result = <-results + } else { + result = executor.NewResult(mp) + result.Batches = []*batch.Batch{cloned} + } + + source.Clean(mp) + snapshot := account.Seal() + require.Zero(t, snapshot.Used) + require.Zero(t, registry.LiveAllocationMetadata()) + _, err = registry.Finalize(account) + require.NoError(t, err) + + require.Equal( + t, + int64(42), + vector.GetFixedAtNoTypeCheck[int64](result.Batches[0].Vecs[0], 0), + ) + require.Equal(t, []byte("result"), result.Batches[0].Vecs[1].GetBytesAt(0)) + result.Close() + require.Zero(t, mp.CurrNB()) + }) + } +} + +func TestInternalExecutorStreamCancellationCleansUnpublishedResult(t *testing.T) { + mp := mpool.MustNewZero() + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + + source := batch.NewWithSchema( + true, + []string{"value"}, + []types.Type{types.T_int64.ToType()}, + ) + require.NoError(t, source.SetAllocationAccount(selection)) + require.NoError(t, vector.AppendFixed(source.Vecs[0], int64(42), false, mp)) + source.SetRowCount(1) + + cloned, err := cloneInternalExecutorResultBatch(source, mp, true) + require.NoError(t, err) + result := executor.NewResult(mp) + result.Batches = []*batch.Batch{cloned} + results := make(chan executor.Result, 1) + results <- executor.NewResult(mp) + procCtx, cancel := context.WithCancel(context.Background()) + cancel() + + err = publishInternalExecutorStreamResult( + procCtx, + context.Background(), + results, + result, + ) + require.Error(t, err) + require.Len(t, results, 1) + + source.Clean(mp) + snapshot := account.Seal() + require.Zero(t, snapshot.Used) + require.Zero(t, registry.LiveAllocationMetadata()) + _, err = registry.Finalize(account) + require.NoError(t, err) + require.Zero(t, mp.CurrNB()) +} From 6d26e8bd44bb7ed1ee31badc6f2934038d055e78 Mon Sep 17 00:00:00 2001 From: aptend Date: Sat, 1 Aug 2026 21:35:59 +0800 Subject: [PATCH 44/61] executor: reconcile physical admission after main rebase --- pkg/sql/colexec/fuzzyfilter/filter_test.go | 4 +- .../hashbuild/allocation_test_helpers_test.go | 19 +++- pkg/sql/colexec/hashbuild/build_test.go | 106 ++++++++---------- pkg/sql/colexec/hashbuild/errors.go | 20 +++- pkg/sql/colexec/hashbuild/errors_test.go | 18 +++ pkg/sql/colexec/indexbuild/build_test.go | 4 +- pkg/sql/colexec/rightdedupjoin/join_test.go | 4 +- .../colexec/runtimefilter/contract_test.go | 4 +- pkg/vm/process/hashbuild_budget.go | 2 +- 9 files changed, 107 insertions(+), 74 deletions(-) diff --git a/pkg/sql/colexec/fuzzyfilter/filter_test.go b/pkg/sql/colexec/fuzzyfilter/filter_test.go index c84ffd84b5d9a..631e485f5f5ff 100644 --- a/pkg/sql/colexec/fuzzyfilter/filter_test.go +++ b/pkg/sql/colexec/fuzzyfilter/filter_test.go @@ -387,7 +387,7 @@ func TestFuzzyRuntimeFilterBudgetErrorPolicy(t *testing.T) { registry, err := mpool.NewAllocationAccountRegistry(1, 16) require.NoError(t, err) account, err := registry.OpenWithController( - generation.Cap(), generation) + 2*generation.Cap(), generation) require.NoError(t, err) require.NoError(t, arg.SetAllocationAccount(account)) prepareFuzzyFilter(t, arg, proc) @@ -398,7 +398,7 @@ func TestFuzzyRuntimeFilterBudgetErrorPolicy(t *testing.T) { if test.closed { generation.Close() } else { - remaining := account.Snapshot().Limit - account.Snapshot().Used + remaining := generation.Cap() - generation.Used() filler, err = proc.Mp().AllocAccounted( int(remaining), account, 63, 255) require.NoError(t, err) diff --git a/pkg/sql/colexec/hashbuild/allocation_test_helpers_test.go b/pkg/sql/colexec/hashbuild/allocation_test_helpers_test.go index 5e7545f95564d..c06b9c0da4696 100644 --- a/pkg/sql/colexec/hashbuild/allocation_test_helpers_test.go +++ b/pkg/sql/colexec/hashbuild/allocation_test_helpers_test.go @@ -39,13 +39,28 @@ func installTestHashBuildBudget( t.Helper() registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) require.NoError(t, err) - account, err := registry.OpenWithController( - generation.Cap(), generation) + account, err := registry.OpenWithController(1<<60, generation) require.NoError(t, err) replaceTestHashBuildAllocation(t, op, account) op.ctr.hashmapBuilder.setBudget(generation) } +func installTestProcessHashBuildBudget( + t testing.TB, + op *HashBuild, + proc *process.Process, +) *process.HashBuildBudgetGeneration { + t.Helper() + generation, err := proc.GetHashBuildBudget() + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.OpenWithController(1<<60, generation) + require.NoError(t, err) + replaceTestHashBuildAllocation(t, op, account) + return generation +} + func newTestHashmapBuilder(t testing.TB) *HashmapBuilder { t.Helper() builder := &HashmapBuilder{} diff --git a/pkg/sql/colexec/hashbuild/build_test.go b/pkg/sql/colexec/hashbuild/build_test.go index 7646953de74ec..664e3b808ac2c 100644 --- a/pkg/sql/colexec/hashbuild/build_test.go +++ b/pkg/sql/colexec/hashbuild/build_test.go @@ -168,6 +168,7 @@ func TestHashBuildRepeatedResetFinalizesRuntimeFilterOnce(t *testing.T) { func TestBroadcastBudgetFailureUnblocksAllConsumers(t *testing.T) { tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) + installTestProcessHashBuildBudget(t, tc.arg, tc.proc) tc.arg.SetChildren([]vm.Operator{tc.marg}) require.NoError(t, tc.marg.Prepare(tc.proc)) require.NoError(t, tc.arg.Prepare(tc.proc)) @@ -232,6 +233,7 @@ func TestHashBuildPrepareConvertsTerminalBudgetAdmission(t *testing.T) { Idx: 0, }}, } + installTestProcessHashBuildBudget(t, arg, proc) var prepareErr error t.Cleanup(func() { arg.Free(proc, true, prepareErr) @@ -254,6 +256,7 @@ func TestHashBuildPrepareConvertsTerminalBudgetAdmission(t *testing.T) { func TestHashBuildWithoutMapStillBudgetsRetainedBatches(t *testing.T) { tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, nil) + installTestProcessHashBuildBudget(t, tc.arg, tc.proc) tc.arg.NeedHashMap = false tc.arg.NeedBatches = true tc.arg.SetChildren([]vm.Operator{tc.marg}) @@ -618,14 +621,13 @@ func TestHashBuildOptionalRuntimeFilterCollectionFallsBackToJoinMap( tc.arg.RuntimeFilterSpec = rawRuntimeFilterSpec( tc.arg.JoinMapTag+500, 100, typ) tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - const capBytes = uint64(64 << 20) aggregate := process.MustNewHashBuildBudget(capBytes, capBytes) generation, err := aggregate.OpenGeneration(1) require.NoError(t, err) installTestHashBuildBudget(t, tc.arg, generation) + require.NoError(t, tc.marg.Prepare(tc.proc)) + require.NoError(t, tc.arg.Prepare(tc.proc)) providerCalls := 0 forcedCollectionReject := false @@ -710,14 +712,13 @@ func TestHashBuildClosedMapBudgetDoesNotRecordCollectionFallback( tc.arg.RuntimeFilterSpec = rawRuntimeFilterSpec( tc.arg.JoinMapTag+501, 100, typ) tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - const capBytes = uint64(64 << 20) aggregate := process.MustNewHashBuildBudget(capBytes, capBytes) generation, err := aggregate.OpenGeneration(1) require.NoError(t, err) installTestHashBuildBudget(t, tc.arg, generation) + require.NoError(t, tc.marg.Prepare(tc.proc)) + require.NoError(t, tc.arg.Prepare(tc.proc)) providerCalls := 0 forcedClosed := false @@ -784,13 +785,12 @@ func TestHashmapBuilderUniqueGrowthFailureAbandonsOptionalKeysInPlace( []types.Type{typ}, []*plan.Expr{newExpr(0, typ)}, ) - require.NoError(t, tc.arg.Prepare(tc.proc)) - const capBytes = uint64(64 << 20) aggregate := process.MustNewHashBuildBudget(capBytes, capBytes) generation, err := aggregate.OpenGeneration(1) require.NoError(t, err) installTestHashBuildBudget(t, tc.arg, generation) + require.NoError(t, tc.arg.Prepare(tc.proc)) const uniqueGrowthRows = hashmap.UnitLimit * 2 input := newBatch( @@ -990,14 +990,13 @@ func TestShuffleDedupAdmissionAfterRewriteDoesNotSpillPartialInput( Tag: tc.arg.JoinMapTag + 700, } tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - const capBytes = uint64(64 << 20) aggregate := process.MustNewHashBuildBudget(capBytes, capBytes) generation, err := aggregate.OpenGeneration(1) require.NoError(t, err) installTestHashBuildBudget(t, tc.arg, generation) + require.NoError(t, tc.marg.Prepare(tc.proc)) + require.NoError(t, tc.arg.Prepare(tc.proc)) // Ingress happens before BuildHashmap initializes this phase. Mark the // retained source safe so the provider rejects only after Dedup crosses // its explicit in-place rewrite boundary. @@ -1206,6 +1205,10 @@ func TestHashBuildFloatRuntimeFilterAllocationFailureFallsBackToPass(t *testing. RuntimeFilterSpec: spec, } arg.OpAnalyzer = process.NewAnalyzer(0, false, false, "hash build") + budget := process.MustNewHashBuildBudget(64<<20, 64<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + installTestHashBuildBudget(t, arg, generation) keyVec := vector.NewOffHeapVecWithType(typ) require.NoError(t, keyVec.PreExtend(256, mp)) @@ -1213,11 +1216,6 @@ func TestHashBuildFloatRuntimeFilterAllocationFailureFallsBackToPass(t *testing. arg.ctr.hashmapBuilder.InputBatchRowCount = keyVec.Length() arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{keyVec} - budget := process.MustNewHashBuildBudget(64<<20, 64<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - installTestHashBuildBudget(t, arg, generation) - var filler []byte defer func() { if filler != nil { @@ -1516,15 +1514,14 @@ func TestRuntimeFilterExplicitDecimalContractProducesIn(t *testing.T) { spec := rawRuntimeFilterSpec(105, 100, decimalType) tc.arg.RuntimeFilterSpec = spec tc.arg.ctr.hashmapBuilder.InputBatchRowCount = 1 - payload := vector.NewVec(decimalType) - require.NoError(t, vector.AppendFixed( - payload, types.Decimal64(1000), false, tc.proc.Mp())) - tc.arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{payload} budget := process.MustNewHashBuildBudget(1<<20, 1<<20) generation, err := budget.OpenGeneration(1) require.NoError(t, err) installTestHashBuildBudget(t, tc.arg, generation) - + payload := vector.NewVec(decimalType) + require.NoError(t, vector.AppendFixed( + payload, types.Decimal64(1000), false, tc.proc.Mp())) + tc.arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{payload} require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) require.True(t, tc.arg.ctr.runtimeFilterDone) require.True(t, tc.arg.ctr.runtimeFilterIn) @@ -1564,15 +1561,14 @@ func TestDirectRuntimeFilterUsesDeclaredHashSlot(t *testing.T) { spec.BuildExpr = newExpr(1, typ) tc.arg.RuntimeFilterSpec = spec tc.arg.ctr.hashmapBuilder.InputBatchRowCount = 2 - tc.arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{ - testutil.MakeInt32Vector([]int32{901, 902}, nil, tc.proc.Mp()), - testutil.MakeInt32Vector([]int32{11, 12}, nil, tc.proc.Mp()), - } budget := process.MustNewHashBuildBudget(1<<20, 1<<20) generation, err := budget.OpenGeneration(1) require.NoError(t, err) installTestHashBuildBudget(t, tc.arg, generation) - + tc.arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{ + testutil.MakeInt32Vector([]int32{901, 902}, nil, tc.proc.Mp()), + testutil.MakeInt32Vector([]int32{11, 12}, nil, tc.proc.Mp()), + } require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) receiver := message.NewMessageReceiver( []int32{spec.Tag}, @@ -1650,6 +1646,10 @@ func TestHashBuildSerializedRuntimeFilterAllocationFailureFallsBackToPass(t *tes RuntimeFilterSpec: spec, } arg.OpAnalyzer = process.NewAnalyzer(0, false, false, "hash build") + budget := process.MustNewHashBuildBudget(64<<20, 64<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + installTestHashBuildBudget(t, arg, generation) arg.ctr.hashmapBuilder.InputBatchRowCount = 1 arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{ testutil.MakeInt32Vector([]int32{1}, nil, mp), @@ -1668,11 +1668,6 @@ func TestHashBuildSerializedRuntimeFilterAllocationFailureFallsBackToPass(t *tes } }) - budget := process.MustNewHashBuildBudget(64<<20, 64<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - installTestHashBuildBudget(t, arg, generation) - var filler []byte defer func() { if filler != nil { @@ -1734,6 +1729,10 @@ func TestSerializedRuntimeFilterUsesTightBudgetAndProducesIn(t *testing.T) { []types.Type{componentType, componentType}, true) tc.arg.RuntimeFilterSpec = spec tc.arg.ctr.hashmapBuilder.InputBatchRowCount = rowCount + budget := process.MustNewHashBuildBudget(512<<10, 512<<10) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + installTestHashBuildBudget(t, tc.arg, generation) first := make([]int32, rowCount) second := make([]int32, rowCount) @@ -1762,11 +1761,6 @@ func TestSerializedRuntimeFilterUsesTightBudgetAndProducesIn(t *testing.T) { // The generic VARCHAR(max) estimator would request roughly 64 MiB for // these tiny tuples. The tuple-specific bound must fit comfortably here. - budget := process.MustNewHashBuildBudget(512<<10, 512<<10) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - installTestHashBuildBudget(t, tc.arg, generation) - require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) require.True(t, tc.arg.ctr.runtimeFilterDone) require.True(t, tc.arg.ctr.runtimeFilterIn) @@ -1982,14 +1976,13 @@ func TestRuntimeFilterMarshalBudgetAdmissionFallsBackToPass(t *testing.T) { tc.arg.RuntimeFilterSpec = spec tc.arg.OpAnalyzer = process.NewAnalyzer(0, false, false, "hash build") tc.arg.ctr.hashmapBuilder.InputBatchRowCount = 1 - tc.arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{ - testutil.MakeInt32Vector([]int32{1}, nil, tc.proc.Mp()), - } - budget := process.MustNewHashBuildBudget(1, 1) generation, err := budget.OpenGeneration(1) require.NoError(t, err) installTestHashBuildBudget(t, tc.arg, generation) + tc.arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{ + testutil.MakeInt32Vector([]int32{1}, nil, tc.proc.Mp()), + } require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) require.True(t, tc.arg.ctr.runtimeFilterDone) @@ -2028,11 +2021,7 @@ func TestRuntimeFilterMarshalUsesSinglePayloadBudget(t *testing.T) { tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) vec := testutil.MakeInt32Vector([]int32{1, 2, 3, 4}, nil, tc.proc.Mp()) - wireBytes := uint64(1+len(types.EncodeType(vec.GetType()))+4*4+1) + - uint64(len(vec.GetData())+len(vec.GetArea())) - projected := wireBytes + 64<<10 - - budget := process.MustNewHashBuildBudget(projected, projected) + budget := process.MustNewHashBuildBudget(1<<20, 1<<20) generation, err := budget.OpenGeneration(1) require.NoError(t, err) installTestHashBuildBudget(t, tc.arg, generation) @@ -2040,8 +2029,8 @@ func TestRuntimeFilterMarshalUsesSinglePayloadBudget(t *testing.T) { data, release, err := tc.arg.ctr.hashmapBuilder.marshalRuntimeFilterVector(vec, tc.proc.Mp()) require.NoError(t, err) require.NotEmpty(t, data) - require.Equal(t, projected, generation.Peak()) - require.LessOrEqual(t, generation.Used(), projected) + require.Equal(t, uint64(cap(data)), generation.Peak()) + require.Equal(t, uint64(cap(data)), generation.Used()) require.NotNil(t, release) release() require.Zero(t, generation.Used()) @@ -2060,11 +2049,7 @@ func TestRuntimeFilterMarshalSinglePayloadCoversVarlenaPeak(t *testing.T) { values[i] = strings.Repeat("x", 1024+i) } vec := testutil.MakeVarcharVector(values, nil, tc.proc.Mp()) - wireBytes := uint64(1+len(types.EncodeType(vec.GetType()))+4*4+1) + - uint64(len(vec.GetData())+len(vec.GetArea())) - projected := wireBytes + 64<<10 - - budget := process.MustNewHashBuildBudget(projected, projected) + budget := process.MustNewHashBuildBudget(1<<20, 1<<20) generation, err := budget.OpenGeneration(1) require.NoError(t, err) installTestHashBuildBudget(t, tc.arg, generation) @@ -2072,8 +2057,8 @@ func TestRuntimeFilterMarshalSinglePayloadCoversVarlenaPeak(t *testing.T) { data, release, err := tc.arg.ctr.hashmapBuilder.marshalRuntimeFilterVector(vec, tc.proc.Mp()) require.NoError(t, err) require.NotEmpty(t, data) - require.Equal(t, projected, generation.Peak()) - require.LessOrEqual(t, generation.Used(), projected) + require.Equal(t, uint64(cap(data)), generation.Peak()) + require.Equal(t, uint64(cap(data)), generation.Used()) release() require.Zero(t, generation.Used()) @@ -2088,7 +2073,7 @@ func TestRuntimeFilterMarshalAccountedPayloadMessageLifecycle(t *testing.T) { []*plan.Expr{newExpr(0, types.T_varchar.ToType())}) vec := testutil.MakeVarcharVector( []string{strings.Repeat("x", 4<<10), strings.Repeat("y", 8<<10)}, - []uint64{1}, + nil, tc.proc.Mp(), ) @@ -2098,7 +2083,7 @@ func TestRuntimeFilterMarshalAccountedPayloadMessageLifecycle(t *testing.T) { require.NoError(t, err) registry, err := mpool.NewAllocationAccountRegistry(1, 16) require.NoError(t, err) - account, err := registry.OpenWithController(limit, generation) + account, err := registry.OpenWithController(2*limit, generation) require.NoError(t, err) tc.arg.NeedHashMap = true replaceTestHashBuildAllocation(t, tc.arg, account) @@ -2154,16 +2139,13 @@ func TestRuntimeFilterMarshalAccountedOneByteShortFallsBackToPass(t *testing.T) require.NoError(t, err) registry, err := mpool.NewAllocationAccountRegistry(1, 16) require.NoError(t, err) - account, err := registry.OpenWithController(limit, generation) + account, err := registry.OpenWithController(2*limit, generation) require.NoError(t, err) tc.arg.NeedHashMap = true replaceTestHashBuildAllocation(t, tc.arg, account) tc.arg.ctr.hashmapBuilder.setBudget(generation) - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{ - Tag: 104, - UpperLimit: 100, - Expr: newExpr(0, types.T_int32.ToType()), - } + tc.arg.RuntimeFilterSpec = rawRuntimeFilterSpec( + 104, 100, types.T_int32.ToType()) tc.arg.OpAnalyzer = process.NewAnalyzer(0, false, false, "hash build") tc.arg.ctr.hashmapBuilder.InputBatchRowCount = 4 tc.arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{ diff --git a/pkg/sql/colexec/hashbuild/errors.go b/pkg/sql/colexec/hashbuild/errors.go index 39f432e147eea..0b9af9faf0c5f 100644 --- a/pkg/sql/colexec/hashbuild/errors.go +++ b/pkg/sql/colexec/hashbuild/errors.go @@ -20,6 +20,7 @@ import ( "strings" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/vm/process" ) @@ -27,7 +28,7 @@ import ( // operator's public Call boundary. Spill and other recovery paths must keep the // typed admission error until they have exhausted every recovery option. func TerminalBudgetError(ctx context.Context, err error) error { - if err == nil || !errors.Is(err, process.ErrHashBuildBudgetAdmission) { + if err == nil { return err } // A joined lifecycle/accounting failure is not a capacity rejection. Keep @@ -40,6 +41,23 @@ func TerminalBudgetError(ctx context.Context, err error) error { var budgetErr *process.HashBuildBudgetError if !errors.As(err, &budgetErr) || budgetErr.Kind != process.HashBuildBudgetErrorAdmission { + switch { + case mpool.AllocationFailureReasonOf(err) == + mpool.AllocationFailureCapacity: + return moerr.NewResourceExhaustedf(ctx, + "hash build memory budget exceeded; reduce join build width or query concurrency, increase processLimitationSize, or lower join_spill_mem for an eligible shuffle join") + case errors.Is(err, process.ErrHashBuildBudgetAdmission): + return moerr.NewResourceExhaustedf(ctx, + "hash build resource budget exceeded; inspect hash-build budget metrics and resource limits") + default: + return err + } + } + if budgetErr.Component == 0 { + reason := terminalBudgetReason(budgetErr.Message) + if reason != "" { + return moerr.NewResourceExhaustedf(ctx, "%s", reason) + } return moerr.NewResourceExhaustedf( ctx, "hash build resource budget exceeded; inspect hash-build budget metrics and resource limits", diff --git a/pkg/sql/colexec/hashbuild/errors_test.go b/pkg/sql/colexec/hashbuild/errors_test.go index d4735e2734c90..2d55a817fa9d0 100644 --- a/pkg/sql/colexec/hashbuild/errors_test.go +++ b/pkg/sql/colexec/hashbuild/errors_test.go @@ -20,6 +20,7 @@ import ( "testing" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/vm/process" "github.com/stretchr/testify/require" ) @@ -91,6 +92,23 @@ func TestTerminalBudgetError(t *testing.T) { require.NotContains(t, err.Error(), process.ErrHashBuildBudgetAdmission.Error()) }) + t.Run("physical capacity is terminal resource exhaustion", func(t *testing.T) { + err := TerminalBudgetError( + context.Background(), mpool.ErrAllocationAccountCapacity) + require.True(t, moerr.IsMoErrCode(err, moerr.ErrOOM)) + require.Contains(t, err.Error(), "hash build memory budget exceeded") + require.Contains(t, err.Error(), "processLimitationSize") + }) + + t.Run("physical lifecycle failure stays fatal", func(t *testing.T) { + joined := errors.Join( + mpool.ErrAllocationAccountCapacity, + mpool.ErrAllocationAccountSealed, + ) + require.Same(t, joined, + TerminalBudgetError(context.Background(), joined)) + }) + for _, lifecycle := range []error{ process.ErrHashBuildBudgetClosed, process.ErrHashBuildBudgetInvalid, diff --git a/pkg/sql/colexec/indexbuild/build_test.go b/pkg/sql/colexec/indexbuild/build_test.go index da49eb2107f32..4da04788d46b0 100644 --- a/pkg/sql/colexec/indexbuild/build_test.go +++ b/pkg/sql/colexec/indexbuild/build_test.go @@ -482,7 +482,7 @@ func TestIndexBuildRuntimeFilterBudgetErrorPolicy(t *testing.T) { registry, err := mpool.NewAllocationAccountRegistry(1, 16) require.NoError(t, err) account, err := registry.OpenWithController( - generation.Cap(), generation) + 2*generation.Cap(), generation) require.NoError(t, err) require.NoError(t, arg.SetAllocationAccount(account)) prepareIndexBuild(t, arg, proc) @@ -495,7 +495,7 @@ func TestIndexBuildRuntimeFilterBudgetErrorPolicy(t *testing.T) { if test.closed { generation.Close() } else { - remaining := account.Snapshot().Limit - account.Snapshot().Used + remaining := generation.Cap() - generation.Used() filler, err = proc.Mp().AllocAccounted( int(remaining), account, 63, 255) require.NoError(t, err) diff --git a/pkg/sql/colexec/rightdedupjoin/join_test.go b/pkg/sql/colexec/rightdedupjoin/join_test.go index 31adeb1c2d067..19702a8f41617 100644 --- a/pkg/sql/colexec/rightdedupjoin/join_test.go +++ b/pkg/sql/colexec/rightdedupjoin/join_test.go @@ -356,8 +356,8 @@ func TestRightDedupEmptyBuildProbeMapHonorsHashBuildBudget(t *testing.T) { require.NotContains(t, callErr.Error(), "convert go error") require.NotContains(t, callErr.Error(), process.ErrHashBuildBudgetAdmission.Error()) require.Contains(t, callErr.Error(), "hash build memory budget exceeded") - require.Equal(t, initialBytes, budget.Used(), - "the admitted initial table remains owned until operator cleanup") + require.Zero(t, budget.Used(), + "failed probe-map construction must roll back its physical allocation") } var ( diff --git a/pkg/sql/colexec/runtimefilter/contract_test.go b/pkg/sql/colexec/runtimefilter/contract_test.go index ae822b573fc3f..5b40139a8fae4 100644 --- a/pkg/sql/colexec/runtimefilter/contract_test.go +++ b/pkg/sql/colexec/runtimefilter/contract_test.go @@ -274,7 +274,7 @@ func TestMarshalExactFilterVectorUsesWireSizedBudget(t *testing.T) { require.NoError(t, err) registry, err := mpool.NewAllocationAccountRegistry(1, 16) require.NoError(t, err) - account, err := registry.OpenWithController(budget.Cap(), budget) + account, err := registry.OpenWithController(1<<20, budget) require.NoError(t, err) data, release, err := MarshalExactFilterVector(vec, mp, account, 1, 1) require.NoError(t, err) @@ -301,7 +301,7 @@ func TestMarshalExactFilterVectorAdmissionFailsBeforeAllocation(t *testing.T) { require.NoError(t, err) registry, err := mpool.NewAllocationAccountRegistry(1, 16) require.NoError(t, err) - account, err := registry.OpenWithController(budget.Cap(), budget) + account, err := registry.OpenWithController(1<<20, budget) require.NoError(t, err) data, release, err := MarshalExactFilterVector(vec, mp, account, 1, 1) require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) diff --git a/pkg/vm/process/hashbuild_budget.go b/pkg/vm/process/hashbuild_budget.go index 266d71b5c43c9..a5f06f216cdb2 100644 --- a/pkg/vm/process/hashbuild_budget.go +++ b/pkg/vm/process/hashbuild_budget.go @@ -152,7 +152,7 @@ const ( type HashBuildBudgetComponent uint8 const ( - HashBuildBudgetComponentMemory HashBuildBudgetComponent = iota + HashBuildBudgetComponentMemory HashBuildBudgetComponent = iota + 1 HashBuildBudgetComponentSpillDisk HashBuildBudgetComponentSpillFD ) From ea5b314cd247d2961128ccfc1c43c9b07921ff46 Mon Sep 17 00:00:00 2001 From: aptend Date: Sat, 1 Aug 2026 22:43:49 +0800 Subject: [PATCH 45/61] executor: close allocation accounting review gaps --- pkg/common/moerr/error.go | 9 ++ pkg/common/moerr/error_no_ctx.go | 4 + pkg/common/moerr/error_test.go | 13 ++ pkg/common/mpool/allocation_account.go | 121 ++------------- .../mpool/allocation_account_mpool_test.go | 113 +------------- pkg/common/mpool/allocation_account_test.go | 50 ++---- pkg/common/mpool/mpool.go | 87 ++++------- pkg/container/types/packer.go | 61 ++++++-- pkg/container/types/packer_test.go | 60 +++++++- .../vector/allocation_account_test.go | 66 ++++++++ pkg/container/vector/vector.go | 44 ++++-- pkg/pb/pipeline/error_test.go | 50 ++++++ pkg/sql/colexec/fuzzyfilter/filter.go | 2 +- pkg/sql/colexec/hashbuild/build.go | 27 +++- pkg/sql/colexec/hashbuild/build_test.go | 44 +++++- pkg/sql/colexec/hashbuild/errors.go | 3 +- pkg/sql/colexec/hashbuild/errors_test.go | 6 + pkg/sql/colexec/hashbuild/hashmap_test.go | 145 ++++++++++-------- pkg/sql/colexec/hashbuild/pressure.go | 34 ++-- pkg/sql/colexec/hashbuild/pressure_test.go | 21 ++- pkg/sql/colexec/hashbuild/types.go | 6 +- .../hashjoin/spill_integration_test.go | 83 ++++++++-- pkg/sql/colexec/indexbuild/build.go | 2 +- pkg/sql/colexec/runtimefilter/contract.go | 30 ++-- .../colexec/runtimefilter/contract_test.go | 46 ++++-- .../remote_allocation_statement_group.go | 4 +- .../remote_allocation_statement_group_test.go | 7 +- pkg/vm/message/message.go | 53 +++++-- pkg/vm/message/message_test.go | 28 ++++ pkg/vm/process/hashbuild_budget_test.go | 77 +++++++--- 30 files changed, 770 insertions(+), 526 deletions(-) create mode 100644 pkg/pb/pipeline/error_test.go diff --git a/pkg/common/moerr/error.go b/pkg/common/moerr/error.go index 729097fcbcaf0..8fe6c67f411e1 100644 --- a/pkg/common/moerr/error.go +++ b/pkg/common/moerr/error.go @@ -65,6 +65,7 @@ const ( ErrQueryInterrupted uint16 = 20104 ErrNotSupported uint16 = 20105 ErrRemoteDispatchNotRegistered uint16 = 20106 + ErrMPoolCapacity uint16 = 20107 // Group 2: numeric and functions ErrDivByZero uint16 = 20200 @@ -392,6 +393,7 @@ var errorMsgRefer = map[uint16]moErrorMsgItem{ ErrQueryInterrupted: {ER_QUERY_INTERRUPTED, []string{MySQLDefaultSqlState}, "query interrupted"}, ErrNotSupported: {ER_UNKNOWN_ERROR, []string{MySQLDefaultSqlState}, "not supported: %s"}, ErrRemoteDispatchNotRegistered: {ER_UNKNOWN_ERROR, []string{MySQLDefaultSqlState}, "remote dispatch receiver %s is not registered yet"}, + ErrMPoolCapacity: {ER_ENGINE_OUT_OF_MEMORY, []string{MySQLDefaultSqlState}, "mpool physical capacity exceeded: %s"}, // Group 2: numeric ErrDivByZero: {ER_DIVISION_BY_ZERO, []string{MySQLDefaultSqlState}, "division by zero"}, @@ -951,6 +953,13 @@ func NewOOM(ctx context.Context) *Error { return newError(ctx, ErrOOM) } +// NewMPoolCapacity reports a physical allocator or MPool capacity failure. +// Its dedicated wire code lets pressure recovery distinguish retryable +// physical capacity from unrelated OOMs without wrapping the MO error. +func NewMPoolCapacity(ctx context.Context, msg string) *Error { + return newError(ctx, ErrMPoolCapacity, msg) +} + // NewResourceExhaustedf preserves the existing resource-exhaustion wire code // while adding bounded, actionable context for guards that reject before the // allocator or operating system itself fails. The formatted message is diff --git a/pkg/common/moerr/error_no_ctx.go b/pkg/common/moerr/error_no_ctx.go index 750ab8b2e6d76..68b07dc7637b4 100644 --- a/pkg/common/moerr/error_no_ctx.go +++ b/pkg/common/moerr/error_no_ctx.go @@ -57,6 +57,10 @@ func NewOOMNoCtx() *Error { return newError(Context(), ErrOOM) } +func NewMPoolCapacityNoCtxf(format string, args ...any) *Error { + return NewMPoolCapacity(Context(), fmt.Sprintf(format, args...)) +} + func NewDivByZeroNoCtx() *Error { return newError(Context(), ErrDivByZero) } diff --git a/pkg/common/moerr/error_test.go b/pkg/common/moerr/error_test.go index 4514c4985a96d..3f96ab3e9d9d3 100644 --- a/pkg/common/moerr/error_test.go +++ b/pkg/common/moerr/error_test.go @@ -173,6 +173,19 @@ func TestResourceExhaustedWithDetailsEncoding(t *testing.T) { require.Equal(t, err, decoded) } +func TestMPoolCapacityEncoding(t *testing.T) { + err := NewMPoolCapacityNoCtxf("alloc %d bytes, cap %d", 8, 4) + require.Equal(t, ErrMPoolCapacity, err.ErrorCode()) + require.Equal(t, ER_ENGINE_OUT_OF_MEMORY, err.MySQLCode()) + require.Contains(t, err.Error(), "alloc 8 bytes, cap 4") + + data, marshalErr := err.MarshalBinary() + require.NoError(t, marshalErr) + decoded := new(Error) + require.NoError(t, decoded.UnmarshalBinary(data)) + require.Equal(t, err, decoded) +} + func TestErrSubqueryNo1RowContract(t *testing.T) { err := NewErrSubqueryNo1Row(context.Background()) require.Equal(t, ErrSubqueryNo1Row, err.ErrorCode()) diff --git a/pkg/common/mpool/allocation_account.go b/pkg/common/mpool/allocation_account.go index 15e31063f1918..b94231526ab44 100644 --- a/pkg/common/mpool/allocation_account.go +++ b/pkg/common/mpool/allocation_account.go @@ -21,6 +21,8 @@ import ( "runtime" "sync" "sync/atomic" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" ) // AllocationOwner and AllocationSite are bounded diagnostic dimensions. @@ -166,6 +168,11 @@ func AllocationFailureReasonOf(err error) AllocationFailureReason { switch { case errors.Is(err, ErrAllocationAccountInvariant): return AllocationFailureInvariant + case errors.Is(err, ErrAllocationAccountInvalid), + errors.Is(err, ErrAllocationAccountStale), + errors.Is(err, ErrAllocationAccountLive), + errors.Is(err, ErrAllocationGenerationSlots): + return AllocationFailureInvariant case errors.Is(err, ErrAllocationAccountMismatch): return AllocationFailureMismatch case errors.Is(err, ErrAllocationAccountSealed): @@ -175,13 +182,21 @@ func AllocationFailureReasonOf(err error) AllocationFailureReason { case errors.Is(err, ErrAllocationAdmissionSuspended): return AllocationFailureSuspended case errors.Is(err, ErrAllocationAccountCapacity), - errors.Is(err, ErrAllocationMetadataSlots): + errors.Is(err, ErrAllocationMetadataSlots), + IsMPoolCapacityFailure(err): return AllocationFailureCapacity default: return AllocationFailureNone } } +// IsMPoolCapacityFailure recognizes both a direct MO error and a contextual +// wrapper retained by an intermediate owner. +func IsMPoolCapacityFailure(err error) bool { + var moErr *moerr.Error + return errors.As(err, &moErr) && moErr.ErrorCode() == moerr.ErrMPoolCapacity +} + func IsRetryableAllocationCapacity(err error) bool { return AllocationFailureReasonOf(err) == AllocationFailureCapacity } @@ -197,15 +212,6 @@ type AllocationAccountTerminalSnapshot struct { LiveAllocations uint64 } -// AllocationAccountCheckpoint records the physical live-byte boundary before -// a retryable logical operation. The owner performs its own private-allocation -// rollback, then ValidateRollback proves that the same generation returned to -// this exact boundary before a retry can begin. -type AllocationAccountCheckpoint struct { - Handle AllocationAccountHandle - Used uint64 -} - // AllocationCapacityController lets an account share a higher-level aggregate // cap. The controller owns cap policy only; physical MPool metadata remains // the sole release owner. @@ -249,78 +255,6 @@ func (a *AllocationAccount) Snapshot() AllocationAccountSnapshot { } } -func (a *AllocationAccount) Checkpoint() (AllocationAccountCheckpoint, error) { - if a == nil || a.registry == nil || a.handle == 0 { - return AllocationAccountCheckpoint{}, ErrAllocationAccountInvalid - } - resolved, ok := a.registry.Resolve(a.handle) - if !ok || resolved != a { - return AllocationAccountCheckpoint{}, ErrAllocationAccountStale - } - snapshot := a.Snapshot() - if snapshot.Sealed { - return AllocationAccountCheckpoint{}, ErrAllocationAccountSealed - } - return AllocationAccountCheckpoint{ - Handle: snapshot.Handle, - Used: snapshot.Used, - }, nil -} - -// ValidateRollback proves that an owner restored its complete physical -// allocation boundary. It never mutates accounting: only physical Free owns a -// release, so a helper cannot hide a leaked allocation by decrementing usage. -func (a *AllocationAccount) ValidateRollback( - checkpoint AllocationAccountCheckpoint, -) error { - if a == nil || checkpoint.Handle == 0 { - return ErrAllocationAccountInvalid - } - if checkpoint.Handle != a.handle { - return wrapAllocationAccountError( - ErrAllocationAccountMismatch, - "checkpoint=%d account=%d", - checkpoint.Handle, - a.handle, - ) - } - snapshot := a.Snapshot() - if snapshot.Sealed { - return ErrAllocationAccountSealed - } - if snapshot.Used != checkpoint.Used { - return wrapAllocationAccountError( - ErrAllocationAccountInvariant, - "checkpoint-used=%d current-used=%d", - checkpoint.Used, - snapshot.Used, - ) - } - return nil -} - -// RollbackToCheckpoint runs the owner's physical cleanup and then proves the -// exact generation boundary. It deliberately does not own or synthesize any -// release: MPool allocation metadata remains the sole release authority. -func (a *AllocationAccount) RollbackToCheckpoint( - checkpoint AllocationAccountCheckpoint, - rollback func() error, -) error { - if rollback == nil { - return ErrAllocationAccountInvalid - } - if a == nil || checkpoint.Handle != a.handle { - return ErrAllocationAccountMismatch - } - if a.Snapshot().Sealed { - return ErrAllocationAccountSealed - } - if err := rollback(); err != nil { - return err - } - return a.ValidateRollback(checkpoint) -} - func (a *AllocationAccount) acquire(capacity uint64) error { if a == nil || a.registry == nil || a.handle == 0 { return ErrAllocationAccountInvalid @@ -859,22 +793,8 @@ type allocationAccountRequest struct { account *AllocationAccount owner AllocationOwner site AllocationSite - // checkpoint is nil for every public caller. Same-package fault tests use - // it to prove rollback at each unpublished transaction boundary. - checkpoint func(allocationCheckpoint) error } -type allocationCheckpoint uint8 - -const ( - allocationAfterAccount allocationCheckpoint = iota + 1 - allocationAfterMetadata - allocationAfterGlobalStats - allocationAfterPoolStats - allocationAfterPhysical - allocationAfterHeader -) - func (r allocationAccountRequest) validate() error { if r.account == nil || r.account.registry == nil || r.owner < AllocationOwnerMin || r.owner > AllocationOwnerMax || @@ -888,15 +808,6 @@ func (r allocationAccountRequest) validate() error { return nil } -func (r allocationAccountRequest) reach( - checkpoint allocationCheckpoint, -) error { - if r.checkpoint == nil { - return nil - } - return r.checkpoint(checkpoint) -} - type allocationLease struct { account *AllocationAccount owner AllocationOwner diff --git a/pkg/common/mpool/allocation_account_mpool_test.go b/pkg/common/mpool/allocation_account_mpool_test.go index df20d08fd8d0f..be552e8e215b6 100644 --- a/pkg/common/mpool/allocation_account_mpool_test.go +++ b/pkg/common/mpool/allocation_account_mpool_test.go @@ -16,7 +16,6 @@ package mpool import ( "bytes" - "errors" "fmt" "sync" "testing" @@ -273,6 +272,9 @@ func TestMPoolAccountedRollback(t *testing.T) { testAllocationSite, ) require.Error(t, err) + require.True(t, IsMPoolCapacityFailure(err)) + require.Equal(t, AllocationFailureCapacity, AllocationFailureReasonOf(err)) + require.True(t, IsRetryableAllocationCapacity(err)) require.Equal(t, uint64(allocationSize), account.Snapshot().Used) require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) require.Equal(t, uint64(2), registry.PeakAllocationMetadata()) @@ -302,6 +304,9 @@ func TestMPoolAccountedRollback(t *testing.T) { testAllocationSite, ) require.Error(t, err) + require.True(t, IsMPoolCapacityFailure(err)) + require.Equal(t, AllocationFailureCapacity, AllocationFailureReasonOf(err)) + require.True(t, IsRetryableAllocationCapacity(err)) require.Equal(t, globalBefore, GlobalStats().NumCurrBytes.Load()) require.Zero(t, mp.CurrNB()) require.Zero(t, account.Snapshot().Used) @@ -604,112 +609,6 @@ func TestMPoolAccountedConcurrentAllocFree(t *testing.T) { finalizeTestAllocationAccount(t, registry, account) } -func TestMPoolAccountedTransactionRollback(t *testing.T) { - countGlobalMetadata := func() (headers int, leases int) { - for i := range globalPtrShards { - shard := &globalPtrShards[i] - shard.mu.Lock() - headers += len(shard.m) - leases += len(shard.leases) - shard.mu.Unlock() - } - return headers, leases - } - - injectedError := errors.New("injected allocation error") - stages := []struct { - name string - checkpoint allocationCheckpoint - }{ - {name: "account", checkpoint: allocationAfterAccount}, - {name: "metadata", checkpoint: allocationAfterMetadata}, - {name: "global-stats", checkpoint: allocationAfterGlobalStats}, - {name: "pool-stats", checkpoint: allocationAfterPoolStats}, - {name: "physical", checkpoint: allocationAfterPhysical}, - {name: "header-publication", checkpoint: allocationAfterHeader}, - } - for _, fault := range []struct { - name string - trigger func() error - wantPanic bool - }{ - { - name: "error", - trigger: func() error { - return injectedError - }, - }, - { - name: "panic", - trigger: func() error { - panic("injected allocation panic") - }, - wantPanic: true, - }, - } { - for _, stage := range stages { - for _, noLock := range []bool{false, true} { - poolKind := "sharded" - if noLock { - poolKind = "no-lock" - } - t.Run(fault.name+"/"+stage.name+"/"+poolKind, func(t *testing.T) { - registry, account := newTestAllocationAccount(t, 64, 1) - var mp *MPool - if noLock { - mp = MustNewNoLock("accounted-rollback-no-lock") - } else { - mp = MustNew("accounted-rollback") - } - defer DeleteMPool(mp) - - globalBytesBefore := GlobalStats().NumCurrBytes.Load() - headersBefore, leasesBefore := countGlobalMetadata() - request := allocationAccountRequest{ - account: account, - owner: testAllocationOwner, - site: testAllocationSite, - checkpoint: func( - reached allocationCheckpoint, - ) error { - if reached != stage.checkpoint { - return nil - } - return fault.trigger() - }, - } - if fault.wantPanic { - require.PanicsWithValue(t, "injected allocation panic", func() { - _, _ = mp.allocAccountedWithDetailK("", 64, request) - }) - } else { - _, err := mp.allocAccountedWithDetailK("", 64, request) - require.ErrorIs(t, err, injectedError) - } - - require.Zero(t, mp.CurrNB()) - require.Equal( - t, - globalBytesBefore, - GlobalStats().NumCurrBytes.Load(), - ) - require.Zero(t, account.Snapshot().Used) - require.Zero(t, registry.LiveAllocationMetadata()) - if noLock { - require.Empty(t, mp.ptrs) - require.Empty(t, mp.leases) - } else { - headersAfter, leasesAfter := countGlobalMetadata() - require.Equal(t, headersBefore, headersAfter) - require.Equal(t, leasesBefore, leasesAfter) - } - finalizeTestAllocationAccount(t, registry, account) - }) - } - } - } -} - func BenchmarkMPoolAccountedAllocation(b *testing.B) { for _, accounted := range []bool{false, true} { mode := "unaccounted" diff --git a/pkg/common/mpool/allocation_account_test.go b/pkg/common/mpool/allocation_account_test.go index c0d07192538d6..d4c44be0f7557 100644 --- a/pkg/common/mpool/allocation_account_test.go +++ b/pkg/common/mpool/allocation_account_test.go @@ -22,6 +22,7 @@ import ( "sync/atomic" "testing" + "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/stretchr/testify/require" ) @@ -434,47 +435,6 @@ func TestAllocationAccountOpenSuspendLinearization(t *testing.T) { require.False(t, registry.AdmissionSuspended()) } -func TestAllocationAccountCheckpointValidation(t *testing.T) { - registry, err := NewAllocationAccountRegistry(2, 2) - require.NoError(t, err) - account, err := registry.Open(8) - require.NoError(t, err) - other, err := registry.Open(8) - require.NoError(t, err) - - checkpoint, err := account.Checkpoint() - require.NoError(t, err) - require.NoError(t, account.acquire(1)) - require.ErrorIs(t, account.ValidateRollback(checkpoint), ErrAllocationAccountInvariant) - require.NoError(t, account.RollbackToCheckpoint(checkpoint, func() error { - account.release(1) - return nil - })) - require.NoError(t, account.ValidateRollback(checkpoint)) - require.ErrorIs( - t, - account.RollbackToCheckpoint(checkpoint, nil), - ErrAllocationAccountInvalid, - ) - - otherCheckpoint, err := other.Checkpoint() - require.NoError(t, err) - require.ErrorIs(t, account.ValidateRollback(otherCheckpoint), ErrAllocationAccountMismatch) - called := false - require.ErrorIs(t, account.RollbackToCheckpoint(otherCheckpoint, func() error { - called = true - return nil - }), ErrAllocationAccountMismatch) - require.False(t, called) - account.Seal() - require.ErrorIs(t, account.ValidateRollback(checkpoint), ErrAllocationAccountSealed) - _, err = registry.Finalize(account) - require.NoError(t, err) - other.Seal() - _, err = registry.Finalize(other) - require.NoError(t, err) -} - func TestAllocationFailureReasonsAreNonOverlapping(t *testing.T) { testCases := []struct { err error @@ -483,10 +443,18 @@ func TestAllocationFailureReasonsAreNonOverlapping(t *testing.T) { }{ {ErrAllocationAccountCapacity, AllocationFailureCapacity, true}, {ErrAllocationMetadataSlots, AllocationFailureCapacity, true}, + {moerr.NewMPoolCapacityNoCtxf("test"), AllocationFailureCapacity, true}, {ErrAllocationAccountSealed, AllocationFailureSealed, false}, {ErrAllocationAccountMismatch, AllocationFailureMismatch, false}, {ErrAllocationAllocatorLimit, AllocationFailureAllocatorLimit, false}, {ErrAllocationAccountInvariant, AllocationFailureInvariant, false}, + {ErrAllocationAccountInvalid, AllocationFailureInvariant, false}, + {ErrAllocationAccountStale, AllocationFailureInvariant, false}, + {ErrAllocationAccountLive, AllocationFailureInvariant, false}, + {ErrAllocationGenerationSlots, AllocationFailureInvariant, false}, + {errors.Join(moerr.NewMPoolCapacityNoCtxf("test"), ErrAllocationAccountInvariant), AllocationFailureInvariant, false}, + {errors.Join(moerr.NewMPoolCapacityNoCtxf("test"), ErrAllocationAccountInvalid), AllocationFailureInvariant, false}, + {errors.Join(moerr.NewMPoolCapacityNoCtxf("test"), ErrAllocationAccountSealed), AllocationFailureSealed, false}, {ErrAllocationAdmissionSuspended, AllocationFailureSuspended, false}, {errors.New("unrelated"), AllocationFailureNone, false}, } diff --git a/pkg/common/mpool/mpool.go b/pkg/common/mpool/mpool.go index ecc96bc4a0abc..14aaf3896ef3d 100644 --- a/pkg/common/mpool/mpool.go +++ b/pkg/common/mpool/mpool.go @@ -361,14 +361,12 @@ func (mp *MPool) recordAccountedPtrMetadata( ptr unsafe.Pointer, pHdr memHdr, lease allocationLease, - request allocationAccountRequest, ) error { if !mp.noLock { return gRecordAccountedPtrMetadata( ptr, pHdr, lease, - request, ) } if _, ok := mp.ptrs[ptr]; ok { @@ -377,22 +375,11 @@ func (mp *MPool) recordAccountedPtrMetadata( if _, ok := mp.leases[ptr]; ok { return moerr.NewInternalErrorNoCtx("account lease already recorded") } - committed := false - defer func() { - if !committed { - delete(mp.ptrs, ptr) - delete(mp.leases, ptr) - } - }() - mp.ptrs[ptr] = pHdr - if err := request.reach(allocationAfterHeader); err != nil { - return err - } if mp.leases == nil { mp.leases = make(map[unsafe.Pointer]allocationLease) } + mp.ptrs[ptr] = pHdr mp.leases[ptr] = lease - committed = true return nil } @@ -837,6 +824,17 @@ func allocationAccountSiteError( request allocationAccountRequest, err error, ) error { + if IsMPoolCapacityFailure(err) { + // Keep a direct MO error at the operator boundary: remote pipeline + // encoding preserves only direct *moerr.Error values. The ownership + // dimensions are folded into the serialized message instead. + return moerr.NewMPoolCapacityNoCtxf( + "allocation owner=%d site=%d: %s", + request.owner, + request.site, + err.Error(), + ) + } return prefixAllocationAccountError( err, "allocation owner=%d site=%d", @@ -866,19 +864,22 @@ func (mp *MPool) alloc( gcurr := globalStats.RecordAlloc("global", sz) if gcurr > GlobalCap() { globalStats.RecordFree("global", sz) - return nil, moerr.NewOOMNoCtx() + return nil, moerr.NewMPoolCapacityNoCtxf( + "global cap exceeded while allocating %d bytes", sz) } mycurr := mp.stats.RecordAlloc(mp.tag, sz) if mycurr > mp.Cap() { mp.stats.RecordFree(mp.tag, sz) globalStats.RecordFree("global", sz) - return nil, moerr.NewInternalErrorNoCtxf("mpool out of space, alloc %d bytes, cap %d", sz, mp.cap) + return nil, moerr.NewMPoolCapacityNoCtxf( + "mpool out of space, alloc %d bytes, cap %d", sz, mp.cap) } bs, err = simpleCAllocator().Allocate(uint64(sz)) if err != nil { mp.stats.RecordFree(mp.tag, sz) globalStats.RecordFree("global", sz) - return nil, err + return nil, moerr.NewMPoolCapacityNoCtxf( + "physical allocator rejected %d bytes: %v", sz, err) } } else { bs = make([]byte, sz) @@ -941,16 +942,10 @@ func (mp *MPool) allocAccounted( return nil, err } accountHeld = true - if err = request.reach(allocationAfterAccount); err != nil { - return nil, err - } if err = request.account.registry.reserveMetadata(); err != nil { return nil, err } metadataHeld = true - if err = request.reach(allocationAfterMetadata); err != nil { - return nil, err - } hdr := memHdr{ poolId: mp.id, @@ -966,39 +961,28 @@ func (mp *MPool) allocAccounted( gcurr := globalStats.RecordAlloc("global", sz) globalHeld = true - if err = request.reach(allocationAfterGlobalStats); err != nil { - return nil, err - } if gcurr > GlobalCap() { - return nil, moerr.NewOOMNoCtx() + return nil, moerr.NewMPoolCapacityNoCtxf( + "global cap exceeded while allocating %d bytes", sz) } mycurr := mp.stats.RecordAlloc(mp.tag, sz) poolHeld = true - if err = request.reach(allocationAfterPoolStats); err != nil { - return nil, err - } if mycurr > mp.Cap() { - return nil, moerr.NewInternalErrorNoCtxf( - "mpool out of space, alloc %d bytes, cap %d", - sz, - mp.cap, - ) + return nil, moerr.NewMPoolCapacityNoCtxf( + "mpool out of space, alloc %d bytes, cap %d", sz, mp.cap) } bs, err = simpleCAllocator().Allocate(uint64(sz)) if err != nil { - return nil, err + return nil, moerr.NewMPoolCapacityNoCtxf( + "physical allocator rejected %d bytes: %v", sz, err) } physicalHeld = true - if err = request.reach(allocationAfterPhysical); err != nil { - return nil, err - } ptr := unsafe.Pointer(&bs[0]) if err = mp.recordAccountedPtrMetadata( ptr, hdr, lease, - request, ); err != nil { return nil, err } @@ -1282,12 +1266,14 @@ func (mp *MPool) ReallocZero(old []byte, sz int, offHeap bool) ([]byte, error) { // retain the new-size charge and release only the old-size charge on success. if globalStats.RecordAlloc("global", int64(sz)) > GlobalCap() { globalStats.RecordFree("global", int64(sz)) - return nil, moerr.NewOOMNoCtx() + return nil, moerr.NewMPoolCapacityNoCtxf( + "global cap exceeded while reallocating %d bytes", sz) } if mp.stats.RecordAlloc(mp.tag, int64(sz)) > mp.Cap() { mp.stats.RecordFree(mp.tag, int64(sz)) globalStats.RecordFree("global", int64(sz)) - return nil, moerr.NewInternalErrorNoCtxf("mpool out of space, realloc %d bytes, cap %d", sz, mp.cap) + return nil, moerr.NewMPoolCapacityNoCtxf( + "mpool out of space, realloc %d bytes, cap %d", sz, mp.cap) } newbs, err := simpleCAllocator().ReallocZero( @@ -1298,7 +1284,8 @@ func (mp *MPool) ReallocZero(old []byte, sz int, offHeap bool) ([]byte, error) { if err != nil { mp.stats.RecordFree(mp.tag, int64(sz)) globalStats.RecordFree("global", int64(sz)) - return nil, err + return nil, moerr.NewMPoolCapacityNoCtxf( + "physical allocator rejected realloc to %d bytes: %v", sz, err) } newptr := unsafe.Pointer(&newbs[0]) var removedLease allocationLease @@ -1497,7 +1484,6 @@ func gRecordAccountedPtrMetadata( ptr unsafe.Pointer, hdr memHdr, lease allocationLease, - request allocationAccountRequest, ) error { shard := getPtrShard(ptr) shard.mu.Lock() @@ -1508,22 +1494,11 @@ func gRecordAccountedPtrMetadata( if _, ok := shard.leases[ptr]; ok { return moerr.NewInternalErrorNoCtx("account lease already recorded") } - committed := false - defer func() { - if !committed { - delete(shard.m, ptr) - delete(shard.leases, ptr) - } - }() - shard.m[ptr] = hdr - if err := request.reach(allocationAfterHeader); err != nil { - return err - } if shard.leases == nil { shard.leases = make(map[unsafe.Pointer]allocationLease) } + shard.m[ptr] = hdr shard.leases[ptr] = lease - committed = true return nil } diff --git a/pkg/container/types/packer.go b/pkg/container/types/packer.go index e6a96e6c7ca0d..3322e2d6c8926 100644 --- a/pkg/container/types/packer.go +++ b/pkg/container/types/packer.go @@ -23,13 +23,20 @@ import ( "unsafe" "github.com/matrixorigin/matrixone/pkg/common/malloc" + "github.com/matrixorigin/matrixone/pkg/common/moerr" ) type Packer struct { buffer []byte bufferDeallocator malloc.Deallocator + fixed bool + overflow bool } +var ErrPackerCapacity = moerr.NewInternalErrorNoCtx( + "packer fixed buffer capacity exceeded", +) + var packerAllocator = malloc.NewShardedAllocator( runtime.GOMAXPROCS(0), func() malloc.Allocator { @@ -43,13 +50,6 @@ func NewPacker() *Packer { return NewPackerWithSize(4096) } -// PackerAllocationSize returns the backing size-class allocation made by -// NewPackerWithSize. It lets memory-governed callers reserve the actual -// allocation, including allocator rounding, before constructing a packer. -func PackerAllocationSize(size uint64) (uint64, bool) { - return malloc.ClassAllocationSize(size) -} - func NewPackerWithSize(size uint64) *Packer { bs, dec, err := packerAllocator.Allocate(size, malloc.NoClear) if err != nil { @@ -61,6 +61,15 @@ func NewPackerWithSize(size uint64) *Packer { } } +// NewPackerWithFixedBuffer uses caller-owned storage and never allocates. +// Err reports an encoding that exceeded the supplied physical capacity. +func NewPackerWithFixedBuffer(buffer []byte) *Packer { + return &Packer{ + buffer: buffer[:0:cap(buffer)], + fixed: true, + } +} + func NewPackerArray(length int) []*Packer { return NewPackerArrayWithSize(length, 4096) } @@ -82,16 +91,14 @@ func (p *Packer) Close() { func (p *Packer) Reset() { p.buffer = p.buffer[:0] + p.overflow = false } -func (p *Packer) ensureSize(n int) { - if len(p.buffer)+n <= cap(p.buffer) { +func (p *Packer) ensureSizeSlow(n int) { + if p.fixed { + p.overflow = true return } - p.ensureSizeSlow(n) -} - -func (p *Packer) ensureSizeSlow(n int) { newBuffer, newDec, err := packerAllocator.Allocate(uint64(cap(p.buffer)+n), malloc.NoClear) if err != nil { panic(err) @@ -106,12 +113,26 @@ func (p *Packer) ensureSizeSlow(n int) { } func (p *Packer) putByte(b byte) { - p.ensureSize(1) + if len(p.buffer) < cap(p.buffer) { + p.buffer = append(p.buffer, b) + return + } + p.ensureSizeSlow(1) + if p.overflow { + return + } p.buffer = append(p.buffer, b) } func (p *Packer) putBytes(bs []byte) { - p.ensureSize(len(bs)) + if len(bs) <= cap(p.buffer)-len(p.buffer) { + p.buffer = append(p.buffer, bs...) + return + } + p.ensureSizeSlow(len(bs)) + if p.overflow { + return + } p.buffer = append(p.buffer, bs...) } @@ -331,6 +352,16 @@ func (p *Packer) GetBuf() []byte { return p.buffer } +func (p *Packer) Err() error { + if p == nil { + return ErrPackerCapacity + } + if p.overflow { + return ErrPackerCapacity + } + return nil +} + func (p *Packer) Bytes() []byte { return slices.Clone(p.buffer) } diff --git a/pkg/container/types/packer_test.go b/pkg/container/types/packer_test.go index 6792c9df48200..ec18ba8120297 100644 --- a/pkg/container/types/packer_test.go +++ b/pkg/container/types/packer_test.go @@ -14,7 +14,10 @@ package types -import "testing" +import ( + "errors" + "testing" +) func TestPacker(t *testing.T) { packer := NewPacker() @@ -41,6 +44,31 @@ func TestClosedPackerIsOK(t *testing.T) { packer.Close() } +func TestFixedBufferPackerNeverAllocatesPastCapacity(t *testing.T) { + storage := make([]byte, 0, 3) + packer := NewPackerWithFixedBuffer(storage) + packer.EncodeBool(true) + packer.EncodeNull() + if err := packer.Err(); err != nil { + t.Fatal(err) + } + if len(packer.GetBuf()) != 2 { + t.Fatalf("encoded length = %d", len(packer.GetBuf())) + } + packer.EncodeInt64(42) + if !errors.Is(packer.Err(), ErrPackerCapacity) { + t.Fatalf("overflow error = %v", packer.Err()) + } + if len(packer.GetBuf()) > cap(storage) { + t.Fatal("fixed packer exceeded caller-owned storage") + } + packer.Reset() + packer.EncodeBool(true) + if err := packer.Err(); err != nil { + t.Fatalf("reset fixed packer error = %v", err) + } +} + func BenchmarkPacker(b *testing.B) { for i := 0; i < b.N; i++ { packer := NewPacker() @@ -50,11 +78,29 @@ func BenchmarkPacker(b *testing.B) { } func BenchmarkPackerEncode(b *testing.B) { - packer := NewPacker() - defer packer.Close() - b.ResetTimer() - for i := 0; i < b.N; i++ { - packer.EncodeInt64(42) - packer.Reset() + for _, fixed := range []bool{false, true} { + mode := "allocator-backed" + if fixed { + mode = "fixed-buffer" + } + b.Run(mode, func(b *testing.B) { + var packer *Packer + if fixed { + packer = NewPackerWithFixedBuffer(make([]byte, 16)) + } else { + packer = NewPacker() + defer packer.Close() + } + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + packer.EncodeInt64(42) + packer.Reset() + } + b.StopTimer() + if err := packer.Err(); err != nil { + b.Fatal(err) + } + }) } } diff --git a/pkg/container/vector/allocation_account_test.go b/pkg/container/vector/allocation_account_test.go index 2ae706b9e241f..462a78e516934 100644 --- a/pkg/container/vector/allocation_account_test.go +++ b/pkg/container/vector/allocation_account_test.go @@ -992,6 +992,72 @@ func BenchmarkVectorAllocationAccount(b *testing.B) { finalizeTestVectorAllocationAccount(b, state) } +func BenchmarkVectorElementAccounting(b *testing.B) { + const rows = 8192 + mp := mpool.MustNewZero() + state := newTestVectorAllocationAccount(b, 1<<40, 64) + source := NewOffHeapVecWithType(types.T_int64.ToType()) + for i := range rows { + if err := AppendFixed(source, int64(i), false, mp); err != nil { + b.Fatal(err) + } + } + b.Cleanup(func() { + source.Free(mp) + finalizeTestVectorAllocationAccount(b, state) + }) + + for _, accounted := range []bool{false, true} { + mode := "unaccounted" + if accounted { + mode = "accounted" + } + b.Run("union-one/"+mode, func(b *testing.B) { + destination := NewOffHeapVecWithType(types.T_int64.ToType()) + if accounted { + if err := destination.SetAllocationAccount(state.selection); err != nil { + b.Fatal(err) + } + } + if err := destination.PreExtend(1, mp); err != nil { + b.Fatal(err) + } + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + destination.ResetWithSameType() + if err := destination.UnionOne(source, int64(i%rows), mp); err != nil { + b.Fatal(err) + } + } + b.StopTimer() + destination.Free(mp) + }) + + b.Run("copy/"+mode, func(b *testing.B) { + destination := NewOffHeapVecWithType(types.T_int64.ToType()) + if accounted { + if err := destination.SetAllocationAccount(state.selection); err != nil { + b.Fatal(err) + } + } + if err := destination.PreExtend(1, mp); err != nil { + b.Fatal(err) + } + destination.SetLength(1) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if err := destination.Copy(source, 0, int64(i%rows), mp); err != nil { + b.Fatal(err) + } + } + b.StopTimer() + destination.Free(mp) + }) + } +} + func TestVectorAllocationAccountErrorsAreTyped(t *testing.T) { state := newTestVectorAllocationAccount(t, 1, 1) mp := mpool.MustNewZero() diff --git a/pkg/container/vector/vector.go b/pkg/container/vector/vector.go index 33bd3b913a8c3..b746b24b4a4c2 100644 --- a/pkg/container/vector/vector.go +++ b/pkg/container/vector/vector.go @@ -2092,17 +2092,19 @@ func (v *Vector) ShuffleWithBuf(sels []int64, mp *mpool.MPool, buf *[]byte) (err // Copy simply does v[vi] = w[wi] func (v *Vector) Copy(w *Vector, vi, wi int64, mp *mpool.MPool) error { sourceGrouping := w.GetGrouping().Contains(uint64(wi)) - if sourceGrouping { + if sourceGrouping && v.allocationAccount != nil { if err := v.ensureGroupingCapacity(int(vi)+1, mp); err != nil { return err } + } + if sourceGrouping { v.GetGrouping().Set(uint64(vi)) } else { v.GetGrouping().Unset(uint64(vi)) } sourceNull := w.IsConstNull() || (!w.IsConst() && w.GetNulls().Contains(uint64(wi))) - if sourceNull { + if sourceNull && v.allocationAccount != nil { if err := v.ensureNullCapacity(int(vi)+1, mp); err != nil { return err } @@ -2121,7 +2123,7 @@ func (v *Vector) Copy(w *Vector, vi, wi int64, mp *mpool.MPool) error { // Non-null constant vectors still share the regular null/data path below. wi = 0 } - if w.GetNulls().Contains(uint64(wi)) { + if sourceNull { if !v.typ.IsFixedLen() { vva := MustFixedColNoTypeCheck[types.Varlena](v) vva[vi] = types.Varlena{} @@ -3366,25 +3368,31 @@ func (v *Vector) UnionNull(mp *mpool.MPool) error { // It is simply append. the purpose of retention is ease of use func (v *Vector) UnionOne(w *Vector, sel int64, mp *mpool.MPool) error { - needGrouping := nulls.Contains(&w.gsp, uint64(sel)) - needNulls := w.IsConstNull() || + sourceGrouping := nulls.Contains(&w.gsp, uint64(sel)) + sourceNull := w.IsConstNull() || (!w.IsConst() && nulls.Contains(&w.nsp, uint64(sel))) - if err := extendWithBitmaps(v, 1, mp, needNulls, needGrouping); err != nil { + if err := extendWithBitmaps( + v, + 1, + mp, + sourceNull && v.allocationAccount != nil, + sourceGrouping && v.allocationAccount != nil, + ); err != nil { return err } oldLen := v.length v.length++ - if nulls.Contains(&w.gsp, uint64(sel)) { + if sourceGrouping { nulls.Add(&v.gsp, uint64(oldLen)) } if w.IsConst() { - if w.IsConstNull() { + if sourceNull { nulls.Add(&v.nsp, uint64(oldLen)) return nil } sel = 0 - } else if nulls.Contains(&w.nsp, uint64(sel)) { + } else if sourceNull { nulls.Add(&v.nsp, uint64(oldLen)) return nil } @@ -3443,25 +3451,31 @@ func (v *Vector) UnionMulti(w *Vector, sel int64, cnt int, mp *mpool.MPool) erro return nil } - needGrouping := nulls.Contains(&w.gsp, uint64(sel)) - needNulls := w.IsConstNull() || + sourceGrouping := nulls.Contains(&w.gsp, uint64(sel)) + sourceNull := w.IsConstNull() || (!w.IsConst() && nulls.Contains(&w.nsp, uint64(sel))) - if err := extendWithBitmaps(v, cnt, mp, needNulls, needGrouping); err != nil { + if err := extendWithBitmaps( + v, + cnt, + mp, + sourceNull && v.allocationAccount != nil, + sourceGrouping && v.allocationAccount != nil, + ); err != nil { return err } oldLen := v.length v.length += cnt - if nulls.Contains(&w.gsp, uint64(sel)) { + if sourceGrouping { nulls.AddRange(&v.gsp, uint64(oldLen), uint64(oldLen+cnt)) } if w.IsConst() { - if w.IsConstNull() { + if sourceNull { nulls.AddRange(&v.nsp, uint64(oldLen), uint64(oldLen+cnt)) return nil } sel = 0 - } else if nulls.Contains(&w.nsp, uint64(sel)) { + } else if sourceNull { nulls.AddRange(&v.nsp, uint64(oldLen), uint64(oldLen+cnt)) return nil } diff --git a/pkg/pb/pipeline/error_test.go b/pkg/pb/pipeline/error_test.go new file mode 100644 index 0000000000000..279d80ea7278a --- /dev/null +++ b/pkg/pb/pipeline/error_test.go @@ -0,0 +1,50 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pipeline + +import ( + "context" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/stretchr/testify/require" +) + +func TestMPoolCapacityErrorPreservesWireIdentity(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 4) + require.NoError(t, err) + account, err := registry.Open(2 << 20) + require.NoError(t, err) + pool, err := mpool.NewMPool("pipeline-mpool-capacity", 1<<20, mpool.NoFixed) + require.NoError(t, err) + defer mpool.DeleteMPool(pool) + + first, err := pool.AllocAccounted(768<<10, account, 1, 1) + require.NoError(t, err) + defer pool.Free(first) + _, capacityErr := pool.AllocAccounted(768<<10, account, 1, 1) + require.Error(t, capacityErr) + require.IsType(t, new(moerr.Error), capacityErr) + require.True(t, mpool.IsMPoolCapacityFailure(capacityErr)) + + message := new(Message) + message.SetMoError(context.Background(), capacityErr) + + wireErr, ok := message.TryToGetMoErr() + require.True(t, ok) + require.True(t, moerr.IsMoErrCode(wireErr, moerr.ErrMPoolCapacity)) + require.Contains(t, wireErr.Error(), "allocation owner=1 site=1") +} diff --git a/pkg/sql/colexec/fuzzyfilter/filter.go b/pkg/sql/colexec/fuzzyfilter/filter.go index aab0610860356..1f6ff802d1252 100644 --- a/pkg/sql/colexec/fuzzyfilter/filter.go +++ b/pkg/sql/colexec/fuzzyfilter/filter.go @@ -402,7 +402,7 @@ func (fuzzyFilter *FuzzyFilter) handleRuntimeFilter(proc *process.Process) error } if encoding == keycodec.ExactRuntimeFilterFloatZeroClosed { if err := runtimefilter.CloseFloatSignedZero( - ctr.pass2RuntimeFilter, proc.Mp(), nil); err != nil { + ctr.pass2RuntimeFilter, proc.Mp()); err != nil { if fuzzyFilter.fallbackRuntimeFilter(proc, err) { return nil } diff --git a/pkg/sql/colexec/hashbuild/build.go b/pkg/sql/colexec/hashbuild/build.go index e36a01df1e4cf..abffc5bad7f6b 100644 --- a/pkg/sql/colexec/hashbuild/build.go +++ b/pkg/sql/colexec/hashbuild/build.go @@ -858,7 +858,6 @@ func (hashBuild *HashBuild) handleRuntimeFilter( if err := runtimefilter.CloseFloatSignedZero( keyVec, proc.Mp(), - nil, ); err != nil { if hashBuild.fallbackOptionalRuntimeFilter(err, &runtimeFilter, spec, proc) { return nil @@ -1015,6 +1014,7 @@ func (hashBuild *HashBuild) materializeSerializedRuntimeFilter( payloadType, ok := planExprType( runtimefilter.BuildKeyExpr(spec)) if !ok || areaBound > uint64(math.MaxInt) || + maxRowBound > uint64(math.MaxInt) || hashBuild.ctr.hashmapBuilder.uniqueKeyAllocation == nil { return nil, nil, 0, false, nil } @@ -1037,8 +1037,23 @@ func (hashBuild *HashBuild) materializeSerializedRuntimeFilter( if packerSize == 0 { packerSize = 1 } - packer := types.NewPackerWithSize(packerSize) - defer packer.Close() + scratch, err := mpool.NewAccountedBuffer( + proc.Mp(), + hashBuild.ctr.hashmapBuilder.mapAllocationAccount, + HashBuildAllocationOwner, + HashBuildAllocationSiteRuntimeFilterScratch, + ) + if err != nil { + return nil, nil, 0, false, err + } + defer scratch.Free() + if err = scratch.Resize(int(packerSize)); err != nil { + if mpool.IsRetryableAllocationCapacity(err) { + err = runtimefilter.MarkOptionalAllocationError(err) + } + return nil, nil, 0, false, err + } + packer := types.NewPackerWithFixedBuffer(scratch.Bytes()) for row := 0; row < rowCount; row++ { if row&8191 == 0 { @@ -1060,6 +1075,12 @@ func (hashBuild *HashBuild) materializeSerializedRuntimeFilter( } encoders[i](component, row, packer) } + if err = packer.Err(); err != nil { + return nil, nil, 0, false, errors.Join( + mpool.ErrAllocationAccountInvariant, + err, + ) + } if rowIsNull { // serial is NULL if any component is NULL. NULL build keys never // match SQL equality, so omit them rather than turning a reset null diff --git a/pkg/sql/colexec/hashbuild/build_test.go b/pkg/sql/colexec/hashbuild/build_test.go index 664e3b808ac2c..bd012abb35732 100644 --- a/pkg/sql/colexec/hashbuild/build_test.go +++ b/pkg/sql/colexec/hashbuild/build_test.go @@ -763,9 +763,6 @@ func TestHashBuildRuntimeFilterFallbackStatsTriggerDiagnostics(t *testing.T) { "HashBuildRuntimeFilterCollectionFallbacks", "HashBuildRuntimeFilterBudgetFallbacks", "HashBuildRuntimeFilterAllocationFallbacks", - "HashBuildSpillScratchReserveRejects", - "HashBuildSpillScratchGrowRejects", - "HashBuildSpillScratchGrowCount", } { t.Run(stat, func(t *testing.T) { require.True(t, hasHashBuildDiagnosticStats( @@ -1808,6 +1805,47 @@ func TestSerializedRuntimeFilterUsesTightBudgetAndProducesIn(t *testing.T) { require.Zero(t, tc.proc.Mp().CurrNB()) } +func TestSerializedRuntimeFilterScratchUsesPhysicalAccount(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + componentType := types.T_varchar.ToType() + spec := makeSerializedRuntimeFilterSpec( + t, proc, 107, 2, []types.Type{componentType}, false) + arg := &HashBuild{RuntimeFilterSpec: spec} + budget := process.MustNewHashBuildBudget(450<<10, 450<<10) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + installTestHashBuildBudget(t, arg, generation) + + value := strings.Repeat("x", 300<<10) + arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{ + testutil.MakeVarcharVector([]string{value}, nil, proc.Mp()), + } + arg.ctr.hashmapBuilder.InputBatchRowCount = 1 + + data, release, rows, usable, err := + arg.materializeSerializedRuntimeFilter( + proc, + spec, + []types.Type{componentType}, + 1, + ) + require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) + require.Equal(t, runtimefilter.OptionalFallbackBudgetAdmission, + runtimefilter.ClassifyOptionalFallback(err)) + require.Nil(t, data) + require.Nil(t, release) + require.Zero(t, rows) + require.False(t, usable) + require.Equal(t, uint64(1), generation.RejectCount()) + require.Zero(t, generation.Used()) + require.NotZero(t, generation.Peak()) + + arg.ctr.hashmapBuilder.Free(proc) + generation.Close() + proc.Free() + require.Zero(t, proc.Mp().CurrNB()) +} + func TestSerializedRuntimeFilterBoundsObserveCancellation(t *testing.T) { tc := newTestCase(t, nil, nil, nil) vec := testutil.MakeInt32Vector([]int32{1}, nil, tc.proc.Mp()) diff --git a/pkg/sql/colexec/hashbuild/errors.go b/pkg/sql/colexec/hashbuild/errors.go index 0b9af9faf0c5f..36957706ebc3b 100644 --- a/pkg/sql/colexec/hashbuild/errors.go +++ b/pkg/sql/colexec/hashbuild/errors.go @@ -43,7 +43,8 @@ func TerminalBudgetError(ctx context.Context, err error) error { if !errors.As(err, &budgetErr) || budgetErr.Kind != process.HashBuildBudgetErrorAdmission { switch { case mpool.AllocationFailureReasonOf(err) == - mpool.AllocationFailureCapacity: + mpool.AllocationFailureCapacity && + !mpool.IsMPoolCapacityFailure(err): return moerr.NewResourceExhaustedf(ctx, "hash build memory budget exceeded; reduce join build width or query concurrency, increase processLimitationSize, or lower join_spill_mem for an eligible shuffle join") case errors.Is(err, process.ErrHashBuildBudgetAdmission): diff --git a/pkg/sql/colexec/hashbuild/errors_test.go b/pkg/sql/colexec/hashbuild/errors_test.go index 2d55a817fa9d0..d7243eee729f7 100644 --- a/pkg/sql/colexec/hashbuild/errors_test.go +++ b/pkg/sql/colexec/hashbuild/errors_test.go @@ -100,6 +100,12 @@ func TestTerminalBudgetError(t *testing.T) { require.Contains(t, err.Error(), "processLimitationSize") }) + t.Run("mpool capacity preserves allocator error", func(t *testing.T) { + capacity := moerr.NewMPoolCapacityNoCtxf("mpool out of space") + require.Same(t, capacity, + TerminalBudgetError(context.Background(), capacity)) + }) + t.Run("physical lifecycle failure stays fatal", func(t *testing.T) { joined := errors.Join( mpool.ErrAllocationAccountCapacity, diff --git a/pkg/sql/colexec/hashbuild/hashmap_test.go b/pkg/sql/colexec/hashbuild/hashmap_test.go index 900123d85d10a..37d6bc517f148 100644 --- a/pkg/sql/colexec/hashbuild/hashmap_test.go +++ b/pkg/sql/colexec/hashbuild/hashmap_test.go @@ -1720,9 +1720,11 @@ func BenchmarkCopyBuildBatchAccounting(b *testing.B) { } } -// BenchmarkResidentHashBuildAccounting compares the complete resident owner -// closure, not only the primitive allocator: copied batches, key expression, -// hash cells/descriptors, and terminal release all run on every iteration. +// BenchmarkResidentHashBuildAccounting compares a local physical account with +// the production shared budget controller across the complete resident owner +// closure: copied batches, key expression, hash cells/descriptors, and terminal +// release all run on every iteration. The builder intentionally has no +// unaccounted mode. // The 32-row case models high-frequency TP statements; 8,192 rows exercises a // full physical batch without entering spill. func BenchmarkResidentHashBuildAccounting(b *testing.B) { @@ -1733,73 +1735,92 @@ func BenchmarkResidentHashBuildAccounting(b *testing.B) { if stringKey { kind = "varchar" } - b.Run(fmt.Sprintf("%s/rows-%d", kind, rows), func(b *testing.B) { - proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) - defer proc.Free() - var input *batch.Batch - var keyType types.Type - if stringKey { - input = makeStrBatch(b, rows, proc) - keyType = types.T_varchar.ToType() - } else { - input = makeIntBatch(b, rows, proc) - keyType = types.T_int32.ToType() - } - defer input.Clean(proc.Mp()) - - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - if err != nil { - b.Fatal(err) - } - registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) - if err != nil { - b.Fatal(err) - } - account, err := registry.OpenWithController(capBytes, generation) - if err != nil { - b.Fatal(err) + for _, controlled := range []bool{false, true} { + mode := "local-account" + if controlled { + mode = "budget-controlled" } + b.Run(fmt.Sprintf("%s/rows-%d/%s", kind, rows, mode), func(b *testing.B) { + proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) + defer proc.Free() + var input *batch.Batch + var keyType types.Type + if stringKey { + input = makeStrBatch(b, rows, proc) + keyType = types.T_varchar.ToType() + } else { + input = makeIntBatch(b, rows, proc) + keyType = types.T_int32.ToType() + } + defer input.Clean(proc.Mp()) - b.ReportAllocs() - b.SetBytes(int64(input.Size())) - b.ResetTimer() - for range b.N { - hb := &HashmapBuilder{} - hb.SetBudget(generation) - if err = hb.SetAllocationAccount(account); err != nil { + var generation *process.HashBuildBudgetGeneration + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + if err != nil { b.Fatal(err) } - if err = hb.Prepare( - []*plan.Expr{newExpr(0, keyType)}, - -1, - -1, - nil, - proc, - ); err != nil { - b.Fatal(err) + var account *mpool.AllocationAccount + if controlled { + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err = budget.OpenGeneration(1) + if err != nil { + b.Fatal(err) + } + account, err = registry.OpenWithController(capBytes, generation) + if err != nil { + b.Fatal(err) + } + } else { + account, err = registry.Open(capBytes) + if err != nil { + b.Fatal(err) + } } - hb.InputBatchRowCount = input.RowCount() - if err = hb.CopyBuildBatch(input, proc); err != nil { - b.Fatal(err) + + b.ReportAllocs() + b.SetBytes(int64(input.Size())) + b.ResetTimer() + for range b.N { + hb := &HashmapBuilder{} + if controlled { + hb.SetBudget(generation) + } + if err := hb.SetAllocationAccount(account); err != nil { + b.Fatal(err) + } + if err := hb.Prepare( + []*plan.Expr{newExpr(0, keyType)}, + -1, + -1, + nil, + proc, + ); err != nil { + b.Fatal(err) + } + hb.InputBatchRowCount = input.RowCount() + if err := hb.CopyBuildBatch(input, proc); err != nil { + b.Fatal(err) + } + if err := hb.BuildHashmap(false, false, false, proc); err != nil { + b.Fatal(err) + } + hb.Free(proc) + } + b.StopTimer() + if account.Snapshot().Used != 0 { + b.Fatalf("account used = %d", account.Snapshot().Used) } - if err = hb.BuildHashmap(false, false, false, proc); err != nil { + if _, _, err := registry.CompleteTerminal(account); err != nil { b.Fatal(err) } - hb.Free(proc) - } - b.StopTimer() - if generation.Used() != 0 { - b.Fatalf("generation used = %d", generation.Used()) - } - if account.Snapshot().Used != 0 { - b.Fatalf("account used = %d", account.Snapshot().Used) - } - if _, _, err = registry.CompleteTerminal(account); err != nil { - b.Fatal(err) - } - generation.Close() - }) + if controlled { + if generation.Used() != 0 { + b.Fatalf("generation used = %d", generation.Used()) + } + generation.Close() + } + }) + } } } } diff --git a/pkg/sql/colexec/hashbuild/pressure.go b/pkg/sql/colexec/hashbuild/pressure.go index aa89f4f271746..851f7f05b9592 100644 --- a/pkg/sql/colexec/hashbuild/pressure.go +++ b/pkg/sql/colexec/hashbuild/pressure.go @@ -53,6 +53,23 @@ func MemoryPressureReasonOf(err error) MemoryPressureReason { return MemoryPressureMinimumUnit } + // Physical-account ownership and lifecycle failures dominate any joined + // logical budget error. Otherwise a sealed or invariant account could be + // mistaken for retryable memory pressure. + switch mpool.AllocationFailureReasonOf(err) { + case mpool.AllocationFailureCapacity: + return MemoryPressureCapacity + case mpool.AllocationFailureSealed, + mpool.AllocationFailureSuspended: + return MemoryPressureSealed + case mpool.AllocationFailureMismatch: + return MemoryPressureMismatch + case mpool.AllocationFailureAllocatorLimit: + return MemoryPressureAllocatorLimit + case mpool.AllocationFailureInvariant: + return MemoryPressureInvariant + } + var budgetErr *process.HashBuildBudgetError if errors.As(err, &budgetErr) { switch budgetErr.Kind { @@ -77,27 +94,10 @@ func MemoryPressureReasonOf(err error) MemoryPressureReason { } } - switch mpool.AllocationFailureReasonOf(err) { - case mpool.AllocationFailureCapacity: - return MemoryPressureCapacity - case mpool.AllocationFailureSealed, - mpool.AllocationFailureSuspended: - return MemoryPressureSealed - case mpool.AllocationFailureMismatch: - return MemoryPressureMismatch - case mpool.AllocationFailureAllocatorLimit: - return MemoryPressureAllocatorLimit - case mpool.AllocationFailureInvariant: - return MemoryPressureInvariant - } - // Resource-ledger helpers may still return a bare lifecycle sentinel. if errors.Is(err, process.ErrHashBuildBudgetClosed) { return MemoryPressureSealed } - if errors.Is(err, process.ErrHashBuildBudgetAdmission) { - return MemoryPressureCapacity - } if errors.Is(err, process.ErrHashBuildBudgetInvalid) || errors.Is(err, process.ErrHashBuildCeilingMissing) { return MemoryPressureInvalid diff --git a/pkg/sql/colexec/hashbuild/pressure_test.go b/pkg/sql/colexec/hashbuild/pressure_test.go index 1ee55c7f2b2de..799a3a08ffc22 100644 --- a/pkg/sql/colexec/hashbuild/pressure_test.go +++ b/pkg/sql/colexec/hashbuild/pressure_test.go @@ -15,9 +15,11 @@ package hashbuild import ( + "errors" "fmt" "testing" + "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/vm/process" "github.com/stretchr/testify/require" @@ -35,8 +37,9 @@ func TestMemoryPressureReasonSeparatesCapacityFromLifecycle(t *testing.T) { {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorAdmission, Component: process.HashBuildBudgetComponentSpillFD}, MemoryPressureSpillFDLimit}, {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorClosed}, MemoryPressureSealed}, {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorInvalid}, MemoryPressureInvalid}, - {fmt.Errorf("wrapped: %w", process.ErrHashBuildBudgetAdmission), MemoryPressureCapacity}, + {fmt.Errorf("wrapped: %w", process.ErrHashBuildBudgetAdmission), MemoryPressureNone}, {mpool.ErrAllocationAccountCapacity, MemoryPressureCapacity}, + {moerr.NewMPoolCapacityNoCtxf("test"), MemoryPressureCapacity}, {mpool.ErrAllocationMetadataSlots, MemoryPressureCapacity}, {mpool.ErrAllocationAccountSealed, MemoryPressureSealed}, {mpool.ErrAllocationAccountMismatch, MemoryPressureMismatch}, @@ -48,6 +51,22 @@ func TestMemoryPressureReasonSeparatesCapacityFromLifecycle(t *testing.T) { require.Equal(t, test.reason, MemoryPressureReasonOf(test.err)) require.Equal(t, test.reason == MemoryPressureCapacity, IsRetryableMemoryCapacity(test.err)) } + + memoryAdmission := &process.HashBuildBudgetError{ + Kind: process.HashBuildBudgetErrorAdmission, + Component: process.HashBuildBudgetComponentMemory, + } + for _, test := range []struct { + err error + reason MemoryPressureReason + }{ + {errors.Join(memoryAdmission, mpool.ErrAllocationAccountInvariant), MemoryPressureInvariant}, + {errors.Join(memoryAdmission, mpool.ErrAllocationAccountInvalid), MemoryPressureInvariant}, + {errors.Join(memoryAdmission, mpool.ErrAllocationAccountSealed), MemoryPressureSealed}, + } { + require.Equal(t, test.reason, MemoryPressureReasonOf(test.err)) + require.False(t, IsRetryableMemoryCapacity(test.err)) + } } func TestPressureRetryGuardRequiresMonotonicProgress(t *testing.T) { diff --git a/pkg/sql/colexec/hashbuild/types.go b/pkg/sql/colexec/hashbuild/types.go index 37724acee6c6e..bbe643d956e17 100644 --- a/pkg/sql/colexec/hashbuild/types.go +++ b/pkg/sql/colexec/hashbuild/types.go @@ -82,6 +82,7 @@ const ( HashBuildAllocationSiteUniqueKeyNulls HashBuildAllocationSiteUniqueKeyGrouping HashBuildAllocationSiteRuntimeFilterPayload + HashBuildAllocationSiteRuntimeFilterScratch HashBuildAllocationSiteDedupIgnoreBitmap HashBuildAllocationSiteDedupDeleteBitmap HashBuildAllocationSiteDedupLastRows @@ -570,10 +571,7 @@ func hasHashBuildDiagnosticStats(extra map[string]int64) bool { extra["QueryHashBudgetRejects"] != 0 || extra["HashBuildRuntimeFilterCollectionFallbacks"] != 0 || extra["HashBuildRuntimeFilterBudgetFallbacks"] != 0 || - extra["HashBuildRuntimeFilterAllocationFallbacks"] != 0 || - extra["HashBuildSpillScratchReserveRejects"] != 0 || - extra["HashBuildSpillScratchGrowRejects"] != 0 || - extra["HashBuildSpillScratchGrowCount"] != 0 + extra["HashBuildRuntimeFilterAllocationFallbacks"] != 0 } func (hashBuild *HashBuild) publishJoinMap(proc *process.Process, jm *message.JoinMap) bool { diff --git a/pkg/sql/colexec/hashjoin/spill_integration_test.go b/pkg/sql/colexec/hashjoin/spill_integration_test.go index 9b152e09db2c7..e97230d4eb63d 100644 --- a/pkg/sql/colexec/hashjoin/spill_integration_test.go +++ b/pkg/sql/colexec/hashjoin/spill_integration_test.go @@ -214,14 +214,19 @@ func TestShuffleJoinFiniteBudgetInitialSpillAndReSpill(t *testing.T) { func TestShuffleJoinSpillUsesCanonicalGroupingPartitionKey(t *testing.T) { for _, test := range []struct { - name string - typ types.Type - probe func(*process.Process) *vector.Vector - build func(*process.Process) *vector.Vector + name string + typ types.Type + rows int + spillThreshold int64 + wantRespill bool + probe func(*process.Process) *vector.Vector + build func(*process.Process) *vector.Vector }{ { - name: "varchar", - typ: types.T_varchar.ToType(), + name: "varchar", + typ: types.T_varchar.ToType(), + rows: 1, + spillThreshold: 1, probe: func(proc *process.Process) *vector.Vector { return testutil.MakeVarcharVector([]string{"probe"}, nil, proc.Mp()) }, @@ -230,8 +235,10 @@ func TestShuffleJoinSpillUsesCanonicalGroupingPartitionKey(t *testing.T) { }, }, { - name: "int32", - typ: types.T_int32.ToType(), + name: "int32", + typ: types.T_int32.ToType(), + rows: 1, + spillThreshold: 1, probe: func(proc *process.Process) *vector.Vector { return testutil.MakeInt32Vector([]int32{222}, nil, proc.Mp()) }, @@ -239,10 +246,45 @@ func TestShuffleJoinSpillUsesCanonicalGroupingPartitionKey(t *testing.T) { return testutil.MakeInt32Vector([]int32{111}, nil, proc.Mp()) }, }, + { + name: "scaled-float32", + rows: 8192, + spillThreshold: 50, + wantRespill: true, + typ: func() types.Type { + typ := types.T_float32.ToType() + typ.Scale = 2 + return typ + }(), + probe: func(proc *process.Process) *vector.Vector { + values := make([]float32, 8192) + values[0] = 1.234 + for i := 1; i < len(values); i++ { + values[i] = float32(i)/10 + 0.001 + } + vec := testutil.MakeFloat32Vector(values, nil, proc.Mp()) + vec.GetType().Scale = 2 + return vec + }, + build: func(proc *process.Process) *vector.Vector { + values := make([]float32, 8192) + values[0] = 9.876 + for i := 1; i < len(values); i++ { + values[i] = float32(i) / 10 + } + vec := testutil.MakeFloat32Vector(values, nil, proc.Mp()) + vec.GetType().Scale = 2 + return vec + }, + }, } { t.Run(test.name, func(t *testing.T) { keyExpr := []*plan.Expr{{ - Typ: plan.Type{Id: int32(test.typ.Oid), Width: test.typ.Width}, + Typ: plan.Type{ + Id: int32(test.typ.Oid), + Width: test.typ.Width, + Scale: test.typ.Scale, + }, Expr: &plan.Expr_Col{Col: &plan.ColRef{ ColPos: 0, }}, @@ -257,10 +299,10 @@ func TestShuffleJoinSpillUsesCanonicalGroupingPartitionKey(t *testing.T) { tc.arg.NonEqCond = nil tc.arg.IsShuffle = true tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 1 + tc.arg.SpillThreshold = test.spillThreshold tc.barg.IsShuffle = true tc.barg.ShuffleIdx = 0 - tc.barg.SpillThreshold = 1 + tc.barg.SpillThreshold = test.spillThreshold tc.barg.NeedBatches = false tc.barg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{ Tag: tc.arg.JoinMapTag + 1_500, @@ -268,15 +310,20 @@ func TestShuffleJoinSpillUsesCanonicalGroupingPartitionKey(t *testing.T) { probe := batch.NewWithSize(1) probe.Vecs[0] = test.probe(tc.proc) - probe.SetRowCount(1) + probe.SetRowCount(test.rows) build := batch.NewWithSize(1) build.Vecs[0] = test.build(tc.proc) - build.SetRowCount(1) + build.SetRowCount(test.rows) probe.Vecs[0].GetGrouping().Add(0) build.Vecs[0].GetGrouping().Add(0) resetChildrenWithBatch(tc.arg, probe) resetHashBuildChildrenWithBatch(tc.barg, build) + spillBefore := promtestutil.ToFloat64( + metricv2.HashBuildSpillDepthCounter.WithLabelValues("spill", "1")) + respillBefore := promtestutil.ToFloat64( + metricv2.HashBuildSpillDepthCounter.WithLabelValues("respill", "2")) + require.NoError(t, tc.arg.Prepare(tc.proc)) require.NoError(t, tc.barg.Prepare(tc.proc)) _, err := vm.Exec(tc.barg, tc.proc) @@ -293,9 +340,13 @@ func TestShuffleJoinSpillUsesCanonicalGroupingPartitionKey(t *testing.T) { break } } - require.Equal(t, 1, rows) - require.Positive(t, - tc.barg.OpAnalyzer.GetOpStats().ExtraStats["HashBuildSpillStarts"]) + require.Equal(t, test.rows, rows) + require.Greater(t, promtestutil.ToFloat64( + metricv2.HashBuildSpillDepthCounter.WithLabelValues("spill", "1")), spillBefore) + if test.wantRespill { + require.Greater(t, promtestutil.ToFloat64( + metricv2.HashBuildSpillDepthCounter.WithLabelValues("respill", "2")), respillBefore) + } tc.arg.Reset(tc.proc, false, nil) tc.barg.Reset(tc.proc, false, nil) diff --git a/pkg/sql/colexec/indexbuild/build.go b/pkg/sql/colexec/indexbuild/build.go index bb2c0f7647c04..949373430e169 100644 --- a/pkg/sql/colexec/indexbuild/build.go +++ b/pkg/sql/colexec/indexbuild/build.go @@ -354,7 +354,7 @@ func (ctr *container) handleRuntimeFilter(ap *IndexBuild, proc *process.Process) return nil } if encoding == keycodec.ExactRuntimeFilterFloatZeroClosed { - if err := runtimefilter.CloseFloatSignedZero(vec, proc.Mp(), nil); err != nil { + if err := runtimefilter.CloseFloatSignedZero(vec, proc.Mp()); err != nil { if ctr.fallbackRuntimeFilter(ap, proc, err) { return nil } diff --git a/pkg/sql/colexec/runtimefilter/contract.go b/pkg/sql/colexec/runtimefilter/contract.go index 7b87a4ca2d2f5..453c09dd7129f 100644 --- a/pkg/sql/colexec/runtimefilter/contract.go +++ b/pkg/sql/colexec/runtimefilter/contract.go @@ -46,7 +46,8 @@ func (e *optionalAllocationError) Unwrap() error { return e.cause } // MarkOptionalAllocationError preserves the allocation error while giving a // runtime-filter producer a narrow fail-open classification. func MarkOptionalAllocationError(err error) error { - if err == nil || IsOptionalAllocationError(err) { + if err == nil || IsOptionalAllocationError(err) || + !mpool.IsRetryableAllocationCapacity(err) { return err } return &optionalAllocationError{cause: err} @@ -78,6 +79,11 @@ func ClassifyOptionalFallback(err error) OptionalFallbackKind { errors.Is(err, context.DeadlineExceeded) { return OptionalFallbackNone } + allocationReason := mpool.AllocationFailureReasonOf(err) + if allocationReason != mpool.AllocationFailureNone && + allocationReason != mpool.AllocationFailureCapacity { + return OptionalFallbackNone + } // Reject every known fatal branch before accepting an admission branch. // This also keeps errors.Join(admission, fatal) fatal regardless of the // traversal order chosen by errors.As below. @@ -90,7 +96,8 @@ func ClassifyOptionalFallback(err error) OptionalFallbackKind { var budgetErr *process.HashBuildBudgetError if errors.As(err, &budgetErr) { if budgetErr != nil && - budgetErr.Kind == process.HashBuildBudgetErrorAdmission { + budgetErr.Kind == process.HashBuildBudgetErrorAdmission && + budgetErr.Component == process.HashBuildBudgetComponentMemory { return OptionalFallbackBudgetAdmission } return OptionalFallbackNone @@ -102,7 +109,8 @@ func ClassifyOptionalFallback(err error) OptionalFallbackKind { return OptionalFallbackNone } - if IsOptionalAllocationError(err) { + if IsOptionalAllocationError(err) && + mpool.IsRetryableAllocationCapacity(err) { return OptionalFallbackAllocation } return OptionalFallbackNone @@ -289,12 +297,10 @@ func planType(typ plan.Type) types.Type { } // CloseFloatSignedZero appends the complementary representation when an exact -// float payload contains only one of +0 and -0. beforeAppend lets budgeted -// producers reserve the vector-growth overlap before the allocation occurs. +// float payload contains only one of +0 and -0. func CloseFloatSignedZero( vec *vector.Vector, mp *mpool.MPool, - beforeAppend func() (release func(), err error), ) error { if vec == nil || mp == nil { return moerr.NewInternalErrorNoCtx("invalid float runtime-filter vector") @@ -344,18 +350,6 @@ func CloseFloatSignedZero( return nil } - var release func() - var err error - if beforeAppend != nil { - release, err = beforeAppend() - if err != nil { - return err - } - } - if release != nil { - defer release() - } - if vec.GetType().Oid == types.T_float32 { value := float32(0) if hasPositiveZero { diff --git a/pkg/sql/colexec/runtimefilter/contract_test.go b/pkg/sql/colexec/runtimefilter/contract_test.go index 5b40139a8fae4..9a1f9bff1329e 100644 --- a/pkg/sql/colexec/runtimefilter/contract_test.go +++ b/pkg/sql/colexec/runtimefilter/contract_test.go @@ -20,6 +20,7 @@ import ( "testing" "github.com/matrixorigin/matrixone/pkg/common/hashmap/keycodec" + "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" @@ -30,13 +31,22 @@ import ( ) func TestClassifyOptionalFallbackFatalFirst(t *testing.T) { - mpoolErr := errors.New("mpool allocation failed") + mpoolErr := mpool.ErrAllocationAccountCapacity providerErr := errors.New("budget provider failed") var nilBudgetErr *process.HashBuildBudgetError - budgetErr := func(kind process.HashBuildBudgetErrorKind) error { - return &process.HashBuildBudgetError{Kind: kind} + budgetErr := func( + kind process.HashBuildBudgetErrorKind, + component process.HashBuildBudgetComponent, + ) error { + return &process.HashBuildBudgetError{Kind: kind, Component: component} } marked := MarkOptionalAllocationError + memoryAdmission := func() error { + return budgetErr( + process.HashBuildBudgetErrorAdmission, + process.HashBuildBudgetComponentMemory, + ) + } tests := []struct { name string @@ -45,19 +55,29 @@ func TestClassifyOptionalFallbackFatalFirst(t *testing.T) { }{ {name: "nil", want: OptionalFallbackNone}, {name: "typed nil budget error", err: nilBudgetErr, want: OptionalFallbackNone}, - {name: "typed admission", err: budgetErr(process.HashBuildBudgetErrorAdmission), want: OptionalFallbackBudgetAdmission}, - {name: "marked typed admission", err: marked(budgetErr(process.HashBuildBudgetErrorAdmission)), want: OptionalFallbackBudgetAdmission}, - {name: "joined admission and closed", err: errors.Join(budgetErr(process.HashBuildBudgetErrorAdmission), budgetErr(process.HashBuildBudgetErrorClosed)), want: OptionalFallbackNone}, + {name: "typed memory admission", err: memoryAdmission(), want: OptionalFallbackBudgetAdmission}, + {name: "marked typed memory admission", err: marked(memoryAdmission()), want: OptionalFallbackBudgetAdmission}, + {name: "typed spill disk admission", err: budgetErr(process.HashBuildBudgetErrorAdmission, process.HashBuildBudgetComponentSpillDisk), want: OptionalFallbackNone}, + {name: "typed spill fd admission", err: budgetErr(process.HashBuildBudgetErrorAdmission, process.HashBuildBudgetComponentSpillFD), want: OptionalFallbackNone}, + {name: "typed admission without component", err: budgetErr(process.HashBuildBudgetErrorAdmission, 0), want: OptionalFallbackNone}, + {name: "joined admission and closed", err: errors.Join(memoryAdmission(), budgetErr(process.HashBuildBudgetErrorClosed, 0)), want: OptionalFallbackNone}, {name: "marked mpool allocation", err: marked(mpoolErr), want: OptionalFallbackAllocation}, {name: "plain mpool error", err: mpoolErr, want: OptionalFallbackNone}, + {name: "marked sealed", err: marked(mpool.ErrAllocationAccountSealed), want: OptionalFallbackNone}, + {name: "marked mismatch", err: marked(mpool.ErrAllocationAccountMismatch), want: OptionalFallbackNone}, + {name: "marked invariant", err: marked(mpool.ErrAllocationAccountInvariant), want: OptionalFallbackNone}, + {name: "marked allocator limit", err: marked(mpool.ErrAllocationAllocatorLimit), want: OptionalFallbackNone}, + {name: "marked suspended", err: marked(mpool.ErrAllocationAdmissionSuspended), want: OptionalFallbackNone}, + {name: "marked joined capacity and invariant", err: marked(errors.Join(mpool.ErrAllocationAccountCapacity, mpool.ErrAllocationAccountInvariant)), want: OptionalFallbackNone}, + {name: "marked joined mpool capacity and invalid", err: marked(errors.Join(moerr.NewMPoolCapacityNoCtxf("test"), mpool.ErrAllocationAccountInvalid)), want: OptionalFallbackNone}, {name: "plain provider error", err: providerErr, want: OptionalFallbackNone}, {name: "raw admission sentinel", err: process.ErrHashBuildBudgetAdmission, want: OptionalFallbackNone}, - {name: "typed closed", err: budgetErr(process.HashBuildBudgetErrorClosed), want: OptionalFallbackNone}, - {name: "marked typed closed", err: marked(budgetErr(process.HashBuildBudgetErrorClosed)), want: OptionalFallbackNone}, - {name: "typed invalid", err: budgetErr(process.HashBuildBudgetErrorInvalid), want: OptionalFallbackNone}, - {name: "marked typed invalid", err: marked(budgetErr(process.HashBuildBudgetErrorInvalid)), want: OptionalFallbackNone}, - {name: "typed ceiling missing", err: budgetErr(process.HashBuildBudgetErrorCeilingMissing), want: OptionalFallbackNone}, - {name: "marked typed ceiling missing", err: marked(budgetErr(process.HashBuildBudgetErrorCeilingMissing)), want: OptionalFallbackNone}, + {name: "typed closed", err: budgetErr(process.HashBuildBudgetErrorClosed, 0), want: OptionalFallbackNone}, + {name: "marked typed closed", err: marked(budgetErr(process.HashBuildBudgetErrorClosed, 0)), want: OptionalFallbackNone}, + {name: "typed invalid", err: budgetErr(process.HashBuildBudgetErrorInvalid, 0), want: OptionalFallbackNone}, + {name: "marked typed invalid", err: marked(budgetErr(process.HashBuildBudgetErrorInvalid, 0)), want: OptionalFallbackNone}, + {name: "typed ceiling missing", err: budgetErr(process.HashBuildBudgetErrorCeilingMissing, 0), want: OptionalFallbackNone}, + {name: "marked typed ceiling missing", err: marked(budgetErr(process.HashBuildBudgetErrorCeilingMissing, 0)), want: OptionalFallbackNone}, {name: "marked canceled", err: marked(context.Canceled), want: OptionalFallbackNone}, {name: "marked deadline", err: marked(context.DeadlineExceeded), want: OptionalFallbackNone}, {name: "marked raw closed", err: marked(process.ErrHashBuildBudgetClosed), want: OptionalFallbackNone}, @@ -279,7 +299,7 @@ func TestMarshalExactFilterVectorUsesWireSizedBudget(t *testing.T) { data, release, err := MarshalExactFilterVector(vec, mp, account, 1, 1) require.NoError(t, err) require.Len(t, data, 34+vec.Length()) - // The retained charge is the actual bytes.Buffer capacity, not a + // The retained charge is the actual caller-owned buffer capacity, not a // row-count-derived metadata estimate. require.LessOrEqual(t, budget.Used(), uint64(2*len(data))) require.NotZero(t, budget.Used()) diff --git a/pkg/sql/compile/remote_allocation_statement_group.go b/pkg/sql/compile/remote_allocation_statement_group.go index 64616d3080cb2..dd024d9b70663 100644 --- a/pkg/sql/compile/remote_allocation_statement_group.go +++ b/pkg/sql/compile/remote_allocation_statement_group.go @@ -286,7 +286,7 @@ func (p *remoteAllocationStatementParticipant) finish(cause error) ( remoteAllocationStatementGroups.Unlock() if abort { abortErr := allocationLifecycleCall(func() error { - group.board.CloseAndDrain() + group.board.Close() return nil }) abortErr = errors.Join( @@ -447,7 +447,7 @@ func expireRemoteAllocationStatementGroup( terminalErr = errors.Join( terminalErr, allocationLifecycleCall(func() error { - group.board.CloseAndDrain() + group.board.Close() return nil }), ) diff --git a/pkg/sql/compile/remote_allocation_statement_group_test.go b/pkg/sql/compile/remote_allocation_statement_group_test.go index 370873c942d71..f6c89a72ed00d 100644 --- a/pkg/sql/compile/remote_allocation_statement_group_test.go +++ b/pkg/sql/compile/remote_allocation_statement_group_test.go @@ -299,8 +299,12 @@ func TestRemoteAllocationStatementRegistrationTimerStartsBeforeFinish(t *testing t.Fatal("registration timeout did not cancel the active fragment") } require.Eventually(t, func() bool { - return destroyed.Load() == 1 && strings.Contains(board.DebugString(), "closed") + return strings.Contains(board.DebugString(), "closed") }, time.Second, time.Millisecond) + // Closing wakes the active fragment but cannot destroy ownership that the + // fragment may still be consuming. The last finish performs the drain. + require.Zero(t, destroyed.Load()) + require.Equal(t, uint64(cap(buffer)), attempt.account.Snapshot().Used) participant.stage(attempt, producer.proc.Mp()) terminal, err := participant.finish(errors.New("active fragment observed cancellation")) @@ -309,6 +313,7 @@ func TestRemoteAllocationStatementRegistrationTimerStartsBeforeFinish(t *testing require.Len(t, terminal.allocation, 1) require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.allocation[0].State) require.Zero(t, terminal.allocation[0].Used) + require.Equal(t, int32(1), destroyed.Load()) require.Zero(t, terminal.memory.LiveBytesAtSeal) require.Zero(t, registry.LiveAllocationMetadata()) require.False(t, remoteAllocationStatementGroupRegistered(board)) diff --git a/pkg/vm/message/message.go b/pkg/vm/message/message.go index 22c2a500c3b40..44eba0cda7299 100644 --- a/pkg/vm/message/message.go +++ b/pkg/vm/message/message.go @@ -76,6 +76,7 @@ type MessageCenter struct { type MessageBoard struct { reset bool // for debug purpose closed bool + drained bool multiCN bool stmtId uuid.UUID messageCenter *MessageCenter @@ -174,13 +175,26 @@ func (m *MessageBoard) Reset() *MessageBoard { return m } +// Close prevents publication and wakes blocked receivers without destroying +// queued ownership. It is safe while consumers are still unwinding; the +// terminal owner must call CloseAndDrain after every producer and consumer is +// quiescent. True identifies the call that first closed the board. +func (m *MessageBoard) Close() bool { + return m.close(false) +} + // CloseAndDrain is the terminal MessageBoard boundary for one execution // attempt. Callers invoke it only after all scope and remote-notifier producers -// are quiescent. It removes a multi-CN registration, destroys every queued -// ownership-bearing message, and prevents a late producer from republishing -// into the closed generation. The operation is idempotent; true identifies -// the call that performed the close. +// and consumers are quiescent. It removes a multi-CN registration, destroys +// every queued ownership-bearing message, and prevents a late producer from +// republishing into the closed generation. The operation is idempotent; true +// identifies the call that performed the drain, including a drain after an +// earlier Close. func (m *MessageBoard) CloseAndDrain() bool { + return m.close(true) +} + +func (m *MessageBoard) close(drain bool) bool { if m == nil || m.rwMutex == nil { return false } @@ -199,20 +213,27 @@ func (m *MessageBoard) CloseAndDrain() bool { m.rwMutex.Lock() defer m.rwMutex.Unlock() - if m.closed { - return false - } - m.closed = true - m.reset = true - for _, waiter := range m.waiters { - if waiter == nil { - continue - } - select { - case waiter <- true: - default: + firstClose := !m.closed + if firstClose { + m.closed = true + m.reset = true + for _, waiter := range m.waiters { + if waiter == nil { + continue + } + select { + case waiter <- true: + default: + } } } + if !drain { + return firstClose + } + if m.drained { + return false + } + m.drained = true m.cleanupQueuedMessagesLocked() return true } diff --git a/pkg/vm/message/message_test.go b/pkg/vm/message/message_test.go index 549bd8153b364..07a87eb520dfd 100644 --- a/pkg/vm/message/message_test.go +++ b/pkg/vm/message/message_test.go @@ -166,6 +166,34 @@ func TestMessageBoardCloseAndDrainRemovesMultiCNRegistration(t *testing.T) { require.False(t, ok) } +func TestMessageBoardCloseDefersQueuedOwnershipDrain(t *testing.T) { + mb := NewMessageBoard() + var destroyed atomic.Int32 + SendMessage(testMessage{tag: 1, destroyed: &destroyed}, mb) + + receiver := NewMessageReceiver( + []int32{2}, + AddrBroadCastOnCurrentCN(), + mb, + ) + waiting := make(chan error, 1) + go func() { + _, _, err := receiver.ReceiveMessage(true, context.Background()) + waiting <- err + }() + + require.True(t, mb.Close()) + require.False(t, mb.Close()) + require.ErrorContains(t, <-waiting, "message board is closed") + require.Zero(t, destroyed.Load()) + require.Len(t, mb.messages, 1) + + require.True(t, mb.CloseAndDrain()) + require.False(t, mb.CloseAndDrain()) + require.Equal(t, int32(1), destroyed.Load()) + require.Empty(t, mb.messages) +} + func TestClosedMessageBoardLatePayloadDrainsOriginalGeneration(t *testing.T) { registry, err := mpool.NewAllocationAccountRegistry(1, 1) require.NoError(t, err) diff --git a/pkg/vm/process/hashbuild_budget_test.go b/pkg/vm/process/hashbuild_budget_test.go index 2a657433c9cc4..31cba8e4e55bf 100644 --- a/pkg/vm/process/hashbuild_budget_test.go +++ b/pkg/vm/process/hashbuild_budget_test.go @@ -16,6 +16,7 @@ package process import ( "errors" + "fmt" "math" "sync" "sync/atomic" @@ -504,27 +505,61 @@ func TestHashBuildBudgetUsesCurrentMemoryInputs(t *testing.T) { } func BenchmarkHashBuildBudgetAllocationAccount(b *testing.B) { - budget := MustNewHashBuildBudget(math.MaxUint64, math.MaxUint64) - generation, err := budget.OpenGeneration(1) - if err != nil { - b.Fatal(err) - } - registry, err := mpool.NewAllocationAccountRegistry(1, uint64(b.N)+1) - if err != nil { - b.Fatal(err) - } - account, err := registry.OpenWithController(math.MaxUint64, generation) - if err != nil { - b.Fatal(err) - } - mp := mpool.MustNewZero() - b.ReportAllocs() - b.ResetTimer() - for i := 0; i < b.N; i++ { - allocation, allocErr := mp.AllocAccounted(1, account, 1, 1) - if allocErr != nil { - b.Fatal(allocErr) + for _, accounted := range []bool{false, true} { + mode := "unaccounted" + if accounted { + mode = "accounted" + } + for _, size := range []int{64, 4 << 10, 64 << 10} { + b.Run(mode+"/alloc-free/"+fmt.Sprint(size), func(b *testing.B) { + mp := mpool.MustNewZero() + var generation *HashBuildBudgetGeneration + var registry *mpool.AllocationAccountRegistry + var account *mpool.AllocationAccount + if accounted { + budget := MustNewHashBuildBudget(math.MaxUint64, math.MaxUint64) + var err error + generation, err = budget.OpenGeneration(1) + if err != nil { + b.Fatal(err) + } + registry, err = mpool.NewAllocationAccountRegistry(1, 2) + if err != nil { + b.Fatal(err) + } + account, err = registry.OpenWithController(math.MaxInt64, generation) + if err != nil { + b.Fatal(err) + } + } + + b.ReportAllocs() + b.SetBytes(int64(size)) + b.ResetTimer() + for range b.N { + var allocation []byte + var err error + if accounted { + allocation, err = mp.AllocAccounted(size, account, 1, 1) + } else { + allocation, err = mp.Alloc(size, true) + } + if err != nil { + b.Fatal(err) + } + mp.Free(allocation) + } + b.StopTimer() + if accounted { + if generation.Used() != 0 || account.Snapshot().Used != 0 { + b.Fatal("physical allocation capacity leaked") + } + if _, _, err := registry.CompleteTerminal(account); err != nil { + b.Fatal(err) + } + generation.Close() + } + }) } - mp.Free(allocation) } } From 258010da171c24b6294a1f606fd4387e5ac5dbf0 Mon Sep 17 00:00:00 2001 From: aptend Date: Sat, 1 Aug 2026 22:53:43 +0800 Subject: [PATCH 46/61] docs: record allocation accounting benchmark evidence --- .../26459_allocation_accounting_bench.txt | 130 ++++++++++++++---- 1 file changed, 101 insertions(+), 29 deletions(-) diff --git a/docs/design/evidence/26459_allocation_accounting_bench.txt b/docs/design/evidence/26459_allocation_accounting_bench.txt index cf1d5b7b50707..c7f55f43bbc25 100644 --- a/docs/design/evidence/26459_allocation_accounting_bench.txt +++ b/docs/design/evidence/26459_allocation_accounting_bench.txt @@ -3,11 +3,12 @@ Artifact Date: 2026-08-01 Branch: feature/26459-statement-lifecycle -Base: origin/main ecc389d420 +Base: origin/main 49f9e33cae48 +Code under test: ea5b314cd247 Go: go1.26.4 linux/amd64 CPU: 11th Gen Intel Core i7-11700, GOMAXPROCS=16 -All commands used the MatrixOne CGo wrapper: +All commands used the same CGo environment as the MatrixOne test wrapper: .agents/skills/mo-dev/scripts/mo-cgo-test @@ -22,19 +23,40 @@ Command: mo-cgo-test ./pkg/common/mpool -run '^$' \ -bench '^BenchmarkMPoolAccountedAllocation$' \ - -benchmem -benchtime=1s -count=5 + -benchmem -benchtime=300ms -count=3 Operation Unaccounted Accounted Delta -alloc/free 64 B 248.6 ns 357.5 ns +43.8% -alloc/free 4 KiB 294.6 ns 410.5 ns +39.3% -alloc/free 16 KiB 341.4 ns 449.3 ns +31.6% -alloc/free 64 KiB 1014 ns 1121 ns +10.6% -grow replacement 1292 ns 1492 ns +15.5% -parallel accounted alloc/free 64K 285.4 ns +alloc/free 64 B 252.1 ns 356.5 ns +41.4% +alloc/free 4 KiB 302.9 ns 405.4 ns +33.8% +alloc/free 16 KiB 343.4 ns 447.1 ns +30.2% +alloc/free 64 KiB 1024 ns 1123 ns +9.7% +grow replacement 1307 ns 1495 ns +14.4% +parallel accounted alloc/free 64K 280.7 ns Every sample reported 0 B/op and 0 allocs/op. The fixed transaction cost is largest for tiny allocations; retained HashBuild cells and batch buffers use -larger capacity changes rather than one admission per row. +larger capacity changes rather than one admission per row. Fault-injection +branches that existed only for tests were removed before this measurement. + + +Production shared-budget admission +================================== + +Command: + + mo-cgo-test ./pkg/vm/process -run '^$' \ + -bench '^BenchmarkHashBuildBudgetAllocationAccount$' \ + -benchmem -benchtime=300ms -count=3 + +Operation Unaccounted Budget-controlled Delta +alloc/free 64 B 253.1 ns 503.5 ns +98.9% +alloc/free 4 KiB 301.4 ns 549.7 ns +82.4% +alloc/free 64 KiB 1029 ns 1304 ns +26.7% + +This deliberately measures the full production controller mutex and aggregate +metrics, not only a local AllocationAccount. Every sample remains 0 B/op and +0 allocs/op. This primitive is the worst-case fixed cost; retained-closure +results below show the operator-level effect after allocation reuse. Vector growth and reuse @@ -43,16 +65,56 @@ Vector growth and reuse Command: mo-cgo-test ./pkg/container/vector -run '^$' \ - -bench '^BenchmarkVectorAllocationAccount$' \ + -bench '^(BenchmarkVectorAllocationAccount|BenchmarkVectorElementAccounting)$' \ -benchmem -benchtime=1s -count=3 Operation Unaccounted Accounted Allocation result -fixed preextend/free 1047 ns 1158 ns 0 B/op, 0 allocs/op -varlen preextend/free 77963 ns 78360 ns 48 B/op, 2 allocs/op -accounted fixed Reset reuse 2.63 ns 0 B/op, 0 allocs/op +fixed preextend/free 1066 ns 1172 ns 0 B/op, 0 allocs/op +varlen preextend/free 79088 ns 78729 ns 48 B/op, 2 allocs/op +accounted fixed Reset reuse 2.65 ns 0 B/op, 0 allocs/op Accounting adds no Go object to vector allocation or Reset reuse. The varlen -difference was 0.5%, inside local run noise. +difference was -0.5%, inside local run noise. + +The per-row UnionOne and Copy hot paths were measured separately after +eliminating duplicate bitmap probes and skipping accounted-capacity checks for +ordinary vectors: + +Operation Unaccounted Accounted Delta +UnionOne, retained capacity 11.58 ns 11.83 ns +2.2% +Copy, retained capacity 8.96 ns 8.95 ns -0.1% + +All samples were 0 B/op and 0 allocs/op. + + +Caller-owned Packer storage +=========================== + +Command: + + # origin/main 49f9e33cae48 + taskset -c 4 env GOMAXPROCS=1 mo-cgo-test \ + ./pkg/container/types -run '^$' \ + -bench '^BenchmarkPacker(Encode)?$' \ + -benchmem -cpu=1 -benchtime=500ms -count=8 + + # code commit ea5b314cd247 + taskset -c 4 env GOMAXPROCS=1 mo-cgo-test \ + ./pkg/container/types -run '^$' \ + -bench '^BenchmarkPacker(Encode)?$' \ + -benchmem -cpu=1 -benchtime=500ms -count=8 + +Operation origin/main Code commit Result +allocator-backed retained encode 9.58 ns 9.23 ns -3.7% +fixed-buffer retained encode n/a 9.15 ns +Packer construction/close 48 B/op 48 B/op unchanged + +All retained-encode samples report 0 B/op and 0 allocs/op. The fixed-buffer +path used by serialized runtime filters therefore adds no retained-encode +allocation or CPU regression. Keeping fixed/overflow as one-byte state also +preserves the baseline 48-byte Packer allocation class; an error interface in +the object would have raised it to 64 B/op. Runtime-filter scratch storage is +admitted once through the statement physical account. Retained HashBuild closure @@ -61,18 +123,19 @@ Retained HashBuild closure Command: mo-cgo-test ./pkg/sql/colexec/hashbuild -run '^$' \ - -bench '^(BenchmarkResidentHashBuildAccounting|BenchmarkCopyBuildBatchAccounting)$' \ - -benchmem -benchtime=1s -count=3 + -bench '^BenchmarkResidentHashBuildAccounting$' \ + -benchmem -benchtime=300ms -count=3 -Benchmark Median B/op allocs/op -copy retained build batch 18919 ns 760 7 -resident int key, 32 rows 4428 ns 2320 24 -resident varchar key, 32 rows 7919 ns 20136 23 -resident int key, 8192 rows 137526 ns 10432 26 -resident varchar key, 8192 rows 411112 ns 20233 27 +Benchmark Local account Budget-controlled Delta +resident int key, 32 rows 4034 ns 4400 ns +9.1% +resident varchar key, 32 rows 6953 ns 7649 ns +10.0% +resident int key, 8192 rows 137067 ns 139062 ns +1.5% +resident varchar key, 8192 rows 422190 ns 423558 ns +0.3% -The retained data itself is MPool-backed; B/op here is bounded Go control -state. Allocations do not scale per input row. +Both modes use the one production physical-account ownership path; the only +difference is whether that account also uses the shared query/CN controller. +B/op and allocs/op are identical within each pair. The full-batch delta is +0.3-1.5%; the 32-row case exposes the fixed cost rather than a row-scaled cost. Pipeline spool ownership transfer @@ -84,7 +147,7 @@ Command: -bench '^BenchmarkCachedBatchReuse$' \ -benchmem -benchtime=2s -count=5 -Median: 1584 ns/op, 224 B/op, 3 allocs/op. +Median: 1603 ns/op, 224 B/op, 3 allocs/op. The same DetachedBuffer path handles both storage outside the controlled HashBuild domain and account-provenance storage. There is no second raw-byte @@ -104,7 +167,8 @@ Ordinary 256-row grouped path: 6075 ns/op, 9344 B/op, 6 allocs/op. Alternating ordinary/GROUPING rows: 26869 ns/op, 19256 B/op, 12 allocs/op. A paired temporary reference containing the pre-fix ordinary hash loop was -measured in the same binary and then removed: reference median 5989 ns/op, +measured earlier on this branch in the same binary and then removed (the +Sample code is unchanged by the final review fixes): reference median 5989 ns/op, final median 6064 ns/op (+1.25%); both were 9344 B/op and 6 allocs/op. The committed alternating-domain test also applies AllocsPerRun and proves the iterator count remains constant per batch rather than per grouping run. @@ -119,7 +183,7 @@ Command: -bench '^BenchmarkSpillScatterAccounting$' \ -benchmem -benchtime=1s -count=5 -Final median: 62490 ns/op, 344 B/op, 4 allocs/op, about 524 MB/s. +Final median: 62709 ns/op, 344 B/op, 4 allocs/op, about 523 MB/s. An allocation profile found that primitive marshal fallbacks took the address of integer parameters even when AccountedBuffer's typed writer fast path was @@ -133,9 +197,17 @@ Acceptance ========== - physical admission/release adds zero Go allocations; +- the production shared controller adds no Go allocations and no per-row + accounting object; - vector and spool reuse retain bounded allocation counts; +- UnionOne and Copy retained-capacity paths add no allocations and show at + most a 2.2% local accounting delta; +- caller-owned Packer encoding matches allocator-backed retained encoding and + adds no Go allocations; - ordinary Sample grouping has unchanged allocation shape and only 1.25% measured CPU delta; - mixed GROUPING input uses a fixed number of iterators per batch; - spill serialization has four bounded Go allocations per 4096-row scatter; -- no benchmark indicates a per-row accounting object or compatibility ledger. +- full-batch resident HashBuild controller overhead is 0.3-1.5%, while the + intentionally worst-case 32-row microbenchmark is 9.1-10.0%; +- no benchmark indicates a compatibility ledger or unaccounted builder path. From 741f07d666281aa0555b2605c805e9431b4319ac Mon Sep 17 00:00:00 2001 From: aptend Date: Sun, 2 Aug 2026 01:11:40 +0800 Subject: [PATCH 47/61] fix: finalize accounted pipeline spools after consumers stop --- pkg/container/pSpool/sender.go | 39 ++++++ pkg/container/pSpool/sender_test.go | 97 ++++++++++++++ pkg/sql/colexec/connector/types.go | 55 +++++++- pkg/sql/colexec/connector/types_test.go | 124 +++++++++++++++++ pkg/sql/colexec/dispatch/dispatch_test.go | 125 ++++++++++++++++++ pkg/sql/colexec/dispatch/types.go | 54 +++++++- .../compile/allocation_account_lifecycle.go | 23 +++- .../allocation_account_lifecycle_test.go | 39 ++++++ .../remote_allocation_statement_group.go | 5 +- .../remote_allocation_statement_group_test.go | 11 ++ 10 files changed, 563 insertions(+), 9 deletions(-) diff --git a/pkg/container/pSpool/sender.go b/pkg/container/pSpool/sender.go index 20d1af7fb46bb..5d78c2001308f 100644 --- a/pkg/container/pSpool/sender.go +++ b/pkg/container/pSpool/sender.go @@ -251,6 +251,45 @@ func (ps *PipelineSpool) ForceCleanupAfterTerminalSignal() { ps.cleanupOnce.Do(ps.forceCleanup) } +// ReleaseReusableCacheAfterProducerQuiesced returns buffers from batches that +// receivers have already released. The producer must have stopped, so no new +// SendBatch call can race to reuse those buffers. Receivers may still return +// current batches afterward; FinalizeAfterConsumersQuiesced performs the final +// cache pass once those receivers have joined. +func (ps *PipelineSpool) ReleaseReusableCacheAfterProducerQuiesced() { + if ps == nil { + return + } + + ps.mu.Lock() + defer ps.mu.Unlock() + ps.cache.free() +} + +// FinalizeAfterConsumersQuiesced releases every batch still retained by the +// transport. The caller must have joined all producer and consumer scopes, so +// no goroutine can subsequently read a queued slot or use a current batch. +// +// A typed terminal signal bypasses the spool queue. Consequently a receiver +// may stop with older GetFromSpool signals still queued even though its scope +// has completed. Those slots are no longer observable and must be reclaimed +// here before the statement allocation account can reach terminal zero. +func (ps *PipelineSpool) FinalizeAfterConsumersQuiesced() { + if ps == nil { + return + } + + ps.mu.Lock() + defer ps.mu.Unlock() + for i := range ps.shardPool { + ps.cleanSlotLocked(uint32(i)) + } + for i := range ps.rs { + ps.rs[i].flagLastPopRelease() + } + ps.cleanupOnce.Do(ps.forceCleanup) +} + // Abort terminates the spool without waiting for receiver acknowledgement. // Pending, not-yet-consumed slots are released immediately. Slots already handed // to receivers stay valid until their receiver calls ReleaseCurrent. The first diff --git a/pkg/container/pSpool/sender_test.go b/pkg/container/pSpool/sender_test.go index ca05143aeed85..b1d43233843e5 100644 --- a/pkg/container/pSpool/sender_test.go +++ b/pkg/container/pSpool/sender_test.go @@ -303,6 +303,103 @@ func TestPipelineSpoolForceCleanupAfterTerminalSignalDoesNotNeedNilEndMessage(t require.Equal(t, int64(0), mp.CurrNB()) } +func TestPipelineSpoolReleaseReusableCacheBeforeFinalization(t *testing.T) { + mp := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(mp) + }) + srcMP := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(srcMP) + }) + src := newSpoolTestBatch(t, srcMP, 1024) + t.Cleanup(func() { + src.Clean(srcMP) + }) + + sp := InitMyPipelineSpool(mp, 2) + for range 2 { + done, err := sp.SendBatch(context.Background(), 0, src, nil) + require.NoError(t, err) + require.False(t, done) + } + got, info := sp.ReceiveBatch(0) + require.NoError(t, info) + require.NotNil(t, got) + sp.ReleaseCurrent(0) + got, info = sp.ReceiveBatch(0) + require.NoError(t, info) + require.NotNil(t, got) + beforeRelease := mp.CurrNB() + + sp.ReleaseReusableCacheAfterProducerQuiesced() + require.Positive(t, mp.CurrNB()) + require.Less(t, mp.CurrNB(), beforeRelease) + + sp.ReleaseCurrent(0) + require.Positive(t, mp.CurrNB()) + sp.FinalizeAfterConsumersQuiesced() + require.Zero(t, mp.CurrNB()) +} + +func TestPipelineSpoolFinalizeAfterConsumersQuiescedReleasesPendingBatch(t *testing.T) { + mp := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(mp) + }) + srcMP := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(srcMP) + }) + src := newSpoolTestBatch(t, srcMP, 1024) + t.Cleanup(func() { + src.Clean(srcMP) + }) + + sp := InitMyPipelineSpool(mp, 1) + done, err := sp.SendBatch(context.Background(), 0, src, nil) + require.NoError(t, err) + require.False(t, done) + require.Greater(t, mp.CurrNB(), int64(0)) + + sp.ForceCleanupAfterTerminalSignal() + require.Greater(t, mp.CurrNB(), int64(0)) + + sp.FinalizeAfterConsumersQuiesced() + require.Equal(t, int64(0), mp.CurrNB()) + sp.FinalizeAfterConsumersQuiesced() + require.Equal(t, int64(0), mp.CurrNB()) +} + +func TestPipelineSpoolFinalizeAfterConsumersQuiescedReleasesBroadcastCurrentAndPending(t *testing.T) { + mp := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(mp) + }) + srcMP := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(srcMP) + }) + src := newSpoolTestBatch(t, srcMP, 1024) + t.Cleanup(func() { + src.Clean(srcMP) + }) + + sp := InitMyPipelineSpool(mp, 2) + done, err := sp.SendBatch(context.Background(), SendToAllLocal, src, nil) + require.NoError(t, err) + require.False(t, done) + got, info := sp.ReceiveBatch(0) + require.NoError(t, info) + require.NotNil(t, got) + require.Greater(t, mp.CurrNB(), int64(0)) + + sp.FinalizeAfterConsumersQuiesced() + require.Equal(t, int64(0), mp.CurrNB()) + sp.FinalizeAfterConsumersQuiesced() + require.Equal(t, int64(0), mp.CurrNB()) +} + func TestPipelineSpoolLateReleaseAfterTerminalCleanupFreesDirectly(t *testing.T) { mp := mpool.MustNewZeroNoFixed() t.Cleanup(func() { diff --git a/pkg/sql/colexec/connector/types.go b/pkg/sql/colexec/connector/types.go index d7e02fb5af4ba..5169fafc3dec5 100644 --- a/pkg/sql/colexec/connector/types.go +++ b/pkg/sql/colexec/connector/types.go @@ -17,6 +17,7 @@ package connector import ( "context" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/reuse" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/pSpool" @@ -30,8 +31,9 @@ var _ vm.Operator = new(Connector) type Connector struct { ctr container - Reg *process.WaitRegister - cleanupSpool *pSpool.PipelineSpool + Reg *process.WaitRegister + cleanupSpool *pSpool.PipelineSpool + allocationAccount *mpool.AllocationAccount vm.OperatorBase } @@ -73,6 +75,45 @@ func (connector *Connector) WithReg(reg *process.WaitRegister) *Connector { return connector } +func (connector *Connector) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if account == nil || account.Handle() == 0 { + return mpool.ErrAllocationAccountInvalid + } + if connector.allocationAccount != nil && connector.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + connector.allocationAccount = account + return nil +} + +// ActivatesAllocationAccountLifecycle reports that Connector only participates +// in an account already required by an allocation-producing operator. +func (connector *Connector) ActivatesAllocationAccountLifecycle() bool { + return false +} + +func (connector *Connector) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if connector.allocationAccount == nil { + return nil + } + if connector.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if connector.ctr.sp != nil { + return mpool.ErrAllocationAccountInvariant + } + if connector.cleanupSpool != nil { + connector.cleanupSpool.FinalizeAfterConsumersQuiesced() + connector.cleanupSpool = nil + } + connector.allocationAccount = nil + return nil +} + func (connector *Connector) Release() { if connector != nil { reuse.Free[Connector](connector, nil) @@ -100,7 +141,11 @@ func (connector *Connector) Reset(proc *process.Process, pipelineFailed bool, er abortErr = fallbackErr } sp.Abort(abortErr) - connector.cleanupSpool = nil + if connector.allocationAccount != nil { + connector.cleanupSpool = sp + } else { + connector.cleanupSpool = nil + } } connector.ctr.sp = nil } else if terminalSignal.EventType == process.EventEnd && !terminalDelivered { @@ -146,6 +191,10 @@ func (connector *Connector) CleanupDeferredSpool() { if connector.cleanupSpool == nil { return } + if connector.allocationAccount != nil { + connector.cleanupSpool.ReleaseReusableCacheAfterProducerQuiesced() + return + } connector.cleanupSpool.ForceCleanupAfterTerminalSignal() connector.cleanupSpool = nil } diff --git a/pkg/sql/colexec/connector/types_test.go b/pkg/sql/colexec/connector/types_test.go index 34dd9b182ad94..824d38ccd8e95 100644 --- a/pkg/sql/colexec/connector/types_test.go +++ b/pkg/sql/colexec/connector/types_test.go @@ -378,6 +378,130 @@ func TestConnectorResetEndPreservesQueuedSpoolBatchUntilDeferredCleanup(t *testi require.Equal(t, int64(0), mp.CurrNB()) } +func TestConnectorAllocationClearFinalizesAbortedSpool(t *testing.T) { + testConnectorAllocationClearFinalizesSpool(t, true) +} + +func TestConnectorAccountedDeferredCleanupReleasesReusableCache(t *testing.T) { + mp := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(mp) + }) + srcMP := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(srcMP) + }) + src := newConnectorSpoolTestBatch(t, srcMP, 1024) + t.Cleanup(func() { + src.Clean(srcMP) + }) + + sp := pSpool.InitMyPipelineSpool(mp, 1) + done, err := sp.SendBatch(context.Background(), 0, src, nil) + require.NoError(t, err) + require.False(t, done) + got, info := sp.ReceiveBatch(0) + require.NoError(t, info) + require.NotNil(t, got) + sp.ReleaseCurrent(0) + require.Positive(t, mp.CurrNB()) + + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + conn := &Connector{cleanupSpool: sp} + require.NoError(t, conn.SetAllocationAccount(account)) + conn.CleanupDeferredSpool() + require.Same(t, sp, conn.cleanupSpool) + require.Zero(t, mp.CurrNB()) + require.NoError(t, conn.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestConnectorAllocationAccountContract(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(2, 1) + require.NoError(t, err) + first, err := registry.Open(1) + require.NoError(t, err) + second, err := registry.Open(1) + require.NoError(t, err) + conn := &Connector{} + require.False(t, conn.ActivatesAllocationAccountLifecycle()) + require.ErrorIs(t, conn.SetAllocationAccount(nil), mpool.ErrAllocationAccountInvalid) + require.NoError(t, conn.SetAllocationAccount(first)) + require.ErrorIs(t, conn.SetAllocationAccount(second), mpool.ErrAllocationAccountMismatch) + require.ErrorIs(t, conn.ClearAllocationAccount(second), mpool.ErrAllocationAccountMismatch) + conn.ctr.sp = &pSpool.PipelineSpool{} + require.ErrorIs(t, conn.ClearAllocationAccount(first), mpool.ErrAllocationAccountInvariant) + conn.ctr.sp = nil + require.NoError(t, conn.ClearAllocationAccount(first)) + require.NoError(t, conn.ClearAllocationAccount(first)) + _, _, err = registry.CompleteTerminal(first) + require.NoError(t, err) + _, _, err = registry.CompleteTerminal(second) + require.NoError(t, err) +} + +func TestConnectorAllocationClearFinalizesTerminalSpoolPending(t *testing.T) { + testConnectorAllocationClearFinalizesSpool(t, false) +} + +func testConnectorAllocationClearFinalizesSpool(t *testing.T, abort bool) { + mp := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(mp) + }) + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection( + account, + 1, + 102, + 103, + 104, + 105, + ) + require.NoError(t, err) + src := batch.NewOffHeapWithSize(1) + require.NoError(t, src.SetAllocationAccount(selection)) + src.SetVector(0, vector.NewOffHeapVecWithType(types.T_int64.ToType())) + require.NoError(t, vector.AppendFixed(src.Vecs[0], int64(1), false, mp)) + src.SetRowCount(1) + + sp := pSpool.InitMyPipelineSpool(mp, 1) + done, err := sp.SendBatch(context.Background(), 0, src, nil) + require.NoError(t, err) + require.False(t, done) + conn := &Connector{} + require.NoError(t, conn.SetAllocationAccount(account)) + if abort { + got, info := sp.ReceiveBatch(0) + require.NoError(t, info) + require.NotNil(t, got) + conn.ctr.sp = sp + conn.Reg = process.NewPipelineEdge(1, 0) + conn.Reset(nil, true, moerr.NewInternalErrorNoCtx("pipeline failed")) + require.Same(t, sp, conn.cleanupSpool) + sp.ReleaseCurrent(0) + } else { + conn.cleanupSpool = sp + sp.ForceCleanupAfterTerminalSignal() + } + conn.CleanupDeferredSpool() + require.Same(t, sp, conn.cleanupSpool) + require.NoError(t, conn.ClearAllocationAccount(account)) + require.Nil(t, conn.cleanupSpool) + src.Clean(mp) + snapshot := account.Seal() + require.Zero(t, snapshot.Used) + _, err = registry.Finalize(account) + require.NoError(t, err) +} + func newConnectorSpoolTestBatch(t *testing.T, mp *mpool.MPool, rows int) *batch.Batch { t.Helper() src := batch.NewWithSize(1) diff --git a/pkg/sql/colexec/dispatch/dispatch_test.go b/pkg/sql/colexec/dispatch/dispatch_test.go index 1b39a02f1ec73..ab3b3ea2068bc 100644 --- a/pkg/sql/colexec/dispatch/dispatch_test.go +++ b/pkg/sql/colexec/dispatch/dispatch_test.go @@ -897,6 +897,131 @@ func TestDispatchResetEndPreservesQueuedBroadcastBatchUntilDeferredCleanup(t *te require.Equal(t, int64(0), mp.CurrNB()) } +func TestDispatchAllocationClearFinalizesTerminalSpoolPending(t *testing.T) { + testDispatchAllocationClearFinalizesSpool(t, false) +} + +func TestDispatchAccountedDeferredCleanupReleasesReusableCache(t *testing.T) { + mp := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(mp) + }) + srcMP := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(srcMP) + }) + src := newDispatchSpoolTestBatch(t, srcMP, 1024) + t.Cleanup(func() { + src.Clean(srcMP) + }) + + sp := pSpool.InitMyPipelineSpool(mp, 1) + done, err := sp.SendBatch(context.Background(), 0, src, nil) + require.NoError(t, err) + require.False(t, done) + got, info := sp.ReceiveBatch(0) + require.NoError(t, info) + require.NotNil(t, got) + sp.ReleaseCurrent(0) + require.Positive(t, mp.CurrNB()) + + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + d := &Dispatch{cleanupSpool: sp} + require.NoError(t, d.SetAllocationAccount(account)) + d.CleanupDeferredSpool() + require.Same(t, sp, d.cleanupSpool) + require.Zero(t, mp.CurrNB()) + require.NoError(t, d.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestDispatchAllocationAccountContract(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(2, 1) + require.NoError(t, err) + first, err := registry.Open(1) + require.NoError(t, err) + second, err := registry.Open(1) + require.NoError(t, err) + d := &Dispatch{} + require.False(t, d.ActivatesAllocationAccountLifecycle()) + require.ErrorIs(t, d.SetAllocationAccount(nil), mpool.ErrAllocationAccountInvalid) + require.NoError(t, d.SetAllocationAccount(first)) + require.ErrorIs(t, d.SetAllocationAccount(second), mpool.ErrAllocationAccountMismatch) + require.ErrorIs(t, d.ClearAllocationAccount(second), mpool.ErrAllocationAccountMismatch) + d.ctr = &container{sp: &pSpool.PipelineSpool{}} + require.ErrorIs(t, d.ClearAllocationAccount(first), mpool.ErrAllocationAccountInvariant) + d.ctr = nil + require.NoError(t, d.ClearAllocationAccount(first)) + require.NoError(t, d.ClearAllocationAccount(first)) + _, _, err = registry.CompleteTerminal(first) + require.NoError(t, err) + _, _, err = registry.CompleteTerminal(second) + require.NoError(t, err) +} + +func TestDispatchAllocationClearFinalizesAbortedSpool(t *testing.T) { + testDispatchAllocationClearFinalizesSpool(t, true) +} + +func testDispatchAllocationClearFinalizesSpool(t *testing.T, abort bool) { + mp := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(mp) + }) + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection( + account, + 1, + 102, + 103, + 104, + 105, + ) + require.NoError(t, err) + src := batch.NewOffHeapWithSize(1) + require.NoError(t, src.SetAllocationAccount(selection)) + src.SetVector(0, vector.NewOffHeapVecWithType(types.T_int64.ToType())) + require.NoError(t, vector.AppendFixed(src.Vecs[0], int64(1), false, mp)) + src.SetRowCount(1) + + sp := pSpool.InitMyPipelineSpool(mp, 1) + done, err := sp.SendBatch(context.Background(), 0, src, nil) + require.NoError(t, err) + require.False(t, done) + + d := &Dispatch{} + require.NoError(t, d.SetAllocationAccount(account)) + if abort { + got, info := sp.ReceiveBatch(0) + require.NoError(t, info) + require.NotNil(t, got) + d.ctr = &container{sp: sp} + d.Reset(nil, true, moerr.NewInternalErrorNoCtx("pipeline failed")) + require.Same(t, sp, d.cleanupSpool) + sp.ReleaseCurrent(0) + } else { + d.cleanupSpool = sp + sp.ForceCleanupAfterTerminalSignal() + } + d.CleanupDeferredSpool() + require.Same(t, sp, d.cleanupSpool) + require.NoError(t, d.ClearAllocationAccount(account)) + require.Nil(t, d.cleanupSpool) + + src.Clean(mp) + snapshot := account.Seal() + require.Zero(t, snapshot.Used) + _, err = registry.Finalize(account) + require.NoError(t, err) +} + // TestReceiverDone_OldBehavior tests the old behavior (kept for backward compatibility verification) func TestReceiverDone_OldBehavior(t *testing.T) { proc := testutil.NewProcess(t) diff --git a/pkg/sql/colexec/dispatch/types.go b/pkg/sql/colexec/dispatch/types.go index 3b73c149ea50c..f9e8b76b3312d 100644 --- a/pkg/sql/colexec/dispatch/types.go +++ b/pkg/sql/colexec/dispatch/types.go @@ -76,8 +76,9 @@ type container struct { } type Dispatch struct { - ctr *container - cleanupSpool *pSpool.PipelineSpool + ctr *container + cleanupSpool *pSpool.PipelineSpool + allocationAccount *mpool.AllocationAccount // MaterializedSource is used by a multi-reference CTE whose consumers can // have execution dependencies on one another. It is local-only and bypasses @@ -109,6 +110,45 @@ func (dispatch *Dispatch) GetOperatorBase() *vm.OperatorBase { return &dispatch.OperatorBase } +func (dispatch *Dispatch) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if account == nil || account.Handle() == 0 { + return mpool.ErrAllocationAccountInvalid + } + if dispatch.allocationAccount != nil && dispatch.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + dispatch.allocationAccount = account + return nil +} + +// ActivatesAllocationAccountLifecycle reports that Dispatch only participates +// in an account already required by an allocation-producing operator. +func (dispatch *Dispatch) ActivatesAllocationAccountLifecycle() bool { + return false +} + +func (dispatch *Dispatch) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if dispatch.allocationAccount == nil { + return nil + } + if dispatch.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if dispatch.ctr != nil && dispatch.ctr.sp != nil { + return mpool.ErrAllocationAccountInvariant + } + if dispatch.cleanupSpool != nil { + dispatch.cleanupSpool.FinalizeAfterConsumersQuiesced() + dispatch.cleanupSpool = nil + } + dispatch.allocationAccount = nil + return nil +} + func init() { reuse.CreatePool[Dispatch]( func() *Dispatch { @@ -286,7 +326,11 @@ func (dispatch *Dispatch) Reset(proc *process.Process, pipelineFailed bool, err abortErr = fallbackErr } sp.Abort(abortErr) - dispatch.cleanupSpool = nil + if dispatch.allocationAccount != nil { + dispatch.cleanupSpool = sp + } else { + dispatch.cleanupSpool = nil + } } dispatch.ctr.sp = nil } else { @@ -308,6 +352,10 @@ func (dispatch *Dispatch) CleanupDeferredSpool() { if dispatch.cleanupSpool == nil { return } + if dispatch.allocationAccount != nil { + dispatch.cleanupSpool.ReleaseReusableCacheAfterProducerQuiesced() + return + } dispatch.cleanupSpool.ForceCleanupAfterTerminalSignal() dispatch.cleanupSpool = nil } diff --git a/pkg/sql/compile/allocation_account_lifecycle.go b/pkg/sql/compile/allocation_account_lifecycle.go index 3638fad8e9341..b8bc959b53e51 100644 --- a/pkg/sql/compile/allocation_account_lifecycle.go +++ b/pkg/sql/compile/allocation_account_lifecycle.go @@ -45,6 +45,22 @@ type executionAllocationAccountOwner interface { ClearAllocationAccount(*mpool.AllocationAccount) error } +type executionAllocationAccountActivationPolicy interface { + ActivatesAllocationAccountLifecycle() bool +} + +func hasAllocationAccountActivator( + owners []executionAllocationAccountOwner, +) bool { + for _, owner := range owners { + policy, ok := owner.(executionAllocationAccountActivationPolicy) + if !ok || policy.ActivatesAllocationAccountLifecycle() { + return true + } + } + return false +} + // statementAllocationAttempt owns one local execution generation. The // MessageBoard pointer is captured at open so prepared/retry Reset cannot make // terminal cleanup drain a newer board. @@ -182,9 +198,12 @@ func (c *Compile) attachRuntimeAllocationOwners(scopes []*Scope) error { } if c.allocationAttempt == nil { owners, err := collectAllocationAccountOwners(scopes) - if err != nil || len(owners) == 0 { + if err != nil { return err } + if !hasAllocationAccountActivator(owners) { + return nil + } return mpool.ErrAllocationAccountInvariant } return c.allocationAttempt.attachRuntimeOwners(scopes) @@ -267,7 +286,7 @@ func (c *Compile) ensureAllocationAccountLifecycle( return err } c.allocationAccountOwners = owners - if len(owners) == 0 { + if !hasAllocationAccountActivator(owners) { c.allocationAccountOwners = nil if c.allocationControllerProvider != nil { c.allocationAccountRegistry = nil diff --git a/pkg/sql/compile/allocation_account_lifecycle_test.go b/pkg/sql/compile/allocation_account_lifecycle_test.go index 5b6462bf83e5c..ceacda7824dfa 100644 --- a/pkg/sql/compile/allocation_account_lifecycle_test.go +++ b/pkg/sql/compile/allocation_account_lifecycle_test.go @@ -28,6 +28,8 @@ import ( "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/pb/txn" "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/connector" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/dispatch" "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" "github.com/matrixorigin/matrixone/pkg/sql/colexec/product" "github.com/matrixorigin/matrixone/pkg/testutil" @@ -429,6 +431,43 @@ func TestAllocationAccountConfiguresEveryStatementOwner(t *testing.T) { require.NoError(t, err) } +func TestAllocationAccountTransportParticipantsDoNotActivateLifecycle(t *testing.T) { + connectorOp := connector.NewArgument() + dispatchOp := dispatch.NewArgument() + t.Cleanup(connectorOp.Release) + t.Cleanup(dispatchOp.Release) + c := &Compile{ + proc: testutil.NewProcess(t), + MessageBoard: message.NewMessageBoard(), + scopes: []*Scope{ + {RootOp: connectorOp}, + {RootOp: dispatchOp}, + }, + } + require.NoError(t, c.ensureAllocationAccountLifecycle(func( + mpool.AllocationAccountTerminalSnapshot, + ) { + })) + require.Nil(t, c.allocationAccountOwners) + require.Nil(t, c.allocationControllerProvider) + require.Nil(t, c.allocationAccountRegistry) + require.NoError(t, c.attachRuntimeAllocationOwners(c.scopes)) + + transportOwners := []executionAllocationAccountOwner{ + connectorOp, + dispatchOp, + } + require.False(t, hasAllocationAccountActivator(transportOwners)) + + active := &allocationLifecycleOwnerOperator{ + MockOperator: colexec.NewMockOperator(), + } + require.True(t, hasAllocationAccountActivator(append(transportOwners, active))) + require.ErrorIs(t, c.attachRuntimeAllocationOwners([]*Scope{{ + RootOp: active, + }}), mpool.ErrAllocationAccountInvariant) +} + func TestAllocationAccountCollectsProductConsumerAndHashBuild(t *testing.T) { consumer := product.NewArgument() producer := hashbuild.NewArgument() diff --git a/pkg/sql/compile/remote_allocation_statement_group.go b/pkg/sql/compile/remote_allocation_statement_group.go index dd024d9b70663..6d58d1a7296ab 100644 --- a/pkg/sql/compile/remote_allocation_statement_group.go +++ b/pkg/sql/compile/remote_allocation_statement_group.go @@ -97,9 +97,12 @@ func validateRemoteAllocationTopologyCapability( return nil } owners, err := collectAllocationAccountOwners(scopes) - if err != nil || len(owners) == 0 { + if err != nil { return err } + if !hasAllocationAccountActivator(owners) { + return nil + } return moerr.NewNotSupportedNoCtx( "remote allocation-accounted execution requires fragment topology metadata", ) diff --git a/pkg/sql/compile/remote_allocation_statement_group_test.go b/pkg/sql/compile/remote_allocation_statement_group_test.go index f6c89a72ed00d..fe39c0ceb23b0 100644 --- a/pkg/sql/compile/remote_allocation_statement_group_test.go +++ b/pkg/sql/compile/remote_allocation_statement_group_test.go @@ -23,6 +23,8 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/connector" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/dispatch" "github.com/matrixorigin/matrixone/pkg/vm/engine" "github.com/matrixorigin/matrixone/pkg/vm/message" "github.com/stretchr/testify/require" @@ -124,6 +126,15 @@ func TestRemoteAllocationTopologyCapabilityIsRequiredForOwners(t *testing.T) { scopes, map[string]uint32{"cn-a:6001": 1}, )) + + connectorOp := connector.NewArgument() + dispatchOp := dispatch.NewArgument() + t.Cleanup(connectorOp.Release) + t.Cleanup(dispatchOp.Release) + require.NoError(t, validateRemoteAllocationTopologyCapability([]*Scope{ + {RootOp: connectorOp}, + {RootOp: dispatchOp}, + }, nil)) } func TestRemoteAllocationStatementGroupDefersSharedBoardTerminal(t *testing.T) { From 60d65446582bb79419241fb697e5096d182e2215 Mon Sep 17 00:00:00 2001 From: aptend Date: Sun, 2 Aug 2026 10:29:14 +0800 Subject: [PATCH 48/61] fix: release accounted prepared result buffers --- pkg/container/batch/batch.go | 6 ++++ pkg/sql/colexec/limit/limit_test.go | 44 ++++++++++++++++++++++++ pkg/sql/colexec/limit/types.go | 6 ++++ pkg/sql/colexec/mergeorder/order_test.go | 29 ++++++++++++++++ pkg/sql/colexec/mergeorder/types.go | 9 ++++- pkg/sql/colexec/offset/offset_test.go | 41 ++++++++++++++++++++++ pkg/sql/colexec/offset/types.go | 9 ++++- pkg/sql/plan/apply_indices_test.go | 35 +++++++++++++++++-- 8 files changed, 175 insertions(+), 4 deletions(-) diff --git a/pkg/container/batch/batch.go b/pkg/container/batch/batch.go index d766f404f3f49..f07dab65ae4e6 100644 --- a/pkg/container/batch/batch.go +++ b/pkg/container/batch/batch.go @@ -1191,6 +1191,12 @@ func (bat *Batch) hasAllocationAccountVector() bool { return false } +// HasAllocationAccount reports whether the batch or one of its vectors owns +// memory charged to an execution-scoped allocation account. +func (bat *Batch) HasAllocationAccount() bool { + return bat != nil && (bat.allocationAccount != nil || bat.hasAllocationAccountVector()) +} + func (bat *Batch) selectedColumnsHaveAllocationAccount(selectCols []int) bool { for _, sourceIdx := range selectCols { vec := bat.Vecs[sourceIdx] diff --git a/pkg/sql/colexec/limit/limit_test.go b/pkg/sql/colexec/limit/limit_test.go index 25cb7316e6355..36d1e439ddfe4 100644 --- a/pkg/sql/colexec/limit/limit_test.go +++ b/pkg/sql/colexec/limit/limit_test.go @@ -22,6 +22,8 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/sql/colexec" plan2 "github.com/matrixorigin/matrixone/pkg/sql/plan" "github.com/matrixorigin/matrixone/pkg/testutil" @@ -184,6 +186,48 @@ func TestLimitDoesNotMutateInputBatch(t *testing.T) { require.Zero(t, proc.Mp().CurrNB()) } +func TestLimitResetReleasesCopiedAllocationAccountData(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + + input := batch.NewOffHeapWithSize(1) + input.SetVector(0, vector.NewOffHeapVecWithType(types.T_int64.ToType())) + require.NoError(t, input.SetAllocationAccount(selection)) + for i := range 32 { + require.NoError(t, vector.AppendFixed(input.Vecs[0], int64(i), false, proc.Mp())) + } + input.SetRowCount(32) + + arg := NewArgument().WithLimit(plan2.MakePlan2Uint64ConstExprWithType(1)) + child := colexec.NewMockOperator().WithBatchs([]*batch.Batch{input}) + arg.AppendChild(child) + require.NoError(t, arg.Prepare(proc)) + result, err := arg.Call(proc) + require.NoError(t, err) + require.Equal(t, 1, result.Batch.RowCount()) + + // Pipeline cleanup resets children before parents. Simulate HashJoin + // releasing its result batch, then verify Limit releases its accounted copy. + input.Clean(proc.Mp()) + require.Positive(t, account.Snapshot().Used) + arg.Reset(proc, false, nil) + require.Nil(t, arg.ctr.buf) + require.Zero(t, account.Snapshot().Used) + + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + arg.Free(proc, false, nil) + child.Free(proc, false, nil) + arg.Release() + proc.Free() + require.Zero(t, proc.Mp().CurrNB()) +} + func BenchmarkLimit(b *testing.B) { for i := 0; i < b.N; i++ { tcs := []limitTestCase{ diff --git a/pkg/sql/colexec/limit/types.go b/pkg/sql/colexec/limit/types.go index 04e143cf327c7..a6b8f75441783 100644 --- a/pkg/sql/colexec/limit/types.go +++ b/pkg/sql/colexec/limit/types.go @@ -78,6 +78,12 @@ func (limit *Limit) Reset(proc *process.Process, pipelineFailed bool, err error) if limit.ctr.limitExecutor != nil { limit.ctr.limitExecutor.ResetForNextQuery() } + if limit.ctr.buf.HasAllocationAccount() { + // Prepared operators may reuse ordinary buffers across executions, but an + // accounted buffer belongs to exactly one execution generation. + limit.ctr.buf.Clean(proc.Mp()) + limit.ctr.buf = nil + } limit.ctr.seen = 0 } diff --git a/pkg/sql/colexec/mergeorder/order_test.go b/pkg/sql/colexec/mergeorder/order_test.go index 635b364a85f50..d36ea0241a1c9 100644 --- a/pkg/sql/colexec/mergeorder/order_test.go +++ b/pkg/sql/colexec/mergeorder/order_test.go @@ -1080,6 +1080,35 @@ func TestMergeOrderResetAndOpType(t *testing.T) { arg.Free(proc, false, nil) } +func TestMergeOrderResetReleasesAccountedResult(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + + result := batch.NewOffHeapWithSize(1) + result.Vecs[0] = vector.NewOffHeapVecWithType(types.T_int64.ToType()) + require.NoError(t, result.SetAllocationAccount(selection)) + require.NoError(t, vector.AppendFixed(result.Vecs[0], int64(1), false, proc.Mp())) + result.SetRowCount(1) + require.Positive(t, account.Snapshot().Used) + + arg := &MergeOrder{} + arg.ctr.buf = result + arg.Reset(proc, false, nil) + require.Nil(t, arg.ctr.buf) + require.Zero(t, account.Snapshot().Used) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + + arg.Free(proc, false, nil) + proc.Free() + require.Zero(t, proc.Mp().CurrNB()) +} + func TestSpillHelperBranches(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer func() { diff --git a/pkg/sql/colexec/mergeorder/types.go b/pkg/sql/colexec/mergeorder/types.go index e9e8f6caac2e1..20e798083b77d 100644 --- a/pkg/sql/colexec/mergeorder/types.go +++ b/pkg/sql/colexec/mergeorder/types.go @@ -174,7 +174,14 @@ func (mergeOrder *MergeOrder) Reset(proc *process.Process, pipelineFailed bool, } } if ctr.buf != nil { - ctr.buf.CleanOnlyData() + if ctr.buf.HasAllocationAccount() { + // The final merge batch may directly own a child execution's result. + // Accounted storage cannot survive that execution's Reset boundary. + ctr.buf.Clean(proc.Mp()) + ctr.buf = nil + } else { + ctr.buf.CleanOnlyData() + } } } diff --git a/pkg/sql/colexec/offset/offset_test.go b/pkg/sql/colexec/offset/offset_test.go index 4c13d023f9d7a..2df732ac35c92 100644 --- a/pkg/sql/colexec/offset/offset_test.go +++ b/pkg/sql/colexec/offset/offset_test.go @@ -23,6 +23,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/sql/colexec" plan2 "github.com/matrixorigin/matrixone/pkg/sql/plan" "github.com/matrixorigin/matrixone/pkg/testutil" @@ -140,6 +141,46 @@ func TestOffset(t *testing.T) { } } +func TestOffsetResetReleasesCopiedAllocationAccountData(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + + input := batch.NewOffHeapWithSize(1) + input.SetVector(0, vector.NewOffHeapVecWithType(types.T_int64.ToType())) + require.NoError(t, input.SetAllocationAccount(selection)) + for i := range 32 { + require.NoError(t, vector.AppendFixed(input.Vecs[0], int64(i), false, proc.Mp())) + } + input.SetRowCount(32) + + arg := NewArgument().WithOffset(plan2.MakePlan2Uint64ConstExprWithType(1)) + child := colexec.NewMockOperator().WithBatchs([]*batch.Batch{input}) + arg.AppendChild(child) + require.NoError(t, arg.Prepare(proc)) + result, err := arg.Call(proc) + require.NoError(t, err) + require.Equal(t, 31, result.Batch.RowCount()) + + input.Clean(proc.Mp()) + require.Positive(t, account.Snapshot().Used) + arg.Reset(proc, false, nil) + require.Nil(t, arg.ctr.buf) + require.Zero(t, account.Snapshot().Used) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + + arg.Free(proc, false, nil) + child.Free(proc, false, nil) + arg.Release() + proc.Free() + require.Zero(t, proc.Mp().CurrNB()) +} + func BenchmarkOffset(b *testing.B) { for i := 0; i < b.N; i++ { tcs := []offsetTestCase{ diff --git a/pkg/sql/colexec/offset/types.go b/pkg/sql/colexec/offset/types.go index c6f3ef638aafc..bd52640ce925f 100644 --- a/pkg/sql/colexec/offset/types.go +++ b/pkg/sql/colexec/offset/types.go @@ -79,7 +79,14 @@ func (offset *Offset) Reset(proc *process.Process, pipelineFailed bool, err erro offset.ctr.offsetExecutor.ResetForNextQuery() } if offset.ctr.buf != nil { - offset.ctr.buf.CleanOnlyData() + if offset.ctr.buf.HasAllocationAccount() { + // Do not carry an execution-scoped allocation selection into the next + // prepared execution. + offset.ctr.buf.Clean(proc.Mp()) + offset.ctr.buf = nil + } else { + offset.ctr.buf.CleanOnlyData() + } } offset.ctr.seen = 0 } diff --git a/pkg/sql/plan/apply_indices_test.go b/pkg/sql/plan/apply_indices_test.go index 1cc81da1e0b21..4d5b4775caf82 100644 --- a/pkg/sql/plan/apply_indices_test.go +++ b/pkg/sql/plan/apply_indices_test.go @@ -752,6 +752,7 @@ func TestUniqueIndexRuntimeFilterUsesSelectedHashSlot(t *testing.T) { } arg.RuntimeFilterSpec = spec arg.AppendChild(child) + registry, account := installIndexPlanHashBuildAllocation(t, arg) require.NoError(t, child.Prepare(proc)) require.NoError(t, arg.Prepare(proc)) result, err := vm.Exec(arg, proc) @@ -778,10 +779,11 @@ func TestUniqueIndexRuntimeFilterUsesSelectedHashSlot(t *testing.T) { payload.Free(proc.Mp()) runtimeFilter.Destroy() arg.Free(proc, false, nil) + proc.GetMessageBoard().Reset() + finishIndexPlanHashBuildAllocation(t, registry, account, arg) child.Free(proc, false, nil) arg.Release() child.Release() - proc.GetMessageBoard().Reset() proc.Free() require.Zero(t, proc.Mp().CurrNB()) } @@ -858,6 +860,7 @@ func TestIndexJoinGeneratedSerializedRuntimeFilterExecutesEndToEnd(t *testing.T) planpb.Type{Id: int32(types.T_int32)}, 0, 0)} arg.RuntimeFilterSpec = spec arg.AppendChild(child) + registry, account := installIndexPlanHashBuildAllocation(t, arg) require.NoError(t, child.Prepare(proc)) require.NoError(t, arg.Prepare(proc)) result, err := vm.Exec(arg, proc) @@ -916,14 +919,42 @@ func TestIndexJoinGeneratedSerializedRuntimeFilterExecutesEndToEnd(t *testing.T) payload.Free(proc.Mp()) runtimeFilter.Destroy() arg.Free(proc, false, nil) + proc.GetMessageBoard().Reset() + finishIndexPlanHashBuildAllocation(t, registry, account, arg) child.Free(proc, false, nil) arg.Release() child.Release() - proc.GetMessageBoard().Reset() proc.Free() require.Zero(t, proc.Mp().CurrNB()) } +func installIndexPlanHashBuildAllocation( + t testing.TB, + arg *hashbuild.HashBuild, +) (*mpool.AllocationAccountRegistry, *mpool.AllocationAccount) { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.Open(1 << 60) + require.NoError(t, err) + require.NoError(t, arg.SetAllocationAccount(account)) + return registry, account +} + +func finishIndexPlanHashBuildAllocation( + t testing.TB, + registry *mpool.AllocationAccountRegistry, + account *mpool.AllocationAccount, + arg *hashbuild.HashBuild, +) { + t.Helper() + require.NoError(t, arg.ClearAllocationAccount(account)) + snapshot, first, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.True(t, first) + require.Zero(t, snapshot.Used) +} + func TestForceIndexForJoinBuildsRightAccessWithoutReorder(t *testing.T) { builder, joinID, leftScanID, _ := makeIndexHintJoinBuilder(t) joinNode := builder.qry.Nodes[joinID] From 74923c49fe1e2ae2c801fda69c08dd393eae21c9 Mon Sep 17 00:00:00 2001 From: aptend Date: Sun, 2 Aug 2026 12:16:04 +0800 Subject: [PATCH 49/61] fix: release consumed merge-order batches --- pkg/sql/colexec/mergeorder/order.go | 8 ++--- pkg/sql/colexec/mergeorder/order_test.go | 39 ++++++++++++++++++++++++ pkg/sql/colexec/mergeorder/types.go | 8 ++--- pkg/sql/colexec/order/order_test.go | 38 +++++++++++++++++++++++ pkg/sql/colexec/order/types.go | 6 +++- 5 files changed, 87 insertions(+), 12 deletions(-) diff --git a/pkg/sql/colexec/mergeorder/order.go b/pkg/sql/colexec/mergeorder/order.go index 9904728afc619..47a4653bc8c3c 100644 --- a/pkg/sql/colexec/mergeorder/order.go +++ b/pkg/sql/colexec/mergeorder/order.go @@ -278,12 +278,8 @@ func (ctr *container) removeInMemoryBatch(proc *process.Process, index int) erro if ctr.inMemoryHeap != nil { heap.Remove(ctr.inMemoryHeap, ctr.inMemoryHeapPos[index]) } - for i := range cols { - if batchContainsVector(bat, cols[i]) { - continue - } - cols[i].Free(proc.GetMPool()) - } + freeOrderColumns(proc.GetMPool(), bat, cols) + bat.Clean(proc.GetMPool()) ctr.batchList[index] = nil ctr.orderCols[index] = nil ctr.indexList[index] = -1 diff --git a/pkg/sql/colexec/mergeorder/order_test.go b/pkg/sql/colexec/mergeorder/order_test.go index d36ea0241a1c9..50549da223a1f 100644 --- a/pkg/sql/colexec/mergeorder/order_test.go +++ b/pkg/sql/colexec/mergeorder/order_test.go @@ -1109,6 +1109,45 @@ func TestMergeOrderResetReleasesAccountedResult(t *testing.T) { require.Zero(t, proc.Mp().CurrNB()) } +func TestRemoveInMemoryBatchReleasesAccountedBatch(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + + input := batch.NewOffHeapWithSize(2) + for i := range input.Vecs { + input.Vecs[i] = vector.NewOffHeapVecWithType(types.T_int64.ToType()) + } + require.NoError(t, input.SetAllocationAccount(selection)) + for i := range 17 { + require.NoError(t, vector.AppendFixed(input.Vecs[0], int64(i), false, proc.Mp())) + require.NoError(t, vector.AppendFixed(input.Vecs[1], int64(i), false, proc.Mp())) + } + input.SetRowCount(17) + require.Positive(t, account.Snapshot().Used) + + ctr := container{ + batchList: []*batch.Batch{input}, + orderCols: [][]*vector.Vector{{input.Vecs[0]}}, + indexList: []int64{17}, + spillMemUsage: int64(input.Size()), + } + require.NoError(t, ctr.removeInMemoryBatch(proc, 0)) + require.Nil(t, ctr.batchList[0]) + require.Nil(t, ctr.orderCols[0]) + require.Zero(t, ctr.spillMemUsage) + require.Zero(t, account.Snapshot().Used) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + + proc.Free() + require.Zero(t, proc.Mp().CurrNB()) +} + func TestSpillHelperBranches(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer func() { diff --git a/pkg/sql/colexec/mergeorder/types.go b/pkg/sql/colexec/mergeorder/types.go index 20e798083b77d..3f944e95a4790 100644 --- a/pkg/sql/colexec/mergeorder/types.go +++ b/pkg/sql/colexec/mergeorder/types.go @@ -215,15 +215,13 @@ func (mergeOrder *MergeOrder) cleanBatchAndCol(proc *process.Process) { mp := proc.Mp() ctr := &mergeOrder.ctr for i := range ctr.batchList { + if ctr.batchList[i] != nil && i < len(ctr.orderCols) && ctr.orderCols[i] != nil { + freeOrderColumns(mp, ctr.batchList[i], ctr.orderCols[i]) + } if ctr.batchList[i] != nil { ctr.batchList[i].Clean(mp) } } - for i := range ctr.orderCols { - if ctr.orderCols[i] != nil { - freeOrderColumns(mp, ctr.batchList[i], ctr.orderCols[i]) - } - } } func (ctr *container) cleanupSpill(proc *process.Process) { diff --git a/pkg/sql/colexec/order/order_test.go b/pkg/sql/colexec/order/order_test.go index 66bdc0fbb8fad..4fe2e60929321 100644 --- a/pkg/sql/colexec/order/order_test.go +++ b/pkg/sql/colexec/order/order_test.go @@ -21,6 +21,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" "github.com/matrixorigin/matrixone/pkg/testutil" @@ -94,6 +95,43 @@ func TestOrder(t *testing.T) { } } +func TestOrderResetReleasesPartiallyAccumulatedAccountedBatch(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + + input := batch.NewOffHeapWithSize(2) + for i := range input.Vecs { + input.Vecs[i] = vector.NewOffHeapVecWithType(types.T_int64.ToType()) + } + require.NoError(t, input.SetAllocationAccount(selection)) + for i := range 64 { + require.NoError(t, vector.AppendFixed(input.Vecs[0], int64(i), false, proc.Mp())) + require.NoError(t, vector.AppendFixed(input.Vecs[1], int64(i), false, proc.Mp())) + } + input.SetRowCount(64) + + arg := &Order{} + _, err = arg.ctr.appendBatch(proc, input) + require.NoError(t, err) + input.Clean(proc.Mp()) + require.Positive(t, account.Snapshot().Used) + + arg.Reset(proc, true, nil) + require.Nil(t, arg.ctr.batWaitForSort) + require.Zero(t, account.Snapshot().Used) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + + arg.Free(proc, true, nil) + proc.Free() + require.Zero(t, proc.Mp().CurrNB()) +} + func BenchmarkOrder(b *testing.B) { for i := 0; i < b.N; i++ { tcs := []orderTestCase{ diff --git a/pkg/sql/colexec/order/types.go b/pkg/sql/colexec/order/types.go index 6055dab28a94c..dbab07cd34823 100644 --- a/pkg/sql/colexec/order/types.go +++ b/pkg/sql/colexec/order/types.go @@ -84,7 +84,11 @@ type container struct { func (order *Order) Reset(proc *process.Process, pipelineFailed bool, err error) { ctr := &order.ctr if ctr.batWaitForSort != nil { - if ctr.batWaitForSort.RowCount() > colexec.DefaultBatchSize { + if ctr.batWaitForSort.HasAllocationAccount() || + ctr.batWaitForSort.RowCount() > colexec.DefaultBatchSize { + // A partially accumulated sort batch can survive when an upstream + // pipeline is stopped before sortAndSend transfers it to rbat. Its + // allocation account belongs to the completed execution generation. ctr.batWaitForSort.Clean(proc.Mp()) ctr.batWaitForSort = nil } else { From 4cdd43c13c95fe8b6de03a8762ffce24d682eb27 Mon Sep 17 00:00:00 2001 From: aptend Date: Sun, 2 Aug 2026 15:43:01 +0800 Subject: [PATCH 50/61] perf: remove accounted hash-build copy amplification --- pkg/common/malloc/profiler.go | 19 ++++ pkg/common/mpool/allocation_account.go | 9 +- pkg/common/mpool/mpool.go | 21 +++- pkg/common/mpool/mpool_profile.go | 50 +++++++- pkg/common/mpool/mpool_profile_test.go | 82 ++++++++++++++ pkg/container/vector/vector.go | 34 ++++++ pkg/container/vector/vector_test.go | 28 +++++ pkg/sql/colexec/join_util.go | 91 +++++++++------ pkg/sql/colexec/join_util_test.go | 151 +++++++++++++++++++++++++ 9 files changed, 436 insertions(+), 49 deletions(-) diff --git a/pkg/common/malloc/profiler.go b/pkg/common/malloc/profiler.go index 148d599ee35b4..fad3f9d5123c7 100644 --- a/pkg/common/malloc/profiler.go +++ b/pkg/common/malloc/profiler.go @@ -103,6 +103,25 @@ func (p *Profiler[T, P]) Sample( return p.getSampleValueFromPCs(pcs, int64(fullStackFraction)) } +// SampleNamed returns a stable synthetic sample without collecting a runtime +// stack. It is intended for allocations that already carry explicit, bounded +// provenance supplied by their owner. +func (p *Profiler[T, P]) SampleNamed(name string) P { + locations := []*profile.Location{p.getMockLocation(name)} + locationsHashSum := hashLocations(locations) + if v, ok := p.locationsToSample.Load(locationsHashSum); ok { + return v.(*SampleInfo[P]).Values + } + + var value T + P(&value).Init() + v, _ := p.locationsToSample.LoadOrStore(locationsHashSum, &SampleInfo[P]{ + Values: &value, + Locations: locations, + }) + return v.(*SampleInfo[P]).Values +} + func (p *Profiler[T, P]) getLocation(frame runtime.Frame) *profile.Location { locationKey := LocationKey{ File: frame.File, diff --git a/pkg/common/mpool/allocation_account.go b/pkg/common/mpool/allocation_account.go index b94231526ab44..ecfcfdc8aee63 100644 --- a/pkg/common/mpool/allocation_account.go +++ b/pkg/common/mpool/allocation_account.go @@ -809,10 +809,11 @@ func (r allocationAccountRequest) validate() error { } type allocationLease struct { - account *AllocationAccount - owner AllocationOwner - site AllocationSite - _ [6]byte + account *AllocationAccount + owner AllocationOwner + site AllocationSite + profiled bool + _ [5]byte } func (l allocationLease) release(capacity uint64) { diff --git a/pkg/common/mpool/mpool.go b/pkg/common/mpool/mpool.go index 14aaf3896ef3d..f65df78c28bb2 100644 --- a/pkg/common/mpool/mpool.go +++ b/pkg/common/mpool/mpool.go @@ -954,9 +954,10 @@ func (mp *MPool) allocAccounted( } hdr.SetGuard() lease := allocationLease{ - account: request.account, - owner: request.owner, - site: request.site, + account: request.account, + owner: request.owner, + site: request.site, + profiled: ProfilingEnabled(), } gcurr := globalStats.RecordAlloc("global", sz) @@ -991,7 +992,7 @@ func (mp *MPool) allocAccounted( if mp.details != nil { mp.details.recordAlloc(detailk, sz) } - profileRecordAlloc(3, uintptr(ptr), sz) + profileRecordAccountedAlloc(lease, sz) return bs, nil } @@ -1030,7 +1031,11 @@ func (mp *MPool) freePtr(detailk string, ptr unsafe.Pointer) { // consistent with freePtrInternal. if hdr.isOffHeap() { sz := int64(hdr.allocSz) - profileRecordFree(uintptr(ptr), sz) + if hdr.isAccounted() { + profileRecordAccountedFree(lease, sz) + } else { + profileRecordFree(uintptr(ptr), sz) + } globalStats.RecordFree("global", sz) simpleCAllocator().Deallocate(unsafe.Slice((*byte)(ptr), sz), uint64(sz)) if hdr.isAccounted() { @@ -1069,7 +1074,11 @@ func (mp *MPool) freePtrInternal( bs[i] = 0xDD } } - profileRecordFree(uintptr(ptr), sz) + if hdr.isAccounted() { + profileRecordAccountedFree(lease, sz) + } else { + profileRecordFree(uintptr(ptr), sz) + } mp.stats.RecordFree(mp.tag, sz) globalStats.RecordFree("global", sz) mp.resource.recordFree(sz) diff --git a/pkg/common/mpool/mpool_profile.go b/pkg/common/mpool/mpool_profile.go index 8736697c458fc..fd2f93a2fa7d7 100644 --- a/pkg/common/mpool/mpool_profile.go +++ b/pkg/common/mpool/mpool_profile.go @@ -15,6 +15,7 @@ package mpool import ( + "fmt" "sync" "sync/atomic" @@ -23,8 +24,9 @@ import ( var profilingEnabled atomic.Bool -// EnableProfiling turns on per-allocation stack tracking for off-heap mpool -// allocations. Tracked allocations appear in the malloc profiler output. +// EnableProfiling turns on tracking for off-heap mpool allocations. Ordinary +// allocations are grouped by sampled call stack; accounted allocations are +// grouped by their explicit owner/site provenance. func EnableProfiling() { profilingEnabled.Store(true) } // DisableProfiling turns off per-allocation stack tracking. @@ -43,6 +45,11 @@ type profileShard struct { var globalProfileShards [numProfileShards]profileShard +// Accounted allocations already carry stable, bounded provenance. Reusing one +// synthetic sample per owner/site avoids collecting and hashing the same +// runtime stack for every vector growth in a hash build. +var accountedProfileSamples [AllocationOwnerMax + 1][256]atomic.Pointer[malloc.HeapSampleValues] + func init() { for i := range globalProfileShards { globalProfileShards[i].m = make(map[uintptr]*malloc.HeapSampleValues, 64) @@ -98,6 +105,45 @@ func profileRecordFree(ptr uintptr, sz int64) { } } +func accountedProfileSample( + owner AllocationOwner, + site AllocationSite, +) *malloc.HeapSampleValues { + slot := &accountedProfileSamples[owner][site] + if values := slot.Load(); values != nil { + return values + } + values := malloc.GlobalProfiler().SampleNamed(fmt.Sprintf( + "| mpool accounted owner=%d site=%d |", + owner, + site, + )) + if slot.CompareAndSwap(nil, values) { + return values + } + return slot.Load() +} + +func profileRecordAccountedAlloc(lease allocationLease, sz int64) { + if !lease.profiled { + return + } + values := accountedProfileSample(lease.owner, lease.site) + values.Bytes.Allocated.Add(uint64(sz)) + values.Objects.Allocated.Add(1) + values.Bytes.Inuse.Add(sz) + values.Objects.Inuse.Add(1) +} + +func profileRecordAccountedFree(lease allocationLease, sz int64) { + if !lease.profiled { + return + } + values := accountedProfileSample(lease.owner, lease.site) + values.Bytes.Inuse.Add(-sz) + values.Objects.Inuse.Add(-1) +} + func profileRecordRealloc(skip int, oldPtr, newPtr uintptr, oldSz, newSz int64) { if !profilingEnabled.Load() { return diff --git a/pkg/common/mpool/mpool_profile_test.go b/pkg/common/mpool/mpool_profile_test.go index 8867cda1e013d..c95750f6a0126 100644 --- a/pkg/common/mpool/mpool_profile_test.go +++ b/pkg/common/mpool/mpool_profile_test.go @@ -103,3 +103,85 @@ func TestProfileWritable(t *testing.T) { mp.Free(bs) } + +func TestAccountedProfileUsesProvenanceAcrossProfilingToggle(t *testing.T) { + DisableProfiling() + defer DisableProfiling() + registry, account := newTestAllocationAccount(t, 1024, 1) + mp := MustNew("accounted-profile") + defer DeleteMPool(mp) + values := accountedProfileSample(testAllocationOwner, testAllocationSite) + before := values.Values() + trackedBefore := ProfileTrackedCount() + + EnableProfiling() + buffer, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + require.Equal(t, trackedBefore, ProfileTrackedCount(), + "accounted provenance does not need a per-pointer stack entry") + afterAlloc := values.Values() + require.Equal(t, int64(1), afterAlloc[0]-before[0]) + require.Equal(t, int64(64), afterAlloc[1]-before[1]) + require.Equal(t, int64(1), afterAlloc[2]-before[2]) + require.Equal(t, int64(64), afterAlloc[3]-before[3]) + + // The allocation lease remembers whether it was profiled, so disabling + // collection cannot strand an existing in-use sample. + DisableProfiling() + mp.Free(buffer) + afterFree := values.Values() + require.Equal(t, before[2], afterFree[2]) + require.Equal(t, before[3], afterFree[3]) + finalizeTestAllocationAccount(t, registry, account) +} + +func BenchmarkProfileAllocFree(b *testing.B) { + EnableProfiling() + defer DisableProfiling() + for _, accounted := range []bool{false, true} { + name := "stack" + if accounted { + name = "accounted-provenance" + } + b.Run(name, func(b *testing.B) { + mp := MustNew("profile-benchmark") + defer DeleteMPool(mp) + var registry *AllocationAccountRegistry + var account *AllocationAccount + if accounted { + registry, account = newTestAllocationAccount(b, 1<<60, 1) + // Initialize the bounded owner/site sample outside the measured loop. + accountedProfileSample(testAllocationOwner, testAllocationSite) + } + b.ReportAllocs() + b.ResetTimer() + for range b.N { + var buffer []byte + var err error + if accounted { + buffer, err = mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + } else { + buffer, err = mp.Alloc(64, true) + } + if err != nil { + b.Fatal(err) + } + mp.Free(buffer) + } + b.StopTimer() + if accounted { + finalizeTestAllocationAccount(b, registry, account) + } + }) + } +} diff --git a/pkg/container/vector/vector.go b/pkg/container/vector/vector.go index b746b24b4a4c2..d86f391f88aff 100644 --- a/pkg/container/vector/vector.go +++ b/pkg/container/vector/vector.go @@ -199,6 +199,40 @@ func (v *Vector) SetLength(n int) { v.length = n } +// AppendCheckpoint captures the logical state changed by append operations. +// Capacity growth is deliberately not rolled back: it remains owned by the +// vector and can be reused by a later append. +type AppendCheckpoint struct { + length int + areaLength int + sorted bool +} + +func (v *Vector) MakeAppendCheckpoint() AppendCheckpoint { + return AppendCheckpoint{ + length: v.length, + areaLength: len(v.area), + sorted: v.sorted, + } +} + +// RollbackAppend restores the logical state captured before an attempted +// append. attemptedRows is needed because grouping bits can be published +// before a varlen copy fails and advances length. +func (v *Vector) RollbackAppend(checkpoint AppendCheckpoint, attemptedRows int) { + if checkpoint.length < 0 || checkpoint.length > v.length || + checkpoint.areaLength < 0 || checkpoint.areaLength > len(v.area) || + attemptedRows < 0 { + panic("invalid vector append checkpoint") + } + end := max(v.length, checkpoint.length+attemptedRows) + nulls.RemoveRange(&v.nsp, uint64(checkpoint.length), uint64(end)) + nulls.RemoveRange(&v.gsp, uint64(checkpoint.length), uint64(end)) + v.length = checkpoint.length + v.area = v.area[:checkpoint.areaLength] + v.sorted = checkpoint.sorted +} + // Size of data, I think this function is inherently broken. This // Size is not meaningful other than used in (approximate) memory accounting. func (v *Vector) Size() int { diff --git a/pkg/container/vector/vector_test.go b/pkg/container/vector/vector_test.go index e54c763b2ad58..236b565827d4b 100644 --- a/pkg/container/vector/vector_test.go +++ b/pkg/container/vector/vector_test.go @@ -64,6 +64,34 @@ func TestLength(t *testing.T) { } } +func TestAppendCheckpointRollback(t *testing.T) { + mp := mpool.MustNewZero() + vec := NewVec(types.T_varchar.ToType()) + defer vec.Free(mp) + first := strings.Repeat("a", 64) + require.NoError(t, AppendBytes(vec, []byte(first), false, mp)) + vec.GetGrouping().Set(0) + vec.SetSorted(true) + checkpoint := vec.MakeAppendCheckpoint() + + require.NoError(t, AppendBytes(vec, []byte(strings.Repeat("b", 96)), false, mp)) + vec.GetNulls().Set(1) + vec.GetGrouping().Set(1) + // Grouping publication can precede a failed varlen copy and therefore can + // extend beyond the length reached by the copy itself. + vec.GetGrouping().Set(2) + vec.SetSorted(false) + vec.RollbackAppend(checkpoint, 2) + + require.Equal(t, 1, vec.Length()) + require.Equal(t, []string{first}, InefficientMustStrCol(vec)) + require.False(t, vec.GetNulls().Contains(1)) + require.True(t, vec.GetGrouping().Contains(0)) + require.False(t, vec.GetGrouping().Contains(1)) + require.False(t, vec.GetGrouping().Contains(2)) + require.True(t, vec.GetSorted()) +} + func TestCapacityForUntypedNull(t *testing.T) { vec := NewVec(types.T_any.ToType()) require.Equal(t, 0, vec.Capacity()) diff --git a/pkg/sql/colexec/join_util.go b/pkg/sql/colexec/join_util.go index d7c42214ff6f6..56b0cae454458 100644 --- a/pkg/sql/colexec/join_util.go +++ b/pkg/sql/colexec/join_util.go @@ -71,9 +71,10 @@ func (bs *Batches) CopyIntoBatches(src *batch.Batch, proc *process.Process) (err // // The append is transactional. In particular, an allocation rejection while // copying a later input must not destroy batches retained from earlier inputs: -// HashBuild needs those batches intact to recover by spilling them. A partial -// tail is copied into the private staging set before it is extended, so an -// error cannot leave the published tail partially mutated either. +// HashBuild needs those batches intact to recover by spilling them. Existing +// vectors are restored to logical append checkpoints on failure; successful +// capacity growth remains owned and reusable. This avoids repeatedly copying +// a partial 8,192-row tail as small input batches arrive. func (bs *Batches) CopyIntoBatchesWithAllocation( src *batch.Batch, proc *process.Process, @@ -84,46 +85,46 @@ func (bs *Batches) CopyIntoBatchesWithAllocation( return mpool.ErrAllocationAccountMismatch } - var staged Batches - defer func() { - if err != nil { - staged.Clean(proc.Mp()) + originalLen := len(bs.Buf) + originalMemSize := bs.MemSize + originalNil := bs.Buf == nil + var originalTail *batch.Batch + originalTailRows := 0 + var localTailCheckpoints [16]vector.AppendCheckpoint + var tailCheckpoints []vector.AppendCheckpoint + if originalLen > 0 && bs.Buf[originalLen-1].RowCount() != DefaultBatchSize { + originalTail = bs.Buf[originalLen-1] + originalTailRows = originalTail.RowCount() + if len(originalTail.Vecs) > len(localTailCheckpoints) { + tailCheckpoints = make([]vector.AppendCheckpoint, len(originalTail.Vecs)) + } else { + tailCheckpoints = localTailCheckpoints[:len(originalTail.Vecs)] } - }() - - replaceTail := len(bs.Buf) > 0 && - bs.Buf[len(bs.Buf)-1].RowCount() != DefaultBatchSize - if replaceTail { - if err = staged.copyIntoBatches( - bs.Buf[len(bs.Buf)-1], - proc, - selection, - ); err != nil { - return err + for i := range originalTail.Vecs { + tailCheckpoints[i] = originalTail.Vecs[i].MakeAppendCheckpoint() } } - if err = staged.copyIntoBatches(src, proc, selection); err != nil { - return err - } - - if replaceTail { - oldTail := bs.Buf[len(bs.Buf)-1] - bs.Buf = bs.Buf[:len(bs.Buf)-1] - bs.Buf = append(bs.Buf, staged.Buf...) - bs.MemSize += staged.MemSize - staged.Buf = nil - staged.MemSize = 0 - oldTail.Clean(proc.Mp()) + if err = bs.copyIntoBatches(src, proc, selection); err == nil { return nil } - if bs.Buf == nil { - bs.Buf = make([]*batch.Batch, 0, max(16, len(staged.Buf))) + for i := originalLen; i < len(bs.Buf); i++ { + bs.Buf[i].Clean(proc.Mp()) } - bs.Buf = append(bs.Buf, staged.Buf...) - bs.MemSize += staged.MemSize - staged.Buf = nil - staged.MemSize = 0 - return nil + bs.Buf = bs.Buf[:originalLen] + if originalNil { + bs.Buf = nil + } + if originalTail != nil { + for i := range originalTail.Vecs { + originalTail.Vecs[i].RollbackAppend( + tailCheckpoints[i], + DefaultBatchSize-originalTailRows, + ) + } + originalTail.SetRowCount(originalTailRows) + } + bs.MemSize = originalMemSize + return err } func (bs *Batches) copyIntoBatches( @@ -150,6 +151,9 @@ func (bs *Batches) copyIntoBatches( } } if err != nil { + if tmp != nil { + tmp.Clean(proc.Mp()) + } return err } bs.MemSize += int64(tmp.Size()) @@ -280,8 +284,21 @@ func appendToFixedSizeFromOffset(dst *batch.Batch, src *batch.Batch, offset int, if length+offset > src.RowCount() { length = src.RowCount() - offset } + var localCheckpoints [16]vector.AppendCheckpoint + var checkpoints []vector.AppendCheckpoint + if len(dst.Vecs) > len(localCheckpoints) { + checkpoints = make([]vector.AppendCheckpoint, len(dst.Vecs)) + } else { + checkpoints = localCheckpoints[:len(dst.Vecs)] + } + for i := range dst.Vecs { + checkpoints[i] = dst.Vecs[i].MakeAppendCheckpoint() + } for i := range dst.Vecs { if err = dst.Vecs[i].UnionBatch(src.Vecs[i], int64(offset), length, nil, proc.Mp()); err != nil { + for j := 0; j <= i; j++ { + dst.Vecs[j].RollbackAppend(checkpoints[j], length) + } return 0, err } dst.Vecs[i].SetSorted(false) diff --git a/pkg/sql/colexec/join_util_test.go b/pkg/sql/colexec/join_util_test.go index 94a91ca78ebff..1dd17a9e1fdec 100644 --- a/pkg/sql/colexec/join_util_test.go +++ b/pkg/sql/colexec/join_util_test.go @@ -15,6 +15,7 @@ package colexec import ( + "sync/atomic" "testing" "github.com/stretchr/testify/require" @@ -26,6 +27,65 @@ import ( "github.com/matrixorigin/matrixone/pkg/testutil" ) +type testAppendCapacityController struct { + limit atomic.Uint64 + used atomic.Uint64 +} + +func (c *testAppendCapacityController) AcquireAllocationCapacity(size uint64) error { + for { + used := c.used.Load() + limit := c.limit.Load() + if size > limit || used > limit-size { + return mpool.ErrAllocationAccountCapacity + } + if c.used.CompareAndSwap(used, used+size) { + return nil + } + } +} + +func (c *testAppendCapacityController) ReleaseAllocationCapacity(size uint64) { + for { + used := c.used.Load() + if size > used { + panic("test allocation capacity underflow") + } + if c.used.CompareAndSwap(used, used-size) { + return + } + } +} + +func BenchmarkCopyIntoBatchesPartialTail(b *testing.B) { + const rowsPerInput = 128 + proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) + defer proc.Free() + input := testutil.NewBatch( + []types.Type{ + types.T_int64.ToType(), + types.T_int64.ToType(), + types.T_int64.ToType(), + types.T_int64.ToType(), + }, + true, + rowsPerInput, + proc.Mp(), + ) + defer input.Clean(proc.Mp()) + + b.ResetTimer() + for range b.N { + var batches Batches + for rows := 0; rows < DefaultBatchSize; rows += rowsPerInput { + if err := batches.CopyIntoBatches(input, proc); err != nil { + b.Fatal(err) + } + } + batches.Clean(proc.Mp()) + } +} + func TestBatches(t *testing.T) { var batches Batches proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) @@ -133,3 +193,94 @@ func TestBatchesShrinkPreservesAllocationAndRollback(t *testing.T) { _, err = measure(1<<20, true) require.NoError(t, err) } + +func TestCopyIntoBatchesAllocationFailureRollsBackPartialTail(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + controller := &testAppendCapacityController{} + controller.limit.Store(1 << 60) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.OpenWithController(1<<60, controller) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + + typesInTail := []types.Type{types.T_int64.ToType(), types.T_int64.ToType()} + initial := testutil.NewBatch(typesInTail, true, 1, proc.Mp()) + defer initial.Clean(proc.Mp()) + var batches Batches + require.NoError(t, batches.CopyIntoBatchesWithAllocation(initial, proc, selection)) + require.Len(t, batches.Buf, 1) + require.Equal(t, 1, batches.RowCount()) + firstBefore := append([]int64(nil), vector.MustFixedColNoTypeCheck[int64](batches.Buf[0].Vecs[0])...) + secondBefore := append([]int64(nil), vector.MustFixedColNoTypeCheck[int64](batches.Buf[0].Vecs[1])...) + + const appendRows = 128 + oldCapacity := cap(batches.Buf[0].Vecs[0].GetData()) + requiredBytes := (batches.RowCount() + appendRows) * types.T_int64.ToType().TypeSize() + newCapacity, ok := mpool.GrowCapacity(int64(oldCapacity), int64(requiredBytes)) + require.True(t, ok) + require.Greater(t, newCapacity, int64(oldCapacity)) + controller.limit.Store(controller.used.Load() + uint64(newCapacity)) + + more := testutil.NewBatch(typesInTail, true, appendRows, proc.Mp()) + defer more.Clean(proc.Mp()) + err = batches.CopyIntoBatchesWithAllocation(more, proc, selection) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Len(t, batches.Buf, 1) + require.Equal(t, 1, batches.RowCount()) + require.Equal(t, firstBefore, vector.MustFixedColNoTypeCheck[int64](batches.Buf[0].Vecs[0])) + require.Equal(t, secondBefore, vector.MustFixedColNoTypeCheck[int64](batches.Buf[0].Vecs[1])) + + controller.limit.Store(1 << 60) + require.NoError(t, batches.CopyIntoBatchesWithAllocation(more, proc, selection)) + require.Equal(t, 1+appendRows, batches.RowCount()) + batches.Clean(proc.Mp()) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, controller.used.Load()) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestCopyIntoBatchesFailureRollsBackEarlierTailChunk(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + controller := &testAppendCapacityController{} + controller.limit.Store(1 << 60) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.OpenWithController(1<<60, controller) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + + typesInTail := []types.Type{types.T_int64.ToType(), types.T_int64.ToType()} + initial := testutil.NewBatch(typesInTail, true, DefaultBatchSize-2, proc.Mp()) + defer initial.Clean(proc.Mp()) + var batches Batches + require.NoError(t, batches.CopyIntoBatchesWithAllocation(initial, proc, selection)) + firstBefore := append([]int64(nil), vector.MustFixedColNoTypeCheck[int64](batches.Buf[0].Vecs[0])...) + secondBefore := append([]int64(nil), vector.MustFixedColNoTypeCheck[int64](batches.Buf[0].Vecs[1])...) + + // Filling the last two rows needs no growth. Reject creation of the next + // batch, after that first chunk has already succeeded. + controller.limit.Store(controller.used.Load()) + more := testutil.NewBatch(typesInTail, true, 128, proc.Mp()) + defer more.Clean(proc.Mp()) + err = batches.CopyIntoBatchesWithAllocation(more, proc, selection) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Len(t, batches.Buf, 1) + require.Equal(t, DefaultBatchSize-2, batches.RowCount()) + require.Equal(t, firstBefore, vector.MustFixedColNoTypeCheck[int64](batches.Buf[0].Vecs[0])) + require.Equal(t, secondBefore, vector.MustFixedColNoTypeCheck[int64](batches.Buf[0].Vecs[1])) + + controller.limit.Store(1 << 60) + require.NoError(t, batches.CopyIntoBatchesWithAllocation(more, proc, selection)) + require.Equal(t, DefaultBatchSize-2+128, batches.RowCount()) + batches.Clean(proc.Mp()) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, controller.used.Load()) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} From 8ce077f808d704577d649cd234479441d1a36111 Mon Sep 17 00:00:00 2001 From: aptend Date: Sun, 2 Aug 2026 17:23:47 +0800 Subject: [PATCH 51/61] perf: coalesce equivalent accounted batches --- .../batch/allocation_account_test.go | 73 +++++++++++++++++++ pkg/container/batch/batch.go | 16 +++- pkg/container/pSpool/copy.go | 5 +- pkg/container/pSpool/sender_test.go | 14 +++- pkg/container/vector/allocation_account.go | 23 +++++- .../vector/allocation_account_test.go | 52 ++++++++++++- pkg/container/vector/pSpoolTools.go | 2 +- pkg/sql/colexec/join_util.go | 15 +++- pkg/sql/colexec/join_util_test.go | 33 +++++++++ 9 files changed, 219 insertions(+), 14 deletions(-) diff --git a/pkg/container/batch/allocation_account_test.go b/pkg/container/batch/allocation_account_test.go index ba710e90a90be..799605335da72 100644 --- a/pkg/container/batch/allocation_account_test.go +++ b/pkg/container/batch/allocation_account_test.go @@ -543,6 +543,79 @@ func TestMixedBatchAllocationBatchSetPreservesVectorProvenance(t *testing.T) { finalizeTestBatchAllocationAccount(t, state) } +func TestBatchSetFillsTailAcrossEquivalentAllocationSelections(t *testing.T) { + state := newTestBatchAllocationAccount(t, 128) + equivalent, err := vector.NewAllocationAccountSelection( + state.account, + 1, + 1, + 2, + 3, + 4, + ) + require.NoError(t, err) + require.NotSame(t, state.selection, equivalent) + mp := mpool.MustNewZero() + set := NewBatchSet(4) + first := newMixedBatchAllocationSource(t, mp, state.selection, 2) + second := newMixedBatchAllocationSource(t, mp, equivalent, 2) + + _, err = set.Extend(mp, first, nil) + require.NoError(t, err) + require.Equal(t, 1, set.ReadyDeltaFor(second, second.RowCount())) + _, err = set.Extend(mp, second, nil) + require.NoError(t, err) + require.Equal(t, 1, set.Length()) + require.Equal(t, 1, set.ReadyCount()) + require.Equal(t, 4, set.Get(0).RowCount()) + require.Same(t, state.selection, set.Get(0).Vecs[0].AllocationAccountSelection()) + + first.Clean(mp) + second.Clean(mp) + set.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchSetCoalescesEquivalentParallelProducerChunks(t *testing.T) { + const ( + batchMaxRows = 8192 + chunkRows = 72 + chunkCount = 114 + ) + state := newTestBatchAllocationAccount(t, 256) + mp := mpool.MustNewZero() + set := NewBatchSet(batchMaxRows) + + for range chunkCount { + selection, err := vector.NewAllocationAccountSelection( + state.account, + 1, + 1, + 2, + 3, + 4, + ) + require.NoError(t, err) + source := newMixedBatchAllocationSource( + t, + mp, + selection, + chunkRows, + ) + _, err = set.Extend(mp, source, nil) + require.NoError(t, err) + source.Clean(mp) + } + + require.Equal(t, 2, set.Length()) + require.Equal(t, 1, set.ReadyCount()) + require.Equal(t, batchMaxRows, set.Get(0).RowCount()) + require.Equal(t, chunkRows*chunkCount-batchMaxRows, set.Get(1).RowCount()) + + set.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + func TestBatchSetStartsNewTailWhenVectorProvenanceChanges(t *testing.T) { state := newTestBatchAllocationAccount(t, 128) mp := mpool.MustNewZero() diff --git a/pkg/container/batch/batch.go b/pkg/container/batch/batch.go index f07dab65ae4e6..9bce06e335050 100644 --- a/pkg/container/batch/batch.go +++ b/pkg/container/batch/batch.go @@ -765,7 +765,10 @@ func (bat *Batch) prepareOwnedDecodeVectors(count int, mp *mpool.MPool) error { vec.Free(mp) } vec.SetOffHeap(bat.offHeap) - if vec.AllocationAccountSelection() != selection { + if !vector.AllocationAccountSelectionsEqual( + vec.AllocationAccountSelection(), + selection, + ) { if err := vec.CanSetAllocationAccount(selection); err != nil { vec.Free(mp) vec.SetOffHeap(bat.offHeap) @@ -1071,7 +1074,10 @@ func vectorAllocationSelectionsMatch(left, right *Batch) bool { if left == nil || right == nil || len(left.Vecs) != len(right.Vecs) { return false } - if left.allocationAccount != right.allocationAccount { + if !vector.AllocationAccountSelectionsEqual( + left.allocationAccount, + right.allocationAccount, + ) { return false } for i := range left.Vecs { @@ -1081,8 +1087,10 @@ func vectorAllocationSelectionsMatch(left, right *Batch) bool { } continue } - if left.Vecs[i].AllocationAccountSelection() != - right.Vecs[i].AllocationAccountSelection() { + if !vector.AllocationAccountSelectionsEqual( + left.Vecs[i].AllocationAccountSelection(), + right.Vecs[i].AllocationAccountSelection(), + ) { return false } } diff --git a/pkg/container/pSpool/copy.go b/pkg/container/pSpool/copy.go index 14898f482bad9..27a2a8d6e994b 100644 --- a/pkg/container/pSpool/copy.go +++ b/pkg/container/pSpool/copy.go @@ -72,7 +72,10 @@ func (cb *cachedBatch) GetCopiedBatch( cacheID, dst = cb.buffer.getCacheID() dst.Recursive = src.Recursive dst.ShuffleIDX = src.ShuffleIDX - if sourceSelection := src.AllocationAccountSelection(); sourceSelection != dst.AllocationAccountSelection() { + if sourceSelection := src.AllocationAccountSelection(); !vector.AllocationAccountSelectionsEqual( + sourceSelection, + dst.AllocationAccountSelection(), + ) { if err = dst.SetAllocationAccount(sourceSelection); err != nil { cb.CacheBatch(true, cacheID, dst) return nil, false, 0, err diff --git a/pkg/container/pSpool/sender_test.go b/pkg/container/pSpool/sender_test.go index b1d43233843e5..6a2d05276d1d3 100644 --- a/pkg/container/pSpool/sender_test.go +++ b/pkg/container/pSpool/sender_test.go @@ -114,6 +114,16 @@ func TestCachedBatchPreservesAllocationProvenance(t *testing.T) { 4, ) require.NoError(t, err) + equivalentSelection, err := vector.NewAllocationAccountSelection( + account, + 1, + 1, + 2, + 3, + 4, + ) + require.NoError(t, err) + require.NotSame(t, selection, equivalentSelection) otherAccount, err := registry.Open(1 << 20) require.NoError(t, err) otherSelection, err := vector.NewAllocationAccountSelection( @@ -145,7 +155,7 @@ func TestCachedBatchPreservesAllocationProvenance(t *testing.T) { return source } firstSource := newSource("first cached allocation payload", selection) - secondSource := newSource("second", selection) + secondSource := newSource("second", equivalentSelection) secondSource.Vecs[0].ToConst() cache := initCachedBatch(mp, 1) @@ -169,7 +179,7 @@ func TestCachedBatchPreservesAllocationProvenance(t *testing.T) { require.Same(t, selection, second.AllocationAccountSelection()) require.Same( t, - selection, + equivalentSelection, second.Vecs[0].AllocationAccountSelection(), ) require.True(t, second.Vecs[0].GetGrouping().Contains(0)) diff --git a/pkg/container/vector/allocation_account.go b/pkg/container/vector/allocation_account.go index 6163657940ddf..bab11ae29546a 100644 --- a/pkg/container/vector/allocation_account.go +++ b/pkg/container/vector/allocation_account.go @@ -47,6 +47,25 @@ type AllocationAccountSelection struct { groupingSite mpool.AllocationSite } +// AllocationAccountSelectionsEqual reports whether two immutable selections +// describe the same physical allocation provenance. Separately constructed +// selections are interchangeable only when they charge the same account, +// owner, and allocation sites. +func AllocationAccountSelectionsEqual( + left, right *AllocationAccountSelection, +) bool { + if left == right { + return true + } + return left != nil && right != nil && + left.account == right.account && + left.owner == right.owner && + left.dataSite == right.dataSite && + left.areaSite == right.areaSite && + left.nullsSite == right.nullsSite && + left.groupingSite == right.groupingSite +} + func NewAllocationAccountSelection( account *mpool.AllocationAccount, owner mpool.AllocationOwner, @@ -209,7 +228,7 @@ func (v *Vector) CanSetAllocationAccount( ) } } - if v.allocationAccount == selection { + if AllocationAccountSelectionsEqual(v.allocationAccount, selection) { return nil } if v.hasBackingStorage() { @@ -247,7 +266,7 @@ func (v *Vector) SetAllocationAccount( if err := v.CanSetAllocationAccount(selection); err != nil { return err } - if v.allocationAccount == selection { + if AllocationAccountSelectionsEqual(v.allocationAccount, selection) { return nil } if v.allocationAccount != nil && selection == nil { diff --git a/pkg/container/vector/allocation_account_test.go b/pkg/container/vector/allocation_account_test.go index 462a78e516934..167bca55eac57 100644 --- a/pkg/container/vector/allocation_account_test.go +++ b/pkg/container/vector/allocation_account_test.go @@ -136,6 +136,46 @@ func TestVectorAllocationAccountConfiguration(t *testing.T) { finalizeTestVectorAllocationAccount(t, state) } +func TestAllocationAccountSelectionsEqual(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 8) + equivalent, err := NewAllocationAccountSelection( + state.account, + testVectorAllocationOwner, + testVectorDataAllocationSite, + testVectorAreaAllocationSite, + testVectorNullAllocationSite, + testVectorGroupAllocationSite, + ) + require.NoError(t, err) + differentSite, err := NewAllocationAccountSelection( + state.account, + testVectorAllocationOwner, + testVectorDataAllocationSite+1, + testVectorAreaAllocationSite, + testVectorNullAllocationSite, + testVectorGroupAllocationSite, + ) + require.NoError(t, err) + + require.NotSame(t, state.selection, equivalent) + require.True(t, AllocationAccountSelectionsEqual(state.selection, equivalent)) + require.False(t, AllocationAccountSelectionsEqual(state.selection, differentSite)) + require.False(t, AllocationAccountSelectionsEqual(state.selection, nil)) + require.True(t, AllocationAccountSelectionsEqual(nil, nil)) + + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, AppendFixed(vec, int64(1), false, mp)) + require.NoError(t, vec.CanSetAllocationAccount(equivalent)) + require.NoError(t, vec.SetAllocationAccount(equivalent)) + // Equivalent provenance is a no-op: existing physical ownership remains + // attached to the original immutable selection. + require.Same(t, state.selection, vec.AllocationAccountSelection()) + vec.Free(mp) + + finalizeTestVectorAllocationAccount(t, state) +} + func TestVectorAllocationAccountFixedResetReuseAndFree(t *testing.T) { state := newTestVectorAllocationAccount(t, 1<<20, 16) mp := mpool.MustNewZero() @@ -1092,10 +1132,20 @@ func TestDetachedBufferPreservesAllocationProvenance(t *testing.T) { source.Free(mp) require.Equal(t, used, state.account.Snapshot().Used) + equivalent, err := NewAllocationAccountSelection( + state.account, + testVectorAllocationOwner, + testVectorDataAllocationSite, + testVectorAreaAllocationSite, + testVectorNullAllocationSite, + testVectorGroupAllocationSite, + ) + require.NoError(t, err) + require.NotSame(t, state.selection, equivalent) destination := newAccountedTestVector( t, types.T_varchar.ToType(), - state.selection, + equivalent, ) require.True(t, data.CanAttachTo(destination, DetachedDataBuffer)) require.False(t, data.CanAttachTo(destination, DetachedAreaBuffer)) diff --git a/pkg/container/vector/pSpoolTools.go b/pkg/container/vector/pSpoolTools.go index f2bdfecddf4b1..3f8beeeee3ab3 100644 --- a/pkg/container/vector/pSpoolTools.go +++ b/pkg/container/vector/pSpoolTools.go @@ -72,7 +72,7 @@ func (b *DetachedBuffer) CanAttachTo( kind DetachedBufferKind, ) bool { if b == nil || v == nil || cap(b.data) == 0 || - b.selection != v.allocationAccount || + !AllocationAccountSelectionsEqual(b.selection, v.allocationAccount) || kind > DetachedAreaBuffer { return false } diff --git a/pkg/sql/colexec/join_util.go b/pkg/sql/colexec/join_util.go index 56b0cae454458..e61fdd3346196 100644 --- a/pkg/sql/colexec/join_util.go +++ b/pkg/sql/colexec/join_util.go @@ -81,7 +81,10 @@ func (bs *Batches) CopyIntoBatchesWithAllocation( selection *vector.AllocationAccountSelection, ) (err error) { if len(bs.Buf) > 0 && - bs.Buf[len(bs.Buf)-1].AllocationAccountSelection() != selection { + !vector.AllocationAccountSelectionsEqual( + bs.Buf[len(bs.Buf)-1].AllocationAccountSelection(), + selection, + ) { return mpool.ErrAllocationAccountMismatch } @@ -136,7 +139,10 @@ func (bs *Batches) copyIntoBatches( bs.Buf = make([]*batch.Batch, 0, 16) } if len(bs.Buf) > 0 && - bs.Buf[len(bs.Buf)-1].AllocationAccountSelection() != selection { + !vector.AllocationAccountSelectionsEqual( + bs.Buf[len(bs.Buf)-1].AllocationAccountSelection(), + selection, + ) { return mpool.ErrAllocationAccountMismatch } @@ -174,7 +180,10 @@ func (bs *Batches) copyIntoBatches( lenBuf := len(bs.Buf) if lenBuf > 0 && bs.Buf[lenBuf-1].RowCount() != DefaultBatchSize { tmp = bs.Buf[lenBuf-1] - if tmp.AllocationAccountSelection() != selection { + if !vector.AllocationAccountSelectionsEqual( + tmp.AllocationAccountSelection(), + selection, + ) { return mpool.ErrAllocationAccountMismatch } } else { diff --git a/pkg/sql/colexec/join_util_test.go b/pkg/sql/colexec/join_util_test.go index 1dd17a9e1fdec..a872de2d164a8 100644 --- a/pkg/sql/colexec/join_util_test.go +++ b/pkg/sql/colexec/join_util_test.go @@ -194,6 +194,39 @@ func TestBatchesShrinkPreservesAllocationAndRollback(t *testing.T) { require.NoError(t, err) } +func TestCopyIntoBatchesAcceptsEquivalentAllocationSelection(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + first, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + second, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + require.NotSame(t, first, second) + input := testutil.NewBatch( + []types.Type{types.T_int64.ToType()}, + true, + 32, + proc.Mp(), + ) + defer input.Clean(proc.Mp()) + + var batches Batches + require.NoError(t, batches.CopyIntoBatchesWithAllocation(input, proc, first)) + require.NoError(t, batches.CopyIntoBatchesWithAllocation(input, proc, second)) + require.Len(t, batches.Buf, 1) + require.Equal(t, 64, batches.RowCount()) + require.Same(t, first, batches.Buf[0].AllocationAccountSelection()) + + batches.Clean(proc.Mp()) + require.Zero(t, account.Snapshot().Used) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + func TestCopyIntoBatchesAllocationFailureRollsBackPartialTail(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() From 1129ec8a274b74da6d92a469f47b904dad80726a Mon Sep 17 00:00:00 2001 From: aptend Date: Sun, 2 Aug 2026 23:29:09 +0800 Subject: [PATCH 52/61] fix: preserve sparse bitmap row capacity on duplicate --- .../vector/allocation_account_test.go | 40 +++++++++++++++++++ pkg/container/vector/vector.go | 27 +++++-------- 2 files changed, 49 insertions(+), 18 deletions(-) diff --git a/pkg/container/vector/allocation_account_test.go b/pkg/container/vector/allocation_account_test.go index 167bca55eac57..917104d404b46 100644 --- a/pkg/container/vector/allocation_account_test.go +++ b/pkg/container/vector/allocation_account_test.go @@ -318,6 +318,46 @@ func TestVectorAllocationAccountBitmapShrinkUsesNoScratch(t *testing.T) { finalizeTestVectorAllocationAccount(t, state) } +func TestVectorAllocationAccountDupPreservesSparseBitmapRowDomain(t *testing.T) { + state := newTestVectorAllocationAccount(t, 8<<20, 16) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, AppendFixed(vec, int64(0), true, mp)) + for i := 1; i < 130; i++ { + require.NoError(t, AppendFixed(vec, int64(i), false, mp)) + } + require.NoError(t, vec.ensureGroupingCapacity(1, mp)) + vec.GetGrouping().Add(0) + require.Equal(t, 1, vec.GetNulls().Count()) + require.Equal(t, 1, vec.GetGrouping().Count()) + dup, err := vec.Dup(mp) + require.NoError(t, err) + require.GreaterOrEqual( + t, + dup.GetNulls().GetBitmap().ExternalStorageCapacity(), + 3, + ) + require.GreaterOrEqual( + t, + dup.GetGrouping().GetBitmap().ExternalStorageCapacity(), + 3, + ) + sels := make([]int64, 129) + for i := range sels { + sels[i] = int64(i + 1) + } + require.NotPanics(t, func() { + dup.Shrink(sels, false) + }) + require.Equal(t, 129, dup.Length()) + require.Zero(t, dup.GetNulls().Count()) + require.Zero(t, dup.GetGrouping().Count()) + + dup.Free(mp) + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + func TestVectorAllocationAccountBitmapShuffleAccountsScratch(t *testing.T) { state := newTestVectorAllocationAccount(t, 8<<20, 16) mp := mpool.MustNewZero() diff --git a/pkg/container/vector/vector.go b/pkg/container/vector/vector.go index d86f391f88aff..cf840e4163764 100644 --- a/pkg/container/vector/vector.go +++ b/pkg/container/vector/vector.go @@ -1643,13 +1643,8 @@ func (v *Vector) dup( if v.IsConstNull() { w.length = v.length if v.HasGrouping() { - groupingRows := v.GetGrouping().GetBitmap().Len() - if groupingRows < 0 || groupingRows > int64(math.MaxInt) { - w.Free(mp) - return nil, mpool.ErrAllocationAccountInvalid - } if err := w.ensureGroupingCapacity( - int(groupingRows), + v.length, mp, ); err != nil { w.Free(mp) @@ -1674,22 +1669,18 @@ func (v *Vector) dup( } dataLen *= v.length } - if nullRows := v.GetNulls().GetBitmap().Len(); !v.GetNulls().EmptyByFlag() { - if nullRows < 0 || nullRows > int64(math.MaxInt) { - w.Free(mp) - return nil, mpool.ErrAllocationAccountInvalid - } - if err := w.ensureNullCapacity(int(nullRows), mp); err != nil { + // Bitmap logical length only reaches the highest set bit, so it can be much + // shorter than a sparse vector. The duplicate must cover the complete row + // domain because allocation-free transforms (for example ordered Shrink) + // may address every output row. + if !v.GetNulls().EmptyByFlag() { + if err := w.ensureNullCapacity(v.length, mp); err != nil { w.Free(mp) return nil, err } } - if groupingRows := v.GetGrouping().GetBitmap().Len(); !v.GetGrouping().EmptyByFlag() { - if groupingRows < 0 || groupingRows > int64(math.MaxInt) { - w.Free(mp) - return nil, mpool.ErrAllocationAccountInvalid - } - if err := w.ensureGroupingCapacity(int(groupingRows), mp); err != nil { + if !v.GetGrouping().EmptyByFlag() { + if err := w.ensureGroupingCapacity(v.length, mp); err != nil { w.Free(mp) return nil, err } From 2aefd528eaa5a2472d94aae954caa13422488895 Mon Sep 17 00:00:00 2001 From: aptend Date: Sun, 2 Aug 2026 23:29:09 +0800 Subject: [PATCH 53/61] fix: reject late remote fragments after abort --- .../remote_allocation_statement_group.go | 61 ++++++++++++- .../remote_allocation_statement_group_test.go | 87 ++++++++++++++++--- pkg/sql/compile/remoterunServer.go | 1 + 3 files changed, 133 insertions(+), 16 deletions(-) diff --git a/pkg/sql/compile/remote_allocation_statement_group.go b/pkg/sql/compile/remote_allocation_statement_group.go index 6d58d1a7296ab..2a2f86734e863 100644 --- a/pkg/sql/compile/remote_allocation_statement_group.go +++ b/pkg/sql/compile/remote_allocation_statement_group.go @@ -59,6 +59,11 @@ func remoteAllocationStatementGroupKey( // statement's execution time. var remoteAllocationStatementRegistrationTimeout = 5 * time.Minute +// A late RPC can carry a new MessageBoard, so keep a bounded record of an +// incomplete execution after its active group has been released. The record +// contains no statement resources. +var remoteAllocationStatementTombstoneTimeout = 5 * time.Minute + // collectRemoteFragmentCounts computes the number of pipeline RPCs that the // complete physical scope graph will send to each CN. The execution address // changes when traversal crosses a Remote scope: nested scopes targeting that @@ -110,9 +115,17 @@ func validateRemoteAllocationTopologyCapability( var remoteAllocationStatementGroups = struct { sync.Mutex - byBoard map[*message.MessageBoard]*remoteAllocationStatementGroup + byBoard map[*message.MessageBoard]*remoteAllocationStatementGroup + byKey map[string]*remoteAllocationStatementGroup + tombstones map[string]*remoteAllocationStatementTombstone }{ - byBoard: make(map[*message.MessageBoard]*remoteAllocationStatementGroup), + byBoard: make(map[*message.MessageBoard]*remoteAllocationStatementGroup), + byKey: make(map[string]*remoteAllocationStatementGroup), + tombstones: make(map[string]*remoteAllocationStatementTombstone), +} + +type remoteAllocationStatementTombstone struct { + timer *time.Timer } // remoteAllocationStatementGroup is the terminal owner for all pipeline RPCs @@ -120,6 +133,7 @@ var remoteAllocationStatementGroups = struct { // no individual fragment may close it or validate transferred allocations // while a sibling can still consume them. type remoteAllocationStatementGroup struct { + key string board *message.MessageBoard expected uint32 registered uint32 @@ -152,11 +166,12 @@ type remoteAllocationStatementTerminal struct { } func acquireRemoteAllocationStatementParticipant( + key string, board *message.MessageBoard, expected uint32, cancel func(error), ) (*remoteAllocationStatementParticipant, error) { - if board == nil { + if key == "" || board == nil { return nil, mpool.ErrAllocationAccountInvariant } if expected == 0 { @@ -165,15 +180,29 @@ func acquireRemoteAllocationStatementParticipant( remoteAllocationStatementGroups.Lock() defer remoteAllocationStatementGroups.Unlock() + if remoteAllocationStatementGroups.tombstones[key] != nil { + return nil, errors.Join( + mpool.ErrAllocationAccountInvariant, + moerr.NewInternalErrorNoCtx("remote allocation statement group already aborted"), + ) + } group := remoteAllocationStatementGroups.byBoard[board] if group == nil { + if remoteAllocationStatementGroups.byKey[key] != nil { + return nil, errors.Join( + mpool.ErrAllocationAccountInvariant, + moerr.NewInternalErrorNoCtx("remote allocation statement group key already registered"), + ) + } group = &remoteAllocationStatementGroup{ + key: key, board: board, expected: expected, } remoteAllocationStatementGroups.byBoard[board] = group + remoteAllocationStatementGroups.byKey[key] = group } - if group.expected != expected || group.finalized || + if group.key != key || group.expected != expected || group.finalized || group.expired || group.registered >= group.expected { return nil, errors.Join( mpool.ErrAllocationAccountInvariant, @@ -325,6 +354,9 @@ func releaseRemoteAllocationStatementGroup(group *remoteAllocationStatementGroup if remoteAllocationStatementGroups.byBoard[group.board] == group { delete(remoteAllocationStatementGroups.byBoard, group.board) } + if remoteAllocationStatementGroups.byKey[group.key] == group { + delete(remoteAllocationStatementGroups.byKey, group.key) + } remoteAllocationStatementGroups.Unlock() } @@ -360,6 +392,9 @@ func takeRemoteAllocationStatementGroupLocked( group *remoteAllocationStatementGroup, ) ([]*statementAllocationAttempt, []*mpool.MPool) { group.finalized = true + if group.expired && group.registered < group.expected { + installRemoteAllocationStatementTombstoneLocked(group.key) + } if group.timer != nil { group.timer.Stop() group.timer = nil @@ -371,6 +406,24 @@ func takeRemoteAllocationStatementGroupLocked( return attempts, pools } +func installRemoteAllocationStatementTombstoneLocked(key string) { + if remoteAllocationStatementGroups.tombstones[key] != nil { + return + } + tombstone := &remoteAllocationStatementTombstone{} + remoteAllocationStatementGroups.tombstones[key] = tombstone + tombstone.timer = time.AfterFunc( + remoteAllocationStatementTombstoneTimeout, + func() { + remoteAllocationStatementGroups.Lock() + if remoteAllocationStatementGroups.tombstones[key] == tombstone { + delete(remoteAllocationStatementGroups.tombstones, key) + } + remoteAllocationStatementGroups.Unlock() + }, + ) +} + func completeRemoteAllocationStatementGroup( group *remoteAllocationStatementGroup, attempts []*statementAllocationAttempt, diff --git a/pkg/sql/compile/remote_allocation_statement_group_test.go b/pkg/sql/compile/remote_allocation_statement_group_test.go index fe39c0ceb23b0..cf148a8f77809 100644 --- a/pkg/sql/compile/remote_allocation_statement_group_test.go +++ b/pkg/sql/compile/remote_allocation_statement_group_test.go @@ -16,6 +16,7 @@ package compile import ( "errors" + "fmt" "strings" "sync/atomic" "testing" @@ -43,6 +44,27 @@ func remoteAllocationStatementGroupRegistered(board *message.MessageBoard) bool return registered } +func acquireRemoteAllocationStatementTestParticipant( + t *testing.T, + board *message.MessageBoard, + expected uint32, + cancel func(error), +) (*remoteAllocationStatementParticipant, error) { + t.Helper() + key := fmt.Sprintf("test@%p", board) + t.Cleanup(func() { clearRemoteAllocationStatementTestTombstone(key) }) + return acquireRemoteAllocationStatementParticipant(key, board, expected, cancel) +} + +func clearRemoteAllocationStatementTestTombstone(key string) { + remoteAllocationStatementGroups.Lock() + if tombstone := remoteAllocationStatementGroups.tombstones[key]; tombstone != nil { + tombstone.timer.Stop() + delete(remoteAllocationStatementGroups.tombstones, key) + } + remoteAllocationStatementGroups.Unlock() +} + func (m *remoteAllocationAccountedMessage) Serialize() []byte { return nil } func (m *remoteAllocationAccountedMessage) Deserialize([]byte) message.Message { @@ -164,7 +186,7 @@ func TestRemoteAllocationStatementGroupDefersSharedBoardTerminal(t *testing.T) { destroyed: &destroyed, }, board) - first, err := acquireRemoteAllocationStatementParticipant(board, 2, nil) + first, err := acquireRemoteAllocationStatementTestParticipant(t, board, 2, nil) require.NoError(t, err) first.stage(attempt, producer.proc.Mp()) terminal, err := first.finish(nil) @@ -176,7 +198,7 @@ func TestRemoteAllocationStatementGroupDefersSharedBoardTerminal(t *testing.T) { require.NotContains(t, board.DebugString(), "closed") require.False(t, registry.AdmissionSuspended()) - second, err := acquireRemoteAllocationStatementParticipant(board, 2, nil) + second, err := acquireRemoteAllocationStatementTestParticipant(t, board, 2, nil) require.NoError(t, err) // The second fragment has no allocation owner. It still participates in // the statement boundary and, as the last fragment, drains the producer's @@ -201,12 +223,12 @@ func TestRemoteAllocationStatementGroupDefersSharedBoardTerminal(t *testing.T) { func TestRemoteAllocationStatementGroupRejectsTopologyMismatch(t *testing.T) { board := message.NewMessageBoard() - participant, err := acquireRemoteAllocationStatementParticipant(board, 2, nil) + participant, err := acquireRemoteAllocationStatementTestParticipant(t, board, 2, nil) require.NoError(t, err) - _, err = acquireRemoteAllocationStatementParticipant(board, 3, nil) + _, err = acquireRemoteAllocationStatementTestParticipant(t, board, 3, nil) require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) - second, err := acquireRemoteAllocationStatementParticipant(board, 2, nil) + second, err := acquireRemoteAllocationStatementTestParticipant(t, board, 2, nil) require.NoError(t, err) _, err = participant.finish(nil) require.NoError(t, err) @@ -247,7 +269,7 @@ func TestRemoteAllocationStatementGroupExpiresMissingFragment(t *testing.T) { destroyed: &destroyed, }, board) - participant, err := acquireRemoteAllocationStatementParticipant(board, 2, nil) + participant, err := acquireRemoteAllocationStatementTestParticipant(t, board, 2, nil) require.NoError(t, err) participant.stage(attempt, producer.proc.Mp()) terminal, err := participant.finish(nil) @@ -296,7 +318,8 @@ func TestRemoteAllocationStatementRegistrationTimerStartsBeforeFinish(t *testing }, board) canceled := make(chan error, 1) - participant, err := acquireRemoteAllocationStatementParticipant( + participant, err := acquireRemoteAllocationStatementTestParticipant( + t, board, 2, func(cause error) { canceled <- cause }, @@ -342,7 +365,7 @@ func TestRemoteAllocationStatementGroupFailureAbortsMissingFragment(t *testing.T producer.MessageBoard = board attempt, err := producer.beginAllocationAccountAttempt() require.NoError(t, err) - participant, err := acquireRemoteAllocationStatementParticipant(board, 2, nil) + participant, err := acquireRemoteAllocationStatementTestParticipant(t, board, 2, nil) require.NoError(t, err) participant.stage(attempt, producer.proc.Mp()) @@ -357,16 +380,54 @@ func TestRemoteAllocationStatementGroupFailureAbortsMissingFragment(t *testing.T require.False(t, remoteAllocationStatementGroupRegistered(board)) } +func TestRemoteAllocationStatementGroupRejectsLateFragmentAfterAbort(t *testing.T) { + key := remoteAllocationStatementGroupKey(newRemoteExecutionID(), "cn-a:6001") + t.Cleanup(func() { clearRemoteAllocationStatementTestTombstone(key) }) + + board := message.NewMessageBoard() + participant, err := acquireRemoteAllocationStatementParticipant( + key, + board, + 2, + nil, + ) + require.NoError(t, err) + terminal, err := participant.finish(errors.New("first fragment aborted")) + require.Error(t, err) + require.True(t, terminal.complete) + require.False(t, remoteAllocationStatementGroupRegistered(board)) + + lateBoard := message.NewMessageBoard() + _, err = acquireRemoteAllocationStatementParticipant(key, lateBoard, 2, nil) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + lateBoard.CloseAndDrain() + + retryKey := remoteAllocationStatementGroupKey(newRemoteExecutionID(), "cn-a:6001") + retryBoard := message.NewMessageBoard() + retry, err := acquireRemoteAllocationStatementParticipant( + retryKey, + retryBoard, + 1, + nil, + ) + require.NoError(t, err) + terminal, err = retry.finish(nil) + require.NoError(t, err) + require.True(t, terminal.complete) +} + func TestRemoteAllocationStatementGroupFailureCancelsActiveSibling(t *testing.T) { board := message.NewMessageBoard() canceled := make(chan error, 2) - first, err := acquireRemoteAllocationStatementParticipant( + first, err := acquireRemoteAllocationStatementTestParticipant( + t, board, 3, func(cause error) { canceled <- cause }, ) require.NoError(t, err) - second, err := acquireRemoteAllocationStatementParticipant( + second, err := acquireRemoteAllocationStatementTestParticipant( + t, board, 3, func(cause error) { canceled <- cause }, @@ -426,11 +487,13 @@ func TestRemoteAllocationStatementGroupExpirationWaitsForActiveFragment(t *testi firstCompile, firstAttempt, firstBuffer := newAttempt() secondCompile, secondAttempt, secondBuffer := newAttempt() canceled := make(chan error, 2) - first, err := acquireRemoteAllocationStatementParticipant( + first, err := acquireRemoteAllocationStatementTestParticipant( + t, board, 3, func(cause error) { canceled <- cause }, ) require.NoError(t, err) - second, err := acquireRemoteAllocationStatementParticipant( + second, err := acquireRemoteAllocationStatementTestParticipant( + t, board, 3, func(cause error) { canceled <- cause }, ) require.NoError(t, err) diff --git a/pkg/sql/compile/remoterunServer.go b/pkg/sql/compile/remoterunServer.go index 5ad79d63db0b3..5a3684331a337 100644 --- a/pkg/sql/compile/remoterunServer.go +++ b/pkg/sql/compile/remoterunServer.go @@ -403,6 +403,7 @@ func handlePipelineMessage(receiver *messageReceiverOnServer) (err error) { runCompile.addr, ) allocationParticipant, runErr = acquireRemoteAllocationStatementParticipant( + allocationGroupKey, runCompile.MessageBoard, expectedFragments, func(cause error) { From d13b9103c8e17f828e48134d35bffcc4b5adf482 Mon Sep 17 00:00:00 2001 From: aptend Date: Mon, 3 Aug 2026 01:07:08 +0800 Subject: [PATCH 54/61] fix: preserve vector bitmap extent on duplicate --- .../vector/allocation_account_test.go | 69 +++++++++++++++++++ pkg/container/vector/vector.go | 23 ++++--- 2 files changed, 83 insertions(+), 9 deletions(-) diff --git a/pkg/container/vector/allocation_account_test.go b/pkg/container/vector/allocation_account_test.go index 917104d404b46..787e063a3e316 100644 --- a/pkg/container/vector/allocation_account_test.go +++ b/pkg/container/vector/allocation_account_test.go @@ -358,6 +358,75 @@ func TestVectorAllocationAccountDupPreservesSparseBitmapRowDomain(t *testing.T) finalizeTestVectorAllocationAccount(t, state) } +func TestVectorAllocationAccountDupPreservesStaleBitmapExtent(t *testing.T) { + t.Run("flat null and grouping", func(t *testing.T) { + state := newTestVectorAllocationAccount(t, 8<<20, 16) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, vec.PreExtend(130, mp)) + require.NoError(t, vec.PreExtendBitmap(130, mp)) + for i := range 130 { + require.NoError(t, AppendFixed(vec, int64(i), false, mp)) + } + vec.SetNull(129) + vec.GetGrouping().Add(128) + vec.SetLength(1) + + dup, err := vec.Dup(mp) + require.NoError(t, err) + require.Equal(t, 1, dup.Length()) + require.True(t, dup.GetNulls().Contains(129)) + require.True(t, dup.GetGrouping().Contains(128)) + + dup.Free(mp) + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) + }) + + t.Run("constant grouping", func(t *testing.T) { + state := newTestVectorAllocationAccount(t, 8<<20, 16) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, AppendFixed(vec, int64(1), false, mp)) + vec.SetClass(CONSTANT) + require.NoError(t, vec.PreExtendGrouping(130, mp)) + vec.GetGrouping().Add(129) + vec.SetLength(1) + + dup, err := vec.Dup(mp) + require.NoError(t, err) + require.Equal(t, 1, dup.Length()) + require.True(t, dup.GetGrouping().Contains(129)) + + dup.Free(mp) + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) + }) + + t.Run("empty stale extent", func(t *testing.T) { + state := newTestVectorAllocationAccount(t, 8<<20, 16) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, AppendFixed(vec, int64(1), false, mp)) + require.NoError(t, vec.PreExtendBitmap(130, mp)) + vec.SetNull(129) + vec.GetGrouping().Add(129) + vec.UnsetNull(129) + vec.GetGrouping().Del(129) + require.Zero(t, vec.GetNulls().Count()) + require.Zero(t, vec.GetGrouping().Count()) + + dup, err := vec.Dup(mp) + require.NoError(t, err) + require.Equal(t, int64(130), dup.GetNulls().GetBitmap().Len()) + require.Equal(t, int64(130), dup.GetGrouping().GetBitmap().Len()) + + dup.Free(mp) + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) + }) +} + func TestVectorAllocationAccountBitmapShuffleAccountsScratch(t *testing.T) { state := newTestVectorAllocationAccount(t, 8<<20, 16) mp := mpool.MustNewZero() diff --git a/pkg/container/vector/vector.go b/pkg/container/vector/vector.go index cf840e4163764..98335a021c56c 100644 --- a/pkg/container/vector/vector.go +++ b/pkg/container/vector/vector.go @@ -1644,7 +1644,7 @@ func (v *Vector) dup( w.length = v.length if v.HasGrouping() { if err := w.ensureGroupingCapacity( - v.length, + max(v.length, int(v.GetGrouping().GetBitmap().Len())), mp, ); err != nil { w.Free(mp) @@ -1669,18 +1669,23 @@ func (v *Vector) dup( } dataLen *= v.length } - // Bitmap logical length only reaches the highest set bit, so it can be much - // shorter than a sparse vector. The duplicate must cover the complete row - // domain because allocation-free transforms (for example ordered Shrink) - // may address every output row. - if !v.GetNulls().EmptyByFlag() { - if err := w.ensureNullCapacity(v.length, mp); err != nil { + // A bitmap may be shorter than a sparse vector or longer than a reused vector + // that was shortened with SetLength. Preserve both the complete row domain + // and the source bitmap extent before InitWith copies its storage. + if v.GetNulls().GetBitmap().Len() > 0 { + if err := w.ensureNullCapacity( + max(v.length, int(v.GetNulls().GetBitmap().Len())), + mp, + ); err != nil { w.Free(mp) return nil, err } } - if !v.GetGrouping().EmptyByFlag() { - if err := w.ensureGroupingCapacity(v.length, mp); err != nil { + if v.GetGrouping().GetBitmap().Len() > 0 { + if err := w.ensureGroupingCapacity( + max(v.length, int(v.GetGrouping().GetBitmap().Len())), + mp, + ); err != nil { w.Free(mp) return nil, err } From 078d05ae6125056529126174bd7a5941f2abfe20 Mon Sep 17 00:00:00 2001 From: aptend Date: Mon, 3 Aug 2026 10:42:39 +0800 Subject: [PATCH 55/61] fix: close allocation and remote generation review gaps --- .../design/evidence/26459_local_validation.md | 52 ++++++++++++- pkg/common/hashmap/strhashmap.go | 14 ++-- pkg/common/hashmap/strhashmap_test.go | 78 +++++++++++++++++++ pkg/common/mpool/accounted_buffer.go | 12 +-- pkg/common/mpool/accounted_buffer_test.go | 48 ++++++++++++ .../remote_allocation_statement_group.go | 6 +- .../remote_allocation_statement_group_test.go | 38 +++++++++ 7 files changed, 233 insertions(+), 15 deletions(-) diff --git a/docs/design/evidence/26459_local_validation.md b/docs/design/evidence/26459_local_validation.md index 081d4400ee31a..0a68681e158f0 100644 --- a/docs/design/evidence/26459_local_validation.md +++ b/docs/design/evidence/26459_local_validation.md @@ -129,4 +129,54 @@ hash-map build/lookup, and spill scatter. The acceptance rule is no new per-row or per-allocation Go object in steady state and no material regression outside measurement noise. -Remote auto-test is deliberately not part of this validation cycle. +The local measurements are complemented by the distributed validation below. + +## Distributed workload evidence + +The final semantic head before the two review counterexample fixes, +`d13b9103c8`, completed the TPCH 100G and 1T TKE run +[`30758186183`](https://github.com/matrixorigin/mo-auto-test/actions/runs/30758186183). +The workflow built that commit, loaded the native fixtures in 11 seconds and +49 seconds, and compared every Q1-Q22 result with its golden result. No query +failed and the run reported no OOM or budget-admission error. + +The measured query-only totals were: + +| Workload | Candidate turns | Candidate average | Recent main average | Delta | +| --- | --- | ---: | ---: | ---: | +| TPCH 100G | 97.739 / 95.005 / 94.982 / 97.469 s | 96.298 s | 98.252 s | -1.99% | +| TPCH 1T | 1045.573 / 1039.371 s | 1042.472 s | 1027.439 s | +1.46% | + +The cited main result is job +[`91396824792`](https://github.com/matrixorigin/mo-nightly-regression/actions/runs/30708854656/job/91396824792): +100G turns were 103.550 / 95.127 / 97.321 / 97.010 seconds and 1T turns +were 1027.322 / 1027.557 seconds. These are adjacent runs of the same TKE +benchmark shape, not a simultaneous same-base A/B; the deltas establish that +the stabilized candidate is within normal workload variance, not a stronger +causal performance claim. Compared with the earlier regressed candidate run +`30738374292` (135.093 seconds for 100G and 1591.049 seconds for 1T), this head +recovered 28.7% and 34.5% respectively. + +The current review fixes after `d13b9103c8` are allocation-boundary and +late-RPC-lifetime corrections. They add no per-row work: existing-buffer grow +now passes the logical requirement to the allocator's single capacity policy, +and aborted remote generations retain only a key and timer for the maximum +possible RPC lifetime. Their focused and package validation is recorded in the +PR review response after the final commit. + +## Incident acceptance matrix + +This matrix separates durable mechanism regressions from workload executions; +one is not presented as a substitute for the other. + +| Incident | Durable regression retained on this branch | Workload evidence | Current-head gap | +| --- | --- | --- | --- | +| #26174 | HashBuild build/hashmap/spill regressions introduced by #26178, plus exact physical batch/vector allocation boundaries in this PR | #26178 TKE BVT: all three 3,840,001-row fulltext inserts succeeded with zero HashBuild rejection | full fulltext workload has not been rerun at the final head | +| #26192 | exact accounted runtime-filter payload, one-byte-short PASS degradation, varlena/null coverage, and spill decode/reuse lifecycle tests | historical LOAD failure shape is covered by #26231/#26318; the current TPCH fixture LOAD path succeeds | the original `ca_comprehensive_dataset` workload has not been rerun at the final head | +| #26413 | segmented `CopyIntoBatches` and accounted hash-map growth/rollback regressions, including large external-batch shapes | #26438 verified the real Parquet self-join with both expected 50,000-row results | the Hive fixture has not been rerun at the final head | +| #26454 | `TestIssue26454ExpressionKeyBuildUsesActualCapacity` exercises the CONCAT/CAST and CASE key shapes under a 16 MiB physical account and validates terminal zero | the exact jinpan SQL/data is not available in this repository | full jinpan workload remains external evidence | +| #25782 | `TestShuffleHashBuildAccountedSpillLifecycle`, `TestHashTableAccountedHighCardinalityResizeReturnsToZero`, broadcast error propagation, recursive spill, and terminal-zero tests | the two-CN 132,096-row harness at `f5cc97efe7` returned the exact count with positive spill and zero OOM; current-head TPCH 1T also completed without OOM/query failure | the private original high-cardinality SQL harness has not been rerun at the final head | + +Accordingly, the current TKE TPCH acceptance is complete, while the unavailable +external-data workloads and the original private #25782 harness remain explicit +follow-up evidence rather than being silently marked complete. diff --git a/pkg/common/hashmap/strhashmap.go b/pkg/common/hashmap/strhashmap.go index e0c14ebefc06c..25c33cbe4b9b1 100644 --- a/pkg/common/hashmap/strhashmap.go +++ b/pkg/common/hashmap/strhashmap.go @@ -208,17 +208,17 @@ func (itr *strHashmapIterator) prepareHashKeys( } if cap(itr.keyBuffer) < total { if allocation := itr.mp.iteratorAllocation; allocation != nil { - capacity, ok := mpool.GrowCapacity( - int64(cap(itr.keyBuffer)), int64(total), - ) - if !ok || int64(int(capacity)) != capacity { - return mpool.ErrAllocationAllocatorLimit - } var next []byte var err error if cap(itr.keyBuffer) > 0 { - next, err = itr.mp.mp.Grow(itr.keyBuffer, int(capacity), true) + // Grow owns the capacity policy. Passing a pre-grown capacity + // would apply that policy twice and falsely inflate admission. + next, err = itr.mp.mp.Grow(itr.keyBuffer, total, true) } else { + capacity, ok := mpool.GrowCapacity(0, int64(total)) + if !ok || int64(int(capacity)) != capacity { + return mpool.ErrAllocationAllocatorLimit + } next, err = itr.mp.mp.AllocAccounted( int(capacity), allocation.account, diff --git a/pkg/common/hashmap/strhashmap_test.go b/pkg/common/hashmap/strhashmap_test.go index b671c9b4ea97f..3359c71152b65 100644 --- a/pkg/common/hashmap/strhashmap_test.go +++ b/pkg/common/hashmap/strhashmap_test.go @@ -264,6 +264,84 @@ func TestStringHashMapCanonicalizesFullyGroupedKeys(t *testing.T) { require.Equal(t, []int64{1}, zValues) } +func TestStringHashIteratorAccountedGrowthCapacityBoundary(t *testing.T) { + const ( + oldCapacity = 10_240 + required = oldCapacity + 1 + payloadSize = required - 4 + ) + newCapacity, ok := mpool.GrowCapacity(oldCapacity, required) + require.True(t, ok) + exactLimit := uint64(oldCapacity) + uint64(newCapacity) + + for _, testCase := range []struct { + name string + limit uint64 + wantError bool + }{ + {name: "exact-old-plus-rounded-new", limit: exactLimit}, + {name: "one-byte-short", limit: exactLimit - 1, wantError: true}, + } { + t.Run(testCase.name, func(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 2) + require.NoError(t, err) + account, err := registry.Open(testCase.limit) + require.NoError(t, err) + allocation, err := NewIteratorAllocation( + account, + mpool.AllocationOwnerMin, + mpool.AllocationSiteMin, + ) + require.NoError(t, err) + mp := mpool.MustNewZero() + hashMap, err := NewStrHashMapWithAllocations( + false, + mp, + nil, + allocation, + ) + require.NoError(t, err) + iterator := hashMap.NewIterator().(*strHashmapIterator) + iterator.keyBuffer, err = mp.AllocAccounted( + oldCapacity, + account, + mpool.AllocationOwnerMin, + mpool.AllocationSiteMin, + ) + require.NoError(t, err) + iterator.keyBuffer = iterator.keyBuffer[:0] + vec := vector.NewVec(types.T_varchar.ToType()) + require.NoError(t, vector.AppendBytes( + vec, + make([]byte, payloadSize), + false, + mp, + )) + + err = iterator.prepareHashKeys([]*vector.Vector{vec}, 0, 1) + if testCase.wantError { + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Equal(t, oldCapacity, cap(iterator.keyBuffer)) + require.Equal(t, uint64(oldCapacity), account.Snapshot().Used) + } else { + require.NoError(t, err) + require.Equal(t, int(newCapacity), cap(iterator.keyBuffer)) + require.Equal(t, uint64(newCapacity), account.Snapshot().Used) + require.Equal(t, exactLimit, account.Snapshot().Peak) + } + + iterator.releaseScratch() + hashMap.Free() + vec.Free(mp) + require.Zero(t, mp.Stats().NumCurrBytes.Load()) + require.Zero(t, account.Seal().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + _, err = registry.Finalize(account) + require.NoError(t, err) + }) + } +} + func TestGroupingAwareStringHashMapSeparatesRawSentinelBytes(t *testing.T) { mp := mpool.MustNewZero() hashMap, err := NewStrHashMap(false, mp) diff --git a/pkg/common/mpool/accounted_buffer.go b/pkg/common/mpool/accounted_buffer.go index c7a1f00550999..eca8325d5590a 100644 --- a/pkg/common/mpool/accounted_buffer.go +++ b/pkg/common/mpool/accounted_buffer.go @@ -89,11 +89,11 @@ func (b *AccountedBuffer) EnsureCapacity(required int) error { } oldLength := len(b.data) - capacity, ok := GrowCapacity(int64(cap(b.data)), int64(required)) - if !ok || capacity > int64(math.MaxInt) { - return ErrAllocationAllocatorLimit - } if cap(b.data) == 0 { + capacity, ok := GrowCapacity(0, int64(required)) + if !ok || capacity > int64(math.MaxInt) { + return ErrAllocationAllocatorLimit + } data, err := b.mp.AllocAccounted( int(capacity), b.account, @@ -107,7 +107,9 @@ func (b *AccountedBuffer) EnsureCapacity(required int) error { return nil } - data, err := b.mp.Grow(b.data, int(capacity), true) + // Grow owns the capacity policy. Pass the caller's requirement instead of + // applying GrowCapacity a second time to the already rounded capacity. + data, err := b.mp.Grow(b.data, required, true) if err != nil { return err } diff --git a/pkg/common/mpool/accounted_buffer_test.go b/pkg/common/mpool/accounted_buffer_test.go index 855cf9e8409a2..c70765a9f48f5 100644 --- a/pkg/common/mpool/accounted_buffer_test.go +++ b/pkg/common/mpool/accounted_buffer_test.go @@ -98,6 +98,54 @@ func TestAccountedBufferFailureRetainsPublishedData(t *testing.T) { finalizeTestAllocationAccount(t, registry, account) } +func TestAccountedBufferGrowthCapacityBoundary(t *testing.T) { + const ( + oldCapacity = 10_240 + required = oldCapacity + 1 + ) + newCapacity, ok := GrowCapacity(oldCapacity, required) + require.True(t, ok) + exactLimit := uint64(oldCapacity) + uint64(newCapacity) + + for _, testCase := range []struct { + name string + limit uint64 + wantError bool + }{ + {name: "exact-old-plus-rounded-new", limit: exactLimit}, + {name: "one-byte-short", limit: exactLimit - 1, wantError: true}, + } { + t.Run(testCase.name, func(t *testing.T) { + registry, account := newTestAllocationAccount(t, testCase.limit, 2) + mp := MustNew("accounted-buffer-growth-boundary") + defer DeleteMPool(mp) + buffer, err := NewAccountedBuffer( + mp, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + require.NoError(t, buffer.EnsureCapacity(oldCapacity)) + + err = buffer.EnsureCapacity(required) + if testCase.wantError { + require.ErrorIs(t, err, ErrAllocationAccountCapacity) + require.Equal(t, oldCapacity, buffer.Cap()) + require.Equal(t, uint64(oldCapacity), account.Snapshot().Used) + } else { + require.NoError(t, err) + require.Equal(t, int(newCapacity), buffer.Cap()) + require.Equal(t, uint64(newCapacity), account.Snapshot().Used) + require.Equal(t, exactLimit, account.Snapshot().Peak) + } + + buffer.Free() + finalizeTestAllocationAccount(t, registry, account) + }) + } +} + func TestAccountedBufferConfiguration(t *testing.T) { _, err := NewAccountedBuffer(nil, nil, 0, 0) require.ErrorIs(t, err, ErrAllocationAccountInvalid) diff --git a/pkg/sql/compile/remote_allocation_statement_group.go b/pkg/sql/compile/remote_allocation_statement_group.go index 2a2f86734e863..1ffc81c6e87f2 100644 --- a/pkg/sql/compile/remote_allocation_statement_group.go +++ b/pkg/sql/compile/remote_allocation_statement_group.go @@ -60,9 +60,11 @@ func remoteAllocationStatementGroupKey( var remoteAllocationStatementRegistrationTimeout = 5 * time.Minute // A late RPC can carry a new MessageBoard, so keep a bounded record of an -// incomplete execution after its active group has been released. The record +// incomplete execution after its active group has been released. A sender +// without a caller deadline can remain in flight for MaxRpcTime; expiring the +// key earlier would let that same physical generation reopen. The record // contains no statement resources. -var remoteAllocationStatementTombstoneTimeout = 5 * time.Minute +const remoteAllocationStatementTombstoneTimeout = MaxRpcTime // collectRemoteFragmentCounts computes the number of pipeline RPCs that the // complete physical scope graph will send to each CN. The execution address diff --git a/pkg/sql/compile/remote_allocation_statement_group_test.go b/pkg/sql/compile/remote_allocation_statement_group_test.go index cf148a8f77809..0ae51c378c9b0 100644 --- a/pkg/sql/compile/remote_allocation_statement_group_test.go +++ b/pkg/sql/compile/remote_allocation_statement_group_test.go @@ -416,6 +416,44 @@ func TestRemoteAllocationStatementGroupRejectsLateFragmentAfterAbort(t *testing. require.True(t, terminal.complete) } +func TestRemoteAllocationStatementGroupExpiryRejectsLateFragment(t *testing.T) { + require.Equal(t, MaxRpcTime, remoteAllocationStatementTombstoneTimeout) + key := remoteAllocationStatementGroupKey(newRemoteExecutionID(), "cn-a:6001") + t.Cleanup(func() { clearRemoteAllocationStatementTestTombstone(key) }) + + board := message.NewMessageBoard() + participant, err := acquireRemoteAllocationStatementParticipant( + key, + board, + 2, + nil, + ) + require.NoError(t, err) + expireRemoteAllocationStatementGroup(participant.group) + terminal, err := participant.finish(nil) + require.Error(t, err) + require.True(t, terminal.complete) + require.False(t, remoteAllocationStatementGroupRegistered(board)) + + lateBoard := message.NewMessageBoard() + _, err = acquireRemoteAllocationStatementParticipant(key, lateBoard, 2, nil) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + lateBoard.CloseAndDrain() + + retryKey := remoteAllocationStatementGroupKey(newRemoteExecutionID(), "cn-a:6001") + retryBoard := message.NewMessageBoard() + retry, err := acquireRemoteAllocationStatementParticipant( + retryKey, + retryBoard, + 1, + nil, + ) + require.NoError(t, err) + terminal, err = retry.finish(nil) + require.NoError(t, err) + require.True(t, terminal.complete) +} + func TestRemoteAllocationStatementGroupFailureCancelsActiveSibling(t *testing.T) { board := message.NewMessageBoard() canceled := make(chan error, 2) From f1bce32098643ae9256ee15597c585947f9447b0 Mon Sep 17 00:00:00 2001 From: aptend Date: Mon, 3 Aug 2026 11:51:51 +0800 Subject: [PATCH 56/61] fix: preserve statement errors through allocation cleanup --- .../compile/allocation_account_lifecycle.go | 30 ++++++++++++++---- .../allocation_account_lifecycle_test.go | 14 +++++++++ pkg/sql/compile/compile2.go | 5 ++- .../remote_allocation_statement_group.go | 31 ++++++++++--------- .../remote_allocation_statement_group_test.go | 19 ++++++++++++ pkg/sql/compile/remoterunServer.go | 14 ++++----- 6 files changed, 83 insertions(+), 30 deletions(-) diff --git a/pkg/sql/compile/allocation_account_lifecycle.go b/pkg/sql/compile/allocation_account_lifecycle.go index b8bc959b53e51..02986242bc587 100644 --- a/pkg/sql/compile/allocation_account_lifecycle.go +++ b/pkg/sql/compile/allocation_account_lifecycle.go @@ -16,6 +16,7 @@ package compile import ( "errors" + "reflect" "sync" "github.com/matrixorigin/matrixone/pkg/common/moerr" @@ -40,6 +41,23 @@ func allocationLifecycleCall(call func() error) (err error) { return call() } +// joinAllocationLifecycleErrors keeps a lone failure's concrete type intact +// and avoids rejoining the same terminal failure. errors.Join wraps even one +// non-nil error, which would turn a statement *moerr.Error into a generic Go +// error before it crosses the pipeline wire. +func joinAllocationLifecycleErrors(primary, secondary error) error { + if primary == nil { + return secondary + } + if secondary == nil { + return primary + } + if reflect.TypeOf(primary).Comparable() && primary == secondary { + return primary + } + return errors.Join(primary, secondary) +} + type executionAllocationAccountOwner interface { SetAllocationAccount(*mpool.AllocationAccount) error ClearAllocationAccount(*mpool.AllocationAccount) error @@ -130,7 +148,7 @@ func (c *Compile) beginAllocationAccountAttempt() ( return terminalErr }) if first { - finalizeErr = errors.Join( + finalizeErr = joinAllocationLifecycleErrors( finalizeErr, allocationLifecycleCall(func() error { c.allocationTerminalExporter(snapshot) @@ -139,7 +157,7 @@ func (c *Compile) beginAllocationAccountAttempt() ( ) } if finalizeErr != nil { - return nil, errors.Join(err, finalizeErr) + return nil, joinAllocationLifecycleErrors(err, finalizeErr) } return nil, err } @@ -216,7 +234,7 @@ func configureAllocationAccountOwners( configured := make([]executionAllocationAccountOwner, 0, len(owners)) rollback := func(cause error) error { for i := len(configured) - 1; i >= 0; i-- { - cause = errors.Join( + cause = joinAllocationLifecycleErrors( cause, allocationLifecycleCall(func() error { return configured[i].ClearAllocationAccount(account) @@ -347,7 +365,7 @@ func (a *statementAllocationAttempt) prepareTerminal(closeBoard bool) error { } a.prepareOnce.Do(func() { if closeBoard { - a.prepareErr = errors.Join( + a.prepareErr = joinAllocationLifecycleErrors( a.prepareErr, allocationLifecycleCall(func() error { a.board.CloseAndDrain() @@ -362,7 +380,7 @@ func (a *statementAllocationAttempt) prepareTerminal(closeBoard bool) error { a.ownerSet = nil a.ownersMu.Unlock() for i := len(owners) - 1; i >= 0; i-- { - a.prepareErr = errors.Join( + a.prepareErr = joinAllocationLifecycleErrors( a.prepareErr, allocationLifecycleCall(func() error { return owners[i].ClearAllocationAccount(a.account) @@ -393,7 +411,7 @@ func (a *statementAllocationAttempt) completeTerminal() ( return terminalErr }) if first && a.exporter != nil { - a.completeErr = errors.Join( + a.completeErr = joinAllocationLifecycleErrors( a.completeErr, allocationLifecycleCall(func() error { a.exporter(a.snapshot) diff --git a/pkg/sql/compile/allocation_account_lifecycle_test.go b/pkg/sql/compile/allocation_account_lifecycle_test.go index ceacda7824dfa..3ffb90ea636a0 100644 --- a/pkg/sql/compile/allocation_account_lifecycle_test.go +++ b/pkg/sql/compile/allocation_account_lifecycle_test.go @@ -16,6 +16,7 @@ package compile import ( "context" + "errors" "sync" "sync/atomic" "testing" @@ -40,6 +41,19 @@ import ( "github.com/stretchr/testify/require" ) +func TestJoinAllocationLifecycleErrorsPreservesSingle(t *testing.T) { + primary := moerr.NewDuplicateEntryNoCtx("duplicate", "primary") + secondary := errors.New("cleanup failed") + + require.Same(t, primary, joinAllocationLifecycleErrors(primary, nil)) + require.Same(t, secondary, joinAllocationLifecycleErrors(nil, secondary)) + require.Same(t, primary, joinAllocationLifecycleErrors(primary, primary)) + + joined := joinAllocationLifecycleErrors(primary, secondary) + require.ErrorIs(t, joined, primary) + require.ErrorIs(t, joined, secondary) +} + type allocationLifecycleErrorOperator struct { *colexec.MockOperator err error diff --git a/pkg/sql/compile/compile2.go b/pkg/sql/compile/compile2.go index 4def9d28eb934..59c82c6939d9b 100644 --- a/pkg/sql/compile/compile2.go +++ b/pkg/sql/compile/compile2.go @@ -17,7 +17,6 @@ package compile import ( "context" "encoding/hex" - "errors" "math" gotrace "runtime/trace" "strings" @@ -356,7 +355,7 @@ func (c *Compile) Run(_ uint64) (queryResult *util2.RunResult, err error) { coordinatorPhaseStart = time.Time{} coordinatorPhaseBase = 0 if terminalErr := finishAllocationAttempt(); terminalErr != nil { - err = errors.Join(err, terminalErr) + err = joinAllocationLifecycleErrors(err, terminalErr) resourceRecorder.finishAttempt( uint64(retryTimes), attemptStart, preRunWall, attemptRemoteWait, stats, attemptScopes, attemptAnal, c.addr, false, @@ -500,7 +499,7 @@ func (c *Compile) Run(_ uint64) (queryResult *util2.RunResult, err error) { // this call returns. c.AnalyzeExecPlan(runC, queryResult, stats, isExplainPhyPlan, option) if terminalErr := finishAllocationAttempt(); terminalErr != nil { - err = errors.Join(err, terminalErr) + err = joinAllocationLifecycleErrors(err, terminalErr) resourceRecorder.finishAttempt( uint64(retryTimes), attemptStart, attemptPreRunWall, attemptRemoteWait, stats, attemptScopes, attemptAnal, c.addr, false, diff --git a/pkg/sql/compile/remote_allocation_statement_group.go b/pkg/sql/compile/remote_allocation_statement_group.go index 1ffc81c6e87f2..5547b53df7bd0 100644 --- a/pkg/sql/compile/remote_allocation_statement_group.go +++ b/pkg/sql/compile/remote_allocation_statement_group.go @@ -289,13 +289,16 @@ func (p *remoteAllocationStatementParticipant) finish(cause error) ( if group.finalized || (!group.expired && group.finished >= group.expected) || (group.expired && group.finished >= group.registered) { - p.err = errors.Join(p.err, mpool.ErrAllocationAccountInvariant) + p.err = joinAllocationLifecycleErrors( + p.err, + mpool.ErrAllocationAccountInvariant, + ) remoteAllocationStatementGroups.Unlock() return } abort := cause != nil && !group.expired if cause != nil { - group.err = errors.Join(group.err, cause) + group.err = joinAllocationLifecycleErrors(group.err, cause) group.expired = true if group.timer != nil { group.timer.Stop() @@ -323,15 +326,15 @@ func (p *remoteAllocationStatementParticipant) finish(cause error) ( group.board.Close() return nil }) - abortErr = errors.Join( + abortErr = joinAllocationLifecycleErrors( abortErr, cancelRemoteAllocationStatementParticipants(cancels, cause), ) if abortErr != nil { remoteAllocationStatementGroups.Lock() - group.err = errors.Join(group.err, abortErr) + group.err = joinAllocationLifecycleErrors(group.err, abortErr) remoteAllocationStatementGroups.Unlock() - p.err = errors.Join(p.err, abortErr) + p.err = joinAllocationLifecycleErrors(p.err, abortErr) } } } else { @@ -345,7 +348,7 @@ func (p *remoteAllocationStatementParticipant) finish(cause error) ( pools, terminalErr, ) - p.err = errors.Join(p.err, terminalErr) + p.err = joinAllocationLifecycleErrors(p.err, terminalErr) } }) return p.terminal, p.err @@ -369,7 +372,7 @@ func cancelRemoteAllocationStatementParticipants( var err error for _, cancel := range cancels { if cancel != nil { - err = errors.Join(err, allocationLifecycleCall(func() error { + err = joinAllocationLifecycleErrors(err, allocationLifecycleCall(func() error { cancel(cause) return nil })) @@ -435,7 +438,7 @@ func completeRemoteAllocationStatementGroup( remoteAllocationStatementTerminal, error, ) { - terminalErr = errors.Join( + terminalErr = joinAllocationLifecycleErrors( terminalErr, allocationLifecycleCall(func() error { group.board.CloseAndDrain() @@ -453,10 +456,10 @@ func completeRemoteAllocationStatementGroup( for _, attempt := range attempts { snapshot, err := attempt.completeTerminal() terminal.allocation = append(terminal.allocation, snapshot) - terminalErr = errors.Join(terminalErr, err) + terminalErr = joinAllocationLifecycleErrors(terminalErr, err) } for _, pool := range pools { - terminalErr = errors.Join( + terminalErr = joinAllocationLifecycleErrors( terminalErr, allocationLifecycleCall(func() error { domain, quality := pool.ResourceSnapshot() @@ -481,7 +484,7 @@ func expireRemoteAllocationStatementGroup( timeoutErr := moerr.NewInternalErrorNoCtx( "remote allocation statement group registration timed out", ) - group.err = errors.Join(group.err, timeoutErr) + group.err = joinAllocationLifecycleErrors(group.err, timeoutErr) expected, registered, finished := group.expected, group.registered, group.finished cancels := activeRemoteAllocationStatementCancelsLocked(group) var attempts []*statementAllocationAttempt @@ -502,7 +505,7 @@ func expireRemoteAllocationStatementGroup( terminalErr, ) } else { - terminalErr = errors.Join( + terminalErr = joinAllocationLifecycleErrors( terminalErr, allocationLifecycleCall(func() error { group.board.Close() @@ -511,11 +514,11 @@ func expireRemoteAllocationStatementGroup( ) } cancelErr := cancelRemoteAllocationStatementParticipants(cancels, timeoutErr) - terminalErr = errors.Join(terminalErr, cancelErr) + terminalErr = joinAllocationLifecycleErrors(terminalErr, cancelErr) if !complete && cancelErr != nil { remoteAllocationStatementGroups.Lock() if !group.finalized { - group.err = errors.Join(group.err, cancelErr) + group.err = joinAllocationLifecycleErrors(group.err, cancelErr) } remoteAllocationStatementGroups.Unlock() } diff --git a/pkg/sql/compile/remote_allocation_statement_group_test.go b/pkg/sql/compile/remote_allocation_statement_group_test.go index 0ae51c378c9b0..e76a6b35705a5 100644 --- a/pkg/sql/compile/remote_allocation_statement_group_test.go +++ b/pkg/sql/compile/remote_allocation_statement_group_test.go @@ -15,6 +15,7 @@ package compile import ( + "context" "errors" "fmt" "strings" @@ -22,6 +23,7 @@ import ( "testing" "time" + "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/sql/colexec" "github.com/matrixorigin/matrixone/pkg/sql/colexec/connector" @@ -221,6 +223,23 @@ func TestRemoteAllocationStatementGroupDefersSharedBoardTerminal(t *testing.T) { require.False(t, remoteAllocationStatementGroupRegistered(board)) } +func TestRemoteAllocationStatementGroupPreservesStatementError(t *testing.T) { + board := message.NewMessageBoard() + participant, err := acquireRemoteAllocationStatementTestParticipant( + t, + board, + 1, + nil, + ) + require.NoError(t, err) + + primary := moerr.NewDuplicateEntryNoCtx("duplicate", "primary") + terminal, err := participant.finish(primary) + require.True(t, terminal.complete) + require.Same(t, primary, err) + require.Same(t, primary, moerr.ConvertGoError(context.Background(), err)) +} + func TestRemoteAllocationStatementGroupRejectsTopologyMismatch(t *testing.T) { board := message.NewMessageBoard() participant, err := acquireRemoteAllocationStatementTestParticipant(t, board, 2, nil) diff --git a/pkg/sql/compile/remoterunServer.go b/pkg/sql/compile/remoterunServer.go index 5a3684331a337..680d31f7050b9 100644 --- a/pkg/sql/compile/remoterunServer.go +++ b/pkg/sql/compile/remoterunServer.go @@ -277,13 +277,13 @@ func handlePipelineMessage(receiver *messageReceiverOnServer) (err error) { allocationParticipant.stage(allocationAttempt, memoryPool) cause := err if recovered != nil { - cause = errors.Join( + cause = joinAllocationLifecycleErrors( cause, moerr.ConvertPanicError(receiver.messageCtx, recovered), ) } _, terminalErr := allocationParticipant.finish(cause) - err = errors.Join(err, terminalErr) + err = joinAllocationLifecycleErrors(err, terminalErr) participantFinished = true } if recovered != nil { @@ -292,7 +292,7 @@ func handlePipelineMessage(receiver *messageReceiverOnServer) (err error) { }() defer func() { if recovered := recover(); recovered != nil { - err = errors.Join( + err = joinAllocationLifecycleErrors( err, moerr.ConvertPanicError(receiver.messageCtx, recovered), ) @@ -305,7 +305,7 @@ func handlePipelineMessage(receiver *messageReceiverOnServer) (err error) { var localDelta resource.Delta var descendant remoteResourceSnapshot var expectedDirect uint64 - err = errors.Join(err, allocationLifecycleCall(func() error { + err = joinAllocationLifecycleErrors(err, allocationLifecycleCall(func() error { localDelta = collectScopeResourceDelta( runCompile.scopes, receiver.cnInformation.cnAddr, @@ -322,7 +322,7 @@ func handlePipelineMessage(receiver *messageReceiverOnServer) (err error) { runCompile.allocationAttempt = nil } allocationParticipant.stage(allocationAttempt, memoryPool) - err = errors.Join(err, allocationLifecycleCall(func() error { + err = joinAllocationLifecycleErrors(err, allocationLifecycleCall(func() error { if statementGroupEnabled { // The remote statement group, rather than any one fragment Compile, // owns the shared multi-CN board. Detach it before clear resets the @@ -339,7 +339,7 @@ func handlePipelineMessage(receiver *messageReceiverOnServer) (err error) { })) terminal, terminalErr := allocationParticipant.finish(err) participantFinished = true - err = errors.Join(err, terminalErr) + err = joinAllocationLifecycleErrors(err, terminalErr) for _, snapshot := range terminal.allocation { localAllocationQuality |= localAllocation.AddGeneration( snapshot.Peak, @@ -350,7 +350,7 @@ func handlePipelineMessage(receiver *messageReceiverOnServer) (err error) { var localMemory resource.MemoryDomainSummary var localMemoryQuality resource.QualityFlags if !statementGroupEnabled { - err = errors.Join(err, allocationLifecycleCall(func() error { + err = joinAllocationLifecycleErrors(err, allocationLifecycleCall(func() error { localMemory, localMemoryQuality = memoryPool.ResourceSnapshot() return nil })) From 60ae274c411dda6b80a47657ba52f8517ca00b8d Mon Sep 17 00:00:00 2001 From: aptend Date: Mon, 3 Aug 2026 13:52:51 +0800 Subject: [PATCH 57/61] fix vector unmarshal nil checks --- pkg/container/vector/vector.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pkg/container/vector/vector.go b/pkg/container/vector/vector.go index 8be712a872604..c7dadb634f3db 100644 --- a/pkg/container/vector/vector.go +++ b/pkg/container/vector/vector.go @@ -1172,10 +1172,10 @@ func decodeVectorBinaryLayout( } func (v *Vector) unmarshalBinary(data []byte, validateValues bool) error { - v.areaDisjoint = false if v == nil { return io.ErrClosedPipe } + v.areaDisjoint = false if v.allocationAccount != nil { return allocationAccountInvalid( "cannot install aliases in an accounted vector", @@ -1330,10 +1330,10 @@ func canonicalVectorTypeSize(typ types.Type) (int, error) { } func (v *Vector) UnmarshalBinaryWithCopy(data []byte, mp *mpool.MPool) error { - v.areaDisjoint = false if v == nil || mp == nil { return io.ErrClosedPipe } + v.areaDisjoint = false if v.hasBackingStorage() { return allocationAccountInvalid( "cannot replace vector storage without Free", @@ -1393,10 +1393,10 @@ func (v *Vector) UnmarshalBinaryWithCopy(data []byte, mp *mpool.MPool) error { } func (v *Vector) UnmarshalWithReader(r io.Reader, mp *mpool.MPool) error { - v.areaDisjoint = false if v == nil || r == nil { return io.ErrClosedPipe } + v.areaDisjoint = false v.ResetWithSameType() var err error From cd36bbbe776e68267218f0df1f8d81c274691bb4 Mon Sep 17 00:00:00 2001 From: aptend Date: Mon, 3 Aug 2026 14:35:15 +0800 Subject: [PATCH 58/61] fix(hashbuild): admit direct spill by allocation --- ...ocation_accounted_memory_admission_impl.md | 15 +++- .../design/evidence/26459_local_validation.md | 6 +- pkg/sql/colexec/hashbuild/build.go | 43 ++++----- pkg/sql/colexec/hashbuild/build_test.go | 87 +++++++++++++++++++ pkg/sql/colexec/hashbuild/spill.go | 12 +++ pkg/sql/colexec/hashbuild/types.go | 8 +- 6 files changed, 140 insertions(+), 31 deletions(-) diff --git a/docs/design/allocation_accounted_memory_admission_impl.md b/docs/design/allocation_accounted_memory_admission_impl.md index 901f94611dce6..315b968459b22 100644 --- a/docs/design/allocation_accounted_memory_admission_impl.md +++ b/docs/design/allocation_accounted_memory_admission_impl.md @@ -17,9 +17,15 @@ The implementation has one production path: 3. the account charges the query generation and the CN aggregate controller; 4. the same allocation releases the charge when MPool frees it. -There is no feature switch, activation gate, estimated-memory reservation, or -parallel compatibility ledger. Cardinality estimates may still select a hash -table capacity, but they never create a separately releasable memory charge. +There is no feature switch, activation gate, or parallel compatibility ledger. +Physical storage is charged only by its allocation. A shuffle HashBuild may +additionally reserve a conservative recovery floor before accepting spillable +retained state; that floor proves the retained state can be drained under a +shared budget. A capacity failure while growing the floor changes the decision +from retain to direct spill; lifecycle and invariant failures remain terminal. +The projection is never reapplied to an upstream-owned direct source as a +query-fatal estimate. Cardinality estimates may likewise select a hash-table +capacity without creating a separately releasable physical allocation charge. ## Scope @@ -216,7 +222,8 @@ one iterator allocation per row. The implementation is complete only when: -- production has no estimated HashBuild memory reservation or activation gate; +- no estimated HashBuild value is a query-fatal admission gate for unretained + work; the recovery projection is limited to the retain-versus-spill decision; - every retained HashBuild/join allocation in the controlled domain has immutable provenance; - runtime parallel clones join the current attempt before `Prepare`; diff --git a/docs/design/evidence/26459_local_validation.md b/docs/design/evidence/26459_local_validation.md index 0a68681e158f0..9b39451504576 100644 --- a/docs/design/evidence/26459_local_validation.md +++ b/docs/design/evidence/26459_local_validation.md @@ -7,7 +7,8 @@ are not retained as validation dimensions. ## Static closure checks - no production allocation-account enable switch; -- no HashBuild logical-size memory reservation token; +- no logical-size estimate is a query-fatal gate for unretained HashBuild work; + the shuffle recovery floor is used only before retaining spillable state; - every join/HashBuild expression-owned MPool vector is constructed with the attempt account, while opaque library Go heap remains an explicit boundary; - SpillEngine construction rejects a missing or closed budget generation; @@ -174,7 +175,8 @@ one is not presented as a substitute for the other. | #26174 | HashBuild build/hashmap/spill regressions introduced by #26178, plus exact physical batch/vector allocation boundaries in this PR | #26178 TKE BVT: all three 3,840,001-row fulltext inserts succeeded with zero HashBuild rejection | full fulltext workload has not been rerun at the final head | | #26192 | exact accounted runtime-filter payload, one-byte-short PASS degradation, varlena/null coverage, and spill decode/reuse lifecycle tests | historical LOAD failure shape is covered by #26231/#26318; the current TPCH fixture LOAD path succeeds | the original `ca_comprehensive_dataset` workload has not been rerun at the final head | | #26413 | segmented `CopyIntoBatches` and accounted hash-map growth/rollback regressions, including large external-batch shapes | #26438 verified the real Parquet self-join with both expected 50,000-row results | the Hive fixture has not been rerun at the final head | -| #26454 | `TestIssue26454ExpressionKeyBuildUsesActualCapacity` exercises the CONCAT/CAST and CASE key shapes under a 16 MiB physical account and validates terminal zero | the exact jinpan SQL/data is not available in this repository | full jinpan workload remains external evidence | +| #26454 | `TestIssue26454ExpressionKeyBuildUsesActualCapacity` exercises the CONCAT/CAST and CASE key shapes under a 16 MiB physical account and validates terminal zero; `TestShuffleHashBuildDirectSpillUsesActualAllocation` proves the conservative recovery projection cannot reject an upstream-owned direct source | the exact jinpan SQL/data is not available in this repository | full jinpan workload remains external evidence | +| #26586 | retained shuffle batches reserve recovery ownership before copy, while direct spill uses allocation-led admission; the focused direct-spill regression and full HashBuild race suite pass | `TestHashBuildSharedBudgetRecoverySQL` passes locally with the 28 MiB shared-budget reproducer | the external TPCH Q9 workload is not rerun for this control-flow-only correction | | #25782 | `TestShuffleHashBuildAccountedSpillLifecycle`, `TestHashTableAccountedHighCardinalityResizeReturnsToZero`, broadcast error propagation, recursive spill, and terminal-zero tests | the two-CN 132,096-row harness at `f5cc97efe7` returned the exact count with positive spill and zero OOM; current-head TPCH 1T also completed without OOM/query failure | the private original high-cardinality SQL harness has not been rerun at the final head | Accordingly, the current TKE TPCH acceptance is complete, while the unavailable diff --git a/pkg/sql/colexec/hashbuild/build.go b/pkg/sql/colexec/hashbuild/build.go index 53cac9f706d13..065b5a04afbb9 100644 --- a/pkg/sql/colexec/hashbuild/build.go +++ b/pkg/sql/colexec/hashbuild/build.go @@ -256,20 +256,6 @@ func (hashBuild *HashBuild) build( analyzer, ) } - ensureDirectRecovery := func(bat *batch.Batch) error { - if bat == nil || bat.RowCount() <= 0 { - return nil - } - selected, err := projectedSelectedRange(bat, 0, bat.RowCount()) - if err != nil { - return err - } - return ensureRecovery(recoveryBatchProjection{ - maxRows: bat.RowCount(), - maxSelected: selected, - columns: len(bat.Vecs), - }) - } defer func() { observeHashBuildBudget(analyzer, ctr.hashmapBuilder.budget) @@ -316,10 +302,6 @@ func (hashBuild *HashBuild) build( hashBuild.JoinMapRefCnt, ) } - execs, err := ctr.initSpillExprExecs(proc, hashBuild.Conditions) - if err != nil { - return err - } if spillFiles == nil { spillFiles = make([]*os.File, spillNumBuckets) } @@ -328,6 +310,14 @@ func (hashBuild *HashBuild) build( // Drain retained copies oldest-first. Each successful partition is // followed immediately by reservation and mpool release, so the source // batch and one partition scratch are the only simultaneous peaks. + execs := ctr.spillExprExecs + if len(ctr.hashmapBuilder.Batches.Buf) > 0 { + var err error + execs, err = ctr.initSpillExprExecs(proc, hashBuild.Conditions) + if err != nil { + return err + } + } for len(ctr.hashmapBuilder.Batches.Buf) > 0 { if err := checkHashBuildCanceled(proc); err != nil { return err @@ -346,6 +336,16 @@ func (hashBuild *HashBuild) build( return err } } + // No retained state remains after the drain. Drop every mandatory + // recovery-class borrower before returning the conservative floor, then + // let direct sources use ordinary allocation-led admission. + ctr.dropMandatorySpillRecoveryScratch() + if err := hashBuild.releaseRecoveryCapacity( + ctr.hashmapBuilder.mapAllocationAccount, + true, + ); err != nil { + return err + } v2.HashBuildSpillDepthCounter.WithLabelValues("spill", "1").Inc() return nil } @@ -353,9 +353,10 @@ func (hashBuild *HashBuild) build( if err := startSpill(); err != nil { return err } - if err := ensureDirectRecovery(bat); err != nil { - return err - } + // Recovery headroom proves that already-retained batches can be drained; + // an upstream-owned direct source cannot strand retained state. Admit its + // scratch at the physical allocation sites instead of applying the + // conservative retained-batch projection as a query-fatal gate. return ctr.spillBatchWithPressure( proc, bat, spillFiles, ctr.spillExprExecs, analyzer, false) } diff --git a/pkg/sql/colexec/hashbuild/build_test.go b/pkg/sql/colexec/hashbuild/build_test.go index bd012abb35732..a8ab28ae0148f 100644 --- a/pkg/sql/colexec/hashbuild/build_test.go +++ b/pkg/sql/colexec/hashbuild/build_test.go @@ -2673,6 +2673,93 @@ func TestShuffleHashBuildAccountedSpillLifecycle(t *testing.T) { require.Zero(t, tc.proc.Mp().CurrNB()) } +func TestShuffleHashBuildDirectSpillUsesActualAllocation(t *testing.T) { + tc := newTestCase( + t, + []bool{false, false}, + []types.Type{types.T_int32.ToType(), types.T_int32.ToType()}, + nil, + ) + tc.arg.Conditions = []*plan.Expr{makeIssue26454ConcatKey(t, tc.proc)} + tc.arg.IsShuffle = true + tc.arg.ShuffleIdx = 0 + tc.arg.SpillThreshold = 1 + tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{ + Tag: tc.arg.JoinMapTag + 4_501, + } + tc.arg.SetChildren([]vm.Operator{tc.marg}) + + const ( + limit = uint64(8 << 20) + rows = 128 + ) + estimated, err := expressionRecoveryBytes( + tc.proc, + tc.arg.Conditions, + rows, + false, + ) + require.NoError(t, err) + require.Greater(t, estimated, limit, + "fixture must exceed the conservative retained-recovery projection") + + budget := process.MustNewHashBuildBudget(limit, limit) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 256) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + replaceTestHashBuildAllocation(t, tc.arg, account) + require.NoError(t, tc.marg.Prepare(tc.proc)) + require.NoError(t, tc.arg.Prepare(tc.proc)) + + build := newBatch(tc.types, tc.proc, rows) + tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly( + build, + nil, + tc.proc.Mp(), + ) + tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly( + nil, + nil, + tc.proc.Mp(), + ) + _, err = vm.Exec(tc.arg, tc.proc) + require.NoError(t, err) + require.Zero(t, generation.RejectCount(), + "an unretained direct source must not hit the recovery projection gate") + require.Less(t, generation.Peak(), limit) + + result, err := message.ReceiveJoinMapResult( + tc.arg.JoinMapTag, + true, + tc.arg.ShuffleIdx, + tc.proc.GetMessageBoard(), + tc.proc.Ctx, + ) + require.NoError(t, err) + require.True(t, result.IsSuccess()) + jm := result.JoinMap() + require.NotNil(t, jm) + require.True(t, jm.IsSpilled()) + require.Equal(t, int64(rows), jm.GetRowCount()) + payload, err := jm.TakeSpillBuildPayload() + require.NoError(t, err) + require.NoError(t, payload.Close()) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + + tc.arg.Reset(tc.proc, false, nil) + tc.marg.Reset(tc.proc, false, nil) + require.NoError(t, tc.arg.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + tc.arg.Free(tc.proc, false, nil) + tc.proc.Free() + require.Zero(t, tc.proc.Mp().CurrNB()) +} + func TestObserveHashBuildBudgetUsesGenerationSnapshot(t *testing.T) { budget := process.MustNewHashBuildBudget(1024, 1024) generation, err := budget.OpenGeneration(1) diff --git a/pkg/sql/colexec/hashbuild/spill.go b/pkg/sql/colexec/hashbuild/spill.go index f2f83aaa28055..73baf2db3e290 100644 --- a/pkg/sql/colexec/hashbuild/spill.go +++ b/pkg/sql/colexec/hashbuild/spill.go @@ -558,6 +558,18 @@ func (ctr *container) releaseSpillComputeScratch() { ctr.spillBucketRowIds = nil } +// dropMandatorySpillRecoveryScratch releases only allocations that borrow the +// retained-state recovery floor. Optional coalescing buffers use ordinary +// admission and may keep their already-produced records across the transition. +func (ctr *container) dropMandatorySpillRecoveryScratch() { + ctr.freeSpillExprExecs() + ctr.releaseSpillComputeScratch() + if ctr.spillAccountedWrite != nil { + ctr.spillAccountedWrite.Free() + ctr.spillAccountedWrite = nil + } +} + // spillBatchWithPressure retries only the unpublished prefix of an exact // spill operation. Hash/expression capacity failures happen before any bucket // write; selected/codec failures are handled transactionally inside diff --git a/pkg/sql/colexec/hashbuild/types.go b/pkg/sql/colexec/hashbuild/types.go index 699f602bbe086..8e19ad938e113 100644 --- a/pkg/sql/colexec/hashbuild/types.go +++ b/pkg/sql/colexec/hashbuild/types.go @@ -384,10 +384,10 @@ func (hashBuild *HashBuild) installRecoveryCapacity( return nil } -// releaseRecoveryCapacity returns unused recovery headroom as soon as build -// reaches a terminal result. restoreDefault keeps direct test/reuse spill -// allocations on the statement's ordinary controller; statement teardown -// passes false and drops the selection immediately afterward. +// releaseRecoveryCapacity returns recovery headroom after retained spill state +// has been drained or build reaches a terminal result. restoreDefault keeps +// later direct/test/reuse allocations on the statement's ordinary controller; +// statement teardown passes false and drops the selection immediately afterward. func (hashBuild *HashBuild) releaseRecoveryCapacity( account *mpool.AllocationAccount, restoreDefault bool, From 95baebc538e0889995f8d8560d8fd3043234b648 Mon Sep 17 00:00:00 2001 From: aptend Date: Mon, 3 Aug 2026 16:02:46 +0800 Subject: [PATCH 59/61] ci: skip onnxruntime for branch validation images --- optools/images/Dockerfile | 4 +++- optools/images/Dockerfile.ci | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/optools/images/Dockerfile b/optools/images/Dockerfile index 678e1b24b7a18..fc666d1915428 100644 --- a/optools/images/Dockerfile +++ b/optools/images/Dockerfile @@ -10,7 +10,9 @@ WORKDIR /go/src/github.com/matrixorigin/matrixone FROM build-base AS native COPY thirdparties thirdparties -RUN make -C thirdparties +# This branch image is only used for non-ONNX regression validation. Build the +# native dependencies those workloads need without fetching ONNX Runtime. +RUN make -C thirdparties init usearch xxhash croaring COPY cgo cgo RUN make -C cgo diff --git a/optools/images/Dockerfile.ci b/optools/images/Dockerfile.ci index 87cf42425fc04..3dc841c850aea 100644 --- a/optools/images/Dockerfile.ci +++ b/optools/images/Dockerfile.ci @@ -11,7 +11,9 @@ RUN go env -w GOPROXY=${GOPROXY} WORKDIR /go/src/github.com/matrixorigin/matrixone COPY thirdparties thirdparties -RUN make -C thirdparties +# This branch image is only used for non-ONNX regression validation. Build the +# native dependencies those workloads need without fetching ONNX Runtime. +RUN make -C thirdparties init usearch xxhash croaring COPY cgo cgo RUN make -C cgo RUN mkdir -p /root/.cache/go-build /go/pkg/mod From 57811ed7fa7cc5ed9234f92a9167c8187abaacdf Mon Sep 17 00:00:00 2001 From: aptend Date: Mon, 3 Aug 2026 16:20:20 +0800 Subject: [PATCH 60/61] Revert "ci: skip onnxruntime for branch validation images" This reverts commit 95baebc538e0889995f8d8560d8fd3043234b648. --- optools/images/Dockerfile | 4 +--- optools/images/Dockerfile.ci | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/optools/images/Dockerfile b/optools/images/Dockerfile index fc666d1915428..678e1b24b7a18 100644 --- a/optools/images/Dockerfile +++ b/optools/images/Dockerfile @@ -10,9 +10,7 @@ WORKDIR /go/src/github.com/matrixorigin/matrixone FROM build-base AS native COPY thirdparties thirdparties -# This branch image is only used for non-ONNX regression validation. Build the -# native dependencies those workloads need without fetching ONNX Runtime. -RUN make -C thirdparties init usearch xxhash croaring +RUN make -C thirdparties COPY cgo cgo RUN make -C cgo diff --git a/optools/images/Dockerfile.ci b/optools/images/Dockerfile.ci index 3dc841c850aea..87cf42425fc04 100644 --- a/optools/images/Dockerfile.ci +++ b/optools/images/Dockerfile.ci @@ -11,9 +11,7 @@ RUN go env -w GOPROXY=${GOPROXY} WORKDIR /go/src/github.com/matrixorigin/matrixone COPY thirdparties thirdparties -# This branch image is only used for non-ONNX regression validation. Build the -# native dependencies those workloads need without fetching ONNX Runtime. -RUN make -C thirdparties init usearch xxhash croaring +RUN make -C thirdparties COPY cgo cgo RUN make -C cgo RUN mkdir -p /root/.cache/go-build /go/pkg/mod From bf244d6b5ca6090b8e29ccbcc550816f554ca42e Mon Sep 17 00:00:00 2001 From: aptend Date: Mon, 3 Aug 2026 19:28:44 +0800 Subject: [PATCH 61/61] fix(compile): bound remote execution tombstones --- .../remote_allocation_statement_group.go | 16 ++++ .../remote_allocation_statement_group_test.go | 75 +++++++++++++++++++ 2 files changed, 91 insertions(+) diff --git a/pkg/sql/compile/remote_allocation_statement_group.go b/pkg/sql/compile/remote_allocation_statement_group.go index 5547b53df7bd0..46a2e321d9765 100644 --- a/pkg/sql/compile/remote_allocation_statement_group.go +++ b/pkg/sql/compile/remote_allocation_statement_group.go @@ -66,6 +66,12 @@ var remoteAllocationStatementRegistrationTimeout = 5 * time.Minute // contains no statement resources. const remoteAllocationStatementTombstoneTimeout = MaxRpcTime +// Reserve tombstone capacity when a remote execution is first admitted. An +// active group can become a tombstone after a partial dispatch, so bounding +// only the tombstone map would admit more generations than can later be +// retained safely. Existing fragments do not consume another reservation. +const remoteAllocationStatementGenerationLimit = 4096 + // collectRemoteFragmentCounts computes the number of pipeline RPCs that the // complete physical scope graph will send to each CN. The execution address // changes when traversal crosses a Remote scope: nested scopes targeting that @@ -196,6 +202,16 @@ func acquireRemoteAllocationStatementParticipant( moerr.NewInternalErrorNoCtx("remote allocation statement group key already registered"), ) } + if len(remoteAllocationStatementGroups.byKey)+ + len(remoteAllocationStatementGroups.tombstones) >= + remoteAllocationStatementGenerationLimit { + return nil, errors.Join( + mpool.ErrAllocationAccountInvariant, + moerr.NewInternalErrorNoCtx( + "remote allocation statement generation capacity reached", + ), + ) + } group = &remoteAllocationStatementGroup{ key: key, board: board, diff --git a/pkg/sql/compile/remote_allocation_statement_group_test.go b/pkg/sql/compile/remote_allocation_statement_group_test.go index e76a6b35705a5..f45d22c4af06e 100644 --- a/pkg/sql/compile/remote_allocation_statement_group_test.go +++ b/pkg/sql/compile/remote_allocation_statement_group_test.go @@ -473,6 +473,81 @@ func TestRemoteAllocationStatementGroupExpiryRejectsLateFragment(t *testing.T) { require.True(t, terminal.complete) } +func TestRemoteAllocationStatementGenerationCapacity(t *testing.T) { + activeKey := remoteAllocationStatementGroupKey(newRemoteExecutionID(), "cn-a:6001") + activeBoard := message.NewMessageBoard() + first, err := acquireRemoteAllocationStatementParticipant( + activeKey, + activeBoard, + 3, + nil, + ) + require.NoError(t, err) + t.Cleanup(func() { + remoteAllocationStatementGroups.Lock() + if group := remoteAllocationStatementGroups.byKey[activeKey]; group != nil { + if group.timer != nil { + group.timer.Stop() + } + delete(remoteAllocationStatementGroups.byBoard, group.board) + delete(remoteAllocationStatementGroups.byKey, activeKey) + } + if tombstone := remoteAllocationStatementGroups.tombstones[activeKey]; tombstone != nil { + if tombstone.timer != nil { + tombstone.timer.Stop() + } + delete(remoteAllocationStatementGroups.tombstones, activeKey) + } + remoteAllocationStatementGroups.Unlock() + activeBoard.CloseAndDrain() + }) + + remoteAllocationStatementGroups.Lock() + for i := 1; i < remoteAllocationStatementGenerationLimit; i++ { + key := fmt.Sprintf("capacity-tombstone-%d", i) + remoteAllocationStatementGroups.tombstones[key] = + &remoteAllocationStatementTombstone{} + } + remoteAllocationStatementGroups.Unlock() + t.Cleanup(func() { + remoteAllocationStatementGroups.Lock() + for i := 1; i < remoteAllocationStatementGenerationLimit; i++ { + delete( + remoteAllocationStatementGroups.tombstones, + fmt.Sprintf("capacity-tombstone-%d", i), + ) + } + remoteAllocationStatementGroups.Unlock() + }) + + // The reservation belongs to the generation, not each fragment. + second, err := acquireRemoteAllocationStatementParticipant( + activeKey, + activeBoard, + 3, + nil, + ) + require.NoError(t, err) + + newKey := remoteAllocationStatementGroupKey(newRemoteExecutionID(), "cn-a:6001") + newBoard := message.NewMessageBoard() + _, err = acquireRemoteAllocationStatementParticipant(newKey, newBoard, 1, nil) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + newBoard.CloseAndDrain() + + terminal, err := first.finish(errors.New("partial dispatch failed")) + require.NoError(t, err) + require.False(t, terminal.complete) + terminal, err = second.finish(errors.New("sibling canceled")) + require.Error(t, err) + require.True(t, terminal.complete) + + remoteAllocationStatementGroups.Lock() + tombstoneCount := len(remoteAllocationStatementGroups.tombstones) + remoteAllocationStatementGroups.Unlock() + require.Equal(t, remoteAllocationStatementGenerationLimit, tombstoneCount) +} + func TestRemoteAllocationStatementGroupFailureCancelsActiveSibling(t *testing.T) { board := message.NewMessageBoard() canceled := make(chan error, 2)