Repository: polarsignals/frostdb
Pull request #526 (2023-09-12 13:57:10)

Fix physical projections
Previously we weren't performing physical projections on returned arrow records which caused discrepancy with the data that was returned for a given query whether it was in arrow or parquet format. This fixes that discrepancy by running the physical projections on arrow records to drop columns that are not included in the query.


Patch:
diff --git a/pqarrow/arrow.go b/pqarrow/arrow.go
index 68e6f8a..4f3973a 100644
--- a/pqarrow/arrow.go
+++ b/pqarrow/arrow.go
@@ -1186,3 +1186,25 @@ func ColToWriter(col int, writers []MultiColumnWriter) writer.ValueWriter {
 
 	return nil
 }
+
+// Project will project the record according to the given projections.
+func Project(r arrow.Record, projections []logicalplan.Expr) arrow.Record {
+	if len(projections) == 0 {
+		r.Retain() // NOTE: we're creating another reference to this record, so retain it
+		return r
+	}
+
+	cols := make([]arrow.Array, 0, r.Schema().NumFields())
+	fields := make([]arrow.Field, 0, r.Schema().NumFields())
+	for i := 0; i < r.Schema().NumFields(); i++ {
+		for _, projection := range projections {
+			if projection.MatchColumn(r.Schema().Field(i).Name) {
+				cols = append(cols, r.Column(i))
+				fields = append(fields, r.Schema().Field(i))
+				break
+			}
+		}
+	}
+
+	return array.NewRecord(arrow.NewSchema(fields, nil), cols, r.NumRows())
+}
diff --git a/table.go b/table.go
index 3e3baab..0f4dcdd 100644
--- a/table.go
+++ b/table.go
@@ -946,7 +946,9 @@ func (t *Table) Iterator(
 					switch t := rg.(type) {
 					case arrow.Record:
 						defer t.Release()
-						err := callback(ctx, t)
+						r := pqarrow.Project(t, iterOpts.PhysicalProjection)
+						defer r.Release()
+						err := callback(ctx, r)
 						if err != nil {
 							return err
 						}


Test patch:
diff --git a/db_test.go b/db_test.go
index 893781e..093d6fd 100644
--- a/db_test.go
+++ b/db_test.go
@@ -13,6 +13,7 @@ import (
 	"time"
 
 	"github.com/apache/arrow/go/v14/arrow"
+	"github.com/apache/arrow/go/v14/arrow/array"
 	"github.com/apache/arrow/go/v14/arrow/memory"
 	"github.com/go-kit/log/level"
 	"github.com/google/uuid"
@@ -338,26 +339,51 @@ func Test_DB_WithStorage(t *testing.T) {
 		Value:     3,
 	}}
 
-	buf, err := samples.ToBuffer(table.Schema())
+	r, err := samples.ToRecord()
 	require.NoError(t, err)
 
 	ctx := context.Background()
-	_, err = table.InsertBuffer(ctx, buf)
+	_, err = table.InsertRecord(ctx, r)
+	require.NoError(t, err)
+
+	pool := memory.NewGoAllocator()
+	engine := query.NewEngine(pool, db.TableProvider())
+	var inMemory arrow.Record
+	err = engine.ScanTable(t.Name()).
+		Filter(logicalplan.Col("timestamp").GtEq(logicalplan.Literal(2))).
+		Execute(context.Background(), func(ctx context.Context, r arrow.Record) error {
+			r.Retain()
+			inMemory = r
+			return nil
+		})
 	require.NoError(t, err)
 
 	// Gracefully close the db to persist blocks
 	c.Close()
 
-	pool := memory.NewGoAllocator()
-	engine := query.NewEngine(pool, db.TableProvider())
+	c, err = New(
+		WithLogger(logger),
+		WithReadWriteStorage(sinksource),
+	)
+	require.NoError(t, err)
+	defer c.Close()
+
+	db, err = c.DB(context.Background(), t.Name())
+	require.NoError(t, err)
+	engine = query.NewEngine(pool, db.TableProvider())
+	var onDisk arrow.Record
 	err = engine.ScanTable(t.Name()).
 		Filter(logicalplan.Col("timestamp").GtEq(logicalplan.Literal(2))).
 		Execute(context.Background(), func(ctx context.Context, r arrow.Record) error {
-			require.Equal(t, int64(1), r.NumCols())
-			require.Equal(t, int64(2), r.NumRows())
+			r.Retain()
+			onDisk = r
 			return nil
 		})
 	require.NoError(t, err)
+
+	require.True(t, array.RecordEqual(inMemory, onDisk))
+	require.Equal(t, int64(1), onDisk.NumCols())
+	require.Equal(t, int64(2), onDisk.NumRows())
 }
 
 func Test_DB_ColdStart(t *testing.T) {
@@ -1315,7 +1341,7 @@ func Test_DB_TableWrite_ArrowRecord(t *testing.T) {
 		"timestamp filter": {
 			filter: logicalplan.Col("timestamp").GtEq(logicalplan.Literal(12)),
 			rows:   1,
-			cols:   7,
+			cols:   1,
 		},
 		"distinct": {
 			distinct: logicalplan.DynCol("labels"),