geoffreyclaude commented on code in PR #26603:
URL: https://github.com/apache/datafusion/pull/26603#discussion_r4412969813
##########
datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs:
##########
@@ -897,4 +915,211 @@ mod tests {
assert_eq!(result.fields()[0].name(), "a");
assert!(consumer.get_outer_schema(2).is_none());
}
+
+ // A custom consumer exercises the trait default independently of the
default
+ // consumer's scheduling policy.
+ struct InputConsumer {
+ extensions: Extensions,
+ state: SessionState,
+ reads: std::sync::Mutex<Vec<String>>,
+ ready: tokio::sync::Semaphore,
+ }
+
+ #[async_trait]
+ impl SubstraitConsumer for InputConsumer {
+ fn get_extensions(&self) -> &Extensions {
+ &self.extensions
+ }
+
+ fn get_function_registry(&self) -> &impl FunctionRegistry {
+ &self.state
+ }
+
+ async fn resolve_table_ref(
+ &self,
+ _: &TableReference,
+ ) -> datafusion::common::Result<Option<Arc<dyn TableProvider>>> {
+ unreachable!("test consumes reads directly")
+ }
+
+ async fn consume_rel(
+ &self,
+ rel: &Rel,
+ ) -> datafusion::common::Result<LogicalPlan> {
+ let Some(proto::rel::RelType::Read(read)) = &rel.rel_type else {
+ panic!("expected a read")
+ };
+ let Some(proto::read_rel::ReadType::NamedTable(table)) =
&read.read_type
+ else {
+ panic!("expected a named table")
+ };
+ self.reads.lock().unwrap().push(table.names[0].clone());
+ self.ready.acquire().await.unwrap().forget();
+ if table.names[0] == "fail" {
+ return substrait_err!("input failed");
+ }
+ let provider =
Review Comment:
You can probably simplify here with the `table_scan` test helper since you
don't need the provider
--
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]
