Repository: segmentio/parquet-go (go)

Issue:
[Bug] parquet.readFile only reads 1024 rows
This file (zipped parquet file):
[debug.zip](https://github.com/segmentio/parquet-go/files/10542463/debug.zip) has more than 1024 rows (around 9K).

But if I read it using the following code:

```
type RowType struct {
	Name string `parquet:",optional"`
	Data []byte `parquet:",optional"`
}

rows, err := parquet.ReadFile[RowType](file)
```

and print out the number of rows, I only see that there are 1024 rows.
Tools like the `pandas` library from Python don't have this issue.

Pull request:
Fix Generic Read methods to read all rows up to length of argument
This fixes the generic Read methods where only one page worth of rows was read instead of all rows up to length of the argument slice. This affected at least `parquet.ReadFile`, `parquet.Read`, and `GenericReader.Read`.

You can see a failed test here https://github.com/segmentio/parquet-go/actions/runs/4515595007/jobs/7953065166 without the fix with a new test case using larger sizes here https://github.com/segmentio/parquet-go/pull/489/files#diff-306797de164c962fe5b992da67ab7ff245b8a8e7c76ca21e350ee08c7071d39eR494.

I added a new generic Read path test with multiple RowGroups & multiple Pages to increase the test coverage around this.

I also noted that the TestWriter writerTests have been skipped since using Github Actions because parquet-tools are not installed in Github CI.

Fixes #469, #471, & possibly others.

Patch:
diff --git a/Makefile b/Makefile
index b58b15c..137660e 100644
--- a/Makefile
+++ b/Makefile
@@ -9,4 +9,4 @@ format:
 	differ gofmt -w .
 
 test:
-	go test -trimpath -race ./...
+	go test -v -trimpath -race -tags= ./...
diff --git a/reader_go18.go b/reader_go18.go
index bbbc225..42130f7 100644
--- a/reader_go18.go
+++ b/reader_go18.go
@@ -103,6 +103,13 @@ func (r *GenericReader[T]) Reset() {
 	r.base.Reset()
 }
 
+// Read reads the next rows from the reader into the given rows slice up to len(rows).
+//
+// The returned values are safe to reuse across Read calls and do not share
+// memory with the reader's underlying page buffers.
+//
+// The method returns the number of rows read and io.EOF when no more rows
+// can be read from the reader.
 func (r *GenericReader[T]) Read(rows []T) (int, error) {
 	return r.read(r, rows)
 }
@@ -127,6 +134,13 @@ func (r *GenericReader[T]) Close() error {
 	return r.base.Close()
 }
 
+// readRows reads the next rows from the reader into the given rows slice up to len(rows).
+//
+// The returned values are safe to reuse across readRows calls and do not share
+// memory with the reader's underlying page buffers.
+//
+// The method returns the number of rows read and io.EOF when no more rows
+// can be read from the reader.
 func (r *GenericReader[T]) readRows(rows []T) (int, error) {
 	if cap(r.base.rowbuf) < len(rows) {
 		r.base.rowbuf = make([]Row, len(rows))
@@ -134,17 +148,26 @@ func (r *GenericReader[T]) readRows(rows []T) (int, error) {
 		r.base.rowbuf = r.base.rowbuf[:len(rows)]
 	}
 
-	n, err := r.base.ReadRows(r.base.rowbuf)
-	if n > 0 {
-		schema := r.base.Schema()
-
-		for i, row := range r.base.rowbuf[:n] {
-			if err := schema.Reconstruct(&rows[i], row); err != nil {
-				return i, err
+	var n, nTotal int
+	var err error
+	for {
+		n, err = r.base.ReadRows(r.base.rowbuf)
+		if n > 0 {
+			schema := r.base.Schema()
+
+			for i, row := range r.base.rowbuf[:n] {
+				if err := schema.Reconstruct(&rows[nTotal+i], row); err != nil {
+					return i, err
+				}
 			}
 		}
+		nTotal += n
+		if n == 0 || nTotal == len(rows) {
+			break
+		}
 	}
-	return n, err
+
+	return nTotal, err
 }
 
 var (


Test patch:
diff --git a/parquet_go18_test.go b/parquet_go18_test.go
index cebf311..840d94d 100644
--- a/parquet_go18_test.go
+++ b/parquet_go18_test.go
@@ -326,6 +326,72 @@ func TestIssue377(t *testing.T) {
 	assertRowsEqual(t, rows, ods)
 }
 
+func TestReadFileGenericMultipleRowGroupsMultiplePages(t *testing.T) {
+	type MyRow struct {
+		ID    [16]byte `parquet:"id,delta,uuid"`
+		File  string   `parquet:"file,dict,zstd"`
+		Index int64    `parquet:"index,delta,zstd"`
+	}
+
+	numRows := 20_000
+	maxPageBytes := 5000
+
+	tmp, err := os.CreateTemp("/tmp", "*.parquet")
+	if err != nil {
+		t.Fatal("os.CreateTemp: ", err)
+	}
+	path := tmp.Name()
+	defer os.Remove(path)
+	t.Log("TestReadFileMultiplePages:", path)
+
+	// The page buffer size ensures we get multiple pages out of this example.
+	w := parquet.NewGenericWriter[MyRow](tmp, parquet.PageBufferSize(maxPageBytes))
+	// Need to write 1 row at a time here as writing many at once disregards PageBufferSize option.
+	for i := 0; i < numRows; i++ {
+		row := MyRow{
+			ID:    [16]byte{15: byte(i)},
+			File:  "hi" + fmt.Sprint(i),
+			Index: int64(i),
+		}
+		_, err := w.Write([]MyRow{row})
+		if err != nil {
+			t.Fatal("w.Write: ", err)
+		}
+		// Flush writes rows as row group. 4 total (20k/5k) in this file.
+		if (i+1)%maxPageBytes == 0 {
+			err = w.Flush()
+			if err != nil {
+				t.Fatal("w.Flush: ", err)
+			}
+		}
+	}
+	err = w.Close()
+	if err != nil {
+		t.Fatal("w.Close: ", err)
+	}
+	err = tmp.Close()
+	if err != nil {
+		t.Fatal("tmp.Close: ", err)
+	}
+
+	rows, err := parquet.ReadFile[MyRow](path)
+	if err != nil {
+		t.Fatal("parquet.ReadFile: ", err)
+	}
+	if len(rows) != numRows {
+		t.Fatal("parquet.ReadFile returned", len(rows), "rows instead of", numRows)
+	}
+	for i, row := range rows {
+		id := [16]byte{15: byte(i)}
+		file := "hi" + fmt.Sprint(i)
+		index := int64(i)
+
+		if row.ID != id || row.File != file || row.Index != index {
+			t.Error("rows mismatch at ", i, "got ", row)
+		}
+	}
+}
+
 func assertRowsEqual[T any](t *testing.T, rows1, rows2 []T) {
 	if !reflect.DeepEqual(rows1, rows2) {
 		t.Error("rows mismatch")
diff --git a/parquet_test.go b/parquet_test.go
index 04d5001..5147b06 100644
--- a/parquet_test.go
+++ b/parquet_test.go
@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+
 	"github.com/segmentio/parquet-go"
 	"github.com/segmentio/parquet-go/deprecated"
 	"github.com/segmentio/parquet-go/internal/quick"
@@ -490,6 +491,7 @@ var quickCheckConfig = quick.Config{
 	Sizes: []int{
 		0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
 		10, 20, 30, 40, 50, 123,
+		4096 + 1,
 	},
 }
 
diff --git a/reader_go18_test.go b/reader_go18_test.go
index d44b4aa..e3902a6 100644
--- a/reader_go18_test.go
+++ b/reader_go18_test.go
@@ -118,7 +118,7 @@ func TestIssue400(t *testing.T) {
 	}
 
 	r := parquet.NewGenericReader[A](bytes.NewReader(b.Bytes()))
-	values := make([]A, 2)
+	values := make([]A, 1)
 	_, err = r.Read(values)
 	if err != nil {
 		t.Fatal(err)
diff --git a/writer_test.go b/writer_test.go
index 4b6b482..348e352 100644
--- a/writer_test.go
+++ b/writer_test.go
@@ -12,6 +12,7 @@ import (
 	"github.com/hexops/gotextdiff"
 	"github.com/hexops/gotextdiff/myers"
 	"github.com/hexops/gotextdiff/span"
+
 	"github.com/segmentio/parquet-go"
 	"github.com/segmentio/parquet-go/compress"
 )
@@ -50,7 +51,7 @@ func generateParquetFile(rows rows, options ...parquet.WriterOption) ([]byte, er
 	defer tmp.Close()
 	path := tmp.Name()
 	defer os.Remove(path)
-	//fmt.Println(path)
+	// fmt.Println(path)
 
 	writerOptions := []parquet.WriterOption{parquet.PageBufferSize(20)}
 	writerOptions = append(writerOptions, options...)
@@ -453,7 +454,7 @@ value 2: R:0 D:0 V:1.0
 
 func TestWriter(t *testing.T) {
 	if !hasParquetTools() {
-		t.Skip("parquet-tools are not installed")
+		t.Skip("Skipping TestWriter writerTests because parquet-tools are not installed in Github CI. FIXME.") // TODO
 	}
 
 	for _, test := range writerTests {