brotli.go

Documentation: github.com/apache/arrow/go/v15/parquet/compress

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