From 6b5be5f76b23d17be527f8db138c16af8524ab02 Mon Sep 17 00:00:00 2001
From: proost <jwalag87@gmail.com>
Date: Mon, 14 Sep 2026 17:33:03 +0900
Subject: [PATCH] feat: CompactTrimmed

---
 theta/union.go              |   6 +-
 theta/update_sketch.go      |  31 +++++++
 theta/update_sketch_test.go | 175 ++++++++++++++++++++++++++++++++++++
 theta/utils.go              |  13 +++
 theta/utils_test.go         |  71 +++++++++++++++
 5 files changed, 291 insertions(+), 5 deletions(-)

diff --git a/theta/union.go b/theta/union.go
index b501639..7162f8d 100644
--- a/theta/union.go
+++ b/theta/union.go
@@ -185,11 +185,7 @@ func (u *Union) Result(ordered bool) (*CompactSketch, error) {
 		}
 	}
 
-	if uint32(len(entries)) > nominalNum {
-		internal.QuickSelect(entries, 0, len(entries)-1, int(nominalNum))
-		theta = entries[nominalNum]
-		entries = entries[:nominalNum]
-	}
+	entries, theta = trimToNominal(entries, nominalNum, theta)
 
 	if ordered {
 		slices.Sort(entries)
diff --git a/theta/update_sketch.go b/theta/update_sketch.go
index a21eae3..bc20592 100644
--- a/theta/update_sketch.go
+++ b/theta/update_sketch.go
@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"iter"
 	"math"
+	"slices"
 	"strings"
 
 	"github.com/apache/datasketches-go/internal"
@@ -415,10 +416,40 @@ func (s *QuickSelectUpdateSketch) All() iter.Seq[uint64] {
 	}
 }
 
+// Compact converts this sketch to a compact sketch (ordered or unordered).
 func (s *QuickSelectUpdateSketch) Compact(ordered bool) *CompactSketch {
 	return NewCompactSketch(s, ordered)
 }
 
+// CompactOrdered converts this sketch to an ordered compact sketch.
 func (s *QuickSelectUpdateSketch) CompactOrdered() *CompactSketch {
 	return s.Compact(true)
 }
+
+// CompactTrimmed converts this sketch to a compact sketch (ordered or unordered)
+// reduced to at most the nominal size k. This sketch is not modified.
+//
+// If this sketch retains more than k entries, theta is lowered to the (k+1)th
+// smallest retained hash and only the k entries below it are kept. Otherwise the
+// result is the same as Compact. This is equivalent to Trim followed by Compact,
+// without mutating this sketch or rebuilding its hash table.
+func (s *QuickSelectUpdateSketch) CompactTrimmed(ordered bool) *CompactSketch {
+	if s.IsEmpty() {
+		return s.Compact(ordered)
+	}
+
+	entries := make([]uint64, 0, s.table.numEntries)
+	for _, entry := range s.table.entries {
+		if entry != 0 {
+			entries = append(entries, entry)
+		}
+	}
+
+	entries, theta := trimToNominal(entries, uint32(1)<<s.table.lgNomSize, s.table.theta)
+	if ordered {
+		slices.Sort(entries)
+	}
+
+	seedHash, _ := s.SeedHash()
+	return newCompactSketchFromEntries(false, ordered, seedHash, theta, entries)
+}
diff --git a/theta/update_sketch_test.go b/theta/update_sketch_test.go
index e1ef9e4..f08e131 100644
--- a/theta/update_sketch_test.go
+++ b/theta/update_sketch_test.go
@@ -18,6 +18,8 @@
 package theta
 
 import (
+	"bytes"
+	"slices"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -662,3 +664,176 @@ func TestUpdateSketch_Compact(t *testing.T) {
 		assert.GreaterOrEqual(t, ub, float64(n))
 	})
 }
+
+func TestQuickSelectUpdateSketch_CompactTrimmed(t *testing.T) {
+	const k = uint32(1) << DefaultLgK
+
+	newSketch := func(t *testing.T, n int) *QuickSelectUpdateSketch {
+		t.Helper()
+		sketch, err := NewQuickSelectUpdateSketch()
+		assert.NoError(t, err)
+		for i := 0; i < n; i++ {
+			assertUpdate(t, sketch.UpdateInt64(int64(i)))
+		}
+		return sketch
+	}
+
+	sortedEntries := func(s *CompactSketch) []uint64 {
+		entries := slices.Collect(s.All())
+		slices.Sort(entries)
+		return entries
+	}
+
+	t.Run("Ordered And Trimmed Combinations", func(t *testing.T) {
+		n := 40000
+		sketch := newSketch(t, n)
+		retainedBefore := sketch.NumRetained()
+		thetaBefore := sketch.Theta64()
+		assert.True(t, sketch.IsEstimationMode())
+		assert.Greater(t, retainedBefore, k)
+
+		reference := newSketch(t, n)
+		reference.Trim()
+		expected := reference.CompactOrdered()
+		expectedEntries := sortedEntries(expected)
+		assert.Equal(t, k, expected.NumRetained())
+
+		// case 1: ordered, not trimmed
+		c1 := sketch.Compact(true)
+		assert.True(t, c1.IsOrdered())
+		assert.Equal(t, retainedBefore, c1.NumRetained())
+		assert.Equal(t, thetaBefore, c1.Theta64())
+
+		// case 2: unordered, not trimmed
+		c2 := sketch.Compact(false)
+		assert.False(t, c2.IsOrdered())
+		assert.Equal(t, retainedBefore, c2.NumRetained())
+		assert.Equal(t, thetaBefore, c2.Theta64())
+
+		// case 3: ordered and trimmed
+		c3 := sketch.CompactTrimmed(true)
+		assert.True(t, c3.IsOrdered())
+		assert.Equal(t, k, c3.NumRetained())
+		assert.Less(t, c3.Theta64(), thetaBefore)
+		assert.Equal(t, expected.Theta64(), c3.Theta64())
+		c3Entries := slices.Collect(c3.All())
+		assert.True(t, slices.IsSorted(c3Entries))
+		assert.Equal(t, expectedEntries, c3Entries)
+		for _, entry := range c3Entries {
+			assert.Less(t, entry, c3.Theta64())
+		}
+
+		// case 4: unordered and trimmed: same set and theta, no sort
+		c4 := sketch.CompactTrimmed(false)
+		assert.False(t, c4.IsOrdered())
+		assert.Equal(t, k, c4.NumRetained())
+		assert.Equal(t, expected.Theta64(), c4.Theta64())
+		assert.Equal(t, expectedEntries, sortedEntries(c4))
+
+		// the source sketch must be untouched by any of the four
+		assert.Equal(t, retainedBefore, sketch.NumRetained())
+		assert.Equal(t, thetaBefore, sketch.Theta64())
+	})
+
+	t.Run("Serialization Matches Trim Then Compact", func(t *testing.T) {
+		n := 40000
+		sketch := newSketch(t, n)
+		reference := newSketch(t, n)
+		reference.Trim()
+
+		for _, compressed := range []bool{false, true} {
+			var trimmedBuf, referenceBuf bytes.Buffer
+			trimmedEncoder := NewEncoder(&trimmedBuf, compressed)
+			assert.NoError(t, trimmedEncoder.Encode(sketch.CompactTrimmed(true)))
+			referenceEncoder := NewEncoder(&referenceBuf, compressed)
+			assert.NoError(t, referenceEncoder.Encode(reference.CompactOrdered()))
+
+			assert.Equal(t, referenceBuf.Bytes(), trimmedBuf.Bytes(), "compressed=%v", compressed)
+		}
+	})
+
+	t.Run("Exact Mode Converts To Estimation", func(t *testing.T) {
+		n := 5000
+		sketch := newSketch(t, n)
+		assert.Greater(t, sketch.NumRetained(), k)
+		assert.False(t, sketch.IsEstimationMode())
+		assert.Equal(t, 1.0, sketch.Theta())
+
+		exact := sketch.Compact(true)
+		assert.False(t, exact.IsEstimationMode())
+		assert.Equal(t, uint32(n), exact.NumRetained())
+		assert.Equal(t, float64(n), exact.Estimate())
+
+		trimmed := sketch.CompactTrimmed(true)
+		assert.True(t, trimmed.IsEstimationMode())
+		assert.Equal(t, k, trimmed.NumRetained())
+		assert.Less(t, trimmed.Theta(), 1.0)
+		assert.NotEqual(t, float64(n), trimmed.Estimate())
+
+		lb, err := trimmed.LowerBound(3)
+		assert.NoError(t, err)
+		assert.LessOrEqual(t, lb, float64(n))
+		ub, err := trimmed.UpperBound(3)
+		assert.NoError(t, err)
+		assert.GreaterOrEqual(t, ub, float64(n))
+	})
+
+	t.Run("Bounds Widen In Estimation Mode", func(t *testing.T) {
+		sketch := newSketch(t, 40000)
+		assert.True(t, sketch.IsEstimationMode())
+		assert.Greater(t, sketch.NumRetained(), k)
+
+		plain := sketch.Compact(true)
+		trimmed := sketch.CompactTrimmed(true)
+		assert.Equal(t, k, trimmed.NumRetained())
+		assert.Greater(t, plain.NumRetained(), trimmed.NumRetained())
+
+		width := func(s *CompactSketch) float64 {
+			lb, err := s.LowerBound(2)
+			assert.NoError(t, err)
+			ub, err := s.UpperBound(2)
+			assert.NoError(t, err)
+			return ub - lb
+		}
+		assert.Greater(t, width(trimmed), width(plain))
+	})
+
+	t.Run("Empty", func(t *testing.T) {
+		sketch := newSketch(t, 0)
+
+		result := sketch.CompactTrimmed(true)
+		assert.True(t, result.IsEmpty())
+		assert.Zero(t, result.NumRetained())
+		assert.Equal(t, 1.0, result.Theta())
+		assert.True(t, result.IsOrdered())
+		assert.True(t, sketch.CompactTrimmed(false).IsOrdered())
+	})
+
+	t.Run("Non Empty No Retained Keys", func(t *testing.T) {
+		sketch, err := NewQuickSelectUpdateSketch(WithUpdateSketchP(0.001))
+		assert.NoError(t, err)
+		assertUpdate(t, sketch.UpdateInt64(1))
+
+		result := sketch.CompactTrimmed(true)
+		assert.False(t, result.IsEmpty())
+		assert.Zero(t, result.NumRetained())
+		assert.Equal(t, sketch.Theta64(), result.Theta64())
+		assert.True(t, result.IsEstimationMode())
+	})
+
+	t.Run("Below K", func(t *testing.T) {
+		sketch := newSketch(t, 100)
+		assert.False(t, sketch.IsEstimationMode())
+
+		result := sketch.CompactTrimmed(true)
+		assert.False(t, result.IsEstimationMode())
+		assert.Equal(t, uint32(100), result.NumRetained())
+		assert.Equal(t, sketch.Theta64(), result.Theta64())
+		assert.Equal(t, 100.0, result.Estimate())
+		assert.True(t, slices.IsSorted(slices.Collect(result.All())))
+		assert.False(t, sketch.CompactTrimmed(false).IsOrdered())
+
+		single := newSketch(t, 1)
+		assert.True(t, single.CompactTrimmed(false).IsOrdered())
+	})
+}
diff --git a/theta/utils.go b/theta/utils.go
index 246ee1e..269bcd9 100644
--- a/theta/utils.go
+++ b/theta/utils.go
@@ -19,6 +19,8 @@ package theta
 
 import (
 	"fmt"
+
+	"github.com/apache/datasketches-go/internal"
 )
 
 func checkEqual[T comparable](actual, expected T, description string) error {
@@ -68,3 +70,14 @@ func startingSubMultiple(lgTgt, lgMin, lgRf uint8) uint8 {
 	}
 	return ((lgTgt - lgMin) % lgRf) + lgMin
 }
+
+func trimToNominal(entries []uint64, nominalSize uint32, theta uint64) ([]uint64, uint64) {
+	if uint32(len(entries)) <= nominalSize {
+		return entries, theta
+	}
+
+	internal.QuickSelect(entries, 0, len(entries)-1, int(nominalSize))
+	trimmed := make([]uint64, nominalSize)
+	copy(trimmed, entries[:nominalSize])
+	return trimmed, entries[nominalSize]
+}
diff --git a/theta/utils_test.go b/theta/utils_test.go
index 548808e..d64f19f 100644
--- a/theta/utils_test.go
+++ b/theta/utils_test.go
@@ -18,6 +18,8 @@
 package theta
 
 import (
+	"math/rand/v2"
+	"slices"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -140,3 +142,72 @@ func TestStartingSubMultiple(t *testing.T) {
 		})
 	}
 }
+
+func TestTrimToNominal(t *testing.T) {
+	shuffledRange := func(n int) []uint64 {
+		r := rand.New(rand.NewPCG(1, 2))
+		entries := make([]uint64, n)
+		for i, v := range r.Perm(n) {
+			entries[i] = uint64(v + 1)
+		}
+		return entries
+	}
+
+	t.Run("Below Nominal", func(t *testing.T) {
+		entries := shuffledRange(5)
+		expected := slices.Clone(entries)
+
+		trimmed, theta := trimToNominal(entries, 8, MaxTheta)
+		assert.Equal(t, expected, trimmed)
+		assert.Equal(t, MaxTheta, theta)
+	})
+
+	t.Run("At Nominal", func(t *testing.T) {
+		entries := shuffledRange(8)
+		expected := slices.Clone(entries)
+
+		trimmed, theta := trimToNominal(entries, 8, MaxTheta)
+		assert.Equal(t, expected, trimmed)
+		assert.Equal(t, MaxTheta, theta)
+	})
+
+	t.Run("Empty", func(t *testing.T) {
+		trimmed, theta := trimToNominal(nil, 8, MaxTheta)
+		assert.Empty(t, trimmed)
+		assert.Equal(t, MaxTheta, theta)
+	})
+
+	t.Run("Above Nominal", func(t *testing.T) {
+		const nominal = 13
+		entries := make([]uint64, 0, 64)
+		entries = append(entries, shuffledRange(40)...)
+
+		trimmed, theta := trimToNominal(entries, nominal, MaxTheta)
+
+		// values are 1..40, so the (nominal+1)th smallest is nominal+1
+		assert.Equal(t, uint64(nominal+1), theta)
+		assert.Len(t, trimmed, nominal)
+		assert.Equal(t, nominal, cap(trimmed), "trimmed result must not keep the untrimmed allocation")
+		for _, entry := range trimmed {
+			assert.Less(t, entry, theta)
+		}
+		slices.Sort(trimmed)
+		expected := make([]uint64, nominal)
+		for i := range expected {
+			expected[i] = uint64(i + 1)
+		}
+		assert.Equal(t, expected, trimmed)
+	})
+
+	t.Run("One Above Nominal", func(t *testing.T) {
+		const nominal = 8
+		entries := shuffledRange(nominal + 1)
+
+		trimmed, theta := trimToNominal(entries, nominal, MaxTheta)
+		assert.Equal(t, uint64(nominal+1), theta)
+		assert.Len(t, trimmed, nominal)
+		for _, entry := range trimmed {
+			assert.Less(t, entry, theta)
+		}
+	})
+}