diff --git a/sstable/colblk_writer.go b/sstable/colblk_writer.go index 5e35eb4740..7e39045f48 100644 --- a/sstable/colblk_writer.go +++ b/sstable/colblk_writer.go @@ -203,11 +203,13 @@ func (w *RawColumnWriter) EstimatedSize() uint64 { if w.rangeKeyBlock.KeyCount() > 0 { sz += uint64(w.rangeKeyBlock.Size()) } - for _, blk := range w.valueBlock.blocks { - sz += uint64(blk.block.LengthWithTrailer()) - } - if w.valueBlock.buf != nil { - sz += uint64(len(w.valueBlock.buf.b)) + if w.valueBlock != nil { + for _, blk := range w.valueBlock.blocks { + sz += uint64(blk.block.LengthWithTrailer()) + } + if w.valueBlock.buf != nil { + sz += uint64(len(w.valueBlock.buf.b)) + } } // TODO(jackson): Include an estimate of the properties, filter and meta // index blocks sizes. diff --git a/sstable/data_test.go b/sstable/data_test.go index 5f218c6eb3..452f17d65b 100644 --- a/sstable/data_test.go +++ b/sstable/data_test.go @@ -103,39 +103,8 @@ func runBuildMemObjCmd( _ = w.Close() } }() - for _, data := range strings.Split(td.Input, "\n") { - err := func() (err error) { - defer func() { - if r := recover(); r != nil { - err = errors.Errorf("%v", r) - } - }() - - switch { - case strings.HasPrefix(data, "EncodeSpan:"): - return w.EncodeSpan(keyspan.ParseSpan(strings.TrimPrefix(data, "EncodeSpan:"))) - default: - forceObsolete := strings.HasPrefix(data, "force-obsolete:") - if forceObsolete { - data = strings.TrimSpace(strings.TrimPrefix(data, "force-obsolete:")) - } - j := strings.Index(data, ":") - key := base.ParseInternalKey(data[:j]) - value := []byte(data[j+1:]) - switch key.Kind() { - case InternalKeyKindRangeDelete: - if forceObsolete { - return errors.Errorf("force-obsolete is not allowed for RANGEDEL") - } - return w.AddWithForceObsolete(key, value, false /* forceObsolete */) - default: - return w.AddWithForceObsolete(key, value, forceObsolete) - } - } - }() - if err != nil { - return nil, nil, err - } + if err := writeKVs(w, td.Input); err != nil { + return nil, nil, err } if err := w.Close(); err != nil { return nil, nil, err @@ -148,6 +117,41 @@ func runBuildMemObjCmd( return meta, obj, nil } +func writeKVs(w RawWriter, input string) (err error) { + defer func() { + if r := recover(); r != nil { + err = errors.Errorf("%v", r) + } + }() + for _, data := range strings.Split(input, "\n") { + switch { + case strings.HasPrefix(data, "EncodeSpan:"): + err = w.EncodeSpan(keyspan.ParseSpan(strings.TrimPrefix(data, "EncodeSpan:"))) + default: + forceObsolete := strings.HasPrefix(data, "force-obsolete:") + if forceObsolete { + data = strings.TrimSpace(strings.TrimPrefix(data, "force-obsolete:")) + } + j := strings.Index(data, ":") + key := base.ParseInternalKey(data[:j]) + value := []byte(data[j+1:]) + switch key.Kind() { + case InternalKeyKindRangeDelete: + if forceObsolete { + return errors.Errorf("force-obsolete is not allowed for RANGEDEL") + } + err = w.AddWithForceObsolete(key, value, false /* forceObsolete */) + default: + err = w.AddWithForceObsolete(key, value, forceObsolete) + } + } + if err != nil { + return err + } + } + return err +} + func runBuildCmd( td *datadriven.TestData, writerOpts *WriterOptions, cacheSize int, ) (*WriterMetadata, *Reader, error) { diff --git a/sstable/testdata/writer_v5 b/sstable/testdata/writer_v5 index f825bdb95c..c220740356 100644 --- a/sstable/testdata/writer_v5 +++ b/sstable/testdata/writer_v5 @@ -327,3 +327,18 @@ rocksdb.property.collectors: [obsolete-key] pebble.raw.range-key.key.size: 6 pebble.raw.range-key.value.size: 9 obsolete-key: hex:0074 + +open-writer disable-value-blocks +---- + +write-kvs +a@2.SET.1:a2 +a@1.SET.1:a1 +b@2.SET.1:b2 +---- +EstimatedSize()=53 + +close +---- +point: [a@2#1,SET-b@2#1,SET] +seqnums: [1-1] diff --git a/sstable/writer_test.go b/sstable/writer_test.go index bac220b59c..13de340e12 100644 --- a/sstable/writer_test.go +++ b/sstable/writer_test.go @@ -120,20 +120,24 @@ func formatWriterMetadata(td *datadriven.TestData, m *WriterMetadata) string { func runDataDriven(t *testing.T, file string, tableFormat TableFormat, parallelism bool) { var r *Reader + var w RawWriter var obj *objstorage.MemObj - defer func() { + closeReaderAndWriter := func() { if r != nil { require.NoError(t, r.Close()) } - }() + if w != nil { + require.NoError(t, w.Close()) + } + r = nil + w = nil + } + defer closeReaderAndWriter() datadriven.RunTest(t, file, func(t *testing.T, td *datadriven.TestData) string { switch td.Cmd { case "build": - if r != nil { - _ = r.Close() - r = nil - } + closeReaderAndWriter() var meta *WriterMetadata var err error wopts := &WriterOptions{ @@ -152,6 +156,7 @@ func runDataDriven(t *testing.T, file string, tableFormat TableFormat, paralleli return formatWriterMetadata(td, meta) case "build-raw": + closeReaderAndWriter() if r != nil { _ = r.Close() r = nil @@ -167,6 +172,41 @@ func runDataDriven(t *testing.T, file string, tableFormat TableFormat, paralleli } return formatWriterMetadata(td, meta) + case "open-writer": + if r != nil { + _ = r.Close() + r = nil + } + wopts := &WriterOptions{ + Comparer: testkeys.Comparer, + DisableValueBlocks: td.HasArg("disable-value-blocks"), + TableFormat: tableFormat, + Parallelism: parallelism, + } + obj := &objstorage.MemObj{} + if err := optsFromArgs(td, wopts); err != nil { + return err.Error() + } + w = NewRawWriter(obj, *wopts) + return "" + + case "write-kvs": + if err := writeKVs(w, td.Input); err != nil { + return err.Error() + } + return fmt.Sprintf("EstimatedSize()=%d", w.EstimatedSize()) + + case "close": + if err := w.Close(); err != nil { + return err.Error() + } + meta, err := w.Metadata() + if err != nil { + return err.Error() + } + w = nil + return formatWriterMetadata(td, meta) + case "scan": iter, err := r.NewIter(NoTransforms, nil /* lower */, nil /* upper */) if err != nil {