Fix JNI variant writes under session-authority edition gating

The writer now rejects encodings outside the session's enabled editions
regardless of the configured strategy, so vortex-jni's strategy-level
allow-set override for `vortex.parquet.variant` no longer worked. JNI
sessions instead enable the latest unstable edition (the family that
encoding belongs to), which restores variant writability through the
session policy. The schema-conditional strategy override is removed:
parquet-variant arrays only arise for variant-typed columns, so the
session-wide policy is equivalent.

Sibling unstable encodings are not registered in the JNI build, so the
footer's pre-populated context may advertise enabled-but-unregistered
ids; a new integration test pins that such files stay readable.

Signed-off-by: "Joe Isaacs" <joe.isaacs@live.co.uk>

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01CK2nCuXnyd2g3mNHQC8Lz2

Author: claude

vortex-file/tests/enabled_unregistered_ids.rs

@@ -0,0 +1,80 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! A session's enabled editions may include encodings the session has no registration for,
+//! e.g. an optional encoding crate that is not compiled in (vortex-jni enables the unstable
+//! edition for `vortex.parquet.variant` without registering its sibling encodings). Files
+//! written by such a session must stay readable.
+
+#![expect(clippy::tests_outside_test_module)]
+
+use futures::TryStreamExt;
+use vortex_array::IntoArray;
+use vortex_array::VortexSessionExecute;
+use vortex_array::arrays::ChunkedArray;
+use vortex_array::assert_arrays_eq;
+use vortex_array::session::ArraySessionExt;
+use vortex_buffer::ByteBuffer;
+use vortex_buffer::ByteBufferMut;
+use vortex_buffer::buffer;
+use vortex_edition::Edition;
+use vortex_edition::EditionId;
+use vortex_edition::EditionInclusion;
+use vortex_edition::EditionSessionExt;
+use vortex_error::VortexResult;
+use vortex_error::vortex_err;
+use vortex_file::OpenOptionsSessionExt;
+use vortex_file::WriteOptionsSessionExt;
+use vortex_io::session::RuntimeSession;
+use vortex_layout::session::LayoutSession;
+use vortex_session::VortexSession;
+
+const TEST_EDITION: EditionId = EditionId::new("test", 2026, 7, 0);
+
+fn session_with_unregistered_edition_member() -> VortexResult<VortexSession> {
+    let session = vortex_array::array_session()
+        .with::<LayoutSession>()
+        .with::<RuntimeSession>();
+    vortex_file::register_default_encodings(&session);
+
+    let editions = session.editions();
+    editions
+        .declare_edition(Edition {
+            id: TEST_EDITION,
+            min_vortex_version: None,
+        })
+        .map_err(|error| vortex_err!("{error}"))?;
+    for id in session.arrays().registry().ids() {
+        editions
+            .declare_inclusion(EditionInclusion::new(&id, TEST_EDITION))
+            .map_err(|error| vortex_err!("{error}"))?;
+    }
+    editions
+        .declare_inclusion(EditionInclusion::new(&"test.not_registered", TEST_EDITION))
+        .map_err(|error| vortex_err!("{error}"))?;
+    session
+        .enable_edition(TEST_EDITION)
+        .map_err(|error| vortex_err!("{error}"))?;
+    Ok(session)
+}
+
+#[tokio::test]
+async fn roundtrip_with_enabled_but_unregistered_encoding() -> VortexResult<()> {
+    let session = session_with_unregistered_edition_member()?;
+
+    let array = buffer![1u32, 2, 3, 4].into_array();
+    let mut output = ByteBufferMut::empty();
+    session
+        .write_options()
+        .write(&mut output, array.clone().to_array_stream())
+        .await?;
+
+    let file = session
+        .open_options()
+        .open_buffer(ByteBuffer::from(output))?;
+    let chunks: Vec<_> = file.scan()?.into_stream()?.try_collect().await?;
+    let actual = ChunkedArray::from_iter(chunks).into_array();
+    let mut ctx = session.create_execution_ctx();
+    assert_arrays_eq!(array, actual, &mut ctx);
+    Ok(())
+}

vortex-jni/src/session.rs

@@ -7,6 +7,10 @@ use jni::EnvUnowned;
 use jni::objects::JClass;
 use jni::sys::jlong;
 use vortex::VortexSessionDefault;
+use vortex::editions::DEFAULT_UNSTABLE_EDITION;
+use vortex::editions::EditionSessionExt;
+use vortex::error::VortexExpect;
+use vortex::error::vortex_err;
 use vortex::io::runtime::BlockingRuntime;
 use vortex::io::session::RuntimeSessionExt;
 use vortex::session::VortexSession;
@@ -15,9 +19,17 @@ use crate::RUNTIME;
 
 /// Constructs a fresh [`VortexSession`] bound to the JNI-shared tokio runtime and returns
 /// an opaque pointer that Java must pass to [`Java_dev_vortex_jni_NativeSession_free`].
+///
+/// JNI sessions enable the latest unstable edition so that `vortex.parquet.variant`, an
+/// unstable-family encoding registered by [`vortex_parquet_variant::initialize`], stays
+/// writable for variant-typed columns.
 pub(crate) fn new_session() -> Box<VortexSession> {
     let session = VortexSession::default().with_handle(RUNTIME.handle());
     vortex_parquet_variant::initialize(&session);
+    session
+        .enable_edition(DEFAULT_UNSTABLE_EDITION)
+        .map_err(|e| vortex_err!("{e}"))
+        .vortex_expect("default unstable edition is registered");
     Box::new(session)
 }
 

vortex-jni/src/writer.rs

@@ -32,26 +32,20 @@ use jni::sys::jlong;
 use jni::sys::jobject;
 use object_store::ObjectStore;
 use object_store::path::Path as ObjectStorePath;
-use vortex::array::ArrayId;
 use vortex::array::ArrayRef;
-use vortex::array::VTable;
 use vortex::array::scalar::PValue;
 use vortex::array::scalar::Scalar;
 use vortex::array::scalar::ScalarValue;
 use vortex::array::stats::StatsSet;
 use vortex::array::stream::ArrayStreamAdapter;
 use vortex::dtype::DType;
-use vortex::dtype::Field as DTypeField;
-use vortex::dtype::FieldPath;
-use vortex::editions::EditionSessionExt;
 use vortex::error::VortexError;
 use vortex::error::VortexResult;
 use vortex::error::vortex_err;
 use vortex::expr::stats::Stat;
 use vortex::expr::stats::StatsProvider;
 use vortex::file::CountingVortexWrite;
 use vortex::file::WriteOptionsSessionExt;
-use vortex::file::WriteStrategyBuilder;
 use vortex::file::WriteSummary;
 use vortex::file::multi::parse_uri_or_path;
 use vortex::io::VortexWrite;
@@ -62,9 +56,7 @@ use vortex::io::runtime::Task;
 use vortex::io::session::RuntimeSessionExt;
 use vortex::session::VortexSession;
 use vortex::utils::aliases::hash_map::HashMap;
-use vortex::utils::aliases::hash_set::HashSet;
 use vortex_arrow::ArrowSessionExt;
-use vortex_parquet_variant::ParquetVariant;
 
 use crate::RUNTIME;
 use crate::dtype::import_arrow_schema;
@@ -102,45 +94,6 @@ fn resolve_store(
     }
 }
 
-fn write_options_for_schema(
-    session: &VortexSession,
-    write_schema: &DType,
-) -> vortex::file::VortexWriteOptions {
-    let variant_paths = variant_field_paths(write_schema);
-    if variant_paths.is_empty() {
-        return session.write_options();
-    }
-
-    let mut allowed: HashSet<ArrayId> = session.enabled_encoding_ids().into_iter().collect();
-    allowed.insert(ParquetVariant.id());
-
-    let strategy = WriteStrategyBuilder::default().with_allow_encodings(allowed);
-
-    session.write_options().with_strategy(strategy.build())
-}
-
-fn variant_field_paths(dtype: &DType) -> Vec<FieldPath> {
-    let mut paths = Vec::new();
-    collect_variant_field_paths(dtype, FieldPath::root(), &mut paths);
-    paths
-}
-
-fn collect_variant_field_paths(dtype: &DType, path: FieldPath, paths: &mut Vec<FieldPath>) {
-    match dtype {
-        DType::Variant(_) => paths.push(path),
-        DType::Struct(fields, _) => {
-            for (name, field_dtype) in fields.names().iter().zip(fields.fields()) {
-                collect_variant_field_paths(
-                    &field_dtype,
-                    path.clone().push(DTypeField::from(name.clone())),
-                    paths,
-                );
-            }
-        }
-        _ => {}
-    }
-}
-
 /// Native writer holding a write-task handle and a sender that Java pushes batches into.
 pub struct NativeWriter {
     handle: Option<Task<VortexResult<WriteSummary>>>,
@@ -416,7 +369,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create(
         let resolved = resolve_store(&file_path, &properties)?;
         let (tx, rx) = mpsc::channel(WRITE_CHANNEL_CAPACITY);
         let stream = ArrayStreamAdapter::new(write_schema.clone(), rx);
-        let write_options = write_options_for_schema(session, &write_schema);
+        let write_options = session.write_options();
 
         let (bytes_written, handle) = match resolved {
             ResolvedStore::Path(path) => {
@@ -494,7 +447,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_createStream(
         let writable = Arc::new(env.new_global_ref(&writable)?);
         let (tx, rx) = mpsc::channel(WRITE_CHANNEL_CAPACITY);
         let stream = ArrayStreamAdapter::new(write_schema.clone(), rx);
-        let write_options = write_options_for_schema(session, &write_schema);
+        let write_options = session.write_options();
 
         let mut write = CountingVortexWrite::new(JavaWrite::new(vm, writable));
         let bytes_written = write.counter();