Skip to content

Commit 5484316

Browse files
authored
fix: check for extension type before built-in arrow types (#109)
The JSON type was not getting selected as the underlying storage type (Binary) is the underlying arrow type. Checking for extension type first will allow us to prefer our defined Arrow extensions over the built-in types.
1 parent 436a826 commit 5484316

File tree

4 files changed

+26
-16
lines changed

4 files changed

+26
-16
lines changed

Diff for: lib/src/main/java/io/cloudquery/memdb/MemDB.java

+4
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,9 @@
1515
import io.cloudquery.transformers.Tables;
1616
import io.cloudquery.transformers.TransformWithClass;
1717
import io.grpc.stub.StreamObserver;
18+
import java.time.LocalDateTime;
1819
import java.util.List;
20+
import java.util.Map;
1921
import java.util.UUID;
2022

2123
public class MemDB extends Plugin {
@@ -32,6 +34,8 @@ public void resolve(
3234
Table1Data.builder()
3335
.id(UUID.fromString("46b2b6e6-8f3e-4340-a721-4aa0786b1cc0"))
3436
.name("name1")
37+
.timestamp(LocalDateTime.now())
38+
.json(Map.of("key1", "value1", "key2", "value2"))
3539
.build());
3640
stream.write(
3741
Table1Data.builder()

Diff for: lib/src/main/java/io/cloudquery/memdb/Table1Data.java

+4
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package io.cloudquery.memdb;
22

3+
import java.time.LocalDateTime;
4+
import java.util.Map;
35
import java.util.UUID;
46
import lombok.Builder;
57
import lombok.Getter;
@@ -9,4 +11,6 @@
911
public class Table1Data {
1012
private UUID id;
1113
private String name;
14+
private LocalDateTime timestamp;
15+
private Map<String, String> json;
1216
}

Diff for: lib/src/main/java/io/cloudquery/scalar/Scalar.java

+16-16
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package io.cloudquery.scalar;
22

3+
import io.cloudquery.types.JSONType;
34
import io.cloudquery.types.UUIDType;
45
import java.util.Objects;
56
import org.apache.arrow.vector.types.pojo.ArrowType;
@@ -77,6 +78,21 @@ public final int hashCode() {
7778
public static final java.lang.String NULL_VALUE_STRING = "(null)";
7879

7980
public static Scalar<?> fromArrowType(ArrowType arrowType) {
81+
if (arrowType instanceof ArrowType.ExtensionType extensionType) {
82+
switch (extensionType.extensionName()) {
83+
case UUIDType.EXTENSION_NAME -> {
84+
return new UUID();
85+
}
86+
case JSONType.EXTENSION_NAME -> {
87+
return new JSON();
88+
}
89+
// TODO: Add support for these types when scalar available
90+
// case INETType.EXTENSION_NAME -> {
91+
// return new INET();
92+
// }
93+
}
94+
}
95+
8096
switch (arrowType.getTypeID()) {
8197
case Timestamp -> {
8298
return new Timestamp();
@@ -110,22 +126,6 @@ public static Scalar<?> fromArrowType(ArrowType arrowType) {
110126
}
111127
}
112128

113-
if (arrowType instanceof ArrowType.ExtensionType extensionType) {
114-
//noinspection SwitchStatementWithTooFewBranches
115-
switch (extensionType.extensionName()) {
116-
case UUIDType.EXTENSION_NAME -> {
117-
return new UUID();
118-
}
119-
// TODO: Add support for these types when scalar available
120-
// case JSONType.EXTENSION_NAME -> {
121-
// return new JSON();
122-
// }
123-
// case INETType.EXTENSION_NAME -> {
124-
// return new INET();
125-
// }
126-
}
127-
}
128-
129129
throw new UnsupportedOperationException("Unsupported type: " + arrowType);
130130
}
131131

Diff for: lib/src/test/java/io/cloudquery/scalar/ScalarTest.java

+2
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
44

5+
import io.cloudquery.types.JSONType;
56
import io.cloudquery.types.UUIDType;
67
import java.time.ZoneOffset;
78
import java.util.stream.Stream;
@@ -64,6 +65,7 @@ public static Stream<Arguments> testDataSource() {
6465

6566
// Extension
6667
Arguments.of(new UUIDType(), UUID.class),
68+
Arguments.of(new JSONType(), JSON.class),
6769

6870
// Date
6971
Arguments.of(new ArrowType.Date(DateUnit.DAY), DateDay.class),

0 commit comments

Comments
 (0)