go/parquet/internal/encoding/levels_test.go - arrow - Git at Google

 // Licensed to the Apache Software Foundation (ASF) under one
 // or more contributor license agreements.  See the NOTICE file
 // distributed with this work for additional information
 // regarding copyright ownership.  The ASF licenses this file
 // to you under the Apache License, Version 2.0 (the
 // "License"); you may not use this file except in compliance
 // with the License.  You may obtain a copy of the License at
 //
 // http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 package encoding_test

 import (
 	"encoding/binary"
 	"strconv"
 	"testing"

 	"github.com/apache/arrow/go/v6/arrow"
 	"github.com/apache/arrow/go/v6/arrow/memory"
 	"github.com/apache/arrow/go/v6/parquet"
 	"github.com/apache/arrow/go/v6/parquet/internal/encoding"
 	"github.com/apache/arrow/go/v6/parquet/internal/utils"
 	"github.com/stretchr/testify/assert"
 )

 func generateLevels(minRepeat, maxRepeat int, maxLevel int16) []int16 {
 	// for each repetition count up to max repeat
 	ret := make([]int16, 0)
 	for rep := minRepeat; rep <= maxRepeat; rep++ {
 		var (
 			repCount       = 1 << rep
 			val      int16 = 0
 			bwidth         = 0
 		)
 		// generate levels for repetition count up to max level
 		for val <= maxLevel {
 			for i := 0; i < repCount; i++ {
 				ret = append(ret, val)
 			}
 			val = int16((2 << bwidth) - 1)
 			bwidth++
 		}
 	}
 	return ret
 }

 func encodeLevels(t *testing.T, enc parquet.Encoding, maxLvl int16, numLevels int, input []int16) []byte {
 	var (
 		encoder  encoding.LevelEncoder
 		lvlCount = 0
 		buf      = encoding.NewBufferWriter(2*numLevels, memory.DefaultAllocator)
 	)

 	if enc == parquet.Encodings.RLE {
 		buf.SetOffset(arrow.Int32SizeBytes)
 		// leave space to write the rle length value
 		encoder.Init(enc, maxLvl, buf)
 		lvlCount, _ = encoder.Encode(input)
 		buf.SetOffset(0)
 		arrow.Int32Traits.CastFromBytes(buf.Bytes())[0] = utils.ToLEInt32(int32(encoder.Len()))
 	} else {
 		encoder.Init(enc, maxLvl, buf)
 		lvlCount, _ = encoder.Encode(input)
 	}

 	assert.Equal(t, numLevels, lvlCount)
 	return buf.Bytes()
 }

 func verifyDecodingLvls(t *testing.T, enc parquet.Encoding, maxLvl int16, input []int16, buf []byte) {
 	var (
 		decoder        encoding.LevelDecoder
 		lvlCount       = 0
 		numLevels      = len(input)
 		output         = make([]int16, numLevels)
 		decodeCount    = 4
 		numInnerLevels = numLevels / decodeCount
 	)

 	// decode levels and test with multiple decode calls
 	_, err := decoder.SetData(enc, maxLvl, numLevels, buf)
 	assert.NoError(t, err)
 	// try multiple decoding on a single setdata call
 	for ct := 0; ct < decodeCount; ct++ {
 		offset := ct * numInnerLevels
 		lvlCount, _ = decoder.Decode(output[:numInnerLevels])
 		assert.Equal(t, numInnerLevels, lvlCount)
 		assert.Equal(t, input[offset:offset+numInnerLevels], output[:numInnerLevels])
 	}

 	// check the remaining levels
 	var (
 		levelsCompleted = decodeCount * (numLevels / decodeCount)
 		remaining       = numLevels - levelsCompleted
 	)

 	if remaining > 0 {
 		lvlCount, _ = decoder.Decode(output[:remaining])
 		assert.Equal(t, remaining, lvlCount)
 		assert.Equal(t, input[levelsCompleted:], output[:remaining])
 	}
 	// test decode zero values
 	lvlCount, _ = decoder.Decode(output[:1])
 	assert.Zero(t, lvlCount)
 }

 func verifyDecodingMultipleSetData(t *testing.T, enc parquet.Encoding, max int16, input []int16, buf [][]byte) {
 	var (
 		decoder      encoding.LevelDecoder
 		lvlCount     = 0
 		setdataCount = len(buf)
 		numLevels    = len(input) / setdataCount
 		output       = make([]int16, numLevels)
 	)

 	for ct := 0; ct < setdataCount; ct++ {
 		offset := ct * numLevels
 		assert.Len(t, output, numLevels)
 		_, err := decoder.SetData(enc, max, numLevels, buf[ct])
 		assert.NoError(t, err)
 		lvlCount, _ = decoder.Decode(output)
 		assert.Equal(t, numLevels, lvlCount)
 		assert.Equal(t, input[offset:offset+numLevels], output)
 	}
 }

 func TestLevelsDecodeMultipleBitWidth(t *testing.T) {
 	t.Parallel()
 	// Test levels with maximum bit-width from 1 to 8
 	// increase the repetition count for each iteration by a factor of 2
 	var (
 		minRepeat   = 0
 		maxRepeat   = 7 // 128
 		maxBitWidth = 8
 		input       []int16
 		buf         []byte
 		encodings   = [2]parquet.Encoding{parquet.Encodings.RLE, parquet.Encodings.BitPacked}
 	)

 	for _, enc := range encodings {
 		t.Run(enc.String(), func(t *testing.T) {
 			// bitpacked requires a sequence of at least 8
 			if enc == parquet.Encodings.BitPacked {
 				minRepeat = 3
 			}
 			// for each max bit width
 			for bitWidth := 1; bitWidth <= maxBitWidth; bitWidth++ {
 				t.Run(strconv.Itoa(bitWidth), func(t *testing.T) {
 					max := int16((1 << bitWidth) - 1)
 					// generate levels
 					input = generateLevels(minRepeat, maxRepeat, max)
 					assert.NotPanics(t, func() {
 						buf = encodeLevels(t, enc, max, len(input), input)
 					})
 					assert.NotPanics(t, func() {
 						verifyDecodingLvls(t, enc, max, input, buf)
 					})
 				})
 			}
 		})
 	}
 }

 func TestLevelsDecodeMultipleSetData(t *testing.T) {
 	t.Parallel()

 	var (
 		minRepeat = 3
 		maxRepeat = 7
 		bitWidth  = 8
 		maxLevel  = int16((1 << bitWidth) - 1)
 		encodings = [2]parquet.Encoding{parquet.Encodings.RLE, parquet.Encodings.BitPacked}
 	)

 	input := generateLevels(minRepeat, maxRepeat, maxLevel)

 	var (
 		numLevels      = len(input)
 		setdataFactor  = 8
 		splitLevelSize = numLevels / setdataFactor
 		buf            = make([][]byte, setdataFactor)
 	)

 	for _, enc := range encodings {
 		t.Run(enc.String(), func(t *testing.T) {
 			for rf := 0; rf < setdataFactor; rf++ {
 				offset := rf * splitLevelSize
 				assert.NotPanics(t, func() {
 					buf[rf] = encodeLevels(t, enc, maxLevel, splitLevelSize, input[offset:offset+splitLevelSize])
 				})
 			}
 			assert.NotPanics(t, func() {
 				verifyDecodingMultipleSetData(t, enc, maxLevel, input, buf)
 			})
 		})
 	}
 }

 func TestMinimumBufferSize(t *testing.T) {
 	t.Parallel()

 	const numToEncode = 1024
 	levels := make([]int16, numToEncode)

 	for idx := range levels {
 		if idx%9 == 0 {
 			levels[idx] = 0
 		} else {
 			levels[idx] = 1
 		}
 	}

 	output := encoding.NewBufferWriter(0, memory.DefaultAllocator)

 	var encoder encoding.LevelEncoder
 	encoder.Init(parquet.Encodings.RLE, 1, output)
 	count, _ := encoder.Encode(levels)
 	assert.Equal(t, numToEncode, count)
 }

 func TestMinimumBufferSize2(t *testing.T) {
 	t.Parallel()

 	// test the worst case for bit_width=2 consisting of
 	// LiteralRun(size=8)
 	// RepeatedRun(size=8)
 	// LiteralRun(size=8)
 	// ...
 	const numToEncode = 1024
 	levels := make([]int16, numToEncode)

 	for idx := range levels {
 		// This forces a literal run of 00000001
 		// followed by eight 1s
 		if (idx % 16) < 7 {
 			levels[idx] = 0
 		} else {
 			levels[idx] = 1
 		}
 	}

 	for bitWidth := int16(1); bitWidth <= 8; bitWidth++ {
 		output := encoding.NewBufferWriter(0, memory.DefaultAllocator)

 		var encoder encoding.LevelEncoder
 		encoder.Init(parquet.Encodings.RLE, bitWidth, output)
 		count, _ := encoder.Encode(levels)
 		assert.Equal(t, numToEncode, count)
 	}
 }

 func TestEncodeDecodeLevels(t *testing.T) {
 	t.Parallel()
 	const numToEncode = 2048
 	levels := make([]int16, numToEncode)
 	numones := 0
 	for idx := range levels {
 		if (idx % 16) < 7 {
 			levels[idx] = 0
 		} else {
 			levels[idx] = 1
 			numones++
 		}
 	}

 	output := encoding.NewBufferWriter(0, memory.DefaultAllocator)

 	var encoder encoding.LevelEncoder
 	encoder.Init(parquet.Encodings.RLE, 1, output)
 	count, _ := encoder.Encode(levels)
 	assert.Equal(t, numToEncode, count)
 	encoder.Flush()

 	buf := output.Bytes()
 	var prefix [4]byte
 	binary.LittleEndian.PutUint32(prefix[:], uint32(len(buf)))

 	var decoder encoding.LevelDecoder
 	_, err := decoder.SetData(parquet.Encodings.RLE, 1, numToEncode, append(prefix[:], buf...))
 	assert.NoError(t, err)

 	var levelOut [numToEncode]int16
 	total, vals := decoder.Decode(levelOut[:])
 	assert.EqualValues(t, numToEncode, total)
 	assert.EqualValues(t, numones, vals)
 	assert.Equal(t, levels, levelOut[:])
 }
	// Licensed to the Apache Software Foundation (ASF) under one
	// or more contributor license agreements. See the NOTICE file
	// distributed with this work for additional information
	// regarding copyright ownership. The ASF licenses this file
	// to you under the Apache License, Version 2.0 (the
	// "License"); you may not use this file except in compliance
	// with the License. You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	package encoding_test

	import (
	"encoding/binary"
	"strconv"
	"testing"

	"github.com/apache/arrow/go/v6/arrow"
	"github.com/apache/arrow/go/v6/arrow/memory"
	"github.com/apache/arrow/go/v6/parquet"
	"github.com/apache/arrow/go/v6/parquet/internal/encoding"
	"github.com/apache/arrow/go/v6/parquet/internal/utils"
	"github.com/stretchr/testify/assert"
	)

	func generateLevels(minRepeat, maxRepeat int, maxLevel int16) []int16 {
	// for each repetition count up to max repeat
	ret := make([]int16, 0)
	for rep := minRepeat; rep <= maxRepeat; rep++ {
	var (
	repCount = 1 << rep
	val int16 = 0
	bwidth = 0
	)
	// generate levels for repetition count up to max level
	for val <= maxLevel {
	for i := 0; i < repCount; i++ {
	ret = append(ret, val)
	}
	val = int16((2 << bwidth) - 1)
	bwidth++
	}
	}
	return ret
	}

	func encodeLevels(t *testing.T, enc parquet.Encoding, maxLvl int16, numLevels int, input []int16) []byte {
	var (
	encoder encoding.LevelEncoder
	lvlCount = 0
	buf = encoding.NewBufferWriter(2*numLevels, memory.DefaultAllocator)
	)

	if enc == parquet.Encodings.RLE {
	buf.SetOffset(arrow.Int32SizeBytes)
	// leave space to write the rle length value
	encoder.Init(enc, maxLvl, buf)
	lvlCount, _ = encoder.Encode(input)
	buf.SetOffset(0)
	arrow.Int32Traits.CastFromBytes(buf.Bytes())[0] = utils.ToLEInt32(int32(encoder.Len()))
	} else {
	encoder.Init(enc, maxLvl, buf)
	lvlCount, _ = encoder.Encode(input)
	}

	assert.Equal(t, numLevels, lvlCount)
	return buf.Bytes()
	}

	func verifyDecodingLvls(t *testing.T, enc parquet.Encoding, maxLvl int16, input []int16, buf []byte) {
	var (
	decoder encoding.LevelDecoder
	lvlCount = 0
	numLevels = len(input)
	output = make([]int16, numLevels)
	decodeCount = 4
	numInnerLevels = numLevels / decodeCount
	)

	// decode levels and test with multiple decode calls
	_, err := decoder.SetData(enc, maxLvl, numLevels, buf)
	assert.NoError(t, err)
	// try multiple decoding on a single setdata call
	for ct := 0; ct < decodeCount; ct++ {
	offset := ct * numInnerLevels
	lvlCount, _ = decoder.Decode(output[:numInnerLevels])
	assert.Equal(t, numInnerLevels, lvlCount)
	assert.Equal(t, input[offset:offset+numInnerLevels], output[:numInnerLevels])
	}

	// check the remaining levels
	var (
	levelsCompleted = decodeCount * (numLevels / decodeCount)
	remaining = numLevels - levelsCompleted
	)

	if remaining > 0 {
	lvlCount, _ = decoder.Decode(output[:remaining])
	assert.Equal(t, remaining, lvlCount)
	assert.Equal(t, input[levelsCompleted:], output[:remaining])
	}
	// test decode zero values
	lvlCount, _ = decoder.Decode(output[:1])
	assert.Zero(t, lvlCount)
	}

	func verifyDecodingMultipleSetData(t *testing.T, enc parquet.Encoding, max int16, input []int16, buf [][]byte) {
	var (
	decoder encoding.LevelDecoder
	lvlCount = 0
	setdataCount = len(buf)
	numLevels = len(input) / setdataCount
	output = make([]int16, numLevels)
	)

	for ct := 0; ct < setdataCount; ct++ {
	offset := ct * numLevels
	assert.Len(t, output, numLevels)
	_, err := decoder.SetData(enc, max, numLevels, buf[ct])
	assert.NoError(t, err)
	lvlCount, _ = decoder.Decode(output)
	assert.Equal(t, numLevels, lvlCount)
	assert.Equal(t, input[offset:offset+numLevels], output)
	}
	}

	func TestLevelsDecodeMultipleBitWidth(t *testing.T) {
	t.Parallel()
	// Test levels with maximum bit-width from 1 to 8
	// increase the repetition count for each iteration by a factor of 2
	var (
	minRepeat = 0
	maxRepeat = 7 // 128
	maxBitWidth = 8
	input []int16
	buf []byte
	encodings = [2]parquet.Encoding{parquet.Encodings.RLE, parquet.Encodings.BitPacked}
	)

	for _, enc := range encodings {
	t.Run(enc.String(), func(t *testing.T) {
	// bitpacked requires a sequence of at least 8
	if enc == parquet.Encodings.BitPacked {
	minRepeat = 3
	}
	// for each max bit width
	for bitWidth := 1; bitWidth <= maxBitWidth; bitWidth++ {
	t.Run(strconv.Itoa(bitWidth), func(t *testing.T) {
	max := int16((1 << bitWidth) - 1)
	// generate levels
	input = generateLevels(minRepeat, maxRepeat, max)
	assert.NotPanics(t, func() {
	buf = encodeLevels(t, enc, max, len(input), input)
	})
	assert.NotPanics(t, func() {
	verifyDecodingLvls(t, enc, max, input, buf)
	})
	})
	}
	})
	}
	}

	func TestLevelsDecodeMultipleSetData(t *testing.T) {
	t.Parallel()

	var (
	minRepeat = 3
	maxRepeat = 7
	bitWidth = 8
	maxLevel = int16((1 << bitWidth) - 1)
	encodings = [2]parquet.Encoding{parquet.Encodings.RLE, parquet.Encodings.BitPacked}
	)

	input := generateLevels(minRepeat, maxRepeat, maxLevel)

	var (
	numLevels = len(input)
	setdataFactor = 8
	splitLevelSize = numLevels / setdataFactor
	buf = make([][]byte, setdataFactor)
	)

	for _, enc := range encodings {
	t.Run(enc.String(), func(t *testing.T) {
	for rf := 0; rf < setdataFactor; rf++ {
	offset := rf * splitLevelSize
	assert.NotPanics(t, func() {
	buf[rf] = encodeLevels(t, enc, maxLevel, splitLevelSize, input[offset:offset+splitLevelSize])
	})
	}
	assert.NotPanics(t, func() {
	verifyDecodingMultipleSetData(t, enc, maxLevel, input, buf)
	})
	})
	}
	}

	func TestMinimumBufferSize(t *testing.T) {
	t.Parallel()

	const numToEncode = 1024
	levels := make([]int16, numToEncode)

	for idx := range levels {
	if idx%9 == 0 {
	levels[idx] = 0
	} else {
	levels[idx] = 1
	}
	}

	output := encoding.NewBufferWriter(0, memory.DefaultAllocator)

	var encoder encoding.LevelEncoder
	encoder.Init(parquet.Encodings.RLE, 1, output)
	count, _ := encoder.Encode(levels)
	assert.Equal(t, numToEncode, count)
	}

	func TestMinimumBufferSize2(t *testing.T) {
	t.Parallel()

	// test the worst case for bit_width=2 consisting of
	// LiteralRun(size=8)
	// RepeatedRun(size=8)
	// LiteralRun(size=8)
	// ...
	const numToEncode = 1024
	levels := make([]int16, numToEncode)

	for idx := range levels {
	// This forces a literal run of 00000001
	// followed by eight 1s
	if (idx % 16) < 7 {
	levels[idx] = 0
	} else {
	levels[idx] = 1
	}
	}

	for bitWidth := int16(1); bitWidth <= 8; bitWidth++ {
	output := encoding.NewBufferWriter(0, memory.DefaultAllocator)

	var encoder encoding.LevelEncoder
	encoder.Init(parquet.Encodings.RLE, bitWidth, output)
	count, _ := encoder.Encode(levels)
	assert.Equal(t, numToEncode, count)
	}
	}

	func TestEncodeDecodeLevels(t *testing.T) {
	t.Parallel()
	const numToEncode = 2048
	levels := make([]int16, numToEncode)
	numones := 0
	for idx := range levels {
	if (idx % 16) < 7 {
	levels[idx] = 0
	} else {
	levels[idx] = 1
	numones++
	}
	}

	output := encoding.NewBufferWriter(0, memory.DefaultAllocator)

	var encoder encoding.LevelEncoder
	encoder.Init(parquet.Encodings.RLE, 1, output)
	count, _ := encoder.Encode(levels)
	assert.Equal(t, numToEncode, count)
	encoder.Flush()

	buf := output.Bytes()
	var prefix [4]byte
	binary.LittleEndian.PutUint32(prefix[:], uint32(len(buf)))

	var decoder encoding.LevelDecoder
	_, err := decoder.SetData(parquet.Encodings.RLE, 1, numToEncode, append(prefix[:], buf...))
	assert.NoError(t, err)

	var levelOut [numToEncode]int16
	total, vals := decoder.Decode(levelOut[:])
	assert.EqualValues(t, numToEncode, total)
	assert.EqualValues(t, numones, vals)
	assert.Equal(t, levels, levelOut[:])
	}