mirror of
https://github.com/featurebasedb/featurebase.git
synced 2026-08-28 02:44:59 +00:00
We had a bunch of different places which had basically the same logic, except that some were testing both RFC3339Nano and RFC3339 formats, and some weren't. This turns out not to matter, because the fractional second part is always permitted and never required, so those two formats are identical. Mostly, though, we now ensure that everything we do that is trying to convert timestamps has the same logic, so if we want to make changes to that logic, we have a central point, which lives in the parser. This came out of an attempt to figure out why the RFC3339 case wasn't getting any test coverage.
1189 lines
33 KiB
Go
1189 lines
33 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 {
|
|
tm, err := parser.ConvertStringToTimestamp(evalValue)
|
|
if err != nil {
|
|
return nil, sql3.NewErrTypeConversionOnMap(0, 0, evalValue, mapColumn.colType.TypeDescription())
|
|
}
|
|
result[idx] = tm
|
|
} 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:
|
|
tm, err := parser.ConvertStringToTimestamp(v)
|
|
if err != nil {
|
|
return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription())
|
|
}
|
|
result[idx] = tm
|
|
|
|
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 {
|
|
tm, err := parser.ConvertStringToTimestamp(stringVal)
|
|
if err != nil {
|
|
return nil, sql3.NewErrTypeConversionOnMap(0, 0, stringVal, mapColumn.colType.TypeDescription())
|
|
}
|
|
result[idx] = tm
|
|
}
|
|
|
|
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()
|
|
}
|
|
}
|