Add helper to count documents

1 files changed, 30 insertions, 0 deletions

diff --git a/dataset/parquet.go b/dataset/parquet.go
index f08c332..6710260 100644
--- a/dataset/parquet.go
+++ b/dataset/parquet.go
@@ -46,6 +46,22 @@ func NewParquetReader(name string) (*ParquetReader, error) {
 	return r, nil
 }
 
+func (p *ParquetReader) Num() (int, error) {
+	r := 0
+
+	for _, name := range p.shards {
+		n, err := numRows(name)
+
+		if err != nil {
+			return 0, err
+		}
+
+		r += int(n)
+	}
+
+	return r, nil
+}
+
 func (p *ParquetReader) Err() error {
 	return p.err
 }
@@ -147,3 +163,17 @@ func (p *ParquetReader) read(name string) iter.Seq[string] {
 		}
 	}
 }
+
+func numRows(name string) (int64, error) {
+	var reader *file.Reader
+
+	if r, err := file.OpenParquetFile(name, false); err != nil {
+		return 0, err
+	} else {
+		reader = r
+
+		defer reader.Close()
+	}
+
+	return reader.NumRows(), nil
+}