Copilot commented on code in PR #23628:
URL: https://github.com/apache/datafusion/pull/23628#discussion_r3592398249
##########
datafusion/functions-aggregate/src/first_last/state.rs:
##########
@@ -459,4 +534,358 @@ mod tests {
Ok(())
}
+
+ // ---------- GenericValueState (nested types) ----------
+
+ /// Build a `List<Utf8>` array with three rows: `["a"]`, `["b", "c"]`,
+ /// `["d", "e", "f"]`. Used by several tests.
+ fn make_list_utf8_array() -> ArrayRef {
+ let mut builder = ListBuilder::new(StringBuilder::new());
+ builder.values().append_value("a");
+ builder.append(true);
+ builder.values().append_value("b");
+ builder.values().append_value("c");
+ builder.append(true);
+ builder.values().append_value("d");
+ builder.values().append_value("e");
+ builder.values().append_value("f");
+ builder.append(true);
+ Arc::new(builder.finish())
+ }
+
+ #[test]
+ fn test_generic_value_state_list_utf8() -> Result<()> {
+ let list_utf8 =
+ DataType::List(Arc::new(Field::new("item", DataType::Utf8, true)));
+ let mut state = GenericValueState::new(list_utf8.clone());
+ state.resize(2);
+
+ let array = make_list_utf8_array();
+
+ // group 0 <- ["a"] ; group 1 <- ["b", "c"]
+ state.update(0, &array, 0)?;
+ state.update(1, &array, 1)?;
+
+ // Overwrite group 0 with the wider ["d", "e", "f"] (size-accounting
+ // must decrement the old value before adding the new one).
+ let size_after_first = state.total_size;
+ state.update(0, &array, 2)?;
+ assert!(
+ state.total_size > 0,
+ "total_size must remain positive after overwrite"
+ );
+ // The overwrite replaced group 0's payload; the delta relative to the
+ // previous state should equal `new.size() - old.size()`. If the caller
+ // forgot to subtract the old size, `total_size` would drift upward.
+ let expected_delta = {
+ let new_scalar = {
+ let mut s = ScalarValue::try_from_array(&array, 2)?;
+ s.compact();
+ s
+ };
+ let old_scalar = {
+ let mut s = ScalarValue::try_from_array(&array, 0)?;
+ s.compact();
+ s
+ };
+ new_scalar.size() as isize - old_scalar.size() as isize
+ };
+ assert_eq!(
+ state.total_size as isize - size_after_first as isize,
+ expected_delta,
+ "size accounting drifted after overwrite"
+ );
+
+ let result = state.take(EmitTo::All)?;
+ let result = result.as_list::<i32>();
+ assert_eq!(result.len(), 2);
+ let g0 = result.value(0);
+ let g0 = g0.as_any().downcast_ref::<StringArray>().unwrap();
+ assert_eq!(g0.value(0), "d");
+ assert_eq!(g0.value(1), "e");
+ assert_eq!(g0.value(2), "f");
+ let g1 = result.value(1);
+ let g1 = g1.as_any().downcast_ref::<StringArray>().unwrap();
+ assert_eq!(g1.value(0), "b");
+ assert_eq!(g1.value(1), "c");
+
+ assert_eq!(state.total_size, 0, "state must be fully drained");
+ Ok(())
+ }
+
+ #[test]
+ fn test_generic_value_state_struct() -> Result<()> {
+ let fields = Fields::from(vec![
+ Field::new("id", DataType::Int32, false),
+ Field::new("name", DataType::Utf8, false),
+ ]);
+ let struct_type = DataType::Struct(fields.clone());
+
+ let id = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
+ let name = Arc::new(StringArray::from(vec!["a", "b", "c"])) as
ArrayRef;
+ let struct_array =
+ Arc::new(StructArray::new(fields, vec![id, name], None)) as
ArrayRef;
+
+ let mut state = GenericValueState::new(struct_type);
+ state.resize(2);
+ state.update(0, &struct_array, 0)?;
+ state.update(1, &struct_array, 2)?;
+
+ let out = state.take(EmitTo::All)?;
+ let out = out.as_any().downcast_ref::<StructArray>().unwrap();
+ assert_eq!(out.len(), 2);
+
+ let out_id =
out.column(0).as_any().downcast_ref::<Int32Array>().unwrap();
+ let out_name = out
+ .column(1)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ assert_eq!(out_id.value(0), 1);
+ assert_eq!(out_id.value(1), 3);
+ assert_eq!(out_name.value(0), "a");
+ assert_eq!(out_name.value(1), "c");
+ Ok(())
+ }
+
+ #[test]
+ fn test_generic_value_state_large_list() -> Result<()> {
+ let field = Arc::new(Field::new("item", DataType::Int32, true));
+ let large_list_type = DataType::LargeList(Arc::clone(&field));
+
+ let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
+ let offsets: OffsetBuffer<i64> =
+ OffsetBuffer::new(ScalarBuffer::from(vec![0_i64, 2, 5, 6]));
+ let array = Arc::new(LargeListArray::new(field, offsets,
Arc::new(values), None))
+ as ArrayRef;
+
+ let mut state = GenericValueState::new(large_list_type);
+ state.resize(2);
+ state.update(0, &array, 0)?; // [1, 2]
+ state.update(1, &array, 2)?; // [6]
+
+ let out = state.take(EmitTo::All)?;
+ let out = out.as_list::<i64>();
+ assert_eq!(out.len(), 2);
+ let g0 = out.value(0);
+ let g0 = g0.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(g0.len(), 2);
+ assert_eq!(g0.value(0), 1);
+ assert_eq!(g0.value(1), 2);
+ let g1 = out.value(1);
+ let g1 = g1.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(g1.len(), 1);
+ assert_eq!(g1.value(0), 6);
+ Ok(())
+ }
+
+ #[test]
+ fn test_generic_value_state_fixed_size_list() -> Result<()> {
+ let field = Arc::new(Field::new("item", DataType::Int32, true));
+ let fsl_type = DataType::FixedSizeList(Arc::clone(&field), 2);
+
+ let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
+ let array = Arc::new(FixedSizeListArray::new(field, 2,
Arc::new(values), None))
+ as ArrayRef;
+
+ let mut state = GenericValueState::new(fsl_type);
+ state.resize(2);
+ state.update(0, &array, 0)?; // [1, 2]
+ state.update(1, &array, 2)?; // [5, 6]
+
+ let out = state.take(EmitTo::All)?;
+ let out = out
+ .as_any()
+ .downcast_ref::<FixedSizeListArray>()
+ .expect("emitted FixedSizeListArray");
+ assert_eq!(out.len(), 2);
+ let g0 = out.value(0);
+ let g0 = g0.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(g0.value(0), 1);
+ assert_eq!(g0.value(1), 2);
+ let g1 = out.value(1);
+ let g1 = g1.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(g1.value(0), 5);
+ assert_eq!(g1.value(1), 6);
+ Ok(())
+ }
+
+ #[test]
+ fn test_generic_value_state_map() -> Result<()> {
+ // Map<Utf8, Int32> with two entries: {"a": 1, "b": 2}, {"c": 3}
+ let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as
ArrayRef;
+ let values = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
+ let entry_fields = Fields::from(vec![
+ Field::new("keys", DataType::Utf8, false),
+ Field::new("values", DataType::Int32, true),
+ ]);
+ let entries = StructArray::new(entry_fields.clone(), vec![keys,
values], None);
+ let offsets: OffsetBuffer<i32> =
+ OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 2, 3]));
+ let map_field =
+ Arc::new(Field::new("entries", DataType::Struct(entry_fields),
false));
+ let map_array = Arc::new(MapArray::new(
+ Arc::clone(&map_field),
+ offsets,
+ entries,
+ None,
+ false,
+ )) as ArrayRef;
+ let map_type = DataType::Map(map_field, false);
+
+ let mut state = GenericValueState::new(map_type);
+ state.resize(2);
+ state.update(0, &map_array, 0)?; // {"a": 1, "b": 2}
+ state.update(1, &map_array, 1)?; // {"c": 3}
+
+ let out = state.take(EmitTo::All)?;
+ let out = out.as_any().downcast_ref::<MapArray>().unwrap();
+ assert_eq!(out.len(), 2);
+ assert_eq!(out.value_length(0), 2);
+ assert_eq!(out.value_length(1), 1);
+ Ok(())
+ }
+
+ #[test]
+ fn test_generic_value_state_emit_first() -> Result<()> {
+ let list_utf8 =
+ DataType::List(Arc::new(Field::new("item", DataType::Utf8, true)));
+ let mut state = GenericValueState::new(list_utf8);
+ state.resize(3);
+
+ let array = make_list_utf8_array();
+ state.update(0, &array, 0)?;
+ state.update(1, &array, 1)?;
+ state.update(2, &array, 2)?;
+
+ let after_all_updates = state.total_size;
+ assert!(after_all_updates > 0);
+
+ // Emit the first 2 groups; remaining group 2 stays.
+ let head = state.take(EmitTo::First(2))?;
+ let head = head.as_list::<i32>();
+ assert_eq!(head.len(), 2);
+ let h0 = head.value(0);
+ let h0 = h0.as_any().downcast_ref::<StringArray>().unwrap();
+ assert_eq!(h0.value(0), "a");
+ let h1 = head.value(1);
+ let h1 = h1.as_any().downcast_ref::<StringArray>().unwrap();
+ assert_eq!(h1.value(0), "b");
+ assert_eq!(h1.value(1), "c");
+
+ // After partial emit, `total_size` shrank but is still positive.
+ assert!(state.total_size > 0);
+ assert!(state.total_size < after_all_updates);
+
+ let tail = state.take(EmitTo::All)?;
+ let tail = tail.as_list::<i32>();
+ assert_eq!(tail.len(), 1);
+ let t0 = tail.value(0);
+ let t0 = t0.as_any().downcast_ref::<StringArray>().unwrap();
+ assert_eq!(t0.value(0), "d");
+ assert_eq!(t0.value(1), "e");
+ assert_eq!(t0.value(2), "f");
+
+ assert_eq!(state.total_size, 0);
+ Ok(())
+ }
+
+ #[test]
+ fn test_generic_value_state_update_null() -> Result<()> {
+ // List<Int32> with rows: [1, 2], NULL
+ let mut builder = ListBuilder::new(Int32Builder::new());
+ builder.values().append_value(1);
+ builder.values().append_value(2);
+ builder.append(true);
+ builder.append(false); // null entry
+ let array: ArrayRef = Arc::new(builder.finish());
+
+ let list_type = array.data_type().clone();
+ let mut state = GenericValueState::new(list_type);
+ state.resize(1);
+
+ // group 0 = [1, 2]
+ state.update(0, &array, 0)?;
+ let size_after_value = state.total_size;
+ assert!(size_after_value > 0);
+
+ // Overwrite group 0 with NULL. The size accounting must subtract the
+ // previous value's size and then add the null-scalar's size; the point
+ // of this test is that `total_size` stays consistent (no drift) and
+ // the null is emitted correctly.
+ state.update(0, &array, 1)?;
+ // Recomputing from scratch must match the cached total_size.
+ let recomputed: usize = state.vals.iter().flatten().map(|v|
v.size()).sum();
+ assert_eq!(
+ state.total_size, recomputed,
+ "total_size drifted after null update"
+ );
+
+ let out = state.take(EmitTo::All)?;
+ let out = out.as_list::<i32>();
+ assert_eq!(out.len(), 1);
+ assert!(out.is_null(0));
+ assert_eq!(state.total_size, 0);
+ Ok(())
+ }
+
+ #[test]
+ fn test_generic_value_state_compact_releases_parent_batch() -> Result<()> {
+ // Regression test for the memory-pinning bug: without compact(),
+ // `ScalarValue::try_from_array` on a List column returns an Arc
+ // slice that keeps the whole parent batch alive. compact() must
+ // copy the referenced bytes into an owned buffer so the parent
+ // array can be dropped.
+ let list_utf8 =
+ DataType::List(Arc::new(Field::new("item", DataType::Utf8, true)));
+ let mut state = GenericValueState::new(list_utf8);
+ state.resize(1);
+
+ let array: ArrayRef = make_list_utf8_array();
+ let weak_before = Arc::downgrade(&array);
+ assert_eq!(weak_before.strong_count(), 1);
+
+ state.update(0, &array, 0)?;
+ drop(array);
+
+ // If compact() did its job, dropping the source array must release
+ // the only strong reference. If the stored ScalarValue kept a slice
+ // of the parent, strong_count would stay >= 1 and the batch would
+ // remain pinned.
+ assert_eq!(
+ weak_before.strong_count(),
+ 0,
+ "compact() failed to break the Arc reference to the source batch"
+ );
Review Comment:
The regression test `test_generic_value_state_compact_releases_parent_batch`
currently checks `Arc` strong_count on the *array object*
(`Arc::downgrade(&array)`), but `ScalarValue::try_from_array` for `List`/nested
types can retain the *underlying buffers* via sliced child arrays without
keeping the original `ArrayRef` alive. That means this test can pass even if
`compact()` does not break buffer sharing, so it doesn’t actually validate the
memory-pinning behavior it claims to cover. Consider asserting that the stored
ScalarValue no longer shares the original Utf8 value-data buffer (e.g., compare
`to_data().buffers()[1].as_ptr()` before/after).
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
