github.com/apache/arrow/go/v7@v7.0.1/parquet/compress/brotli.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 compress
    18  
    19  import (
    20  	"bytes"
    21  	"io"
    22  	"io/ioutil"
    23  
    24  	"github.com/andybalholm/brotli"
    25  	"github.com/apache/arrow/go/v7/parquet/internal/debug"
    26  )
    27  
    28  type brotliCodec struct{}
    29  
    30  func (brotliCodec) NewReader(r io.Reader) io.ReadCloser {
    31  	return ioutil.NopCloser(brotli.NewReader(r))
    32  }
    33  
    34  func (b brotliCodec) EncodeLevel(dst, src []byte, level int) []byte {
    35  	if level == DefaultCompressionLevel {
    36  		level = brotli.DefaultCompression
    37  	}
    38  
    39  	maxlen := int(b.CompressBound(int64(len(src))))
    40  	if dst == nil || cap(dst) < maxlen {
    41  		dst = make([]byte, 0, maxlen)
    42  	}
    43  	buf := bytes.NewBuffer(dst[:0])
    44  	w := brotli.NewWriterLevel(buf, level)
    45  	_, err := w.Write(src)
    46  	if err != nil {
    47  		panic(err)
    48  	}
    49  	if err := w.Close(); err != nil {
    50  		panic(err)
    51  	}
    52  	return buf.Bytes()
    53  }
    54  
    55  func (b brotliCodec) Encode(dst, src []byte) []byte {
    56  	return b.EncodeLevel(dst, src, brotli.DefaultCompression)
    57  }
    58  
    59  func (brotliCodec) Decode(dst, src []byte) []byte {
    60  	rdr := brotli.NewReader(bytes.NewReader(src))
    61  	if dst != nil {
    62  		var (
    63  			sofar       = 0
    64  			n           = -1
    65  			err   error = nil
    66  		)
    67  		for n != 0 && err == nil {
    68  			n, err = rdr.Read(dst[sofar:])
    69  			sofar += n
    70  		}
    71  		if err != nil && err != io.EOF {
    72  			panic(err)
    73  		}
    74  		return dst[:sofar]
    75  	}
    76  
    77  	dst, err := ioutil.ReadAll(rdr)
    78  	if err != nil {
    79  		panic(err)
    80  	}
    81  
    82  	return dst
    83  }
    84  
    85  // taken from brotli/enc/encode.c:1426
    86  // BrotliEncoderMaxCompressedSize
    87  func (brotliCodec) CompressBound(len int64) int64 {
    88  	// [window bits / empty metadata] + N * [uncompressed] + [last empty]
    89  	debug.Assert(len > 0, "brotli compressbound should be > 0")
    90  	nlarge := len >> 14
    91  	overhead := 2 + (4 * nlarge) + 3 + 1
    92  	result := len + overhead
    93  	if len == 0 {
    94  		return 2
    95  	}
    96  	if result < len {
    97  		return 0
    98  	}
    99  	return len
   100  }
   101  
   102  func (brotliCodec) NewWriter(w io.Writer) io.WriteCloser {
   103  	return brotli.NewWriter(w)
   104  }
   105  
   106  func (brotliCodec) NewWriterLevel(w io.Writer, level int) (io.WriteCloser, error) {
   107  	if level == DefaultCompressionLevel {
   108  		level = brotli.DefaultCompression
   109  	}
   110  	return brotli.NewWriterLevel(w, level), nil
   111  }
   112  
   113  func init() {
   114  	codecs[Codecs.Brotli] = brotliCodec{}
   115  }