Support writing RunEndEncoded as Parquet #8069
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -222,6 +222,10 @@ impl LevelInfoBuilder { | |
| _ => unreachable!(), | ||
| }) | ||
| } | ||
| DataType::RunEndEncoded(_, v) if is_leaf(v.data_type()) => { | ||
| let levels = ArrayLevels::new(parent_ctx, is_nullable, array.clone()); | ||
| Ok(Self::Primitive(levels)) | ||
| } | ||
vegarsti marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| d => Err(nyi_err!("Datatype {} is not yet supported", d)), | ||
| } | ||
| } | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -1122,6 +1122,17 @@ impl ArrowColumnWriterFactory { | |||||
| ArrowDataType::FixedSizeBinary(_) => out.push(bytes(leaves.next().unwrap())?), | ||||||
| _ => out.push(col(leaves.next().unwrap())?), | ||||||
| }, | ||||||
| ArrowDataType::RunEndEncoded(_, value_type) => match value_type.data_type() { | ||||||
| ArrowDataType::Utf8 | ||||||
| | ArrowDataType::LargeUtf8 | ||||||
| | ArrowDataType::Binary | ||||||
| | ArrowDataType::LargeBinary => out.push(bytes(leaves.next().unwrap())?), | ||||||
| ArrowDataType::Utf8View | ArrowDataType::BinaryView => { | ||||||
| out.push(bytes(leaves.next().unwrap())?) | ||||||
| } | ||||||
| ArrowDataType::FixedSizeBinary(_) => out.push(bytes(leaves.next().unwrap())?), | ||||||
| _ => out.push(col(leaves.next().unwrap())?), | ||||||
| }, | ||||||
| _ => { | ||||||
| return Err(ParquetError::NYI(format!( | ||||||
| "Attempting to write an Arrow type {data_type} to parquet that is not yet implemented" | ||||||
|
|
@@ -1215,6 +1226,41 @@ fn write_leaf(writer: &mut ColumnWriter<'_>, levels: &ArrayLevels) -> Result<usi | |||||
| write_primitive(typed, array.values(), levels) | ||||||
| } | ||||||
| }, | ||||||
| ArrowDataType::RunEndEncoded(_, value_type) => match value_type.data_type() { | ||||||
| ArrowDataType::Decimal32(_, _) => { | ||||||
| let array = arrow_cast::cast(column, value_type.data_type())?; | ||||||
| let array = array | ||||||
| .as_primitive::<Decimal32Type>() | ||||||
| .unary::<_, Int32Type>(|v| v); | ||||||
| write_primitive(typed, array.values(), levels) | ||||||
| } | ||||||
| ArrowDataType::Decimal64(_, _) => { | ||||||
| let array = arrow_cast::cast(column, value_type.data_type())?; | ||||||
| let array = array | ||||||
| .as_primitive::<Decimal64Type>() | ||||||
| .unary::<_, Int32Type>(|v| v as i32); | ||||||
| write_primitive(typed, array.values(), levels) | ||||||
| } | ||||||
| ArrowDataType::Decimal128(_, _) => { | ||||||
| let array = arrow_cast::cast(column, value_type.data_type())?; | ||||||
| let array = array | ||||||
| .as_primitive::<Decimal128Type>() | ||||||
| .unary::<_, Int32Type>(|v| v as i32); | ||||||
| write_primitive(typed, array.values(), levels) | ||||||
| } | ||||||
| ArrowDataType::Decimal256(_, _) => { | ||||||
| let array = arrow_cast::cast(column, value_type.data_type())?; | ||||||
| let array = array | ||||||
| .as_primitive::<Decimal256Type>() | ||||||
| .unary::<_, Int32Type>(|v| v.as_i128() as i32); | ||||||
| write_primitive(typed, array.values(), levels) | ||||||
| } | ||||||
| _ => { | ||||||
| let array = arrow_cast::cast(column, &ArrowDataType::Int32)?; | ||||||
| let array = array.as_primitive::<Int32Type>(); | ||||||
| write_primitive(typed, array.values(), levels) | ||||||
| } | ||||||
| }, | ||||||
| _ => { | ||||||
| let array = arrow_cast::cast(column, &ArrowDataType::Int32)?; | ||||||
| let array = array.as_primitive::<Int32Type>(); | ||||||
|
|
@@ -1297,6 +1343,12 @@ fn write_leaf(writer: &mut ColumnWriter<'_>, levels: &ArrayLevels) -> Result<usi | |||||
| write_primitive(typed, array.values(), levels) | ||||||
| } | ||||||
| }, | ||||||
| ArrowDataType::RunEndEncoded(_run_ends, _values) => { | ||||||
| Err(ParquetError::NYI( | ||||||
| "Int64ColumnWriter: Attempting to write an Arrow REE type that is not yet implemented" | ||||||
|
Member
There was a problem hiding this comment.
Suggested change
nit: IMO it would be easier for the end user to debug/understand the error if it does not use acronyms. |
||||||
| .to_string(), | ||||||
| )) | ||||||
| } | ||||||
| _ => { | ||||||
| let array = arrow_cast::cast(column, &ArrowDataType::Int64)?; | ||||||
| let array = array.as_primitive::<Int64Type>(); | ||||||
|
|
@@ -1371,6 +1423,12 @@ fn write_leaf(writer: &mut ColumnWriter<'_>, levels: &ArrayLevels) -> Result<usi | |||||
| let array = column.as_primitive::<Float16Type>(); | ||||||
| get_float_16_array_slice(array, indices) | ||||||
| } | ||||||
| ArrowDataType::RunEndEncoded(_run_ends, _values) => { | ||||||
| return Err(ParquetError::NYI( | ||||||
| "FixedLenByteArrayColumnWriter: Attempting to write an Arrow REE type that is not yet implemented" | ||||||
|
Member
There was a problem hiding this comment.
Suggested change
nit: IMO it would be easier for the end user to debug/understand the error if it does not use acronyms. |
||||||
| .to_string(), | ||||||
| )); | ||||||
| } | ||||||
| _ => { | ||||||
| return Err(ParquetError::NYI( | ||||||
| "Attempting to write an Arrow type that is not yet implemented".to_string(), | ||||||
|
|
@@ -4481,4 +4539,153 @@ mod tests { | |||||
| assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024); | ||||||
| assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4); | ||||||
| } | ||||||
|
|
||||||
| #[test] | ||||||
| fn arrow_writer_run_end_encoded_string() { | ||||||
| // Create a run array of strings | ||||||
| let mut builder = StringRunBuilder::<Int32Type>::new(); | ||||||
| builder.extend( | ||||||
| vec![Some("alpha"); 100000] | ||||||
| .into_iter() | ||||||
| .chain(vec![Some("beta"); 100000]), | ||||||
| ); | ||||||
| let run_array: RunArray<Int32Type> = builder.finish(); | ||||||
| let schema = Arc::new(Schema::new(vec![Field::new( | ||||||
| "ree", | ||||||
| run_array.data_type().clone(), | ||||||
| run_array.is_nullable(), | ||||||
| )])); | ||||||
|
|
||||||
| // Write to parquet | ||||||
| let mut parquet_bytes: Vec<u8> = Vec::new(); | ||||||
| let mut writer = ArrowWriter::try_new(&mut parquet_bytes, schema.clone(), None).unwrap(); | ||||||
| let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(run_array)]).unwrap(); | ||||||
| writer.write(&batch).unwrap(); | ||||||
| writer.close().unwrap(); | ||||||
|
|
||||||
| // Read back and verify | ||||||
| let bytes = Bytes::from(parquet_bytes); | ||||||
| let reader = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap(); | ||||||
|
|
||||||
| // Check if dictionary was used by examining the metadata | ||||||
| let metadata = reader.metadata(); | ||||||
| let row_group = &metadata.row_groups()[0]; | ||||||
| let col_meta = &row_group.columns()[0]; | ||||||
|
|
||||||
| // If dictionary encoding worked, we should see RLE_DICTIONARY encoding | ||||||
| // and have a dictionary page offset | ||||||
| let has_dict_encoding = col_meta.encodings().any(|e| e == Encoding::RLE_DICTIONARY); | ||||||
| let has_dict_page = col_meta.dictionary_page_offset().is_some(); | ||||||
|
|
||||||
| // Verify the schema is REE encoded when we read it back | ||||||
| let expected_schema = Arc::new(Schema::new(vec![Field::new( | ||||||
| "ree", | ||||||
| DataType::RunEndEncoded( | ||||||
| Arc::new(Field::new("run_ends", arrow_schema::DataType::Int32, false)), | ||||||
| Arc::new(Field::new("values", arrow_schema::DataType::Utf8, true)), | ||||||
| ), | ||||||
| false, | ||||||
| )])); | ||||||
| assert_eq!(&expected_schema, reader.schema()); | ||||||
|
|
||||||
| // Read the data back | ||||||
| let batches: Vec<_> = reader | ||||||
| .build() | ||||||
| .unwrap() | ||||||
| .collect::<ArrowResult<Vec<_>>>() | ||||||
| .unwrap(); | ||||||
| assert_eq!(batches.len(), 196); | ||||||
| // Count rows in total | ||||||
| let total_rows = batches.iter().map(|b| b.num_rows()).sum::<usize>(); | ||||||
| assert_eq!(total_rows, 200000); | ||||||
|
|
||||||
| // Ensure dictionary encoding | ||||||
| assert!(has_dict_encoding, "RunArray should be dictionary encoded"); | ||||||
| assert!(has_dict_page, "RunArray should have dictionary page"); | ||||||
| } | ||||||
|
|
||||||
| #[test] | ||||||
| fn arrow_writer_run_end_encoded_int() { | ||||||
| // Create a run array of strings | ||||||
| let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new(); | ||||||
| builder.extend( | ||||||
| vec![Some(1); 100000] | ||||||
| .into_iter() | ||||||
| .chain(vec![Some(2); 100000]), | ||||||
| ); | ||||||
| let run_array: RunArray<Int32Type> = builder.finish(); | ||||||
| let schema = Arc::new(Schema::new(vec![Field::new( | ||||||
| "ree", | ||||||
| run_array.data_type().clone(), | ||||||
| run_array.is_nullable(), | ||||||
| )])); | ||||||
|
|
||||||
| // Write to parquet | ||||||
| let mut parquet_bytes: Vec<u8> = Vec::new(); | ||||||
| let mut writer = ArrowWriter::try_new(&mut parquet_bytes, schema.clone(), None).unwrap(); | ||||||
| let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(run_array)]).unwrap(); | ||||||
| writer.write(&batch).unwrap(); | ||||||
| writer.close().unwrap(); | ||||||
|
|
||||||
| // Read back and verify | ||||||
| let bytes = Bytes::from(parquet_bytes); | ||||||
| let reader = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap(); | ||||||
|
|
||||||
| // Check if dictionary was used by examining the metadata | ||||||
| let metadata = reader.metadata(); | ||||||
| let row_group = &metadata.row_groups()[0]; | ||||||
| let col_meta = &row_group.columns()[0]; | ||||||
| let has_dict_encoding = col_meta.encodings().any(|e| e == Encoding::RLE_DICTIONARY); | ||||||
|
|
||||||
| // If dictionary encoding worked, we should see RLE_DICTIONARY encoding | ||||||
| // and have a dictionary page offset | ||||||
| // let has_dict_encoding = col_meta.encodings().contains(&Encoding::RLE_DICTIONARY); | ||||||
| let has_dict_page = col_meta.dictionary_page_offset().is_some(); | ||||||
|
|
||||||
| // Verify the schema is REE encoded when we read it back | ||||||
| let expected_schema = Arc::new(Schema::new(vec![Field::new( | ||||||
| "ree", | ||||||
| DataType::RunEndEncoded( | ||||||
| Arc::new(Field::new("run_ends", arrow_schema::DataType::Int32, false)), | ||||||
| Arc::new(Field::new("values", arrow_schema::DataType::Int32, true)), | ||||||
| ), | ||||||
| false, | ||||||
| )])); | ||||||
| assert_eq!(&expected_schema, reader.schema()); | ||||||
|
|
||||||
| // Read the data back | ||||||
| let batches: Vec<_> = reader | ||||||
| .build() | ||||||
| .unwrap() | ||||||
| .collect::<ArrowResult<Vec<_>>>() | ||||||
| .unwrap(); | ||||||
| assert_eq!(batches.len(), 196); | ||||||
| // Count rows in total | ||||||
| let total_rows = batches.iter().map(|b| b.num_rows()).sum::<usize>(); | ||||||
| assert_eq!(total_rows, 200000); | ||||||
|
|
||||||
| // Ensure dictionary encoding | ||||||
| assert!(has_dict_encoding, "RunArray should be dictionary encoded"); | ||||||
| assert!(has_dict_page, "RunArray should have dictionary page"); | ||||||
| } | ||||||
|
|
||||||
| #[test] | ||||||
| fn arrow_writer_round_trip_run_end_encoded_string() { | ||||||
| // Create a run array of strings (cannot have more than 1024 values per record batch) | ||||||
| let mut builder = StringRunBuilder::<Int32Type>::new(); | ||||||
| builder.extend( | ||||||
| vec![Some("alpha"); 512] | ||||||
| .into_iter() | ||||||
| .chain(vec![Some("beta"); 512]), | ||||||
| ); | ||||||
| let run_array: RunArray<Int32Type> = builder.finish(); | ||||||
| let schema = Arc::new(Schema::new(vec![Field::new( | ||||||
| "ree", | ||||||
| run_array.data_type().clone(), | ||||||
| run_array.is_nullable(), | ||||||
| )])); | ||||||
|
|
||||||
| let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(run_array)]).unwrap(); | ||||||
| roundtrip(batch, None); | ||||||
| } | ||||||
|
Member
There was a problem hiding this comment. Tests for decimal REE arrays are missing |
||||||
| } | ||||||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.