Implement TSDB->Parquet RowReader

convert/convert.go

@@ -0,0 +1,159 @@
+package convert
+
+import (
+	"context"
+	"fmt"
+	"github.com/hashicorp/go-multierror"
+	"github.com/parquet-go/parquet-go"
+	"github.com/prometheus-community/parquet-common/schema"
+	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/storage"
+	"github.com/prometheus/prometheus/tsdb"
+	"github.com/prometheus/prometheus/tsdb/chunks"
+	"github.com/prometheus/prometheus/tsdb/tombstones"
+	"io"
+)
+
+type Convertible interface {
+	Index() (tsdb.IndexReader, error)
+	Chunks() (tsdb.ChunkReader, error)
+	Tombstones() (tombstones.Reader, error)
+	Meta() tsdb.BlockMeta
+}
+
+type tsdbRowReader struct {
+	ctx context.Context
+
+	closers []io.Closer
+
+	seriesSet storage.ChunkSeriesSet
+
+	rowBuilder *parquet.RowBuilder
+	schema     *schema.TSDBSchema
+
+	encoder *schema.PrometheusParquetChunksEncoder
+}
+
+func newTsdbRowReader(ctx context.Context, mint, maxt, colDuration int64, blks []Convertible) (*tsdbRowReader, error) {
+	var (
+		seriesSets = make([]storage.ChunkSeriesSet, 0, len(blks))
+		closers    = make([]io.Closer, 0, len(blks))
+	)
+
+	b := schema.NewBuilder(mint, maxt, colDuration)
+
+	for _, blk := range blks {
+		indexr, err := blk.Index()
+		if err != nil {
+			return nil, fmt.Errorf("unable to get index reader from block: %s", err)
+		}
+		closers = append(closers, indexr)
+
+		chunkr, err := blk.Chunks()
+		if err != nil {
+			return nil, fmt.Errorf("unable to get chunk reader from block: %s", err)
+		}
+		closers = append(closers, chunkr)
+
+		tombsr, err := blk.Tombstones()
+		if err != nil {
+			return nil, fmt.Errorf("unable to get tombstone reader from block: %s", err)
+		}
+		closers = append(closers, tombsr)
+
+		lblns, err := indexr.LabelNames(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("unable to get label names from block: %s", err)
+		}
+
+		postings := tsdb.AllSortedPostings(ctx, indexr)
+		seriesSet := tsdb.NewBlockChunkSeriesSet(blk.Meta().ULID, indexr, chunkr, tombsr, postings, mint, maxt, false)
+		seriesSets = append(seriesSets, seriesSet)
+
+		b.AddLabelNameColumn(lblns...)
+	}
+
+	cseriesSet := storage.NewMergeChunkSeriesSet(seriesSets, 0, storage.NewConcatenatingChunkSeriesMerger())
+
+	s, err := b.Build()
+	if err != nil {
+		return nil, fmt.Errorf("unable to build index reader from block: %s", err)
+	}
+
+	return &tsdbRowReader{
+		ctx:       ctx,
+		seriesSet: cseriesSet,
+		closers:   closers,
+		schema:    s,
+
+		rowBuilder: parquet.NewRowBuilder(s.Schema),
+		encoder:    schema.NewPrometheusParquetChunksEncoder(s),
+	}, nil
+}
+
+func (rr *tsdbRowReader) Close() error {
+	err := &multierror.Error{}
+	for i := range rr.closers {
+		err = multierror.Append(err, rr.closers[i].Close())
+	}
+	return err.ErrorOrNil()
+}
+
+func (rr *tsdbRowReader) Schema() *parquet.Schema {
+	return rr.schema.Schema
+}
+
+func (rr *tsdbRowReader) ReadRows(buf []parquet.Row) (int, error) {
+	select {
+	case <-rr.ctx.Done():
+		return 0, rr.ctx.Err()
+	default:
+	}
+
+	var it chunks.Iterator
+
+	i := 0
+	for i < len(buf) && rr.seriesSet.Next() {
+		rr.rowBuilder.Reset()
+		s := rr.seriesSet.At()
+		it = s.Iterator(it)
+
+		chkBytes, err := rr.encoder.Encode(it)
+		if err != nil {
+			return i, fmt.Errorf("unable to collect chunks: %s", err)
+		}
+
+		// skip series that have no chunks in the requested time
+		if allChunksEmpty(chkBytes) {
+			continue
+		}
+
+		s.Labels().Range(func(l labels.Label) {
+			colName := schema.LabelToColumn(l.Name)
+			lc, _ := rr.schema.Schema.Lookup(colName)
+			rr.rowBuilder.Add(lc.ColumnIndex, parquet.ValueOf(l.Value))
+		})
+
+		for idx, chk := range chkBytes {
+			if len(chk) == 0 {
+				continue
+			}
+			rr.rowBuilder.Add(rr.schema.DataColsIndexes[idx], parquet.ValueOf(chk))
+		}
+		buf[i] = rr.rowBuilder.AppendRow(buf[i][:0])
+		i++
+	}
+	if i < len(buf) {
+		return i, io.EOF
+	}
+	return i, rr.seriesSet.Err()
+}
+
+func allChunksEmpty(chkBytes [][]byte) bool {
+	for _, chk := range chkBytes {
+		if len(chk) != 0 {
+			return false
+		}
+	}
+	return true
+}

convert/convert_test.go

@@ -0,0 +1,185 @@
+package convert
+
+import (
+	"context"
+	"fmt"
+	"math"
+	"testing"
+	"time"
+
+	"github.com/parquet-go/parquet-go"
+	"github.com/prometheus-community/parquet-common/schema"
+	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/tsdb/chunkenc"
+	"github.com/prometheus/prometheus/tsdb/chunks"
+	"github.com/prometheus/prometheus/util/teststorage"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_Convert_TSDB(t *testing.T) {
+	ctx := context.Background()
+
+	tc := []struct {
+		dataColDurationMs      time.Duration
+		step                   time.Duration
+		numberOfSamples        int
+		expectedNumberOfChunks int
+		expectedPointsPerChunk int
+	}{
+		{
+			dataColDurationMs:      time.Hour,
+			step:                   time.Hour,
+			numberOfSamples:        3,
+			expectedNumberOfChunks: 3,
+			expectedPointsPerChunk: 1,
+		},
+		{
+			dataColDurationMs:      time.Hour,
+			step:                   time.Hour,
+			numberOfSamples:        48,
+			expectedNumberOfChunks: 48,
+			expectedPointsPerChunk: 1,
+		},
+		{
+			dataColDurationMs:      8 * time.Hour,
+			step:                   time.Hour / 2,
+			numberOfSamples:        10,
+			expectedNumberOfChunks: 1,
+			expectedPointsPerChunk: 10,
+		},
+		{
+			dataColDurationMs:      8 * time.Hour,
+			step:                   time.Hour / 2,
+			numberOfSamples:        32,
+			expectedNumberOfChunks: 2,
+			expectedPointsPerChunk: 16,
+		},
+	}
+
+	for _, tt := range tc {
+		t.Run(fmt.Sprintf("dataColDurationMs:%v,step:%v,numberOfSamples:%v", tt.dataColDurationMs.Hours(), tt.step.Seconds(), tt.numberOfSamples), func(t *testing.T) {
+			st := teststorage.New(t)
+			t.Cleanup(func() { st.Close() })
+
+			app := st.Appender(ctx)
+			seriesHash := make(map[uint64]struct{})
+			for i := 0; i != 1_000; i++ {
+				for j := 0; j < tt.numberOfSamples; j++ {
+					lbls := labels.FromStrings("__name__", "foo", "bar", fmt.Sprintf("%d", 2*i))
+					seriesHash[lbls.Hash()] = struct{}{}
+					_, err := app.Append(0, lbls, (tt.step * time.Duration(j)).Milliseconds(), float64(i))
+					require.NoError(t, err)
+				}
+			}
+
+			require.NoError(t, app.Commit())
+
+			h := st.DB.Head()
+			rr, err := newTsdbRowReader(ctx, h.MinTime(), h.MaxTime(), tt.dataColDurationMs.Milliseconds(), []Convertible{h})
+			require.NoError(t, err)
+
+			defer func() { _ = rr.Close() }()
+
+			buf := make([]parquet.Row, 100)
+			chunksDecoder := schema.NewPrometheusParquetChunksDecoder(chunkenc.NewPool())
+			total := 0
+
+			for {
+				n, _ := rr.ReadRows(buf)
+				if n == 0 {
+					break
+				}
+
+				total += n
+				series, chunks, err := rowToSeries(rr.schema, chunksDecoder, buf[:n])
+				require.NoError(t, err)
+				require.Len(t, series, n)
+				for i, s := range series {
+					require.Contains(t, seriesHash, s.Hash())
+					require.Len(t, chunks[i], tt.expectedNumberOfChunks)
+					totalSamples := 0
+					for _, c := range chunks[i] {
+						require.Equal(t, tt.expectedPointsPerChunk, c.Chunk.NumSamples())
+						totalSamples += c.Chunk.NumSamples()
+					}
+					require.Equal(t, tt.numberOfSamples, totalSamples)
+				}
+			}
+
+			require.Equal(t, st.DB.Head().NumSeries(), uint64(total))
+		})
+	}
+}
+
+func Test_CreateParquetWithReducedTimestampSamples(t *testing.T) {
+	ctx := context.Background()
+	st := teststorage.New(t)
+	t.Cleanup(func() { st.Close() })
+
+	app := st.Appender(ctx)
+
+	// 240 samples * 30 seconds = 2 hours
+	step := (30 * time.Second).Milliseconds()
+	for i := 0; i < 240; i++ {
+		_, err := app.Append(0, labels.FromStrings("__name__", "foo"), int64(i)*step, float64(i))
+		require.NoError(t, err)
+	}
+
+	require.NoError(t, app.Commit())
+
+	h := st.DB.Head()
+	mint, maxt := (time.Minute * 30).Milliseconds(), (time.Minute*90).Milliseconds()-1
+	rr, err := newTsdbRowReader(ctx, mint, maxt, (time.Minute * 10).Milliseconds(), []Convertible{h})
+	require.NoError(t, err)
+	defer func() { _ = rr.Close() }()
+	// 6 data cols with 10 min duration
+	require.Len(t, rr.schema.DataColsIndexes, 6)
+
+	chunksDecoder := schema.NewPrometheusParquetChunksDecoder(chunkenc.NewPool())
+	buf := make([]parquet.Row, 100)
+	n, _ := rr.ReadRows(buf)
+	require.Equal(t, 1, n)
+
+	series, chunks, err := rowToSeries(rr.schema, chunksDecoder, buf[:n])
+	require.NoError(t, err)
+	require.Len(t, series, 1)
+	require.Len(t, chunks, 1)
+	require.Equal(t, labels.FromStrings("__name__", "foo").Hash(), series[0].Hash())
+
+	totalSamples := 0
+	for _, c := range chunks[0] {
+		totalSamples += c.Chunk.NumSamples()
+		require.LessOrEqual(t, c.MaxTime, maxt)
+		require.GreaterOrEqual(t, c.MinTime, mint)
+	}
+	require.Equal(t, 120, totalSamples)
+}
+
+func rowToSeries(s *schema.TSDBSchema, dec *schema.PrometheusParquetChunksDecoder, rows []parquet.Row) ([]labels.Labels, [][]chunks.Meta, error) {
+	cols := s.Schema.Columns()
+	b := labels.NewScratchBuilder(10)
+	series := make([]labels.Labels, len(rows))
+	chunksMetas := make([][]chunks.Meta, len(rows))
+
+	for i, row := range rows {
+		b.Reset()
+		for colIdx, colVal := range row {
+			col := cols[colIdx][0]
+			label, ok := schema.ExtractLabelFromColumn(col)
+			if ok {
+				b.Add(label, colVal.String())
+			}
+
+			if schema.IsDataColumn(col) {
+				c, err := dec.Decode(colVal.ByteArray(), 0, math.MaxInt64)
+				if err != nil {
+					return nil, nil, err
+				}
+				chunksMetas[i] = append(chunksMetas[i], c...)
+			}
+		}
+		series[i] = b.Labels()
+	}
+
+	return series, chunksMetas, nil
+}