Skip to content

Commit 30a1db9

Browse files
AdamGSAdamGS
committed
cleanup tests
1 parent 18fee37 commit 30a1db9

2 files changed

Lines changed: 53 additions & 62 deletions

File tree

‎encodings/parquet-variant/src/variant_get/tests.rs‎

Lines changed: 51 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ use parquet_variant_compute::GetOptions;
1919
use parquet_variant_compute::VariantArray as ArrowVariantArray;
2020
use parquet_variant_compute::VariantArrayBuilder;
2121
use parquet_variant_compute::json_to_variant;
22+
use rstest::fixture;
2223
use rstest::rstest;
2324
use vortex_array::ArrayRef;
2425
use vortex_array::LEGACY_SESSION;
@@ -40,24 +41,23 @@ use crate::ParquetVariant;
4041
use crate::ParquetVariantArrayExt;
4142
use crate::ParquetVariantData;
4243

43-
/// Apply variant_get and execute through the full pipeline (including execute_parent).
44-
fn apply_variant_get(arr: &ArrayRef, path: impl Into<VortexVariantPath>) -> VortexResult<ArrayRef> {
45-
let expr = variant_get(path, root());
46-
let array = arr.clone().apply(&expr)?;
47-
let mut ctx = LEGACY_SESSION.create_execution_ctx();
48-
array.execute::<ArrayRef>(&mut ctx)
49-
}
50-
51-
/// Apply typed variant_get and execute through the full pipeline.
52-
fn apply_variant_get_as(
53-
arr: &ArrayRef,
54-
path: impl Into<VortexVariantPath>,
55-
as_dtype: DType,
56-
) -> VortexResult<ArrayRef> {
57-
let expr = variant_get_as(path, as_dtype, root());
58-
let array = arr.clone().apply(&expr)?;
59-
let mut ctx = LEGACY_SESSION.create_execution_ctx();
60-
array.execute::<ArrayRef>(&mut ctx)
44+
macro_rules! apply_variant_get {
45+
($arr:expr, $path:expr) => {{
46+
(|| -> VortexResult<ArrayRef> {
47+
let expr = variant_get($path, root());
48+
let array = $arr.clone().apply(&expr)?;
49+
let mut ctx = LEGACY_SESSION.create_execution_ctx();
50+
array.execute::<ArrayRef>(&mut ctx)
51+
})()
52+
}};
53+
($arr:expr, $path:expr, $as_dtype:expr) => {{
54+
(|| -> VortexResult<ArrayRef> {
55+
let expr = variant_get_as($path, $as_dtype, root());
56+
let array = $arr.clone().apply(&expr)?;
57+
let mut ctx = LEGACY_SESSION.create_execution_ctx();
58+
array.execute::<ArrayRef>(&mut ctx)
59+
})()
60+
}};
6161
}
6262

6363
/// Convert a Vortex result back to an Arrow VariantArray for comparison.
@@ -86,7 +86,7 @@ fn assert_matches_arrow(json_rows: &[&str], field: &str) {
8686
.unwrap();
8787

8888
let vortex_input = ParquetVariantData::from_arrow_variant(&arrow_variant).unwrap();
89-
let vortex_result = apply_variant_get(&vortex_input, field).unwrap();
89+
let vortex_result = apply_variant_get!(&vortex_input, field).unwrap();
9090
let vortex_as_arrow = vortex_to_arrow_variant(&vortex_result);
9191

9292
assert_eq!(
@@ -136,7 +136,7 @@ fn assert_matches_arrow_with_path(
136136
.unwrap();
137137

138138
let vortex_input = ParquetVariantData::from_arrow_variant(&arrow_variant).unwrap();
139-
let vortex_result = apply_variant_get(&vortex_input, path).unwrap();
139+
let vortex_result = apply_variant_get!(&vortex_input, path).unwrap();
140140
let vortex_as_arrow = vortex_to_arrow_variant(&vortex_result);
141141

142142
assert_eq!(
@@ -189,7 +189,7 @@ fn assert_typed_matches_arrow_with_path(
189189
let expected = ArrayRef::from_arrow(arrow_result.as_ref(), true)?;
190190

191191
let vortex_input = ParquetVariantData::from_arrow_variant(&arrow_variant).unwrap();
192-
let vortex_result = apply_variant_get_as(&vortex_input, path, as_dtype.clone())?;
192+
let vortex_result = apply_variant_get!(&vortex_input, path, as_dtype.clone())?;
193193

194194
assert_eq!(vortex_result.dtype(), &as_dtype.as_nullable());
195195
assert_arrays_eq!(vortex_result, expected);
@@ -223,7 +223,7 @@ fn assert_matches_arrow_nullable(json_rows: &[&str], validity: &[bool], field: &
223223
.unwrap();
224224

225225
let vortex_input = ParquetVariantData::from_arrow_variant(&arrow_variant).unwrap();
226-
let vortex_result = apply_variant_get(&vortex_input, field).unwrap();
226+
let vortex_result = apply_variant_get!(&vortex_input, field).unwrap();
227227
let vortex_as_arrow = vortex_to_arrow_variant(&vortex_result);
228228

229229
assert_eq!(
@@ -374,14 +374,9 @@ fn test_variant_get_matches_arrow_typed_path() -> VortexResult<()> {
374374
)
375375
}
376376

377-
// ---------------------------------------------------------------------------
378-
// Original standalone tests
379-
// ---------------------------------------------------------------------------
380-
381-
#[test]
382-
fn test_variant_get_basic() -> VortexResult<()> {
383-
let arr = make_object_array()?;
384-
let result = apply_variant_get(&arr, "a")?;
377+
#[rstest]
378+
fn test_variant_get_basic(object_array: ArrayRef) -> VortexResult<()> {
379+
let result = apply_variant_get!(&object_array, "a")?;
385380

386381
assert_eq!(result.len(), 3);
387382

@@ -404,10 +399,9 @@ fn test_variant_get_basic() -> VortexResult<()> {
404399
Ok(())
405400
}
406401

407-
#[test]
408-
fn test_variant_get_missing_field() -> VortexResult<()> {
409-
let arr = make_object_array()?;
410-
let result = apply_variant_get(&arr, "nonexistent")?;
402+
#[rstest]
403+
fn test_variant_get_missing_field(object_array: ArrayRef) -> VortexResult<()> {
404+
let result = apply_variant_get!(&object_array, "nonexistent")?;
411405

412406
assert_eq!(result.len(), 3);
413407
for i in 0..3 {
@@ -417,10 +411,9 @@ fn test_variant_get_missing_field() -> VortexResult<()> {
417411
Ok(())
418412
}
419413

420-
#[test]
421-
fn test_variant_get_null_input() -> VortexResult<()> {
422-
let arr = make_nullable_object_array()?;
423-
let result = apply_variant_get(&arr, "a")?;
414+
#[rstest]
415+
fn test_variant_get_null_input(nullable_object_array: ArrayRef) -> VortexResult<()> {
416+
let result = apply_variant_get!(&nullable_object_array, "a")?;
424417

425418
assert_eq!(result.len(), 3);
426419
assert!(!result.scalar_at(0)?.is_null());
@@ -437,7 +430,7 @@ fn test_variant_get_non_object() -> VortexResult<()> {
437430
builder.append_variant(PqVariant::from("hello"));
438431
let arr = ParquetVariantData::from_arrow_variant(&builder.build())?;
439432

440-
let result = apply_variant_get(&arr, "a")?;
433+
let result = apply_variant_get!(&arr, "a")?;
441434

442435
assert_eq!(result.len(), 2);
443436
assert!(result.scalar_at(0)?.is_null());
@@ -446,10 +439,9 @@ fn test_variant_get_non_object() -> VortexResult<()> {
446439
Ok(())
447440
}
448441

449-
#[test]
450-
fn test_variant_get_different_field() -> VortexResult<()> {
451-
let arr = make_object_array()?;
452-
let result = apply_variant_get(&arr, "b")?;
442+
#[rstest]
443+
fn test_variant_get_different_field(object_array: ArrayRef) -> VortexResult<()> {
444+
let result = apply_variant_get!(&object_array, "b")?;
453445

454446
assert_eq!(result.len(), 3);
455447
assert!(!result.scalar_at(0)?.is_null());
@@ -459,37 +451,34 @@ fn test_variant_get_different_field() -> VortexResult<()> {
459451
Ok(())
460452
}
461453

462-
#[test]
463-
fn test_variant_get_through_slice_wrapper() -> VortexResult<()> {
464-
let arr = make_object_array()?;
465-
454+
#[rstest]
455+
fn test_variant_get_through_slice_wrapper(object_array: ArrayRef) -> VortexResult<()> {
466456
let mut ctx = LEGACY_SESSION.create_execution_ctx();
467457

468458
let expr = variant_get("a", root());
469-
let actual = arr
459+
let actual = object_array
470460
.slice(1..3)?
471461
.apply(&expr)?
472462
.execute::<ArrayRef>(&mut ctx)?;
473463

474-
let expected = apply_variant_get(&arr, "a")?;
464+
let expected = apply_variant_get!(&object_array, "a")?;
475465

476466
assert_eq!(actual.len(), 2);
477467
assert_eq!(actual.scalar_at(0)?, expected.scalar_at(1)?);
478468
assert_eq!(actual.scalar_at(1)?, expected.scalar_at(2)?);
479469
Ok(())
480470
}
481471

482-
#[test]
483-
fn test_variant_get_through_filter_wrapper() -> VortexResult<()> {
484-
let arr = make_object_array()?;
472+
#[rstest]
473+
fn test_variant_get_through_filter_wrapper(object_array: ArrayRef) -> VortexResult<()> {
485474
let mask = Mask::from_iter([true, false, true]);
486475

487476
let expr = variant_get("a", root());
488477
let mut ctx = LEGACY_SESSION.create_execution_ctx();
489478

490-
let array = arr.filter(mask.clone())?.apply(&expr)?;
479+
let array = object_array.filter(mask.clone())?.apply(&expr)?;
491480
let actual = array.execute::<ArrayRef>(&mut ctx)?;
492-
let expected = apply_variant_get(&arr, "a")?;
481+
let expected = apply_variant_get!(&object_array, "a")?;
493482

494483
assert_eq!(mask.true_count(), 2);
495484
assert_eq!(actual.len(), 2);
@@ -502,8 +491,9 @@ fn test_variant_get_through_filter_wrapper() -> VortexResult<()> {
502491
// Test data helpers
503492
// ---------------------------------------------------------------------------
504493

505-
/// Build a small non-null object variant array used by the standalone tests.
506-
fn make_object_array() -> VortexResult<ArrayRef> {
494+
/// Small non-null object variant array used by the standalone tests.
495+
#[fixture]
496+
fn object_array() -> ArrayRef {
507497
let mut builder = VariantArrayBuilder::new(3);
508498

509499
builder
@@ -520,11 +510,12 @@ fn make_object_array() -> VortexResult<ArrayRef> {
520510

521511
builder.new_object().with_field("b", "y").finish();
522512

523-
ParquetVariantData::from_arrow_variant(&builder.build())
513+
ParquetVariantData::from_arrow_variant(&builder.build()).unwrap()
524514
}
525515

526-
/// Build the same object array shape with an explicit top-level validity bitmap.
527-
fn make_nullable_object_array() -> VortexResult<ArrayRef> {
516+
/// The same object array shape with an explicit top-level validity bitmap.
517+
#[fixture]
518+
fn nullable_object_array() -> ArrayRef {
528519
let mut builder = VariantArrayBuilder::new(3);
529520

530521
builder.new_object().with_field("a", 10i32).finish();
@@ -539,5 +530,5 @@ fn make_nullable_object_array() -> VortexResult<ArrayRef> {
539530
)
540531
.unwrap();
541532
let arrow_variant = ArrowVariantArray::try_new(&null_struct).unwrap();
542-
ParquetVariantData::from_arrow_variant(&arrow_variant)
533+
ParquetVariantData::from_arrow_variant(&arrow_variant).unwrap()
543534
}

‎vortex-array/public-api.lock‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17424,7 +17424,7 @@ pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::coerce_args(&self,
1742417424

1742517425
pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::deserialize(&self, metadata: &[u8], session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
1742617426

17427-
pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::execute(&self, _options: &vortex_array::scalar_fn::fns::variant_get::VariantGetOptions, _args: &dyn vortex_array::scalar_fn::ExecutionArgs, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ArrayRef>
17427+
pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::execute(&self, options: &vortex_array::scalar_fn::fns::variant_get::VariantGetOptions, args: &dyn vortex_array::scalar_fn::ExecutionArgs, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ArrayRef>
1742817428

1742917429
pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::fmt_sql(&self, options: &vortex_array::scalar_fn::fns::variant_get::VariantGetOptions, expr: &vortex_array::expr::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
1743017430

@@ -18774,7 +18774,7 @@ pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::coerce_args(&self,
1877418774

1877518775
pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::deserialize(&self, metadata: &[u8], session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
1877618776

18777-
pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::execute(&self, _options: &vortex_array::scalar_fn::fns::variant_get::VariantGetOptions, _args: &dyn vortex_array::scalar_fn::ExecutionArgs, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ArrayRef>
18777+
pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::execute(&self, options: &vortex_array::scalar_fn::fns::variant_get::VariantGetOptions, args: &dyn vortex_array::scalar_fn::ExecutionArgs, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ArrayRef>
1877818778

1877918779
pub fn vortex_array::scalar_fn::fns::variant_get::VariantGet::fmt_sql(&self, options: &vortex_array::scalar_fn::fns::variant_get::VariantGetOptions, expr: &vortex_array::expr::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
1878018780

0 commit comments

Comments
 (0)