feat(post-processing): allow write_tables to compact output LevelDB (…

…#3215)

kythe/go/serving/tools/write_tables/write_tables.go

@@ -23,6 +23,7 @@ import (
 	"errors"
 	"flag"
 	"log"
+	"time"
 
 	"kythe.io/kythe/go/platform/vfs"
 	"kythe.io/kythe/go/services/graphstore"
@@ -61,6 +62,7 @@ var (
 
 	experimentalBeamPipeline = flag.Bool("experimental_beam_pipeline", false, "Whether to use the Beam experimental pipeline implementation")
 	experimentalColumnarData = flag.Bool("experimental_beam_columnar_data", false, "Whether to emit columnar data from the Beam pipeline implementation")
+	compactTable             = flag.Bool("compact_table", false, "Whether to compact the output LevelDB after its creation")
 )
 
 func init() {
@@ -78,6 +80,11 @@ func main() {
 		if err := runExperimentalBeamPipeline(ctx); err != nil {
 			log.Fatalf("Pipeline error: %v", err)
 		}
+		if *compactTable {
+			if err := compactLevelDB(*tablePath); err != nil {
+				log.Fatalf("Error compacting LevelDB: %v", err)
+			}
+		}
 		return
 	}
 
@@ -123,6 +130,17 @@ func main() {
 	}); err != nil {
 		log.Fatal("FATAL ERROR: ", err)
 	}
+
+	if *compactTable {
+		if err := compactLevelDB(*tablePath); err != nil {
+			log.Fatalf("Error compacting LevelDB: %v", err)
+		}
+	}
+}
+
+func compactLevelDB(path string) error {
+	defer func(start time.Time) { log.Printf("Compaction completed in %s", time.Since(start)) }(time.Now())
+	return leveldb.CompactRange(*tablePath, nil)
 }
 
 func runExperimentalBeamPipeline(ctx context.Context) error {

kythe/go/storage/leveldb/leveldb.go

@@ -48,6 +48,25 @@ type levelDB struct {
 	writeOpts     *levigo.WriteOptions
 }
 
+// CompactRange runs a manual compaction on the Range of keys given.
+// If r == nil, the entire table will be compacted.
+func CompactRange(path string, r *keyvalue.Range) error {
+	options := levigo.NewOptions()
+	defer options.Close()
+	db, err := levigo.Open(path, options)
+	if err != nil {
+		return err
+	}
+	lr := levigo.Range{}
+	if r != nil {
+		lr.Start = r.Start
+		lr.Limit = r.End
+	}
+	db.CompactRange(lr)
+	db.Close()
+	return nil
+}
+
 // DefaultOptions is the default Options struct passed to Open when not
 // otherwise given one.
 var DefaultOptions = &Options{