diff options
Diffstat (limited to 'dataset/parquet.go')
| -rw-r--r-- | dataset/parquet.go | 16 |
1 files changed, 9 insertions, 7 deletions
diff --git a/dataset/parquet.go b/dataset/parquet.go index c80758f..8b91458 100644 --- a/dataset/parquet.go +++ b/dataset/parquet.go @@ -16,9 +16,10 @@ import ( ) type ParquetReader struct { - shards []string - batchSize int64 - err error + shards []string + textColumn string + batchSize int64 + err error } func NewParquetReader(name string) (*ParquetReader, error) { @@ -37,8 +38,9 @@ func NewParquetReader(name string) (*ParquetReader, error) { slices.Sort(shards) r := &ParquetReader{ - shards: shards, - batchSize: 1024, + shards: shards, + textColumn: "text", + batchSize: 1024, } return r, nil @@ -48,10 +50,10 @@ func (r *ParquetReader) Err() error { return r.err } -func (r *ParquetReader) Texts(column string) iter.Seq[string] { +func (r *ParquetReader) Texts() iter.Seq[string] { return func(yield func(string) bool) { for _, name := range r.shards { - err := read(name, column, r.batchSize, yield) + err := read(name, r.textColumn, r.batchSize, yield) if errors.Is(err, stop) { return |