diff --git a/arrow.go b/arrow.go index 3ddaa4178..3e268d08f 100644 --- a/arrow.go +++ b/arrow.go @@ -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 { diff --git a/executor.go b/executor.go index 24dc1d3a3..7bec99f58 100644 --- a/executor.go +++ b/executor.go @@ -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: diff --git a/sql3/planner/compilebulkinsert.go b/sql3/planner/compilebulkinsert.go index a72cfdf99..117812189 100644 --- a/sql3/planner/compilebulkinsert.go +++ b/sql3/planner/compilebulkinsert.go @@ -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 diff --git a/sql3/planner/opbulkinsert.go b/sql3/planner/opbulkinsert.go index d8bd33413..eeff56b2b 100644 --- a/sql3/planner/opbulkinsert.go +++ b/sql3/planner/opbulkinsert.go @@ -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() + } +} diff --git a/sql3/sql_complex_test.go b/sql3/sql_complex_test.go index 03ca4700b..5e496907f 100644 --- a/sql3/sql_complex_test.go +++ b/sql3/sql_complex_test.go @@ -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]) + }) +} diff --git a/sql3/test/helpers.go b/sql3/test/helpers.go index 1208f18cf..2cdef180e 100644 --- a/sql3/test/helpers.go +++ b/sql3/test/helpers.go @@ -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{