Copy bytes as needed

document/field/value.go

@@ -266,6 +266,20 @@ func (v ValueUnion) MarshalJSON() ([]byte, error) {
 	}
 }
 
+// ValueCloneOptions controls how a value should be cloned.
+type ValueCloneOptions struct {
+	DeepCloneBytes bool
+}
+
+// Clone clones a value union.
+func (v *ValueUnion) Clone(opts ValueCloneOptions) ValueUnion {
+	cloned := *v
+	if v.Type == BytesType && opts.DeepCloneBytes {
+		cloned.BytesVal = bytes.NewImmutableBytes(v.BytesVal.SafeBytes())
+	}
+	return cloned
+}
+
 // ToProto converts a value to a value proto message.
 func (v *ValueUnion) ToProto() (servicepb.FieldValue, error) {
 	var fb servicepb.FieldValue
@@ -491,14 +505,13 @@ func (v Values) Hash() uint64 {
 }
 
 // Clone clones the values.
-func (v Values) Clone() Values {
+func (v Values) Clone(opts ValueCloneOptions) Values {
 	if len(v) == 0 {
 		return nil
 	}
 	cloned := make(Values, 0, len(v))
 	for i := 0; i < len(v); i++ {
-		// NB: This is fine as each value union does not contain reference types.
-		cloned = append(cloned, v[i])
+		cloned = append(cloned, v[i].Clone(opts))
 	}
 	return cloned
 }

query/executor/doc_id_values.go

@@ -10,17 +10,18 @@ type docIDValues struct {
 // cloneDocIDValues clones the incoming (doc ID, values) pair.
 func cloneDocIDValues(v docIDValues) docIDValues {
 	// TODO(xichen): Should pool and reuse the value array here.
-	v.Values = v.Values.Clone()
+	v.Values = v.Values.Clone(field.ValueCloneOptions{DeepCloneBytes: true})
 	return v
 }
 
 // cloneDocIDValuesTo clones the (doc ID, values) pair from `src` to `target`.
 // Precondition: `target` values are guaranteed to have the same length as `src` and can be reused.
 func cloneDocIDValuesTo(src docIDValues, target *docIDValues) {
-	reusedValues := target.Values
-	copy(reusedValues, src.Values)
 	target.DocID = src.DocID
-	target.Values = reusedValues
+	for i := 0; i < len(src.Values); i++ {
+		cloned := src.Values[i].Clone(field.ValueCloneOptions{DeepCloneBytes: true})
+		target.Values[i] = cloned
+	}
 }
 
 type docIDValuesByDocIDAsc []docIDValues

query/single_key_result_groups.go

@@ -303,7 +303,10 @@ func (m *SingleKeyResultGroups) getOrInsertBytes(
 		return nil, RejectedDueToLimit
 	}
 	arr = m.resultArrayProtoType.New()
-	m.bytesResults.Set(v.SafeBytes(), arr)
+	m.bytesResults.SetUnsafe(v.SafeBytes(), arr, SetUnsafeBytesOptions{
+		NoCopyKey:     true,
+		NoFinalizeKey: true,
+	})
 	return arr, Inserted
 }
 
@@ -423,7 +426,10 @@ func (m *SingleKeyResultGroups) mergeBytesGroups(other *SingleKeyResultGroups) {
 			// Limit reached, do nothing.
 			continue
 		}
-		m.bytesResults.Set(k, v)
+		m.bytesResults.SetUnsafe(k, v, SetUnsafeBytesOptions{
+			NoCopyKey:     true,
+			NoFinalizeKey: true,
+		})
 	}
 }
 
@@ -891,7 +897,10 @@ func (m *SingleKeyResultGroups) trimBytesToTopN(targetSize int) {
 	})
 	data := m.topNBytes.RawData()
 	for i := 0; i < len(data); i++ {
-		m.bytesResults.Set(data[i].Key, data[i].Values)
+		m.bytesResults.SetUnsafe(data[i].Key, data[i].Values, SetUnsafeBytesOptions{
+			NoCopyKey:     true,
+			NoFinalizeKey: true,
+		})
 		data[i] = emptyBytesResultGroup
 	}
 	m.topNBytes.Reset()

query/values_result_array_new_map.go

@@ -12,7 +12,7 @@ func NewValuesResultArrayMap(initialSize int) *ValuesResultArrayHash {
 			return v1.Equal(v2)
 		},
 		copy: func(v field.Values) field.Values {
-			return v.Clone()
+			return v.Clone(field.ValueCloneOptions{DeepCloneBytes: true})
 		},
 		finalize:    nil, // No op on key removal
 		initialSize: initialSize,