GH-1205: Fix byte-array element leak in FromSchemaByteArray

dataset/src/main/cpp/jni_util.cc

@@ -352,13 +352,17 @@ arrow::Result<std::shared_ptr<arrow::Schema>> 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<jbyte, decltype(release_elements)> elements_guard(schemaBytes_data,
+                                                                    release_elements);
   auto serialized_schema = std::make_shared<arrow::Buffer>(
       reinterpret_cast<uint8_t*>(schemaBytes_data), schemaBytes_len);
   arrow::io::BufferReader buf_reader(serialized_schema);
-  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Schema> 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<RecordBatch>& batch,
                                 jlong struct_array) {

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();
+    }
+  }
+}