Repository: polarsignals/frostdb
Pull request #530 (2023-09-13 19:53:54)

Project all
Since fixing the projection pushdowns to also apply to in-memory arrow records we no longer have a way to wildcard select all columns.

This adds a new logicalplan Expr that allows projection of all columns.


Patch:
diff --git a/query/logicalplan/expr.go b/query/logicalplan/expr.go
index 185c6a7..6035e2e 100644
--- a/query/logicalplan/expr.go
+++ b/query/logicalplan/expr.go
@@ -690,3 +690,27 @@ func (a *AverageExpr) Accept(visitor Visitor) bool {
 
 	return visitor.PostVisit(a)
 }
+
+type AllExpr struct{}
+
+func All() *AllExpr {
+	return &AllExpr{}
+}
+
+func (a *AllExpr) DataType(*parquet.Schema) (arrow.DataType, error) { return nil, nil }
+func (a *AllExpr) Accept(visitor Visitor) bool {
+	continu := visitor.PreVisit(a)
+	if !continu {
+		return false
+	}
+
+	return visitor.PostVisit(a)
+}
+func (a *AllExpr) Name() string { return "all" }
+func (a *AllExpr) ColumnsUsedExprs() []Expr {
+	return []Expr{&AllExpr{}}
+}
+func (a *AllExpr) MatchColumn(columnName string) bool { return true }
+func (a *AllExpr) MatchPath(path string) bool         { return true }
+func (a *AllExpr) Computed() bool                     { return false }
+func (a *AllExpr) Clone() Expr                        { return &AllExpr{} }
diff --git a/query/physicalplan/physicalplan.go b/query/physicalplan/physicalplan.go
index 164abf7..5e2da04 100644
--- a/query/physicalplan/physicalplan.go
+++ b/query/physicalplan/physicalplan.go
@@ -351,6 +351,11 @@ func Build(
 			prev = append(prev[:0], plans...)
 			oInfo.nodeMaintainsOrdering()
 		case plan.Projection != nil:
+			for _, e := range plan.Projection.Exprs { // Don't build the projection if it's a wildcard, the projection pushdown optimization will handle it.
+				if e.Name() == "all" {
+					return true
+				}
+			}
 			// For each previous physical plan create one Projection
 			for i := range prev {
 				p, err := Project(pool, tracer, plan.Projection.Exprs)
diff --git a/query/physicalplan/project.go b/query/physicalplan/project.go
index 09b1a96..1dbf693 100644
--- a/query/physicalplan/project.go
+++ b/query/physicalplan/project.go
@@ -168,6 +168,8 @@ func (p dynamicProjection) Project(mem memory.Allocator, ar arrow.Record) ([]arr
 
 func projectionFromExpr(expr logicalplan.Expr) (columnProjection, error) {
 	switch e := expr.(type) {
+	case *logicalplan.AllExpr:
+		return allProjection{}, nil
 	case *logicalplan.Column:
 		return plainProjection{
 			expr: e,
@@ -356,3 +358,11 @@ func avgInt64arrays(pool memory.Allocator, sums, counts arrow.Array) arrow.Array
 
 	return res.NewArray()
 }
+
+type allProjection struct{}
+
+func (a allProjection) Name() string { return "all" }
+
+func (a allProjection) Project(mem memory.Allocator, ar arrow.Record) ([]arrow.Field, []arrow.Array, error) {
+	return ar.Schema().Fields(), ar.Columns(), nil
+}


Test patch:
diff --git a/db_test.go b/db_test.go
index 093d6fd..b01c63c 100644
--- a/db_test.go
+++ b/db_test.go
@@ -1986,3 +1986,82 @@ func Test_DB_SnapshotOnClose(t *testing.T) {
 	}))
 	require.True(t, found)
 }
+
+func Test_DB_All(t *testing.T) {
+	config := NewTableConfig(
+		dynparquet.SampleDefinition(),
+	)
+
+	logger := newTestLogger(t)
+
+	c, err := New(
+		WithLogger(logger),
+	)
+	require.NoError(t, err)
+	defer c.Close()
+
+	db, err := c.DB(context.Background(), t.Name())
+	require.NoError(t, err)
+	defer os.RemoveAll(t.Name())
+	table, err := db.Table(t.Name(), config)
+	require.NoError(t, err)
+
+	samples := dynparquet.Samples{{
+		ExampleType: "test",
+		Labels: []dynparquet.Label{
+			{Name: "label1", Value: "value1"},
+			{Name: "label2", Value: "value2"},
+		},
+		Stacktrace: []uuid.UUID{
+			{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1},
+			{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x2},
+		},
+		Timestamp: 1,
+		Value:     1,
+	}, {
+		ExampleType: "test",
+		Labels: []dynparquet.Label{
+			{Name: "label1", Value: "value2"},
+			{Name: "label2", Value: "value2"},
+			{Name: "label3", Value: "value3"},
+		},
+		Stacktrace: []uuid.UUID{
+			{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1},
+			{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x2},
+		},
+		Timestamp: 2,
+		Value:     2,
+	}, {
+		ExampleType: "test",
+		Labels: []dynparquet.Label{
+			{Name: "label1", Value: "value3"},
+			{Name: "label2", Value: "value2"},
+			{Name: "label4", Value: "value4"},
+		},
+		Stacktrace: []uuid.UUID{
+			{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1},
+			{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x2},
+		},
+		Timestamp: 3,
+		Value:     3,
+	}}
+
+	r, err := samples.ToRecord()
+	require.NoError(t, err)
+
+	ctx := context.Background()
+	_, err = table.InsertRecord(ctx, r)
+	require.NoError(t, err)
+
+	pool := memory.NewGoAllocator()
+	engine := query.NewEngine(pool, db.TableProvider())
+	err = engine.ScanTable(t.Name()).
+		Project(logicalplan.All()).
+		Filter(logicalplan.Col("timestamp").GtEq(logicalplan.Literal(2))).
+		Execute(context.Background(), func(ctx context.Context, r arrow.Record) error {
+			require.Equal(t, int64(2), r.NumRows())
+			require.Equal(t, int64(8), r.NumCols())
+			return nil
+		})
+	require.NoError(t, err)
+}