feat: support zstandard in ocf (#343)

go.mod

@@ -6,6 +6,7 @@ require (
  github.com/ettle/strcase v0.2.0
  github.com/golang/snappy v0.0.4
  github.com/json-iterator/go v1.1.12
+ github.com/klauspost/compress v1.17.4
  github.com/mitchellh/mapstructure v1.5.0
  github.com/modern-go/reflect2 v1.0.2
  github.com/stretchr/testify v1.7.1

go.sum

@@ -8,6 +8,8 @@ github.com/golang/snappy v0.0.4/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEW
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
 github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
+github.com/klauspost/compress v1.17.4 h1:Ej5ixsIri7BrIjBkRZLTo6ghwrEtHFk7ijlczPW4fZ4=
+github.com/klauspost/compress v1.17.4/go.mod h1:/dCuZOvVtNoHsyb+cuJD3itjs3NbnF6KH9zAO4BDxPM=
 github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY=
 github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo=
 github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=

ocf/codec.go

@@ -10,16 +10,18 @@ import (
  "io"
 
  "github.com/golang/snappy"
+ "github.com/klauspost/compress/zstd"
 )
 
 // CodecName represents a compression codec name.
 type CodecName string
 
 // Supported compression codecs.
 const (
- Null CodecName = "null"
- Deflate CodecName = "deflate"
- Snappy CodecName = "snappy"
+ Null CodecName = "null"
+ Deflate CodecName = "deflate"
+ Snappy CodecName = "snappy"
+ ZStandard CodecName = "zstandard"
 )
 
 func resolveCodec(name CodecName, lvl int) (Codec, error) {
@@ -33,6 +35,9 @@ func resolveCodec(name CodecName, lvl int) (Codec, error) {
  case Snappy:
  return &SnappyCodec{}, nil
 
+ case ZStandard:
+ return &ZStandardCodec{}, nil
+
  default:
  return nil, fmt.Errorf("unknown codec %s", name)
  }
@@ -120,3 +125,22 @@ func (*SnappyCodec) Encode(b []byte) []byte {
 
  return dst
 }
+
+// ZStandardCodec is a zstandard compression codec.
+type ZStandardCodec struct{}
+
+// Decode decodes the given bytes.
+func (*ZStandardCodec) Decode(b []byte) ([]byte, error) {
+ dec, _ := zstd.NewReader(nil)
+ defer dec.Close()
+
+ return dec.DecodeAll(b, nil)
+}
+
+// Encode encodes the given bytes.
+func (*ZStandardCodec) Encode(b []byte) []byte {
+ enc, _ := zstd.NewWriter(nil)
+ defer func() { _ = enc.Close() }()
+
+ return enc.EncodeAll(b, nil)
+}

ocf/ocf_test.go

@@ -311,6 +311,65 @@ func TestDecoder_WithSnappyHandlesInvalidCRC(t *testing.T) {
  assert.Error(t, dec.Error())
 }
 
+func TestDecoder_WithZStandard(t *testing.T) {
+ unionStr := "union value"
+ want := FullRecord{
+ Strings: []string{"string1", "string2", "string3", "string4", "string5"},
+ Longs: []int64{1, 2, 3, 4, 5},
+ Enum: "C",
+ Map: map[string]int{
+ "key1": 1,
+ "key2": 2,
+ "key3": 3,
+ "key4": 4,
+ "key5": 5,
+ },
+ Nullable: &unionStr,
+ Fixed: [16]byte{0x01, 0x02, 0x03, 0x04, 0x01, 0x02, 0x03, 0x04, 0x01, 0x02, 0x03, 0x04, 0x01, 0x02, 0x03, 0x04},
+ Record: &TestRecord{
+ Long: 1925639126735,
+ String: "I am a test record",
+ Int: 666,
+ Float: 7171.17,
+ Double: 916734926348163.01973408746523,
+ Bool: true,
+ },
+ }
+
+ f, err := os.Open("testdata/full-zstd.avro")
+ require.NoError(t, err)
+ t.Cleanup(func() { _ = f.Close() })
+
+ dec, err := ocf.NewDecoder(f)
+ require.NoError(t, err)
+
+ var count int
+ for dec.HasNext() {
+ count++
+ var got FullRecord
+ err = dec.Decode(&got)
+
+ require.NoError(t, err)
+ assert.Equal(t, want, got)
+ }
+
+ require.NoError(t, dec.Error())
+ assert.Equal(t, 1, count)
+}
+
+func TestDecoder_WithZStandardHandlesInvalidData(t *testing.T) {
+ f, err := os.Open("testdata/zstd-invalid-data.avro")
+ require.NoError(t, err)
+ t.Cleanup(func() { _ = f.Close() })
+
+ dec, err := ocf.NewDecoder(f)
+ require.NoError(t, err)
+
+ dec.HasNext()
+
+ assert.Error(t, dec.Error())
+}
+
 func TestDecoder_DecodeAvroError(t *testing.T) {
  data := []byte{'O', 'b', 'j', 0x01, 0x01, 0x26, 0x16, 'a', 'v', 'r', 'o', '.', 's', 'c', 'h', 'e', 'm', 'a',
  0x0c, '"', 'l', 'o', 'n', 'g', '"', 0x00, 0xfb, 0x2b, 0x0f, 0x1a, 0xdd, 0xfd, 0x90, 0x7d, 0x87, 0x12,
@@ -681,6 +740,43 @@ func TestEncoder_EncodeCompressesSnappy(t *testing.T) {
  assert.Equal(t, 938, buf.Len())
 }
 
+func TestEncoder_EncodeCompressesZStandard(t *testing.T) {
+ unionStr := "union value"
+ record := FullRecord{
+ Strings: []string{"string1", "string2", "string3", "string4", "string5"},
+ Longs: []int64{1, 2, 3, 4, 5},
+ Enum: "C",
+ Map: map[string]int{
+ "key1": 1,
+ "key2": 2,
+ "key3": 3,
+ "key4": 4,
+ "key5": 5,
+ },
+ Nullable: &unionStr,
+ Fixed: [16]byte{0x01, 0x02, 0x03, 0x04, 0x01, 0x02, 0x03, 0x04, 0x01, 0x02, 0x03, 0x04, 0x01, 0x02, 0x03, 0x04},
+ Record: &TestRecord{
+ Long: 1925639126735,
+ String: "I am a test record",
+ Int: 666,
+ Float: 7171.17,
+ Double: 916734926348163.01973408746523,
+ Bool: true,
+ },
+ }
+
+ buf := &bytes.Buffer{}
+ enc, _ := ocf.NewEncoder(schema, buf, ocf.WithCodec(ocf.ZStandard))
+
+ err := enc.Encode(record)
+ assert.NoError(t, err)
+
+ err = enc.Close()
+
+ require.NoError(t, err)
+ assert.Equal(t, 951, buf.Len())
+}
+
 func TestEncoder_EncodeError(t *testing.T) {
  buf := &bytes.Buffer{}
  enc, err := ocf.NewEncoder(`"long"`, buf)