Repository: pingcap/tidb-lightning
Pull request #363 (2020-08-03 05:53:35)

restore: support split csv source file with header
<!--
Thank you for contributing to TiDB-Lightning! Please read the [CONTRIBUTING](https://github.com/pingcap/tidb-lightning/blob/master/CONTRIBUTING.md) document **BEFORE** filing this PR.
-->

### What problem does this PR solve? <!--add issue link with summary if exists-->
Support split large csv files with header

### What is changed and how it works?
publish header to each chunk when split csv files, so when restore chunk, we can reader column names from chunk info.

### Benchmark result:
Bench with single source datafile size in 1.1GB:

use master commit, restore cost: 8.6s total cost: 30.5s
```
[2020/08/03 13:34:59.210 +08:00] [INFO] [restore.go:843] ["restore engine completed"] [table=`noheader`.`customer`] [engineNumber=0] [takeTime=8.573293547s] []
[2020/08/03 13:35:10.670 +08:00] [INFO] [restore.go:1109] ["local checksum"] [table=`noheader`.`customer`] [checksum="{cksum=15306146695610273290,size=1479995146,kvs=5760000}"]
[2020/08/03 13:35:21.104 +08:00] [INFO] [restore.go:605] ["restore table completed"] [table=`noheader`.`customer`] [takeTime=30.47492601s] []
```

use this comit: restore cost: 38s total cost: 59s
```
[2020/08/03 13:47:21.646 +08:00] [INFO] [restore.go:843] ["restore engine completed"] [table=`test`.`customer`] [engineNumber=0] [takeTime=38.462242082s] []
[2020/08/03 13:47:35.041 +08:00] [INFO] [restore.go:1446] ["checksum pass"] [table=`test`.`customer`] [local="{cksum=12990959430449404358,size=1479995146,kvs=5760000}"]
[2020/08/03 13:47:42.276 +08:00] [INFO] [restore.go:605] ["restore table completed"] [table=`test`.`customer`] [takeTime=59.096241639s] []
```

NOTE: this pr is base on the commits in #362 .

### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->

 - Unit test
 - Integration test
 - Manual test (add detailed scripts or steps below)
 - No code

Side effects

Related changes

 - Need to update the documentation
 - Need to be included in the release note

kennytm: (WIP depends on #362)
glorv: @kennytm @3pointer @XuHuaiyu  PTAL
kennytm: /run-all-tests

Patch:
diff --git a/lightning/mydump/csv/split_large_file.csv b/lightning/mydump/csv/split_large_file.csv
index 5b7f8ad..7b6512d 100644
--- a/lightning/mydump/csv/split_large_file.csv
+++ b/lightning/mydump/csv/split_large_file.csv
@@ -1,3 +1,4 @@
+a,b,c
 1,1,2
 2,2,1
 3,2,2
diff --git a/lightning/mydump/csv_parser.go b/lightning/mydump/csv_parser.go
index 1c52a9b..e5f34b0 100644
--- a/lightning/mydump/csv_parser.go
+++ b/lightning/mydump/csv_parser.go
@@ -342,15 +342,10 @@ func (parser *CSVParser) ReadRow() error {
 
 	// skip the header first
 	if parser.shouldParseHeader {
-		columns, err := parser.readRecord(nil)
+		err := parser.ReadColumns()
 		if err != nil {
 			return errors.Trace(err)
 		}
-		parser.columns = make([]string, 0, len(columns))
-		for _, colName := range columns {
-			colName, _ = parser.unescapeString(colName)
-			parser.columns = append(parser.columns, strings.ToLower(colName))
-		}
 		parser.shouldParseHeader = false
 	}
 
@@ -385,6 +380,19 @@ func (parser *CSVParser) ReadRow() error {
 	return nil
 }
 
+func (parser *CSVParser) ReadColumns() error {
+	columns, err := parser.readRecord(nil)
+	if err != nil {
+		return errors.Trace(err)
+	}
+	parser.columns = make([]string, 0, len(columns))
+	for _, colName := range columns {
+		colName, _ = parser.unescapeString(colName)
+		parser.columns = append(parser.columns, strings.ToLower(colName))
+	}
+	return nil
+}
+
 var newLineAsciiSet, _ = makeASCIISet("\r\n")
 
 func indexOfNewLine(b []byte) int {
diff --git a/lightning/mydump/parser.go b/lightning/mydump/parser.go
index ba4aeb3..bcea22f 100644
--- a/lightning/mydump/parser.go
+++ b/lightning/mydump/parser.go
@@ -85,6 +85,7 @@ type Chunk struct {
 	EndOffset    int64
 	PrevRowIDMax int64
 	RowIDMax     int64
+	Columns      []string
 }
 
 // Row is the content of a row.
diff --git a/lightning/mydump/region.go b/lightning/mydump/region.go
index cdbd4b4..9e16b0a 100644
--- a/lightning/mydump/region.go
+++ b/lightning/mydump/region.go
@@ -163,7 +163,7 @@ func MakeTableRegions(
 		}
 		// If a csv file is overlarge, we need to split it into mutiple regions.
 		// Note: We can only split a csv file whose format is strict and header is empty.
-		if isCsvFile && dataFileSize > cfg.Mydumper.MaxRegionSize && cfg.Mydumper.StrictFormat && !cfg.Mydumper.CSV.Header {
+		if isCsvFile && dataFileSize > cfg.Mydumper.MaxRegionSize && cfg.Mydumper.StrictFormat {
 			var (
 				regions      []*TableRegion
 				subFileSizes []float64
@@ -226,6 +226,20 @@ func SplitLargeFile(
 	maxRegionSize := cfg.Mydumper.MaxRegionSize
 	dataFileSizes = make([]float64, 0, dataFileSize/maxRegionSize+1)
 	startOffset, endOffset := int64(0), maxRegionSize
+	var columns []string
+	if cfg.Mydumper.CSV.Header {
+		reader, err := os.Open(dataFilePath)
+		if err != nil {
+			return 0, nil, nil, err
+		}
+		parser := NewCSVParser(&cfg.Mydumper.CSV, reader, cfg.Mydumper.ReadBlockSize, ioWorker, true)
+		if err = parser.ReadColumns(); err != nil {
+			return 0, nil, nil, err
+		}
+		columns = parser.Columns()
+		startOffset, _ = parser.Pos()
+		endOffset = startOffset + maxRegionSize
+	}
 	for {
 		curRowsCnt := (endOffset - startOffset) / divisor
 		rowIDMax := prevRowIdxMax + curRowsCnt
@@ -253,6 +267,7 @@ func SplitLargeFile(
 					EndOffset:    endOffset,
 					PrevRowIDMax: prevRowIdxMax,
 					RowIDMax:     rowIDMax,
+					Columns:      columns,
 				},
 			})
 		dataFileSizes = append(dataFileSizes, float64(endOffset-startOffset))
diff --git a/lightning/restore/restore.go b/lightning/restore/restore.go
index 2a8fd4a..9507f32 100644
--- a/lightning/restore/restore.go
+++ b/lightning/restore/restore.go
@@ -1348,7 +1348,7 @@ func (t *TableRestore) populateChunks(rc *RestoreController, cp *TableCheckpoint
 				}
 				cp.Engines[chunk.EngineID] = engine
 			}
-			engine.Chunks = append(engine.Chunks, &ChunkCheckpoint{
+			ccp := &ChunkCheckpoint{
 				Key: ChunkCheckpointKey{
 					Path:   chunk.File,
 					Offset: chunk.Chunk.Offset,
@@ -1356,7 +1356,11 @@ func (t *TableRestore) populateChunks(rc *RestoreController, cp *TableCheckpoint
 				ColumnPermutation: nil,
 				Chunk:             chunk.Chunk,
 				Timestamp:         timestamp,
-			})
+			}
+			if len(chunk.Chunk.Columns) > 0 {
+				ccp.ColumnPermutation = t.parseColumnPermutations(chunk.Chunk.Columns)
+			}
+			engine.Chunks = append(engine.Chunks, ccp)
 		}
 
 		// Add index engine checkpoint
@@ -1383,10 +1387,11 @@ func (t *TableRestore) populateChunks(rc *RestoreController, cp *TableCheckpoint
 //
 // The argument `columns` _must_ be in lower case.
 func (t *TableRestore) initializeColumns(columns []string, ccp *ChunkCheckpoint) {
-	colPerm := make([]int, 0, len(t.tableInfo.Core.Columns)+1)
-	shouldIncludeRowID := common.TableHasAutoRowID(t.tableInfo.Core)
-
+	var colPerm []int
 	if len(columns) == 0 {
+		colPerm = make([]int, 0, len(t.tableInfo.Core.Columns)+1)
+		shouldIncludeRowID := common.TableHasAutoRowID(t.tableInfo.Core)
+
 		// no provided columns, so use identity permutation.
 		for i := range t.tableInfo.Core.Columns {
 			colPerm = append(colPerm, i)
@@ -1395,30 +1400,37 @@ func (t *TableRestore) initializeColumns(columns []string, ccp *ChunkCheckpoint)
 			colPerm = append(colPerm, -1)
 		}
 	} else {
-		columnMap := make(map[string]int)
-		for i, column := range columns {
-			columnMap[column] = i
-		}
-		for _, colInfo := range t.tableInfo.Core.Columns {
-			if i, ok := columnMap[colInfo.Name.L]; ok {
-				colPerm = append(colPerm, i)
-			} else {
-				t.logger.Warn("column missing from data file, going to fill with default value",
-					zap.Stringer("path", &ccp.Key),
-					zap.String("colName", colInfo.Name.O),
-					zap.Stringer("colType", &colInfo.FieldType),
-				)
-				colPerm = append(colPerm, -1)
-			}
-		}
-		if i, ok := columnMap[model.ExtraHandleName.L]; ok {
+		colPerm = t.parseColumnPermutations(columns)
+	}
+
+	ccp.ColumnPermutation = colPerm
+}
+
+func (t *TableRestore) parseColumnPermutations(columns []string) []int {
+	colPerm := make([]int, 0, len(t.tableInfo.Core.Columns)+1)
+
+	columnMap := make(map[string]int)
+	for i, column := range columns {
+		columnMap[column] = i
+	}
+	for _, colInfo := range t.tableInfo.Core.Columns {
+		if i, ok := columnMap[colInfo.Name.L]; ok {
 			colPerm = append(colPerm, i)
-		} else if shouldIncludeRowID {
+		} else {
+			t.logger.Warn("column missing from data file, going to fill with default value",
+				zap.String("colName", colInfo.Name.O),
+				zap.Stringer("colType", &colInfo.FieldType),
+			)
 			colPerm = append(colPerm, -1)
 		}
 	}
+	if i, ok := columnMap[model.ExtraHandleName.L]; ok {
+		colPerm = append(colPerm, i)
+	} else if common.TableHasAutoRowID(t.tableInfo.Core) {
+		colPerm = append(colPerm, -1)
+	}
 
-	ccp.ColumnPermutation = colPerm
+	return colPerm
 }
 
 func getColumnNames(tableInfo *model.TableInfo, permutation []int) []string {


Test patch:
diff --git a/lightning/mydump/region_test.go b/lightning/mydump/region_test.go
index 20aba61..4b54108 100644
--- a/lightning/mydump/region_test.go
+++ b/lightning/mydump/region_test.go
@@ -209,7 +209,7 @@ func (s *testMydumpRegionSuite) TestSplitLargeFile(c *C) {
 			CSV: config.CSVConfig{
 				Separator:       ",",
 				Delimiter:       "",
-				Header:          false,
+				Header:          true,
 				TrimLastSep:     false,
 				NotNull:         false,
 				Null:            "NULL",
@@ -226,18 +226,19 @@ func (s *testMydumpRegionSuite) TestSplitLargeFile(c *C) {
 	}
 	fileSize := dataFileInfo.Size()
 	colCnt := int64(3)
+	columns := []string{"a", "b", "c"}
 	for _, tc := range []struct {
 		maxRegionSize int64
 		chkCnt        int
 		offsets       [][]int64
 	}{
-		{1, 4, [][]int64{{0, 6}, {6, 12}, {12, 18}, {18, 24}}},
-		{6, 2, [][]int64{{0, 12}, {12, 24}}},
-		{8, 2, [][]int64{{0, 12}, {12, 24}}},
-		{12, 2, [][]int64{{0, 18}, {18, 24}}},
-		{13, 2, [][]int64{{0, 18}, {18, 24}}},
-		{18, 1, [][]int64{{0, 24}}},
-		{19, 1, [][]int64{{0, 24}}},
+		{1, 4, [][]int64{{6, 12}, {12, 18}, {18, 24}, {24, 30}}},
+		{6, 2, [][]int64{{6, 18}, {18, 30}}},
+		{8, 2, [][]int64{{6, 18}, {18, 30}}},
+		{12, 2, [][]int64{{6, 24}, {24, 30}}},
+		{13, 2, [][]int64{{6, 24}, {24, 30}}},
+		{18, 1, [][]int64{{6, 30}}},
+		{19, 1, [][]int64{{6, 30}}},
 	} {
 		cfg.Mydumper.MaxRegionSize = tc.maxRegionSize
 		prevRowIdxMax := int64(0)
@@ -248,6 +249,8 @@ func (s *testMydumpRegionSuite) TestSplitLargeFile(c *C) {
 		for i := range tc.offsets {
 			c.Assert(regions[i].Chunk.Offset, Equals, tc.offsets[i][0])
 			c.Assert(regions[i].Chunk.EndOffset, Equals, tc.offsets[i][1])
+			c.Assert(len(regions[i].Chunk.Columns), Equals, len(columns))
+			c.Assert(regions[i].Chunk.Columns, DeepEquals, columns)
 		}
 	}
 }