From 47f15b07080d62cd912bfbfd5d067cf70dfe6960 Mon Sep 17 00:00:00 2001
From: Yan Zhou <zhouyan@users.noreply.github.com>
Date: Wed, 21 Feb 2024 00:27:23 +0800
Subject: [PATCH] GH-40113 [Go][Parquet] New RegisterCodec function (#40114)

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>
---
 go/parquet/compress/brotli.go   |  2 +-
 go/parquet/compress/compress.go | 20 ++++++++++++++++++++
 go/parquet/compress/gzip.go     |  2 +-
 go/parquet/compress/snappy.go   |  2 +-
 go/parquet/compress/zstd.go     |  2 +-
 5 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/go/parquet/compress/brotli.go b/go/parquet/compress/brotli.go
index 8a7e92a1403c3..3b1575a70cfc8 100644
--- 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
index dc45b6ee9311f..2798defca9444 100644
--- 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
index 31f1729e9b3af..4b43f8e906599 100644
--- 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
index b7fa1142c3a6c..5c82a2c8dc33e 100644
--- 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
index 02ffd2eae568a..be3fb507262d4 100644
--- 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{})
 }