// Copyright 2021 Molecula Corp. All rights reserved. package planner import ( "bufio" "bytes" "context" "encoding/csv" "encoding/json" "fmt" "io" "net/http" "os" "strconv" "strings" "time" "github.com/PaesslerAG/gval" "github.com/PaesslerAG/jsonpath" "github.com/apache/arrow/go/v10/arrow" "github.com/apache/arrow/go/v10/arrow/memory" "github.com/apache/arrow/go/v10/parquet/file" "github.com/apache/arrow/go/v10/parquet/pqarrow" pilosa "github.com/featurebasedb/featurebase/v3" "github.com/featurebasedb/featurebase/v3/pql" "github.com/featurebasedb/featurebase/v3/sql3" "github.com/featurebasedb/featurebase/v3/sql3/parser" "github.com/featurebasedb/featurebase/v3/sql3/planner/types" ) type bulkInsertMapColumn struct { name string expr types.PlanExpression colType parser.ExprDataType } // bulkInsertOptions contains options for bulk insert type bulkInsertOptions struct { // name of the file we're going to read sourceData string // number of rows in a batch batchSize int // stop after this many rows rowsLimit int // format specifier (CSV is the only one right now) 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 // target columns targetColumns []*qualifiedRefPlanExpression // transformations transformExpressions []types.PlanExpression // map expressions mapExpressions []*bulkInsertMapColumn } // PlanOpBulkInsert plan operator to handle INSERT. type PlanOpBulkInsert struct { planner *ExecutionPlanner tableName string options *bulkInsertOptions warnings []string } func NewPlanOpBulkInsert(p *ExecutionPlanner, tableName string, options *bulkInsertOptions) *PlanOpBulkInsert { return &PlanOpBulkInsert{ planner: p, tableName: tableName, options: options, warnings: make([]string, 0), } } func (p *PlanOpBulkInsert) Plan() map[string]interface{} { result := make(map[string]interface{}) result["_op"] = fmt.Sprintf("%T", p) result["_schema"] = p.Schema().Plan() result["tableName"] = p.tableName options := make(map[string]interface{}) options["sourceData"] = p.options.sourceData options["batchSize"] = p.options.batchSize options["rowsLimit"] = p.options.rowsLimit 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 { colMap = append(colMap, m.Plan()) } options["targetColumns"] = colMap mapList := make([]interface{}, 0) for _, m := range p.options.mapExpressions { mapItem := make(map[string]interface{}) options["name"] = m.name options["type"] = m.colType.TypeDescription() options["expr"] = m.expr.Plan() mapList = append(mapList, mapItem) } options["mapExpressions"] = mapList if p.options.transformExpressions != nil && len(p.options.transformExpressions) > 0 { transformList := make([]interface{}, 0) for _, m := range p.options.transformExpressions { transformList = append(transformList, m.Plan()) } options["transformExpressions"] = transformList } result["options"] = options return result } func (p *PlanOpBulkInsert) String() string { return "" } func (p *PlanOpBulkInsert) AddWarning(warning string) { p.warnings = append(p.warnings, warning) } func (p *PlanOpBulkInsert) Warnings() []string { return p.warnings } func (p *PlanOpBulkInsert) Schema() types.Schema { return types.Schema{} } func (p *PlanOpBulkInsert) Children() []types.PlanOperator { return []types.PlanOperator{} } func (p *PlanOpBulkInsert) Iterator(ctx context.Context, row types.Row) (types.RowIterator, error) { switch strings.ToUpper(p.options.format) { case "CSV": return &bulkInsertLineRowIter{ planner: p.planner, tableName: p.tableName, options: p.options, sourceIter: &bulkInsertSourceCSVRowIter{ planner: p.planner, options: p.options, }, }, nil case "NDJSON": return &bulkInsertNDJsonRowIter{ planner: p.planner, tableName: p.tableName, options: p.options, sourceIter: &bulkInsertSourceNDJsonRowIter{ planner: p.planner, options: p.options, }, }, nil case "PARQUET": return &bulkInsertLineRowIter{ planner: p.planner, tableName: p.tableName, options: p.options, sourceIter: &bulkInsertSourceParquetRowIter{ planner: p.planner, options: p.options, }, }, nil default: return nil, sql3.NewErrInternalf("unexpected format '%s'", p.options.format) } } func (p *PlanOpBulkInsert) WithChildren(children ...types.PlanOperator) (types.PlanOperator, error) { return NewPlanOpBulkInsert(p.planner, p.tableName, p.options), nil } type bulkInsertSourceCSVRowIter struct { planner *ExecutionPlanner options *bulkInsertOptions csvReader *csv.Reader closeFunc func() mapValues []int64 hasStarted *struct{} } var _ types.RowIterator = (*bulkInsertSourceCSVRowIter)(nil) func (i *bulkInsertSourceCSVRowIter) Next(ctx context.Context) (types.Row, error) { if i.hasStarted == nil { i.hasStarted = &struct{}{} // pre-calculate map values since these represent column offsets and will be constant for csv i.mapValues = []int64{} for _, mc := range i.options.mapExpressions { // this is csv so map value will be an int rawMapValue, err := mc.expr.Evaluate(nil) if err != nil { return nil, err } mapValue, ok := rawMapValue.(int64) if !ok { return nil, sql3.NewErrInternalf("unexpected type for mapValue '%T'", rawMapValue) } i.mapValues = append(i.mapValues, mapValue) } switch strings.ToUpper(i.options.input) { case "FILE": f, err := os.Open(i.options.sourceData) if err != nil { return nil, err } i.closeFunc = func() { f.Close() } i.csvReader = csv.NewReader(f) case "URL": response, err := http.Get(i.options.sourceData) if err != nil { return nil, err } i.closeFunc = func() { response.Body.Close() } if response.StatusCode != 200 { return nil, sql3.NewErrReadingDatasource(0, 0, i.options.sourceData, fmt.Sprintf("unexpected response %d", response.StatusCode)) } i.csvReader = csv.NewReader(response.Body) case "STREAM": i.csvReader = csv.NewReader(strings.NewReader(i.options.sourceData)) default: return nil, sql3.NewErrInternalf("unexpected input specification type '%s'", i.options.input) } i.csvReader.LazyQuotes = true i.csvReader.TrimLeadingSpace = true // skip header row if necessary if i.options.hasHeaderRow { _, err := i.csvReader.Read() if err == io.EOF { return nil, types.ErrNoMoreRows } else if err != nil { return nil, err } } } rec, err := i.csvReader.Read() if err == io.EOF { return nil, types.ErrNoMoreRows } // err == csv.ParseError is impossible if LazyQuotes is true // we should uncomment the code below if we ever dispable LazyQuotes // so there is no need check for any other error /* else if err != nil { pe, ok := err.(*csv.ParseError) if ok { return nil, sql3.NewErrReadingDatasource(0, 0, i.options.sourceData, fmt.Sprintf("csv parse error on line %d: %s", pe.Line, pe.Error())) } return nil, err } */ // now we do the mapping to the output row result := make([]interface{}, len(i.options.mapExpressions)) for idx := range i.options.mapExpressions { mapExpressionResult := i.mapValues[idx] if !(mapExpressionResult >= 0 && int(mapExpressionResult) < len(rec)) { return nil, sql3.NewErrMappingFromDatasource(0, 0, i.options.sourceData, fmt.Sprintf("map index %d out of range", mapExpressionResult)) } evalValue := rec[mapExpressionResult] mapColumn := i.options.mapExpressions[idx] // if an empty string and the map type is not a string, treat it as a null if len(evalValue) == 0 && !typeIsString(mapColumn.colType) { result[idx] = nil continue } switch mapColumn.colType.(type) { case *parser.DataTypeID, *parser.DataTypeInt: intVal, err := strconv.ParseInt(evalValue, 10, 64) if err != nil { return nil, sql3.NewErrTypeConversionOnMap(0, 0, evalValue, mapColumn.colType.TypeDescription()) } result[idx] = intVal case *parser.DataTypeIDSet: intVal, err := strconv.ParseInt(evalValue, 10, 64) if err != nil { return nil, sql3.NewErrTypeConversionOnMap(0, 0, evalValue, mapColumn.colType.TypeDescription()) } result[idx] = []int64{intVal} case *parser.DataTypeStringSet: result[idx] = []string{evalValue} case *parser.DataTypeTimestamp: intVal, err := strconv.ParseInt(evalValue, 10, 64) if err != nil { if tm, err := time.ParseInLocation(time.RFC3339Nano, evalValue, time.UTC); err == nil { result[idx] = tm } else if tm, err := time.ParseInLocation("2006-01-02", evalValue, time.UTC); err == nil { result[idx] = tm } else { return nil, sql3.NewErrTypeConversionOnMap(0, 0, evalValue, mapColumn.colType.TypeDescription()) } } else { // implicit conversion of int to timestamp will treat int as seconds since unix epoch result[idx] = time.Unix(intVal, 0).UTC() } case *parser.DataTypeString: result[idx] = evalValue case *parser.DataTypeBool: bval, err := strconv.ParseBool(evalValue) if err != nil { return nil, sql3.NewErrTypeConversionOnMap(0, 0, evalValue, mapColumn.colType.TypeDescription()) } result[idx] = bval case *parser.DataTypeDecimal: dval, err := pql.ParseDecimal(evalValue) if err != nil { return nil, sql3.NewErrTypeConversionOnMap(0, 0, evalValue, mapColumn.colType.TypeDescription()) } result[idx] = dval default: return nil, sql3.NewErrInternalf("unhandled type '%T'", mapColumn.colType) } } return result, nil } func (i *bulkInsertSourceCSVRowIter) Close(ctx context.Context) { if i.closeFunc != nil { i.closeFunc() } } type bulkInsertBasicRowIter interface { Next(context.Context) (types.Row, error) Close(context.Context) } type bulkInsertLineRowIter struct { planner *ExecutionPlanner tableName string options *bulkInsertOptions linesRead int currentBatch [][]interface{} sourceIter bulkInsertBasicRowIter } var _ types.RowIterator = (*bulkInsertLineRowIter)(nil) func (i *bulkInsertLineRowIter) Next(ctx context.Context) (types.Row, error) { defer i.sourceIter.Close(ctx) for { row, err := i.sourceIter.Next(ctx) if err != nil && err != types.ErrNoMoreRows { return nil, err } if err == types.ErrNoMoreRows { break } i.linesRead++ if i.currentBatch == nil { i.currentBatch = make([][]interface{}, 0) } i.currentBatch = append(i.currentBatch, row) if len(i.currentBatch) >= i.options.batchSize { err := processBatch(ctx, i.planner, i.tableName, i.currentBatch, i.options) if err != nil { return nil, err } i.currentBatch = nil // update the counter for bulk insert batches pilosa.PerfCounterSQLBulkInsertBatchesSec.Add(1) } if i.options.rowsLimit > 0 && i.linesRead >= i.options.rowsLimit { break } } if len(i.currentBatch) > 0 { err := processBatch(ctx, i.planner, i.tableName, i.currentBatch, i.options) if err != nil { return nil, err } i.currentBatch = nil // update the counter for bulk insert batches pilosa.PerfCounterSQLBulkInsertBatchesSec.Add(1) } return nil, types.ErrNoMoreRows } type bulkInsertSourceNDJsonRowIter struct { planner *ExecutionPlanner options *bulkInsertOptions reader *bufio.Scanner closeFunc func() mapExpressionResults []string pathExpressions []gval.Evaluable hasStarted *struct{} } var _ types.RowIterator = (*bulkInsertSourceNDJsonRowIter)(nil) func (i *bulkInsertSourceNDJsonRowIter) Next(ctx context.Context) (types.Row, error) { if i.hasStarted == nil { i.hasStarted = &struct{}{} builder := gval.Full(jsonpath.PlaceholderExtension()) // pre-calculate map values since these represent ndjson expressions and will be constant i.mapExpressionResults = []string{} i.pathExpressions = []gval.Evaluable{} for _, mc := range i.options.mapExpressions { rawMapValue, err := mc.expr.Evaluate(nil) if err != nil { return nil, err } mapValue, ok := rawMapValue.(string) if !ok { return nil, sql3.NewErrInternalf("unexpected type for mapValue '%T'", rawMapValue) } i.mapExpressionResults = append(i.mapExpressionResults, mapValue) path, err := builder.NewEvaluable(mapValue) if err != nil { return nil, err } i.pathExpressions = append(i.pathExpressions, path) } switch strings.ToUpper(i.options.input) { case "FILE": f, err := os.Open(i.options.sourceData) if err != nil { return nil, err } i.closeFunc = func() { f.Close() } i.reader = bufio.NewScanner(f) case "URL": response, err := http.Get(i.options.sourceData) if err != nil { return nil, err } i.closeFunc = func() { response.Body.Close() } if response.StatusCode != 200 { return nil, sql3.NewErrReadingDatasource(0, 0, i.options.sourceData, fmt.Sprintf("unexpected response %d", response.StatusCode)) } i.reader = bufio.NewScanner(response.Body) case "STREAM": i.reader = bufio.NewScanner(strings.NewReader(i.options.sourceData)) default: return nil, sql3.NewErrInternalf("unexpected input specification type '%s'", i.options.input) } } for { if i.reader.Scan() { if err := i.reader.Err(); err != nil { return nil, err } jsonValue := i.reader.Text() jsonValue = strings.TrimSpace(jsonValue) if len(jsonValue) == 0 { continue } // now we do the mapping to the output row result := make([]interface{}, len(i.options.mapExpressions)) // parse the json v := interface{}(nil) dec := json.NewDecoder(bytes.NewReader([]byte(jsonValue))) dec.UseNumber() err := dec.Decode(&v) if err != nil { return nil, sql3.NewErrParsingJSON(0, 0, jsonValue, err.Error()) } // type check against the output type of the map operation for idx, expr := range i.pathExpressions { evalValue, err := expr(ctx, v) if err != nil { if i.options.allowMissingValues && (strings.HasPrefix(err.Error(), "unknown key") || strings.HasPrefix(err.Error(), "unknown parameter")) { evalValue = nil } else { return nil, sql3.NewErrEvaluatingJSONPathExpr(0, 0, i.mapExpressionResults[idx], jsonValue, err.Error()) } } // if nil (null) then return nil if evalValue == nil { result[idx] = nil continue } mapColumn := i.options.mapExpressions[idx] switch mapColumn.colType.(type) { case *parser.DataTypeID, *parser.DataTypeInt: switch v := evalValue.(type) { case json.Number: n, err := v.Int64() if err == nil { result[idx] = n } else { return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) } case []interface{}: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) case string: intVal, err := strconv.ParseInt(v, 10, 64) if err != nil { return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) } result[idx] = intVal case bool: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) case interface{}: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) default: return nil, sql3.NewErrInternalf("unhandled type '%T'", evalValue) } case *parser.DataTypeIDSet: switch v := evalValue.(type) { case json.Number: n, err := v.Int64() if err == nil { result[idx] = []int64{n} } else { return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) } case []interface{}: setValue := make([]int64, 0) for _, i := range v { switch v := i.(type) { case json.Number: i, e := v.Int64() if e == nil { setValue = append(setValue, i) } else { return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) } case string: intVal, err := strconv.ParseInt(v, 10, 64) if err != nil { return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) } setValue = append(setValue, int64(intVal)) default: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) } } result[idx] = setValue case string: intVal, err := strconv.ParseInt(v, 10, 64) if err != nil { return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) } result[idx] = []int64{int64(intVal)} case bool: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) case interface{}: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) default: return nil, sql3.NewErrInternalf("unhandled type '%T'", evalValue) } case *parser.DataTypeStringSet: switch v := evalValue.(type) { case json.Number: result[idx] = []string{v.String()} case []interface{}: setValue := make([]string, 0) for _, i := range v { f, ok := i.(string) if !ok { return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) } setValue = append(setValue, f) } result[idx] = setValue case string: result[idx] = []string{v} case bool: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) case interface{}: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) default: return nil, sql3.NewErrInternalf("unhandled type '%T'", evalValue) } case *parser.DataTypeTimestamp: switch v := evalValue.(type) { case json.Number: n, err := v.Int64() // if v is a whole number then make it an int if err == nil { // implicit conversion of int to timestamp will treat int as seconds since unix epoch result[idx] = time.Unix(n, 0).UTC() } else { return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) } case []interface{}: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) case string: if tm, err := time.ParseInLocation(time.RFC3339Nano, v, time.UTC); err == nil { result[idx] = tm } else if tm, err := time.ParseInLocation("2006-01-02", v, time.UTC); err == nil { result[idx] = tm } else { return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) } case bool: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) case interface{}: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) default: return nil, sql3.NewErrInternalf("unhandled type '%T'", evalValue) } case *parser.DataTypeString: switch v := evalValue.(type) { case json.Number: result[idx] = v.String() case []interface{}: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) case string: result[idx] = v case bool: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) case interface{}: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) default: return nil, sql3.NewErrInternalf("unhandled type '%T'", evalValue) } case *parser.DataTypeBool: switch v := evalValue.(type) { case json.Number: // if a whole number make it an int, and convert to a bool n, err := v.Int64() if err == nil { result[idx] = n > 0 } else { return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) } case []interface{}: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) case string: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) case bool: result[idx] = v case interface{}: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) default: return nil, sql3.NewErrInternalf("unhandled type '%T'", evalValue) } case *parser.DataTypeDecimal: switch v := evalValue.(type) { case json.Number: f, err := v.Float64() if err != nil { return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) } result[idx] = pql.FromFloat64(f) case []interface{}: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) case string: // try to parse from a string dv, err := pql.ParseDecimal(v) if err != nil { return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) } result[idx] = dv case bool: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) case interface{}: return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription()) default: return nil, sql3.NewErrInternalf("unhandled type '%T'", evalValue) } default: return nil, sql3.NewErrInternalf("unhandled type '%T'", mapColumn.colType) } } return result, nil } return nil, types.ErrNoMoreRows } } func (i *bulkInsertSourceNDJsonRowIter) Close(ctx context.Context) { if i.closeFunc != nil { i.closeFunc() } } type bulkInsertNDJsonRowIter struct { planner *ExecutionPlanner tableName string options *bulkInsertOptions linesRead int currentBatch [][]interface{} sourceIter *bulkInsertSourceNDJsonRowIter } var _ types.RowIterator = (*bulkInsertNDJsonRowIter)(nil) func (i *bulkInsertNDJsonRowIter) Next(ctx context.Context) (types.Row, error) { defer i.sourceIter.Close(ctx) for { row, err := i.sourceIter.Next(ctx) if err != nil && err != types.ErrNoMoreRows { return nil, err } if err == types.ErrNoMoreRows { break } i.linesRead++ if i.currentBatch == nil { i.currentBatch = make([][]interface{}, 0) } i.currentBatch = append(i.currentBatch, row) if len(i.currentBatch) >= i.options.batchSize { err := processBatch(ctx, i.planner, i.tableName, i.currentBatch, i.options) if err != nil { return nil, err } i.currentBatch = nil // update the counter for bulk insert batches pilosa.PerfCounterSQLBulkInsertBatchesSec.Add(1) } if i.options.rowsLimit > 0 && i.linesRead >= i.options.rowsLimit { break } } if len(i.currentBatch) > 0 { err := processBatch(ctx, i.planner, i.tableName, i.currentBatch, i.options) if err != nil { return nil, err } i.currentBatch = nil // update the counter for bulk insert batches pilosa.PerfCounterSQLBulkInsertBatchesSec.Add(1) } return nil, types.ErrNoMoreRows } func processColumnValue(rawValue interface{}, targetType parser.ExprDataType) (types.PlanExpression, error) { if rawValue == nil { return newNullLiteralPlanExpression(), nil } switch targetType.(type) { case *parser.DataTypeID, *parser.DataTypeInt: ival, ok := rawValue.(int64) if !ok { return nil, sql3.NewErrInternalf("unexpected value type '%T'", rawValue) } return newIntLiteralPlanExpression(ival), nil case *parser.DataTypeIDSet: val, ok := rawValue.([]int64) if !ok { return nil, sql3.NewErrInternalf("unable to convert '%s", rawValue) } members := make([]types.PlanExpression, 0) for _, m := range val { members = append(members, newIntLiteralPlanExpression(m)) } return newExprSetLiteralPlanExpression(members, parser.NewDataTypeIDSet()), nil case *parser.DataTypeStringSet: val, ok := rawValue.([]string) if !ok { return nil, sql3.NewErrInternalf("unable to convert '%s", rawValue) } members := make([]types.PlanExpression, 0) for _, m := range val { members = append(members, newStringLiteralPlanExpression(m)) } return newExprSetLiteralPlanExpression(members, parser.NewDataTypeStringSet()), nil case *parser.DataTypeTimestamp: tval, ok := rawValue.(time.Time) if !ok { return nil, sql3.NewErrInternalf("unable to convert '%s", rawValue) } return newTimestampLiteralPlanExpression(tval), nil case *parser.DataTypeString: sval, ok := rawValue.(string) if !ok { return nil, sql3.NewErrInternalf("unable to convert '%s", rawValue) } return newStringLiteralPlanExpression(sval), nil case *parser.DataTypeBool: bval, ok := rawValue.(bool) if !ok { return nil, sql3.NewErrInternalf("unable to convert '%s", rawValue) } return newBoolLiteralPlanExpression(bval), nil case *parser.DataTypeDecimal: dval, ok := rawValue.(pql.Decimal) if !ok { return nil, sql3.NewErrInternalf("unable to convert '%s", rawValue) } return newFloatLiteralPlanExpression(fmt.Sprintf("%f", dval.Float64())), nil default: return nil, sql3.NewErrInternalf("unhandled type '%T'", targetType) } } func processBatch(ctx context.Context, planner *ExecutionPlanner, tableName string, currentBatch [][]interface{}, options *bulkInsertOptions) error { insertValues := [][]types.PlanExpression{} // we're going to take a different path if transforms are specified // mostly for performmance reasons if len(options.transformExpressions) > 0 { // we have transformations so we are going to evaluate them and then build the insert tuple for _, row := range currentBatch { tupleValues := []types.PlanExpression{} // handle each transform for idx, mc := range options.transformExpressions { rawValue, err := mc.Evaluate(row) if err != nil { return err } // handle nulls if rawValue == nil { tupleValues = append(tupleValues, newNullLiteralPlanExpression()) continue } tupleExpr, err := processColumnValue(rawValue, options.targetColumns[idx].dataType) if err != nil { return err } tupleValues = append(tupleValues, tupleExpr) } insertValues = append(insertValues, tupleValues) } } else { // we are just going to take the values from the source row and copy pasta them across // for each row in the batch add value to each mapped column for _, row := range currentBatch { tupleValues := []types.PlanExpression{} // handle each column for idx, rawValue := range row { tupleExpr, err := processColumnValue(rawValue, options.targetColumns[idx].dataType) if err != nil { return err } tupleValues = append(tupleValues, tupleExpr) } insertValues = append(insertValues, tupleValues) } } insert := &insertRowIter{ planner: planner, tableName: tableName, targetColumns: options.targetColumns, insertValues: insertValues, } _, err := insert.Next(ctx) if err != nil && err != types.ErrNoMoreRows { return err } // update the counter for bulk inserts pilosa.PerfCounterSQLBulkInsertsSec.Add(int64(len(insertValues))) return nil } // / // TODO(twg) 2023/01/23 need to refactor this type colOrder struct { realColumn int } type parquetReader struct { table *pilosa.BasicTable rowOffset int columnOrder []colOrder row []interface{} } func (pr *parquetReader) Read() ([]interface{}, error) { // need to read the row // and package it up according to the mappings if pr.rowOffset >= int(pr.table.NumRows()) { return nil, io.EOF // done } for i, col := range pr.columnOrder { pr.row[i] = pr.table.Get(col.realColumn, pr.rowOffset) } pr.rowOffset++ return pr.row, nil } func process(typeMappings []*bulkInsertMapColumn, schema *arrow.Schema) ([]colOrder, error) { ret := make([]colOrder, len(typeMappings)) find := func(n string) int { for i, x := range schema.Fields() { if n == x.Name { return i } } return -1 } for i, column := range typeMappings { iname, err := column.expr.Evaluate(nil) if err != nil { return nil, err } columnIdx := find(iname.(string)) if columnIdx < 0 { return nil, sql3.NewErrInternalf("unexpected type for mapping '%v' not found in parquet", iname) } ret[i] = colOrder{columnIdx} } return ret, nil } func NewParquetReader(ctx context.Context, mappings []*bulkInsertMapColumn, r *os.File, mem memory.Allocator) (*parquetReader, error) { pf, err := file.NewParquetReader(r) if err != nil { return nil, err } reader, err := pqarrow.NewFileReader(pf, pqarrow.ArrowReadProperties{}, mem) if err != nil { return nil, err } table, err := reader.ReadTable(ctx) if err != nil { return nil, err } m := &parquetReader{} m.table = pilosa.BasicTableFromArrow(table, mem) m.columnOrder, err = process(mappings, table.Schema()) if err != nil { return nil, err } m.row = make([]interface{}, len(m.columnOrder)) return m, nil } type bulkInsertSourceParquetRowIter struct { planner *ExecutionPlanner options *bulkInsertOptions parquetReader *parquetReader closeFunc func() hasStarted *struct{} pool memory.Allocator } var _ types.RowIterator = (*bulkInsertSourceCSVRowIter)(nil) func (i *bulkInsertSourceParquetRowIter) Next(ctx context.Context) (types.Row, error) { if i.hasStarted == nil { i.hasStarted = &struct{}{} i.pool = memory.NewGoAllocator() switch strings.ToUpper(i.options.input) { case "FILE": f, err := os.Open(i.options.sourceData) if err != nil { return nil, err } i.closeFunc = func() { f.Close() } i.parquetReader, err = NewParquetReader(ctx, i.options.mapExpressions, f, i.pool) if err != nil { return nil, sql3.NewErrInternalf("problems with parquet file '%v' '%v'", i.options.sourceData, err) } case "URL": response, err := http.Get(i.options.sourceData) if err != nil { return nil, err } if response.StatusCode != 200 { return nil, sql3.NewErrReadingDatasource(0, 0, i.options.sourceData, fmt.Sprintf("unexpected response %d", response.StatusCode)) } defer response.Body.Close() // download to temp file first tmpFile, err := os.CreateTemp("", "BulkParquetFile.parquet") if err != nil { return nil, sql3.NewErrReadingDatasource(0, 0, i.options.sourceData, fmt.Sprintf("error creating tempfile %v", err)) } i.closeFunc = func() { tmpFile.Close() } _, err = io.Copy(tmpFile, response.Body) if err != nil { return nil, sql3.NewErrReadingDatasource(0, 0, i.options.sourceData, fmt.Sprintf("error downloading url %v %v", i.options.sourceData, err)) } defer os.Remove(tmpFile.Name()) _, err = tmpFile.Seek(0, io.SeekStart) if err != nil { return nil, sql3.NewErrReadingDatasource(0, 0, i.options.sourceData, fmt.Sprintf("error reseting file for reading %v ", err)) } i.parquetReader, err = NewParquetReader(ctx, i.options.mapExpressions, tmpFile, i.pool) if err != nil { return nil, sql3.NewErrReadingDatasource(0, 0, i.options.sourceData, fmt.Sprintf("reading parquet file %v ", err)) } default: return nil, sql3.NewErrInternalf("unexpected input specification type '%s'", i.options.input) } } rec, err := i.parquetReader.Read() if err == io.EOF { return nil, types.ErrNoMoreRows } else if err != nil { pe, ok := err.(*csv.ParseError) if ok { return nil, sql3.NewErrReadingDatasource(0, 0, i.options.sourceData, fmt.Sprintf("csv parse error on line %d: %s", pe.Line, pe.Error())) } return nil, err } // now we do the mapping to the output row // current assumption is float--> DECIMAL(n) result := make([]interface{}, len(i.options.mapExpressions)) for idx := range i.options.mapExpressions { evalValue := rec[idx] mapColumn := i.options.mapExpressions[idx] switch mapColumn.colType.(type) { case *parser.DataTypeID, *parser.DataTypeInt: if intVal, ok := evalValue.(int64); ok { result[idx] = intVal } else { return nil, sql3.NewErrTypeConversionOnMap(0, 0, evalValue, mapColumn.colType.TypeDescription()) } case *parser.DataTypeIDSet: if intVal, ok := evalValue.(int64); ok { result[idx] = []int64{intVal} } else { return nil, sql3.NewErrTypeConversionOnMap(0, 0, evalValue, mapColumn.colType.TypeDescription()) } case *parser.DataTypeStringSet: if stringVal, ok := evalValue.(string); ok { result[idx] = []string{stringVal} } else { return nil, sql3.NewErrTypeConversionOnMap(0, 0, evalValue, mapColumn.colType.TypeDescription()) } case *parser.DataTypeTimestamp: if intVal, ok := evalValue.(int64); ok { // implicit conversion of int to timestamp will treat int as seconds since unix epoch result[idx] = time.Unix(intVal, 0).UTC() } else if stringVal, ok := evalValue.(string); ok { if tm, err := time.ParseInLocation(time.RFC3339Nano, stringVal, time.UTC); err == nil { result[idx] = tm } else if tm, err := time.ParseInLocation("2006-01-02", stringVal, time.UTC); err == nil { result[idx] = tm } else { return nil, sql3.NewErrTypeConversionOnMap(0, 0, stringVal, mapColumn.colType.TypeDescription()) } } case *parser.DataTypeString: if stringVal, ok := evalValue.(string); ok { result[idx] = stringVal } else { return nil, sql3.NewErrTypeConversionOnMap(0, 0, evalValue, mapColumn.colType.TypeDescription()) } case *parser.DataTypeBool: if boolVal, ok := evalValue.(bool); ok { result[idx] = boolVal } else { return nil, sql3.NewErrTypeConversionOnMap(0, 0, evalValue, mapColumn.colType.TypeDescription()) } case *parser.DataTypeDecimal: if floatVal, ok := evalValue.(float64); ok { result[idx] = pql.FromFloat64(floatVal) } else { return nil, sql3.NewErrTypeConversionOnMap(0, 0, evalValue, mapColumn.colType.TypeDescription()) } default: return nil, sql3.NewErrInternalf("unhandled type '%T'", mapColumn.colType) } } return result, nil } func (i *bulkInsertSourceParquetRowIter) Close(ctx context.Context) { if i.closeFunc != nil { i.closeFunc() } }