scriptoriumadmin/featurebase
1
0
Fork 0
mirror of https://github.com/featurebasedb/featurebase.git synced 2026-08-28 02:44:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

bulk import support for parquet files (#2226)

Browse source 
* bulk import support form parquet files
This commit is contained in:
tgruben 2023-01-26 08:32:39 -06:00 • committed by GitHub
parent 49ef905b89
commit 90e2808f52
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 475 additions and 50 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

67
arrow.go
View file

@ -60,7 +60,7 @@ func (e *executor) executeArrow(ctx context.Context, qcx *Qcx, index string, c *
mu.Unlock()
return e.executeArrowShard(ctx, qcx, index, c, shard, pool, columnFilter)
}
tables := make([]*basicTable, 0)
tables := make([]*BasicTable, 0)
reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
mu.Lock()
@ -70,7 +70,7 @@ func (e *executor) executeArrow(ctx context.Context, qcx *Qcx, index string, c *
return prev
}
switch t := v.(type) {
case *basicTable:
case *BasicTable:
if t.resolver != nil {
mu.Lock()
@ -93,104 +93,103 @@ func (e *executor) executeArrow(ctx context.Context, qcx *Qcx, index string, c *
return nil, err
}
if len(tables) == 0 {
return &basicTable{name: "empty"}, nil
return &BasicTable{name: "empty"}, nil
}
tbl := Concat(tables[0].Schema(), tables, pool)
r := dataframe.NewChunkResolver(tbl.Column(0))
return &basicTable{resolver: &r, table: tbl}, nil
return &BasicTable{resolver: &r, table: tbl}, nil
}
type basicTable struct {
type BasicTable struct {
resolver dataframe.Resolver
table arrow.Table
filtered bool
name string
}
func (st *basicTable) Name() string {
func (st *BasicTable) Name() string {
return st.name
}
func (st *basicTable) Schema() *arrow.Schema {
func (st *BasicTable) Schema() *arrow.Schema {
if st.table != nil {
return st.table.Schema()
}
return &arrow.Schema{}
}
func (st *basicTable) IsFiltered() bool {
func (st *BasicTable) IsFiltered() bool {
return st.filtered
}
func (st *basicTable) NumRows() int64 {
func (st *BasicTable) NumRows() int64 {
if st.resolver == nil {
return 0
}
return int64(st.resolver.NumRows())
}
func (st *basicTable) NumCols() int64 {
func (st *BasicTable) NumCols() int64 {
if st.table != nil {
return st.table.NumCols()
}
return 0
}
func (st *basicTable) Column(i int) *arrow.Column {
func (st *BasicTable) Column(i int) *arrow.Column {
if st.table != nil {
return st.table.Column(i)
}
return nil
}
func (st *basicTable) Retain() {
func (st *BasicTable) Retain() {
if st.table != nil {
st.table.Retain()
}
}
func (st *basicTable) Release() {
func (st *BasicTable) Release() {
if st.table != nil {
st.table.Retain()
}
}
func (st *basicTable) Get(column, row int) interface{} {
func (st *BasicTable) Get(column, row int) interface{} {
field := st.Schema().Field(column)
c, i := st.resolver.Resolve(row)
chunk := st.Column(column).Data().Chunk(c)
switch field.Type.(type) {
// case *arrow.BooleanType:
// v := chunk.(*array.Boolean).BooleanValues()
// return v[i]
case *arrow.BooleanType:
return chunk.(*array.Boolean).Value(i)
case *arrow.Int8Type:
v := chunk.(*array.Int8).Int8Values()
return v[i]
return int64(v[i])
case *arrow.Int16Type:
v := chunk.(*array.Int16).Int16Values()
return v[i]
return int64(v[i])
case *arrow.Int32Type:
v := chunk.(*array.Int32).Int32Values()
return v[i]
return int64(v[i])
case *arrow.Int64Type:
v := chunk.(*array.Int64).Int64Values()
return v[i]
return int64(v[i])
case *arrow.Uint8Type:
v := chunk.(*array.Uint8).Uint8Values()
return v[i]
return uint64(v[i])
case *arrow.Uint16Type:
v := chunk.(*array.Uint16).Uint16Values()
return v[i]
return uint64(v[i])
case *arrow.Uint32Type:
v := chunk.(*array.Uint32).Uint32Values()
return v[i]
return uint64(v[i])
case *arrow.Uint64Type:
v := chunk.(*array.Uint64).Uint64Values()
return v[i]
case *arrow.Float32Type:
v := chunk.(*array.Float32).Float32Values()
return v[i]
return float64(v[i])
case *arrow.Float64Type:
v := chunk.(*array.Float64).Float64Values()
return v[i]
@ -265,7 +264,7 @@ func appendData(bldr array.Builder, v interface{}) {
}
}
func Concat(schema *arrow.Schema, tables []*basicTable, mem memory.Allocator) arrow.Table {
func Concat(schema *arrow.Schema, tables []*BasicTable, mem memory.Allocator) arrow.Table {
if len(tables) == 1 {
if !tables[0].IsFiltered() {
return tables[0]
@ -307,7 +306,7 @@ func Concat(schema *arrow.Schema, tables []*basicTable, mem memory.Allocator) ar
return array.NewTable(schema, cols, -1)
}
func (st *basicTable) MarshalJSON() ([]byte, error) {
func (st *BasicTable) MarshalJSON() ([]byte, error) {
results := make(map[string]interface{})
n := 0
if st.table != nil {
@ -326,10 +325,10 @@ func (st *basicTable) MarshalJSON() ([]byte, error) {
return json.Marshal(results)
}
func BasicTableFromArrow(table arrow.Table, mem memory.Allocator) *basicTable {
func BasicTableFromArrow(table arrow.Table, mem memory.Allocator) *BasicTable {
col := table.Column(0)
r := dataframe.NewChunkResolver(col)
return &basicTable{resolver: &r, table: table}
return &BasicTable{resolver: &r, table: table}
}
func filterColumns(filters []string, table arrow.Table) arrow.Table {
@ -359,7 +358,7 @@ func filterColumns(filters []string, table arrow.Table) arrow.Table {
return array.NewTable(filterdSchema, cols, table.NumRows())
}
func (e *executor) executeArrowShard(ctx context.Context, qcx *Qcx, index string, c *pql.Call, shard uint64, pool memory.Allocator, columnFilter []string) (*basicTable, error) {
func (e *executor) executeArrowShard(ctx context.Context, qcx *Qcx, index string, c *pql.Call, shard uint64, pool memory.Allocator, columnFilter []string) (*BasicTable, error) {
name := fmt.Sprintf("a. %v", shard)
span, _ := tracing.StartSpanFromContext(ctx, "Executor.executeArrowShard")
defer span.Finish()
@ -373,7 +372,7 @@ func (e *executor) executeArrowShard(ctx context.Context, qcx *Qcx, index string
filter = row
if !filter.Any() {
// no need to actuall run the query for its not operating against any values
return &basicTable{name: name}, nil
return &BasicTable{name: name}, nil
}
}
//
@ -387,7 +386,7 @@ func (e *executor) executeArrowShard(ctx context.Context, qcx *Qcx, index string
fname := idx.GetDataFramePath(shard)
if !e.dataFrameExists(fname) {
return &basicTable{name: name}, nil
return &BasicTable{name: name}, nil
}
table, err := e.getDataTable(ctx, fname, pool)
@ -407,7 +406,7 @@ func (e *executor) executeArrowShard(ctx context.Context, qcx *Qcx, index string
resolver = &p
if filter != nil {
if len(ids) == 0 {
return &basicTable{name: name}, nil
return &BasicTable{name: name}, nil
}
resolver, err = filterDataframe(resolver, pool, ids)
if err != nil {
@ -415,7 +414,7 @@ func (e *executor) executeArrowShard(ctx context.Context, qcx *Qcx, index string
}
}
table.Retain()
return &basicTable{resolver: resolver, table: table, filtered: filter != nil, name: name}, nil
return &BasicTable{resolver: resolver, table: table, filtered: filter != nil, name: name}, nil
}
func (e *executor) dataFrameExists(fname string) bool {

2
executor.go
View file

@ -347,7 +347,7 @@ func safeCopy(resp QueryResponse) (out QueryResponse) {
case *dataframe.DataFrame:
// dumpTable(x)
out.Results = append(out.Results, x)
case *basicTable:
case *BasicTable:
// dumpTable(x)
out.Results = append(out.Results, x)
case ExtractedIDMatrixSorted:

10
sql3/planner/compilebulkinsert.go
View file

@ -158,7 +158,7 @@ func (p *ExecutionPlanner) compileBulkInsertStatement(stmt *parser.BulkInsertSta
// analyzeBulkInsertStatement analyzes a BULK INSERT statement and returns an
// error if anything is invalid.
func (p *ExecutionPlanner) analyzeBulkInsertStatement(stmt *parser.BulkInsertStatement) error {
//check referred to table exists
// check referred to table exists
tableName := parser.IdentName(stmt.Table)
tname := dax.TableName(tableName)
tbl, err := p.schemaAPI.TableByName(context.Background(), tname)
@ -203,6 +203,14 @@ func (p *ExecutionPlanner) analyzeBulkInsertStatement(stmt *parser.BulkInsertSta
return sql3.NewErrIntegerLiteral(im.MapExpr.Pos().Line, im.MapExpr.Pos().Column)
}
}
case "PARQUET":
// for parquet the map expressions need to be string values
// that represent the offsets in the source file
for _, im := range stmt.MapList {
if !(im.MapExpr.IsLiteral() && typeIsString(im.MapExpr.DataType())) {
return sql3.NewErrStringLiteral(im.MapExpr.Pos().Line, im.MapExpr.Pos().Column)
}
}
case "NDJSON":
// for ndjson the map expressions need to be string values
// that represent json path expressions

267
sql3/planner/opbulkinsert.go
View file

@ -17,6 +17,10 @@ import (
"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"
@ -139,7 +143,7 @@ func (p *PlanOpBulkInsert) Children() []types.PlanOperator {
func (p *PlanOpBulkInsert) Iterator(ctx context.Context, row types.Row) (types.RowIterator, error) {
switch strings.ToUpper(p.options.format) {
case "CSV":
return &bulkInsertCSVRowIter{
return &bulkInsertLineRowIter{
planner: p.planner,
tableName: p.tableName,
options: p.options,
@ -159,6 +163,16 @@ func (p *PlanOpBulkInsert) Iterator(ctx context.Context, row types.Row) (types.R
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)
@ -184,7 +198,6 @@ type bulkInsertSourceCSVRowIter struct {
var _ types.RowIterator = (*bulkInsertSourceCSVRowIter)(nil)
func (i *bulkInsertSourceCSVRowIter) Next(ctx context.Context) (types.Row, error) {
if i.hasStarted == nil {
i.hasStarted = &struct{}{}
@ -334,7 +347,12 @@ func (i *bulkInsertSourceCSVRowIter) Close(ctx context.Context) {
}
}
type bulkInsertCSVRowIter struct {
type bulkInsertBasicRowIter interface {
Next(context.Context) (types.Row, error)
Close(context.Context)
}
type bulkInsertLineRowIter struct {
planner *ExecutionPlanner
tableName string
options *bulkInsertOptions
@ -342,12 +360,12 @@ type bulkInsertCSVRowIter struct {
currentBatch [][]interface{}
sourceIter *bulkInsertSourceCSVRowIter
sourceIter bulkInsertBasicRowIter
}
var _ types.RowIterator = (*bulkInsertCSVRowIter)(nil)
var _ types.RowIterator = (*bulkInsertLineRowIter)(nil)
func (i *bulkInsertCSVRowIter) Next(ctx context.Context) (types.Row, error) {
func (i *bulkInsertLineRowIter) Next(ctx context.Context) (types.Row, error) {
defer i.sourceIter.Close(ctx)
for {
row, err := i.sourceIter.Next(ctx)
@ -404,7 +422,6 @@ type bulkInsertSourceNDJsonRowIter struct {
var _ types.RowIterator = (*bulkInsertSourceNDJsonRowIter)(nil)
func (i *bulkInsertSourceNDJsonRowIter) Next(ctx context.Context) (types.Row, error) {
if i.hasStarted == nil {
i.hasStarted = &struct{}{}
@ -484,7 +501,6 @@ func (i *bulkInsertSourceNDJsonRowIter) Next(ctx context.Context) (types.Row, er
// parse the json
v := interface{}(nil)
err := json.Unmarshal([]byte(jsonValue), &v)
if err != nil {
return nil, sql3.NewErrParsingJSON(0, 0, jsonValue, err.Error())
}
@ -876,7 +892,6 @@ func processColumnValue(rawValue interface{}, targetType parser.ExprDataType) (t
}
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
@ -888,7 +903,7 @@ func processBatch(ctx context.Context, planner *ExecutionPlanner, tableName stri
for _, row := range currentBatch {
tupleValues := []types.PlanExpression{}
//handle each transform
// handle each transform
for idx, mc := range options.transformExpressions {
rawValue, err := mc.Evaluate(row)
if err != nil {
@ -909,7 +924,6 @@ func processBatch(ctx context.Context, planner *ExecutionPlanner, tableName stri
}
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
@ -928,7 +942,6 @@ func processBatch(ctx context.Context, planner *ExecutionPlanner, tableName stri
}
insertValues = append(insertValues, tupleValues)
}
}
insert := &insertRowIter{
@ -948,3 +961,233 @@ func processBatch(ctx context.Context, planner *ExecutionPlanner, tableName stri
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 {
// vprint.VV("check row:%v col:%v", pr.rowOffset, col.realColumn)
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 {
result[idx] = time.UnixMilli(intVal).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(time.RFC3339, 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()
}
}

177
sql3/sql_complex_test.go
View file

@ -4,6 +4,8 @@ package sql3_test
import (
"context"
"fmt"
"net/http"
"net/http/httptest"
"os"
"reflect"
"sort"
@ -11,6 +13,11 @@ import (
"testing"
"time"
"github.com/apache/arrow/go/v10/arrow"
"github.com/apache/arrow/go/v10/arrow/array"
"github.com/apache/arrow/go/v10/arrow/memory"
"github.com/apache/arrow/go/v10/parquet"
"github.com/apache/arrow/go/v10/parquet/pqarrow"
pilosa "github.com/featurebasedb/featurebase/v3"
"github.com/featurebasedb/featurebase/v3/dax"
"github.com/featurebasedb/featurebase/v3/pql"
@ -119,7 +126,6 @@ func TestPlanner_SystemTableFanout(t *testing.T) {
t.Fatal(diff)
}
})
}
func TestPlanner_Show(t *testing.T) {
@ -2631,3 +2637,172 @@ func TestPlanner_BulkInsert_FB1831(t *testing.T) {
t.Fatal("Expecting to be equal")
}
}
type tb struct {
Name string
Type arrow.DataType
Value interface{}
}
func simpleParquetMaker(t *testing.T, f *os.File, numRows int64, input []tb) {
t.Helper()
mem := memory.NewGoAllocator()
chunks := make([]arrow.Array, 0)
for i := range input {
switch input[i].Type {
case arrow.PrimitiveTypes.Int64:
idbuild := array.NewInt64Builder(mem)
idbuild.AppendValues(input[i].Value.([]int64), nil)
newChunk := idbuild.NewArray()
chunks = append(chunks, newChunk)
case arrow.PrimitiveTypes.Float64:
fbuild := array.NewFloat64Builder(mem)
fbuild.AppendValues(input[i].Value.([]float64), nil)
newChunk := fbuild.NewArray()
chunks = append(chunks, newChunk)
case arrow.BinaryTypes.String:
sbuild := array.NewStringBuilder(mem)
sbuild.AppendValues(input[i].Value.([]string), nil)
newChunk := sbuild.NewArray()
chunks = append(chunks, newChunk)
}
}
// make schema
fields := make([]arrow.Field, len(input))
for i := range input {
fields[i].Name = input[i].Name
fields[i].Type = input[i].Type
}
schema := arrow.NewSchema(fields, nil)
rec := array.NewRecord(schema, chunks, numRows)
table := array.NewTableFromRecords(schema, []arrow.Record{rec})
props := parquet.NewWriterProperties(parquet.WithDictionaryDefault(false))
arrProps := pqarrow.DefaultWriterProps()
pqarrow.WriteTable(table, f, 4096, props, arrProps)
}
func TestPlanner_BulkInsertParquet(t *testing.T) {
c := test.MustRunUnsharedCluster(t, 1)
defer c.Close()
t.Run("BulkFromLocalFile", func(t *testing.T) {
// check that can pull parquet file from local file
_, _, err := sql_test.MustQueryRows(t, c.GetNode(0).Server, `create table j1 (_id ID, a INT, b DECIMAL(2), c STRING);`)
assert.NoError(t, err)
tmpfile, err := os.CreateTemp("", "BulkParquetFile.parquet")
assert.NoError(t, err)
defer os.Remove(tmpfile.Name())
// create a parquet file with all the example data
simpleParquetMaker(t, tmpfile, 2, []tb{
{Name: "id", Type: arrow.PrimitiveTypes.Int64, Value: []int64{1, 2}},
{Name: "int64V", Type: arrow.PrimitiveTypes.Int64, Value: []int64{42, 7}},
{Name: "float64V", Type: arrow.PrimitiveTypes.Float64, Value: []float64{3.14159, 1.61803}},
{Name: "stringV", Type: arrow.BinaryTypes.String, Value: []string{"pi", "goldenratio"}},
})
_, _, err = sql_test.MustQueryRows(t, c.GetNode(0).Server, fmt.Sprintf(`bulk insert
into j1 (_id,a,b,c )
map(
'id' id,
'int64V' INT,
'float64V' DECIMAL(2),
'stringV' STRING)
from
'%s'
WITH FORMAT 'PARQUET'
INPUT 'FILE';`, tmpfile.Name()))
assert.NoError(t, err)
results, _, err := sql_test.MustQueryRows(t, c.GetNode(0).Server, `select _id, a,c from j1`)
assert.NoError(t, err)
if diff := cmp.Diff([][]interface{}{
{int64(1), int64(42), "pi"},
{int64(2), int64(7), "goldenratio"},
}, results); diff != "" {
t.Fatal(diff)
}
results, _, err = sql_test.MustQueryRows(t, c.GetNode(0).Server, `select b from j1`)
assert.NoError(t, err)
d, _ := pql.FromFloat64WithScale(3.14159, 2)
if !pql.Decimal.EqualTo(d, results[0][0].(pql.Decimal)) {
t.Fatal("Should be equal")
}
})
t.Run("BulkFromUrl", func(t *testing.T) {
// check that can pull parquet file from URL and load
_, _, err := sql_test.MustQueryRows(t, c.GetNode(0).Server, `create table j2 (_id ID, a INT, b STRING);`)
assert.NoError(t, err)
tmpfile, err := os.CreateTemp("", "BulkParquetFile.parquet")
assert.NoError(t, err)
defer os.Remove(tmpfile.Name())
simpleParquetMaker(t, tmpfile, 1, []tb{
{Name: "id", Type: arrow.PrimitiveTypes.Int64, Value: []int64{1}},
{Name: "int64V", Type: arrow.PrimitiveTypes.Int64, Value: []int64{42}},
{Name: "stringV", Type: arrow.BinaryTypes.String, Value: []string{"pi"}},
})
// create a parquet file with all the example data
mux := http.NewServeMux()
ts := httptest.NewServer(mux)
defer ts.Close()
mux.HandleFunc("/static", func(w http.ResponseWriter, r *http.Request) {
payload, _ := os.ReadFile(tmpfile.Name())
w.Header().Set("Content-Type", "application/octet-stream")
w.WriteHeader(200)
w.Write(payload)
})
_, _, err = sql_test.MustQueryRows(t, c.GetNode(0).Server, fmt.Sprintf(`bulk insert
into j2 (_id,a,b )
map(
'id' id,
'int64V' INT,
'stringV' STRING)
from
'%s'
WITH FORMAT 'PARQUET'
INPUT 'URL';`, ts.URL+"/static"))
assert.NoError(t, err)
results, _, err := sql_test.MustQueryRows(t, c.GetNode(0).Server, `select _id, a,b from j2`)
assert.NoError(t, err)
if diff := cmp.Diff([][]interface{}{
{int64(1), int64(42), "pi"},
}, results); diff != "" {
t.Fatal(diff)
}
})
t.Run("FileTimeStamp", func(t *testing.T) {
_, _, err := sql_test.MustQueryRows(t, c.GetNode(0).Server, `create table continuum (_id ID, created timestamp, updated timestamp);`)
assert.NoError(t, err)
tmpfile, err := os.CreateTemp("", "BulkParquetFile.parquet")
assert.NoError(t, err)
defer os.Remove(tmpfile.Name())
// create a parquet file with all the example data
now := time.Now()
simpleParquetMaker(t, tmpfile, 1, []tb{
{Name: "id", Type: arrow.PrimitiveTypes.Int64, Value: []int64{1}},
{Name: "unixtime", Type: arrow.PrimitiveTypes.Int64, Value: []int64{now.UnixMilli()}},
{Name: "stringtime", Type: arrow.BinaryTypes.String, Value: []string{now.Format(time.RFC3339)}},
})
_, _, err = sql_test.MustQueryRows(t, c.GetNode(0).Server, fmt.Sprintf(`bulk insert
into continuum (_id,created,updated )
map(
'id' id,
'unixtime' timestamp,
'stringtime' timestamp )
from
'%s'
WITH FORMAT 'PARQUET'
INPUT 'FILE';`, tmpfile.Name()))
assert.NoError(t, err)
results, _, err := sql_test.MustQueryRows(t, c.GetNode(0).Server, `select _id, created,updated from continuum`)
assert.NoError(t, err)
row := results[0]
assert.Equal(t, row[1], row[2])
})
}

2
sql3/test/helpers.go
View file

@ -53,7 +53,7 @@ func MustQueryRows(tb testing.TB, svr *featurebase.Server, q string) ([][]interf
return nil, nil, err
}
}
//temporarily transform to Columns()
// temporarily transform to Columns()
cols := make([]*featurebase.WireQueryField, 0)
for _, oc := range ocolumns {
cols = append(cols, &featurebase.WireQueryField{