github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/col/colserde/file_test.go (about)

     1  // Copyright 2019 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package colserde_test
    12  
    13  import (
    14  	"bytes"
    15  	"os"
    16  	"path/filepath"
    17  	"testing"
    18  
    19  	"github.com/cockroachdb/cockroach/pkg/col/coldata"
    20  	"github.com/cockroachdb/cockroach/pkg/col/coldatatestutils"
    21  	"github.com/cockroachdb/cockroach/pkg/col/colserde"
    22  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    23  	"github.com/cockroachdb/cockroach/pkg/testutils"
    24  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    25  	"github.com/stretchr/testify/require"
    26  )
    27  
    28  func TestFileRoundtrip(t *testing.T) {
    29  	defer leaktest.AfterTest(t)()
    30  	typs, b := randomBatch(testAllocator)
    31  
    32  	t.Run(`mem`, func(t *testing.T) {
    33  		// Make a copy of the original batch because the converter modifies and
    34  		// casts data without copying for performance reasons.
    35  		original := coldatatestutils.CopyBatch(b, typs, testColumnFactory)
    36  
    37  		var buf bytes.Buffer
    38  		s, err := colserde.NewFileSerializer(&buf, typs)
    39  		require.NoError(t, err)
    40  		require.NoError(t, s.AppendBatch(b))
    41  		require.NoError(t, s.Finish())
    42  
    43  		// Parts of the deserialization modify things (null bitmaps) in place, so
    44  		// run it twice to make sure those modifications don't leak back to the
    45  		// buffer.
    46  		for i := 0; i < 2; i++ {
    47  			func() {
    48  				roundtrip := testAllocator.NewMemBatchWithSize(typs, b.Length())
    49  				d, err := colserde.NewFileDeserializerFromBytes(typs, buf.Bytes())
    50  				require.NoError(t, err)
    51  				defer func() { require.NoError(t, d.Close()) }()
    52  				require.Equal(t, typs, d.Typs())
    53  				require.Equal(t, 1, d.NumBatches())
    54  				require.NoError(t, d.GetBatch(0, roundtrip))
    55  
    56  				coldata.AssertEquivalentBatches(t, original, roundtrip)
    57  			}()
    58  		}
    59  	})
    60  
    61  	t.Run(`disk`, func(t *testing.T) {
    62  		dir, cleanup := testutils.TempDir(t)
    63  		defer cleanup()
    64  		path := filepath.Join(dir, `rng.arrow`)
    65  
    66  		// Make a copy of the original batch because the converter modifies and
    67  		// casts data without copying for performance reasons.
    68  		original := coldatatestutils.CopyBatch(b, typs, testColumnFactory)
    69  
    70  		f, err := os.Create(path)
    71  		require.NoError(t, err)
    72  		defer func() { require.NoError(t, f.Close()) }()
    73  		s, err := colserde.NewFileSerializer(f, typs)
    74  		require.NoError(t, err)
    75  		require.NoError(t, s.AppendBatch(b))
    76  		require.NoError(t, s.Finish())
    77  		require.NoError(t, f.Sync())
    78  
    79  		// Parts of the deserialization modify things (null bitmaps) in place, so
    80  		// run it twice to make sure those modifications don't leak back to the
    81  		// file.
    82  		for i := 0; i < 2; i++ {
    83  			func() {
    84  				roundtrip := testAllocator.NewMemBatchWithSize(typs, b.Length())
    85  				d, err := colserde.NewFileDeserializerFromPath(typs, path)
    86  				require.NoError(t, err)
    87  				defer func() { require.NoError(t, d.Close()) }()
    88  				require.Equal(t, typs, d.Typs())
    89  				require.Equal(t, 1, d.NumBatches())
    90  				require.NoError(t, d.GetBatch(0, roundtrip))
    91  
    92  				coldata.AssertEquivalentBatches(t, original, roundtrip)
    93  			}()
    94  		}
    95  	})
    96  }
    97  
    98  func TestFileIndexing(t *testing.T) {
    99  	defer leaktest.AfterTest(t)()
   100  
   101  	const numInts = 10
   102  	typs := []*types.T{types.Int}
   103  	batchSize := 1
   104  
   105  	var buf bytes.Buffer
   106  	s, err := colserde.NewFileSerializer(&buf, typs)
   107  	require.NoError(t, err)
   108  
   109  	for i := 0; i < numInts; i++ {
   110  		b := testAllocator.NewMemBatchWithSize(typs, batchSize)
   111  		b.SetLength(batchSize)
   112  		b.ColVec(0).Int64()[0] = int64(i)
   113  		require.NoError(t, s.AppendBatch(b))
   114  	}
   115  	require.NoError(t, s.Finish())
   116  
   117  	d, err := colserde.NewFileDeserializerFromBytes(typs, buf.Bytes())
   118  	require.NoError(t, err)
   119  	defer func() { require.NoError(t, d.Close()) }()
   120  	require.Equal(t, typs, d.Typs())
   121  	require.Equal(t, numInts, d.NumBatches())
   122  	for batchIdx := numInts - 1; batchIdx >= 0; batchIdx-- {
   123  		b := testAllocator.NewMemBatchWithSize(typs, batchSize)
   124  		require.NoError(t, d.GetBatch(batchIdx, b))
   125  		require.Equal(t, batchSize, b.Length())
   126  		require.Equal(t, 1, b.Width())
   127  		require.Equal(t, int64(batchIdx), b.ColVec(0).Int64()[0])
   128  	}
   129  }