Skip to content

Commit 04ef310

Browse files
dankingdanking
authored
chore: use assert_arrays_eq in more places (#4988)
I'm just trying to climb the deletions leader-board. Signed-off-by: Daniel King <[email protected]>
1 parent 369defb commit 04ef310

File tree

22 files changed

+299
-632
lines changed

22 files changed

+299
-632
lines changed

‎encodings/alp/src/alp/compress.rs‎

Lines changed: 8 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -133,10 +133,10 @@ mod tests {
133133
use core::f64;
134134

135135
use f64::consts::{E, PI};
136+
use vortex_array::assert_arrays_eq;
136137
use vortex_array::validity::Validity;
137138
use vortex_buffer::{Buffer, buffer};
138139
use vortex_dtype::NativePType;
139-
use vortex_scalar::Scalar;
140140

141141
use super::*;
142142

@@ -202,10 +202,7 @@ mod tests {
202202
assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 });
203203

204204
let decoded = decompress(&encoded);
205-
assert_eq!(decoded.scalar_at(0), array.scalar_at(0));
206-
assert_eq!(decoded.scalar_at(1), array.scalar_at(1));
207-
assert!(!decoded.is_valid(2));
208-
assert_eq!(decoded.scalar_at(3), array.scalar_at(3));
205+
assert_arrays_eq!(decoded, array);
209206
}
210207

211208
#[test]
@@ -223,14 +220,7 @@ mod tests {
223220

224221
assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 });
225222

226-
for idx in 0..3 {
227-
let s = encoded.scalar_at(idx);
228-
assert!(s.is_valid());
229-
}
230-
231-
assert!(!encoded.is_valid(4));
232-
let s = encoded.scalar_at(4);
233-
assert!(s.is_null());
223+
assert_arrays_eq!(&encoded, array);
234224

235225
let _decoded = decompress(&encoded);
236226
}
@@ -245,21 +235,18 @@ mod tests {
245235

246236
#[test]
247237
fn roundtrips_all_null() {
248-
let original = PrimitiveArray::new(
249-
Buffer::from_iter([195.26274f64, PI, -48.815685]),
250-
Validity::AllInvalid,
251-
);
238+
let original =
239+
PrimitiveArray::new(buffer![195.26274f64, PI, -48.815685], Validity::AllInvalid);
252240
let alp_arr = alp_encode(&original, None).unwrap();
253241
let decompressed = alp_arr.to_primitive();
242+
254243
assert_eq!(
255244
// The second and third values become exceptions and are replaced
256245
[195.26274, 195.26274, 195.26274],
257246
decompressed.as_slice::<f64>()
258247
);
259-
assert_eq!(original.validity(), decompressed.validity());
260-
assert_eq!(original.scalar_at(0), Scalar::null_typed::<f64>());
261-
assert_eq!(original.scalar_at(1), Scalar::null_typed::<f64>());
262-
assert_eq!(original.scalar_at(2), Scalar::null_typed::<f64>());
248+
249+
assert_arrays_eq!(decompressed, original);
263250
}
264251

265252
#[test]

‎encodings/alp/src/alp_rd/ops.rs‎

Lines changed: 6 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -78,9 +78,8 @@ impl OperationsVTable<ALPRDVTable> for ALPRDVTable {
7878
#[cfg(test)]
7979
mod test {
8080
use rstest::rstest;
81-
use vortex_array::ToCanonical;
8281
use vortex_array::arrays::PrimitiveArray;
83-
use vortex_dtype::Nullability;
82+
use vortex_array::assert_arrays_eq;
8483
use vortex_scalar::Scalar;
8584

8685
use crate::{ALPRDFloat, RDEncoder};
@@ -93,10 +92,7 @@ mod test {
9392
let encoded = RDEncoder::new(&[a, b]).encode(&array);
9493

9594
assert!(encoded.left_parts_patches().is_some());
96-
97-
let decoded = encoded.slice(1..3).to_primitive();
98-
99-
assert_eq!(decoded.as_slice::<T>(), &[b, outlier]);
95+
assert_arrays_eq!(encoded, array);
10096
}
10197

10298
#[rstest]
@@ -109,16 +105,8 @@ mod test {
109105
) {
110106
let array = PrimitiveArray::from_iter([a, b, outlier]);
111107
let encoded = RDEncoder::new(&[a, b]).encode(&array);
112-
113-
// Make sure that we're testing the exception pathway.
114108
assert!(encoded.left_parts_patches().is_some());
115-
116-
// The first two values need no patching
117-
assert_eq!(encoded.scalar_at(0), a.into());
118-
assert_eq!(encoded.scalar_at(1), b.into());
119-
120-
// The right value hits the left_part_exceptions
121-
assert_eq!(encoded.scalar_at(2), outlier.into());
109+
assert_arrays_eq!(encoded, array);
122110
}
123111

124112
#[test]
@@ -128,24 +116,10 @@ mod test {
128116
let outlier = 3e100f64;
129117
let array = PrimitiveArray::from_option_iter([Some(a), Some(b), Some(outlier)]);
130118
let encoded = RDEncoder::new(&[a, b]).encode(&array);
131-
132-
// Make sure that we're testing the exception pathway.
133119
assert!(encoded.left_parts_patches().is_some());
134-
135-
// The first two values need no patching
136-
assert_eq!(
137-
encoded.scalar_at(0),
138-
Scalar::primitive(a, Nullability::Nullable)
139-
);
140-
assert_eq!(
141-
encoded.scalar_at(1),
142-
Scalar::primitive(b, Nullability::Nullable)
143-
);
144-
145-
// The right value hits the left_part_exceptions
146-
assert_eq!(
147-
encoded.scalar_at(2),
148-
Scalar::primitive(outlier, Nullability::Nullable)
120+
assert_arrays_eq!(
121+
encoded,
122+
PrimitiveArray::from_option_iter([Some(a), Some(b), Some(outlier)])
149123
);
150124
}
151125
}

‎encodings/runend/src/arrow.rs‎

Lines changed: 30 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,10 @@ where
5252
mod tests {
5353
use arrow_array::types::{Int32Type, Int64Type};
5454
use arrow_array::{Float64Array, Int32Array, Int64Array, RunArray};
55+
use vortex_array::arrays::PrimitiveArray;
5556
use vortex_array::arrow::FromArrowArray;
57+
use vortex_array::{IntoArray as _, assert_arrays_eq};
58+
use vortex_buffer::buffer;
5659
use vortex_dtype::{DType, Nullability, PType};
5760

5861
use crate::RunEndArray;
@@ -69,20 +72,10 @@ mod tests {
6972
// Convert to Vortex
7073
let vortex_array = RunEndArray::from_arrow(&arrow_run_array, false);
7174

72-
// Verify basic properties
73-
assert_eq!(vortex_array.len(), 8);
74-
assert_eq!(
75-
vortex_array.dtype(),
76-
&DType::Primitive(PType::I32, Nullability::NonNullable)
75+
assert_arrays_eq!(
76+
vortex_array.as_ref(),
77+
buffer![10i32, 10, 10, 20, 20, 30, 30, 30].into_array()
7778
);
78-
79-
// Verify the values at different positions
80-
assert_eq!(vortex_array.scalar_at(0), 10.into()); // First run
81-
assert_eq!(vortex_array.scalar_at(2), 10.into()); // Still first run
82-
assert_eq!(vortex_array.scalar_at(3), 20.into()); // Second run
83-
assert_eq!(vortex_array.scalar_at(4), 20.into()); // Still second run
84-
assert_eq!(vortex_array.scalar_at(5), 30.into()); // Third run
85-
assert_eq!(vortex_array.scalar_at(7), 30.into()); // Still third run
8679
}
8780

8881
#[test]
@@ -95,20 +88,17 @@ mod tests {
9588
// Convert to Vortex with nullable=true
9689
let vortex_array = RunEndArray::from_arrow(&arrow_run_array, true);
9790

98-
// Verify basic properties
99-
assert_eq!(vortex_array.len(), 6);
100-
assert_eq!(
101-
vortex_array.dtype(),
102-
&DType::Primitive(PType::I32, Nullability::Nullable)
91+
assert_arrays_eq!(
92+
vortex_array.as_ref(),
93+
PrimitiveArray::from_option_iter([
94+
Some(100i32),
95+
Some(100i32),
96+
None,
97+
None,
98+
Some(300i32),
99+
Some(300i32)
100+
])
103101
);
104-
105-
// Verify the values
106-
assert_eq!(vortex_array.scalar_at(0), 100.into());
107-
assert_eq!(vortex_array.scalar_at(1), 100.into());
108-
assert!(vortex_array.scalar_at(2).is_null()); // Null value
109-
assert!(vortex_array.scalar_at(3).is_null()); // Null value
110-
assert_eq!(vortex_array.scalar_at(4), 300.into());
111-
assert_eq!(vortex_array.scalar_at(5), 300.into());
112102
}
113103

114104
#[test]
@@ -121,18 +111,7 @@ mod tests {
121111
// Convert to Vortex
122112
let vortex_array = RunEndArray::from_arrow(&arrow_run_array, false);
123113

124-
// Verify properties
125-
assert_eq!(vortex_array.len(), 4);
126-
assert_eq!(
127-
vortex_array.dtype(),
128-
&DType::Primitive(PType::F64, Nullability::NonNullable)
129-
);
130-
131-
// Verify values
132-
assert_eq!(vortex_array.scalar_at(0), 1.5.into());
133-
assert_eq!(vortex_array.scalar_at(1), 2.5.into());
134-
assert_eq!(vortex_array.scalar_at(2), 2.5.into());
135-
assert_eq!(vortex_array.scalar_at(3), 3.5.into());
114+
assert_arrays_eq!(vortex_array, buffer![1.5f64, 2.5, 2.5, 3.5].into_array());
136115
}
137116

138117
#[test]
@@ -150,21 +129,10 @@ mod tests {
150129

151130
// Convert the sliced array to Vortex
152131
let vortex_array = RunEndArray::from_arrow(&sliced_array, false);
153-
154-
// Verify the sliced array properties
155-
assert_eq!(vortex_array.len(), 6);
156-
assert_eq!(
157-
vortex_array.dtype(),
158-
&DType::Primitive(PType::I32, Nullability::NonNullable)
132+
assert_arrays_eq!(
133+
vortex_array,
134+
buffer![100, 200, 200, 200, 300, 300].into_array()
159135
);
160-
161-
// Verify the values in the sliced array
162-
assert_eq!(vortex_array.scalar_at(0), 100.into()); // Index 1 of original
163-
assert_eq!(vortex_array.scalar_at(1), 200.into()); // Index 2 of original
164-
assert_eq!(vortex_array.scalar_at(2), 200.into()); // Index 3 of original
165-
assert_eq!(vortex_array.scalar_at(3), 200.into()); // Index 4 of original
166-
assert_eq!(vortex_array.scalar_at(4), 300.into()); // Index 5 of original
167-
assert_eq!(vortex_array.scalar_at(5), 300.into()); // Index 6 of original
168136
}
169137

170138
#[test]
@@ -184,20 +152,17 @@ mod tests {
184152
// Convert to Vortex with nullable=true
185153
let vortex_array = RunEndArray::from_arrow(&sliced_array, true);
186154

187-
// Verify properties
188-
assert_eq!(vortex_array.len(), 6);
189-
assert_eq!(
190-
vortex_array.dtype(),
191-
&DType::Primitive(PType::I64, Nullability::Nullable)
155+
assert_arrays_eq!(
156+
vortex_array,
157+
PrimitiveArray::from_option_iter([
158+
None,
159+
None,
160+
Some(30i64),
161+
Some(30),
162+
Some(30),
163+
Some(40),
164+
])
192165
);
193-
194-
// Verify the values in the sliced array
195-
assert!(vortex_array.scalar_at(0).is_null());
196-
assert!(vortex_array.scalar_at(1).is_null());
197-
assert_eq!(vortex_array.scalar_at(2), 30i64.into());
198-
assert_eq!(vortex_array.scalar_at(3), 30i64.into());
199-
assert_eq!(vortex_array.scalar_at(4), 30i64.into());
200-
assert_eq!(vortex_array.scalar_at(5), 40i64.into());
201166
}
202167

203168
#[test]

0 commit comments

Comments
 (0)