[Go][Parquet] unpack32Avx2 returns 0 if n is equal to 0

apache · Aug 9, 2024 · 1af4ee9 · 1af4ee9
1 parent 8e034c7
commit 1af4ee9
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 4 deletions.
diff --git a/go/parquet/internal/utils/bit_packing_avx2_amd64.go b/go/parquet/internal/utils/bit_packing_avx2_amd64.go
@@ -39,6 +39,10 @@ func unpack32Avx2(in io.Reader, out []uint32, nbits int) int {
 
 	n := batch * nbits / 8
 
+	if n == 0 {
+		return 0
+	}
+
 	buffer := bufferPool.Get().(*bytes.Buffer)
 	defer bufferPool.Put(buffer)
 	buffer.Reset()

diff --git a/go/parquet/pqarrow/file_reader_test.go b/go/parquet/pqarrow/file_reader_test.go
@@ -401,16 +401,28 @@ func TestReadParquetFile(t *testing.T) {
 	}()
 
 	arrowRdr, err := pqarrow.NewFileReader(rdr, pqarrow.ArrowReadProperties{
-		Parallel:  false,
+		Parallel:  true,
 		BatchSize: 0,
 	}, mem)
 	if err != nil {
 		t.Errorf("unexpected error: %v", err)
 	}
 
-	_, err = arrowRdr.ReadTable(ctx)
+	table, err := arrowRdr.ReadTable(ctx)
+	defer table.Release()
 
-	if err == nil {
-		t.Errorf("expected error: %v", err)
+	if err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+
+	assert.Equal(t, table.NumCols(), int64(1))
+	assert.Equal(t, table.NumRows(), int64(21186))
+
+	for i := 0; i < int(table.NumCols()); i++ {
+		col := table.Column(i)
+		assert.Equal(t, col.Len(), int(table.NumRows()))
+		for j := 0; j < col.Len(); j++ {
+			assert.Equal(t, col.Data().Chunk(0).(*array.Uint16).Value(j), uint16(0))
+		}
 	}
 }