zeroshade commented on code in PR #1376:
URL: https://github.com/apache/arrow-go/pull/1376#discussion_r4315343528
##########
arrow/ipc/writer_test.go:
##########
@@ -958,3 +957,230 @@ func TestRecordEncoderResetRestoresDepth(t *testing.T) {
require.NoError(t, enc.encode(&simplePayload, simpleRecord))
simplePayload.Release()
}
+
+func TestWithZstdLevel(t *testing.T) {
+ rec := benchmarkRecordBatch(4, 4096)
+ defer rec.Release()
+
+ // WithZstdLevel should be sufficient on its own — it enables ZSTD.
+ t.Run("NewWriter roundtrip with level", func(t *testing.T) {
+ var buf bytes.Buffer
+ w := NewWriter(&buf, WithSchema(rec.Schema()), WithZstdLevel(1))
+ require.NoError(t, w.Write(rec))
+ require.NoError(t, w.Close())
+
+ r, err := NewReader(bytes.NewReader(buf.Bytes()))
+ require.NoError(t, err)
+ defer r.Release()
+
+ require.True(t, r.Next())
+ assert.True(t, array.RecordEqual(rec, r.RecordBatch()),
"records differ")
+ assert.False(t, r.Next())
+ assert.NoError(t, r.Err())
+ })
+
+ t.Run("NewFileWriter roundtrip with level", func(t *testing.T) {
+ var buf bytes.Buffer
+ fw, err := NewFileWriter(&buf, WithSchema(rec.Schema()),
WithZstdLevel(1))
+ require.NoError(t, err)
+ require.NoError(t, fw.Write(rec))
+ require.NoError(t, fw.Close())
+
+ r, err := NewFileReader(bytes.NewReader(buf.Bytes()))
+ require.NoError(t, err)
+ defer r.Close()
+
+ require.Equal(t, 1, r.NumRecords())
+ got, err := r.RecordBatchAt(0)
+ require.NoError(t, err)
+ defer got.Release()
+ assert.True(t, array.RecordEqual(rec, got), "records differ")
+ })
+
+ t.Run("different levels produce different output sizes", func(t
*testing.T) {
+ // Build a record with sequential integers so the four encoder
presets
+ // (SpeedFastest through SpeedBestCompression) yield visibly
different
+ // compressed sizes. All-zeros data is trivially compressible
at every
+ // level, but sequential values give zstd actual work to do.
+ const ncols, nrows = 8, 8192
Review Comment:
Nit: this sequential-int32 record builder is repeated verbatim at L1096 and
L1132; a small helper would trim it. Also, "non-positive level falls back to
default" is already covered by the `WithZstd`/`WithZstdLevel(0)` byte-equality
subtest, and the two `WithZstd` round-trip subtests don't exercise the new
option. Either could be dropped.
##########
arrow/ipc/ipc.go:
##########
@@ -165,6 +166,28 @@ func WithZstd() Option {
}
}
+// WithZstdLevel tells the writer to use ZSTD compression with a given
+// compression level and implicitly enables ZSTD compression (equivalent to
+// calling WithZstd). The specified value is passed to the underlying
+// klauspost/compress/zstd encoder via EncoderLevelFromZstd, which collapses
+// the full Zstandard level range (1-22) into four encoder speed presets:
+//
+// level <= 0 -> default (same as WithZstd alone)
+// <3 -> SpeedFastest
Review Comment:
Nit: `<3` overlaps the `level <= 0` row above it. Suggest:
```suggestion
// 1-2 -> SpeedFastest
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
