Repository: polarsignals/frostdb
Pull request #750 (2024-02-20 15:17:40)

*: change WithInMemory physicalplan option to WithReadMode
This allows introducing different read modes such as ReadModeObjectStorage which reads only from object storage. WithInMemory is now WithReadMode(ReadModeInMemory).


Patch:
diff --git a/query/logicalplan/logicalplan.go b/query/logicalplan/logicalplan.go
index e9a3ce5..6ab845f 100644
--- a/query/logicalplan/logicalplan.go
+++ b/query/logicalplan/logicalplan.go
@@ -36,14 +36,14 @@ type IterOptions struct {
 	Projection         []Expr
 	Filter             Expr
 	DistinctColumns    []Expr
-	InMemoryOnly       bool
+	ReadMode           ReadMode
 }
 
 type Option func(opts *IterOptions)
 
-func WithInMemoryOnly() Option {
+func WithReadMode(m ReadMode) Option {
 	return func(opts *IterOptions) {
-		opts.InMemoryOnly = true
+		opts.ReadMode = m
 	}
 }
 
@@ -250,8 +250,8 @@ 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
+	// ReadMode indicates the mode to use when reading.
+	ReadMode ReadMode
 }
 
 func (scan *TableScan) DataTypeForExpr(expr Expr) (arrow.DataType, error) {
@@ -322,6 +322,18 @@ func (scan *TableScan) String() string {
 		" Distinct: " + fmt.Sprint(scan.Distinct)
 }
 
+type ReadMode int
+
+const (
+	// ReadModeDefault is the default read mode. Reads from in-memory and object
+	// storage.
+	ReadModeDefault ReadMode = iota
+	// ReadModeInMemoryOnly reads from in-memory storage only.
+	ReadModeInMemoryOnly
+	// ReadModeDataSourcesOnly reads from data sources only.
+	ReadModeDataSourcesOnly
+)
+
 type SchemaScan struct {
 	TableProvider TableProvider
 	TableName     string
@@ -340,8 +352,8 @@ 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
+	// ReadMode indicates the mode to use when reading.
+	ReadMode ReadMode
 }
 
 func (s *SchemaScan) String() string {
diff --git a/query/physicalplan/physicalplan.go b/query/physicalplan/physicalplan.go
index d43eed9..b91b4c1 100644
--- a/query/physicalplan/physicalplan.go
+++ b/query/physicalplan/physicalplan.go
@@ -134,10 +134,7 @@ func (s *TableScan) Execute(ctx context.Context, pool memory.Allocator) error {
 		logicalplan.WithProjection(s.options.Projection...),
 		logicalplan.WithFilter(s.options.Filter),
 		logicalplan.WithDistinctColumns(s.options.Distinct...),
-	}
-
-	if s.options.SkipSources {
-		opts = append(opts, logicalplan.WithInMemoryOnly())
+		logicalplan.WithReadMode(s.options.ReadMode),
 	}
 
 	errg, _ := errgroup.WithContext(ctx)
@@ -197,10 +194,7 @@ func (s *SchemaScan) Execute(ctx context.Context, pool memory.Allocator) error {
 		logicalplan.WithProjection(s.options.Projection...),
 		logicalplan.WithFilter(s.options.Filter),
 		logicalplan.WithDistinctColumns(s.options.Distinct...),
-	}
-
-	if s.options.SkipSources {
-		opts = append(opts, logicalplan.WithInMemoryOnly())
+		logicalplan.WithReadMode(s.options.ReadMode),
 	}
 
 	errg, _ := errgroup.WithContext(ctx)
@@ -260,14 +254,14 @@ func (p *noopOperator) Draw() *Diagram {
 type execOptions struct {
 	orderedAggregations bool
 	overrideInput       []PhysicalPlan
-	skipSources         bool
+	readMode            logicalplan.ReadMode
 }
 
 type Option func(o *execOptions)
 
-func WithInMemoryOnly() Option {
+func WithReadMode(m logicalplan.ReadMode) Option {
 	return func(o *execOptions) {
-		o.skipSources = true
+		o.readMode = m
 	}
 }
 
@@ -324,7 +318,7 @@ func Build(
 			for i := range plans {
 				plans[i] = &noopOperator{}
 			}
-			plan.SchemaScan.SkipSources = execOpts.skipSources
+			plan.SchemaScan.ReadMode = execOpts.readMode
 			outputPlan.scan = &SchemaScan{
 				tracer:  tracer,
 				options: plan.SchemaScan,
@@ -339,7 +333,7 @@ func Build(
 			for i := range plans {
 				plans[i] = &noopOperator{}
 			}
-			plan.TableScan.SkipSources = execOpts.skipSources
+			plan.TableScan.ReadMode = execOpts.readMode
 			outputPlan.scan = &TableScan{
 				tracer:  tracer,
 				options: plan.TableScan,
diff --git a/table.go b/table.go
index 820adbb..7d2048d 100644
--- a/table.go
+++ b/table.go
@@ -882,7 +882,7 @@ func (t *Table) Iterator(
 
 	errg.Go(func() error {
 		defer close(rowGroups)
-		return t.collectRowGroups(ctx, tx, iterOpts.Filter, iterOpts.InMemoryOnly, rowGroups)
+		return t.collectRowGroups(ctx, tx, iterOpts.Filter, iterOpts.ReadMode, rowGroups)
 	})
 
 	return errg.Wait()
@@ -1003,7 +1003,7 @@ func (t *Table) SchemaIterator(
 	}
 
 	errg.Go(func() error {
-		if err := t.collectRowGroups(ctx, tx, iterOpts.Filter, iterOpts.InMemoryOnly, rowGroups); err != nil {
+		if err := t.collectRowGroups(ctx, tx, iterOpts.Filter, iterOpts.ReadMode, rowGroups); err != nil {
 			return err
 		}
 		close(rowGroups)
@@ -1201,7 +1201,7 @@ func (t *Table) collectRowGroups(
 	ctx context.Context,
 	tx uint64,
 	filterExpr logicalplan.Expr,
-	skipSources bool,
+	readMode logicalplan.ReadMode,
 	rowGroups chan<- any,
 ) error {
 	ctx, span := t.tracer.Start(ctx, "Table/collectRowGroups")
@@ -1211,29 +1211,33 @@ func (t *Table) collectRowGroups(
 	// to avoid to iterate on them again while reading the block file
 	// we keep the last block timestamp to be read from the bucket and pass it to the IterateBucketBlocks() function
 	// so that every block with a timestamp >= lastReadBlockTimestamp is discarded while being read.
-	memoryBlocks, lastBlockTimestamp := t.memoryBlocks()
-	defer func() {
+	var lastBlockTimestamp uint64
+	if readMode != logicalplan.ReadModeDataSourcesOnly {
+		memoryBlocks, lbt := t.memoryBlocks()
+		lastBlockTimestamp = lbt
+		defer func() {
+			for _, block := range memoryBlocks {
+				block.pendingReadersWg.Done()
+			}
+		}()
 		for _, block := range memoryBlocks {
-			block.pendingReadersWg.Done()
-		}
-	}()
-	for _, block := range memoryBlocks {
-		if err := block.index.Scan(ctx, "", t.schema, filterExpr, tx, func(ctx context.Context, v any) error {
-			select {
-			case <-ctx.Done():
-				if rg, ok := v.(index.ReleaseableRowGroup); ok {
-					rg.Release()
+			if err := block.index.Scan(ctx, "", t.schema, filterExpr, tx, func(ctx context.Context, v any) error {
+				select {
+				case <-ctx.Done():
+					if rg, ok := v.(index.ReleaseableRowGroup); ok {
+						rg.Release()
+					}
+					return ctx.Err()
+				case rowGroups <- v:
+					return nil
 				}
-				return ctx.Err()
-			case rowGroups <- v:
-				return nil
+			}); err != nil {
+				return err
 			}
-		}); err != nil {
-			return err
 		}
 	}
 
-	if skipSources {
+	if readMode == logicalplan.ReadModeInMemoryOnly {
 		return nil
 	}
 


Test patch:
diff --git a/db_test.go b/db_test.go
index 3fe1c51..df7e06e 100644
--- a/db_test.go
+++ b/db_test.go
@@ -1713,7 +1713,7 @@ func Test_DB_EngineInMemory(t *testing.T) {
 
 	pool := memory.NewCheckedAllocator(memory.DefaultAllocator)
 	defer pool.AssertSize(t, 0)
-	engine := query.NewEngine(pool, db.TableProvider(), query.WithPhysicalplanOptions(physicalplan.WithInMemoryOnly()))
+	engine := query.NewEngine(pool, db.TableProvider(), query.WithPhysicalplanOptions(physicalplan.WithReadMode(logicalplan.ReadModeInMemoryOnly)))
 	err = engine.ScanTable("test").
 		Aggregate(
 			[]*logicalplan.AggregationFunction{logicalplan.Sum(logicalplan.Col("value"))},