Skip to content

Commit 2cd311b

Browse files
connortsui20connortsui20
committed
fix some tests
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent 16953eb commit 2cd311b

7 files changed

Lines changed: 32 additions & 125 deletions

File tree

‎encodings/bytebool/src/array.rs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -352,7 +352,6 @@ impl From<Vec<Option<bool>>> for ByteBoolData {
352352
mod tests {
353353
use vortex_array::ArrayContext;
354354
use vortex_array::IntoArray;
355-
use vortex_array::LEGACY_SESSION;
356355
use vortex_array::assert_arrays_eq;
357356
use vortex_array::serde::SerializeOptions;
358357
use vortex_array::serde::SerializedArray;
@@ -407,7 +406,7 @@ mod tests {
407406
let serialized = array
408407
.clone()
409408
.into_array()
410-
.serialize(&ctx, &LEGACY_SESSION, &SerializeOptions::default())
409+
.serialize(&ctx, &session, &SerializeOptions::default())
411410
.unwrap();
412411

413412
let mut concat = ByteBufferMut::empty();

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -253,6 +253,7 @@ mod tests {
253253
use vortex_array::dtype::PType;
254254
use vortex_array::serde::SerializeOptions;
255255
use vortex_array::serde::SerializedArray;
256+
use vortex_array::session::ArraySession;
256257
use vortex_array::session::ArraySessionExt;
257258
use vortex_array::validity::Validity;
258259
use vortex_buffer::BitBuffer;
@@ -267,7 +268,10 @@ mod tests {
267268
let dtype = array.dtype().clone();
268269
let len = array.len();
269270

270-
let session = VortexSession::empty().with::<vortex_array::session::ArraySession>();
271+
let session = VortexSession::empty().with::<ArraySession>();
272+
session.arrays().register(ParquetVariant);
273+
session.arrays().register(Variant);
274+
271275
let ctx = ArrayContext::empty();
272276
let serialized = array
273277
.serialize(&ctx, &session, &SerializeOptions::default())

‎encodings/pco/src/test.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -184,7 +184,7 @@ fn test_serde() -> VortexResult<()> {
184184
let bytes = pco
185185
.serialize(
186186
&context,
187-
&LEGACY_SESSION,
187+
&SESSION,
188188
&SerializeOptions {
189189
offset: 0,
190190
include_padding: true,

‎encodings/runend/src/arrow.rs‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ mod tests {
9090
use vortex_array::search_sorted::SearchSorted;
9191
use vortex_array::search_sorted::SearchSortedSide;
9292
use vortex_array::session::ArraySession;
93+
use vortex_array::session::ArraySessionExt;
9394
use vortex_array::validity::Validity;
9495
use vortex_buffer::Buffer;
9596
use vortex_buffer::buffer;
@@ -99,8 +100,11 @@ mod tests {
99100
use crate::RunEnd;
100101
use crate::ops::find_slice_end_index;
101102

102-
static SESSION: LazyLock<VortexSession> =
103-
LazyLock::new(|| VortexSession::empty().with::<ArraySession>());
103+
static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
104+
let session = VortexSession::empty().with::<ArraySession>();
105+
session.arrays().register(RunEnd);
106+
session
107+
});
104108

105109
fn decode_run_array<R: RunEndIndexType>(
106110
array: &RunArray<R>,

‎vortex-array/src/serde.rs‎

Lines changed: 0 additions & 100 deletions
Original file line numberDiff line numberDiff line change
@@ -691,103 +691,3 @@ impl TryFrom<BufferHandle> for SerializedArray {
691691
Self::try_from(value.try_to_host_sync()?)
692692
}
693693
}
694-
695-
#[cfg(test)]
696-
mod tests {
697-
use std::sync::LazyLock;
698-
699-
use flatbuffers::FlatBufferBuilder;
700-
use vortex_session::VortexSession;
701-
use vortex_session::registry::ReadContext;
702-
703-
use super::SerializeOptions;
704-
use super::SerializedArray;
705-
use crate::ArrayContext;
706-
use crate::array::ArrayId;
707-
use crate::dtype::DType;
708-
use crate::dtype::Nullability;
709-
use crate::flatbuffers as fba;
710-
use crate::session::ArraySession;
711-
use crate::session::ArraySessionExt;
712-
713-
static SESSION: LazyLock<VortexSession> = LazyLock::new(VortexSession::empty);
714-
715-
#[test]
716-
fn unknown_array_encoding_allow_unknown() {
717-
let mut fbb = FlatBufferBuilder::new();
718-
719-
let child_metadata = fbb.create_vector(&[9u8]);
720-
let child = fba::ArrayNode::create(
721-
&mut fbb,
722-
&fba::ArrayNodeArgs {
723-
encoding: 1,
724-
metadata: Some(child_metadata),
725-
children: None,
726-
buffers: None,
727-
stats: None,
728-
},
729-
);
730-
731-
let children = fbb.create_vector(&[child]);
732-
let metadata = fbb.create_vector(&[1u8, 2, 3]);
733-
let root = fba::ArrayNode::create(
734-
&mut fbb,
735-
&fba::ArrayNodeArgs {
736-
encoding: 0,
737-
metadata: Some(metadata),
738-
children: Some(children),
739-
buffers: None,
740-
stats: None,
741-
},
742-
);
743-
let array = fba::Array::create(
744-
&mut fbb,
745-
&fba::ArrayArgs {
746-
root: Some(root),
747-
buffers: None,
748-
},
749-
);
750-
fbb.finish_minimal(array);
751-
let (buf, start) = fbb.collapse();
752-
let tree = vortex_buffer::ByteBuffer::from(buf).slice(start..);
753-
754-
let ser = SerializedArray::from_array_tree(tree).unwrap();
755-
let ctx = ReadContext::new([
756-
ArrayId::new_ref("vortex.test.foreign_array"),
757-
ArrayId::new_ref("vortex.test.foreign_child"),
758-
]);
759-
let session = VortexSession::empty()
760-
.with::<ArraySession>()
761-
.allow_unknown();
762-
763-
let decoded = ser
764-
.decode(&DType::Variant(Nullability::Nullable), 5, &ctx, &session)
765-
.unwrap();
766-
assert_eq!(decoded.encoding_id().as_ref(), "vortex.test.foreign_array");
767-
assert_eq!(decoded.nchildren(), 1);
768-
assert_eq!(
769-
decoded.nth_child(0).unwrap().encoding_id().as_ref(),
770-
"vortex.test.foreign_child"
771-
);
772-
assert_eq!(
773-
SESSION.array_serialize(&decoded).unwrap().unwrap(),
774-
vec![1, 2, 3]
775-
);
776-
assert_eq!(
777-
SESSION
778-
.array_serialize(&decoded.nth_child(0).unwrap())
779-
.unwrap()
780-
.unwrap(),
781-
vec![9]
782-
);
783-
784-
let serialized = decoded
785-
.serialize(
786-
&ArrayContext::default(),
787-
&SESSION,
788-
&SerializeOptions::default(),
789-
)
790-
.unwrap();
791-
assert!(!serialized.is_empty());
792-
}
793-
}

‎vortex-array/src/session/mod.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ use crate::arrays::Bool;
1515
use crate::arrays::Chunked;
1616
use crate::arrays::Constant;
1717
use crate::arrays::Decimal;
18+
use crate::arrays::Dict;
1819
use crate::arrays::Extension;
1920
use crate::arrays::FixedSizeList;
2021
use crate::arrays::List;
@@ -26,6 +27,7 @@ use crate::arrays::Primitive;
2627
use crate::arrays::Struct;
2728
use crate::arrays::VarBin;
2829
use crate::arrays::VarBinView;
30+
use crate::arrays::Variant;
2931

3032
pub type ArrayRegistry = Registry<ArrayPluginRef>;
3133

@@ -68,11 +70,13 @@ impl Default for ArraySession {
6870
this.register(ListView);
6971
this.register(FixedSizeList);
7072
this.register(Struct);
73+
this.register(Variant);
7174
this.register(Extension);
7275

7376
// Register the utility encodings.
7477
this.register(Chunked);
7578
this.register(Constant);
79+
this.register(Dict);
7680
this.register(List);
7781
this.register(Masked);
7882
this.register(Patched);

‎vortex-layout/src/layouts/flat/writer.rs‎

Lines changed: 15 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -468,27 +468,23 @@ mod tests {
468468
let ctx = ArrayContext::empty();
469469

470470
// Write the array into a byte buffer.
471-
let (layout, _segments) = {
472-
let segments = Arc::new(TestSegments::default());
473-
let (ptr, eof) = SequenceId::root().split();
474-
// Only allow the dict encoding; canonical primitive children remain permitted.
475-
let mut allowed = HashSet::default();
476-
allowed.insert(Dict.id());
477-
let layout = FlatLayoutStrategy::default()
478-
.with_allow_encodings(allowed)
479-
.write_stream(
480-
ctx,
481-
Arc::<TestSegments>::clone(&segments),
482-
dict.into_array().to_array_stream().sequenced(ptr),
483-
eof,
484-
&session,
485-
)
486-
.await;
487471

488-
(layout, segments)
489-
};
472+
let segments = Arc::new(TestSegments::default());
473+
let (ptr, eof) = SequenceId::root().split();
474+
// Only allow the dict encoding; canonical primitive children remain permitted.
475+
let mut allowed = HashSet::default();
476+
allowed.insert(Dict.id());
490477

491-
assert!(layout.is_ok());
478+
FlatLayoutStrategy::default()
479+
.with_allow_encodings(allowed)
480+
.write_stream(
481+
ctx,
482+
Arc::<TestSegments>::clone(&segments),
483+
dict.into_array().to_array_stream().sequenced(ptr),
484+
eof,
485+
&session,
486+
)
487+
.await?;
492488

493489
Ok(())
494490
})

0 commit comments

Comments
 (0)