diff --git a/arrow/extensions/extensions.go b/arrow/extensions/extensions.go index 04566c75..3859bb3d 100644 --- a/arrow/extensions/extensions.go +++ b/arrow/extensions/extensions.go @@ -26,6 +26,8 @@ var canonicalExtensionTypes = []arrow.ExtensionType{ &OpaqueType{}, &JSONType{}, &VariantType{}, + // GeoArrow extension types + NewPointType(), } func init() { diff --git a/arrow/extensions/geoarrow.go b/arrow/extensions/geoarrow.go new file mode 100644 index 00000000..ee3c0909 --- /dev/null +++ b/arrow/extensions/geoarrow.go @@ -0,0 +1,186 @@ +// 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 extensions + +import ( + "encoding/json" + "fmt" + + "github.com/apache/arrow-go/v18/arrow" +) + +// CoordinateDimension represents the dimensionality of coordinates +type CoordinateDimension int + +const ( + DimensionXY CoordinateDimension = iota + DimensionXYZ + DimensionXYM + DimensionXYZM +) + +// String returns the string representation of the coordinate dimension +func (d CoordinateDimension) String() string { + switch d { + case DimensionXY: + return "xy" + case DimensionXYZ: + return "xyz" + case DimensionXYM: + return "xym" + case DimensionXYZM: + return "xyzm" + default: + return "unknown" + } +} + +// Size returns the number of coordinate values per point +func (d CoordinateDimension) Size() int { + switch d { + case DimensionXY: + return 2 + case DimensionXYZ, DimensionXYM: + return 3 + case DimensionXYZM: + return 4 + default: + return 2 + } +} + +// GeometryEncoding represents the encoding method for geometry data +type GeometryEncoding int + +const ( + EncodingGeoArrow GeometryEncoding = iota +) + +// String returns the string representation of the encoding +func (e GeometryEncoding) String() string { + switch e { + case EncodingGeoArrow: + return "geoarrow" + default: + return "unknown" + } +} + +// EdgeType represents the edge interpretation for geometry data +type EdgeType string + +const ( + EdgePlanar EdgeType = "planar" + EdgeSpherical EdgeType = "spherical" +) + +// String returns the string representation of the edge type +func (e EdgeType) String() string { + return string(e) +} + +// CoordType represents the coordinate layout type +type CoordType string + +const ( + CoordSeparate CoordType = "separate" + CoordInterleaved CoordType = "interleaved" +) + +// String returns the string representation of the coordinate type +func (c CoordType) String() string { + return string(c) +} + +// GeometryMetadata contains metadata for GeoArrow geometry types +type GeometryMetadata struct { + // Encoding specifies the geometry encoding format + Encoding GeometryEncoding `json:"encoding,omitempty"` + + // CRS contains PROJJSON coordinate reference system information + CRS json.RawMessage `json:"crs,omitempty"` + + // Edges specifies the edge interpretation for the geometry + Edges EdgeType `json:"edges,omitempty"` + + // CoordType specifies the coordinate layout (separate vs interleaved) + CoordType CoordType `json:"coord_type,omitempty"` +} + +// NewGeometryMetadata creates a new GeometryMetadata with default values +func NewGeometryMetadata() *GeometryMetadata { + return &GeometryMetadata{ + Encoding: EncodingGeoArrow, + Edges: EdgePlanar, + CoordType: CoordSeparate, + } +} + +// Serialize serializes the metadata to a JSON string +func (gm *GeometryMetadata) Serialize() (string, error) { + if gm == nil { + return "", nil + } + data, err := json.Marshal(gm) + if err != nil { + return "", fmt.Errorf("failed to serialize geometry metadata: %w", err) + } + return string(data), nil +} + +// DeserializeGeometryMetadata deserializes geometry metadata from a JSON string +func DeserializeGeometryMetadata(data string) (*GeometryMetadata, error) { + if data == "" { + return NewGeometryMetadata(), nil + } + + var gm GeometryMetadata + if err := json.Unmarshal([]byte(data), &gm); err != nil { + return nil, fmt.Errorf("failed to deserialize geometry metadata: %w", err) + } + + return &gm, nil +} + +// createCoordinateType creates an Arrow data type for coordinates based on dimension +func createCoordinateType(dim CoordinateDimension) arrow.DataType { + var fieldNames []string + + switch dim { + case DimensionXY: + fieldNames = []string{"x", "y"} + case DimensionXYZ: + fieldNames = []string{"x", "y", "z"} + case DimensionXYM: + fieldNames = []string{"x", "y", "m"} + case DimensionXYZM: + fieldNames = []string{"x", "y", "z", "m"} + default: + fieldNames = []string{"x", "y"} + } + + fields := make([]arrow.Field, len(fieldNames)) + for i, name := range fieldNames { + fields[i] = arrow.Field{ + Name: name, + Type: arrow.PrimitiveTypes.Float64, + Nullable: false, + } + } + + return arrow.StructOf(fields...) +} diff --git a/arrow/extensions/geoarrow_point.go b/arrow/extensions/geoarrow_point.go new file mode 100644 index 00000000..35e47f4a --- /dev/null +++ b/arrow/extensions/geoarrow_point.go @@ -0,0 +1,407 @@ +// 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 extensions + +import ( + "encoding/json" + "fmt" + "math" + "reflect" + "strings" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +// Point represents a point geometry with coordinates +type Point struct { + X, Y, Z, M float64 + Dimension CoordinateDimension +} + +// NewPoint creates a new 2D point +func NewPoint(x, y float64) Point { + return Point{X: x, Y: y, Dimension: DimensionXY} +} + +// NewPointZ creates a new 3D point with Z coordinate +func NewPointZ(x, y, z float64) Point { + return Point{X: x, Y: y, Z: z, Dimension: DimensionXYZ} +} + +// NewPointM creates a new 2D point with M coordinate +func NewPointM(x, y, m float64) Point { + return Point{X: x, Y: y, M: m, Dimension: DimensionXYM} +} + +// NewPointZM creates a new 3D point with Z and M coordinates +func NewPointZM(x, y, z, m float64) Point { + return Point{X: x, Y: y, Z: z, M: m, Dimension: DimensionXYZM} +} + +// NewEmptyPoint creates a new empty point +func NewEmptyPoint() Point { + return Point{X: math.NaN(), Y: math.NaN(), Z: math.NaN(), M: math.NaN(), Dimension: DimensionXY} +} + +// IsEmpty returns true if this is an empty point (all coordinates are NaN or zero with no dimension) +func (p Point) IsEmpty() bool { + return math.IsNaN(p.X) && math.IsNaN(p.Y) && math.IsNaN(p.Z) && math.IsNaN(p.M) +} + +// String returns a string representation of the point +func (p Point) String() string { + if p.IsEmpty() { + return "POINT EMPTY" + } + + switch p.Dimension { + case DimensionXY: + return fmt.Sprintf("POINT(%.6f %.6f)", p.X, p.Y) + case DimensionXYZ: + return fmt.Sprintf("POINT Z(%.6f %.6f %.6f)", p.X, p.Y, p.Z) + case DimensionXYM: + return fmt.Sprintf("POINT M(%.6f %.6f %.6f)", p.X, p.Y, p.M) + case DimensionXYZM: + return fmt.Sprintf("POINT ZM(%.6f %.6f %.6f %.6f)", p.X, p.Y, p.Z, p.M) + default: + return fmt.Sprintf("POINT(%.6f %.6f)", p.X, p.Y) + } +} + +// PointType is the extension type for Point geometries +type PointType struct { + arrow.ExtensionBase + metadata *GeometryMetadata + dimension CoordinateDimension +} + +// NewPointType creates a new Point extension type for 2D points +func NewPointType() *PointType { + return NewPointTypeWithDimension(DimensionXY) +} + +// NewPointTypeWithDimension creates a new Point extension type with specified dimension +func NewPointTypeWithDimension(dim CoordinateDimension) *PointType { + metadata := NewGeometryMetadata() + coordType := createCoordinateType(dim) + + return &PointType{ + ExtensionBase: arrow.ExtensionBase{Storage: coordType}, + metadata: metadata, + dimension: dim, + } +} + +// NewPointTypeWithMetadata creates a new Point extension type with custom metadata +func NewPointTypeWithMetadata(dim CoordinateDimension, metadata *GeometryMetadata) *PointType { + if metadata == nil { + metadata = NewGeometryMetadata() + } + coordType := createCoordinateType(dim) + + return &PointType{ + ExtensionBase: arrow.ExtensionBase{Storage: coordType}, + metadata: metadata, + dimension: dim, + } +} + +// ArrayType returns the array type for Point arrays +func (*PointType) ArrayType() reflect.Type { + return reflect.TypeOf(PointArray{}) +} + +// ExtensionName returns the name of the extension +func (*PointType) ExtensionName() string { + return "geoarrow.point" +} + +// String returns a string representation of the type +func (p *PointType) String() string { + return fmt.Sprintf("extension<%s[%s]>", p.ExtensionName(), p.dimension.String()) +} + +// ExtensionEquals checks if two extension types are equal +func (p *PointType) ExtensionEquals(other arrow.ExtensionType) bool { + if p.ExtensionName() != other.ExtensionName() { + return false + } + if otherPoint, ok := other.(*PointType); ok { + return p.dimension == otherPoint.dimension + } + return arrow.TypeEqual(p.Storage, other.StorageType()) +} + +// Serialize serializes the extension type metadata +func (p *PointType) Serialize() string { + if p.metadata == nil { + return "" + } + serialized, _ := p.metadata.Serialize() + return serialized +} + +// Deserialize deserializes the extension type metadata +func (*PointType) Deserialize(storageType arrow.DataType, data string) (arrow.ExtensionType, error) { + metadata, err := DeserializeGeometryMetadata(data) + if err != nil { + return nil, err + } + + // Determine dimension from storage type + dim := DimensionXY + if structType, ok := storageType.(*arrow.StructType); ok { + numFields := structType.NumFields() + switch numFields { + case 2: + dim = DimensionXY + case 3: + // Check if it's XYZ or XYM by looking at field names + if structType.Field(2).Name == "z" { + dim = DimensionXYZ + } else { + dim = DimensionXYM + } + case 4: + dim = DimensionXYZM + default: + dim = DimensionXY + } + } + + return &PointType{ + ExtensionBase: arrow.ExtensionBase{Storage: storageType}, + metadata: metadata, + dimension: dim, + }, nil +} + +// NewBuilder creates a new array builder for this type +func (p *PointType) NewBuilder(mem memory.Allocator) array.Builder { + return NewPointBuilder(mem, p) +} + +// Metadata returns the geometry metadata +func (p *PointType) Metadata() *GeometryMetadata { + return p.metadata +} + +// Dimension returns the coordinate dimension +func (p *PointType) Dimension() CoordinateDimension { + return p.dimension +} + +// PointArray represents an array of Point geometries +type PointArray struct { + array.ExtensionArrayBase +} + +// String returns a string representation of the array +func (p *PointArray) String() string { + o := new(strings.Builder) + o.WriteString("PointArray[") + for i := 0; i < p.Len(); i++ { + if i > 0 { + o.WriteString(" ") + } + if p.IsNull(i) { + o.WriteString(array.NullValueStr) + } else { + point := p.Value(i) + o.WriteString(point.String()) + } + } + o.WriteString("]") + return o.String() +} + +// Value returns the Point at the given index +func (p *PointArray) Value(i int) Point { + pointType := p.ExtensionType().(*PointType) + structArray := p.Storage().(*array.Struct) + + point := Point{Dimension: pointType.dimension} + + if p.IsNull(i) { + return point + } + + // Get X coordinate + xArray := structArray.Field(0).(*array.Float64) + point.X = xArray.Value(i) + + // Get Y coordinate + yArray := structArray.Field(1).(*array.Float64) + point.Y = yArray.Value(i) + + // Get Z coordinate if present + if pointType.dimension == DimensionXYZ || pointType.dimension == DimensionXYZM { + zArray := structArray.Field(2).(*array.Float64) + point.Z = zArray.Value(i) + } + + // Get M coordinate if present + switch pointType.dimension { + case DimensionXYM: + mArray := structArray.Field(2).(*array.Float64) + point.M = mArray.Value(i) + case DimensionXYZM: + mArray := structArray.Field(3).(*array.Float64) + point.M = mArray.Value(i) + } + + return point +} + +// Values returns all Point values as a slice +func (p *PointArray) Values() []Point { + values := make([]Point, p.Len()) + for i := range values { + values[i] = p.Value(i) + } + return values +} + +// ValueStr returns a string representation of the value at index i +func (p *PointArray) ValueStr(i int) string { + if p.IsNull(i) { + return array.NullValueStr + } + return p.Value(i).String() +} + +// GetOneForMarshal returns the value at index i for JSON marshaling +func (p *PointArray) GetOneForMarshal(i int) any { + if p.IsNull(i) { + return nil + } + point := p.Value(i) + switch point.Dimension { + case DimensionXY: + return []float64{point.X, point.Y} + case DimensionXYZ: + return []float64{point.X, point.Y, point.Z} + case DimensionXYM: + return []float64{point.X, point.Y, point.M} + case DimensionXYZM: + return []float64{point.X, point.Y, point.Z, point.M} + default: + // Should never happen but defensive programming + panic(fmt.Sprintf("unknown coordinate dimension: %v", point.Dimension)) + } +} + +// MarshalJSON implements json.Marshaler +func (p *PointArray) MarshalJSON() ([]byte, error) { + vals := make([]any, p.Len()) + for i := range vals { + vals[i] = p.GetOneForMarshal(i) + } + return json.Marshal(vals) +} + +// PointBuilder is an array builder for Point geometries +type PointBuilder struct { + *array.ExtensionBuilder +} + +// NewPointBuilder creates a new Point array builder +func NewPointBuilder(mem memory.Allocator, dtype *PointType) *PointBuilder { + return &PointBuilder{ + ExtensionBuilder: array.NewExtensionBuilder(mem, dtype), + } +} + +// Append appends a Point to the array +func (b *PointBuilder) Append(point Point) { + pointType := b.Type().(*PointType) + structBuilder := b.Builder.(*array.StructBuilder) + structBuilder.Append(true) + + // X coordinate + xBuilder := structBuilder.FieldBuilder(0).(*array.Float64Builder) + xBuilder.Append(point.X) + + // Y coordinate + yBuilder := structBuilder.FieldBuilder(1).(*array.Float64Builder) + yBuilder.Append(point.Y) + + // Z coordinate if present + if pointType.dimension == DimensionXYZ || pointType.dimension == DimensionXYZM { + zBuilder := structBuilder.FieldBuilder(2).(*array.Float64Builder) + zBuilder.Append(point.Z) + } + + // M coordinate if present + switch pointType.dimension { + case DimensionXYM: + mBuilder := structBuilder.FieldBuilder(2).(*array.Float64Builder) + mBuilder.Append(point.M) + case DimensionXYZM: + mBuilder := structBuilder.FieldBuilder(3).(*array.Float64Builder) + mBuilder.Append(point.M) + } +} + +// AppendXY appends a 2D point to the array +func (b *PointBuilder) AppendXY(x, y float64) { + b.Append(NewPoint(x, y)) +} + +// AppendXYZ appends a 3D point to the array +func (b *PointBuilder) AppendXYZ(x, y, z float64) { + b.Append(NewPointZ(x, y, z)) +} + +// AppendXYM appends a 2D point with M coordinate to the array +func (b *PointBuilder) AppendXYM(x, y, m float64) { + b.Append(NewPointM(x, y, m)) +} + +// AppendXYZM appends a 3D point with M coordinate to the array +func (b *PointBuilder) AppendXYZM(x, y, z, m float64) { + b.Append(NewPointZM(x, y, z, m)) +} + +// AppendNull appends a null value to the array +func (b *PointBuilder) AppendNull() { + b.ExtensionBuilder.Builder.(*array.StructBuilder).AppendNull() +} + +// AppendValues appends multiple Point values to the array +func (b *PointBuilder) AppendValues(v []Point, valid []bool) { + if len(v) != len(valid) && len(valid) != 0 { + panic("len(v) != len(valid) && len(valid) != 0") + } + + for i, point := range v { + if len(valid) > 0 && !valid[i] { + b.AppendNull() + } else { + b.Append(point) + } + } +} + +// NewArray creates a new array from the builder +func (b *PointBuilder) NewArray() arrow.Array { + storage := b.Builder.(*array.StructBuilder).NewArray() + defer storage.Release() + return array.NewExtensionArrayWithStorage(b.Type().(*PointType), storage) +} diff --git a/arrow/extensions/geoarrow_point_test.go b/arrow/extensions/geoarrow_point_test.go new file mode 100644 index 00000000..79744a44 --- /dev/null +++ b/arrow/extensions/geoarrow_point_test.go @@ -0,0 +1,435 @@ +// 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 extensions_test + +import ( + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/extensions" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestPointTypeBasics(t *testing.T) { + // Test extensions.NewPointType (default XY) + pointType := extensions.NewPointType() + assert.Equal(t, "geoarrow.point", pointType.ExtensionName()) + assert.Equal(t, extensions.DimensionXY, pointType.Dimension()) + + // Test extensions.NewPointTypeWithDimension for different dimensions + pointTypeXYZ := extensions.NewPointTypeWithDimension(extensions.DimensionXYZ) + pointTypeXYM := extensions.NewPointTypeWithDimension(extensions.DimensionXYM) + pointTypeXYZM := extensions.NewPointTypeWithDimension(extensions.DimensionXYZM) + + assert.Equal(t, "geoarrow.point", pointType.ExtensionName()) + assert.Equal(t, "geoarrow.point", pointTypeXYZ.ExtensionName()) + assert.Equal(t, "geoarrow.point", pointTypeXYM.ExtensionName()) + assert.Equal(t, "geoarrow.point", pointTypeXYZM.ExtensionName()) + + assert.True(t, pointType.ExtensionEquals(pointType)) + assert.True(t, pointTypeXYZ.ExtensionEquals(pointTypeXYZ)) + assert.True(t, pointTypeXYM.ExtensionEquals(pointTypeXYM)) + assert.True(t, pointTypeXYZM.ExtensionEquals(pointTypeXYZM)) + + // Different dimensions should not be equal + assert.False(t, pointType.ExtensionEquals(pointTypeXYZ)) + assert.False(t, pointTypeXYZ.ExtensionEquals(pointTypeXYM)) + assert.False(t, pointTypeXYM.ExtensionEquals(pointTypeXYZM)) + + // Test string representations + assert.Equal(t, "extension", pointType.String()) + assert.Equal(t, "extension", pointTypeXYZ.String()) + assert.Equal(t, "extension", pointTypeXYM.String()) + assert.Equal(t, "extension", pointTypeXYZM.String()) +} + +func TestGeometryMetadata(t *testing.T) { + t.Run("default_metadata", func(t *testing.T) { + metadata := extensions.NewGeometryMetadata() + assert.Equal(t, extensions.EncodingGeoArrow, metadata.Encoding) + assert.Equal(t, extensions.EdgePlanar, metadata.Edges) + assert.Equal(t, extensions.CoordSeparate, metadata.CoordType) + assert.Empty(t, metadata.CRS) + }) + + t.Run("edge_types", func(t *testing.T) { + assert.Equal(t, "planar", extensions.EdgePlanar.String()) + assert.Equal(t, "spherical", extensions.EdgeSpherical.String()) + }) + + t.Run("coord_types", func(t *testing.T) { + assert.Equal(t, "separate", extensions.CoordSeparate.String()) + assert.Equal(t, "interleaved", extensions.CoordInterleaved.String()) + }) + + t.Run("custom_metadata", func(t *testing.T) { + metadata := extensions.NewGeometryMetadata() + metadata.Edges = extensions.EdgeSpherical + metadata.CoordType = extensions.CoordInterleaved + metadata.CRS = []byte(`{"type": "name", "properties": {"name": "EPSG:4326"}}`) + + // Test serialization + serialized, err := metadata.Serialize() + require.NoError(t, err) + assert.Contains(t, serialized, "spherical") + assert.Contains(t, serialized, "interleaved") + assert.Contains(t, serialized, "EPSG:4326") + + // Test deserialization + deserialized, err := extensions.DeserializeGeometryMetadata(serialized) + require.NoError(t, err) + assert.Equal(t, extensions.EdgeSpherical, deserialized.Edges) + assert.Equal(t, extensions.CoordInterleaved, deserialized.CoordType) + assert.JSONEq(t, string(metadata.CRS), string(deserialized.CRS)) + }) + + t.Run("point_type_with_custom_metadata", func(t *testing.T) { + metadata := extensions.NewGeometryMetadata() + metadata.Edges = extensions.EdgeSpherical + metadata.CoordType = extensions.CoordInterleaved + + pointType := extensions.NewPointTypeWithMetadata(extensions.DimensionXYZ, metadata) + assert.Equal(t, extensions.DimensionXYZ, pointType.Dimension()) + + retrievedMetadata := pointType.Metadata() + assert.Equal(t, extensions.EdgeSpherical, retrievedMetadata.Edges) + assert.Equal(t, extensions.CoordInterleaved, retrievedMetadata.CoordType) + }) +} + +func TestPointExtensionRegistration(t *testing.T) { + // Test that Point type is registered with Arrow + extType := arrow.GetExtensionType("geoarrow.point") + require.NotNil(t, extType, "Point extension type should be registered") + assert.Equal(t, "geoarrow.point", extType.ExtensionName()) +} + +func TestPointBuilderExtensive(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + t.Run("basic_building", func(t *testing.T) { + builder := extensions.NewPointBuilder(mem, extensions.NewPointType()) + defer builder.Release() + + // Test all append methods + builder.AppendXY(1.0, 2.0) + builder.AppendNull() + builder.Append(extensions.NewPoint(3.0, 4.0)) + + // Test AppendValues + points := []extensions.Point{extensions.NewPoint(5.0, 6.0), extensions.NewPoint(7.0, 8.0)} + valid := []bool{true, true} + builder.AppendValues(points, valid) + + arr := builder.NewArray() + defer arr.Release() + + pointArray := arr.(*extensions.PointArray) + assert.Equal(t, 5, pointArray.Len()) + assert.Equal(t, 1, pointArray.NullN()) + + // Test Values() method + allPoints := pointArray.Values() + assert.Equal(t, 5, len(allPoints)) + assert.Equal(t, 1.0, allPoints[0].X) + assert.Equal(t, 2.0, allPoints[0].Y) + + // Test null handling + assert.True(t, pointArray.IsNull(1)) + assert.False(t, pointArray.IsNull(0)) + }) + + t.Run("multidimensional_building", func(t *testing.T) { + // Test XYZ + builderXYZ := extensions.NewPointBuilder(mem, extensions.NewPointTypeWithDimension(extensions.DimensionXYZ)) + defer builderXYZ.Release() + builderXYZ.AppendXYZ(1.0, 2.0, 3.0) + builderXYZ.Append(extensions.NewPointZ(4.0, 5.0, 6.0)) + + arrXYZ := builderXYZ.NewArray() + defer arrXYZ.Release() + pointArrayXYZ := arrXYZ.(*extensions.PointArray) + + point := pointArrayXYZ.Value(0) + assert.Equal(t, extensions.DimensionXYZ, point.Dimension) + assert.Equal(t, 3.0, point.Z) + + // Test XYM + builderXYM := extensions.NewPointBuilder(mem, extensions.NewPointTypeWithDimension(extensions.DimensionXYM)) + defer builderXYM.Release() + builderXYM.AppendXYM(1.0, 2.0, 100.0) + builderXYM.Append(extensions.NewPointM(4.0, 5.0, 200.0)) + + arrXYM := builderXYM.NewArray() + defer arrXYM.Release() + pointArrayXYM := arrXYM.(*extensions.PointArray) + + pointM := pointArrayXYM.Value(0) + assert.Equal(t, extensions.DimensionXYM, pointM.Dimension) + assert.Equal(t, 100.0, pointM.M) + + // Test XYZM + builderXYZM := extensions.NewPointBuilder(mem, extensions.NewPointTypeWithDimension(extensions.DimensionXYZM)) + defer builderXYZM.Release() + builderXYZM.AppendXYZM(1.0, 2.0, 3.0, 100.0) + + arrXYZM := builderXYZM.NewArray() + defer arrXYZM.Release() + pointArrayXYZM := arrXYZM.(*extensions.PointArray) + + pointZM := pointArrayXYZM.Value(0) + assert.Equal(t, extensions.DimensionXYZM, pointZM.Dimension) + assert.Equal(t, 3.0, pointZM.Z) + assert.Equal(t, 100.0, pointZM.M) + }) +} + +func TestPointTypeCreateFromArray(t *testing.T) { + pointType := extensions.NewPointType() + + // Use the storage type from the point type + structType := pointType.StorageType().(*arrow.StructType) + structBuilder := array.NewStructBuilder(memory.DefaultAllocator, structType) + defer structBuilder.Release() + + // Add some points manually to struct + structBuilder.Append(true) + structBuilder.FieldBuilder(0).(*array.Float64Builder).Append(1.5) + structBuilder.FieldBuilder(1).(*array.Float64Builder).Append(2.5) + + structBuilder.AppendNull() + + structBuilder.Append(true) + structBuilder.FieldBuilder(0).(*array.Float64Builder).Append(3.5) + structBuilder.FieldBuilder(1).(*array.Float64Builder).Append(4.5) + + storage := structBuilder.NewArray() + defer storage.Release() + + arr := array.NewExtensionArrayWithStorage(pointType, storage) + defer arr.Release() + + assert.Equal(t, 3, arr.Len()) + assert.Equal(t, 1, arr.NullN()) + + pointArray, ok := arr.(*extensions.PointArray) + require.True(t, ok) + + point0 := pointArray.Value(0) + assert.Equal(t, 1.5, point0.X) + assert.Equal(t, 2.5, point0.Y) + + // Check null + assert.True(t, pointArray.IsNull(1)) + + point2 := pointArray.Value(2) + assert.Equal(t, 3.5, point2.X) + assert.Equal(t, 4.5, point2.Y) +} + +func TestPointRecordBuilder(t *testing.T) { + schema := arrow.NewSchema([]arrow.Field{ + {Name: "point", Type: extensions.NewPointType()}, + }, nil) + builder := array.NewRecordBuilder(memory.DefaultAllocator, schema) + defer builder.Release() + + pointBuilder := builder.Field(0).(*extensions.PointBuilder) + pointBuilder.AppendXY(1.0, 2.0) + pointBuilder.AppendNull() + pointBuilder.AppendXY(3.0, 4.0) + + record := builder.NewRecordBatch() + defer record.Release() + + // Test JSON marshaling of record + b, err := record.MarshalJSON() + require.NoError(t, err) + + // Should contain coordinate arrays + jsonStr := string(b) + assert.Contains(t, jsonStr, "[1,2]") + assert.Contains(t, jsonStr, "null") + assert.Contains(t, jsonStr, "[3,4]") +} + +func TestMarshalPointArray(t *testing.T) { + mem := memory.NewGoAllocator() + + // Test different dimensions + testCases := []struct { + name string + pointType *extensions.PointType + points []extensions.Point + expected string + }{ + { + name: "XY", + pointType: extensions.NewPointType(), + points: []extensions.Point{extensions.NewPoint(1.5, 2.5), extensions.NewPoint(3.0, 4.0)}, + expected: "[[1.5,2.5],[3,4]]", + }, + { + name: "XYZ", + pointType: extensions.NewPointTypeWithDimension(extensions.DimensionXYZ), + points: []extensions.Point{extensions.NewPointZ(1.0, 2.0, 3.0), extensions.NewPointZ(4.0, 5.0, 6.0)}, + expected: "[[1,2,3],[4,5,6]]", + }, + { + name: "XYM", + pointType: extensions.NewPointTypeWithDimension(extensions.DimensionXYM), + points: []extensions.Point{extensions.NewPointM(1.0, 2.0, 100.0), extensions.NewPointM(4.0, 5.0, 200.0)}, + expected: "[[1,2,100],[4,5,200]]", + }, + { + name: "XYZM", + pointType: extensions.NewPointTypeWithDimension(extensions.DimensionXYZM), + points: []extensions.Point{extensions.NewPointZM(1.0, 2.0, 3.0, 100.0)}, + expected: "[[1,2,3,100]]", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + builder := extensions.NewPointBuilder(mem, tc.pointType) + defer builder.Release() + + for _, point := range tc.points { + builder.Append(point) + } + + arr := builder.NewArray() + defer arr.Release() + + pointArray := arr.(*extensions.PointArray) + b, err := pointArray.MarshalJSON() + require.NoError(t, err) + + require.JSONEq(t, tc.expected, string(b)) + }) + } +} + +func TestPointConstructors(t *testing.T) { + // Test coordinate dimension enum first + t.Run("coordinate_dimension", func(t *testing.T) { + assert.Equal(t, "xy", extensions.DimensionXY.String()) + assert.Equal(t, "xyz", extensions.DimensionXYZ.String()) + assert.Equal(t, "xym", extensions.DimensionXYM.String()) + assert.Equal(t, "xyzm", extensions.DimensionXYZM.String()) + + assert.Equal(t, 2, extensions.DimensionXY.Size()) + assert.Equal(t, 3, extensions.DimensionXYZ.Size()) + assert.Equal(t, 3, extensions.DimensionXYM.Size()) + assert.Equal(t, 4, extensions.DimensionXYZM.Size()) + }) + + // Test all Point constructor functions + t.Run("extensions.NewPoint", func(t *testing.T) { + point := extensions.NewPoint(1.5, 2.5) + assert.Equal(t, 1.5, point.X) + assert.Equal(t, 2.5, point.Y) + assert.Equal(t, 0.0, point.Z) + assert.Equal(t, 0.0, point.M) + assert.Equal(t, extensions.DimensionXY, point.Dimension) + assert.False(t, point.IsEmpty()) + assert.Equal(t, "POINT(1.500000 2.500000)", point.String()) + }) + + t.Run("extensions.NewPointZ", func(t *testing.T) { + point := extensions.NewPointZ(1.5, 2.5, 3.5) + assert.Equal(t, 1.5, point.X) + assert.Equal(t, 2.5, point.Y) + assert.Equal(t, 3.5, point.Z) + assert.Equal(t, 0.0, point.M) + assert.Equal(t, extensions.DimensionXYZ, point.Dimension) + assert.False(t, point.IsEmpty()) + assert.Equal(t, "POINT Z(1.500000 2.500000 3.500000)", point.String()) + }) + + t.Run("extensions.NewPointM", func(t *testing.T) { + point := extensions.NewPointM(1.5, 2.5, 100.0) + assert.Equal(t, 1.5, point.X) + assert.Equal(t, 2.5, point.Y) + assert.Equal(t, 0.0, point.Z) + assert.Equal(t, 100.0, point.M) + assert.Equal(t, extensions.DimensionXYM, point.Dimension) + assert.False(t, point.IsEmpty()) + assert.Equal(t, "POINT M(1.500000 2.500000 100.000000)", point.String()) + }) + + t.Run("extensions.NewPointZM", func(t *testing.T) { + point := extensions.NewPointZM(1.5, 2.5, 3.5, 100.0) + assert.Equal(t, 1.5, point.X) + assert.Equal(t, 2.5, point.Y) + assert.Equal(t, 3.5, point.Z) + assert.Equal(t, 100.0, point.M) + assert.Equal(t, extensions.DimensionXYZM, point.Dimension) + assert.False(t, point.IsEmpty()) + assert.Equal(t, "POINT ZM(1.500000 2.500000 3.500000 100.000000)", point.String()) + }) + + t.Run("NewEmptyPoint", func(t *testing.T) { + point := extensions.NewEmptyPoint() + assert.True(t, point.IsEmpty()) + assert.Equal(t, "POINT EMPTY", point.String()) + assert.Equal(t, extensions.DimensionXY, point.Dimension) + }) +} + +func TestPointSerialization(t *testing.T) { + t.Run("type_roundtrip", func(t *testing.T) { + originalType := extensions.NewPointTypeWithDimension(extensions.DimensionXYM) + serialized := originalType.Serialize() + + deserializedInterface, err := originalType.Deserialize(originalType.StorageType(), serialized) + require.NoError(t, err) + + deserializedType, ok := deserializedInterface.(*extensions.PointType) + require.True(t, ok) + + assert.True(t, originalType.ExtensionEquals(deserializedType)) + assert.Equal(t, originalType.Dimension(), deserializedType.Dimension()) + }) + + t.Run("json_marshaling", func(t *testing.T) { + mem := memory.NewGoAllocator() + pointType := extensions.NewPointType() + builder := extensions.NewPointBuilder(mem, pointType) + + builder.AppendXY(1.5, 2.5) + builder.AppendNull() + + arr := builder.NewArray() + defer arr.Release() + + pointArray := arr.(*extensions.PointArray) + + // Test GetOneForMarshal + jsonVal0 := pointArray.GetOneForMarshal(0) + coords0, ok := jsonVal0.([]float64) + require.True(t, ok) + assert.Equal(t, []float64{1.5, 2.5}, coords0) + + jsonVal1 := pointArray.GetOneForMarshal(1) + assert.Nil(t, jsonVal1) // null value + }) +}