Jefffrey commented on code in PR #11764:
URL: https://github.com/apache/arrow-rs/pull/11764#discussion_r4442061137
##########
arrow-row/src/run.rs:
##########
@@ -140,25 +142,45 @@ pub unsafe fn decode<R: RunEndIndexType>(
// Convert the unique decoded values using the row converter
let mut unique_rows: Vec<&[u8]> = decoded_values.iter().map(|v|
v.as_slice()).collect();
- let values = if unique_rows.is_empty() {
- unsafe { converter.convert_raw(&mut [], validate_utf8) }?
- } else {
- unsafe { converter.convert_raw(&mut unique_rows, validate_utf8) }?
- };
+ let values = unsafe { converter.convert_raw(&mut unique_rows,
validate_utf8) }?;
- // Create run ends array
- let run_ends_array =
PrimitiveArray::<R>::try_new(ScalarBuffer::from(run_ends), None)?;
+ let run_ends = PrimitiveArray::<R>::try_new(ScalarBuffer::from(run_ends),
None)?;
+ let (_, run_ends, values) = RunArray::<R>::try_new(&run_ends,
&values[0])?.into_parts();
+ let data_type = decoded_run_type(&field.data_type, values.data_type());
+ // SAFETY: run ends and value count are validated above; fields match the
decoded types.
+ Ok(unsafe { RunArray::<R>::new_unchecked(data_type, run_ends, values) })
Review Comment:
it seems confusing we build the run array for validation, then deconstruct
it so we can replace the type; is there a reason we cant calculate the correct
type upfront?
##########
arrow-row/src/run.rs:
##########
@@ -140,25 +142,45 @@ pub unsafe fn decode<R: RunEndIndexType>(
// Convert the unique decoded values using the row converter
let mut unique_rows: Vec<&[u8]> = decoded_values.iter().map(|v|
v.as_slice()).collect();
- let values = if unique_rows.is_empty() {
- unsafe { converter.convert_raw(&mut [], validate_utf8) }?
- } else {
- unsafe { converter.convert_raw(&mut unique_rows, validate_utf8) }?
- };
+ let values = unsafe { converter.convert_raw(&mut unique_rows,
validate_utf8) }?;
- // Create run ends array
- let run_ends_array =
PrimitiveArray::<R>::try_new(ScalarBuffer::from(run_ends), None)?;
+ let run_ends = PrimitiveArray::<R>::try_new(ScalarBuffer::from(run_ends),
None)?;
+ let (_, run_ends, values) = RunArray::<R>::try_new(&run_ends,
&values[0])?.into_parts();
+ let data_type = decoded_run_type(&field.data_type, values.data_type());
+ // SAFETY: run ends and value count are validated above; fields match the
decoded types.
+ Ok(unsafe { RunArray::<R>::new_unchecked(data_type, run_ends, values) })
+}
- // Create the RunEndEncodedArray
- RunArray::<R>::try_new(&run_ends_array, &values[0])
+fn decoded_run_type(data_type: &DataType, values_type: &DataType) -> DataType {
+ match data_type {
+ DataType::RunEndEncoded(run_ends_field, values_field)
+ if values_field.data_type() != values_type =>
Review Comment:
is this function supposed to be like what we do for lists?
https://github.com/apache/arrow-rs/blob/e2c48ef0dd83fa464d643e25f8e9f9177ca7316e/arrow-row/src/list.rs#L258-L269
if so, why do we have this equality check here whereas for lists we always
overwrite the type?
##########
arrow-row/src/run.rs:
##########
@@ -730,6 +793,67 @@ mod tests {
assert_eq!(result_ree.len(), 0);
}
+ #[test]
+ fn test_run_end_encoded_dictionary_values() {
Review Comment:
this test code doesnt read well, with how it constructs the test data via
mapping the iterators. can we simplify it; maybe its easier to just have them
as separate tests, etc.
##########
arrow-row/src/run.rs:
##########
@@ -730,6 +793,67 @@ mod tests {
assert_eq!(result_ree.len(), 0);
}
+ #[test]
+ fn test_run_end_encoded_dictionary_values() {
+ let dictionary: DictionaryArray<Int32Type> = ["a",
"b"].into_iter().collect();
+ let ends = PrimitiveArray::<Int32Type>::from(vec![2, 4]).to_data();
+ for nested in [false, true] {
+ let values = [
+ Arc::new(dictionary.clone()) as ArrayRef,
+ Arc::new(StringArray::from(vec!["a", "b"])) as ArrayRef,
+ ]
+ .map(|values| {
+ if nested {
+ Arc::new(StructArray::from(vec![(
+ Arc::new(Field::new("child",
values.data_type().clone(), false)),
+ values,
+ )])) as ArrayRef
+ } else {
+ values
+ }
+ });
+ let [input, expected] = values.map(|values| {
+ let data_type = DataType::RunEndEncoded(
+ Arc::new(Field::new("ends", DataType::Int32, false)),
+ Arc::new(
+ Field::new("labels", values.data_type().clone(), false)
+ .with_metadata([("source".to_string(),
"test".to_string())]),
+ ),
+ );
+ RunArray::<Int32Type>::from(
+ ArrayData::builder(data_type)
+ .len(4)
+ .child_data(vec![ends.clone(), values.to_data()])
+ .build()
+ .unwrap(),
+ )
+ });
+ let converter =
+
RowConverter::new(vec![SortField::new(input.data_type().clone())]).unwrap();
+ for (offset, len) in [(0, 4), (1, 3), (4, 0)] {
+ let rows = converter
+ .convert_columns(&[Arc::new(input.slice(offset, len))])
+ .unwrap();
+ let actual = converter.convert_rows(&rows).unwrap();
+ assert_eq!(actual[0].to_data(), expected.slice(offset,
len).to_data());
+ actual[0].to_data().validate_full().unwrap();
+ }
+ }
+ }
+
+ #[test]
+ fn test_run_end_encoded_decode_overflow() {
+ let array = RunArray::<Int16Type>::try_new(
+ &PrimitiveArray::from(vec![1]),
+ &Int64Array::from(vec![42]),
+ )
+ .unwrap();
+ let converter =
RowConverter::new(vec![SortField::new(array.data_type().clone())]).unwrap();
+ let rows = converter.convert_columns(&[Arc::new(array)]).unwrap();
+ let too_many_rows = std::iter::repeat_with(||
rows.row(0)).take(i16::MAX as usize + 1);
+ assert!(converter.convert_rows(too_many_rows).is_err());
Review Comment:
how much memory would this test allocate if we're trying to push the bounds?
##########
arrow-select/src/interleave.rs:
##########
@@ -1004,6 +1003,23 @@ mod tests {
use arrow_buffer::ScalarBuffer;
use arrow_schema::{Field, Schema, TimeUnit};
+ #[test]
+ fn test_interleave_preserves_run_end_type() {
+ for array in crate::test_util::run_arrays() {
+ let selected: Vec<u32> = (0..array.len() as
u32).step_by(2).collect();
+ let indices: Vec<_> = selected
+ .iter()
+ .enumerate()
+ .map(|(i, row)| (i % 2, *row as usize))
+ .collect();
+ let expected = crate::take::take(&array,
&UInt32Array::from(selected), None).unwrap();
Review Comment:
im hesistant on having take as the oracle; can we have proper expected data
created, instead of relying on take being correct here? same for filter
##########
arrow-select/src/coalesce.rs:
##########
@@ -943,14 +943,37 @@ mod tests {
use arrow_array::types::{Int32Type, UInt8Type};
use arrow_array::{
BinaryViewArray, DictionaryArray, Int32Array, Int64Array,
RecordBatchOptions, StringArray,
- StringViewArray, TimestampNanosecondArray, UInt8Array, UInt32Array,
UInt64Array,
- make_array, record_batch,
+ StringViewArray, StructArray, TimestampNanosecondArray, UInt8Array,
UInt32Array,
+ UInt64Array, make_array, record_batch,
};
use arrow_buffer::BooleanBufferBuilder;
use arrow_schema::{DataType, Field, Schema};
use rand::{RngExt, SeedableRng};
use std::ops::Range;
+ #[test]
+ fn test_filter_preserves_nested_run_end_type() {
Review Comment:
is this necessary if we're testing filter? isn't that covered in filter.rs?
##########
arrow-select/src/interleave.rs:
##########
@@ -701,12 +701,11 @@ fn interleave_run_end<R: RunEndIndexType>(
run_ends_buf.push(R::Native::from_usize(n).unwrap());
let taken_values = interleave(&value_arrays, &dedup_pairs)?;
- let run_ends = PrimitiveArray::<R>::from_iter_values(run_ends_buf);
-
- Ok(Arc::new(RunArray::<R>::try_new(
- &run_ends,
- taken_values.as_ref(),
- )?))
+ let run_ends = RunEndBuffer::new(run_ends_buf.into(), 0, n);
+ // SAFETY: inputs have the same type and interleaving produces one value
per run.
Review Comment:
`and interleaving produces one value per run.` isnt a strong enough safety
comment i feel; maybe we are better off having a better api for run array to
allow creating with fields we want, in a safe api, instead of relying on the
unchecked constructor
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
