Support raw query against mutable segments (#75)

document/field/field.go

@@ -18,6 +18,11 @@ const (
 	TimeType
 )
 
+var (
+	// NumValidFieldTypes returns the number of valid field types.
+	NumValidFieldTypes = len(validTypes)
+)
+
 // IsValid returns true if this is a valid value type.
 func (t ValueType) IsValid() bool {
 	_, exists := validTypes[t]
@@ -58,6 +63,16 @@ func (m ValueTypeSet) Clone() ValueTypeSet {
 	return cloned
 }
 
+// MergeInPlace merges another value type set into the current set in place.
+func (m ValueTypeSet) MergeInPlace(other ValueTypeSet) {
+	if len(other) == 0 {
+		return
+	}
+	for k := range other {
+		m[k] = struct{}{}
+	}
+}
+
 var (
 	// validTypes is a list of valid value types.
 	validTypes = ValueTypeSet{

index/doc_id_set.go

@@ -126,7 +126,8 @@ func NewBitmapBasedDocIDSetBuilder(bm *roaring.Bitmap) DocIDSetBuilder {
 func (s *bitmapBasedDocIDSetBuilder) Add(docID int32) { s.bm.DirectAdd(uint64(docID)) }
 
 // NB(xichen): Clone the internal bitmap so the builder can be mutated independently
-// of the snapshot. In the future we can look into the bitmap implementation to see
+// of the snapshot.
+// TODO(xichen): We should look into the bitmap implementation to see
 // if there are more efficient ways of doing this without requiring a full copy.
 func (s *bitmapBasedDocIDSetBuilder) Snapshot() DocIDSet {
 	return newBitmapBasedDocIDSet(s.bm.Clone())

index/field/docs_field.go

@@ -113,9 +113,10 @@ type DocsFieldBuilder interface {
 	// Add adds a value with its document ID.
 	Add(docID int32, v field.ValueUnion) error
 
-	// Snapshot returns an immutable snapshot of the doc IDs an the field values
-	// contained in the field.
-	Snapshot() DocsField
+	// Snapshot return an immutable snapshot of the typed fields for the given
+	// value type set, returning the new docs field and any types remaining in
+	// the given value type set that's not available in the builder.
+	SnapshotFor(fieldTypes field.ValueTypeSet) (DocsField, field.ValueTypeSet, error)
 
 	// Seal seals and closes the builder and returns an immutable docs field that contains (and
 	// owns) all the doc IDs and the field values accummulated across `numTotalDocs`
@@ -634,43 +635,102 @@ func (b *docsFieldBuilder) Add(docID int32, v field.ValueUnion) error {
 	}
 }
 
-// Snapshot return an immutable snapshot of the builder state.
-func (b *docsFieldBuilder) Snapshot() DocsField {
+func (b *docsFieldBuilder) SnapshotFor(
+	fieldTypes field.ValueTypeSet,
+) (DocsField, field.ValueTypeSet, error) {
 	var (
-		fieldTypes = make([]field.ValueType, 0, 6)
-		nf         CloseableNullField
-		bf         CloseableBoolField
-		intf       CloseableIntField
-		df         CloseableDoubleField
-		sf         CloseableStringField
-		tf         CloseableTimeField
+		nf             CloseableNullField
+		bf             CloseableBoolField
+		intf           CloseableIntField
+		df             CloseableDoubleField
+		sf             CloseableStringField
+		tf             CloseableTimeField
+		availableTypes = make([]field.ValueType, 0, len(fieldTypes))
+		remainingTypes field.ValueTypeSet
+		err            error
 	)
-	if b.nfb != nil {
-		fieldTypes = append(fieldTypes, field.NullType)
-		nf = b.nfb.Snapshot()
+	for t := range fieldTypes {
+		hasType := true
+
+		switch t {
+		case field.NullType:
+			if b.nfb != nil {
+				nf = b.nfb.Snapshot()
+				break
+			}
+			hasType = false
+		case field.BoolType:
+			if b.bfb != nil {
+				bf = b.bfb.Snapshot()
+				break
+			}
+			hasType = false
+		case field.IntType:
+			if b.ifb != nil {
+				intf = b.ifb.Snapshot()
+				break
+			}
+			hasType = false
+		case field.DoubleType:
+			if b.dfb != nil {
+				df = b.dfb.Snapshot()
+				break
+			}
+			hasType = false
+		case field.StringType:
+			if b.sfb != nil {
+				sf = b.sfb.Snapshot()
+				break
+			}
+			hasType = false
+		case field.TimeType:
+			if b.tfb != nil {
+				tf = b.tfb.Snapshot()
+				break
+			}
+			hasType = false
+		default:
+			err = fmt.Errorf("unknown field type %v", t)
+		}
+
+		if err != nil {
+			break
+		}
+		if hasType {
+			availableTypes = append(availableTypes, t)
+			continue
+		}
+		if remainingTypes == nil {
+			remainingTypes = make(field.ValueTypeSet, len(fieldTypes))
+		}
+		remainingTypes[t] = struct{}{}
 	}
-	if b.bfb != nil {
-		fieldTypes = append(fieldTypes, field.BoolType)
-		bf = b.bfb.Snapshot()
+
+	if err == nil {
+		return NewDocsField(b.fieldPath, availableTypes, nf, bf, intf, df, sf, tf), remainingTypes, nil
 	}
-	if b.ifb != nil {
-		fieldTypes = append(fieldTypes, field.IntType)
-		intf = b.ifb.Snapshot()
+
+	// Close all resources on error.
+	if nf != nil {
+		nf.Close()
 	}
-	if b.dfb != nil {
-		fieldTypes = append(fieldTypes, field.DoubleType)
-		df = b.dfb.Snapshot()
+	if bf != nil {
+		bf.Close()
 	}
-	if b.sfb != nil {
-		fieldTypes = append(fieldTypes, field.StringType)
-		sf = b.sfb.Snapshot()
+	if intf != nil {
+		intf.Close()
 	}
-	if b.tfb != nil {
-		fieldTypes = append(fieldTypes, field.TimeType)
-		tf = b.tfb.Snapshot()
+	if df != nil {
+		df.Close()
+	}
+	if sf != nil {
+		sf.Close()
+	}
+	if tf != nil {
+		tf.Close()
 	}
 
-	return NewDocsField(b.fieldPath, fieldTypes, nf, bf, intf, df, sf, tf)
+	return nil, nil, err
 }
 
 // Seal seals the builder.

index/field/docs_field_test.go

@@ -17,7 +17,9 @@ func TestDocsFieldBuilderStringSealAndSnapshot(t *testing.T) {
 	}
 	stringArrayPool := pool.NewBucketizedStringArrayPool(stringArrayBuckets, nil)
 	stringArrayPool.Init(func(capacity int) []string { return make([]string, 0, capacity) })
-
+	fieldTypes := field.ValueTypeSet{
+		field.StringType: struct{}{},
+	}
 	opts := NewDocsFieldBuilderOptions().
 		SetStringArrayPool(stringArrayPool)
 	builder := NewDocsFieldBuilder([]string{"testPath"}, opts)
@@ -37,7 +39,10 @@ func TestDocsFieldBuilderStringSealAndSnapshot(t *testing.T) {
 	})
 
 	// Take a snapshot.
-	snapshot1 := builder.Snapshot()
+	snapshot1, remainderTypes, err := builder.SnapshotFor(fieldTypes)
+	require.Equal(t, 0, len(remainderTypes))
+	require.NoError(t, err)
+
 	snapshotField1, _ := snapshot1.StringField()
 	metadata1 := snapshotField1.Values().Metadata()
 	require.Equal(t, 3, metadata1.Size)
@@ -52,7 +57,10 @@ func TestDocsFieldBuilderStringSealAndSnapshot(t *testing.T) {
 	require.Equal(t, 3, metadata1.Size)
 
 	// Take another snapshot.
-	snapshot2 := builder.Snapshot()
+	snapshot2, remainderTypes, err := builder.SnapshotFor(fieldTypes)
+	require.Equal(t, 0, len(remainderTypes))
+	require.NoError(t, err)
+
 	snapshotField2, _ := snapshot2.StringField()
 	metadata2 := snapshotField2.Values().Metadata()
 	require.Equal(t, 4, metadata2.Size)