fix: Inet extension MarshalJSON (#894)

Right now there is an inconsistency in our extensions: Some of them are
using `BinaryType` as the underlying storage and some `StringType`.

There is another issue in the parquet reader that seems not to like
extensions with `StringType` but we can solve this bug if we want our
Extensions to be of `StringType` rather then `BinaryType`.

Right now I changed `Inet` to `BinaryType` as well for consistency but
we have a chance to change all of them to `StringType` before we migrate
sources without breaking backward compatibility because we didn't send
anything over the wire yet - so right it's just an implementation change
rather then a breaking change.

Author: yevgenypats

types/inet.go

@@ -25,26 +25,26 @@ func (b *InetBuilder) Append(v *net.IPNet) {
 		b.AppendNull()
 		return
 	}
-	b.ExtensionBuilder.Builder.(*array.StringBuilder).Append(v.String())
+	b.ExtensionBuilder.Builder.(*array.BinaryBuilder).Append([]byte(v.String()))
 }
 
 func (b *InetBuilder) UnsafeAppend(v *net.IPNet) {
-	b.ExtensionBuilder.Builder.(*array.StringBuilder).UnsafeAppend([]byte(v.String()))
+	b.ExtensionBuilder.Builder.(*array.BinaryBuilder).UnsafeAppend([]byte(v.String()))
 }
 
 func (b *InetBuilder) AppendValues(v []*net.IPNet, valid []bool) {
 	if len(v) != len(valid) && len(valid) != 0 {
 		panic("len(v) != len(valid) && len(valid) != 0")
 	}
 
-	data := make([]string, len(v))
+	data := make([][]byte, len(v))
 	for i, v := range v {
 		if len(valid) > 0 && !valid[i] {
 			continue
 		}
-		data[i] = v.String()
+		data[i] = []byte(v.String())
 	}
-	b.ExtensionBuilder.Builder.(*array.StringBuilder).AppendValues(data, valid)
+	b.ExtensionBuilder.Builder.(*array.BinaryBuilder).AppendValues(data, valid)
 }
 
 func (b *InetBuilder) AppendValueFromString(s string) error {
@@ -127,7 +127,7 @@ type InetArray struct {
 }
 
 func (a *InetArray) String() string {
-	arr := a.Storage().(*array.String)
+	arr := a.Storage().(*array.Binary)
 	o := new(strings.Builder)
 	o.WriteString("[")
 	for i := 0; i < arr.Len(); i++ {
@@ -149,7 +149,7 @@ func (a *InetArray) Value(i int) *net.IPNet {
 	if a.IsNull(i) {
 		return nil
 	}
-	_, ipnet, err := net.ParseCIDR(a.Storage().(*array.String).Value(i))
+	_, ipnet, err := net.ParseCIDR(string(a.Storage().(*array.Binary).Value(i)))
 	if err != nil {
 		panic(fmt.Errorf("invalid ip+net: %w", err))
 	}
@@ -167,13 +167,26 @@ func (a *InetArray) ValueStr(i int) string {
 }
 
 func (a *InetArray) GetOneForMarshal(i int) any {
-	if val := a.Value(i); val != nil {
-		return val.String()
+	if a.IsNull(i) {
+		return nil
 	}
-	return nil
+	return string(a.Storage().(*array.Binary).Value(i))
+}
+
+func (a *InetArray) MarshalJSON() ([]byte, error) {
+	arr := a.Storage().(*array.Binary)
+	values := make([]any, a.Len())
+	for i := 0; i < a.Len(); i++ {
+		if a.IsValid(i) {
+			values[i] = string(arr.Value(i))
+		} else {
+			values[i] = nil
+		}
+	}
+	return json.Marshal(values)
 }
 
-// InetType is a simple extension type that represents a StringType
+// InetType is a simple extension type that represents a BinaryType
 // to be used for representing IP Addresses and CIDRs
 type InetType struct {
 	arrow.ExtensionBase
@@ -182,7 +195,7 @@ type InetType struct {
 // NewInetType is a convenience function to create an instance of InetType
 // with the correct storage type
 func NewInetType() *InetType {
-	return &InetType{ExtensionBase: arrow.ExtensionBase{Storage: &arrow.StringType{}}}
+	return &InetType{ExtensionBase: arrow.ExtensionBase{Storage: &arrow.BinaryType{}}}
 }
 
 // ArrayType returns TypeOf(InetArray{}) for constructing Inet arrays
@@ -199,13 +212,13 @@ func (*InetType) Serialize() string {
 	return "inet-serialized"
 }
 
-// Deserialize expects storageType to be StringType and the data to be
+// Deserialize expects storageType to be BinaryType and the data to be
 // "inet-serialized" in order to correctly create a InetType for testing deserialize.
 func (*InetType) Deserialize(storageType arrow.DataType, data string) (arrow.ExtensionType, error) {
 	if data != "inet-serialized" {
 		return nil, fmt.Errorf("type identifier did not match: '%s'", data)
 	}
-	if !arrow.TypeEqual(storageType, &arrow.StringType{}) {
+	if !arrow.TypeEqual(storageType, &arrow.BinaryType{}) {
 		return nil, fmt.Errorf("invalid storage type for InetType: %s", storageType.Name())
 	}
 	return NewInetType(), nil