feat(parquet): add experimental VECTOR repetition for Arrow FixedSizeList

parquet/file/column_reader.go

@@ -118,6 +118,11 @@ type ColumnChunkReader interface {
 	Err() error
 
 	SeekToRow(rowIdx int64) error
+	// seekToRowWithValueStride seeks to rowIdx using row-group row ordinals but
+	// positions the value decoder at rowIdx*valueStride physical values. This is
+	// used by VECTOR columns where one logical row contributes a fixed number of
+	// primitive leaf values.
+	seekToRowWithValueStride(rowIdx, valueStride int64) error
 
 	// Skip buffered values
 	consumeBufferedValues(int64)
@@ -571,6 +576,36 @@ func (c *columnChunkReader) determineNumToRead(batchLen int64, defLvls, repLvls
 	return
 }
 
+// seekToPageForRow seeks the page reader to the page containing rowIdx and
+// initializes the current data page. If seeking replays an already-read
+// dictionary page, it skips reconfiguration and advances to the next data page.
+func (c *columnChunkReader) seekToPageForRow(rowIdx int64) error {
+	if err := c.pager().SeekToPageWithRow(rowIdx); err != nil {
+		return err
+	}
+
+	c.numBuffered, c.numDecoded = 0, 0
+	c.curPage = c.rdr.Page()
+	if c.curPage == nil {
+		c.err = c.rdr.Err()
+		return c.err
+	}
+
+	gotDataPage := false
+	if _, isDict := c.curPage.(*DictionaryPage); !isDict || c.dictState == dictNotRead {
+		var err error
+		if gotDataPage, err = c.processPage(); err != nil {
+			c.err = err
+			return err
+		}
+	}
+
+	if !gotDataPage {
+		c.readNewPage()
+	}
+	return c.err
+}
+
 // SeekToRow will seek to the row index provided in the column chunk. If
 // the metadata contains an OffsetIndex for skipping pages based on row indexes
 // then the pager will use that to skip to the correct page.
@@ -580,28 +615,109 @@ func (c *columnChunkReader) determineNumToRead(batchLen int64, defLvls, repLvls
 // reader will discard values until it reaches the desired row index according
 // to the definition and repetition levels.
 func (c *columnChunkReader) SeekToRow(rowIdx int64) error {
-	if err := c.pager().SeekToPageWithRow(rowIdx); err != nil {
+	if c.descr.InVectorColumn() {
+		// VECTOR columns contribute EffectiveVectorLength leaf values
+		// per row and carry no per-element levels, so a row ordinal maps to a
+		// fixed value stride. skipRows would otherwise treat one leaf slot as one
+		// row and undershoot the seek by a factor of the vector length.
+		return c.seekToRowWithValueStride(rowIdx, int64(c.descr.EffectiveVectorLength()))
+	}
+	if err := c.seekToPageForRow(rowIdx); err != nil {
 		return err
 	}
+	return c.skipRows(rowIdx - c.curPage.(DataPage).FirstRowIndex())
+}
 
-	c.numBuffered, c.numDecoded = 0, 0
-	c.curPage = c.rdr.Page()
-	if c.curPage == nil {
-		c.err = c.rdr.Err()
-		return c.err
+func (c *columnChunkReader) hasOffsetIndex() bool {
+	spr, ok := c.rdr.(*serializedPageReader)
+	if !ok || spr.pgIndexReader == nil {
+		return false
 	}
+	oidx, err := spr.pgIndexReader.GetOffsetIndex(spr.colIdx)
+	return err == nil && oidx != nil
+}
 
-	gotDataPage, err := c.processPage()
-	if err != nil {
-		c.err = err
+func (c *columnChunkReader) seekToRowWithValueStride(rowIdx, valueStride int64) error {
+	if rowIdx < 0 {
+		return fmt.Errorf("parquet: cannot seek column reader to row index %d", rowIdx)
+	}
+	if valueStride <= 0 {
+		return fmt.Errorf("parquet: invalid value stride %d", valueStride)
+	}
+	valuesToSkip, ok := utils.Mul64(rowIdx, valueStride)
+	if !ok {
+		return errors.New("parquet: seek value offset overflow")
+	}
+	if spr, ok := c.rdr.(*serializedPageReader); ok {
+		if spr.nrows%valueStride != 0 {
+			return fmt.Errorf("parquet: VECTOR column chunk has %d values, not a multiple of vector length %d", spr.nrows, valueStride)
+		}
+		if spr.parentRows > 0 {
+			if rowIdx >= spr.parentRows {
+				return fmt.Errorf("parquet: cannot seek column reader to row index %d", rowIdx)
+			}
+		} else if valuesToSkip < 0 || valuesToSkip >= spr.nrows {
+			return fmt.Errorf("parquet: cannot seek column reader to row index %d", rowIdx)
+		}
+	}
+
+	// Use page-level seeking when page row ordinals are known to be parent rows:
+	// DataPageV2 stores num_rows, and V1 is reliable when an offset index supplies
+	// first_row_index. V1 pages without an offset index only expose num_values
+	// (leaf slots), so fall back to row-group-local value skipping below.
+	if err := c.seekToPageForRow(rowIdx); err != nil {
 		return err
 	}
+	page := c.curPage.(DataPage)
+	pageFirstRow := page.FirstRowIndex()
+	if _, ok := page.(*DataPageV2); ok || c.hasOffsetIndex() {
+		if pageFirstRow <= rowIdx {
+			pageValuesToSkip, ok := utils.Mul64(rowIdx-pageFirstRow, valueStride)
+			if !ok {
+				return errors.New("parquet: seek page value offset overflow")
+			}
+			return c.skipValues(pageValuesToSkip)
+		}
+	}
 
-	if !gotDataPage {
-		c.readNewPage()
+	// The caller has already selected the row group by parent row. Reset to the
+	// start of that row group before skipping leaf values. This is the universally
+	// correct path for V1 pages without an offset index.
+	if err := c.seekToPageForRow(0); err != nil {
+		return err
 	}
+	return c.skipValues(valuesToSkip)
+}
 
-	return c.skipRows(rowIdx - c.curPage.(DataPage).FirstRowIndex())
+func (c *columnChunkReader) skipValues(nvalues int64) error {
+	toSkip := nvalues
+	for c.HasNext() && toSkip > 0 {
+		available := c.numBuffered - c.numDecoded
+		if available <= 0 {
+			continue
+		}
+		if toSkip >= available {
+			toSkip -= available
+			c.consumeBufferedValues(available)
+			continue
+		}
+
+		skipped, err := c.curDecoder.Discard(int(toSkip))
+		if err != nil {
+			c.err = err
+			return err
+		}
+		if skipped <= 0 {
+			c.err = errors.New("parquet: decoder did not discard any values")
+			return c.err
+		}
+		toSkip -= int64(skipped)
+		c.consumeBufferedValues(int64(skipped))
+	}
+	if toSkip > 0 {
+		return fmt.Errorf("parquet: cannot seek past end of column chunk, %d values remaining", toSkip)
+	}
+	return c.err
 }
 
 func (c *columnChunkReader) skipRows(nrows int64) error {

parquet/file/column_writer.go

@@ -490,8 +490,9 @@ func (w *columnWriter) writeLevels(numValues int64, defLevels, repLevels []int16
 
 		w.WriteRepetitionLevels(repLevels[:numValues])
 	} else {
-		// each value is exactly 1 row
-		w.numBufferedRows += int(numValues)
+		// each value is exactly 1 row, except for VECTOR columns where each row
+		// spans effectiveVectorLength leaf slots
+		w.numBufferedRows += w.rowsForLeafValues(numValues)
 	}
 	return toWrite
 }
@@ -509,7 +510,7 @@ func (w *columnWriter) writeLevelsSpaced(numLevels int64, defLevels, repLevels [
 		}
 		w.WriteRepetitionLevels(repLevels[:numLevels])
 	} else {
-		w.numBufferedRows += int(numLevels)
+		w.numBufferedRows += w.rowsForLeafValues(numLevels)
 	}
 }
 
@@ -654,8 +655,84 @@ func (w *columnWriter) Close() (err error) {
 	return err
 }
 
+// rowsForLeafValues converts a count of physical leaf values (slots) to the
+// number of parent rows they represent. For VECTOR columns each row contributes
+// effectiveVectorLength leaf slots; for every other column each leaf value is
+// exactly one row at this nesting level. VECTOR page/column statistics
+// intentionally remain flattened element-level statistics, not vector-level or
+// lexicographic-vector ordering statistics.
+func (w *columnWriter) rowsForLeafValues(numLeafValues int64) int {
+	if vectorLen := w.vectorLengthForBatch(numLeafValues); vectorLen > 0 {
+		return int(numLeafValues / vectorLen)
+	}
+	return int(numLeafValues)
+}
+
+// ErrVectorBatchMisaligned is returned by the typed column writers when a VECTOR
+// column is handed a batch whose physical leaf-value count is not a whole
+// multiple of the column's vector_length. A vector value must be written in one
+// piece, so partial vectors are rejected up front.
+var ErrVectorBatchMisaligned = errors.New("parquet: VECTOR columns must be written in whole-vector batches")
+
+// validateVectorBatch checks, before any values are written, that n physical
+// leaf values form whole vectors for a VECTOR column. It is a no-op for every
+// other column. Validating here lets the public WriteBatch/WriteBatchSpaced
+// methods return ErrVectorBatchMisaligned directly instead of surfacing the
+// internal vectorLengthForBatch panic as a recovered "unknown error type".
+func (w *columnWriter) validateVectorBatch(n int64) error {
+	if !w.descr.InVectorColumn() {
+		return nil
+	}
+	vectorLen := int64(w.descr.EffectiveVectorLength())
+	if vectorLen <= 0 || n%vectorLen != 0 {
+		return fmt.Errorf("%w: %d values is not a multiple of vector_length %d", ErrVectorBatchMisaligned, n, vectorLen)
+	}
+	return nil
+}
+
+func (w *columnWriter) vectorLengthForBatch(total int64) int64 {
+	if !w.descr.InVectorColumn() {
+		return 0
+	}
+	vectorLen := int64(w.descr.EffectiveVectorLength())
+	if vectorLen <= 0 || total%vectorLen != 0 {
+		// Unreachable on the public path: WriteBatch/WriteBatchSpaced call
+		// validateVectorBatch up front. Kept as an internal invariant guard.
+		panic(ErrVectorBatchMisaligned)
+	}
+	return vectorLen
+}
+
+func alignBatchToVector(batch, vectorLen int64) int64 {
+	if vectorLen <= 0 {
+		return batch
+	}
+	if aligned := batch - batch%vectorLen; aligned >= vectorLen {
+		return aligned
+	}
+	return vectorLen
+}
+
+func (w *columnWriter) doBatchesAlignedForVector(total int64, repLevels []int16, action func(offset, batch int64)) {
+	if w.descr.InVectorColumn() {
+		w.doBatches(total, repLevels, action)
+		return
+	}
+	doBatches(total, w.props.WriteBatchSize(), action)
+}
+
 func (w *columnWriter) doBatches(total int64, repLevels []int16, action func(offset, batch int64)) {
 	batchSize := w.props.WriteBatchSize()
+	// VECTOR columns: a vector value occupies effectiveVectorLength contiguous
+	// leaf slots and must not be split across data pages. Page flushes are
+	// evaluated once per batch, so batching on whole-vector boundaries keeps
+	// every page boundary on a whole-vector boundary. (VECTOR columns have
+	// max repetition level 0, so without this they would fall into the plain
+	// batching path below and could split a vector across pages.)
+	if vectorLen := w.vectorLengthForBatch(total); vectorLen > 0 {
+		doBatches(total, alignBatchToVector(batchSize, vectorLen), action)
+		return
+	}
 	// if we're writing V1 data pages, have no replevels or the max replevel is 0 then just
 	// use the regular doBatches function
 	if w.props.DataPageVersion() == parquet.DataPageV1 || repLevels == nil || w.descr.MaxRepetitionLevel() == 0 {