Add example of external sort using BigQueue

Author: mangalaman93

examples/extsort/.gitignore

@@ -0,0 +1,3 @@
+extsort
+*.dat
+bq/*

examples/extsort/main.go

@@ -0,0 +1,107 @@
+package main
+
+import (
+	"bufio"
+	"flag"
+	"fmt"
+	"io"
+	"log"
+	"math/rand"
+	"os"
+	"strconv"
+	"strings"
+	"time"
+)
+
+const (
+	cMaxRandomNum = 10000
+)
+
+func main() {
+	var numDataElements, maxMemSortSize int
+	var inputPath, tempPath, outputPath string
+	flag.IntVar(&numDataElements, "num", 1000, "# elements to generate & sort")
+	flag.IntVar(&maxMemSortSize, "nummem", 100, "# elements to sort in mem")
+	flag.StringVar(&inputPath, "i", "input.dat", "input file (overwritten)")
+	flag.StringVar(&tempPath, "t", "bq", "path to write intermediate data")
+	flag.StringVar(&outputPath, "o", "output.dat", "output file (overwritten)")
+	flag.Parse()
+
+	// generate random data for testing
+	log.Println("generating random dataset")
+	if err := generateData(numDataElements, cMaxRandomNum, inputPath); err != nil {
+		panic(err)
+	}
+
+	// sort the data
+	if err := ExternalSort(inputPath, tempPath, outputPath, maxMemSortSize); err != nil {
+		panic(err)
+	}
+
+	// validate the sorted data
+	log.Println("validating sort")
+	if err := validateSort(outputPath); err != nil {
+		panic(err)
+	}
+}
+
+// generateData generates random data and writes them to a file with newline delimiter
+func generateData(size, maxRandomNum int, dataFilePath string) error {
+	// open file
+	fd, err := os.Create(dataFilePath)
+	if err != nil {
+		return fmt.Errorf("unable to open file :: %v", err)
+	}
+	defer fd.Close()
+	writer := bufio.NewWriter(fd)
+
+	// generate random numbers and write to file
+	rand.Seed(time.Now().UnixNano())
+	for i := 0; i < size; i++ {
+		num := (rand.Int() % (2 * maxRandomNum)) - maxRandomNum
+		str := strconv.Itoa(num) + "\n"
+		writer.WriteString(str)
+	}
+
+	return writer.Flush()
+}
+
+func validateSort(outputPath string) error {
+	// open input file
+	fd, err := os.Open(outputPath)
+	if err != nil {
+		return fmt.Errorf("error in opening input file :: %v", err)
+	}
+	defer fd.Close()
+	reader := bufio.NewReader(fd)
+
+	prevNum := 0
+	firstIter := true
+	for {
+		// each line contains 1 element in the file
+		str, err := reader.ReadString('\n')
+		if err == io.EOF {
+			break
+		} else if err != nil {
+			return fmt.Errorf("error in reading input file :: %v", err)
+		}
+
+		// convert the element into integer
+		str = strings.TrimSpace(str)
+		num, err := strconv.Atoi(str)
+		if err != nil {
+			return fmt.Errorf("error in converting {%s} :: %v", str, err)
+		}
+
+		// check the order of prev and cur element
+		if firstIter {
+			firstIter = false
+		} else if prevNum > num {
+			return fmt.Errorf("incorrect sort")
+		}
+
+		prevNum = num
+	}
+
+	return nil
+}

examples/extsort/sort.go

@@ -0,0 +1,254 @@
+package main
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"log"
+	"os"
+	"path"
+	"sort"
+	"strconv"
+	"strings"
+
+	"github.com/grandecola/bigqueue"
+)
+
+// TODO:
+//  - release resources for bigqueue
+//  - ensure that we do not cross the available memory
+//  - use fewer queues
+
+var (
+	bqFileIndexCount = 0
+)
+
+// ExternalSort perform external sort https://en.wikipedia.org/wiki/External_sorting
+// The inputPath should be a path to a file containing integers in each line
+// The outputPath is similarly formatted file with sorted integers
+// The tempPath is used to write intermediate files
+// maxMemSortSize is number of elements that can be sorted directly in memory
+func ExternalSort(inputPath, tempPath, outputPath string, maxMemSortSize int) error {
+	files, err := ioutil.ReadDir(tempPath)
+	if err != nil {
+		return fmt.Errorf("unable to read temp directory :: %v", err)
+	}
+	if len(files) != 0 {
+		return fmt.Errorf("non-empty temp directory")
+	}
+
+	log.Println("starting divide step")
+	iqs, err := divide(inputPath, tempPath, maxMemSortSize)
+	if err != nil {
+		return fmt.Errorf("error in divide step :: %v", err)
+	}
+
+	log.Println("starting merge step")
+	oq, err := merge(tempPath, iqs)
+	if err != nil {
+		return fmt.Errorf("error in merge step :: %v", err)
+	}
+
+	if err := writeToFile(oq, outputPath); err != nil {
+		return fmt.Errorf("error in writing output to file :: %v", err)
+	}
+
+	return nil
+}
+
+// divide step divides all the input data into sorted group of elements.
+// Each group is persisted to disk using bigqueue interface.
+func divide(inputPath, tempPath string, maxMemSortSize int) ([]*bigqueue.BigQueue, error) {
+	log.Println("reading input file")
+	queues := make([]*bigqueue.BigQueue, 0)
+
+	// open input file
+	fd, err := os.Open(inputPath)
+	if err != nil {
+		return nil, fmt.Errorf("error in opening input file :: %v", err)
+	}
+	defer fd.Close()
+	reader := bufio.NewReader(fd)
+
+	// read all the data from input file and divide it in multiple queues
+	// such that each queue has data sorted and has maximum size of maxMemSortSize
+	elemCount := 0
+	data := make([]int, 0, maxMemSortSize)
+	for {
+		// each line contains 1 element in the file
+		str, err := reader.ReadString('\n')
+		if err == io.EOF {
+			break
+		} else if err != nil {
+			return nil, fmt.Errorf("error in reading input file :: %v", err)
+		}
+
+		// convert the element into integer
+		str = strings.TrimSpace(str)
+		num, err := strconv.Atoi(str)
+		if err != nil {
+			return nil, fmt.Errorf("error in converting {%s} :: %v", str, err)
+		}
+		elemCount++
+		data = append(data, num)
+
+		// check whether we have enough element to perform in memory sort
+		if elemCount < maxMemSortSize {
+			continue
+		}
+
+		// if yes, add the sorted elements into the queue
+		sort.Ints(data)
+		bq, err := buildBigQueue(tempPath, data)
+		if err != nil {
+			return nil, fmt.Errorf("error in building bigqueue :: %v", err)
+		}
+
+		// add the queue in the list and truncate the slice that holds data in memory
+		queues = append(queues, bq)
+		elemCount = 0
+		data = data[:0]
+	}
+
+	// write the final list of elements to bigqueue
+	if len(data) != 0 {
+		sort.Ints(data)
+		bq, err := buildBigQueue(tempPath, data)
+		if err != nil {
+			return nil, fmt.Errorf("error in building bigqueue :: %v", err)
+		}
+
+		queues = append(queues, bq)
+	}
+
+	return queues, nil
+}
+
+// merge step merges the sorted group of elements stored in bigqueue using bigqueue
+func merge(tempPath string, queues []*bigqueue.BigQueue) (*bigqueue.BigQueue, error) {
+	currentQueues := queues
+	nextQueues := make([]*bigqueue.BigQueue, 0)
+	for iteration := 0; len(currentQueues) != 1; iteration++ {
+		log.Printf("iteration %d, # queues %d\n", iteration, len(currentQueues))
+
+		for i := 0; i < len(currentQueues); i += 2 {
+			// if only one queue is left, just add this queue
+			q1 := currentQueues[i]
+			if i+1 >= len(currentQueues) {
+				nextQueues = append(nextQueues, q1)
+				continue
+			}
+
+			// otherwise, merge the two queues
+			q2 := currentQueues[i+1]
+			mq, err := mergeQueues(q1, q2, tempPath)
+			if err != nil {
+				return nil, fmt.Errorf("error in merging two queues :: %v", err)
+			}
+
+			nextQueues = append(nextQueues, mq)
+		}
+
+		currentQueues = nextQueues
+		nextQueues = make([]*bigqueue.BigQueue, 0)
+	}
+
+	return currentQueues[0], nil
+}
+
+func buildBigQueue(tempPath string, data []int) (*bigqueue.BigQueue, error) {
+	bq, err := bigqueue.NewBigQueue(getTempDir(tempPath))
+	if err != nil {
+		return nil, fmt.Errorf("unable to init bigqueue :: %v", err)
+	}
+
+	// write all the data to bigqueue
+	for _, e := range data {
+		if err := bq.Enqueue([]byte(strconv.Itoa(e))); err != nil {
+			return nil, fmt.Errorf("unable to write to bigqueue :: %v", err)
+		}
+	}
+
+	return bq, nil
+}
+
+func getTempDir(tempPath string) string {
+	queueDir := "q" + strconv.Itoa(bqFileIndexCount)
+	bqFileIndexCount++
+
+	queuePath := path.Join(tempPath, queueDir)
+	if err := os.MkdirAll(queuePath, 0700); err != nil {
+		panic(err)
+	}
+
+	return queuePath
+}
+
+func mergeQueues(q1, q2 *bigqueue.BigQueue, tempPath string) (*bigqueue.BigQueue, error) {
+	mq, err := bigqueue.NewBigQueue(getTempDir(tempPath))
+	if err != nil {
+		return nil, fmt.Errorf("unable to create bigqueue :: %v", err)
+	}
+
+	for !q1.IsEmpty() && !q2.IsEmpty() {
+		e1, err1 := q1.Peek()
+		e2, err2 := q2.Peek()
+		if err1 != nil || err2 != nil {
+			return nil, fmt.Errorf("unable to dequeue :: %v || %v", err1, err2)
+		}
+
+		num1, err1 := strconv.Atoi(string(e1))
+		num2, err2 := strconv.Atoi(string(e2))
+		if err1 != nil || err2 != nil {
+			return nil, fmt.Errorf("error in conversion :: %v || %v", err1, err2)
+		}
+
+		if num1 < num2 {
+			q1.Dequeue()
+			mq.Enqueue([]byte(strconv.Itoa(num1)))
+		} else {
+			q2.Dequeue()
+			mq.Enqueue([]byte(strconv.Itoa(num2)))
+		}
+	}
+
+	// add elements from the non-empty queue
+	var lq *bigqueue.BigQueue
+	if q1.IsEmpty() {
+		lq = q1
+	} else {
+		lq = q2
+	}
+	for !lq.IsEmpty() {
+		e1, err := lq.Dequeue()
+		if err != nil {
+			return nil, fmt.Errorf("unable to dequeue :: %v", err)
+		}
+
+		mq.Enqueue(e1)
+	}
+
+	return mq, nil
+}
+
+func writeToFile(oq *bigqueue.BigQueue, outputPath string) error {
+	// write the final output to file
+	od, err := os.Create(outputPath)
+	if err != nil {
+		return fmt.Errorf("error in opening input file :: %v", err)
+	}
+	defer od.Close()
+
+	w := bufio.NewWriter(od)
+	for !oq.IsEmpty() {
+		v, err := oq.Dequeue()
+		if err != nil {
+			return fmt.Errorf("unable to dequeue from bigqueue :: %v", err)
+		}
+
+		w.WriteString(string(v) + "\n")
+	}
+
+	return w.Flush()
+}