From 9f64b779d82c5713db1b10b3827319c97735ef40 Mon Sep 17 00:00:00 2001 From: PG1204 Date: Fri, 24 Jul 2026 10:23:45 -0700 Subject: [PATCH 1/2] GH-1205: Fix byte-array element leak in FromSchemaByteArray by releasing elements via RAII guard --- dataset/src/main/cpp/jni_util.cc | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/dataset/src/main/cpp/jni_util.cc b/dataset/src/main/cpp/jni_util.cc index 35bfb328f0..23c9db06c5 100644 --- a/dataset/src/main/cpp/jni_util.cc +++ b/dataset/src/main/cpp/jni_util.cc @@ -352,13 +352,17 @@ arrow::Result> FromSchemaByteArray( arrow::ipc::DictionaryMemo in_memo; int schemaBytes_len = env->GetArrayLength(schemaBytes); jbyte* schemaBytes_data = env->GetByteArrayElements(schemaBytes, nullptr); + // Ensure the pinned/copied Java array elements are always released, even when + // schema parsing fails and an error Result is returned early below. + auto release_elements = [&](jbyte* data) { + env->ReleaseByteArrayElements(schemaBytes, data, JNI_ABORT); + }; + std::unique_ptr elements_guard(schemaBytes_data, + release_elements); auto serialized_schema = std::make_shared( reinterpret_cast(schemaBytes_data), schemaBytes_len); arrow::io::BufferReader buf_reader(serialized_schema); - ARROW_ASSIGN_OR_RAISE(std::shared_ptr schema, - arrow::ipc::ReadSchema(&buf_reader, &in_memo)) - env->ReleaseByteArrayElements(schemaBytes, schemaBytes_data, JNI_ABORT); - return schema; + return arrow::ipc::ReadSchema(&buf_reader, &in_memo); } arrow::Status ExportRecordBatch(JNIEnv* env, const std::shared_ptr& batch, jlong struct_array) { From 8667b46c06c65bb31306c689fcec62f64e247098 Mon Sep 17 00:00:00 2001 From: PG1204 Date: Sat, 25 Jul 2026 10:14:11 -0700 Subject: [PATCH 2/2] GH-1205: Add regression test for byte-array element leak in FromSchemaByteArray --- .../dataset/jni/TestFromSchemaByteArray.java | 74 +++++++++++++++++++ 1 file changed, 74 insertions(+) create mode 100644 dataset/src/test/java/org/apache/arrow/dataset/jni/TestFromSchemaByteArray.java diff --git a/dataset/src/test/java/org/apache/arrow/dataset/jni/TestFromSchemaByteArray.java b/dataset/src/test/java/org/apache/arrow/dataset/jni/TestFromSchemaByteArray.java new file mode 100644 index 0000000000..37471e9201 --- /dev/null +++ b/dataset/src/test/java/org/apache/arrow/dataset/jni/TestFromSchemaByteArray.java @@ -0,0 +1,74 @@ +/* + * 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.dataset.jni; + +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.io.File; +import java.lang.reflect.Field; +import org.apache.arrow.dataset.ParquetWriteSupport; +import org.apache.arrow.dataset.file.FileFormat; +import org.apache.arrow.dataset.file.FileSystemDatasetFactory; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +/** + * Regression test for the native {@code FromSchemaByteArray} helper (GH-1205). Passing malformed + * serialized schema bytes to {@code JniWrapper#createDataset} must surface a Java exception rather + * than crash, and repeated failures must not leak the pinned/copied Java byte-array elements + * acquired via {@code GetByteArrayElements}. + */ +public class TestFromSchemaByteArray extends TestNativeDataset { + + @TempDir public File TMP; + + public static final String AVRO_SCHEMA_USER = "user.avsc"; + + private static long factoryId(NativeDatasetFactory factory) throws Exception { + Field field = NativeDatasetFactory.class.getDeclaredField("datasetFactoryId"); + field.setAccessible(true); + return field.getLong(factory); + } + + @Test + public void testCreateDatasetWithMalformedSchemaBytes() throws Exception { + ParquetWriteSupport writeSupport = + ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a"); + FileSystemDatasetFactory factory = + new FileSystemDatasetFactory( + rootAllocator(), + NativeMemoryPool.getDefault(), + FileFormat.PARQUET, + writeSupport.getOutputURI()); + try { + final long datasetFactoryId = factoryId(factory); + // Bytes that are not a valid serialized Arrow schema, so native ReadSchema fails and the + // error path of FromSchemaByteArray is taken. + final byte[] malformedSchemaBytes = new byte[] {0, 1, 2, 3, 4, 5, 6, 7}; + + // Repeat many times: before the fix each failed call leaked the acquired array elements. + // The loop keeps the test meaningful as a leak regression while asserting graceful failure. + for (int i = 0; i < 1000; i++) { + assertThrows( + RuntimeException.class, + () -> JniWrapper.get().createDataset(datasetFactoryId, malformedSchemaBytes)); + } + } finally { + factory.close(); + } + } +}