Repository: planetlabs/gpq
Pull request #31 (2023-06-05 20:49:22)

Convert Parquet without metadata to GeoParquet
This adds support for converting Parquet files without "geo" file metadata to GeoParquet files.
```bash
gpq convert non-geo.parquet geo.parquet
```

For this to work, the non-geo Parquet file must have a column named "geometry" with WKT or WKB encoded geometries.


Patch:
diff --git a/cmd/gpq/convert.go b/cmd/gpq/convert.go
index 9e90872..cb0431a 100644
--- a/cmd/gpq/convert.go
+++ b/cmd/gpq/convert.go
@@ -15,17 +15,19 @@
 package main
 
 import (
+	"errors"
 	"fmt"
 	"os"
 	"strings"
 
 	"github.com/planetlabs/gpq/internal/geojson"
+	"github.com/planetlabs/gpq/internal/geoparquet"
 	"github.com/segmentio/parquet-go"
 )
 
 type ConvertCmd struct {
 	Input       string `arg:"" name:"input" help:"Input file." type:"existingfile"`
-	From        string `help:"Input file format.  Possible values: ${enum}." enum:"auto, geojson, geoparquet" default:"auto"`
+	From        string `help:"Input file format.  Possible values: ${enum}." enum:"auto, geojson, geoparquet, parquet" default:"auto"`
 	Output      string `arg:"" name:"output" help:"Output file." type:"path"`
 	To          string `help:"Output file format.  Possible values: ${enum}." enum:"auto, geojson, geoparquet" default:"auto"`
 	Min         int    `help:"Minimum number of features to consider when building a schema." default:"10"`
@@ -38,6 +40,7 @@ type FormatType string
 const (
 	AutoType       FormatType = "auto"
 	GeoParquetType FormatType = "geoparquet"
+	ParquetType    FormatType = "parquet"
 	GeoJSONType    FormatType = "geojson"
 	UnknownType    FormatType = "unknown"
 )
@@ -45,6 +48,7 @@ const (
 var validTypes = map[FormatType]bool{
 	AutoType:       true,
 	GeoParquetType: true,
+	ParquetType:    true,
 	GeoJSONType:    true,
 }
 
@@ -60,9 +64,12 @@ func getFormatType(filename string) FormatType {
 	if strings.HasSuffix(filename, ".json") || strings.HasSuffix(filename, ".geojson") {
 		return GeoJSONType
 	}
-	if strings.HasSuffix(filename, ".pq") || strings.HasSuffix(filename, ".parquet") || strings.HasSuffix(filename, ".geoparquet") {
+	if strings.HasSuffix(filename, ".gpq") || strings.HasSuffix(filename, ".geoparquet") {
 		return GeoParquetType
 	}
+	if strings.HasSuffix(filename, ".pq") || strings.HasSuffix(filename, ".parquet") {
+		return ParquetType
+	}
 	return UnknownType
 }
 
@@ -83,10 +90,6 @@ func (c *ConvertCmd) Run() error {
 		return fmt.Errorf("could not determine input format for %s", c.Input)
 	}
 
-	if inputFormat == outputFormat {
-		return fmt.Errorf("input and output are both the same type: %s", inputFormat)
-	}
-
 	input, readErr := os.Open(c.Input)
 	if readErr != nil {
 		return fmt.Errorf("failed to read from %q: %w", c.Input, readErr)
@@ -99,20 +102,27 @@ func (c *ConvertCmd) Run() error {
 	}
 	defer output.Close()
 
-	if inputFormat == GeoParquetType {
-		stat, statErr := os.Stat(c.Input)
-		if statErr != nil {
-			return fmt.Errorf("failed to get size of %q: %w", c.Input, statErr)
+	if inputFormat == GeoJSONType {
+		if outputFormat != ParquetType && outputFormat != GeoParquetType {
+			return errors.New("GeoJSON input can only be converted to GeoParquet")
 		}
+		convertOptions := &geojson.ConvertOptions{MinFeatures: c.Min, MaxFeatures: c.Max, Compression: c.Compression}
+		return geojson.ToParquet(input, output, convertOptions)
+	}
 
-		file, fileErr := parquet.OpenFile(input, stat.Size())
-		if fileErr != nil {
-			return fileErr
-		}
+	stat, statErr := os.Stat(c.Input)
+	if statErr != nil {
+		return fmt.Errorf("failed to get size of %q: %w", c.Input, statErr)
+	}
+
+	file, fileErr := parquet.OpenFile(input, stat.Size())
+	if fileErr != nil {
+		return fileErr
+	}
 
+	if outputFormat == GeoJSONType {
 		return geojson.FromParquet(file, output)
 	}
 
-	convertOptions := &geojson.ConvertOptions{MinFeatures: c.Min, MaxFeatures: c.Max, Compression: c.Compression}
-	return geojson.ToParquet(input, output, convertOptions)
+	return geoparquet.FromParquet(file, output, nil)
 }
diff --git a/internal/geojson/geojson.go b/internal/geojson/geojson.go
index 7121fa0..9ae8b8b 100644
--- a/internal/geojson/geojson.go
+++ b/internal/geojson/geojson.go
@@ -25,7 +25,6 @@ import (
 	orbjson "github.com/paulmach/orb/geojson"
 	"github.com/planetlabs/gpq/internal/geoparquet"
 	"github.com/segmentio/parquet-go"
-	"github.com/segmentio/parquet-go/compress"
 )
 
 type FeatureWriter struct {
@@ -65,7 +64,7 @@ func toFeature(row parquet.Row, schema *parquet.Schema, metadata *geoparquet.Met
 			continue
 		}
 
-		geometry, err := geoparquet.Geometry(value, name, metadata, schema)
+		geometry, _, err := geoparquet.Geometry(value, name, metadata, schema)
 		if err != nil {
 			return nil, err
 		}
@@ -523,25 +522,6 @@ var defaultOptions = &ConvertOptions{
 	Compression: "gzip",
 }
 
-func getCodec(codec string) (compress.Codec, error) {
-	switch codec {
-	case "uncompressed":
-		return &parquet.Uncompressed, nil
-	case "snappy":
-		return &parquet.Snappy, nil
-	case "gzip":
-		return &parquet.Gzip, nil
-	case "brotli":
-		return &parquet.Brotli, nil
-	case "zstd":
-		return &parquet.Zstd, nil
-	case "lz4raw":
-		return &parquet.Lz4Raw, nil
-	default:
-		return nil, fmt.Errorf("invalid compression codec %s", codec)
-	}
-}
-
 func ToParquet(input io.Reader, output io.Writer, convertOptions *ConvertOptions) error {
 	reader := NewFeatureReader(input)
 
@@ -571,7 +551,7 @@ func ToParquet(input io.Reader, output io.Writer, convertOptions *ConvertOptions
 		compression = defaultOptions.Compression
 	}
 
-	codec, codecErr := getCodec(compression)
+	codec, codecErr := geoparquet.GetCodec(compression)
 	if codecErr != nil {
 		return codecErr
 	}
diff --git a/internal/geoparquet/geoparquet.go b/internal/geoparquet/geoparquet.go
index 66d04e1..f37699b 100644
--- a/internal/geoparquet/geoparquet.go
+++ b/internal/geoparquet/geoparquet.go
@@ -24,6 +24,7 @@ import (
 	"github.com/paulmach/orb/encoding/wkb"
 	"github.com/paulmach/orb/encoding/wkt"
 	"github.com/segmentio/parquet-go"
+	"github.com/segmentio/parquet-go/compress"
 )
 
 const (
@@ -105,15 +106,19 @@ func (col *GeometryColumn) GetGeometryTypes() []string {
 	return types
 }
 
+func getDefaultGeometryColumn() *GeometryColumn {
+	return &GeometryColumn{
+		Encoding:      EncodingWKB,
+		GeometryTypes: []string{},
+	}
+}
+
 func DefaultMetadata() *Metadata {
 	return &Metadata{
 		Version:       Version,
 		PrimaryColumn: defaultGeometryColumn,
 		Columns: map[string]*GeometryColumn{
-			defaultGeometryColumn: {
-				Encoding:      EncodingWKB,
-				GeometryTypes: []string{},
-			},
+			defaultGeometryColumn: getDefaultGeometryColumn(),
 		},
 	}
 }
@@ -243,17 +248,17 @@ func (w *GenericWriter[T]) Close() error {
 
 var stringType = parquet.String().Type()
 
-func Geometry(value any, name string, metadata *Metadata, schema *parquet.Schema) (orb.Geometry, error) {
+func Geometry(value any, name string, metadata *Metadata, schema *parquet.Schema) (orb.Geometry, string, error) {
 	geometryString, ok := value.(string)
 	if !ok {
-		return nil, fmt.Errorf("unexpected geometry type: %t", value)
+		return nil, "", fmt.Errorf("unexpected geometry type: %t", value)
 	}
 
 	encoding := metadata.Columns[name].Encoding
 	if encoding == "" {
 		column, ok := schema.Lookup(name)
 		if !ok {
-			return nil, fmt.Errorf("missing column: %s", name)
+			return nil, "", fmt.Errorf("missing column: %s", name)
 		}
 		nodeType := column.Node.Type()
 		if nodeType == stringType {
@@ -261,7 +266,7 @@ func Geometry(value any, name string, metadata *Metadata, schema *parquet.Schema
 		} else if nodeType == parquet.ByteArrayType {
 			encoding = EncodingWKB
 		} else {
-			return nil, fmt.Errorf("unsupported geometry type: %s", nodeType)
+			return nil, "", fmt.Errorf("unsupported geometry type: %s", nodeType)
 		}
 	}
 
@@ -271,18 +276,173 @@ func Geometry(value any, name string, metadata *Metadata, schema *parquet.Schema
 	case EncodingWKB:
 		g, err := wkb.Unmarshal([]byte(geometryString))
 		if err != nil {
-			return nil, fmt.Errorf("trouble reading geometry: %w", err)
+			return nil, "", fmt.Errorf("trouble reading geometry as WKB: %w", err)
 		}
 		geometry = g
 	case EncodingWKT:
 		g, err := wkt.Unmarshal(geometryString)
 		if err != nil {
-			return nil, fmt.Errorf("trouble reading geometry: %w", err)
+			return nil, "", fmt.Errorf("trouble reading geometry as WKT: %w", err)
 		}
 		geometry = g
 	default:
-		return nil, fmt.Errorf("unsupported encoding: %s", encoding)
+		return nil, "", fmt.Errorf("unsupported encoding: %s", encoding)
+	}
+
+	return geometry, strings.ToUpper(encoding), nil
+}
+
+func GetCodec(codec string) (compress.Codec, error) {
+	switch codec {
+	case "uncompressed":
+		return &parquet.Uncompressed, nil
+	case "snappy":
+		return &parquet.Snappy, nil
+	case "gzip":
+		return &parquet.Gzip, nil
+	case "brotli":
+		return &parquet.Brotli, nil
+	case "zstd":
+		return &parquet.Zstd, nil
+	case "lz4raw":
+		return &parquet.Lz4Raw, nil
+	default:
+		return nil, fmt.Errorf("invalid compression codec %s", codec)
+	}
+}
+
+type ConvertOptions struct {
+	Compression string
+}
+
+func FromParquet(file *parquet.File, output io.Writer, convertOptions *ConvertOptions) error {
+	if convertOptions == nil {
+		convertOptions = &ConvertOptions{}
+	}
+	reader := NewRowReader(file)
+
+	schema := file.Schema()
+
+	codec := schema.Compression()
+	if convertOptions.Compression != "" {
+		candidate, codecErr := GetCodec(convertOptions.Compression)
+		if codecErr != nil {
+			return codecErr
+		}
+		codec = candidate
+	}
+
+	options := []parquet.WriterOption{
+		parquet.Compression(codec),
+		schema,
+	}
+
+	writerConfig, configErr := parquet.NewWriterConfig(options...)
+	if configErr != nil {
+		return configErr
+	}
+
+	writer := parquet.NewGenericWriter[any](output, writerConfig)
+
+	boundsLookup := map[string]*orb.Bound{}
+	geometryTypeLookup := map[string]map[string]bool{}
+
+	metadata, metadataErr := GetMetadata(file)
+	if metadataErr != nil {
+		metadata = &Metadata{
+			PrimaryColumn: defaultGeometryColumn,
+			Columns: map[string]*GeometryColumn{
+				defaultGeometryColumn: {},
+			},
+		}
+	}
+
+	for {
+		row, err := reader.Next()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return err
+		}
+
+		properties := map[string]any{}
+		if err := schema.Reconstruct(&properties, row); err != nil {
+			return err
+		}
+
+		for name, columnMetadata := range metadata.Columns {
+			value, ok := properties[name]
+			if !ok {
+				return fmt.Errorf("missing geometry column: %s", name)
+			}
+			geometry, encoding, err := Geometry(value, name, metadata, schema)
+			if err != nil {
+				return err
+			}
+
+			if encoding != EncodingWKB {
+				column, ok := schema.Lookup(name)
+				if !ok {
+					return fmt.Errorf("missing geometry column: %s", name)
+				}
+				geomBytes, wkbErr := wkb.Marshal(geometry)
+				if wkbErr != nil {
+					return fmt.Errorf("failed to encode %q geometry as wkb: %w", name, wkbErr)
+				}
+				row[column.ColumnIndex] = parquet.ValueOf(geomBytes)
+			}
+
+			if columnMetadata.Encoding != EncodingWKB {
+				columnMetadata.Encoding = EncodingWKB
+			}
+
+			bounds := geometry.Bound()
+			if boundsLookup[name] != nil {
+				bounds = bounds.Union(*boundsLookup[name])
+			}
+			boundsLookup[name] = &bounds
+
+			if geometryTypeLookup[name] == nil {
+				geometryTypeLookup[name] = map[string]bool{}
+			}
+			geometryTypeLookup[name][geometry.GeoJSONType()] = true
+		}
+
+		_, writeErr := writer.WriteRows([]parquet.Row{row})
+		if writeErr != nil {
+			return writeErr
+		}
+	}
+
+	for name, bounds := range boundsLookup {
+		if bounds != nil {
+			if metadata.Columns[name] == nil {
+				metadata.Columns[name] = getDefaultGeometryColumn()
+			}
+			metadata.Columns[name].Bounds = []float64{
+				bounds.Left(), bounds.Bottom(), bounds.Right(), bounds.Top(),
+			}
+		}
 	}
 
-	return geometry, nil
+	for name, types := range geometryTypeLookup {
+		geometryTypes := []string{}
+		if len(types) > 0 {
+			for geometryType := range types {
+				geometryTypes = append(geometryTypes, geometryType)
+			}
+		}
+		if metadata.Columns[name] == nil {
+			metadata.Columns[name] = getDefaultGeometryColumn()
+		}
+		metadata.Columns[name].GeometryTypes = geometryTypes
+	}
+
+	metadataBytes, jsonErr := json.Marshal(metadata)
+	if jsonErr != nil {
+		return fmt.Errorf("failed to serialize geo metadata: %w", jsonErr)
+	}
+	writer.SetKeyValueMetadata(MetadataKey, string(metadataBytes))
+	return writer.Close()
 }
diff --git a/internal/validator/rules.go b/internal/validator/rules.go
index 257ec9f..63567f7 100644
--- a/internal/validator/rules.go
+++ b/internal/validator/rules.go
@@ -505,7 +505,7 @@ func GeometryEncoding() Rule {
 			metadata := info.Metadata
 
 			for name, encoded := range geometries {
-				_, err := geoparquet.Geometry(encoded, name, metadata, schema)
+				_, _, err := geoparquet.Geometry(encoded, name, metadata, schema)
 				if err != nil {
 					return fatal("invalid geometry in column %q: %s", name, err)
 				}
diff --git a/internal/validator/validator.go b/internal/validator/validator.go
index 5573cbc..19a9d00 100644
--- a/internal/validator/validator.go
+++ b/internal/validator/validator.go
@@ -235,7 +235,7 @@ func (v *Validator) Report(ctx context.Context, file *parquet.File) (*Report, er
 
 		decodedGeometryMap := DecodedGeometryMap{}
 		for name, value := range encodedGeometryMap {
-			decoded, err := geoparquet.Geometry(value, name, metadata, schema)
+			decoded, _, err := geoparquet.Geometry(value, name, metadata, schema)
 			if err != nil {
 				return nil, fmt.Errorf("failed to decode geometry: %w", err)
 			}


Test patch:
diff --git a/internal/geoparquet/geoparquet_test.go b/internal/geoparquet/geoparquet_test.go
index a30d7e0..8b33981 100644
--- a/internal/geoparquet/geoparquet_test.go
+++ b/internal/geoparquet/geoparquet_test.go
@@ -15,10 +15,16 @@
 package geoparquet_test
 
 import (
+	"bytes"
+	"encoding/json"
+	"fmt"
 	"io"
 	"os"
 	"testing"
 
+	"github.com/paulmach/orb"
+	"github.com/paulmach/orb/encoding/wkb"
+	"github.com/paulmach/orb/encoding/wkt"
 	"github.com/planetlabs/gpq/internal/geoparquet"
 	"github.com/segmentio/parquet-go"
 	"github.com/stretchr/testify/assert"
@@ -152,3 +158,119 @@ func TestRowReaderV100Beta1(t *testing.T) {
 	name := firstRow[nameCol.ColumnIndex]
 	assert.Equal(t, "Fiji", name.String())
 }
+
+func makeParquet[T any](rows []T, metadata *geoparquet.Metadata) (*parquet.File, error) {
+	data := &bytes.Buffer{}
+
+	writer := parquet.NewGenericWriter[T](data)
+
+	_, writeErr := writer.Write(rows)
+	if writeErr != nil {
+		return nil, fmt.Errorf("trouble writing rows: %w", writeErr)
+	}
+
+	if metadata != nil {
+		jsonMetadata, jsonErr := json.Marshal(metadata)
+		if jsonErr != nil {
+			return nil, fmt.Errorf("trouble encoding metadata as json: %w", jsonErr)
+		}
+
+		writer.SetKeyValueMetadata(geoparquet.MetadataKey, string(jsonMetadata))
+	}
+
+	closeErr := writer.Close()
+	if closeErr != nil {
+		return nil, fmt.Errorf("trouble closing writer: %w", closeErr)
+	}
+
+	return parquet.OpenFile(bytes.NewReader(data.Bytes()), int64(data.Len()))
+}
+
+func TestFromParquetWithoutMetadata(t *testing.T) {
+	type Row struct {
+		Name     string `parquet:"name"`
+		Geometry []byte `parquet:"geometry"`
+	}
+
+	point, pointErr := wkb.Marshal(orb.Point{1, 2})
+	require.NoError(t, pointErr)
+
+	rows := []*Row{
+		{
+			Name:     "test-point",
+			Geometry: point,
+		},
+	}
+
+	parquetFile, inputErr := makeParquet(rows, nil)
+	require.NoError(t, inputErr)
+
+	output := &bytes.Buffer{}
+	convertErr := geoparquet.FromParquet(parquetFile, output, nil)
+	require.NoError(t, convertErr)
+
+	geoparquetInput := bytes.NewReader(output.Bytes())
+	geoparquetFile, outputErr := parquet.OpenFile(geoparquetInput, geoparquetInput.Size())
+	require.NoError(t, outputErr)
+
+	metadata, geoErr := geoparquet.GetMetadata(geoparquetFile)
+	require.NoError(t, geoErr)
+
+	assert.Len(t, metadata.Columns, 1)
+
+	primaryColumnMetadata := metadata.Columns[metadata.PrimaryColumn]
+
+	geometryTypes := primaryColumnMetadata.GetGeometryTypes()
+	assert.Len(t, geometryTypes, 1)
+	assert.Contains(t, geometryTypes, "Point")
+
+	bounds := primaryColumnMetadata.Bounds
+	assert.Equal(t, []float64{1, 2, 1, 2}, bounds)
+
+	assert.Equal(t, geoparquet.EncodingWKB, primaryColumnMetadata.Encoding)
+
+	assert.Equal(t, int64(1), geoparquetFile.NumRows())
+}
+
+func TestFromParquetWithWKT(t *testing.T) {
+	type Row struct {
+		Name     string `parquet:"name"`
+		Geometry string `parquet:"geometry"`
+	}
+
+	rows := []*Row{
+		{
+			Name:     "test-point",
+			Geometry: string(wkt.Marshal(orb.Point{1, 2})),
+		},
+	}
+
+	parquetFile, inputErr := makeParquet(rows, nil)
+	require.NoError(t, inputErr)
+
+	output := &bytes.Buffer{}
+	convertErr := geoparquet.FromParquet(parquetFile, output, nil)
+	require.NoError(t, convertErr)
+
+	geoparquetInput := bytes.NewReader(output.Bytes())
+	geoparquetFile, outputErr := parquet.OpenFile(geoparquetInput, geoparquetInput.Size())
+	require.NoError(t, outputErr)
+
+	metadata, geoErr := geoparquet.GetMetadata(geoparquetFile)
+	require.NoError(t, geoErr)
+
+	assert.Len(t, metadata.Columns, 1)
+
+	primaryColumnMetadata := metadata.Columns[metadata.PrimaryColumn]
+
+	geometryTypes := primaryColumnMetadata.GetGeometryTypes()
+	assert.Len(t, geometryTypes, 1)
+	assert.Contains(t, geometryTypes, "Point")
+
+	bounds := primaryColumnMetadata.Bounds
+	assert.Equal(t, []float64{1, 2, 1, 2}, bounds)
+
+	assert.Equal(t, geoparquet.EncodingWKB, primaryColumnMetadata.Encoding)
+
+	assert.Equal(t, int64(1), geoparquetFile.NumRows())
+}