GH-40113 [Go][Parquet] New RegisterCodec function (#40114)

zhouyan · web-flow · commit 47f15b07080d · 2024-02-20T11:27:23.000-05:00
This is to allow addition/overwrite of custom codec implementation This allows other modules to provide alternative implementations for the compression algorithms, such as using libdeflate for Gzip, or CGO version of ZSTD. In addition, it allows others to supply codecs that cannot be easily supported by this library such as LZO due to license reasons or LZ4. ### Rationale for this change See #40113 ### What changes are included in this PR? A new RegisterCodec function added ### Are these changes tested? yes ### Are there any user-facing changes? It's an addition more targeted towards library writers. * Closes: #40113 Authored-by: Yan Zhou <zhouyan@me.com> Signed-off-by: Matt Topol <zotthewizard@gmail.com>
diff --git a/go/parquet/compress/brotli.go b/go/parquet/compress/brotli.go
@@ -110,5 +110,5 @@ func (brotliCodec) NewWriterLevel(w io.Writer, level int) (io.WriteCloser, error
 }
 
 func init() {
-	codecs[Codecs.Brotli] = brotliCodec{}
+	RegisterCodec(Codecs.Brotli, brotliCodec{})
 }
diff --git a/go/parquet/compress/compress.go b/go/parquet/compress/compress.go
@@ -92,6 +92,26 @@ type Codec interface {
 
 var codecs = map[Compression]Codec{}
 
+// RegisterCodec adds or overrides a codec implementation for a given compression algorithm.
+// The intended use case is within the init() section of a package. For example,
+//
+//	// inside a custom codec package, say czstd
+//
+//	func init() {
+//	    RegisterCodec(compress.Codecs.Zstd, czstdCodec{})
+//	}
+//
+//	type czstdCodec struct{} // implementing Codec interface using CGO based ZSTD wrapper
+//
+// And user of the custom codec can import the above package like below,
+//
+//	package main
+//
+//	import _ "package/path/to/czstd"
+func RegisterCodec(compression Compression, codec Codec) {
+	codecs[compression] = codec
+}
+
 type nocodec struct{}
 
 func (nocodec) NewReader(r io.Reader) io.ReadCloser {
diff --git a/go/parquet/compress/gzip.go b/go/parquet/compress/gzip.go
@@ -93,5 +93,5 @@ func (gzipCodec) NewWriterLevel(w io.Writer, level int) (io.WriteCloser, error)
 }
 
 func init() {
-	codecs[Codecs.Gzip] = gzipCodec{}
+	RegisterCodec(Codecs.Gzip, gzipCodec{})
 }
diff --git a/go/parquet/compress/snappy.go b/go/parquet/compress/snappy.go
@@ -57,5 +57,5 @@ func (s snappyCodec) NewWriterLevel(w io.Writer, _ int) (io.WriteCloser, error)
 }
 
 func init() {
-	codecs[Codecs.Snappy] = snappyCodec{}
+	RegisterCodec(Codecs.Snappy, snappyCodec{})
 }
diff --git a/go/parquet/compress/zstd.go b/go/parquet/compress/zstd.go
@@ -108,5 +108,5 @@ func (zstdCodec) CompressBound(len int64) int64 {
 }
 
 func init() {
-	codecs[Codecs.Zstd] = zstdCodec{}
+	RegisterCodec(Codecs.Zstd, zstdCodec{})
 }

Original file line number	Diff line number	Diff line change
`@@ -110,5 +110,5 @@ func (brotliCodec) NewWriterLevel(w io.Writer, level int) (io.WriteCloser, error`
`110`	`110`	`}`
`111`	`111`
`112`	`112`	`func init() {`
`113`		`- codecs[Codecs.Brotli] = brotliCodec{}`
	`113`	`+ RegisterCodec(Codecs.Brotli, brotliCodec{})`
`114`	`114`	`}`
Original file line number	Diff line number	Diff line change
`@@ -93,5 +93,5 @@ func (gzipCodec) NewWriterLevel(w io.Writer, level int) (io.WriteCloser, error)`
`93`	`93`	`}`
`94`	`94`
`95`	`95`	`func init() {`
`96`		`- codecs[Codecs.Gzip] = gzipCodec{}`
	`96`	`+ RegisterCodec(Codecs.Gzip, gzipCodec{})`
`97`	`97`	`}`
Original file line number	Diff line number	Diff line change
`@@ -57,5 +57,5 @@ func (s snappyCodec) NewWriterLevel(w io.Writer, _ int) (io.WriteCloser, error)`
`57`	`57`	`}`
`58`	`58`
`59`	`59`	`func init() {`
`60`		`- codecs[Codecs.Snappy] = snappyCodec{}`
	`60`	`+ RegisterCodec(Codecs.Snappy, snappyCodec{})`
`61`	`61`	`}`
Original file line number	Diff line number	Diff line change
`@@ -108,5 +108,5 @@ func (zstdCodec) CompressBound(len int64) int64 {`
`108`	`108`	`}`
`109`	`109`
`110`	`110`	`func init() {`
`111`		`- codecs[Codecs.Zstd] = zstdCodec{}`
	`111`	`+ RegisterCodec(Codecs.Zstd, zstdCodec{})`
`112`	`112`	`}`