Skip to content

Commit ff676f7

Browse files
AdamGSAdamGS
committed
more
Signed-off-by: Adam Gutglick <adam@spiraldb.com>
1 parent ae094af commit ff676f7

35 files changed

Lines changed: 275 additions & 110 deletions

File tree

‎encodings/datetime-parts/src/compute/mod.rs‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,11 @@ mod tests {
2727
use crate::DateTimePartsArray;
2828

2929
fn dtp_from_temporal(temporal: TemporalArray) -> DateTimePartsArray {
30-
DateTimeParts::try_from_temporal(temporal, &mut array_session().build().create_execution_ctx())
31-
.unwrap()
30+
DateTimeParts::try_from_temporal(
31+
temporal,
32+
&mut array_session().build().create_execution_ctx(),
33+
)
34+
.unwrap()
3235
}
3336

3437
#[rstest]

‎encodings/datetime-parts/src/compute/rules.rs‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -230,8 +230,11 @@ mod tests {
230230
time_unit,
231231
None,
232232
);
233-
DateTimeParts::try_from_temporal(temporal, &mut array_session().build().create_execution_ctx())
234-
.vortex_expect("TemporalArray must produce valid DateTimeParts")
233+
DateTimeParts::try_from_temporal(
234+
temporal,
235+
&mut array_session().build().create_execution_ctx(),
236+
)
237+
.vortex_expect("TemporalArray must produce valid DateTimeParts")
235238
}
236239

237240
/// Create a constant timestamp scalar at midnight for the given day.

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ use vortex_array::arrow::ArrowSessionExt;
3737
use vortex_array::arrow::FromArrowArray;
3838
use vortex_array::dtype::DType;
3939
use vortex_array::kernel::ExecuteParentKernel;
40-
use vortex_array::optimizer::kernels::ArrayKernelsExt;
40+
use vortex_array::optimizer::kernels::builder_kernels;
4141
use vortex_array::scalar_fn::ScalarFnVTable;
4242
use vortex_array::scalar_fn::fns::variant_get::VariantGet;
4343
use vortex_array::scalar_fn::fns::variant_get::VariantPath;
@@ -48,13 +48,13 @@ use vortex_error::vortex_err;
4848
use vortex_json::Json;
4949
use vortex_json::JsonToVariant;
5050
use vortex_mask::Mask;
51-
use vortex_session::VortexSession;
51+
use vortex_session::VortexSessionBuilder;
5252

5353
use crate::ParquetVariant;
5454
use crate::ParquetVariantArrayExt;
5555

56-
pub(crate) fn initialize(session: &VortexSession) {
57-
let kernels = session.kernels();
56+
pub(crate) fn initialize(session: &mut VortexSessionBuilder) {
57+
let kernels = builder_kernels(session);
5858
kernels.register_execute_parent_kernel(
5959
Filter.id(),
6060
ParquetVariant,

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

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -36,13 +36,13 @@ mod vtable;
3636
use std::sync::Arc;
3737

3838
pub use array::ParquetVariantArrayExt;
39-
use vortex_array::arrow::ArrowSessionExt;
40-
use vortex_array::session::ArraySessionExt;
39+
use vortex_array::arrow::ArrowSession;
40+
use vortex_array::session::ArraySession;
4141
pub use vortex_json::JsonToVariant;
4242
pub use vortex_json::JsonToVariantOptions;
4343
pub use vortex_json::ShreddingSpec;
4444
pub use vortex_json::json_to_variant;
45-
use vortex_session::VortexSession;
45+
use vortex_session::VortexSessionBuilder;
4646
pub use vtable::ParquetVariant;
4747
pub use vtable::ParquetVariantArray;
4848

@@ -51,10 +51,12 @@ pub use vtable::ParquetVariantArray;
5151
///
5252
/// This also initializes [`vortex_json`], registering the `Json` extension dtype and the
5353
/// `json_to_variant` scalar function whose execution this crate provides.
54-
pub fn initialize(session: &VortexSession) {
54+
pub fn initialize(session: &mut VortexSessionBuilder) {
5555
vortex_json::initialize(session);
56-
session.arrays().register(ParquetVariant);
56+
session.get_mut::<ArraySession>().register(ParquetVariant);
5757
kernel::initialize(session);
58-
session.arrow().register_exporter(Arc::new(ParquetVariant));
59-
session.arrow().register_importer(Arc::new(ParquetVariant));
58+
59+
let arrow = session.get_mut::<ArrowSession>();
60+
arrow.register_exporter(Arc::new(ParquetVariant));
61+
arrow.register_importer(Arc::new(ParquetVariant));
6062
}

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

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -420,7 +420,8 @@ mod tests {
420420
Some(ScalarValue::Variant(Box::new(expected_inner))),
421421
)?;
422422
assert_eq!(
423-
vortex_arr.execute_scalar(index, &mut array_session().build().create_execution_ctx())?,
423+
vortex_arr
424+
.execute_scalar(index, &mut array_session().build().create_execution_ctx())?,
424425
expected
425426
);
426427
}
@@ -673,7 +674,8 @@ mod tests {
673674
let arrow_variant = ArrowVariantArray::try_new(&struct_array)?;
674675
let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?;
675676

676-
let row0 = vortex_arr.execute_scalar(0, &mut array_session().build().create_execution_ctx())?;
677+
let row0 =
678+
vortex_arr.execute_scalar(0, &mut array_session().build().create_execution_ctx())?;
677679
let row0 = row0.as_variant().value().unwrap().as_list();
678680
assert_eq!(row0.len(), 2);
679681
assert_eq!(
@@ -697,7 +699,8 @@ mod tests {
697699
Some(20)
698700
);
699701

700-
let row1 = vortex_arr.execute_scalar(1, &mut array_session().build().create_execution_ctx())?;
702+
let row1 =
703+
vortex_arr.execute_scalar(1, &mut array_session().build().create_execution_ctx())?;
701704
let row1 = row1.as_variant().value().unwrap().as_list();
702705
assert_eq!(row1.len(), 1);
703706
assert_eq!(
@@ -764,7 +767,8 @@ mod tests {
764767

765768
let arrow_variant = ArrowVariantArray::try_new(&struct_array)?;
766769
let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?;
767-
let object = vortex_arr.execute_scalar(0, &mut array_session().build().create_execution_ctx())?;
770+
let object =
771+
vortex_arr.execute_scalar(0, &mut array_session().build().create_execution_ctx())?;
768772
let object = object.as_variant().value().unwrap().as_struct();
769773

770774
assert_eq!(

‎encodings/zstd/src/compute/mod.rs‎

Lines changed: 56 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -18,52 +18,100 @@ mod tests {
1818

1919
fn zstd_i32() -> ZstdArray {
2020
let values = PrimitiveArray::from_iter([100i32, 200, 300, 400, 500]);
21-
Zstd::from_primitive(&values, 0, 0, &mut array_session().build().create_execution_ctx()).unwrap()
21+
Zstd::from_primitive(
22+
&values,
23+
0,
24+
0,
25+
&mut array_session().build().create_execution_ctx(),
26+
)
27+
.unwrap()
2228
}
2329

2430
fn zstd_f64() -> ZstdArray {
2531
let values = PrimitiveArray::from_iter([1.1f64, 2.2, 3.3, 4.4, 5.5]);
26-
Zstd::from_primitive(&values, 0, 0, &mut array_session().build().create_execution_ctx()).unwrap()
32+
Zstd::from_primitive(
33+
&values,
34+
0,
35+
0,
36+
&mut array_session().build().create_execution_ctx(),
37+
)
38+
.unwrap()
2739
}
2840

2941
fn zstd_u32() -> ZstdArray {
3042
let values = PrimitiveArray::from_iter([10u32, 20, 30, 40, 50]);
31-
Zstd::from_primitive(&values, 0, 0, &mut array_session().build().create_execution_ctx()).unwrap()
43+
Zstd::from_primitive(
44+
&values,
45+
0,
46+
0,
47+
&mut array_session().build().create_execution_ctx(),
48+
)
49+
.unwrap()
3250
}
3351

3452
fn zstd_nullable_i64() -> ZstdArray {
3553
let values =
3654
PrimitiveArray::from_option_iter([Some(1000i64), None, Some(3000), Some(4000), None]);
37-
Zstd::from_primitive(&values, 0, 0, &mut array_session().build().create_execution_ctx()).unwrap()
55+
Zstd::from_primitive(
56+
&values,
57+
0,
58+
0,
59+
&mut array_session().build().create_execution_ctx(),
60+
)
61+
.unwrap()
3862
}
3963

4064
fn zstd_single() -> ZstdArray {
4165
let values = PrimitiveArray::new(
4266
buffer![42i64],
4367
vortex_array::validity::Validity::NonNullable,
4468
);
45-
Zstd::from_primitive(&values, 0, 0, &mut array_session().build().create_execution_ctx()).unwrap()
69+
Zstd::from_primitive(
70+
&values,
71+
0,
72+
0,
73+
&mut array_session().build().create_execution_ctx(),
74+
)
75+
.unwrap()
4676
}
4777

4878
fn zstd_large() -> ZstdArray {
4979
let values = PrimitiveArray::new(
5080
buffer![0u32..1000],
5181
vortex_array::validity::Validity::NonNullable,
5282
);
53-
Zstd::from_primitive(&values, 3, 0, &mut array_session().build().create_execution_ctx()).unwrap()
83+
Zstd::from_primitive(
84+
&values,
85+
3,
86+
0,
87+
&mut array_session().build().create_execution_ctx(),
88+
)
89+
.unwrap()
5490
}
5591

5692
fn zstd_all_same() -> ZstdArray {
5793
let values = PrimitiveArray::new(
5894
buffer![42i32; 100],
5995
vortex_array::validity::Validity::NonNullable,
6096
);
61-
Zstd::from_primitive(&values, 0, 0, &mut array_session().build().create_execution_ctx()).unwrap()
97+
Zstd::from_primitive(
98+
&values,
99+
0,
100+
0,
101+
&mut array_session().build().create_execution_ctx(),
102+
)
103+
.unwrap()
62104
}
63105

64106
fn zstd_negative() -> ZstdArray {
65107
let values = PrimitiveArray::from_iter([-100i32, -50, 0, 50, 100]);
66-
Zstd::from_primitive(&values, 0, 0, &mut array_session().build().create_execution_ctx()).unwrap()
108+
Zstd::from_primitive(
109+
&values,
110+
0,
111+
0,
112+
&mut array_session().build().create_execution_ctx(),
113+
)
114+
.unwrap()
67115
}
68116

69117
#[rstest]

‎encodings/zstd/src/test.rs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,10 @@ fn test_validity_vtable() {
164164
assert_eq!(
165165
arr.validity()
166166
.unwrap()
167-
.execute_mask(arr.len(), &mut array_session().build().create_execution_ctx())
167+
.execute_mask(
168+
arr.len(),
169+
&mut array_session().build().create_execution_ctx()
170+
)
168171
.unwrap(),
169172
Mask::from_iter(mask_bools)
170173
);
@@ -173,7 +176,10 @@ fn test_validity_vtable() {
173176
sliced
174177
.validity()
175178
.unwrap()
176-
.execute_mask(sliced.len(), &mut array_session().build().create_execution_ctx())
179+
.execute_mask(
180+
sliced.len(),
181+
&mut array_session().build().create_execution_ctx()
182+
)
177183
.unwrap(),
178184
Mask::from_iter(vec![true, true, false])
179185
);

‎vortex-array/benches/aggregate_grouped.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ fn main() {
3131
divan::main();
3232
}
3333

34-
static SESSION: LazyLock<VortexSession> = LazyLock::new(vortex_array::array_session);
34+
static SESSION: LazyLock<VortexSession> = LazyLock::new(|| vortex_array::array_session().build());
3535

3636
const GROUP_COUNT: usize = 128;
3737
const GROUP_SIZE_SEED: u64 = 42;

‎vortex-array/benches/bool_zip.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ fn main() {
1919
divan::main();
2020
}
2121

22-
static SESSION: LazyLock<VortexSession> = LazyLock::new(vortex_array::array_session);
22+
static SESSION: LazyLock<VortexSession> = LazyLock::new(|| vortex_array::array_session().build());
2323

2424
const LEN: usize = 65_536;
2525

‎vortex-array/benches/dict_compare.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ fn main() {
2828
divan::main();
2929
}
3030

31-
static SESSION: LazyLock<VortexSession> = LazyLock::new(vortex_array::array_session);
31+
static SESSION: LazyLock<VortexSession> = LazyLock::new(|| vortex_array::array_session().build());
3232

3333
const LENGTH_AND_UNIQUE_VALUES: &[(usize, usize)] = &[
3434
// length, unique_values

0 commit comments

Comments
 (0)