refactor(query engine): json as bytes

api/seqproxyapi/v1/seq_proxy_api.proto

@@ -491,7 +491,7 @@ message Typing {
 enum DataType {
   DATA_TYPE_BYTES = 0;
   DATA_TYPE_SEQ_ID = 1;
-  DATA_TYPE_RAW_DOCUMENT = 2;
+  DATA_TYPE_JSON = 2;
   DATA_TYPE_STRING = 3;
   DATA_TYPE_UINT32 = 4;
   DATA_TYPE_UINT64 = 5;

api/storeapi/store_api.proto

@@ -321,7 +321,7 @@ message Typing {
 enum DataType {
   DATA_TYPE_BYTES = 0;
   DATA_TYPE_SEQ_ID = 1;
-  DATA_TYPE_RAW_DOCUMENT = 2;
+  DATA_TYPE_JSON = 2;
   DATA_TYPE_STRING = 3;
   DATA_TYPE_UINT32 = 4;
   DATA_TYPE_UINT64 = 5;

go.mod

@@ -12,6 +12,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/credentials v1.20.4
 	github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.19.13
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.113.1
+	github.com/buger/jsonparser v1.6.1
 	github.com/c2h5oh/datasize v0.0.0-20200112174442-28bbd4740fee
 	github.com/cep21/circuit/v3 v3.2.2
 	github.com/felixge/fgprof v0.9.5

go.sum

@@ -75,6 +75,8 @@ github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
 github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
 github.com/bits-and-blooms/bitset v1.24.2 h1:M7/NzVbsytmtfHbumG+K2bremQPMJuqv1JD3vOaFxp0=
 github.com/bits-and-blooms/bitset v1.24.2/go.mod h1:7hO7Gc7Pp1vODcmWvKMRA9BNmbv6a/7QIWpPxHddWR8=
+github.com/buger/jsonparser v1.6.1 h1:I0phFv0PlbLHnM7TZAVjZ2MJ2/eWRTDyuO7GLR98IEs=
+github.com/buger/jsonparser v1.6.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
 github.com/c2h5oh/datasize v0.0.0-20200112174442-28bbd4740fee h1:BnPxIde0gjtTnc9Er7cxvBk8DHLWhEux0SxayC8dP6I=
 github.com/c2h5oh/datasize v0.0.0-20200112174442-28bbd4740fee/go.mod h1:S/7n9copUssQ56c7aAgHqftWO4LTf4xY6CGWt8Bc+3M=
 github.com/cactus/go-statsd-client v3.1.1+incompatible/go.mod h1:cMRcwZDklk7hXp+Law83urTHUiHMzCev/r4JMYr/zU0=

pkg/seqproxyapi/v1/mappings.go

@@ -133,7 +133,7 @@ func AsyncSearchStatusFromString(s string) (AsyncSearchStatus, error) {
 var typeMappings = []DataType{
 	query.DataTypeBytes:        DataType_DATA_TYPE_BYTES,
 	query.DataTypeSeqID:        DataType_DATA_TYPE_SEQ_ID,
-	query.DataTypeDocument:     DataType_DATA_TYPE_RAW_DOCUMENT,
+	query.DataTypeJSON:         DataType_DATA_TYPE_JSON,
 	query.DataTypeString:       DataType_DATA_TYPE_STRING,
 	query.DataTypeUint32:       DataType_DATA_TYPE_UINT32,
 	query.DataTypeUint64:       DataType_DATA_TYPE_UINT64,

pkg/seqproxyapi/v1/stream_render.go

@@ -135,7 +135,7 @@ func (t *Typing) renderCell(cell []byte) any {
 		return string(cell)
 	case DataType_DATA_TYPE_BYTES:
 		return cell
-	case DataType_DATA_TYPE_RAW_DOCUMENT:
+	case DataType_DATA_TYPE_JSON:
 		return json.RawMessage(cell)
 	case DataType_DATA_TYPE_UINT32:
 		return encoding.Uint32FromBytes(cell)

pkg/seqproxyapi/v1/stream_render_test.go

@@ -31,7 +31,7 @@ func TestRowSchemaRender(t *testing.T) {
 			[]*Typing{
 				{Title: "id", Type: DataType_DATA_TYPE_SEQ_ID},
 				{Title: "time", Type: DataType_DATA_TYPE_UINT64},
-				{Title: "data", Type: DataType_DATA_TYPE_RAW_DOCUMENT},
+				{Title: "data", Type: DataType_DATA_TYPE_JSON},
 			},
 			[][]byte{
 				[]byte("46e48be997010000-e70163d0fa7582e4"),

pkg/storeapi/mappings.go

@@ -147,7 +147,7 @@ func MustProtoAsyncSearchStatus(s asyncsearcher.AsyncSearchStatus) AsyncSearchSt
 var typeMappings = []DataType{
 	query.DataTypeBytes:        DataType_DATA_TYPE_BYTES,
 	query.DataTypeSeqID:        DataType_DATA_TYPE_SEQ_ID,
-	query.DataTypeDocument:     DataType_DATA_TYPE_RAW_DOCUMENT,
+	query.DataTypeJSON:         DataType_DATA_TYPE_JSON,
 	query.DataTypeString:       DataType_DATA_TYPE_STRING,
 	query.DataTypeUint32:       DataType_DATA_TYPE_UINT32,
 	query.DataTypeUint64:       DataType_DATA_TYPE_UINT64,

proxy/search/stream_search_test.go

@@ -7,7 +7,6 @@ import (
 	"sync"
 	"testing"
 
-	insaneJSON "github.com/ozontech/insane-json"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc/metadata"
@@ -189,7 +188,7 @@ func TestNewStreamSearchIteratorEOFOnPrefetch(t *testing.T) {
 func testTyping() []*storeapi.Typing {
 	return []*storeapi.Typing{
 		{Title: "id", Type: storeapi.DataType_DATA_TYPE_SEQ_ID},
-		{Title: "data", Type: storeapi.DataType_DATA_TYPE_RAW_DOCUMENT},
+		{Title: "data", Type: storeapi.DataType_DATA_TYPE_JSON},
 	}
 }
 
@@ -314,15 +313,15 @@ func (f *fakeStream) RecvMsg(any) error            { return nil }
 func TestSchemaFromTyping(t *testing.T) {
 	typing := []*storeapi.Typing{
 		{Title: "id", Type: storeapi.DataType_DATA_TYPE_SEQ_ID},
-		{Title: "data", Type: storeapi.DataType_DATA_TYPE_RAW_DOCUMENT},
+		{Title: "data", Type: storeapi.DataType_DATA_TYPE_JSON},
 	}
 
 	schema, types, err := schemaFromTyping(typing)
 	require.NoError(t, err)
 	assert.Equal(t, 2, schema.Len())
-	assert.Equal(t, []query.DataType{query.DataTypeSeqID, query.DataTypeDocument}, types)
+	assert.Equal(t, []query.DataType{query.DataTypeSeqID, query.DataTypeJSON}, types)
 	assert.Equal(t, 0, schema.MustColumn[seq.ID]("id").Idx())
-	assert.Equal(t, query.DataTypeDocument, schema.MustColumn[*insaneJSON.Root]("data").DataType())
+	assert.Equal(t, query.DataTypeJSON, schema.MustColumn[[]byte]("data").DataType())
 }
 
 func TestSchemaFromTypingDuplicateName(t *testing.T) {
@@ -343,11 +342,11 @@ func TestValidateShardSchemas(t *testing.T) {
 
 	docsTyping := []*storeapi.Typing{
 		{Title: "id", Type: storeapi.DataType_DATA_TYPE_SEQ_ID},
-		{Title: "data", Type: storeapi.DataType_DATA_TYPE_RAW_DOCUMENT},
+		{Title: "data", Type: storeapi.DataType_DATA_TYPE_JSON},
 	}
 	otherTyping := []*storeapi.Typing{
 		{Title: "id", Type: storeapi.DataType_DATA_TYPE_SEQ_ID},
-		{Title: "payload", Type: storeapi.DataType_DATA_TYPE_RAW_DOCUMENT},
+		{Title: "payload", Type: storeapi.DataType_DATA_TYPE_JSON},
 	}
 	aggsTyping := []*storeapi.Typing{
 		{Title: "token", Type: storeapi.DataType_DATA_TYPE_STRING},

proxyapi/grpc_stream_search.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"io"
 
-	insaneJSON "github.com/ozontech/insane-json"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
@@ -304,7 +303,7 @@ func docsTyping() []*seqproxyapi.Typing {
 	return []*seqproxyapi.Typing{
 		{Title: "id", Type: seqproxyapi.DataType_DATA_TYPE_SEQ_ID},
 		{Title: "time", Type: seqproxyapi.DataType_DATA_TYPE_UINT64}, // extract time from id
-		{Title: "data", Type: seqproxyapi.DataType_DATA_TYPE_RAW_DOCUMENT},
+		{Title: "data", Type: seqproxyapi.DataType_DATA_TYPE_JSON},
 	}
 }
 
@@ -321,7 +320,7 @@ func schemaToTyping(out *query.Schema) []*seqproxyapi.Typing {
 
 var (
 	docIDCol   = query.DocsSchema.MustColumn[seq.ID](query.DocsIDCol)
-	docDataCol = query.DocsSchema.MustColumn[*insaneJSON.Root](query.DocsDataCol)
+	docDataCol = query.DocsSchema.MustColumn[[]byte](query.DocsDataCol)
 )
 
 // converts *query.Record to *seqproxyapi.Record according to hardcoded schemas from both store and proxy

query/column.go

@@ -3,8 +3,6 @@ package query
 import (
 	"cmp"
 
-	insaneJSON "github.com/ozontech/insane-json"
-
 	"github.com/ozontech/seq-db/seq"
 )
 
@@ -74,11 +72,11 @@ func StringColumn(idx int) Column[string] {
 	}
 }
 
-func DocColumn(idx int) Column[*insaneJSON.Root] {
-	return Column[*insaneJSON.Root]{
+func JSONColumn(idx int) Column[[]byte] {
+	return Column[[]byte]{
 		idx:      idx,
-		dataType: DataTypeDocument,
-		get:      (*RecordVals).AsDoc,
+		dataType: DataTypeJSON,
+		get:      (*RecordVals).AsJSON,
 	}
 }
 

query/exec/docs_datasource.go

@@ -165,7 +165,7 @@ func makeDocumentRecord(id seq.ID, payload []byte) *query.Record {
 	return &query.Record{
 		Vals: []*query.RecordVals{
 			query.NewSeqIDVals(id),
-			query.NewDocVals(payload),
+			query.NewJSONVals(payload),
 		},
 	}
 }

query/exec/extractor.go

@@ -1,56 +1,80 @@
 package exec
 
 import (
-	insaneJSON "github.com/ozontech/insane-json"
+	"slices"
+
+	"github.com/buger/jsonparser"
 
 	"github.com/ozontech/seq-db/query"
 )
 
-type DocFieldsExtractor struct {
-	input  query.RecordProducer
-	docCol query.Column[*insaneJSON.Root]
+var nullFieldValue = []byte("null")
 
-	// extractFields lists scalar JSON fields to extract out of the document column as string values.
-	extractFields []string
+type JSONFieldsExtractor struct {
+	input query.RecordProducer
+	col   query.Column[[]byte]
 
-	// roots holds every record whose val has been decoded (spawned an insaneJSON root).
-	// They are released back to the library pool in Finalize.
-	roots []*query.Record
+	extractFields []string
+	paths         [][]string
 }
 
-func NewDocFieldsExtractor(
+func NewJSONFieldsExtractor(
 	input query.RecordProducer,
-	docCol query.Column[*insaneJSON.Root],
+	col query.Column[[]byte],
 	extractFields ...string,
-) *DocFieldsExtractor {
-	return &DocFieldsExtractor{
+) *JSONFieldsExtractor {
+	paths := make([][]string, len(extractFields))
+	for i, f := range extractFields {
+		paths[i] = []string{f}
+	}
+	return &JSONFieldsExtractor{
 		input:         input,
-		docCol:        docCol,
+		col:           col,
 		extractFields: extractFields,
+		paths:         paths,
 	}
 }
 
-func (e *DocFieldsExtractor) Next() *query.Record {
+func (e *JSONFieldsExtractor) Next() *query.Record {
 	r := e.input.Next()
 	if r == nil {
 		return nil
 	}
 
-	root := e.docCol.Val(r)
-	for _, field := range e.extractFields {
-		// now we treat everything as strings, but later will support other types as well
-		fieldVal := root.Dig(field).AsBytes()
-		r.Vals = append(r.Vals, query.NewRecordVals(query.DataTypeString, fieldVal))
+	doc := e.col.Val(r)
+	base := len(r.Vals)
+	r.Vals = slices.Grow(r.Vals, len(e.extractFields))
+
+	// pre fill vals in case of not found paths
+	for range e.extractFields {
+		// now treat everything as strings, but later will support other types as well
+		r.Vals = append(r.Vals, query.NewRecordVals(query.DataTypeString, nil))
 	}
-	e.roots = append(e.roots, r)
+
+	jsonparser.EachKey(doc, func(i int, value []byte, dataType jsonparser.ValueType, _ error) {
+		if i >= 0 {
+			r.Vals[base+i] = query.NewRecordVals(query.DataTypeString, jsonFieldValue(value, dataType))
+		}
+	}, e.paths...)
 
 	return r
 }
 
-func (e *DocFieldsExtractor) Finalize() *query.Summary {
-	for _, r := range e.roots {
-		r.Release()
-	}
-	e.roots = nil
+func (e *JSONFieldsExtractor) Finalize() *query.Summary {
 	return e.input.Finalize()
 }
+
+func jsonFieldValue(value []byte, dataType jsonparser.ValueType) []byte {
+	switch dataType {
+	case jsonparser.String:
+		unescaped, err := jsonparser.Unescape(value, nil)
+		if err != nil {
+			return nil
+		}
+		return unescaped
+	case jsonparser.Null:
+		return nullFieldValue
+	default:
+		return value
+	}
+}