feat: add logical projection node

Author: trevorwhitney

pkg/engine/internal/executor/unwrap.go renamed to pkg/engine/internal/executor/cast.go

@@ -10,34 +10,41 @@ import (
 	"github.com/apache/arrow-go/v18/arrow/memory"
 	"github.com/dustin/go-humanize"
 
+	"github.com/grafana/loki/v3/pkg/engine/internal/semconv"
 	"github.com/grafana/loki/v3/pkg/engine/internal/types"
 )
 
-func unwrapFn(operation types.UnaryOp) UnaryFunction {
-	//TODO: change signature of Evaluate to add this alloc parameter
-	alloc := memory.DefaultAllocator
-	return UnaryFunc(func(input ColumnVector) (ColumnVector, error) {
-		sourceCol, ok := input.ToArray().(*array.String)
+func castFn(operation types.UnaryOp) UnaryFunction {
+	return UnaryFunc(func(input ColumnVector, allocator memory.Allocator) (ColumnVector, error) {
+		arr := input.ToArray()
+		defer arr.Release()
+
+		sourceCol, ok := arr.(*array.String)
 		if !ok {
-			return nil, fmt.Errorf("expected column to be of type string, got %T", input.ToArray())
+			return nil, fmt.Errorf("expected column to be of type string, got %T", arr)
 		}
 
 		// Get conversion function and process values
 		conversionFn := getConversionFunction(operation)
-		unwrappedCol, errTracker := convertValues(sourceCol, conversionFn, alloc)
+		castCol, errTracker := castValues(sourceCol, conversionFn, allocator)
+		defer castCol.Release()
 
 		// Build error columns if needed
 		errorCol, errorDetailsCol := errTracker.buildArrays()
 		defer errTracker.releaseBuilders()
+		if errTracker.hasErrors {
+			defer errorCol.Release()
+			defer errorDetailsCol.Release()
+		}
 
 		// Build output schema and record
 		fields := buildOutputFields(errTracker.hasErrors)
-		arr, err := buildResult(unwrappedCol, errorCol, errorDetailsCol, fields)
+		result, err := buildResult(castCol, errorCol, errorDetailsCol, fields)
 		if err != nil {
 			return nil, err
 		}
 		return &ArrayStruct{
-			array: arr,
+			array: result,
 			ct:    types.ColumnTypeGenerated,
 			rows:  input.Len(),
 		}, nil
@@ -48,43 +55,43 @@ type conversionFn func(value string) (float64, error)
 
 func getConversionFunction(operation types.UnaryOp) conversionFn {
 	switch operation {
-	case types.UnaryOpUnwrapBytes:
+	case types.UnaryOpCastBytes:
 		return convertBytes
-	case types.UnaryOpUnwrapDuration:
+	case types.UnaryOpCastDuration:
 		return convertDuration
 	default:
 		return convertFloat
 	}
 }
 
-func convertValues(
+func castValues(
 	sourceCol *array.String,
 	conversionFn conversionFn,
 	allocator memory.Allocator,
 ) (arrow.Array, *errorTracker) {
-	unwrappedBuilder := array.NewFloat64Builder(allocator)
-	defer unwrappedBuilder.Release()
+	castBuilder := array.NewFloat64Builder(allocator)
+	defer castBuilder.Release()
 
 	tracker := newErrorTracker(allocator)
 
 	for i := 0; i < sourceCol.Len(); i++ {
 		if sourceCol.IsNull(i) {
-			unwrappedBuilder.AppendNull()
+			castBuilder.AppendNull()
 			tracker.recordSuccess()
 		} else {
 			valueStr := sourceCol.Value(i)
 			if val, err := conversionFn(valueStr); err == nil {
-				unwrappedBuilder.Append(val)
+				castBuilder.Append(val)
 				tracker.recordSuccess()
 			} else {
 				// Use 0.0 as default for errors, for backwards compatibility with old engine
-				unwrappedBuilder.Append(0.0)
+				castBuilder.Append(0.0)
 				tracker.recordError(i, err)
 			}
 		}
 	}
 
-	return unwrappedBuilder.NewArray(), tracker
+	return castBuilder.NewArray(), tracker
 }
 
 func buildOutputFields(
@@ -93,45 +100,21 @@ func buildOutputFields(
 	fields := make([]arrow.Field, 0, 3)
 
 	// Add value field
-	fields = append(fields, arrow.Field{
-		Name: types.ColumnNameGeneratedValue,
-		Type: arrow.PrimitiveTypes.Float64,
-		Metadata: types.ColumnMetadata(
-			types.ColumnTypeGenerated,
-			types.Loki.Float,
-		),
-		Nullable: true,
-	})
+	fields = append(fields, semconv.FieldFromIdent(semconv.ColumnIdentValue, false))
 
 	// Add error fields if needed
 	if hasErrors {
 		fields = append(fields,
-			arrow.Field{
-				Name: types.ColumnNameError,
-				Type: arrow.BinaryTypes.String,
-				Metadata: types.ColumnMetadata(
-					types.ColumnTypeParsed,
-					types.Loki.String,
-				),
-				Nullable: true,
-			},
-			arrow.Field{
-				Name: types.ColumnNameErrorDetails,
-				Type: arrow.BinaryTypes.String,
-				Metadata: types.ColumnMetadata(
-					types.ColumnTypeParsed,
-					types.Loki.String,
-				),
-				Nullable: true,
-			},
+			semconv.FieldFromIdent(semconv.ColumnIdentError, true),
+			semconv.FieldFromIdent(semconv.ColumnIdentErrorDetails, true),
 		)
 	}
 
 	return fields
 }
 
 func buildResult(
-	unwrappedCol, errorCol, errorDetailsCol arrow.Array,
+	castCol, errorCol, errorDetailsCol arrow.Array,
 	fields []arrow.Field,
 ) (*array.Struct, error) {
 	hasErrors := errorCol != nil
@@ -144,23 +127,14 @@ func buildResult(
 	columns := make([]arrow.Array, totalCols)
 
 	// Add new columns - these are newly created so don't need extra retain
-	columns[0] = unwrappedCol
+	columns[0] = castCol
 	if hasErrors {
 		columns[1] = errorCol
 		columns[2] = errorDetailsCol
 	}
 
 	// NewStructArrayWithFields will retain all columns
-	result, err := array.NewStructArrayWithFields(columns, fields)
-
-	// Release our references to newly created columns (result now owns them)
-	unwrappedCol.Release()
-	if hasErrors {
-		errorCol.Release()
-		errorDetailsCol.Release()
-	}
-
-	return result, err
+	return array.NewStructArrayWithFields(columns, fields)
 }
 
 func convertFloat(v string) (float64, error) {

pkg/engine/internal/executor/expressions.go

@@ -134,7 +134,7 @@ func (e expressionEvaluator) eval(expr physical.Expression, input arrow.Record)
 		if err != nil {
 			return nil, fmt.Errorf("failed to lookup unary function: %w", err)
 		}
-		return fn.Evaluate(lhr)
+		return fn.Evaluate(lhr, e.allocator)
 
 	case *physical.BinaryExpr:
 		lhs, err := e.eval(expr.Left, input)
@@ -339,6 +339,7 @@ func NewArrayStruct(arr *array.Struct, ct types.ColumnType) *ArrayStruct {
 // ToArray implements ColumnVector.
 // Returns the underlying struct array.
 func (a *ArrayStruct) ToArray() arrow.Array {
+	a.array.Retain()
 	return a.array
 }
 
@@ -402,6 +403,11 @@ func (a *ArrayStruct) Len() int64 {
 	return a.rows
 }
 
+// Release decreases the reference count by 1 on underlying Arrow array
+func (a *ArrayStruct) Release() {
+	a.array.Release()
+}
+
 // CoalesceVector represents multiple columns with the same name but different [types.ColumnType]
 // Vectors are ordered by precedence (highest precedence first).
 type CoalesceVector struct {

pkg/engine/internal/executor/expressions_test.go

@@ -397,7 +397,12 @@ label_2
 	})
 }
 
-func TestEvaluateUnwrapExpression(t *testing.T) {
+func TestEvaluateUnaryCastExpression(t *testing.T) {
+	colMsg := semconv.ColumnIdentMessage
+	colStatusCode := semconv.NewIdentifier("status_code", types.ColumnTypeMetadata, types.Loki.String)
+	colTimeout := semconv.NewIdentifier("timeout", types.ColumnTypeParsed, types.Loki.String)
+	colMixedValues := semconv.NewIdentifier("mixed_values", types.ColumnTypeMetadata, types.Loki.String)
+
 	t.Run("unknown column", func(t *testing.T) {
 		alloc := memory.NewCheckedAllocator(memory.DefaultAllocator)
 		defer alloc.AssertSize(t, 0) // Assert empty on test exit
@@ -409,13 +414,14 @@ func TestEvaluateUnwrapExpression(t *testing.T) {
 					Type:   types.ColumnTypeAmbiguous,
 				},
 			},
-			Op: types.UnaryOpUnwrap,
+			Op: types.UnaryOpCastFloat,
 		}
 
 		n := len(words)
 		rec := batch(n, time.Now())
 		colVec, err := e.eval(expr, rec)
 		require.NoError(t, err)
+		defer colVec.Release()
 
 		id := colVec.Type().ArrowType().ID()
 		require.Equal(t, arrow.STRUCT, id)
@@ -442,37 +448,41 @@ func TestEvaluateUnwrapExpression(t *testing.T) {
 		}
 	})
 
-	t.Run("unwrap column generates a value", func(t *testing.T) {
+	t.Run("cast column generates a value", func(t *testing.T) {
 		expr := &physical.UnaryExpr{
 			Left: &physical.ColumnExpr{
 				Ref: types.ColumnRef{
 					Column: "status_code",
 					Type:   types.ColumnTypeAmbiguous,
 				},
 			},
-			Op: types.UnaryOpUnwrap,
+			Op: types.UnaryOpCastBytes,
 		}
 
 		alloc := memory.NewCheckedAllocator(memory.DefaultAllocator)
 		defer alloc.AssertSize(t, 0) // Assert empty on test exit
 		e := newExpressionEvaluator(alloc)
 
-		schema := arrow.NewSchema([]arrow.Field{
-			{Name: types.ColumnNameBuiltinMessage, Type: arrow.BinaryTypes.String, Metadata: types.ColumnMetadata(types.ColumnTypeBuiltin, types.Loki.String)},
-			{Name: "status_code", Type: arrow.BinaryTypes.String, Metadata: types.ColumnMetadata(types.ColumnTypeMetadata, types.Loki.String)},
-			{Name: "timeout", Type: arrow.BinaryTypes.String, Metadata: types.ColumnMetadata(types.ColumnTypeParsed, types.Loki.String)},
-		}, nil)
+		schema := arrow.NewSchema(
+			[]arrow.Field{
+				semconv.FieldFromIdent(colMsg, false),
+				semconv.FieldFromIdent(colStatusCode, true),
+				semconv.FieldFromIdent(colTimeout, true),
+			},
+			nil,
+		)
 		rows := arrowtest.Rows{
-			{"message": "timeout set", "status_code": "200", "timeout": "2m"},
-			{"message": "short timeout", "status_code": "204", "timeout": "10s"},
-			{"message": "long timeout", "status_code": "404", "timeout": "1h"},
+			{"utf8.builtin.message": "timeout set", "utf8.metadata.status_code": "200", "utf8.parsed.timeout": "2m"},
+			{"utf8.builtin.message": "short timeout", "utf8.metadata.status_code": "204", "utf8.parsed.timeout": "10s"},
+			{"utf8.builtin.message": "long timeout", "utf8.metadata.status_code": "404", "utf8.parsed.timeout": "1h"},
 		}
 
 		record := rows.Record(alloc, schema)
 		defer record.Release()
 
 		colVec, err := e.eval(expr, record)
 		require.NoError(t, err)
+		defer colVec.Release()
 		id := colVec.Type().ArrowType().ID()
 		require.Equal(t, arrow.STRUCT, id)
 
@@ -490,37 +500,41 @@ func TestEvaluateUnwrapExpression(t *testing.T) {
 		require.Equal(t, 404.0, value.Value(2))
 	})
 
-	t.Run("unwrap column generates a value from a parsed column", func(t *testing.T) {
+	t.Run("cast column generates a value from a parsed column", func(t *testing.T) {
 		expr := &physical.UnaryExpr{
 			Left: &physical.ColumnExpr{
 				Ref: types.ColumnRef{
 					Column: "timeout",
 					Type:   types.ColumnTypeAmbiguous,
 				},
 			},
-			Op: types.UnaryOpUnwrapDuration,
+			Op: types.UnaryOpCastDuration,
 		}
 
 		alloc := memory.NewCheckedAllocator(memory.DefaultAllocator)
 		defer alloc.AssertSize(t, 0) // Assert empty on test exit
 		e := newExpressionEvaluator(alloc)
 
-		schema := arrow.NewSchema([]arrow.Field{
-			{Name: types.ColumnNameBuiltinMessage, Type: arrow.BinaryTypes.String, Metadata: types.ColumnMetadata(types.ColumnTypeBuiltin, types.Loki.String)},
-			{Name: "status_code", Type: arrow.BinaryTypes.String, Metadata: types.ColumnMetadata(types.ColumnTypeMetadata, types.Loki.String)},
-			{Name: "timeout", Type: arrow.BinaryTypes.String, Metadata: types.ColumnMetadata(types.ColumnTypeParsed, types.Loki.String)},
-		}, nil)
+		schema := arrow.NewSchema(
+			[]arrow.Field{
+				semconv.FieldFromIdent(colMsg, false),
+				semconv.FieldFromIdent(colStatusCode, true),
+				semconv.FieldFromIdent(colTimeout, true),
+			},
+			nil,
+		)
 		rows := arrowtest.Rows{
-			{"message": "timeout set", "status_code": "200", "timeout": "2m"},
-			{"message": "short timeout", "status_code": "204", "timeout": "10s"},
-			{"message": "long timeout", "status_code": "404", "timeout": "1h"},
+			{"utf8.builtin.message": "timeout set", "utf8.metadata.status_code": "200", "utf8.parsed.timeout": "2m"},
+			{"utf8.builtin.message": "short timeout", "utf8.metadata.status_code": "204", "utf8.parsed.timeout": "10s"},
+			{"utf8.builtin.message": "long timeout", "utf8.metadata.status_code": "404", "utf8.parsed.timeout": "1h"},
 		}
 
 		record := rows.Record(alloc, schema)
 		defer record.Release()
 
 		colVec, err := e.eval(expr, record)
 		require.NoError(t, err)
+		defer colVec.Release()
 		id := colVec.Type().ArrowType().ID()
 		require.Equal(t, arrow.STRUCT, id)
 
@@ -537,38 +551,42 @@ func TestEvaluateUnwrapExpression(t *testing.T) {
 		require.Equal(t, 3600.0, value.Value(2))
 	})
 
-	t.Run("unwrap tracks errors", func(t *testing.T) {
+	t.Run("cast operation tracks errors", func(t *testing.T) {
 		colExpr := &physical.UnaryExpr{
 			Left: &physical.ColumnExpr{
 				Ref: types.ColumnRef{
 					Column: "mixed_values",
 					Type:   types.ColumnTypeAmbiguous,
 				},
 			},
-			Op: types.UnaryOpUnwrap,
+			Op: types.UnaryOpCastFloat,
 		}
 
 		alloc := memory.NewCheckedAllocator(memory.DefaultAllocator)
 		defer alloc.AssertSize(t, 0) // Assert empty on test exit
 		e := newExpressionEvaluator(alloc)
 
-		schema := arrow.NewSchema([]arrow.Field{
-			{Name: types.ColumnNameBuiltinMessage, Type: arrow.BinaryTypes.String, Metadata: types.ColumnMetadata(types.ColumnTypeBuiltin, types.Loki.String)},
-			{Name: "mixed_values", Type: arrow.BinaryTypes.String, Metadata: types.ColumnMetadata(types.ColumnTypeMetadata, types.Loki.String)},
-		}, nil)
+		schema := arrow.NewSchema(
+			[]arrow.Field{
+				semconv.FieldFromIdent(colMsg, false),
+				semconv.FieldFromIdent(colMixedValues, true),
+			},
+			nil,
+		)
 		rows := arrowtest.Rows{
-			{"message": "valid numeric", "mixed_values": "42.5"},
-			{"message": "invalid numeric", "mixed_values": "not_a_number"},
-			{"message": "valid bytes", "mixed_values": "1KB"},
-			{"message": "invalid bytes", "mixed_values": "invalid_bytes"},
-			{"message": "empty string", "mixed_values": ""},
+			{"utf8.builtin.message": "valid numeric", "utf8.metadata.mixed_values": "42.5"},
+			{"utf8.builtin.message": "invalid numeric", "utf8.metadata.mixed_values": "not_a_number"},
+			{"utf8.builtin.message": "valid bytes", "utf8.metadata.mixed_values": "1KB"},
+			{"utf8.builtin.message": "invalid bytes", "utf8.metadata.mixed_values": "invalid_bytes"},
+			{"utf8.builtin.message": "empty string", "utf8.metadata.mixed_values": ""},
 		}
 
 		record := rows.Record(alloc, schema)
 		defer record.Release()
 
 		colVec, err := e.eval(colExpr, record)
 		require.NoError(t, err)
+		defer colVec.Release()
 		id := colVec.Type().ArrowType().ID()
 		require.Equal(t, arrow.STRUCT, id)