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

fb-1915 Support large id's in NDJSON (#2290)

Browse source 
* uses json.Decoder to allow for large integer values in ndjson format in bulk import
This commit is contained in:
tgruben 2023-03-03 15:37:56 -06:00 • committed by GitHub
parent 7b8b3d8e4f
commit dba15669c6
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 99 additions and 84 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
arrow.go
View file

@ -478,7 +478,7 @@ func readTableArrow(filename string, mem memory.Allocator) (arrow.Table, error)
return nil, err
}
defer rr.Close()
records := make([]arrow.Record, rr.NumRecords(), rr.NumRecords())
records := make([]arrow.Record, rr.NumRecords())
i := 0
for {
rec, err := rr.Read()

81
sql3/planner/opbulkinsert.go
View file

@ -4,6 +4,7 @@ package planner
import (
"bufio"
"bytes"
"context"
"encoding/csv"
"encoding/json"
@ -320,7 +321,7 @@ func (i *bulkInsertSourceCSVRowIter) Next(ctx context.Context) (types.Row, error
return nil, sql3.NewErrTypeConversionOnMap(0, 0, evalValue, mapColumn.colType.TypeDescription())
}
} else {
//implicit conversion of int to timestamp will treat int as seconds since unix epoch
// implicit conversion of int to timestamp will treat int as seconds since unix epoch
result[idx] = time.Unix(intVal, 0).UTC()
}
@ -448,7 +449,6 @@ func (i *bulkInsertSourceNDJsonRowIter) Next(ctx context.Context) (types.Row, er
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
@ -507,13 +507,14 @@ func (i *bulkInsertSourceNDJsonRowIter) Next(ctx context.Context) (types.Row, er
// parse the json
v := interface{}(nil)
err := json.Unmarshal([]byte(jsonValue), &v)
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)
@ -534,16 +535,14 @@ func (i *bulkInsertSourceNDJsonRowIter) Next(ctx context.Context) (types.Row, er
mapColumn := i.options.mapExpressions[idx]
switch mapColumn.colType.(type) {
case *parser.DataTypeID, *parser.DataTypeInt:
switch v := evalValue.(type) {
case float64:
// if v is a whole number then make it an int
if v == float64(int64(v)) {
result[idx] = int64(v)
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())
@ -566,21 +565,21 @@ func (i *bulkInsertSourceNDJsonRowIter) Next(ctx context.Context) (types.Row, er
case *parser.DataTypeIDSet:
switch v := evalValue.(type) {
case float64:
// if v is a whole number then make it an int, and then turn that into an idset
if v == float64(int64(v)) {
result[idx] = []int64{int64(v)}
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 float64:
if v == float64(int64(v)) {
setValue = append(setValue, int64(v))
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())
}
@ -616,13 +615,8 @@ func (i *bulkInsertSourceNDJsonRowIter) Next(ctx context.Context) (types.Row, er
case *parser.DataTypeStringSet:
switch v := evalValue.(type) {
case float64:
if v == float64(int64(v)) {
result[idx] = []string{fmt.Sprintf("%d", int64(v))}
} else {
result[idx] = []string{fmt.Sprintf("%f", v)}
}
case json.Number:
result[idx] = []string{v.String()}
case []interface{}:
setValue := make([]string, 0)
for _, i := range v {
@ -649,11 +643,12 @@ func (i *bulkInsertSourceNDJsonRowIter) Next(ctx context.Context) (types.Row, er
case *parser.DataTypeTimestamp:
switch v := evalValue.(type) {
case float64:
case json.Number:
n, err := v.Int64()
// if v is a whole number then make it an int
if v == float64(int64(v)) {
//implicit conversion of int to timestamp will treat int as seconds since unix epoch
result[idx] = time.Unix(int64(v), 0).UTC()
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())
}
@ -684,14 +679,8 @@ func (i *bulkInsertSourceNDJsonRowIter) Next(ctx context.Context) (types.Row, er
case *parser.DataTypeString:
switch v := evalValue.(type) {
case float64:
// if a whole number make it an int
if v == float64(int64(v)) {
result[idx] = fmt.Sprintf("%d", int64(v))
} else {
result[idx] = fmt.Sprintf("%f", v)
}
case json.Number:
result[idx] = v.String()
case []interface{}:
return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription())
@ -710,10 +699,11 @@ func (i *bulkInsertSourceNDJsonRowIter) Next(ctx context.Context) (types.Row, er
case *parser.DataTypeBool:
switch v := evalValue.(type) {
case float64:
case json.Number:
// if a whole number make it an int, and convert to a bool
if v == float64(int64(v)) {
result[idx] = v > 0
n, err := v.Int64()
if err == nil {
result[idx] = n > 0
} else {
return nil, sql3.NewErrTypeConversionOnMap(0, 0, v, mapColumn.colType.TypeDescription())
}
@ -736,8 +726,12 @@ func (i *bulkInsertSourceNDJsonRowIter) Next(ctx context.Context) (types.Row, er
case *parser.DataTypeDecimal:
switch v := evalValue.(type) {
case float64:
result[idx] = pql.FromFloat64(v)
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())
@ -989,7 +983,6 @@ func (pr *parquetReader) Read() ([]interface{}, error) {
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++
@ -1155,7 +1148,7 @@ func (i *bulkInsertSourceParquetRowIter) Next(ctx context.Context) (types.Row, e
case *parser.DataTypeTimestamp:
if intVal, ok := evalValue.(int64); ok {
//implicit conversion of int to timestamp will treat int as seconds since unix epoch
// implicit conversion of int to timestamp will treat int as seconds since unix epoch
result[idx] = time.Unix(intVal, 0).UTC()
} else if stringVal, ok := evalValue.(string); ok {
if tm, err := time.ParseInLocation(time.RFC3339Nano, stringVal, time.UTC); err == nil {

37
sql3/sql_complex_test.go
View file

@ -23,6 +23,7 @@ import (
"github.com/featurebasedb/featurebase/v3/pql"
sql_test "github.com/featurebasedb/featurebase/v3/sql3/test"
"github.com/featurebasedb/featurebase/v3/test"
"github.com/featurebasedb/featurebase/v3/vprint"
"github.com/google/go-cmp/cmp"
"github.com/stretchr/testify/assert"
)
@ -3038,3 +3039,39 @@ func TestPlanner_BulkInsert_FP1916(t *testing.T) {
expected := int64(8924809397503602651)
assert.Equal(t, got, expected)
}
func TestPlanner_BulkInsert_FP1915(t *testing.T) {
c := test.MustRunCluster(t, 3)
defer c.Close()
// 8924809397503602651 is larger than 2^53 which is the largest integer value representable in float64
node := c.GetNode(0).Server
_, _, _, err := sql_test.MustQueryRows(t, node, `create table ids (
_id id,
a int,
b int);`)
assert.NoError(t, err)
_, _, _, err = sql_test.MustQueryRows(t, node, `BULK INSERT INTO ids (_id, a, b)
map ('$._id' id, '$.a' int, '$.b' int)
from x'{ "_id":8924809397503602651 , "a": 10, "b": 20 }
{ "_id":"8924809397503602652" , "a": 10, "b": 20 }'
WITH
FORMAT 'NDJSON'
INPUT 'STREAM';`)
assert.NoError(t, err)
results, _, _, err := sql_test.MustQueryRows(t, node, `select _id from ids`)
assert.NoError(t, err)
got := make([]int64, 0)
for i := range results {
got = append(got, results[i][0].(int64))
}
sort.Slice(got, func(i, j int) bool {
return got[i] < got[j]
})
vprint.VV("results %#v", results)
if diff := cmp.Diff([]int64{
8924809397503602651,
8924809397503602652,
}, got); diff != "" {
t.Fatal(diff)
}
}

13
translate_boltdb.go
View file

@ -10,14 +10,13 @@ import (
"io"
"os"
"path/filepath"
"runtime/pprof"
"sync"
"time"
"github.com/featurebasedb/featurebase/v3/roaring"
"github.com/pkg/errors"
bolt "go.etcd.io/bbolt"
"runtime/pprof"
)
var _ = pprof.StartCPUProfile
@ -102,7 +101,6 @@ func NewBoltTranslateStore(index, field string, partitionID, partitionN int, fsy
// Open opens the translate file.
func (s *BoltTranslateStore) Open() (err error) {
// add the path to the problem database if we panic handling it.
defer func() {
r := recover()
@ -111,9 +109,9 @@ func (s *BoltTranslateStore) Open() (err error) {
}
}()
if err := os.MkdirAll(filepath.Dir(s.Path), 0750); err != nil {
if err := os.MkdirAll(filepath.Dir(s.Path), 0o750); err != nil {
return errors.Wrapf(err, "mkdir %s", filepath.Dir(s.Path))
} else if s.db, err = bolt.Open(s.Path, 0600, &bolt.Options{Timeout: 1 * time.Second, NoSync: !s.fsyncEnabled, InitialMmapSize: 0}); err != nil {
} else if s.db, err = bolt.Open(s.Path, 0o600, &bolt.Options{Timeout: 1 * time.Second, NoSync: !s.fsyncEnabled, InitialMmapSize: 0}); err != nil {
return errors.Wrapf(err, "open file: %s", err)
}
@ -513,7 +511,6 @@ func (r *BoltTranslateEntryReader) ReadEntry(entry *TranslateEntry) error {
type boltWrapper struct {
tx *bolt.Tx
db *bolt.DB
}
func (w *boltWrapper) Commit() error {
@ -528,6 +525,7 @@ func (w *boltWrapper) Rollback() {
w.tx.Rollback()
}
}
func (s *BoltTranslateStore) FreeIDs() (*roaring.Bitmap, error) {
result := roaring.NewBitmap()
err := s.db.View(func(tx *bolt.Tx) error {
@ -544,11 +542,12 @@ func (s *BoltTranslateStore) FreeIDs() (*roaring.Bitmap, error) {
})
return result, err
}
func (s *BoltTranslateStore) MergeFree(tx *bolt.Tx, newIDs *roaring.Bitmap) error {
bkt := tx.Bucket(bucketFree)
b := bkt.Get(freeKey)
buf := new(bytes.Buffer)
if b != nil { //if existing combine with newIDs
if b != nil { // if existing combine with newIDs
before := roaring.NewBitmap()
err := before.UnmarshalBinary(b)
if err != nil {

50
wire_response.go
View file

@ -45,24 +45,6 @@ func (s *WireQueryResponse) UnmarshalJSON(in []byte) error {
return s.UnmarshalJSONTyped(in, false)
}
// decodeDataWithNumber allows access to integers larger than 2^53 which is normally
// not supported in JSON, this extra parse is inefficient and should be re-examined
// once all the typing is settled
func (s *WireQueryResponse) decodeDataWithNumber(in []byte) ([]interface{}, error) {
dat := make(map[string]interface{})
d := json.NewDecoder(bytes.NewBuffer(in))
d.UseNumber()
if err := d.Decode(&dat); err != nil {
return nil, err
}
if a, ok := dat["data"].([]interface{}); ok {
return a, nil
}
return nil, noArrayPresent
}
var noArrayPresent = errors.New("no data in response")
// UnmarshalJSONTyped is a temporary until we send typed values back in sql
// responses. At that point, we can get rid of the typed=false path. In order to
// do that, we need sql3 to return typed values, and we need the sql3/test/defs
@ -72,16 +54,12 @@ func (s *WireQueryResponse) UnmarshalJSONTyped(in []byte, typed bool) error {
type Alias WireQueryResponse
var aux Alias
if err := json.Unmarshal(in, &aux); err != nil {
dec := json.NewDecoder(bytes.NewReader(in))
dec.UseNumber()
err := dec.Decode(&aux)
if err != nil {
return err
}
bypass, err := s.decodeDataWithNumber(in)
if err != nil {
if err != noArrayPresent { // no data is not an error but just allows the compiler
return err
}
}
*s = WireQueryResponse(aux)
// If the SQLResponse contains an error, don't bother doing any conversions
@ -100,7 +78,12 @@ func (s *WireQueryResponse) UnmarshalJSONTyped(in []byte, typed bool) error {
for k, v := range fld.TypeInfo {
switch k {
case "scale":
fld.TypeInfo[k] = int64(v.(float64))
switch n := v.(type) {
case float64:
fld.TypeInfo[k] = int64(n)
case json.Number:
fld.TypeInfo[k], _ = n.Int64()
}
}
}
}
@ -110,7 +93,7 @@ func (s *WireQueryResponse) UnmarshalJSONTyped(in []byte, typed bool) error {
for j, hdr := range s.Schema.Fields {
switch hdr.BaseType {
case dax.BaseTypeID, dax.BaseTypeInt:
jn := bypass[i].([]interface{})[j] // forced to do this to handle ints bigger than 2^53
jn := s.Data[i][j]
if v, ok := jn.(json.Number); ok {
if x, err := v.Int64(); err == nil {
s.Data[i][j] = x
@ -120,8 +103,7 @@ func (s *WireQueryResponse) UnmarshalJSONTyped(in []byte, typed bool) error {
}
case dax.BaseTypeIDSet:
if _, ok := s.Data[i][j].([]interface{}); ok {
src := bypass[i].([]interface{})[j].([]interface{})
if src, ok := s.Data[i][j].([]interface{}); ok {
if typed {
val := make(IDSet, len(src))
for k := range src {
@ -148,7 +130,7 @@ func (s *WireQueryResponse) UnmarshalJSONTyped(in []byte, typed bool) error {
}
case dax.BaseTypeDecimal:
if _, ok := s.Data[i][j].(float64); ok {
if jn, ok := s.Data[i][j].(json.Number); ok {
var scale int64
if scaleVal, ok := hdr.TypeInfo["scale"]; !ok {
return errors.New("decimal does not have a scale")
@ -159,7 +141,11 @@ func (s *WireQueryResponse) UnmarshalJSONTyped(in []byte, typed bool) error {
}
format := fmt.Sprintf("%%.%df", scale)
dec, err := pql.ParseDecimal(fmt.Sprintf(format, s.Data[i][j]))
f, err := jn.Float64()
if err != nil {
return errors.Wrap(err, "parsing decimal")
}
dec, err := pql.ParseDecimal(fmt.Sprintf(format, f))
if err != nil {
return errors.Wrap(err, "parsing decimal")
}