Add helper to count documents

4 files changed, 39 insertions, 0 deletions

diff --git a/dataset/file.go b/dataset/file.go
index 081fbdf..242172e 100644
--- a/dataset/file.go
+++ b/dataset/file.go
@@ -37,6 +37,10 @@ func NewFileReader(name, pattern string) (*FileReader, error) {
 	return f, nil
 }
 
+func (f *FileReader) Num() (int, error) {
+	return len(f.shards), nil
+}
+
 func (f *FileReader) Err() error {
 	return f.err
 }
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
+}
diff --git a/dataset/reader.go b/dataset/reader.go
index 6458010..31a25b8 100644
--- a/dataset/reader.go
+++ b/dataset/reader.go
@@ -4,5 +4,6 @@ import "iter"
 
 type Reader interface {
 	Texts() iter.Seq2[int, string]
+	Num() (int, error)
 	Err() error
 }
diff --git a/dataset/string.go b/dataset/string.go
index c1f7600..76f047b 100644
--- a/dataset/string.go
+++ b/dataset/string.go
@@ -20,6 +20,10 @@ func (s *StringReader) Texts() iter.Seq2[int, string] {
 	}
 }
 
+func (s *StringReader) Num() (int, error) {
+	return 1, nil
+}
+
 func (s *StringReader) Err() error {
 	return nil
 }