Visitar URL original
avros · apache/arrow-java@63f93f3 · GitHub
Skip to content

Commit 63f93f3

Browse files
robert3005robert3005
committed
avros
1 parent 4fb64d7 commit 63f93f3

11 files changed

Lines changed: 488 additions & 35 deletions

File tree

‎adapter/avro/src/main/java/org/apache/arrow/adapter/avro/AvroToArrowUtils.java‎

Lines changed: 34 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,8 @@
5151
import org.apache.arrow.adapter.avro.consumers.SkipFunction;
5252
import org.apache.arrow.adapter.avro.consumers.logical.AvroDateConsumer;
5353
import org.apache.arrow.adapter.avro.consumers.logical.AvroDecimal256Consumer;
54+
import org.apache.arrow.adapter.avro.consumers.logical.AvroDecimal32Consumer;
55+
import org.apache.arrow.adapter.avro.consumers.logical.AvroDecimal64Consumer;
5456
import org.apache.arrow.adapter.avro.consumers.logical.AvroDecimalConsumer;
5557
import org.apache.arrow.adapter.avro.consumers.logical.AvroTimeMicroConsumer;
5658
import org.apache.arrow.adapter.avro.consumers.logical.AvroTimeMillisConsumer;
@@ -67,6 +69,8 @@
6769
import org.apache.arrow.vector.BitVector;
6870
import org.apache.arrow.vector.DateDayVector;
6971
import org.apache.arrow.vector.Decimal256Vector;
72+
import org.apache.arrow.vector.Decimal32Vector;
73+
import org.apache.arrow.vector.Decimal64Vector;
7074
import org.apache.arrow.vector.DecimalVector;
7175
import org.apache.arrow.vector.FieldVector;
7276
import org.apache.arrow.vector.FixedSizeBinaryVector;
@@ -226,15 +230,7 @@ private static Consumer createConsumer(
226230
/* dictionary= */ null,
227231
getMetaData(schema, extProps, config));
228232
vector = createVector(consumerVector, fieldType, name, allocator);
229-
if (schema.getFixedSize() <= 16) {
230-
consumer =
231-
new AvroDecimalConsumer.FixedDecimalConsumer(
232-
(DecimalVector) vector, schema.getFixedSize());
233-
} else {
234-
consumer =
235-
new AvroDecimal256Consumer.FixedDecimal256Consumer(
236-
(Decimal256Vector) vector, schema.getFixedSize());
237-
}
233+
consumer = createFixedDecimalConsumer(vector, schema.getFixedSize());
238234
} else {
239235
arrowType = new ArrowType.FixedSizeBinary(schema.getFixedSize());
240236
fieldType =
@@ -407,12 +403,11 @@ private static ArrowType createDecimalArrowType(LogicalTypes.Decimal logicalType
407403
scale,
408404
precision);
409405

406+
// A fixed decimal maps to the narrowest Arrow decimal that holds its bytes, so an Arrow
407+
// decimal written as fixed(bitWidth / 8) reads back with the same bit width.
410408
if (schema.getType() == Schema.Type.FIXED) {
411-
if (schema.getFixedSize() <= 16) {
412-
return new ArrowType.Decimal(precision, scale, 128);
413-
} else {
414-
return new ArrowType.Decimal(precision, scale, 256);
415-
}
409+
return new ArrowType.Decimal(
410+
precision, scale, decimalBitWidthForFixed(schema.getFixedSize()));
416411
} else {
417412
if (precision <= 38) {
418413
return new ArrowType.Decimal(precision, scale, 128);
@@ -422,6 +417,31 @@ private static ArrowType createDecimalArrowType(LogicalTypes.Decimal logicalType
422417
}
423418
}
424419

420+
private static int decimalBitWidthForFixed(int fixedSize) {
421+
if (fixedSize <= Decimal32Vector.TYPE_WIDTH) {
422+
return Decimal32Vector.TYPE_WIDTH * 8;
423+
} else if (fixedSize <= Decimal64Vector.TYPE_WIDTH) {
424+
return Decimal64Vector.TYPE_WIDTH * 8;
425+
} else if (fixedSize <= DecimalVector.TYPE_WIDTH) {
426+
return DecimalVector.TYPE_WIDTH * 8;
427+
} else {
428+
return Decimal256Vector.TYPE_WIDTH * 8;
429+
}
430+
}
431+
432+
private static Consumer createFixedDecimalConsumer(FieldVector vector, int fixedSize) {
433+
if (vector instanceof Decimal32Vector) {
434+
return new AvroDecimal32Consumer.FixedDecimal32Consumer((Decimal32Vector) vector, fixedSize);
435+
} else if (vector instanceof Decimal64Vector) {
436+
return new AvroDecimal64Consumer.FixedDecimal64Consumer((Decimal64Vector) vector, fixedSize);
437+
} else if (vector instanceof DecimalVector) {
438+
return new AvroDecimalConsumer.FixedDecimalConsumer((DecimalVector) vector, fixedSize);
439+
} else {
440+
return new AvroDecimal256Consumer.FixedDecimal256Consumer(
441+
(Decimal256Vector) vector, fixedSize);
442+
}
443+
}
444+
425445
private static Consumer createSkipConsumer(Schema schema) {
426446

427447
SkipFunction skipFunction;

‎adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroDecimal32Consumer.java‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
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.avro.consumers.logical;
18+
19+
import java.io.IOException;
20+
import java.nio.ByteBuffer;
21+
import org.apache.arrow.adapter.avro.consumers.BaseAvroConsumer;
22+
import org.apache.arrow.util.Preconditions;
23+
import org.apache.arrow.vector.Decimal32Vector;
24+
import org.apache.avro.io.Decoder;
25+
26+
/**
27+
* Consumer which consume 32-bit decimal type values from avro decoder. Write the data to {@link
28+
* Decimal32Vector}.
29+
*/
30+
public abstract class AvroDecimal32Consumer extends BaseAvroConsumer<Decimal32Vector> {
31+
32+
private static final String LENGTH_ERROR =
33+
"Decimal bytes length should <= " + Decimal32Vector.TYPE_WIDTH + ".";
34+
35+
protected AvroDecimal32Consumer(Decimal32Vector vector) {
36+
super(vector);
37+
}
38+
39+
/** Consumer for decimal logical type with 32 bit width and original bytes type. */
40+
public static class BytesDecimal32Consumer extends AvroDecimal32Consumer {
41+
42+
private ByteBuffer cacheBuffer;
43+
44+
/** Instantiate a BytesDecimal32Consumer. */
45+
public BytesDecimal32Consumer(Decimal32Vector vector) {
46+
super(vector);
47+
}
48+
49+
@Override
50+
public void consume(Decoder decoder) throws IOException {
51+
cacheBuffer = decoder.readBytes(cacheBuffer);
52+
byte[] bytes = new byte[cacheBuffer.limit()];
53+
Preconditions.checkArgument(bytes.length <= Decimal32Vector.TYPE_WIDTH, LENGTH_ERROR);
54+
cacheBuffer.get(bytes);
55+
vector.setBigEndian(currentIndex++, bytes);
56+
}
57+
}
58+
59+
/** Consumer for decimal logical type with 32 bit width and original fixed type. */
60+
public static class FixedDecimal32Consumer extends AvroDecimal32Consumer {
61+
62+
private final byte[] reuseBytes;
63+
64+
/** Instantiate a FixedDecimal32Consumer. */
65+
public FixedDecimal32Consumer(Decimal32Vector vector, int size) {
66+
super(vector);
67+
Preconditions.checkArgument(size <= Decimal32Vector.TYPE_WIDTH, LENGTH_ERROR);
68+
reuseBytes = new byte[size];
69+
}
70+
71+
@Override
72+
public void consume(Decoder decoder) throws IOException {
73+
decoder.readFixed(reuseBytes);
74+
vector.setBigEndian(currentIndex++, reuseBytes);
75+
}
76+
}
77+
}

‎adapter/avro/src/main/java/org/apache/arrow/adapter/avro/consumers/logical/AvroDecimal64Consumer.java‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
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.avro.consumers.logical;
18+
19+
import java.io.IOException;
20+
import java.nio.ByteBuffer;
21+
import org.apache.arrow.adapter.avro.consumers.BaseAvroConsumer;
22+
import org.apache.arrow.util.Preconditions;
23+
import org.apache.arrow.vector.Decimal64Vector;
24+
import org.apache.avro.io.Decoder;
25+
26+
/**
27+
* Consumer which consume 64-bit decimal type values from avro decoder. Write the data to {@link
28+
* Decimal64Vector}.
29+
*/
30+
public abstract class AvroDecimal64Consumer extends BaseAvroConsumer<Decimal64Vector> {
31+
32+
private static final String LENGTH_ERROR =
33+
"Decimal bytes length should <= " + Decimal64Vector.TYPE_WIDTH + ".";
34+
35+
protected AvroDecimal64Consumer(Decimal64Vector vector) {
36+
super(vector);
37+
}
38+
39+
/** Consumer for decimal logical type with 64 bit width and original bytes type. */
40+
public static class BytesDecimal64Consumer extends AvroDecimal64Consumer {
41+
42+
private ByteBuffer cacheBuffer;
43+
44+
/** Instantiate a BytesDecimal64Consumer. */
45+
public BytesDecimal64Consumer(Decimal64Vector vector) {
46+
super(vector);
47+
}
48+
49+
@Override
50+
public void consume(Decoder decoder) throws IOException {
51+
cacheBuffer = decoder.readBytes(cacheBuffer);
52+
byte[] bytes = new byte[cacheBuffer.limit()];
53+
Preconditions.checkArgument(bytes.length <= Decimal64Vector.TYPE_WIDTH, LENGTH_ERROR);
54+
cacheBuffer.get(bytes);
55+
vector.setBigEndian(currentIndex++, bytes);
56+
}
57+
}
58+
59+
/** Consumer for decimal logical type with 64 bit width and original fixed type. */
60+
public static class FixedDecimal64Consumer extends AvroDecimal64Consumer {
61+
62+
private final byte[] reuseBytes;
63+
64+
/** Instantiate a FixedDecimal64Consumer. */
65+
public FixedDecimal64Consumer(Decimal64Vector vector, int size) {
66+
super(vector);
67+
Preconditions.checkArgument(size <= Decimal64Vector.TYPE_WIDTH, LENGTH_ERROR);
68+
reuseBytes = new byte[size];
69+
}
70+
71+
@Override
72+
public void consume(Decoder decoder) throws IOException {
73+
decoder.readFixed(reuseBytes);
74+
vector.setBigEndian(currentIndex++, reuseBytes);
75+
}
76+
}
77+
}

‎adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripDataTest.java‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@
3737
import org.apache.arrow.vector.BitVector;
3838
import org.apache.arrow.vector.DateDayVector;
3939
import org.apache.arrow.vector.Decimal256Vector;
40+
import org.apache.arrow.vector.Decimal32Vector;
41+
import org.apache.arrow.vector.Decimal64Vector;
4042
import org.apache.arrow.vector.DecimalVector;
4143
import org.apache.arrow.vector.FieldVector;
4244
import org.apache.arrow.vector.FixedSizeBinaryVector;
@@ -621,6 +623,82 @@ public void testRoundTripNullableBinary() throws Exception {
621623

622624
// Data round trip for logical types, nullable and non-nullable
623625

626+
@Test
627+
public void testRoundTripNarrowDecimals() throws Exception {
628+
629+
// Field definitions
630+
FieldType decimal32Field = new FieldType(false, new ArrowType.Decimal(9, 3, 32), null);
631+
FieldType decimal64Field = new FieldType(false, new ArrowType.Decimal(18, 6, 64), null);
632+
633+
// Create empty vectors
634+
BufferAllocator allocator = new RootAllocator();
635+
Decimal32Vector decimal32Vector =
636+
new Decimal32Vector(new Field("decimal32", decimal32Field, null), allocator);
637+
Decimal64Vector decimal64Vector =
638+
new Decimal64Vector(new Field("decimal64", decimal64Field, null), allocator);
639+
640+
// Set up VSR
641+
List<FieldVector> vectors = Arrays.asList(decimal32Vector, decimal64Vector);
642+
int rowCount = 3;
643+
644+
try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) {
645+
646+
root.setRowCount(rowCount);
647+
root.allocateNew();
648+
649+
// Set test data
650+
decimal32Vector.setSafe(0, new BigDecimal("123456.789"));
651+
decimal32Vector.setSafe(1, new BigDecimal("-999999.999"));
652+
decimal32Vector.setSafe(2, new BigDecimal("0.001"));
653+
654+
decimal64Vector.setSafe(0, new BigDecimal("123456789012.345678"));
655+
decimal64Vector.setSafe(1, new BigDecimal("-999999999999.999999"));
656+
decimal64Vector.setSafe(2, new BigDecimal("0.000001"));
657+
658+
File dataFile = new File(TMP, "testRoundTripNarrowDecimals.avro");
659+
660+
roundTripTest(root, allocator, dataFile, rowCount);
661+
}
662+
}
663+
664+
@Test
665+
public void testRoundTripNullableNarrowDecimals() throws Exception {
666+
667+
// Field definitions
668+
FieldType decimal32Field = new FieldType(true, new ArrowType.Decimal(9, 3, 32), null);
669+
FieldType decimal64Field = new FieldType(true, new ArrowType.Decimal(18, 6, 64), null);
670+
671+
// Create empty vectors
672+
BufferAllocator allocator = new RootAllocator();
673+
Decimal32Vector decimal32Vector =
674+
new Decimal32Vector(new Field("decimal32", decimal32Field, null), allocator);
675+
Decimal64Vector decimal64Vector =
676+
new Decimal64Vector(new Field("decimal64", decimal64Field, null), allocator);
677+
678+
// Set up VSR
679+
List<FieldVector> vectors = Arrays.asList(decimal32Vector, decimal64Vector);
680+
int rowCount = 3;
681+
682+
try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) {
683+
684+
root.setRowCount(rowCount);
685+
root.allocateNew();
686+
687+
// Set test data
688+
decimal32Vector.setNull(0);
689+
decimal32Vector.setSafe(1, BigDecimal.ZERO.setScale(3, RoundingMode.UNNECESSARY));
690+
decimal32Vector.setSafe(2, new BigDecimal("-1.500"));
691+
692+
decimal64Vector.setNull(0);
693+
decimal64Vector.setSafe(1, BigDecimal.ZERO.setScale(6, RoundingMode.UNNECESSARY));
694+
decimal64Vector.setSafe(2, new BigDecimal("-1.500000"));
695+
696+
File dataFile = new File(TMP, "testRoundTripNullableNarrowDecimals.avro");
697+
698+
roundTripTest(root, allocator, dataFile, rowCount);
699+
}
700+
}
701+
624702
@Test
625703
public void testRoundTripDecimals() throws Exception {
626704

‎adapter/avro/src/test/java/org/apache/arrow/adapter/avro/RoundTripSchemaTest.java‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,27 @@ public void testRoundTripDecimalTypes() {
217217
doRoundTripTest(fields);
218218
}
219219

220+
@Test
221+
public void testRoundTripNarrowDecimalTypes() {
222+
223+
List<Field> fields =
224+
Arrays.asList(
225+
new Field(
226+
"nullableDecimal32", FieldType.nullable(new ArrowType.Decimal(9, 2, 32)), null),
227+
new Field(
228+
"nonNullableDecimal32",
229+
FieldType.notNullable(new ArrowType.Decimal(5, 0, 32)),
230+
null),
231+
new Field(
232+
"nullableDecimal64", FieldType.nullable(new ArrowType.Decimal(18, 6, 64)), null),
233+
new Field(
234+
"nonNullableDecimal64",
235+
FieldType.notNullable(new ArrowType.Decimal(12, 4, 64)),
236+
null));
237+
238+
doRoundTripTest(fields);
239+
}
240+
220241
@Test
221242
public void testRoundTripDateTypes() {
222243

‎docs/source/jdbc.rst‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,10 @@ mapping, with additional support for the UUID extension type noted below.
227227
+----------------------------+----------------------------+-------+
228228
| Date64 | DATE (setDate) | |
229229
+----------------------------+----------------------------+-------+
230+
| Decimal32 | DECIMAL (setBigDecimal) | |
231+
+----------------------------+----------------------------+-------+
232+
| Decimal64 | DECIMAL (setBigDecimal) | |
233+
+----------------------------+----------------------------+-------+
230234
| Decimal128 | DECIMAL (setBigDecimal) | |
231235
+----------------------------+----------------------------+-------+
232236
| Decimal256 | DECIMAL (setBigDecimal) | |

0 commit comments

Comments
 (0)