ethan-tyler commented on code in PR #19633:
URL: https://github.com/apache/datafusion/pull/19633#discussion_r2662853240
##########
datafusion/core/tests/custom_sources_cases/dml_planning.rs:
##########
@@ -165,6 +165,66 @@ impl TableProvider for CaptureUpdateProvider {
}
}
+/// A TableProvider that captures whether truncate() was called.
+struct CaptureTruncateProvider {
+ schema: SchemaRef,
+ truncate_called: Arc<Mutex<bool>>,
+}
+
+impl CaptureTruncateProvider {
+ fn new(schema: SchemaRef) -> Self {
+ Self {
+ schema,
+ truncate_called: Arc::new(Mutex::new(false)),
+ }
+ }
+
+ fn was_truncated(&self) -> bool {
+ *self.truncate_called.lock().unwrap()
+ }
+}
+
+impl std::fmt::Debug for CaptureTruncateProvider {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ f.debug_struct("CaptureTruncateProvider")
+ .field("schema", &self.schema)
+ .finish()
+ }
+}
+
+#[async_trait]
+impl TableProvider for CaptureTruncateProvider {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+
+ fn table_type(&self) -> TableType {
+ TableType::Base
+ }
+
+ async fn scan(
+ &self,
+ _state: &dyn Session,
+ _projection: Option<&Vec<usize>>,
+ _filters: &[Expr],
+ _limit: Option<usize>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(Arc::new(EmptyExec::new(Arc::clone(&self.schema))))
+ }
+
+ async fn truncate(&self, _state: &dyn Session) -> Result<Arc<dyn
ExecutionPlan>> {
+ *self.truncate_called.lock().unwrap() = true;
+
+ Ok(Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![
+ Field::new("count", DataType::UInt64, false),
+ ])))))
+ }
+}
Review Comment:
TableProvider::truncate docs say it returns an ExecutionPlan producing
{count: UInt64}. This test returns EmptyExec which produces zero rows, it
validates the hook was called, but not the contract.
Could tighten this to return an actual row and assert on it:
```suggestion
// Add these imports:
use arrow::array::UInt64Array;
use arrow::record_batch::RecordBatch;
use datafusion_physical_plan::test::TestMemoryExec;
// Then update the truncate impl:
async fn truncate(&self, _state: &dyn Session) -> Result<Arc<dyn
ExecutionPlan>> {
*self.truncate_called.lock().unwrap() = true;
let schema = Arc::new(Schema::new(vec![
Field::new("count", DataType::UInt64, false),
]));
let batch = RecordBatch::try_new(
Arc::clone(&schema),
vec![Arc::new(UInt64Array::from(vec![0u64]))],
)?;
Ok(Arc::new(TestMemoryExec::try_new(&[vec![batch]], schema, None)?))
}
```
--
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]
