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 5bc241e8 fix(parquet): build bloom filters from dictionary entries
(#1164)
5bc241e8 is described below
commit 5bc241e801ed76a1c5226bf479c84b794d2bcfd5
Author: Minh Vu <[email protected]>
AuthorDate: Wed Aug 26 21:30:01 2026 +0200
fix(parquet): build bloom filters from dictionary entries (#1164)
## Summary
- track dictionary entries referenced by encoded values
- populate Bloom filters from those entries once per column chunk
- cover direct Arrow dictionary arrays, multiple data pages, null
indices, and dictionary fallback
- keep the existing per-value path after fallback to plain encoding
## Why
Bloom filter updates currently happen in the ordinary value-writing
paths. Direct dictionary-array writes bypass those paths and write
indices directly, which leaves their Bloom filters empty. This can
produce false negatives when readers use the filter.
Ordinary dictionary encoding also hashes every logical value even though
inserting the same hash repeatedly does not change the filter. Reusing
the dictionary reduces hashing from the row count to the number of
referenced dictionary entries while dictionary encoding remains active.
## Benchmark
Apple M1 Pro, 100,000 int32 values, median of 10 runs with a 500 ms
benchmark time:
| Cardinality | Before | After | Speedup | B/op | allocs/op |
|---:|---:|---:|---:|---:|---:|
| 1 | 3.568 ms | 1.745 ms | 2.04x as fast | -75.0% | -51.8% |
| 10 | 3.902 ms | 2.136 ms | 1.83x as fast | -71.7% | -50.8% |
| 100 | 3.956 ms | 2.118 ms | 1.87x as fast | -69.5% | -49.0% |
Speedups are calculated directly from the displayed median times.
Higher-cardinality cases that trigger dictionary fallback remain close
to the existing path.
## Tests
- full repository test suite
- race tests for parquet/file, parquet/internal/encoding, and
parquet/pqarrow
- vet for the changed packages
---
parquet/file/column_writer.go | 86 +++++++++-
parquet/file/column_writer_bloom_internal_test.go | 117 ++++++++++++++
parquet/file/column_writer_test.go | 47 ++++++
parquet/file/column_writer_types.gen.go | 112 ++++++++-----
parquet/file/column_writer_types.gen.go.tmpl | 14 +-
parquet/file/file_writer_test.go | 81 ++++++++++
parquet/file/writer_performance_test.go | 80 +++++++++
parquet/internal/encoding/encoder.go | 189 ++++++++++++++++------
parquet/internal/encoding/typed_encoder.go | 7 +-
parquet/internal/encoding/typed_encoder_test.go | 98 +++++++++++
parquet/internal/encoding/types.go | 15 +-
parquet/metadata/bloom_filter.go | 66 ++++++++
parquet/metadata/bloom_filter_test.go | 61 +++++++
parquet/pqarrow/encode_dict_compute.go | 2 +-
parquet/pqarrow/encode_dictionary_test.go | 145 +++++++++++++++++
15 files changed, 1021 insertions(+), 99 deletions(-)
diff --git a/parquet/file/column_writer.go b/parquet/file/column_writer.go
index b19d1ce9..5da1663b 100644
--- a/parquet/file/column_writer.go
+++ b/parquet/file/column_writer.go
@@ -115,9 +115,10 @@ type columnWriter struct {
repEncoder encoding.LevelEncoder
mem memory.Allocator
- pageStatistics metadata.TypedStatistics
- chunkStatistics metadata.TypedStatistics
- bloomFilter metadata.BloomFilterBuilder
+ pageStatistics metadata.TypedStatistics
+ chunkStatistics metadata.TypedStatistics
+ bloomFilter metadata.BloomFilterBuilder
+ dictBloomPopulated bool
// total number of values stored in the current data page. this is the
maximum
// of the number of encoded def levels or encoded values. for
@@ -519,6 +520,9 @@ func (w *columnWriter) WriteDictionaryPage() error {
buffer.Resize(dictEncoder.DictEncodedSize())
dictEncoder.WriteDict(buffer.Bytes())
defer buffer.Release()
+ if err := w.populateBloomFilterFromDictionary(dictEncoder,
buffer.Bytes()); err != nil {
+ return err
+ }
page := NewDictionaryPage(buffer, int32(dictEncoder.NumEntries()),
w.props.DictionaryPageEncoding())
written, err := w.pager.WriteDictionaryPage(page)
@@ -527,6 +531,82 @@ func (w *columnWriter) WriteDictionaryPage() error {
return err
}
+func (w *columnWriter) populateBloomFilterFromEncoder(dictEncoder
encoding.DictEncoder) error {
+ if w.bloomFilter == nil || w.dictBloomPopulated {
+ return nil
+ }
+
+ buffer := memory.NewResizableBuffer(w.mem)
+ defer buffer.Release()
+ buffer.Resize(dictEncoder.DictEncodedSize())
+ dictEncoder.WriteDict(buffer.Bytes())
+ return w.populateBloomFilterFromDictionary(dictEncoder, buffer.Bytes())
+}
+
+// populateBloomFilterFromDictionary hashes the PLAIN-encoded dictionary
entries
+// that are referenced by data pages. Repeated logical values therefore
contribute
+// one hash for the entire column chunk.
+func (w *columnWriter) populateBloomFilterFromDictionary(dictEncoder
encoding.DictEncoder, dictionary []byte) error {
+ if w.bloomFilter == nil || w.dictBloomPopulated {
+ return nil
+ }
+
+ referenced := dictEncoder.ReferencedDictionaryIndices()
+ if len(referenced) == 0 {
+ w.dictBloomPopulated = true
+ return nil
+ }
+
+ hasher := w.bloomFilter.Hasher()
+ insertFixedWidth := func(width int) error {
+ for _, index := range referenced {
+ start := int(index) * width
+ end := start + width
+ if index < 0 || start < 0 || end > len(dictionary) {
+ return fmt.Errorf("parquet: referenced
dictionary index %d out of bounds", index)
+ }
+
w.bloomFilter.InsertHash(hasher.Sum64(dictionary[start:end]))
+ }
+ return nil
+ }
+
+ var err error
+ switch w.descr.PhysicalType() {
+ case parquet.Types.Int32, parquet.Types.Float:
+ err = insertFixedWidth(arrow.Int32SizeBytes)
+ case parquet.Types.Int64, parquet.Types.Double:
+ err = insertFixedWidth(arrow.Int64SizeBytes)
+ case parquet.Types.Int96:
+ err = insertFixedWidth(12)
+ case parquet.Types.FixedLenByteArray:
+ err = insertFixedWidth(w.descr.TypeLength())
+ case parquet.Types.ByteArray:
+ offset := 0
+ for index := 0; index < dictEncoder.NumEntries(); index++ {
+ if len(dictionary)-offset < arrow.Uint32SizeBytes {
+ return fmt.Errorf("parquet: truncated
byte-array dictionary length at index %d", index)
+ }
+ valueLen :=
int(binary.LittleEndian.Uint32(dictionary[offset:]))
+ offset += arrow.Uint32SizeBytes
+ if valueLen < 0 || valueLen > len(dictionary)-offset {
+ return fmt.Errorf("parquet: truncated
byte-array dictionary value at index %d", index)
+ }
+ if dictEncoder.DictionaryIndexReferenced(index) {
+
w.bloomFilter.InsertHash(hasher.Sum64(dictionary[offset : offset+valueLen]))
+ }
+ offset += valueLen
+ }
+ default:
+ err = fmt.Errorf("parquet: bloom filters are not supported for
dictionary type %s", w.descr.PhysicalType())
+ }
+ if err != nil {
+ return err
+ }
+
+ w.dictBloomPopulated = true
+ return nil
+}
+
type batchWriteInfo struct {
batchNum int64
nullCount int64
diff --git a/parquet/file/column_writer_bloom_internal_test.go
b/parquet/file/column_writer_bloom_internal_test.go
new file mode 100644
index 00000000..19335eb3
--- /dev/null
+++ b/parquet/file/column_writer_bloom_internal_test.go
@@ -0,0 +1,117 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package file
+
+import (
+ "math"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/apache/arrow-go/v18/parquet"
+ "github.com/apache/arrow-go/v18/parquet/internal/encoding"
+ "github.com/apache/arrow-go/v18/parquet/metadata"
+ "github.com/apache/arrow-go/v18/parquet/schema"
+ "github.com/cespare/xxhash/v2"
+ "github.com/stretchr/testify/require"
+)
+
+func TestDictionaryBloomHashesUsePlainEncoding(t *testing.T) {
+ tests := []struct {
+ name string
+ node *schema.PrimitiveNode
+ putValue func(encoding.TypedEncoder)
+ expected []byte
+ }{
+ {
+ name: "int32",
+ node: schema.NewInt32Node("value",
parquet.Repetitions.Required, -1),
+ putValue: func(enc encoding.TypedEncoder) {
+
enc.(encoding.Int32Encoder).Put([]int32{0x01020304})
+ },
+ expected: []byte{0x04, 0x03, 0x02, 0x01},
+ },
+ {
+ name: "int64",
+ node: schema.NewInt64Node("value",
parquet.Repetitions.Required, -1),
+ putValue: func(enc encoding.TypedEncoder) {
+
enc.(encoding.Int64Encoder).Put([]int64{0x0102030405060708})
+ },
+ expected: []byte{0x08, 0x07, 0x06, 0x05, 0x04, 0x03,
0x02, 0x01},
+ },
+ {
+ name: "float",
+ node: schema.NewFloat32Node("value",
parquet.Repetitions.Required, -1),
+ putValue: func(enc encoding.TypedEncoder) {
+
enc.(encoding.Float32Encoder).Put([]float32{math.Float32frombits(0x01020304)})
+ },
+ expected: []byte{0x04, 0x03, 0x02, 0x01},
+ },
+ {
+ name: "double",
+ node: schema.NewFloat64Node("value",
parquet.Repetitions.Required, -1),
+ putValue: func(enc encoding.TypedEncoder) {
+
enc.(encoding.Float64Encoder).Put([]float64{math.Float64frombits(0x0102030405060708)})
+ },
+ expected: []byte{0x08, 0x07, 0x06, 0x05, 0x04, 0x03,
0x02, 0x01},
+ },
+ {
+ name: "int96",
+ node: schema.NewInt96Node("value",
parquet.Repetitions.Required, -1),
+ putValue: func(enc encoding.TypedEncoder) {
+
enc.(encoding.Int96Encoder).Put([]parquet.Int96{{0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
10, 11}})
+ },
+ expected: []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11},
+ },
+ {
+ name: "byte array",
+ node: schema.NewByteArrayNode("value",
parquet.Repetitions.Required, -1),
+ putValue: func(enc encoding.TypedEncoder) {
+
enc.(encoding.ByteArrayEncoder).Put([]parquet.ByteArray{[]byte("plain")})
+ },
+ expected: []byte("plain"),
+ },
+ {
+ name: "fixed length byte array",
+ node: schema.NewFixedLenByteArrayNode("value",
parquet.Repetitions.Required, 5, -1),
+ putValue: func(enc encoding.TypedEncoder) {
+
enc.(encoding.FixedLenByteArrayEncoder).Put([]parquet.FixedLenByteArray{[]byte("fixed")})
+ },
+ expected: []byte("fixed"),
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ descr := schema.NewColumn(tt.node, 0, 0)
+ enc := encoding.NewEncoder(descr.PhysicalType(),
parquet.Encodings.PlainDict, true, descr, memory.DefaultAllocator)
+ defer enc.Release()
+ dictEnc := enc.(encoding.DictEncoder)
+ dictEnc.EnableDictionaryReferenceTracking()
+ tt.putValue(enc)
+
+ dictionary := make([]byte, dictEnc.DictEncodedSize())
+ dictEnc.WriteDict(dictionary)
+ builder := metadata.NewBloomFilter(32, 32,
memory.DefaultAllocator)
+ writer := columnWriter{descr: descr, bloomFilter:
builder}
+ require.NoError(t,
writer.populateBloomFilterFromDictionary(dictEnc, dictionary))
+
+ filter, ok := builder.(metadata.BloomFilter)
+ require.True(t, ok)
+ require.True(t,
filter.CheckHash(xxhash.Sum64(tt.expected)))
+ })
+ }
+}
diff --git a/parquet/file/column_writer_test.go
b/parquet/file/column_writer_test.go
index 7f66f79f..37491a52 100644
--- a/parquet/file/column_writer_test.go
+++ b/parquet/file/column_writer_test.go
@@ -43,6 +43,7 @@ import (
"github.com/apache/arrow-go/v18/parquet/schema"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
+ "github.com/stretchr/testify/require"
"github.com/stretchr/testify/suite"
)
@@ -170,6 +171,52 @@ func TestWriteDataPageV2NumRows(t *testing.T) {
assert.EqualValues(t, 3, wr.RowsWritten())
}
+func TestWriteDictIndicesRejectsBeforeWritingLevels(t *testing.T) {
+ root, err := schema.NewGroupNode("schema",
parquet.Repetitions.Required, schema.FieldList{
+ schema.NewInt32Node("v", parquet.Repetitions.Required, -1),
+ }, -1)
+ require.NoError(t, err)
+
+ props := parquet.NewWriterProperties(
+ parquet.WithDictionaryDefault(true),
+ parquet.WithStats(false),
+ )
+ var buf bytes.Buffer
+ w := file.NewParquetWriter(&buf, root, file.WithWriterProps(props))
+ rgw := w.AppendRowGroup()
+ cw, err := rgw.NextColumn()
+ require.NoError(t, err)
+ writer := cw.(*file.Int32ColumnChunkWriter)
+
+ dictionary := array.NewInt32Builder(memory.DefaultAllocator)
+ dictionary.AppendValues([]int32{10, 20}, nil)
+ dictValues := dictionary.NewArray()
+ dictionary.Release()
+ defer dictValues.Release()
+ indices := array.NewInt8Builder(memory.DefaultAllocator)
+ indices.Append(2)
+ invalidIndices := indices.NewArray()
+ indices.Release()
+ defer invalidIndices.Release()
+
+ dictEncoder := cw.CurrentEncoder().(encoding.DictEncoder)
+ require.NoError(t, dictEncoder.PutDictionary(dictValues))
+ require.ErrorIs(t, writer.WriteDictIndices(invalidIndices, nil, nil),
arrow.ErrInvalid)
+ assert.Equal(t, 0, writer.RowsWritten())
+
+ validBuilder := array.NewInt8Builder(memory.DefaultAllocator)
+ validBuilder.Append(0)
+ validIndices := validBuilder.NewArray()
+ validBuilder.Release()
+ defer validIndices.Release()
+ require.NoError(t, writer.WriteDictIndices(validIndices, nil, nil))
+ assert.Equal(t, 1, writer.RowsWritten())
+
+ require.NoError(t, cw.Close())
+ require.NoError(t, rgw.Close())
+ require.NoError(t, w.Close())
+}
+
func TestDataPageV2RowBoundaries(t *testing.T) {
sc := schema.NewSchema(schema.MustGroup(schema.NewGroupNode("schema",
parquet.Repetitions.Required, schema.FieldList{
schema.Must(schema.ListOf(
diff --git a/parquet/file/column_writer_types.gen.go
b/parquet/file/column_writer_types.gen.go
index bab348a3..fd5446df 100644
--- a/parquet/file/column_writer_types.gen.go
+++ b/parquet/file/column_writer_types.gen.go
@@ -45,6 +45,9 @@ type Int32ColumnChunkWriter struct {
func NewInt32ColumnChunkWriter(meta *metadata.ColumnChunkMetaDataBuilder,
pager PageWriter, useDict bool, enc parquet.Encoding, props
*parquet.WriterProperties) *Int32ColumnChunkWriter {
ret := &Int32ColumnChunkWriter{columnWriter: newColumnWriterBase(meta,
pager, useDict, enc, props)}
ret.currentEncoder =
encoding.Int32EncoderTraits.Encoder(format.Encoding(enc), useDict,
meta.Descr(), props.Allocator())
+ if useDict && ret.bloomFilter != nil {
+
ret.currentEncoder.(encoding.DictEncoder).EnableDictionaryReferenceTracking()
+ }
return ret
}
@@ -180,7 +183,6 @@ func (w *Int32ColumnChunkWriter) WriteDictIndices(indices
arrow.Array, defLevels
w.doBatches(int64(length), repLevels, func(offset, batch int64) {
info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels,
offset, batch), batch)
- w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset,
batch), levelSliceOrNil(repLevels, offset, batch))
writeableIndices := array.NewSlice(indices, valueOffset,
valueOffset+info.numSpaced())
defer writeableIndices.Release()
@@ -191,6 +193,8 @@ func (w *Int32ColumnChunkWriter) WriteDictIndices(indices
arrow.Array, defLevels
panic(err) // caught above
}
+ w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset,
batch), levelSliceOrNil(repLevels, offset, batch))
+
if err := w.commitWriteAndCheckPageLimit(batch, info.batchNum);
err != nil {
panic(err)
}
@@ -206,8 +210,7 @@ func (w *Int32ColumnChunkWriter) writeValues(values
[]int32, numNulls int64) {
if w.pageStatistics != nil {
w.pageStatistics.(*metadata.Int32Statistics).Update(values,
numNulls)
}
- if w.bloomFilter != nil {
- // TODO: optimize for Dictionary Encoding case
+ if w.bloomFilter != nil && w.currentEncoder.Encoding() !=
parquet.Encodings.PlainDict {
w.bloomFilter.InsertBulk(metadata.GetHashes(w.bloomFilter.Hasher(), values))
}
}
@@ -222,8 +225,7 @@ func (w *Int32ColumnChunkWriter)
writeValuesSpaced(spacedValues []int32, numRead
nulls := numValues - numRead
w.pageStatistics.(*metadata.Int32Statistics).UpdateSpaced(spacedValues,
validBits, validBitsOffset, nulls)
}
- if w.bloomFilter != nil {
- // TODO: optimize for Dictionary Encoding case
+ if w.bloomFilter != nil && w.currentEncoder.Encoding() !=
parquet.Encodings.PlainDict {
w.bloomFilter.InsertBulk(metadata.GetSpacedHashes(w.bloomFilter.Hasher(),
numRead, spacedValues, validBits, validBitsOffset))
}
}
@@ -285,6 +287,8 @@ func (w *Int32ColumnChunkWriter) FallbackToPlain() {
if err := w.drainBufferedDataPages(); err != nil {
panic(err)
}
+ } else if err := w.populateBloomFilterFromEncoder(dictEnc); err != nil {
+ panic(err)
}
if err := dictEnc.FallBackTo(plainEnc); err != nil {
@@ -312,6 +316,9 @@ type Int64ColumnChunkWriter struct {
func NewInt64ColumnChunkWriter(meta *metadata.ColumnChunkMetaDataBuilder,
pager PageWriter, useDict bool, enc parquet.Encoding, props
*parquet.WriterProperties) *Int64ColumnChunkWriter {
ret := &Int64ColumnChunkWriter{columnWriter: newColumnWriterBase(meta,
pager, useDict, enc, props)}
ret.currentEncoder =
encoding.Int64EncoderTraits.Encoder(format.Encoding(enc), useDict,
meta.Descr(), props.Allocator())
+ if useDict && ret.bloomFilter != nil {
+
ret.currentEncoder.(encoding.DictEncoder).EnableDictionaryReferenceTracking()
+ }
return ret
}
@@ -447,7 +454,6 @@ func (w *Int64ColumnChunkWriter) WriteDictIndices(indices
arrow.Array, defLevels
w.doBatches(int64(length), repLevels, func(offset, batch int64) {
info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels,
offset, batch), batch)
- w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset,
batch), levelSliceOrNil(repLevels, offset, batch))
writeableIndices := array.NewSlice(indices, valueOffset,
valueOffset+info.numSpaced())
defer writeableIndices.Release()
@@ -458,6 +464,8 @@ func (w *Int64ColumnChunkWriter) WriteDictIndices(indices
arrow.Array, defLevels
panic(err) // caught above
}
+ w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset,
batch), levelSliceOrNil(repLevels, offset, batch))
+
if err := w.commitWriteAndCheckPageLimit(batch, info.batchNum);
err != nil {
panic(err)
}
@@ -473,8 +481,7 @@ func (w *Int64ColumnChunkWriter) writeValues(values
[]int64, numNulls int64) {
if w.pageStatistics != nil {
w.pageStatistics.(*metadata.Int64Statistics).Update(values,
numNulls)
}
- if w.bloomFilter != nil {
- // TODO: optimize for Dictionary Encoding case
+ if w.bloomFilter != nil && w.currentEncoder.Encoding() !=
parquet.Encodings.PlainDict {
w.bloomFilter.InsertBulk(metadata.GetHashes(w.bloomFilter.Hasher(), values))
}
}
@@ -489,8 +496,7 @@ func (w *Int64ColumnChunkWriter)
writeValuesSpaced(spacedValues []int64, numRead
nulls := numValues - numRead
w.pageStatistics.(*metadata.Int64Statistics).UpdateSpaced(spacedValues,
validBits, validBitsOffset, nulls)
}
- if w.bloomFilter != nil {
- // TODO: optimize for Dictionary Encoding case
+ if w.bloomFilter != nil && w.currentEncoder.Encoding() !=
parquet.Encodings.PlainDict {
w.bloomFilter.InsertBulk(metadata.GetSpacedHashes(w.bloomFilter.Hasher(),
numRead, spacedValues, validBits, validBitsOffset))
}
}
@@ -552,6 +558,8 @@ func (w *Int64ColumnChunkWriter) FallbackToPlain() {
if err := w.drainBufferedDataPages(); err != nil {
panic(err)
}
+ } else if err := w.populateBloomFilterFromEncoder(dictEnc); err != nil {
+ panic(err)
}
if err := dictEnc.FallBackTo(plainEnc); err != nil {
@@ -579,6 +587,9 @@ type Int96ColumnChunkWriter struct {
func NewInt96ColumnChunkWriter(meta *metadata.ColumnChunkMetaDataBuilder,
pager PageWriter, useDict bool, enc parquet.Encoding, props
*parquet.WriterProperties) *Int96ColumnChunkWriter {
ret := &Int96ColumnChunkWriter{columnWriter: newColumnWriterBase(meta,
pager, useDict, enc, props)}
ret.currentEncoder =
encoding.Int96EncoderTraits.Encoder(format.Encoding(enc), useDict,
meta.Descr(), props.Allocator())
+ if useDict && ret.bloomFilter != nil {
+
ret.currentEncoder.(encoding.DictEncoder).EnableDictionaryReferenceTracking()
+ }
return ret
}
@@ -714,7 +725,6 @@ func (w *Int96ColumnChunkWriter) WriteDictIndices(indices
arrow.Array, defLevels
w.doBatches(int64(length), repLevels, func(offset, batch int64) {
info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels,
offset, batch), batch)
- w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset,
batch), levelSliceOrNil(repLevels, offset, batch))
writeableIndices := array.NewSlice(indices, valueOffset,
valueOffset+info.numSpaced())
defer writeableIndices.Release()
@@ -725,6 +735,8 @@ func (w *Int96ColumnChunkWriter) WriteDictIndices(indices
arrow.Array, defLevels
panic(err) // caught above
}
+ w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset,
batch), levelSliceOrNil(repLevels, offset, batch))
+
if err := w.commitWriteAndCheckPageLimit(batch, info.batchNum);
err != nil {
panic(err)
}
@@ -740,8 +752,7 @@ func (w *Int96ColumnChunkWriter) writeValues(values
[]parquet.Int96, numNulls in
if w.pageStatistics != nil {
w.pageStatistics.(*metadata.Int96Statistics).Update(values,
numNulls)
}
- if w.bloomFilter != nil {
- // TODO: optimize for Dictionary Encoding case
+ if w.bloomFilter != nil && w.currentEncoder.Encoding() !=
parquet.Encodings.PlainDict {
w.bloomFilter.InsertBulk(metadata.GetHashes(w.bloomFilter.Hasher(), values))
}
}
@@ -756,8 +767,7 @@ func (w *Int96ColumnChunkWriter)
writeValuesSpaced(spacedValues []parquet.Int96,
nulls := numValues - numRead
w.pageStatistics.(*metadata.Int96Statistics).UpdateSpaced(spacedValues,
validBits, validBitsOffset, nulls)
}
- if w.bloomFilter != nil {
- // TODO: optimize for Dictionary Encoding case
+ if w.bloomFilter != nil && w.currentEncoder.Encoding() !=
parquet.Encodings.PlainDict {
w.bloomFilter.InsertBulk(metadata.GetSpacedHashes(w.bloomFilter.Hasher(),
numRead, spacedValues, validBits, validBitsOffset))
}
}
@@ -819,6 +829,8 @@ func (w *Int96ColumnChunkWriter) FallbackToPlain() {
if err := w.drainBufferedDataPages(); err != nil {
panic(err)
}
+ } else if err := w.populateBloomFilterFromEncoder(dictEnc); err != nil {
+ panic(err)
}
if err := dictEnc.FallBackTo(plainEnc); err != nil {
@@ -846,6 +858,9 @@ type Float32ColumnChunkWriter struct {
func NewFloat32ColumnChunkWriter(meta *metadata.ColumnChunkMetaDataBuilder,
pager PageWriter, useDict bool, enc parquet.Encoding, props
*parquet.WriterProperties) *Float32ColumnChunkWriter {
ret := &Float32ColumnChunkWriter{columnWriter:
newColumnWriterBase(meta, pager, useDict, enc, props)}
ret.currentEncoder =
encoding.Float32EncoderTraits.Encoder(format.Encoding(enc), useDict,
meta.Descr(), props.Allocator())
+ if useDict && ret.bloomFilter != nil {
+
ret.currentEncoder.(encoding.DictEncoder).EnableDictionaryReferenceTracking()
+ }
return ret
}
@@ -981,7 +996,6 @@ func (w *Float32ColumnChunkWriter) WriteDictIndices(indices
arrow.Array, defLeve
w.doBatches(int64(length), repLevels, func(offset, batch int64) {
info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels,
offset, batch), batch)
- w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset,
batch), levelSliceOrNil(repLevels, offset, batch))
writeableIndices := array.NewSlice(indices, valueOffset,
valueOffset+info.numSpaced())
defer writeableIndices.Release()
@@ -992,6 +1006,8 @@ func (w *Float32ColumnChunkWriter)
WriteDictIndices(indices arrow.Array, defLeve
panic(err) // caught above
}
+ w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset,
batch), levelSliceOrNil(repLevels, offset, batch))
+
if err := w.commitWriteAndCheckPageLimit(batch, info.batchNum);
err != nil {
panic(err)
}
@@ -1007,8 +1023,7 @@ func (w *Float32ColumnChunkWriter) writeValues(values
[]float32, numNulls int64)
if w.pageStatistics != nil {
w.pageStatistics.(*metadata.Float32Statistics).Update(values,
numNulls)
}
- if w.bloomFilter != nil {
- // TODO: optimize for Dictionary Encoding case
+ if w.bloomFilter != nil && w.currentEncoder.Encoding() !=
parquet.Encodings.PlainDict {
w.bloomFilter.InsertBulk(metadata.GetHashes(w.bloomFilter.Hasher(), values))
}
}
@@ -1023,8 +1038,7 @@ func (w *Float32ColumnChunkWriter)
writeValuesSpaced(spacedValues []float32, num
nulls := numValues - numRead
w.pageStatistics.(*metadata.Float32Statistics).UpdateSpaced(spacedValues,
validBits, validBitsOffset, nulls)
}
- if w.bloomFilter != nil {
- // TODO: optimize for Dictionary Encoding case
+ if w.bloomFilter != nil && w.currentEncoder.Encoding() !=
parquet.Encodings.PlainDict {
w.bloomFilter.InsertBulk(metadata.GetSpacedHashes(w.bloomFilter.Hasher(),
numRead, spacedValues, validBits, validBitsOffset))
}
}
@@ -1086,6 +1100,8 @@ func (w *Float32ColumnChunkWriter) FallbackToPlain() {
if err := w.drainBufferedDataPages(); err != nil {
panic(err)
}
+ } else if err := w.populateBloomFilterFromEncoder(dictEnc); err != nil {
+ panic(err)
}
if err := dictEnc.FallBackTo(plainEnc); err != nil {
@@ -1113,6 +1129,9 @@ type Float64ColumnChunkWriter struct {
func NewFloat64ColumnChunkWriter(meta *metadata.ColumnChunkMetaDataBuilder,
pager PageWriter, useDict bool, enc parquet.Encoding, props
*parquet.WriterProperties) *Float64ColumnChunkWriter {
ret := &Float64ColumnChunkWriter{columnWriter:
newColumnWriterBase(meta, pager, useDict, enc, props)}
ret.currentEncoder =
encoding.Float64EncoderTraits.Encoder(format.Encoding(enc), useDict,
meta.Descr(), props.Allocator())
+ if useDict && ret.bloomFilter != nil {
+
ret.currentEncoder.(encoding.DictEncoder).EnableDictionaryReferenceTracking()
+ }
return ret
}
@@ -1248,7 +1267,6 @@ func (w *Float64ColumnChunkWriter)
WriteDictIndices(indices arrow.Array, defLeve
w.doBatches(int64(length), repLevels, func(offset, batch int64) {
info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels,
offset, batch), batch)
- w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset,
batch), levelSliceOrNil(repLevels, offset, batch))
writeableIndices := array.NewSlice(indices, valueOffset,
valueOffset+info.numSpaced())
defer writeableIndices.Release()
@@ -1259,6 +1277,8 @@ func (w *Float64ColumnChunkWriter)
WriteDictIndices(indices arrow.Array, defLeve
panic(err) // caught above
}
+ w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset,
batch), levelSliceOrNil(repLevels, offset, batch))
+
if err := w.commitWriteAndCheckPageLimit(batch, info.batchNum);
err != nil {
panic(err)
}
@@ -1274,8 +1294,7 @@ func (w *Float64ColumnChunkWriter) writeValues(values
[]float64, numNulls int64)
if w.pageStatistics != nil {
w.pageStatistics.(*metadata.Float64Statistics).Update(values,
numNulls)
}
- if w.bloomFilter != nil {
- // TODO: optimize for Dictionary Encoding case
+ if w.bloomFilter != nil && w.currentEncoder.Encoding() !=
parquet.Encodings.PlainDict {
w.bloomFilter.InsertBulk(metadata.GetHashes(w.bloomFilter.Hasher(), values))
}
}
@@ -1290,8 +1309,7 @@ func (w *Float64ColumnChunkWriter)
writeValuesSpaced(spacedValues []float64, num
nulls := numValues - numRead
w.pageStatistics.(*metadata.Float64Statistics).UpdateSpaced(spacedValues,
validBits, validBitsOffset, nulls)
}
- if w.bloomFilter != nil {
- // TODO: optimize for Dictionary Encoding case
+ if w.bloomFilter != nil && w.currentEncoder.Encoding() !=
parquet.Encodings.PlainDict {
w.bloomFilter.InsertBulk(metadata.GetSpacedHashes(w.bloomFilter.Hasher(),
numRead, spacedValues, validBits, validBitsOffset))
}
}
@@ -1353,6 +1371,8 @@ func (w *Float64ColumnChunkWriter) FallbackToPlain() {
if err := w.drainBufferedDataPages(); err != nil {
panic(err)
}
+ } else if err := w.populateBloomFilterFromEncoder(dictEnc); err != nil {
+ panic(err)
}
if err := dictEnc.FallBackTo(plainEnc); err != nil {
@@ -1383,6 +1403,9 @@ func NewBooleanColumnChunkWriter(meta
*metadata.ColumnChunkMetaDataBuilder, page
}
ret := &BooleanColumnChunkWriter{columnWriter:
newColumnWriterBase(meta, pager, useDict, enc, props)}
ret.currentEncoder =
encoding.BooleanEncoderTraits.Encoder(format.Encoding(enc), useDict,
meta.Descr(), props.Allocator())
+ if useDict && ret.bloomFilter != nil {
+
ret.currentEncoder.(encoding.DictEncoder).EnableDictionaryReferenceTracking()
+ }
return ret
}
@@ -1598,7 +1621,6 @@ func (w *BooleanColumnChunkWriter)
WriteDictIndices(indices arrow.Array, defLeve
w.doBatches(int64(length), repLevels, func(offset, batch int64) {
info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels,
offset, batch), batch)
- w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset,
batch), levelSliceOrNil(repLevels, offset, batch))
writeableIndices := array.NewSlice(indices, valueOffset,
valueOffset+info.numSpaced())
defer writeableIndices.Release()
@@ -1609,6 +1631,8 @@ func (w *BooleanColumnChunkWriter)
WriteDictIndices(indices arrow.Array, defLeve
panic(err) // caught above
}
+ w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset,
batch), levelSliceOrNil(repLevels, offset, batch))
+
if err := w.commitWriteAndCheckPageLimit(batch, info.batchNum);
err != nil {
panic(err)
}
@@ -1624,8 +1648,7 @@ func (w *BooleanColumnChunkWriter) writeValues(values
[]bool, numNulls int64) {
if w.pageStatistics != nil {
w.pageStatistics.(*metadata.BooleanStatistics).Update(values,
numNulls)
}
- if w.bloomFilter != nil {
- // TODO: optimize for Dictionary Encoding case
+ if w.bloomFilter != nil && w.currentEncoder.Encoding() !=
parquet.Encodings.PlainDict {
w.bloomFilter.InsertBulk(metadata.GetHashes(w.bloomFilter.Hasher(), values))
}
}
@@ -1640,8 +1663,7 @@ func (w *BooleanColumnChunkWriter)
writeValuesSpaced(spacedValues []bool, numRea
nulls := numValues - numRead
w.pageStatistics.(*metadata.BooleanStatistics).UpdateSpaced(spacedValues,
validBits, validBitsOffset, nulls)
}
- if w.bloomFilter != nil {
- // TODO: optimize for Dictionary Encoding case
+ if w.bloomFilter != nil && w.currentEncoder.Encoding() !=
parquet.Encodings.PlainDict {
w.bloomFilter.InsertBulk(metadata.GetSpacedHashes(w.bloomFilter.Hasher(),
numRead, spacedValues, validBits, validBitsOffset))
}
}
@@ -1766,6 +1788,8 @@ func (w *BooleanColumnChunkWriter) FallbackToPlain() {
if err := w.drainBufferedDataPages(); err != nil {
panic(err)
}
+ } else if err := w.populateBloomFilterFromEncoder(dictEnc); err != nil {
+ panic(err)
}
if err := dictEnc.FallBackTo(plainEnc); err != nil {
@@ -1793,6 +1817,9 @@ type ByteArrayColumnChunkWriter struct {
func NewByteArrayColumnChunkWriter(meta *metadata.ColumnChunkMetaDataBuilder,
pager PageWriter, useDict bool, enc parquet.Encoding, props
*parquet.WriterProperties) *ByteArrayColumnChunkWriter {
ret := &ByteArrayColumnChunkWriter{columnWriter:
newColumnWriterBase(meta, pager, useDict, enc, props)}
ret.currentEncoder =
encoding.ByteArrayEncoderTraits.Encoder(format.Encoding(enc), useDict,
meta.Descr(), props.Allocator())
+ if useDict && ret.bloomFilter != nil {
+
ret.currentEncoder.(encoding.DictEncoder).EnableDictionaryReferenceTracking()
+ }
return ret
}
@@ -2038,7 +2065,6 @@ func (w *ByteArrayColumnChunkWriter)
WriteDictIndices(indices arrow.Array, defLe
w.doBatches(int64(length), repLevels, func(offset, batch int64) {
info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels,
offset, batch), batch)
- w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset,
batch), levelSliceOrNil(repLevels, offset, batch))
writeableIndices := array.NewSlice(indices, valueOffset,
valueOffset+info.numSpaced())
defer writeableIndices.Release()
@@ -2049,6 +2075,8 @@ func (w *ByteArrayColumnChunkWriter)
WriteDictIndices(indices arrow.Array, defLe
panic(err) // caught above
}
+ w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset,
batch), levelSliceOrNil(repLevels, offset, batch))
+
if err := w.commitWriteAndCheckPageLimit(batch, info.batchNum);
err != nil {
panic(err)
}
@@ -2064,8 +2092,7 @@ func (w *ByteArrayColumnChunkWriter) writeValues(values
[]parquet.ByteArray, num
if w.pageStatistics != nil {
w.pageStatistics.(*metadata.ByteArrayStatistics).Update(values,
numNulls)
}
- if w.bloomFilter != nil {
- // TODO: optimize for Dictionary Encoding case
+ if w.bloomFilter != nil && w.currentEncoder.Encoding() !=
parquet.Encodings.PlainDict {
w.bloomFilter.InsertBulk(metadata.GetHashes(w.bloomFilter.Hasher(), values))
}
}
@@ -2080,8 +2107,7 @@ func (w *ByteArrayColumnChunkWriter)
writeValuesSpaced(spacedValues []parquet.By
nulls := numValues - numRead
w.pageStatistics.(*metadata.ByteArrayStatistics).UpdateSpaced(spacedValues,
validBits, validBitsOffset, nulls)
}
- if w.bloomFilter != nil {
- // TODO: optimize for Dictionary Encoding case
+ if w.bloomFilter != nil && w.currentEncoder.Encoding() !=
parquet.Encodings.PlainDict {
w.bloomFilter.InsertBulk(metadata.GetSpacedHashes(w.bloomFilter.Hasher(),
numRead, spacedValues, validBits, validBitsOffset))
}
}
@@ -2143,6 +2169,8 @@ func (w *ByteArrayColumnChunkWriter) FallbackToPlain() {
if err := w.drainBufferedDataPages(); err != nil {
panic(err)
}
+ } else if err := w.populateBloomFilterFromEncoder(dictEnc); err != nil {
+ panic(err)
}
if err := dictEnc.FallBackTo(plainEnc); err != nil {
@@ -2170,6 +2198,9 @@ type FixedLenByteArrayColumnChunkWriter struct {
func NewFixedLenByteArrayColumnChunkWriter(meta
*metadata.ColumnChunkMetaDataBuilder, pager PageWriter, useDict bool, enc
parquet.Encoding, props *parquet.WriterProperties)
*FixedLenByteArrayColumnChunkWriter {
ret := &FixedLenByteArrayColumnChunkWriter{columnWriter:
newColumnWriterBase(meta, pager, useDict, enc, props)}
ret.currentEncoder =
encoding.FixedLenByteArrayEncoderTraits.Encoder(format.Encoding(enc), useDict,
meta.Descr(), props.Allocator())
+ if useDict && ret.bloomFilter != nil {
+
ret.currentEncoder.(encoding.DictEncoder).EnableDictionaryReferenceTracking()
+ }
return ret
}
@@ -2415,7 +2446,6 @@ func (w *FixedLenByteArrayColumnChunkWriter)
WriteDictIndices(indices arrow.Arra
w.doBatches(int64(length), repLevels, func(offset, batch int64) {
info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels,
offset, batch), batch)
- w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset,
batch), levelSliceOrNil(repLevels, offset, batch))
writeableIndices := array.NewSlice(indices, valueOffset,
valueOffset+info.numSpaced())
defer writeableIndices.Release()
@@ -2426,6 +2456,8 @@ func (w *FixedLenByteArrayColumnChunkWriter)
WriteDictIndices(indices arrow.Arra
panic(err) // caught above
}
+ w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset,
batch), levelSliceOrNil(repLevels, offset, batch))
+
if err := w.commitWriteAndCheckPageLimit(batch, info.batchNum);
err != nil {
panic(err)
}
@@ -2445,8 +2477,7 @@ func (w *FixedLenByteArrayColumnChunkWriter)
writeValues(values []parquet.FixedL
w.pageStatistics.(*metadata.FixedLenByteArrayStatistics).Update(values,
numNulls)
}
}
- if w.bloomFilter != nil {
- // TODO: optimize for Dictionary Encoding case
+ if w.bloomFilter != nil && w.currentEncoder.Encoding() !=
parquet.Encodings.PlainDict {
w.bloomFilter.InsertBulk(metadata.GetHashes(w.bloomFilter.Hasher(), values))
}
}
@@ -2465,8 +2496,7 @@ func (w *FixedLenByteArrayColumnChunkWriter)
writeValuesSpaced(spacedValues []pa
w.pageStatistics.(*metadata.FixedLenByteArrayStatistics).UpdateSpaced(spacedValues,
validBits, validBitsOffset, nulls)
}
}
- if w.bloomFilter != nil {
- // TODO: optimize for Dictionary Encoding case
+ if w.bloomFilter != nil && w.currentEncoder.Encoding() !=
parquet.Encodings.PlainDict {
w.bloomFilter.InsertBulk(metadata.GetSpacedHashes(w.bloomFilter.Hasher(),
numRead, spacedValues, validBits, validBitsOffset))
}
}
@@ -2528,6 +2558,8 @@ func (w *FixedLenByteArrayColumnChunkWriter)
FallbackToPlain() {
if err := w.drainBufferedDataPages(); err != nil {
panic(err)
}
+ } else if err := w.populateBloomFilterFromEncoder(dictEnc); err != nil {
+ panic(err)
}
if err := dictEnc.FallBackTo(plainEnc); err != nil {
diff --git a/parquet/file/column_writer_types.gen.go.tmpl
b/parquet/file/column_writer_types.gen.go.tmpl
index 6988c468..ba898bec 100644
--- a/parquet/file/column_writer_types.gen.go.tmpl
+++ b/parquet/file/column_writer_types.gen.go.tmpl
@@ -49,6 +49,9 @@ func New{{.Name}}ColumnChunkWriter(meta
*metadata.ColumnChunkMetaDataBuilder, pa
{{- end}}
ret := &{{.Name}}ColumnChunkWriter{columnWriter: newColumnWriterBase(meta,
pager, useDict, enc, props)}
ret.currentEncoder =
encoding.{{.Name}}EncoderTraits.Encoder(format.Encoding(enc), useDict,
meta.Descr(), props.Allocator())
+ if useDict && ret.bloomFilter != nil {
+
ret.currentEncoder.(encoding.DictEncoder).EnableDictionaryReferenceTracking()
+ }
return ret
}
@@ -428,7 +431,6 @@ func (w *{{.Name}}ColumnChunkWriter)
WriteDictIndices(indices arrow.Array, defLe
w.doBatches(int64(length), repLevels, func(offset, batch int64) {
info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset,
batch), batch)
- w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset, batch),
levelSliceOrNil(repLevels, offset, batch))
writeableIndices := array.NewSlice(indices, valueOffset,
valueOffset+info.numSpaced())
defer writeableIndices.Release()
@@ -439,6 +441,8 @@ func (w *{{.Name}}ColumnChunkWriter)
WriteDictIndices(indices arrow.Array, defLe
panic(err) // caught above
}
+ w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset, batch),
levelSliceOrNil(repLevels, offset, batch))
+
if err := w.commitWriteAndCheckPageLimit(batch, info.batchNum); err != nil
{
panic(err)
}
@@ -462,8 +466,7 @@ func (w *{{.Name}}ColumnChunkWriter) writeValues(values
[]{{.name}}, numNulls in
}
{{- end}}
}
- if w.bloomFilter != nil {
- // TODO: optimize for Dictionary Encoding case
+ if w.bloomFilter != nil && w.currentEncoder.Encoding() !=
parquet.Encodings.PlainDict {
w.bloomFilter.InsertBulk(metadata.GetHashes(w.bloomFilter.Hasher(),
values))
}
}
@@ -486,8 +489,7 @@ func (w *{{.Name}}ColumnChunkWriter)
writeValuesSpaced(spacedValues []{{.name}},
}
{{- end}}
}
- if w.bloomFilter != nil {
- // TODO: optimize for Dictionary Encoding case
+ if w.bloomFilter != nil && w.currentEncoder.Encoding() !=
parquet.Encodings.PlainDict {
w.bloomFilter.InsertBulk(metadata.GetSpacedHashes(w.bloomFilter.Hasher(),
numRead, spacedValues, validBits, validBitsOffset))
}
}
@@ -614,6 +616,8 @@ func (w *{{.Name}}ColumnChunkWriter) FallbackToPlain() {
if err := w.drainBufferedDataPages(); err != nil {
panic(err)
}
+ } else if err := w.populateBloomFilterFromEncoder(dictEnc); err != nil {
+ panic(err)
}
if err := dictEnc.FallBackTo(plainEnc); err != nil {
diff --git a/parquet/file/file_writer_test.go b/parquet/file/file_writer_test.go
index 9e1abd08..2a28abde 100644
--- a/parquet/file/file_writer_test.go
+++ b/parquet/file/file_writer_test.go
@@ -1211,6 +1211,87 @@ func TestWriteBloomFilters(t *testing.T) {
assert.False(t, byteArrayFilter.Check(parquet.ByteArray("baz")))
}
+func TestWriteDictionaryBloomFilterAfterFallback(t *testing.T) {
+ values := []parquet.ByteArray{
+ parquet.ByteArray("alpha"),
+ parquet.ByteArray("beta"),
+ parquet.ByteArray("alpha"),
+ }
+ props := parquet.NewWriterProperties(
+ parquet.WithDictionaryDefault(true),
+ parquet.WithDictionaryPageSizeLimit(1),
+ parquet.WithBloomFilterEnabledFor("values", true),
+ parquet.WithBloomFilterNDVFor("values", 2),
+ )
+ field := schema.NewByteArrayNode("values",
parquet.Repetitions.Required, -1)
+ sc, err := schema.NewGroupNode("schema", parquet.Repetitions.Required,
schema.FieldList{field}, -1)
+ require.NoError(t, err)
+
+ sink := encoding.NewBufferWriter(0, memory.DefaultAllocator)
+ defer sink.Release()
+ writer := file.NewParquetWriter(sink, sc, file.WithWriterProps(props))
+ rowGroup := writer.AppendRowGroup()
+ column, err := rowGroup.NextColumn()
+ require.NoError(t, err)
+ _, err = column.(*file.ByteArrayColumnChunkWriter).WriteBatch(values,
nil, nil)
+ require.NoError(t, err)
+ require.NoError(t, column.Close())
+ require.NoError(t, rowGroup.Close())
+ require.NoError(t, writer.Close())
+
+ reader, err := file.NewParquetReader(bytes.NewReader(sink.Bytes()))
+ require.NoError(t, err)
+ defer reader.Close()
+ bloomReader := reader.GetBloomFilterReader()
+ rowGroupBloom, err := bloomReader.RowGroup(0)
+ require.NoError(t, err)
+ require.NoError(t, rowGroupBloom.VisitColumnBloomFilter(0, func(filter
metadata.BloomFilter) error {
+ typedFilter :=
metadata.TypedBloomFilter[parquet.ByteArray]{BloomFilter: filter}
+ assert.True(t, typedFilter.Check(parquet.ByteArray("alpha")))
+ assert.True(t, typedFilter.Check(parquet.ByteArray("beta")))
+ return nil
+ }))
+}
+
+func TestWriteDictionaryBloomFilterAcrossPages(t *testing.T) {
+ values := []int32{7, 42, 7, 99, 42, 7}
+ props := parquet.NewWriterProperties(
+ parquet.WithDictionaryDefault(true),
+ parquet.WithDataPageSize(1),
+ parquet.WithBatchSize(2),
+ parquet.WithBloomFilterEnabledFor("values", true),
+ parquet.WithBloomFilterNDVFor("values", 3),
+ )
+ field := schema.NewInt32Node("values", parquet.Repetitions.Required, -1)
+ sc, err := schema.NewGroupNode("schema", parquet.Repetitions.Required,
schema.FieldList{field}, -1)
+ require.NoError(t, err)
+
+ sink := encoding.NewBufferWriter(0, memory.DefaultAllocator)
+ defer sink.Release()
+ writer := file.NewParquetWriter(sink, sc, file.WithWriterProps(props))
+ rowGroup := writer.AppendRowGroup()
+ column, err := rowGroup.NextColumn()
+ require.NoError(t, err)
+ _, err = column.(*file.Int32ColumnChunkWriter).WriteBatch(values, nil,
nil)
+ require.NoError(t, err)
+ require.NoError(t, column.Close())
+ require.NoError(t, rowGroup.Close())
+ require.NoError(t, writer.Close())
+
+ reader, err := file.NewParquetReader(bytes.NewReader(sink.Bytes()))
+ require.NoError(t, err)
+ defer reader.Close()
+ rowGroupBloom, err := reader.GetBloomFilterReader().RowGroup(0)
+ require.NoError(t, err)
+ require.NoError(t, rowGroupBloom.VisitColumnBloomFilter(0, func(filter
metadata.BloomFilter) error {
+ typedFilter := metadata.TypedBloomFilter[int32]{BloomFilter:
filter}
+ assert.True(t, typedFilter.Check(7))
+ assert.True(t, typedFilter.Check(42))
+ assert.True(t, typedFilter.Check(99))
+ return nil
+ }))
+}
+
// TestBufferedStreamDictionaryCompressed tests the fix for issue #619
// where BufferedStreamEnabled=true with dictionary encoding and compression
// caused "dict spaced eof exception" and "snappy: corrupt input" errors.
diff --git a/parquet/file/writer_performance_test.go
b/parquet/file/writer_performance_test.go
index f3283ae6..0d823769 100644
--- a/parquet/file/writer_performance_test.go
+++ b/parquet/file/writer_performance_test.go
@@ -18,13 +18,93 @@ package file_test
import (
"bytes"
+ "fmt"
"testing"
+ "github.com/apache/arrow-go/v18/arrow"
"github.com/apache/arrow-go/v18/parquet"
"github.com/apache/arrow-go/v18/parquet/file"
"github.com/apache/arrow-go/v18/parquet/schema"
)
+func BenchmarkWriteDictionaryBloomFilter(b *testing.B) {
+ const numValues = 100_000
+
+ tests := []struct {
+ name string
+ cardinality int
+ }{
+ {name: "cardinality=1", cardinality: 1},
+ {name: "cardinality=10", cardinality: 10},
+ {name: "cardinality=100", cardinality: 100},
+ {name: "cardinality=1000", cardinality: 1_000},
+ {name: "cardinality=10000", cardinality: 10_000},
+ }
+
+ for _, tt := range tests {
+ b.Run("int32/"+tt.name, func(b *testing.B) {
+ values := make([]int32, numValues)
+ for i := range values {
+ values[i] = int32(i % tt.cardinality)
+ }
+ sc :=
schema.NewSchema(schema.MustGroup(schema.NewGroupNode("schema",
parquet.Repetitions.Required, schema.FieldList{
+ schema.NewInt32Node("values",
parquet.Repetitions.Required, -1),
+ }, -1)))
+ props := parquet.NewWriterProperties(
+ parquet.WithDictionaryDefault(true),
+ parquet.WithStats(false),
+ parquet.WithBloomFilterEnabledFor("values",
true),
+ parquet.WithBloomFilterNDVFor("values",
int64(tt.cardinality)),
+ )
+
+ b.ReportAllocs()
+ b.SetBytes(int64(numValues * arrow.Int32SizeBytes))
+ for range b.N {
+ var output bytes.Buffer
+ writer := file.NewParquetWriter(&output,
sc.Root(), file.WithWriterProps(props))
+ rowGroup := writer.AppendRowGroup()
+ column, _ := rowGroup.NextColumn()
+ _, _ =
column.(*file.Int32ColumnChunkWriter).WriteBatch(values, nil, nil)
+ _ = column.Close()
+ _ = rowGroup.Close()
+ _ = writer.Close()
+ }
+ })
+
+ b.Run("byte_array/"+tt.name, func(b *testing.B) {
+ dictionary := make([]parquet.ByteArray, tt.cardinality)
+ for i := range dictionary {
+ dictionary[i] =
parquet.ByteArray(fmt.Sprintf("value-%08d", i))
+ }
+ values := make([]parquet.ByteArray, numValues)
+ for i := range values {
+ values[i] = dictionary[i%tt.cardinality]
+ }
+ sc :=
schema.NewSchema(schema.MustGroup(schema.NewGroupNode("schema",
parquet.Repetitions.Required, schema.FieldList{
+ schema.NewByteArrayNode("values",
parquet.Repetitions.Required, -1),
+ }, -1)))
+ props := parquet.NewWriterProperties(
+ parquet.WithDictionaryDefault(true),
+ parquet.WithStats(false),
+ parquet.WithBloomFilterEnabledFor("values",
true),
+ parquet.WithBloomFilterNDVFor("values",
int64(tt.cardinality)),
+ )
+
+ b.ReportAllocs()
+ for range b.N {
+ var output bytes.Buffer
+ writer := file.NewParquetWriter(&output,
sc.Root(), file.WithWriterProps(props))
+ rowGroup := writer.AppendRowGroup()
+ column, _ := rowGroup.NextColumn()
+ _, _ =
column.(*file.ByteArrayColumnChunkWriter).WriteBatch(values, nil, nil)
+ _ = column.Close()
+ _ = rowGroup.Close()
+ _ = writer.Close()
+ }
+ })
+ }
+}
+
// Benchmark writing small ByteArray values (typical case)
// This tests the common scenario where values are small (< 1KB)
func BenchmarkWriteSmallByteArrayValues(b *testing.B) {
diff --git a/parquet/internal/encoding/encoder.go
b/parquet/internal/encoding/encoder.go
index e6561d69..58851b30 100644
--- a/parquet/internal/encoding/encoder.go
+++ b/parquet/internal/encoding/encoder.go
@@ -107,10 +107,13 @@ func (e *encoder) Reset() { e.sink.Reset(0) }
type dictEncoder struct {
encoder
- dictEncodedSize int
- idxBuffer *memory.Buffer
- idxValues []int32
- memo MemoTable
+ dictEncodedSize int
+ idxBuffer *memory.Buffer
+ idxValues []int32
+ memo MemoTable
+ trackReferences bool
+ referencedBitmap []byte
+ referencedIndices []int32
// rawDataSize is the number of bytes of input values observed since
// the last page flush. Mirrors parquet-mr's rawDataByteSize and is
@@ -140,6 +143,8 @@ func (d *dictEncoder) Reset() {
d.idxValues = d.idxValues[:0]
d.idxBuffer.ResizeNoShrink(0)
d.rawDataSize = 0
+ clear(d.referencedBitmap)
+ d.referencedIndices = d.referencedIndices[:0]
d.memo.Reset()
if d.preservedDict != nil {
d.preservedDict.Release()
@@ -147,6 +152,38 @@ func (d *dictEncoder) Reset() {
}
}
+func (d *dictEncoder) EnableDictionaryReferenceTracking() {
+ d.trackReferences = true
+}
+
+func (d *dictEncoder) ReferencedDictionaryIndices() []int32 {
+ return d.referencedIndices
+}
+
+func (d *dictEncoder) DictionaryIndexReferenced(index int) bool {
+ if index < 0 || index>>3 >= len(d.referencedBitmap) {
+ return false
+ }
+ return d.referencedBitmap[index>>3]&(1<<uint(index&7)) != 0
+}
+
+func (d *dictEncoder) recordDictionaryReference(index int32) {
+ if !d.trackReferences || index < 0 {
+ return
+ }
+
+ byteIndex := int(index) >> 3
+ if byteIndex >= len(d.referencedBitmap) {
+ d.referencedBitmap = append(d.referencedBitmap,
+ make([]byte, byteIndex-len(d.referencedBitmap)+1)...)
+ }
+ mask := byte(1 << uint(index&7))
+ if d.referencedBitmap[byteIndex]&mask == 0 {
+ d.referencedBitmap[byteIndex] |= mask
+ d.referencedIndices = append(d.referencedIndices, index)
+ }
+}
+
// ObservedRawSize returns the number of raw input bytes accumulated since
// the last data page flush. Used with DictEncodedSize and
// EstimatedDataEncodedSize to evaluate whether dictionary encoding is
@@ -182,56 +219,115 @@ func (d *dictEncoder) expandBuffer(newCap int) {
d.idxValues = arrow.Int32Traits.CastFromBytes(d.idxBuffer.Buf())[:
curLen : d.idxBuffer.Len()/arrow.Int32SizeBytes]
}
+type signedDictionaryIndex interface {
+ int8 | int16 | int32 | int64
+}
+
+type unsignedDictionaryIndex interface {
+ uint8 | uint16 | uint32 | uint64
+}
+
+const maxDictionaryIndex = uint64(1<<31 - 1)
+
+func (d *dictEncoder) invalidDictionaryIndex(index any) error {
+ return fmt.Errorf("%w: dictionary index %v out of bounds for dictionary
of length %d",
+ arrow.ErrInvalid, index, d.NumEntries())
+}
+
+func putSignedDictionaryIndices[T signedDictionaryIndex](d *dictEncoder, data
arrow.Array, values []T, start int) error {
+ dictSize := uint64(d.NumEntries())
+ curPos := start
+ return bitutils.VisitSetBitRuns(data.NullBitmapBytes(),
+ int64(data.Data().Offset()), int64(data.Len()),
+ func(pos, length int64) error {
+ for i := int64(0); i < length; i++ {
+ index := values[i+pos]
+ if index < 0 || uint64(index) >= dictSize ||
uint64(index) > maxDictionaryIndex {
+ return d.invalidDictionaryIndex(index)
+ }
+ d.idxValues[curPos] = int32(index)
+ d.recordDictionaryReference(int32(index))
+ curPos++
+ }
+ return nil
+ })
+}
+
+func putUnsignedDictionaryIndices[T unsignedDictionaryIndex](d *dictEncoder,
data arrow.Array, values []T, start int) error {
+ dictSize := uint64(d.NumEntries())
+ curPos := start
+ return bitutils.VisitSetBitRuns(data.NullBitmapBytes(),
+ int64(data.Data().Offset()), int64(data.Len()),
+ func(pos, length int64) error {
+ for i := int64(0); i < length; i++ {
+ index := values[i+pos]
+ if uint64(index) >= dictSize || uint64(index) >
maxDictionaryIndex {
+ return d.invalidDictionaryIndex(index)
+ }
+ d.idxValues[curPos] = int32(index)
+ d.recordDictionaryReference(int32(index))
+ curPos++
+ }
+ return nil
+ })
+}
+
+func (d *dictEncoder) rollbackDictionaryReferences(start, bitmapLen int) {
+ for _, index := range d.referencedIndices[start:] {
+ d.referencedBitmap[index>>3] &^= byte(1 << uint(index&7))
+ }
+ d.referencedIndices = d.referencedIndices[:start]
+ d.referencedBitmap = d.referencedBitmap[:bitmapLen]
+}
+
func (d *dictEncoder) PutIndices(data arrow.Array) error {
+ switch data.DataType().ID() {
+ case arrow.INT8, arrow.UINT8, arrow.INT16, arrow.UINT16,
+ arrow.INT32, arrow.UINT32, arrow.INT64, arrow.UINT64:
+ default:
+ return fmt.Errorf("%w: passed non-integer array to PutIndices",
arrow.ErrInvalid)
+ }
+
newValues := data.Len() - data.NullN()
curPos := len(d.idxValues)
newLen := newValues + curPos
d.expandBuffer(newLen)
d.idxValues = d.idxValues[:newLen:cap(d.idxValues)]
+ referenceStart := len(d.referencedIndices)
+ bitmapLen := len(d.referencedBitmap)
+ valueOffset := data.Data().Offset()
+ var err error
switch data.DataType().ID() {
- case arrow.UINT8, arrow.INT8:
- values :=
arrow.Uint8Traits.CastFromBytes(data.Data().Buffers()[1].Bytes())[data.Data().Offset():]
- bitutils.VisitSetBitRunsNoErr(data.NullBitmapBytes(),
- int64(data.Data().Offset()), int64(data.Len()),
- func(pos, length int64) {
- for i := int64(0); i < length; i++ {
- d.idxValues[curPos] =
int32(values[i+pos])
- curPos++
- }
- })
- case arrow.UINT16, arrow.INT16:
- values :=
arrow.Uint16Traits.CastFromBytes(data.Data().Buffers()[1].Bytes())[data.Data().Offset():]
- bitutils.VisitSetBitRunsNoErr(data.NullBitmapBytes(),
- int64(data.Data().Offset()), int64(data.Len()),
- func(pos, length int64) {
- for i := int64(0); i < length; i++ {
- d.idxValues[curPos] =
int32(values[i+pos])
- curPos++
- }
- })
- case arrow.UINT32, arrow.INT32:
- values :=
arrow.Uint32Traits.CastFromBytes(data.Data().Buffers()[1].Bytes())[data.Data().Offset():]
- bitutils.VisitSetBitRunsNoErr(data.NullBitmapBytes(),
- int64(data.Data().Offset()), int64(data.Len()),
- func(pos, length int64) {
- for i := int64(0); i < length; i++ {
- d.idxValues[curPos] =
int32(values[i+pos])
- curPos++
- }
- })
- case arrow.UINT64, arrow.INT64:
- values :=
arrow.Uint64Traits.CastFromBytes(data.Data().Buffers()[1].Bytes())[data.Data().Offset():]
- bitutils.VisitSetBitRunsNoErr(data.NullBitmapBytes(),
- int64(data.Data().Offset()), int64(data.Len()),
- func(pos, length int64) {
- for i := int64(0); i < length; i++ {
- d.idxValues[curPos] =
int32(values[i+pos])
- curPos++
- }
- })
- default:
- return fmt.Errorf("%w: passed non-integer array to PutIndices",
arrow.ErrInvalid)
+ case arrow.INT8:
+ values :=
arrow.Int8Traits.CastFromBytes(data.Data().Buffers()[1].Bytes())[valueOffset:]
+ err = putSignedDictionaryIndices(d, data, values, curPos)
+ case arrow.UINT8:
+ values :=
arrow.Uint8Traits.CastFromBytes(data.Data().Buffers()[1].Bytes())[valueOffset:]
+ err = putUnsignedDictionaryIndices(d, data, values, curPos)
+ case arrow.INT16:
+ values :=
arrow.Int16Traits.CastFromBytes(data.Data().Buffers()[1].Bytes())[valueOffset:]
+ err = putSignedDictionaryIndices(d, data, values, curPos)
+ case arrow.UINT16:
+ values :=
arrow.Uint16Traits.CastFromBytes(data.Data().Buffers()[1].Bytes())[valueOffset:]
+ err = putUnsignedDictionaryIndices(d, data, values, curPos)
+ case arrow.INT32:
+ values :=
arrow.Int32Traits.CastFromBytes(data.Data().Buffers()[1].Bytes())[valueOffset:]
+ err = putSignedDictionaryIndices(d, data, values, curPos)
+ case arrow.UINT32:
+ values :=
arrow.Uint32Traits.CastFromBytes(data.Data().Buffers()[1].Bytes())[valueOffset:]
+ err = putUnsignedDictionaryIndices(d, data, values, curPos)
+ case arrow.INT64:
+ values :=
arrow.Int64Traits.CastFromBytes(data.Data().Buffers()[1].Bytes())[valueOffset:]
+ err = putSignedDictionaryIndices(d, data, values, curPos)
+ case arrow.UINT64:
+ values :=
arrow.Uint64Traits.CastFromBytes(data.Data().Buffers()[1].Bytes())[valueOffset:]
+ err = putUnsignedDictionaryIndices(d, data, values, curPos)
+ }
+ if err != nil {
+ d.idxValues = d.idxValues[:curPos]
+ d.rollbackDictionaryReferences(referenceStart, bitmapLen)
+ return err
}
return nil
@@ -242,6 +338,7 @@ func (d *dictEncoder) addIndex(idx int) {
curLen := len(d.idxValues)
d.expandBuffer(curLen + 1)
d.idxValues = append(d.idxValues, int32(idx))
+ d.recordDictionaryReference(int32(idx))
}
// FlushValues dumps all the currently buffered indexes that would become the
data page to a buffer and
diff --git a/parquet/internal/encoding/typed_encoder.go
b/parquet/internal/encoding/typed_encoder.go
index 79b887a3..d547137a 100644
--- a/parquet/internal/encoding/typed_encoder.go
+++ b/parquet/internal/encoding/typed_encoder.go
@@ -168,6 +168,7 @@ func (enc *typedDictEncoder[T]) Put(in []T) {
enc.dictEncodedSize += int(unsafe.Sizeof(T(0)))
}
enc.idxValues[curPos+i] = int32(memoIdx)
+ enc.recordDictionaryReference(int32(memoIdx))
}
enc.AddRawSize(int64(len(in)) * int64(unsafe.Sizeof(T(0))))
}
@@ -458,7 +459,7 @@ func (enc *DictInt96Encoder) WriteDict(out []byte) {
// Put encodes the values passed in, adding to the index as needed
func (enc *DictInt96Encoder) Put(in []parquet.Int96) {
for _, v := range in {
- memoIdx, found, err := enc.memo.GetOrInsert(v)
+ memoIdx, found, err := enc.memo.GetOrInsert(v[:])
if err != nil {
panic(err)
}
@@ -487,6 +488,10 @@ func (enc *DictInt96Encoder) PutDictionary(arrow.Array)
error {
return fmt.Errorf("%w: direct PutDictionary to Int96",
arrow.ErrNotImplemented)
}
+func (enc *DictInt96Encoder) NormalizeDict(arrow.Array) (arrow.Array, error) {
+ return nil, fmt.Errorf("%w: direct PutDictionary to Int96",
arrow.ErrNotImplemented)
+}
+
// FallBackTo drains buffered indices through the dictionary into the
// fallback plain encoder and clears the index buffer.
func (enc *DictInt96Encoder) FallBackTo(fallback TypedEncoder) error {
diff --git a/parquet/internal/encoding/typed_encoder_test.go
b/parquet/internal/encoding/typed_encoder_test.go
index b9af99cb..d20b5d63 100644
--- a/parquet/internal/encoding/typed_encoder_test.go
+++ b/parquet/internal/encoding/typed_encoder_test.go
@@ -17,6 +17,8 @@
package encoding
import (
+ "math"
+ "strings"
"testing"
"unsafe"
@@ -26,6 +28,7 @@ import (
"github.com/apache/arrow-go/v18/parquet"
"github.com/apache/arrow-go/v18/parquet/schema"
"github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
)
func assertTypedDictEncoderPut[T int32 | int64 | float32 | float64](t
*testing.T, values, expectedDict []T, expectedIndices []int32) {
@@ -87,3 +90,98 @@ func TestPutDictionary(t *testing.T) {
err := enc.PutDictionary(arr)
assert.NoError(t, err)
}
+
+func TestDictionaryReferenceTracking(t *testing.T) {
+ dictionary, _, err := array.FromJSON(memory.DefaultAllocator,
arrow.PrimitiveTypes.Int32,
+ strings.NewReader(`[10, 20, 30]`))
+ require.NoError(t, err)
+ defer dictionary.Release()
+
+ indices, _, err := array.FromJSON(memory.DefaultAllocator,
arrow.PrimitiveTypes.Int32,
+ strings.NewReader(`[2, null, 2, 1]`))
+ require.NoError(t, err)
+ defer indices.Release()
+
+ typ := schema.NewInt32Node("a", parquet.Repetitions.Required, -1)
+ descr := schema.NewColumn(typ, 0, 0)
+ enc := &typedDictEncoder[int32]{newDictEncoderBase(descr,
NewDictionary[int32](), memory.DefaultAllocator)}
+ defer enc.Release()
+ enc.EnableDictionaryReferenceTracking()
+
+ require.NoError(t, enc.PutDictionary(dictionary))
+ require.NoError(t, enc.PutIndices(indices))
+ assert.Equal(t, []int32{2, 1}, enc.ReferencedDictionaryIndices())
+ assert.False(t, enc.DictionaryIndexReferenced(0))
+ assert.True(t, enc.DictionaryIndexReferenced(1))
+ assert.True(t, enc.DictionaryIndexReferenced(2))
+}
+
+func TestPutIndicesRejectsOutOfBoundsIndices(t *testing.T) {
+ dictionary, _, err := array.FromJSON(memory.DefaultAllocator,
arrow.PrimitiveTypes.Int32,
+ strings.NewReader(`[10]`))
+ require.NoError(t, err)
+ defer dictionary.Release()
+
+ tests := []struct {
+ name string
+ newIndices func() arrow.Array
+ }{
+ {
+ name: "negative signed index",
+ newIndices: func() arrow.Array {
+ builder :=
array.NewInt8Builder(memory.DefaultAllocator)
+ defer builder.Release()
+ builder.Append(-1)
+ return builder.NewArray()
+ },
+ },
+ {
+ name: "index equal to dictionary length",
+ newIndices: func() arrow.Array {
+ builder :=
array.NewInt16Builder(memory.DefaultAllocator)
+ defer builder.Release()
+ builder.Append(0)
+ builder.Append(1)
+ return builder.NewArray()
+ },
+ },
+ {
+ name: "maximum int32 index",
+ newIndices: func() arrow.Array {
+ builder :=
array.NewInt32Builder(memory.DefaultAllocator)
+ defer builder.Release()
+ builder.Append(math.MaxInt32)
+ return builder.NewArray()
+ },
+ },
+ {
+ name: "large unsigned index",
+ newIndices: func() arrow.Array {
+ builder :=
array.NewUint64Builder(memory.DefaultAllocator)
+ defer builder.Release()
+ builder.Append(math.MaxUint64)
+ return builder.NewArray()
+ },
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ indices := tt.newIndices()
+ defer indices.Release()
+
+ typ := schema.NewInt32Node("a",
parquet.Repetitions.Required, -1)
+ descr := schema.NewColumn(typ, 0, 0)
+ enc :=
&typedDictEncoder[int32]{newDictEncoderBase(descr, NewDictionary[int32](),
memory.DefaultAllocator)}
+ defer enc.Release()
+ enc.EnableDictionaryReferenceTracking()
+ require.NoError(t, enc.PutDictionary(dictionary))
+
+ err := enc.PutIndices(indices)
+ require.ErrorIs(t, err, arrow.ErrInvalid)
+ assert.Empty(t, enc.idxValues)
+ assert.Empty(t, enc.referencedBitmap)
+ assert.Empty(t, enc.ReferencedDictionaryIndices())
+ })
+ }
+}
diff --git a/parquet/internal/encoding/types.go
b/parquet/internal/encoding/types.go
index afbda0d6..9c940f6b 100644
--- a/parquet/internal/encoding/types.go
+++ b/parquet/internal/encoding/types.go
@@ -84,6 +84,16 @@ type TypedEncoder interface {
// encoding.
type DictEncoder interface {
TypedEncoder
+ // EnableDictionaryReferenceTracking records which dictionary entries
are
+ // referenced by encoded values. Tracking is disabled by default so
writers
+ // that do not need it do not pay its per-value cost.
+ EnableDictionaryReferenceTracking()
+ // ReferencedDictionaryIndices returns dictionary indices in the order
they
+ // were first referenced. The returned slice is owned by the encoder
and must
+ // not be modified.
+ ReferencedDictionaryIndices() []int32
+ // DictionaryIndexReferenced reports whether an entry has been
referenced.
+ DictionaryIndexReferenced(index int) bool
// WriteIndices populates the byte slice with the final indexes of data
and returns
// the number of bytes written
WriteIndices(out []byte) (int, error)
@@ -109,9 +119,8 @@ type DictEncoder interface {
// from PutDictionary or nil.
PreservedDictionary() arrow.Array
// PutIndices adds the indices from the passed in integral array to
- // the column data. It is assumed that the indices are within the bounds
- // of [0,dictSize) and is not validated. Returns an error if a
non-integral
- // array is passed.
+ // the column data. Returns an error if an index is outside [0,
dictSize)
+ // or a non-integral array is passed.
PutIndices(arrow.Array) error
// NormalizeDict takes an arrow array and normalizes it to a parquet
// native type. e.g. a dictionary of type int8 will be cast to an int32
diff --git a/parquet/metadata/bloom_filter.go b/parquet/metadata/bloom_filter.go
index 100786e1..16ded8b2 100644
--- a/parquet/metadata/bloom_filter.go
+++ b/parquet/metadata/bloom_filter.go
@@ -17,6 +17,7 @@
package metadata
import (
+ "encoding/binary"
"errors"
"fmt"
"io"
@@ -26,6 +27,7 @@ import (
"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/arrow-go/v18/arrow/bitutil"
+ "github.com/apache/arrow-go/v18/arrow/endian"
"github.com/apache/arrow-go/v18/arrow/memory"
"github.com/apache/arrow-go/v18/internal/bitutils"
"github.com/apache/arrow-go/v18/parquet"
@@ -194,6 +196,30 @@ func GetSpacedHashesFromBitmap(h Hasher, numValid int64,
bitmap []byte, bitmapOf
func getBytes[T parquet.ColumnTypes](v T) []byte {
switch v := any(v).(type) {
+ case int32:
+ if endian.IsBigEndian {
+ var out [arrow.Int32SizeBytes]byte
+ binary.LittleEndian.PutUint32(out[:], uint32(v))
+ return out[:]
+ }
+ case int64:
+ if endian.IsBigEndian {
+ var out [arrow.Int64SizeBytes]byte
+ binary.LittleEndian.PutUint64(out[:], uint64(v))
+ return out[:]
+ }
+ case float32:
+ if endian.IsBigEndian {
+ var out [arrow.Float32SizeBytes]byte
+ binary.LittleEndian.PutUint32(out[:],
math.Float32bits(v))
+ return out[:]
+ }
+ case float64:
+ if endian.IsBigEndian {
+ var out [arrow.Float64SizeBytes]byte
+ binary.LittleEndian.PutUint64(out[:],
math.Float64bits(v))
+ return out[:]
+ }
case parquet.ByteArray:
return v
case parquet.FixedLenByteArray:
@@ -208,6 +234,46 @@ func getBytes[T parquet.ColumnTypes](v T) []byte {
func getBytesSlice[T parquet.ColumnTypes](v []T) [][]byte {
b := make([][]byte, len(v))
switch v := any(v).(type) {
+ case []int32:
+ if endian.IsBigEndian {
+ raw := make([]byte, arrow.Int32SizeBytes*len(v))
+ for i, vv := range v {
+ value := raw[i*arrow.Int32SizeBytes :
(i+1)*arrow.Int32SizeBytes]
+ binary.LittleEndian.PutUint32(value, uint32(vv))
+ b[i] = value
+ }
+ return b
+ }
+ case []int64:
+ if endian.IsBigEndian {
+ raw := make([]byte, arrow.Int64SizeBytes*len(v))
+ for i, vv := range v {
+ value := raw[i*arrow.Int64SizeBytes :
(i+1)*arrow.Int64SizeBytes]
+ binary.LittleEndian.PutUint64(value, uint64(vv))
+ b[i] = value
+ }
+ return b
+ }
+ case []float32:
+ if endian.IsBigEndian {
+ raw := make([]byte, arrow.Float32SizeBytes*len(v))
+ for i, vv := range v {
+ value := raw[i*arrow.Float32SizeBytes :
(i+1)*arrow.Float32SizeBytes]
+ binary.LittleEndian.PutUint32(value,
math.Float32bits(vv))
+ b[i] = value
+ }
+ return b
+ }
+ case []float64:
+ if endian.IsBigEndian {
+ raw := make([]byte, arrow.Float64SizeBytes*len(v))
+ for i, vv := range v {
+ value := raw[i*arrow.Float64SizeBytes :
(i+1)*arrow.Float64SizeBytes]
+ binary.LittleEndian.PutUint64(value,
math.Float64bits(vv))
+ b[i] = value
+ }
+ return b
+ }
case []parquet.ByteArray:
for i, vv := range v {
b[i] = vv
diff --git a/parquet/metadata/bloom_filter_test.go
b/parquet/metadata/bloom_filter_test.go
index 4206ec56..720de535 100644
--- a/parquet/metadata/bloom_filter_test.go
+++ b/parquet/metadata/bloom_filter_test.go
@@ -17,6 +17,7 @@
package metadata
import (
+ "bytes"
"fmt"
"math"
"math/rand/v2"
@@ -29,9 +30,37 @@ import (
"github.com/apache/arrow-go/v18/arrow/memory"
"github.com/apache/arrow-go/v18/parquet"
"github.com/apache/arrow-go/v18/parquet/schema"
+ "github.com/cespare/xxhash/v2"
"github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
)
+type recordingHasher struct {
+ inputs [][]byte
+}
+
+func (h *recordingHasher) Sum64(value []byte) uint64 {
+ h.inputs = append(h.inputs, bytes.Clone(value))
+ return xxhash.Sum64(value)
+}
+
+func (h *recordingHasher) Sum64s(values [][]byte) []uint64 {
+ hashes := make([]uint64, len(values))
+ for i, value := range values {
+ hashes[i] = h.Sum64(value)
+ }
+ return hashes
+}
+
+type expectedHashBloomFilter struct {
+ hasher Hasher
+ expected uint64
+}
+
+func (b expectedHashBloomFilter) Hasher() Hasher { return b.hasher
}
+func (b expectedHashBloomFilter) CheckHash(hash uint64) bool { return hash ==
b.expected }
+func (expectedHashBloomFilter) Size() int64 { return 0 }
+
func TestSplitBlockFilter(t *testing.T) {
const N = 1000
const S = 3
@@ -114,6 +143,38 @@ func TestGetHashes(t *testing.T) {
testHash(t, h, valsI32)
}
+func assertPlainEncodedBloomHash[T parquet.ColumnTypes](t *testing.T, value T,
expected []byte) {
+ t.Helper()
+
+ expectedHash := xxhash.Sum64(expected)
+ hasher := &recordingHasher{}
+ assert.Equal(t, expectedHash, GetHash[T](hasher, value))
+ require.Equal(t, [][]byte{expected}, hasher.inputs)
+
+ hasher.inputs = nil
+ assert.Equal(t, []uint64{expectedHash}, GetHashes(hasher, []T{value}))
+ require.Equal(t, [][]byte{expected}, hasher.inputs)
+
+ hasher.inputs = nil
+ assert.Equal(t, []uint64{expectedHash}, GetSpacedHashes(hasher, 1,
[]T{value}, []byte{1}, 0))
+ require.Equal(t, [][]byte{expected}, hasher.inputs)
+
+ hasher.inputs = nil
+ filter := TypedBloomFilter[T]{BloomFilter:
expectedHashBloomFilter{hasher: hasher, expected: expectedHash}}
+ assert.True(t, filter.Check(value))
+ require.Equal(t, [][]byte{expected}, hasher.inputs)
+}
+
+func TestBloomHashesUsePlainEncoding(t *testing.T) {
+ assertPlainEncodedBloomHash(t, int32(0x01020304), []byte{0x04, 0x03,
0x02, 0x01})
+ assertPlainEncodedBloomHash(t, int64(0x0102030405060708), []byte{0x08,
0x07, 0x06, 0x05, 0x04, 0x03, 0x02, 0x01})
+ assertPlainEncodedBloomHash(t, math.Float32frombits(0x01020304),
[]byte{0x04, 0x03, 0x02, 0x01})
+ assertPlainEncodedBloomHash(t,
math.Float64frombits(0x0102030405060708), []byte{0x08, 0x07, 0x06, 0x05, 0x04,
0x03, 0x02, 0x01})
+ assertPlainEncodedBloomHash(t, parquet.Int96{0, 1, 2, 3, 4, 5, 6, 7, 8,
9, 10, 11}, []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11})
+ assertPlainEncodedBloomHash(t, parquet.ByteArray("plain"),
[]byte("plain"))
+ assertPlainEncodedBloomHash(t, parquet.FixedLenByteArray("fixed"),
[]byte("fixed"))
+}
+
func TestNewBloomFilter(t *testing.T) {
tests := []struct {
ndv uint32
diff --git a/parquet/pqarrow/encode_dict_compute.go
b/parquet/pqarrow/encode_dict_compute.go
index 184ec0fa..be020c1d 100644
--- a/parquet/pqarrow/encode_dict_compute.go
+++ b/parquet/pqarrow/encode_dict_compute.go
@@ -113,7 +113,7 @@ func writeDictionaryArrow(ctx *arrowWriteContext, cw
file.ColumnChunkWriter, lea
if referencedIndices.Len() == normalized.Len() {
referencedDict = normalized
} else {
- referencedDict, err =
compute.TakeArrayOpts(ctx, normalized, referencedIndices,
compute.TakeOptions{BoundsCheck: false})
+ referencedDict, err =
compute.TakeArrayOpts(ctx, normalized, referencedIndices,
compute.TakeOptions{BoundsCheck: true})
if err != nil {
return err
}
diff --git a/parquet/pqarrow/encode_dictionary_test.go
b/parquet/pqarrow/encode_dictionary_test.go
index b027f7f0..7f9a0a09 100644
--- a/parquet/pqarrow/encode_dictionary_test.go
+++ b/parquet/pqarrow/encode_dictionary_test.go
@@ -34,6 +34,7 @@ import (
"github.com/apache/arrow-go/v18/parquet"
"github.com/apache/arrow-go/v18/parquet/file"
"github.com/apache/arrow-go/v18/parquet/internal/testutils"
+ "github.com/apache/arrow-go/v18/parquet/metadata"
"github.com/apache/arrow-go/v18/parquet/pqarrow"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
@@ -81,6 +82,150 @@ func TestWriteColumnChunkedPropagatesLevelBuilderError(t
*testing.T) {
require.NoError(t, writer.Close())
}
+func TestDictionaryArrayBloomFilter(t *testing.T) {
+ mem := memory.DefaultAllocator
+
+ dictionary, _, err := array.FromJSON(mem, arrow.BinaryTypes.String,
+ strings.NewReader(`["unused", "alpha", "beta"]`))
+ require.NoError(t, err)
+ defer dictionary.Release()
+ indices, _, err := array.FromJSON(mem, arrow.PrimitiveTypes.Int32,
+ strings.NewReader(`[1, null, 1, 2, 1]`))
+ require.NoError(t, err)
+ defer indices.Release()
+
+ dictType := &arrow.DictionaryType{IndexType:
arrow.PrimitiveTypes.Int32, ValueType: arrow.BinaryTypes.String}
+ values := array.NewDictionaryArray(dictType, indices, dictionary)
+ defer values.Release()
+ sc := arrow.NewSchema([]arrow.Field{{Name: "values", Type: dictType,
Nullable: true}}, nil)
+
+ var output bytes.Buffer
+ writer, err := pqarrow.NewFileWriter(sc, &output,
parquet.NewWriterProperties(
+ parquet.WithAllocator(mem),
+ parquet.WithDictionaryDefault(true),
+ parquet.WithStats(false),
+ parquet.WithDataPageSize(1),
+ parquet.WithBatchSize(2),
+ parquet.WithBloomFilterEnabledFor("values", true),
+ parquet.WithBloomFilterNDVFor("values", 2),
+ ), pqarrow.NewArrowWriterProperties(pqarrow.WithAllocator(mem)))
+ require.NoError(t, err)
+ require.NoError(t, writer.NewRowGroupChecked())
+ require.NoError(t, writer.WriteColumnData(values))
+ require.NoError(t, writer.Close())
+
+ reader, err := file.NewParquetReader(bytes.NewReader(output.Bytes()))
+ require.NoError(t, err)
+ defer reader.Close()
+ bloomReader := reader.GetBloomFilterReader()
+ rowGroupBloom, err := bloomReader.RowGroup(0)
+ require.NoError(t, err)
+ require.NoError(t, rowGroupBloom.VisitColumnBloomFilter(0, func(filter
metadata.BloomFilter) error {
+ typedFilter :=
metadata.TypedBloomFilter[parquet.ByteArray]{BloomFilter: filter}
+ assert.True(t, typedFilter.Check(parquet.ByteArray("alpha")))
+ assert.True(t, typedFilter.Check(parquet.ByteArray("beta")))
+ return nil
+ }))
+}
+
+func TestWriteColumnDataRejectsOutOfBoundsDictionaryIndex(t *testing.T) {
+ mem := memory.DefaultAllocator
+
+ dictionary, _, err := array.FromJSON(mem, arrow.PrimitiveTypes.Int32,
+ strings.NewReader(`[10, 20]`))
+ require.NoError(t, err)
+ defer dictionary.Release()
+ indices, _, err := array.FromJSON(mem, arrow.PrimitiveTypes.Int8,
+ strings.NewReader(`[2]`))
+ require.NoError(t, err)
+ defer indices.Release()
+
+ dictType := &arrow.DictionaryType{IndexType: arrow.PrimitiveTypes.Int8,
ValueType: arrow.PrimitiveTypes.Int32}
+ values := array.NewDictionaryArray(dictType, indices, dictionary)
+ defer values.Release()
+ sc := arrow.NewSchema([]arrow.Field{{Name: "values", Type: dictType,
Nullable: false}}, nil)
+
+ var output bytes.Buffer
+ writer, err := pqarrow.NewFileWriter(sc, &output,
parquet.NewWriterProperties(
+ parquet.WithAllocator(mem),
+ parquet.WithDictionaryDefault(true),
+ parquet.WithStats(true),
+ ), pqarrow.NewArrowWriterProperties(pqarrow.WithAllocator(mem)))
+ require.NoError(t, err)
+ require.NoError(t, writer.NewRowGroupChecked())
+
+ err = writer.WriteColumnData(values)
+ require.ErrorIs(t, err, arrow.ErrIndex)
+ require.NoError(t, writer.Close())
+}
+
+func BenchmarkDictionaryArrayBloomFilter(b *testing.B) {
+ tests := []struct {
+ name string
+ dictionaryEntries int
+ referencedEntries int
+ }{
+ {name: "dictionary=100k/referenced=10", dictionaryEntries:
100_000, referencedEntries: 10},
+ {name: "dictionary=1m/referenced=100", dictionaryEntries:
1_000_000, referencedEntries: 100},
+ }
+
+ for _, tt := range tests {
+ b.Run(tt.name, func(b *testing.B) {
+ const numRows = 100_000
+ mem := memory.DefaultAllocator
+
+ dictionaryBuilder := array.NewStringBuilder(mem)
+ dictionaryBuilder.Reserve(tt.dictionaryEntries)
+ for i := 0; i < tt.dictionaryEntries; i++ {
+
dictionaryBuilder.Append(fmt.Sprintf("value-%08d", i))
+ }
+ dictionary := dictionaryBuilder.NewArray()
+ dictionaryBuilder.Release()
+ defer dictionary.Release()
+
+ indicesBuilder := array.NewInt32Builder(mem)
+ indicesBuilder.Reserve(numRows)
+ for i := 0; i < numRows; i++ {
+ indicesBuilder.Append(int32(i %
tt.referencedEntries))
+ }
+ indices := indicesBuilder.NewArray()
+ indicesBuilder.Release()
+ defer indices.Release()
+
+ dictType := &arrow.DictionaryType{IndexType:
arrow.PrimitiveTypes.Int32, ValueType: arrow.BinaryTypes.String}
+ values := array.NewDictionaryArray(dictType, indices,
dictionary)
+ defer values.Release()
+ sc := arrow.NewSchema([]arrow.Field{{Name: "values",
Type: dictType}}, nil)
+ props := parquet.NewWriterProperties(
+ parquet.WithDictionaryDefault(true),
+ parquet.WithStats(false),
+ parquet.WithBloomFilterEnabledFor("values",
true),
+ parquet.WithBloomFilterNDVFor("values",
int64(tt.referencedEntries)),
+ )
+ arrowProps :=
pqarrow.NewArrowWriterProperties(pqarrow.WithAllocator(mem))
+
+ b.ReportAllocs()
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ var output bytes.Buffer
+ writer, err := pqarrow.NewFileWriter(sc,
&output, props, arrowProps)
+ if err != nil {
+ b.Fatal(err)
+ }
+ if err := writer.NewRowGroupChecked(); err !=
nil {
+ b.Fatal(err)
+ }
+ if err := writer.WriteColumnData(values); err
!= nil {
+ b.Fatal(err)
+ }
+ if err := writer.Close(); err != nil {
+ b.Fatal(err)
+ }
+ }
+ })
+ }
+}
+
func (ps *ParquetIOTestSuite) TestSingleColumnOptionalDictionaryWrite() {
for _, dt := range fullTypeList {
// skip tests for bool as we don't do dictionaries for it
