Skip to content

Commit 970746c

Browse files
connortsui20connortsui20
committed
Export to Arrow with the session of the execution context
An Arrow export took an ArrowSession as its receiver and an ExecutionCtx that carries its own VortexSession. The top-level dispatch and the target checks used the receiver, while export plugins and the inference of nested fields used the session of the context. So `ArrowSession::default().execute_arrow(array, None, &mut other.create_execution_ctx())` compiled and mixed the plugins of two sessions in one export. Exports now go through the new `ExecuteArrowExt` extension trait on arrays (`execute_arrow` and `execute_arrow_with_options`), which reads the ArrowSession from the context and takes no other session. This replaces `ArrowSession::execute_arrow`, `ArrowSession::exporter`, `ArrowExporter`, and the deprecated `ArrowArrayExecutor` trait. ArrowSession has no public constructor or Default anymore, so every ArrowSession belongs to a VortexSession. A lazily created ArrowSession was a bare default with an empty extension dtype registry. It exported an extension type without a plugin under its Vortex ID, but imported the same field back as plain storage. Only `vortex_arrow::initialize` wired the registry, and it replaced any existing ArrowSession, which dropped plugins that other crates had registered earlier. `ArrowSessionExt::arrow` now builds the session from the dtype registry on first use through the new `SessionExt::get_or_insert_with`, so `vortex_arrow::initialize` is no longer needed and is removed. Signed-off-by: Connor Tsui <[email protected]>
1 parent 19ae853 commit 970746c

67 files changed

Lines changed: 850 additions & 957 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎benchmarks/compress-bench/src/gpu/vortex.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ use vortex::file::WriteOptionsSessionExt;
2929
use vortex::layout::layouts::chunked::writer::ChunkedLayoutStrategy;
3030
use vortex::layout::layouts::compressed::CompressingStrategy;
3131
use vortex::layout::scan::split_by::SplitBy;
32-
use vortex_arrow::ArrowSessionExt;
32+
use vortex_arrow::ExecuteArrowExt;
3333
use vortex_bench::Format;
3434
use vortex_bench::SESSION;
3535
use vortex_bench::compress::Compressed;
@@ -303,11 +303,11 @@ fn verify_field(
303303
batch_index: usize,
304304
field_index: usize,
305305
) -> Result<()> {
306-
let expected = SESSION.arrow().execute_arrow(host.clone(), None, ctx)?;
306+
let expected = host.clone().execute_arrow(None, ctx)?;
307307
// Pin the Arrow target type so the two sides cannot land on different but equivalent
308308
// encodings of the same logical values.
309309
let target = Field::new("", expected.data_type().clone(), gpu.dtype().is_nullable());
310-
let actual = SESSION.arrow().execute_arrow(gpu, Some(&target), ctx)?;
310+
let actual = gpu.execute_arrow(Some(&target), ctx)?;
311311

312312
if expected.to_data() == actual.to_data() {
313313
return Ok(());

‎encodings/parquet-variant/src/array.rs‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ use vortex_array::vtable::child_to_validity;
3838
use vortex_array::vtable::validity_to_child;
3939
use vortex_arrow::ArrowExportOptions;
4040
use vortex_arrow::ArrowSession;
41-
use vortex_arrow::ArrowSessionExt;
41+
use vortex_arrow::ExecuteArrowExt;
4242
use vortex_arrow::to_arrow_null_buffer;
4343
use vortex_buffer::BitBuffer;
4444
use vortex_error::VortexExpect;
@@ -449,17 +449,16 @@ pub trait ParquetVariantArrayExt:
449449
options: &ArrowExportOptions,
450450
ctx: &mut ExecutionCtx,
451451
) -> VortexResult<ArrowVariantArray> {
452-
let session = ctx.session().clone();
453-
let arrow = session.arrow();
454-
let exporter = arrow.exporter(options);
455452
let metadata = self.metadata();
456453
let len = metadata.len();
457454
let nulls = to_arrow_null_buffer(self.parquet_variant_validity(), len, ctx)?;
458455

459456
let mut fields = Vec::with_capacity(3);
460457
let mut arrays: Vec<ArrowArrayRef> = Vec::with_capacity(3);
461458

462-
let metadata_arrow = exporter.execute_arrow(metadata.clone(), None, ctx)?;
459+
let metadata_arrow = metadata
460+
.clone()
461+
.execute_arrow_with_options(None, options, ctx)?;
463462
fields.push(Arc::new(Field::new(
464463
"metadata",
465464
metadata_arrow.data_type().clone(),
@@ -468,7 +467,9 @@ pub trait ParquetVariantArrayExt:
468467
arrays.push(metadata_arrow);
469468

470469
if let Some(value) = self.value() {
471-
let value_arrow = exporter.execute_arrow(value.clone(), None, ctx)?;
470+
let value_arrow = value
471+
.clone()
472+
.execute_arrow_with_options(None, options, ctx)?;
472473
fields.push(Arc::new(Field::new(
473474
"value",
474475
value_arrow.data_type().clone(),
@@ -478,7 +479,9 @@ pub trait ParquetVariantArrayExt:
478479
}
479480

480481
if let Some(typed_value) = self.typed_value() {
481-
let tv_arrow = exporter.execute_arrow(typed_value.clone(), None, ctx)?;
482+
let tv_arrow = typed_value
483+
.clone()
484+
.execute_arrow_with_options(None, options, ctx)?;
482485
fields.push(Arc::new(Field::new(
483486
"typed_value",
484487
tv_arrow.data_type().clone(),

‎encodings/parquet-variant/src/arrow.rs‎

Lines changed: 13 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ use vortex_arrow::ArrowImport;
4545
use vortex_arrow::ArrowImportVTable;
4646
use vortex_arrow::ArrowSession;
4747
use vortex_arrow::ArrowSessionExt;
48+
use vortex_arrow::ExecuteArrowExt;
4849
use vortex_arrow::to_arrow_null_buffer;
4950
use vortex_error::VortexResult;
5051
use vortex_error::vortex_bail;
@@ -84,9 +85,6 @@ pub(crate) fn export_storage_to_target<T: ParquetVariantArrayExt>(
8485
options: &ArrowExportOptions,
8586
ctx: &mut ExecutionCtx,
8687
) -> VortexResult<ArrowArrayRef> {
87-
let session = ctx.session().clone();
88-
let arrow = session.arrow();
89-
let exporter = arrow.exporter(options);
9088
let mut arrays = Vec::with_capacity(target_fields.len());
9189

9290
for field in target_fields {
@@ -104,7 +102,7 @@ pub(crate) fn export_storage_to_target<T: ParquetVariantArrayExt>(
104102
);
105103
};
106104

107-
arrays.push(exporter.execute_arrow(child, Some(field.as_ref()), ctx)?);
105+
arrays.push(child.execute_arrow_with_options(Some(field.as_ref()), options, ctx)?);
108106
}
109107

110108
let nulls = to_arrow_null_buffer(
@@ -471,6 +469,7 @@ mod tests {
471469
use vortex_array::scalar::Scalar;
472470
use vortex_array::validity::Validity;
473471
use vortex_arrow::ArrowSessionExt;
472+
use vortex_arrow::ExecuteArrowExt;
474473
use vortex_buffer::BitBuffer;
475474
use vortex_buffer::buffer;
476475
use vortex_error::VortexResult;
@@ -490,9 +489,7 @@ mod tests {
490489
fn export_variant(session: &VortexSession, array: ArrayRef) -> VortexResult<ArrowVariantArray> {
491490
let mut ctx = session.create_execution_ctx();
492491
let field = session.arrow().to_arrow_field("variant", array.dtype())?;
493-
let result = session
494-
.arrow()
495-
.execute_arrow(array, Some(&field), &mut ctx)?;
492+
let result = array.execute_arrow(Some(&field), &mut ctx)?;
496493

497494
Ok(ArrowVariantArray::try_new(result.as_ref())?)
498495
}
@@ -784,9 +781,7 @@ mod tests {
784781
assert!(imported.as_opt::<ParquetVariant>().is_some());
785782

786783
let mut ctx = session.create_execution_ctx();
787-
let exported = session
788-
.arrow()
789-
.execute_arrow(imported, Some(&field), &mut ctx)?;
784+
let exported = imported.execute_arrow(Some(&field), &mut ctx)?;
790785
let exported = exported.as_struct();
791786

792787
assert_struct_arrays_eq(exported, &storage);
@@ -805,9 +800,7 @@ mod tests {
805800
let canonical = VariantArray::try_new(core_storage, None)?.into_array();
806801

807802
let mut ctx = session.create_execution_ctx();
808-
let exported = session
809-
.arrow()
810-
.execute_arrow(canonical, Some(&field), &mut ctx)?;
803+
let exported = canonical.execute_arrow(Some(&field), &mut ctx)?;
811804
let exported = exported.as_struct();
812805

813806
assert_struct_arrays_eq(exported, &storage);
@@ -855,9 +848,7 @@ mod tests {
855848
.into(),
856849
);
857850

858-
let exported = session
859-
.arrow()
860-
.execute_arrow(canonical, Some(&field), &mut ctx)?;
851+
let exported = canonical.execute_arrow(Some(&field), &mut ctx)?;
861852
assert_eq!(exported.data_type(), field.data_type());
862853

863854
let actual = session.arrow().from_arrow_array(exported, &field)?;
@@ -883,9 +874,7 @@ mod tests {
883874
let expected = array.clone();
884875

885876
let mut ctx = session.create_execution_ctx();
886-
let exported = session
887-
.arrow()
888-
.execute_arrow(array.clone(), None, &mut ctx)?;
877+
let exported = array.clone().execute_arrow(None, &mut ctx)?;
889878

890879
let field = Field::new(
891880
"",
@@ -926,9 +915,7 @@ mod tests {
926915
.into_array();
927916

928917
let mut ctx = session.create_execution_ctx();
929-
let exported = session
930-
.arrow()
931-
.execute_arrow(expected.clone(), None, &mut ctx)?;
918+
let exported = expected.clone().execute_arrow(None, &mut ctx)?;
932919

933920
assert_eq!(
934921
exported.data_type(),
@@ -999,9 +986,7 @@ mod tests {
999986
);
1000987

1001988
let mut ctx = session.create_execution_ctx();
1002-
let exported = session
1003-
.arrow()
1004-
.execute_arrow(expected.clone(), Some(&field), &mut ctx)?;
989+
let exported = expected.clone().execute_arrow(Some(&field), &mut ctx)?;
1005990
assert_eq!(exported.data_type(), field.data_type());
1006991

1007992
let actual = session.arrow().from_arrow_array(exported, &field)?;
@@ -1044,9 +1029,8 @@ mod tests {
10441029
);
10451030

10461031
let mut ctx = session.create_execution_ctx();
1047-
let err = session
1048-
.arrow()
1049-
.execute_arrow(array, Some(&field), &mut ctx)
1032+
let err = array
1033+
.execute_arrow(Some(&field), &mut ctx)
10501034
.unwrap_err()
10511035
.to_string();
10521036
assert!(
@@ -1094,9 +1078,7 @@ mod tests {
10941078
.into(),
10951079
);
10961080
let mut ctx = session.create_execution_ctx();
1097-
let exported = session
1098-
.arrow()
1099-
.execute_arrow(array, Some(&field), &mut ctx)?;
1081+
let exported = array.execute_arrow(Some(&field), &mut ctx)?;
11001082
assert_eq!(
11011083
exported.data_type(),
11021084
field.data_type(),

‎encodings/parquet-variant/src/kernel.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ use vortex_array::scalar_fn::fns::variant_get::VariantPathElement;
4646
use vortex_arrow::ArrowExportOptions;
4747
use vortex_arrow::ArrowSession;
4848
use vortex_arrow::ArrowSessionExt;
49+
use vortex_arrow::ExecuteArrowExt;
4950
use vortex_error::VortexResult;
5051
use vortex_error::vortex_ensure_eq;
5152
use vortex_error::vortex_err;
@@ -154,7 +155,7 @@ fn json_strings_to_variant(
154155
) -> VortexResult<ArrayRef> {
155156
let nullable = parent.dtype().is_nullable();
156157
let session = ctx.session().clone();
157-
let arrow_strings = session.arrow().execute_arrow(strings, None, ctx)?;
158+
let arrow_strings = strings.execute_arrow(None, ctx)?;
158159
// Any row that fails to parse as JSON fails the whole conversion.
159160
let arrow_variant = parquet_variant_compute::json_to_variant(&arrow_strings)?;
160161

‎encodings/pco/src/tests.rs‎

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,8 @@ static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
3535
session
3636
});
3737

38+
use vortex_arrow::ExecuteArrowExt;
39+
3840
use crate::Pco;
3941
#[test]
4042
fn test_compress_decompress() {
@@ -207,16 +209,11 @@ fn test_serde() -> VortexResult<()> {
207209
&SESSION,
208210
)?;
209211
let data_type = SESSION.arrow().to_arrow_field("", data.dtype())?;
210-
let pco_arrow = SESSION
211-
.arrow()
212-
.execute_arrow(pco, Some(&data_type), &mut ctx)?;
213-
let decoded_arrow = SESSION
214-
.arrow()
215-
.execute_arrow(decoded, Some(&data_type), &mut ctx)?;
216-
let expected_arrow =
217-
SESSION
218-
.arrow()
219-
.execute_arrow(data.into_array(), Some(&data_type), &mut ctx)?;
212+
let pco_arrow = pco.execute_arrow(Some(&data_type), &mut ctx)?;
213+
let decoded_arrow = decoded.execute_arrow(Some(&data_type), &mut ctx)?;
214+
let expected_arrow = data
215+
.into_array()
216+
.execute_arrow(Some(&data_type), &mut ctx)?;
220217
assert_eq!(pco_arrow.as_ref(), expected_arrow.as_ref());
221218
assert_eq!(decoded_arrow.as_ref(), expected_arrow.as_ref());
222219
Ok(())

‎encodings/sparse/src/canonical.rs‎

Lines changed: 19 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -910,6 +910,7 @@ mod test {
910910
use vortex_array::scalar::Scalar;
911911
use vortex_array::validity::Validity;
912912
use vortex_arrow::ArrowSessionExt;
913+
use vortex_arrow::ExecuteArrowExt;
913914
use vortex_buffer::ByteBuffer;
914915
use vortex_buffer::buffer;
915916
use vortex_buffer::buffer_mut;
@@ -1136,33 +1137,23 @@ mod test {
11361137
let sparse_struct =
11371138
Sparse::try_new(indices, patch_values, len, fill_scalar, &mut ctx).unwrap();
11381139

1139-
let expected = SESSION
1140-
.arrow()
1141-
.execute_arrow(
1142-
DecimalArray::new(
1143-
buffer![100i128, 200, 123, 123, 123, 123, 123, 300, 4000, 123],
1144-
decimal_dtype,
1145-
// NB: patch indices: [0, 1, 7, 8]; patch validity: [Valid, Valid, Valid, Invalid]; ergo 0, 1, 7 are valid.
1146-
Validity::from_mask(Mask::from_excluded_indices(10, vec![8]), Nullable),
1147-
)
1148-
.into_array(),
1149-
None,
1150-
&mut ctx,
1151-
)
1152-
.unwrap();
1140+
let expected = DecimalArray::new(
1141+
buffer![100i128, 200, 123, 123, 123, 123, 123, 300, 4000, 123],
1142+
decimal_dtype,
1143+
// NB: patch indices: [0, 1, 7, 8]; patch validity: [Valid, Valid, Valid, Invalid]; ergo 0, 1, 7 are valid.
1144+
Validity::from_mask(Mask::from_excluded_indices(10, vec![8]), Nullable),
1145+
)
1146+
.into_array()
1147+
.execute_arrow(None, &mut ctx)
1148+
.unwrap();
11531149

1154-
let actual = SESSION
1155-
.arrow()
1156-
.execute_arrow(
1157-
sparse_struct
1158-
.as_array()
1159-
.clone()
1160-
.execute::<DecimalArray>(&mut ctx)
1161-
.unwrap()
1162-
.into_array(),
1163-
None,
1164-
&mut ctx,
1165-
)
1150+
let actual = sparse_struct
1151+
.as_array()
1152+
.clone()
1153+
.execute::<DecimalArray>(&mut ctx)
1154+
.unwrap()
1155+
.into_array()
1156+
.execute_arrow(None, &mut ctx)
11661157
.unwrap();
11671158

11681159
assert_eq!(expected.data_type(), actual.data_type());
@@ -1931,12 +1922,8 @@ mod test {
19311922

19321923
// Note that the preferred arrow list representation is `List` (not `ListView`).
19331924
let arrow_dtype = SESSION.arrow().to_arrow_field("", expected.dtype())?;
1934-
let actual = SESSION
1935-
.arrow()
1936-
.execute_arrow(actual, Some(&arrow_dtype), &mut ctx)?;
1937-
let expected = SESSION
1938-
.arrow()
1939-
.execute_arrow(expected, Some(&arrow_dtype), &mut ctx)?;
1925+
let actual = actual.execute_arrow(Some(&arrow_dtype), &mut ctx)?;
1926+
let expected = expected.execute_arrow(Some(&arrow_dtype), &mut ctx)?;
19401927

19411928
assert_eq!(actual.data_type(), expected.data_type());
19421929
Ok(())

0 commit comments

Comments
 (0)