From dba15669c6ed732e7cf972f58e496fb5ceb6f44b Mon Sep 17 00:00:00 2001 From: tgruben Date: Fri, 3 Mar 2023 15:37:56 -0600 Subject: [PATCH] fb-1915 Support large id's in NDJSON (#2290) * uses json.Decoder to allow for large integer values in ndjson format in bulk import --- arrow.go | 2 +- sql3/planner/opbulkinsert.go | 81 ++++++++++++++++-------------------- sql3/sql_complex_test.go | 37 ++++++++++++++++ translate_boltdb.go | 13 +++--- wire_response.go | 50 ++++++++-------------- 5 files changed, 99 insertions(+), 84 deletions(-) diff --git a/arrow.go b/arrow.go index 4b416351e..348dec94d 100644 --- a/arrow.go +++ b/arrow.go @@ -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() diff --git a/sql3/planner/opbulkinsert.go b/sql3/planner/opbulkinsert.go index 223aa66ac..fb062fda3 100644 --- a/sql3/planner/opbulkinsert.go +++ b/sql3/planner/opbulkinsert.go @@ -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 { diff --git a/sql3/sql_complex_test.go b/sql3/sql_complex_test.go index ff8bc5386..69955408b 100644 --- a/sql3/sql_complex_test.go +++ b/sql3/sql_complex_test.go @@ -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) + } +} diff --git a/translate_boltdb.go b/translate_boltdb.go index 7a3ccfbe7..92139e87a 100644 --- a/translate_boltdb.go +++ b/translate_boltdb.go @@ -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 { diff --git a/wire_response.go b/wire_response.go index 1eff5351e..2aa91ff2f 100644 --- a/wire_response.go +++ b/wire_response.go @@ -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") }