Skip to content

Commit 0c9d879

Browse files
caicancaiILuffZhe
caicancai
authored and
ILuffZhe
committed
[CALCITE-6760] Support time data type in Arrow adapter
1 parent 901aadb commit 0c9d879

File tree

4 files changed

+55
-19
lines changed

4 files changed

+55
-19
lines changed

Diff for: ‎arrow/src/main/java/org/apache/calcite/adapter/arrow/ArrowFieldTypeFactory.java

+2
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,8 @@ private static RelDataType of(ArrowType arrowType, JavaTypeFactory typeFactory)
8080
return typeFactory.createSqlType(SqlTypeName.DECIMAL,
8181
((ArrowType.Decimal) arrowType).getPrecision(),
8282
((ArrowType.Decimal) arrowType).getScale());
83+
case Time:
84+
return typeFactory.createSqlType(SqlTypeName.TIME);
8385
default:
8486
throw new IllegalArgumentException("Unsupported type: " + arrowType);
8587
}

Diff for: ‎arrow/src/test/java/org/apache/calcite/adapter/arrow/ArrowAdapterDataTypesTest.java

+25-10
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
6868
String sql = "select \"tinyIntField\" from arrowdatatype";
6969
String plan = "PLAN=ArrowToEnumerableConverter\n"
7070
+ " ArrowProject(tinyIntField=[$0])\n"
71-
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n";
71+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
7272
String result = "tinyIntField=0\ntinyIntField=1\n";
7373
CalciteAssert.that()
7474
.with(arrow)
@@ -82,7 +82,7 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
8282
String sql = "select \"smallIntField\" from arrowdatatype";
8383
String plan = "PLAN=ArrowToEnumerableConverter\n"
8484
+ " ArrowProject(smallIntField=[$1])\n"
85-
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n";
85+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
8686
String result = "smallIntField=0\nsmallIntField=1\n";
8787
CalciteAssert.that()
8888
.with(arrow)
@@ -96,7 +96,7 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
9696
String sql = "select \"intField\" from arrowdatatype";
9797
String plan = "PLAN=ArrowToEnumerableConverter\n"
9898
+ " ArrowProject(intField=[$2])\n"
99-
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n";
99+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
100100
String result = "intField=0\nintField=1\n";
101101
CalciteAssert.that()
102102
.with(arrow)
@@ -110,7 +110,7 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
110110
String sql = "select \"longField\" from arrowdatatype";
111111
String plan = "PLAN=ArrowToEnumerableConverter\n"
112112
+ " ArrowProject(longField=[$5])\n"
113-
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n";
113+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
114114
String result = "longField=0\nlongField=1\n";
115115
CalciteAssert.that()
116116
.with(arrow)
@@ -124,7 +124,7 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
124124
String sql = "select \"floatField\" from arrowdatatype";
125125
String plan = "PLAN=ArrowToEnumerableConverter\n"
126126
+ " ArrowProject(floatField=[$4])\n"
127-
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n";
127+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
128128
String result = "floatField=0.0\nfloatField=1.0\n";
129129
CalciteAssert.that()
130130
.with(arrow)
@@ -138,7 +138,7 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
138138
String sql = "select \"doubleField\" from arrowdatatype";
139139
String plan = "PLAN=ArrowToEnumerableConverter\n"
140140
+ " ArrowProject(doubleField=[$6])\n"
141-
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n";
141+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
142142
String result = "doubleField=0.0\ndoubleField=1.0\n";
143143
CalciteAssert.that()
144144
.with(arrow)
@@ -152,7 +152,7 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
152152
String sql = "select \"decimalField\" from arrowdatatype";
153153
String plan = "PLAN=ArrowToEnumerableConverter\n"
154154
+ " ArrowProject(decimalField=[$8])\n"
155-
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n";
155+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
156156
String result = "decimalField=0.00\ndecimalField=1.00\n";
157157
CalciteAssert.that()
158158
.with(arrow)
@@ -166,7 +166,7 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
166166
String sql = "select \"dateField\" from arrowdatatype";
167167
String plan = "PLAN=ArrowToEnumerableConverter\n"
168168
+ " ArrowProject(dateField=[$9])\n"
169-
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n";
169+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
170170
String result = "dateField=1970-01-01\n"
171171
+ "dateField=1970-01-02\n";
172172
CalciteAssert.that()
@@ -181,7 +181,7 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
181181
String sql = "select \"booleanField\" from arrowdatatype";
182182
String plan = "PLAN=ArrowToEnumerableConverter\n"
183183
+ " ArrowProject(booleanField=[$7])\n"
184-
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n";
184+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
185185
String result = "booleanField=null\nbooleanField=true\nbooleanField=false\n";
186186
CalciteAssert.that()
187187
.with(arrow)
@@ -198,7 +198,7 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
198198
String sql = "select \"decimalField2\" from arrowdatatype";
199199
String plan = "PLAN=ArrowToEnumerableConverter\n"
200200
+ " ArrowProject(decimalField2=[$10])\n"
201-
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n";
201+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
202202
String result = "decimalField2=20.000\ndecimalField2=21.000\n";
203203
CalciteAssert.that()
204204
.with(arrow)
@@ -207,4 +207,19 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
207207
.returns(result)
208208
.explainContains(plan);
209209
}
210+
211+
@Test void testTimeProject() {
212+
String sql = "select \"timeField\" from arrowdatatype";
213+
String plan = "PLAN=ArrowToEnumerableConverter\n"
214+
+ " ArrowProject(timeField=[$11])\n"
215+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
216+
String result = "timeField=00:00:00\n"
217+
+ "timeField=00:00:01\n";
218+
CalciteAssert.that()
219+
.with(arrow)
220+
.query(sql)
221+
.limit(2)
222+
.returns(result)
223+
.explainContains(plan);
224+
}
210225
}

Diff for: ‎arrow/src/test/java/org/apache/calcite/adapter/arrow/ArrowAdapterTest.java

+9-9
Original file line numberDiff line numberDiff line change
@@ -860,7 +860,7 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
860860
String plan = "PLAN=ArrowToEnumerableConverter\n"
861861
+ " ArrowProject(booleanField=[$7])\n"
862862
+ " ArrowFilter(condition=[$7])\n"
863-
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n";
863+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
864864
String result = "booleanField=true\nbooleanField=true\n";
865865

866866
CalciteAssert.that()
@@ -878,7 +878,7 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
878878
String plan = "PLAN=ArrowToEnumerableConverter\n"
879879
+ " ArrowProject(intField=[$2])\n"
880880
+ " ArrowFilter(condition=[>($2, 10)])\n"
881-
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n";
881+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
882882
String result = "intField=11\nintField=12\n";
883883

884884
CalciteAssert.that()
@@ -896,7 +896,7 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
896896
String plan = "PLAN=ArrowToEnumerableConverter\n"
897897
+ " ArrowProject(booleanField=[$7])\n"
898898
+ " ArrowFilter(condition=[NOT($7)])\n"
899-
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n";
899+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
900900
String result = "booleanField=false\nbooleanField=false\n";
901901

902902
CalciteAssert.that()
@@ -915,7 +915,7 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
915915
String plan = "PLAN=ArrowToEnumerableConverter\n"
916916
+ " ArrowProject(booleanField=[$7])\n"
917917
+ " ArrowFilter(condition=[IS NOT TRUE($7)])\n"
918-
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n";
918+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
919919
String result = "booleanField=null\nbooleanField=false\n";
920920

921921
CalciteAssert.that()
@@ -933,7 +933,7 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
933933
String plan = "PLAN=ArrowToEnumerableConverter\n"
934934
+ " ArrowProject(booleanField=[$7])\n"
935935
+ " ArrowFilter(condition=[IS NOT FALSE($7)])\n"
936-
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n";
936+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
937937
String result = "booleanField=null\nbooleanField=true\n";
938938

939939
CalciteAssert.that()
@@ -951,7 +951,7 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
951951
String plan = "PLAN=ArrowToEnumerableConverter\n"
952952
+ " ArrowProject(booleanField=[$7])\n"
953953
+ " ArrowFilter(condition=[IS NULL($7)])\n"
954-
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n";
954+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
955955
String result = "booleanField=null\n";
956956

957957
CalciteAssert.that()
@@ -972,7 +972,7 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
972972
String plan = "PLAN=ArrowToEnumerableConverter\n"
973973
+ " ArrowProject(decimalField=[$8])\n"
974974
+ " ArrowFilter(condition=[=($8, 1.00)])\n"
975-
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n";
975+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
976976
String result = "decimalField=1.00\n";
977977

978978
CalciteAssert.that()
@@ -989,7 +989,7 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
989989
String plan = "PLAN=ArrowToEnumerableConverter\n"
990990
+ " ArrowProject(doubleField=[$6])\n"
991991
+ " ArrowFilter(condition=[=($6, 1.0E0)])\n"
992-
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n";
992+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
993993
String result = "doubleField=1.0\n";
994994

995995
CalciteAssert.that()
@@ -1006,7 +1006,7 @@ static void initializeArrowState(@TempDir Path sharedTempDir)
10061006
String plan = "PLAN=ArrowToEnumerableConverter\n"
10071007
+ " ArrowProject(stringField=[$3])\n"
10081008
+ " ArrowFilter(condition=[=($3, '1')])\n"
1009-
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])\n\n";
1009+
+ " ArrowTableScan(table=[[ARROW, ARROWDATATYPE]], fields=[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])\n\n";
10101010
String result = "stringField=1\n";
10111011

10121012
CalciteAssert.that()

Diff for: ‎arrow/src/test/java/org/apache/calcite/adapter/arrow/ArrowDataTest.java

+19
Original file line numberDiff line numberDiff line change
@@ -31,12 +31,14 @@
3131
import org.apache.arrow.vector.FloatingPointVector;
3232
import org.apache.arrow.vector.IntVector;
3333
import org.apache.arrow.vector.SmallIntVector;
34+
import org.apache.arrow.vector.TimeSecVector;
3435
import org.apache.arrow.vector.TinyIntVector;
3536
import org.apache.arrow.vector.VarCharVector;
3637
import org.apache.arrow.vector.VectorSchemaRoot;
3738
import org.apache.arrow.vector.ipc.ArrowFileWriter;
3839
import org.apache.arrow.vector.types.DateUnit;
3940
import org.apache.arrow.vector.types.FloatingPointPrecision;
41+
import org.apache.arrow.vector.types.TimeUnit;
4042
import org.apache.arrow.vector.types.pojo.ArrowType;
4143
import org.apache.arrow.vector.types.pojo.Field;
4244
import org.apache.arrow.vector.types.pojo.FieldType;
@@ -76,6 +78,7 @@ public class ArrowDataTest {
7678
private double doubleValue;
7779
private boolean booleanValue;
7880
private BigDecimal decimalValue;
81+
private int timeValue;
7982

8083
public ArrowDataTest() {
8184
this.batchSize = 20;
@@ -89,6 +92,7 @@ public ArrowDataTest() {
8992
this.doubleValue = 0;
9093
this.booleanValue = false;
9194
this.decimalValue = BigDecimal.ZERO;
95+
this.timeValue = 0;
9296
}
9397

9498
private Schema makeArrowDateTypeSchema() {
@@ -106,6 +110,7 @@ private Schema makeArrowDateTypeSchema() {
106110
FieldType decimalType = FieldType.nullable(new ArrowType.Decimal(12, 2, 128));
107111
FieldType decimalType2 = FieldType.nullable(new ArrowType.Decimal(12, 3, 128));
108112
FieldType dateType = FieldType.nullable(new ArrowType.Date(DateUnit.DAY));
113+
FieldType timeType = FieldType.nullable(new ArrowType.Time(TimeUnit.SECOND, 32));
109114

110115
childrenBuilder.add(new Field("tinyIntField", tinyIntType, null));
111116
childrenBuilder.add(new Field("smallIntField", smallIntType, null));
@@ -118,6 +123,7 @@ private Schema makeArrowDateTypeSchema() {
118123
childrenBuilder.add(new Field("decimalField", decimalType, null));
119124
childrenBuilder.add(new Field("dateField", dateType, null));
120125
childrenBuilder.add(new Field("decimalField2", decimalType2, null));
126+
childrenBuilder.add(new Field("timeField", timeType, null));
121127

122128
return new Schema(childrenBuilder.build(), null);
123129
}
@@ -273,6 +279,9 @@ public void writeArrowDataType(File file) throws IOException {
273279
case "dateField":
274280
dateField(vector, numRows);
275281
break;
282+
case "timeField":
283+
timeField(vector, numRows);
284+
break;
276285
default:
277286
throw new IllegalStateException("Not supported type yet: " + vector.getMinorType());
278287
}
@@ -411,4 +420,14 @@ private void dateField(FieldVector fieldVector, int rowCount) {
411420
}
412421
fieldVector.setValueCount(rowCount);
413422
}
423+
424+
private void timeField(FieldVector fieldVector, int rowCount) {
425+
TimeSecVector timeVector = (TimeSecVector) fieldVector;
426+
timeVector.setInitialCapacity(rowCount);
427+
timeVector.allocateNew();
428+
for (int i = 0; i < rowCount; i++) {
429+
timeVector.set(i, i * 1000);
430+
}
431+
fieldVector.setValueCount(rowCount);
432+
}
414433
}

0 commit comments

Comments
 (0)