Re: [PR] Propagate .execute() calls immediately in `RepartitionExec` [datafusion]

["gabotechs (via GitHub)" <gi...@apache.org>]

gabotechs commented on code in PR #16093:
URL: https://github.com/apache/datafusion/pull/16093#discussion_r2095320415


##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -1298,15 +1347,9 @@ mod tests {
         let partitioning = Partitioning::RoundRobinBatch(1);
         let exec = RepartitionExec::try_new(Arc::new(input), partitioning).unwrap();
 
-        // Note: this should pass (the stream can be created) but the
-        // error when the input is executed should get passed back
-        let output_stream = exec.execute(0, task_ctx).unwrap();
-
         // Expect that an error is returned
-        let result_string = crate::common::collect(output_stream)
-            .await
-            .unwrap_err()
-            .to_string();
+        let result_string = exec.execute(0, task_ctx).err().unwrap().to_string();

Review Comment:
   The error is now propagated immediately, rather than lazily upon the first message poll



##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -1496,7 +1539,14 @@ mod tests {
         });
         let batches_with_drop = crate::common::collect(output_stream1).await.unwrap();
 
-        assert_eq!(batches_without_drop, batches_with_drop);
+        fn sort(batch: Vec<RecordBatch>) -> Vec<RecordBatch> {
+            batch
+                .into_iter()
+                .sorted_by_key(|b| format!("{b:?}"))
+                .collect()
+        }
+
+        assert_eq!(sort(batches_without_drop), sort(batches_with_drop));

Review Comment:
   Found this test to be flaky after my changes. Just ordering the record batches removes the flakiness. Not sure why this started being flaky now though...



-- 
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-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org