1package fileio23import (4 "context"5 "errors"6 "io"7 "os"8 "slices"9 "strings"10 "sync"1112 "github.com/parquet-go/parquet-go"1314 "src.rest/cat/012/internal/sheet"15)1617// parquetSource is a Parquet file read in place. Each read opens its18// own readers of the column chunks it needs (the file is read with19// ReadAt), so reads on several goroutines don't meet.20type parquetSource struct {21 f *os.File22 pf *parquet.File23 groups []parquet.RowGroup24 starts []int64 // each row group's first row, and the rows in all last25 cols []parquetColumn26 meta []SourceColumn27 dir string2829 mu sync.Mutex30 views []*rowFile // what views wrote, removed by Close31}3233func openParquetSource(spec SourceSpec) (*parquetSource, error) {34 f, err := os.Open(spec.Path)35 if err != nil {36 return nil, err37 }38 st, err := f.Stat()39 if err != nil {40 f.Close()41 return nil, err42 }43 pf, err := parquet.OpenFile(f, st.Size())44 if err != nil {45 f.Close()46 return nil, err47 }48 s := &parquetSource{f: f, pf: pf, groups: pf.RowGroups(), dir: spec.TempDir}49 s.starts = make([]int64, len(s.groups)+1)50 for i, g := range s.groups {51 s.starts[i+1] = s.starts[i] + g.NumRows()52 }53 schema := pf.Schema()54 for _, path := range schema.Columns() {55 var c parquetColumn56 if leaf, ok := schema.Lookup(path...); ok && leaf.Node != nil {57 c = parquetColumnOf(leaf.Node.Type())58 c.repeated = leaf.MaxRepetitionLevel > 059 }60 f, num := parquetFormat(c)61 s.cols = append(s.cols, c)62 s.meta = append(s.meta, SourceColumn{Name: strings.Join(path, "."), Format: f, Numeric: num})63 }64 return s, nil65}6667func (s *parquetSource) Columns() []SourceColumn { return slices.Clone(s.meta) }6869func (s *parquetSource) Rows() int64 { return s.starts[len(s.starts)-1] }7071func (s *parquetSource) Close() error {72 s.mu.Lock()73 views := s.views74 s.views = nil75 s.mu.Unlock()76 var errs []error77 for _, v := range views {78 errs = append(errs, v.close())79 }80 return errors.Join(append(errs, s.f.Close())...)81}8283// group is the row group holding row.84func (s *parquetSource) group(row int64) int {85 i, found := slices.BinarySearch(s.starts, row)86 if !found {87 i--88 }89 return i90}9192// Scan streams rows from row from on.93func (s *parquetSource) Scan(ctx context.Context, from int64, cols []int, fn func(int64, []sheet.LiveCell) bool) error {94 cols = allCols(cols, len(s.cols))95 if err := checkCols(cols, len(s.cols)); err != nil {96 return err97 }98 rs := s.newRows(cols)99 defer rs.close()100 vals := make([]sheet.LiveCell, len(cols))101 for row := max(from, 0); row < s.Rows(); row++ {102 if row%1024 == 0 {103 if err := ctx.Err(); err != nil {104 return err105 }106 }107 if err := rs.read(row, vals); err != nil {108 return err109 }110 if !fn(row, vals) {111 return nil112 }113 }114 return nil115}116117// Fetch reads the rows numbered rows.118func (s *parquetSource) Fetch(ctx context.Context, rows []int64, cols []int) ([][]sheet.LiveCell, error) {119 cols = allCols(cols, len(s.cols))120 if err := checkCols(cols, len(s.cols)); err != nil {121 return nil, err122 }123 rs := s.newRows(cols)124 defer rs.close()125 out := make([][]sheet.LiveCell, len(rows))126 for _, i := range fetchOrder(rows) {127 if err := ctx.Err(); err != nil {128 return nil, err129 }130 row := rows[i]131 if row < 0 || row >= s.Rows() {132 return nil, errors.New("no such row in the source")133 }134 vals := make([]sheet.LiveCell, len(cols))135 if err := rs.read(row, vals); err != nil {136 return nil, err137 }138 out[i] = vals139 }140 return out, nil141}142143// pqRows reads rows of some columns, each through a cursor of its own.144type pqRows struct {145 s *parquetSource146 curs []*pqCursor147 buf []parquet.Value148}149150func (s *parquetSource) newRows(cols []int) *pqRows {151 rs := &pqRows{s: s}152 for _, c := range cols {153 rs.curs = append(rs.curs, &pqCursor{s: s, col: c, repeated: s.cols[c].repeated, group: -1})154 }155 return rs156}157158// read reads row's values into vals.159func (rs *pqRows) read(row int64, vals []sheet.LiveCell) error {160 for i, c := range rs.curs {161 var err error162 if rs.buf, err = c.row(row, rs.buf[:0]); err != nil {163 return err164 }165 cell, _ := parquetCells(rs.buf, rs.s.cols[c.col])166 vals[i] = cell.typed()167 }168 return nil169}170171func (rs *pqRows) close() {172 for _, c := range rs.curs {173 c.close()174 }175}176177// pqCursor reads one column's values a row at a time: from the page178// readers of its chunk in the row group being read, sought to a row179// when a read isn't the next one.180type pqCursor struct {181 s *parquetSource182 col int183 repeated bool184 group int // the row group being read, -1 before any185 next int64 // the row the next value starts186 pages parquet.Pages187 vals parquet.ValueReader188 buf []parquet.Value189 i, n int190}191192// skipRows is how far ahead a cursor reads its way rather than seeking.193const skipRows = 4096194195// row appends the values of row to out: one, or a repeated column's196// several (none for an empty list, read as null).197func (c *pqCursor) row(row int64, out []parquet.Value) ([]parquet.Value, error) {198 if err := c.seek(row); err != nil {199 return out, err200 }201 for first := true; ; first = false {202 if c.i == c.n {203 if err := c.fill(); err != nil {204 if errors.Is(err, io.EOF) && !first {205 break206 }207 return out, err208 }209 }210 v := c.buf[c.i]211 if !first && v.RepetitionLevel() == 0 {212 break213 }214 if c.repeated {215 v = v.Clone() // the next fill may reuse what it points at216 }217 out = append(out, v)218 c.i++219 }220 c.next = row + 1221 return out, nil222}223224// seek places the cursor at row: reading forward to it when it is a225// little ahead in the same group, else seeking the group's pages.226func (c *pqCursor) seek(row int64) error {227 if row == c.next && c.group >= 0 && row < c.s.starts[c.group+1] {228 return nil // the next row, as a scan reads them229 }230 g := c.s.group(row)231 if g == c.group && row >= c.next && row-c.next <= skipRows {232 var skip []parquet.Value233 for c.next < row {234 var err error235 if skip, err = c.row(c.next, skip[:0]); err != nil {236 return err237 }238 }239 return nil240 }241 c.close()242 pages := c.s.groups[g].ColumnChunks()[c.col].Pages()243 if err := pages.SeekToRow(row - c.s.starts[g]); err != nil {244 pages.Close()245 return err246 }247 c.pages, c.group, c.next, c.vals, c.i, c.n = pages, g, row, nil, 0, 0248 return nil249}250251// fill reads the next values of the group's pages, io.EOF past them.252func (c *pqCursor) fill() error {253 if c.buf == nil {254 c.buf = make([]parquet.Value, 1024)255 }256 for {257 if c.vals != nil {258 n, err := c.vals.ReadValues(c.buf)259 if n > 0 {260 c.i, c.n = 0, n261 return nil262 }263 if err != nil && !errors.Is(err, io.EOF) {264 return err265 }266 c.vals = nil267 }268 if c.pages == nil {269 return io.EOF270 }271 p, err := c.pages.ReadPage()272 if err != nil {273 return err274 }275 c.vals = p.Values()276 }277}278279func (c *pqCursor) close() {280 if c.pages != nil {281 c.pages.Close()282 }283 c.pages, c.vals, c.group, c.i, c.n = nil, nil, -1, 0, 0284}