Amogh-2404 commented on code in PR #23906:
URL: https://github.com/apache/datafusion/pull/23906#discussion_r3766188856
##########
datafusion/ffi/src/table_provider.rs:
##########
@@ -336,6 +378,133 @@ unsafe extern "C" fn insert_into_fn_wrapper(
.into_ffi()
}
+unsafe extern "C" fn delete_from_fn_wrapper(
+ provider: &FFI_TableProvider,
+ session: FFI_SessionRef,
+ filters_serialized: SVec<u8>,
+) -> FfiFuture<FFI_Result<FFI_ExecutionPlan>> {
+ let task_ctx: Result<Arc<TaskContext>, DataFusionError> =
+ (&provider.logical_codec.task_ctx_provider).try_into();
+ let runtime = provider.runtime().clone();
+ let logical_codec: Arc<dyn LogicalExtensionCodec> =
(&provider.logical_codec).into();
+ let internal_provider = Arc::clone(provider.inner());
+
+ async move {
+ let mut foreign_session = None;
+ let session = sresult_return!(
+ session
+ .as_local()
+ .map(Ok::<&(dyn Session + Send + Sync), DataFusionError>)
+ .unwrap_or_else(|| {
+ foreign_session =
Some(ForeignSession::try_from(&session)?);
+ Ok(foreign_session.as_ref().unwrap())
+ })
+ );
+
+ let task_ctx = sresult_return!(task_ctx);
+ let filters = sresult_return!(parse_serialized_exprs(
+ &filters_serialized,
+ &task_ctx,
+ logical_codec.as_ref(),
+ ));
+
+ let plan = sresult_return!(internal_provider.delete_from(session,
filters).await);
+
+ FFI_Result::Ok(FFI_ExecutionPlan::new(plan, runtime))
+ }
+ .into_ffi()
+}
+
+unsafe extern "C" fn update_fn_wrapper(
+ provider: &FFI_TableProvider,
+ session: FFI_SessionRef,
+ assignments: SVec<FFI_TableProviderUpdateAssignment>,
+ filters_serialized: SVec<u8>,
+) -> FfiFuture<FFI_Result<FFI_ExecutionPlan>> {
+ let task_ctx: Result<Arc<TaskContext>, DataFusionError> =
+ (&provider.logical_codec.task_ctx_provider).try_into();
+ let runtime = provider.runtime().clone();
+ let logical_codec: Arc<dyn LogicalExtensionCodec> =
(&provider.logical_codec).into();
+ let internal_provider = Arc::clone(provider.inner());
+
+ async move {
+ let mut foreign_session = None;
+ let session = sresult_return!(
+ session
+ .as_local()
+ .map(Ok::<&(dyn Session + Send + Sync), DataFusionError>)
+ .unwrap_or_else(|| {
+ foreign_session =
Some(ForeignSession::try_from(&session)?);
+ Ok(foreign_session.as_ref().unwrap())
+ })
+ );
+
+ let task_ctx = sresult_return!(task_ctx);
+ let assignments = sresult_return!(
+ assignments
+ .into_iter()
+ .map(|assignment| {
+ let mut exprs = parse_serialized_exprs(
+ &assignment.expr_serialized,
+ &task_ctx,
+ logical_codec.as_ref(),
+ )?;
+ let expr = match exprs.len() {
+ 1 => exprs.remove(0),
+ _ => {
+ return Err(DataFusionError::Plan(
+ "Expected exactly one expression for update
assignment"
+ .to_string(),
+ ));
+ }
+ };
Review Comment:
Updated the error to include the assignment column and the decoded
expression count.
##########
datafusion/ffi/src/table_provider.rs:
##########
@@ -336,6 +378,133 @@ unsafe extern "C" fn insert_into_fn_wrapper(
.into_ffi()
}
+unsafe extern "C" fn delete_from_fn_wrapper(
+ provider: &FFI_TableProvider,
+ session: FFI_SessionRef,
+ filters_serialized: SVec<u8>,
+) -> FfiFuture<FFI_Result<FFI_ExecutionPlan>> {
+ let task_ctx: Result<Arc<TaskContext>, DataFusionError> =
+ (&provider.logical_codec.task_ctx_provider).try_into();
+ let runtime = provider.runtime().clone();
+ let logical_codec: Arc<dyn LogicalExtensionCodec> =
(&provider.logical_codec).into();
+ let internal_provider = Arc::clone(provider.inner());
+
+ async move {
+ let mut foreign_session = None;
+ let session = sresult_return!(
+ session
+ .as_local()
+ .map(Ok::<&(dyn Session + Send + Sync), DataFusionError>)
+ .unwrap_or_else(|| {
+ foreign_session =
Some(ForeignSession::try_from(&session)?);
+ Ok(foreign_session.as_ref().unwrap())
+ })
+ );
Review Comment:
I kept these wrappers explicit. The local `ForeignSession` guard keeps the
borrowed session alive for the call, so extracting this would make the lifetime
relationship less clear for very little reduction.
##########
datafusion/ffi/src/table_provider.rs:
##########
@@ -671,11 +906,161 @@ mod tests {
Ok(())
}
+ #[derive(Debug, Default)]
+ struct DmlCalls {
+ delete_filters: Option<Vec<Expr>>,
+ update_assignments: Option<Vec<(String, Expr)>>,
+ update_filters: Option<Vec<Expr>>,
+ truncated: bool,
+ }
+
+ #[derive(Debug)]
+ struct DmlTableProvider {
+ calls: Arc<Mutex<DmlCalls>>,
+ schema: SchemaRef,
+ }
+
+ fn dml_count_plan() -> Arc<dyn ExecutionPlan> {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "count",
+ DataType::UInt64,
+ false,
+ )]));
+ Arc::new(crate::execution_plan::tests::EmptyExec::new(schema))
+ }
+
+ #[async_trait]
+ impl TableProvider for DmlTableProvider {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+
+ fn table_type(&self) -> TableType {
+ TableType::Base
+ }
+
+ async fn scan(
+ &self,
+ _session: &dyn Session,
+ _projection: Option<&Vec<usize>>,
+ _filters: &[Expr],
+ _limit: Option<usize>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Err(DataFusionError::Internal(
+ "DmlTableProvider scan should not be called".to_string(),
+ ))
+ }
Review Comment:
Reworked this to extend the existing `TableWithStats` fixture and removed
the separate DML provider and call-state type.
##########
datafusion/ffi/src/table_provider.rs:
##########
@@ -202,21 +228,45 @@ unsafe extern "C" fn table_type_fn_wrapper(
provider.inner().table_type().into()
}
-fn supports_filters_pushdown_internal(
- provider: &Arc<dyn TableProvider>,
- filters_serialized: &[u8],
+fn parse_serialized_exprs(
+ exprs_serialized: &[u8],
task_ctx: &Arc<TaskContext>,
codec: &dyn LogicalExtensionCodec,
-) -> Result<SVec<FFI_TableProviderFilterPushDown>> {
- let filters = match filters_serialized.is_empty() {
- true => vec![],
+) -> Result<Vec<Expr>> {
+ match exprs_serialized.is_empty() {
+ true => Ok(vec![]),
false => {
- let proto_filters = LogicalExprList::decode(filters_serialized)
+ let proto_exprs = LogicalExprList::decode(exprs_serialized)
.map_err(|e| DataFusionError::Plan(e.to_string()))?;
- parse_exprs(proto_filters.expr.iter(), task_ctx.as_ref(), codec)?
+ Ok(parse_exprs(
+ proto_exprs.expr.iter(),
+ task_ctx.as_ref(),
+ codec,
+ )?)
}
- };
+ }
+}
+
+fn serialize_expr_list<'a>(
+ exprs: impl IntoIterator<Item = &'a Expr>,
+ codec: &dyn LogicalExtensionCodec,
+) -> Result<SVec<u8>> {
+ Ok(LogicalExprList {
+ expr: serialize_exprs(exprs, codec)?,
+ }
+ .encode_to_vec()
+ .into_iter()
+ .collect())
+}
Review Comment:
Moved the protobuf expression-list encoding and decoding into
`datafusion-proto`. The FFI helper left here only converts `Bytes` to
`SVec<u8>` at the ABI boundary. I also reused the proto helpers in the UDTF
paths.
##########
datafusion/ffi/src/table_provider.rs:
##########
@@ -202,21 +228,45 @@ unsafe extern "C" fn table_type_fn_wrapper(
provider.inner().table_type().into()
}
-fn supports_filters_pushdown_internal(
- provider: &Arc<dyn TableProvider>,
- filters_serialized: &[u8],
+fn parse_serialized_exprs(
+ exprs_serialized: &[u8],
task_ctx: &Arc<TaskContext>,
codec: &dyn LogicalExtensionCodec,
-) -> Result<SVec<FFI_TableProviderFilterPushDown>> {
- let filters = match filters_serialized.is_empty() {
- true => vec![],
+) -> Result<Vec<Expr>> {
Review Comment:
Done. The protobuf decode and parse logic now lives in `datafusion-proto`;
this local wrapper only supplies the task context used at the FFI boundary.
--
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]
