lightning: optimize CSV field unescaping

Project: pingcap/tidb | Language: Go | Concepts: algorithmic-optimization, data-structures, refactoring, testing-strategy | Difficulty: mid
Link: https://github.com/pingcap/tidb/pull/70384

The problem

The CSV parser in TiDB Lightning used a regular expression to unescape every field with escape sequences. This approach was slow, especially for CSV files with many escape sequences, causing significant CPU usage during imports.

Why it's non-trivial

A naive fix might just replace the regular expression with a simple string replace, but this would miss several details:

• The escape character is configurable and could be a regexp metacharacter (like *).
• MySQL-style CSV escaping supports several special mappings (e.g., \n, \t, \0, etc.).
• The parser must handle edge cases like trailing escape characters, custom escape bytes, and correct null value recognition.
• The solution must preserve all existing behaviors and not break compatibility.

The approach

1. Remove regexp from parser state and construction:

-	"regexp"
...
-	unescapeRegexp *regexp.Regexp

The parser no longer stores or compiles a regexp for unescaping.

2. Remove regexp construction in parser initialization:

-	var r *regexp.Regexp
...
-		r, err = regexp.Compile(`(?s)` + regexp.QuoteMeta(cfg.FieldsEscapedBy) + `.`)
-		if err != nil {
-			return nil, errors.Trace(err)
-		}

No regexp is compiled, even if the escape byte is a regexp metacharacter.

3. Use a custom unescape function in the parser:

-		unescaped = unescape(unescaped, "", parser.escFlavor, parser.escapedBy[0], parser.unescapeRegexp)
+		unescaped = unescape(unescaped, "", parser.escFlavor, parser.escapedBy[0])

The parser now calls a new unescape function that does not use regexps.

4. Replace regexp-based unescaping with a single-pass byte scanner:

func unescape(
	input string,
	delim string,
	escFlavor escapeFlavor,
	escChar byte,
) string {
	if len(delim) > 0 {
		delim2 := delim + delim
		if strings.Contains(input, delim2) {
			input = strings.ReplaceAll(input, delim2, delim)
		}
	}
	if escFlavor != escapeFlavorNone {
		input = unescapeByChar(input, escChar)
	}
	return input
}

This function first handles doubled delimiters, then delegates to unescapeByChar if escaping is enabled.

5. Implement unescapeByChar as a manual scanner:

func unescapeByChar(input string, escChar byte) string {
	first := strings.IndexByte(input, escChar)
	if first < 0 || first+1 == len(input) {
		return input
	}

	var result strings.Builder
	result.Grow(len(input) - 1)
	result.WriteString(input[:first])

	remaining := input[first:]
	for len(remaining) > 1 {
		result.WriteByte(escapedCharToByte(remaining[1]))
		remaining = remaining[2:]

		next := strings.IndexByte(remaining, escChar)
		if next < 0 {
			result.WriteString(remaining)
			return result.String()
		}
		result.WriteString(remaining[:next])
		remaining = remaining[next:]
	}

	result.WriteString(remaining)
	return result.String()
}

This function scans the input for the escape character, handles each escape sequence, and writes the result efficiently.

6. Correctly map MySQL escape sequences:

func escapedCharToByte(input byte) byte {
	switch input {
	case '0':
		return 0
	case 'b':
		return '\b'
	case 'n':
		return '\n'
	case 'r':
		return '\r'
	case 't':
		return '\t'
	case 'Z':
		return '\x1a'
	default:
		return input
	}
}

This function ensures that only known escape sequences are replaced; all others are left as-is.

7. Add comprehensive tests and benchmarks:

func TestCustomEscapeChar(t *testing.T) {
	// ...
	cfg.FieldsEscapedBy = `*`
	cfg.FieldNullDefinedBy = []string{`*N`}
	parser, err = mydump.NewCSVParser(
		context.Background(),
		&cfg,
		mydump.NewStringReader(`"*0*b*n*r*t*Z***q",*N,**N`),
		// ...
	)
	// ...
	require.Equal(t, parsedef.Row{
		RowID: 1,
		Row: []types.Datum{
			types.NewStringDatum(string([]byte{0, '\b', '\n', '\r', '\t', 26, '*', 'q'})),
			nullDatum,
			types.NewStringDatum(`*N`),
		},
		Length: 21,
	}, parser.LastRow())
	// ...
}

Test cases cover custom escape bytes (including regexp metacharacters), all supported escape mappings, and edge cases.

Benchmarks are also added to measure performance improvements for both escape-heavy and unescaped data.

From the review thread

No notable review discussion on this PR.

Takeaways

• Avoid regular expressions in hot paths when a single-pass scanner is sufficient; this can yield significant performance improvements.
• Always handle configurable delimiters and escape bytes carefully, especially when they can be special characters in other contexts (like regular expressions).
• Comprehensive tests—including edge cases and performance benchmarks—are critical when refactoring parsing logic to ensure correctness and measure improvements.