SIMDJSON S3 select input (#8401)

pkg/s3select/sql/record.go

@@ -17,9 +17,10 @@
 package sql
 
 import (
+	"fmt"
 	"io"
 
-	"github.com/bcicen/jstream"
+	"github.com/minio/simdjson-go"
 )
 
 // SelectObjectFormat specifies the format of the underlying data
@@ -32,6 +33,8 @@ const (
 	SelectFmtCSV
 	// SelectFmtJSON - JSON format
 	SelectFmtJSON
+	// SelectFmtSIMDJSON - SIMD JSON format
+	SelectFmtSIMDJSON
 	// SelectFmtParquet - Parquet format
 	SelectFmtParquet
 )
@@ -39,7 +42,10 @@ const (
 // Record - is a type containing columns and their values.
 type Record interface {
 	Get(name string) (*Value, error)
-	Set(name string, value *Value) error
+
+	// Set a value.
+	// Can return a different record type.
+	Set(name string, value *Value) (Record, error)
 	WriteCSV(writer io.Writer, fieldDelimiter rune) error
 	WriteJSON(writer io.Writer) error
 
@@ -51,5 +57,77 @@ type Record interface {
 	Raw() (SelectObjectFormat, interface{})
 
 	// Replaces the underlying data
-	Replace(k jstream.KVS) error
+	Replace(k interface{}) error
+}
+
+// IterToValue converts a simdjson Iter to its underlying value.
+// Objects are returned as simdjson.Object
+// Arrays are returned as []interface{} with parsed values.
+func IterToValue(iter simdjson.Iter) (interface{}, error) {
+	switch iter.Type() {
+	case simdjson.TypeString:
+		v, err := iter.String()
+		if err != nil {
+			return nil, err
+		}
+		return v, nil
+	case simdjson.TypeFloat:
+		v, err := iter.Float()
+		if err != nil {
+			return nil, err
+		}
+		return v, nil
+	case simdjson.TypeInt:
+		v, err := iter.Int()
+		if err != nil {
+			return nil, err
+		}
+		return 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 v, err
+		}
+		return v, nil
+	case simdjson.TypeBool:
+		v, err := iter.Bool()
+		if err != nil {
+			return nil, err
+		}
+		return v, nil
+	case simdjson.TypeObject:
+		obj, err := iter.Object(nil)
+		if err != nil {
+			return nil, err
+		}
+		return *obj, err
+	case simdjson.TypeArray:
+		arr, err := iter.Array(nil)
+		if err != nil {
+			return nil, err
+		}
+		iter := arr.Iter()
+		var dst []interface{}
+		var next simdjson.Iter
+		for {
+			typ, err := iter.AdvanceIter(&next)
+			if err != nil {
+				return nil, err
+			}
+			if typ == simdjson.TypeNone {
+				break
+			}
+			v, err := IterToValue(next)
+			if err != nil {
+				return nil, err
+			}
+			dst = append(dst, v)
+		}
+		return dst, err
+	case simdjson.TypeNull:
+		return nil, nil
+	}
+	return nil, fmt.Errorf("IterToValue: unknown JSON type: %s", iter.Type().String())
 }