Repository: segmentio/parquet-go
Pull request #126 (2022-04-06 18:22:21)

add parquet.ColumnMapping
This PR adds a new `parquet.ColumnMapping` type to the library, which intends to help applications perform conversions of column paths to their indexes in a row, which is useful when performing conversions of Go values into parquet rows for example.


Patch:
diff --git a/column_mapping.go b/column_mapping.go
new file mode 100644
index 0000000..7d69cae
--- /dev/null
+++ b/column_mapping.go
@@ -0,0 +1,106 @@
+package parquet
+
+import (
+	"fmt"
+	"strings"
+)
+
+// ColumnMapping is an interface representing types which map column paths to
+// their zero-based index in a parquet schema.
+//
+// ColumnMapping is a read-only type, and therefore can be used concurrently
+// from multiple goroutines.
+type ColumnMapping struct {
+	mapping columnMapping
+	columns [][]string
+}
+
+// ColumnIndex returns the column index of the column at the given path.
+// The path is the sequence of column names identifying a leaf column, starting
+// from the root column of the schema.
+//
+// If the path was not found in the mapping, or if it did not represent a
+// leaf column of the parquet schema, the method returns a negative value.
+func (m *ColumnMapping) ColumnIndex(path ...string) (columnIndex int) {
+	return int(m.mapping.lookup(path))
+}
+
+// ColumnPaths returns the list of column paths available in the mapping.
+//
+// The method always returns the same slice value across calls to ColumnPaths,
+// applications should treat it as immutable.
+func (m *ColumnMapping) ColumnPaths() [][]string {
+	return m.columns
+}
+
+// String returns a human-readable representation of the column mapping.
+func (m *ColumnMapping) String() string {
+	s := new(strings.Builder)
+	s.WriteByte('{')
+
+	if len(m.columns) > 0 {
+		for _, path := range m.columns {
+			fmt.Fprintf(s, "\n  % 2d => %q", m.ColumnIndex(path...), columnPath(path))
+		}
+		s.WriteByte('\n')
+	}
+
+	s.WriteByte('}')
+	return s.String()
+}
+
+// ColumnMappingOf constructs the column mapping of the given schema.
+func ColumnMappingOf(schema Node) *ColumnMapping {
+	mapping := make(columnMappingGroup)
+	columns := make([][]string, 0, 16)
+
+	forEachLeafColumnOf(schema, func(leaf leafColumn) {
+		column := make([]string, len(leaf.path))
+		copy(column, leaf.path)
+		columns = append(columns, column)
+
+		group, path := mapping, leaf.path
+
+		for len(path) > 1 {
+			columnName := path[0]
+			g, ok := group[columnName].(columnMappingGroup)
+			if !ok {
+				g = make(columnMappingGroup)
+				group[columnName] = g
+			}
+			group, path = g, path[1:]
+		}
+
+		group[path[0]] = columnMappingLeaf(leaf.columnIndex)
+	})
+
+	return &ColumnMapping{
+		mapping: mapping,
+		columns: columns,
+	}
+}
+
+type columnMapping interface {
+	lookup(path columnPath) (columnIndex int16)
+}
+
+type columnMappingGroup map[string]columnMapping
+
+func (group columnMappingGroup) lookup(path columnPath) int16 {
+	if len(path) > 0 {
+		c, ok := group[path[0]]
+		if ok {
+			return c.lookup(path[1:])
+		}
+	}
+	return -1
+}
+
+type columnMappingLeaf int16
+
+func (leaf columnMappingLeaf) lookup(path columnPath) int16 {
+	if len(path) == 0 {
+		return int16(leaf)
+	}
+	return -1
+}


Test patch:
diff --git a/column_mapping_test.go b/column_mapping_test.go
new file mode 100644
index 0000000..a19214c
--- /dev/null
+++ b/column_mapping_test.go
@@ -0,0 +1,29 @@
+package parquet_test
+
+import (
+	"fmt"
+
+	"github.com/segmentio/parquet-go"
+)
+
+func ExampleColumnMapping() {
+	schema := parquet.SchemaOf(struct {
+		FirstName  string `parquet:"first_name"`
+		LastName   string `parquet:"last_name"`
+		Attributes []struct {
+			Name  string `parquet:"name"`
+			Value string `parquet:"value"`
+		} `parquet:"attributes"`
+	}{})
+
+	mapping := parquet.ColumnMappingOf(schema)
+	fmt.Println(mapping)
+
+	// Output:
+	// {
+	//    0 => "attributes.name"
+	//    1 => "attributes.value"
+	//    2 => "first_name"
+	//    3 => "last_name"
+	// }
+}