fix(parquet): build bloom filters from dictionary entries

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

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)))
+		})
+	}
+}