storj.io/minio@v0.0.0-20230509071714-0cbc90f649b1/pkg/s3select/internal/parquet-go/common/common.go (about)

     1  package common
     2  
     3  import (
     4  	"bytes"
     5  	"compress/gzip"
     6  	"fmt"
     7  	"io/ioutil"
     8  
     9  	"github.com/klauspost/compress/snappy"
    10  	"github.com/pierrec/lz4"
    11  
    12  	"storj.io/minio/pkg/s3select/internal/parquet-go/gen-go/parquet"
    13  )
    14  
    15  // ToSliceValue converts values to a slice value.
    16  func ToSliceValue(values []interface{}, parquetType parquet.Type) interface{} {
    17  	switch parquetType {
    18  	case parquet.Type_BOOLEAN:
    19  		bs := make([]bool, len(values))
    20  		for i := range values {
    21  			bs[i] = values[i].(bool)
    22  		}
    23  		return bs
    24  	case parquet.Type_INT32:
    25  		i32s := make([]int32, len(values))
    26  		for i := range values {
    27  			i32s[i] = values[i].(int32)
    28  		}
    29  		return i32s
    30  	case parquet.Type_INT64:
    31  		i64s := make([]int64, len(values))
    32  		for i := range values {
    33  			i64s[i] = values[i].(int64)
    34  		}
    35  		return i64s
    36  	case parquet.Type_FLOAT:
    37  		f32s := make([]float32, len(values))
    38  		for i := range values {
    39  			f32s[i] = values[i].(float32)
    40  		}
    41  		return f32s
    42  	case parquet.Type_DOUBLE:
    43  		f64s := make([]float64, len(values))
    44  		for i := range values {
    45  			f64s[i] = values[i].(float64)
    46  		}
    47  		return f64s
    48  	case parquet.Type_BYTE_ARRAY:
    49  		array := make([][]byte, len(values))
    50  		for i := range values {
    51  			array[i] = values[i].([]byte)
    52  		}
    53  		return array
    54  	}
    55  
    56  	return nil
    57  }
    58  
    59  // BitWidth returns bits count required to accommodate given value.
    60  func BitWidth(ui64 uint64) (width int32) {
    61  	for ; ui64 != 0; ui64 >>= 1 {
    62  		width++
    63  	}
    64  
    65  	return width
    66  }
    67  
    68  // Compress compresses given data.
    69  func Compress(compressionType parquet.CompressionCodec, data []byte) ([]byte, error) {
    70  	switch compressionType {
    71  	case parquet.CompressionCodec_UNCOMPRESSED:
    72  		return data, nil
    73  
    74  	case parquet.CompressionCodec_SNAPPY:
    75  		return snappy.Encode(nil, data), nil
    76  
    77  	case parquet.CompressionCodec_GZIP:
    78  		buf := new(bytes.Buffer)
    79  		writer := gzip.NewWriter(buf)
    80  		n, err := writer.Write(data)
    81  		if err != nil {
    82  			return nil, err
    83  		}
    84  		if n != len(data) {
    85  			return nil, fmt.Errorf("short writes")
    86  		}
    87  
    88  		if err = writer.Flush(); err != nil {
    89  			return nil, err
    90  		}
    91  
    92  		if err = writer.Close(); err != nil {
    93  			return nil, err
    94  		}
    95  
    96  		return buf.Bytes(), nil
    97  
    98  	case parquet.CompressionCodec_LZ4:
    99  		buf := new(bytes.Buffer)
   100  		writer := lz4.NewWriter(buf)
   101  		n, err := writer.Write(data)
   102  		if err != nil {
   103  			return nil, err
   104  		}
   105  		if n != len(data) {
   106  			return nil, fmt.Errorf("short writes")
   107  		}
   108  
   109  		if err = writer.Flush(); err != nil {
   110  			return nil, err
   111  		}
   112  
   113  		if err = writer.Close(); err != nil {
   114  			return nil, err
   115  		}
   116  
   117  		return buf.Bytes(), nil
   118  	}
   119  
   120  	return nil, fmt.Errorf("unsupported compression codec %v", compressionType)
   121  }
   122  
   123  // Uncompress uncompresses given data.
   124  func Uncompress(compressionType parquet.CompressionCodec, data []byte) ([]byte, error) {
   125  	switch compressionType {
   126  	case parquet.CompressionCodec_UNCOMPRESSED:
   127  		return data, nil
   128  
   129  	case parquet.CompressionCodec_SNAPPY:
   130  		return snappy.Decode(nil, data)
   131  
   132  	case parquet.CompressionCodec_GZIP:
   133  		reader, err := gzip.NewReader(bytes.NewReader(data))
   134  		if err != nil {
   135  			return nil, err
   136  		}
   137  		defer reader.Close()
   138  		return ioutil.ReadAll(reader)
   139  
   140  	case parquet.CompressionCodec_LZ4:
   141  		return ioutil.ReadAll(lz4.NewReader(bytes.NewReader(data)))
   142  	}
   143  
   144  	return nil, fmt.Errorf("unsupported compression codec %v", compressionType)
   145  }