mirror of
https://github.com/featurebasedb/featurebase.git
synced 2026-08-28 10:54:59 +00:00
* added bulk insert tests * test idset,stringset,bool in parquet * bulk insert time coverage
1195 lines
34 KiB
Go
1195 lines
34 KiB
Go
// 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 newDateLiteralPlanExpression(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()
|
|
}
|
|
}
|