add allow_missing_values option to bulk insert (fb-1823) (#2372)

* add allow_missing_values option to bulk insert

* test coverage

* review feedback
This commit is contained in:
pokeeffe-molecula 2022-12-14 18:32:27 -06:00 committed by GitHub
parent 843312dfc9
commit 15d2ee8b07
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 96 additions and 20 deletions

View file

@ -2803,14 +2803,15 @@ type BulkInsertStatement struct {
TransformList []Expr // source to column map
TransformRparen Pos // position of column list right paren
From Pos // position of FROM keyword
DataSource Expr // data source
With Pos // position of WITH keyword
BatchSize Expr
RowsLimit Expr
Format Expr
Input Expr
HeaderRow Expr // has header row (that needs to be skipped)
From Pos // position of FROM keyword
DataSource Expr // data source
With Pos // position of WITH keyword
BatchSize Expr
RowsLimit Expr
Format Expr
Input Expr
HeaderRow Expr // has header row (that needs to be skipped)
AllowMissingValues Expr // allows missing values
}
func (s *BulkInsertStatement) String() string {

View file

@ -1523,7 +1523,7 @@ func (p *Parser) parseBulkInsertStatement() (_ *BulkInsertStatement, err error)
}
stmt.With, _, _ = p.scan()
if !isBulkInsertOptionStartToken(p.peek(), p) {
return nil, p.errorExpected(p.pos, p.tok, "BATCHSIZE, ROWSLIMIT, FORMAT, INPUT or HEADER_ROW")
return nil, p.errorExpected(p.pos, p.tok, "BATCHSIZE, ROWSLIMIT, FORMAT, INPUT, ALLOW_MISSING_VALUES or HEADER_ROW")
}
for {
err := p.parseBulkInsertOption(&stmt)
@ -1577,6 +1577,10 @@ func (p *Parser) parseBulkInsertOption(stmt *BulkInsertStatement) error {
} else {
return p.errorExpected(p.pos, p.tok, "literal")
}
case "ALLOW_MISSING_VALUES":
stmt.AllowMissingValues = ident
return nil
case "HEADER_ROW":
stmt.HeaderRow = ident
return nil
@ -3436,7 +3440,7 @@ func isBulkInsertOptionStartToken(tok Token, p *Parser) bool {
return false
}
switch strings.ToUpper(ident.Name) {
case "BATCHSIZE", "ROWSLIMIT", "FORMAT", "INPUT", "HEADER_ROW":
case "BATCHSIZE", "ROWSLIMIT", "FORMAT", "INPUT", "HEADER_ROW", "ALLOW_MISSING_VALUES":
return true
}
}

View file

@ -80,6 +80,13 @@ func (p *ExecutionPlanner) compileBulkInsertStatement(stmt *parser.BulkInsertSta
}
options.hasHeaderRow = bliteral.Value
// ALLOW_MISSING_VALUES
bliteral, sok = stmt.AllowMissingValues.(*parser.BoolLit)
if !sok {
return nil, sql3.NewErrBoolLiteral(stmt.AllowMissingValues.Pos().Line, stmt.AllowMissingValues.Pos().Column)
}
options.allowMissingValues = bliteral.Value
// batchsize
literal, ok := stmt.BatchSize.(*parser.IntegerLit)
if !ok {
@ -263,14 +270,13 @@ func (p *ExecutionPlanner) analyzeBulkInsertStatement(stmt *parser.BulkInsertSta
}
// header row is true if specified, false if not
if stmt.HeaderRow != nil {
stmt.HeaderRow = &parser.BoolLit{
Value: true,
}
} else {
stmt.HeaderRow = &parser.BoolLit{
Value: false,
}
stmt.HeaderRow = &parser.BoolLit{
Value: stmt.HeaderRow != nil,
}
// allow missing values is true if specified, false if not
stmt.AllowMissingValues = &parser.BoolLit{
Value: stmt.AllowMissingValues != nil,
}
// analyze map expressions

View file

@ -41,6 +41,8 @@ type bulkInsertOptions struct {
format string
// whether the source has a header row
hasHeaderRow bool
// whether we allow missing values for NDJSON jsonpath expressions
allowMissingValues bool
// input specifier (FILE is the only one right now)
input string
@ -88,6 +90,7 @@ func (p *PlanOpBulkInsert) Plan() map[string]interface{} {
options["format"] = p.options.format
options["input"] = p.options.input
options["hasHeaderRow"] = p.options.hasHeaderRow
options["allowMissingValues"] = p.options.allowMissingValues
colMap := make([]interface{}, 0)
for _, m := range p.options.targetColumns {
@ -485,7 +488,11 @@ func (i *bulkInsertSourceNDJsonRowIter) Next(ctx context.Context) (types.Row, er
evalValue, err := expr(ctx, v)
if err != nil {
return nil, sql3.NewErrEvaluatingJSONPathExpr(0, 0, i.mapExpressionResults[idx], jsonValue, err.Error())
if i.options.allowMissingValues && strings.HasPrefix(err.Error(), "unknown key") {
evalValue = nil
} else {
return nil, sql3.NewErrEvaluatingJSONPathExpr(0, 0, i.mapExpressionResults[idx], jsonValue, err.Error())
}
}
// if nil (null) then return nil

View file

@ -1392,7 +1392,7 @@ func TestPlanner_BulkInsert(t *testing.T) {
t.Run("BulkBadWith", func(t *testing.T) {
_, _, err = sql_test.MustQueryRows(t, c.GetNode(0).Server, `bulk insert into j (_id, a, b) map (0 id, 1 int, 2 int) from '/Users/bar/foo.csv' WITH UNICORNS AND RAINBOWS;`)
if err == nil || !strings.Contains(err.Error(), `expected BATCHSIZE, ROWSLIMIT, FORMAT, INPUT or HEADER_ROW, found UNICORNS`) {
if err == nil || !strings.Contains(err.Error(), `expected BATCHSIZE, ROWSLIMIT, FORMAT, INPUT, ALLOW_MISSING_VALUES or HEADER_ROW, found UNICORNS`) {
t.Fatalf("unexpected error: %v", err)
}
})
@ -1897,6 +1897,64 @@ func TestPlanner_BulkInsert(t *testing.T) {
}
})
t.Run("BulkInsertAllowMissingValues", func(t *testing.T) {
_, _, err = sql_test.MustQueryRows(t, c.GetNode(0).Server, `create table greg-test-amv (
_id STRING,
id_col ID,
string_col STRING cachetype ranked size 1000,
int_col int,
decimal_col DECIMAL(2),
bool_col BOOL
time_col TIMESTAMP,
stringset_col STRINGSET,
ideset_col IDSET
);`)
if err != nil {
t.Fatal(err)
}
_, _, err = sql_test.MustQueryRows(t, c.GetNode(0).Server, `BULK INSERT INTO greg-test-amv (
_id,
id_col,
string_col,
int_col,
decimal_col,
bool_col,
time_col,
stringset_col,
ideset_col)
map (
'$.id_col' ID,
'$.string_col' STRING,
'$.int_col' INT,
'$.decimal_col' DECIMAL(2),
'$.bool_col' BOOL,
'$.time_col' TIMESTAMP,
'$.stringset_col' STRINGSET,
'$.ideset_col' IDSET)
transform(
@1,
@0,
@1,
@2,
@3,
@4,
@5,
@6,
@7)
FROM '{"id_col": "3", "string_col": "TEST", "decimal_col": "1.12", "bool_col": false, "time_col": "2013-07-15T01:18:46Z", "stringset_col": "stringset1","ideset_col": 1}
{"id_col": "4", "string_col": "TEST2", "decimal_col": "1.12", "bool_col": false, "time_col": "2013-07-15T01:18:46Z", "stringset_col": ["stringset1","stringset3"],"ideset_col": [1,2]}
{"id_col": "5", "string_col": "TEST", "int_col": "321", "decimal_col": "12.1", "bool_col": 1, "time_col": "2014-07-15T01:18:46Z", "stringset_col": "stringset2","ideset_col": [1,3]}'
with
BATCHSIZE 10000
format 'NDJSON'
input 'STREAM'
allow_missing_values;`)
if err != nil {
t.Fatal(err)
}
})
t.Run("BulkInsertNDJSONStringIDSet", func(t *testing.T) {
_, _, err = sql_test.MustQueryRows(t, c.GetNode(0).Server, `create table greg-test-01 (