This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git
The following commit(s) were added to refs/heads/main by this push:
new 0a28b231 fix(parquet): preserve dictionary index types (#1098)
0a28b231 is described below
commit 0a28b2316a2b456bdbb5316ea78e9169fdb70967
Author: Minh Vu <[email protected]>
AuthorDate: Fri Aug 7 18:40:43 2026 +0200
fix(parquet): preserve dictionary index types (#1098)
### Rationale for this change
Parquet dictionary encoding uses int32 indexes, but stored Arrow schema
metadata can specify a different dictionary index type. The schema
restoration path currently hardcodes int32 and changes the dictionary
type and index buffer width on read.
### What changes are included in this PR?
Restore the stored index type and safely cast the reader indexes before
rebuilding each dictionary array. Preserve dictionary values and the
ordered flag.
### Are these changes tested?
- `go test ./parquet/pqarrow -run
TestArrowDictionaryTypePreservesIndexType -count=1`
### Are there any user-facing changes?
No API changes. This corrects the reported behavior while preserving the
existing ownership and compatibility contracts.
---
parquet/pqarrow/column_readers.go | 47 +++++++++++++++++++-----
parquet/pqarrow/dictionary_multipage_test.go | 55 ++++++++++++++++++++++++++++
parquet/pqarrow/schema.go | 2 +-
3 files changed, 93 insertions(+), 11 deletions(-)
diff --git a/parquet/pqarrow/column_readers.go
b/parquet/pqarrow/column_readers.go
index 5c216946..9064d018 100644
--- a/parquet/pqarrow/column_readers.go
+++ b/parquet/pqarrow/column_readers.go
@@ -17,6 +17,7 @@
package pqarrow
import (
+ "context"
"encoding/binary"
"errors"
"fmt"
@@ -29,6 +30,7 @@ import (
"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/arrow-go/v18/arrow/array"
"github.com/apache/arrow-go/v18/arrow/bitutil"
+ "github.com/apache/arrow-go/v18/arrow/compute"
"github.com/apache/arrow-go/v18/arrow/decimal128"
"github.com/apache/arrow-go/v18/arrow/decimal256"
"github.com/apache/arrow-go/v18/arrow/memory"
@@ -120,7 +122,7 @@ func (lr *leafReader) LoadBatch(nrecords int64) (err error)
{
}
}
}
- lr.out, err = transferColumnData(lr.recordRdr, lr.field.Type, lr.descr)
+ lr.out, err = transferColumnData(lr.recordRdr, lr.field.Type, lr.descr,
lr.rctx.mem)
return
}
@@ -587,7 +589,7 @@ func chunksToSingle(chunked *arrow.Chunked, mem
memory.Allocator) (arrow.ArrayDa
}
// create a chunked arrow array from the raw record data
-func transferColumnData(rdr file.RecordReader, valueType arrow.DataType, descr
*schema.Column) (*arrow.Chunked, error) {
+func transferColumnData(rdr file.RecordReader, valueType arrow.DataType, descr
*schema.Column, mem memory.Allocator) (*arrow.Chunked, error) {
dt := valueType
if valueType.ID() == arrow.EXTENSION {
dt = valueType.(arrow.ExtensionType).StorageType()
@@ -596,7 +598,7 @@ func transferColumnData(rdr file.RecordReader, valueType
arrow.DataType, descr *
var data arrow.ArrayData
switch dt.ID() {
case arrow.DICTIONARY:
- return transferDictionary(rdr, valueType), nil
+ return transferDictionary(rdr, valueType, mem)
case arrow.NULL:
return arrow.NewChunked(arrow.Null,
[]arrow.Array{array.NewNull(rdr.ValuesWritten())}), nil
case arrow.INT32, arrow.INT64, arrow.FLOAT32, arrow.FLOAT64:
@@ -616,7 +618,7 @@ func transferColumnData(rdr file.RecordReader, valueType
arrow.DataType, descr *
case arrow.DATE64:
data = transferDate64(rdr, valueType)
case arrow.FIXED_SIZE_BINARY, arrow.BINARY, arrow.STRING,
arrow.LARGE_BINARY, arrow.LARGE_STRING:
- return transferBinary(rdr, valueType), nil
+ return transferBinary(rdr, valueType, mem)
case arrow.DECIMAL, arrow.DECIMAL256:
switch descr.PhysicalType() {
case parquet.Types.Int32, parquet.Types.Int64:
@@ -647,7 +649,7 @@ func transferColumnData(rdr file.RecordReader, valueType
arrow.DataType, descr *
if len := arrow.Float16SizeBytes; descr.TypeLength() != len {
return nil, fmt.Errorf("fixed len byte array length for
float16 must be %d", len)
}
- return transferBinary(rdr, valueType), nil
+ return transferBinary(rdr, valueType, mem)
default:
return nil, fmt.Errorf("no support for reading columns of type:
%s", valueType.Name())
}
@@ -675,10 +677,10 @@ func transferZeroCopy(rdr file.RecordReader, dt
arrow.DataType) arrow.ArrayData
nil, int(rdr.NullCount()), 0)
}
-func transferBinary(rdr file.RecordReader, dt arrow.DataType) *arrow.Chunked {
+func transferBinary(rdr file.RecordReader, dt arrow.DataType, mem
memory.Allocator) (*arrow.Chunked, error) {
brdr := rdr.(file.BinaryRecordReader)
if brdr.ReadDictionary() {
- return transferDictionary(brdr,
&arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: dt})
+ return transferDictionary(brdr,
&arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int32, ValueType: dt},
mem)
}
chunks := brdr.GetBuilderChunks()
defer releaseArrays(chunks)
@@ -703,7 +705,7 @@ func transferBinary(rdr file.RecordReader, dt
arrow.DataType) *arrow.Chunked {
chunk.Release()
}
}
- return arrow.NewChunked(dt, chunks)
+ return arrow.NewChunked(dt, chunks), nil
}
func transferInt(rdr file.RecordReader, dt arrow.DataType) arrow.ArrayData {
@@ -1080,9 +1082,34 @@ func transferDecimalBytes(rdr file.BinaryRecordReader,
dt arrow.DataType) (*arro
return arrow.NewChunked(dt, chunks), nil
}
-func transferDictionary(rdr file.RecordReader, logicalValueType
arrow.DataType) *arrow.Chunked {
+func transferDictionary(rdr file.RecordReader, logicalValueType
arrow.DataType, mem memory.Allocator) (*arrow.Chunked, error) {
brdr := rdr.(file.BinaryRecordReader)
chunks := brdr.GetBuilderChunks()
defer releaseArrays(chunks)
- return arrow.NewChunked(logicalValueType, chunks)
+
+ dictType, ok := logicalValueType.(*arrow.DictionaryType)
+ if !ok || dictType.IndexType.ID() == arrow.INT32 {
+ return arrow.NewChunked(logicalValueType, chunks), nil
+ }
+
+ ctx := compute.WithAllocator(context.Background(), mem)
+ for idx, chunk := range chunks {
+ dictArr, ok := chunk.(*array.Dictionary)
+ if !ok {
+ return nil, fmt.Errorf("expected dictionary array, got
%T", chunk)
+ }
+
+ indices, err := compute.CastArray(ctx, dictArr.Indices(),
compute.SafeCastOptions(dictType.IndexType))
+ if err != nil {
+ return nil, err
+ }
+ converted, err := array.NewValidatedDictionaryArray(dictType,
indices, dictArr.Dictionary())
+ indices.Release()
+ if err != nil {
+ return nil, err
+ }
+ chunk.Release()
+ chunks[idx] = converted
+ }
+ return arrow.NewChunked(logicalValueType, chunks), nil
}
diff --git a/parquet/pqarrow/dictionary_multipage_test.go
b/parquet/pqarrow/dictionary_multipage_test.go
index 8acd72f2..b6e9db5c 100644
--- a/parquet/pqarrow/dictionary_multipage_test.go
+++ b/parquet/pqarrow/dictionary_multipage_test.go
@@ -145,3 +145,58 @@ func TestArrowDictionaryTypeMultiplePages(t *testing.T) {
require.Equal(t, int64(numRows), totalRows, "Should read all rows")
t.Logf("Successfully read %d rows", totalRows)
}
+
+func TestArrowDictionaryTypePreservesIndexType(t *testing.T) {
+ for _, indexType := range []arrow.DataType{
+ arrow.PrimitiveTypes.Int8,
+ arrow.PrimitiveTypes.Int16,
+ arrow.PrimitiveTypes.Int32,
+ arrow.PrimitiveTypes.Int64,
+ arrow.PrimitiveTypes.Uint8,
+ arrow.PrimitiveTypes.Uint16,
+ arrow.PrimitiveTypes.Uint32,
+ arrow.PrimitiveTypes.Uint64,
+ } {
+ t.Run(indexType.Name(), func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ dictType := &arrow.DictionaryType{IndexType: indexType,
ValueType: arrow.BinaryTypes.String}
+ sc := arrow.NewSchema([]arrow.Field{{Name: "dict_col",
Type: dictType}}, nil)
+ builder := array.NewDictionaryBuilder(mem,
dictType).(*array.BinaryDictionaryBuilder)
+ defer builder.Release()
+ for _, value := range []string{"a", "b", "a", "c"} {
+ require.NoError(t, builder.AppendString(value))
+ }
+ values := builder.NewDictionaryArray()
+ defer values.Release()
+
+ var buf bytes.Buffer
+ writer, err := pqarrow.NewFileWriter(sc, &buf,
+
parquet.NewWriterProperties(parquet.WithAllocator(mem)),
+
pqarrow.NewArrowWriterProperties(pqarrow.WithStoreSchema(),
pqarrow.WithAllocator(mem)))
+ require.NoError(t, err)
+ rec := array.NewRecordBatch(sc, []arrow.Array{values},
int64(values.Len()))
+ require.NoError(t, writer.Write(rec))
+ require.NoError(t, writer.Close())
+ rec.Release()
+
+ pf, err :=
file.NewParquetReader(bytes.NewReader(buf.Bytes()),
+
file.WithReadProps(parquet.NewReaderProperties(mem)))
+ require.NoError(t, err)
+ defer pf.Close()
+ reader, err := pqarrow.NewFileReader(pf,
pqarrow.ArrowReadProperties{}, mem)
+ require.NoError(t, err)
+
+ readSchema, err := reader.Schema()
+ require.NoError(t, err)
+ require.True(t, arrow.TypeEqual(dictType,
readSchema.Field(0).Type))
+
+ tbl, err := reader.ReadTable(context.Background())
+ require.NoError(t, err)
+ defer tbl.Release()
+ require.True(t, arrow.TypeEqual(dictType,
tbl.Column(0).DataType()))
+ require.True(t, array.Equal(values,
tbl.Column(0).Data().Chunk(0)))
+ })
+ }
+}
diff --git a/parquet/pqarrow/schema.go b/parquet/pqarrow/schema.go
index ef87e663..d2a97559 100644
--- a/parquet/pqarrow/schema.go
+++ b/parquet/pqarrow/schema.go
@@ -1240,7 +1240,7 @@ func applyOriginalStorageMetadata(origin arrow.Field,
inferred *SchemaField) (mo
// direct dictionary reads are only supported for a few
primitive types
// so no need to recurse on value types
dictOriginType := origin.Type.(*arrow.DictionaryType)
- inferred.Field.Type = &arrow.DictionaryType{IndexType:
arrow.PrimitiveTypes.Int32,
+ inferred.Field.Type = &arrow.DictionaryType{IndexType:
dictOriginType.IndexType,
ValueType: inferred.Field.Type, Ordered:
dictOriginType.Ordered}
modified = true
case arrow.DECIMAL256:
