chore(arrow): Simplify with arrow.ListLikeType (#186)

Follow-up after apache/arrow#35885
cloudquery · Jun 2, 2023 · 5553242 · 5553242
1 parent 9dd0523
commit 5553242
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 27 deletions.
diff --git a/go.mod b/go.mod
@@ -12,7 +12,7 @@ require (
 )
 
 // TODO: remove once all updates are merged
-replace github.com/apache/arrow/go/v13 => github.com/cloudquery/arrow/go/v13 v13.0.0-20230602001318-a7aad4c5365c
+replace github.com/apache/arrow/go/v13 => github.com/cloudquery/arrow/go/v13 v13.0.0-20230602155531-6d34568d4501
 
 require (
 	github.com/JohnCGriffin/overflow v0.0.0-20211019200055-46fa312c352c // indirect

diff --git a/go.sum b/go.sum
@@ -6,8 +6,8 @@ github.com/apache/thrift v0.18.1 h1:lNhK/1nqjbwbiOPDBPFJVKxgDEGSepKuTh6OLiXW8kg=
 github.com/apache/thrift v0.18.1/go.mod h1:rdQn/dCcDKEWjjylUeueum4vQEjG2v8v2PqriUnbr+I=
 github.com/bradleyjkemp/cupaloy/v2 v2.8.0 h1:any4BmKE+jGIaMpnU8YgH/I2LPiLBufr6oMMlVBbn9M=
 github.com/bradleyjkemp/cupaloy/v2 v2.8.0/go.mod h1:bm7JXdkRd4BHJk9HpwqAI8BoAY1lps46Enkdqw6aRX0=
-github.com/cloudquery/arrow/go/v13 v13.0.0-20230602001318-a7aad4c5365c h1:z/inZxI0DTuu4DELOZQN5CjlSx0jwdVcn4XILXWDB+o=
-github.com/cloudquery/arrow/go/v13 v13.0.0-20230602001318-a7aad4c5365c/go.mod h1:/XatdE3kDIBqZKhZ7OBUHwP2jaASDFZHqF4puOWM8po=
+github.com/cloudquery/arrow/go/v13 v13.0.0-20230602155531-6d34568d4501 h1:vZJyF3/vy6nhP7guV0I5aB1eK3CGXZb7vg7UQPPnqqo=
+github.com/cloudquery/arrow/go/v13 v13.0.0-20230602155531-6d34568d4501/go.mod h1:/XatdE3kDIBqZKhZ7OBUHwP2jaASDFZHqF4puOWM8po=
 github.com/cloudquery/plugin-pb-go v1.0.8 h1:wn3GXhcNItcP+6wUUZuzUFbvdL59liKBO37/izMi+FQ=
 github.com/cloudquery/plugin-pb-go v1.0.8/go.mod h1:vAGA27psem7ZZNAY4a3S9TKuA/JDQWstjKcHPJX91Mc=
 github.com/cloudquery/plugin-sdk/v3 v3.10.1 h1:Jb7UAt8hU7zs3A+5nIHO7LQmYfz3hDfFfHixGKYDmTU=

diff --git a/parquet/list.go b/parquet/list.go
diff --git a/parquet/read.go b/parquet/read.go
@@ -82,20 +82,11 @@ func reverseTransformArray(dt arrow.DataType, arr arrow.Array) arrow.Array {
 	case *array.Struct:
 		return reverseTransformStruct(dt.(*arrow.StructType), arr)
 
-	case array.ListLike:
-		var child arrow.ArrayData
-		switch dt := dt.(type) {
-		case *arrow.MapType:
-			child = reverseTransformArray(dt.ValueType(), arr.ListValues()).Data()
-		case listLikeType:
-			child = reverseTransformArray(dt.Elem(), arr.ListValues()).Data()
-		default:
-			panic("unsupported list like conv to " + dt.String())
-		}
+	case array.ListLike: // this also handles maps
 		return array.MakeFromData(array.NewData(
 			dt, arr.Len(),
 			arr.Data().Buffers(),
-			[]arrow.ArrayData{child},
+			[]arrow.ArrayData{reverseTransformArray(dt.(arrow.ListLikeType).Elem(), arr.ListValues()).Data()},
 			arr.NullN(), arr.Data().Offset(),
 		))
 

diff --git a/parquet/write.go b/parquet/write.go
@@ -53,7 +53,6 @@ func transformDataType(t arrow.DataType) arrow.DataType {
 		return arrow.BinaryTypes.String
 
 	case *arrow.LargeBinaryType,
-		*arrow.LargeListType,
 		*arrow.LargeStringType: // not yet implemented in arrow
 		return arrow.BinaryTypes.String
 
@@ -69,8 +68,9 @@ func transformDataType(t arrow.DataType) arrow.DataType {
 	case *arrow.MapType:
 		return arrow.MapOf(transformDataType(dt.KeyType()), transformDataType(dt.ItemType()))
 
-	case listLikeType:
+	case arrow.ListLikeType:
 		return arrow.ListOf(transformDataType(dt.Elem()))
+
 	default:
 		return t
 	}
@@ -107,7 +107,7 @@ func transformArray(arr arrow.Array) arrow.Array {
 			arr.NullN(), arr.Data().Offset(),
 		))
 
-	case array.ListLike:
+	case array.ListLike: // this also handles maps
 		return array.MakeFromData(array.NewData(
 			transformDataType(arr.DataType()), arr.Len(),
 			arr.Data().Buffers(),