1package fileio23import (4 "context"5 "errors"6 "io"7 "math"8 "math/big"9 "os"10 "strings"11 "time"1213 "github.com/parquet-go/parquet-go"14 "github.com/parquet-go/parquet-go/deprecated"15 "github.com/parquet-go/parquet-go/format"1617 "src.rest/cat/012/internal/sheet"18)1920// parquetColumn is how a leaf column's values become cells.21type parquetColumn struct {22 logical format.LogicalTypeValue23 kind parquet.Kind24 // repeated is set on a list's column, whose rows hold several25 // values; sources read it a row at a time (sourceparquet.go).26 repeated bool27 // interval is set on a column of the INTERVAL converted type: 1228 // bytes of months, days and milliseconds.29 interval bool30}3132// parquetColumnOf is how a column of type t becomes cells.33func parquetColumnOf(t parquet.Type) parquetColumn {34 c := parquetColumn{kind: t.Kind()}35 if lt := t.LogicalType(); lt != nil {36 c.logical = lt.Value37 }38 if ct := t.ConvertedType(); ct != nil && *ct == deprecated.Interval {39 c.interval = true40 }41 return c42}4344// importParquet reads every leaf column, named by its path, into a45// sheet with a header row. Repeated values in a row are joined with46// commas. Rows past the sheet's last are counted from the file's47// metadata, not read.48func importParquet(ctx context.Context, name string, opt Options) (*Result, error) {49 prog := opt.Progress50 f, err := os.Open(name)51 if err != nil {52 return nil, err53 }54 defer f.Close()55 defer context.AfterFunc(ctx, func() { f.Close() })() // unblocks a read from a pipe56 st, err := f.Stat()57 if err != nil {58 return nil, err59 }60 pf, err := parquet.OpenFile(f, st.Size())61 if err != nil {62 return nil, err63 }64 b := newBuilder(ctx, opt.MaxCells)65 r := &parquetReader{66 b: b, cols: parquetHeader(b, pf.Schema()), prog: prog,67 row: 1, total: pf.NumRows(), buf: make([]parquet.Row, 256),68 }69 for _, rg := range pf.RowGroups() {70 if r.row >= sheet.MaxRows || b.isFull() {71 break72 }73 if err := r.group(ctx, rg); err != nil {74 return nil, err75 }76 }77 // The rows left out, as far as the file says there are.78 r.row = max(r.row, int(r.total)+1)79 b.fits(sheet.Addr{Row: r.row - 1})80 prog.setRows(r.row)81 var notes []string82 if r.lists > 0 {83 notes = append(notes, "repeated values joined with commas")84 }85 s, notes := b.finish(notes)86 return &Result{Sheet: s, Rows: r.row, Notes: notes}, nil87}8889// parquetHeader writes the columns' paths as a bold header row and90// returns how each column's values become cells.91func parquetHeader(b *builder, schema *parquet.Schema) []parquetColumn {92 paths := schema.Columns()93 cols := make([]parquetColumn, len(paths))94 header := sheet.Style{Bold: true}95 for i, path := range paths {96 b.text(sheet.Addr{Col: i}, strings.Join(path, "."), sheet.Format{}, header)97 if leaf, ok := schema.Lookup(path...); ok && leaf.Node != nil {98 cols[i] = parquetColumnOf(leaf.Node.Type())99 }100 }101 return cols102}103104// parquetReader holds the state of one import.105type parquetReader struct {106 b *builder107 cols []parquetColumn108 prog *Progress109 row int // the sheet row of the next file row110 total int64 // rows in the file111 lists int // rows with repeated values112 buf []parquet.Row113}114115// group reads a row group, a batch at a time, until it ends or the116// sheet is full.117func (r *parquetReader) group(ctx context.Context, rg parquet.RowGroup) error {118 rows := rg.Rows()119 defer rows.Close()120 for r.row < sheet.MaxRows && !r.b.isFull() {121 if err := ctx.Err(); err != nil {122 return err123 }124 n, err := rows.ReadRows(r.buf)125 for _, row := range r.buf[:n] {126 if putParquetRow(r.b, r.row, row, r.cols) {127 r.lists++128 }129 r.row++130 }131 r.prog.setRows(r.row)132 r.prog.setFrac(int64(r.row-1), r.total)133 if errors.Is(err, io.EOF) {134 return nil135 }136 if err != nil {137 return err138 }139 if n == 0 {140 return nil141 }142 }143 return nil144}145146// putParquetRow stores one row, reporting whether a column repeated.147func putParquetRow(b *builder, row int, r parquet.Row, cols []parquetColumn) (repeated bool) {148 if row >= sheet.MaxRows {149 b.fits(sheet.Addr{Row: row})150 return false151 }152 eachParquetColumn(r, func(col int, vals []parquet.Value) {153 if col < 0 || col >= len(cols) {154 return155 }156 c, rep := parquetCells(vals, cols[col])157 c.put(b, sheet.Addr{Col: col, Row: row})158 repeated = repeated || rep159 })160 return repeated161}162163// eachParquetColumn calls fn with each column's values in a row: they164// arrive grouped by column, and a repeated column has several.165func eachParquetColumn(r parquet.Row, fn func(col int, vals []parquet.Value)) {166 for i := 0; i < len(r); {167 col := r[i].Column()168 j := i + 1169 for j < len(r) && r[j].Column() == col {170 j++171 }172 fn(col, r[i:j])173 i = j174 }175}176177// parquetText is a value as text, for byte arrays and joined lists.178func parquetText(v parquet.Value, c parquetColumn) string {179 switch c.kind {180 case parquet.ByteArray, parquet.FixedLenByteArray:181 if _, isDec := c.logical.(*format.DecimalType); !isDec {182 return string(v.ByteArray())183 }184 }185 return v.String()186}187188func timeUnit(u format.TimeUnit) time.Duration {189 if u.Value == nil {190 return time.Millisecond191 }192 return u.Value.Duration()193}194195func timestamp(n int64, u format.TimeUnit) time.Time {196 d := timeUnit(u)197 return time.Unix(0, 0).UTC().Add(time.Duration(n) * d)198}199200func timeInt(v parquet.Value) int64 {201 if v.Kind() == parquet.Int32 {202 return int64(v.Int32())203 }204 return v.Int64()205}206207// decimal reads a DECIMAL stored as an integer or as big-endian bytes.208func decimal(v parquet.Value, scale int32) (float64, bool) {209 var unscaled float64210 switch v.Kind() {211 case parquet.Int32:212 unscaled = float64(v.Int32())213 case parquet.Int64:214 unscaled = float64(v.Int64())215 case parquet.ByteArray, parquet.FixedLenByteArray:216 raw := v.ByteArray()217 if len(raw) == 0 {218 return 0, false219 }220 n := new(big.Int).SetBytes(raw)221 if raw[0]&0x80 != 0 { // two's complement222 n.Sub(n, new(big.Int).Lsh(big.NewInt(1), uint(8*len(raw))))223 }224 unscaled, _ = new(big.Float).SetInt(n).Float64()225 default:226 return 0, false227 }228 return unscaled / math.Pow10(int(scale)), true229}