From 810f9d964061ebe252ae3c0d37c0fcceef141f2d Mon Sep 17 00:00:00 2001 From: Wei Yan Date: Tue, 25 Aug 2026 21:01:59 -0700 Subject: [PATCH 1/4] perf(spark): reduce map deduplication allocations --- datafusion/spark/src/function/map/utils.rs | 67 ++++++++++++++++++++-- 1 file changed, 61 insertions(+), 6 deletions(-) diff --git a/datafusion/spark/src/function/map/utils.rs b/datafusion/spark/src/function/map/utils.rs index fa6b2a960dabb..9fe9e56985d47 100644 --- a/datafusion/spark/src/function/map/utils.rs +++ b/datafusion/spark/src/function/map/utils.rs @@ -161,6 +161,9 @@ fn map_deduplicate_keys( values_nulls: Option<&NullBuffer>, last_value_wins: bool, ) -> Result<(ArrayRef, ArrayRef, OffsetBuffer)> { + const MIN_RETAINED_LOOKUP_CAPACITY: usize = 16; + const MAX_RETAINED_LOOKUP_CAPACITY_RATIO: usize = 4; + let offsets_len = keys_offsets.len(); let mut new_offsets = Vec::with_capacity(offsets_len); @@ -178,11 +181,15 @@ fn map_deduplicate_keys( // Mirror Spark's `ArrayBasedMapBuilder`: the first occurrence of a key // fixes its position in the output; under LAST_WIN a later duplicate - // overwrites that slot's value. `keys_mask` selects the first-seen keys, - // `value_indices` records the source index in `flat_values` to materialize - // for each output slot (updated in place on overwrite). + // overwrites that slot's value. `keys_mask` selects the first-seen keys. + // Keep a matching `values_mask` for the common case where no overwrite + // occurs, so Arrow's all-true filter path can share the original value + // buffers. Only materialize `value_indices` with `take` after an actual + // LAST_WIN overwrite changes the value selection order. let mut keys_mask_builder = BooleanBuilder::new(); + let mut values_mask_builder = BooleanBuilder::new(); let mut value_indices: Vec = Vec::new(); + let mut needs_value_take = false; let mut key_to_output_idx: HashMap = HashMap::new(); for (row_idx, (next_keys_offset, next_values_offset)) in keys_offsets .iter() @@ -202,7 +209,15 @@ fn map_deduplicate_keys( "map_deduplicate_keys: keys and values lists in the same row must have equal lengths" ); } + // Reuse normal-sized tables. Only shrink when a prior row left a + // table much larger than both the current row and a small floor. + let target_capacity = num_keys_entries.max(MIN_RETAINED_LOOKUP_CAPACITY); key_to_output_idx.clear(); + if key_to_output_idx.capacity() + > target_capacity.saturating_mul(MAX_RETAINED_LOOKUP_CAPACITY_RATIO) + { + key_to_output_idx.shrink_to(target_capacity); + } for cur_entry_idx in 0..num_keys_entries { let key = ScalarValue::try_from_array( &flat_keys, @@ -213,8 +228,10 @@ fn map_deduplicate_keys( if let Some(&output_idx) = key_to_output_idx.get(&key) { if last_value_wins { + needs_value_take = true; value_indices[output_idx] = abs_value_idx; keys_mask_builder.append_value(false); + values_mask_builder.append_value(false); continue; } return exec_err!( @@ -225,23 +242,30 @@ fn map_deduplicate_keys( ); } keys_mask_builder.append_value(true); + values_mask_builder.append_value(true); key_to_output_idx.insert(key, value_indices.len()); value_indices.push(abs_value_idx); new_last_offset += 1; } } else { // The result entry is NULL — no keys/values emitted. Still pad the - // mask so it stays aligned with `flat_keys`. + // masks so they stay aligned with their respective flat arrays. keys_mask_builder.append_n(num_keys_entries, false); + values_mask_builder.append_n(num_values_entries, false); } new_offsets.push(new_last_offset); cur_keys_offset += num_keys_entries; cur_values_offset += num_values_entries; } let keys_mask = keys_mask_builder.finish(); + let values_mask = values_mask_builder.finish(); let needed_keys = filter(&flat_keys, &keys_mask)?; - let value_indices_array = Int32Array::from(value_indices); - let needed_values = take(&flat_values, &value_indices_array, None)?; + let needed_values = if needs_value_take { + let value_indices_array = Int32Array::from(value_indices); + take(&flat_values, &value_indices_array, None)? + } else { + filter(&flat_values, &values_mask)? + }; let offsets = OffsetBuffer::new(new_offsets.into()); Ok((needed_keys, needed_values, offsets)) } @@ -276,6 +300,37 @@ mod tests { assert_eq!(map.value_offsets(), &[0, 2, 3]); } + #[test] + fn distinct_keys_reuse_value_payload_buffer() { + let (keys, values) = + int32_utf8_inputs(vec![1, 2, 3], vec![Some("a"), Some("b"), Some("c")]); + let offsets = [0i32, 3]; + let source_values = values.as_any().downcast_ref::().unwrap(); + + for last_value_wins in [false, true] { + let (_, needed_values, _) = map_deduplicate_keys( + &keys, + &values, + &offsets, + &offsets, + None, + None, + last_value_wins, + ) + .unwrap(); + let needed_values = needed_values + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!( + source_values.value_data().as_ptr(), + needed_values.value_data().as_ptr(), + "distinct keys should not copy values under last_value_wins={last_value_wins}" + ); + } + } + #[test] fn single_row_duplicate_errors_under_exception() { let (keys, values) = From f89176ef1b67482e400dd5966a11fe74e345995c Mon Sep 17 00:00:00 2001 From: Wei Yan Date: Thu, 27 Aug 2026 10:12:37 -0700 Subject: [PATCH 2/4] fix(spark): preserve sliced map values and expand dedup coverage Align the value filter with the starting list offset while keeping the zero-offset path free of extra slice allocations. Add regression coverage for wide-to-small rows, LAST_WIN ordering, sliced values, and null rows with unequal child lengths. --- datafusion/spark/src/function/map/utils.rs | 182 ++++++++++++++++++++- 1 file changed, 180 insertions(+), 2 deletions(-) diff --git a/datafusion/spark/src/function/map/utils.rs b/datafusion/spark/src/function/map/utils.rs index 9fe9e56985d47..2b51917824dd2 100644 --- a/datafusion/spark/src/function/map/utils.rs +++ b/datafusion/spark/src/function/map/utils.rs @@ -171,10 +171,11 @@ fn map_deduplicate_keys( .first() .map(|offset| *offset as usize) .unwrap_or(0); - let mut cur_values_offset = values_offsets + let values_start_offset = values_offsets .first() .map(|offset| *offset as usize) .unwrap_or(0); + let mut cur_values_offset = values_start_offset; let mut new_last_offset = 0; new_offsets.push(new_last_offset); @@ -264,7 +265,15 @@ fn map_deduplicate_keys( let value_indices_array = Int32Array::from(value_indices); take(&flat_values, &value_indices_array, None)? } else { - filter(&flat_values, &values_mask)? + // Values lists can be sliced independently of keys lists. Align the + // relative mask with their first offset, without allocating a slice + // wrapper for the common case where values start at zero. + let flat_values = if values_start_offset == 0 { + Cow::Borrowed(flat_values) + } else { + Cow::Owned(flat_values.slice(values_start_offset, values_mask.len())) + }; + filter(flat_values.as_ref(), &values_mask)? }; let offsets = OffsetBuffer::new(new_offsets.into()); Ok((needed_keys, needed_values, offsets)) @@ -331,6 +340,175 @@ mod tests { } } + #[test] + fn wide_row_followed_by_small_rows() { + // The wide row grows the lookup table beyond the retention threshold + // for the following small and empty rows. + let wide_len = 512; + let keys = (0..wide_len).chain([0, 1, 0]).collect(); + let values = std::iter::repeat_n(Some("wide"), wide_len as usize) + .chain([Some("small"), None, Some("last")]) + .collect(); + let (keys, values) = int32_utf8_inputs(keys, values); + let offsets = [0, wide_len, wide_len + 2, wide_len + 2, wide_len + 3]; + + for last_value_wins in [false, true] { + let result = map_from_keys_values_offsets_nulls( + &keys, + &values, + &offsets, + &offsets, + None, + None, + last_value_wins, + ) + .unwrap(); + + let map = result.as_map(); + assert_eq!(map.value_offsets(), &offsets); + assert_eq!(map.keys().to_data(), keys.to_data()); + assert_eq!(map.values().to_data(), values.to_data()); + } + } + + #[test] + fn last_win_after_wide_row_preserves_key_order_and_values() { + let wide_len = 512; + let keys = (0..wide_len).chain([1, 0, 1, 2, 0, 1]).collect(); + let values = std::iter::repeat_n(Some("wide"), wide_len as usize) + .chain([ + Some("old-a"), + Some("old-b"), + Some("new-a"), + Some("c"), + None, + Some("next-row"), + ]) + .collect(); + let (keys, values) = int32_utf8_inputs(keys, values); + let offsets = [0, wide_len, wide_len + 5, wide_len + 6]; + + let result = map_from_keys_values_offsets_nulls( + &keys, &values, &offsets, &offsets, None, None, true, + ) + .unwrap(); + + let expected_keys = (0..wide_len).chain([1, 0, 2, 1]).collect(); + let expected_values = std::iter::repeat_n(Some("wide"), wide_len as usize) + .chain([Some("new-a"), None, Some("c"), Some("next-row")]) + .collect(); + let (expected_keys, expected_values) = + int32_utf8_inputs(expected_keys, expected_values); + let map = result.as_map(); + assert_eq!( + map.value_offsets(), + &[0, wide_len, wide_len + 3, wide_len + 4] + ); + assert_eq!(map.keys().to_data(), expected_keys.to_data()); + assert_eq!(map.values().to_data(), expected_values.to_data()); + } + + #[test] + fn distinct_keys_with_nonzero_values_offset_reuse_payload_buffer() { + let (keys, values) = int32_utf8_inputs( + vec![1, 2, 3], + vec![Some("prefix"), Some("a"), None, Some("c"), Some("suffix")], + ); + let keys_offsets = [0, 2, 3]; + let values_offsets = [1, 3, 4]; + let expected_values = StringArray::from(vec![Some("a"), None, Some("c")]); + let source_values = values.as_any().downcast_ref::().unwrap(); + + for last_value_wins in [false, true] { + let result = map_from_keys_values_offsets_nulls( + &keys, + &values, + &keys_offsets, + &values_offsets, + None, + None, + last_value_wins, + ) + .unwrap(); + + let map = result.as_map(); + assert_eq!(map.value_offsets(), &keys_offsets); + assert_eq!(map.keys().to_data(), keys.to_data()); + let needed_values = + map.values().as_any().downcast_ref::().unwrap(); + assert_eq!(needed_values, &expected_values); + assert_eq!( + needed_values.value_data().as_ptr(), + source_values.value_data().as_ptr(), + ); + } + } + + #[test] + fn last_win_with_nonzero_values_offset() { + let (keys, values) = int32_utf8_inputs( + vec![1, 2, 1], + vec![Some("prefix"), Some("a"), Some("b"), None, Some("suffix")], + ); + let result = map_from_keys_values_offsets_nulls( + &keys, + &values, + &[0, 3], + &[1, 4], + None, + None, + true, + ) + .unwrap(); + + let map = result.as_map(); + assert_eq!(map.value_offsets(), &[0, 2]); + assert_eq!(map.keys().to_data(), Int32Array::from(vec![1, 2]).to_data()); + assert_eq!( + map.values().to_data(), + StringArray::from(vec![None, Some("b")]).to_data(), + ); + } + + #[test] + fn null_row_with_unequal_child_lengths_keeps_values_aligned() { + let (keys, values) = int32_utf8_inputs( + vec![1, 9, 9, 3, 4], + vec![Some("first"), Some("ignored"), Some("last"), None], + ); + let keys_offsets = [0, 1, 3, 5]; + let values_offsets = [0, 1, 2, 4]; + let nulls = NullBuffer::from(vec![true, false, true]); + + for last_value_wins in [false, true] { + for (keys_nulls, values_nulls) in [(Some(&nulls), None), (None, Some(&nulls))] + { + let result = map_from_keys_values_offsets_nulls( + &keys, + &values, + &keys_offsets, + &values_offsets, + keys_nulls, + values_nulls, + last_value_wins, + ) + .unwrap(); + + let map = result.as_map(); + assert_eq!(map.value_offsets(), &[0, 1, 1, 3]); + assert_eq!(map.nulls(), Some(&nulls)); + assert_eq!( + map.keys().to_data(), + Int32Array::from(vec![1, 3, 4]).to_data(), + ); + assert_eq!( + map.values().to_data(), + StringArray::from(vec![Some("first"), Some("last"), None]).to_data(), + ); + } + } + } + #[test] fn single_row_duplicate_errors_under_exception() { let (keys, values) = From dd83e70097aa9d3b68e3590379e60bca3929c8f3 Mon Sep 17 00:00:00 2001 From: Wei Yan Date: Thu, 27 Aug 2026 20:10:53 -0700 Subject: [PATCH 3/4] fix(spark): handle large map value offsets --- datafusion/spark/src/function/map/utils.rs | 104 +++++++++++++++++++-- 1 file changed, 98 insertions(+), 6 deletions(-) diff --git a/datafusion/spark/src/function/map/utils.rs b/datafusion/spark/src/function/map/utils.rs index 2b51917824dd2..5d2ece5b1a1b9 100644 --- a/datafusion/spark/src/function/map/utils.rs +++ b/datafusion/spark/src/function/map/utils.rs @@ -185,12 +185,14 @@ fn map_deduplicate_keys( // overwrites that slot's value. `keys_mask` selects the first-seen keys. // Keep a matching `values_mask` for the common case where no overwrite // occurs, so Arrow's all-true filter path can share the original value - // buffers. Only materialize `value_indices` with `take` after an actual - // LAST_WIN overwrite changes the value selection order. + // buffers. Use `take` for value reordering after a LAST_WIN overwrite or + // for narrowed large child offsets that cannot be used directly by `slice`. let mut keys_mask_builder = BooleanBuilder::new(); let mut values_mask_builder = BooleanBuilder::new(); let mut value_indices: Vec = Vec::new(); - let mut needs_value_take = false; + // LargeList offsets can narrow to negative i32 values. Keep take's index + // handling for those offsets instead of sign-extending them for slice. + let mut needs_value_take = values_offsets.first().is_some_and(|offset| *offset < 0); let mut key_to_output_idx: HashMap = HashMap::new(); for (row_idx, (next_keys_offset, next_values_offset)) in keys_offsets .iter() @@ -200,6 +202,7 @@ fn map_deduplicate_keys( { let num_keys_entries = *next_keys_offset as usize - cur_keys_offset; let num_values_entries = *next_values_offset as usize - cur_values_offset; + needs_value_take |= *next_values_offset < 0; let key_is_valid = keys_nulls.is_none_or(|buf| buf.is_valid(row_idx)); let value_is_valid = values_nulls.is_none_or(|buf| buf.is_valid(row_idx)); @@ -250,9 +253,11 @@ fn map_deduplicate_keys( } } else { // The result entry is NULL — no keys/values emitted. Still pad the - // masks so they stay aligned with their respective flat arrays. + // masks used by filter so they stay aligned with their flat arrays. keys_mask_builder.append_n(num_keys_entries, false); - values_mask_builder.append_n(num_values_entries, false); + if !needs_value_take { + values_mask_builder.append_n(num_values_entries, false); + } } new_offsets.push(new_last_offset); cur_keys_offset += num_keys_entries; @@ -282,7 +287,7 @@ fn map_deduplicate_keys( #[cfg(test)] mod tests { use super::*; - use arrow::array::{Int32Array, StringArray}; + use arrow::array::{Int32Array, LargeListArray, NullArray, StringArray}; fn int32_utf8_inputs( keys: Vec, @@ -444,6 +449,93 @@ mod tests { } } + #[test] + fn large_list_values_near_i32_offset_limit() { + let keys: ArrayRef = Arc::new(Int32Array::from(vec![7])); + keys.to_data().validate_full().unwrap(); + + for start in [(1_i64 << 31) - 2, 1_i64 << 31] { + for row_valid in [true, false] { + // NullArray represents the large child without a large allocation. + let values: ArrayRef = Arc::new(LargeListArray::new( + Arc::new(Field::new("item", DataType::Null, true)), + OffsetBuffer::new(vec![start, start + 1].into()), + Arc::new(NullArray::new(start as usize + 1)), + Some(NullBuffer::from(vec![row_valid])), + )); + values.to_data().validate_full().unwrap(); + let values_offsets = get_list_offsets(&values).unwrap(); + let expected_keys = if row_valid { vec![7] } else { vec![] }; + + for last_value_wins in [false, true] { + let result = map_from_keys_values_offsets_nulls( + &keys, + get_list_values(&values).unwrap(), + &[0, 1], + &values_offsets, + None, + values.nulls(), + last_value_wins, + ) + .unwrap(); + + result.to_data().validate_full().unwrap(); + let map = result.as_map(); + assert_eq!(map.len(), 1); + assert_eq!(map.is_null(0), !row_valid); + assert_eq!(map.value_offsets(), &[0, expected_keys.len() as i32]); + assert_eq!( + map.keys().to_data(), + Int32Array::from(expected_keys.clone()).to_data(), + ); + assert_eq!( + map.values().to_data(), + NullArray::new(expected_keys.len()).to_data(), + ); + } + } + } + } + + #[test] + fn null_large_list_values_cross_i32_offset_limit() { + let start = (1_i64 << 31) - 1; + let keys: ArrayRef = Arc::new(Int32Array::from(vec![7, 8])); + // The first values row is NULL and crosses the narrowing boundary. + // The second row must still produce its map entry at the large offset. + let values: ArrayRef = Arc::new(LargeListArray::new( + Arc::new(Field::new("item", DataType::Null, true)), + OffsetBuffer::new(vec![start, start + 1, start + 2].into()), + Arc::new(NullArray::new(start as usize + 2)), + Some(NullBuffer::from(vec![false, true])), + )); + keys.to_data().validate_full().unwrap(); + values.to_data().validate_full().unwrap(); + let values_offsets = get_list_offsets(&values).unwrap(); + + for last_value_wins in [false, true] { + let result = map_from_keys_values_offsets_nulls( + &keys, + get_list_values(&values).unwrap(), + &[0, 1, 2], + &values_offsets, + None, + values.nulls(), + last_value_wins, + ) + .unwrap(); + + result.to_data().validate_full().unwrap(); + let map = result.as_map(); + assert_eq!(map.len(), 2); + assert!(map.is_null(0)); + assert!(map.is_valid(1)); + assert_eq!(map.value_offsets(), &[0, 0, 1]); + assert_eq!(map.keys().to_data(), Int32Array::from(vec![8]).to_data()); + assert_eq!(map.values().to_data(), NullArray::new(1).to_data()); + } + } + #[test] fn last_win_with_nonzero_values_offset() { let (keys, values) = int32_utf8_inputs( From b0b359d1ea69beafc1c33a6db5be59cb1fe8ddf0 Mon Sep 17 00:00:00 2001 From: Wei Yan Date: Mon, 31 Aug 2026 13:33:42 -0700 Subject: [PATCH 4/4] fix(spark): avoid allocations for discarded map values --- datafusion/spark/src/function/map/utils.rs | 251 +++++++++++++++++++-- 1 file changed, 227 insertions(+), 24 deletions(-) diff --git a/datafusion/spark/src/function/map/utils.rs b/datafusion/spark/src/function/map/utils.rs index 5d2ece5b1a1b9..1a9bbe69a07f2 100644 --- a/datafusion/spark/src/function/map/utils.rs +++ b/datafusion/spark/src/function/map/utils.rs @@ -21,6 +21,7 @@ use std::sync::Arc; use arrow::array::{ Array, ArrayRef, AsArray, BooleanBuilder, Int32Array, MapArray, StructArray, + new_empty_array, }; use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::compute::{filter, take}; @@ -161,6 +162,9 @@ fn map_deduplicate_keys( values_nulls: Option<&NullBuffer>, last_value_wins: bool, ) -> Result<(ArrayRef, ArrayRef, OffsetBuffer)> { + // Round-number heuristics balance retaining space for small rows against + // resizing after a wide row. HashMap::capacity counts insertable entries, + // not buckets, so the ratio is an approximate retention limit. const MIN_RETAINED_LOOKUP_CAPACITY: usize = 16; const MAX_RETAINED_LOOKUP_CAPACITY_RATIO: usize = 4; @@ -183,12 +187,10 @@ fn map_deduplicate_keys( // Mirror Spark's `ArrayBasedMapBuilder`: the first occurrence of a key // fixes its position in the output; under LAST_WIN a later duplicate // overwrites that slot's value. `keys_mask` selects the first-seen keys. - // Keep a matching `values_mask` for the common case where no overwrite - // occurs, so Arrow's all-true filter path can share the original value - // buffers. Use `take` for value reordering after a LAST_WIN overwrite or - // for narrowed large child offsets that cannot be used directly by `slice`. + // Share the value buffers with a slice only if every value is retained in + // order. Once a value is skipped or overwritten, use `take`: a filter mask + // would allocate for skipped spans and can materialize NULL list children. let mut keys_mask_builder = BooleanBuilder::new(); - let mut values_mask_builder = BooleanBuilder::new(); let mut value_indices: Vec = Vec::new(); // LargeList offsets can narrow to negative i32 values. Keep take's index // handling for those offsets instead of sign-extending them for slice. @@ -235,7 +237,6 @@ fn map_deduplicate_keys( needs_value_take = true; value_indices[output_idx] = abs_value_idx; keys_mask_builder.append_value(false); - values_mask_builder.append_value(false); continue; } return exec_err!( @@ -246,39 +247,32 @@ fn map_deduplicate_keys( ); } keys_mask_builder.append_value(true); - values_mask_builder.append_value(true); key_to_output_idx.insert(key, value_indices.len()); value_indices.push(abs_value_idx); new_last_offset += 1; } } else { - // The result entry is NULL — no keys/values emitted. Still pad the - // masks used by filter so they stay aligned with their flat arrays. + // The result entry is NULL — no keys/values emitted. Keep the key + // mask aligned, but do not allocate a mask for ignored values. keys_mask_builder.append_n(num_keys_entries, false); - if !needs_value_take { - values_mask_builder.append_n(num_values_entries, false); - } + needs_value_take |= num_values_entries != 0; } new_offsets.push(new_last_offset); cur_keys_offset += num_keys_entries; cur_values_offset += num_values_entries; } let keys_mask = keys_mask_builder.finish(); - let values_mask = values_mask_builder.finish(); let needed_keys = filter(&flat_keys, &keys_mask)?; - let needed_values = if needs_value_take { + let needed_values = if value_indices.is_empty() { + // An empty slice could retain unused nested child buffers. + new_empty_array(flat_values.data_type()) + } else if needs_value_take { let value_indices_array = Int32Array::from(value_indices); take(&flat_values, &value_indices_array, None)? } else { - // Values lists can be sliced independently of keys lists. Align the - // relative mask with their first offset, without allocating a slice - // wrapper for the common case where values start at zero. - let flat_values = if values_start_offset == 0 { - Cow::Borrowed(flat_values) - } else { - Cow::Owned(flat_values.slice(values_start_offset, values_mask.len())) - }; - filter(flat_values.as_ref(), &values_mask)? + // No values were skipped or overwritten, so the selected range is + // contiguous. Its first offset can differ from the keys' first offset. + flat_values.slice(values_start_offset, value_indices.len()) }; let offsets = OffsetBuffer::new(new_offsets.into()); Ok((needed_keys, needed_values, offsets)) @@ -287,7 +281,7 @@ fn map_deduplicate_keys( #[cfg(test)] mod tests { use super::*; - use arrow::array::{Int32Array, LargeListArray, NullArray, StringArray}; + use arrow::array::{Int32Array, LargeListArray, ListArray, NullArray, StringArray}; fn int32_utf8_inputs( keys: Vec, @@ -601,6 +595,215 @@ mod tests { } } + #[test] + fn null_rows_with_large_values_spans_are_skipped() { + // A NullArray has no payload buffer, even when a skipped row spans + // many values. Selection should only allocate for retained entries. + let skipped_len = 1 << 20; + for null_row in 0..3 { + let nulls = + NullBuffer::from((0..3).map(|row| row != null_row).collect::>()); + let mut keys_offsets = vec![0]; + let mut values_offsets = vec![0]; + for row in 0..3 { + keys_offsets.push(keys_offsets[row] + i32::from(row != null_row)); + values_offsets.push( + values_offsets[row] + if row == null_row { skipped_len } else { 1 }, + ); + } + + for null_keys in [false, true] { + let keys: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new("item", DataType::Int32, false)), + OffsetBuffer::new(keys_offsets.clone().into()), + Arc::new(Int32Array::from(vec![7, 8])), + null_keys.then(|| nulls.clone()), + )); + let values: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new("item", DataType::Null, true)), + OffsetBuffer::new(values_offsets.clone().into()), + Arc::new(NullArray::new(skipped_len as usize + 2)), + (!null_keys).then(|| nulls.clone()), + )); + keys.to_data().validate_full().unwrap(); + values.to_data().validate_full().unwrap(); + + for last_value_wins in [false, true] { + let result = map_from_keys_values_offsets_nulls( + get_list_values(&keys).unwrap(), + get_list_values(&values).unwrap(), + &get_list_offsets(&keys).unwrap(), + &get_list_offsets(&values).unwrap(), + keys.nulls(), + values.nulls(), + last_value_wins, + ) + .unwrap(); + + result.to_data().validate_full().unwrap(); + let map = result.as_map(); + assert_eq!(map.value_offsets(), keys_offsets.as_slice()); + assert_eq!(map.nulls(), Some(&nulls)); + assert_eq!( + map.keys().to_data(), + Int32Array::from(vec![7, 8]).to_data(), + ); + assert_eq!(map.values().to_data(), NullArray::new(2).to_data()); + } + } + } + } + + #[test] + fn partial_selection_does_not_materialize_null_list_children() { + let child_len = 1 << 16; + let list_field = Arc::new(Field::new("item", DataType::Null, true)); + let child_values: ArrayRef = Arc::new(NullArray::new(child_len + 1)); + let inner_nulls = NullBuffer::from(vec![false, true]); + let inner_lists: Vec = vec![ + Arc::new(ListArray::new( + Arc::clone(&list_field), + OffsetBuffer::new(vec![0, child_len as i32, child_len as i32 + 1].into()), + Arc::clone(&child_values), + Some(inner_nulls.clone()), + )), + Arc::new(LargeListArray::new( + list_field, + OffsetBuffer::new(vec![0, child_len as i64, child_len as i64 + 1].into()), + Arc::clone(&child_values), + Some(inner_nulls), + )), + ]; + let keys: ArrayRef = Arc::new(Int32Array::from(vec![7, 8])); + let offsets = [0, 1, 2]; + let outer_nulls = NullBuffer::from(vec![true, false]); + + for values in inner_lists { + values.to_data().validate_full().unwrap(); + for last_value_wins in [false, true] { + for (keys_nulls, values_nulls) in + [(Some(&outer_nulls), None), (None, Some(&outer_nulls))] + { + let result = map_from_keys_values_offsets_nulls( + &keys, + &values, + &offsets, + &offsets, + keys_nulls, + values_nulls, + last_value_wins, + ) + .unwrap(); + + result.to_data().validate_full().unwrap(); + let map = result.as_map(); + assert_eq!(map.value_offsets(), &[0, 1, 1]); + assert_eq!(map.nulls(), Some(&outer_nulls)); + assert_eq!(map.keys().to_data(), Int32Array::from(vec![7]).to_data()); + assert!(map.values().is_null(0)); + assert_eq!(get_list_values(map.values()).unwrap().len(), 0); + } + + // Retaining every value should still share the nested child, + // even when a NULL list has a nonempty child span. + let result = map_from_keys_values_offsets_nulls( + &keys, + &values, + &offsets, + &offsets, + None, + None, + last_value_wins, + ) + .unwrap(); + result.to_data().validate_full().unwrap(); + assert_eq!(result.as_map().values().to_data(), values.to_data()); + assert!(Arc::ptr_eq( + get_list_values(result.as_map().values()).unwrap(), + &child_values, + )); + } + } + } + + #[test] + fn empty_null_rows_preserve_value_buffer_sharing() { + let (keys, values) = int32_utf8_inputs( + vec![1, 2], + vec![Some("prefix"), Some("a"), Some("b"), Some("suffix")], + ); + let source_values = values.as_any().downcast_ref::().unwrap(); + let nulls = NullBuffer::from(vec![true, false, true]); + for last_value_wins in [false, true] { + for (keys_nulls, values_nulls) in [(Some(&nulls), None), (None, Some(&nulls))] + { + let result = map_from_keys_values_offsets_nulls( + &keys, + &values, + &[0, 1, 1, 2], + &[1, 2, 2, 3], + keys_nulls, + values_nulls, + last_value_wins, + ) + .unwrap(); + result.to_data().validate_full().unwrap(); + let map = result.as_map(); + let needed_values = + map.values().as_any().downcast_ref::().unwrap(); + assert_eq!(map.value_offsets(), &[0, 1, 1, 2]); + assert_eq!(map.nulls(), Some(&nulls)); + assert_eq!(needed_values, &StringArray::from(vec!["a", "b"])); + assert_eq!( + needed_values.value_data().as_ptr(), + source_values.value_data().as_ptr(), + ); + } + } + } + + #[test] + fn empty_maps_do_not_retain_nested_values() { + let keys: ArrayRef = Arc::new(Int32Array::from(Vec::::new())); + let values: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new("item", DataType::Null, true)), + OffsetBuffer::new(vec![0, 1 << 16].into()), + Arc::new(NullArray::new(1 << 16)), + None, + )); + values.to_data().validate_full().unwrap(); + let nulls = NullBuffer::from(vec![false]); + + for last_value_wins in [false, true] { + // Cover an empty batch, an empty map, and a NULL map that skips + // either an empty or a nonempty values row. + for (keys_offsets, values_offsets, keys_nulls) in [ + (&[0][..], &[0][..], None), + (&[0, 0][..], &[0, 0][..], None), + (&[0, 0][..], &[1, 1][..], Some(&nulls)), + (&[0, 0][..], &[0, 1][..], Some(&nulls)), + ] { + let result = map_from_keys_values_offsets_nulls( + &keys, + &values, + keys_offsets, + values_offsets, + keys_nulls, + None, + last_value_wins, + ) + .unwrap(); + + result.to_data().validate_full().unwrap(); + let map = result.as_map(); + assert_eq!(map.value_offsets(), keys_offsets); + assert_eq!(map.nulls(), keys_nulls); + assert_eq!(map.values().len(), 0); + assert_eq!(get_list_values(map.values()).unwrap().len(), 0); + } + } + } + #[test] fn single_row_duplicate_errors_under_exception() { let (keys, values) =