Repository: polarsignals/frostdb
Pull request #522 (2023-09-06 21:00:15)

Query engine in memory only
Adds an option to the query engine to perform queries only against the in memory portion of the db.


Patch:
diff --git a/db.go b/db.go
index db22eed..3ac45fa 100644
--- a/db.go
+++ b/db.go
@@ -1144,7 +1144,9 @@ func (db *DB) dropStorage() error {
 	trashDir := db.trashDir()
 
 	if moveErr := func() error {
-		_ = os.Mkdir(trashDir, os.FileMode(0o755))
+		if err := os.MkdirAll(trashDir, os.FileMode(0o755)); err != nil {
+			return fmt.Errorf("making trash dir: %w", err)
+		}
 		// Create a temporary directory in the trash dir to avoid clashing
 		// with other wal/snapshot dirs that might not have been removed
 		// previously.
diff --git a/query/logicalplan/logicalplan.go b/query/logicalplan/logicalplan.go
index 5586bb8..d3c7ba1 100644
--- a/query/logicalplan/logicalplan.go
+++ b/query/logicalplan/logicalplan.go
@@ -35,10 +35,17 @@ type IterOptions struct {
 	Projection         []Expr
 	Filter             Expr
 	DistinctColumns    []Expr
+	InMemoryOnly       bool
 }
 
 type Option func(opts *IterOptions)
 
+func WithInMemoryOnly() Option {
+	return func(opts *IterOptions) {
+		opts.InMemoryOnly = true
+	}
+}
+
 func WithPhysicalProjection(e ...Expr) Option {
 	return func(opts *IterOptions) {
 		opts.PhysicalProjection = append(opts.PhysicalProjection, e...)
@@ -177,6 +184,9 @@ type TableScan struct {
 
 	// Projection is the list of columns that are to be projected.
 	Projection []Expr
+
+	// SkipSources indicates to skip scanning the tables sources.
+	SkipSources bool
 }
 
 func (scan *TableScan) String() string {
@@ -204,6 +214,9 @@ type SchemaScan struct {
 
 	// Projection is the list of columns that are to be projected.
 	Projection []Expr
+
+	// SkipSources indicates to skip scanning the tables sources.
+	SkipSources bool
 }
 
 func (s *SchemaScan) String() string {
diff --git a/query/physicalplan/physicalplan.go b/query/physicalplan/physicalplan.go
index 4c54955..164abf7 100644
--- a/query/physicalplan/physicalplan.go
+++ b/query/physicalplan/physicalplan.go
@@ -128,6 +128,17 @@ func (s *TableScan) Execute(ctx context.Context, pool memory.Allocator) error {
 		}
 	}()
 
+	opts := []logicalplan.Option{
+		logicalplan.WithPhysicalProjection(s.options.PhysicalProjection...),
+		logicalplan.WithProjection(s.options.Projection...),
+		logicalplan.WithFilter(s.options.Filter),
+		logicalplan.WithDistinctColumns(s.options.Distinct...),
+	}
+
+	if s.options.SkipSources {
+		opts = append(opts, logicalplan.WithInMemoryOnly())
+	}
+
 	errg, _ := errgroup.WithContext(ctx)
 	errg.Go(recovery.Do(func() error {
 		return table.View(ctx, func(ctx context.Context, tx uint64) error {
@@ -136,10 +147,7 @@ func (s *TableScan) Execute(ctx context.Context, pool memory.Allocator) error {
 				tx,
 				pool,
 				callbacks,
-				logicalplan.WithPhysicalProjection(s.options.PhysicalProjection...),
-				logicalplan.WithProjection(s.options.Projection...),
-				logicalplan.WithFilter(s.options.Filter),
-				logicalplan.WithDistinctColumns(s.options.Distinct...),
+				opts...,
 			)
 		})
 	}))
@@ -187,6 +195,17 @@ func (s *SchemaScan) Execute(ctx context.Context, pool memory.Allocator) error {
 		callbacks = append(callbacks, plan.Callback)
 	}
 
+	opts := []logicalplan.Option{
+		logicalplan.WithPhysicalProjection(s.options.PhysicalProjection...),
+		logicalplan.WithProjection(s.options.Projection...),
+		logicalplan.WithFilter(s.options.Filter),
+		logicalplan.WithDistinctColumns(s.options.Distinct...),
+	}
+
+	if s.options.SkipSources {
+		opts = append(opts, logicalplan.WithInMemoryOnly())
+	}
+
 	errg, _ := errgroup.WithContext(ctx)
 	errg.Go(recovery.Do(func() error {
 		return table.View(ctx, func(ctx context.Context, tx uint64) error {
@@ -195,10 +214,7 @@ func (s *SchemaScan) Execute(ctx context.Context, pool memory.Allocator) error {
 				tx,
 				pool,
 				callbacks,
-				logicalplan.WithPhysicalProjection(s.options.PhysicalProjection...),
-				logicalplan.WithProjection(s.options.Projection...),
-				logicalplan.WithFilter(s.options.Filter),
-				logicalplan.WithDistinctColumns(s.options.Distinct...),
+				opts...,
 			)
 		})
 	}))
@@ -247,10 +263,17 @@ func (p *noopOperator) Draw() *Diagram {
 type execOptions struct {
 	orderedAggregations bool
 	overrideInput       []PhysicalPlan
+	skipSources         bool
 }
 
 type Option func(o *execOptions)
 
+func WithInMemoryOnly() Option {
+	return func(o *execOptions) {
+		o.skipSources = true
+	}
+}
+
 func WithOrderedAggregations() Option {
 	return func(o *execOptions) {
 		o.orderedAggregations = true
@@ -304,6 +327,7 @@ func Build(
 			for i := range plans {
 				plans[i] = &noopOperator{}
 			}
+			plan.SchemaScan.SkipSources = execOpts.skipSources
 			outputPlan.scan = &SchemaScan{
 				tracer:  tracer,
 				options: plan.SchemaScan,
@@ -318,6 +342,7 @@ func Build(
 			for i := range plans {
 				plans[i] = &noopOperator{}
 			}
+			plan.TableScan.SkipSources = execOpts.skipSources
 			outputPlan.scan = &TableScan{
 				tracer:  tracer,
 				options: plan.TableScan,
diff --git a/table.go b/table.go
index d0dc468..3e3baab 100644
--- a/table.go
+++ b/table.go
@@ -977,7 +977,7 @@ func (t *Table) Iterator(
 	}
 
 	errg.Go(func() error {
-		if err := t.collectRowGroups(ctx, tx, iterOpts.Filter, rowGroups); err != nil {
+		if err := t.collectRowGroups(ctx, tx, iterOpts.Filter, iterOpts.InMemoryOnly, rowGroups); err != nil {
 			return err
 		}
 		close(rowGroups)
@@ -1073,7 +1073,7 @@ func (t *Table) SchemaIterator(
 	}
 
 	errg.Go(func() error {
-		if err := t.collectRowGroups(ctx, tx, iterOpts.Filter, rowGroups); err != nil {
+		if err := t.collectRowGroups(ctx, tx, iterOpts.Filter, iterOpts.InMemoryOnly, rowGroups); err != nil {
 			return err
 		}
 		close(rowGroups)
@@ -1741,6 +1741,7 @@ func (t *Table) collectRowGroups(
 	ctx context.Context,
 	tx uint64,
 	filterExpr logicalplan.Expr,
+	skipSources bool,
 	rowGroups chan<- any,
 ) error {
 	ctx, span := t.tracer.Start(ctx, "Table/collectRowGroups")
@@ -1768,6 +1769,10 @@ func (t *Table) collectRowGroups(
 		}
 	}
 
+	if skipSources {
+		return nil
+	}
+
 	// Collect from all other data sources.
 	for _, source := range t.db.sources {
 		span.AddEvent(fmt.Sprintf("source/%s", source.String()))


Test patch:
diff --git a/db_test.go b/db_test.go
index c0fb181..797d04e 100644
--- a/db_test.go
+++ b/db_test.go
@@ -23,6 +23,7 @@ import (
 	schemapb "github.com/polarsignals/frostdb/gen/proto/go/frostdb/schema/v1alpha1"
 	"github.com/polarsignals/frostdb/query"
 	"github.com/polarsignals/frostdb/query/logicalplan"
+	"github.com/polarsignals/frostdb/query/physicalplan"
 )
 
 func TestDBWithWALAndBucket(t *testing.T) {
@@ -1840,3 +1841,77 @@ func Test_DB_DropStorage(t *testing.T) {
 	require.NoError(t, err)
 	countRows(0)
 }
+
+func Test_DB_EngineInMemory(t *testing.T) {
+	config := NewTableConfig(
+		dynparquet.SampleDefinition(),
+	)
+
+	logger := newTestLogger(t)
+
+	dir := t.TempDir()
+	bucket := objstore.NewInMemBucket()
+
+	sinksource := NewDefaultObjstoreBucket(bucket)
+
+	c, err := New(
+		WithLogger(logger),
+		WithStoragePath(dir),
+		WithReadWriteStorage(sinksource),
+		WithActiveMemorySize(100*1024),
+	)
+	require.NoError(t, err)
+	db, err := c.DB(context.Background(), "test")
+	require.NoError(t, err)
+	table, err := db.Table("test", config)
+	require.NoError(t, err)
+
+	samples := dynparquet.NewTestSamples()
+
+	ctx := context.Background()
+	for i := 0; i < 100; i++ {
+		buf, err := samples.ToBuffer(table.Schema())
+		require.NoError(t, err)
+		_, err = table.InsertBuffer(ctx, buf)
+		require.NoError(t, err)
+	}
+	require.NoError(t, c.Close())
+
+	c, err = New(
+		WithLogger(logger),
+		WithStoragePath(dir),
+		WithReadWriteStorage(sinksource),
+		WithActiveMemorySize(100*1024),
+	)
+	require.NoError(t, err)
+	defer c.Close()
+
+	db, err = c.DB(context.Background(), "test")
+	require.NoError(t, err)
+
+	pool := memory.NewCheckedAllocator(memory.DefaultAllocator)
+	defer pool.AssertSize(t, 0)
+	engine := query.NewEngine(pool, db.TableProvider(), query.WithPhysicalplanOptions(physicalplan.WithInMemoryOnly()))
+	err = engine.ScanTable("test").
+		Aggregate(
+			[]logicalplan.Expr{logicalplan.Sum(logicalplan.Col("value"))},
+			[]logicalplan.Expr{logicalplan.Col("labels.namespace")},
+		).
+		Execute(context.Background(), func(ctx context.Context, r arrow.Record) error {
+			t.FailNow() // should not be called
+			return nil
+		})
+	require.NoError(t, err)
+
+	engine = query.NewEngine(pool, db.TableProvider())
+	err = engine.ScanTable("test").
+		Aggregate(
+			[]logicalplan.Expr{logicalplan.Sum(logicalplan.Col("value"))},
+			[]logicalplan.Expr{logicalplan.Col("labels.namespace")},
+		).
+		Execute(context.Background(), func(ctx context.Context, r arrow.Record) error {
+			require.Equal(t, int64(2), r.NumRows())
+			return nil
+		})
+	require.NoError(t, err)
+}