1package fileio23import (4 "bufio"5 "context"6 "errors"7 "os"8 "slices"910 "src.rest/cat/012/internal/sheet"11)1213// Views of a Parquet file stream: building one reads the columns it14// sorts and filters by once, writing the rows it keeps, in its order,15// to a file of row numbers (sourcesort.go); a page reads its row16// numbers there and fetches those rows.1718// View orders the file's rows as o says.19func (s *parquetSource) View(ctx context.Context, o SourceOrder) (SourceView, error) {20 if o.IsZero() {21 return pqPlainView{s}, nil22 }23 plan, err := newViewPlan(o, len(s.cols))24 if err != nil {25 return nil, err26 }27 f, err := newRowFile(s.dir)28 if err != nil {29 return nil, err30 }31 rf := &rowFile{f: f}32 if rf.n, err = plan.write(ctx, s, f); err != nil {33 rf.close()34 return nil, err35 }36 s.mu.Lock()37 s.views = append(s.views, rf)38 s.mu.Unlock()39 return &pqView{s: s, rf: rf}, nil40}4142// viewPlan is what building a streamed view reads: the columns it sorts43// and filters by, once each, and where each key and test finds its44// column among them.45type viewPlan struct {46 cols []int47 keys []int // positions in cols of the sort's columns48 desc []bool49 tests []viewTest50}5152type viewTest struct {53 at int // the column's position in cols54 pass func(sheet.Value, string) bool55}5657func newViewPlan(o SourceOrder, ncols int) (*viewPlan, error) {58 p := &viewPlan{}59 at := func(col int) (int, error) {60 if col < 0 || col >= ncols {61 return 0, errors.New("the source has no such column")62 }63 if i := slices.Index(p.cols, col); i >= 0 {64 return i, nil65 }66 p.cols = append(p.cols, col)67 return len(p.cols) - 1, nil68 }69 for _, k := range o.Sort {70 i, err := at(k.Col)71 if err != nil {72 return nil, err73 }74 p.keys, p.desc = append(p.keys, i), append(p.desc, k.Desc)75 }76 for _, f := range o.Filter {77 i, err := at(f.Col)78 if err != nil {79 return nil, err80 }81 p.tests = append(p.tests, viewTest{at: i, pass: f.Cond.Test()})82 }83 return p, nil84}8586// passes reports whether a row's values pass every test, each on the87// value and the text it shows, as a sheet's filter tests them.88func (p *viewPlan) passes(vals []sheet.LiveCell) bool {89 for _, t := range p.tests {90 c := vals[t.at]91 if !t.pass(c.V, sheet.FormatText(c.V, c.F)) {92 return false93 }94 }95 return true96}9798// write streams the source, writing the numbers of the rows the view99// keeps to f in its order, and returns how many.100func (p *viewPlan) write(ctx context.Context, s *parquetSource, f *os.File) (int64, error) {101 w := bufio.NewWriterSize(f, 1<<20)102 var n int64103 var sorter *extSorter104 if len(p.keys) > 0 {105 sorter = newExtSorter(s.dir, p.desc)106 defer sorter.remove()107 }108 keys := make([]sheet.LiveCell, len(p.keys))109 var ferr error110 err := s.Scan(ctx, 0, p.cols, func(row int64, vals []sheet.LiveCell) bool {111 if !p.passes(vals) {112 return true113 }114 if sorter == nil {115 ferr = writeRow(w, row)116 n++117 return ferr == nil118 }119 for i, k := range p.keys {120 keys[i] = vals[k]121 }122 ferr = sorter.add(row, keys)123 return ferr == nil124 })125 if err = errors.Join(err, ferr); err != nil {126 return 0, err127 }128 if sorter != nil {129 if n, err = sorter.finish(w); err != nil {130 return 0, err131 }132 }133 return n, w.Flush()134}135136// pqView is a sorted or filtered view of a Parquet file.137type pqView struct {138 s *parquetSource139 rf *rowFile140}141142func (v *pqView) Rows() int64 { return v.rf.n }143144func (v *pqView) Page(ctx context.Context, from int64, n int) ([]int64, [][]sheet.LiveCell, error) {145 nums, err := v.rf.read(from, n)146 if err != nil || len(nums) == 0 {147 return nil, nil, err148 }149 rows, err := v.s.Fetch(ctx, nums, nil)150 if err != nil {151 return nil, nil, err152 }153 return nums, rows, nil154}155156func (v *pqView) Close() error {157 v.s.mu.Lock()158 i := slices.Index(v.s.views, v.rf)159 if i < 0 {160 v.s.mu.Unlock()161 return nil // closed with the source162 }163 v.s.views = slices.Delete(v.s.views, i, i+1)164 v.s.mu.Unlock()165 return v.rf.close()166}167168// pqPlainView is the file's rows in its own order.169type pqPlainView struct{ s *parquetSource }170171func (v pqPlainView) Rows() int64 { return v.s.Rows() }172173func (v pqPlainView) Page(ctx context.Context, from int64, n int) ([]int64, [][]sheet.LiveCell, error) {174 return scanPage(ctx, v.s, from, n)175}176177func (pqPlainView) Close() error { return nil }178179// scanPage reads n rows from row from of a source in its own order.180func scanPage(ctx context.Context, s Source, from int64, n int) ([]int64, [][]sheet.LiveCell, error) {181 if from < 0 || from >= s.Rows() || n <= 0 {182 return nil, nil, nil183 }184 var nums []int64185 var rows [][]sheet.LiveCell186 err := s.Scan(ctx, from, nil, func(row int64, vals []sheet.LiveCell) bool {187 nums, rows = append(nums, row), append(rows, slices.Clone(vals))188 return len(nums) < n189 })190 return nums, rows, err191}