lidavidm commented on a change in pull request #11982:
URL: https://github.com/apache/arrow/pull/11982#discussion_r823881866
##########
File path:
java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/example/FlightSqlExample.java
##########
@@ -375,29 +388,66 @@ private static VectorSchemaRoot getSchemasRoot(final
ResultSet data, final Buffe
private static <T extends FieldVector> int saveToVectors(final Map<T,
String> vectorToColumnName,
final ResultSet
data, boolean emptyToNull)
throws SQLException {
+ Predicate<ResultSet> alwaysTrue = (resultSet) -> true;
+ return saveToVectors(vectorToColumnName, data, emptyToNull, alwaysTrue);
+ }
+
+ private static <T extends FieldVector> int saveToVectors(final Map<T,
String> vectorToColumnName,
+ final ResultSet
data, boolean emptyToNull,
+
Predicate<ResultSet> resultSetPredicate)
+ throws SQLException {
Objects.requireNonNull(vectorToColumnName, "vectorToColumnName cannot be
null.");
Objects.requireNonNull(data, "data cannot be null.");
final Set<Entry<T, String>> entrySet = vectorToColumnName.entrySet();
int rows = 0;
- for (; data.next(); rows++) {
+
+ while (data.next()) {
+ if (!resultSetPredicate.test(data)) {
+ continue;
+ }
for (final Entry<T, String> vectorToColumn : entrySet) {
final T vector = vectorToColumn.getKey();
final String columnName = vectorToColumn.getValue();
if (vector instanceof VarCharVector) {
String thisData = data.getString(columnName);
saveToVector(emptyToNull ? emptyToNull(thisData) : thisData,
(VarCharVector) vector, rows);
- continue;
} else if (vector instanceof IntVector) {
final int intValue = data.getInt(columnName);
saveToVector(data.wasNull() ? null : intValue, (IntVector) vector,
rows);
- continue;
} else if (vector instanceof UInt1Vector) {
final byte byteValue = data.getByte(columnName);
saveToVector(data.wasNull() ? null : byteValue, (UInt1Vector)
vector, rows);
- continue;
+ } else if (vector instanceof BitVector) {
+ final byte byteValue = data.getByte(columnName);
+ saveToVector(data.wasNull() ? null : byteValue, (BitVector) vector,
rows);
+ } else if (vector instanceof ListVector) {
+ String createParamsValues = data.getString(columnName);
+
+ UnionListWriter writer = ((ListVector) vector).getWriter();
+
+ BufferAllocator allocator = vector.getAllocator();
+ final ArrowBuf buf = allocator.buffer(1024);
+
+ writer.setPosition(rows);
+ writer.startList();
+
+ if (createParamsValues != null) {
+ String[] split = createParamsValues.split(",");
+
+ range(0, split.length)
+ .forEach(i -> {
+ byte[] bytes = split[i].getBytes(UTF_8);
+ buf.setBytes(0, bytes);
+ buf.writerIndex(1);
+ writer.varChar().writeVarChar(0, bytes.length, buf);
Review comment:
Filed https://issues.apache.org/jira/browse/ARROW-15907
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
