1package fileio23import (4 "context"5 "fmt"6 "os"7 "path/filepath"8 "slices"9 "strings"10 "sync"11 "testing"12 "time"1314 "github.com/parquet-go/parquet-go"1516 "src.rest/cat/012/internal/sheet"17)1819// sale is a row of the Parquet sources tests read.20type sale struct {21 ID int64 `parquet:"id"`22 Cat *string `parquet:"cat,optional"`23 Amount *float64 `parquet:"amount,optional"`24 Day int32 `parquet:"day,date"`25 Tags []string `parquet:"tags,list"`26}2728var saleCats = []string{"North", "south", "East", "west", "Éire", "", "north"}2930// writeSales writes n rows in row groups of group rows.31func writeSales(t testing.TB, dir string, n, group int) string {32 t.Helper()33 name := filepath.Join(dir, "sales.parquet")34 f, err := os.Create(name)35 if err != nil {36 t.Fatal(err)37 }38 w := parquet.NewGenericWriter[sale](f, parquet.MaxRowsPerRowGroup(int64(group)), parquet.PageBufferSize(512))39 day := int32(time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC).Unix() / 86400)40 rows := make([]sale, n)41 for i := range rows {42 r := &rows[i]43 r.ID, r.Day = int64(i), day+int32(i%40)44 if i%9 != 4 {45 c := saleCats[i%len(saleCats)]46 r.Cat = &c47 }48 if i%11 != 7 {49 a := float64((i*37)%101) - 20.550 r.Amount = &a51 }52 if i%5 == 0 {53 r.Tags = []string{"a", fmt.Sprint(i)}54 }55 }56 if _, err := w.Write(rows); err != nil {57 t.Fatal(err)58 }59 if err := w.Close(); err != nil {60 t.Fatal(err)61 }62 f.Close()63 return name64}6566// cellsText is a row as its cells show, separated by |.67func cellsText(row []sheet.LiveCell) string {68 parts := make([]string, len(row))69 for i, c := range row {70 parts[i] = sheet.FormatText(c.V, c.F)71 }72 return strings.Join(parts, "|")73}7475// imported is a file imported as a sheet, which sources must read as76// the import does, and the rows under its header.77func imported(t *testing.T, name string, opt Options) (*sheet.Sheet, sheet.Rect) {78 t.Helper()79 res, err := Import(context.Background(), name, opt)80 if err != nil {81 t.Fatal(err)82 }83 used, _ := res.Sheet.UsedRange()84 used.From.Row = 185 return res.Sheet, used86}8788// sheetRows are the rows of r a sheet shows, the filter's hidden ones89// left out.90func sheetRows(s *sheet.Sheet, r sheet.Rect) []string {91 var out []string92 for row := r.From.Row; row <= r.To.Row; row++ {93 if s.RowHidden(row) {94 continue95 }96 parts := make([]string, 0, r.To.Col+1)97 for c := r.From.Col; c <= r.To.Col; c++ {98 parts = append(parts, shown(s, sheet.Addr{Col: c, Row: row}))99 }100 out = append(out, strings.Join(parts, "|"))101 }102 return out103}104105// viewRows pages through a whole view, n rows at a time.106func viewRows(t *testing.T, v SourceView, n int) ([]int64, []string) {107 t.Helper()108 var nums []int64109 var out []string110 for from := int64(0); from < v.Rows(); from += int64(n) {111 ns, rows, err := v.Page(context.Background(), from, n)112 if err != nil {113 t.Fatal(err)114 }115 if len(ns) != min(n, int(v.Rows()-from)) {116 t.Fatalf("page at %d has %d rows", from, len(ns))117 }118 for i, r := range rows {119 nums, out = append(nums, ns[i]), append(out, cellsText(r))120 }121 }122 return nums, out123}124125// orders are the sorts and filters every source is checked with,126// against a sheet sorting and filtering the same rows.127var orders = []SourceOrder{128 {},129 {Sort: []SourceSort{{Col: 1}}},130 {Sort: []SourceSort{{Col: 1, Desc: true}}},131 {Sort: []SourceSort{{Col: 2, Desc: true}, {Col: 1}}},132 {Filter: []SourceFilter{{Col: 2, Cond: sheet.Condition{Op: sheet.CondGreater, Arg: "30"}}}},133 {Filter: []SourceFilter{{Col: 1, Cond: sheet.Condition{Op: sheet.CondContains, Arg: "OR"}}}},134 {Filter: []SourceFilter{{Col: 1, Cond: sheet.Condition{Op: sheet.CondEmpty}}}},135 {Filter: []SourceFilter{{Col: 1, Cond: sheet.Condition{Op: sheet.CondNotEqual, Arg: "north"}}}},136 {Filter: []SourceFilter{{Col: 1, Cond: sheet.Condition{Op: sheet.CondStartsWith, Arg: "e"}}}},137 {Filter: []SourceFilter{{Col: 1, Cond: sheet.Condition{Op: sheet.CondEndsWith, Arg: "TH"}}}},138 {Filter: []SourceFilter{{Col: 1, Cond: sheet.Condition{Op: sheet.CondExactly, Arg: "EAST"}}}},139 {Filter: []SourceFilter{{Col: 2, Cond: sheet.Condition{Op: sheet.CondLessEq, Arg: "-10"}}},140 Sort: []SourceSort{{Col: 2}}},141 {Filter: []SourceFilter{{Col: 1, Cond: sheet.Condition{Op: sheet.CondNotEmpty}},142 {Col: 2, Cond: sheet.Condition{Op: sheet.CondGreaterEq, Arg: "0"}}},143 Sort: []SourceSort{{Col: 1, Desc: true}}},144}145146// checkOrders checks each order's view against the sheet the file147// imports as, sorted and filtered the same way.148func checkOrders(t *testing.T, src Source, name string, opt Options, list []SourceOrder) {149 t.Helper()150 for _, o := range list {151 s, r := imported(t, name, opt)152 if len(o.Sort) > 0 {153 var keys []sheet.SortKey154 for _, k := range o.Sort {155 keys = append(keys, sheet.SortKey{Col: k.Col, Desc: k.Desc})156 }157 s.SortRange(r, keys)158 }159 if len(o.Filter) > 0 {160 s.CreateFilter(sheet.Rect{To: r.To})161 for _, f := range o.Filter {162 s.FilterColumn(f.Col, sheet.Criteria{Cond: f.Cond})163 }164 }165 want := sheetRows(s, r)166 v, err := src.View(context.Background(), o)167 if err != nil {168 t.Fatalf("%+v: %v", o, err)169 }170 _, got := viewRows(t, v, 37)171 if !slices.Equal(got, want) {172 i := 0173 for i < min(len(got), len(want)) && got[i] == want[i] {174 i++175 }176 t.Errorf("%+v: %d rows, want %d; from row %d\n got %q\nwant %q", o, len(got), len(want), i, head(got[i:]), head(want[i:]))177 }178 if err := v.Close(); err != nil {179 t.Error(err)180 }181 }182}183184func head(rows []string) []string { return rows[:min(len(rows), 12)] }185186func TestParquetSource(t *testing.T) {187 dir := t.TempDir()188 name := writeSales(t, dir, 1000, 128)189 tmp := t.TempDir()190 src, err := OpenSource(context.Background(), SourceSpec{Path: name, TempDir: tmp})191 if err != nil {192 t.Fatal(err)193 }194 defer src.Close()195 if src.Rows() != 1000 {196 t.Fatalf("%d rows", src.Rows())197 }198 var cols []string199 for _, c := range src.Columns() {200 cols = append(cols, fmt.Sprintf("%s:%v:%v", c.Name, c.Numeric, c.Format.Kind == sheet.FmtDate))201 }202 if got := strings.Join(cols, " "); got != "id:true:false cat:false:false amount:true:false day:true:true tags.list.element:false:false" {203 t.Errorf("columns %s", got)204 }205 s, r := imported(t, name, Options{})206 want := sheetRows(s, r)207208 // Scan from part way, two columns in another order.209 var got []string210 err = src.Scan(context.Background(), 500, []int{2, 0}, func(row int64, vals []sheet.LiveCell) bool {211 got = append(got, fmt.Sprintf("%d:%s", row, cellsText(vals)))212 return row < 510213 })214 if err != nil {215 t.Fatal(err)216 }217 for i, g := range got {218 parts := strings.Split(want[500+i], "|")219 if w := fmt.Sprintf("%d:%s|%s", 500+i, parts[2], parts[0]); g != w {220 t.Errorf("scan %q, want %q", g, w)221 }222 }223 if len(got) != 11 {224 t.Errorf("scan read %d rows", len(got))225 }226227 // Fetch rows in any order, across row groups.228 nums := []int64{999, 3, 500, 128, 127, 3}229 rows, err := src.Fetch(context.Background(), nums, nil)230 if err != nil {231 t.Fatal(err)232 }233 for i, n := range nums {234 if g := cellsText(rows[i]); g != want[n] {235 t.Errorf("row %d is %q, want %q", n, g, want[n])236 }237 }238 checkOrders(t, src, name, Options{}, orders)239 if left, _ := os.ReadDir(tmp); len(left) != 0 {240 t.Errorf("views left %d files", len(left))241 }242}243244// TestSourceSortSpills sorts in runs spilled to files and merged.245func TestSourceSortSpills(t *testing.T) {246 defer func(n int) { sortRunBytes = n }(sortRunBytes)247 sortRunBytes = 4096248 name := writeSales(t, t.TempDir(), 3000, 1000)249 tmp := t.TempDir()250 src, err := OpenSource(context.Background(), SourceSpec{Path: name, TempDir: tmp})251 if err != nil {252 t.Fatal(err)253 }254 checkOrders(t, src, name, Options{}, orders[1:4])255 v, err := src.View(context.Background(), orders[2])256 if err != nil {257 t.Fatal(err)258 }259 if left, _ := os.ReadDir(tmp); len(left) != 1 {260 t.Errorf("a view keeps %d files, want its row order", len(left))261 }262 if err := src.Close(); err != nil {263 t.Fatal(err)264 }265 if left, _ := os.ReadDir(tmp); len(left) != 0 {266 t.Errorf("closing left %d files", len(left))267 }268 if err := v.Close(); err != nil {269 t.Errorf("closing a view after its source: %v", err)270 }271}272273// TestSourceConcurrent reads a source from several goroutines at once.274func TestSourceConcurrent(t *testing.T) {275 dir := t.TempDir()276 for _, name := range []string{writeSales(t, dir, 2000, 300), writeSalesDB(t, dir, 2000, false)} {277 src, err := OpenSource(context.Background(), SourceSpec{Path: name, Table: "sales", TempDir: dir})278 if err != nil {279 t.Fatal(err)280 }281 sum := func() float64 {282 total := 0.0283 err := src.Scan(context.Background(), 0, []int{2}, func(_ int64, v []sheet.LiveCell) bool {284 total += v[0].V.Num285 return true286 })287 if err != nil {288 t.Error(err)289 }290 return total291 }292 want := sum()293 var wg sync.WaitGroup294 for i := range 8 {295 wg.Go(func() {296 if got := sum(); got != want {297 t.Errorf("%s: sum %v, want %v", name, got, want)298 }299 if _, err := src.Fetch(context.Background(), []int64{int64(i * 200), 1999}, nil); err != nil {300 t.Error(err)301 }302 v, err := src.View(context.Background(), SourceOrder{Sort: []SourceSort{{Col: 2}}})303 if err != nil {304 t.Error(err)305 return306 }307 if _, _, err := v.Page(context.Background(), 1500, 60); err != nil {308 t.Error(err)309 }310 v.Close()311 })312 }313 wg.Wait()314 src.Close()315 }316}317318func TestSourceKind(t *testing.T) {319 for spec, ok := range map[SourceSpec]bool{320 {Path: "a.parquet"}: true, {Path: "a.sqlite"}: true, {Path: "a.db"}: true,321 {Path: "a.csv"}: false, {Path: "a.bin", Format: "Parquet"}: true, {Path: "a.parquet", Format: "CSV"}: false,322 } {323 if _, err := SourceKind(spec); (err == nil) != ok {324 t.Errorf("%+v: %v", spec, err)325 }326 }327 if _, err := OpenSource(context.Background(), SourceSpec{Path: filepath.Join(t.TempDir(), "none.parquet")}); err == nil {328 t.Error("opened a file that isn't there")329 }330}