From 1798fbe3881ec0f4abde416771b88285191cbd16 Mon Sep 17 00:00:00 2001
From: Jonas Knobloch <jonas.knobloch@t-online.de>
Date: Fri, 3 Apr 2026 20:32:02 +0200
Subject: Add helper to count documents

---
 dataset/parquet.go | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

(limited to 'dataset/parquet.go')

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
+}
-- 
cgit v1.3.1