wiedld commented on code in PR #7798:
URL: https://github.com/apache/arrow-datafusion/pull/7798#discussion_r1360039344
##########
datafusion/physical-plan/src/sorts/cursor.rs:
##########
@@ -432,4 +501,119 @@ mod tests {
b.advance();
assert_eq!(a.cmp(&b), Ordering::Less);
}
+
+ #[test]
+ fn test_slice_primitive() {
+ let options = SortOptions {
+ descending: false,
+ nulls_first: true,
+ };
+
+ let buffer = ScalarBuffer::from(vec![0, 1, 2]);
+ let mut cursor = new_primitive(options, buffer, 0);
+
+ // from start
+ let sliced = cursor.slice(0, 1);
+ assert_eq!(sliced.num_rows(), 1);
+ let expected = new_primitive(options, ScalarBuffer::from(vec![0]), 0);
+ assert_eq!(
+ sliced.cmp(&expected),
+ Ordering::Equal,
+ "should slice from start"
+ );
+
+ // with offset
+ let sliced = cursor.slice(1, 2);
+ assert_eq!(sliced.num_rows(), 2);
+ let expected = new_primitive(options, ScalarBuffer::from(vec![1]), 0);
+ assert_eq!(
+ sliced.cmp(&expected),
+ Ordering::Equal,
+ "should slice with offset"
+ );
+
+ // cursor current position != start
+ cursor.advance();
+ let sliced = cursor.slice(0, 1);
+ assert_eq!(sliced.num_rows(), 1);
+ let expected = new_primitive(options, ScalarBuffer::from(vec![0]), 0);
+ assert_eq!(
+ sliced.cmp(&expected),
+ Ordering::Equal,
+ "should ignore current cursor position when sliced"
+ );
+ }
+
+ #[test]
+ #[should_panic(expected = "cursor slice out of bounds")]
+ fn test_slice_panic_can_panic() {
+ let options = SortOptions {
+ descending: false,
+ nulls_first: true,
+ };
+
+ let buffer = ScalarBuffer::from(vec![0, 1, 2]);
+ let cursor = new_primitive(options, buffer, 0);
+
+ cursor.slice(42, 1);
+ }
+
+ #[test]
+ fn test_slice_nulls_first() {
+ let options = SortOptions {
+ descending: false,
+ nulls_first: true,
+ };
+
+ let buffer = ScalarBuffer::from(vec![2, i32::MIN, 2]);
+ let mut a = new_primitive(options, buffer, 1);
+ let buffer = ScalarBuffer::from(vec![3, 2, i32::MIN, 2]);
+ let mut b = new_primitive(options, buffer, 1);
+
+ // NULL == NULL
+ assert_eq!(a, b);
+ assert_eq!(a.cmp(&b), Ordering::Equal);
+
+ // 2 > NULL
Review Comment:
Fixed the null mask to work properly. Explicit test cases pushed. Let me
know if it's correct this time.
--
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: github-unsubscr...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
