github.com/vc42/parquet-go@v0.0.0-20240320194221-1a9adb5f23f5/compress.go (about)

     1  package parquet
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"github.com/vc42/parquet-go/compress"
     7  	"github.com/vc42/parquet-go/compress/brotli"
     8  	"github.com/vc42/parquet-go/compress/gzip"
     9  	"github.com/vc42/parquet-go/compress/lz4"
    10  	"github.com/vc42/parquet-go/compress/snappy"
    11  	"github.com/vc42/parquet-go/compress/uncompressed"
    12  	"github.com/vc42/parquet-go/compress/zstd"
    13  	"github.com/vc42/parquet-go/format"
    14  )
    15  
    16  var (
    17  	// Uncompressed is a parquet compression codec representing uncompressed
    18  	// pages.
    19  	Uncompressed uncompressed.Codec
    20  
    21  	// Snappy is the SNAPPY parquet compression codec.
    22  	Snappy snappy.Codec
    23  
    24  	// Gzip is the GZIP parquet compression codec.
    25  	Gzip = gzip.Codec{
    26  		Level: gzip.DefaultCompression,
    27  	}
    28  
    29  	// Brotli is the BROTLI parquet compression codec.
    30  	Brotli = brotli.Codec{
    31  		Quality: brotli.DefaultQuality,
    32  		LGWin:   brotli.DefaultLGWin,
    33  	}
    34  
    35  	// Zstd is the ZSTD parquet compression codec.
    36  	Zstd = zstd.Codec{
    37  		Level: zstd.DefaultLevel,
    38  	}
    39  
    40  	// Lz4Raw is the LZ4_RAW parquet compression codec.
    41  	Lz4Raw = lz4.Codec{
    42  		Level: lz4.DefaultLevel,
    43  	}
    44  
    45  	// Table of compression codecs indexed by their code in the parquet format.
    46  	compressionCodecs = [...]compress.Codec{
    47  		format.Uncompressed: &Uncompressed,
    48  		format.Snappy:       &Snappy,
    49  		format.Gzip:         &Gzip,
    50  		format.Brotli:       &Brotli,
    51  		format.Zstd:         &Zstd,
    52  		format.Lz4Raw:       &Lz4Raw,
    53  	}
    54  )
    55  
    56  // LookupCompressionCodec returns the compression codec associated with the
    57  // given code.
    58  //
    59  // The function never returns nil. If the encoding is not supported,
    60  // an "unsupported" codec is returned.
    61  func LookupCompressionCodec(codec format.CompressionCodec) compress.Codec {
    62  	if codec >= 0 && int(codec) < len(compressionCodecs) {
    63  		if c := compressionCodecs[codec]; c != nil {
    64  			return c
    65  		}
    66  	}
    67  	return &unsupported{codec}
    68  }
    69  
    70  type unsupported struct {
    71  	codec format.CompressionCodec
    72  }
    73  
    74  func (u *unsupported) String() string {
    75  	return "UNSUPPORTED"
    76  }
    77  
    78  func (u *unsupported) CompressionCodec() format.CompressionCodec {
    79  	return u.codec
    80  }
    81  
    82  func (u *unsupported) Encode(dst, src []byte) ([]byte, error) {
    83  	return dst[:0], u.error()
    84  }
    85  
    86  func (u *unsupported) Decode(dst, src []byte) ([]byte, error) {
    87  	return dst[:0], u.error()
    88  }
    89  
    90  func (u *unsupported) error() error {
    91  	return fmt.Errorf("unsupported compression codec: %s", u.codec)
    92  }
    93  
    94  func isCompressed(c compress.Codec) bool {
    95  	return c != nil && c.CompressionCodec() != format.Uncompressed
    96  }