pkg/query: Set line as null when using frame filter

pkg/query/columnquery.go

@@ -475,16 +475,9 @@ func filterRecord(
 		}
 	}
 
-	// To keep track of which rows and frames to keep, we will build a list of row indices
-	// and the corresponding frames to keep within those rows.
-	type rowInfo struct {
-		rowIndex     int
-		framesToKeep []int // indices of frames to keep within the location list
-	}
-
-	rowsInfo := make([]rowInfo, 0)
 	originalValueSum := math.Int64.Sum(r.Value)
 
+	rowsToKeep := make([]int64, 0, int(rec.NumRows()))
 	for i := 0; i < int(rec.NumRows()); i++ {
 		lOffsetStart, lOffsetEnd := r.Locations.ValueOffsets(i)
 
@@ -509,104 +502,35 @@ func filterRecord(
 		if !keepRow {
 			continue
 		}
+		rowsToKeep = append(rowsToKeep, int64(i))
 
 		// Apply frame filters - determine which frames to keep
-		framesToKeep := make([]int, 0)
 		if lOffsetEnd-lOffsetStart > 0 {
 			for j := int(lOffsetStart); j < int(lOffsetEnd); j++ {
+				keepLocation := false
 				lineStart, lineEnd := r.Lines.ValueOffsets(j)
 				if lineStart >= lineEnd {
 					// For Unsymbolized location, check at location level only
 					if matchesAllFrameFilters(r, j, -1, frameFilters) {
-						framesToKeep = append(framesToKeep, -1)
+						keepLocation = true
 					}
 				} else {
 					// For Symbolized location, check each line/frame
 					for lineIdx := int(lineStart); lineIdx < int(lineEnd); lineIdx++ {
 						if matchesAllFrameFilters(r, j, lineIdx, frameFilters) {
-							framesToKeep = append(framesToKeep, lineIdx)
+							keepLocation = true
+						} else {
+							bitutil.ClearBit(r.Line.NullBitmapBytes(), lineIdx)
 						}
 					}
 				}
-			}
-		}
-
-		// Only keep rows that have at least one frame after filtering
-		if len(framesToKeep) > 0 {
-			rowsInfo = append(rowsInfo, rowInfo{
-				rowIndex:     i,
-				framesToKeep: framesToKeep,
-			})
-		}
-	}
-
-	if len(rowsInfo) == 0 {
-		// No rows match the filters
-		return []arrow.RecordBatch{}, originalValueSum, 0, nil
-	}
-
-	// Now apply frame filtering by nulling out non-matching frames within rows
-	for _, info := range rowsInfo {
-		lOffsetStart, lOffsetEnd := r.Locations.ValueOffsets(info.rowIndex)
-
-		// Create a set of frames to keep for quick lookup
-		keepSet := make(map[int]bool)
-		for _, frameIdx := range info.framesToKeep {
-			keepSet[frameIdx] = true
-		}
-
-		// Null out frames that don't match the filters
-		for j := int(lOffsetStart); j < int(lOffsetEnd); j++ {
-			lineStart, lineEnd := r.Lines.ValueOffsets(j)
-			allLinesFiltered := true
-
-			// Check if any line in this location is kept (or if this is unsymbolized and kept)
-			for lineIdx := int(lineStart); lineIdx < int(lineEnd); lineIdx++ {
-				if keepSet[lineIdx] {
-					allLinesFiltered = false
-					break
-				}
-			}
-			// Also check for unsymbolized case
-			if keepSet[-1] {
-				allLinesFiltered = false
-			}
-
-			// If all lines in this location are filtered out, invalidate the location
-			if allLinesFiltered {
-				bitutil.ClearBit(r.Locations.ListValues().NullBitmapBytes(), j)
-			}
-
-			// Null out individual line fields that don't match
-			for lineIdx := int(lineStart); lineIdx < int(lineEnd); lineIdx++ {
-				if !keepSet[lineIdx] {
-					// Null out the frame fields for this line
-					if r.LineFunctionNameIndices.Len() > 0 {
-						bitutil.ClearBit(r.LineFunctionNameIndices.NullBitmapBytes(), lineIdx)
-					}
-					if r.LineFunctionSystemNameIndices.Len() > 0 {
-						bitutil.ClearBit(r.LineFunctionSystemNameIndices.NullBitmapBytes(), lineIdx)
-					}
-					if r.LineFunctionFilenameIndices.Len() > 0 {
-						bitutil.ClearBit(r.LineFunctionFilenameIndices.NullBitmapBytes(), lineIdx)
-					}
-					if r.LineFunctionStartLine.Len() > 0 {
-						bitutil.ClearBit(r.LineFunctionStartLine.NullBitmapBytes(), lineIdx)
-					}
-					if r.LineNumber.Len() > 0 {
-						bitutil.ClearBit(r.LineNumber.NullBitmapBytes(), lineIdx)
-					}
+				if !keepLocation {
+					bitutil.ClearBit(r.Locations.ListValues().NullBitmapBytes(), j)
 				}
 			}
 		}
 	}
 
-	// Extract the rows we want to keep
-	rowsToKeep := make([]int64, len(rowsInfo))
-	for i, info := range rowsInfo {
-		rowsToKeep[i] = int64(info.rowIndex)
-	}
-
 	// Split the record into slices based on the rowsToKeep
 	recs := sliceRecord(rec, rowsToKeep)
 

pkg/query/flamegraph_arrow.go

@@ -263,8 +263,8 @@ func generateFlamegraphArrowRecord(ctx context.Context, mem memory.Allocator, tr
 
 				// just like locations, pprof stores lines in reverse order.
 				for k := int(llOffsetEnd - 1); k >= int(llOffsetStart); k-- {
-					// Skip lines that have been filtered out (null function names)
-					if !r.LineFunctionNameIndices.IsValid(k) {
+					// Skip lines that have been filtered out (null)
+					if !r.Line.IsValid(k) {
 						continue
 					}
 

pkg/query/multiple_filters_test.go

@@ -761,29 +761,25 @@ func TestFrameFilterFunctionNameSingle(t *testing.T) {
 		}
 	}()
 
-	// Should have 3 samples remaining with database frames
-	// Record 1 sample 1: database.query
-	// Record 2 sample 1: database.connect
-	// Record 3 sample 2: database.execute
 	totalSamples := int64(0)
 	for _, rec := range recs {
 		totalSamples += rec.NumRows()
 	}
-	require.Equal(t, int64(3), totalSamples, "Should have 3 samples with database frames")
+	require.Equal(t, int64(5), totalSamples, "Should have 5 samples with database frames")
 
 	// Validate each remaining sample
 	sampleIndex := 0
 	for _, rec := range recs {
 		r := profile.NewRecordReader(rec)
+		validFrameCount := 0
 		for i := 0; i < int(rec.NumRows()); i++ {
 			lOffsetStart, lOffsetEnd := r.Locations.ValueOffsets(i)
 			firstStart, _ := r.Lines.ValueOffsets(int(lOffsetStart))
 			_, lastEnd := r.Lines.ValueOffsets(int(lOffsetEnd - 1))
 
 			// Count valid frames (non-null) and verify they all contain "database"
-			validFrameCount := 0
 			for k := int(firstStart); k < int(lastEnd); k++ {
-				if r.LineFunctionNameIndices.IsValid(k) {
+				if r.Line.IsValid(k) {
 					fnIndex := r.LineFunctionNameIndices.Value(k)
 					functionName := string(r.LineFunctionNameDict.Value(int(fnIndex)))
 
@@ -793,11 +789,11 @@ func TestFrameFilterFunctionNameSingle(t *testing.T) {
 					validFrameCount++
 				}
 			}
-
-			// Each sample should have exactly one frame remaining
-			require.Equal(t, 1, validFrameCount, "Sample %d should have exactly 1 database frame remaining", sampleIndex)
-			sampleIndex++
 		}
+
+		// Each sample should have exactly one frame remaining
+		require.Equal(t, 1, validFrameCount, "Sample %d should have exactly 1 database frame remaining", sampleIndex)
+		sampleIndex++
 	}
 }
 
@@ -862,7 +858,7 @@ func TestFrameFilterFunctionNameNotContains(t *testing.T) {
 			// Count valid frames (non-null) and verify none contain "database"
 			validFrameCount := 0
 			for k := int(firstStart); k < int(lastEnd); k++ {
-				if r.LineFunctionNameIndices.IsValid(k) {
+				if r.Line.IsValid(k) {
 					fnIndex := r.LineFunctionNameIndices.Value(k)
 					functionName := string(r.LineFunctionNameDict.Value(int(fnIndex)))
 
@@ -1279,34 +1275,33 @@ func TestFrameFilterAddress(t *testing.T) {
 	for _, rec := range recs {
 		totalSamples += rec.NumRows()
 	}
-	require.Equal(t, int64(1), totalSamples, "Should have 1 sample containing frame with address 0x3000")
+	require.Equal(t, int64(5), totalSamples, "Should have unchanged 5 sample")
 
 	// Validate that the returned sample contains the correct frame with address 0x3000
 	// and that other frames in the stack are filtered out (set to null)
 	sampleIndex := 0
+	foundTargetAddress := false
+	validFrameCount := 0
 	for _, rec := range recs {
 		r := profile.NewRecordReader(rec)
 		for i := 0; i < int(rec.NumRows()); i++ {
 			lOffsetStart, lOffsetEnd := r.Locations.ValueOffsets(i)
 
-			foundTargetAddress := false
-			validFrameCount := 0
 			for j := int(lOffsetStart); j < int(lOffsetEnd); j++ {
 				if r.Address.IsValid(j) {
 					address := r.Address.Value(j)
 					validFrameCount++
 					if address == 0x3000 {
 						foundTargetAddress = true
+						sampleIndex++
 					}
 				}
 			}
-
-			require.True(t, foundTargetAddress, "Sample %d should contain frame with address 0x3000", sampleIndex)
-			// Frame filtering keeps the sample but may have multiple valid frames
-			require.Greater(t, validFrameCount, 0, "Sample %d should have at least 1 valid frame after filtering", sampleIndex)
-			sampleIndex++
 		}
 	}
+	require.True(t, foundTargetAddress, "Sample %d should contain frame with address 0x3000", sampleIndex)
+	// Frame filtering keeps the sample but may have multiple valid frames
+	require.Greater(t, validFrameCount, 0, "Sample %d should have at least 1 valid frame after filtering", sampleIndex)
 }
 
 func TestStackFilterAndLogicValidation(t *testing.T) {