You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by ze...@apache.org on 2025/01/09 19:59:43 UTC
(arrow-go) branch main updated: feat(parquet): Add support for Page Indexes (#223)
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 365a84a feat(parquet): Add support for Page Indexes (#223)
365a84a is described below
commit 365a84ac084ad7daae6c5c758235042d6960b96e
Author: Matt Topol <zo...@gmail.com>
AuthorDate: Thu Jan 9 14:59:37 2025 -0500
feat(parquet): Add support for Page Indexes (#223)
### Rationale for this change
Adds support for writing and retrieving Page Indexes from parquet files
### What changes are included in this PR?
Added new writer properties to enable writing PageIndexes either as
default or for specific columns.
Added `RowGroupPageIndexReader`, `ColumnIndexReader`,
`OffsetIndexReader` to access page index information.
### Are these changes tested?
Yes, extensive unit tests have been added including round trip tests and
encryption/decryption tests.
### Are there any user-facing changes?
no. Just new functions available.
closes #33
---
arrow/flight/flightsql/column_metadata.go | 6 +-
parquet/encryption_read_config_test.go | 119 +++-
parquet/encryption_write_config_test.go | 5 +-
parquet/file/column_writer.go | 31 +-
parquet/file/column_writer_test.go | 3 +
parquet/file/file_reader.go | 21 +-
parquet/file/file_writer.go | 46 +-
parquet/file/file_writer_test.go | 471 +++++++++++++
parquet/file/page_reader.go | 57 +-
parquet/file/page_writer.go | 67 +-
parquet/file/row_group_writer.go | 68 +-
parquet/file/size_statistics.go | 28 +
parquet/internal/encryption/decryptor.go | 38 ++
parquet/metadata/column_chunk.go | 31 +
parquet/metadata/file.go | 41 +-
parquet/metadata/metadata_test.go | 49 ++
parquet/metadata/page_index.go | 928 ++++++++++++++++++++++++++
parquet/metadata/page_index_internal_test.go | 197 ++++++
parquet/metadata/page_index_test.go | 802 ++++++++++++++++++++++
parquet/metadata/statistics.go | 454 ++++++++++++-
parquet/metadata/statistics_types.gen.go | 9 +
parquet/metadata/statistics_types.gen.go.tmpl | 1 +
parquet/pqarrow/schema.go | 8 +-
parquet/types.go | 4 +
parquet/writer_properties.go | 63 +-
25 files changed, 3483 insertions(+), 64 deletions(-)
diff --git a/arrow/flight/flightsql/column_metadata.go b/arrow/flight/flightsql/column_metadata.go
index 66962fd..d46fab3 100644
--- a/arrow/flight/flightsql/column_metadata.go
+++ b/arrow/flight/flightsql/column_metadata.go
@@ -63,7 +63,7 @@ type ColumnMetadata struct {
}
func (c *ColumnMetadata) findStrVal(key string) (string, bool) {
- idx := c.Data.FindKey(CatalogNameKey)
+ idx := c.Data.FindKey(key)
if idx == -1 {
return "", false
}
@@ -71,7 +71,7 @@ func (c *ColumnMetadata) findStrVal(key string) (string, bool) {
}
func (c *ColumnMetadata) findBoolVal(key string) (bool, bool) {
- idx := c.Data.FindKey(CatalogNameKey)
+ idx := c.Data.FindKey(key)
if idx == -1 {
return false, false
}
@@ -79,7 +79,7 @@ func (c *ColumnMetadata) findBoolVal(key string) (bool, bool) {
}
func (c *ColumnMetadata) findInt32Val(key string) (int32, bool) {
- idx := c.Data.FindKey(CatalogNameKey)
+ idx := c.Data.FindKey(key)
if idx == -1 {
return 0, false
}
diff --git a/parquet/encryption_read_config_test.go b/parquet/encryption_read_config_test.go
index fae1843..9d4046c 100644
--- a/parquet/encryption_read_config_test.go
+++ b/parquet/encryption_read_config_test.go
@@ -17,6 +17,7 @@
package parquet_test
import (
+ "bytes"
"encoding/binary"
"fmt"
"os"
@@ -27,6 +28,9 @@ import (
"github.com/apache/arrow-go/v18/parquet"
"github.com/apache/arrow-go/v18/parquet/file"
"github.com/apache/arrow-go/v18/parquet/internal/encryption"
+ "github.com/apache/arrow-go/v18/parquet/metadata"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
"github.com/stretchr/testify/suite"
)
@@ -402,8 +406,121 @@ func (d *TestDecryptionSuite) decryptFile(filename string, decryptConfigNum int)
}
}
+func assertColumnIndex[T parquet.ColumnTypes](t *testing.T, colIndex metadata.ColumnIndex, expectedNullCounts []int64, expectedMinVals, expectedMaxVals []T) {
+ typedColIndex, ok := colIndex.(*metadata.TypedColumnIndex[T])
+ require.True(t, ok)
+ assert.Equal(t, typedColIndex.GetNullCounts(), expectedNullCounts)
+
+ assert.Equal(t, expectedMinVals, typedColIndex.MinValues())
+ assert.Equal(t, expectedMaxVals, typedColIndex.MaxValues())
+}
+
+func (d *TestDecryptionSuite) decryptPageIndex(fileName string, decryptConfigNum int) {
+ // if we get decryption_config_num = x then it means the actual number is x+1
+ // and since we want decryption_config_num=4 we set the condition to 3
+ props := parquet.NewReaderProperties(memory.DefaultAllocator)
+ if decryptConfigNum != 3 {
+ props.FileDecryptProps = d.decryptionConfigs[decryptConfigNum].Clone("")
+ }
+
+ fileReader, err := file.OpenParquetFile(fileName, false, file.WithReadProps(props))
+ if err != nil {
+ panic(err)
+ }
+ defer fileReader.Close()
+
+ pageIndexReader := fileReader.GetPageIndexReader()
+ d.Require().NotNil(pageIndexReader)
+
+ fileMetadata := fileReader.MetaData()
+ numRowGroups := fileMetadata.NumRowGroups()
+ numCols := fileMetadata.NumColumns()
+ d.Equal(8, numCols)
+
+ // we cannot read page index of encrypted columns in the plaintext mode
+ needRowGroups := make([]int32, numRowGroups)
+ for i := range needRowGroups {
+ needRowGroups[i] = int32(i)
+ }
+ var needColumns []int32
+ if props.FileDecryptProps == nil {
+ needColumns = []int32{0, 1, 2, 3, 6, 7}
+ } else {
+ needColumns = []int32{0, 1, 2, 3, 4, 5, 6, 7}
+ }
+
+ // provide hint of requested columns to avoid accessing encrypted columns
+ // without decryption properties
+ d.Require().NoError(pageIndexReader.WillNeed(needRowGroups, needColumns,
+ metadata.PageIndexSelection{ColumnIndex: true, OffsetIndex: true}))
+
+ // iterate over all row groups in the file
+ for r := range numRowGroups {
+ rowGrpPageIndexRdr, err := pageIndexReader.RowGroup(r)
+ d.Require().NoError(err)
+ d.Require().NotNil(rowGrpPageIndexRdr)
+
+ for c := range numCols {
+ // skip reading encrypted columns without decryption props
+ if props.FileDecryptProps == nil && (c == 4 || c == 5) {
+ continue
+ }
+
+ const expectedNumPages = 1
+
+ // check offset index
+ offsetIndex, err := rowGrpPageIndexRdr.GetOffsetIndex(c)
+ d.Require().NoError(err)
+ d.Require().NotNil(offsetIndex)
+ d.Len(offsetIndex.GetPageLocations(), expectedNumPages)
+ firstPage := offsetIndex.GetPageLocations()[0]
+ d.EqualValues(0, firstPage.FirstRowIndex)
+ d.Greater(firstPage.CompressedPageSize, int32(0))
+
+ // int96 column does not have column index
+ if c == 3 {
+ continue
+ }
+
+ // check column index
+ colIndex, err := rowGrpPageIndexRdr.GetColumnIndex(c)
+ d.Require().NoError(err)
+ d.Require().NotNil(colIndex)
+ nullPages := colIndex.GetNullPages()
+ d.Len(nullPages, expectedNumPages)
+ d.False(nullPages[0])
+ d.Len(colIndex.GetMinValues(), expectedNumPages)
+ d.Len(colIndex.GetMaxValues(), expectedNumPages)
+ d.True(colIndex.IsSetNullCounts())
+
+ switch c {
+ case 0:
+ assertColumnIndex(d.T(), colIndex, []int64{0}, []bool{false}, []bool{true})
+ case 1:
+ assertColumnIndex(d.T(), colIndex, []int64{0}, []int32{0}, []int32{49})
+ case 2:
+ assertColumnIndex(d.T(), colIndex, []int64{0}, []int64{0}, []int64{99000000000000})
+ case 4:
+ assertColumnIndex(d.T(), colIndex, []int64{0}, []float32{0}, []float32{53.9})
+ case 5:
+ assertColumnIndex(d.T(), colIndex, []int64{0}, []float64{0}, []float64{54.4444439})
+ case 6:
+ assertColumnIndex(d.T(), colIndex, []int64{25},
+ []parquet.ByteArray{[]byte("parquet000")}, []parquet.ByteArray{[]byte("parquet048")})
+ case 7:
+ assertColumnIndex(d.T(), colIndex, []int64{0},
+ []parquet.FixedLenByteArray{make(parquet.FixedLenByteArray, 10)},
+ []parquet.FixedLenByteArray{bytes.Repeat([]byte{49}, 10)})
+ }
+ }
+ }
+}
+
func (d *TestDecryptionSuite) checkResults(fileName string, decryptionConfig, encryptionConfig uint) {
- decFn := func() { d.decryptFile(fileName, int(decryptionConfig-1)) }
+ decFn := func() {
+ d.decryptFile(fileName, int(decryptionConfig-1))
+ d.decryptPageIndex(fileName, int(decryptionConfig-1))
+ }
// Encryption configuration number 5 contains aad_prefix and disable_aad_prefix_storage
// an exception is expected to be thrown if the file is not decrypted with aad_prefix
diff --git a/parquet/encryption_write_config_test.go b/parquet/encryption_write_config_test.go
index d2fe761..b2508a2 100644
--- a/parquet/encryption_write_config_test.go
+++ b/parquet/encryption_write_config_test.go
@@ -82,7 +82,10 @@ type EncryptionConfigTestSuite struct {
func (en *EncryptionConfigTestSuite) encryptFile(configs *parquet.FileEncryptionProperties, filename string) {
filename = filepath.Join(tempdir, filename)
- props := parquet.NewWriterProperties(parquet.WithCompression(compress.Codecs.Snappy), parquet.WithEncryptionProperties(configs))
+ props := parquet.NewWriterProperties(
+ parquet.WithPageIndexEnabled(true),
+ parquet.WithCompression(compress.Codecs.Snappy),
+ parquet.WithEncryptionProperties(configs))
outFile, err := os.Create(filename)
en.Require().NoError(err)
en.Require().NotNil(outFile)
diff --git a/parquet/file/column_writer.go b/parquet/file/column_writer.go
index 22f8fe2..f608cd0 100755
--- a/parquet/file/column_writer.go
+++ b/parquet/file/column_writer.go
@@ -331,16 +331,30 @@ func (w *columnWriter) buildDataPageV1(defLevelsRLESize, repLevelsRLESize, uncom
data = w.uncompressedData.Bytes()
}
+ firstRowIndex := int64(w.rowsWritten)
+
// write the page to sink eagerly if there's no dictionary or if dictionary encoding has fallen back
if w.hasDict && !w.fallbackToNonDict {
pageSlice := make([]byte, len(data))
copy(pageSlice, data)
- page := NewDataPageV1WithStats(memory.NewBufferBytes(pageSlice), int32(w.numBufferedValues), w.encoding, parquet.Encodings.RLE, parquet.Encodings.RLE, uncompressed, pageStats)
+ page := NewDataPageV1WithConfig(memory.NewBufferBytes(pageSlice), parquet.Encodings.RLE, parquet.Encodings.RLE, DataPageConfig{
+ Num: int32(w.numBufferedValues),
+ Encoding: w.encoding,
+ UncompressedSize: uncompressed,
+ Stats: pageStats,
+ FirstRowIndex: firstRowIndex,
+ })
w.totalCompressedBytes += int64(page.buf.Len()) // + size of Pageheader
w.pages = append(w.pages, page)
} else {
w.totalCompressedBytes += int64(len(data))
- dp := NewDataPageV1WithStats(memory.NewBufferBytes(data), int32(w.numBufferedValues), w.encoding, parquet.Encodings.RLE, parquet.Encodings.RLE, uncompressed, pageStats)
+ dp := NewDataPageV1WithConfig(memory.NewBufferBytes(data), parquet.Encodings.RLE, parquet.Encodings.RLE, DataPageConfig{
+ Num: int32(w.numBufferedValues),
+ Encoding: w.encoding,
+ UncompressedSize: uncompressed,
+ Stats: pageStats,
+ FirstRowIndex: firstRowIndex,
+ })
defer dp.Release()
return w.WriteDataPage(dp)
}
@@ -374,9 +388,16 @@ func (w *columnWriter) buildDataPageV2(defLevelsRLESize, repLevelsRLESize, uncom
nullCount := int32(pageStats.NullCount)
defLevelsByteLen := int32(defLevelsRLESize)
repLevelsByteLen := int32(repLevelsRLESize)
-
- page := NewDataPageV2WithStats(memory.NewBufferBytes(combined.Bytes()), numValues, nullCount, numRows, w.encoding,
- defLevelsByteLen, repLevelsByteLen, uncompressed, w.pager.HasCompressor(), pageStats)
+ firstRowIndex := int64(w.rowsWritten)
+
+ page := NewDataPageV2WithConfig(memory.NewBufferBytes(combined.Bytes()), nullCount, numRows, defLevelsByteLen, repLevelsByteLen,
+ w.pager.HasCompressor(), DataPageConfig{
+ Num: numValues,
+ Encoding: w.encoding,
+ UncompressedSize: uncompressed,
+ Stats: pageStats,
+ FirstRowIndex: firstRowIndex,
+ })
if w.hasDict && !w.fallbackToNonDict {
w.totalCompressedBytes += int64(page.buf.Len()) // + sizeof pageheader
w.pages = append(w.pages, page)
diff --git a/parquet/file/column_writer_test.go b/parquet/file/column_writer_test.go
index 718cd30..68bffeb 100755
--- a/parquet/file/column_writer_test.go
+++ b/parquet/file/column_writer_test.go
@@ -79,6 +79,9 @@ func (m *mockpagewriter) Compress(buf *bytes.Buffer, src []byte) []byte {
func (m *mockpagewriter) Reset(sink utils.WriterTell, codec compress.Compression, compressionLevel int, metadata *metadata.ColumnChunkMetaDataBuilder, rgOrdinal, columnOrdinal int16, metaEncryptor, dataEncryptor encryption.Encryptor) error {
return m.Called().Error(0)
}
+func (m *mockpagewriter) SetIndexBuilders(colidx metadata.ColumnIndexBuilder, offsetidx *metadata.OffsetIndexBuilder) {
+ m.Called(colidx, offsetidx)
+}
func TestWriteDataPageV1NumValues(t *testing.T) {
sc := schema.NewSchema(schema.MustGroup(schema.NewGroupNode("schema", parquet.Repetitions.Required, schema.FieldList{
diff --git a/parquet/file/file_reader.go b/parquet/file/file_reader.go
index c7ba813..7807ac6 100644
--- a/parquet/file/file_reader.go
+++ b/parquet/file/file_reader.go
@@ -44,10 +44,11 @@ var (
// Reader is the main interface for reading a parquet file
type Reader struct {
- r parquet.ReaderAtSeeker
- props *parquet.ReaderProperties
- metadata *metadata.FileMetaData
- fileDecryptor encryption.FileDecryptor
+ r parquet.ReaderAtSeeker
+ props *parquet.ReaderProperties
+ metadata *metadata.FileMetaData
+ fileDecryptor encryption.FileDecryptor
+ pageIndexReader *metadata.PageIndexReader
bufferPool sync.Pool
}
@@ -312,3 +313,15 @@ func (f *Reader) RowGroup(i int) *RowGroupReader {
bufferPool: &f.bufferPool,
}
}
+
+func (f *Reader) GetPageIndexReader() *metadata.PageIndexReader {
+ if f.pageIndexReader == nil {
+ f.pageIndexReader = &metadata.PageIndexReader{
+ Input: f.r,
+ Props: f.props,
+ FileMetadata: f.metadata,
+ Decryptor: f.fileDecryptor,
+ }
+ }
+ return f.pageIndexReader
+}
diff --git a/parquet/file/file_writer.go b/parquet/file/file_writer.go
index 6f35ac3..f616831 100644
--- a/parquet/file/file_writer.go
+++ b/parquet/file/file_writer.go
@@ -30,15 +30,16 @@ import (
// Writer is the primary interface for writing a parquet file
type Writer struct {
- sink utils.WriteCloserTell
- open bool
- footerFlushed bool
- props *parquet.WriterProperties
- rowGroups int
- nrows int
- metadata metadata.FileMetaDataBuilder
- fileEncryptor encryption.FileEncryptor
- rowGroupWriter *rowGroupWriter
+ sink utils.WriteCloserTell
+ open bool
+ footerFlushed bool
+ props *parquet.WriterProperties
+ rowGroups int
+ nrows int
+ metadata metadata.FileMetaDataBuilder
+ fileEncryptor encryption.FileEncryptor
+ rowGroupWriter *rowGroupWriter
+ pageIndexBuilder *metadata.PageIndexBuilder
// The Schema of this writer
Schema *schema.Schema
@@ -128,7 +129,12 @@ func (fw *Writer) appendRowGroup(buffered bool) *rowGroupWriter {
fw.rowGroups++
fw.footerFlushed = false
rgMeta := fw.metadata.AppendRowGroup()
- fw.rowGroupWriter = newRowGroupWriter(fw.sink, rgMeta, int16(fw.rowGroups)-1, fw.props, buffered, fw.fileEncryptor)
+ if fw.pageIndexBuilder != nil {
+ fw.pageIndexBuilder.AppendRowGroup()
+ }
+
+ fw.rowGroupWriter = newRowGroupWriter(fw.sink, rgMeta, int16(fw.rowGroups)-1,
+ fw.props, buffered, fw.fileEncryptor, fw.pageIndexBuilder)
return fw.rowGroupWriter
}
@@ -160,6 +166,24 @@ func (fw *Writer) startFile() {
if n != 4 || err != nil {
panic("failed to write magic number")
}
+
+ if fw.props.PageIndexEnabled() {
+ fw.pageIndexBuilder = &metadata.PageIndexBuilder{
+ Schema: fw.Schema,
+ Encryptor: fw.fileEncryptor,
+ }
+ }
+}
+
+func (fw *Writer) writePageIndex() {
+ if fw.pageIndexBuilder != nil {
+ // serialize page index after all row groups have been written and report
+ // location to the file metadata
+ fw.pageIndexBuilder.Finish()
+ var pageIndexLocation metadata.PageIndexLocation
+ fw.pageIndexBuilder.WriteTo(fw.sink, &pageIndexLocation)
+ fw.metadata.SetPageIndexLocation(pageIndexLocation)
+ }
}
// AppendKeyValueMetadata appends a key/value pair to the existing key/value metadata
@@ -206,6 +230,8 @@ func (fw *Writer) FlushWithFooter() error {
}
fw.rowGroupWriter = nil
+ fw.writePageIndex()
+
fileMetadata, err := fw.metadata.Snapshot()
if err != nil {
return err
diff --git a/parquet/file/file_writer_test.go b/parquet/file/file_writer_test.go
index 5c658b0..7bcbe29 100644
--- a/parquet/file/file_writer_test.go
+++ b/parquet/file/file_writer_test.go
@@ -19,15 +19,22 @@ package file_test
import (
"bytes"
"fmt"
+ "math"
"reflect"
+ "slices"
"testing"
+ "unsafe"
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
"github.com/apache/arrow-go/v18/arrow/memory"
"github.com/apache/arrow-go/v18/parquet"
"github.com/apache/arrow-go/v18/parquet/compress"
"github.com/apache/arrow-go/v18/parquet/file"
"github.com/apache/arrow-go/v18/parquet/internal/encoding"
"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/apache/arrow-go/v18/parquet/schema"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
@@ -596,3 +603,467 @@ func TestLZ4RawFileRoundtrip(t *testing.T) {
require.NoError(t, rdr.Close())
}
+
+type ColumnIndexObject struct {
+ nullPages []bool
+ minValues [][]byte
+ maxValues [][]byte
+ boundaryOrder metadata.BoundaryOrder
+ nullCounts []int64
+}
+
+func NewColumnIndexObject(colIdx metadata.ColumnIndex) (ret ColumnIndexObject) {
+ if colIdx == nil {
+ return
+ }
+
+ ret.nullPages = colIdx.GetNullPages()
+ ret.minValues = colIdx.GetMinValues()
+ ret.maxValues = colIdx.GetMaxValues()
+ ret.boundaryOrder = colIdx.GetBoundaryOrder()
+ if colIdx.IsSetNullCounts() {
+ ret.nullCounts = colIdx.GetNullCounts()
+ }
+ return
+}
+
+func simpleEncode[T int32 | int64 | float32 | float64](val T) []byte {
+ return unsafe.Slice((*byte)(unsafe.Pointer(&val)), unsafe.Sizeof(val))
+}
+
+type PageIndexRoundTripSuite struct {
+ suite.Suite
+
+ mem *memory.CheckedAllocator
+ buf bytes.Buffer
+ columnIndexes []ColumnIndexObject
+ pageVersion parquet.DataPageVersion
+}
+
+func (t *PageIndexRoundTripSuite) SetupTest() {
+ t.mem = memory.NewCheckedAllocator(memory.NewGoAllocator())
+
+ t.buf.Reset()
+ t.columnIndexes = make([]ColumnIndexObject, 0)
+}
+
+func (t *PageIndexRoundTripSuite) TearDownTest() {
+ t.mem.AssertSize(t.T(), 0)
+}
+
+func (t *PageIndexRoundTripSuite) writeFile(props *parquet.WriterProperties, tbl arrow.Table) {
+ schema := tbl.Schema()
+ arrWriterProps := pqarrow.NewArrowWriterProperties(pqarrow.WithAllocator(t.mem))
+
+ wr, err := pqarrow.NewFileWriter(schema, &t.buf, props, arrWriterProps)
+ t.Require().NoError(err)
+
+ t.Require().NoError(wr.WriteTable(tbl, tbl.NumRows()))
+ t.Require().NoError(wr.Close())
+}
+
+func (t *PageIndexRoundTripSuite) readPageIndexes(expectNumRG, expectNumPages int, expectColsWithoutIdx []int) {
+ rdr, err := file.NewParquetReader(bytes.NewReader(t.buf.Bytes()),
+ file.WithReadProps(parquet.NewReaderProperties(t.mem)))
+ t.Require().NoError(err)
+ defer rdr.Close()
+
+ fileMeta := rdr.MetaData()
+ t.Equal(expectNumRG, fileMeta.NumRowGroups())
+
+ pageIdxReader := rdr.GetPageIndexReader()
+ t.Require().NotNil(pageIdxReader)
+
+ offsetLowerBound := int64(0)
+ for rg := range fileMeta.NumRowGroups() {
+ rgIdxRdr, err := pageIdxReader.RowGroup(rg)
+ t.Require().NoError(err)
+ t.Require().NotNil(rgIdxRdr)
+
+ rgReader := rdr.RowGroup(rg)
+ t.Require().NotNil(rgReader)
+
+ for col := range fileMeta.Schema.NumColumns() {
+ colIdx, err := rgIdxRdr.GetColumnIndex(col)
+ t.Require().NoError(err)
+
+ t.columnIndexes = append(t.columnIndexes, NewColumnIndexObject(colIdx))
+ expectNoPageIndex := slices.Contains(expectColsWithoutIdx, col)
+
+ offsetIdx, err := rgIdxRdr.GetOffsetIndex(col)
+ t.Require().NoError(err)
+ if expectNoPageIndex {
+ t.Nil(offsetIdx)
+ } else {
+ t.checkOffsetIndex(offsetIdx, expectNumPages, &offsetLowerBound)
+ }
+
+ // verify page stats are not written to page header if page index
+ // is enabled
+ pgRdr, err := rgReader.GetColumnPageReader(col)
+ t.Require().NoError(err)
+ t.Require().NotNil(pgRdr)
+
+ for pgRdr.Next() {
+ page := pgRdr.Page()
+ t.Require().NotNil(page)
+
+ if page.Type() == file.PageTypeDataPage || page.Type() == file.PageTypeDataPageV2 {
+ stats := page.(file.DataPage).Statistics()
+ t.Equalf(expectNoPageIndex, stats.IsSet(), "rg: %d, col: %d", rg, col)
+ }
+ }
+
+ t.Require().NoError(pgRdr.Err())
+ }
+ }
+}
+
+func (t *PageIndexRoundTripSuite) checkOffsetIndex(offsetIdx metadata.OffsetIndex, expectNumPages int, offsetLowerBoundInOut *int64) {
+ t.Require().NotNil(offsetIdx)
+ locations := offsetIdx.GetPageLocations()
+ t.Len(locations, expectNumPages)
+
+ prevFirstRowIdx := int64(-1)
+ for _, loc := range locations {
+ // make sure first row index is ascending within a row group
+ t.Greater(loc.FirstRowIndex, prevFirstRowIdx)
+ // make sure page offset is ascending across the file
+ t.GreaterOrEqual(loc.Offset, *offsetLowerBoundInOut)
+ // make sure page size is positive
+ t.Greater(loc.CompressedPageSize, int32(0))
+ prevFirstRowIdx = loc.FirstRowIndex
+ *offsetLowerBoundInOut = loc.Offset + int64(loc.CompressedPageSize)
+ }
+}
+
+func (t *PageIndexRoundTripSuite) TestSimpleRoundTrip() {
+ props := parquet.NewWriterProperties(parquet.WithAllocator(t.mem),
+ parquet.WithDataPageVersion(t.pageVersion),
+ parquet.WithPageIndexEnabled(true), parquet.WithMaxRowGroupLength(4))
+
+ sc := arrow.NewSchema([]arrow.Field{
+ {Name: "c0", Type: arrow.PrimitiveTypes.Int64, Nullable: true},
+ {Name: "c1", Type: arrow.BinaryTypes.String, Nullable: true},
+ {Name: "c2", Type: arrow.ListOf(arrow.PrimitiveTypes.Int64), Nullable: true},
+ }, nil)
+
+ tbl, err := array.TableFromJSON(t.mem, sc, []string{`[
+ {"c0": 1, "c1": "a", "c2": [1]},
+ {"c0": 2, "c1": "b", "c2": [1, 2]},
+ {"c0": 3, "c1": "c", "c2": [null]},
+ {"c0": null, "c1": "d", "c2": []},
+ {"c0": 5, "c1": null, "c2": [3, 3, 3]},
+ {"c0": 6, "c1": "f", "c2": null}
+ ]`})
+ t.Require().NoError(err)
+ defer tbl.Release()
+
+ t.writeFile(props, tbl)
+ t.readPageIndexes(2, 1, []int{})
+
+ t.Equal(t.columnIndexes, []ColumnIndexObject{
+ {
+ nullPages: []bool{false},
+ minValues: [][]byte{simpleEncode(int64(1))},
+ maxValues: [][]byte{simpleEncode(int64(3))},
+ boundaryOrder: metadata.Ascending, nullCounts: []int64{1},
+ },
+ {
+ nullPages: []bool{false},
+ minValues: [][]byte{[]byte("a")},
+ maxValues: [][]byte{[]byte("d")},
+ boundaryOrder: metadata.Ascending, nullCounts: []int64{0},
+ },
+ {
+ nullPages: []bool{false},
+ minValues: [][]byte{simpleEncode(int64(1))},
+ maxValues: [][]byte{simpleEncode(int64(2))},
+ boundaryOrder: metadata.Ascending, nullCounts: []int64{2},
+ },
+ {
+ nullPages: []bool{false},
+ minValues: [][]byte{simpleEncode(int64(5))},
+ maxValues: [][]byte{simpleEncode(int64(6))},
+ boundaryOrder: metadata.Ascending, nullCounts: []int64{0},
+ },
+ {
+ nullPages: []bool{false},
+ minValues: [][]byte{[]byte("f")},
+ maxValues: [][]byte{[]byte("f")},
+ boundaryOrder: metadata.Ascending, nullCounts: []int64{1},
+ },
+ {
+ nullPages: []bool{false},
+ minValues: [][]byte{simpleEncode(int64(3))},
+ maxValues: [][]byte{simpleEncode(int64(3))},
+ boundaryOrder: metadata.Ascending, nullCounts: []int64{1},
+ },
+ })
+}
+
+func (t *PageIndexRoundTripSuite) TestRoundTripWithStatsDisabled() {
+ props := parquet.NewWriterProperties(parquet.WithAllocator(t.mem),
+ parquet.WithDataPageVersion(t.pageVersion),
+ parquet.WithPageIndexEnabled(true), parquet.WithStats(false))
+
+ sc := arrow.NewSchema([]arrow.Field{
+ {Name: "c0", Type: arrow.PrimitiveTypes.Int64, Nullable: true},
+ {Name: "c1", Type: arrow.BinaryTypes.String, Nullable: true},
+ {Name: "c2", Type: arrow.ListOf(arrow.PrimitiveTypes.Int64), Nullable: true},
+ }, nil)
+
+ tbl, err := array.TableFromJSON(t.mem, sc, []string{`[
+ {"c0": 1, "c1": "a", "c2": [1]},
+ {"c0": 2, "c1": "b", "c2": [1, 2]},
+ {"c0": 3, "c1": "c", "c2": [null]},
+ {"c0": null, "c1": "d", "c2": []},
+ {"c0": 5, "c1": null, "c2": [3, 3, 3]},
+ {"c0": 6, "c1": "f", "c2": null}
+ ]`})
+ t.Require().NoError(err)
+ defer tbl.Release()
+
+ t.writeFile(props, tbl)
+ t.readPageIndexes(1, 1, []int{})
+
+ for _, colidx := range t.columnIndexes {
+ t.Zero(colidx)
+ }
+}
+
+func (t *PageIndexRoundTripSuite) TestRoundTripWithColumnStatsDisabled() {
+ props := parquet.NewWriterProperties(parquet.WithAllocator(t.mem),
+ parquet.WithDataPageVersion(t.pageVersion),
+ parquet.WithPageIndexEnabled(true), parquet.WithStatsFor("c0", false),
+ parquet.WithMaxRowGroupLength(4))
+
+ sc := arrow.NewSchema([]arrow.Field{
+ {Name: "c0", Type: arrow.PrimitiveTypes.Int64, Nullable: true},
+ {Name: "c1", Type: arrow.BinaryTypes.String, Nullable: true},
+ {Name: "c2", Type: arrow.ListOf(arrow.PrimitiveTypes.Int64), Nullable: true},
+ }, nil)
+
+ tbl, err := array.TableFromJSON(t.mem, sc, []string{`[
+ {"c0": 1, "c1": "a", "c2": [1]},
+ {"c0": 2, "c1": "b", "c2": [1, 2]},
+ {"c0": 3, "c1": "c", "c2": [null]},
+ {"c0": null, "c1": "d", "c2": []},
+ {"c0": 5, "c1": null, "c2": [3, 3, 3]},
+ {"c0": 6, "c1": "f", "c2": null}
+ ]`})
+ t.Require().NoError(err)
+ defer tbl.Release()
+
+ t.writeFile(props, tbl)
+ t.readPageIndexes(2, 1, []int{})
+
+ var emptyColIndex ColumnIndexObject
+ t.Equal(t.columnIndexes, []ColumnIndexObject{
+ emptyColIndex,
+ {
+ nullPages: []bool{false},
+ minValues: [][]byte{[]byte("a")},
+ maxValues: [][]byte{[]byte("d")},
+ boundaryOrder: metadata.Ascending, nullCounts: []int64{0},
+ },
+ {
+ nullPages: []bool{false},
+ minValues: [][]byte{simpleEncode(int64(1))},
+ maxValues: [][]byte{simpleEncode(int64(2))},
+ boundaryOrder: metadata.Ascending, nullCounts: []int64{2},
+ },
+ emptyColIndex,
+ {
+ nullPages: []bool{false},
+ minValues: [][]byte{[]byte("f")},
+ maxValues: [][]byte{[]byte("f")},
+ boundaryOrder: metadata.Ascending, nullCounts: []int64{1},
+ },
+ {
+ nullPages: []bool{false},
+ minValues: [][]byte{simpleEncode(int64(3))},
+ maxValues: [][]byte{simpleEncode(int64(3))},
+ boundaryOrder: metadata.Ascending, nullCounts: []int64{1},
+ },
+ })
+}
+
+func (t *PageIndexRoundTripSuite) TestDropLargeStats() {
+ props := parquet.NewWriterProperties(parquet.WithAllocator(t.mem),
+ parquet.WithDataPageVersion(t.pageVersion),
+ parquet.WithPageIndexEnabled(true), parquet.WithMaxRowGroupLength(1),
+ parquet.WithMaxStatsSize(20))
+
+ sc := arrow.NewSchema([]arrow.Field{
+ {Name: "c0", Type: arrow.BinaryTypes.String, Nullable: true},
+ }, nil)
+
+ tbl, err := array.TableFromJSON(t.mem, sc, []string{`[
+ {"c0": "short_string"},
+ {"c0": "very_large_string_to_drop_stats"}
+ ]`})
+ t.Require().NoError(err)
+ defer tbl.Release()
+
+ t.writeFile(props, tbl)
+ t.readPageIndexes(2, 1, []int{})
+
+ t.Equal(t.columnIndexes, []ColumnIndexObject{
+ {
+ nullPages: []bool{false},
+ minValues: [][]byte{[]byte("short_string")},
+ maxValues: [][]byte{[]byte("short_string")},
+ boundaryOrder: metadata.Ascending, nullCounts: []int64{0},
+ },
+ {},
+ })
+}
+
+func (t *PageIndexRoundTripSuite) TestMultiplePages() {
+ props := parquet.NewWriterProperties(parquet.WithAllocator(t.mem),
+ parquet.WithDataPageVersion(t.pageVersion),
+ parquet.WithPageIndexEnabled(true), parquet.WithDataPageSize(1))
+
+ sc := arrow.NewSchema([]arrow.Field{
+ {Name: "c0", Type: arrow.PrimitiveTypes.Int64, Nullable: true},
+ {Name: "c1", Type: arrow.BinaryTypes.String, Nullable: true},
+ }, nil)
+
+ tbl, err := array.TableFromJSON(t.mem, sc, []string{
+ `[{"c0": 1, "c1": "a"},{"c0": 2, "c1": "b"}]`,
+ `[{"c0": 3, "c1": "c"}, {"c0": 4, "c1": "d"}]`,
+ `[{"c0": null, "c1": null}, {"c0": 6, "c1": "f"}]`,
+ `[{"c0": null, "c1": null}, {"c0": null, "c1": null}]`,
+ })
+ t.Require().NoError(err)
+ defer tbl.Release()
+
+ t.writeFile(props, tbl)
+ t.readPageIndexes(1, 4, []int{})
+
+ t.Equal(t.columnIndexes, []ColumnIndexObject{
+ {
+ nullPages: []bool{false, false, false, true},
+ minValues: [][]byte{simpleEncode(int64(1)), simpleEncode(int64(3)),
+ simpleEncode(int64(6)), {}},
+ maxValues: [][]byte{simpleEncode(int64(2)), simpleEncode(int64(4)),
+ simpleEncode(int64(6)), {}},
+ boundaryOrder: metadata.Ascending, nullCounts: []int64{0, 0, 1, 2},
+ },
+ {
+ nullPages: []bool{false, false, false, true},
+ minValues: [][]byte{[]byte("a"), []byte("c"), []byte("f"), {}},
+ maxValues: [][]byte{[]byte("b"), []byte("d"), []byte("f"), {}},
+ boundaryOrder: metadata.Ascending, nullCounts: []int64{0, 0, 1, 2},
+ },
+ })
+}
+
+func (t *PageIndexRoundTripSuite) TestDoubleWithNaNs() {
+ props := parquet.NewWriterProperties(parquet.WithAllocator(t.mem),
+ parquet.WithDataPageVersion(t.pageVersion),
+ parquet.WithPageIndexEnabled(true), parquet.WithMaxRowGroupLength(3))
+
+ sc := arrow.NewSchema([]arrow.Field{
+ {Name: "c0", Type: arrow.PrimitiveTypes.Float64, Nullable: true},
+ }, nil)
+
+ bldr := array.NewFloat64Builder(t.mem)
+ defer bldr.Release()
+
+ chunks := make([]arrow.Array, 4)
+ defer func() {
+ for _, c := range chunks {
+ if c != nil {
+ c.Release()
+ }
+ }
+ }()
+
+ nan := math.NaN()
+ bldr.AppendValues([]float64{1, nan, 0.1}, nil)
+ chunks[0] = bldr.NewArray()
+ bldr.AppendValues([]float64{0, nan, 0}, nil)
+ chunks[1] = bldr.NewArray()
+ bldr.AppendValues([]float64{math.Copysign(0, -1), nan, math.Copysign(0, -1)}, nil)
+ chunks[2] = bldr.NewArray()
+ bldr.AppendValues([]float64{nan, nan, nan}, nil)
+ chunks[3] = bldr.NewArray()
+
+ tbl := array.NewTableFromSlice(sc, [][]arrow.Array{chunks})
+ defer tbl.Release()
+
+ t.writeFile(props, tbl)
+ t.readPageIndexes(4, 1, []int{})
+
+ t.Equal(t.columnIndexes, []ColumnIndexObject{
+ {
+ nullPages: []bool{false},
+ minValues: [][]byte{simpleEncode(float64(0.1))},
+ maxValues: [][]byte{simpleEncode(float64(1))},
+ boundaryOrder: metadata.Ascending, nullCounts: []int64{0},
+ },
+ {
+ nullPages: []bool{false},
+ minValues: [][]byte{simpleEncode(math.Copysign(0, -1))},
+ maxValues: [][]byte{simpleEncode(float64(0))},
+ boundaryOrder: metadata.Ascending, nullCounts: []int64{0},
+ },
+ {
+ nullPages: []bool{false},
+ minValues: [][]byte{simpleEncode(math.Copysign(0, -1))},
+ maxValues: [][]byte{simpleEncode(float64(0))},
+ boundaryOrder: metadata.Ascending, nullCounts: []int64{0},
+ },
+ {}, // page with only NaN values does not have column index built
+ })
+}
+
+func (t *PageIndexRoundTripSuite) TestEnablePerColumn() {
+ props := parquet.NewWriterProperties(parquet.WithAllocator(t.mem),
+ parquet.WithDataPageVersion(t.pageVersion),
+ parquet.WithPageIndexEnabled(true),
+ parquet.WithPageIndexEnabledFor("c0", true),
+ parquet.WithPageIndexEnabledFor("c1", false))
+
+ sc := arrow.NewSchema([]arrow.Field{
+ {Name: "c0", Type: arrow.PrimitiveTypes.Int64, Nullable: true},
+ {Name: "c1", Type: arrow.PrimitiveTypes.Int64, Nullable: true},
+ {Name: "c2", Type: arrow.PrimitiveTypes.Int64, Nullable: true},
+ }, nil)
+
+ tbl, err := array.TableFromJSON(t.mem, sc, []string{
+ `[{"c0": 0, "c1": 1, "c2": 2}]`,
+ })
+ t.Require().NoError(err)
+ defer tbl.Release()
+
+ t.writeFile(props, tbl)
+ t.readPageIndexes(1, 1, []int{1})
+
+ t.Equal(t.columnIndexes, []ColumnIndexObject{
+ {
+ nullPages: []bool{false},
+ minValues: [][]byte{simpleEncode(int64(0))},
+ maxValues: [][]byte{simpleEncode(int64(0))},
+ boundaryOrder: metadata.Ascending, nullCounts: []int64{0},
+ },
+ {}, // page index of c1 is disabled
+ {
+ nullPages: []bool{false},
+ minValues: [][]byte{simpleEncode(int64(2))},
+ maxValues: [][]byte{simpleEncode(int64(2))},
+ boundaryOrder: metadata.Ascending, nullCounts: []int64{0},
+ },
+ })
+}
+
+func TestPageIndexRoundTripSuite(t *testing.T) {
+ t.Run("datapagev1", func(t *testing.T) {
+ suite.Run(t, &PageIndexRoundTripSuite{pageVersion: parquet.DataPageV1})
+ })
+ t.Run("datapagev2", func(t *testing.T) {
+ suite.Run(t, &PageIndexRoundTripSuite{pageVersion: parquet.DataPageV2})
+ })
+}
diff --git a/parquet/file/page_reader.go b/parquet/file/page_reader.go
index f42dd38..2226b84 100644
--- a/parquet/file/page_reader.go
+++ b/parquet/file/page_reader.go
@@ -49,10 +49,19 @@ type PageReader interface {
Reset(r parquet.BufferedReader, nrows int64, compressType compress.Compression, ctx *CryptoContext)
}
+type PageType = format.PageType
+
+const (
+ PageTypeDataPage PageType = format.PageType_DATA_PAGE
+ PageTypeDataPageV2 PageType = format.PageType_DATA_PAGE_V2
+ PageTypeDictionaryPage PageType = format.PageType_DICTIONARY_PAGE
+ PageTypeIndexPage PageType = format.PageType_INDEX_PAGE
+)
+
// Page is an interface for handling DataPages or Dictionary Pages
type Page interface {
// Returns which kind of page this is
- Type() format.PageType
+ Type() PageType
// Get the raw bytes of this page
Data() []byte
// return the encoding used for this page, Plain/RLE, etc.
@@ -71,17 +80,31 @@ type page struct {
encoding format.Encoding
}
-func (p *page) Type() format.PageType { return p.typ }
+func (p *page) Type() PageType { return p.typ }
func (p *page) Data() []byte { return p.buf.Bytes() }
func (p *page) NumValues() int32 { return p.nvals }
func (p *page) Encoding() format.Encoding { return p.encoding }
+// DataPageConfig is a struct for passing configuration params to data page creation
+// which can be expanded in the future without causing any breaking changes.
+type DataPageConfig struct {
+ Num int32
+ Encoding parquet.Encoding
+ UncompressedSize int32
+ Stats metadata.EncodedStatistics
+ FirstRowIndex int64
+ SizeStats SizeStatistics
+}
+
// DataPage is the base interface for both DataPageV1 and DataPageV2 of the
// parquet spec.
type DataPage interface {
Page
UncompressedSize() int32
Statistics() metadata.EncodedStatistics
+ // FirstRowIndex returns the row ordinal within the row group
+ // to the first row in the data page, or -1 if not set.
+ FirstRowIndex() int64
}
// Create some pools to use for reusing the data page objects themselves so that
@@ -110,6 +133,8 @@ type DataPageV1 struct {
repLvlEncoding format.Encoding
uncompressedSize int32
statistics metadata.EncodedStatistics
+ firstRowIndex int64
+ sizeStatistics SizeStatistics
}
// NewDataPageV1 returns a V1 data page with the given buffer as its data and the specified encoding information
@@ -125,6 +150,7 @@ func NewDataPageV1(buffer *memory.Buffer, num int32, encoding, defEncoding, repE
defLvlEncoding: format.Encoding(defEncoding),
repLvlEncoding: format.Encoding(repEncoding),
uncompressedSize: uncompressedSize,
+ firstRowIndex: -1,
}
}
@@ -134,6 +160,8 @@ func NewDataPageV1(buffer *memory.Buffer, num int32, encoding, defEncoding, repE
dp.statistics.HasMax, dp.statistics.HasMin = false, false
dp.statistics.HasNullCount, dp.statistics.HasDistinctCount = false, false
dp.uncompressedSize = uncompressedSize
+ dp.firstRowIndex = -1
+ dp.sizeStatistics = SizeStatistics{}
return dp
}
@@ -144,6 +172,14 @@ func NewDataPageV1WithStats(buffer *memory.Buffer, num int32, encoding, defEncod
return ret
}
+// NewDataPageV1WithConfig uses a DataPageConfig object to encapsulate some parameters for future expansion
+// rather than continuing to add new functions to the API to avoid breaking changes.
+func NewDataPageV1WithConfig(buffer *memory.Buffer, defEncoding, repEncoding parquet.Encoding, cfg DataPageConfig) *DataPageV1 {
+ ret := NewDataPageV1WithStats(buffer, cfg.Num, cfg.Encoding, defEncoding, repEncoding, cfg.UncompressedSize, cfg.Stats)
+ ret.firstRowIndex, ret.sizeStatistics = cfg.FirstRowIndex, cfg.SizeStats
+ return ret
+}
+
// Release this page back into the DataPage object pool so that it can be reused.
//
// After calling this function, the object should not be utilized anymore, otherwise
@@ -154,6 +190,8 @@ func (d *DataPageV1) Release() {
dataPageV1Pool.Put(d)
}
+func (d *DataPageV1) FirstRowIndex() int64 { return d.firstRowIndex }
+
// UncompressedSize returns the size of the data in this data page when uncompressed
func (d *DataPageV1) UncompressedSize() int32 { return d.uncompressedSize }
@@ -181,6 +219,8 @@ type DataPageV2 struct {
compressed bool
uncompressedSize int32
statistics metadata.EncodedStatistics
+ firstRowIndex int64
+ sizeStatistics SizeStatistics
}
// NewDataPageV2 constructs a new V2 data page with the provided information and a buffer of the raw data.
@@ -195,6 +235,7 @@ func NewDataPageV2(buffer *memory.Buffer, numValues, numNulls, numRows int32, en
repLvlByteLen: repLvlsByteLen,
compressed: isCompressed,
uncompressedSize: uncompressed,
+ firstRowIndex: -1,
}
}
@@ -205,6 +246,8 @@ func NewDataPageV2(buffer *memory.Buffer, numValues, numNulls, numRows int32, en
dp.compressed, dp.uncompressedSize = isCompressed, uncompressed
dp.statistics.HasMax, dp.statistics.HasMin = false, false
dp.statistics.HasNullCount, dp.statistics.HasDistinctCount = false, false
+ dp.firstRowIndex = -1
+ dp.sizeStatistics = SizeStatistics{}
return dp
}
@@ -215,6 +258,14 @@ func NewDataPageV2WithStats(buffer *memory.Buffer, numValues, numNulls, numRows
return ret
}
+// NewDataPageV2WithConfig uses a DataPageConfig object to encapsulate some parameters for future expansion
+// rather than continuing to add new functions to the API to avoid breaking changes.
+func NewDataPageV2WithConfig(buffer *memory.Buffer, numNulls, numRows int32, defLvlsByteLen, repLvlsByteLen int32, isCompressed bool, cfg DataPageConfig) *DataPageV2 {
+ ret := NewDataPageV2WithStats(buffer, cfg.Num, numNulls, numRows, cfg.Encoding, defLvlsByteLen, repLvlsByteLen, cfg.UncompressedSize, isCompressed, cfg.Stats)
+ ret.firstRowIndex, ret.sizeStatistics = cfg.FirstRowIndex, cfg.SizeStats
+ return ret
+}
+
// Release this page back into the DataPage object pool so that it can be reused.
//
// After calling this function, the object should not be utilized anymore, otherwise
@@ -225,6 +276,8 @@ func (d *DataPageV2) Release() {
dataPageV2Pool.Put(d)
}
+func (d *DataPageV2) FirstRowIndex() int64 { return d.firstRowIndex }
+
// UncompressedSize is the size of the raw page when uncompressed. If `IsCompressed` is true, then
// the raw data in the buffer is expected to be compressed.
func (d *DataPageV2) UncompressedSize() int32 { return d.uncompressedSize }
diff --git a/parquet/file/page_writer.go b/parquet/file/page_writer.go
index 5576a68..1718292 100644
--- a/parquet/file/page_writer.go
+++ b/parquet/file/page_writer.go
@@ -18,6 +18,8 @@ package file
import (
"bytes"
+ "fmt"
+ "math"
"sync"
"github.com/apache/arrow-go/v18/arrow/memory"
@@ -51,6 +53,8 @@ type PageWriter interface {
// Allow reuse of the pagewriter object by resetting it using these values instead
// of having to create a new object.
Reset(sink utils.WriterTell, codec compress.Compression, compressionLevel int, metadata *metadata.ColumnChunkMetaDataBuilder, rgOrdinal, columnOrdinal int16, metaEncryptor, dataEncryptor encryption.Encryptor) error
+
+ SetIndexBuilders(metadata.ColumnIndexBuilder, *metadata.OffsetIndexBuilder)
}
type serializedPageWriter struct {
@@ -79,7 +83,9 @@ type serializedPageWriter struct {
dictEncodingStats map[parquet.Encoding]int32
dataEncodingStats map[parquet.Encoding]int32
- thriftSerializer *thrift.Serializer
+ columnIndexBuilder metadata.ColumnIndexBuilder
+ offsetIndexBuilder *metadata.OffsetIndexBuilder
+ thriftSerializer *thrift.Serializer
}
func createSerializedPageWriter(sink utils.WriterTell, codec compress.Compression, compressionLevel int, metadata *metadata.ColumnChunkMetaDataBuilder, rowGroupOrdinal, columnChunkOrdinal int16, mem memory.Allocator, metaEncryptor, dataEncryptor encryption.Encryptor) (PageWriter, error) {
@@ -122,6 +128,11 @@ func NewPageWriter(sink utils.WriterTell, codec compress.Compression, compressio
return createSerializedPageWriter(sink, codec, compressionLevel, metadata, rowGroupOrdinal, columnChunkOrdinal, mem, metaEncryptor, dataEncryptor)
}
+func (pw *serializedPageWriter) SetIndexBuilders(colIdxBldr metadata.ColumnIndexBuilder, offsetIdxBldr *metadata.OffsetIndexBuilder) {
+ pw.columnIndexBuilder = colIdxBldr
+ pw.offsetIndexBuilder = offsetIdxBldr
+}
+
// Reset allows reusing the pagewriter object instead of creating a new one.
func (pw *serializedPageWriter) Reset(sink utils.WriterTell, codec compress.Compression, compressionLevel int, metadata *metadata.ColumnChunkMetaDataBuilder, rowGroupOrdinal, columnChunkOrdinal int16, metaEncryptor, dataEncryptor encryption.Encryptor) error {
var (
@@ -156,6 +167,9 @@ func (pw *serializedPageWriter) Reset(sink utils.WriterTell, codec compress.Comp
if metaEncryptor != nil || dataEncryptor != nil {
pw.initEncryption()
}
+
+ pw.columnIndexBuilder, pw.offsetIndexBuilder = nil, nil
+
return nil
}
@@ -205,6 +219,7 @@ func (pw *serializedPageWriter) Close(hasDict, fallback bool) error {
DictEncodingStats: pw.dictEncodingStats,
DataEncodingStats: pw.dataEncodingStats,
}
+ pw.FinishPageIndexes(0)
pw.metaData.Finish(chunkInfo, hasDict, fallback, encodingStats, pw.metaEncryptor)
_, err := pw.metaData.WriteTo(pw.sink)
return err
@@ -227,7 +242,11 @@ func (pw *serializedPageWriter) setDataPageHeader(pageHdr *format.PageHeader, pa
hdr.Encoding = page.encoding
hdr.DefinitionLevelEncoding = page.defLvlEncoding
hdr.RepetitionLevelEncoding = page.repLvlEncoding
- hdr.Statistics = page.statistics.ToThrift()
+ if pw.columnIndexBuilder == nil {
+ hdr.Statistics = page.statistics.ToThrift()
+ } else {
+ hdr.Statistics = nil
+ }
pageHdr.DataPageHeader = hdr
pageHdr.DataPageHeaderV2 = nil
pageHdr.DictionaryPageHeader = nil
@@ -247,7 +266,11 @@ func (pw *serializedPageWriter) setDataPageV2Header(pageHdr *format.PageHeader,
hdr.DefinitionLevelsByteLength = page.defLvlByteLen
hdr.RepetitionLevelsByteLength = page.repLvlByteLen
hdr.IsCompressed = page.compressed
- hdr.Statistics = page.statistics.ToThrift()
+ if pw.columnIndexBuilder == nil {
+ hdr.Statistics = page.statistics.ToThrift()
+ } else {
+ hdr.Statistics = nil
+ }
pageHdr.DataPageHeaderV2 = hdr
pageHdr.DataPageHeader = nil
pageHdr.DictionaryPageHeader = nil
@@ -378,6 +401,29 @@ func (pw *serializedPageWriter) WriteDataPage(page DataPage) (int64, error) {
}
written += headerSize
+ // collect page index
+ if pw.columnIndexBuilder != nil {
+ stats := page.Statistics()
+ pw.columnIndexBuilder.AddPage(&stats)
+ }
+
+ if pw.offsetIndexBuilder != nil {
+ if written > math.MaxInt32 {
+ return int64(written), fmt.Errorf("parquet: compressed page size %d overflows INT32_MAX", written)
+ }
+
+ if page.FirstRowIndex() == -1 {
+ return int64(written), fmt.Errorf("parquet: first row index is not set in data page for offset index")
+ }
+
+ // startPos is a relative offset in the buffered mode, it should be
+ // adjusted via OffsetIndexBuilder.Finish after BufferedPageWriter
+ // has flushed all data pages
+ if err = pw.offsetIndexBuilder.AddPage(startPos, page.FirstRowIndex(), int32(written)); err != nil {
+ return int64(written), err
+ }
+ }
+
pw.totalUncompressed += int64(uncompressed) + int64(headerSize)
pw.totalCompressed += int64(written)
pw.nvalues += int64(page.NumValues())
@@ -386,6 +432,15 @@ func (pw *serializedPageWriter) WriteDataPage(page DataPage) (int64, error) {
return int64(written), nil
}
+func (pw *serializedPageWriter) FinishPageIndexes(finalPos int64) {
+ if pw.columnIndexBuilder != nil {
+ pw.columnIndexBuilder.Finish()
+ }
+ if pw.offsetIndexBuilder != nil {
+ pw.offsetIndexBuilder.Finish(finalPos)
+ }
+}
+
type bufferedPageWriter struct {
finalSink utils.WriterTell
inMemSink *encoding.BufferWriter
@@ -409,6 +464,10 @@ func newBufferedPageWriter(sink utils.WriterTell, codec compress.Compression, co
return wr, nil
}
+func (bw *bufferedPageWriter) SetIndexBuilders(colIdxBldr metadata.ColumnIndexBuilder, offsetIdxBldr *metadata.OffsetIndexBuilder) {
+ bw.pager.SetIndexBuilders(colIdxBldr, offsetIdxBldr)
+}
+
func (bw *bufferedPageWriter) Reset(sink utils.WriterTell, codec compress.Compression, compressionLevel int, metadata *metadata.ColumnChunkMetaDataBuilder, rgOrdinal, columnOrdinal int16, metaEncryptor, dataEncryptor encryption.Encryptor) error {
bw.finalSink = sink
bw.metadata = metadata
@@ -447,6 +506,8 @@ func (bw *bufferedPageWriter) Close(hasDict, fallback bool) error {
DataEncodingStats: bw.pager.dataEncodingStats,
}
bw.metadata.Finish(chunkInfo, hasDict, fallback, encodingStats, bw.pager.metaEncryptor)
+ bw.pager.FinishPageIndexes(position)
+
bw.metadata.WriteTo(bw.inMemSink)
buf := bw.inMemSink.Finish()
diff --git a/parquet/file/row_group_writer.go b/parquet/file/row_group_writer.go
index 4556d46..c335a10 100644
--- a/parquet/file/row_group_writer.go
+++ b/parquet/file/row_group_writer.go
@@ -66,29 +66,31 @@ type BufferedRowGroupWriter interface {
}
type rowGroupWriter struct {
- sink utils.WriterTell
- metadata *metadata.RowGroupMetaDataBuilder
- props *parquet.WriterProperties
- bytesWritten int64
- closed bool
- ordinal int16
- nextColumnIdx int
- nrows int
- buffered bool
- fileEncryptor encryption.FileEncryptor
+ sink utils.WriterTell
+ metadata *metadata.RowGroupMetaDataBuilder
+ props *parquet.WriterProperties
+ bytesWritten int64
+ closed bool
+ ordinal int16
+ nextColumnIdx int
+ nrows int
+ buffered bool
+ fileEncryptor encryption.FileEncryptor
+ pageIndexBuilder *metadata.PageIndexBuilder
columnWriters []ColumnChunkWriter
pager PageWriter
}
-func newRowGroupWriter(sink utils.WriterTell, metadata *metadata.RowGroupMetaDataBuilder, ordinal int16, props *parquet.WriterProperties, buffered bool, fileEncryptor encryption.FileEncryptor) *rowGroupWriter {
+func newRowGroupWriter(sink utils.WriterTell, metadata *metadata.RowGroupMetaDataBuilder, ordinal int16, props *parquet.WriterProperties, buffered bool, fileEncryptor encryption.FileEncryptor, pageIdxBldr *metadata.PageIndexBuilder) *rowGroupWriter {
ret := &rowGroupWriter{
- sink: sink,
- metadata: metadata,
- props: props,
- ordinal: ordinal,
- buffered: buffered,
- fileEncryptor: fileEncryptor,
+ sink: sink,
+ metadata: metadata,
+ props: props,
+ ordinal: ordinal,
+ buffered: buffered,
+ fileEncryptor: fileEncryptor,
+ pageIndexBuilder: pageIdxBldr,
}
if buffered {
ret.initColumns()
@@ -152,22 +154,36 @@ func (rg *rowGroupWriter) NextColumn() (ColumnChunkWriter, error) {
path := colMeta.Descr().Path()
var (
+ columnOrdinal = rg.nextColumnIdx - 1
metaEncryptor encryption.Encryptor
dataEncryptor encryption.Encryptor
+ colIdxBldr metadata.ColumnIndexBuilder
+ offsetIdxBldr *metadata.OffsetIndexBuilder
)
if rg.fileEncryptor != nil {
metaEncryptor = rg.fileEncryptor.GetColumnMetaEncryptor(path)
dataEncryptor = rg.fileEncryptor.GetColumnDataEncryptor(path)
}
+ if rg.pageIndexBuilder != nil && rg.props.PageIndexEnabledFor(path) {
+ var err error
+ if colIdxBldr, err = rg.pageIndexBuilder.GetColumnIndexBuilder(columnOrdinal); err != nil {
+ return nil, err
+ }
+ if offsetIdxBldr, err = rg.pageIndexBuilder.GetOffsetIndexBuilder(columnOrdinal); err != nil {
+ return nil, err
+ }
+ }
if rg.pager == nil {
var err error
- rg.pager, err = NewPageWriter(rg.sink, rg.props.CompressionFor(path), rg.props.CompressionLevelFor(path), colMeta, rg.ordinal, int16(rg.nextColumnIdx-1), rg.props.Allocator(), false, metaEncryptor, dataEncryptor)
+ rg.pager, err = NewPageWriter(rg.sink, rg.props.CompressionFor(path), rg.props.CompressionLevelFor(path), colMeta, rg.ordinal, int16(columnOrdinal), rg.props.Allocator(), false, metaEncryptor, dataEncryptor)
if err != nil {
return nil, err
}
+ rg.pager.SetIndexBuilders(colIdxBldr, offsetIdxBldr)
} else {
- rg.pager.Reset(rg.sink, rg.props.CompressionFor(path), rg.props.CompressionLevelFor(path), colMeta, rg.ordinal, int16(rg.nextColumnIdx-1), metaEncryptor, dataEncryptor)
+ rg.pager.Reset(rg.sink, rg.props.CompressionFor(path), rg.props.CompressionLevelFor(path), colMeta, rg.ordinal, int16(columnOrdinal), metaEncryptor, dataEncryptor)
+ rg.pager.SetIndexBuilders(colIdxBldr, offsetIdxBldr)
}
rg.columnWriters[0] = NewColumnChunkWriter(colMeta, rg.pager, rg.props)
@@ -239,15 +255,29 @@ func (rg *rowGroupWriter) initColumns() error {
var (
metaEncryptor encryption.Encryptor
dataEncryptor encryption.Encryptor
+ colIdxBldr metadata.ColumnIndexBuilder
+ offsetIdxBldr *metadata.OffsetIndexBuilder
)
if rg.fileEncryptor != nil {
metaEncryptor = rg.fileEncryptor.GetColumnMetaEncryptor(path)
dataEncryptor = rg.fileEncryptor.GetColumnDataEncryptor(path)
}
+ if rg.pageIndexBuilder != nil && rg.props.PageIndexEnabledFor(path) {
+ var err error
+ if colIdxBldr, err = rg.pageIndexBuilder.GetColumnIndexBuilder(rg.nextColumnIdx); err != nil {
+ return err
+ }
+ if offsetIdxBldr, err = rg.pageIndexBuilder.GetOffsetIndexBuilder(rg.nextColumnIdx); err != nil {
+ return err
+ }
+ }
+
pager, err := NewPageWriter(rg.sink, rg.props.CompressionFor(path), rg.props.CompressionLevelFor(path), colMeta, rg.ordinal, int16(rg.nextColumnIdx), rg.props.Allocator(), rg.buffered, metaEncryptor, dataEncryptor)
if err != nil {
return err
}
+ pager.SetIndexBuilders(colIdxBldr, offsetIdxBldr)
+
rg.nextColumnIdx++
rg.columnWriters = append(rg.columnWriters, NewColumnChunkWriter(colMeta, pager, rg.props))
}
diff --git a/parquet/file/size_statistics.go b/parquet/file/size_statistics.go
new file mode 100644
index 0000000..74c5210
--- /dev/null
+++ b/parquet/file/size_statistics.go
@@ -0,0 +1,28 @@
+// 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 (
+ format "github.com/apache/arrow-go/v18/parquet/internal/gen-go/parquet"
+)
+
+type SizeStatistics format.SizeStatistics
+
+func (s *SizeStatistics) IsSet() bool {
+ return len(s.RepetitionLevelHistogram) > 0 || len(s.DefinitionLevelHistogram) > 0 ||
+ s.UnencodedByteArrayDataBytes != nil
+}
diff --git a/parquet/internal/encryption/decryptor.go b/parquet/internal/encryption/decryptor.go
index 3bb43ac..50a1a6f 100644
--- a/parquet/internal/encryption/decryptor.go
+++ b/parquet/internal/encryption/decryptor.go
@@ -17,10 +17,14 @@
package encryption
import (
+ "fmt"
"io"
+ "github.com/apache/arrow-go/v18/arrow"
"github.com/apache/arrow-go/v18/arrow/memory"
"github.com/apache/arrow-go/v18/parquet"
+ "github.com/apache/arrow-go/v18/parquet/internal/debug"
+ format "github.com/apache/arrow-go/v18/parquet/internal/gen-go/parquet"
)
// FileDecryptor is an interface used by the filereader for decrypting an
@@ -266,3 +270,37 @@ func (d *decryptor) Decrypt(src []byte) []byte {
func (d *decryptor) DecryptFrom(r io.Reader) []byte {
return d.decryptor.DecryptFrom(r, d.key, d.aad)
}
+
+func getColumnDecryptor(cryptoMetadata *format.ColumnCryptoMetaData, fileDecryptor FileDecryptor, metadata bool) (Decryptor, error) {
+ if cryptoMetadata == nil {
+ return nil, nil
+ }
+
+ if fileDecryptor == nil {
+ return nil, fmt.Errorf("%w: row group is noted as encrypted but no file decryptor", arrow.ErrNotFound)
+ }
+
+ if cryptoMetadata.IsSetENCRYPTION_WITH_FOOTER_KEY() {
+ if metadata {
+ return fileDecryptor.GetFooterDecryptorForColumnMeta(fileDecryptor.FileAad()), nil
+ }
+ return fileDecryptor.GetFooterDecryptorForColumnData(fileDecryptor.FileAad()), nil
+ }
+
+ // column is encrypted with its own key
+ columnKeyMetadata := cryptoMetadata.ENCRYPTION_WITH_COLUMN_KEY.KeyMetadata
+ colPath := parquet.ColumnPath(cryptoMetadata.ENCRYPTION_WITH_COLUMN_KEY.PathInSchema).String()
+ return fileDecryptor.GetColumnMetaDecryptor(colPath, string(columnKeyMetadata), ""), nil
+}
+
+func GetColumnMetaDecryptor(cryptoMetadata *format.ColumnCryptoMetaData, fileDecryptor FileDecryptor) (Decryptor, error) {
+ return getColumnDecryptor(cryptoMetadata, fileDecryptor, true)
+}
+
+const NonPageOrdinal int16 = -1
+
+func UpdateDecryptor(decryptor Decryptor, rgOrdinal, colOrdinal int16, moduleType int8) {
+ debug.Assert(decryptor.FileAad() != "", "file decryptor has no file aad")
+ aad := CreateModuleAad(decryptor.FileAad(), moduleType, rgOrdinal, colOrdinal, NonPageOrdinal)
+ decryptor.UpdateAad(aad)
+}
diff --git a/parquet/metadata/column_chunk.go b/parquet/metadata/column_chunk.go
index 5a79c9b..22c9b9e 100644
--- a/parquet/metadata/column_chunk.go
+++ b/parquet/metadata/column_chunk.go
@@ -59,6 +59,13 @@ func makeColumnStats(metadata *format.ColumnMetaData, descr *schema.Column, mem
metadata.Statistics)
}
+type IndexLocation struct {
+ // file offset of the given index, in bytes
+ Offset int64
+ // length of the given index, in bytes
+ Length int32
+}
+
// ColumnChunkMetaData is a proxy around format.ColumnChunkMetaData
// containing all of the information and metadata for a given column chunk
// and it's associated Column
@@ -174,6 +181,26 @@ func (c *ColumnChunkMetaData) HasIndexPage() bool { return c.columnMeta.IsSetInd
// IndexPageOffset is the location in the file where the index page starts.
func (c *ColumnChunkMetaData) IndexPageOffset() int64 { return c.columnMeta.GetIndexPageOffset() }
+func (c *ColumnChunkMetaData) GetColumnIndexLocation() *IndexLocation {
+ if c.column.IsSetColumnIndexOffset() {
+ return &IndexLocation{
+ Offset: c.column.GetColumnIndexOffset(),
+ Length: c.column.GetColumnIndexLength(),
+ }
+ }
+ return nil
+}
+
+func (c *ColumnChunkMetaData) GetOffsetIndexLocation() *IndexLocation {
+ if c.column.IsSetOffsetIndexOffset() {
+ return &IndexLocation{
+ Offset: c.column.GetOffsetIndexOffset(),
+ Length: c.column.GetOffsetIndexLength(),
+ }
+ }
+ return nil
+}
+
// TotalCompressedSize will be equal to TotalUncompressedSize if the data is not compressed.
// Otherwise this will be the size of the actual data in the file.
func (c *ColumnChunkMetaData) TotalCompressedSize() int64 {
@@ -432,3 +459,7 @@ func (c *ColumnChunkMetaDataBuilder) Finish(info ChunkMetaInfo, hasDict, dictFal
func (c *ColumnChunkMetaDataBuilder) WriteTo(w io.Writer) (int64, error) {
return 0, thrift.SerializeThriftStream(c.chunk, w)
}
+
+type PageIndexLocation struct {
+ ColIndexLocation, OffsetIndexLocation map[uint64][]*IndexLocation
+}
diff --git a/parquet/metadata/file.go b/parquet/metadata/file.go
index c526a85..39a3192 100644
--- a/parquet/metadata/file.go
+++ b/parquet/metadata/file.go
@@ -173,6 +173,42 @@ func (f *FileMetaDataBuilder) Snapshot() (*FileMetaData, error) {
return out, nil
}
+func (f *FileMetaDataBuilder) SetPageIndexLocation(loc PageIndexLocation) error {
+ setIndexLoc := func(rgOrdinal int64, fileIdxLoc map[uint64][]*IndexLocation, colIndex bool) error {
+ rgMeta := f.rowGroups[rgOrdinal]
+ iter, ok := fileIdxLoc[uint64(rgOrdinal)]
+ if ok {
+ for i, idxLoc := range iter {
+ if i >= len(rgMeta.Columns) {
+ return fmt.Errorf("cannot find metadata for column ordinal %d", i)
+ }
+ colMeta := rgMeta.Columns[i]
+ if idxLoc != nil {
+ if colIndex {
+ colMeta.ColumnIndexOffset = &idxLoc.Offset
+ colMeta.ColumnIndexLength = &idxLoc.Length
+ } else {
+ colMeta.OffsetIndexOffset = &idxLoc.Offset
+ colMeta.OffsetIndexLength = &idxLoc.Length
+ }
+ }
+ }
+ }
+ return nil
+ }
+
+ for rgOrdinal := range f.rowGroups {
+ if err := setIndexLoc(int64(rgOrdinal), loc.ColIndexLocation, true); err != nil {
+ return err
+ }
+ if err := setIndexLoc(int64(rgOrdinal), loc.OffsetIndexLocation, false); err != nil {
+ return err
+ }
+ }
+
+ return nil
+}
+
// Clears out this filemetadatabuilder so it can be re-used
func (f *FileMetaDataBuilder) Clear() {
f.metadata = format.NewFileMetaData()
@@ -332,7 +368,7 @@ func (f *FileMetaData) initColumnOrders() {
orders = append(orders, parquet.ColumnOrders.Undefined)
}
}
- } else {
+ } else if f.Schema.NumColumns() > 0 {
orders = orders[:f.Schema.NumColumns()]
orders[0] = parquet.ColumnOrders.Undefined
for i := 1; i < len(orders); i *= 2 {
@@ -500,6 +536,9 @@ func (f *FileMetaData) Version() parquet.Version {
}
}
+func (f *FileMetaData) NumRowGroups() int { return len(f.RowGroups) }
+func (f *FileMetaData) NumColumns() int { return f.Schema.NumColumns() }
+
// FileCryptoMetadata is a proxy for the thrift fileCryptoMetadata object
type FileCryptoMetadata struct {
metadata *format.FileCryptoMetaData
diff --git a/parquet/metadata/metadata_test.go b/parquet/metadata/metadata_test.go
index 106307c..2f35d63 100644
--- a/parquet/metadata/metadata_test.go
+++ b/parquet/metadata/metadata_test.go
@@ -18,10 +18,13 @@ package metadata_test
import (
"context"
+ "os"
+ "path/filepath"
"testing"
"unsafe"
"github.com/apache/arrow-go/v18/parquet"
+ "github.com/apache/arrow-go/v18/parquet/file"
"github.com/apache/arrow-go/v18/parquet/metadata"
"github.com/apache/arrow-go/v18/parquet/schema"
"github.com/stretchr/testify/assert"
@@ -379,3 +382,49 @@ func TestCheckBadDecimalStats(t *testing.T) {
assert.False(t, version2.HasCorrectStatistics(parquet.Types.FixedLenByteArray, schema.NewDecimalLogicalType(5, 0), stats, schema.SortSIGNED))
assert.True(t, version3.HasCorrectStatistics(parquet.Types.FixedLenByteArray, schema.NewDecimalLogicalType(5, 0), stats, schema.SortSIGNED))
}
+
+func TestReadPageIndex(t *testing.T) {
+ dir := os.Getenv("PARQUET_TEST_DATA")
+ if dir == "" {
+ t.Skip("PARQUET_TEST_DATA not set")
+ }
+ require.DirExists(t, dir)
+
+ path := filepath.Join(dir, "alltypes_tiny_pages.parquet")
+ rdr, err := file.OpenParquetFile(path, false)
+ require.NoError(t, err)
+ defer rdr.Close()
+
+ fileMetadata := rdr.MetaData()
+ assert.Equal(t, 1, fileMetadata.NumRowGroups())
+ rgMeta := fileMetadata.RowGroup(0)
+ assert.Equal(t, 13, rgMeta.NumColumns())
+
+ ciOffsets := []int64{323583, 327502, 328009, 331928, 335847,
+ 339766, 350345, 354264, 364843, 384342, -1, 386473, 390392}
+ ciLengths := []int64{3919, 507, 3919, 3919, 3919, 10579, 3919,
+ 10579, 19499, 2131, -1, 3919, 3919}
+ oiOffsets := []int64{394311, 397814, 398637, 401888, 405139,
+ 408390, 413670, 416921, 422201, 431936, 435457, 446002, 449253}
+ oiLengths := []int64{3503, 823, 3251, 3251, 3251, 5280, 3251, 5280,
+ 9735, 3521, 10545, 3251, 3251}
+
+ for i := 0; i < rgMeta.NumColumns(); i++ {
+ colMeta, err := rgMeta.ColumnChunk(i)
+ require.NoError(t, err)
+
+ ciLoc := colMeta.GetColumnIndexLocation()
+ if i == 10 {
+ assert.Nil(t, ciLoc)
+ } else {
+ assert.NotNil(t, ciLoc)
+ assert.EqualValues(t, ciOffsets[i], ciLoc.Offset)
+ assert.EqualValues(t, ciLengths[i], ciLoc.Length)
+ }
+
+ oiLoc := colMeta.GetOffsetIndexLocation()
+ assert.NotNil(t, oiLoc)
+ assert.EqualValues(t, oiOffsets[i], oiLoc.Offset)
+ assert.EqualValues(t, oiLengths[i], oiLoc.Length)
+ }
+}
diff --git a/parquet/metadata/page_index.go b/parquet/metadata/page_index.go
new file mode 100644
index 0000000..5402544
--- /dev/null
+++ b/parquet/metadata/page_index.go
@@ -0,0 +1,928 @@
+// 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 metadata
+
+import (
+ "fmt"
+ "io"
+ "math"
+
+ "github.com/JohnCGriffin/overflow"
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/parquet"
+ "github.com/apache/arrow-go/v18/parquet/internal/debug"
+ "github.com/apache/arrow-go/v18/parquet/internal/encoding"
+ "github.com/apache/arrow-go/v18/parquet/internal/encryption"
+ format "github.com/apache/arrow-go/v18/parquet/internal/gen-go/parquet"
+ "github.com/apache/arrow-go/v18/parquet/internal/thrift"
+ "github.com/apache/arrow-go/v18/parquet/internal/utils"
+ "github.com/apache/arrow-go/v18/parquet/schema"
+)
+
+// BoundaryOrder identifies whether the min and max values are ordered and
+// if so, which direction it is stored in.
+type BoundaryOrder = format.BoundaryOrder
+
+const (
+ Unordered BoundaryOrder = format.BoundaryOrder_UNORDERED
+ Ascending BoundaryOrder = format.BoundaryOrder_ASCENDING
+ Descending BoundaryOrder = format.BoundaryOrder_DESCENDING
+)
+
+// ColumnIndex is an interface for reading optional statistics for
+// each data page of a column chunk. Along with the OffsetIndex this
+// forms the PageIndex for a column chunk.
+type ColumnIndex interface {
+ // GetNullPages returns a list of bools to determine the validity of the
+ // corresponding min/max values. If the value is true, then that page
+ // contains only null values.
+ GetNullPages() []bool
+ // IsSetNullCounts returns true if the null counts are set.
+ IsSetNullCounts() bool
+ // GetNullCounts returns the number of null values in each page. This is
+ // only valid if IsSetNullCounts returns true.
+ GetNullCounts() []int64
+ // GetBoundaryOrder returns if the min/max values are ordered and if so,
+ // what direction.
+ GetBoundaryOrder() BoundaryOrder
+ // GetMinValues returns the encoded minimum value for each page
+ GetMinValues() [][]byte
+ // GetMaxValues returns the encoded max value for each page
+ GetMaxValues() [][]byte
+}
+
+// TypedColumnIndex expands the ColumnIndex interface to provide a
+// type-safe accessor for the min/max values.
+type TypedColumnIndex[T parquet.ColumnTypes] struct {
+ ColumnIndex
+
+ minvals []T
+ maxvals []T
+ nonNullPageIndices []int32
+}
+
+type typedDecoder[T parquet.ColumnTypes] interface {
+ encoding.TypedDecoder
+ Decode([]T) (int, error)
+ DecodeSpaced([]T, int, []byte, int64) (int, error)
+}
+
+func must(err error) {
+ if err != nil {
+ panic(err)
+ }
+}
+
+func mustArg[T any](val T, err error) T {
+ must(err)
+ return val
+}
+
+// NewColumnIndex uses the thrift serialized bytes to deserialize a column index, optionally decrypting it.
+//
+// The column descriptor is used to determine the physical type of the column to create a proper
+// TypedColumnIndex.
+func NewColumnIndex(descr *schema.Column, serializedIndex []byte, props *parquet.ReaderProperties, decryptor encryption.Decryptor) ColumnIndex {
+ if decryptor != nil {
+ serializedIndex = decryptor.Decrypt(serializedIndex)
+ }
+
+ var colidx format.ColumnIndex
+ if _, err := thrift.DeserializeThrift(&colidx, serializedIndex); err != nil {
+ panic(err)
+ }
+
+ switch descr.PhysicalType() {
+ case parquet.Types.Boolean:
+ return newTypedColumnIndex[bool](descr, &colidx)
+ case parquet.Types.Int32:
+ return newTypedColumnIndex[int32](descr, &colidx)
+ case parquet.Types.Int64:
+ return newTypedColumnIndex[int64](descr, &colidx)
+ case parquet.Types.Int96:
+ return newTypedColumnIndex[parquet.Int96](descr, &colidx)
+ case parquet.Types.Float:
+ return newTypedColumnIndex[float32](descr, &colidx)
+ case parquet.Types.Double:
+ return newTypedColumnIndex[float64](descr, &colidx)
+ case parquet.Types.ByteArray:
+ return newTypedColumnIndex[parquet.ByteArray](descr, &colidx)
+ case parquet.Types.FixedLenByteArray:
+ return newTypedColumnIndex[parquet.FixedLenByteArray](descr, &colidx)
+ }
+
+ panic("unreachable: cannot make columnindex of unknown type")
+}
+
+func getDecoder[T parquet.ColumnTypes](descr *schema.Column) func([]byte) T {
+ switch descr.PhysicalType() {
+ case parquet.Types.ByteArray:
+ var f any = func(data []byte) parquet.ByteArray {
+ return parquet.ByteArray(data)
+ }
+ return f.(func([]byte) T)
+ case parquet.Types.FixedLenByteArray:
+ var f any = func(data []byte) parquet.FixedLenByteArray {
+ return parquet.FixedLenByteArray(data)
+ }
+ return f.(func([]byte) T)
+ default:
+ decoder := encoding.NewDecoder(descr.PhysicalType(), parquet.Encodings.Plain, descr, nil).(typedDecoder[T])
+ var buf [1]T
+ return func(data []byte) T {
+ must(decoder.SetData(1, data))
+ mustArg(decoder.Decode(buf[:]))
+ return buf[0]
+ }
+ }
+}
+
+func newTypedColumnIndex[T parquet.ColumnTypes](descr *schema.Column, colIdx *format.ColumnIndex) *TypedColumnIndex[T] {
+ numPages := len(colIdx.NullPages)
+ if numPages >= math.MaxInt32 ||
+ len(colIdx.MinValues) != numPages ||
+ len(colIdx.MaxValues) != numPages ||
+ (colIdx.IsSetNullCounts() && len(colIdx.NullCounts) != numPages) {
+ panic("invalid column index")
+ }
+
+ numNonNullPages := 0
+ for _, page := range colIdx.NullPages {
+ if !page {
+ numNonNullPages++
+ }
+ }
+ debug.Assert(numNonNullPages <= numPages, "invalid column index")
+
+ minvals, maxvals := make([]T, numPages), make([]T, numPages)
+ nonNullPageIndices := make([]int32, 0, numNonNullPages)
+
+ dec := getDecoder[T](descr)
+ for i := 0; i < numPages; i++ {
+ if !colIdx.NullPages[i] {
+ nonNullPageIndices = append(nonNullPageIndices, int32(i))
+ minvals[i] = dec(colIdx.MinValues[i])
+ maxvals[i] = dec(colIdx.MaxValues[i])
+ }
+ }
+ debug.Assert(len(nonNullPageIndices) == numNonNullPages, "invalid column index")
+
+ return &TypedColumnIndex[T]{
+ ColumnIndex: colIdx,
+ minvals: minvals,
+ maxvals: maxvals,
+ nonNullPageIndices: nonNullPageIndices,
+ }
+}
+
+func (idx *TypedColumnIndex[T]) MinValues() []T {
+ return idx.minvals
+}
+
+func (idx *TypedColumnIndex[T]) MaxValues() []T {
+ return idx.maxvals
+}
+
+func (idx *TypedColumnIndex[T]) NonNullPageIndices() []int32 {
+ return idx.nonNullPageIndices
+}
+
+type (
+ // PageLocation describes where in a file a particular page can be found,
+ // along with the index within the rowgroup of the first row in the page
+ PageLocation = format.PageLocation
+ PageIndexSelection struct {
+ // specifies whether to read the column index
+ ColumnIndex bool
+ // specifies whether to read the offset index
+ OffsetIndex bool
+ }
+)
+
+// OffsetIndex forms the page index alongside a ColumnIndex,
+// the OffsetIndex may be present even if a ColumnIndex is not.
+type OffsetIndex interface {
+ GetPageLocations() []*PageLocation
+}
+
+func (p PageIndexSelection) String() string {
+ return fmt.Sprintf("PageIndexSelection{column_index = %t, offset_index = %t}",
+ p.ColumnIndex, p.OffsetIndex)
+}
+
+// NewOffsetIndex constructs an OffsetIndex object from the thrift serialized bytes,
+// optionally decrypting it if it was encrypted.
+func NewOffsetIndex(serializedIndex []byte, _ *parquet.ReaderProperties, decryptor encryption.Decryptor) OffsetIndex {
+ if decryptor != nil {
+ serializedIndex = decryptor.Decrypt(serializedIndex)
+ }
+
+ var offsetIndex format.OffsetIndex
+ if _, err := thrift.DeserializeThrift(&offsetIndex, serializedIndex); err != nil {
+ panic(err)
+ }
+
+ return &offsetIndex
+}
+
+type readRange struct {
+ Offset, Length int64
+}
+
+func (r readRange) Contains(other readRange) bool {
+ return r.Offset <= other.Offset && other.Offset+other.Length <= r.Offset+r.Length
+}
+
+func checkReadRange(loc IndexLocation, idxRange *readRange, rgOrdinal int32) error {
+ if idxRange == nil {
+ return fmt.Errorf("%w: missing page index read range of row group %d, it may not exist or has not been requested",
+ arrow.ErrInvalid, rgOrdinal)
+ }
+
+ // coalesced read range is invalid
+ if idxRange.Offset < 0 || idxRange.Length <= 0 {
+ return fmt.Errorf("%w: invalid page index read range: offset %d, length %d",
+ arrow.ErrInvalid, idxRange.Offset, idxRange.Length)
+ }
+
+ if loc.Offset < 0 || loc.Length <= 0 {
+ return fmt.Errorf("%w: invalid page index location: offset %d, length %d",
+ arrow.ErrInvalid, loc.Offset, loc.Length)
+ }
+
+ if loc.Offset < idxRange.Offset || loc.Offset+int64(loc.Length) > idxRange.Offset+idxRange.Length {
+ return fmt.Errorf("%w: Page index location [offset:%d,length:%d] is out of range from previous WillNeed request [offset:%d,length:%d], row group: %d",
+ arrow.ErrInvalid, loc.Offset, loc.Length, idxRange.Offset, idxRange.Length, rgOrdinal)
+ }
+
+ return nil
+}
+
+type rgIndexReadRange struct {
+ ColIndex, OffsetIndex *readRange
+}
+
+// RowGroupPageIndexReader is a read-only object for retrieving column and offset
+// indexes for a given row group.
+type RowGroupPageIndexReader struct {
+ input parquet.ReaderAtSeeker
+ rowGroupMetadata *RowGroupMetaData
+ props *parquet.ReaderProperties
+ rgOrdinal int32
+ idxReadRange rgIndexReadRange
+ fileDecryptor encryption.FileDecryptor
+
+ // buffers to hold raw bytes of page index
+ // will be lazily set when the corresponding page index is accessed
+ colIndexBuffer, offsetIndexBuffer []byte
+}
+
+func (r *RowGroupPageIndexReader) GetColumnIndex(i int) (ColumnIndex, error) {
+ if i < 0 || i >= r.rowGroupMetadata.NumColumns() {
+ return nil, fmt.Errorf("%w: invalid column index at column ordinal %d",
+ arrow.ErrInvalid, i)
+ }
+
+ colChunk, err := r.rowGroupMetadata.ColumnChunk(i)
+ if err != nil {
+ return nil, err
+ }
+
+ colIndexLocation := colChunk.GetColumnIndexLocation()
+ if colIndexLocation == nil {
+ return nil, nil
+ }
+
+ if err := checkReadRange(*colIndexLocation, r.idxReadRange.ColIndex, r.rgOrdinal); err != nil {
+ return nil, err
+ }
+
+ if r.colIndexBuffer == nil {
+ r.colIndexBuffer = make([]byte, r.idxReadRange.ColIndex.Length)
+ if _, err := r.input.ReadAt(r.colIndexBuffer, r.idxReadRange.ColIndex.Offset); err != nil {
+ return nil, err
+ }
+ }
+
+ bufferOffset := colIndexLocation.Offset - r.idxReadRange.ColIndex.Offset
+ descr := r.rowGroupMetadata.Schema.Column(i)
+ decryptor, err := encryption.GetColumnMetaDecryptor(colChunk.CryptoMetadata(), r.fileDecryptor)
+ if err != nil {
+ return nil, err
+ }
+
+ if decryptor != nil {
+ encryption.UpdateDecryptor(decryptor, int16(r.rgOrdinal),
+ int16(i), encryption.ColumnIndexModule)
+ }
+
+ return NewColumnIndex(descr, r.colIndexBuffer[bufferOffset:], r.props, decryptor), nil
+}
+
+func (r *RowGroupPageIndexReader) GetOffsetIndex(i int) (OffsetIndex, error) {
+ if i < 0 || i >= r.rowGroupMetadata.NumColumns() {
+ return nil, fmt.Errorf("%w: invalid column index at column ordinal %d",
+ arrow.ErrInvalid, i)
+ }
+
+ colChunk, err := r.rowGroupMetadata.ColumnChunk(i)
+ if err != nil {
+ return nil, err
+ }
+
+ offsetIndexLocation := colChunk.GetOffsetIndexLocation()
+ if offsetIndexLocation == nil {
+ return nil, nil
+ }
+
+ if err := checkReadRange(*offsetIndexLocation, r.idxReadRange.OffsetIndex, r.rgOrdinal); err != nil {
+ return nil, err
+ }
+
+ if r.offsetIndexBuffer == nil {
+ r.offsetIndexBuffer = make([]byte, r.idxReadRange.OffsetIndex.Length)
+ if _, err := r.input.ReadAt(r.offsetIndexBuffer, r.idxReadRange.OffsetIndex.Offset); err != nil {
+ return nil, err
+ }
+ }
+
+ bufferOffset := offsetIndexLocation.Offset - r.idxReadRange.OffsetIndex.Offset
+ decryptor, err := encryption.GetColumnMetaDecryptor(colChunk.CryptoMetadata(), r.fileDecryptor)
+ if err != nil {
+ return nil, err
+ }
+
+ if decryptor != nil {
+ encryption.UpdateDecryptor(decryptor, int16(r.rgOrdinal),
+ int16(i), encryption.OffsetIndexModule)
+ }
+
+ return NewOffsetIndex(r.offsetIndexBuffer[bufferOffset:], r.props, decryptor), nil
+}
+
+// PageIndexReader is a read-only object for retrieving the Column and Offset indexes
+// for a particular parquet file.
+type PageIndexReader struct {
+ Input parquet.ReaderAtSeeker
+ FileMetadata *FileMetaData
+ Props *parquet.ReaderProperties
+ Decryptor encryption.FileDecryptor
+
+ // coalesced read ranges of page index of row groups that have
+ // been suggested by WillNeed(). key is the row group ordinal
+ idxReadRanges map[int32]rgIndexReadRange
+}
+
+func determinePageIndexRangesInRowGroup(rgMeta *RowGroupMetaData, cols []int32) (rng rgIndexReadRange, err error) {
+ ciStart, oiStart := int64(math.MaxInt64), int64(math.MaxInt64)
+ ciEnd, oiEnd := int64(-1), int64(-1)
+
+ mergeRange := func(idxLocation *IndexLocation, start, end *int64) error {
+ if idxLocation == nil {
+ return nil
+ }
+
+ indexEnd, ok := overflow.Add64(idxLocation.Offset, int64(idxLocation.Length))
+ if idxLocation.Offset < 0 || idxLocation.Length <= 0 || !ok {
+ return fmt.Errorf("%w: invalid page index location: offset %d length %d",
+ arrow.ErrIndex, idxLocation.Offset, idxLocation.Length)
+ }
+ *start = min(*start, idxLocation.Offset)
+ *end = max(*end, indexEnd)
+ return nil
+ }
+
+ var colChunk *ColumnChunkMetaData
+ if len(cols) == 0 {
+ for i := 0; i < rgMeta.NumColumns(); i++ {
+ if colChunk, err = rgMeta.ColumnChunk(i); err != nil {
+ return
+ }
+
+ if err = mergeRange(colChunk.GetColumnIndexLocation(), &ciStart, &ciEnd); err != nil {
+ return
+ }
+
+ if err = mergeRange(colChunk.GetOffsetIndexLocation(), &oiStart, &oiEnd); err != nil {
+ return
+ }
+ }
+ } else {
+ for _, i := range cols {
+ if i < 0 || i >= int32(rgMeta.NumColumns()) {
+ return rng, fmt.Errorf("%w: invalid column ordinal %d", arrow.ErrIndex, i)
+ }
+
+ if colChunk, err = rgMeta.ColumnChunk(int(i)); err != nil {
+ return
+ }
+
+ if err = mergeRange(colChunk.GetColumnIndexLocation(), &ciStart, &ciEnd); err != nil {
+ return
+ }
+
+ if err = mergeRange(colChunk.GetOffsetIndexLocation(), &oiStart, &oiEnd); err != nil {
+ return
+ }
+ }
+ }
+
+ if ciEnd != -1 {
+ rng.ColIndex = &readRange{Offset: ciStart, Length: ciEnd - ciStart}
+ }
+
+ if oiEnd != -1 {
+ rng.OffsetIndex = &readRange{Offset: oiStart, Length: oiEnd - oiStart}
+ }
+ return
+}
+
+func (r *PageIndexReader) RowGroup(i int) (*RowGroupPageIndexReader, error) {
+ if i < 0 || i >= r.FileMetadata.NumRowGroups() {
+ return nil, fmt.Errorf("%w: invalid row group ordinal %d", arrow.ErrInvalid, i)
+ }
+
+ var err error
+ rgmeta := r.FileMetadata.RowGroup(i)
+ idxReadRange, ok := r.idxReadRanges[int32(i)]
+ if !ok {
+ // row group has not been requested by WillNeed(), by default both
+ // column index and offset index of all column chunks for the row group
+ // can be read.
+ if idxReadRange, err = determinePageIndexRangesInRowGroup(rgmeta, nil); err != nil {
+ return nil, err
+ }
+ }
+
+ if idxReadRange.ColIndex != nil || idxReadRange.OffsetIndex != nil {
+ return &RowGroupPageIndexReader{
+ input: r.Input,
+ rowGroupMetadata: rgmeta,
+ props: r.Props,
+ rgOrdinal: int32(i),
+ idxReadRange: idxReadRange,
+ fileDecryptor: r.Decryptor,
+ }, nil
+ }
+
+ // the row group does not have a page index or has not been requested by willneed
+ // simply return a nil pointer
+ return nil, nil
+}
+
+func (r *PageIndexReader) WillNeed(rgIndices, colIndices []int32, selection PageIndexSelection) error {
+ if r.idxReadRanges == nil {
+ r.idxReadRanges = make(map[int32]rgIndexReadRange)
+ }
+
+ for _, ordinal := range rgIndices {
+ readRange, err := determinePageIndexRangesInRowGroup(r.FileMetadata.RowGroup(int(ordinal)), colIndices)
+ if err != nil {
+ return err
+ }
+
+ if !selection.ColumnIndex || readRange.ColIndex == nil {
+ // mark column index as not requested
+ readRange.ColIndex = nil
+ }
+
+ if !selection.OffsetIndex || readRange.OffsetIndex == nil {
+ // mark offset index as not requested
+ readRange.OffsetIndex = nil
+ }
+ r.idxReadRanges[int32(ordinal)] = readRange
+ }
+ // TODO: possibly use read ranges to prefetch data of the input
+ return nil
+}
+
+func (r *PageIndexReader) WillNotNeed(rgIndices []int32) {
+ if r.idxReadRanges == nil {
+ return
+ }
+
+ for _, i := range rgIndices {
+ delete(r.idxReadRanges, i)
+ }
+}
+
+type builderState int8
+
+const (
+ stateCreated builderState = iota
+ stateStarted
+ stateFinished
+ stateDiscarded
+)
+
+// ColumnIndexBuilder is an interface for constructing column indexes,
+// with the concrete implementations being fully typed.
+type ColumnIndexBuilder interface {
+ AddPage(stats *EncodedStatistics) error
+ Finish() error
+ WriteTo(w io.Writer, encryptor encryption.Encryptor) (int, error)
+ Build() ColumnIndex
+}
+
+type columnIndexBuilder[T parquet.ColumnTypes] struct {
+ descr *schema.Column
+ colIndex format.ColumnIndex
+ nonNullPageIndices []int64
+ state builderState
+}
+
+// NewColumnIndexBuilder creates a new typed ColumnIndexBuilder for the given column descriptor.
+func NewColumnIndexBuilder(descr *schema.Column) ColumnIndexBuilder {
+ switch descr.PhysicalType() {
+ case parquet.Types.Boolean:
+ return newColumnIndexBuilder[bool](descr)
+ case parquet.Types.Int32:
+ return newColumnIndexBuilder[int32](descr)
+ case parquet.Types.Int64:
+ return newColumnIndexBuilder[int64](descr)
+ case parquet.Types.Int96:
+ return newColumnIndexBuilder[parquet.Int96](descr)
+ case parquet.Types.Float:
+ return newColumnIndexBuilder[float32](descr)
+ case parquet.Types.Double:
+ return newColumnIndexBuilder[float64](descr)
+ case parquet.Types.ByteArray:
+ return newColumnIndexBuilder[parquet.ByteArray](descr)
+ case parquet.Types.FixedLenByteArray:
+ return newColumnIndexBuilder[parquet.FixedLenByteArray](descr)
+ case parquet.Types.Undefined:
+ return nil
+ }
+ panic("unreachable: cannot make column index builder of unknown type")
+}
+
+func newColumnIndexBuilder[T parquet.ColumnTypes](descr *schema.Column) *columnIndexBuilder[T] {
+ return &columnIndexBuilder[T]{
+ descr: descr,
+ colIndex: format.ColumnIndex{
+ NullCounts: make([]int64, 0),
+ BoundaryOrder: Unordered,
+ },
+ nonNullPageIndices: make([]int64, 0),
+ state: stateCreated,
+ }
+}
+
+func (b *columnIndexBuilder[T]) AddPage(stats *EncodedStatistics) error {
+ if b.state == stateFinished {
+ return fmt.Errorf("%w: cannot add page to finished ColumnIndexBuilder", arrow.ErrInvalid)
+ } else if b.state == stateDiscarded {
+ return nil
+ }
+
+ b.state = stateStarted
+
+ switch {
+ case stats.AllNullValue:
+ b.colIndex.NullPages = append(b.colIndex.NullPages, true)
+ // thrift deserializes nil byte slice or empty byte slice both as
+ // an empty byte slice. So we should append an empty byte slice
+ // instead of nil so that round trip comparisons are consistent.
+ b.colIndex.MinValues = append(b.colIndex.MinValues, []byte{})
+ b.colIndex.MaxValues = append(b.colIndex.MaxValues, []byte{})
+ case stats.HasMin && stats.HasMax:
+ pageOrdinal := len(b.colIndex.NullPages)
+ b.nonNullPageIndices = append(b.nonNullPageIndices, int64(pageOrdinal))
+ b.colIndex.MinValues = append(b.colIndex.MinValues, stats.Min)
+ b.colIndex.MaxValues = append(b.colIndex.MaxValues, stats.Max)
+ b.colIndex.NullPages = append(b.colIndex.NullPages, false)
+ default:
+ // this is a non-null page but it lacks meaningful min/max values
+ // discard the column index
+ b.state = stateDiscarded
+ return nil
+ }
+
+ if b.colIndex.IsSetNullCounts() && stats.HasNullCount {
+ b.colIndex.NullCounts = append(b.colIndex.NullCounts, stats.NullCount)
+ } else {
+ b.colIndex.NullCounts = nil
+ }
+
+ return nil
+}
+
+func (b *columnIndexBuilder[T]) Finish() error {
+ switch b.state {
+ case stateCreated:
+ // no page added, discard the column index
+ b.state = stateDiscarded
+ return nil
+ case stateFinished:
+ return fmt.Errorf("%w: ColumnIndexBuilder is already finished", arrow.ErrInvalid)
+ case stateDiscarded:
+ // column index is discarded, do nothing
+ return nil
+ case stateStarted:
+ }
+
+ b.state = stateFinished
+ // clear null counts vector because at least one page does not provide it
+ if !b.colIndex.IsSetNullCounts() {
+ b.colIndex.NullCounts = nil
+ }
+
+ // decode min/max values according to data type
+ nonNullPageCnt := len(b.nonNullPageIndices)
+ minVals, maxVals := make([]T, nonNullPageCnt), make([]T, nonNullPageCnt)
+ dec := getDecoder[T](b.descr)
+ for i, pageOrdinal := range b.nonNullPageIndices {
+ minVals[i] = dec(b.colIndex.MinValues[pageOrdinal])
+ maxVals[i] = dec(b.colIndex.MaxValues[pageOrdinal])
+ }
+
+ // decode the boundary order from decoded min/max vals
+ b.colIndex.BoundaryOrder = b.determineBoundaryOrder(minVals, maxVals)
+ return nil
+}
+
+func (b *columnIndexBuilder[T]) Build() ColumnIndex {
+ if b.state != stateFinished {
+ return nil
+ }
+ return newTypedColumnIndex[T](b.descr, &b.colIndex)
+}
+
+func (b *columnIndexBuilder[T]) WriteTo(w io.Writer, encryptor encryption.Encryptor) (int, error) {
+ if b.state == stateFinished {
+ return thrift.NewThriftSerializer().Serialize(&b.colIndex, w, encryptor)
+ }
+ return 0, nil
+}
+
+func (b *columnIndexBuilder[T]) determineBoundaryOrder(minVals, maxVals []T) BoundaryOrder {
+ debug.Assert(len(minVals) == len(maxVals), "min/max values length mismatch")
+ if len(minVals) == 0 {
+ return Unordered
+ }
+
+ comp, err := NewTypedComparator[T](b.descr)
+ if err != nil {
+ return Unordered
+ }
+
+ // check if both minVals and maxVals are in ascending order
+ isAsc := true
+ for i := 1; i < len(minVals); i++ {
+ if comp.Compare(minVals[i], minVals[i-1]) ||
+ comp.Compare(maxVals[i], maxVals[i-1]) {
+ isAsc = false
+ break
+ }
+ }
+
+ if isAsc {
+ return Ascending
+ }
+
+ // check if both minVals and maxVals are in descending order
+ isDesc := true
+ for i := 1; i < len(minVals); i++ {
+ if comp.Compare(minVals[i-1], minVals[i]) ||
+ comp.Compare(maxVals[i-1], maxVals[i]) {
+ isDesc = false
+ break
+ }
+ }
+ if isDesc {
+ return Descending
+ }
+
+ return Unordered
+}
+
+// OffsetIndexBuilder provides a way to construct new OffsetIndexes while writing
+// a parquet file.
+type OffsetIndexBuilder struct {
+ offsetIndex format.OffsetIndex
+ state builderState
+}
+
+func (o *OffsetIndexBuilder) AddPageLoc(pgloc PageLocation) error {
+ return o.AddPage(pgloc.Offset, pgloc.FirstRowIndex, pgloc.CompressedPageSize)
+}
+
+func (o *OffsetIndexBuilder) AddPage(offset, firstRowIdx int64, compressedPgSize int32) error {
+ if o.state == stateFinished {
+ return fmt.Errorf("%w: cannot add page to finished OffsetIndexBuilder", arrow.ErrInvalid)
+ } else if o.state == stateDiscarded {
+ // offset index is discarded, do nothing
+ return nil
+ }
+
+ o.state = stateStarted
+ o.offsetIndex.PageLocations = append(o.offsetIndex.PageLocations, &PageLocation{
+ Offset: offset,
+ FirstRowIndex: firstRowIdx,
+ CompressedPageSize: compressedPgSize,
+ })
+ return nil
+}
+
+func (o *OffsetIndexBuilder) Finish(finalPos int64) error {
+ switch o.state {
+ case stateCreated:
+ o.state = stateDiscarded
+ case stateStarted:
+ // adjust page offsets according to final position
+ if finalPos > 0 {
+ for _, loc := range o.offsetIndex.PageLocations {
+ loc.Offset += finalPos
+ }
+ }
+ o.state = stateFinished
+ case stateFinished, stateDiscarded:
+ return fmt.Errorf("%w: OffsetIndexBuilder is already finished", arrow.ErrInvalid)
+ }
+ return nil
+}
+
+func (o *OffsetIndexBuilder) WriteTo(w io.Writer, encryptor encryption.Encryptor) (int, error) {
+ if o.state == stateFinished {
+ return thrift.NewThriftSerializer().Serialize(&o.offsetIndex, w, encryptor)
+ }
+ return 0, nil
+}
+
+func (o *OffsetIndexBuilder) Build() OffsetIndex {
+ if o.state != stateFinished {
+ return nil
+ }
+
+ return &o.offsetIndex
+}
+
+// PageIndexBuilder manages the creation of the entire PageIndex for a parquet file,
+// managing the builders for each row group as they are added and providing getters
+// to retrieve the particular builders for specific columns and row groups.
+type PageIndexBuilder struct {
+ Schema *schema.Schema
+ Encryptor encryption.FileEncryptor
+
+ colIndexBuilders [][]ColumnIndexBuilder
+ offsetIndexBuilders [][]*OffsetIndexBuilder
+ finished bool
+}
+
+func (b *PageIndexBuilder) AppendRowGroup() error {
+ if b.finished {
+ return fmt.Errorf("%w: cannot append row group to finished PageIndexBuilder", arrow.ErrInvalid)
+ }
+
+ if b.Schema == nil {
+ return fmt.Errorf("%w: schema is not set in PageIndexBuilder", arrow.ErrInvalid)
+ }
+
+ numColumns := b.Schema.NumColumns()
+ b.colIndexBuilders = append(b.colIndexBuilders, make([]ColumnIndexBuilder, numColumns))
+ b.offsetIndexBuilders = append(b.offsetIndexBuilders, make([]*OffsetIndexBuilder, numColumns))
+
+ debug.Assert(len(b.colIndexBuilders) == len(b.offsetIndexBuilders), "column and offset index builders mismatch")
+ return nil
+}
+
+func (b *PageIndexBuilder) GetColumnIndexBuilder(i int) (ColumnIndexBuilder, error) {
+ if err := b.checkState(i); err != nil {
+ return nil, err
+ }
+
+ bldr := &b.colIndexBuilders[len(b.colIndexBuilders)-1][i]
+ if *bldr == nil {
+ *bldr = NewColumnIndexBuilder(b.Schema.Column(i))
+ }
+ return *bldr, nil
+}
+
+func (b *PageIndexBuilder) GetOffsetIndexBuilder(i int) (*OffsetIndexBuilder, error) {
+ if err := b.checkState(i); err != nil {
+ return nil, err
+ }
+
+ bldr := &b.offsetIndexBuilders[len(b.offsetIndexBuilders)-1][i]
+ if *bldr == nil {
+ *bldr = &OffsetIndexBuilder{}
+ }
+ return *bldr, nil
+}
+
+func (b *PageIndexBuilder) Finish() { b.finished = true }
+
+func (b *PageIndexBuilder) WriteTo(w utils.WriterTell, location *PageIndexLocation) error {
+ if !b.finished {
+ return fmt.Errorf("%w: PageIndexBuilder is not finished", arrow.ErrInvalid)
+ }
+
+ location.ColIndexLocation = make(map[uint64][]*IndexLocation)
+ location.OffsetIndexLocation = make(map[uint64][]*IndexLocation)
+
+ // serialize column index
+ if err := serializeIndex(b.Schema, b.colIndexBuilders, w, location.ColIndexLocation, encryption.ColumnIndexModule, b.getColumnMetaEncryptor); err != nil {
+ return err
+ }
+
+ // serialize offset index
+ if err := serializeIndex(b.Schema, b.offsetIndexBuilders, w, location.OffsetIndexLocation, encryption.OffsetIndexModule, b.getColumnMetaEncryptor); err != nil {
+ return err
+ }
+
+ return nil
+}
+
+func (b *PageIndexBuilder) checkState(col int) error {
+ if b.finished {
+ return fmt.Errorf("%w: cannot add page to finished PageIndexBuilder", arrow.ErrInvalid)
+ }
+
+ if col < 0 || col >= b.Schema.NumColumns() {
+ return fmt.Errorf("%w: invalid column ordinal %d", arrow.ErrInvalid, col)
+ }
+
+ if len(b.colIndexBuilders) == 0 || len(b.offsetIndexBuilders) == 0 {
+ return fmt.Errorf("%w: No row group appended to PageIndexBuilder", arrow.ErrInvalid)
+ }
+ return nil
+}
+
+func (b *PageIndexBuilder) getColumnMetaEncryptor(rgOrdinal, colOrdinal int, moduleType int8) encryption.Encryptor {
+ if b.Encryptor == nil {
+ return nil
+ }
+
+ colPath := b.Schema.Column(colOrdinal).Path()
+ encryptor := b.Encryptor.GetColumnMetaEncryptor(colPath)
+ if encryptor != nil {
+ encryptor.UpdateAad(encryption.CreateModuleAad(
+ encryptor.FileAad(), moduleType, int16(rgOrdinal),
+ int16(colOrdinal), encryption.NonPageOrdinal))
+ }
+ return encryptor
+}
+
+func serializeIndex[T interface {
+ comparable
+ WriteTo(io.Writer, encryption.Encryptor) (int, error)
+}](s *schema.Schema, bldrs [][]T, w utils.WriterTell, location map[uint64][]*IndexLocation, moduleType int8, encFn func(int, int, int8) encryption.Encryptor) error {
+ var (
+ z T
+ numCols = s.NumColumns()
+ )
+
+ // serialize the same kind of page index, row group by row group
+ for rg, idxBldrs := range bldrs {
+ debug.Assert(len(idxBldrs) == numCols, "column index builders length mismatch")
+
+ hasValidIndex := false
+ locations := make([]*IndexLocation, numCols)
+
+ // in the same row group, serialize the same kind of page index column by column
+ for col, bldr := range idxBldrs {
+ if bldr == z {
+ continue
+ }
+
+ encryptor := encFn(rg, col, moduleType)
+ posBefore := w.Tell()
+
+ n, err := bldr.WriteTo(w, encryptor)
+ if err != nil {
+ return err
+ }
+
+ if n == 0 {
+ continue
+ }
+
+ if n > math.MaxInt32 {
+ return fmt.Errorf("%w: serialized page index size overflows INT32_MAX", arrow.ErrInvalid)
+ }
+
+ locations[col] = &IndexLocation{Offset: posBefore, Length: int32(n)}
+ hasValidIndex = true
+ }
+
+ if hasValidIndex {
+ location[uint64(rg)] = locations
+ }
+ }
+
+ return nil
+}
diff --git a/parquet/metadata/page_index_internal_test.go b/parquet/metadata/page_index_internal_test.go
new file mode 100644
index 0000000..49a7417
--- /dev/null
+++ b/parquet/metadata/page_index_internal_test.go
@@ -0,0 +1,197 @@
+// 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 metadata
+
+import (
+ "bytes"
+ "strconv"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/parquet"
+ format "github.com/apache/arrow-go/v18/parquet/internal/gen-go/parquet"
+ "github.com/apache/arrow-go/v18/parquet/internal/thrift"
+ "github.com/apache/arrow-go/v18/parquet/schema"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+type PageIndexRanges struct {
+ ColIndexOffset, OffsetIndexOffset int64
+ ColIndexLen, OffsetIndexLen int32
+}
+
+func constructFakeMetadata(rowGroupRanges []PageIndexRanges) *FileMetaData {
+ var rg format.RowGroup
+ for _, r := range rowGroupRanges {
+ var colChunk format.ColumnChunk
+ colChunk.MetaData = &format.ColumnMetaData{}
+ if r.ColIndexOffset != -1 {
+ colChunk.ColumnIndexOffset = &r.ColIndexOffset
+ }
+ if r.ColIndexLen != -1 {
+ colChunk.ColumnIndexLength = &r.ColIndexLen
+ }
+ if r.OffsetIndexOffset != -1 {
+ colChunk.OffsetIndexOffset = &r.OffsetIndexOffset
+ }
+ if r.OffsetIndexLen != -1 {
+ colChunk.OffsetIndexLength = &r.OffsetIndexLen
+ }
+ rg.Columns = append(rg.Columns, &colChunk)
+ }
+
+ var fileMeta format.FileMetaData
+ fileMeta.RowGroups = append(fileMeta.RowGroups, &rg)
+
+ fields := make(schema.FieldList, 0)
+ for i := range rowGroupRanges {
+ fields = append(fields, schema.NewInt64Node(strconv.Itoa(i), parquet.Repetitions.Optional, -1))
+ }
+
+ grpNode, _ := schema.NewGroupNode("schema", parquet.Repetitions.Repeated, fields, -1)
+ fileMeta.Schema = schema.ToThrift(grpNode)
+
+ var buf bytes.Buffer
+ thrift.NewThriftSerializer().Serialize(&fileMeta, &buf, nil)
+
+ meta, _ := NewFileMetaData(buf.Bytes(), nil)
+ return meta
+}
+
+// validates that determinePagteIndexRangesInRowGroup selects the expected
+// file offsets and sizes or returns false when the row group doesn't have
+// a page index
+func validatePageIndexRange(t *testing.T, ranges []PageIndexRanges, colIndices []int32, expectedHasColIndex, expectedHasOffsetIndex bool, expCIStart, expCISize, expOIStart, expOISize int) {
+ fileMeta := constructFakeMetadata(ranges)
+ readRange, err := determinePageIndexRangesInRowGroup(fileMeta.RowGroup(0), colIndices)
+ require.NoError(t, err)
+
+ if expectedHasColIndex {
+ assert.NotNil(t, readRange.ColIndex)
+ assert.EqualValues(t, expCIStart, readRange.ColIndex.Offset)
+ assert.EqualValues(t, expCISize, readRange.ColIndex.Length)
+ } else {
+ assert.Nil(t, readRange.ColIndex)
+ }
+
+ if expectedHasOffsetIndex {
+ assert.NotNil(t, readRange.OffsetIndex)
+ assert.EqualValues(t, expOIStart, readRange.OffsetIndex.Offset)
+ assert.EqualValues(t, expOISize, readRange.OffsetIndex.Length)
+ } else {
+ assert.Nil(t, readRange.OffsetIndex)
+ }
+}
+
+// construct artificial row groups with page index offsets in them.
+// then validate if determinePageIndexRangesInRowGroup properly
+// computes the file range that contains the whole page index
+func TestDeterminePageIndexRangesInRowGroup(t *testing.T) {
+ tests := []struct {
+ name string
+ ranges []PageIndexRanges
+ cols []int32
+ expectHasCol, expectHasOffset bool
+ expectCIStart, expectCISize int
+ expectOIStart, expectOISize int
+ }{
+ {"no col chunks", nil, nil, false, false, -1, -1, -1, -1},
+ {"no page index", []PageIndexRanges{{-1, -1, -1, -1}}, nil, false, false, -1, -1, -1, -1},
+ {"page index for single col", []PageIndexRanges{{10, 15, 5, 5}}, nil, true, true, 10, 5, 15, 5},
+ {"page index for two columns", []PageIndexRanges{{10, 30, 5, 25}, {15, 50, 15, 20}},
+ nil, true, true, 10, 20, 30, 40},
+ {"page index for second column", []PageIndexRanges{{-1, -1, -1, -1}, {20, 30, 10, 25}}, nil,
+ true, true, 20, 10, 30, 25},
+ {"page index for first column", []PageIndexRanges{{10, 15, 5, 5}, {-1, -1, -1, -1}}, nil,
+ true, true, 10, 5, 15, 5},
+ {"missing offset index for first column, gap in col index",
+ []PageIndexRanges{{10, -1, 5, -1}, {20, 30, 10, 25}}, nil, true, true, 10, 20, 30, 25},
+ {"missing offset index for second column", []PageIndexRanges{{10, 25, 5, 5}, {20, -1, 10, -1}},
+ nil, true, true, 10, 20, 25, 5},
+ {"four col chunks",
+ []PageIndexRanges{{100, 220, 10, 30}, {110, 250, 25, 10}, {140, 260, 30, 40}, {200, 300, 10, 100}}, nil,
+ true, true, 100, 110, 220, 180},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ validatePageIndexRange(t, tt.ranges, tt.cols, tt.expectHasCol, tt.expectHasOffset,
+ tt.expectCIStart, tt.expectCISize, tt.expectOIStart, tt.expectOISize)
+ })
+ }
+}
+
+// properly compute file range that contains the page index of selected columns
+func TestDeterminePageIndexRangesInRowGroupWithPartialColumnsSelected(t *testing.T) {
+ tests := []struct {
+ name string
+ ranges []PageIndexRanges
+ cols []int32
+ expectHasCol, expectHasOffset bool
+ expectCIStart, expectCISize int
+ expectOIStart, expectOISize int
+ }{
+ {"no page index", []PageIndexRanges{{-1, -1, -1, -1}}, []int32{0}, false, false, -1, -1, -1, -1},
+ {"page index for single col", []PageIndexRanges{{10, 15, 5, 5}}, []int32{0}, true, true, 10, 5, 15, 5},
+ {"page index for the 1st col", []PageIndexRanges{{10, 30, 5, 25}, {15, 50, 15, 20}}, []int32{0},
+ true, true, 10, 5, 30, 25},
+ {"page index for the 2nd col", []PageIndexRanges{{10, 30, 5, 25}, {15, 50, 15, 20}}, []int32{0},
+ true, true, 10, 5, 30, 25},
+ {"only 2nd col is selected from 4 chunks",
+ []PageIndexRanges{{100, 220, 10, 30}, {110, 250, 25, 10}, {140, 260, 30, 40}, {200, 300, 10, 100}},
+ []int32{1}, true, true, 110, 25, 250, 10},
+ {"2nd and 3rd cols selected", []PageIndexRanges{{100, 220, 10, 30}, {110, 250, 25, 10}, {140, 260, 30, 40}, {200, 300, 10, 100}},
+ []int32{1, 2}, true, true, 110, 60, 250, 50},
+ {"2nd and 4th cols selected", []PageIndexRanges{{100, 220, 10, 30}, {110, 250, 25, 10}, {140, 260, 30, 40}, {200, 300, 10, 100}},
+ []int32{1, 3}, true, true, 110, 100, 250, 150},
+ {"1st, 2nd, and 4th cols selected", []PageIndexRanges{{100, 220, 10, 30}, {110, 250, 25, 10}, {140, 260, 30, 40}, {200, 300, 10, 100}},
+ []int32{0, 1, 3}, true, true, 100, 110, 220, 180},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ validatePageIndexRange(t, tt.ranges, tt.cols, tt.expectHasCol, tt.expectHasOffset,
+ tt.expectCIStart, tt.expectCISize, tt.expectOIStart, tt.expectOISize)
+ })
+ }
+}
+
+// validate that a column index or offset index is missing
+func TestDeterminePageIndexRangesInRowGroupWithMissingPageIndex(t *testing.T) {
+ tests := []struct {
+ name string
+ ranges []PageIndexRanges
+ cols []int32
+ expectHasCol, expectHasOffset bool
+ expectCIStart, expectCISize int
+ expectOIStart, expectOISize int
+ }{
+ {"no page index", []PageIndexRanges{{-1, 15, -1, 5}}, nil, false, true, -1, -1, 15, 5},
+ {"no offset index", []PageIndexRanges{{10, -1, 5, -1}}, nil, true, false, 10, 5, -1, -1},
+ {"no column index at all among two cols", []PageIndexRanges{{-1, 30, -1, 25}, {-1, 50, -1, 20}}, nil,
+ false, true, -1, -1, 30, 40},
+ {"no offset index at all among two cols", []PageIndexRanges{{10, -1, 5, -1}, {15, -1, 15, -1}}, nil,
+ true, false, 10, 20, -1, -1},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ validatePageIndexRange(t, tt.ranges, tt.cols, tt.expectHasCol, tt.expectHasOffset,
+ tt.expectCIStart, tt.expectCISize, tt.expectOIStart, tt.expectOISize)
+ })
+ }
+}
diff --git a/parquet/metadata/page_index_test.go b/parquet/metadata/page_index_test.go
new file mode 100644
index 0000000..fda7397
--- /dev/null
+++ b/parquet/metadata/page_index_test.go
@@ -0,0 +1,802 @@
+// 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 metadata_test
+
+import (
+ "bytes"
+ "fmt"
+ "math"
+ "os"
+ "path/filepath"
+ "slices"
+ "strconv"
+ "testing"
+ "unsafe"
+
+ "github.com/apache/arrow-go/v18/arrow/float16"
+ "github.com/apache/arrow-go/v18/parquet"
+ "github.com/apache/arrow-go/v18/parquet/file"
+ "github.com/apache/arrow-go/v18/parquet/internal/utils"
+ "github.com/apache/arrow-go/v18/parquet/metadata"
+ "github.com/apache/arrow-go/v18/parquet/schema"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+ "github.com/stretchr/testify/suite"
+)
+
+func TestReadOffsetIndex(t *testing.T) {
+ dir := os.Getenv("PARQUET_TEST_DATA")
+ if dir == "" {
+ t.Skip("PARQUET_TEST_DATA not set")
+ }
+ require.DirExists(t, dir)
+
+ path := filepath.Join(dir, "alltypes_tiny_pages.parquet")
+ rdr, err := file.OpenParquetFile(path, false)
+ require.NoError(t, err)
+ defer rdr.Close()
+
+ const (
+ rgID, colID int = 0, 0
+ )
+
+ fileMetadata := rdr.MetaData()
+ colMeta, err := fileMetadata.RowGroup(rgID).ColumnChunk(colID)
+ require.NoError(t, err)
+ indexLoc := colMeta.GetOffsetIndexLocation()
+ require.NotNil(t, indexLoc)
+
+ // read serialized offset index from file
+ src, err := os.Open(path)
+ require.NoError(t, err)
+ buf := make([]byte, indexLoc.Length)
+ n, err := src.ReadAt(buf, indexLoc.Offset)
+ require.NoError(t, err)
+ assert.EqualValues(t, indexLoc.Length, n)
+ require.NoError(t, src.Close())
+
+ // deserialize offset index
+ offsetIdx := metadata.NewOffsetIndex(buf, nil, nil)
+
+ // verify only partial data as it contains 325 pages in total
+ const numPages = 325
+ var (
+ pageIndices = [...]uint64{0, 100, 200, 300}
+ pageLocations = []metadata.PageLocation{
+ {Offset: 4, CompressedPageSize: 109, FirstRowIndex: 0},
+ {Offset: 11480, CompressedPageSize: 133, FirstRowIndex: 2244},
+ {Offset: 22980, CompressedPageSize: 133, FirstRowIndex: 4494},
+ {Offset: 34480, CompressedPageSize: 133, FirstRowIndex: 6744},
+ }
+ )
+
+ offsetPageLocs := offsetIdx.GetPageLocations()
+ assert.Len(t, offsetPageLocs, numPages)
+
+ for i, pageID := range pageIndices {
+ readPageLocation := offsetPageLocs[pageID]
+ expectedPageLoc := pageLocations[i]
+
+ assert.Equal(t, expectedPageLoc, *readPageLocation)
+ }
+}
+
+func ulps64(actual, expected float64) int64 {
+ ulp := math.Nextafter(actual, math.Inf(1)) - actual
+ return int64(math.Abs((expected - actual) / ulp))
+}
+
+func ulps32(actual, expected float32) int64 {
+ ulp := math.Nextafter32(actual, float32(math.Inf(1))) - actual
+ return int64(math.Abs(float64((expected - actual) / ulp)))
+}
+
+func assertFloat32Approx(t assert.TestingT, expected, actual any, _ ...any) bool {
+ const maxulps int64 = 4
+ ulps := ulps32(actual.(float32), expected.(float32))
+ return assert.LessOrEqualf(t, ulps, maxulps, "%f not equal to %f (%d ulps)", actual, expected, ulps)
+}
+
+func assertFloat64Approx(t assert.TestingT, expected, actual any, _ ...any) bool {
+ const maxulps int64 = 4
+ ulps := ulps64(actual.(float64), expected.(float64))
+ return assert.LessOrEqualf(t, ulps, maxulps, "%f not equal to %f (%d ulps)", actual, expected, ulps)
+}
+
+type testcase struct {
+ fileName string
+ colID int
+ numPages uint64
+ boundaryOrder metadata.BoundaryOrder
+ pageIndices []uint64
+ nullPages []bool
+ hasNullCounts bool
+ nullCounts []int64
+}
+
+func testReadTypedColIndex[T parquet.ColumnTypes](t *testing.T, tc testcase, minValues, maxValues []T) {
+ dir := os.Getenv("PARQUET_TEST_DATA")
+ if dir == "" {
+ t.Skip("PARQUET_TEST_DATA not set")
+ }
+ require.DirExists(t, dir)
+
+ path := filepath.Join(dir, tc.fileName)
+ rdr, err := file.OpenParquetFile(path, false)
+ require.NoError(t, err)
+ defer rdr.Close()
+
+ fileMetadata := rdr.MetaData()
+
+ // Get column index location to a specific column chunk
+ const rgID = 0
+ colMeta, err := fileMetadata.RowGroup(rgID).ColumnChunk(tc.colID)
+ require.NoError(t, err)
+
+ indexLoc := colMeta.GetColumnIndexLocation()
+ require.NotNil(t, indexLoc)
+
+ // read serialized column index from the file
+ src, err := os.Open(path)
+ require.NoError(t, err)
+ buf := make([]byte, indexLoc.Length)
+ n, err := src.ReadAt(buf, indexLoc.Offset)
+ require.NoError(t, err)
+ assert.EqualValues(t, indexLoc.Length, n)
+ require.NoError(t, src.Close())
+
+ descr := fileMetadata.Schema.Column(tc.colID)
+ colIndex := metadata.NewColumnIndex(descr, buf, nil, nil)
+ assert.IsType(t, (*metadata.TypedColumnIndex[T])(nil), colIndex)
+
+ typedColIndex := colIndex.(*metadata.TypedColumnIndex[T])
+ // verify only partial data as there are too many pages
+ assert.Len(t, typedColIndex.GetNullPages(), int(tc.numPages))
+ assert.Equal(t, tc.hasNullCounts, typedColIndex.IsSetNullCounts())
+ assert.Equal(t, tc.boundaryOrder, typedColIndex.GetBoundaryOrder())
+
+ var cmp assert.ComparisonAssertionFunc
+ switch any(minValues).(type) {
+ case []float64:
+ cmp = assertFloat64Approx
+ case []float32:
+ cmp = assertFloat32Approx
+ default:
+ cmp = assert.Equal
+ }
+
+ for i, pageID := range tc.pageIndices {
+ assert.Equal(t, tc.nullPages[i], typedColIndex.GetNullPages()[pageID])
+ if tc.hasNullCounts {
+ assert.Equal(t, tc.nullCounts[i], typedColIndex.GetNullCounts()[pageID])
+ }
+
+ // min/max values are only meaningful for non-null pages
+ if !tc.nullPages[i] {
+ cmp(t, minValues[i], typedColIndex.MinValues()[pageID])
+ cmp(t, maxValues[i], typedColIndex.MaxValues()[pageID])
+ }
+ }
+}
+
+func TestReadTypedColumnIndex(t *testing.T) {
+ tests := []struct {
+ tc testcase
+ minVals, maxVals any
+ }{
+ {testcase{"alltypes_tiny_pages.parquet", 5, 528, metadata.Unordered,
+ []uint64{0, 99, 426, 520}, []bool{false, false, false, false}, true, []int64{0, 0, 0, 0}},
+ []int64{0, 10, 0, 0}, []int64{90, 90, 80, 70}},
+ {testcase{"alltypes_tiny_pages.parquet", 7, 528, metadata.Unordered,
+ []uint64{0, 51, 212, 527}, []bool{false, false, false, false}, true, []int64{0, 0, 0, 0}},
+ []float64{-0, 30.3, 10.1, 40.4}, []float64{90.9, 90.9, 90.9, 60.6}},
+ {testcase{"alltypes_tiny_pages.parquet", 9, 352, metadata.Ascending,
+ []uint64{0, 128, 256}, []bool{false, false, false}, true, []int64{0, 0, 0}},
+ []parquet.ByteArray{{'0'}, {'0'}, {'0'}}, []parquet.ByteArray{{'9'}, {'9'}, {'9'}}},
+ {testcase{"alltypes_tiny_pages.parquet", 1, 82, metadata.Ascending,
+ []uint64{0, 16, 64}, []bool{false, false, false}, true, []int64{0, 0, 0}},
+ []bool{false, false, false}, []bool{true, true, true}},
+ {testcase{"fixed_length_byte_array.parquet", 0, 10, metadata.Descending,
+ []uint64{0, 4, 8}, []bool{false, false, false}, true, []int64{9, 13, 9}},
+ []parquet.FixedLenByteArray{{0x00, 0x00, 0x03, 0x85}, {0x00, 0x00, 0x01, 0xF5}, {0x00, 0x00, 0x00, 0x65}},
+ []parquet.FixedLenByteArray{{0x00, 0x00, 0x03, 0xE8}, {0x00, 0x00, 0x02, 0x58}, {0x00, 0x00, 0x00, 0xC8}}},
+ {testcase{"int32_with_null_pages.parquet", 0, 10, metadata.Unordered,
+ []uint64{2, 4, 8}, []bool{true, false, false}, true, []int64{100, 16, 8}},
+ []int32{0, -2048691758, -2046900272}, []int32{0, 2143189382, 2087168549}},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.tc.fileName+" "+strconv.Itoa(tt.tc.colID), func(t *testing.T) {
+ switch mv := tt.minVals.(type) {
+ case []bool:
+ testReadTypedColIndex(t, tt.tc, mv, tt.maxVals.([]bool))
+ case []int32:
+ testReadTypedColIndex(t, tt.tc, mv, tt.maxVals.([]int32))
+ case []int64:
+ testReadTypedColIndex(t, tt.tc, mv, tt.maxVals.([]int64))
+ case []float64:
+ testReadTypedColIndex(t, tt.tc, mv, tt.maxVals.([]float64))
+ case []parquet.ByteArray:
+ testReadTypedColIndex(t, tt.tc, mv, tt.maxVals.([]parquet.ByteArray))
+ case []parquet.FixedLenByteArray:
+ testReadTypedColIndex(t, tt.tc, mv, tt.maxVals.([]parquet.FixedLenByteArray))
+ }
+ })
+ }
+}
+
+func TestWriteOffsetIndex(t *testing.T) {
+ var bldr metadata.OffsetIndexBuilder
+
+ const (
+ numPages = 5
+ finalPosition = 4096
+ )
+ var (
+ offsets = []int64{100, 200, 300, 400, 500}
+ pageSizes = []int32{1024, 2048, 3072, 4096, 8192}
+ firstRowIndices = []int64{0, 10000, 20000, 30000, 40000}
+ )
+
+ for i := 0; i < numPages; i++ {
+ bldr.AddPage(offsets[i], firstRowIndices[i], pageSizes[i])
+ }
+ bldr.Finish(finalPosition)
+
+ offsetIndexes := make([]metadata.OffsetIndex, 0)
+ // 1st element is the offset index just built
+ offsetIndexes = append(offsetIndexes, bldr.Build())
+ // 2nd element is the offset index restored by serialize-then-deserialize round trip
+ var buf bytes.Buffer
+ bldr.WriteTo(&buf, nil)
+
+ offsetIndexes = append(offsetIndexes, metadata.NewOffsetIndex(buf.Bytes(), nil, nil))
+
+ // verify the data in the offset index
+ for _, offsetIndex := range offsetIndexes {
+ pageLocations := offsetIndex.GetPageLocations()
+ assert.Len(t, pageLocations, numPages)
+
+ for i := 0; i < numPages; i++ {
+ pageLoc := pageLocations[i]
+ assert.Equal(t, offsets[i]+finalPosition, pageLoc.Offset)
+ assert.Equal(t, pageSizes[i], pageLoc.CompressedPageSize)
+ assert.Equal(t, firstRowIndices[i], pageLoc.FirstRowIndex)
+ }
+ }
+}
+
+func simpleEncode[T int32 | int64 | float32 | float64](val T) []byte {
+ return unsafe.Slice((*byte)(unsafe.Pointer(&val)), unsafe.Sizeof(val))
+}
+
+func TestWriteTypedColumnIndex(t *testing.T) {
+ tests := []struct {
+ name string
+ n *schema.PrimitiveNode
+ pageStats []*metadata.EncodedStatistics
+ boundaryOrder metadata.BoundaryOrder
+ hasNullCounts bool
+ }{
+ {"write_int32", schema.NewInt32Node("c1", parquet.Repetitions.Optional, -1),
+ []*metadata.EncodedStatistics{
+ (&metadata.EncodedStatistics{}).SetNullCount(1).SetMin(simpleEncode(int32(1))).SetMax(simpleEncode(int32(2))),
+ (&metadata.EncodedStatistics{}).SetNullCount(2).SetMin(simpleEncode(int32(2))).SetMax(simpleEncode(int32(3))),
+ (&metadata.EncodedStatistics{}).SetNullCount(3).SetMin(simpleEncode(int32(3))).SetMax(simpleEncode(int32(4))),
+ }, metadata.Ascending, true},
+ {"write_int64", schema.NewInt64Node("c1", parquet.Repetitions.Optional, -1),
+ []*metadata.EncodedStatistics{
+ (&metadata.EncodedStatistics{}).SetNullCount(4).SetMin(simpleEncode(int64(-1))).SetMax(simpleEncode(int64(-2))),
+ (&metadata.EncodedStatistics{}).SetNullCount(0).SetMin(simpleEncode(int64(-2))).SetMax(simpleEncode(int64(-3))),
+ (&metadata.EncodedStatistics{}).SetNullCount(4).SetMin(simpleEncode(int64(-3))).SetMax(simpleEncode(int64(-4))),
+ }, metadata.Descending, true},
+ {"write_float32", schema.NewFloat32Node("c1", parquet.Repetitions.Optional, -1),
+ []*metadata.EncodedStatistics{
+ (&metadata.EncodedStatistics{}).SetNullCount(0).SetMin(simpleEncode(float32(2.2))).SetMax(simpleEncode(float32(4.4))),
+ (&metadata.EncodedStatistics{}).SetNullCount(0).SetMin(simpleEncode(float32(1.1))).SetMax(simpleEncode(float32(5.5))),
+ (&metadata.EncodedStatistics{}).SetNullCount(0).SetMin(simpleEncode(float32(3.3))).SetMax(simpleEncode(float32(6.6))),
+ }, metadata.Unordered, true},
+ {"write_float64", schema.NewFloat64Node("c1", parquet.Repetitions.Optional, -1),
+ []*metadata.EncodedStatistics{
+ (&metadata.EncodedStatistics{}).SetMin(simpleEncode(float64(1.2))).SetMax(simpleEncode(float64(4.4))),
+ (&metadata.EncodedStatistics{}).SetMin(simpleEncode(float64(2.2))).SetMax(simpleEncode(float64(5.5))),
+ (&metadata.EncodedStatistics{}).SetMin(simpleEncode(float64(3.3))).SetMax(simpleEncode(float64(-6.6))),
+ }, metadata.Unordered, false},
+ {"write_byte_array", schema.NewByteArrayNode("c1", parquet.Repetitions.Optional, -1),
+ []*metadata.EncodedStatistics{
+ (&metadata.EncodedStatistics{}).SetMin([]byte("bar")).SetMax([]byte("foo")),
+ (&metadata.EncodedStatistics{}).SetMin([]byte("bar")).SetMax([]byte("foo")),
+ (&metadata.EncodedStatistics{}).SetMin([]byte("bar")).SetMax([]byte("foo")),
+ }, metadata.Ascending, false},
+ {"write_fixed_len_byte_array", schema.NewFixedLenByteArrayNode("c1", parquet.Repetitions.Optional, 3, -1),
+ []*metadata.EncodedStatistics{
+ (&metadata.EncodedStatistics{}).SetMin([]byte("abc")).SetMax([]byte("ABC")),
+ (&metadata.EncodedStatistics{AllNullValue: true, Min: []byte{}, Max: []byte{}}),
+ (&metadata.EncodedStatistics{}).SetMin([]byte("foo")).SetMax([]byte("FOO")),
+ (&metadata.EncodedStatistics{AllNullValue: true, Min: []byte{}, Max: []byte{}}),
+ (&metadata.EncodedStatistics{}).SetMin([]byte("xyz")).SetMax([]byte("XYZ")),
+ }, metadata.Ascending, false},
+ {"write_float16", schema.MustPrimitive(
+ schema.NewPrimitiveNodeLogical("c1", parquet.Repetitions.Optional,
+ schema.Float16LogicalType{}, parquet.Types.FixedLenByteArray, 2, -1)),
+ []*metadata.EncodedStatistics{
+ (&metadata.EncodedStatistics{}).SetMin(float16.New(-1.3).ToLEBytes()).SetMax(float16.New(3.6).ToLEBytes()),
+ (&metadata.EncodedStatistics{}).SetMin(float16.New(-0.2).ToLEBytes()).SetMax(float16.New(4.5).ToLEBytes()),
+ (&metadata.EncodedStatistics{}).SetMin(float16.New(1.1).ToLEBytes()).SetMax(float16.New(5.4).ToLEBytes()),
+ (&metadata.EncodedStatistics{}).SetMin(float16.New(2.0).ToLEBytes()).SetMax(float16.New(6.3).ToLEBytes()),
+ }, metadata.Ascending, false},
+ {"write_all_nulls", schema.NewInt32Node("c1", parquet.Repetitions.Optional, -1),
+ []*metadata.EncodedStatistics{
+ (&metadata.EncodedStatistics{AllNullValue: true, Min: []byte{}, Max: []byte{}}).SetNullCount(100),
+ (&metadata.EncodedStatistics{AllNullValue: true, Min: []byte{}, Max: []byte{}}).SetNullCount(100),
+ (&metadata.EncodedStatistics{AllNullValue: true, Min: []byte{}, Max: []byte{}}).SetNullCount(100),
+ }, metadata.Unordered, true},
+ {"invalid_null_counts", schema.NewInt32Node("c1", parquet.Repetitions.Optional, -1),
+ // one page doesn't provide null count
+ []*metadata.EncodedStatistics{
+ (&metadata.EncodedStatistics{}).SetMin(simpleEncode(int32(1))).SetMax(simpleEncode(int32(2))).SetNullCount(0),
+ (&metadata.EncodedStatistics{}).SetMin(simpleEncode(int32(1))).SetMax(simpleEncode(int32(3))),
+ (&metadata.EncodedStatistics{}).SetMin(simpleEncode(int32(2))).SetMax(simpleEncode(int32(3))).SetNullCount(0),
+ }, metadata.Ascending, false},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ descr := schema.NewColumn(tt.n, 1, 0)
+
+ bldr := metadata.NewColumnIndexBuilder(descr)
+ for _, stats := range tt.pageStats {
+ bldr.AddPage(stats)
+ }
+ require.NoError(t, bldr.Finish())
+
+ colIndexes := make([]metadata.ColumnIndex, 0)
+ // 1st element is the column index just built
+ colIndexes = append(colIndexes, bldr.Build())
+ // 2nd element is the column index restored by serialize-then-deserialize round trip
+ var buf bytes.Buffer
+ bldr.WriteTo(&buf, nil)
+ colIndexes = append(colIndexes, metadata.NewColumnIndex(descr, buf.Bytes(), nil, nil))
+
+ for _, colIndex := range colIndexes {
+ assert.Equal(t, tt.hasNullCounts, colIndex.IsSetNullCounts())
+ assert.Equal(t, tt.boundaryOrder, colIndex.GetBoundaryOrder())
+ numPages := len(colIndex.GetNullPages())
+ for i := 0; i < numPages; i++ {
+ assert.Equal(t, tt.pageStats[i].AllNullValue, colIndex.GetNullPages()[i])
+ assert.Equalf(t, tt.pageStats[i].Min, colIndex.GetMinValues()[i], "colIndex: %d, page: %d", colIndex, i)
+ assert.Equalf(t, tt.pageStats[i].Max, colIndex.GetMaxValues()[i], "colIndex: %d, page: %d", colIndex, i)
+ if tt.hasNullCounts {
+ assert.Equal(t, tt.pageStats[i].NullCount, colIndex.GetNullCounts()[i])
+ }
+ }
+ }
+ })
+ }
+}
+
+func TestWriteColumnIndexCorruptedStats(t *testing.T) {
+ pageStats := []*metadata.EncodedStatistics{
+ (&metadata.EncodedStatistics{}).SetMin(simpleEncode(int32(1))).SetMax(simpleEncode(int32(2))),
+ (&metadata.EncodedStatistics{}),
+ (&metadata.EncodedStatistics{}).SetMin(simpleEncode(int32(3))).SetMax(simpleEncode(int32(4))),
+ }
+
+ descr := schema.NewColumn(schema.NewInt32Node("c1", parquet.Repetitions.Optional, -1), 1, 0)
+ bldr := metadata.NewColumnIndexBuilder(descr)
+ for _, stats := range pageStats {
+ bldr.AddPage(stats)
+ }
+ require.NoError(t, bldr.Finish())
+ assert.Nil(t, bldr.Build())
+
+ var buf bytes.Buffer
+ bldr.WriteTo(&buf, nil)
+ assert.Zero(t, buf.Len())
+}
+
+func TestPageIndexBuilderWithZeroRowGroup(t *testing.T) {
+ fields := schema.FieldList{
+ schema.NewInt32Node("c1", parquet.Repetitions.Optional, -1),
+ schema.NewByteArrayNode("c2", parquet.Repetitions.Optional, -1),
+ }
+
+ root := schema.MustGroup(schema.NewGroupNode("schema", parquet.Repetitions.Repeated, fields, -1))
+ sc := schema.NewSchema(root)
+
+ bldr := metadata.PageIndexBuilder{Schema: sc}
+
+ // AppendRowGroup not called
+ cb, err := bldr.GetColumnIndexBuilder(0)
+ assert.Error(t, err)
+ assert.Nil(t, cb)
+ ob, err := bldr.GetOffsetIndexBuilder(0)
+ assert.Error(t, err)
+ assert.Nil(t, ob)
+ bldr.Finish()
+
+ var (
+ buf bytes.Buffer
+ loc metadata.PageIndexLocation
+ )
+
+ assert.NoError(t, bldr.WriteTo(&utils.TellWrapper{Writer: &buf}, &loc))
+ assert.Zero(t, buf.Len())
+ assert.Len(t, loc.ColIndexLocation, 0)
+ assert.Len(t, loc.OffsetIndexLocation, 0)
+}
+
+type PageIndexBuilderSuite struct {
+ suite.Suite
+
+ sc *schema.Schema
+ pageIndexLoc metadata.PageIndexLocation
+ buf bytes.Buffer
+}
+
+func (s *PageIndexBuilderSuite) SetupTest() {
+ s.buf.Reset()
+}
+
+func (s *PageIndexBuilderSuite) writePageIndexes(numRGs, numCols int, pageStats [][]*metadata.EncodedStatistics, pageLocs [][]metadata.PageLocation, finalPos int) {
+ bldr := metadata.PageIndexBuilder{Schema: s.sc}
+ for rg := range numRGs {
+ s.Require().NoError(bldr.AppendRowGroup())
+
+ for col := range numCols {
+ if col < len(pageStats[rg]) {
+ colIdxBldr, err := bldr.GetColumnIndexBuilder(col)
+ s.Require().NoError(err)
+ s.Require().NoError(colIdxBldr.AddPage(pageStats[rg][col]))
+ s.Require().NoError(colIdxBldr.Finish())
+ }
+
+ if col < len(pageLocs[rg]) {
+ offsetIdxBldr, err := bldr.GetOffsetIndexBuilder(col)
+ s.Require().NoError(err)
+ s.Require().NoError(offsetIdxBldr.AddPageLoc(pageLocs[rg][col]))
+ s.Require().NoError(offsetIdxBldr.Finish(int64(finalPos)))
+ }
+ }
+ }
+ bldr.Finish()
+
+ s.Require().NoError(bldr.WriteTo(&utils.TellWrapper{Writer: &s.buf}, &s.pageIndexLoc))
+ s.Len(s.pageIndexLoc.ColIndexLocation, numRGs)
+ s.Len(s.pageIndexLoc.OffsetIndexLocation, numRGs)
+ for rg := range numRGs {
+ s.Len(s.pageIndexLoc.ColIndexLocation[uint64(rg)], numCols)
+ s.Len(s.pageIndexLoc.OffsetIndexLocation[uint64(rg)], numCols)
+ }
+}
+
+func (s *PageIndexBuilderSuite) checkColumnIndex(rg, col int, stats *metadata.EncodedStatistics) {
+ colIndex := s.readColumnIndex(rg, col)
+ s.Require().NotNil(colIndex)
+ s.Len(colIndex.GetNullPages(), 1)
+ s.Equal(stats.AllNullValue, colIndex.GetNullPages()[0])
+ s.Equal(stats.Min, colIndex.GetMinValues()[0])
+ s.Equal(stats.Max, colIndex.GetMaxValues()[0])
+ s.Equal(stats.HasNullCount, colIndex.IsSetNullCounts())
+ if stats.HasNullCount {
+ s.Equal(stats.NullCount, colIndex.GetNullCounts()[0])
+ }
+}
+
+func (s *PageIndexBuilderSuite) checkOffsetIndex(rg, col int, loc metadata.PageLocation, finalLoc int64) {
+ offsetIndex := s.readOffsetIndex(rg, col)
+ s.Require().NotNil(offsetIndex)
+ s.Len(offsetIndex.GetPageLocations(), 1)
+ location := offsetIndex.GetPageLocations()[0]
+ s.Equal(loc.Offset+finalLoc, location.Offset)
+ s.Equal(loc.CompressedPageSize, location.CompressedPageSize)
+ s.Equal(loc.FirstRowIndex, location.FirstRowIndex)
+}
+
+func (s *PageIndexBuilderSuite) readColumnIndex(rg, col int) metadata.ColumnIndex {
+ location := s.pageIndexLoc.ColIndexLocation[uint64(rg)][col]
+ if location == nil {
+ return nil
+ }
+
+ start, end := location.Offset, location.Offset+int64(location.Length)
+ return metadata.NewColumnIndex(s.sc.Column(col), s.buf.Bytes()[start:end], nil, nil)
+}
+
+func (s *PageIndexBuilderSuite) readOffsetIndex(rg, col int) metadata.OffsetIndex {
+ location := s.pageIndexLoc.OffsetIndexLocation[uint64(rg)][col]
+ if location == nil {
+ return nil
+ }
+
+ start, end := location.Offset, location.Offset+int64(location.Length)
+ return metadata.NewOffsetIndex(s.buf.Bytes()[start:end], nil, nil)
+}
+
+func (s *PageIndexBuilderSuite) TestSingleRowGroup() {
+ s.sc = schema.NewSchema(schema.MustGroup(
+ schema.NewGroupNode("schema", parquet.Repetitions.Repeated,
+ schema.FieldList{
+ schema.NewByteArrayNode("c1", parquet.Repetitions.Optional, -1),
+ schema.NewByteArrayNode("c2", parquet.Repetitions.Optional, -1),
+ schema.NewByteArrayNode("c3", parquet.Repetitions.Optional, -1),
+ }, -1)))
+
+ // prepare page stats and page locations
+ // note the 3rd column doesnt have any stats
+ const (
+ numRowGroups = 1
+ numCols = 3
+ finalPos = 200
+ )
+
+ var (
+ pageStats = [][]*metadata.EncodedStatistics{
+ // row group id = 0
+ {
+ (&metadata.EncodedStatistics{}).SetNullCount(0).SetMin([]byte("a")).SetMax([]byte("b")),
+ (&metadata.EncodedStatistics{}).SetNullCount(0).SetMin([]byte("A")).SetMax([]byte("B")),
+ },
+ }
+
+ pageLocations = [][]metadata.PageLocation{
+ // row group id = 0
+ {
+ {Offset: 128, CompressedPageSize: 512, FirstRowIndex: 0},
+ {Offset: 1024, CompressedPageSize: 512, FirstRowIndex: 0},
+ },
+ }
+ )
+
+ s.writePageIndexes(numRowGroups, numCols, pageStats, pageLocations, finalPos)
+
+ // verify the column index and offset index
+ for col := 0; col < 2; col++ {
+ s.checkColumnIndex(0, col, pageStats[0][col])
+ s.checkOffsetIndex(0, col, pageLocations[0][col], finalPos)
+ }
+
+ // 3rd column should not have any column index
+ s.Nil(s.readColumnIndex(0, 2))
+ s.Nil(s.readOffsetIndex(0, 2))
+}
+
+func (s *PageIndexBuilderSuite) TestTwoRowGroups() {
+ s.sc = schema.NewSchema(schema.MustGroup(
+ schema.NewGroupNode("schema", parquet.Repetitions.Repeated,
+ schema.FieldList{
+ schema.NewByteArrayNode("c1", parquet.Repetitions.Optional, -1),
+ schema.NewByteArrayNode("c2", parquet.Repetitions.Optional, -1),
+ }, -1)))
+
+ // prepare page stats and page locations
+ // note the 3rd column doesnt have any stats
+ const (
+ numRowGroups = 2
+ numCols = 2
+ finalPos = 200
+ )
+
+ var (
+ pageStats = [][]*metadata.EncodedStatistics{
+ { // row group id = 0
+ (&metadata.EncodedStatistics{}).SetMin([]byte("a")).SetMax([]byte("b")),
+ (&metadata.EncodedStatistics{}).SetNullCount(0).SetMin([]byte("A")).SetMax([]byte("B")),
+ },
+ { // row group id = 1
+ (&metadata.EncodedStatistics{}), // corrupted stats
+ (&metadata.EncodedStatistics{}).SetNullCount(0).SetMin([]byte("bar")).SetMax([]byte("foo")),
+ },
+ }
+
+ pageLocations = [][]metadata.PageLocation{
+ { // row group id = 0
+ {Offset: 128, CompressedPageSize: 512, FirstRowIndex: 0},
+ {Offset: 1024, CompressedPageSize: 512, FirstRowIndex: 0},
+ },
+ { // row group id = 1
+ {Offset: 128, CompressedPageSize: 512, FirstRowIndex: 0},
+ {Offset: 1024, CompressedPageSize: 512, FirstRowIndex: 0},
+ },
+ }
+ )
+
+ s.writePageIndexes(numRowGroups, numCols, pageStats, pageLocations, finalPos)
+
+ // verify that all columns have good column indexes except the 2nd column
+ // in the 2nd row group
+ s.checkColumnIndex(0, 0, pageStats[0][0])
+ s.checkColumnIndex(0, 1, pageStats[0][1])
+ s.checkColumnIndex(1, 1, pageStats[1][1])
+ s.Nil(s.readColumnIndex(1, 0))
+
+ // verify that two columns have good offset indexes
+ s.checkOffsetIndex(0, 0, pageLocations[0][0], finalPos)
+ s.checkOffsetIndex(0, 1, pageLocations[0][1], finalPos)
+ s.checkOffsetIndex(1, 0, pageLocations[1][0], finalPos)
+ s.checkOffsetIndex(1, 1, pageLocations[1][1], finalPos)
+}
+
+func TestPageIndexBuilder(t *testing.T) {
+ suite.Run(t, new(PageIndexBuilderSuite))
+}
+
+func TestPageIndexReader(t *testing.T) {
+ dir := os.Getenv("PARQUET_TEST_DATA")
+ if dir == "" {
+ t.Skip("PARQUET_TEST_DATA not set")
+ }
+ require.DirExists(t, dir)
+
+ path := filepath.Join(dir, "alltypes_tiny_pages.parquet")
+ rdr, err := file.OpenParquetFile(path, false)
+ require.NoError(t, err)
+ defer rdr.Close()
+
+ fileMeta := rdr.MetaData()
+ assert.EqualValues(t, 1, fileMeta.NumRowGroups())
+ assert.EqualValues(t, 13, fileMeta.Schema.NumColumns())
+
+ pageIdxReader := rdr.GetPageIndexReader()
+ require.NotNil(t, pageIdxReader)
+
+ idxSelectionTests := []metadata.PageIndexSelection{
+ {ColumnIndex: true, OffsetIndex: true},
+ {ColumnIndex: true, OffsetIndex: false},
+ {ColumnIndex: false, OffsetIndex: true},
+ {ColumnIndex: false, OffsetIndex: false},
+ }
+
+ tests := []struct {
+ rowGroupIndices []int32
+ colIndices []int32
+ }{
+ {[]int32{}, []int32{}},
+ {[]int32{0}, []int32{}},
+ {[]int32{0}, []int32{0}},
+ {[]int32{0}, []int32{5}},
+ {[]int32{0}, []int32{0, 5}},
+ }
+
+ for _, tt := range tests {
+ t.Run(fmt.Sprintf("rowGroupIndices=%v,colIndices=%v", tt.rowGroupIndices, tt.colIndices), func(t *testing.T) {
+ for _, idxSelection := range idxSelectionTests {
+ t.Run(fmt.Sprintf("indexSelection=%v", idxSelection), func(t *testing.T) {
+ callWillNeed := len(tt.rowGroupIndices) > 0
+ if callWillNeed {
+ pageIdxReader.WillNeed(tt.rowGroupIndices, tt.colIndices, idxSelection)
+ }
+
+ rgIdxReader, err := pageIdxReader.RowGroup(0)
+ assert.NoError(t, err)
+ if !callWillNeed || idxSelection.OffsetIndex || idxSelection.ColumnIndex {
+ assert.NotNil(t, rgIdxReader)
+ } else {
+ assert.Nil(t, rgIdxReader)
+ return
+ }
+
+ colIdxRequested := func(colID int32) bool {
+ return !callWillNeed || (idxSelection.ColumnIndex &&
+ (len(tt.colIndices) == 0 || slices.Contains(tt.colIndices, colID)))
+ }
+
+ offsetIdxRequested := func(colID int32) bool {
+ return !callWillNeed || (idxSelection.OffsetIndex &&
+ (len(tt.colIndices) == 0 || slices.Contains(tt.colIndices, colID)))
+ }
+
+ if offsetIdxRequested(0) {
+ // verify offset index of column 0 and only partial data as it contains 325 pages
+ const numPages = 325
+ var (
+ pageIndices = []uint64{0, 100, 200, 300}
+ pageLocations = []metadata.PageLocation{
+ {Offset: 4, CompressedPageSize: 109, FirstRowIndex: 0},
+ {Offset: 11480, CompressedPageSize: 133, FirstRowIndex: 2244},
+ {Offset: 22980, CompressedPageSize: 133, FirstRowIndex: 4494},
+ {Offset: 34480, CompressedPageSize: 133, FirstRowIndex: 6744},
+ }
+ )
+
+ offsetIdx, err := rgIdxReader.GetOffsetIndex(0)
+ require.NoError(t, err)
+ assert.NotNil(t, offsetIdx)
+
+ assert.Len(t, offsetIdx.GetPageLocations(), numPages)
+ for i, pageID := range pageIndices {
+ readPageLoc, expectedPageLoc := offsetIdx.GetPageLocations()[pageID], pageLocations[i]
+ assert.EqualValues(t, expectedPageLoc.Offset, readPageLoc.Offset)
+ assert.EqualValues(t, expectedPageLoc.CompressedPageSize, readPageLoc.CompressedPageSize)
+ assert.EqualValues(t, expectedPageLoc.FirstRowIndex, readPageLoc.FirstRowIndex)
+ }
+ } else {
+ _, err := rgIdxReader.GetOffsetIndex(0)
+ assert.Error(t, err)
+ }
+
+ if colIdxRequested(5) {
+ // verify column index of column 5 and only partial data as it contains 528 pages
+ const (
+ numPages = 528
+ boundaryOrder = metadata.Unordered
+ hasNullCounts = true
+ )
+
+ var (
+ pageIndices = []uint64{0, 99, 426, 520}
+ nullPages = []bool{false, false, false, false}
+ nullCounts = []int64{0, 0, 0, 0}
+ minValues = []int64{0, 10, 0, 0}
+ maxValues = []int64{90, 90, 80, 70}
+ )
+
+ colIdx, err := rgIdxReader.GetColumnIndex(5)
+ require.NoError(t, err)
+ assert.NotNil(t, colIdx)
+
+ assert.IsType(t, (*metadata.TypedColumnIndex[int64])(nil), colIdx)
+ typedColIdx := colIdx.(*metadata.TypedColumnIndex[int64])
+
+ assert.Len(t, colIdx.GetNullPages(), numPages)
+ assert.Equal(t, hasNullCounts, colIdx.IsSetNullCounts())
+ assert.Equal(t, boundaryOrder, colIdx.GetBoundaryOrder())
+ for i, pageID := range pageIndices {
+ assert.Equal(t, nullPages[i], colIdx.GetNullPages()[pageID])
+ if hasNullCounts {
+ assert.Equal(t, nullCounts[i], colIdx.GetNullCounts()[pageID])
+ }
+ if !nullPages[i] {
+ assert.Equal(t, minValues[i], typedColIdx.MinValues()[pageID])
+ assert.Equal(t, maxValues[i], typedColIdx.MaxValues()[pageID])
+ }
+ }
+ } else {
+ _, err := rgIdxReader.GetColumnIndex(5)
+ assert.Error(t, err)
+ }
+
+ // verify null is returned if column index does not exist
+ colIdx, err := rgIdxReader.GetColumnIndex(10)
+ assert.NoError(t, err)
+ assert.Nil(t, colIdx)
+ })
+ }
+ })
+ }
+}
+
+func TestReadFileWithoutPageIndex(t *testing.T) {
+ dir := os.Getenv("PARQUET_TEST_DATA")
+ if dir == "" {
+ t.Skip("PARQUET_TEST_DATA not set")
+ }
+ require.DirExists(t, dir)
+
+ path := filepath.Join(dir, "int32_decimal.parquet")
+ rdr, err := file.OpenParquetFile(path, false)
+ require.NoError(t, err)
+ defer rdr.Close()
+
+ fileMeta := rdr.MetaData()
+ assert.EqualValues(t, 1, fileMeta.NumRowGroups())
+ pageIdxReader := rdr.GetPageIndexReader()
+ assert.NotNil(t, pageIdxReader)
+ rgIdxReader, err := pageIdxReader.RowGroup(0)
+ assert.NoError(t, err)
+ assert.Nil(t, rgIdxReader)
+}
diff --git a/parquet/metadata/statistics.go b/parquet/metadata/statistics.go
index 65fd8f5..e7d8cfc 100644
--- a/parquet/metadata/statistics.go
+++ b/parquet/metadata/statistics.go
@@ -19,13 +19,15 @@ package metadata
import (
"bytes"
"encoding/binary"
+ "fmt"
"math"
"unsafe"
"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/arrow-go/v18/arrow/float16"
"github.com/apache/arrow-go/v18/arrow/memory"
- "github.com/apache/arrow-go/v18/internal/utils"
+ "github.com/apache/arrow-go/v18/internal/bitutils"
+ shared_utils "github.com/apache/arrow-go/v18/internal/utils"
"github.com/apache/arrow-go/v18/parquet"
"github.com/apache/arrow-go/v18/parquet/internal/debug"
"github.com/apache/arrow-go/v18/parquet/internal/encoding"
@@ -58,6 +60,8 @@ type EncodedStatistics struct {
NullCount int64
HasDistinctCount bool
DistinctCount int64
+
+ AllNullValue bool
}
// ApplyStatSizeLimits sets the maximum size of the min/max values.
@@ -466,8 +470,8 @@ func (s *Int96Statistics) less(a, b parquet.Int96) bool {
i96a := arrow.Uint32Traits.CastFromBytes(a[:])
i96b := arrow.Uint32Traits.CastFromBytes(b[:])
- a0, a1, a2 := utils.ToLEUint32(i96a[0]), utils.ToLEUint32(i96a[1]), utils.ToLEUint32(i96a[2])
- b0, b1, b2 := utils.ToLEUint32(i96b[0]), utils.ToLEUint32(i96b[1]), utils.ToLEUint32(i96b[2])
+ a0, a1, a2 := shared_utils.ToLEUint32(i96a[0]), shared_utils.ToLEUint32(i96a[1]), shared_utils.ToLEUint32(i96a[2])
+ b0, b1, b2 := shared_utils.ToLEUint32(i96b[0]), shared_utils.ToLEUint32(i96b[1]), shared_utils.ToLEUint32(i96b[2])
if a2 != b2 {
// only the msb bit is by signed comparison
@@ -615,3 +619,447 @@ func GetStatValue(typ parquet.Type, val []byte) interface{} {
}
return nil
}
+
+type Comparator[T parquet.ColumnTypes] interface {
+ // return true if a is strictly less than b
+ Compare(a, b T) bool
+ GetMinMax(vals []T) (min, max T)
+ GetMinMaxSpaced(vals []T, validBits []byte, validBitsOffset int64) (min, max T)
+}
+
+func NewComparator(descr *schema.Column) (any, error) {
+ if descr.SortOrder() == schema.SortSIGNED {
+ switch descr.PhysicalType() {
+ case parquet.Types.Boolean:
+ return &booleanComparator{}, nil
+ case parquet.Types.Int32:
+ return &intComparator[int32]{sortOrder: schema.SortSIGNED}, nil
+ case parquet.Types.Int64:
+ return &intComparator[int64]{sortOrder: schema.SortSIGNED}, nil
+ case parquet.Types.Int96:
+ return &int96Comparator{sortOrder: schema.SortSIGNED}, nil
+ case parquet.Types.Float:
+ return &floatComparator[float32]{}, nil
+ case parquet.Types.Double:
+ return &floatComparator[float64]{}, nil
+ case parquet.Types.ByteArray:
+ return &byteArrayComparator[parquet.ByteArray]{sortOrder: schema.SortSIGNED}, nil
+ case parquet.Types.FixedLenByteArray:
+ if descr.LogicalType().Equals(schema.Float16LogicalType{}) {
+ return &float16Comparator{}, nil
+ }
+ return &byteArrayComparator[parquet.FixedLenByteArray]{sortOrder: schema.SortSIGNED}, nil
+ default:
+ return nil, fmt.Errorf("%w: signed compare not implemented for %s",
+ arrow.ErrNotImplemented, descr.PhysicalType())
+ }
+ } else if descr.SortOrder() == schema.SortUNKNOWN {
+ return nil, fmt.Errorf("%w: sort order is unknown", arrow.ErrNotImplemented)
+ }
+
+ switch descr.PhysicalType() {
+ case parquet.Types.Int32:
+ return &intComparator[int32]{sortOrder: schema.SortUNSIGNED}, nil
+ case parquet.Types.Int64:
+ return &intComparator[int64]{sortOrder: schema.SortUNSIGNED}, nil
+ case parquet.Types.Int96:
+ return &int96Comparator{sortOrder: schema.SortUNSIGNED}, nil
+ case parquet.Types.ByteArray:
+ return &byteArrayComparator[parquet.ByteArray]{sortOrder: schema.SortUNSIGNED}, nil
+ case parquet.Types.FixedLenByteArray:
+ return &byteArrayComparator[parquet.FixedLenByteArray]{sortOrder: schema.SortUNSIGNED}, nil
+ default:
+ return nil, fmt.Errorf("%w: unsigned compare not implemented for %s", arrow.ErrNotImplemented, descr.PhysicalType())
+ }
+}
+
+func NewTypedComparator[T parquet.ColumnTypes](descr *schema.Column) (Comparator[T], error) {
+ comp, err := NewComparator(descr)
+ if err != nil {
+ return nil, err
+ }
+
+ c, ok := comp.(Comparator[T])
+ if !ok {
+ return nil, fmt.Errorf("unexpected comparator type %T", comp)
+ }
+ return c, nil
+}
+
+type intComparator[T int32 | int64] struct {
+ sortOrder schema.SortOrder
+ bitSetReader bitutils.SetBitRunReader
+}
+
+func (c *intComparator[T]) defaultMin() T {
+ var z T
+ switch any(z).(type) {
+ case int32:
+ if c.sortOrder == schema.SortUNSIGNED {
+ val := uint32(math.MaxUint32)
+ return T(val)
+ }
+ return math.MaxInt32
+ case int64:
+ if c.sortOrder == schema.SortUNSIGNED {
+ val := uint64(math.MaxUint64)
+ return T(val)
+ }
+ var v int64 = math.MaxInt64
+ return T(v)
+ }
+ panic("unreachable")
+}
+
+func (c *intComparator[T]) defaultMax() T {
+ if c.sortOrder == schema.SortUNSIGNED {
+ return T(0)
+ }
+
+ var z T
+ switch any(z).(type) {
+ case int32:
+ return math.MinInt32
+ case int64:
+ var v int64 = math.MinInt64
+ return T(v)
+ }
+ panic("unreachable")
+}
+
+func (c *intComparator[T]) Compare(a, b T) bool {
+ if c.sortOrder == schema.SortUNSIGNED {
+ return uint64(a) < uint64(b)
+ }
+ return a < b
+}
+
+func (c *intComparator[T]) GetMinMax(vals []T) (min, max T) {
+ if c.sortOrder == schema.SortSIGNED {
+ switch v := any(vals).(type) {
+ case []int32:
+ minv, maxv := shared_utils.GetMinMaxInt32(v)
+ return T(minv), T(maxv)
+ case []int64:
+ minv, maxv := shared_utils.GetMinMaxInt64(v)
+ return T(minv), T(maxv)
+ }
+ panic("unreachable")
+ }
+
+ switch v := any(vals).(type) {
+ case []int32:
+ minv, maxv := shared_utils.GetMinMaxUint32(arrow.GetData[uint32](arrow.GetBytes(v)))
+ return T(minv), T(maxv)
+ case []int64:
+ minv, maxv := shared_utils.GetMinMaxUint64(arrow.GetData[uint64](arrow.GetBytes(v)))
+ return T(minv), T(maxv)
+ }
+ panic("unreachable")
+}
+
+func (c *intComparator[T]) GetMinMaxSpaced(vals []T, validBits []byte, validBitsOffset int64) (minv, maxv T) {
+ minv, maxv = c.defaultMin(), c.defaultMax()
+ var fn func([]T) (T, T)
+ switch any(vals).(type) {
+ case []int32:
+ if c.sortOrder == schema.SortSIGNED {
+ fn = func(t []T) (T, T) {
+ minv, maxv := shared_utils.GetMinMaxInt32(any(t).([]int32))
+ return T(minv), T(maxv)
+ }
+ } else {
+ fn = func(t []T) (T, T) {
+ minv, maxv := shared_utils.GetMinMaxUint32(arrow.GetData[uint32](arrow.GetBytes(any(t).([]int32))))
+ return T(minv), T(maxv)
+ }
+ }
+ case []int64:
+ if c.sortOrder == schema.SortSIGNED {
+ fn = func(t []T) (T, T) {
+ minv, maxv := shared_utils.GetMinMaxInt64(any(t).([]int64))
+ return T(minv), T(maxv)
+ }
+ } else {
+ fn = func(t []T) (T, T) {
+ minv, maxv := shared_utils.GetMinMaxUint64(arrow.GetData[uint64](arrow.GetBytes(any(t).([]int64))))
+ return T(minv), T(maxv)
+ }
+ }
+ }
+
+ if c.bitSetReader == nil {
+ c.bitSetReader = bitutils.NewSetBitRunReader(validBits, validBitsOffset, int64(len(vals)))
+ } else {
+ c.bitSetReader.Reset(validBits, validBitsOffset, int64(len(vals)))
+ }
+
+ for {
+ run := c.bitSetReader.NextRun()
+ if run.Length == 0 {
+ break
+ }
+
+ localMin, localMax := fn(vals[int(run.Pos):int(run.Pos+run.Length)])
+ minv, maxv = min(minv, localMin), max(maxv, localMax)
+ }
+ return
+}
+
+func basicMinMax[T any](vals []T, defMin, defMax T, coalesce, min, max func(T, T) T) (minv, maxv T) {
+ minv, maxv = defMin, defMax
+
+ for _, v := range vals {
+ minv = min(minv, coalesce(v, defMin))
+ maxv = max(maxv, coalesce(v, defMax))
+ }
+ return
+}
+
+func basicMinMaxSpaced[T any](vals []T, validBits []byte, validOffset int64, bitSetReader *bitutils.SetBitRunReader,
+ defMin, defMax T, coalesce, min, max func(T, T) T) (minv, maxv T) {
+
+ minv, maxv = defMin, defMax
+ if *bitSetReader == nil {
+ *bitSetReader = bitutils.NewSetBitRunReader(validBits, validOffset, int64(len(vals)))
+ } else {
+ (*bitSetReader).Reset(validBits, validOffset, int64(len(vals)))
+ }
+
+ for {
+ run := (*bitSetReader).NextRun()
+ if run.Length == 0 {
+ break
+ }
+
+ for _, v := range vals[int(run.Pos):int(run.Pos+run.Length)] {
+ minv = min(minv, coalesce(v, defMin))
+ maxv = max(maxv, coalesce(v, defMax))
+ }
+ }
+ return
+}
+
+type floatComparator[T float32 | float64] struct {
+ bitSetReader bitutils.SetBitRunReader
+}
+
+func (floatComparator[T]) coalesce(val, fallback T) T {
+ if math.IsNaN(float64(val)) {
+ return fallback
+ }
+ return val
+}
+
+func (c *floatComparator[T]) defaultMin() T {
+ var z T
+ switch any(z).(type) {
+ case float32:
+ return math.MaxFloat32
+ case float64:
+ v := math.MaxFloat64
+ return T(v)
+ }
+ panic("unreachable")
+}
+
+func (c *floatComparator[T]) defaultMax() T {
+ return -c.defaultMin()
+}
+
+func (c *floatComparator[T]) Compare(a, b T) bool {
+ return a < b
+}
+
+func (c *floatComparator[T]) GetMinMax(vals []T) (minv, maxv T) {
+ return basicMinMax(vals, c.defaultMin(), c.defaultMax(), c.coalesce,
+ func(a, b T) T { return min(a, b) }, func(a, b T) T { return max(a, b) })
+}
+
+func (c *floatComparator[T]) GetMinMaxSpaced(vals []T, validBits []byte, validBitsOffset int64) (minv, maxv T) {
+ return basicMinMaxSpaced(vals, validBits, validBitsOffset, &c.bitSetReader, c.defaultMin(), c.defaultMax(), c.coalesce,
+ func(a, b T) T { return min(a, b) }, func(a, b T) T { return max(a, b) })
+}
+
+func byteArrMinVal[T parquet.ByteArray | parquet.FixedLenByteArray](cmpFn func(T, T) bool) func(a, b T) T {
+ return func(a, b T) T {
+ switch {
+ case a == nil:
+ return b
+ case b == nil:
+ return a
+ case cmpFn(a, b):
+ return a
+ default:
+ return b
+ }
+ }
+}
+
+func byteArrMaxVal[T parquet.ByteArray | parquet.FixedLenByteArray](cmpFn func(T, T) bool) func(a, b T) T {
+ return func(a, b T) T {
+ switch {
+ case a == nil:
+ return b
+ case b == nil:
+ return a
+ case cmpFn(a, b):
+ return b
+ default:
+ return a
+ }
+ }
+}
+
+type float16Comparator struct {
+ bitSetReader bitutils.SetBitRunReader
+}
+
+func (*float16Comparator) coalesce(v, fallback parquet.FixedLenByteArray) parquet.FixedLenByteArray {
+ if float16.FromLEBytes(v).IsNaN() {
+ return fallback
+ }
+ return v
+}
+
+func (*float16Comparator) Compare(a, b parquet.FixedLenByteArray) bool {
+ return float16.FromLEBytes(a).Less(float16.FromLEBytes(b))
+}
+
+func (c *float16Comparator) GetMinMax(vals []parquet.FixedLenByteArray) (minv, maxv parquet.FixedLenByteArray) {
+ return basicMinMax(vals, defaultMinFloat16, defaultMaxFloat16, c.coalesce, byteArrMinVal(c.Compare), byteArrMaxVal(c.Compare))
+}
+
+func (c *float16Comparator) GetMinMaxSpaced(vals []parquet.FixedLenByteArray, validBits []byte, validBitsOffset int64) (minv, maxv parquet.FixedLenByteArray) {
+ return basicMinMaxSpaced(vals, validBits, validBitsOffset, &c.bitSetReader, defaultMinFloat16, defaultMaxFloat16, c.coalesce, byteArrMinVal(c.Compare), byteArrMaxVal(c.Compare))
+}
+
+type int96Comparator struct {
+ sortOrder schema.SortOrder
+ bitSetReader bitutils.SetBitRunReader
+}
+
+func (c *int96Comparator) defaultMin() parquet.Int96 {
+ if c.sortOrder == schema.SortUNSIGNED {
+ return defaultMinUInt96
+ }
+ return defaultMinInt96
+}
+
+func (c *int96Comparator) defaultMax() parquet.Int96 {
+ if c.sortOrder == schema.SortUNSIGNED {
+ return defaultMaxUInt96
+ }
+ return defaultMaxInt96
+}
+
+func (c *int96Comparator) Compare(a, b parquet.Int96) bool {
+ i96a := arrow.Uint32Traits.CastFromBytes(a[:])
+ i96b := arrow.Uint32Traits.CastFromBytes(b[:])
+
+ a0, a1, a2 := shared_utils.ToLEUint32(i96a[0]), shared_utils.ToLEUint32(i96a[1]), shared_utils.ToLEUint32(i96a[2])
+ b0, b1, b2 := shared_utils.ToLEUint32(i96b[0]), shared_utils.ToLEUint32(i96b[1]), shared_utils.ToLEUint32(i96b[2])
+
+ if a2 != b2 {
+ // only the msb bit is by signed comparison
+ if c.sortOrder == schema.SortSIGNED {
+ return int32(a2) < int32(b2)
+ }
+ return a2 < b2
+ } else if a1 != b1 {
+ return a1 < b1
+ }
+ return a0 < b0
+}
+
+func (c *int96Comparator) GetMinMax(vals []parquet.Int96) (minv, maxv parquet.Int96) {
+ return basicMinMax(vals, c.defaultMin(), c.defaultMax(), func(v, _ parquet.Int96) parquet.Int96 { return v },
+ func(a, b parquet.Int96) parquet.Int96 {
+ if c.Compare(a, b) {
+ return a
+ }
+ return b
+ }, func(a, b parquet.Int96) parquet.Int96 {
+ if c.Compare(a, b) {
+ return b
+ }
+ return a
+ })
+}
+
+func (c *int96Comparator) GetMinMaxSpaced(vals []parquet.Int96, validBits []byte, validBitsOffset int64) (minv, maxv parquet.Int96) {
+ return basicMinMaxSpaced(vals, validBits, validBitsOffset, &c.bitSetReader, c.defaultMin(), c.defaultMax(),
+ func(v, _ parquet.Int96) parquet.Int96 { return v },
+ func(a, b parquet.Int96) parquet.Int96 {
+ if c.Compare(a, b) {
+ return a
+ }
+ return b
+ }, func(a, b parquet.Int96) parquet.Int96 {
+ if c.Compare(a, b) {
+ return b
+ }
+ return a
+ })
+}
+
+type byteArrayComparator[T parquet.ByteArray | parquet.FixedLenByteArray] struct {
+ sortOrder schema.SortOrder
+ bitSetReader bitutils.SetBitRunReader
+}
+
+func (c *byteArrayComparator[T]) Compare(a, b T) bool {
+ if c.sortOrder == schema.SortUNSIGNED {
+ return bytes.Compare(a, b) == -1
+ }
+
+ return signedByteLess(a, b)
+}
+
+func (c *byteArrayComparator[T]) GetMinMax(vals []T) (minv, maxv T) {
+ return basicMinMax(vals, nil, nil, func(v, _ T) T { return v },
+ byteArrMinVal(c.Compare), byteArrMaxVal(c.Compare))
+}
+
+func (c *byteArrayComparator[T]) GetMinMaxSpaced(vals []T, validBits []byte, validBitsOffset int64) (minv, maxv T) {
+ return basicMinMaxSpaced(vals, validBits, validBitsOffset, &c.bitSetReader, nil, nil,
+ func(v, _ T) T { return v }, byteArrMinVal(c.Compare), byteArrMaxVal(c.Compare))
+}
+
+type booleanComparator struct {
+ bitSetReader bitutils.SetBitRunReader
+}
+
+func (*booleanComparator) Compare(a, b bool) bool {
+ return !a && b
+}
+
+func (c *booleanComparator) GetMinMax(vals []bool) (minv, maxv bool) {
+ return basicMinMax(vals, true, false, func(v, _ bool) bool { return v },
+ func(a, b bool) bool {
+ if c.Compare(a, b) {
+ return a
+ }
+ return b
+ }, func(a, b bool) bool {
+ if c.Compare(a, b) {
+ return b
+ }
+ return a
+ })
+}
+
+func (c *booleanComparator) GetMinMaxSpaced(vals []bool, validBits []byte, validBitsOffset int64) (minv, maxv bool) {
+ return basicMinMaxSpaced(vals, validBits, validBitsOffset, &c.bitSetReader, true, false, func(v, _ bool) bool { return v },
+ func(a, b bool) bool {
+ if c.Compare(a, b) {
+ return a
+ }
+ return b
+ }, func(a, b bool) bool {
+ if c.Compare(a, b) {
+ return b
+ }
+ return a
+ })
+}
diff --git a/parquet/metadata/statistics_types.gen.go b/parquet/metadata/statistics_types.gen.go
index 9f44ada..73988e5 100644
--- a/parquet/metadata/statistics_types.gen.go
+++ b/parquet/metadata/statistics_types.gen.go
@@ -326,6 +326,7 @@ func (s *Int32Statistics) Encode() (enc EncodedStatistics, err error) {
}
if s.HasNullCount() {
enc.SetNullCount(s.NullCount())
+ enc.AllNullValue = s.NumValues() == 0
}
if s.HasDistinctCount() {
enc.SetDistinctCount(s.DistinctCount())
@@ -626,6 +627,7 @@ func (s *Int64Statistics) Encode() (enc EncodedStatistics, err error) {
}
if s.HasNullCount() {
enc.SetNullCount(s.NullCount())
+ enc.AllNullValue = s.NumValues() == 0
}
if s.HasDistinctCount() {
enc.SetDistinctCount(s.DistinctCount())
@@ -904,6 +906,7 @@ func (s *Int96Statistics) Encode() (enc EncodedStatistics, err error) {
}
if s.HasNullCount() {
enc.SetNullCount(s.NullCount())
+ enc.AllNullValue = s.NumValues() == 0
}
if s.HasDistinctCount() {
enc.SetDistinctCount(s.DistinctCount())
@@ -1196,6 +1199,7 @@ func (s *Float32Statistics) Encode() (enc EncodedStatistics, err error) {
}
if s.HasNullCount() {
enc.SetNullCount(s.NullCount())
+ enc.AllNullValue = s.NumValues() == 0
}
if s.HasDistinctCount() {
enc.SetDistinctCount(s.DistinctCount())
@@ -1488,6 +1492,7 @@ func (s *Float64Statistics) Encode() (enc EncodedStatistics, err error) {
}
if s.HasNullCount() {
enc.SetNullCount(s.NullCount())
+ enc.AllNullValue = s.NumValues() == 0
}
if s.HasDistinctCount() {
enc.SetDistinctCount(s.DistinctCount())
@@ -1766,6 +1771,7 @@ func (s *BooleanStatistics) Encode() (enc EncodedStatistics, err error) {
}
if s.HasNullCount() {
enc.SetNullCount(s.NullCount())
+ enc.AllNullValue = s.NumValues() == 0
}
if s.HasDistinctCount() {
enc.SetDistinctCount(s.DistinctCount())
@@ -2073,6 +2079,7 @@ func (s *ByteArrayStatistics) Encode() (enc EncodedStatistics, err error) {
}
if s.HasNullCount() {
enc.SetNullCount(s.NullCount())
+ enc.AllNullValue = s.NumValues() == 0
}
if s.HasDistinctCount() {
enc.SetDistinctCount(s.DistinctCount())
@@ -2383,6 +2390,7 @@ func (s *FixedLenByteArrayStatistics) Encode() (enc EncodedStatistics, err error
}
if s.HasNullCount() {
enc.SetNullCount(s.NullCount())
+ enc.AllNullValue = s.NumValues() == 0
}
if s.HasDistinctCount() {
enc.SetDistinctCount(s.DistinctCount())
@@ -2684,6 +2692,7 @@ func (s *Float16Statistics) Encode() (enc EncodedStatistics, err error) {
}
if s.HasNullCount() {
enc.SetNullCount(s.NullCount())
+ enc.AllNullValue = s.NumValues() == 0
}
if s.HasDistinctCount() {
enc.SetDistinctCount(s.DistinctCount())
diff --git a/parquet/metadata/statistics_types.gen.go.tmpl b/parquet/metadata/statistics_types.gen.go.tmpl
index fcd8a69..fb14ad4 100644
--- a/parquet/metadata/statistics_types.gen.go.tmpl
+++ b/parquet/metadata/statistics_types.gen.go.tmpl
@@ -480,6 +480,7 @@ func (s *{{.Name}}Statistics) Encode() (enc EncodedStatistics, err error) {
}
if s.HasNullCount() {
enc.SetNullCount(s.NullCount())
+ enc.AllNullValue = s.NumValues() == 0
}
if s.HasDistinctCount() {
enc.SetDistinctCount(s.DistinctCount())
diff --git a/parquet/pqarrow/schema.go b/parquet/pqarrow/schema.go
index efb9551..8eccd57 100644
--- a/parquet/pqarrow/schema.go
+++ b/parquet/pqarrow/schema.go
@@ -696,7 +696,7 @@ func listToSchemaField(n *schema.GroupNode, currentLevels file.LevelInfo, ctx *s
return err
}
} else {
- if err := groupToStructField(listGroup, currentLevels, ctx, out, &out.Children[0]); err != nil {
+ if err := groupToStructField(listGroup, currentLevels, ctx, &out.Children[0]); err != nil {
return err
}
}
@@ -741,7 +741,7 @@ func listToSchemaField(n *schema.GroupNode, currentLevels file.LevelInfo, ctx *s
return nil
}
-func groupToStructField(n *schema.GroupNode, currentLevels file.LevelInfo, ctx *schemaTree, parent, out *SchemaField) error {
+func groupToStructField(n *schema.GroupNode, currentLevels file.LevelInfo, ctx *schemaTree, out *SchemaField) error {
arrowFields := make([]arrow.Field, 0, n.NumFields())
out.Children = make([]SchemaField, n.NumFields())
@@ -852,7 +852,7 @@ func groupToSchemaField(n *schema.GroupNode, currentLevels file.LevelInfo, ctx *
// }
out.Children = make([]SchemaField, 1)
repeatedAncestorDef := currentLevels.IncrementRepeated()
- if err := groupToStructField(n, currentLevels, ctx, out, &out.Children[0]); err != nil {
+ if err := groupToStructField(n, currentLevels, ctx, &out.Children[0]); err != nil {
return err
}
@@ -865,7 +865,7 @@ func groupToSchemaField(n *schema.GroupNode, currentLevels file.LevelInfo, ctx *
}
currentLevels.Increment(n)
- return groupToStructField(n, currentLevels, ctx, parent, out)
+ return groupToStructField(n, currentLevels, ctx, out)
}
func createFieldMeta(fieldID int) arrow.Metadata {
diff --git a/parquet/types.go b/parquet/types.go
index 8cc0a94..d4a6850 100644
--- a/parquet/types.go
+++ b/parquet/types.go
@@ -371,3 +371,7 @@ func (t Type) ByteSize() int {
}
panic("no bytesize info for type")
}
+
+type ColumnTypes interface {
+ bool | int32 | int64 | float32 | float64 | Int96 | ByteArray | FixedLenByteArray
+}
diff --git a/parquet/writer_properties.go b/parquet/writer_properties.go
index 06b92ca..3ee5f79 100644
--- a/parquet/writer_properties.go
+++ b/parquet/writer_properties.go
@@ -47,8 +47,10 @@ const (
DefaultStatsEnabled = true
// If the stats are larger than 4K the writer will skip writing them out anyways.
DefaultMaxStatsSize int64 = 4096
- DefaultCreatedBy = "parquet-go version 18.0.0-SNAPSHOT"
- DefaultRootName = "schema"
+ // Default is to not write page indexes for columns
+ DefaultPageIndexEnabled = false
+ DefaultCreatedBy = "parquet-go version 18.0.0-SNAPSHOT"
+ DefaultRootName = "schema"
)
// ColumnProperties defines the encoding, codec, and so on for a given column.
@@ -57,6 +59,7 @@ type ColumnProperties struct {
Codec compress.Compression
DictionaryEnabled bool
StatsEnabled bool
+ PageIndexEnabled bool
MaxStatsSize int64
CompressionLevel int
}
@@ -65,18 +68,20 @@ type ColumnProperties struct {
//
// The default column properties are the following constants:
//
-// Encoding: Encodings.Plain
-// Codec: compress.Codecs.Uncompressed
-// DictionaryEnabled: DefaultDictionaryEnabled
-// StatsEnabled: DefaultStatsEnabled
-// MaxStatsSize: DefaultMaxStatsSize
-// CompressionLevel: compress.DefaultCompressionLevel
+// Encoding: Encodings.Plain
+// Codec: compress.Codecs.Uncompressed
+// DictionaryEnabled: DefaultDictionaryEnabled
+// StatsEnabled: DefaultStatsEnabled
+// PageIndexEnabled: DefaultPageIndexEnabled
+// MaxStatsSize: DefaultMaxStatsSize
+// CompressionLevel: compress.DefaultCompressionLevel
func DefaultColumnProperties() ColumnProperties {
return ColumnProperties{
Encoding: Encodings.Plain,
Codec: compress.Codecs.Uncompressed,
DictionaryEnabled: DefaultDictionaryEnabled,
StatsEnabled: DefaultStatsEnabled,
+ PageIndexEnabled: DefaultPageIndexEnabled,
MaxStatsSize: DefaultMaxStatsSize,
CompressionLevel: compress.DefaultCompressionLevel,
}
@@ -92,6 +97,7 @@ type writerPropConfig struct {
compressLevel map[string]int
dictEnabled map[string]bool
statsEnabled map[string]bool
+ indexEnabled map[string]bool
}
// WriterProperty is used as the options for building a writer properties instance
@@ -312,6 +318,25 @@ func WithSortingColumns(cols []SortingColumn) WriterProperty {
}
}
+// WithPageIndexEnabled specifies the default value for whether or not to write page indexes for columns
+func WithPageIndexEnabled(enabled bool) WriterProperty {
+ return func(cfg *writerPropConfig) {
+ cfg.wr.defColumnProps.PageIndexEnabled = enabled
+ }
+}
+
+// WithPageIndexEnabled specifies a per column value as to enable or disable writing page indexes for the column
+func WithPageIndexEnabledFor(path string, enabled bool) WriterProperty {
+ return func(cfg *writerPropConfig) {
+ cfg.indexEnabled[path] = enabled
+ }
+}
+
+// WithPageIndexEnabledPath is like WithPageIndexEnabledFor, but takes a ColumnPath
+func WithPageIndexEnabledPath(path ColumnPath, enabled bool) WriterProperty {
+ return WithPageIndexEnabledFor(path.String(), enabled)
+}
+
// WriterProperties is the collection of properties to use for writing a parquet file. The values are
// read only once it has been constructed.
type WriterProperties struct {
@@ -372,6 +397,7 @@ func NewWriterProperties(opts ...WriterProperty) *WriterProperties {
compressLevel: make(map[string]int),
dictEnabled: make(map[string]bool),
statsEnabled: make(map[string]bool),
+ indexEnabled: make(map[string]bool),
}
for _, o := range opts {
o(&cfg)
@@ -406,6 +432,10 @@ func NewWriterProperties(opts ...WriterProperty) *WriterProperties {
for key, value := range cfg.statsEnabled {
get(key).StatsEnabled = value
}
+
+ for key, value := range cfg.indexEnabled {
+ get(key).PageIndexEnabled = value
+ }
return cfg.wr
}
@@ -535,6 +565,23 @@ func (w *WriterProperties) StatisticsEnabledPath(path ColumnPath) bool {
return w.StatisticsEnabledFor(path.String())
}
+// PageIndexEnabled returns the default value for whether or not page indexes will be written
+func (w *WriterProperties) PageIndexEnabled() bool { return w.defColumnProps.PageIndexEnabled }
+
+// PageIndexEnabledFor returns whether page index writing is enabled for the given column path, or
+// the default value if it wasn't specified separately.
+func (w *WriterProperties) PageIndexEnabledFor(path string) bool {
+ if p, ok := w.columnProps[path]; ok {
+ return p.PageIndexEnabled
+ }
+ return w.defColumnProps.PageIndexEnabled
+}
+
+// PageIndexEnabledPath is the same as PageIndexEnabledFor but takes a ColumnPath object
+func (w *WriterProperties) PageIndexEnabledPath(path ColumnPath) bool {
+ return w.PageIndexEnabledFor(path.String())
+}
+
// MaxStatsSize returns the default maximum size for stats
func (w *WriterProperties) MaxStatsSize() int64 { return w.defColumnProps.MaxStatsSize }