martin-g commented on code in PR #25040:
URL: https://github.com/apache/datafusion/pull/25040#discussion_r4002829153
##########
datafusion/catalog/src/memory/table.rs:
##########
@@ -701,28 +813,139 @@ impl ExecutionPlan for DmlResultExec {
fn execute(
&self,
- _partition: usize,
- _context: Arc<datafusion_execution::TaskContext>,
- ) -> Result<datafusion_execution::SendableRecordBatchStream> {
- // Create a single batch with the count
- let count_array = UInt64Array::from(vec![self.rows_affected]);
- let batch = ArrowRecordBatch::try_new(
- Arc::clone(&self.schema),
- vec![Arc::new(count_array) as ArrayRef],
- )?;
+ partition: usize,
+ _context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ if partition != 0 {
+ return internal_err!(
+ "MemUpdateExec has one partition, but partition {partition}
was requested"
+ );
+ }
- // Create a stream that yields just this one batch
- let stream = futures::stream::iter(vec![Ok(batch)]);
- Ok(Box::pin(RecordBatchStreamAdapter::new(
- Arc::clone(&self.schema),
- stream,
- )))
+ let partitions = self.partitions.clone();
+ let sort_order = Arc::clone(&self.sort_order);
+ let table_schema = Arc::clone(&self.table_schema);
+ let set_exprs = self.set_exprs.clone();
+ let predicates = self.predicates.clone();
+ let schema = Arc::clone(&self.schema);
+ let count_schema = Arc::clone(&self.schema);
Review Comment:
I think there is no need of two clones of the schema here.
##########
datafusion/catalog/src/memory/table.rs:
##########
@@ -356,299 +362,401 @@ impl MemTable {
Ok(Arc::new(DataSinkExec::new(input, Arc::new(sink), None)))
}
- fn delete_from_boxed<'a>(
- &'a self,
- state: &'a dyn Session,
- filters: Vec<Expr>,
- ) -> BoxFuture<'a, Result<Arc<dyn ExecutionPlan>>> {
- Box::pin(self.delete_from_inner(state, filters))
- }
-
- async fn delete_from_inner(
+ /// Plan a `DELETE`. The rows change when the returned plan runs, not here,
+ /// so `EXPLAIN DELETE` prints the plan and the table keeps its rows.
+ fn delete_from_inner(
&self,
state: &dyn Session,
- filters: Vec<Expr>,
+ filters: &[Expr],
) -> Result<Arc<dyn ExecutionPlan>> {
- // Early exit if table has no partitions
- if self.batches.is_empty() {
- return Ok(Arc::new(DmlResultExec::new(0)));
- }
-
- *self.sort_order.lock() = vec![];
-
- let mut total_deleted: u64 = 0;
let df_schema = DFSchema::try_from(Arc::clone(&self.schema))?;
+ let predicates = create_predicates(filters, &df_schema, state)?;
- for partition_data in &self.batches {
- let mut partition = partition_data.write().await;
- let mut new_batches = Vec::with_capacity(partition.len());
-
- for batch in partition.iter() {
- if batch.num_rows() == 0 {
- continue;
- }
-
- // Evaluate filters - None means "match all rows"
- let filter_mask = evaluate_filters_to_mask(
- &filters,
- batch,
- &df_schema,
- state.execution_props(),
- )?;
-
- let (delete_count, keep_mask) = match filter_mask {
- Some(mask) => {
- // Count rows where mask is true (will be deleted)
- let count = mask.iter().filter(|v| v ==
&Some(true)).count();
- // Keep rows where predicate is false or NULL (SQL
three-valued logic)
- let keep: BooleanArray =
- mask.iter().map(|v| Some(v !=
Some(true))).collect();
- (count, keep)
- }
- None => {
- // No filters = delete all rows
- (
- batch.num_rows(),
- BooleanArray::from(vec![false; batch.num_rows()]),
- )
- }
- };
-
- total_deleted += delete_count as u64;
-
- let filtered_batch = filter_record_batch(batch, &keep_mask)?;
- if filtered_batch.num_rows() > 0 {
- new_batches.push(filtered_batch);
- }
- }
-
- *partition = new_batches;
- }
-
- Ok(Arc::new(DmlResultExec::new(total_deleted)))
- }
-
- fn update_boxed<'a>(
- &'a self,
- state: &'a dyn Session,
- assignments: Vec<(String, Expr)>,
- filters: Vec<Expr>,
- ) -> BoxFuture<'a, Result<Arc<dyn ExecutionPlan>>> {
- Box::pin(self.update_inner(state, assignments, filters))
+ Ok(Arc::new(MemDeleteExec::new(
+ self.batches.clone(),
+ Arc::clone(&self.sort_order),
+ predicates,
+ )))
}
- async fn update_inner(
+ /// Plan an `UPDATE`. The rows change when the returned plan runs, not
here,
+ /// so `EXPLAIN UPDATE` prints the plan and the table keeps its rows.
+ fn update_inner(
&self,
state: &dyn Session,
- assignments: Vec<(String, Expr)>,
- filters: Vec<Expr>,
+ assignments: &[(String, Expr)],
+ filters: &[Expr],
) -> Result<Arc<dyn ExecutionPlan>> {
- // Early exit if table has no partitions
- if self.batches.is_empty() {
- return Ok(Arc::new(DmlResultExec::new(0)));
- }
+ let df_schema = DFSchema::try_from(Arc::clone(&self.schema))?;
- // Validate column names upfront with clear error messages
- let available_columns: Vec<&str> = self
- .schema
- .fields()
- .iter()
- .map(|f| f.name().as_str())
- .collect();
- for (column_name, _) in &assignments {
- if self.schema.field_with_name(column_name).is_err() {
+ // One entry for each field of the table, in field order. A `Some`
entry
+ // holds the expression of the `SET` clause for that field.
+ let mut set_exprs: Vec<Option<Arc<dyn PhysicalExpr>>> =
+ vec![None; self.schema.fields().len()];
+ for (column_name, expr) in assignments {
+ let Ok(index) = self.schema.index_of(column_name) else {
+ let available_columns: Vec<&str> = self
+ .schema
+ .fields()
+ .iter()
+ .map(|f| f.name().as_str())
+ .collect();
return plan_err!(
"UPDATE failed: column '{}' does not exist. Available
columns: {}",
column_name,
available_columns.join(", ")
);
- }
+ };
+ set_exprs[index] = Some(create_physical_expr(
+ expr,
+ &df_schema,
+ state.execution_props(),
+ &PhysicalPlanningContext::default(),
+ )?);
}
- let df_schema = DFSchema::try_from(Arc::clone(&self.schema))?;
-
- // Create physical expressions for assignments upfront (outside batch
loop)
- let physical_assignments: HashMap<String, Arc<dyn PhysicalExpr>> =
assignments
- .iter()
- .map(|(name, expr)| {
- let physical_expr = create_physical_expr(
- expr,
- &df_schema,
- state.execution_props(),
- &PhysicalPlanningContext::default(),
- )?;
- Ok((name.clone(), physical_expr))
- })
- .collect::<Result<_>>()?;
+ let predicates = create_predicates(filters, &df_schema, state)?;
- *self.sort_order.lock() = vec![];
+ Ok(Arc::new(MemUpdateExec::new(
+ self.batches.clone(),
+ Arc::clone(&self.sort_order),
+ Arc::clone(&self.schema),
+ set_exprs,
+ predicates,
+ )))
+ }
+}
- let mut total_updated: u64 = 0;
+/// Build one physical predicate for each expression of the `WHERE` clause.
+///
+/// The planner calls this, so a predicate that cannot be planned raises its
+/// error while the plan is built and `EXPLAIN` reports it.
+fn create_predicates(
+ filters: &[Expr],
+ df_schema: &DFSchema,
+ state: &dyn Session,
+) -> Result<Vec<Arc<dyn PhysicalExpr>>> {
+ filters
+ .iter()
+ .map(|filter| {
+ create_physical_expr(
+ filter,
+ df_schema,
+ state.execution_props(),
+ &PhysicalPlanningContext::default(),
+ )
+ })
+ .collect()
+}
- for partition_data in &self.batches {
- let mut partition = partition_data.write().await;
- let mut new_batches = Vec::with_capacity(partition.len());
+/// Combine the predicates into one mask over the rows of `batch`. The mask is
+/// true for a row that every predicate matches. `None` means there is no
+/// `WHERE` clause, which matches every row.
+fn evaluate_predicates(
+ predicates: &[Arc<dyn PhysicalExpr>],
+ batch: &RecordBatch,
+) -> Result<Option<BooleanArray>> {
+ let mut combined_mask: Option<BooleanArray> = None;
- for batch in partition.iter() {
- if batch.num_rows() == 0 {
- continue;
- }
+ for predicate in predicates {
+ let result = predicate.evaluate(batch)?;
+ let array = result.into_array(batch.num_rows())?;
+ let bool_array = array
+ .as_any()
+ .downcast_ref::<BooleanArray>()
+ .ok_or_else(|| {
+ internal_datafusion_err!("Filter did not evaluate to boolean")
+ })?
+ .clone();
- // Evaluate filters - None means "match all rows"
- let filter_mask = evaluate_filters_to_mask(
- &filters,
- batch,
- &df_schema,
- state.execution_props(),
- )?;
+ combined_mask = Some(match combined_mask {
+ Some(existing) => and(&existing, &bool_array)?,
+ None => bool_array,
+ });
+ }
- let (update_count, update_mask) = match filter_mask {
- Some(mask) => {
- // Count rows where mask is true (will be updated)
- let count = mask.iter().filter(|v| v ==
&Some(true)).count();
- // Normalize mask: only true (not NULL) triggers update
- let normalized: BooleanArray =
- mask.iter().map(|v| Some(v ==
Some(true))).collect();
- (count, normalized)
- }
- None => {
- // No filters = update all rows
- (
- batch.num_rows(),
- BooleanArray::from(vec![true; batch.num_rows()]),
- )
- }
- };
+ Ok(combined_mask)
+}
- total_updated += update_count as u64;
+/// Schema of the single `count` column that a DML plan emits.
+fn dml_count_schema() -> SchemaRef {
+ Arc::new(Schema::new(vec![Field::new(
+ "count",
+ DataType::UInt64,
+ false,
+ )]))
+}
- if update_count == 0 {
- new_batches.push(batch.clone());
- continue;
- }
+/// Properties of a DML plan: one partition, one final batch.
+fn dml_plan_properties(schema: &SchemaRef) -> Arc<PlanProperties> {
+ Arc::new(PlanProperties::new(
+ EquivalenceProperties::new(Arc::clone(schema)),
+ Partitioning::UnknownPartitioning(1),
+ EmissionType::Final,
+ Boundedness::Bounded,
+ ))
+}
- let mut new_columns: Vec<ArrayRef> =
- Vec::with_capacity(batch.num_columns());
-
- for field in self.schema.fields() {
- let column_name = field.name();
- let original_column =
- batch.column_by_name(column_name).ok_or_else(|| {
-
datafusion_common::DataFusionError::Internal(format!(
- "Column '{column_name}' not found in batch"
- ))
- })?;
-
- let new_column = if let Some(physical_expr) =
- physical_assignments.get(column_name.as_str())
- {
- // Use evaluate_selection to only evaluate on matching
rows.
- // This avoids errors (e.g., divide-by-zero) on rows
that won't
- // be updated. The result is scattered back with nulls
for
- // non-matching rows, which zip() will replace with
originals.
- let new_values =
- physical_expr.evaluate_selection(batch,
&update_mask)?;
- let new_array =
new_values.into_array(batch.num_rows())?;
+/// A single row holding the number of rows the statement changed.
+fn count_batch(schema: &SchemaRef, rows_affected: u64) -> Result<RecordBatch> {
+ let count_array = UInt64Array::from(vec![rows_affected]);
+ Ok(RecordBatch::try_new(
+ Arc::clone(schema),
+ vec![Arc::new(count_array) as ArrayRef],
+ )?)
+}
- // Convert to &dyn Array which implements Datum
- let new_arr: &dyn Array = new_array.as_ref();
- let orig_arr: &dyn Array = original_column.as_ref();
- zip(&update_mask, &new_arr, &orig_arr)?
- } else {
- Arc::clone(original_column)
- };
+/// Render the predicates as a comma separated list.
+fn format_predicates(predicates: &[Arc<dyn PhysicalExpr>]) -> String {
+ predicates
+ .iter()
+ .map(|predicate| predicate.to_string())
+ .collect::<Vec<_>>()
+ .join(", ")
+}
- new_columns.push(new_column);
- }
+/// Deletes the matching rows of a [`MemTable`] when it runs, and emits the
count.
+///
+/// The provider hook that builds this node changes no row, so `EXPLAIN DELETE`
+/// prints the plan and the table keeps its rows. Each run of the plan applies
+/// the delete once, as [`DataSinkExec`] does for an `INSERT`.
+#[derive(Debug)]
+struct MemDeleteExec {
+ /// Partitions of the target table, shared with the [`MemTable`].
+ partitions: Vec<PartitionData>,
+ /// Declared sort order of the target table. A delete clears it.
+ sort_order: Arc<Mutex<Vec<Vec<SortExpr>>>>,
+ /// Predicates of the `WHERE` clause. An empty list matches every row.
+ predicates: Vec<Arc<dyn PhysicalExpr>>,
+ /// Single `count` column of the output.
+ schema: SchemaRef,
+ properties: Arc<PlanProperties>,
+}
- let updated_batch =
- ArrowRecordBatch::try_new(Arc::clone(&self.schema),
new_columns)?;
- new_batches.push(updated_batch);
- }
+impl MemDeleteExec {
+ fn new(
+ partitions: Vec<PartitionData>,
+ sort_order: Arc<Mutex<Vec<Vec<SortExpr>>>>,
+ predicates: Vec<Arc<dyn PhysicalExpr>>,
+ ) -> Self {
+ let schema = dml_count_schema();
+ let properties = dml_plan_properties(&schema);
- *partition = new_batches;
+ Self {
+ partitions,
+ sort_order,
+ predicates,
+ schema,
+ properties,
}
+ }
+}
- Ok(Arc::new(DmlResultExec::new(total_updated)))
+impl DisplayAs for MemDeleteExec {
+ fn fmt_as(
+ &self,
+ t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ match t {
+ DisplayFormatType::Default
+ | DisplayFormatType::Verbose
+ | DisplayFormatType::TreeRender => {
+ write!(f, "MemDeleteExec")?;
+ if !self.predicates.is_empty() {
+ write!(f, ": predicate=[{}]",
format_predicates(&self.predicates))?;
+ }
+ Ok(())
+ }
+ }
}
}
-/// Evaluate filter expressions against a batch and return a combined boolean
mask.
-/// Returns None if filters is empty (meaning "match all rows").
-/// The returned mask has true for rows that match the filter predicates.
-fn evaluate_filters_to_mask(
- filters: &[Expr],
- batch: &RecordBatch,
- df_schema: &DFSchema,
- execution_props: &datafusion_expr::execution_props::ExecutionProps,
-) -> Result<Option<BooleanArray>> {
- if filters.is_empty() {
- return Ok(None);
+impl ExecutionPlan for MemDeleteExec {
+ fn name(&self) -> &str {
+ "MemDeleteExec"
}
- let mut combined_mask: Option<BooleanArray> = None;
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
- for filter_expr in filters {
- let physical_expr = create_physical_expr(
- filter_expr,
- df_schema,
- execution_props,
- &PhysicalPlanningContext::default(),
- )?;
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.properties
+ }
- let result = physical_expr.evaluate(batch)?;
- let array = result.into_array(batch.num_rows())?;
- let bool_array = array
- .as_any()
- .downcast_ref::<BooleanArray>()
- .ok_or_else(|| {
- datafusion_common::DataFusionError::Internal(
- "Filter did not evaluate to boolean".to_string(),
- )
- })?
- .clone();
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![]
+ }
- combined_mask = Some(match combined_mask {
- Some(existing) => and(&existing, &bool_array)?,
- None => bool_array,
+ fn replace_children(
+ self: Arc<Self>,
+ _: Vec<Arc<dyn ExecutionPlan>>,
+ _: ReplaceChildrenOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(self)
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.replace_children(
+ children,
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )
+ }
+
+ fn execute(
+ &self,
+ partition: usize,
+ _context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ if partition != 0 {
+ return internal_err!(
+ "MemDeleteExec has one partition, but partition {partition}
was requested"
+ );
+ }
+
+ let partitions = self.partitions.clone();
+ let sort_order = Arc::clone(&self.sort_order);
+ let predicates = self.predicates.clone();
+ let schema = Arc::clone(&self.schema);
+ let count_schema = Arc::clone(&self.schema);
Review Comment:
I think there is no need of two clones of the schema here.
--
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]
