Repository: segmentio/parquet-go (go)

Issue:
Add ability to set properties on key and value of a Parquet map 
Currently when defining a Parquet map:

```
type A struct {
  B map[string]string `parquet:","`
}
```

We don't have the ability to set the properties of the key or value for that map.

One possibility is to add something similar to the following:

```
type A struct {
  B map[string]string `parquet:",key(dict,zstd)"`
}
```

Pull request description:
Add the ability to pass Parquet tags on map key and values
Fixes #256 

Patch:
diff --git a/column_buffer.go b/column_buffer.go
index 9981831..f35932a 100644
--- a/column_buffer.go
+++ b/column_buffer.go
@@ -722,10 +722,12 @@ func (col *repeatedColumnBuffer) ReadValuesAt(values []Value, offset int64) (int
 type booleanColumnBuffer struct{ booleanPage }
 
 func newBooleanColumnBuffer(typ Type, columnIndex int16, numValues int32) *booleanColumnBuffer {
+	// Boolean values are bit-packed, we can fit up to 8 values per byte.
+	bufferSize := (numValues + 7) / 8
 	return &booleanColumnBuffer{
 		booleanPage: booleanPage{
 			typ:         typ,
-			bits:        make([]byte, 0, numValues),
+			bits:        make([]byte, 0, bufferSize),
 			columnIndex: ^columnIndex,
 		},
 	}
diff --git a/column_path.go b/column_path.go
index 06f154a..681d135 100644
--- a/column_path.go
+++ b/column_path.go
@@ -1,6 +1,8 @@
 package parquet
 
-import "strings"
+import (
+	"strings"
+)
 
 type columnPath []string
 
diff --git a/schema.go b/schema.go
index bf87605..6026317 100644
--- a/schema.go
+++ b/schema.go
@@ -82,6 +82,15 @@ type Schema struct {
 // and the data will not be written into the parquet file(s).
 // Note that a field with name "-" can still be generated using the tag "-,".
 //
+// Map key configuration can be pass using the `parquet-key` tag. This tag will follow the
+// same properties than the `parquet` tag.
+//
+// For example, the following will set the int64 key of the map to be a timestamp:
+//
+//  type Actions struct {
+//    Action map[int64]string `parquet:"," parquet-key:",timestamp"`
+//  }
+//
 // The schema name is the Go type name of the value.
 func SchemaOf(model interface{}) *Schema {
 	return schemaOf(dereference(reflect.TypeOf(model)))
@@ -98,7 +107,7 @@ func schemaOf(model reflect.Type) *Schema {
 	if model.Kind() != reflect.Struct {
 		panic("cannot construct parquet schema from value of type " + model.String())
 	}
-	schema = NewSchema(model.Name(), nodeOf(model))
+	schema = NewSchema(model.Name(), nodeOf(model, nil))
 	if actual, loaded := cachedSchemas.LoadOrStore(model, schema); loaded {
 		schema = actual.(*Schema)
 	}
@@ -302,7 +311,12 @@ func structNodeOf(t reflect.Type) *structNode {
 	}
 
 	for i := range fields {
-		s.fields[i] = makeStructField(fields[i])
+		field := structField{name: fields[i].Name, index: fields[i].Index}
+		field.Node = makeNodeOf(fields[i].Type, fields[i].Name, []string{
+			fields[i].Tag.Get("parquet"),
+			fields[i].Tag.Get("parquet-key"),
+		})
+		s.fields[i] = field
 	}
 
 	return s
@@ -418,68 +432,249 @@ func (f *structField) Value(base reflect.Value) reflect.Value {
 	}
 }
 
-func structFieldString(f reflect.StructField) string {
-	return f.Name + " " + f.Type.String() + " " + string(f.Tag)
+func nodeString(t reflect.Type, name string, tag ...string) string {
+	return fmt.Sprintf("%s %s %v", name, t.String(), tag)
+}
+
+func throwInvalidTag(t reflect.Type, name string, tag string) {
+	panic(tag + " is an invalid parquet tag: " + nodeString(t, name, tag))
+}
+
+func throwUnknownTag(t reflect.Type, name string, tag string) {
+	panic(tag + " is an unrecognized parquet tag: " + nodeString(t, name, tag))
+}
+
+func throwInvalidNode(t reflect.Type, msg, name string, tag ...string) {
+	panic(msg + ": " + nodeString(t, name, tag...))
+}
+
+// FixedLenByteArray decimals are sized based on precision
+// this function calculates the necessary byte array size.
+func decimalFixedLenByteArraySize(precision int) int {
+	return int(math.Ceil((math.Log10(2) + float64(precision)) / math.Log10(256)))
+}
+
+func forEachStructTagOption(sf reflect.StructField, do func(t reflect.Type, option, args string)) {
+	if tag := sf.Tag.Get("parquet"); tag != "" {
+		_, tag = split(tag) // skip the field name
+		for tag != "" {
+			option := ""
+			args := ""
+			option, tag = split(tag)
+			option, args = splitOptionArgs(option)
+			ft := sf.Type
+			if ft.Kind() == reflect.Ptr {
+				ft = ft.Elem()
+			}
+			do(ft, option, args)
+		}
+	}
+}
+
+func nodeOf(t reflect.Type, tag []string) Node {
+	switch t {
+	case reflect.TypeOf(deprecated.Int96{}):
+		return Leaf(Int96Type)
+	case reflect.TypeOf(uuid.UUID{}):
+		return UUID()
+	}
+
+	var n Node
+	switch t.Kind() {
+	case reflect.Bool:
+		n = Leaf(BooleanType)
+
+	case reflect.Int, reflect.Int64:
+		n = Int(64)
+
+	case reflect.Int8, reflect.Int16, reflect.Int32:
+		n = Int(t.Bits())
+
+	case reflect.Uint, reflect.Uintptr, reflect.Uint64:
+		n = Uint(64)
+
+	case reflect.Uint8, reflect.Uint16, reflect.Uint32:
+		n = Uint(t.Bits())
+
+	case reflect.Float32:
+		n = Leaf(FloatType)
+
+	case reflect.Float64:
+		n = Leaf(DoubleType)
+
+	case reflect.String:
+		n = String()
+
+	case reflect.Ptr:
+		n = Optional(nodeOf(t.Elem(), nil))
+
+	case reflect.Slice:
+		if elem := t.Elem(); elem.Kind() == reflect.Uint8 { // []byte?
+			n = Leaf(ByteArrayType)
+		} else {
+			n = Repeated(nodeOf(elem, nil))
+		}
+
+	case reflect.Array:
+		if t.Elem().Kind() == reflect.Uint8 {
+			n = Leaf(FixedLenByteArrayType(t.Len()))
+		}
+
+	case reflect.Map:
+		var valueTag, keyTag string
+		if len(tag) > 0 {
+			valueTag = tag[0]
+			if len(tag) > 1 {
+				keyTag = tag[1]
+			}
+		}
+		n = Map(
+			makeNodeOf(t.Key(), t.Name(), []string{keyTag}),
+			makeNodeOf(t.Elem(), t.Name(), []string{valueTag}),
+		)
+
+	case reflect.Struct:
+		return structNodeOf(t)
+	}
+
+	if n == nil {
+		panic("cannot create parquet node from go value of type " + t.String())
+	}
+
+	return &goNode{Node: n, gotype: t}
+}
+
+func split(s string) (head, tail string) {
+	if i := strings.IndexByte(s, ','); i < 0 {
+		head = s
+	} else {
+		head, tail = s[:i], s[i+1:]
+	}
+	return
 }
 
-func throwInvalidFieldTag(f reflect.StructField, tag string) {
-	panic("struct has invalid '" + tag + "' parquet tag: " + structFieldString(f))
+func splitOptionArgs(s string) (option, args string) {
+	if i := strings.IndexByte(s, '('); i >= 0 {
+		option = s[:i]
+		args = s[i:]
+	} else {
+		option = s
+		args = "()"
+	}
+	return
 }
 
-func throwUnknownFieldTag(f reflect.StructField, tag string) {
-	panic("struct has unrecognized '" + tag + "' parquet tag: " + structFieldString(f))
+func parseDecimalArgs(args string) (scale, precision int, err error) {
+	if !strings.HasPrefix(args, "(") || !strings.HasSuffix(args, ")") {
+		return 0, 0, fmt.Errorf("malformed decimal args: %s", args)
+	}
+	args = strings.TrimPrefix(args, "(")
+	args = strings.TrimSuffix(args, ")")
+	parts := strings.Split(args, ":")
+	if len(parts) != 2 {
+		return 0, 0, fmt.Errorf("malformed decimal args: (%s)", args)
+	}
+	s, err := strconv.ParseInt(parts[0], 10, 32)
+	if err != nil {
+		return 0, 0, err
+	}
+	p, err := strconv.ParseInt(parts[1], 10, 32)
+	if err != nil {
+		return 0, 0, err
+	}
+	return int(s), int(p), nil
+}
+
+func parseTimestampArgs(args string) (TimeUnit, error) {
+	if !strings.HasPrefix(args, "(") || !strings.HasSuffix(args, ")") {
+		return nil, fmt.Errorf("malformed timestamp args: %s", args)
+	}
+
+	args = strings.TrimPrefix(args, "(")
+	args = strings.TrimSuffix(args, ")")
+
+	if len(args) == 0 {
+		return Millisecond, nil
+	}
+
+	switch args {
+	case "millisecond":
+		return Millisecond, nil
+	case "microsecond":
+		return Microsecond, nil
+	case "nanosecond":
+		return Nanosecond, nil
+	default:
+	}
+
+	return nil, fmt.Errorf("unknown time unit: %s", args)
 }
 
-func throwInvalidStructField(msg string, field reflect.StructField) {
-	panic(msg + ": " + structFieldString(field))
+type goNode struct {
+	Node
+	gotype reflect.Type
 }
 
-func makeStructField(f reflect.StructField) structField {
+func (n *goNode) GoType() reflect.Type { return n.gotype }
+
+var (
+	_ RowGroupOption = (*Schema)(nil)
+	_ ReaderOption   = (*Schema)(nil)
+	_ WriterOption   = (*Schema)(nil)
+)
+
+func makeNodeOf(t reflect.Type, name string, tag []string) Node {
 	var (
-		field      = structField{name: f.Name, index: f.Index}
+		node       Node
 		optional   bool
 		list       bool
 		encoded    encoding.Encoding
 		compressed compress.Codec
 	)
 
-	setNode := func(node Node) {
-		if field.Node != nil {
-			throwInvalidStructField("struct field has multiple logical parquet types declared", f)
+	setNode := func(n Node) {
+		if node != nil {
+			throwInvalidNode(t, "struct field has multiple logical parquet types declared", name, tag...)
 		}
-		field.Node = node
+		node = n
 	}
 
 	setOptional := func() {
 		if optional {
-			throwInvalidStructField("struct field has multiple declaration of the optional tag", f)
+			throwInvalidNode(t, "struct field has multiple declaration of the optional tag", name, tag...)
 		}
 		optional = true
 	}
 
 	setList := func() {
 		if list {
-			throwInvalidStructField("struct field has multiple declaration of the list tag", f)
+			throwInvalidNode(t, "struct field has multiple declaration of the list tag", name, tag...)
 		}
 		list = true
 	}
 
 	setEncoding := func(e encoding.Encoding) {
 		if encoded != nil {
-			throwInvalidStructField("struct field has encoding declared multiple times", f)
+			throwInvalidNode(t, "struct field has encoding declared multiple time", name, tag...)
 		}
 		encoded = e
 	}
 
 	setCompression := func(c compress.Codec) {
 		if compressed != nil {
-			throwInvalidStructField("struct field has compression codecs declared multiple times", f)
+			throwInvalidNode(t, "struct field has compression codecs declared multiple times", name, tag...)
 		}
 		compressed = c
 	}
 
-	forEachStructTagOption(f, func(t reflect.Type, option, args string) {
+	forEachTagOption(tag, func(option, args string) {
+		if t.Kind() == reflect.Map {
+			node = nodeOf(t, tag)
+			return
+		}
 		switch option {
+		case "":
+			return
 		case "optional":
 			setOptional()
 
@@ -520,16 +715,16 @@ func makeStructField(f reflect.StructField) structField {
 				if t.Elem().Kind() == reflect.Uint8 { // []byte?
 					setEncoding(&DeltaByteArray)
 				} else {
-					throwInvalidFieldTag(f, option)
+					throwInvalidTag(t, name, option)
 				}
 			case reflect.Array:
 				if t.Elem().Kind() == reflect.Uint8 { // [N]byte?
 					setEncoding(&DeltaByteArray)
 				} else {
-					throwInvalidFieldTag(f, option)
+					throwInvalidTag(t, name, option)
 				}
 			default:
-				throwInvalidFieldTag(f, option)
+				throwInvalidTag(t, name, option)
 			}
 
 		case "split":
@@ -537,17 +732,17 @@ func makeStructField(f reflect.StructField) structField {
 			case reflect.Float32, reflect.Float64:
 				setEncoding(&ByteStreamSplit)
 			default:
-				throwInvalidFieldTag(f, option)
+				throwInvalidTag(t, name, option)
 			}
 
 		case "list":
 			switch t.Kind() {
 			case reflect.Slice:
-				element := nodeOf(t.Elem())
+				element := nodeOf(t.Elem(), nil)
 				setNode(element)
 				setList()
 			default:
-				throwInvalidFieldTag(f, option)
+				throwInvalidTag(t, name, option)
 			}
 
 		case "enum":
@@ -555,23 +750,23 @@ func makeStructField(f reflect.StructField) structField {
 			case reflect.String:
 				setNode(Enum())
 			default:
-				throwInvalidFieldTag(f, option)
+				throwInvalidTag(t, name, option)
 			}
 
 		case "uuid":
 			switch t.Kind() {
 			case reflect.Array:
 				if t.Elem().Kind() != reflect.Uint8 || t.Len() != 16 {
-					throwInvalidFieldTag(f, option)
+					throwInvalidTag(t, name, option)
 				}
 			default:
-				throwInvalidFieldTag(f, option)
+				throwInvalidTag(t, name, option)
 			}
 
 		case "decimal":
 			scale, precision, err := parseDecimalArgs(args)
 			if err != nil {
-				throwInvalidFieldTag(f, option+args)
+				throwInvalidTag(t, name, option+args)
 			}
 			var baseType Type
 			switch t.Kind() {
@@ -582,7 +777,7 @@ func makeStructField(f reflect.StructField) structField {
 			case reflect.Array, reflect.Slice:
 				baseType = FixedLenByteArrayType(decimalFixedLenByteArraySize(precision))
 			default:
-				throwInvalidFieldTag(f, option)
+				throwInvalidTag(t, name, option)
 			}
 
 			setNode(Decimal(scale, precision, baseType))
@@ -591,204 +786,56 @@ func makeStructField(f reflect.StructField) structField {
 			case reflect.Int32:
 				setNode(Date())
 			default:
-				throwInvalidFieldTag(f, option)
+				throwInvalidTag(t, name, option)
 			}
 		case "timestamp":
 			switch t.Kind() {
 			case reflect.Int64:
 				timeUnit, err := parseTimestampArgs(args)
 				if err != nil {
-					throwInvalidFieldTag(f, args)
+					throwInvalidTag(t, name, option)
 				}
 				setNode(Timestamp(timeUnit))
 			default:
-				throwInvalidFieldTag(f, option)
+				throwInvalidTag(t, name, option)
 			}
 		default:
-			throwUnknownFieldTag(f, option)
+			throwUnknownTag(t, name, option)
 		}
 	})
 
-	if field.Node == nil {
-		field.Node = nodeOf(f.Type)
+	if node == nil {
+		node = nodeOf(t, tag)
 	}
 
 	if compressed != nil {
-		field.Node = Compressed(field.Node, compressed)
+		node = Compressed(node, compressed)
 	}
 
 	if encoded != nil {
-		field.Node = Encoded(field.Node, encoded)
+		node = Encoded(node, encoded)
 	}
 
 	if list {
-		field.Node = List(field.Node)
+		node = List(node)
 	}
 
 	if optional {
-		field.Node = Optional(field.Node)
+		node = Optional(node)
 	}
 
-	return field
-}
-
-// FixedLenByteArray decimals are sized based on precision
-// this function calculates the necessary byte array size.
-func decimalFixedLenByteArraySize(precision int) int {
-	return int(math.Ceil((math.Log10(2) + float64(precision)) / math.Log10(256)))
+	return node
 }
 
-func forEachStructTagOption(sf reflect.StructField, do func(t reflect.Type, option, args string)) {
-	if tag := sf.Tag.Get("parquet"); tag != "" {
+func forEachTagOption(tags []string, do func(option, args string)) {
+	for _, tag := range tags {
 		_, tag = split(tag) // skip the field name
 		for tag != "" {
 			option := ""
 			option, tag = split(tag)
-			option, args := splitOptionArgs(option)
-			ft := sf.Type
-			if ft.Kind() == reflect.Ptr {
-				ft = ft.Elem()
-			}
-			do(ft, option, args)
+			var args string
+			option, args = splitOptionArgs(option)
+			do(option, args)
 		}
 	}
 }
-
-func nodeOf(t reflect.Type) Node {
-	switch t {
-	case reflect.TypeOf(deprecated.Int96{}):
-		return Leaf(Int96Type)
-	case reflect.TypeOf(uuid.UUID{}):
-		return UUID()
-	}
-
-	var n Node
-	switch t.Kind() {
-	case reflect.Bool:
-		n = Leaf(BooleanType)
-
-	case reflect.Int, reflect.Int64:
-		n = Int(64)
-
-	case reflect.Int8, reflect.Int16, reflect.Int32:
-		n = Int(t.Bits())
-
-	case reflect.Uint, reflect.Uintptr, reflect.Uint64:
-		n = Uint(64)
-
-	case reflect.Uint8, reflect.Uint16, reflect.Uint32:
-		n = Uint(t.Bits())
-
-	case reflect.Float32:
-		n = Leaf(FloatType)
-
-	case reflect.Float64:
-		n = Leaf(DoubleType)
-
-	case reflect.String:
-		n = String()
-
-	case reflect.Ptr:
-		n = Optional(nodeOf(t.Elem()))
-
-	case reflect.Slice:
-		if elem := t.Elem(); elem.Kind() == reflect.Uint8 { // []byte?
-			n = Leaf(ByteArrayType)
-		} else {
-			n = Repeated(nodeOf(elem))
-		}
-
-	case reflect.Array:
-		if t.Elem().Kind() == reflect.Uint8 {
-			n = Leaf(FixedLenByteArrayType(t.Len()))
-		}
-
-	case reflect.Map:
-		n = Map(nodeOf(t.Key()), nodeOf(t.Elem()))
-
-	case reflect.Struct:
-		return structNodeOf(t)
-	}
-
-	if n == nil {
-		panic("cannot create parquet node from go value of type " + t.String())
-	}
-
-	return &goNode{Node: n, gotype: t}
-}
-
-func split(s string) (head, tail string) {
-	if i := strings.IndexByte(s, ','); i < 0 {
-		head = s
-	} else {
-		head, tail = s[:i], s[i+1:]
-	}
-	return
-}
-
-func splitOptionArgs(s string) (option, args string) {
-	if i := strings.IndexByte(s, '('); i >= 0 {
-		return s[:i], s[i:]
-	} else {
-		return s, "()"
-	}
-}
-
-func parseDecimalArgs(args string) (scale, precision int, err error) {
-	if !strings.HasPrefix(args, "(") || !strings.HasSuffix(args, ")") {
-		return 0, 0, fmt.Errorf("malformed decimal args: %s", args)
-	}
-	args = strings.TrimPrefix(args, "(")
-	args = strings.TrimSuffix(args, ")")
-	parts := strings.Split(args, ":")
-	if len(parts) != 2 {
-		return 0, 0, fmt.Errorf("malformed decimal args: (%s)", args)
-	}
-	s, err := strconv.ParseInt(parts[0], 10, 32)
-	if err != nil {
-		return 0, 0, err
-	}
-	p, err := strconv.ParseInt(parts[1], 10, 32)
-	if err != nil {
-		return 0, 0, err
-	}
-	return int(s), int(p), nil
-}
-
-func parseTimestampArgs(args string) (TimeUnit, error) {
-	if !strings.HasPrefix(args, "(") || !strings.HasSuffix(args, ")") {
-		return nil, fmt.Errorf("malformed timestamp args: %s", args)
-	}
-
-	args = strings.TrimPrefix(args, "(")
-	args = strings.TrimSuffix(args, ")")
-
-	if len(args) == 0 {
-		return Millisecond, nil
-	}
-
-	switch args {
-	case "millisecond":
-		return Millisecond, nil
-	case "microsecond":
-		return Microsecond, nil
-	case "nanosecond":
-		return Nanosecond, nil
-	default:
-	}
-
-	return nil, fmt.Errorf("unknown time unit: %s", args)
-}
-
-type goNode struct {
-	Node
-	gotype reflect.Type
-}
-
-func (n *goNode) GoType() reflect.Type { return n.gotype }
-
-var (
-	_ RowGroupOption = (*Schema)(nil)
-	_ ReaderOption   = (*Schema)(nil)
-	_ WriterOption   = (*Schema)(nil)
-)


Test patch:
diff --git a/schema_test.go b/schema_test.go
index 2ee9a9b..333854f 100644
--- a/schema_test.go
+++ b/schema_test.go
@@ -116,6 +116,27 @@ func TestSchemaOf(t *testing.T) {
 			}),
 			print: `message {
 	required binary Name (JSON);
+}`,
+		},
+
+		{
+			value: new(struct {
+				A map[int64]string `parquet:"," parquet-key:",timestamp"`
+				B map[int64]string
+			}),
+			print: `message {
+	required group A (MAP) {
+		repeated group key_value {
+			required int64 key (TIMESTAMP(isAdjustedToUTC=true,unit=MILLIS));
+			required binary value (STRING);
+		}
+	}
+	required group B (MAP) {
+		repeated group key_value {
+			required int64 key (INT(64,true));
+			required binary value (STRING);
+		}
+	}
 }`,
 		},
 	}