Skip to content

Commit f86fe9c

Browse files
robert3005robert3005
committed
Decimal32/Decimal64 support
Signed-off-by: Robert Kruszewski <github@robertk.io>
1 parent cb24576 commit f86fe9c

31 files changed

Lines changed: 3234 additions & 156 deletions

File tree

‎adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/JdbcToArrowUtils.java‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,8 @@
4646
import org.apache.arrow.adapter.jdbc.consumer.CompositeJdbcConsumer;
4747
import org.apache.arrow.adapter.jdbc.consumer.DateConsumer;
4848
import org.apache.arrow.adapter.jdbc.consumer.Decimal256Consumer;
49+
import org.apache.arrow.adapter.jdbc.consumer.Decimal32Consumer;
50+
import org.apache.arrow.adapter.jdbc.consumer.Decimal64Consumer;
4951
import org.apache.arrow.adapter.jdbc.consumer.DecimalConsumer;
5052
import org.apache.arrow.adapter.jdbc.consumer.DoubleConsumer;
5153
import org.apache.arrow.adapter.jdbc.consumer.FloatConsumer;
@@ -66,6 +68,8 @@
6668
import org.apache.arrow.vector.BitVector;
6769
import org.apache.arrow.vector.DateDayVector;
6870
import org.apache.arrow.vector.Decimal256Vector;
71+
import org.apache.arrow.vector.Decimal32Vector;
72+
import org.apache.arrow.vector.Decimal64Vector;
6973
import org.apache.arrow.vector.DecimalVector;
7074
import org.apache.arrow.vector.FieldVector;
7175
import org.apache.arrow.vector.Float4Vector;
@@ -510,10 +514,21 @@ public static JdbcConsumer getConsumer(
510514
}
511515
case Decimal:
512516
final RoundingMode bigDecimalRoundingMode = config.getBigDecimalRoundingMode();
513-
if (((ArrowType.Decimal) arrowType).getBitWidth() == 256) {
517+
final int decimalBitWidth = ((ArrowType.Decimal) arrowType).getBitWidth();
518+
if (decimalBitWidth == 256) {
514519
return Decimal256Consumer.createConsumer(
515520
(Decimal256Vector) vector, columnIndex, nullable, bigDecimalRoundingMode);
521+
} else if (decimalBitWidth == 128) {
522+
return DecimalConsumer.createConsumer(
523+
(DecimalVector) vector, columnIndex, nullable, bigDecimalRoundingMode);
524+
} else if (decimalBitWidth == 64) {
525+
return Decimal64Consumer.createConsumer(
526+
(Decimal64Vector) vector, columnIndex, nullable, bigDecimalRoundingMode);
527+
} else if (decimalBitWidth == 32) {
528+
return Decimal32Consumer.createConsumer(
529+
(Decimal32Vector) vector, columnIndex, nullable, bigDecimalRoundingMode);
516530
} else {
531+
// Any other bit width maps to MinorType.DECIMAL, so the root created a DecimalVector.
517532
return DecimalConsumer.createConsumer(
518533
(DecimalVector) vector, columnIndex, nullable, bigDecimalRoundingMode);
519534
}

‎adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/Decimal32Consumer.java‎

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.apache.arrow.adapter.jdbc.consumer;
18+
19+
import java.math.BigDecimal;
20+
import java.math.RoundingMode;
21+
import java.sql.ResultSet;
22+
import java.sql.SQLException;
23+
import org.apache.arrow.vector.Decimal32Vector;
24+
25+
/**
26+
* Consumer which consume decimal type values from {@link ResultSet}. Write the data to {@link
27+
* org.apache.arrow.vector.Decimal32Vector}.
28+
*/
29+
public abstract class Decimal32Consumer extends BaseConsumer<Decimal32Vector> {
30+
private final RoundingMode bigDecimalRoundingMode;
31+
private final int scale;
32+
33+
/**
34+
* Constructs a new consumer.
35+
*
36+
* @param vector the underlying vector for the consumer.
37+
* @param index the column id for the consumer.
38+
*/
39+
public Decimal32Consumer(Decimal32Vector vector, int index) {
40+
this(vector, index, null);
41+
}
42+
43+
/**
44+
* Constructs a new consumer, with optional coercibility.
45+
*
46+
* @param vector the underlying vector for the consumer.
47+
* @param index the column index for the consumer.
48+
* @param bigDecimalRoundingMode java.math.RoundingMode to be applied if the BigDecimal scale does
49+
* not match that of the target vector. Set to null to retain strict matching behavior (scale
50+
* of source and target vector must match exactly).
51+
*/
52+
public Decimal32Consumer(Decimal32Vector vector, int index, RoundingMode bigDecimalRoundingMode) {
53+
super(vector, index);
54+
this.bigDecimalRoundingMode = bigDecimalRoundingMode;
55+
this.scale = vector.getScale();
56+
}
57+
58+
/** Creates a consumer for {@link Decimal32Vector}. */
59+
public static JdbcConsumer<Decimal32Vector> createConsumer(
60+
Decimal32Vector vector, int index, boolean nullable, RoundingMode bigDecimalRoundingMode) {
61+
if (nullable) {
62+
return new NullableDecimal32Consumer(vector, index, bigDecimalRoundingMode);
63+
} else {
64+
return new NonNullableDecimal32Consumer(vector, index, bigDecimalRoundingMode);
65+
}
66+
}
67+
68+
protected void set(BigDecimal value) {
69+
if (bigDecimalRoundingMode != null && value.scale() != scale) {
70+
value = value.setScale(scale, bigDecimalRoundingMode);
71+
}
72+
vector.set(currentIndex, value);
73+
}
74+
75+
/** Consumer for nullable decimal. */
76+
static class NullableDecimal32Consumer extends Decimal32Consumer {
77+
78+
/** Instantiate a Decimal32Consumer. */
79+
public NullableDecimal32Consumer(
80+
Decimal32Vector vector, int index, RoundingMode bigDecimalRoundingMode) {
81+
super(vector, index, bigDecimalRoundingMode);
82+
}
83+
84+
@Override
85+
public void consume(ResultSet resultSet) throws SQLException {
86+
BigDecimal value = resultSet.getBigDecimal(columnIndexInResultSet);
87+
if (!resultSet.wasNull()) {
88+
// for fixed width vectors, we have allocated enough memory proactively,
89+
// so there is no need to call the setSafe method here.
90+
set(value);
91+
}
92+
currentIndex++;
93+
}
94+
}
95+
96+
/** Consumer for non-nullable decimal. */
97+
static class NonNullableDecimal32Consumer extends Decimal32Consumer {
98+
99+
/** Instantiate a Decimal32Consumer. */
100+
public NonNullableDecimal32Consumer(
101+
Decimal32Vector vector, int index, RoundingMode bigDecimalRoundingMode) {
102+
super(vector, index, bigDecimalRoundingMode);
103+
}
104+
105+
@Override
106+
public void consume(ResultSet resultSet) throws SQLException {
107+
BigDecimal value = resultSet.getBigDecimal(columnIndexInResultSet);
108+
// for fixed width vectors, we have allocated enough memory proactively,
109+
// so there is no need to call the setSafe method here.
110+
set(value);
111+
currentIndex++;
112+
}
113+
}
114+
}

‎adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/Decimal64Consumer.java‎

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.apache.arrow.adapter.jdbc.consumer;
18+
19+
import java.math.BigDecimal;
20+
import java.math.RoundingMode;
21+
import java.sql.ResultSet;
22+
import java.sql.SQLException;
23+
import org.apache.arrow.vector.Decimal64Vector;
24+
25+
/**
26+
* Consumer which consume decimal type values from {@link ResultSet}. Write the data to {@link
27+
* org.apache.arrow.vector.Decimal64Vector}.
28+
*/
29+
public abstract class Decimal64Consumer extends BaseConsumer<Decimal64Vector> {
30+
private final RoundingMode bigDecimalRoundingMode;
31+
private final int scale;
32+
33+
/**
34+
* Constructs a new consumer.
35+
*
36+
* @param vector the underlying vector for the consumer.
37+
* @param index the column id for the consumer.
38+
*/
39+
public Decimal64Consumer(Decimal64Vector vector, int index) {
40+
this(vector, index, null);
41+
}
42+
43+
/**
44+
* Constructs a new consumer, with optional coercibility.
45+
*
46+
* @param vector the underlying vector for the consumer.
47+
* @param index the column index for the consumer.
48+
* @param bigDecimalRoundingMode java.math.RoundingMode to be applied if the BigDecimal scale does
49+
* not match that of the target vector. Set to null to retain strict matching behavior (scale
50+
* of source and target vector must match exactly).
51+
*/
52+
public Decimal64Consumer(Decimal64Vector vector, int index, RoundingMode bigDecimalRoundingMode) {
53+
super(vector, index);
54+
this.bigDecimalRoundingMode = bigDecimalRoundingMode;
55+
this.scale = vector.getScale();
56+
}
57+
58+
/** Creates a consumer for {@link Decimal64Vector}. */
59+
public static JdbcConsumer<Decimal64Vector> createConsumer(
60+
Decimal64Vector vector, int index, boolean nullable, RoundingMode bigDecimalRoundingMode) {
61+
if (nullable) {
62+
return new NullableDecimal64Consumer(vector, index, bigDecimalRoundingMode);
63+
} else {
64+
return new NonNullableDecimal64Consumer(vector, index, bigDecimalRoundingMode);
65+
}
66+
}
67+
68+
protected void set(BigDecimal value) {
69+
if (bigDecimalRoundingMode != null && value.scale() != scale) {
70+
value = value.setScale(scale, bigDecimalRoundingMode);
71+
}
72+
vector.set(currentIndex, value);
73+
}
74+
75+
/** Consumer for nullable decimal. */
76+
static class NullableDecimal64Consumer extends Decimal64Consumer {
77+
78+
/** Instantiate a Decimal64Consumer. */
79+
public NullableDecimal64Consumer(
80+
Decimal64Vector vector, int index, RoundingMode bigDecimalRoundingMode) {
81+
super(vector, index, bigDecimalRoundingMode);
82+
}
83+
84+
@Override
85+
public void consume(ResultSet resultSet) throws SQLException {
86+
BigDecimal value = resultSet.getBigDecimal(columnIndexInResultSet);
87+
if (!resultSet.wasNull()) {
88+
// for fixed width vectors, we have allocated enough memory proactively,
89+
// so there is no need to call the setSafe method here.
90+
set(value);
91+
}
92+
currentIndex++;
93+
}
94+
}
95+
96+
/** Consumer for non-nullable decimal. */
97+
static class NonNullableDecimal64Consumer extends Decimal64Consumer {
98+
99+
/** Instantiate a Decimal64Consumer. */
100+
public NonNullableDecimal64Consumer(
101+
Decimal64Vector vector, int index, RoundingMode bigDecimalRoundingMode) {
102+
super(vector, index, bigDecimalRoundingMode);
103+
}
104+
105+
@Override
106+
public void consume(ResultSet resultSet) throws SQLException {
107+
BigDecimal value = resultSet.getBigDecimal(columnIndexInResultSet);
108+
// for fixed width vectors, we have allocated enough memory proactively,
109+
// so there is no need to call the setSafe method here.
110+
set(value);
111+
currentIndex++;
112+
}
113+
}
114+
}

0 commit comments

Comments
 (0)