diff --git a/go.mod b/go.mod index f484f9a79..3e0d4dce5 100644 --- a/go.mod +++ b/go.mod @@ -47,6 +47,7 @@ require ( github.com/testcontainers/testcontainers-go/modules/compose v0.42.0 github.com/twmb/avro v1.7.2 github.com/twmb/murmur3 v1.1.8 + github.com/twpayne/go-geom v1.6.1 github.com/uptrace/bun v1.2.18 github.com/uptrace/bun/dialect/mssqldialect v1.2.18 github.com/uptrace/bun/dialect/mysqldialect v1.2.18 diff --git a/go.sum b/go.sum index a20f133d7..0f7a0347a 100644 --- a/go.sum +++ b/go.sum @@ -56,6 +56,8 @@ github.com/AzureAD/microsoft-authentication-extensions-for-go/cache v0.1.1 h1:WJ github.com/AzureAD/microsoft-authentication-extensions-for-go/cache v0.1.1/go.mod h1:tCcJZ0uHAmvjsVYzEFivsRTN00oz5BEsRgQHu5JZ9WE= github.com/AzureAD/microsoft-authentication-library-for-go v1.7.2 h1:RHK7bS+HQMslb1sZpAokUt+zTVmue0hKSs2C791hhzU= github.com/AzureAD/microsoft-authentication-library-for-go v1.7.2/go.mod h1:HKpQxkWaGLJ+D/5H8QRpyQXA1eKjxkFlOMwck5+33Jk= +github.com/DATA-DOG/go-sqlmock v1.5.2 h1:OcvFkGmslmlZibjAjaHm3L//6LiuBgolP7OputlJIzU= +github.com/DATA-DOG/go-sqlmock v1.5.2/go.mod h1:88MAG/4G7SMwSE3CeA0ZKzrT5CiOU3OJ+JlNzwDqpNU= github.com/DefangLabs/secret-detector v0.0.0-20250403165618-22662109213e h1:rd4bOvKmDIx0WeTv9Qz+hghsgyjikFiPrseXHlKepO0= github.com/DefangLabs/secret-detector v0.0.0-20250403165618-22662109213e/go.mod h1:blbwPQh4DTlCZEfk1BLU4oMIhLda2U+A840Uag9DsZw= github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.31.0 h1:DHa2U07rk8syqvCge0QIGMCE1WxGj9njT44GH7zNJLQ= @@ -85,6 +87,10 @@ github.com/RoaringBitmap/roaring/v2 v2.18.2 h1:oPq3Cgx//iDuJQVp6xSInAKW34J9CEwE5 github.com/RoaringBitmap/roaring/v2 v2.18.2/go.mod h1:eq4wdNXxtJIS/oikeCzdX1rBzek7ANzbth041hrU8Q4= github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d h1:licZJFw2RwpHMqeKTCYkitsPqHNxTmd4SNR5r94FGM8= github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d/go.mod h1:asat636LX7Bqt5lYEZ27JNDcqxfjdBQuJ/MM4CN/Lzo= +github.com/alecthomas/assert/v2 v2.10.0 h1:jjRCHsj6hBJhkmhznrCzoNpbA3zqy0fYiUcYZP/GkPY= +github.com/alecthomas/assert/v2 v2.10.0/go.mod h1:Bze95FyfUr7x34QZrjL+XP+0qgp/zg8yS+TtBj1WA3k= +github.com/alecthomas/repr v0.4.0 h1:GhI2A8MACjfegCPVq9f1FLvIBS+DrQ2KQBFZP1iFzXc= +github.com/alecthomas/repr v0.4.0/go.mod h1:Fr0507jx4eOXV7AlPV6AVZLYrLIuIeSOWtW57eE/O/4= github.com/alexflint/go-arg v1.6.1 h1:uZogJ6VDBjcuosydKgvYYRhh9sRCusjOvoOLZopBlnA= github.com/alexflint/go-arg v1.6.1/go.mod h1:nQ0LFYftLJ6njcaee0sU+G0iS2+2XJQfA8I062D0LGc= github.com/alexflint/go-scalar v1.2.0 h1:WR7JPKkeNpnYIOfHRa7ivM21aWAdHD0gEWHCx+WQBRw= @@ -392,6 +398,8 @@ github.com/hashicorp/go-version v1.9.0 h1:CeOIz6k+LoN3qX9Z0tyQrPtiB1DFYRPfCIBtaX github.com/hashicorp/go-version v1.9.0/go.mod h1:fltr4n8CU8Ke44wwGCBoEymUuxUHl09ZGVZPK5anwXA= github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= +github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM= +github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg= github.com/in-toto/attestation v1.1.2 h1:MBFn6lsMq6dptQZJBhalXTcWMb/aJy3V+GX3VYj/V1E= github.com/in-toto/attestation v1.1.2/go.mod h1:gYFddHMZj3DiQ0b62ltNi1Vj5rC879bTmBbrv9CRHpM= github.com/in-toto/in-toto-golang v0.11.0 h1:nfidMYBFx+E0lnmX5KUnN2Pdm8zdNKal1ayjJuzzRoA= @@ -624,6 +632,8 @@ github.com/twmb/avro v1.7.2 h1:cmrEBRSbELRqsg/dRkQvVWuOaR2EfGifHIt/2iJ9lfI= github.com/twmb/avro v1.7.2/go.mod h1:X0fT1dY2xcbV4YuCE4mYro+qljHl4kUF5uA/2z1rgSk= github.com/twmb/murmur3 v1.1.8 h1:8Yt9taO/WN3l08xErzjeschgZU2QSrwm1kclYq+0aRg= github.com/twmb/murmur3 v1.1.8/go.mod h1:Qq/R7NUyOfr65zD+6Q5IHKsJLwP7exErjN6lyyq3OSQ= +github.com/twpayne/go-geom v1.6.1 h1:iLE+Opv0Ihm/ABIcvQFGIiFBXd76oBIar9drAwHFhR4= +github.com/twpayne/go-geom v1.6.1/go.mod h1:Kr+Nly6BswFsKM5sd31YaoWS5PeDDH2NftJTK7Gd028= github.com/uptrace/bun v1.2.18 h1:3HnRcMfS6OBPMG1eSOzlbFJ/X/AyMEJb7rMxE6VQvDU= github.com/uptrace/bun v1.2.18/go.mod h1:wNltaKJk4JtOt4SG5I5zmA7v0/Mzjh1+/S906Rayd3Y= github.com/uptrace/bun/dialect/mssqldialect v1.2.18 h1:nYzHoyJKJlIyl5i95Exi8ZTK8ooKWG+o3z3f404d/yQ= diff --git a/table/arrow_utils.go b/table/arrow_utils.go index dcc1ff2b5..128196764 100644 --- a/table/arrow_utils.go +++ b/table/arrow_utils.go @@ -18,12 +18,15 @@ package table import ( + "bytes" "context" "encoding/base64" "encoding/json" + "errors" "fmt" "iter" "log/slog" + "regexp" "slices" "strconv" "strings" @@ -375,6 +378,17 @@ func (c convertToIceberg) Primitive(dt arrow.DataType) (result iceberg.NestedFie result.Type = iceberg.PrimitiveTypes.UUID case "parquet.variant": result.Type = iceberg.VariantType{} + case "geoarrow.wkb": + wkb, ok := dt.(*geoarrow.WKBType) + if !ok { + panic(fmt.Errorf("%w: unsupported arrow type for conversion - %s", iceberg.ErrInvalidSchema, dt)) + } + + iceType, err := geoArrowMetadataToIcebergType(wkb.Metadata()) + if err != nil { + panic(fmt.Errorf("%w: converting geoarrow metadata: %w", iceberg.ErrInvalidSchema, err)) + } + result.Type = iceType default: panic(fmt.Errorf("%w: unsupported arrow type for conversion - %s", iceberg.ErrInvalidSchema, dt)) } @@ -658,20 +672,26 @@ func (c convertToArrow) VisitVariant() arrow.Field { return arrow.Field{Type: extensions.NewDefaultVariantType()} } -func (c convertToArrow) VisitGeometry(iceberg.GeometryType) arrow.Field { +func (c convertToArrow) VisitGeometry(g iceberg.GeometryType) arrow.Field { + meta := icebergCRSToGeoArrowMetadata(g.CRS()) if c.useLargeTypes { - return arrow.Field{Type: geoarrow.NewWKBType(geoarrow.WKBWithLargeBinaryStorage())} + return arrow.Field{Type: geoarrow.NewWKBType(geoarrow.WKBWithLargeBinaryStorage(), geoarrow.WKBWithMetadata(meta))} } - return arrow.Field{Type: geoarrow.NewWKBType(geoarrow.WKBWithBinaryStorage())} + return arrow.Field{Type: geoarrow.NewWKBType(geoarrow.WKBWithBinaryStorage(), geoarrow.WKBWithMetadata(meta))} } -func (c convertToArrow) VisitGeography(iceberg.GeographyType) arrow.Field { +func (c convertToArrow) VisitGeography(g iceberg.GeographyType) arrow.Field { + meta := icebergCRSToGeoArrowMetadata(g.CRS()) + // Always add an edge to differentiate between Geography and Geometry arrow fields. + // Note that the edge convention is a best-effort hint and planar geography from other clients won't round-trip through Arrow alone. + meta.Edges = geoarrow.EdgeInterpolation(g.Algorithm()) + if c.useLargeTypes { - return arrow.Field{Type: geoarrow.NewWKBType(geoarrow.WKBWithLargeBinaryStorage())} + return arrow.Field{Type: geoarrow.NewWKBType(geoarrow.WKBWithLargeBinaryStorage(), geoarrow.WKBWithMetadata(meta))} } - return arrow.Field{Type: geoarrow.NewWKBType(geoarrow.WKBWithBinaryStorage())} + return arrow.Field{Type: geoarrow.NewWKBType(geoarrow.WKBWithBinaryStorage(), geoarrow.WKBWithMetadata(meta))} } var _ iceberg.SchemaVisitorPerPrimitiveType[arrow.Field] = convertToArrow{} @@ -1826,3 +1846,150 @@ func positionDeleteRecordsToDataFilesDV(ctx context.Context, rootLocation string } } } + +func checkCRSString(rawCrs json.RawMessage) bool { + b := bytes.TrimSpace(rawCrs) + + return len(b) > 0 && b[0] == '"' +} + +func checkCRSJSON(rawCrs json.RawMessage) bool { + b := bytes.TrimSpace(rawCrs) + + return len(b) > 0 && b[0] == '{' +} + +func geoArrowCRSToIcebergCRS(meta geoarrow.Metadata) (string, error) { + if len(meta.CRS) == 0 { + return "srid:0", nil + } + + switch { + case checkCRSString(meta.CRS): + var crs string + + if err := json.Unmarshal(meta.CRS, &crs); err != nil { + return "", fmt.Errorf("invalid geoarrow CRS metadata: %w", err) + } + + if strings.EqualFold(crs, "OGC:CRS84") || strings.EqualFold(crs, "EPSG:4326") { + return "OGC:CRS84", nil + } + + switch meta.CRSType { + case geoarrow.CRSTypeSRID: + return "srid:" + crs, nil + case geoarrow.CRSTypeWKT22019: + return "", errors.New("CRS type wkt2:2019 not supported") + default: + if len(crs) <= 32 { + return crs, nil + } + + return "", errors.New("crs length too long") + } + case checkCRSJSON(meta.CRS): + var crs map[string]json.RawMessage + + if err := json.Unmarshal(meta.CRS, &crs); err != nil { + return "", fmt.Errorf("invalid geoarrow CRS metadata: %w", err) + } + + idRaw, ok := crs["id"] + if !ok || len(idRaw) == 0 { + return "", errors.New("unsupported CRS") + } + + var id map[string]json.RawMessage + if err := json.Unmarshal(idRaw, &id); err != nil { + return "", errors.New("unsupported CRS") + } + + var authority string + if err := json.Unmarshal(id["authority"], &authority); err != nil || authority == "" { + return "", errors.New("unsupported CRS") + } + + codeRaw := id["code"] + if len(codeRaw) == 0 { + return "", errors.New("unsupported CRS") + } + + var code string + if err := json.Unmarshal(codeRaw, &code); err != nil { + var codeNum json.Number + if err := json.Unmarshal(codeRaw, &codeNum); err != nil { + return "", errors.New("unsupported CRS") + } + code = codeNum.String() + } + if code == "" { + return "", errors.New("unsupported CRS") + } + + authorityCode := authority + ":" + code + if strings.EqualFold(authorityCode, "OGC:CRS84") || strings.EqualFold(authorityCode, "EPSG:4326") { + return "OGC:CRS84", nil + } + + return authorityCode, nil + default: + return "", errors.New("unsupported CRS: CRS must either be omitted, a string, or a JSON object") + } +} + +func geoArrowMetadataToIcebergType(meta geoarrow.Metadata) (iceberg.Type, error) { + crs, err := geoArrowCRSToIcebergCRS(meta) + if err != nil { + return nil, err + } + + switch meta.Edges { + case geoarrow.EdgePlanar: + return iceberg.GeometryTypeOf(crs) + case geoarrow.EdgeVincenty, geoarrow.EdgeKarney, geoarrow.EdgeThomas, geoarrow.EdgeAndoyer, geoarrow.EdgeSpherical: + return iceberg.GeographyTypeOf(crs, string(meta.Edges)) + default: + return nil, fmt.Errorf("unsupported geoarrow edges %q", meta.Edges) + } +} + +var authorityCodeCRS = regexp.MustCompile(`^[A-Za-z][A-Za-z0-9_-]*:[A-Za-z0-9_.-]+$`) + +func icebergCRSToGeoArrowMetadata(crs string) geoarrow.Metadata { + lowerCRS := strings.ToLower(crs) + if strings.HasPrefix(lowerCRS, "srid:") { + id := crs[len("srid:"):] + + if id == "0" { + return geoarrow.NewMetadata() // srid:0 maps to omitted GeoArrow CRS + } + + raw, _ := json.Marshal(id) //nolint:errcheck // Marshalling a string can't fail + + return geoarrow.Metadata{ + CRS: raw, + CRSType: geoarrow.CRSTypeSRID, + } + } + + if strings.HasPrefix(lowerCRS, "projjson:") { + panic(fmt.Errorf("%w: projjson CRS not supported yet", iceberg.ErrInvalidSchema)) + } + + var raw []byte + + if lowerCRS == "epsg:4326" { + // collapse EPSG:4326 to OGC:CRS84 + raw, _ = json.Marshal("OGC:CRS84") //nolint:errcheck // Marshalling a string can't fail + } else { + raw, _ = json.Marshal(crs) //nolint:errcheck // Marshalling a string can't fail + } + + meta := geoarrow.Metadata{CRS: raw} + if authorityCodeCRS.MatchString(crs) { + meta.CRSType = geoarrow.CRSTypeAuthorityCode + } + + return meta +} diff --git a/table/arrow_utils_test.go b/table/arrow_utils_test.go index 07d17fd74..3e5f99262 100644 --- a/table/arrow_utils_test.go +++ b/table/arrow_utils_test.go @@ -35,6 +35,7 @@ import ( "github.com/apache/arrow-go/v18/parquet/variant" "github.com/apache/iceberg-go" "github.com/apache/iceberg-go/table" + "github.com/geoarrow/geoarrow-go" "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -430,6 +431,462 @@ func TestVariantArrowConversion(t *testing.T) { }) } +func assertGeoArrowWKB(t *testing.T, dt arrow.DataType, storage arrow.DataType, want geoarrow.Metadata) { + t.Helper() + + wkb, ok := dt.(*geoarrow.WKBType) + require.True(t, ok, "expected *geoarrow.WKBType, got %T", dt) + assert.Equal(t, "geoarrow.wkb", wkb.ExtensionName()) + assert.True(t, arrow.TypeEqual(storage, wkb.StorageType())) + assert.Equal(t, want.CRS, wkb.Metadata().CRS) + assert.Equal(t, want.Edges, wkb.Metadata().Edges) +} + +func assertGeoArrowWKBMetadataJSON(t *testing.T, dt arrow.DataType, storage arrow.DataType, wantJSON string) { + t.Helper() + + wkb, ok := dt.(*geoarrow.WKBType) + require.True(t, ok, "expected *geoarrow.WKBType, got %T", dt) + assert.Equal(t, "geoarrow.wkb", wkb.ExtensionName()) + assert.True(t, arrow.TypeEqual(storage, wkb.StorageType())) + + var want map[string]json.RawMessage + require.NoError(t, json.Unmarshal([]byte(wantJSON), &want)) + + meta := wkb.Metadata() + + if wantCRS, ok := want["crs"]; ok { + assert.JSONEq(t, string(wantCRS), string(meta.CRS), "crs mismatch") + } else { + assert.Empty(t, meta.CRS, "expected omitted crs") + } + + if wantEdges, ok := want["edges"]; ok { + var edges string + require.NoError(t, json.Unmarshal(wantEdges, &edges)) + assert.Equal(t, geoarrow.EdgeInterpolation(edges), meta.Edges, "edges mismatch") + } else { + assert.Empty(t, meta.Edges, "expected omitted edges") + } +} + +func jsonCRS(s string) json.RawMessage { + raw, _ := json.Marshal(s) + + return raw +} + +func TestIcebergGeoTypesToArrowSchema(t *testing.T) { + geomSRID, err := iceberg.GeometryTypeOf("srid:4326") + require.NoError(t, err) + geogKarney, err := iceberg.GeographyTypeOf("srid:4269", "karney") + require.NoError(t, err) + geomSRID0, err := iceberg.GeometryTypeOf("srid:0") + require.NoError(t, err) + geomEPSG4267, err := iceberg.GeometryTypeOf("EPSG:4267") + require.NoError(t, err) + geogSpherical, err := iceberg.GeographyTypeOf("OGC:CRS84", "spherical") + require.NoError(t, err) + geogKarneyDefaultCRS, err := iceberg.GeographyTypeOf("OGC:CRS84", "karney") + require.NoError(t, err) + geogVincenty, err := iceberg.GeographyTypeOf("OGC:CRS84", "vincenty") + require.NoError(t, err) + geogAndoyer, err := iceberg.GeographyTypeOf("OGC:CRS84", "andoyer") + require.NoError(t, err) + geogThomas, err := iceberg.GeographyTypeOf("OGC:CRS84", "thomas") + require.NoError(t, err) + geogSRID0, err := iceberg.GeographyTypeOf("srid:0", "spherical") + require.NoError(t, err) + geogEPSG4267, err := iceberg.GeographyTypeOf("EPSG:4267", "spherical") + require.NoError(t, err) + defaultGeometry, err := iceberg.GeometryTypeOf("OGC:CRS84") + require.NoError(t, err) + geomEPSG3857, err := iceberg.GeometryTypeOf("EPSG:3857") + require.NoError(t, err) + geogEPSG4267Karney, err := iceberg.GeographyTypeOf("EPSG:4267", "karney") + require.NoError(t, err) + + typeCases := []struct { + name string + ice iceberg.Type + geoarrowMetaJSON string + }{ + // Note that these tests below are based on arrow-rs tests (https://github.com/apache/arrow-rs/pull/10065) + // Geometry with default CRS (defaults to OGC:CRS84 per Parquet spec) + { + name: "geometry_default_crs", + ice: iceberg.GeometryType{}, + geoarrowMetaJSON: `{"crs":"OGC:CRS84"}`, + }, + // Geometry with srid:0 should result in an unset (omitted) CRS + { + name: "geometry_srid_0", + ice: geomSRID0, + geoarrowMetaJSON: `{}`, + }, + // Geometry with custom CRSes (authority:code and partial projjson) + { + name: "geometry_epsg_4267", + ice: geomEPSG4267, + geoarrowMetaJSON: `{"crs":"EPSG:4267"}`, + }, + // Geography with default CRS (default OGC:CRS84, spherical edges) + { + name: "geography_default_crs", + ice: iceberg.GeographyType{}, + geoarrowMetaJSON: `{"crs":"OGC:CRS84","edges":"spherical"}`, + }, + // Geography with explicit edges + { + name: "geography_explicit_spherical", + ice: geogSpherical, + geoarrowMetaJSON: `{"crs":"OGC:CRS84","edges":"spherical"}`, + }, + { + name: "geography_karney", + ice: geogKarneyDefaultCRS, + geoarrowMetaJSON: `{"crs":"OGC:CRS84","edges":"karney"}`, + }, + { + name: "geography_vincenty", + ice: geogVincenty, + geoarrowMetaJSON: `{"crs":"OGC:CRS84","edges":"vincenty"}`, + }, + { + name: "geography_andoyer", + ice: geogAndoyer, + geoarrowMetaJSON: `{"crs":"OGC:CRS84","edges":"andoyer"}`, + }, + { + name: "geography_thomas", + ice: geogThomas, + geoarrowMetaJSON: `{"crs":"OGC:CRS84","edges":"thomas"}`, + }, + // Geography with srid:0 should result in an unset (omitted) CRS and spherical edges + { + name: "geography_srid_0", + ice: geogSRID0, + geoarrowMetaJSON: `{"edges":"spherical"}`, + }, + // Geography with custom CRSes (authority:code and partial projjson) + { + name: "geography_epsg_4267", + ice: geogEPSG4267, + geoarrowMetaJSON: `{"crs":"EPSG:4267","edges":"spherical"}`, + }, + // Happy path SRID + { + name: "geometry_srid_4326", + ice: geomSRID, + geoarrowMetaJSON: `{"crs":"4326", "crs_type":"srid"}`, + }, + } + + // The following tests focus on edge cases and pinning specific behavior for read case + readOnlyCases := []struct { + name string + ice iceberg.Type + geoarrowMetaJSON string + }{ + { + name: "geometry_srid_0_with_crs_type", + ice: geomSRID0, + geoarrowMetaJSON: `{"crs_type":"authority_code"}`, + }, + { + name: "case_insensitive_default_geometry", + ice: defaultGeometry, + geoarrowMetaJSON: `{"crs":"OgC:cRs84"}`, + }, + { + name: "case_insensitive_default_geometry_epsg_4326", + ice: defaultGeometry, + geoarrowMetaJSON: `{"crs":"EpSg:4326"}`, + }, + { + name: "geometry_epsg_4326", + ice: defaultGeometry, + geoarrowMetaJSON: `{"crs":"epsg:4326"}`, + }, + { + name: "geometry_epsg_4326_incorrect_type", + ice: defaultGeometry, + geoarrowMetaJSON: `{"crs":"epsg:4326", "crs_type":"projjson"}`, + }, + + // Translated from arrow-rs geo logical type read tests (https://github.com/apache/arrow-rs/pull/10065) + // Geometry with no CRS should be GEOMETRY(srid:0) + { + name: "geometry_no_crs", + ice: geomSRID0, + geoarrowMetaJSON: `{}`, + }, + // Geometry with string CRS + { + name: "geometry_epsg_4267_from_crs", + ice: geomEPSG4267, + geoarrowMetaJSON: `{"crs":"EPSG:4267"}`, + }, + // Geometry with PROJJSON CRS + { + name: "geometry_projjson_epsg_3857", + ice: geomEPSG3857, + geoarrowMetaJSON: `{"crs":{"id":{"authority":"EPSG","code":3857}}}`, + }, + // Geometry with lon/lat CRSes (canonically removed because lon/lat is the + // default Iceberg CRS) + { + name: "geometry_ogc_crs84_canonical", + ice: iceberg.GeometryType{}, + geoarrowMetaJSON: `{"crs":"OGC:CRS84"}`, + }, + { + name: "geometry_epsg_4326_canonical", + ice: iceberg.GeometryType{}, + geoarrowMetaJSON: `{"crs":"EPSG:4326"}`, + }, + { + name: "geometry_projjson_epsg_4326_canonical", + ice: iceberg.GeometryType{}, + geoarrowMetaJSON: `{"crs":{"id":{"authority":"EPSG","code":4326}}}`, + }, + { + name: "geometry_projjson_epsg_4326_string_code_canonical", + ice: iceberg.GeometryType{}, + geoarrowMetaJSON: `{"crs":{"id":{"authority":"EPSG","code":"4326"}}}`, + }, + { + name: "geometry_projjson_ogc_crs84_canonical", + ice: iceberg.GeometryType{}, + geoarrowMetaJSON: `{"crs":{"id":{"authority":"OGC","code":"CRS84"}}}`, + }, + // Geography with no CRS, spherical edges + { + name: "geography_no_crs_spherical", + ice: geogSRID0, + geoarrowMetaJSON: `{"edges":"spherical"}`, + }, + // Geography with OGC:CRS84 and spherical edges + { + name: "geography_ogc_crs84_spherical_canonical", + ice: iceberg.GeographyType{}, + geoarrowMetaJSON: `{"crs":"OGC:CRS84","edges":"spherical"}`, + }, + // Geography with different edge algorithms + { + name: "geography_ogc_crs84_karney", + ice: geogKarneyDefaultCRS, + geoarrowMetaJSON: `{"crs":"OGC:CRS84","edges":"karney"}`, + }, + { + name: "geography_ogc_crs84_vincenty", + ice: geogVincenty, + geoarrowMetaJSON: `{"crs":"OGC:CRS84","edges":"vincenty"}`, + }, + { + name: "geography_ogc_crs84_andoyer", + ice: geogAndoyer, + geoarrowMetaJSON: `{"crs":"OGC:CRS84","edges":"andoyer"}`, + }, + { + name: "geography_ogc_crs84_thomas", + ice: geogThomas, + geoarrowMetaJSON: `{"crs":"OGC:CRS84","edges":"thomas"}`, + }, + // Geography with custom CRS and edges + { + name: "geography_epsg_4267_karney", + ice: geogEPSG4267Karney, + geoarrowMetaJSON: `{"crs":"EPSG:4267","edges":"karney"}`, + }, + // Geography with PROJJSON CRS + { + name: "geography_projjson_epsg_4267_spherical", + ice: geogEPSG4267, + geoarrowMetaJSON: `{"crs":{"id":{"authority":"EPSG","code":4267}},"edges":"spherical"}`, + }, + } + + for _, tt := range typeCases { + t.Run("iceberg_to_arrow/"+tt.name, func(t *testing.T) { + result, err := table.TypeToArrowType(tt.ice, true, false) + require.NoError(t, err) + assertGeoArrowWKBMetadataJSON(t, result, arrow.BinaryTypes.Binary, tt.geoarrowMetaJSON) + }) + } + + for _, tt := range typeCases { + t.Run("iceberg_to_arrow_round_trip/"+tt.name, func(t *testing.T) { + arrowType, err := table.TypeToArrowType(tt.ice, true, false) + require.NoError(t, err) + + iceType, err := table.ArrowTypeToIceberg(arrowType, false) + require.NoError(t, err) + assert.True(t, tt.ice.Equals(iceType), "expected %s, got %s", tt.ice, iceType) + }) + } + + for _, tt := range append(typeCases, readOnlyCases...) { + t.Run("arrow_to_iceberg/"+tt.name, func(t *testing.T) { + arrowType, err := geoarrow.NewWKBType().Deserialize(arrow.BinaryTypes.Binary, tt.geoarrowMetaJSON) + require.NoError(t, err) + + iceType, err := table.ArrowTypeToIceberg(arrowType, false) + require.NoError(t, err) + assert.True(t, tt.ice.Equals(iceType), "expected %s, got %s", tt.ice, iceType) + }) + } + + t.Run("type/geometry_vs_geography_edge_differentiation", func(t *testing.T) { + geomResult, err := table.TypeToArrowType(iceberg.GeometryType{}, true, false) + require.NoError(t, err) + geogResult, err := table.TypeToArrowType(iceberg.GeographyType{}, true, false) + require.NoError(t, err) + + geomWKB, ok := geomResult.(*geoarrow.WKBType) + require.True(t, ok) + geogWKB, ok := geogResult.(*geoarrow.WKBType) + require.True(t, ok) + + assert.Equal(t, geoarrow.EdgePlanar, geomWKB.Metadata().Edges) + assert.Equal(t, geoarrow.EdgeSpherical, geogWKB.Metadata().Edges) + }) + + t.Run("large_binary_storage", func(t *testing.T) { + result, err := table.TypeToArrowType(iceberg.GeometryType{}, true, true) + require.NoError(t, err) + assertGeoArrowWKB(t, result, arrow.BinaryTypes.LargeBinary, geoarrow.Metadata{CRS: jsonCRS("OGC:CRS84")}) + }) + + t.Run("epsg_4326_behavior", func(t *testing.T) { + geom, err := iceberg.GeometryTypeOf("EPSG:4326") + require.NoError(t, err) + + geomResult, err := table.TypeToArrowType(geom, true, false) + require.NoError(t, err) + + geomWKB, ok := geomResult.(*geoarrow.WKBType) + assert.True(t, ok) + + meta := geomWKB.Metadata() + assert.Equal(t, meta.CRS, jsonCRS("OGC:CRS84")) + assert.Equal(t, meta.Edges, geoarrow.EdgePlanar) + assert.Equal(t, meta.CRSType, geoarrow.CRSTypeAuthorityCode) + + roundTripGeom, err := table.ArrowTypeToIceberg(geomWKB, false) + require.NoError(t, err) + + g, ok := roundTripGeom.(iceberg.GeometryType) + require.True(t, ok) + + assert.Equal(t, g.CRS(), "OGC:CRS84") + assert.Equal(t, g, defaultGeometry) + }) + + t.Run("projjson_error_behavior", func(t *testing.T) { + geom, err := iceberg.GeometryTypeOf("projjson:my-custom-crs") + require.NoError(t, err) + + _, err = table.TypeToArrowType(geom, true, false) + require.Error(t, err) + require.ErrorContains(t, err, "projjson CRS not supported yet") + + arrowType, err := geoarrow.NewWKBType().Deserialize(arrow.BinaryTypes.Binary, + `{"crs_type":"projjson","crs":{"id":{"authority":"OGC", "code":"CRS84"}}}`) + require.NoError(t, err) + + g, err := table.ArrowTypeToIceberg(arrowType, false) + require.NoError(t, err) + assert.Equal(t, g, defaultGeometry) + }) + + t.Run("wkt2:2019_error_behavior", func(t *testing.T) { + arrowType, err := geoarrow.NewWKBType().Deserialize(arrow.BinaryTypes.Binary, + `{"crs_type":"wkt2:2019","crs":"GEOGCRS[\"WGS 84\",DATUM[\"World Geodetic System 1984\",ELLIPSOID[\"WGS 84\",6378137,298.257223563]],CS[ellipsoidal,2],AXIS[\"geodetic latitude (Lat)\",north],AXIS[\"geodetic longitude (Lon)\",east],ID[\"EPSG\",4326]]"}`) + require.NoError(t, err) + + _, err = table.ArrowTypeToIceberg(arrowType, false) + require.Error(t, err) + require.ErrorContains(t, err, "CRS type wkt2:2019 not supported") + + arrowTypeWithoutCRSType, err := geoarrow.NewWKBType().Deserialize(arrow.BinaryTypes.Binary, + `{"crs":"GEOGCRS[\"WGS 84\",DATUM[\"World Geodetic System 1984\",ELLIPSOID[\"WGS 84\",6378137,298.257223563]],CS[ellipsoidal,2],AXIS[\"geodetic latitude (Lat)\",north],AXIS[\"geodetic longitude (Lon)\",east],ID[\"EPSG\",4326]]"}`) + require.NoError(t, err) + + _, err = table.ArrowTypeToIceberg(arrowTypeWithoutCRSType, false) + require.Error(t, err) + require.ErrorContains(t, err, "crs length too long") + }) + + t.Run("schema", func(t *testing.T) { + geomList, err := iceberg.GeometryTypeOf("srid:4326") + require.NoError(t, err) + + iceSchema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "point", Type: iceberg.GeometryType{}, Required: false}, + iceberg.NestedField{ID: 2, Name: "loc", Type: geomSRID, Required: true}, + iceberg.NestedField{ID: 3, Name: "area", Type: iceberg.GeographyType{}, Required: false}, + iceberg.NestedField{ID: 4, Name: "region", Type: geogKarney, Required: false}, + iceberg.NestedField{ + ID: 5, + Name: "locations", + Type: &iceberg.ListType{ + ElementID: 6, + Element: geomList, + ElementRequired: true, + }, + Required: true, + }, + ) + + arrowSc, err := table.SchemaToArrowSchema(iceSchema, nil, true, false) + require.NoError(t, err) + require.Equal(t, 5, arrowSc.NumFields()) + + wantFields := []struct { + name string + nullable bool + fieldID string + wantMeta geoarrow.Metadata + }{ + {"point", true, "1", geoarrow.Metadata{CRS: jsonCRS("OGC:CRS84")}}, + {"loc", false, "2", geoarrow.Metadata{CRS: jsonCRS("4326"), CRSType: geoarrow.CRSTypeSRID}}, + {"area", true, "3", geoarrow.Metadata{CRS: jsonCRS("OGC:CRS84"), Edges: geoarrow.EdgeSpherical}}, + {"region", true, "4", geoarrow.Metadata{CRS: jsonCRS("4269"), CRSType: geoarrow.CRSTypeSRID, Edges: geoarrow.EdgeKarney}}, + } + + for i, want := range wantFields { + field := arrowSc.Field(i) + assert.Equal(t, want.name, field.Name) + assert.Equal(t, want.nullable, field.Nullable) + fieldID, ok := field.Metadata.GetValue(table.ArrowParquetFieldIDKey) + require.True(t, ok) + assert.Equal(t, want.fieldID, fieldID) + assertGeoArrowWKB(t, field.Type, arrow.BinaryTypes.Binary, want.wantMeta) + } + + listField := arrowSc.Field(4) + assert.Equal(t, "locations", listField.Name) + assert.False(t, listField.Nullable) + listFieldID, ok := listField.Metadata.GetValue(table.ArrowParquetFieldIDKey) + require.True(t, ok) + assert.Equal(t, "5", listFieldID) + listType, ok := listField.Type.(*arrow.ListType) + require.True(t, ok, "expected list type, got %T", listField.Type) + + elemField := listType.ElemField() + assert.Equal(t, "element", elemField.Name) + assert.False(t, elemField.Nullable) + elemFieldID, ok := elemField.Metadata.GetValue(table.ArrowParquetFieldIDKey) + require.True(t, ok) + assert.Equal(t, "6", elemFieldID) + assertGeoArrowWKB(t, elemField.Type, arrow.BinaryTypes.Binary, geoarrow.Metadata{ + CRS: jsonCRS("4326"), + CRSType: geoarrow.CRSTypeSRID, + }) + }) +} + func TestVariantArrayBuilderLargeValues(t *testing.T) { mem := memory.NewCheckedAllocator(memory.DefaultAllocator) defer mem.AssertSize(t, 0) diff --git a/table/internal/geo_codec_test.go b/table/internal/geo_codec_test.go new file mode 100644 index 000000000..e8d430678 --- /dev/null +++ b/table/internal/geo_codec_test.go @@ -0,0 +1,101 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package internal_test + +import ( + "errors" + "fmt" + "testing" + + "github.com/geoarrow/geoarrow-go" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/twpayne/go-geom/encoding/wkb" + "github.com/twpayne/go-geom/encoding/wkt" +) + +// wktToWKB is a helper which converts Well Known Text (WKT) to Well Known Bytes (WKB). +// Note that return bytes are little endian. +func wktToWKB(s string) (geoarrow.WKBBytes, error) { + geometry, err := wkt.Unmarshal(s) + if err != nil { + return nil, fmt.Errorf("parse WKT: %w", err) + } + + wkbBytes, err := wkb.Marshal(geometry, wkb.NDR) // little endian + if err != nil { + return nil, fmt.Errorf("parse WKT: %w", err) + } + + return geoarrow.WKBBytes(wkbBytes), nil +} + +func TestWKTToWKB(t *testing.T) { + tests := []struct { + name string + wkt string + wantWKB string + wantErr error + }{ + { + name: "point", + wkt: "POINT (30 10)", + wantWKB: "01010000000000000000003e400000000000002440", + }, + { + name: "linestring", + wkt: "LINESTRING (30 10, 10 30, 40 40)", + wantWKB: "0102000000030000000000000000003e40000000000000244000000000000024400000000000003e4000000000000044400000000000004440", + }, + { + name: "polygon", + wkt: "POLYGON ((30 10, 40 40, 20 40, 10 20, 30 10))", + wantWKB: "010300000001000000050000000000000000003e4000000000000024400000000000004440000000000000444000000000000034400000000000004440000000000000244000000000000034400000000000003e400000000000002440", + }, + { + name: "multipoint", + wkt: "MULTIPOINT ((10 40), (40 30), (20 20), (30 20))", + wantWKB: "010400000004000000010100000000000000000024400000000000004440010100000000000000000044400000000000003e4001010000000000000000003440000000000000344001010000000000000000003e400000000000003440", + }, + { + name: "geometry_collection", + wkt: "GEOMETRYCOLLECTION (POINT (4 6), LINESTRING (4 6, 7 10))", + wantWKB: "010700000002000000010100000000000000000010400000000000001840010200000002000000000000000000104000000000000018400000000000001c400000000000002440", + }, + { + name: "unknown wkt", + wkt: "POINT (30 10", + wantErr: errors.New("parse WKT: syntax error: unexpected $end, expecting ')'"), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := wktToWKB(tt.wkt) + if tt.wantErr != nil { + require.Error(t, err) + assert.ErrorContains(t, err, tt.wantErr.Error()) + + return + } + + require.NoError(t, err) + assert.Equal(t, tt.wantWKB, got.String()) + }) + } +} diff --git a/table/internal/parquet_files_test.go b/table/internal/parquet_files_test.go index 4e9d381a1..d7c0d3011 100644 --- a/table/internal/parquet_files_test.go +++ b/table/internal/parquet_files_test.go @@ -40,6 +40,7 @@ import ( internal2 "github.com/apache/iceberg-go/internal" "github.com/apache/iceberg-go/table" "github.com/apache/iceberg-go/table/internal" + "github.com/geoarrow/geoarrow-go" "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -603,6 +604,101 @@ func TestDecimalPhysicalTypes(t *testing.T) { } } +func TestParquetGeoArrowExtensionMetadataRoundTrip(t *testing.T) { + geomType, err := iceberg.GeometryTypeOf("srid:4326") + require.NoError(t, err) + + geogType, err := iceberg.GeographyTypeOf("srid:4326", "vincenty") + require.NoError(t, err) + + iceSchema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int32, Required: false}, + iceberg.NestedField{ID: 2, Name: "geom", Type: geomType, Required: false}, + iceberg.NestedField{ID: 3, Name: "geog", Type: geogType, Required: false}, + ) + + arrowSchema, err := table.SchemaToArrowSchema(iceSchema, nil, true, false) + require.NoError(t, err) + + geomWKB, err := wktToWKB("POINT (30 10)") + require.NoError(t, err) + geogWKB, err := wktToWKB("POINT (20 5)") + require.NoError(t, err) + + rec, _, err := array.RecordFromJSON(memory.DefaultAllocator, arrowSchema, strings.NewReader(`[ + {"id": 1, "geom": "`+geomWKB.String()+`", "geog": "`+geogWKB.String()+`"}, + {"id": null, "geom": null, "geog": "`+geogWKB.String()+`"} + ]`)) + require.NoError(t, err) + defer rec.Release() + + arrProps := pqarrow.NewArrowWriterProperties( + pqarrow.WithAllocator(memory.DefaultAllocator), + pqarrow.WithStoreSchema(), + ) + var buf bytes.Buffer + wr, err := pqarrow.NewFileWriter(arrowSchema, &buf, + parquet.NewWriterProperties(parquet.WithStats(true)), + arrProps) + require.NoError(t, err) + require.NoError(t, wr.Write(rec)) + require.NoError(t, wr.Close()) + + pqRdr, err := file.NewParquetReader(bytes.NewReader(buf.Bytes())) + require.NoError(t, err) + defer pqRdr.Close() + + arrRdr, err := pqarrow.NewFileReader(pqRdr, pqarrow.ArrowReadProperties{}, memory.DefaultAllocator) + require.NoError(t, err) + fr := internal.WrapParquetFileReader(arrRdr) + + tbl, err := fr.ReadTable(context.Background()) + require.NoError(t, err) + defer tbl.Release() + + roundTripSchema := tbl.Schema() + + origGeomType := arrowSchema.Field(1).Type.(*geoarrow.WKBType) + origGeogType := arrowSchema.Field(2).Type.(*geoarrow.WKBType) + rtGeomField := roundTripSchema.Field(1) + rtGeogField := roundTripSchema.Field(2) + + rtGeomExt := rtGeomField.Type.(arrow.ExtensionType) + rtGeogExt := rtGeogField.Type.(arrow.ExtensionType) + rtGeomType := rtGeomExt.(*geoarrow.WKBType) + rtGeogType := rtGeogExt.(*geoarrow.WKBType) + + assert.Equal(t, origGeomType.ExtensionName(), rtGeomExt.ExtensionName()) + assert.Equal(t, origGeogType.ExtensionName(), rtGeogExt.ExtensionName()) + assert.True(t, arrow.TypeEqual(origGeomType.StorageType(), rtGeomType.StorageType())) + assert.True(t, arrow.TypeEqual(origGeogType.StorageType(), rtGeogType.StorageType())) + assert.Equal(t, origGeomType.Metadata(), rtGeomType.Metadata()) + assert.Equal(t, origGeogType.Metadata(), rtGeogType.Metadata()) + assert.Equal(t, geoarrow.CRSTypeSRID, rtGeomType.Metadata().CRSType) + assert.Equal(t, geoarrow.EdgePlanar, rtGeomType.Metadata().Edges) + assert.Equal(t, geoarrow.CRSTypeSRID, rtGeogType.Metadata().CRSType) + assert.Equal(t, geoarrow.EdgeVincenty, rtGeogType.Metadata().Edges) + + rr, err := fr.GetRecords(context.Background(), nil, nil) + require.NoError(t, err) + defer rr.Release() + require.True(t, rr.Next()) + + batch := rr.RecordBatch() + defer batch.Release() + require.Equal(t, int64(2), batch.NumRows()) + + geomArr, ok := batch.Column(1).(*geoarrow.WKBArray) + assert.True(t, ok) + assert.Equal(t, geomWKB, geomArr.Value(0)) + assert.Nil(t, geomArr.Value(1)) + + geogArr, ok := batch.Column(2).(*geoarrow.WKBArray) + assert.True(t, ok) + assert.Equal(t, geogWKB, geogArr.Value(0)) + assert.Equal(t, geogWKB, geogArr.Value(1)) +} + func TestWriteDataFileErrOnClose(t *testing.T) { ctx := context.Background() fm := internal.GetFileFormat(iceberg.ParquetFile)