GH-39309: [Go][Parquet] handle nil bitWriter for DeltaBinaryPacked (#39347)

zeroshade · web-flow · commit 0aadd5a0e7fd · 2024-01-08T15:28:15.000-05:00
### Rationale for this change If using the DeltaBinaryPacked encoding, we end up with a nil pointer dereference if we end up with an empty column. ### What changes are included in this PR? Add a nil check in `EstimatedDataEncodedSize` for the base `deltaBitPackEncoder`. This should only ever occur if we have an empty column with this encoding when closing a row group. ### Are these changes tested? Yes a unit test was added to verify the fix. * Closes: #39309 Authored-by: Matt Topol <zotthewizard@gmail.com> Signed-off-by: Matt Topol <zotthewizard@gmail.com>
diff --git a/go/parquet/internal/encoding/delta_bit_packing.go b/go/parquet/internal/encoding/delta_bit_packing.go
@@ -466,6 +466,10 @@ func (enc *deltaBitPackEncoder) FlushValues() (Buffer, error) {
 
 // EstimatedDataEncodedSize returns the current amount of data actually flushed out and written
 func (enc *deltaBitPackEncoder) EstimatedDataEncodedSize() int64 {
+	if enc.bitWriter == nil {
+		return 0
+	}
+
 	return int64(enc.bitWriter.Written())
 }
 
diff --git a/go/parquet/pqarrow/encode_arrow_test.go b/go/parquet/pqarrow/encode_arrow_test.go
@@ -1983,3 +1983,40 @@ func TestWriteTableMemoryAllocation(t *testing.T) {
 
 	require.Zero(t, mem.CurrentAlloc())
 }
+
+func TestEmptyListDeltaBinaryPacked(t *testing.T) {
+	schema := arrow.NewSchema([]arrow.Field{
+		{Name: "ts", Type: arrow.ListOf(arrow.PrimitiveTypes.Uint64),
+			Metadata: arrow.NewMetadata([]string{"PARQUET:field_id"}, []string{"-1"})}}, nil)
+	builder := array.NewRecordBuilder(memory.DefaultAllocator, schema)
+	defer builder.Release()
+
+	listBuilder := builder.Field(0).(*array.ListBuilder)
+	listBuilder.Append(true)
+	arrowRec := builder.NewRecord()
+	defer arrowRec.Release()
+
+	var buf bytes.Buffer
+	wr, err := pqarrow.NewFileWriter(schema, &buf,
+		parquet.NewWriterProperties(
+			parquet.WithDictionaryFor("ts.list.element", false),
+			parquet.WithEncodingFor("ts.list.element", parquet.Encodings.DeltaBinaryPacked)),
+		pqarrow.DefaultWriterProps())
+	require.NoError(t, err)
+
+	require.NoError(t, wr.WriteBuffered(arrowRec))
+	require.NoError(t, wr.Close())
+
+	rdr, err := file.NewParquetReader(bytes.NewReader(buf.Bytes()))
+	require.NoError(t, err)
+	reader, err := pqarrow.NewFileReader(rdr, pqarrow.ArrowReadProperties{}, memory.DefaultAllocator)
+	require.NoError(t, err)
+	defer rdr.Close()
+
+	tbl, err := reader.ReadTable(context.Background())
+	require.NoError(t, err)
+	defer tbl.Release()
+
+	assert.True(t, schema.Equal(tbl.Schema()))
+	assert.EqualValues(t, 1, tbl.NumRows())
+}

Original file line number	Diff line number	Diff line change
`@@ -466,6 +466,10 @@ func (enc *deltaBitPackEncoder) FlushValues() (Buffer, error) {`
`466`	`466`
`467`	`467`	`// EstimatedDataEncodedSize returns the current amount of data actually flushed out and written`
`468`	`468`	`func (enc *deltaBitPackEncoder) EstimatedDataEncodedSize() int64 {`
	`469`	`+ if enc.bitWriter == nil {`
	`470`	`+ return 0`
	`471`	`+ }`
	`472`	`+`
`469`	`473`	`return int64(enc.bitWriter.Written())`
`470`	`474`	`}`
`471`	`475`