Skip to content

Commit 4f54a2c

Browse files
wecharyuwecharyu
committed
[GLUTEN-12447] Add Arrow large var types support
1 parent 065c073 commit 4f54a2c

5 files changed

Lines changed: 240 additions & 19 deletions

File tree

‎backends-velox/src/test/scala/org/apache/gluten/expression/UDFPartialProjectSuite.scala‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -158,8 +158,13 @@ class UDFPartialProjectSuite extends WholeStageTransformerSuite {
158158
}
159159

160160
test("test concat with string") {
161-
runQueryAndCompare("SELECT concat_concat(l_comment), hash(l_partkey) from lineitem") {
162-
checkGlutenPlan[ColumnarPartialProjectExec]
161+
Seq("false", "true").foreach {
162+
useLargeVarTypes =>
163+
withSQLConf("spark.sql.execution.arrow.useLargeVarTypes" -> useLargeVarTypes) {
164+
runQueryAndCompare("SELECT concat_concat(l_comment), hash(l_partkey) from lineitem") {
165+
checkGlutenPlan[ColumnarPartialProjectExec]
166+
}
167+
}
163168
}
164169
}
165170

‎gluten-arrow/src/main/java/org/apache/gluten/vectorized/ArrowColumnVector.java‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import org.apache.arrow.vector.complex.ListVector;
2121
import org.apache.arrow.vector.complex.MapVector;
2222
import org.apache.arrow.vector.complex.StructVector;
23+
import org.apache.arrow.vector.holders.NullableLargeVarCharHolder;
2324
import org.apache.arrow.vector.holders.NullableVarCharHolder;
2425
import org.apache.spark.annotation.DeveloperApi;
2526
import org.apache.spark.sql.types.DataType;
@@ -166,8 +167,12 @@ void initAccessor(ValueVector vector) {
166167
accessor = new DecimalAccessor((DecimalVector) vector);
167168
} else if (vector instanceof VarCharVector) {
168169
accessor = new StringAccessor((VarCharVector) vector);
170+
} else if (vector instanceof LargeVarCharVector) {
171+
accessor = new LargeStringAccessor((LargeVarCharVector) vector);
169172
} else if (vector instanceof VarBinaryVector) {
170173
accessor = new BinaryAccessor((VarBinaryVector) vector);
174+
} else if (vector instanceof LargeVarBinaryVector) {
175+
accessor = new LargeBinaryAccessor((LargeVarBinaryVector) vector);
171176
} else if (vector instanceof DateDayVector) {
172177
accessor = new DateAccessor((DateDayVector) vector);
173178
} else if (vector instanceof TimeStampMicroTZVector) {
@@ -414,6 +419,31 @@ final UTF8String getUTF8String(int rowId) {
414419
}
415420
}
416421

422+
static class LargeStringAccessor extends ArrowVectorAccessor {
423+
424+
private final LargeVarCharVector accessor;
425+
private final NullableLargeVarCharHolder stringResult = new NullableLargeVarCharHolder();
426+
427+
LargeStringAccessor(LargeVarCharVector vector) {
428+
super(vector);
429+
this.accessor = vector;
430+
}
431+
432+
@Override
433+
final UTF8String getUTF8String(int rowId) {
434+
accessor.get(rowId, stringResult);
435+
if (stringResult.isSet == 0) {
436+
return null;
437+
} else {
438+
return UTF8String.fromAddress(
439+
null,
440+
stringResult.buffer.memoryAddress() + stringResult.start,
441+
// A single string cannot be larger than the max integer size, so the conversion is safe
442+
(int) (stringResult.end - stringResult.start));
443+
}
444+
}
445+
}
446+
417447
static class BinaryAccessor extends ArrowVectorAccessor {
418448

419449
private final VarBinaryVector accessor;
@@ -429,6 +459,21 @@ final byte[] getBinary(int rowId) {
429459
}
430460
}
431461

462+
static class LargeBinaryAccessor extends ArrowVectorAccessor {
463+
464+
private final LargeVarBinaryVector accessor;
465+
466+
LargeBinaryAccessor(LargeVarBinaryVector vector) {
467+
super(vector);
468+
this.accessor = vector;
469+
}
470+
471+
@Override
472+
final byte[] getBinary(int rowId) {
473+
return accessor.getObject(rowId);
474+
}
475+
}
476+
432477
static class DateAccessor extends ArrowVectorAccessor {
433478

434479
private final DateDayVector accessor;

‎gluten-arrow/src/main/java/org/apache/gluten/vectorized/ArrowWritableColumnVector.java‎

Lines changed: 152 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@
2323
import org.apache.arrow.vector.complex.ListVector;
2424
import org.apache.arrow.vector.complex.MapVector;
2525
import org.apache.arrow.vector.complex.StructVector;
26+
import org.apache.arrow.vector.holders.NullableLargeVarBinaryHolder;
27+
import org.apache.arrow.vector.holders.NullableLargeVarCharHolder;
2628
import org.apache.arrow.vector.holders.NullableVarBinaryHolder;
2729
import org.apache.arrow.vector.holders.NullableVarCharHolder;
2830
import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
@@ -78,8 +80,7 @@ public final class ArrowWritableColumnVector extends WritableColumnVectorShim {
7880
* elements, not number of bytes.
7981
*/
8082
public static ArrowWritableColumnVector[] allocateColumns(int capacity, StructType schema) {
81-
String timeZoneId = SparkSchemaUtil.getLocalTimezoneID();
82-
Schema arrowSchema = SparkArrowUtil.toArrowSchema(schema, timeZoneId);
83+
Schema arrowSchema = SparkSchemaUtil.toArrowSchema(schema);
8384
VectorSchemaRoot newRoot =
8485
VectorSchemaRoot.create(arrowSchema, ArrowBufferAllocators.contextInstance());
8586

@@ -94,8 +95,7 @@ public static ArrowWritableColumnVector[] allocateColumns(int capacity, StructTy
9495

9596
public static ArrowWritableColumnVector[] allocateColumns(
9697
int valueCount, long[] totalSizes, StructType schema) {
97-
String timeZoneId = SparkSchemaUtil.getLocalTimezoneID();
98-
Schema arrowSchema = SparkArrowUtil.toArrowSchema(schema, timeZoneId);
98+
Schema arrowSchema = SparkSchemaUtil.toArrowSchema(schema);
9999
VectorSchemaRoot newRoot =
100100
VectorSchemaRoot.create(arrowSchema, ArrowBufferAllocators.contextInstance());
101101

@@ -208,8 +208,10 @@ public ArrowWritableColumnVector(int capacity, DataType dataType) {
208208
vectorCount.getAndIncrement();
209209
refCnt.getAndIncrement();
210210
String timeZoneId = SparkSchemaUtil.getLocalTimezoneID();
211+
boolean useLargeVarTypes = SparkSchemaUtil.enableLargeVarTypes();
211212
List<Field> fields =
212-
Arrays.asList(SparkArrowUtil.toArrowField("col", dataType, true, timeZoneId));
213+
Arrays.asList(
214+
SparkArrowUtil.toArrowField("col", dataType, true, timeZoneId, useLargeVarTypes));
213215
Schema arrowSchema = new Schema(fields);
214216
VectorSchemaRoot root =
215217
VectorSchemaRoot.create(arrowSchema, ArrowBufferAllocators.contextInstance());
@@ -265,8 +267,12 @@ private void createVectorAccessor(ValueVector vector, ValueVector dictionary) {
265267
accessor = new DecimalAccessor((DecimalVector) vector);
266268
} else if (vector instanceof VarCharVector) {
267269
accessor = new StringAccessor((VarCharVector) vector);
270+
} else if (vector instanceof LargeVarCharVector) {
271+
accessor = new LargeStringAccessor((LargeVarCharVector) vector);
268272
} else if (vector instanceof VarBinaryVector) {
269273
accessor = new BinaryAccessor((VarBinaryVector) vector);
274+
} else if (vector instanceof LargeVarBinaryVector) {
275+
accessor = new LargeBinaryAccessor((LargeVarBinaryVector) vector);
270276
} else if (vector instanceof DateDayVector) {
271277
accessor = new DateAccessor((DateDayVector) vector);
272278
} else if (vector instanceof TimeStampMicroVector || vector instanceof TimeStampMicroTZVector) {
@@ -336,8 +342,12 @@ private ArrowVectorWriter createVectorWriter(ValueVector vector) {
336342
return new DecimalWriter((DecimalVector) vector);
337343
} else if (vector instanceof VarCharVector) {
338344
return new StringWriter((VarCharVector) vector);
345+
} else if (vector instanceof LargeVarCharVector) {
346+
return new LargeStringWriter((LargeVarCharVector) vector);
339347
} else if (vector instanceof VarBinaryVector) {
340348
return new BinaryWriter((VarBinaryVector) vector);
349+
} else if (vector instanceof LargeVarBinaryVector) {
350+
return new LargeBinaryWriter((LargeVarBinaryVector) vector);
341351
} else if (vector instanceof DateDayVector) {
342352
return new DateWriter((DateDayVector) vector);
343353
} else if (vector instanceof TimeStampMicroVector || vector instanceof TimeStampMicroTZVector) {
@@ -1115,6 +1125,30 @@ final UTF8String getUTF8String(int rowId) {
11151125
}
11161126
}
11171127

1128+
private static class LargeStringAccessor extends ArrowVectorAccessor {
1129+
private final LargeVarCharVector accessor;
1130+
private final NullableLargeVarCharHolder stringResult = new NullableLargeVarCharHolder();
1131+
1132+
LargeStringAccessor(LargeVarCharVector vector) {
1133+
super(vector);
1134+
this.accessor = vector;
1135+
}
1136+
1137+
@Override
1138+
final UTF8String getUTF8String(int rowId) {
1139+
accessor.get(rowId, stringResult);
1140+
if (stringResult.isSet == 0) {
1141+
return null;
1142+
} else {
1143+
return UTF8String.fromAddress(
1144+
null,
1145+
stringResult.buffer.memoryAddress() + stringResult.start,
1146+
// A single string cannot be larger than the max integer size, so the conversion is safe
1147+
(int) (stringResult.end - stringResult.start));
1148+
}
1149+
}
1150+
}
1151+
11181152
private static class DictionaryEncodedStringAccessor extends ArrowVectorAccessor {
11191153
private final IntVector index;
11201154
private final VarCharVector dictionary;
@@ -1169,6 +1203,35 @@ final UTF8String getUTF8String(int rowId) {
11691203
}
11701204
}
11711205

1206+
private static class LargeBinaryAccessor extends ArrowVectorAccessor {
1207+
private final LargeVarBinaryVector accessor;
1208+
private final NullableLargeVarBinaryHolder stringResult = new NullableLargeVarBinaryHolder();
1209+
1210+
LargeBinaryAccessor(LargeVarBinaryVector vector) {
1211+
super(vector);
1212+
this.accessor = vector;
1213+
}
1214+
1215+
@Override
1216+
final byte[] getBinary(int rowId) {
1217+
return accessor.getObject(rowId);
1218+
}
1219+
1220+
@Override
1221+
final UTF8String getUTF8String(int rowId) {
1222+
accessor.get(rowId, stringResult);
1223+
if (stringResult.isSet == 0) {
1224+
return null;
1225+
} else {
1226+
return UTF8String.fromAddress(
1227+
null,
1228+
stringResult.buffer.memoryAddress() + stringResult.start,
1229+
// A single string cannot be larger than the max integer size, so the conversion is safe
1230+
(int) (stringResult.end - stringResult.start));
1231+
}
1232+
}
1233+
}
1234+
11721235
private static class DictionaryEncodedBinaryAccessor extends ArrowVectorAccessor {
11731236
private final IntVector index;
11741237
private final VarBinaryVector dictionary;
@@ -2014,6 +2077,52 @@ void unsafeSetValueNullSafe(SpecializedGetters input, int ordinal) {
20142077
}
20152078
}
20162079

2080+
private static class LargeStringWriter extends ArrowVectorWriter {
2081+
private final LargeVarCharVector writer;
2082+
private int rowId;
2083+
2084+
LargeStringWriter(LargeVarCharVector vector) {
2085+
super(vector);
2086+
this.writer = vector;
2087+
this.rowId = 0;
2088+
}
2089+
2090+
@Override
2091+
final void setNull(int rowId) {
2092+
writer.setNull(rowId);
2093+
}
2094+
2095+
@Override
2096+
final void setNulls(int rowId, int count) {
2097+
for (int i = 0; i < count; i++) {
2098+
writer.setNull(rowId + i);
2099+
}
2100+
}
2101+
2102+
@Override
2103+
final void setBytes(int rowId, int count, byte[] src, int srcIndex) {
2104+
writer.setSafe(rowId, src, srcIndex, count);
2105+
}
2106+
2107+
@Override
2108+
final void appendBytes(byte[] value, int offset, int length) {
2109+
writer.setSafe(rowId, value, offset, length);
2110+
rowId++;
2111+
}
2112+
2113+
@Override
2114+
void setValueNullSafe(SpecializedGetters input, int ordinal) {
2115+
UTF8String value = input.getUTF8String(ordinal);
2116+
setBytes(count, value.numBytes(), value.getBytes(), 0);
2117+
}
2118+
2119+
@Override
2120+
void unsafeSetValueNullSafe(SpecializedGetters input, int ordinal) {
2121+
UTF8String value = input.getUTF8String(ordinal);
2122+
writer.set(count, value.getBytes(), 0, value.numBytes());
2123+
}
2124+
}
2125+
20172126
private static class BinaryWriter extends ArrowVectorWriter {
20182127
private final VarBinaryVector writer;
20192128

@@ -2052,6 +2161,44 @@ void unsafeSetValueNullSafe(SpecializedGetters input, int ordinal) {
20522161
}
20532162
}
20542163

2164+
private static class LargeBinaryWriter extends ArrowVectorWriter {
2165+
private final LargeVarBinaryVector writer;
2166+
2167+
LargeBinaryWriter(LargeVarBinaryVector vector) {
2168+
super(vector);
2169+
this.writer = vector;
2170+
}
2171+
2172+
@Override
2173+
final void setNull(int rowId) {
2174+
writer.setNull(rowId);
2175+
}
2176+
2177+
@Override
2178+
final void setNulls(int rowId, int count) {
2179+
for (int i = 0; i < count; i++) {
2180+
writer.setNull(rowId + i);
2181+
}
2182+
}
2183+
2184+
@Override
2185+
final void setBytes(int rowId, int count, byte[] src, int srcIndex) {
2186+
writer.setSafe(rowId, src, srcIndex, count);
2187+
}
2188+
2189+
@Override
2190+
void setValueNullSafe(SpecializedGetters input, int ordinal) {
2191+
byte[] value = input.getBinary(ordinal);
2192+
setBytes(count, value.length, value, 0);
2193+
}
2194+
2195+
@Override
2196+
void unsafeSetValueNullSafe(SpecializedGetters input, int ordinal) {
2197+
byte[] value = input.getBinary(ordinal);
2198+
writer.set(count, value, 0, value.length);
2199+
}
2200+
}
2201+
20552202
private static class DateWriter extends ArrowVectorWriter {
20562203
private final DateDayVector writer;
20572204

0 commit comments

Comments
 (0)