featurebase/sql3/planner/opbulkinsert.go
Seebs a158bba125 Simplify/centralize time parsing
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.
2023-04-07 17:02:37 -05:00

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()
}
}