Skip to content

Commit a5e0c3d

Browse files
committed
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: Claude <noreply@anthropic.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01CK2nCuXnyd2g3mNHQC8Lz2
1 parent e99ba67 commit a5e0c3d

3 files changed

Lines changed: 94 additions & 49 deletions

File tree

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
//! A session's enabled editions may include encodings the session has no registration for,
5+
//! e.g. an optional encoding crate that is not compiled in (vortex-jni enables the unstable
6+
//! edition for `vortex.parquet.variant` without registering its sibling encodings). Files
7+
//! written by such a session must stay readable.
8+
9+
#![expect(clippy::tests_outside_test_module)]
10+
11+
use futures::TryStreamExt;
12+
use vortex_array::IntoArray;
13+
use vortex_array::VortexSessionExecute;
14+
use vortex_array::arrays::ChunkedArray;
15+
use vortex_array::assert_arrays_eq;
16+
use vortex_array::session::ArraySessionExt;
17+
use vortex_buffer::ByteBuffer;
18+
use vortex_buffer::ByteBufferMut;
19+
use vortex_buffer::buffer;
20+
use vortex_edition::Edition;
21+
use vortex_edition::EditionId;
22+
use vortex_edition::EditionInclusion;
23+
use vortex_edition::EditionSessionExt;
24+
use vortex_error::VortexResult;
25+
use vortex_error::vortex_err;
26+
use vortex_file::OpenOptionsSessionExt;
27+
use vortex_file::WriteOptionsSessionExt;
28+
use vortex_io::session::RuntimeSession;
29+
use vortex_layout::session::LayoutSession;
30+
use vortex_session::VortexSession;
31+
32+
const TEST_EDITION: EditionId = EditionId::new("test", 2026, 7, 0);
33+
34+
fn session_with_unregistered_edition_member() -> VortexResult<VortexSession> {
35+
let session = vortex_array::array_session()
36+
.with::<LayoutSession>()
37+
.with::<RuntimeSession>();
38+
vortex_file::register_default_encodings(&session);
39+
40+
let editions = session.editions();
41+
editions
42+
.declare_edition(Edition {
43+
id: TEST_EDITION,
44+
min_vortex_version: None,
45+
})
46+
.map_err(|error| vortex_err!("{error}"))?;
47+
for id in session.arrays().registry().ids() {
48+
editions
49+
.declare_inclusion(EditionInclusion::new(&id, TEST_EDITION))
50+
.map_err(|error| vortex_err!("{error}"))?;
51+
}
52+
editions
53+
.declare_inclusion(EditionInclusion::new(&"test.not_registered", TEST_EDITION))
54+
.map_err(|error| vortex_err!("{error}"))?;
55+
session
56+
.enable_edition(TEST_EDITION)
57+
.map_err(|error| vortex_err!("{error}"))?;
58+
Ok(session)
59+
}
60+
61+
#[tokio::test]
62+
async fn roundtrip_with_enabled_but_unregistered_encoding() -> VortexResult<()> {
63+
let session = session_with_unregistered_edition_member()?;
64+
65+
let array = buffer![1u32, 2, 3, 4].into_array();
66+
let mut output = ByteBufferMut::empty();
67+
session
68+
.write_options()
69+
.write(&mut output, array.clone().to_array_stream())
70+
.await?;
71+
72+
let file = session
73+
.open_options()
74+
.open_buffer(ByteBuffer::from(output))?;
75+
let chunks: Vec<_> = file.scan()?.into_stream()?.try_collect().await?;
76+
let actual = ChunkedArray::from_iter(chunks).into_array();
77+
let mut ctx = session.create_execution_ctx();
78+
assert_arrays_eq!(array, actual, &mut ctx);
79+
Ok(())
80+
}

vortex-jni/src/session.rs

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,10 @@ use jni::EnvUnowned;
77
use jni::objects::JClass;
88
use jni::sys::jlong;
99
use vortex::VortexSessionDefault;
10+
use vortex::editions::DEFAULT_UNSTABLE_EDITION;
11+
use vortex::editions::EditionSessionExt;
12+
use vortex::error::VortexExpect;
13+
use vortex::error::vortex_err;
1014
use vortex::io::runtime::BlockingRuntime;
1115
use vortex::io::session::RuntimeSessionExt;
1216
use vortex::session::VortexSession;
@@ -15,9 +19,17 @@ use crate::RUNTIME;
1519

1620
/// Constructs a fresh [`VortexSession`] bound to the JNI-shared tokio runtime and returns
1721
/// an opaque pointer that Java must pass to [`Java_dev_vortex_jni_NativeSession_free`].
22+
///
23+
/// JNI sessions enable the latest unstable edition so that `vortex.parquet.variant`, an
24+
/// unstable-family encoding registered by [`vortex_parquet_variant::initialize`], stays
25+
/// writable for variant-typed columns.
1826
pub(crate) fn new_session() -> Box<VortexSession> {
1927
let session = VortexSession::default().with_handle(RUNTIME.handle());
2028
vortex_parquet_variant::initialize(&session);
29+
session
30+
.enable_edition(DEFAULT_UNSTABLE_EDITION)
31+
.map_err(|e| vortex_err!("{e}"))
32+
.vortex_expect("default unstable edition is registered");
2133
Box::new(session)
2234
}
2335

vortex-jni/src/writer.rs

Lines changed: 2 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -32,26 +32,20 @@ use jni::sys::jlong;
3232
use jni::sys::jobject;
3333
use object_store::ObjectStore;
3434
use object_store::path::Path as ObjectStorePath;
35-
use vortex::array::ArrayId;
3635
use vortex::array::ArrayRef;
37-
use vortex::array::VTable;
3836
use vortex::array::scalar::PValue;
3937
use vortex::array::scalar::Scalar;
4038
use vortex::array::scalar::ScalarValue;
4139
use vortex::array::stats::StatsSet;
4240
use vortex::array::stream::ArrayStreamAdapter;
4341
use vortex::dtype::DType;
44-
use vortex::dtype::Field as DTypeField;
45-
use vortex::dtype::FieldPath;
46-
use vortex::editions::EditionSessionExt;
4742
use vortex::error::VortexError;
4843
use vortex::error::VortexResult;
4944
use vortex::error::vortex_err;
5045
use vortex::expr::stats::Stat;
5146
use vortex::expr::stats::StatsProvider;
5247
use vortex::file::CountingVortexWrite;
5348
use vortex::file::WriteOptionsSessionExt;
54-
use vortex::file::WriteStrategyBuilder;
5549
use vortex::file::WriteSummary;
5650
use vortex::file::multi::parse_uri_or_path;
5751
use vortex::io::VortexWrite;
@@ -62,9 +56,7 @@ use vortex::io::runtime::Task;
6256
use vortex::io::session::RuntimeSessionExt;
6357
use vortex::session::VortexSession;
6458
use vortex::utils::aliases::hash_map::HashMap;
65-
use vortex::utils::aliases::hash_set::HashSet;
6659
use vortex_arrow::ArrowSessionExt;
67-
use vortex_parquet_variant::ParquetVariant;
6860

6961
use crate::RUNTIME;
7062
use crate::dtype::import_arrow_schema;
@@ -102,45 +94,6 @@ fn resolve_store(
10294
}
10395
}
10496

105-
fn write_options_for_schema(
106-
session: &VortexSession,
107-
write_schema: &DType,
108-
) -> vortex::file::VortexWriteOptions {
109-
let variant_paths = variant_field_paths(write_schema);
110-
if variant_paths.is_empty() {
111-
return session.write_options();
112-
}
113-
114-
let mut allowed: HashSet<ArrayId> = session.enabled_encoding_ids().into_iter().collect();
115-
allowed.insert(ParquetVariant.id());
116-
117-
let strategy = WriteStrategyBuilder::default().with_allow_encodings(allowed);
118-
119-
session.write_options().with_strategy(strategy.build())
120-
}
121-
122-
fn variant_field_paths(dtype: &DType) -> Vec<FieldPath> {
123-
let mut paths = Vec::new();
124-
collect_variant_field_paths(dtype, FieldPath::root(), &mut paths);
125-
paths
126-
}
127-
128-
fn collect_variant_field_paths(dtype: &DType, path: FieldPath, paths: &mut Vec<FieldPath>) {
129-
match dtype {
130-
DType::Variant(_) => paths.push(path),
131-
DType::Struct(fields, _) => {
132-
for (name, field_dtype) in fields.names().iter().zip(fields.fields()) {
133-
collect_variant_field_paths(
134-
&field_dtype,
135-
path.clone().push(DTypeField::from(name.clone())),
136-
paths,
137-
);
138-
}
139-
}
140-
_ => {}
141-
}
142-
}
143-
14497
/// Native writer holding a write-task handle and a sender that Java pushes batches into.
14598
pub struct NativeWriter {
14699
handle: Option<Task<VortexResult<WriteSummary>>>,
@@ -416,7 +369,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create(
416369
let resolved = resolve_store(&file_path, &properties)?;
417370
let (tx, rx) = mpsc::channel(WRITE_CHANNEL_CAPACITY);
418371
let stream = ArrayStreamAdapter::new(write_schema.clone(), rx);
419-
let write_options = write_options_for_schema(session, &write_schema);
372+
let write_options = session.write_options();
420373

421374
let (bytes_written, handle) = match resolved {
422375
ResolvedStore::Path(path) => {
@@ -494,7 +447,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_createStream(
494447
let writable = Arc::new(env.new_global_ref(&writable)?);
495448
let (tx, rx) = mpsc::channel(WRITE_CHANNEL_CAPACITY);
496449
let stream = ArrayStreamAdapter::new(write_schema.clone(), rx);
497-
let write_options = write_options_for_schema(session, &write_schema);
450+
let write_options = session.write_options();
498451

499452
let mut write = CountingVortexWrite::new(JavaWrite::new(vm, writable));
500453
let bytes_written = write.counter();

0 commit comments

Comments
 (0)