Rows past the limit were decoded only to be counted; the file's metadata and the table's row count already say how many there are. Queries, whose size isn't known, still step to the end but no longer scan the rows they drop. importParquet and importSQLite are split into a header, a query and a row-reading helper.
- commit
- ad230063629c3a4588b4b9ff9101c7fd4cb5a90b
- parent
- 7fcd326e053e39bada37f305ed65f973fb046a30
- author
- Dami <dami.etoile@gmail.com>
| M | internal/ | 108 | +++++++++++++++++++----------- |
| M | internal/ | 42 | ++++++++++++ |
| M | internal/ | 97 | +++++++++++++++++---------- |
| M | internal/ | 28 | ++++++++ |
@@ -24,7 +24,8 @@ type parquetColumn struct {
2424 2525 // importParquet reads every leaf column, named by its path, into a2626 // sheet with a header row. Repeated values in a row are joined with27-// commas.27+// commas. Rows past the sheet's last are counted from the file's28+// metadata, not read.2829 func importParquet(ctx context.Context, name string, prog *Progress) (*Result, error) {2930 f, err := os.Open(name)3031 if err != nil {@@ -40,10 +41,36 @@ func importParquet(ctx context.Context, name string, prog *Progress) (*Result, e
4041 if err != nil {4142 return nil, err4243 }43- schema := pf.Schema()44+ b := newBuilder()45+ r := &parquetReader{46+ b: b, cols: parquetHeader(b, pf.Schema()), prog: prog,47+ row: 1, total: pf.NumRows(), buf: make([]parquet.Row, 256),48+ }49+ for _, rg := range pf.RowGroups() {50+ if r.row >= sheet.MaxRows {51+ break52+ }53+ if err := r.group(ctx, rg); err != nil {54+ return nil, err55+ }56+ }57+ // The rows left out, as far as the file says there are.58+ r.row = max(r.row, int(r.total)+1)59+ b.fits(sheet.Addr{Row: r.row - 1})60+ prog.setRows(r.row)61+ var notes []string62+ if r.lists > 0 {63+ notes = append(notes, "repeated values joined with commas")64+ }65+ s, notes := b.finish(notes)66+ return &Result{Sheet: s, Rows: r.row, Notes: notes}, nil67+}68+69+// parquetHeader writes the columns' paths as a bold header row and70+// returns how each column's values become cells.71+func parquetHeader(b *builder, schema *parquet.Schema) []parquetColumn {4472 paths := schema.Columns()4573 cols := make([]parquetColumn, len(paths))46- b := newBuilder()4774 header := sheet.Style{Bold: true}4875 for i, path := range paths {4976 b.text(sheet.Addr{Col: i}, strings.Join(path, "."), sheet.Format{}, header)@@ -55,46 +82,49 @@ func importParquet(ctx context.Context, name string, prog *Progress) (*Result, e
5582 }5683 }5784 }85+ return cols86+}5887 59- total := pf.NumRows()60- row := 161- buf := make([]parquet.Row, 256)62- lists := 063- for _, rg := range pf.RowGroups() {64- rows := rg.Rows()65- for {66- if err := ctx.Err(); err != nil {67- rows.Close()68- return nil, err69- }70- n, err := rows.ReadRows(buf)71- for _, r := range buf[:n] {72- if putParquetRow(b, row, r, cols) {73- lists++74- }75- row++76- }77- prog.setRows(row)78- prog.setFrac(int64(row-1), total)79- if errors.Is(err, io.EOF) {80- break81- }82- if err != nil {83- rows.Close()84- return nil, err85- }86- if n == 0 {87- break88+// parquetReader holds the state of one import.89+type parquetReader struct {90+ b *builder91+ cols []parquetColumn92+ prog *Progress93+ row int // the sheet row of the next file row94+ total int64 // rows in the file95+ lists int // rows with repeated values96+ buf []parquet.Row97+}98+99+// group reads a row group, a batch at a time, until it ends or the100+// sheet is full.101+func (r *parquetReader) group(ctx context.Context, rg parquet.RowGroup) error {102+ rows := rg.Rows()103+ defer rows.Close()104+ for r.row < sheet.MaxRows {105+ if err := ctx.Err(); err != nil {106+ return err107+ }108+ n, err := rows.ReadRows(r.buf)109+ for _, row := range r.buf[:n] {110+ if putParquetRow(r.b, r.row, row, r.cols) {111+ r.lists++88112 }113+ r.row++114+ }115+ r.prog.setRows(r.row)116+ r.prog.setFrac(int64(r.row-1), r.total)117+ if errors.Is(err, io.EOF) {118+ return nil119+ }120+ if err != nil {121+ return err122+ }123+ if n == 0 {124+ return nil89125 }90- rows.Close()91- }92- var notes []string93- if lists > 0 {94- notes = append(notes, "repeated values joined with commas")95126 }96- s, notes := b.finish(notes)97- return &Result{Sheet: s, Rows: row, Notes: notes}, nil127+ return nil98128 }99129 100130 // putParquetRow stores one row, reporting whether a column repeated.@@ -8,6 +8,8 @@ import (
88 "time"99 1010 "github.com/parquet-go/parquet-go"11+12+ "github.com/FelineStateMachine/012/internal/sheet"1113 )1214 1315 type trip struct {@@ -66,6 +68,46 @@ func TestParquetImport(t *testing.T) {
6668 }6769 }6870 71+// TestParquetPastLimit reads rows up to the sheet's last and counts the72+// rest from the file's metadata.73+func TestParquetPastLimit(t *testing.T) {74+ type num struct {75+ N int32 `parquet:"n"`76+ }77+ rows := make([]num, sheet.MaxRows+4)78+ for i := range rows {79+ rows[i].N = int32(i + 1)80+ }81+ name := filepath.Join(t.TempDir(), "nums.parquet")82+ f, err := os.Create(name)83+ if err != nil {84+ t.Fatal(err)85+ }86+ w := parquet.NewGenericWriter[num](f, parquet.MaxRowsPerRowGroup(1000))87+ if _, err := w.Write(rows); err != nil {88+ t.Fatal(err)89+ }90+ if err := w.Close(); err != nil {91+ t.Fatal(err)92+ }93+ f.Close()94+95+ res, err := Import(context.Background(), name, Options{})96+ if err != nil {97+ t.Fatal(err)98+ }99+ last := sheet.Addr{Row: sheet.MaxRows - 1}100+ if g := shown(res.Sheet, last); g != "8191" {101+ t.Errorf("last row shows %q, want 8191", g)102+ }103+ if res.Rows != len(rows)+1 {104+ t.Errorf("rows %d, want %d", res.Rows, len(rows)+1)105+ }106+ if len(res.Notes) != 1 || res.Notes[0] != "only the first 8,192 rows fit; 5 rows left out" {107+ t.Errorf("notes %q", res.Notes)108+ }109+}110+69111 func TestParquetNotParquet(t *testing.T) {70112 name := filepath.Join(t.TempDir(), "x.parquet")71113 os.WriteFile(name, []byte("not parquet"), 0o644)@@ -107,25 +107,9 @@ func importSQLite(ctx context.Context, name string, opt Options) (*Result, error
107107 return nil, err108108 }109109 defer db.Close()110- query, total := opt.Query, 0111- if query == "" {112- table := opt.Table113- if table == "" {114- ts, err := tables(ctx, db)115- switch {116- case err != nil:117- return nil, err118- case len(ts) == 0:119- return nil, errors.New("the database has no tables")120- case len(ts) > 1:121- return nil, &ErrNeedTable{Tables: ts}122- }123- table = ts[0].Name124- }125- query = "SELECT * FROM " + quoteIdent(table)126- if err := db.QueryRowContext(ctx, "SELECT count(*) FROM "+quoteIdent(table)).Scan(&total); err != nil {127- return nil, sqliteErr(err)128- }110+ query, total, err := sqliteQuery(ctx, db, opt)111+ if err != nil {112+ return nil, err129113 }130114 rows, err := db.QueryContext(ctx, query)131115 if err != nil {@@ -142,28 +126,73 @@ func importSQLite(ctx context.Context, name string, opt Options) (*Result, error
142126 for c, name := range cols {143127 b.text(sheet.Addr{Col: c}, name, sheet.Format{}, header)144128 }145- vals := make([]any, len(cols))146- ptrs := make([]any, len(cols))129+ row, blobs, err := readSQLiteRows(ctx, rows, b, len(cols), total, opt.Progress)130+ if err != nil {131+ return nil, err132+ }133+ var notes []string134+ if blobs > 0 {135+ notes = append(notes, count(blobs, "binary value", "binary values")+" shown as sizes")136+ }137+ s, notes := b.finish(notes)138+ return &Result{Sheet: s, Rows: row, Notes: notes}, nil139+}140+141+// sqliteQuery is the query an import runs: the options' own, or all of142+// the table they name, or of the only table, with its row count; the143+// count is -1 for a query.144+func sqliteQuery(ctx context.Context, db *sql.DB, opt Options) (string, int, error) {145+ if opt.Query != "" {146+ return opt.Query, -1, nil147+ }148+ table := opt.Table149+ if table == "" {150+ ts, err := tables(ctx, db)151+ switch {152+ case err != nil:153+ return "", 0, err154+ case len(ts) == 0:155+ return "", 0, errors.New("the database has no tables")156+ case len(ts) > 1:157+ return "", 0, &ErrNeedTable{Tables: ts}158+ }159+ table = ts[0].Name160+ }161+ total := 0162+ if err := db.QueryRowContext(ctx, "SELECT count(*) FROM "+quoteIdent(table)).Scan(&total); err != nil {163+ return "", 0, sqliteErr(err)164+ }165+ return "SELECT * FROM " + quoteIdent(table), total, nil166+}167+168+// readSQLiteRows stores the rows below the header and returns the rows169+// read, counting the header, and how many values were binary. With a170+// known total, rows past the sheet's last are counted, not read.171+func readSQLiteRows(ctx context.Context, rows *sql.Rows, b *builder, ncols, total int, prog *Progress) (row, blobs int, err error) {172+ vals := make([]any, ncols)173+ ptrs := make([]any, ncols)147174 for i := range vals {148175 ptrs[i] = &vals[i]149176 }150- row, blobs := 1, 0151- for ; rows.Next(); row++ {177+ for row = 1; rows.Next(); row++ {152178 if row%256 == 0 {153179 if err := ctx.Err(); err != nil {154- return nil, err180+ return 0, 0, err155181 }156- prog := opt.Progress157182 prog.setRows(row)158183 prog.setFrac(int64(row), int64(total))159184 }160- if err := rows.Scan(ptrs...); err != nil {161- return nil, sqliteErr(err)162- }163185 if row >= sheet.MaxRows {186+ if total >= 0 {187+ row = max(row, total+1)188+ break189+ }164190 b.fits(sheet.Addr{Row: row})165191 continue166192 }193+ if err := rows.Scan(ptrs...); err != nil {194+ return 0, 0, sqliteErr(err)195+ }167196 for c, v := range vals {168197 if putValue(b, sheet.Addr{Col: c, Row: row}, v) {169198 blobs++@@ -171,15 +200,11 @@ func importSQLite(ctx context.Context, name string, opt Options) (*Result, error
171200 }172201 }173202 if err := rows.Err(); err != nil {174- return nil, sqliteErr(err)203+ return 0, 0, sqliteErr(err)175204 }176- opt.Progress.setRows(row)177- var notes []string178- if blobs > 0 {179- notes = append(notes, count(blobs, "binary value", "binary values")+" shown as sizes")180- }181- s, notes := b.finish(notes)182- return &Result{Sheet: s, Rows: row, Notes: notes}, nil205+ b.fits(sheet.Addr{Row: row - 1})206+ prog.setRows(row)207+ return row, blobs, nil183208 }184209 185210 // putValue stores a database value, reporting whether it was binary data@@ -116,6 +116,34 @@ func TestSQLiteImport(t *testing.T) {
116116 }117117 }118118 119+// TestSQLitePastLimit keeps the rows that fit and says how many didn't,120+// for a table (counted, not read) and for a query (read to the end).121+func TestSQLitePastLimit(t *testing.T) {122+ name := filepath.Join(t.TempDir(), "nums.sqlite")123+ db, err := sql.Open("sqlite", name)124+ if err != nil {125+ t.Fatal(err)126+ }127+ _, err = db.Exec(`CREATE TABLE nums AS WITH RECURSIVE c(n) AS128+ (SELECT 1 UNION ALL SELECT n+1 FROM c WHERE n < 8195) SELECT n FROM c`)129+ db.Close()130+ if err != nil {131+ t.Fatal(err)132+ }133+ for _, opt := range []Options{{Table: "nums"}, {Query: "SELECT n FROM nums"}} {134+ res, err := Import(context.Background(), name, opt)135+ if err != nil {136+ t.Fatal(err)137+ }138+ if g := shown(res.Sheet, sheet.Addr{Row: sheet.MaxRows - 1}); g != "8191" {139+ t.Errorf("%+v: last row shows %q, want 8191", opt, g)140+ }141+ if res.Rows != 8196 || len(res.Notes) != 1 || res.Notes[0] != "only the first 8,192 rows fit; 4 rows left out" {142+ t.Errorf("%+v: rows %d, notes %q", opt, res.Rows, res.Notes)143+ }144+ }145+}146+119147 func TestSQLiteExportRoundTrip(t *testing.T) {120148 src := build(t, map[string]string{121149 "A1": "Item", "B1": "Qty", "C1": "Price", "D1": "Due", "E1": "", "F1": "Item",