Skip to content

Commit 25733a8

Browse files
AdamGSAdamGS
committed
Actually write Variant to files
Signed-off-by: Adam Gutglick <adam@spiraldb.com>
1 parent 3006be6 commit 25733a8

7 files changed

Lines changed: 85 additions & 30 deletions

File tree

‎Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -289,6 +289,7 @@ vortex-ipc = { version = "0.1.0", path = "./vortex-ipc", default-features = fals
289289
vortex-layout = { version = "0.1.0", path = "./vortex-layout", default-features = false }
290290
vortex-mask = { version = "0.1.0", path = "./vortex-mask", default-features = false }
291291
vortex-metrics = { version = "0.1.0", path = "./vortex-metrics", default-features = false }
292+
vortex-parquet-variant = { version = "0.1.0", path = "./encodings/parquet-variant" }
292293
vortex-pco = { version = "0.1.0", path = "./encodings/pco", default-features = false }
293294
vortex-proto = { version = "0.1.0", path = "./vortex-proto", default-features = false }
294295
vortex-runend = { version = "0.1.0", path = "./encodings/runend", default-features = false }

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

Lines changed: 66 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -309,6 +309,7 @@ mod tests {
309309
use vortex_array::IntoArray;
310310
use vortex_array::LEGACY_SESSION;
311311
use vortex_array::Precision;
312+
use vortex_array::VTable;
312313
use vortex_array::VortexSessionExecute;
313314
use vortex_array::arrays::PrimitiveArray;
314315
use vortex_array::arrays::VarBinViewArray;
@@ -364,6 +365,44 @@ mod tests {
364365
.unwrap()
365366
}
366367

368+
fn typed_value_variant_array() -> VortexResult<ArrayRef> {
369+
let mut metadata = BinaryViewBuilder::new();
370+
for _ in 0..3 {
371+
metadata.append_value(b"\x01\x00");
372+
}
373+
let metadata: ArrowArrayRef = Arc::new(metadata.finish());
374+
let typed_value: ArrowArrayRef = Arc::new(Int32Array::from(vec![10, 20, 30]));
375+
let arrow_storage = StructArray::try_new(
376+
vec![
377+
Arc::new(Field::new("metadata", DataType::BinaryView, false)),
378+
Arc::new(Field::new("typed_value", DataType::Int32, false)),
379+
]
380+
.into(),
381+
vec![metadata, typed_value],
382+
None,
383+
)?;
384+
385+
ParquetVariant::from_arrow_variant(&ArrowVariantArray::try_new(&arrow_storage)?)
386+
}
387+
388+
fn parquet_variant_file_session() -> VortexSession {
389+
let session = VortexSession::empty()
390+
.with::<ArraySession>()
391+
.with::<LayoutSession>()
392+
.with::<RuntimeSession>();
393+
vortex_file::register_default_encodings(&session);
394+
session.arrays().register(ParquetVariant);
395+
session
396+
}
397+
398+
fn write_strategy_with_parquet_variant() -> Arc<dyn vortex_layout::LayoutStrategy> {
399+
let mut allowed = vortex_file::ALLOWED_ENCODINGS.clone();
400+
allowed.insert(ParquetVariant.id());
401+
vortex_file::WriteStrategyBuilder::default()
402+
.with_allow_encodings(allowed)
403+
.build()
404+
}
405+
367406
#[test]
368407
fn test_execute_exposes_typed_value_as_canonical_shredded() -> VortexResult<()> {
369408
let metadata =
@@ -405,37 +444,38 @@ mod tests {
405444
}
406445

407446
#[tokio::test]
408-
async fn test_file_roundtrip_typed_value_variant() -> VortexResult<()> {
409-
let mut metadata = BinaryViewBuilder::new();
410-
for _ in 0..3 {
411-
metadata.append_value(b"\x01\x00");
412-
}
413-
let metadata: ArrowArrayRef = Arc::new(metadata.finish());
414-
let typed_value: ArrowArrayRef = Arc::new(Int32Array::from(vec![10, 20, 30]));
415-
let arrow_storage = StructArray::try_new(
416-
vec![
417-
Arc::new(Field::new("metadata", DataType::BinaryView, false)),
418-
Arc::new(Field::new("typed_value", DataType::Int32, false)),
419-
]
420-
.into(),
421-
vec![metadata, typed_value],
422-
None,
423-
)?;
424-
let expected =
425-
ParquetVariant::from_arrow_variant(&ArrowVariantArray::try_new(&arrow_storage)?)?;
426-
427-
let session = VortexSession::empty()
428-
.with::<ArraySession>()
429-
.with::<LayoutSession>()
430-
.with::<RuntimeSession>();
431-
vortex_file::register_default_encodings(&session);
432-
session.arrays().register(ParquetVariant);
447+
async fn test_file_roundtrip_typed_value_variant_with_statistics() -> VortexResult<()> {
448+
let expected = typed_value_variant_array()?;
449+
let session = parquet_variant_file_session();
433450

434451
let mut bytes = ByteBufferMut::empty();
435452
session
436453
.write_options()
437454
.with_strategy(Arc::new(FlatLayoutStrategy::default()))
438-
.with_file_statistics(Vec::new())
455+
.write(&mut bytes, expected.to_array_stream())
456+
.await?;
457+
458+
let actual = session
459+
.open_options()
460+
.open_buffer(bytes)?
461+
.scan()?
462+
.into_array_stream()?
463+
.read_all()
464+
.await?;
465+
466+
assert_arrays_eq!(expected, actual);
467+
Ok(())
468+
}
469+
470+
#[tokio::test]
471+
async fn test_file_roundtrip_typed_value_variant_with_zoned_strategy() -> VortexResult<()> {
472+
let expected = typed_value_variant_array()?;
473+
let session = parquet_variant_file_session();
474+
475+
let mut bytes = ByteBufferMut::empty();
476+
session
477+
.write_options()
478+
.with_strategy(write_strategy_with_parquet_variant())
439479
.write(&mut bytes, expected.to_array_stream())
440480
.await?;
441481

‎vortex-compressor/src/compressor.rs‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@ use vortex_array::dtype::DType;
2727
use vortex_array::dtype::Nullability;
2828
use vortex_array::scalar::Scalar;
2929
use vortex_error::VortexResult;
30-
use vortex_error::vortex_bail;
3130

3231
use crate::builtins::IntDictScheme;
3332
use crate::ctx::CompressorContext;
@@ -250,9 +249,7 @@ impl CascadingCompressor {
250249
.into_array(),
251250
)
252251
}
253-
Canonical::Variant(_) => {
254-
vortex_bail!("Variant arrays can not be compressed")
255-
}
252+
Canonical::Variant(variant_array) => Ok(variant_array.into_array()),
256253
}
257254
}
258255

‎vortex-file/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ vortex-io = { workspace = true }
4646
vortex-layout = { workspace = true }
4747
vortex-mask = { workspace = true }
4848
vortex-metrics = { workspace = true }
49+
vortex-parquet-variant = { workspace = true }
4950
vortex-pco = { workspace = true }
5051
vortex-runend = { workspace = true }
5152
vortex-scan = { workspace = true }

‎vortex-file/src/strategy.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ use vortex_array::arrays::Primitive;
2626
use vortex_array::arrays::Struct;
2727
use vortex_array::arrays::VarBin;
2828
use vortex_array::arrays::VarBinView;
29+
use vortex_array::arrays::Variant;
2930
use vortex_array::arrays::patched::use_experimental_patches;
3031
use vortex_array::dtype::FieldPath;
3132
use vortex_btrblocks::BtrBlocksCompressorBuilder;
@@ -52,6 +53,7 @@ use vortex_layout::layouts::repartition::RepartitionWriterOptions;
5253
use vortex_layout::layouts::table::TableStrategy;
5354
use vortex_layout::layouts::zoned::writer::ZonedLayoutOptions;
5455
use vortex_layout::layouts::zoned::writer::ZonedStrategy;
56+
use vortex_parquet_variant::ParquetVariant;
5557
use vortex_pco::Pco;
5658
use vortex_runend::RunEnd;
5759
use vortex_sequence::Sequence;
@@ -89,6 +91,7 @@ pub static ALLOWED_ENCODINGS: LazyLock<HashSet<ArrayId>> = LazyLock::new(|| {
8991
allowed.insert(Constant.id());
9092
allowed.insert(Masked.id());
9193
allowed.insert(Dict.id());
94+
allowed.insert(Variant.id());
9295

9396
// Compressed encodings from encoding crates
9497
allowed.insert(ALP.id());
@@ -111,6 +114,7 @@ pub static ALLOWED_ENCODINGS: LazyLock<HashSet<ArrayId>> = LazyLock::new(|| {
111114

112115
if use_experimental_patches() {
113116
allowed.insert(Patched.id());
117+
allowed.insert(ParquetVariant.id());
114118
}
115119

116120
#[cfg(feature = "zstd")]

‎vortex-layout/src/layouts/zoned/builder.rs‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,13 @@ pub struct StatsAccumulator {
4747

4848
impl StatsAccumulator {
4949
pub fn new(dtype: &DType, stats: &[Stat], max_variable_length_statistics_size: usize) -> Self {
50+
if !supports_file_stats(dtype) {
51+
return Self {
52+
builders: Vec::new(),
53+
length: 0,
54+
};
55+
}
56+
5057
let builders = stats
5158
.iter()
5259
.filter_map(|&stat| {
@@ -165,6 +172,10 @@ impl StatsAccumulator {
165172
}
166173
}
167174

175+
fn supports_file_stats(dtype: &DType) -> bool {
176+
!matches!(dtype, DType::Variant(_))
177+
}
178+
168179
fn stats_builder_with_capacity(
169180
stat: Stat,
170181
dtype: &DType,

0 commit comments

Comments
 (0)