SIMDJSON S3 select input (#8401)

pkg/s3select/simdj/record.go

@@ -0,0 +1,228 @@
+/*
+ * MinIO Cloud Storage, (C) 2019 MinIO, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package simdj
+
+import (
+	"encoding/csv"
+	"fmt"
+	"io"
+
+	"github.com/bcicen/jstream"
+	"github.com/minio/minio/pkg/s3select/json"
+	"github.com/minio/minio/pkg/s3select/sql"
+	"github.com/minio/simdjson-go"
+)
+
+// Record - is JSON record.
+type Record struct {
+	// object
+	object simdjson.Object
+}
+
+// Get - gets the value for a column name.
+func (r *Record) Get(name string) (*sql.Value, error) {
+	elem := r.object.FindKey(name, nil)
+	if elem == nil {
+		return nil, nil
+	}
+	return iterToValue(elem.Iter)
+}
+
+func iterToValue(iter simdjson.Iter) (*sql.Value, error) {
+	switch iter.Type() {
+	case simdjson.TypeString:
+		v, err := iter.String()
+		if err != nil {
+			return nil, err
+		}
+		return sql.FromString(v), nil
+	case simdjson.TypeFloat:
+		v, err := iter.Float()
+		if err != nil {
+			return nil, err
+		}
+		return sql.FromFloat(v), nil
+	case simdjson.TypeInt:
+		v, err := iter.Int()
+		if err != nil {
+			return nil, err
+		}
+		return sql.FromInt(v), nil
+	case simdjson.TypeUint:
+		v, err := iter.Int()
+		if err != nil {
+			// Can't fit into int, convert to float.
+			v, err := iter.Float()
+			return sql.FromFloat(v), err
+		}
+		return sql.FromInt(v), nil
+	case simdjson.TypeBool:
+		v, err := iter.Bool()
+		if err != nil {
+			return nil, err
+		}
+		return sql.FromBool(v), nil
+	case simdjson.TypeNull:
+		return sql.FromNull(), nil
+	case simdjson.TypeObject, simdjson.TypeArray:
+		b, err := iter.MarshalJSON()
+		return sql.FromBytes(b), err
+	}
+	return nil, fmt.Errorf("iterToValue: unknown JSON type: %s", iter.Type().String())
+}
+
+// Reset the record.
+func (r *Record) Reset() {
+	r.object = simdjson.Object{}
+}
+
+// Clone the record and if possible use the destination provided.
+func (r *Record) Clone(dst sql.Record) sql.Record {
+	other, ok := dst.(*Record)
+	if !ok {
+		other = &Record{}
+	}
+	other.object = r.object
+	return other
+}
+
+// CloneTo clones the record to a json Record.
+// Values are only unmashaled on object level.
+func (r *Record) CloneTo(dst *json.Record) (sql.Record, error) {
+	if dst == nil {
+		dst = &json.Record{SelectFormat: sql.SelectFmtJSON}
+	}
+	dst.Reset()
+	elems, err := r.object.Parse(nil)
+	if err != nil {
+		return nil, err
+	}
+	if cap(dst.KVS) < len(elems.Elements) {
+		dst.KVS = make(jstream.KVS, 0, len(elems.Elements))
+	}
+	for _, elem := range elems.Elements {
+		v, err := sql.IterToValue(elem.Iter)
+		if err != nil {
+			v, err = elem.Iter.Interface()
+			if err != nil {
+				panic(err)
+			}
+		}
+		dst.KVS = append(dst.KVS, jstream.KV{
+			Key:   elem.Name,
+			Value: v,
+		})
+	}
+	return dst, nil
+}
+
+// Set - sets the value for a column name.
+func (r *Record) Set(name string, value *sql.Value) (sql.Record, error) {
+	dst, err := r.CloneTo(nil)
+	if err != nil {
+		return nil, err
+	}
+	return dst.Set(name, value)
+}
+
+// WriteCSV - encodes to CSV data.
+func (r *Record) WriteCSV(writer io.Writer, fieldDelimiter rune) error {
+	csvRecord := make([]string, 0, 10)
+	var tmp simdjson.Iter
+	obj := r.object
+allElems:
+	for {
+		_, typ, err := obj.NextElement(&tmp)
+		if err != nil {
+			return err
+		}
+		var columnValue string
+		switch typ {
+		case simdjson.TypeNull, simdjson.TypeFloat, simdjson.TypeUint, simdjson.TypeInt, simdjson.TypeBool, simdjson.TypeString:
+			val, err := tmp.StringCvt()
+			if err != nil {
+				return err
+			}
+			columnValue = val
+		case simdjson.TypeObject, simdjson.TypeArray:
+			b, err := tmp.MarshalJSON()
+			if err != nil {
+				return err
+			}
+			columnValue = string(b)
+		case simdjson.TypeNone:
+			break allElems
+		default:
+			return fmt.Errorf("cannot marshal unhandled type: %s", typ.String())
+		}
+		csvRecord = append(csvRecord, columnValue)
+	}
+	w := csv.NewWriter(writer)
+	w.Comma = fieldDelimiter
+	if err := w.Write(csvRecord); err != nil {
+		return err
+	}
+	w.Flush()
+	if err := w.Error(); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// Raw - returns the underlying representation.
+func (r *Record) Raw() (sql.SelectObjectFormat, interface{}) {
+	return sql.SelectFmtSIMDJSON, r.object
+}
+
+// WriteJSON - encodes to JSON data.
+func (r *Record) WriteJSON(writer io.Writer) error {
+	o := r.object
+	elems, err := o.Parse(nil)
+	if err != nil {
+		return err
+	}
+	b, err := elems.MarshalJSON()
+	if err != nil {
+		return err
+	}
+	n, err := writer.Write(b)
+	if err != nil {
+		return err
+	}
+	if n != len(b) {
+		return io.ErrShortWrite
+	}
+	return nil
+}
+
+// Replace the underlying buffer of json data.
+func (r *Record) Replace(k interface{}) error {
+	v, ok := k.(simdjson.Object)
+	if !ok {
+		return fmt.Errorf("cannot replace internal data in simd json record with type %T", k)
+	}
+	r.object = v
+	return nil
+}
+
+// NewRecord - creates new empty JSON record.
+func NewRecord(f sql.SelectObjectFormat, obj simdjson.Object) *Record {
+	return &Record{
+		object: obj,
+	}
+}