From 43d60c578ef4187c4114efe035e085a2a52e1eab Mon Sep 17 00:00:00 2001 From: Jia-Xuan Liu Date: Fri, 19 Apr 2024 15:03:28 +0800 Subject: [PATCH] handle hugeint for duckdb connector --- .../wren/base/client/duckdb/DuckdbTypes.java | 4 +++- .../java/io/wren/base/type/NumericType.java | 19 +++++++++++++++---- .../testing/duckdb/TestWrenWithDuckDB.java | 3 ++- 3 files changed, 20 insertions(+), 6 deletions(-) diff --git a/wren-base/src/main/java/io/wren/base/client/duckdb/DuckdbTypes.java b/wren-base/src/main/java/io/wren/base/client/duckdb/DuckdbTypes.java index 87b00d6b9..12b46f93b 100644 --- a/wren-base/src/main/java/io/wren/base/client/duckdb/DuckdbTypes.java +++ b/wren-base/src/main/java/io/wren/base/client/duckdb/DuckdbTypes.java @@ -46,9 +46,10 @@ public final class DuckdbTypes { - // other types LIST, ENUM, HUGEINT, UTINYINT, USMALLINT, STRUCT, UUID, JSON, UINTEGER, UBIGINT, INTERVAL, MAP + // other types LIST, ENUM, UTINYINT, USMALLINT, STRUCT, UUID, JSON, UINTEGER, UBIGINT, INTERVAL, MAP public static final DuckdbType BOOLEAN = new DuckdbType(Types.BOOLEAN, "BOOLEAN"); public static final DuckdbType BIGINT = new DuckdbType(Types.BIGINT, "BIGINT"); + public static final DuckdbType HUGEINT = new DuckdbType(Types.DECIMAL, "HUGEINT"); public static final DuckdbType BIT = new DuckdbType(Types.BIT, "BIT"); public static final DuckdbType BLOB = new DuckdbType(Types.BLOB, "BLOB"); public static final DuckdbType DATE = new DuckdbType(Types.DATE, "DATE"); @@ -86,6 +87,7 @@ public final class DuckdbTypes .put(TIMESTAMP_WITH_TIMEZONE.getName(), TIMESTAMP_WITH_TIMEZONE) .put(VARCHAR.getName(), VARCHAR) .put(JSON.getName(), JSON) + .put(HUGEINT.getName(), HUGEINT) .build(); private static final Map> duckdbTypeToPgTypeMap = ImmutableMap.>builder() diff --git a/wren-base/src/main/java/io/wren/base/type/NumericType.java b/wren-base/src/main/java/io/wren/base/type/NumericType.java index 8afb7cd27..a728d40ab 100644 --- a/wren-base/src/main/java/io/wren/base/type/NumericType.java +++ b/wren-base/src/main/java/io/wren/base/type/NumericType.java @@ -19,11 +19,12 @@ import javax.annotation.Nonnull; import java.math.BigDecimal; +import java.math.BigInteger; import java.math.MathContext; import java.nio.charset.StandardCharsets; public class NumericType - extends PGType + extends PGType { static final int OID = 1700; @@ -61,7 +62,17 @@ public String type() } @Override - public int writeAsBinary(ByteBuf buffer, @Nonnull BigDecimal value) + public int writeAsBinary(ByteBuf buffer, @Nonnull Number value) + { + return switch (value) { + case BigDecimal bigDecimal -> writeAsBinary(buffer, bigDecimal); + // TODO: customize the handling of BigInteger for performance. + case BigInteger bigInteger -> writeAsBinary(buffer, new BigDecimal(bigInteger)); + default -> throw new IllegalArgumentException("Unsupported numeric type: " + value.getClass().getName()); + }; + } + + private int writeAsBinary(ByteBuf buffer, @Nonnull BigDecimal value) { // Taken from https://github.com/cockroachdb/cockroach/blob/master/pkg/sql/pgwire/types.go#L336 // and https://github.com/postgres/postgres/blob/master/src/backend/utils/adt/numeric.c#L6760. @@ -174,13 +185,13 @@ public BigDecimal readBinaryValue(ByteBuf buffer, int valueLength) } @Override - public byte[] encodeAsUTF8Text(@Nonnull BigDecimal value) + public byte[] encodeAsUTF8Text(@Nonnull Number value) { return value.toString().getBytes(StandardCharsets.UTF_8); } @Override - public BigDecimal decodeUTF8Text(byte[] bytes) + public Number decodeUTF8Text(byte[] bytes) { return new BigDecimal(new String(bytes, StandardCharsets.UTF_8)); } diff --git a/wren-tests/src/test/java/io/wren/testing/duckdb/TestWrenWithDuckDB.java b/wren-tests/src/test/java/io/wren/testing/duckdb/TestWrenWithDuckDB.java index c5b248a3c..a3b22f7f1 100644 --- a/wren-tests/src/test/java/io/wren/testing/duckdb/TestWrenWithDuckDB.java +++ b/wren-tests/src/test/java/io/wren/testing/duckdb/TestWrenWithDuckDB.java @@ -35,7 +35,8 @@ public Object[][] queryModel() {"SELECT * FROM Orders"}, {"SELECT * FROM Orders WHERE orderkey > 100"}, {"SELECT * FROM Orders a JOIN Customer b ON a.custkey = b.custkey"}, - {"SELECT * FROM Orders WHERE nation_name IS NOT NULL"} + {"SELECT * FROM Orders WHERE nation_name IS NOT NULL"}, + {"SELECT sum(orderkey) FROM Orders"}, // DuckDB always returns HUGEINT when aggregating integers }; }