github.com/apache/arrow/go/v14@v14.0.2/parquet/reader_writer_properties_test.go (about)

     1  // Licensed to the Apache Software Foundation (ASF) under one
     2  // or more contributor license agreements.  See the NOTICE file
     3  // distributed with this work for additional information
     4  // regarding copyright ownership.  The ASF licenses this file
     5  // to you under the Apache License, Version 2.0 (the
     6  // "License"); you may not use this file except in compliance
     7  // with the License.  You may obtain a copy of the License at
     8  //
     9  // http://www.apache.org/licenses/LICENSE-2.0
    10  //
    11  // Unless required by applicable law or agreed to in writing, software
    12  // distributed under the License is distributed on an "AS IS" BASIS,
    13  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  // See the License for the specific language governing permissions and
    15  // limitations under the License.
    16  
    17  package parquet_test
    18  
    19  import (
    20  	"bytes"
    21  	"testing"
    22  
    23  	"github.com/apache/arrow/go/v14/arrow/memory"
    24  	"github.com/apache/arrow/go/v14/parquet"
    25  	"github.com/apache/arrow/go/v14/parquet/compress"
    26  	"github.com/stretchr/testify/assert"
    27  )
    28  
    29  func TestReaderPropBasics(t *testing.T) {
    30  	props := parquet.NewReaderProperties(nil)
    31  	assert.Equal(t, parquet.DefaultBufSize, props.BufferSize)
    32  	assert.False(t, props.BufferedStreamEnabled)
    33  }
    34  
    35  func TestWriterPropBasics(t *testing.T) {
    36  	props := parquet.NewWriterProperties()
    37  
    38  	assert.Equal(t, parquet.DefaultDataPageSize, props.DataPageSize())
    39  	assert.Equal(t, parquet.DefaultDictionaryPageSizeLimit, props.DictionaryPageSizeLimit())
    40  	assert.Equal(t, parquet.V2_LATEST, props.Version())
    41  	assert.Equal(t, parquet.DataPageV1, props.DataPageVersion())
    42  }
    43  
    44  func TestWriterPropAdvanced(t *testing.T) {
    45  	props := parquet.NewWriterProperties(
    46  		parquet.WithCompressionFor("gzip", compress.Codecs.Gzip),
    47  		parquet.WithCompressionFor("zstd", compress.Codecs.Zstd),
    48  		parquet.WithCompression(compress.Codecs.Snappy),
    49  		parquet.WithEncoding(parquet.Encodings.DeltaBinaryPacked),
    50  		parquet.WithEncodingFor("delta-length", parquet.Encodings.DeltaLengthByteArray),
    51  		parquet.WithDataPageVersion(parquet.DataPageV2),
    52  		parquet.WithRootName("test2"),
    53  		parquet.WithRootRepetition(parquet.Repetitions.Required))
    54  
    55  	assert.Equal(t, compress.Codecs.Gzip, props.CompressionPath(parquet.ColumnPathFromString("gzip")))
    56  	assert.Equal(t, compress.Codecs.Zstd, props.CompressionFor("zstd"))
    57  	assert.Equal(t, compress.Codecs.Snappy, props.CompressionPath(parquet.ColumnPathFromString("delta-length")))
    58  	assert.Equal(t, parquet.Encodings.DeltaBinaryPacked, props.EncodingFor("gzip"))
    59  	assert.Equal(t, parquet.Encodings.DeltaLengthByteArray, props.EncodingPath(parquet.ColumnPathFromString("delta-length")))
    60  	assert.Equal(t, parquet.DataPageV2, props.DataPageVersion())
    61  	assert.Equal(t, "test2", props.RootName())
    62  	assert.Equal(t, parquet.Repetitions.Required, props.RootRepetition())
    63  }
    64  
    65  func TestReaderPropsGetStreamInsufficient(t *testing.T) {
    66  	data := "shorter than expected"
    67  	buf := memory.NewBufferBytes([]byte(data))
    68  	rdr := bytes.NewReader(buf.Bytes())
    69  
    70  	props := parquet.NewReaderProperties(nil)
    71  	_, err := props.GetStream(rdr, 12, 15)
    72  	assert.Error(t, err)
    73  }