summaryrefslogtreecommitdiff
path: root/dataset/parquet.go
diff options
context:
space:
mode:
Diffstat (limited to 'dataset/parquet.go')
-rw-r--r--dataset/parquet.go30
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
+}