Refactor s3select to support parquet. (#7023)

Also handle pretty formatted JSON documents.
2026-07-20 12:40:24 +03:00 · 2019-01-09 06:23:04 +05:30
parent e98d89274f
commit b0deea27df
124 changed files with 27376 additions and 4152 deletions
@@ -0,0 +1,95 @@
+/*
+ * 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 json
+
+import (
+	"encoding/xml"
+	"fmt"
+	"strings"
+)
+
+const (
+	document = "document"
+	lines    = "lines"
+
+	defaultRecordDelimiter = "\n"
+)
+
+// ReaderArgs - represents elements inside <InputSerialization><JSON/> in request XML.
+type ReaderArgs struct {
+	ContentType string `xml:"Type"`
+	unmarshaled bool
+}
+
+// IsEmpty - returns whether reader args is empty or not.
+func (args *ReaderArgs) IsEmpty() bool {
+	return !args.unmarshaled
+}
+
+// UnmarshalXML - decodes XML data.
+func (args *ReaderArgs) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {
+	// Make subtype to avoid recursive UnmarshalXML().
+	type subReaderArgs ReaderArgs
+	parsedArgs := subReaderArgs{}
+	if err := d.DecodeElement(&parsedArgs, &start); err != nil {
+		return err
+	}
+
+	parsedArgs.ContentType = strings.ToLower(parsedArgs.ContentType)
+	switch parsedArgs.ContentType {
+	case document, lines:
+	default:
+		return errInvalidJSONType(fmt.Errorf("invalid ContentType '%v'", parsedArgs.ContentType))
+	}
+
+	*args = ReaderArgs(parsedArgs)
+	args.unmarshaled = true
+	return nil
+}
+
+// WriterArgs - represents elements inside <OutputSerialization><JSON/> in request XML.
+type WriterArgs struct {
+	RecordDelimiter string `xml:"RecordDelimiter"`
+	unmarshaled     bool
+}
+
+// IsEmpty - returns whether writer args is empty or not.
+func (args *WriterArgs) IsEmpty() bool {
+	return !args.unmarshaled
+}
+
+// UnmarshalXML - decodes XML data.
+func (args *WriterArgs) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {
+	// Make subtype to avoid recursive UnmarshalXML().
+	type subWriterArgs WriterArgs
+	parsedArgs := subWriterArgs{}
+	if err := d.DecodeElement(&parsedArgs, &start); err != nil {
+		return err
+	}
+
+	switch len(parsedArgs.RecordDelimiter) {
+	case 0:
+		parsedArgs.RecordDelimiter = defaultRecordDelimiter
+	case 1, 2:
+	default:
+		return fmt.Errorf("invalid RecordDelimiter '%v'", parsedArgs.RecordDelimiter)
+	}
+
+	*args = WriterArgs(parsedArgs)
+	args.unmarshaled = true
+	return nil
+}
@@ -0,0 +1,62 @@
+/*
+ * 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 json
+
+type s3Error struct {
+	code       string
+	message    string
+	statusCode int
+	cause      error
+}
+
+func (err *s3Error) Cause() error {
+	return err.cause
+}
+
+func (err *s3Error) ErrorCode() string {
+	return err.code
+}
+
+func (err *s3Error) ErrorMessage() string {
+	return err.message
+}
+
+func (err *s3Error) HTTPStatusCode() int {
+	return err.statusCode
+}
+
+func (err *s3Error) Error() string {
+	return err.message
+}
+
+func errInvalidJSONType(err error) *s3Error {
+	return &s3Error{
+		code:       "InvalidJsonType",
+		message:    "The JsonType is invalid. Only DOCUMENT and LINES are supported.",
+		statusCode: 400,
+		cause:      err,
+	}
+}
+
+func errJSONParsingError(err error) *s3Error {
+	return &s3Error{
+		code:       "JSONParsingError",
+		message:    "Encountered an error parsing the JSON file. Check the file and try again.",
+		statusCode: 400,
+		cause:      err,
+	}
+}
@@ -0,0 +1,217 @@
+/*
+ * 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 json
+
+import (
+	"bytes"
+	"io"
+	"io/ioutil"
+	"strconv"
+
+	"github.com/minio/minio/pkg/s3select/sql"
+	"github.com/tidwall/gjson"
+	"github.com/tidwall/sjson"
+)
+
+func toSingleLineJSON(input string, currentKey string, result gjson.Result) (output string, err error) {
+	switch {
+	case result.IsObject():
+		result.ForEach(func(key, value gjson.Result) bool {
+			jsonKey := key.String()
+			if currentKey != "" {
+				jsonKey = currentKey + "." + key.String()
+			}
+			output, err = toSingleLineJSON(input, jsonKey, value)
+			input = output
+			return err == nil
+		})
+	case result.IsArray():
+		i := 0
+		result.ForEach(func(key, value gjson.Result) bool {
+			if currentKey == "" {
+				panic("currentKey is empty")
+			}
+
+			indexKey := currentKey + "." + strconv.Itoa(i)
+			output, err = toSingleLineJSON(input, indexKey, value)
+			input = output
+			i++
+			return err == nil
+		})
+	default:
+		output, err = sjson.Set(input, currentKey, result.Value())
+	}
+
+	return output, err
+}
+
+type objectReader struct {
+	reader io.Reader
+	err    error
+
+	p     []byte
+	start int
+	end   int
+
+	escaped     bool
+	quoteOpened bool
+	curlyCount  uint64
+	endOfObject bool
+}
+
+func (or *objectReader) objectEndIndex(p []byte, length int) int {
+	for i := 0; i < length; i++ {
+		if p[i] == '\\' {
+			or.escaped = !or.escaped
+			continue
+		}
+
+		if p[i] == '"' && !or.escaped {
+			or.quoteOpened = !or.quoteOpened
+		}
+
+		or.escaped = false
+
+		switch p[i] {
+		case '{':
+			if !or.quoteOpened {
+				or.curlyCount++
+			}
+		case '}':
+			if or.quoteOpened || or.curlyCount == 0 {
+				break
+			}
+
+			if or.curlyCount--; or.curlyCount == 0 {
+				return i + 1
+			}
+		}
+	}
+
+	return -1
+}
+
+func (or *objectReader) Read(p []byte) (n int, err error) {
+	if or.endOfObject {
+		return 0, io.EOF
+	}
+
+	if or.p != nil {
+		n = copy(p, or.p[or.start:or.end])
+		or.start += n
+		if or.start == or.end {
+			// made full copy.
+			or.p = nil
+			or.start = 0
+			or.end = 0
+		}
+	} else {
+		if or.err != nil {
+			return 0, or.err
+		}
+
+		n, err = or.reader.Read(p)
+		or.err = err
+		switch err {
+		case nil:
+		case io.EOF, io.ErrUnexpectedEOF, io.ErrClosedPipe:
+			or.err = io.EOF
+		default:
+			return 0, err
+		}
+	}
+
+	index := or.objectEndIndex(p, n)
+	if index == -1 || index == n {
+		return n, nil
+	}
+
+	or.endOfObject = true
+	if or.p == nil {
+		or.p = p
+		or.start = index
+		or.end = n
+	} else {
+		or.start -= index
+	}
+
+	return index, nil
+}
+
+func (or *objectReader) Reset() error {
+	or.endOfObject = false
+
+	if or.p != nil {
+		return nil
+	}
+
+	return or.err
+}
+
+// Reader - JSON record reader for S3Select.
+type Reader struct {
+	args         *ReaderArgs
+	objectReader *objectReader
+	readCloser   io.ReadCloser
+}
+
+// Read - reads single record.
+func (r *Reader) Read() (sql.Record, error) {
+	if err := r.objectReader.Reset(); err != nil {
+		return nil, err
+	}
+
+	data, err := ioutil.ReadAll(r.objectReader)
+	if err != nil {
+		return nil, errJSONParsingError(err)
+	}
+
+	data = bytes.TrimSpace(data)
+	if len(data) == 0 {
+		return nil, io.EOF
+	}
+
+	if !gjson.ValidBytes(data) {
+		return nil, errJSONParsingError(err)
+	}
+
+	if bytes.Count(data, []byte("\n")) > 0 {
+		var s string
+		if s, err = toSingleLineJSON("", "", gjson.ParseBytes(data)); err != nil {
+			return nil, errJSONParsingError(err)
+		}
+		data = []byte(s)
+	}
+
+	return &Record{
+		data: data,
+	}, nil
+}
+
+// Close - closes underlaying reader.
+func (r *Reader) Close() error {
+	return r.readCloser.Close()
+}
+
+// NewReader - creates new JSON reader using readCloser.
+func NewReader(readCloser io.ReadCloser, args *ReaderArgs) *Reader {
+	return &Reader{
+		args:         args,
+		objectReader: &objectReader{reader: readCloser},
+		readCloser:   readCloser,
+	}
+}
@@ -0,0 +1,107 @@
+/*
+ * 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 json
+
+import (
+	"bytes"
+	"encoding/csv"
+	"fmt"
+	"strings"
+
+	"github.com/minio/minio/pkg/s3select/sql"
+	"github.com/tidwall/gjson"
+	"github.com/tidwall/sjson"
+)
+
+// Record - is JSON record.
+type Record struct {
+	data []byte
+}
+
+// Get - gets the value for a column name.
+func (r *Record) Get(name string) (*sql.Value, error) {
+	result := gjson.GetBytes(r.data, name)
+	switch result.Type {
+	case gjson.False:
+		return sql.NewBool(false), nil
+	case gjson.Number:
+		return sql.NewFloat(result.Float()), nil
+	case gjson.String:
+		return sql.NewString(result.String()), nil
+	case gjson.True:
+		return sql.NewBool(true), nil
+	}
+
+	return nil, fmt.Errorf("unsupported gjson value %v; %v", result, result.Type)
+}
+
+// Set - sets the value for a column name.
+func (r *Record) Set(name string, value *sql.Value) (err error) {
+	var v interface{}
+	switch value.Type() {
+	case sql.Bool:
+		v = value.BoolValue()
+	case sql.Int:
+		v = value.IntValue()
+	case sql.Float:
+		v = value.FloatValue()
+	case sql.String:
+		v = value.StringValue()
+	default:
+		return fmt.Errorf("unsupported sql value %v and type %v", value, value.Type())
+	}
+
+	name = strings.Replace(name, "*", "__ALL__", -1)
+	r.data, err = sjson.SetBytes(r.data, name, v)
+	return err
+}
+
+// MarshalCSV - encodes to CSV data.
+func (r *Record) MarshalCSV(fieldDelimiter rune) ([]byte, error) {
+	var csvRecord []string
+	result := gjson.ParseBytes(r.data)
+	result.ForEach(func(key, value gjson.Result) bool {
+		csvRecord = append(csvRecord, value.String())
+		return true
+	})
+
+	buf := new(bytes.Buffer)
+	w := csv.NewWriter(buf)
+	w.Comma = fieldDelimiter
+	if err := w.Write(csvRecord); err != nil {
+		return nil, err
+	}
+	w.Flush()
+	if err := w.Error(); err != nil {
+		return nil, err
+	}
+
+	data := buf.Bytes()
+	return data[:len(data)-1], nil
+}
+
+// MarshalJSON - encodes to JSON data.
+func (r *Record) MarshalJSON() ([]byte, error) {
+	return r.data, nil
+}
+
+// NewRecord - creates new empty JSON record.
+func NewRecord() *Record {
+	return &Record{
+		data: []byte("{}"),
+	}
+}