Repository: segmentio/parquet-go
Pull request #138 (2022-04-13 16:57:35)

Core multi-page column index support
This PR was extracted from #134. It pulls the multi-page tests without the need for adding an equality state. As such, it focuses only on expanding test coverage, achieving order parity between buffer and file indexes, and fixing a bug in ordering string slices.


Patch:
diff --git a/column_index.go b/column_index.go
index df8b91b..3f36bcb 100644
--- a/column_index.go
+++ b/column_index.go
@@ -1,8 +1,6 @@
 package parquet
 
 import (
-	"bytes"
-
 	"github.com/segmentio/parquet-go/encoding"
 	"github.com/segmentio/parquet-go/format"
 	"github.com/segmentio/parquet-go/internal/bits"
@@ -145,8 +143,8 @@ func (i byteArrayColumnIndex) NullCount(int) int64 { return 0 }
 func (i byteArrayColumnIndex) NullPage(int) bool   { return false }
 func (i byteArrayColumnIndex) MinValue(int) Value  { return makeValueBytes(ByteArray, i.page.min()) }
 func (i byteArrayColumnIndex) MaxValue(int) Value  { return makeValueBytes(ByteArray, i.page.max()) }
-func (i byteArrayColumnIndex) IsAscending() bool   { return bytes.Compare(i.page.bounds()) < 0 }
-func (i byteArrayColumnIndex) IsDescending() bool  { return bytes.Compare(i.page.bounds()) > 0 }
+func (i byteArrayColumnIndex) IsAscending() bool   { return false }
+func (i byteArrayColumnIndex) IsDescending() bool  { return false }
 
 type fixedLenByteArrayColumnIndex struct{ page *fixedLenByteArrayPage }
 
@@ -159,8 +157,8 @@ func (i fixedLenByteArrayColumnIndex) MinValue(int) Value {
 func (i fixedLenByteArrayColumnIndex) MaxValue(int) Value {
 	return makeValueBytes(FixedLenByteArray, i.page.max())
 }
-func (i fixedLenByteArrayColumnIndex) IsAscending() bool  { return bytes.Compare(i.page.bounds()) < 0 }
-func (i fixedLenByteArrayColumnIndex) IsDescending() bool { return bytes.Compare(i.page.bounds()) > 0 }
+func (i fixedLenByteArrayColumnIndex) IsAscending() bool  { return false }
+func (i fixedLenByteArrayColumnIndex) IsDescending() bool { return false }
 
 // The ColumnIndexer interface is implemented by types that support generating
 // parquet column indexes.
diff --git a/column_index_default.go b/column_index_default.go
index d9bace2..a1f7c83 100644
--- a/column_index_default.go
+++ b/column_index_default.go
@@ -15,8 +15,8 @@ func (i booleanColumnIndex) NullCount(int) int64 { return 0 }
 func (i booleanColumnIndex) NullPage(int) bool   { return false }
 func (i booleanColumnIndex) MinValue(int) Value  { return makeValueBoolean(i.page.min()) }
 func (i booleanColumnIndex) MaxValue(int) Value  { return makeValueBoolean(i.page.max()) }
-func (i booleanColumnIndex) IsAscending() bool   { return compareBool(i.page.bounds()) < 0 }
-func (i booleanColumnIndex) IsDescending() bool  { return compareBool(i.page.bounds()) > 0 }
+func (i booleanColumnIndex) IsAscending() bool   { return false }
+func (i booleanColumnIndex) IsDescending() bool  { return false }
 
 type int32ColumnIndex struct{ page *int32Page }
 
@@ -25,8 +25,8 @@ func (i int32ColumnIndex) NullCount(int) int64 { return 0 }
 func (i int32ColumnIndex) NullPage(int) bool   { return false }
 func (i int32ColumnIndex) MinValue(int) Value  { return makeValueInt32(i.page.min()) }
 func (i int32ColumnIndex) MaxValue(int) Value  { return makeValueInt32(i.page.max()) }
-func (i int32ColumnIndex) IsAscending() bool   { return compareInt32(i.page.bounds()) < 0 }
-func (i int32ColumnIndex) IsDescending() bool  { return compareInt32(i.page.bounds()) > 0 }
+func (i int32ColumnIndex) IsAscending() bool   { return false }
+func (i int32ColumnIndex) IsDescending() bool  { return false }
 
 type int64ColumnIndex struct{ page *int64Page }
 
@@ -35,8 +35,8 @@ func (i int64ColumnIndex) NullCount(int) int64 { return 0 }
 func (i int64ColumnIndex) NullPage(int) bool   { return false }
 func (i int64ColumnIndex) MinValue(int) Value  { return makeValueInt64(i.page.min()) }
 func (i int64ColumnIndex) MaxValue(int) Value  { return makeValueInt64(i.page.max()) }
-func (i int64ColumnIndex) IsAscending() bool   { return compareInt64(i.page.bounds()) < 0 }
-func (i int64ColumnIndex) IsDescending() bool  { return compareInt64(i.page.bounds()) > 0 }
+func (i int64ColumnIndex) IsAscending() bool   { return false }
+func (i int64ColumnIndex) IsDescending() bool  { return false }
 
 type int96ColumnIndex struct{ page *int96Page }
 
@@ -45,8 +45,8 @@ func (i int96ColumnIndex) NullCount(int) int64 { return 0 }
 func (i int96ColumnIndex) NullPage(int) bool   { return false }
 func (i int96ColumnIndex) MinValue(int) Value  { return makeValueInt96(i.page.min()) }
 func (i int96ColumnIndex) MaxValue(int) Value  { return makeValueInt96(i.page.max()) }
-func (i int96ColumnIndex) IsAscending() bool   { return compareInt96(i.page.bounds()) < 0 }
-func (i int96ColumnIndex) IsDescending() bool  { return compareInt96(i.page.bounds()) > 0 }
+func (i int96ColumnIndex) IsAscending() bool   { return false }
+func (i int96ColumnIndex) IsDescending() bool  { return false }
 
 type floatColumnIndex struct{ page *floatPage }
 
@@ -55,8 +55,8 @@ func (i floatColumnIndex) NullCount(int) int64 { return 0 }
 func (i floatColumnIndex) NullPage(int) bool   { return false }
 func (i floatColumnIndex) MinValue(int) Value  { return makeValueFloat(i.page.min()) }
 func (i floatColumnIndex) MaxValue(int) Value  { return makeValueFloat(i.page.max()) }
-func (i floatColumnIndex) IsAscending() bool   { return compareFloat32(i.page.bounds()) < 0 }
-func (i floatColumnIndex) IsDescending() bool  { return compareFloat32(i.page.bounds()) > 0 }
+func (i floatColumnIndex) IsAscending() bool   { return false }
+func (i floatColumnIndex) IsDescending() bool  { return false }
 
 type doubleColumnIndex struct{ page *doublePage }
 
@@ -65,8 +65,8 @@ func (i doubleColumnIndex) NullCount(int) int64 { return 0 }
 func (i doubleColumnIndex) NullPage(int) bool   { return false }
 func (i doubleColumnIndex) MinValue(int) Value  { return makeValueDouble(i.page.min()) }
 func (i doubleColumnIndex) MaxValue(int) Value  { return makeValueDouble(i.page.max()) }
-func (i doubleColumnIndex) IsAscending() bool   { return compareFloat64(i.page.bounds()) < 0 }
-func (i doubleColumnIndex) IsDescending() bool  { return compareFloat64(i.page.bounds()) > 0 }
+func (i doubleColumnIndex) IsAscending() bool   { return false }
+func (i doubleColumnIndex) IsDescending() bool  { return false }
 
 type uint32ColumnIndex struct{ page uint32Page }
 
@@ -75,8 +75,8 @@ func (i uint32ColumnIndex) NullCount(int) int64 { return 0 }
 func (i uint32ColumnIndex) NullPage(int) bool   { return false }
 func (i uint32ColumnIndex) MinValue(int) Value  { return makeValueInt32(int32(i.page.min())) }
 func (i uint32ColumnIndex) MaxValue(int) Value  { return makeValueInt32(int32(i.page.max())) }
-func (i uint32ColumnIndex) IsAscending() bool   { return compareUint32(i.page.bounds()) < 0 }
-func (i uint32ColumnIndex) IsDescending() bool  { return compareUint32(i.page.bounds()) > 0 }
+func (i uint32ColumnIndex) IsAscending() bool   { return false }
+func (i uint32ColumnIndex) IsDescending() bool  { return false }
 
 type uint64ColumnIndex struct{ page uint64Page }
 
@@ -85,8 +85,8 @@ func (i uint64ColumnIndex) NullCount(int) int64 { return 0 }
 func (i uint64ColumnIndex) NullPage(int) bool   { return false }
 func (i uint64ColumnIndex) MinValue(int) Value  { return makeValueInt64(int64(i.page.min())) }
 func (i uint64ColumnIndex) MaxValue(int) Value  { return makeValueInt64(int64(i.page.max())) }
-func (i uint64ColumnIndex) IsAscending() bool   { return compareUint64(i.page.bounds()) < 0 }
-func (i uint64ColumnIndex) IsDescending() bool  { return compareUint64(i.page.bounds()) > 0 }
+func (i uint64ColumnIndex) IsAscending() bool   { return false }
+func (i uint64ColumnIndex) IsDescending() bool  { return false }
 
 type booleanColumnIndexer struct {
 	baseColumnIndexer
diff --git a/column_index_go18.go b/column_index_go18.go
index 9f593b1..46885ce 100644
--- a/column_index_go18.go
+++ b/column_index_go18.go
@@ -14,8 +14,8 @@ func (i columnIndex[T]) NullCount(int) int64 { return 0 }
 func (i columnIndex[T]) NullPage(int) bool   { return false }
 func (i columnIndex[T]) MinValue(int) Value  { return i.page.class.makeValue(i.page.min()) }
 func (i columnIndex[T]) MaxValue(int) Value  { return i.page.class.makeValue(i.page.max()) }
-func (i columnIndex[T]) IsAscending() bool   { return i.page.class.compare(i.page.bounds()) < 0 }
-func (i columnIndex[T]) IsDescending() bool  { return i.page.class.compare(i.page.bounds()) > 0 }
+func (i columnIndex[T]) IsAscending() bool   { return false }
+func (i columnIndex[T]) IsDescending() bool  { return false }
 
 type columnIndexer[T primitive] struct {
 	class      *class[T]
diff --git a/deprecated/int96.go b/deprecated/int96.go
index e242f3c..f82d2d4 100644
--- a/deprecated/int96.go
+++ b/deprecated/int96.go
@@ -132,7 +132,7 @@ func MinMaxInt96(data []Int96) (min, max Int96) {
 }
 
 func OrderOfInt96(data []Int96) int {
-	if len(data) > 0 {
+	if len(data) > 1 {
 		if int96AreInAscendingOrder(data) {
 			return +1
 		}
diff --git a/internal/bits/order.go b/internal/bits/order.go
index ca38dd5..567145d 100644
--- a/internal/bits/order.go
+++ b/internal/bits/order.go
@@ -6,10 +6,8 @@ import (
 
 func OrderOfBool(data []bool) int {
 	switch len(data) {
-	case 0:
+	case 0, 1:
 		return 0
-	case 1:
-		return 1
 	default:
 		k := 0
 		i := 0
@@ -51,10 +49,8 @@ func streakOfFalse(data []bool) int {
 
 func OrderOfInt32(data []int32) int {
 	switch len(data) {
-	case 0:
+	case 0, 1:
 		return 0
-	case 1:
-		return 1
 	default:
 		return orderOfInt32(data)
 	}
@@ -62,10 +58,8 @@ func OrderOfInt32(data []int32) int {
 
 func OrderOfInt64(data []int64) int {
 	switch len(data) {
-	case 0:
+	case 0, 1:
 		return 0
-	case 1:
-		return 1
 	default:
 		return orderOfInt64(data)
 	}
@@ -73,10 +67,8 @@ func OrderOfInt64(data []int64) int {
 
 func OrderOfUint32(data []uint32) int {
 	switch len(data) {
-	case 0:
+	case 0, 1:
 		return 0
-	case 1:
-		return 1
 	default:
 		return orderOfUint32(data)
 	}
@@ -84,10 +76,8 @@ func OrderOfUint32(data []uint32) int {
 
 func OrderOfUint64(data []uint64) int {
 	switch len(data) {
-	case 0:
+	case 0, 1:
 		return 0
-	case 1:
-		return 1
 	default:
 		return orderOfUint64(data)
 	}
@@ -95,10 +85,8 @@ func OrderOfUint64(data []uint64) int {
 
 func OrderOfFloat32(data []float32) int {
 	switch len(data) {
-	case 0:
+	case 0, 1:
 		return 0
-	case 1:
-		return 1
 	default:
 		return orderOfFloat32(data)
 	}
@@ -106,33 +94,29 @@ func OrderOfFloat32(data []float32) int {
 
 func OrderOfFloat64(data []float64) int {
 	switch len(data) {
-	case 0:
+	case 0, 1:
 		return 0
-	case 1:
-		return 1
 	default:
 		return orderOfFloat64(data)
 	}
 }
 
 func OrderOfBytes(data [][]byte) int {
-	if len(data) == 0 {
+	switch len(data) {
+	case 0, 1:
 		return 0
 	}
-	if len(data) == 1 {
-		return 1
-	}
 	data = skipBytesStreak(data)
 	if len(data) < 2 {
 		return 1
 	}
 	ordering := bytes.Compare(data[0], data[1])
 	switch {
-	case ordering > 0:
+	case ordering < 0:
 		if bytesAreInAscendingOrder(data[1:]) {
 			return +1
 		}
-	case ordering < 0:
+	case ordering > 0:
 		if bytesAreInDescendingOrder(data[1:]) {
 			return -1
 		}


Test patch:
diff --git a/column_test.go b/column_test.go
index ca5f63c..347e483 100644
--- a/column_test.go
+++ b/column_test.go
@@ -2,8 +2,10 @@ package parquet_test
 
 import (
 	"fmt"
+	"math/rand"
 	"testing"
 	"testing/quick"
+	"time"
 
 	"github.com/google/uuid"
 	"github.com/segmentio/parquet-go"
@@ -142,6 +144,7 @@ func checkColumnChunkColumnIndex(columnChunk parquet.ColumnChunk) error {
 	columnIndex := columnChunk.ColumnIndex()
 	numPages := columnIndex.NumPages()
 	pagesRead := 0
+	stats := newColumnStats(columnType)
 	err := forEachPage(columnChunk.Pages(), func(page parquet.Page) error {
 		pageMin, pageMax := page.Bounds()
 		indexMin := columnIndex.MinValue(pagesRead)
@@ -157,6 +160,7 @@ func checkColumnChunkColumnIndex(columnChunk parquet.ColumnChunk) error {
 		numNulls := int64(0)
 		numValues := int64(0)
 		err := forEachValue(page.Values(), func(value parquet.Value) error {
+			stats.observe(value)
 			if value.IsNull() {
 				numNulls++
 			}
@@ -177,17 +181,7 @@ func checkColumnChunkColumnIndex(columnChunk parquet.ColumnChunk) error {
 			return fmt.Errorf("page only contained null values but the index did not categorize it as a null page: nulls=%d", numNulls)
 		}
 
-		switch {
-		case columnIndex.IsAscending():
-			if columnType.Compare(pageMin, pageMax) > 0 {
-				return fmt.Errorf("column index is declared to be in ascending order but %v > %v", pageMin, pageMax)
-			}
-		case columnIndex.IsDescending():
-			if columnType.Compare(pageMin, pageMax) < 0 {
-				return fmt.Errorf("column index is declared to be in desending order but %v < %v", pageMin, pageMax)
-			}
-		}
-
+		stats.pageRead()
 		pagesRead++
 		return nil
 	})
@@ -197,6 +191,23 @@ func checkColumnChunkColumnIndex(columnChunk parquet.ColumnChunk) error {
 	if pagesRead != numPages {
 		return fmt.Errorf("number of pages found in column index differs from the number of pages read: index=%d read=%d", numPages, pagesRead)
 	}
+
+	actualOrder := columnIndexOrder(columnIndex)
+	observedOrder := observedIndexOrder(columnType, stats.minValues, stats.maxValues)
+	xorAscending := (columnIndex.IsAscending() || observedOrder == ascendingIndexOrder) &&
+		!(columnIndex.IsAscending() && observedOrder == ascendingIndexOrder)
+	xorDescending := (columnIndex.IsDescending() || observedOrder == descendingIndexOrder) &&
+		!(columnIndex.IsDescending() && observedOrder == descendingIndexOrder)
+
+	if xorAscending || xorDescending {
+		return fmt.Errorf("column index is declared to be %s while observed values %s (min values %s, max values %s)",
+			actualOrder,
+			observedOrder,
+			valueOrder(columnType, stats.minValues),
+			valueOrder(columnType, stats.maxValues),
+		)
+	}
+
 	return nil
 }
 
@@ -236,7 +247,9 @@ func checkColumnChunkOffsetIndex(columnChunk parquet.ColumnChunk) error {
 
 func testColumnPageIndexWithFile(t *testing.T, rows rows) bool {
 	if len(rows) > 0 {
-		f, err := createParquetFile(rows)
+		r := rand.New(rand.NewSource(time.Now().UnixNano()))
+		size := parquet.PageBufferSize(r.Intn(49) + 1)
+		f, err := createParquetFile(rows, size)
 		if err != nil {
 			t.Error(err)
 			return false
@@ -397,3 +410,109 @@ func (i *fileOffsetIndex) CompressedPageSize(j int) int64 {
 	return int64(i.PageLocations[j].CompressedPageSize)
 }
 func (i *fileOffsetIndex) FirstRowIndex(j int) int64 { return i.PageLocations[j].FirstRowIndex }
+
+type columnStats struct {
+	page       int
+	columnType parquet.Type
+	minValues  []parquet.Value
+	maxValues  []parquet.Value
+}
+
+func newColumnStats(columnType parquet.Type) *columnStats {
+	return &columnStats{columnType: columnType}
+}
+
+func (c *columnStats) observe(value parquet.Value) {
+	if c.page >= len(c.minValues) {
+		c.minValues = append(c.minValues, value)
+	} else if c.columnType.Compare(c.minValues[c.page], value) > 0 {
+		c.minValues[c.page] = value
+	}
+
+	if c.page >= len(c.maxValues) {
+		c.maxValues = append(c.maxValues, value)
+	} else if c.columnType.Compare(c.maxValues[c.page], value) < 0 {
+		c.maxValues[c.page] = value
+	}
+}
+
+func (c *columnStats) pageRead() {
+	c.page++
+}
+
+type indexOrder int
+
+const (
+	invalidIndexOrder indexOrder = iota
+	unorderedIndexOrder
+	ascendingIndexOrder
+	descendingIndexOrder
+)
+
+func (o indexOrder) String() string {
+	switch o {
+	case unorderedIndexOrder:
+		return "unordered"
+	case ascendingIndexOrder:
+		return "ascending"
+	case descendingIndexOrder:
+		return "descending"
+	default:
+		return "invalid"
+	}
+}
+
+func columnIndexOrder(index parquet.ColumnIndex) indexOrder {
+	switch {
+	case index.IsAscending() && index.IsDescending():
+		return invalidIndexOrder
+	case index.IsAscending():
+		return ascendingIndexOrder
+	case index.IsDescending():
+		return descendingIndexOrder
+	default:
+		return unorderedIndexOrder
+	}
+}
+
+func observedIndexOrder(columnType parquet.Type, minValues []parquet.Value, maxValues []parquet.Value) indexOrder {
+	a := valueOrder(columnType, minValues)
+	b := valueOrder(columnType, maxValues)
+
+	switch {
+	case a == ascendingIndexOrder && b == ascendingIndexOrder:
+		return ascendingIndexOrder
+	case a == descendingIndexOrder && b == descendingIndexOrder:
+		return descendingIndexOrder
+	default:
+		return unorderedIndexOrder
+	}
+}
+
+func valueOrder(columnType parquet.Type, values []parquet.Value) indexOrder {
+	switch len(values) {
+	case 0, 1:
+		return unorderedIndexOrder
+	}
+
+	var order int
+	for i := 1; i < len(values); i++ {
+		next := columnType.Compare(values[i-1], values[i])
+		if next == 0 {
+			continue
+		}
+		if order == 0 {
+			order = next
+			continue
+		}
+		if order != next {
+			return unorderedIndexOrder
+		}
+	}
+
+	if order > 0 {
+		return descendingIndexOrder
+	}
+
+	return ascendingIndexOrder
+}
diff --git a/internal/bits/order_test.go b/internal/bits/order_test.go
index 11cd78d..6807dae 100644
--- a/internal/bits/order_test.go
+++ b/internal/bits/order_test.go
@@ -85,19 +85,19 @@ func isUndefined(ordering int) bool {
 	return ordering == 0
 }
 
-func isSorted(set sort.Interface) bool {
-	return set.Len() > 0 && sort.IsSorted(set)
+func isOrdered(set sort.Interface) bool {
+	return set.Len() > 1 && sort.IsSorted(set)
 }
 
 func checkOrdering(t *testing.T, set sort.Interface, ordering int) bool {
 	t.Helper()
 	switch {
-	case isSorted(set):
+	case isOrdered(set):
 		if !isAscending(ordering) {
 			t.Errorf("got=%s want=%s", orderingName(ordering), ascending)
 			return false
 		}
-	case isSorted(sort.Reverse(set)):
+	case isOrdered(sort.Reverse(set)):
 		if !isDescending(ordering) {
 			t.Errorf("got=%s want=%s", orderingName(ordering), descending)
 			return false
@@ -348,8 +348,8 @@ func TestOrderOfBytes(t *testing.T) {
 	}
 	err := quickCheck(func(values [][16]byte) bool {
 		slices := make([][]byte, len(values))
-		for i, v := range values {
-			slices[i] = v[:]
+		for i := range values {
+			slices[i] = values[i][:]
 		}
 		if !check(slices) {
 			return false