chore(dataobj): Wire dataset.Reader into dataobj.LogsReader and dataobj.StreamsReader (#16701)

Co-authored-by: Ashwanth Goli <[email protected]>

Author: rfratto

pkg/dataobj/internal/dataset/predicate.go

@@ -15,6 +15,13 @@ type (
 	// included if either the Left or Right Predicate are true.
 	OrPredicate struct{ Left, Right Predicate }
 
+	// A NotePredicate is a [Predicate] which asserts that a row may only be
+	// included if the inner Predicate is false.
+	NotPredicate struct{ Inner Predicate }
+
+	// FalsePredicate is a [Predicate] which always returns false.
+	FalsePredicate struct{}
+
 	// An EqualPredicate is a [Predicate] which asserts that a row may only be
 	// included if the Value of the Column is equal to the Value.
 	EqualPredicate struct {
@@ -55,6 +62,8 @@ type (
 
 func (AndPredicate) isPredicate()         {}
 func (OrPredicate) isPredicate()          {}
+func (NotPredicate) isPredicate()         {}
+func (FalsePredicate) isPredicate()       {}
 func (EqualPredicate) isPredicate()       {}
 func (GreaterThanPredicate) isPredicate() {}
 func (LessThanPredicate) isPredicate()    {}
@@ -78,6 +87,10 @@ func WalkPredicate(p Predicate, fn func(p Predicate) bool) {
 		WalkPredicate(p.Left, fn)
 		WalkPredicate(p.Right, fn)
 
+	case NotPredicate:
+		WalkPredicate(p.Inner, fn)
+
+	case FalsePredicate: // No children.
 	case EqualPredicate: // No children.
 	case GreaterThanPredicate: // No children.
 	case LessThanPredicate: // No children.

pkg/dataobj/internal/dataset/reader.go

@@ -126,6 +126,8 @@ func (r *Reader) Read(ctx context.Context, s []Row) (n int, err error) {
 	count, err := r.inner.ReadColumns(ctx, r.dl.PrimaryColumns(), s[:readSize])
 	if err != nil && !errors.Is(err, io.EOF) {
 		return n, err
+	} else if count == 0 && errors.Is(err, io.EOF) {
+		return 0, io.EOF
 	}
 
 	var passCount int // passCount tracks how many rows pass the predicate.
@@ -196,6 +198,12 @@ func checkPredicate(p Predicate, lookup map[Column]int, row Row) bool {
 	case OrPredicate:
 		return checkPredicate(p.Left, lookup, row) || checkPredicate(p.Right, lookup, row)
 
+	case NotPredicate:
+		return !checkPredicate(p.Inner, lookup, row)
+
+	case FalsePredicate:
+		return false
+
 	case EqualPredicate:
 		columnIndex, ok := lookup[p.Column]
 		if !ok {
@@ -350,7 +358,7 @@ func (r *Reader) validatePredicate() error {
 			err = process(p.Column)
 		case FuncPredicate:
 			err = process(p.Column)
-		case AndPredicate, OrPredicate, nil:
+		case AndPredicate, OrPredicate, NotPredicate, FalsePredicate, nil:
 			// No columns to process.
 		default:
 			panic(fmt.Sprintf("dataset.Reader.validatePredicate: unsupported predicate type %T", p))
@@ -422,7 +430,7 @@ func (r *Reader) fillPrimaryMask(mask *bitmask.Mask) {
 			process(p.Column)
 		case FuncPredicate:
 			process(p.Column)
-		case AndPredicate, OrPredicate, nil:
+		case AndPredicate, OrPredicate, NotPredicate, FalsePredicate, nil:
 			// No columns to process.
 		default:
 			panic(fmt.Sprintf("dataset.Reader.fillPrimaryMask: unsupported predicate type %T", p))
@@ -463,6 +471,25 @@ func (r *Reader) buildPredicateRanges(ctx context.Context, p Predicate) (rowRang
 		}
 		return unionRanges(nil, left, right), nil
 
+	case NotPredicate:
+		// De Morgan's laws must be applied to reduce the NotPredicate to a set of
+		// predicates that can be applied to pages.
+		//
+		// See comment on [simplifyNotPredicate] for more information.
+		simplified, err := simplifyNotPredicate(p)
+		if err != nil {
+			// Predicate can't be simplfied, so we permit the full range.
+			var rowsCount uint64
+			for _, column := range r.dl.AllColumns() {
+				rowsCount = max(rowsCount, uint64(column.ColumnInfo().RowsCount))
+			}
+			return rowRanges{{Start: 0, End: rowsCount - 1}}, nil
+		}
+		return r.buildPredicateRanges(ctx, simplified)
+
+	case FalsePredicate:
+		return nil, nil // No valid ranges.
+
 	case EqualPredicate:
 		return r.buildColumnPredicateRanges(ctx, p.Column, p)
 
@@ -489,6 +516,67 @@ func (r *Reader) buildPredicateRanges(ctx context.Context, p Predicate) (rowRang
 	}
 }
 
+// simplifyNotPredicate applies De Morgan's laws to a NotPredicate to permit
+// page filtering.
+//
+// While during evaluation, a NotPredicate inverts the result of the inner
+// predicate, the same can't be done for page filtering. For example, imagine
+// that a page is included from a rule "a > 10." If we inverted that inclusion,
+// we may be incorrectly filtering out that page, as that page may also have
+// values less than 10.
+//
+// To correctly apply page filtering to a NotPredicate, we reduce the
+// NotPredicate to a set of predicates that can be applied to pages. This may
+// result in other NotPredicates that also need to be simplified.
+//
+// If the NotPredicate can't be simplified, simplifyNotPredicate returns an
+// error.
+func simplifyNotPredicate(p NotPredicate) (Predicate, error) {
+	switch inner := p.Inner.(type) {
+	case AndPredicate: // De Morgan's law: !(A && B) == !A || !B
+		return OrPredicate{
+			Left:  NotPredicate{Inner: inner.Left},
+			Right: NotPredicate{Inner: inner.Right},
+		}, nil
+
+	case OrPredicate: // De Morgan's law: !(A || B) == !A && !B
+		return AndPredicate{
+			Left:  NotPredicate{Inner: inner.Left},
+			Right: NotPredicate{Inner: inner.Right},
+		}, nil
+
+	case NotPredicate: // De Morgan's law: !!A == A
+		return inner.Inner, nil
+
+	case FalsePredicate:
+		return nil, fmt.Errorf("can't simplify FalsePredicate")
+
+	case EqualPredicate: // De Morgan's law: !(A == B) == A != B == A < B || A > B
+		return OrPredicate{
+			Left:  LessThanPredicate(inner),
+			Right: GreaterThanPredicate(inner),
+		}, nil
+
+	case GreaterThanPredicate: // De Morgan's law: !(A > B) == A <= B
+		return OrPredicate{
+			Left:  EqualPredicate(inner),
+			Right: LessThanPredicate(inner),
+		}, nil
+
+	case LessThanPredicate: // De Morgan's law: !(A < B) == A >= B
+		return OrPredicate{
+			Left:  EqualPredicate(inner),
+			Right: GreaterThanPredicate(inner),
+		}, nil
+
+	case FuncPredicate:
+		return nil, fmt.Errorf("can't simplify FuncPredicate")
+
+	default:
+		panic(fmt.Sprintf("unsupported predicate type %T", inner))
+	}
+}
+
 // buildColumnPredicateRanges returns a set of rowRanges that are valid based
 // on whether EqualPredicate, GreaterThanPredicate, or LessThanPredicate may be
 // true for each page in a column.
@@ -536,10 +624,10 @@ func (r *Reader) buildColumnPredicateRanges(ctx context.Context, c Column, p Pre
 		switch p := p.(type) {
 		case EqualPredicate: // EqualPredicate may be true if p.Value is inside the range of the page.
 			include = CompareValues(p.Value, minValue) >= 0 && CompareValues(p.Value, maxValue) <= 0
-		case GreaterThanPredicate: // GreaterThanPredicate may be true if p.Value is greater than the min value.
-			include = CompareValues(p.Value, minValue) > 0
-		case LessThanPredicate: // LessThanPredicate may be true if p.Value is less than the max value.
-			include = CompareValues(p.Value, maxValue) < 0
+		case GreaterThanPredicate: // GreaterThanPredicate may be true if maxValue of a page is greater than p.Value
+			include = CompareValues(maxValue, p.Value) > 0
+		case LessThanPredicate: // LessThanPredicate may be true if minValue of a page is less than p.Value
+			include = CompareValues(minValue, p.Value) < 0
 		default:
 			panic(fmt.Sprintf("unsupported predicate type %T", p))
 		}

pkg/dataobj/internal/dataset/reader_basic.go

@@ -90,7 +90,7 @@ func (pr *basicReader) ReadColumns(ctx context.Context, columns []Column, s []Ro
 // Fill does not advance the offset of the basicReader.
 func (pr *basicReader) Fill(ctx context.Context, columns []Column, s []Row) (n int, err error) {
 	if len(columns) == 0 {
-		return 0, fmt.Errorf("no columns to read")
+		return 0, fmt.Errorf("no columns to fill")
 	}
 
 	for partition := range partitionRows(s) {

pkg/dataobj/internal/dataset/reader_test.go

@@ -8,6 +8,9 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/loki/v3/pkg/dataobj/internal/metadata/datasetmd"
+	"github.com/grafana/loki/v3/pkg/dataobj/internal/result"
 )
 
 func Test_Reader_ReadAll(t *testing.T) {
@@ -228,3 +231,159 @@ func readDataset(br *Reader, batchSize int) ([]Row, error) {
 		}
 	}
 }
+
+func Test_BuildPredicateRanges(t *testing.T) {
+	ds, cols := buildMemDatasetWithStats(t)
+	tt := []struct {
+		name      string
+		predicate Predicate
+		want      rowRanges
+	}{
+		{
+			name:      "nil predicate returns full range",
+			predicate: nil,
+			want:      rowRanges{{Start: 0, End: 999}}, // Full dataset range
+		},
+		{
+			name:      "equal predicate in range",
+			predicate: EqualPredicate{Column: cols[1], Value: Int64Value(50)},
+			want:      rowRanges{{Start: 0, End: 249}}, // Page 1 of Timestamp column
+		},
+		{
+			name:      "equal predicate not in any range",
+			predicate: EqualPredicate{Column: cols[1], Value: Int64Value(1500)},
+			want:      nil, // No ranges should match
+		},
+		{
+			name:      "greater than predicate",
+			predicate: GreaterThanPredicate{Column: cols[1], Value: Int64Value(400)},
+			want:      rowRanges{{Start: 250, End: 749}, {Start: 750, End: 999}}, // Pages 2 and 3 of Timestamp column
+		},
+		{
+			name:      "less than predicate",
+			predicate: LessThanPredicate{Column: cols[1], Value: Int64Value(300)},
+			want:      rowRanges{{Start: 0, End: 249}, {Start: 250, End: 749}}, // Pages 1 and 2 of Timestamp column
+		},
+		{
+			name: "and predicate",
+			predicate: AndPredicate{
+				Left:  EqualPredicate{Column: cols[0], Value: Int64Value(1)},      // Rows 0 - 299 of stream column
+				Right: LessThanPredicate{Column: cols[1], Value: Int64Value(600)}, // Rows 0 - 249, 250 - 749 of timestamp column
+			},
+			want: rowRanges{{Start: 0, End: 249}, {Start: 250, End: 299}},
+		},
+		{
+			name: "or predicate",
+			predicate: OrPredicate{
+				Left:  EqualPredicate{Column: cols[0], Value: Int64Value(1)},         // Rows 0 - 299 of stream column
+				Right: GreaterThanPredicate{Column: cols[1], Value: Int64Value(800)}, // Rows 750 - 999 of timestamp column
+			},
+			want: rowRanges{{Start: 0, End: 299}, {Start: 750, End: 999}}, // Rows 0 - 299, 750 - 999
+		},
+	}
+
+	ctx := context.Background()
+	for _, tc := range tt {
+		t.Run(tc.name, func(t *testing.T) {
+			r := NewReader(ReaderOptions{
+				Dataset:   ds,
+				Columns:   cols,
+				Predicate: tc.predicate,
+			})
+			defer r.Close()
+
+			// Initialize downloader
+			require.NoError(t, r.initDownloader(ctx))
+
+			got, err := r.buildPredicateRanges(ctx, tc.predicate)
+			require.NoError(t, err)
+
+			require.Equal(t, tc.want, got, "row ranges should match expected ranges")
+		})
+	}
+}
+
+// buildMemDatasetWithStats creates a test dataset with only column and page stats.
+func buildMemDatasetWithStats(t *testing.T) (Dataset, []Column) {
+	t.Helper()
+
+	dset := FromMemory([]*MemColumn{
+		{
+			Info: ColumnInfo{
+				Name:      "stream",
+				Type:      datasetmd.VALUE_TYPE_INT64,
+				RowsCount: 1000, // 0 - 999
+			},
+			Pages: []*MemPage{
+				{
+					Info: PageInfo{
+						RowCount: 300, // 0 - 299
+						Stats: &datasetmd.Statistics{
+							MinValue: encodeInt64Value(t, 1),
+							MaxValue: encodeInt64Value(t, 2),
+						},
+					},
+				},
+				{
+					Info: PageInfo{
+						RowCount: 700, // 300 - 999
+						Stats: &datasetmd.Statistics{
+							MinValue: encodeInt64Value(t, 2),
+							MaxValue: encodeInt64Value(t, 2),
+						},
+					},
+				},
+			},
+		},
+		{
+			Info: ColumnInfo{
+				Name:      "timestamp",
+				Type:      datasetmd.VALUE_TYPE_INT64,
+				RowsCount: 1000, // 0 - 999
+			},
+			Pages: []*MemPage{
+				{
+					Info: PageInfo{
+						RowCount: 250, // 0 - 249
+						Stats: &datasetmd.Statistics{
+							MinValue: encodeInt64Value(t, 0),
+							MaxValue: encodeInt64Value(t, 100),
+						},
+					},
+				},
+				{
+					Info: PageInfo{
+						RowCount: 500, // 249 - 749
+						Stats: &datasetmd.Statistics{
+							MinValue: encodeInt64Value(t, 200),
+							MaxValue: encodeInt64Value(t, 500),
+						},
+					},
+				},
+				{
+					Info: PageInfo{
+						RowCount: 250, // 750 - 999
+						Stats: &datasetmd.Statistics{
+							MinValue: encodeInt64Value(t, 800),
+							MaxValue: encodeInt64Value(t, 1000),
+						},
+					},
+				},
+			},
+		},
+	})
+
+	cols, err := result.Collect(dset.ListColumns(context.Background()))
+	require.NoError(t, err)
+
+	return dset, cols
+}
+
+// Helper function to encode an integer value for statistics
+func encodeInt64Value(t *testing.T, v int64) []byte {
+	t.Helper()
+
+	data, err := Int64Value(v).MarshalBinary()
+	require.NoError(t, err)
+	return data
+}

pkg/dataobj/internal/sections/logs/iter.go

@@ -80,7 +80,7 @@ func IterSection(ctx context.Context, dec encoding.LogsDecoder, section *filemd.
 			}
 
 			for _, row := range rows[:n] {
-				record, err := decodeRecord(streamsColumns, row)
+				record, err := Decode(streamsColumns, row)
 				if err != nil || !yield(record) {
 					return err
 				}
@@ -89,7 +89,10 @@ func IterSection(ctx context.Context, dec encoding.LogsDecoder, section *filemd.
 	})
 }
 
-func decodeRecord(columns []*logsmd.ColumnDesc, row dataset.Row) (Record, error) {
+// Decode decodes a record from a [dataset.Row], using the provided columns to
+// determine the column type. The list of columns must match the columns used
+// to create the row.
+func Decode(columns []*logsmd.ColumnDesc, row dataset.Row) (Record, error) {
 	record := Record{
 		// Preallocate metadata to exact number of metadata columns to avoid
 		// oversizing.