nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: planetlabs/gpq
Pull request #92 (2023-10-04 20:13:56)

Check if geometry columns exist when converting
This makes it so the `convert` command fails if the data is missing columns in the geo metadata.


Patch:
diff --git a/internal/geoparquet/geoparquet.go b/internal/geoparquet/geoparquet.go
index 9a784b7..23f121c 100644
--- a/internal/geoparquet/geoparquet.go
+++ b/internal/geoparquet/geoparquet.go
@@ -2,6 +2,7 @@ package geoparquet
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 
@@ -62,9 +63,20 @@ func FromParquet(input parquet.ReaderAtSeeker, output io.Writer, convertOptions
 	datasetInfo := geo.NewDatasetStats(true)
 	transformSchema := func(fileReader *file.Reader) (*schema.Schema, error) {
 		inputSchema := fileReader.MetaData().Schema
+		inputRoot := inputSchema.Root()
 		metadata := getMetadata(fileReader, convertOptions)
-		for fieldNum := 0; fieldNum < inputSchema.Root().NumFields(); fieldNum += 1 {
-			field := inputSchema.Root().Field(fieldNum)
+		for geomColName := range metadata.Columns {
+			if inputRoot.FieldIndexByName(geomColName) < 0 {
+				message := fmt.Sprintf(
+					"expected a geometry column named %q,"+
+						" use the --input-primary-column to supply a different primary geometry",
+					geomColName,
+				)
+				return nil, errors.New(message)
+			}
+		}
+		for fieldNum := 0; fieldNum < inputRoot.NumFields(); fieldNum += 1 {
+			field := inputRoot.Field(fieldNum)
 			name := field.Name()
 			if _, ok := metadata.Columns[name]; !ok {
 				continue
@@ -78,9 +90,7 @@ func FromParquet(input parquet.ReaderAtSeeker, output io.Writer, convertOptions
 			return inputSchema, nil
 		}
 
-		inputRoot := inputSchema.Root()
 		numFields := inputRoot.NumFields()
-
 		fields := make([]schema.Node, numFields)
 		for fieldNum := 0; fieldNum < numFields; fieldNum += 1 {
 			inputField := inputRoot.Field(fieldNum)


Test patch:
diff --git a/internal/geoparquet/geoparquet_test.go b/internal/geoparquet/geoparquet_test.go
index 2df7d63..86803b8 100644
--- a/internal/geoparquet/geoparquet_test.go
+++ b/internal/geoparquet/geoparquet_test.go
@@ -192,6 +192,26 @@ func TestFromParquetWithoutMetadata(t *testing.T) {
 	assert.Equal(t, int64(1), reader.NumRows())
 }
 
+func TestFromParquetWithoutDefaultGeometryColumn(t *testing.T) {
+	type Row struct {
+		Name string `parquet:"name=name, logical=String" json:"name"`
+		Geom []byte `parquet:"name=geom" json:"geom"`
+	}
+
+	rows := []*Row{
+		{
+			Name: "test-point",
+			Geom: toWKB(t, orb.Point{1, 2}),
+		},
+	}
+
+	input := test.ParquetFromStructs(t, rows)
+
+	output := &bytes.Buffer{}
+	convertErr := geoparquet.FromParquet(input, output, nil)
+	require.ErrorContains(t, convertErr, "expected a geometry column named \"geometry\"")
+}
+
 func TestMetadataClone(t *testing.T) {
 	metadata := geoparquet.DefaultMetadata()
 	clone := metadata.Clone()