Repository: segmentio/parquet-go
Pull request #75 (2022-02-19 02:13:10)

change parquet.BloomFilter.Check to accept a parquet.Value
This PR modifies the `parquet.BloomFilter` interface to accept a `parquet.Value` instead of a `[]byte`. The intent is to reduce reliance on a particular encoding, which results in more readable code in general in my opinion. It also creates better compile and runtime type checking, since the API does not accept any arbitrary byte slice anymore.

**Before**
```go
b := [8]byte{}
binary.LittleEndian.PutUint64(b[:], v)

if ok, err := bloomFilter.Check(b[:]); err != nil {
    ...
} else if ok {
    ...
}
```

**After**
```go
x := parquet.ValueOf(v)

if ok, err := bloomFilter.Check(x); err != nil {
    ...
} else if ok {
    ...
}
```


Patch:
diff --git a/bloom.go b/bloom.go
index 17754db..b123e4f 100644
--- a/bloom.go
+++ b/bloom.go
@@ -22,15 +22,12 @@ type BloomFilter interface {
 
 	// Tests whether the given value is present in the filter.
 	//
-	// The parquet value is expected to be PLAIN-encoded, as it would be by
-	// calling the parquet.Value.Bytes method.
-	//
 	// A non-nil error may be returned if reading the filter failed. This may
 	// happen if the filter was lazily loaded from a storage medium during the
 	// call to Check for example. Applications that can guarantee that the
 	// filter was in memory at the time Check was called can safely ignore the
 	// error, which would always be nil in this case.
-	Check(value []byte) (bool, error)
+	Check(value Value) (bool, error)
 }
 
 type bloomFilter struct {
@@ -39,8 +36,23 @@ type bloomFilter struct {
 	check func(io.ReaderAt, int64, uint64) (bool, error)
 }
 
-func (f *bloomFilter) Check(value []byte) (bool, error) {
-	return f.check(&f.SectionReader, f.Size(), f.hash.Sum64(value))
+func (f *bloomFilter) Check(v Value) (bool, error) {
+	return f.check(&f.SectionReader, f.Size(), v.hash(f.hash))
+}
+
+func (v Value) hash(h bloom.Hash) uint64 {
+	switch v.Kind() {
+	case Boolean:
+		return h.Sum64Uint8(uint8(v.u64))
+	case Int32, Float:
+		return h.Sum64Uint32(uint32(v.u64))
+	case Int64, Double:
+		return h.Sum64Uint64(v.u64)
+	case Int96:
+		return h.Sum64(v.Bytes())
+	default:
+		return h.Sum64(v.ByteArray())
+	}
 }
 
 func newBloomFilter(file io.ReaderAt, offset int64, header *format.BloomFilterHeader) *bloomFilter {
@@ -124,10 +136,6 @@ func newBloomFilterEncoder(filter bloom.MutableFilter, hash bloom.Hash) *bloomFi
 	return &bloomFilterEncoder{filter: filter, hash: hash}
 }
 
-func (e *bloomFilterEncoder) Check(value []byte) bool {
-	return e.filter.Check(e.hash.Sum64(value))
-}
-
 func (e *bloomFilterEncoder) Bytes() []byte {
 	return e.filter.Bytes()
 }
diff --git a/bloom/hash.go b/bloom/hash.go
index d8b7b5c..90df4e0 100644
--- a/bloom/hash.go
+++ b/bloom/hash.go
@@ -8,6 +8,14 @@ import "github.com/segmentio/parquet-go/bloom/xxhash"
 type Hash interface {
 	// Returns the 64 bit hash of the value passed as argument.
 	Sum64(value []byte) uint64
+
+	// Compute hashes of individual values of primitive types.
+	Sum64Uint8(value uint8) uint64
+	Sum64Uint16(value uint16) uint64
+	Sum64Uint32(value uint32) uint64
+	Sum64Uint64(value uint64) uint64
+	Sum64Uint128(value [16]byte) uint64
+
 	// Compute hashes of the array of fixed size values passed as arguments,
 	// returning the number of hashes written to the destination buffer.
 	MultiSum64Uint8(dst []uint64, src []uint8) int
@@ -24,6 +32,26 @@ func (XXH64) Sum64(b []byte) uint64 {
 	return xxhash.Sum64(b)
 }
 
+func (XXH64) Sum64Uint8(v uint8) uint64 {
+	return xxhash.Sum64Uint8(v)
+}
+
+func (XXH64) Sum64Uint16(v uint16) uint64 {
+	return xxhash.Sum64Uint16(v)
+}
+
+func (XXH64) Sum64Uint32(v uint32) uint64 {
+	return xxhash.Sum64Uint32(v)
+}
+
+func (XXH64) Sum64Uint64(v uint64) uint64 {
+	return xxhash.Sum64Uint64(v)
+}
+
+func (XXH64) Sum64Uint128(v [16]byte) uint64 {
+	return xxhash.Sum64Uint128(v)
+}
+
 func (XXH64) MultiSum64Uint8(h []uint64, v []uint8) int {
 	return xxhash.MultiSum64Uint8(h, v)
 }
diff --git a/concat.go b/concat.go
index a35921e..33ae2df 100644
--- a/concat.go
+++ b/concat.go
@@ -145,10 +145,10 @@ func (f concatenatedBloomFilter) Size() int64 {
 	return size
 }
 
-func (f concatenatedBloomFilter) Check(key []byte) (bool, error) {
+func (f concatenatedBloomFilter) Check(v Value) (bool, error) {
 	for _, c := range f.chunks {
 		if b := c.BloomFilter(); b != nil {
-			if ok, err := b.Check(key); ok || err != nil {
+			if ok, err := b.Check(v); ok || err != nil {
 				return ok, err
 			}
 		}
diff --git a/convert.go b/convert.go
index 9441490..47ce0aa 100644
--- a/convert.go
+++ b/convert.go
@@ -497,7 +497,7 @@ type missingBloomFilter struct{}
 
 func (missingBloomFilter) ReadAt([]byte, int64) (int, error) { return 0, io.EOF }
 func (missingBloomFilter) Size() int64                       { return 0 }
-func (missingBloomFilter) Check([]byte) (bool, error)        { return false, nil }
+func (missingBloomFilter) Check(Value) (bool, error)         { return false, nil }
 
 type missingPage struct{ *missingColumnChunk }
 
diff --git a/dictionary.go b/dictionary.go
index dfce3b3..5d01164 100644
--- a/dictionary.go
+++ b/dictionary.go
@@ -762,14 +762,18 @@ type indexedPageReader struct {
 }
 
 func (r *indexedPageReader) ReadValues(values []Value) (n int, err error) {
+	var v Value
 	for n < len(values) && r.offset < len(r.page.values) {
-		values[n] = r.page.dict.Index(int(r.page.values[r.offset]))
+		v = r.page.dict.Index(int(r.page.values[r.offset]))
+		v.columnIndex = r.page.columnIndex
+		values[n] = v
 		r.offset++
 		n++
 	}
 	if r.offset == len(r.page.values) {
 		err = io.EOF
 	}
+
 	return n, err
 }
 
diff --git a/row_group.go b/row_group.go
index 0024a3b..7f2aaa2 100644
--- a/row_group.go
+++ b/row_group.go
@@ -222,6 +222,10 @@ func (r *rowGroup) SortingColumns() []SortingColumn { return r.sorting }
 func (r *rowGroup) Schema() *Schema                 { return r.schema }
 func (r *rowGroup) Rows() Rows                      { return &rowGroupRowReader{rowGroup: r} }
 
+func NewRowGroupRowReader(rowGroup RowGroup) Rows {
+	return &rowGroupRowReader{rowGroup: rowGroup}
+}
+
 type rowGroupRowReader struct {
 	rowGroup RowGroup
 	schema   *Schema
@@ -438,7 +442,7 @@ type emptyBloomFilter struct{}
 
 func (emptyBloomFilter) ReadAt([]byte, int64) (int, error) { return 0, io.EOF }
 func (emptyBloomFilter) Size() int64                       { return 0 }
-func (emptyBloomFilter) Check([]byte) (bool, error)        { return false, nil }
+func (emptyBloomFilter) Check(Value) (bool, error)         { return false, nil }
 
 type emptyRowReader struct{ schema *Schema }
 
diff --git a/value.go b/value.go
index ae02e83..807cbb1 100644
--- a/value.go
+++ b/value.go
@@ -326,6 +326,8 @@ func ValueOf(v interface{}) Value {
 		return Value{}
 	case uuid.UUID:
 		return makeValueBytes(FixedLenByteArray, value[:])
+	case deprecated.Int96:
+		return makeValueInt96(value)
 	}
 
 	k := Kind(-1)


Test patch:
diff --git a/bloom_test.go b/bloom_test.go
index 6fd803a..760a832 100644
--- a/bloom_test.go
+++ b/bloom_test.go
@@ -1,8 +1,6 @@
 package parquet
 
 import (
-	"encoding/binary"
-	"math"
 	"math/rand"
 	"testing"
 	"testing/quick"
@@ -20,8 +18,8 @@ func TestBloomFilterEncoder(t *testing.T) {
 		)
 	}
 
-	check := func(e *bloomFilterEncoder, b []byte) bool {
-		return e.filter.Check(e.hash.Sum64(b))
+	check := func(e *bloomFilterEncoder, v Value) bool {
+		return e.filter.Check(v.hash(e.hash))
 	}
 
 	tests := []struct {
@@ -34,43 +32,7 @@ func TestBloomFilterEncoder(t *testing.T) {
 				f := newFilter(len(values))
 				f.EncodeBoolean(values)
 				for _, v := range values {
-					b := make([]byte, 1)
-					if v {
-						b[0] = 1
-					}
-					if !check(f, b) {
-						return false
-					}
-				}
-				return true
-			},
-		},
-
-		{
-			scenario: "INT8",
-			function: func(values []int8) bool {
-				f := newFilter(len(values))
-				f.EncodeInt8(values)
-				for _, v := range values {
-					b := make([]byte, 1)
-					b[0] = byte(v)
-					if !check(f, b) {
-						return false
-					}
-				}
-				return true
-			},
-		},
-
-		{
-			scenario: "INT16",
-			function: func(values []int16) bool {
-				f := newFilter(len(values))
-				f.EncodeInt16(values)
-				for _, v := range values {
-					b := make([]byte, 2)
-					binary.LittleEndian.PutUint16(b, uint16(v))
-					if !check(f, b) {
+					if !check(f, ValueOf(v)) {
 						return false
 					}
 				}
@@ -84,9 +46,7 @@ func TestBloomFilterEncoder(t *testing.T) {
 				f := newFilter(len(values))
 				f.EncodeInt32(values)
 				for _, v := range values {
-					b := make([]byte, 4)
-					binary.LittleEndian.PutUint32(b, uint32(v))
-					if !check(f, b) {
+					if !check(f, ValueOf(v)) {
 						return false
 					}
 				}
@@ -100,9 +60,7 @@ func TestBloomFilterEncoder(t *testing.T) {
 				f := newFilter(len(values))
 				f.EncodeInt64(values)
 				for _, v := range values {
-					b := make([]byte, 8)
-					binary.LittleEndian.PutUint64(b, uint64(v))
-					if !check(f, b) {
+					if !check(f, ValueOf(v)) {
 						return false
 					}
 				}
@@ -116,11 +74,7 @@ func TestBloomFilterEncoder(t *testing.T) {
 				f := newFilter(len(values))
 				f.EncodeInt96(values)
 				for _, v := range values {
-					b := make([]byte, 12)
-					binary.LittleEndian.PutUint32(b[0:4], v[0])
-					binary.LittleEndian.PutUint32(b[4:8], v[1])
-					binary.LittleEndian.PutUint32(b[8:12], v[2])
-					if !check(f, b) {
+					if !check(f, ValueOf(v)) {
 						return false
 					}
 				}
@@ -134,9 +88,7 @@ func TestBloomFilterEncoder(t *testing.T) {
 				f := newFilter(len(values))
 				f.EncodeFloat(values)
 				for _, v := range values {
-					b := make([]byte, 4)
-					binary.LittleEndian.PutUint32(b, math.Float32bits(v))
-					if !check(f, b) {
+					if !check(f, ValueOf(v)) {
 						return false
 					}
 				}
@@ -150,9 +102,7 @@ func TestBloomFilterEncoder(t *testing.T) {
 				f := newFilter(len(values))
 				f.EncodeDouble(values)
 				for _, v := range values {
-					b := make([]byte, 8)
-					binary.LittleEndian.PutUint64(b, math.Float64bits(v))
-					if !check(f, b) {
+					if !check(f, ValueOf(v)) {
 						return false
 					}
 				}
@@ -170,7 +120,7 @@ func TestBloomFilterEncoder(t *testing.T) {
 				f := newFilter(len(values))
 				f.EncodeByteArray(a)
 				for _, v := range values {
-					if !f.Check(v) {
+					if !check(f, ValueOf(v)) {
 						return false
 					}
 				}
@@ -180,15 +130,11 @@ func TestBloomFilterEncoder(t *testing.T) {
 
 		{
 			scenario: "FIXED_LEN_BYTE_ARRAY",
-			function: func(values [][]byte) bool {
-				a := encoding.ByteArrayList{}
-				for _, v := range values {
-					a.Push(v)
-				}
+			function: func(values []byte) bool {
 				f := newFilter(len(values))
-				f.EncodeByteArray(a)
+				f.EncodeFixedLenByteArray(1, values)
 				for _, v := range values {
-					if !f.Check(v) {
+					if !check(f, ValueOf([1]byte{v})) {
 						return false
 					}
 				}
diff --git a/buffer_test.go b/buffer_test.go
index d770f11..c364566 100644
--- a/buffer_test.go
+++ b/buffer_test.go
@@ -2,7 +2,6 @@ package parquet_test
 
 import (
 	"bytes"
-	"encoding/binary"
 	"io"
 	"math"
 	"sort"
@@ -399,9 +398,7 @@ func TestBufferGenerateBloomFilters(t *testing.T) {
 		fz := z.BloomFilter()
 
 		test := func(f parquet.BloomFilter, v float64) bool {
-			b := [8]byte{}
-			binary.LittleEndian.PutUint64(b[:], math.Float64bits(v))
-			if ok, err := f.Check(b[:]); err != nil {
+			if ok, err := f.Check(parquet.ValueOf(v)); err != nil {
 				t.Errorf("unexpected error checking bloom filter: %v", err)
 				return false
 			} else if !ok {
diff --git a/writer_test.go b/writer_test.go
index 1af09de..b4ab7a2 100644
--- a/writer_test.go
+++ b/writer_test.go
@@ -511,7 +511,7 @@ func TestWriterGenerateBloomFilters(t *testing.T) {
 		}
 
 		for i, row := range rows {
-			if ok, err := bloomFilter.Check([]byte(row.LastName)); err != nil {
+			if ok, err := bloomFilter.Check(parquet.ValueOf(row.LastName)); err != nil {
 				t.Errorf("unexpected error checking bloom filter: %v", err)
 				return false
 			} else if !ok {