Skip to content

Commit 1641216

Browse files
aptendaptend
committed
fix: refine load external stats estimates
1 parent 183ca26 commit 1641216

4 files changed

Lines changed: 271 additions & 37 deletions

File tree

‎pkg/sql/plan/bind_load.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -184,7 +184,7 @@ func (builder *QueryBuilder) bindExternalScan(
184184

185185
externalScanNode := &plan.Node{
186186
NodeType: plan.Node_EXTERNAL_SCAN,
187-
Stats: makeLoadExternalStats(stmt.Param, tableDef, offset),
187+
Stats: makeLoadExternalStats(stmt.Param, tableDef, offset, ctx.GetContext()),
188188
ObjRef: objRef,
189189
TableDef: tableDef,
190190
ExternScan: &plan.ExternScan{

‎pkg/sql/plan/build_load.go‎

Lines changed: 98 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ package plan
1616

1717
import (
1818
"bufio"
19+
"context"
1920
"encoding/json"
2021
"io"
2122
"math"
@@ -287,7 +288,7 @@ func hasLoadUserVariable(cols []tree.LoadColumn) bool {
287288
return false
288289
}
289290

290-
func makeLoadExternalStats(param *tree.ExternParam, tableDef *TableDef, offset int64) *plan.Stats {
291+
func makeLoadExternalStats(param *tree.ExternParam, tableDef *TableDef, offset int64, ctx context.Context) *plan.Stats {
291292
// LOAD external scan parallelism is currently sized by
292293
// getParallelSizeForExternalScan as Cost*Rowsize/WriteS3Threshold.
293294
// Keep Cost*Rowsize close to input bytes, but express Cost/Outcnt/BlockNum
@@ -309,7 +310,7 @@ func makeLoadExternalStats(param *tree.ExternParam, tableDef *TableDef, offset i
309310
return stats
310311
}
311312

312-
rowSize := estimateLoadRowsize(param, tableDef, inputSize)
313+
rowSize := estimateLoadRowsize(param, tableDef, inputSize, offset, ctx)
313314
rowCount := math.Ceil(float64(inputSize) / rowSize)
314315
if rowCount < 1 {
315316
rowCount = 1
@@ -319,16 +320,19 @@ func makeLoadExternalStats(param *tree.ExternParam, tableDef *TableDef, offset i
319320
stats.TableCnt = rowCount
320321
stats.Rowsize = rowSize
321322
stats.Selectivity = 1
322-
stats.BlockNum = int32(rowCount/float64(options.DefaultBlockMaxRows)) + 1
323+
stats.BlockNum = int32(math.Ceil(rowCount / float64(options.DefaultBlockMaxRows)))
323324
return stats
324325
}
325326

326-
func estimateLoadRowsize(param *tree.ExternParam, tableDef *TableDef, inputSize int64) float64 {
327+
func estimateLoadRowsize(param *tree.ExternParam, tableDef *TableDef, inputSize int64, offset int64, ctx context.Context) float64 {
327328
if param != nil && param.ScanType == tree.INLINE && param.Format == tree.CSV {
328-
if idx := strings.Index(param.Data, "\n"); idx > 0 {
329-
return clampLoadRowsize(float64(idx), inputSize)
329+
if rowSize := inlineCSVRowsize(param.Data, loadLinesTerminatedBy(param)); rowSize > 0 {
330+
return clampLoadRowsize(rowSize, inputSize)
330331
}
331332
}
333+
if rowSize := estimateLoadRowsizeFromFirstLine(param, inputSize, offset, ctx); rowSize > 0 {
334+
return rowSize
335+
}
332336
if tableDef != nil {
333337
if rowSize := GetRowSizeFromTableDef(tableDef, true) * 0.8; rowSize > 0 {
334338
return clampLoadRowsize(rowSize, inputSize)
@@ -337,6 +341,93 @@ func estimateLoadRowsize(param *tree.ExternParam, tableDef *TableDef, inputSize
337341
return clampLoadRowsize(1, inputSize)
338342
}
339343

344+
func inlineCSVRowsize(data string, terminatedBy string) float64 {
345+
if terminatedBy == "" {
346+
terminatedBy = "\n"
347+
}
348+
if idx := strings.Index(data, terminatedBy); idx >= 0 {
349+
return float64(idx + len(terminatedBy))
350+
}
351+
return float64(len(data))
352+
}
353+
354+
func loadLinesTerminatedBy(param *tree.ExternParam) string {
355+
if param != nil && param.Tail != nil && param.Tail.Lines != nil {
356+
if terminated := param.Tail.Lines.TerminatedBy; terminated != nil && terminated.Value != "" {
357+
return terminated.Value
358+
}
359+
}
360+
return "\n"
361+
}
362+
363+
func estimateLoadRowsizeFromFirstLine(param *tree.ExternParam, inputSize int64, offset int64, ctx context.Context) float64 {
364+
lineTerminator := loadLinesTerminatedBy(param)
365+
if param == nil ||
366+
param.ScanType == tree.INLINE ||
367+
param.Local ||
368+
param.Format == tree.PARQUET ||
369+
getCompressType(param, param.Filepath) != tree.NOCOMPRESS ||
370+
(lineTerminator != "\n" && lineTerminator != "\r\n") ||
371+
strings.HasPrefix(param.Filepath, "SHARED:/query_result/") {
372+
return 0
373+
}
374+
375+
if size := readExternalFirstLineSize(param, offset, ctx); size > 0 {
376+
return clampLoadRowsize(float64(size), inputSize)
377+
}
378+
return 0
379+
}
380+
381+
func readExternalFirstLineSize(param *tree.ExternParam, offset int64, ctx context.Context) int {
382+
if param == nil {
383+
return 0
384+
}
385+
if ctx == nil {
386+
ctx = param.Ctx
387+
}
388+
if ctx == nil {
389+
return 0
390+
}
391+
392+
fs, readPath, err := GetForETLWithType(param, param.Filepath)
393+
if err != nil {
394+
return 0
395+
}
396+
var r io.ReadCloser
397+
vec := fileservice.IOVector{
398+
FilePath: readPath,
399+
Entries: []fileservice.IOEntry{
400+
0: {
401+
Offset: offset,
402+
Size: -1,
403+
ReadCloserForRead: &r,
404+
},
405+
},
406+
}
407+
if err = fs.Read(ctx, &vec); err != nil {
408+
return 0
409+
}
410+
if r == nil {
411+
return 0
412+
}
413+
defer r.Close()
414+
415+
reader := bufio.NewReader(r)
416+
if offset == 0 && param.Tail != nil {
417+
for i := uint64(0); i < param.Tail.IgnoredLines; i++ {
418+
if _, err := reader.ReadString('\n'); err != nil {
419+
return 0
420+
}
421+
}
422+
}
423+
424+
line, err := reader.ReadString('\n')
425+
if len(line) == 0 && err != nil {
426+
return 0
427+
}
428+
return len(line)
429+
}
430+
340431
func clampLoadRowsize(rowSize float64, inputSize int64) float64 {
341432
if rowSize < 1 {
342433
return 1
@@ -453,7 +544,7 @@ func buildLoad(stmt *tree.Load, ctx CompilerContext, isPrepareStmt bool) (*Plan,
453544

454545
externalScanNode := &plan.Node{
455546
NodeType: plan.Node_EXTERNAL_SCAN,
456-
Stats: makeLoadExternalStats(stmt.Param, tableDef, offset),
547+
Stats: makeLoadExternalStats(stmt.Param, tableDef, offset, ctx.GetContext()),
457548
ProjectList: externalProject,
458549
ObjRef: objRef,
459550
TableDef: tableDef,

‎pkg/sql/plan/build_load_parquet_test.go‎

Lines changed: 169 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020

2121
"github.com/matrixorigin/matrixone/pkg/common/moerr"
2222
"github.com/matrixorigin/matrixone/pkg/container/types"
23+
"github.com/matrixorigin/matrixone/pkg/fileservice"
2324
pbplan "github.com/matrixorigin/matrixone/pkg/pb/plan"
2425
"github.com/matrixorigin/matrixone/pkg/sql/parsers/tree"
2526
"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/options"
@@ -292,7 +293,7 @@ func TestMakeLoadExternalStatsUsesInputBytes(t *testing.T) {
292293
}
293294
stats := makeLoadExternalStats(&tree.ExternParam{
294295
ExParamConst: tree.ExParamConst{FileSize: 100},
295-
}, tableDef, 25)
296+
}, tableDef, 25, context.Background())
296297
require.GreaterOrEqual(t, stats.Cost, float64(1))
297298
require.GreaterOrEqual(t, stats.Rowsize, float64(1))
298299
requireLoadByteHint(t, stats, 75)
@@ -301,7 +302,7 @@ func TestMakeLoadExternalStatsUsesInputBytes(t *testing.T) {
301302

302303
stats = makeLoadExternalStats(&tree.ExternParam{
303304
ExParamConst: tree.ExParamConst{FileSize: 10},
304-
}, tableDef, 20)
305+
}, tableDef, 20, context.Background())
305306
require.Equal(t, float64(0), stats.Cost)
306307
require.Equal(t, float64(1), stats.Rowsize)
307308
require.Equal(t, int32(0), stats.BlockNum)
@@ -312,9 +313,43 @@ func TestMakeLoadExternalStatsUsesInputBytes(t *testing.T) {
312313
Format: tree.CSV,
313314
Data: "1,2\n3,4\n",
314315
},
315-
}, tableDef, 0)
316-
require.Equal(t, float64(3), stats.Rowsize)
316+
}, tableDef, 0, context.Background())
317+
require.Equal(t, float64(4), stats.Rowsize)
318+
require.Equal(t, float64(2), stats.Cost)
317319
requireLoadByteHint(t, stats, 8)
320+
321+
stats = makeLoadExternalStats(&tree.ExternParam{
322+
ExParamConst: tree.ExParamConst{
323+
ScanType: tree.INLINE,
324+
Format: tree.CSV,
325+
Data: "a\nb\nc\n",
326+
},
327+
}, tableDef, 0, context.Background())
328+
require.Equal(t, float64(2), stats.Rowsize)
329+
require.Equal(t, float64(3), stats.Cost)
330+
requireLoadByteHint(t, stats, 6)
331+
332+
stats = makeLoadExternalStats(&tree.ExternParam{
333+
ExParamConst: tree.ExParamConst{
334+
ScanType: tree.INLINE,
335+
Format: tree.CSV,
336+
Data: "a|b|c|",
337+
Tail: &tree.TailParameter{
338+
Lines: &tree.Lines{
339+
TerminatedBy: &tree.Terminated{Value: "|"},
340+
},
341+
},
342+
},
343+
}, tableDef, 0, context.Background())
344+
require.Equal(t, float64(2), stats.Rowsize)
345+
require.Equal(t, float64(3), stats.Cost)
346+
requireLoadByteHint(t, stats, 6)
347+
348+
rowSize := GetRowSizeFromTableDef(tableDef, true) * 0.8
349+
stats = makeLoadExternalStats(&tree.ExternParam{
350+
ExParamConst: tree.ExParamConst{FileSize: int64(float64(options.DefaultBlockMaxRows) * rowSize)},
351+
}, tableDef, 0, context.Background())
352+
require.Equal(t, int32(1), stats.BlockNum)
318353
}
319354

320355
func TestMakeLoadExternalStatsKeepsLargeLoadMultiCN(t *testing.T) {
@@ -327,7 +362,7 @@ func TestMakeLoadExternalStatsKeepsLargeLoadMultiCN(t *testing.T) {
327362
inputSize := int64(float64(options.DefaultBlockMaxRows) * GetRowSizeFromTableDef(tableDef, true) * 0.8 * float64(BlockThresholdForOneCN+1))
328363
stats := makeLoadExternalStats(&tree.ExternParam{
329364
ExParamConst: tree.ExParamConst{FileSize: inputSize},
330-
}, tableDef, 0)
365+
}, tableDef, 0, context.Background())
331366
require.Greater(t, stats.BlockNum, int32(BlockThresholdForOneCN))
332367
require.Greater(t, stats.Cost, float64(costThresholdForOneCN))
333368
require.Equal(t, ExecTypeAP_MULTICN, GetExecType(&Query{
@@ -339,6 +374,135 @@ func TestMakeLoadExternalStatsKeepsLargeLoadMultiCN(t *testing.T) {
339374
}, false, false))
340375
}
341376

377+
func TestMakeLoadExternalStatsUsesFirstLineForTextLoad(t *testing.T) {
378+
ctx := context.Background()
379+
fs, err := fileservice.NewMemoryFS("memory", fileservice.DisabledCacheConfig, nil)
380+
require.NoError(t, err)
381+
filePath := fileservice.JoinPath(fs.Name(), "wide.csv")
382+
require.NoError(t, fs.Write(ctx, fileservice.IOVector{
383+
FilePath: filePath,
384+
Entries: []fileservice.IOEntry{{
385+
Offset: 0,
386+
Size: int64(len("1,2\n3,4\n")),
387+
Data: []byte("1,2\n3,4\n"),
388+
}},
389+
}))
390+
391+
tableDef := &TableDef{
392+
Cols: []*ColDef{
393+
{Name: "a", Typ: Type{Id: int32(types.T_varchar), Width: 65535}},
394+
{Name: "b", Typ: Type{Id: int32(types.T_varchar), Width: 65535}},
395+
},
396+
}
397+
inputSize := int64(4 * options.DefaultBlockMaxRows * (BlockThresholdForOneCN + 1))
398+
stats := makeLoadExternalStats(&tree.ExternParam{
399+
ExParamConst: tree.ExParamConst{
400+
Filepath: filePath,
401+
FileSize: inputSize,
402+
Format: tree.CSV,
403+
},
404+
ExParam: tree.ExParam{
405+
FileService: fs,
406+
Ctx: ctx,
407+
},
408+
}, tableDef, 0, context.Background())
409+
410+
require.Equal(t, float64(4), stats.Rowsize)
411+
require.Greater(t, stats.BlockNum, int32(BlockThresholdForOneCN))
412+
require.Equal(t, ExecTypeAP_MULTICN, GetExecType(&Query{
413+
Nodes: []*Node{{
414+
NodeType: pbplan.Node_EXTERNAL_SCAN,
415+
Stats: stats,
416+
}},
417+
Steps: []int32{0},
418+
}, false, false))
419+
420+
crlfPath := fileservice.JoinPath(fs.Name(), "crlf.csv")
421+
require.NoError(t, fs.Write(ctx, fileservice.IOVector{
422+
FilePath: crlfPath,
423+
Entries: []fileservice.IOEntry{{
424+
Offset: 0,
425+
Size: int64(len("1,2\r\n3,4\r\n")),
426+
Data: []byte("1,2\r\n3,4\r\n"),
427+
}},
428+
}))
429+
stats = makeLoadExternalStats(&tree.ExternParam{
430+
ExParamConst: tree.ExParamConst{
431+
Filepath: crlfPath,
432+
FileSize: inputSize,
433+
Format: tree.CSV,
434+
Tail: &tree.TailParameter{
435+
Lines: &tree.Lines{
436+
TerminatedBy: &tree.Terminated{Value: "\r\n"},
437+
},
438+
},
439+
},
440+
ExParam: tree.ExParam{
441+
FileService: fs,
442+
Ctx: ctx,
443+
},
444+
}, tableDef, 0, context.Background())
445+
require.Equal(t, float64(5), stats.Rowsize)
446+
447+
ignoredPath := fileservice.JoinPath(fs.Name(), "ignored.csv")
448+
require.NoError(t, fs.Write(ctx, fileservice.IOVector{
449+
FilePath: ignoredPath,
450+
Entries: []fileservice.IOEntry{{
451+
Offset: 0,
452+
Size: int64(len("long_header_value\n1,2\n3,4\n")),
453+
Data: []byte("long_header_value\n1,2\n3,4\n"),
454+
}},
455+
}))
456+
stats = makeLoadExternalStats(&tree.ExternParam{
457+
ExParamConst: tree.ExParamConst{
458+
Filepath: ignoredPath,
459+
FileSize: inputSize,
460+
Format: tree.CSV,
461+
Tail: &tree.TailParameter{
462+
IgnoredLines: 1,
463+
},
464+
},
465+
ExParam: tree.ExParam{
466+
FileService: fs,
467+
Ctx: ctx,
468+
},
469+
}, tableDef, 0, context.Background())
470+
require.Equal(t, float64(4), stats.Rowsize)
471+
472+
schemaRowSize := GetRowSizeFromTableDef(tableDef, true) * 0.8
473+
stats = makeLoadExternalStats(&tree.ExternParam{
474+
ExParamConst: tree.ExParamConst{
475+
Filepath: filePath,
476+
FileSize: inputSize,
477+
Format: tree.CSV,
478+
CompressType: tree.GZIP,
479+
},
480+
ExParam: tree.ExParam{
481+
FileService: fs,
482+
Ctx: ctx,
483+
},
484+
}, tableDef, 0, context.Background())
485+
require.Equal(t, schemaRowSize, stats.Rowsize)
486+
487+
stats = makeLoadExternalStats(&tree.ExternParam{
488+
ExParamConst: tree.ExParamConst{
489+
Filepath: filePath,
490+
FileSize: inputSize,
491+
Format: tree.CSV,
492+
Tail: &tree.TailParameter{
493+
Lines: &tree.Lines{
494+
TerminatedBy: &tree.Terminated{Value: "|"},
495+
},
496+
},
497+
},
498+
ExParam: tree.ExParam{
499+
FileService: fs,
500+
Ctx: ctx,
501+
},
502+
}, tableDef, 0, context.Background())
503+
require.Equal(t, schemaRowSize, stats.Rowsize)
504+
}
505+
342506
func requireLoadByteHint(t *testing.T, stats *Stats, inputSize int64) {
343507
t.Helper()
344508
estimatedBytes := stats.Cost * stats.Rowsize

0 commit comments

Comments
 (0)