Repository: segmentio/parquet-go
Pull request #190 (2022-05-23 23:05:30)

fix slicing of repeated pages
This PR adds a test provided by @mdisibio which highlighted a bug, and fixes it.

The problem was a misuse of the repetition levels to detect the start of new rows. New rows start where the repetition level is zero, we were testing for `def != page.maxRepetitionLevel` instead of `def == 0`.



Patch:
diff --git a/bloom/block_amd64.go b/bloom/block_amd64.go
index 16fb17f..1113010 100644
--- a/bloom/block_amd64.go
+++ b/bloom/block_amd64.go
@@ -2,6 +2,8 @@
 
 package bloom
 
+import "golang.org/x/sys/cpu"
+
 // The functions in this file are SIMD-optimized versions of the functions
 // declared in block_optimized.go for x86 targets.
 //
@@ -24,6 +26,7 @@ package bloom
 // block_optimized.go; the delta comparing to functions in block_default.go is
 // significantly larger but not very interesting since those functions have no
 // practical use cases.
+var hasAVX2 = cpu.X86.HasAVX2
 
 //go:noescape
 func blockInsert(b *Block, x uint32)
diff --git a/bloom/block_amd64.s b/bloom/block_amd64.s
index b8ade61..713aea2 100644
--- a/bloom/block_amd64.s
+++ b/bloom/block_amd64.s
@@ -2,6 +2,15 @@
 
 #include "textflag.h"
 
+#define salt0 0x47b6137b
+#define salt1 0x44974d91
+#define salt2 0x8824ad5b
+#define salt3 0xa2b7289d
+#define salt4 0x705495c7
+#define salt5 0x2df1424b
+#define salt6 0x9efc4947
+#define salt7 0x5c6bfb31
+
 DATA ones+0(SB)/4, $1
 DATA ones+4(SB)/4, $1
 DATA ones+8(SB)/4, $1
@@ -12,14 +21,14 @@ DATA ones+24(SB)/4, $1
 DATA ones+28(SB)/4, $1
 GLOBL ones(SB), RODATA|NOPTR, $32
 
-DATA salt+0(SB)/4, $0x47b6137b
-DATA salt+4(SB)/4, $0x44974d91
-DATA salt+8(SB)/4, $0x8824ad5b
-DATA salt+12(SB)/4, $0xa2b7289d
-DATA salt+16(SB)/4, $0x705495c7
-DATA salt+20(SB)/4, $0x2df1424b
-DATA salt+24(SB)/4, $0x9efc4947
-DATA salt+28(SB)/4, $0x5c6bfb31
+DATA salt+0(SB)/4, $salt0
+DATA salt+4(SB)/4, $salt1
+DATA salt+8(SB)/4, $salt2
+DATA salt+12(SB)/4, $salt3
+DATA salt+16(SB)/4, $salt4
+DATA salt+20(SB)/4, $salt5
+DATA salt+24(SB)/4, $salt6
+DATA salt+28(SB)/4, $salt7
 GLOBL salt(SB), RODATA|NOPTR, $32
 
 // This initial block is a SIMD implementation of the mask function declared in
@@ -46,19 +55,52 @@ GLOBL salt(SB), RODATA|NOPTR, $32
     VPSRLD $27, tmp, tmp \
     VPSLLVD tmp, dst, dst
 
+#define insert(salt, src, dst) \
+    MOVL src, CX \
+    IMULL salt, CX \
+    SHRL $27, CX \
+    MOVL $1, DX \
+    SHLL CX, DX \
+    ORL DX, dst
+
+#define check(salt, b, x) \
+    MOVL b, CX \
+    MOVL x, DX \
+    IMULL salt, DX \
+    SHRL $27, DX \
+    BTL DX, CX \
+    JAE notfound
+
 // func blockInsert(b *Block, x uint32)
 TEXT ·blockInsert(SB), NOSPLIT, $0-16
     MOVQ b+0(FP), AX
+    CMPB ·hasAVX2(SB), $0
+    JE fallback
+avx2:
     generateMask(x+8(FP), Y1, Y0)
     // Set all 1 bits of the mask in the bloom filter block.
     VPOR (AX), Y0, Y0
     VMOVDQU Y0, (AX)
     VZEROUPPER
     RET
+fallback:
+    MOVL x+8(FP), BX
+    insert($salt0, BX, 0(AX))
+    insert($salt1, BX, 4(AX))
+    insert($salt2, BX, 8(AX))
+    insert($salt3, BX, 12(AX))
+    insert($salt4, BX, 16(AX))
+    insert($salt5, BX, 20(AX))
+    insert($salt6, BX, 24(AX))
+    insert($salt7, BX, 28(AX))
+    RET
 
 // func blockCheck(b *Block, x uint32) bool
 TEXT ·blockCheck(SB), NOSPLIT, $0-17
     MOVQ b+0(FP), AX
+    CMPB ·hasAVX2(SB), $0
+    JE fallback
+avx2:
     generateMask(x+8(FP), Y1, Y0)
     // Compare the 1 bits of the mask with the bloom filter block, then compare
     // the result with the mask, expecting equality if the value `x` was present
@@ -68,3 +110,20 @@ TEXT ·blockCheck(SB), NOSPLIT, $0-17
     SETCS ret+16(FP)   // return CF == 1
     VZEROUPPER
     RET
+fallback:
+    MOVL x+8(FP), BX
+    check($salt0, 0(AX), BX)
+    check($salt1, 4(AX), BX)
+    check($salt2, 8(AX), BX)
+    check($salt3, 12(AX), BX)
+    check($salt4, 16(AX), BX)
+    check($salt5, 20(AX), BX)
+    check($salt6, 24(AX), BX)
+    check($salt7, 28(AX), BX)
+    MOVB $1, CX
+    JMP done
+notfound:
+    XORB CX, CX
+done:
+    MOVB CX, ret+16(FP)
+    RET
diff --git a/bloom/block_optimized.go b/bloom/block_optimized.go
index e879d15..d67bec1 100644
--- a/bloom/block_optimized.go
+++ b/bloom/block_optimized.go
@@ -24,41 +24,26 @@ package bloom
 // The benchmarks measure throughput based on the byte size of a bloom filter
 // block.
 
-func mask(x uint32) Block {
-	return Block{
-		0: 1 << ((x * salt0) >> 27),
-		1: 1 << ((x * salt1) >> 27),
-		2: 1 << ((x * salt2) >> 27),
-		3: 1 << ((x * salt3) >> 27),
-		4: 1 << ((x * salt4) >> 27),
-		5: 1 << ((x * salt5) >> 27),
-		6: 1 << ((x * salt6) >> 27),
-		7: 1 << ((x * salt7) >> 27),
-	}
-}
-
 func (b *Block) Insert(x uint32) {
-	masked := mask(x)
-	b[0] |= masked[0]
-	b[1] |= masked[1]
-	b[2] |= masked[2]
-	b[3] |= masked[3]
-	b[4] |= masked[4]
-	b[5] |= masked[5]
-	b[6] |= masked[6]
-	b[7] |= masked[7]
+	b[0] |= 1 << ((x * salt0) >> 27)
+	b[1] |= 1 << ((x * salt1) >> 27)
+	b[2] |= 1 << ((x * salt2) >> 27)
+	b[3] |= 1 << ((x * salt3) >> 27)
+	b[4] |= 1 << ((x * salt4) >> 27)
+	b[5] |= 1 << ((x * salt5) >> 27)
+	b[6] |= 1 << ((x * salt6) >> 27)
+	b[7] |= 1 << ((x * salt7) >> 27)
 }
 
 func (b *Block) Check(x uint32) bool {
-	masked := mask(x)
-	return ((b[0] & masked[0]) != 0) &&
-		((b[1] & masked[1]) != 0) &&
-		((b[2] & masked[2]) != 0) &&
-		((b[3] & masked[3]) != 0) &&
-		((b[4] & masked[4]) != 0) &&
-		((b[5] & masked[5]) != 0) &&
-		((b[6] & masked[6]) != 0) &&
-		((b[7] & masked[7]) != 0)
+	return ((b[0] & (1 << ((x * salt0) >> 27))) != 0) &&
+		((b[1] & (1 << ((x * salt1) >> 27))) != 0) &&
+		((b[2] & (1 << ((x * salt2) >> 27))) != 0) &&
+		((b[3] & (1 << ((x * salt3) >> 27))) != 0) &&
+		((b[4] & (1 << ((x * salt4) >> 27))) != 0) &&
+		((b[5] & (1 << ((x * salt5) >> 27))) != 0) &&
+		((b[6] & (1 << ((x * salt6) >> 27))) != 0) &&
+		((b[7] & (1 << ((x * salt7) >> 27))) != 0)
 }
 
 func (f SplitBlockFilter) insertBulk(x []uint64) {
diff --git a/bloom/filter_amd64.s b/bloom/filter_amd64.s
index e216210..cfa75e6 100644
--- a/bloom/filter_amd64.s
+++ b/bloom/filter_amd64.s
@@ -2,6 +2,15 @@
 
 #include "textflag.h"
 
+#define salt0 0x47b6137b
+#define salt1 0x44974d91
+#define salt2 0x8824ad5b
+#define salt3 0xa2b7289d
+#define salt4 0x705495c7
+#define salt5 0x2df1424b
+#define salt6 0x9efc4947
+#define salt7 0x5c6bfb31
+
 // See block_amd64.s for a description of this algorithm.
 #define generateMask(src, dst) \
     VMOVDQA ones(SB), dst \
@@ -32,23 +41,41 @@
     MOVQ tmpXMM, r2 \
     VPEXTRQ $1, tmpXMM, r3
 
+#define insert(salt, src, dst) \
+    MOVL src, CX \
+    IMULL salt, CX \
+    SHRL $27, CX \
+    MOVL $1, DX \
+    SHLL CX, DX \
+    ORL DX, dst
+
+#define check(salt, b, x) \
+    MOVL b, CX \
+    MOVL x, DX \
+    IMULL salt, DX \
+    SHRL $27, DX \
+    BTL DX, CX \
+    JAE notfound
+
 // func filterInsertBulk(f []Block, x []uint64)
 TEXT ·filterInsertBulk(SB), NOSPLIT, $0-48
     MOVQ f_base+0(FP), AX
     MOVQ f_len+8(FP), CX
     MOVQ x_base+24(FP), BX
     MOVQ x_len+32(FP), DX
+    CMPB ·hasAVX2(SB), $0
+    JE fallback
+avx2:
     VPBROADCASTQ f_base+8(FP), Y0
-
     // Loop initialization, SI holds the current index in `x`, DI is the number
     // of elements in `x` rounded down to the nearest multiple of 4.
     XORQ SI, SI
     MOVQ DX, DI
     SHRQ $2, DI
     SHLQ $2, DI
-loop4x64:
+avx2loop4x64:
     CMPQ SI, DI
-    JAE loop
+    JAE avx2loop1x64
 
     // The masks and indexes for 4 input hashes are computed in each loop
     // iteration. The hashes are loaded in Y1 so we can use vector instructions
@@ -84,22 +111,44 @@ loop4x64:
     applyMask(Y9, (AX)(R11*1))
 
     ADDQ $4, SI
-    JMP loop4x64
-loop:
+    JMP avx2loop4x64
+avx2loop1x64:
     // Compute trailing elements in `x` if the length was not a multiple of 4.
     // This is the same algorithm as the one in the loop4x64 section, working
     // on a single mask/block pair at a time.
     CMPQ SI, DX
-    JE done
+    JE avx2done
     MOVQ (BX)(SI*8), R8
     VPBROADCASTD (BX)(SI*8), Y0
     fasthash1x64(CX, R8)
     generateMask(Y0, Y1)
     applyMask(Y1, (AX)(R8*1))
     INCQ SI
+    JMP avx2loop1x64
+avx2done:
+    VZEROUPPER
+    JMP done
+fallback:
+    XORQ SI, SI
+    MOVQ DX, DI
+    MOVQ CX, R10
+loop:
+    CMPQ SI, DI
+    JE done
+    MOVLQZX (BX)(SI*8), R8
+    MOVQ (BX)(SI*8), R9
+    fasthash1x64(R10, R9)
+    insert($salt0, R8, 0(AX)(R9*1))
+    insert($salt1, R8, 4(AX)(R9*1))
+    insert($salt2, R8, 8(AX)(R9*1))
+    insert($salt3, R8, 12(AX)(R9*1))
+    insert($salt4, R8, 16(AX)(R9*1))
+    insert($salt5, R8, 20(AX)(R9*1))
+    insert($salt6, R8, 24(AX)(R9*1))
+    insert($salt7, R8, 28(AX)(R9*1))
+    INCQ SI
     JMP loop
 done:
-    VZEROUPPER
     RET
 
 // func filterInsert(f []Block, x uint64)
@@ -107,23 +156,59 @@ TEXT ·filterInsert(SB), NOSPLIT, $0-32
     MOVQ f_base+0(FP), AX
     MOVQ f_len+8(FP), BX
     MOVQ x+24(FP), CX
-    VPBROADCASTD x+24(FP), Y1
     fasthash1x64(BX, CX)
+    CMPB ·hasAVX2(SB), $0
+    JE fallback
+avx2:
+    VPBROADCASTD x+24(FP), Y1
     generateMask(Y1, Y0)
     applyMask(Y0, (AX)(CX*1))
     VZEROUPPER
     RET
+fallback:
+    ADDQ CX, AX
+    MOVL x+24(FP), BX
+    insert($salt0, BX, 0(AX))
+    insert($salt1, BX, 4(AX))
+    insert($salt2, BX, 8(AX))
+    insert($salt3, BX, 12(AX))
+    insert($salt4, BX, 16(AX))
+    insert($salt5, BX, 20(AX))
+    insert($salt6, BX, 24(AX))
+    insert($salt7, BX, 28(AX))
+    RET
 
 // func filterCheck(f []Block, x uint64) bool
 TEXT ·filterCheck(SB), NOSPLIT, $0-33
     MOVQ f_base+0(FP), AX
     MOVQ f_len+8(FP), BX
     MOVQ x+24(FP), CX
-    VPBROADCASTD x+24(FP), Y1
     fasthash1x64(BX, CX)
+    CMPB ·hasAVX2(SB), $0
+    JE fallback
+avx2:
+    VPBROADCASTD x+24(FP), Y1
     generateMask(Y1, Y0)
     VPAND (AX)(CX*1), Y0, Y1
     VPTEST Y0, Y1
     SETCS ret+32(FP)
     VZEROUPPER
     RET
+fallback:
+    ADDQ CX, AX
+    MOVL x+24(FP), BX
+    check($salt0, 0(AX), BX)
+    check($salt1, 4(AX), BX)
+    check($salt2, 8(AX), BX)
+    check($salt3, 12(AX), BX)
+    check($salt4, 16(AX), BX)
+    check($salt5, 20(AX), BX)
+    check($salt6, 24(AX), BX)
+    check($salt7, 28(AX), BX)
+    MOVB $1, CX
+    JMP done
+notfound:
+    XORB CX, CX
+done:
+    MOVB CX, ret+32(FP)
+    RET
diff --git a/page.go b/page.go
index 913d8b0..3b4d4f4 100644
--- a/page.go
+++ b/page.go
@@ -444,27 +444,35 @@ func (page *repeatedPage) Slice(i, j int64) BufferedPage {
 		panic(errPageBoundsOutOfRange(i, j, numRows))
 	}
 
-	rowIndex0 := int64(0)
-	rowIndex1 := int64(len(page.repetitionLevels))
-	rowIndex2 := int64(len(page.repetitionLevels))
+	rowIndex0 := 0
+	rowIndex1 := len(page.repetitionLevels)
+	rowIndex2 := len(page.repetitionLevels)
 
 	for k, def := range page.repetitionLevels {
-		if def != page.maxRepetitionLevel {
-			if rowIndex0 == i {
-				rowIndex1 = int64(k)
+		if def == 0 {
+			if rowIndex0 == int(i) {
+				rowIndex1 = k
+				break
 			}
-			if rowIndex0 == j {
-				rowIndex2 = int64(k)
+			rowIndex0++
+		}
+	}
+
+	for k, def := range page.repetitionLevels[rowIndex1:] {
+		if def == 0 {
+			if rowIndex0 == int(j) {
+				rowIndex2 = rowIndex1 + k
+				break
 			}
 			rowIndex0++
 		}
 	}
 
-	numNulls1 := int64(countLevelsNotEqual(page.definitionLevels[:rowIndex1], page.maxDefinitionLevel))
-	numNulls2 := int64(countLevelsNotEqual(page.definitionLevels[rowIndex1:rowIndex2], page.maxDefinitionLevel))
+	numNulls1 := countLevelsNotEqual(page.definitionLevels[:rowIndex1], page.maxDefinitionLevel)
+	numNulls2 := countLevelsNotEqual(page.definitionLevels[rowIndex1:rowIndex2], page.maxDefinitionLevel)
 
-	i = rowIndex1 - numNulls1
-	j = rowIndex2 - (numNulls1 + numNulls2)
+	i = int64(rowIndex1 - numNulls1)
+	j = int64(rowIndex2 - (numNulls1 + numNulls2))
 
 	return newRepeatedPage(
 		page.base.Slice(i, j),
diff --git a/print.go b/print.go
index ea2b930..92d47d5 100644
--- a/print.go
+++ b/print.go
@@ -223,7 +223,6 @@ func PrintRowGroup(w io.Writer, rowGroup RowGroup) error {
 
 	rowbuf := make([]Row, defaultRowBufferSize)
 	cells := make([]string, 0, len(columns))
-	parts := make([]string, 0)
 	rows := rowGroup.Rows()
 	defer rows.Close()
 
@@ -233,25 +232,19 @@ func PrintRowGroup(w io.Writer, rowGroup RowGroup) error {
 		for _, row := range rowbuf[:n] {
 			cells = cells[:0]
 
-			for i := 0; i < len(row); {
-				j := i + 1
+			for _, value := range row {
+				columnIndex := value.Column()
 
-				for j < len(row) && row[j].Column() == row[i].Column() {
-					j++
+				for len(cells) <= columnIndex {
+					cells = append(cells, "")
 				}
 
-				if (j - i) == 1 {
-					cells = append(cells, row[i].String())
+				if cells[columnIndex] == "" {
+					cells[columnIndex] = value.String()
 				} else {
-					parts = parts[:0]
-					for k := i; k < j; k++ {
-						parts = append(parts, row[k].String())
-					}
-					alignment[len(cells)] = tablewriter.ALIGN_LEFT
-					cells = append(cells, strings.Join(parts, ","))
+					cells[columnIndex] += "," + value.String()
+					alignment[columnIndex] = tablewriter.ALIGN_LEFT
 				}
-
-				i = j
 			}
 
 			tw.Append(cells)
diff --git a/row.go b/row.go
index bf051ea..adf2a90 100644
--- a/row.go
+++ b/row.go
@@ -568,7 +568,7 @@ func reconstructFuncOfRepeated(columnIndex int16, node Node) (int16, reconstruct
 
 func reconstructRepeated(columnIndex, rowLength int16, levels levels, row Row, do func(levels, Row) (Row, error)) (Row, error) {
 	if !row.startsWith(columnIndex) {
-		return row, fmt.Errorf("row is missing repeated column %d", columnIndex)
+		return row, fmt.Errorf("row is missing repeated column %d: %+v", columnIndex, row)
 	}
 	if len(row) < int(rowLength) {
 		return row, fmt.Errorf("expected repeated column %d to have at least %d values but got %d", columnIndex, rowLength, len(row))


Test patch:
diff --git a/buffer_test.go b/buffer_test.go
index faaffa0..9b5154e 100644
--- a/buffer_test.go
+++ b/buffer_test.go
@@ -8,6 +8,7 @@ import (
 	"math/rand"
 	"reflect"
 	"sort"
+	"strconv"
 	"testing"
 	"testing/quick"
 	"time"
@@ -508,6 +509,50 @@ func TestBufferRoundtripNestedRepeatedPointer(t *testing.T) {
 	}
 }
 
+func TestBufferSeekToRow(t *testing.T) {
+	type B struct {
+		I int
+		C []string
+	}
+	type A struct {
+		B []B
+	}
+
+	buffer := parquet.NewBuffer()
+	var objs []A
+	for i := 0; i < 2; i++ {
+		o := A{
+			B: []B{
+				{I: i, C: []string{"foo", strconv.Itoa(i)}},
+				{I: i + 1, C: []string{"bar", strconv.Itoa(i + 1)}},
+			},
+		}
+		buffer.Write(&o)
+		objs = append(objs, o)
+	}
+
+	buf := new(bytes.Buffer)
+	w := parquet.NewWriter(buf)
+	w.WriteRowGroup(buffer)
+	w.Flush()
+	w.Close()
+
+	file := bytes.NewReader(buf.Bytes())
+	r := parquet.NewReader(file)
+
+	i := 1
+	o := new(A)
+	if err := r.SeekToRow(int64(i)); err != nil {
+		t.Fatal(err)
+	}
+	if err := r.Read(o); err != nil {
+		t.Fatal(err)
+	}
+	if !reflect.DeepEqual(*o, objs[i]) {
+		t.Errorf("points mismatch at row index %d: want=%v got=%v", i, objs[i], o)
+	}
+}
+
 type benchmarkBufferRowType struct {
 	ID    [16]byte `parquet:"id,uuid"`
 	Value float64  `parquet:"value"`