alamb commented on code in PR #24655:
URL: https://github.com/apache/datafusion/pull/24655#discussion_r3874201066
##########
datafusion/catalog/src/memory/table.rs:
##########
@@ -469,133 +427,216 @@ impl MemTable {
// Create physical expressions for assignments upfront (outside batch
loop)
let physical_assignments: HashMap<String, Arc<dyn PhysicalExpr>> =
assignments
- .iter()
+ .into_iter()
.map(|(name, expr)| {
let physical_expr = create_physical_expr(
- expr,
+ &expr,
&df_schema,
state.execution_props(),
&PhysicalPlanningContext::default(),
)?;
- Ok((name.clone(), physical_expr))
+ Ok((name, physical_expr))
})
.collect::<Result<_>>()?;
- *self.sort_order.lock() = vec![];
+ let filters = compile_filters(filters, &df_schema,
state.execution_props())?;
- let mut total_updated: u64 = 0;
+ Ok(self.dml_exec(
+ self.batches.clone(),
+ filters,
+ MemDmlOp::Update(physical_assignments),
+ ))
+ }
- for partition_data in &self.batches {
- let mut partition = partition_data.write().await;
- let mut new_batches = Vec::with_capacity(partition.len());
+ /// Build the plan node that applies `op` to `partitions`.
+ fn dml_exec(
+ &self,
+ partitions: Vec<PartitionData>,
+ filters: Vec<Arc<dyn PhysicalExpr>>,
+ op: MemDmlOp,
+ ) -> Arc<dyn ExecutionPlan> {
+ Arc::new(MemDmlExec::new(MemDmlState {
+ partitions,
+ table_schema: Arc::clone(&self.schema),
+ sort_order: Arc::clone(&self.sort_order),
+ filters,
+ op,
+ }))
+ }
+}
- for batch in partition.iter() {
- if batch.num_rows() == 0 {
- continue;
- }
+/// Compile the `WHERE` clause of a DELETE or an UPDATE into physical
+/// expressions. An empty result means "match all rows".
+fn compile_filters(
+ filters: Vec<Expr>,
+ df_schema: &DFSchema,
+ execution_props: &datafusion_expr::execution_props::ExecutionProps,
+) -> Result<Vec<Arc<dyn PhysicalExpr>>> {
+ filters
+ .into_iter()
+ .map(|filter_expr| {
+ create_physical_expr(
+ &filter_expr,
+ df_schema,
+ execution_props,
+ &PhysicalPlanningContext::default(),
+ )
+ })
+ .collect()
+}
- // Evaluate filters - None means "match all rows"
- let filter_mask = evaluate_filters_to_mask(
- &filters,
- batch,
- &df_schema,
- state.execution_props(),
- )?;
+/// Delete the rows of `state` that its filters match, and return the number of
+/// rows deleted.
+async fn apply_delete(state: &MemDmlState) -> Result<u64> {
Review Comment:
minor: would this be more natural as a method of `MemDmlState` rather than
taking it as an argument ?
```rust
impl MemDmlState {
fn apply_delete(&self) -> Result<u64> {
...
}
}
```
##########
datafusion/catalog/src/memory/table.rs:
##########
@@ -469,133 +427,216 @@ impl MemTable {
// Create physical expressions for assignments upfront (outside batch
loop)
let physical_assignments: HashMap<String, Arc<dyn PhysicalExpr>> =
assignments
- .iter()
+ .into_iter()
.map(|(name, expr)| {
let physical_expr = create_physical_expr(
- expr,
+ &expr,
&df_schema,
state.execution_props(),
&PhysicalPlanningContext::default(),
)?;
- Ok((name.clone(), physical_expr))
+ Ok((name, physical_expr))
})
.collect::<Result<_>>()?;
- *self.sort_order.lock() = vec![];
+ let filters = compile_filters(filters, &df_schema,
state.execution_props())?;
- let mut total_updated: u64 = 0;
+ Ok(self.dml_exec(
+ self.batches.clone(),
+ filters,
+ MemDmlOp::Update(physical_assignments),
+ ))
+ }
- for partition_data in &self.batches {
- let mut partition = partition_data.write().await;
- let mut new_batches = Vec::with_capacity(partition.len());
+ /// Build the plan node that applies `op` to `partitions`.
+ fn dml_exec(
+ &self,
+ partitions: Vec<PartitionData>,
+ filters: Vec<Arc<dyn PhysicalExpr>>,
+ op: MemDmlOp,
+ ) -> Arc<dyn ExecutionPlan> {
+ Arc::new(MemDmlExec::new(MemDmlState {
+ partitions,
+ table_schema: Arc::clone(&self.schema),
+ sort_order: Arc::clone(&self.sort_order),
+ filters,
+ op,
+ }))
+ }
+}
- for batch in partition.iter() {
- if batch.num_rows() == 0 {
- continue;
- }
+/// Compile the `WHERE` clause of a DELETE or an UPDATE into physical
+/// expressions. An empty result means "match all rows".
+fn compile_filters(
+ filters: Vec<Expr>,
+ df_schema: &DFSchema,
+ execution_props: &datafusion_expr::execution_props::ExecutionProps,
+) -> Result<Vec<Arc<dyn PhysicalExpr>>> {
+ filters
+ .into_iter()
+ .map(|filter_expr| {
+ create_physical_expr(
+ &filter_expr,
+ df_schema,
+ execution_props,
+ &PhysicalPlanningContext::default(),
+ )
+ })
+ .collect()
+}
- // Evaluate filters - None means "match all rows"
- let filter_mask = evaluate_filters_to_mask(
- &filters,
- batch,
- &df_schema,
- state.execution_props(),
- )?;
+/// Delete the rows of `state` that its filters match, and return the number of
+/// rows deleted.
+async fn apply_delete(state: &MemDmlState) -> Result<u64> {
+ let mut total_deleted: u64 = 0;
- 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()]),
- )
- }
- };
+ for partition_data in &state.partitions {
+ let mut partition = partition_data.write().await;
+ let mut new_batches = Vec::with_capacity(partition.len());
- total_updated += update_count as u64;
+ for batch in partition.iter() {
+ if batch.num_rows() == 0 {
+ continue;
+ }
- if update_count == 0 {
- new_batches.push(batch.clone());
- continue;
+ // Evaluate filters - None means "match all rows"
+ let filter_mask = evaluate_filters_to_mask(&state.filters, batch)?;
+
+ 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 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())?;
-
- // 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)
- };
-
- new_columns.push(new_column);
+ let filtered_batch = filter_record_batch(batch, &keep_mask)?;
+ if filtered_batch.num_rows() > 0 {
+ new_batches.push(filtered_batch);
+ }
+ }
+
+ *partition = new_batches;
+ }
+
+ Ok(total_deleted)
+}
+
+/// Assign a new value to each row of `state` that its filters match, and
return
+/// the number of rows updated.
+async fn apply_update(
+ state: &MemDmlState,
Review Comment:
similar question above -- maybe this would be nicer as a
MemDmlState::apply_update method
##########
datafusion/catalog/src/memory/table.rs:
##########
@@ -361,71 +363,24 @@ impl MemTable {
state: &'a dyn Session,
filters: Vec<Expr>,
) -> BoxFuture<'a, Result<Arc<dyn ExecutionPlan>>> {
- Box::pin(self.delete_from_inner(state, filters))
+ Box::pin(ready(self.plan_delete(state, filters)))
}
- async fn delete_from_inner(
+ /// Build the plan of a DELETE. The rows change when the plan runs, not
here.
+ fn plan_delete(
&self,
state: &dyn Session,
filters: Vec<Expr>,
) -> Result<Arc<dyn ExecutionPlan>> {
// Early exit if table has no partitions
if self.batches.is_empty() {
- return Ok(Arc::new(DmlResultExec::new(0)));
+ return Ok(self.dml_exec(vec![], vec![], MemDmlOp::Delete));
Review Comment:
As in the new code will error and the old didn't? I think that is a good
change if so
--
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]
