diff options
Diffstat (limited to 'dataset')
| -rw-r--r-- | dataset/cmd/dataset/main.go | 4 | ||||
| -rw-r--r-- | dataset/parquet.go (renamed from dataset/reader.go) | 10 |
2 files changed, 7 insertions, 7 deletions
diff --git a/dataset/cmd/dataset/main.go b/dataset/cmd/dataset/main.go index 5fd633b..8bd892d 100644 --- a/dataset/cmd/dataset/main.go +++ b/dataset/cmd/dataset/main.go @@ -19,9 +19,9 @@ func main() { split := filepath.Join(root, "train") - var reader *dataset.Reader + var reader *dataset.ParquetReader - if r, err := dataset.NewReader(split); err != nil { + if r, err := dataset.NewParquetReader(split); err != nil { log.Fatal(err) } else { reader = r diff --git a/dataset/reader.go b/dataset/parquet.go index e882332..c80758f 100644 --- a/dataset/reader.go +++ b/dataset/parquet.go @@ -15,13 +15,13 @@ import ( "github.com/apache/arrow-go/v18/parquet/pqarrow" ) -type Reader struct { +type ParquetReader struct { shards []string batchSize int64 err error } -func NewReader(name string) (*Reader, error) { +func NewParquetReader(name string) (*ParquetReader, error) { var shards []string if matches, err := filepath.Glob(filepath.Join(name, "*.parquet")); err != nil { @@ -36,7 +36,7 @@ func NewReader(name string) (*Reader, error) { slices.Sort(shards) - r := &Reader{ + r := &ParquetReader{ shards: shards, batchSize: 1024, } @@ -44,11 +44,11 @@ func NewReader(name string) (*Reader, error) { return r, nil } -func (r *Reader) Err() error { +func (r *ParquetReader) Err() error { return r.err } -func (r *Reader) Texts(column string) iter.Seq[string] { +func (r *ParquetReader) Texts(column string) iter.Seq[string] { return func(yield func(string) bool) { for _, name := range r.shards { err := read(name, column, r.batchSize, yield) |