Expand nest field to parquet path

Author: monsterxx03

engine/engine.go

@@ -96,7 +96,12 @@ func (p *ParquetEngine) executeSelect(stmt *parser.Select) (*RecordSet, error) {
 		case *parser.ColExpr:
 			cols = append(cols, field.(*parser.ColExpr).Name)
 		case *parser.NestColExpr:
-			fmt.Println(field.(*parser.NestColExpr).Subs)
+			subs :=field.(*parser.NestColExpr).Subs
+			_p, err := p.schema.GetFieldPath(subs)
+			if err != nil {
+				return nil, err
+			}
+			cols = append(cols, _p)
 		}
 	}
 	filterCols := make([]string, 0)
@@ -111,7 +116,7 @@ func (p *ParquetEngine) executeSelect(stmt *parser.Select) (*RecordSet, error) {
 	if stmt.Limit != nil {
 		limit = stmt.Limit.Rowcount
 	}
-	cr, err := p.GetColumnReader()
+	cr, err := p.GetReader()
 	if err != nil {
 		return nil, err
 	}
@@ -162,15 +167,15 @@ func filter(expr parser.Expr, result *RecordSet) (*RecordSet, error) {
 }
 
 func (p *ParquetEngine) GetTotalRowCount() (int64, error) {
-	cr, err := p.GetColumnReader()
+	cr, err := p.GetReader()
 	if err != nil {
 		return 0, err
 	}
 	return cr.GetNumRows(), nil
 }
 
-func (p *ParquetEngine) GetColumnReader() (*reader.ParquetReader, error) {
-	return reader.NewParquetColumnReader(p.fr, 2)
+func (p *ParquetEngine) GetReader() (*reader.ParquetReader, error) {
+	return reader.NewParquetReader(p.fr,nil, 2)
 }
 
 func (p *ParquetEngine) FetchRows(cr *reader.ParquetReader, cols []string, limit int) (result *RecordSet, err error) {

engine/schema.go

@@ -3,9 +3,16 @@ package engine
 import (
 	"fmt"
 	"github.com/xitongsys/parquet-go/parquet"
+	"strconv"
 	"strings"
 )
 
+const (
+	TYPE_STRUCT = "STRUCT"
+	TYPE_LIST   = "LIST"
+	TYPE_MAP    = "MAP"
+)
+
 type Schema interface {
 }
 
@@ -18,32 +25,21 @@ type Field struct {
 }
 
 func (f *Field) String() string {
-	if len(f.Fields) == 0 {
-		if f.SE.ConvertedType != nil {
-			return f.SE.ConvertedType.String()
-		}
-		if f.SE.Type != nil {
-			return f.SE.Type.String()
-		}
-	}
-	if f.SE.ConvertedType == nil {
+	switch f.Type {
+	case TYPE_LIST:
+		// LIST.ELEMENT
+		return fmt.Sprintf("array<%s>", f.Fields[0].Fields[0])
+	case TYPE_MAP:
+		// map<KEY_VALUE.KEY, KEY_VALUE.VALUE>
+		return fmt.Sprintf("map<%s,%s>", f.Fields[0].Fields[0], f.Fields[0].Fields[1])
+	case TYPE_STRUCT:
 		fields := make([]string, 0, len(f.Fields))
 		for _, _f := range f.Fields {
 			fields = append(fields, fmt.Sprintf("%s:%s", _f.Name, _f))
 		}
 		return fmt.Sprintf("struct<%s>", strings.Join(fields, ","))
-	}
-	switch *f.SE.ConvertedType {
-	case parquet.ConvertedType_LIST:
-		// LIST.ELEMENT
-		return fmt.Sprintf("array<%s>", f.Fields[0].Fields[0])
-	case parquet.ConvertedType_MAP:
-		// map<KEY_VALUE.KEY, KEY_VALUE.VALUE>
-		return fmt.Sprintf("map<%s,%s>", f.Fields[0].Fields[0], f.Fields[0].Fields[1])
-	case parquet.ConvertedType_JSON:
-		return "json"
 	default:
-		return f.SE.ConvertedType.String()
+		return f.Type
 	}
 }
 
@@ -53,6 +49,47 @@ type ParquetSchema struct {
 	Fields  []*Field
 }
 
+func (s *ParquetSchema) GetFieldPath(subs []string) (string, error) {
+	paths := make([]string, 0)
+	f := s.MFields[subs[0]]
+	paths = append(paths, subs[0])
+	subs = subs[1:]
+	for len(subs) > 0 {
+		sub := subs[0]
+		idx, err := strconv.Atoi(sub)
+		if err != nil {
+			switch f.Type{
+			case TYPE_MAP:
+				paths = append(paths, []string{"Key_value", "Value"}...)
+				f = f.Fields[0].Fields[1]
+			case TYPE_STRUCT:
+				found := false
+				for _, _f := range f.Fields{
+					if _f.Name == sub {
+						found = true
+						paths = append(paths, _f.Name)
+						f = _f
+						break
+					}
+				}
+				if !found {
+					return "", fmt.Errorf("can't find %s on %s:%s", sub, f.Name, f)
+				}
+			default:
+				return "", fmt.Errorf("unsuport to get `%s` on %s:%s", sub, f.Name, f)
+			}
+		} else {
+			if f.Type != TYPE_LIST {
+				return "", fmt.Errorf("unsupport to retrive index `%d` on %s:%s", idx, f.Name, f)
+			}
+			paths = append(paths, "List", "Element")
+			f = f.Fields[0].Fields[0]
+		}
+		subs = subs[1:]
+	}
+	return strings.Join(paths, "."), nil
+}
+
 func (s *ParquetSchema) GetName() string {
 	return s.Name
 }
@@ -86,11 +123,33 @@ func NewParquetSchema(schemas []*parquet.SchemaElement) *ParquetSchema {
 	return &ParquetSchema{Name: root.Name, MFields: root.MFields, Fields: root.Fields}
 }
 
-func NewField(schema *parquet.SchemaElement) *Field {
+func NewField(se *parquet.SchemaElement) *Field {
+	TYPE := ""
+	if se.GetNumChildren() == 0 {
+		if se.ConvertedType != nil {
+			TYPE = se.ConvertedType.String()
+		} else if se.Type != nil {
+			TYPE = se.Type.String()
+		}
+	} else {
+		if se.ConvertedType == nil {
+			TYPE = TYPE_STRUCT
+		} else {
+			switch *se.ConvertedType {
+			case parquet.ConvertedType_MAP:
+				TYPE = TYPE_MAP
+			case parquet.ConvertedType_LIST:
+				TYPE = TYPE_LIST
+			default:
+				TYPE = se.ConvertedType.String()
+			}
+		}
+	}
 	return &Field{
-		Name:    schema.Name,
-		SE:      schema,
+		Name:    se.Name,
+		SE:      se,
+		Type:    TYPE,
 		MFields: make(map[string]*Field),
-		Fields:  make([]*Field, 0, schema.GetNumChildren()),
+		Fields:  make([]*Field, 0, se.GetNumChildren()),
 	}
 }