Merge pull request #2 from SkyAPM/external-seg

Add external segment streaming support with deduplication

Author: hanahmily

.github/workflows/lint.yml

@@ -13,4 +13,4 @@ jobs:
       - name: golangci-lint
         uses: golangci/[email protected]
         with:
-          version: v1.61.0
+          version: v1.64.8

.github/workflows/tests.yml

@@ -8,7 +8,7 @@ jobs:
   test:
     strategy:
       matrix:
-        go-version: [1.23.x]
+        go-version: [1.24.x]
         platform: [ubuntu-latest, macos-latest, windows-latest]
     runs-on: ${{ matrix.platform }}
     steps:

config.go

@@ -39,6 +39,11 @@ type Config struct {
 
 	SearchStartFunc func(size uint64) error
 	SearchEndFunc   func(size uint64)
+
+	// External segment streaming options
+	EnableExternalSegments bool
+	EnableDeduplication    bool
+	ExternalSegmentTempDir string
 }
 
 // WithVirtualField allows you to describe a field that
@@ -87,6 +92,24 @@ func (config Config) WithPrepareMergeCallback(f func(src []*roaring.Bitmap, segm
 	return config
 }
 
+func (config Config) WithExternalSegments(tempDir string, enableDeduplication bool) Config {
+	config.EnableExternalSegments = true
+	config.EnableDeduplication = enableDeduplication
+	config.ExternalSegmentTempDir = tempDir
+
+	// Propagate to index config
+	config.indexConfig.EnableExternalSegments = true
+	config.indexConfig.EnableDeduplication = enableDeduplication
+	config.indexConfig.ExternalSegmentTempDir = tempDir
+
+	return config
+}
+
+// DefaultExternalSegmentConfig provides sensible defaults for external segment streaming
+func DefaultExternalSegmentConfig(tempDir string) Config {
+	return Config{}.WithExternalSegments(tempDir, true)
+}
+
 func DefaultConfig(path string) Config {
 	indexConfig := index.DefaultConfig(path)
 	return defaultConfig(indexConfig)
@@ -124,6 +147,12 @@ func defaultConfig(indexConfig index.Config) Config {
 		}
 		return rv.DefaultSimilarity.ComputeNorm(length)
 	})
+
+	// Propagate external segment settings from main config to index config
+	indexConfig.EnableExternalSegments = rv.EnableExternalSegments
+	indexConfig.EnableDeduplication = rv.EnableDeduplication
+	indexConfig.ExternalSegmentTempDir = rv.ExternalSegmentTempDir
+
 	rv.indexConfig = indexConfig
 
 	return rv

external_segment_example.go

@@ -0,0 +1,141 @@
+//  Copyright (c) 2020 The Bluge Authors.
+//
+// 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 bluge
+
+import (
+	"fmt"
+	"io"
+	"log"
+)
+
+// ExampleExternalSegmentStreaming demonstrates how to use the external segment streaming feature
+func ExampleExternalSegmentStreaming() {
+	// Configure writer with external segment support
+	config := DefaultConfig("./index").WithExternalSegments("./temp_segments", true)
+
+	writer, err := OpenWriter(config)
+	if err != nil {
+		log.Printf("failed to open writer: %v", err)
+		return
+	}
+	defer writer.Close()
+
+	// Enable external segment receiver
+	receiver, err := writer.EnableExternalSegments()
+	if err != nil {
+		log.Printf("failed to enable external segments: %v", err)
+		return
+	}
+
+	// Start receiving a segment
+	err = receiver.StartSegment()
+	if err != nil {
+		log.Printf("failed to start segment: %v", err)
+		return
+	}
+
+	// Simulate streaming chunks from an external source
+	// In a real scenario, this would come from a network stream, file, etc.
+	var externalSegmentData []byte // This would be your actual segment data
+	chunkSize := 64 * 1024         // 64KB chunks
+
+	for len(externalSegmentData) > 0 {
+		chunkEnd := chunkSize
+		if chunkEnd > len(externalSegmentData) {
+			chunkEnd = len(externalSegmentData)
+		}
+
+		chunk := externalSegmentData[:chunkEnd]
+		externalSegmentData = externalSegmentData[chunkEnd:]
+
+		err = receiver.WriteChunk(chunk)
+		if err != nil {
+			log.Printf("failed to write chunk: %v", err)
+			return
+		}
+	}
+
+	// Signal completion of segment streaming
+	err = receiver.CompleteSegment()
+	if err != nil {
+		log.Printf("failed to complete segment: %v", err)
+		return
+	}
+
+	fmt.Printf("External segment streamed successfully\n")
+	fmt.Printf("Status: %v\n", receiver.Status())
+	fmt.Printf("Bytes received: %d\n", receiver.BytesReceived())
+}
+
+// ExampleExternalSegmentStreamingFromReader demonstrates streaming from an io.Reader
+func ExampleExternalSegmentStreamingFromReader(segmentReader io.Reader) error {
+	// Configure writer with external segment support
+	config := DefaultConfig("./index").WithExternalSegments("./temp_segments", true)
+
+	writer, err := OpenWriter(config)
+	if err != nil {
+		return err
+	}
+	defer writer.Close()
+
+	// Enable external segment receiver
+	receiver, err := writer.EnableExternalSegments()
+	if err != nil {
+		return err
+	}
+
+	// Start receiving a segment
+	err = receiver.StartSegment()
+	if err != nil {
+		return err
+	}
+
+	// Stream data from reader
+	chunkSize := 64 * 1024 // 64KB chunks
+	buffer := make([]byte, chunkSize)
+
+	for {
+		n, err := segmentReader.Read(buffer)
+		if err != nil && err != io.EOF {
+			if recoverErr := receiver.RecoverFromFailure(); recoverErr != nil {
+				return fmt.Errorf("failed to recover from failure: %w, original error: %w", recoverErr, err)
+			}
+			return err
+		}
+
+		if n == 0 {
+			break // End of data
+		}
+
+		// Write chunk
+		err = receiver.WriteChunk(buffer[:n])
+		if err != nil {
+			return err
+		}
+
+		if err == io.EOF {
+			break
+		}
+	}
+
+	// Signal completion
+	err = receiver.CompleteSegment()
+	if err != nil {
+		return err
+	}
+
+	fmt.Printf("External segment streamed successfully from reader\n")
+	return nil
+}