fix LOAD DATA 64K buffer limit (#2709)

Author: jycor

.gitattributes

@@ -7,4 +7,5 @@ enginetest/testdata/test5.txt binary
 enginetest/testdata/test6.csv binary
 enginetest/testdata/test7.txt binary
 enginetest/testdata/test8.txt binary
-enginetest/testdata/test9.txt binary
+enginetest/testdata/test9.txt binary
+enginetest/testdata/test10.txt binary

enginetest/queries/load_queries.go

@@ -539,6 +539,23 @@ var LoadDataScripts = []ScriptTest{
 			},
 		},
 	},
+	{
+		Name: "LOAD DATA with column data larger than 64KB",
+		SetUpScript: []string{
+			"create table t(id int primary key, lt longtext);",
+			"load data infile './testdata/test10.txt' into table t fields terminated by ',' lines terminated by '\n';",
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query: "select id, length(lt) from t order by id",
+				Expected: []sql.Row{
+					{1, 65535},
+					{2, 100000},
+					{3, 1000000},
+				},
+			},
+		},
+	},
 }
 
 var LoadDataErrorScripts = []ScriptTest{

enginetest/testdata/test10.txt

@@ -15,7 +15,7 @@
 package plan
 
 import (
-	"strings"
+	"bytes"
 
 	"github.com/dolthub/vitess/go/vt/sqlparser"
 
@@ -76,7 +76,7 @@ func (l *LoadData) SplitLines(data []byte, atEOF bool) (advance int, token []byt
 	}
 
 	// Find the index of the LINES TERMINATED BY delim.
-	if i := strings.Index(string(data), l.LinesTerminatedBy); i >= 0 {
+	if i := bytes.Index(data, []byte(l.LinesTerminatedBy)); i >= 0 {
 		return i + len(l.LinesTerminatedBy), data[0:i], nil
 	}
 

sql/plan/load_data.go

@@ -94,19 +94,9 @@ func (b *BaseBuilder) buildLoadData(ctx *sql.Context, n *plan.LoadData, row sql.
 	}
 
 	scanner := bufio.NewScanner(reader)
+	scanner.Buffer(nil, int(types.LongTextBlobMax))
 	scanner.Split(n.SplitLines)
 
-	// Skip through the lines that need to be ignored.
-	for n.IgnoreNum > 0 && scanner.Scan() {
-		scanner.Text()
-		n.IgnoreNum--
-	}
-
-	if scanner.Err() != nil {
-		reader.Close()
-		return nil, scanner.Err()
-	}
-
 	sch := n.Schema()
 	source := sch[0].Source // Schema will always have at least one column
 	colNames := n.ColNames
@@ -136,6 +126,8 @@ func (b *BaseBuilder) buildLoadData(ctx *sql.Context, n *plan.LoadData, row sql.
 		setExprs:      n.SetExprs,
 		userVars:      n.UserVars,
 
+		ignoreNum: n.IgnoreNum,
+
 		fieldsTerminatedBy:  n.FieldsTerminatedBy,
 		fieldsEnclosedBy:    n.FieldsEnclosedBy,
 		fieldsEnclosedByOpt: n.FieldsEnclosedByOpt,

sql/rowexec/ddl.go

@@ -47,6 +47,8 @@ type loadDataIter struct {
 	setExprs      []sql.Expression
 	userVars      []sql.Expression
 
+	ignoreNum int64
+
 	fieldsTerminatedBy  string
 	fieldsEnclosedBy    string
 	fieldsEnclosedByOpt bool
@@ -56,21 +58,30 @@ type loadDataIter struct {
 	linesTerminatedBy string
 }
 
-func (l loadDataIter) Next(ctx *sql.Context) (returnRow sql.Row, returnErr error) {
-	var exprs []sql.Expression
+var _ sql.RowIter = (*loadDataIter)(nil)
+var _ sql.Closer = (*loadDataIter)(nil)
+
+func (l *loadDataIter) Next(ctx *sql.Context) (returnRow sql.Row, returnErr error) {
+	// skip first ignoreNum lines
 	var err error
+	for ; l.ignoreNum > 0 && l.scanner.Scan(); l.ignoreNum-- {
+		if err = l.scanner.Err(); err != nil {
+			l.reader.Close()
+			return nil, err
+		}
+	}
+
 	// If exprs is nil then this is a skipped line (see test cases). Keep skipping
 	// until exprs != nil
+	var exprs []sql.Expression
 	for exprs == nil {
 		if keepGoing := l.scanner.Scan(); !keepGoing {
-			if l.scanner.Err() != nil {
-				return nil, l.scanner.Err()
+			if err = l.scanner.Err(); err != nil {
+				return nil, err
 			}
 			return nil, io.EOF
 		}
-
-		line := l.scanner.Text()
-		exprs, err = l.parseFields(ctx, line)
+		exprs, err = l.parseFields(ctx, l.scanner.Text())
 		if err != nil {
 			return nil, err
 		}
@@ -101,12 +112,12 @@ func (l loadDataIter) Next(ctx *sql.Context) (returnRow sql.Row, returnErr error
 	return sql.NewRow(row...), nil
 }
 
-func (l loadDataIter) Close(ctx *sql.Context) error {
+func (l *loadDataIter) Close(ctx *sql.Context) error {
 	return l.reader.Close()
 }
 
 // parseLinePrefix searches for the delim defined by linesStartingByDelim.
-func (l loadDataIter) parseLinePrefix(line string) string {
+func (l *loadDataIter) parseLinePrefix(line string) string {
 	if l.linesStartingBy == "" {
 		return line
 	}
@@ -121,7 +132,7 @@ func (l loadDataIter) parseLinePrefix(line string) string {
 	}
 }
 
-func (l loadDataIter) parseFields(ctx *sql.Context, line string) ([]sql.Expression, error) {
+func (l *loadDataIter) parseFields(ctx *sql.Context, line string) ([]sql.Expression, error) {
 	// Step 1. Start by Searching for prefix if there is one
 	line = l.parseLinePrefix(line)
 	if line == "" {

sql/rowexec/ddl_iters.go

@@ -17,8 +17,9 @@ package types
 import (
 	"testing"
 
-	"github.com/dolthub/go-mysql-server/sql"
 	"github.com/stretchr/testify/assert"
+
+	"github.com/dolthub/go-mysql-server/sql"
 )
 
 func TestIsGeometry(t *testing.T) {