org.codehaus.mojo
diff --git a/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java b/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java
index cf50f9417b..f51fb25105 100644
--- a/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java
+++ b/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java
@@ -64,13 +64,18 @@ void release() {
*/
ArrowBuf unsafeAssociateAllocation(
BufferAllocator trackingAllocator, long capacity, long memoryAddress) {
+ // Retain only after wrapForeignAllocation succeeds. On the allocator-limit OOM path,
+ // wrapForeignAllocation throws before the ForeignAllocation is associated, so release0()
+ // is not called; retaining first would leave the count elevated with no matching release0().
+ ArrowBuf buf =
+ trackingAllocator.wrapForeignAllocation(
+ new ForeignAllocation(capacity, memoryAddress) {
+ @Override
+ protected void release0() {
+ ReferenceCountedArrowArray.this.release();
+ }
+ });
retain();
- return trackingAllocator.wrapForeignAllocation(
- new ForeignAllocation(capacity, memoryAddress) {
- @Override
- protected void release0() {
- ReferenceCountedArrowArray.this.release();
- }
- });
+ return buf;
}
}
diff --git a/c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java b/c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java
new file mode 100644
index 0000000000..7c099f2ef0
--- /dev/null
+++ b/c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java
@@ -0,0 +1,139 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.arrow.c;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.OutOfMemoryException;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.types.pojo.Schema;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Regression test: a mid-import {@link OutOfMemoryException} must not leak the imported array.
+ *
+ * A "producer" allocator owns the exported batch; if the C Data release callback fires, the
+ * producer drains to zero. A too-small consumer allocator forces an OOM part-way through the
+ * import. The test asserts the producer drains, confirming the release callback fired despite the
+ * failure.
+ */
+final class ImportOutOfMemoryTest {
+ private static final int ROWS = 1024;
+ private static final int VALUE_BYTES = 256;
+ private static final int COLUMNS = 4;
+ // Far smaller than the exported batch, so the import OOMs part-way through the buffers.
+ private static final long TINY_LIMIT = 16 * 1024;
+
+ private RootAllocator root;
+
+ @BeforeEach
+ public void setUp() {
+ root = new RootAllocator(Long.MAX_VALUE);
+ }
+
+ @AfterEach
+ public void tearDown() {
+ root.close();
+ }
+
+ @Test
+ public void importOomDoesNotLeakExportedArray() {
+ // "producer" owns only the exported batch buffers; the C Data struct containers live on a
+ // separate allocator (they are consumed/closed by import, which would otherwise muddy the
+ // producer's balance). So producer draining to zero is an exact signal that the array's release
+ // callback fired.
+ try (BufferAllocator producer = root.newChildAllocator("producer", 0, Long.MAX_VALUE);
+ BufferAllocator structs = root.newChildAllocator("structs", 0, Long.MAX_VALUE)) {
+ try (ArrowArray array = ArrowArray.allocateNew(structs);
+ ArrowSchema schema = ArrowSchema.allocateNew(structs)) {
+ exportBatch(producer, array, schema);
+ assertTrue(
+ producer.getAllocatedMemory() > 0, "producer holds the exported batch before import");
+
+ // A consumer allocator far too small to hold the batch: the import throws part-way through.
+ try (BufferAllocator consumer = root.newChildAllocator("consumer", 0, TINY_LIMIT);
+ CDataDictionaryProvider provider = new CDataDictionaryProvider()) {
+ Schema importSchema = Data.importSchema(consumer, schema, provider);
+ try (VectorSchemaRoot importRoot = VectorSchemaRoot.create(importSchema, consumer)) {
+ Exception thrown =
+ assertThrows(
+ Exception.class,
+ () -> Data.importIntoVectorSchemaRoot(consumer, array, importRoot, provider));
+ assertTrue(
+ hasOutOfMemoryCause(thrown),
+ "mid-import failure must be an allocator OOM: " + thrown);
+ }
+ }
+
+ // The array's release callback must have fired despite the mid-import OOM, freeing the
+ // whole exported batch. On the unfixed retain-before-wrap code the batch is stranded.
+ assertEquals(
+ 0L,
+ producer.getAllocatedMemory(),
+ "import OOM leaked the exported batch (producer not drained)");
+ }
+ }
+ }
+
+ /** True if {@code t} is, or is caused by, an Arrow {@link OutOfMemoryException}. */
+ private static boolean hasOutOfMemoryCause(Throwable t) {
+ for (Throwable cause = t; cause != null; cause = cause.getCause()) {
+ if (cause instanceof OutOfMemoryException) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ /**
+ * Builds a wide multi-column VarChar batch on {@code alloc} and exports it into the C structs.
+ */
+ private void exportBatch(BufferAllocator alloc, ArrowArray array, ArrowSchema schema) {
+ byte[] value = new byte[VALUE_BYTES];
+ for (int i = 0; i < value.length; i++) {
+ value[i] = (byte) 'x';
+ }
+ List vectors = new ArrayList<>(COLUMNS);
+ for (int c = 0; c < COLUMNS; c++) {
+ VarCharVector vector = new VarCharVector("col" + c, alloc);
+ vector.allocateNew((long) ROWS * VALUE_BYTES, ROWS);
+ for (int r = 0; r < ROWS; r++) {
+ vector.setSafe(r, value);
+ }
+ vector.setValueCount(ROWS);
+ vectors.add(vector);
+ }
+ try (VectorSchemaRoot source = new VectorSchemaRoot(vectors)) {
+ long total = 0;
+ for (FieldVector vector : source.getFieldVectors()) {
+ total += vector.getBufferSize();
+ }
+ assertTrue(total > TINY_LIMIT, "test setup: batch must exceed the consumer limit");
+ Data.exportVectorSchemaRoot(alloc, source, null, array, schema);
+ }
+ }
+}
diff --git a/ci/scripts/test.sh b/ci/scripts/test.sh
index cacc20034e..8061ee455d 100755
--- a/ci/scripts/test.sh
+++ b/ci/scripts/test.sh
@@ -34,10 +34,11 @@ fi
mvn="mvn -B -Dorg.slf4j.simpleLogger.log.org.apache.maven.cli.transfer.Slf4jMavenTransferListener=warn"
# Use `2 * ncores` threads
mvn="${mvn} -T 2C"
+mvn="${mvn} -Denforcer.skip=true"
pushd "${build_dir}"
-${mvn} -Darrow.test.dataRoot="${source_dir}/testing/data" clean test
+${mvn} -Darrow.test.dataRoot="${source_dir}/testing/data" test
projects=()
if [ "${ARROW_JAVA_JNI}" = "ON" ]; then
@@ -46,7 +47,7 @@ if [ "${ARROW_JAVA_JNI}" = "ON" ]; then
projects+=(gandiva)
fi
if [ "${#projects[@]}" -gt 0 ]; then
- ${mvn} clean test \
+ ${mvn} test \
-Parrow-jni \
-pl "$(
IFS=,
@@ -56,7 +57,7 @@ if [ "${#projects[@]}" -gt 0 ]; then
fi
if [ "${ARROW_JAVA_CDATA}" = "ON" ]; then
- ${mvn} clean test -Parrow-c-data -pl c -Darrow.c.jni.dist.dir="${java_jni_dist_dir}"
+ ${mvn} test -Parrow-c-data -pl c -Darrow.c.jni.dist.dir="${java_jni_dist_dir}"
fi
popd
diff --git a/compression/pom.xml b/compression/pom.xml
index f3de6fc248..aa7dee6f89 100644
--- a/compression/pom.xml
+++ b/compression/pom.xml
@@ -55,7 +55,7 @@ under the License.
com.github.luben
zstd-jni
- 1.5.7-8
+ 1.5.7-11
diff --git a/compression/src/main/java/org/apache/arrow/compression/Lz4CompressionCodec.java b/compression/src/main/java/org/apache/arrow/compression/Lz4CompressionCodec.java
index 91cefc2a9e..f268e815fe 100644
--- a/compression/src/main/java/org/apache/arrow/compression/Lz4CompressionCodec.java
+++ b/compression/src/main/java/org/apache/arrow/compression/Lz4CompressionCodec.java
@@ -80,6 +80,13 @@ protected ArrowBuf doDecompress(BufferAllocator allocator, ArrowBuf compressedBu
}
byte[] outBytes = out.toByteArray();
+ if (outBytes.length != decompressedLength) {
+ throw new RuntimeException(
+ "Expected != actual decompressed length: "
+ + decompressedLength
+ + " != "
+ + outBytes.length);
+ }
ArrowBuf decompressedBuffer = allocator.buffer(outBytes.length);
decompressedBuffer.setBytes(/* index= */ 0, outBytes);
decompressedBuffer.writerIndex(decompressedLength);
diff --git a/compression/src/main/java/org/apache/arrow/compression/ZstdCompressionCodec.java b/compression/src/main/java/org/apache/arrow/compression/ZstdCompressionCodec.java
index 290723608d..ed46fe81b4 100644
--- a/compression/src/main/java/org/apache/arrow/compression/ZstdCompressionCodec.java
+++ b/compression/src/main/java/org/apache/arrow/compression/ZstdCompressionCodec.java
@@ -44,7 +44,7 @@ protected ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf uncompressedBu
long bytesWritten =
Zstd.compressUnsafe(
compressedBuffer.memoryAddress() + CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH,
- dstSize,
+ maxSize,
/*src*/ uncompressedBuffer.memoryAddress(),
/* srcSize= */ uncompressedBuffer.writerIndex(),
/* level= */ this.compressionLevel);
diff --git a/compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java b/compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java
index b8fb4e28b9..d2d2921649 100644
--- a/compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java
+++ b/compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java
@@ -20,6 +20,7 @@
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import java.io.ByteArrayOutputStream;
@@ -59,6 +60,7 @@
import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
@@ -231,6 +233,26 @@ void testEmptyBuffer(int vectorLength, CompressionCodec codec) throws Exception
AutoCloseables.close(decompressedBuffers);
}
+ @Test
+ void testLz4DecompressRejectsWrongLength() {
+ byte[] data = new byte[512]; // all zeros, highly compressible
+ ArrowBuf orig = allocator.buffer(data.length);
+ orig.setBytes(0, data);
+ orig.writerIndex(data.length);
+
+ CompressionCodec codec = new Lz4CompressionCodec();
+ ArrowBuf compressed = codec.compress(allocator, orig);
+
+ // tamper with the 8-byte uncompressed-length prefix so it no longer matches
+ // the real decompressed size
+ compressed.setLong(0, 1_000_000L);
+
+ RuntimeException e =
+ assertThrows(RuntimeException.class, () -> codec.decompress(allocator, compressed));
+ assertTrue(e.getMessage().contains("decompressed length"));
+ compressed.close();
+ }
+
private static Stream codecTypes() {
return Arrays.stream(CompressionUtil.CodecType.values());
}
diff --git a/flight/flight-core/pom.xml b/flight/flight-core/pom.xml
index 15b870905d..9ae402cdc4 100644
--- a/flight/flight-core/pom.xml
+++ b/flight/flight-core/pom.xml
@@ -134,7 +134,7 @@ under the License.
com.google.api.grpc
proto-google-common-protos
- 2.71.0
+ 2.72.0
test
diff --git a/flight/flight-sql-jdbc-core/pom.xml b/flight/flight-sql-jdbc-core/pom.xml
index 96a11f2b9a..be2ee32868 100644
--- a/flight/flight-sql-jdbc-core/pom.xml
+++ b/flight/flight-sql-jdbc-core/pom.xml
@@ -123,19 +123,19 @@ under the License.
com.squareup.okhttp3
mockwebserver3
- 5.3.2
+ 5.4.0
test
com.squareup.okhttp3
mockwebserver3-junit5
- 5.3.2
+ 5.4.0
test
com.squareup.okhttp3
okhttp-jvm
- 5.3.2
+ 5.4.0
test
@@ -182,7 +182,7 @@ under the License.
com.nimbusds
oauth2-oidc-sdk
- 11.37.1
+ 11.37.2
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java
index 2885f7895b..376e5b11e7 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java
@@ -106,7 +106,7 @@ static ArrowFlightJdbcFlightStreamResultSet fromFlightInfo(
final TimeZone timeZone = TimeZone.getDefault();
final QueryState state = new QueryState();
- final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null);
+ final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null, null);
final AvaticaResultSetMetaData resultSetMetaData =
new AvaticaResultSetMetaData(null, null, signature);
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java
index 5d02d6e843..ad6670a001 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java
@@ -73,7 +73,7 @@ public static ArrowFlightJdbcVectorSchemaRootResultSet fromVectorSchemaRoot(
final TimeZone timeZone = TimeZone.getDefault();
final QueryState state = new QueryState();
- final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null);
+ final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null, null);
final AvaticaResultSetMetaData resultSetMetaData =
new AvaticaResultSetMetaData(null, null, signature);
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java
index 64529b50c8..0d85b5eddb 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java
@@ -53,7 +53,8 @@ public ArrowFlightMetaImpl(final AvaticaConnection connection) {
}
/** Construct a signature. */
- static Signature newSignature(final String sql, Schema resultSetSchema, Schema parameterSchema) {
+ static Signature newSignature(
+ final String sql, Schema resultSetSchema, Schema parameterSchema, Boolean isUpdate) {
List columnMetaData =
resultSetSchema == null
? new ArrayList<>()
@@ -62,10 +63,17 @@ static Signature newSignature(final String sql, Schema resultSetSchema, Schema p
parameterSchema == null
? new ArrayList<>()
: ConvertUtils.convertArrowFieldsToAvaticaParameters(parameterSchema.getFields());
- StatementType statementType =
- resultSetSchema == null || resultSetSchema.getFields().isEmpty()
- ? StatementType.IS_DML
- : StatementType.SELECT;
+ // If the server provided the is_update field, use it to determine the statement type
+ StatementType statementType;
+ if (isUpdate != null) {
+ statementType = isUpdate ? StatementType.IS_DML : StatementType.SELECT;
+ } else {
+ // Fall back to the legacy logic: check if the result set schema is empty
+ statementType =
+ resultSetSchema == null || resultSetSchema.getFields().isEmpty()
+ ? StatementType.IS_DML
+ : StatementType.SELECT;
+ }
return new Signature(
columnMetaData,
sql,
@@ -178,7 +186,10 @@ private PreparedStatement prepareForHandle(final String query, StatementHandle h
((ArrowFlightConnection) connection).getClientHandler().prepare(query);
handle.signature =
newSignature(
- query, preparedStatement.getDataSetSchema(), preparedStatement.getParameterSchema());
+ query,
+ preparedStatement.getDataSetSchema(),
+ preparedStatement.getParameterSchema(),
+ preparedStatement.isUpdate());
statementHandlePreparedStatementMap.put(new StatementHandleKey(handle), preparedStatement);
return preparedStatement;
}
diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java
index f0ea284239..719cc38a2b 100644
--- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java
+++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java
@@ -388,6 +388,14 @@ public interface PreparedStatement extends AutoCloseable {
*/
Schema getParameterSchema();
+ /**
+ * Gets whether this {@link PreparedStatement} is an update statement.
+ *
+ * @return {@code true} if this is an update statement, {@code false} if it's a query, or {@code
+ * null} if the server did not provide this information.
+ */
+ @Nullable Boolean isUpdate();
+
void setParameters(VectorSchemaRoot parameters);
@Override
@@ -456,6 +464,12 @@ public long executeUpdate() {
@Override
public StatementType getType() {
+ // If the server provided the is_update field, use it to determine the statement type
+ final Boolean isUpdate = preparedStatement.isUpdate();
+ if (isUpdate != null) {
+ return isUpdate ? StatementType.UPDATE : StatementType.SELECT;
+ }
+ // Fall back to the legacy logic: check if the result set schema is empty
final Schema schema = preparedStatement.getResultSetSchema();
return schema.getFields().isEmpty() ? StatementType.UPDATE : StatementType.SELECT;
}
@@ -475,6 +489,11 @@ public void setParameters(VectorSchemaRoot parameters) {
preparedStatement.setParameters(parameters);
}
+ @Override
+ public Boolean isUpdate() {
+ return preparedStatement.isUpdate();
+ }
+
@Override
public void close() {
try {
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java
index 0369c3a162..078837adf3 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java
@@ -98,6 +98,40 @@ public void testSimpleQueryNoParameterBindingWithExecute() throws SQLException {
}
}
+ @Test
+ public void testSimpleQueryNoParameterBindingWithExecuteV2() throws SQLException {
+ final String query = "SELECT * FROM TEST_V2";
+ final Schema schema =
+ new Schema(Collections.singletonList(Field.nullable("", Types.MinorType.INT.getType())));
+ PRODUCER.addSelectQuery(
+ query,
+ schema,
+ Collections.singletonList(
+ listener -> {
+ try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
+ final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
+ root.allocateNew();
+ ((IntVector) root.getVector(0)).setSafe(0, 123);
+ root.setRowCount(1);
+ listener.start(root);
+ listener.putNext();
+ } finally {
+ listener.completed();
+ }
+ }),
+ false);
+ try (final PreparedStatement preparedStatement = connection.prepareStatement(query)) {
+ boolean isResultSet = preparedStatement.execute();
+ assertTrue(isResultSet);
+ final ResultSet resultSet = preparedStatement.getResultSet();
+ assertTrue(resultSet.next());
+ assertEquals(123, resultSet.getInt(1));
+ assertFalse(resultSet.next());
+ assertFalse(preparedStatement.getMoreResults());
+ assertEquals(-1, preparedStatement.getUpdateCount());
+ }
+ }
+
@Test
public void testQueryWithParameterBinding() throws SQLException {
final String query = "Fake query with parameters";
@@ -203,6 +237,20 @@ public void testUpdateQueryWithExecute() throws SQLException {
}
}
+ @Test
+ public void testUpdateQueryWithExecuteV2() throws SQLException {
+ String query = "Fake update with execute V2";
+ PRODUCER.addUpdateQuery(query, /*updatedRows*/ 99, true);
+ try (final PreparedStatement stmt = connection.prepareStatement(query)) {
+ boolean isResultSet = stmt.execute();
+ assertFalse(isResultSet);
+ int updated = stmt.getUpdateCount();
+ assertEquals(99, updated);
+ assertFalse(stmt.getMoreResults());
+ assertEquals(-1, stmt.getUpdateCount());
+ }
+ }
+
@Test
public void testUpdateQueryWithParameters() throws SQLException {
String query = "Fake update with parameters";
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteTest.java
index 632cb0ba56..6acce9c2a6 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteTest.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteTest.java
@@ -62,6 +62,9 @@ public class ArrowFlightStatementExecuteTest {
private static final String SAMPLE_LARGE_UPDATE_QUERY =
"UPDATE this_large_table SET this_large_field = that_large_field FROM this_large_test WHERE this_large_condition";
private static final long SAMPLE_LARGE_UPDATE_COUNT = Long.MAX_VALUE;
+ private static final String SAMPLE_QUERY_CMD_V2 = "SELECT * FROM this_test_v2";
+ private static final String SAMPLE_LARGE_UPDATE_QUERY_V2 =
+ "UPDATE this_large_table_v2 SET this_large_field = that_large_field FROM this_large_test WHERE this_large_condition";
private static final MockFlightSqlProducer PRODUCER = new MockFlightSqlProducer();
@RegisterExtension
@@ -96,6 +99,31 @@ public static void setUpBeforeClass() {
}));
PRODUCER.addUpdateQuery(SAMPLE_UPDATE_QUERY, SAMPLE_UPDATE_COUNT);
PRODUCER.addUpdateQuery(SAMPLE_LARGE_UPDATE_QUERY, SAMPLE_LARGE_UPDATE_COUNT);
+
+ // V2 queries with is_update field set
+ PRODUCER.addSelectQuery(
+ SAMPLE_QUERY_CMD_V2,
+ SAMPLE_QUERY_SCHEMA,
+ Collections.singletonList(
+ listener -> {
+ try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
+ final VectorSchemaRoot root =
+ VectorSchemaRoot.create(SAMPLE_QUERY_SCHEMA, allocator)) {
+ final UInt1Vector vector = (UInt1Vector) root.getVector(VECTOR_NAME);
+ IntStream.range(0, SAMPLE_QUERY_ROWS)
+ .forEach(index -> vector.setSafe(index, index));
+ vector.setValueCount(SAMPLE_QUERY_ROWS);
+ root.setRowCount(SAMPLE_QUERY_ROWS);
+ listener.start(root);
+ listener.putNext();
+ } catch (final Throwable throwable) {
+ listener.error(throwable);
+ } finally {
+ listener.completed();
+ }
+ }),
+ false);
+ PRODUCER.addUpdateQuery(SAMPLE_LARGE_UPDATE_QUERY_V2, SAMPLE_LARGE_UPDATE_COUNT, true);
}
@BeforeEach
@@ -168,4 +196,42 @@ public void testUpdateCountShouldStartOnZero() throws SQLException {
is(allOf(equalTo(statement.getLargeUpdateCount()), equalTo(0L))));
assertThat(statement.getResultSet(), is(nullValue()));
}
+
+ @Test
+ public void testExecuteShouldRunSelectQueryV2() throws SQLException {
+ assertThat(statement.execute(SAMPLE_QUERY_CMD_V2), is(true));
+ final Set numbers =
+ IntStream.range(0, SAMPLE_QUERY_ROWS)
+ .boxed()
+ .map(Integer::byteValue)
+ .collect(Collectors.toCollection(HashSet::new));
+ try (final ResultSet resultSet = statement.getResultSet()) {
+ final int columnCount = resultSet.getMetaData().getColumnCount();
+ assertThat(columnCount, is(1));
+ int rowCount = 0;
+ for (; resultSet.next(); rowCount++) {
+ assertThat(numbers.remove(resultSet.getByte(1)), is(true));
+ }
+ assertThat(rowCount, is(equalTo(SAMPLE_QUERY_ROWS)));
+ }
+ assertThat(numbers, is(Collections.emptySet()));
+ assertThat(
+ (long) statement.getUpdateCount(),
+ is(allOf(equalTo(statement.getLargeUpdateCount()), equalTo(-1L))));
+ }
+
+ @Test
+ public void testExecuteShouldRunUpdateQueryForLargeUpdateV2() throws SQLException {
+ assertThat(statement.execute(SAMPLE_LARGE_UPDATE_QUERY_V2), is(false)); // UPDATE query.
+ final long updateCountSmall = statement.getUpdateCount();
+ final long updateCountLarge = statement.getLargeUpdateCount();
+ assertThat(updateCountLarge, is(equalTo(SAMPLE_LARGE_UPDATE_COUNT)));
+ assertThat(
+ updateCountSmall,
+ is(
+ allOf(
+ equalTo((long) AvaticaUtils.toSaturatedInt(updateCountLarge)),
+ not(equalTo(updateCountLarge)))));
+ assertThat(statement.getResultSet(), is(nullValue()));
+ }
}
diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java
index 45c2a96404..6627d91ab6 100644
--- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java
+++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java
@@ -87,6 +87,7 @@
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.arrow.vector.util.JsonStringArrayList;
import org.apache.calcite.avatica.Meta.StatementType;
+import org.checkerframework.checker.nullness.qual.Nullable;
/** An ad-hoc {@link FlightSqlProducer} for tests. */
public final class MockFlightSqlProducer implements FlightSqlProducer {
@@ -101,6 +102,7 @@ public final class MockFlightSqlProducer implements FlightSqlProducer {
private final SqlInfoBuilder sqlInfoBuilder = new SqlInfoBuilder();
private final Map parameterSchemas = new HashMap<>();
private final Map>> expectedParameterValues = new HashMap<>();
+ private final Map isUpdateMap = new HashMap<>();
private final Map actionTypeCounter = new HashMap<>();
@@ -176,6 +178,40 @@ public void addUpdateQuery(final String sqlCommand, final long updatedRows) {
});
}
+ /**
+ * Registers a new {@link StatementType#SELECT} SQL query, optionally setting the is_update field.
+ *
+ * @param sqlCommand the SQL command under which to register the new query.
+ * @param schema the schema to use for the query result.
+ * @param resultProviders the result provider for this query.
+ * @param isUpdate value to report for the is_update field, or {@code null} to leave it unset.
+ */
+ public void addSelectQuery(
+ final String sqlCommand,
+ final Schema schema,
+ final List> resultProviders,
+ final @Nullable Boolean isUpdate) {
+ addSelectQuery(sqlCommand, schema, resultProviders);
+ if (isUpdate != null) {
+ isUpdateMap.put(sqlCommand, isUpdate);
+ }
+ }
+
+ /**
+ * Registers a new {@link StatementType#UPDATE} SQL query, optionally setting the is_update field.
+ *
+ * @param sqlCommand the SQL command.
+ * @param updatedRows the number of rows affected.
+ * @param isUpdate value to report for the is_update field, or {@code null} to leave it unset.
+ */
+ public void addUpdateQuery(
+ final String sqlCommand, final long updatedRows, final @Nullable Boolean isUpdate) {
+ addUpdateQuery(sqlCommand, updatedRows);
+ if (isUpdate != null) {
+ isUpdateMap.put(sqlCommand, isUpdate);
+ }
+ }
+
/**
* Adds a catalog query to the results.
*
@@ -247,6 +283,12 @@ public void createPreparedStatement(
resultBuilder.setParameterSchema(ByteString.copyFrom(outputStream.toByteArray()));
}
+ // Set is_update field if present
+ final Boolean isUpdate = isUpdateMap.get(query);
+ if (isUpdate != null) {
+ resultBuilder.setIsUpdate(isUpdate);
+ }
+
listener.onNext(new Result(pack(resultBuilder.build()).toByteArray()));
} catch (final Throwable t) {
listener.onError(t);
diff --git a/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java b/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java
index 623f9311e8..0af09faee1 100644
--- a/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java
+++ b/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java
@@ -1284,6 +1284,19 @@ public Schema getParameterSchema() {
return parameterSchema;
}
+ /**
+ * Returns whether the server indicated this prepared statement is an update query.
+ *
+ * @return true if the server indicated this is an update query, false if the server indicated
+ * this is a select query, or null if the server did not provide this information.
+ */
+ public Boolean isUpdate() {
+ if (preparedStatementResult.hasIsUpdate()) {
+ return preparedStatementResult.getIsUpdate();
+ }
+ return null;
+ }
+
/** Get the schema of the result set (should be identical to {@link #getResultSetSchema()}). */
public SchemaResult fetchSchema(CallOption... options) {
checkOpen();
diff --git a/pom.xml b/pom.xml
index 0ce71f74cc..78137eb4e9 100644
--- a/pom.xml
+++ b/pom.xml
@@ -95,13 +95,13 @@ under the License.
1773644827
${project.build.directory}/generated-sources
1.9.0
- 5.12.2
+ 6.1.1
2.0.18
33.6.0-jre
- 4.2.13.Final
- 1.81.0
- 4.34.1
- 2.21.3
+ 4.2.15.Final
+ 1.82.1
+ 4.35.1
+ 2.22.0
3.5.0
25.2.10
1.12.1
@@ -112,8 +112,8 @@ under the License.
10.23.0
true
2.42.0
- 4.1.0
- 1.5.32
+ 4.2.1
+ 1.5.37
none
-Xdoclint:none
@@ -314,7 +314,7 @@ under the License.
org.immutables
value
- 2.12.1
+ 2.12.2
@@ -352,7 +352,7 @@ under the License.
org.jacoco
jacoco-maven-plugin
- 0.8.14
+ 0.8.15