diff --git a/parquet/tests/arrow_reader/parquet_testing.rs b/parquet/tests/arrow_reader/parquet_testing.rs index daa3ff79d202..cbc7aa3eaef5 100644 --- a/parquet/tests/arrow_reader/parquet_testing.rs +++ b/parquet/tests/arrow_reader/parquet_testing.rs @@ -20,7 +20,7 @@ //! [parquet-testing]: https://github.com/apache/parquet-testing use arrow_array::cast::AsArray; -use arrow_array::{Array, Int64Array, types}; +use arrow_array::{Array, BinaryArray, Int64Array, StringArray, types}; use arrow_schema::{Field, Schema, TimeUnit}; use parquet::arrow::arrow_reader::{ArrowReaderOptions, ParquetRecordBatchReaderBuilder}; use parquet::basic::{LogicalType, Type as PhysicalType}; @@ -190,3 +190,43 @@ fn test_read_unknown_logical_type() { assert_eq!(out.num_rows(), 3); assert_eq!(out.num_columns(), 2); } + +#[test] +fn test_json_and_bson_logical_types() { + let test_data = arrow::util::test_util::parquet_test_data(); + + let json_file = File::open(format!("{test_data}/json.parquet")).unwrap(); + let mut json_reader = ParquetRecordBatchReaderBuilder::try_new(json_file) + .unwrap() + .build() + .unwrap(); + let json_batch = json_reader.next().unwrap().unwrap(); + assert!(json_reader.next().is_none()); + let json = json_batch.column(0).as_string::(); + assert_eq!( + json, + &StringArray::from(vec![ + Some(r#"{"a":1}"#), + Some(r#"{"a":1,"b":null}"#), + Some("[1,null,3]"), + None, + ]) + ); + + let bson_file = File::open(format!("{test_data}/bson.parquet")).unwrap(); + let mut bson_reader = ParquetRecordBatchReaderBuilder::try_new(bson_file) + .unwrap() + .build() + .unwrap(); + let bson_batch = bson_reader.next().unwrap().unwrap(); + assert!(bson_reader.next().is_none()); + let bson = bson_batch.column(0).as_binary::(); + assert_eq!( + bson, + &BinaryArray::from(vec![ + Some(&[12, 0, 0, 0, 16, 97, 0, 1, 0, 0, 0, 0][..]), + Some(&[15, 0, 0, 0, 16, 97, 0, 1, 0, 0, 0, 10, 98, 0, 0][..]), + None, + ]) + ); +} diff --git a/parquet/tests/variant_integration.rs b/parquet/tests/variant_integration.rs index a6e56b3af682..a80428b9617c 100644 --- a/parquet/tests/variant_integration.rs +++ b/parquet/tests/variant_integration.rs @@ -24,6 +24,7 @@ //! Inspired by the arrow-go implementation: use arrow::util::test_util::parquet_test_data; +use arrow::{array::BinaryArray, record_batch::RecordBatch}; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use parquet_variant::{Variant, VariantMetadata}; use parquet_variant_compute::{VariantArray, unshred_variant}; @@ -421,3 +422,61 @@ static ALL_CASES: LazyLock>> = LazyLock::new(|| { fn all_cases() -> &'static [VariantTestCase] { ALL_CASES.as_ref().unwrap() } + +#[test] +fn test_variant_validation_files() { + const TEST_CASES: &[(&str, bool)] = &[ + ("duplicate_field_offsets.parquet", true), + ("field_id_out_of_range.parquet", false), + ("int_overflow_in_bounds_check.parquet", false), + ("malformed_child_inside_well_formed_parent.parquet", false), + ("negative_dictionary_size.parquet", false), + ("out_of_range_child_offset.parquet", false), + ("out_of_range_dictionary_size.parquet", false), + ("out_of_range_element_count.parquet", false), + ("over_deep_nested_children.parquet", false), + ("oversized_primitive_size.parquet", false), + ("short_string_length_exceeds_buffer.parquet", false), + ("truncated_primitive_size.parquet", false), + ("unknown_primitive_type.parquet", false), + ("variant_version_2_header.parquet", false), + ]; + + let test_dir = PathBuf::from(parquet_test_data()) + .join("..") + .join("bad_data") + .join("variants"); + for (filename, expected_valid) in TEST_CASES { + let path = test_dir.join(filename); + let file = fs::File::open(&path).unwrap(); + let mut reader = ParquetRecordBatchReaderBuilder::try_new(file) + .unwrap() + .build() + .unwrap(); + let batch = reader.next().unwrap().unwrap(); + assert!(reader.next().is_none()); + assert_eq!( + variant_is_valid(&batch, filename), + *expected_valid, + "unexpected validation result for {filename}" + ); + } +} + +fn variant_is_valid(batch: &RecordBatch, filename: &str) -> bool { + assert_eq!(batch.num_rows(), 1, "unexpected row count in {filename}"); + let metadata = batch + .column_by_name("metadata") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let value = batch + .column_by_name("value") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + + Variant::try_new(metadata.value(0), value.value(0)).is_ok() +}