From 1cca82ceb99f4f21e931901d004e581db377cfcc Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Mon, 5 Apr 2021 17:26:43 -0400 Subject: [PATCH 1/2] ARROW-12214: [Rust][DataFusion] Add tests for limit --- rust/datafusion/src/execution/context.rs | 54 ++++++++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/rust/datafusion/src/execution/context.rs b/rust/datafusion/src/execution/context.rs index a22b5d3c1e7..6b94d470e05 100644 --- a/rust/datafusion/src/execution/context.rs +++ b/rust/datafusion/src/execution/context.rs @@ -1721,6 +1721,60 @@ mod tests { Ok(()) } + #[tokio::test] + async fn limit() -> Result<()> { + let tmp_dir = TempDir::new()?; + let mut ctx = create_ctx(&tmp_dir, 1)?; + ctx.register_table("t", table_with_sequence(1, 1000).unwrap()) + .unwrap(); + + let results = plan_and_collect(&mut ctx, "SELECT i FROM t ORDER BY i DESC limit 3") + .await + .unwrap(); + + let expected = vec! + [ + "+------+", + "| i |", + "+------+", + "| 1000 |", + "| 999 |", + "| 998 |", + "+------+", + ]; + + assert_batches_eq!(expected, &results); + + + let results = plan_and_collect(&mut ctx, "SELECT i FROM t ORDER BY i limit 3") + .await + .unwrap(); + + let expected = vec! + [ + "+---+", + "| i |", + "+---+", + "| 1 |", + "| 2 |", + "| 3 |", + "+---+", + ]; + + assert_batches_eq!(expected, &results); + + let results = plan_and_collect(&mut ctx, "SELECT i FROM t limit 3") + .await + .unwrap(); + + // the actual rows are not guaranteed, so only check the count (should be 3) + let num_rows: usize = results.into_iter().map(|b| b.num_rows()).sum(); + assert_eq!(num_rows, 3); + + Ok(()) + } + + #[tokio::test] async fn case_sensitive_identifiers_functions() { let mut ctx = ExecutionContext::new(); From ba55fcb8c8e9b5e91db9696520cbac97181fccd7 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Mon, 5 Apr 2021 17:56:05 -0400 Subject: [PATCH 2/2] fix fmt --- rust/datafusion/src/execution/context.rs | 36 ++++++++---------------- 1 file changed, 11 insertions(+), 25 deletions(-) diff --git a/rust/datafusion/src/execution/context.rs b/rust/datafusion/src/execution/context.rs index 6b94d470e05..50b3a7a1add 100644 --- a/rust/datafusion/src/execution/context.rs +++ b/rust/datafusion/src/execution/context.rs @@ -1728,38 +1728,25 @@ mod tests { ctx.register_table("t", table_with_sequence(1, 1000).unwrap()) .unwrap(); - let results = plan_and_collect(&mut ctx, "SELECT i FROM t ORDER BY i DESC limit 3") - .await - .unwrap(); + let results = + plan_and_collect(&mut ctx, "SELECT i FROM t ORDER BY i DESC limit 3") + .await + .unwrap(); - let expected = vec! - [ - "+------+", - "| i |", - "+------+", - "| 1000 |", - "| 999 |", - "| 998 |", - "+------+", - ]; + let expected = vec![ + "+------+", "| i |", "+------+", "| 1000 |", "| 999 |", "| 998 |", + "+------+", + ]; assert_batches_eq!(expected, &results); - let results = plan_and_collect(&mut ctx, "SELECT i FROM t ORDER BY i limit 3") .await .unwrap(); - let expected = vec! - [ - "+---+", - "| i |", - "+---+", - "| 1 |", - "| 2 |", - "| 3 |", - "+---+", - ]; + let expected = vec![ + "+---+", "| i |", "+---+", "| 1 |", "| 2 |", "| 3 |", "+---+", + ]; assert_batches_eq!(expected, &results); @@ -1774,7 +1761,6 @@ mod tests { Ok(()) } - #[tokio::test] async fn case_sensitive_identifiers_functions() { let mut ctx = ExecutionContext::new();