gabotechs commented on code in PR #27167:
URL: https://github.com/apache/datafusion/pull/27167#discussion_r4522718114
##########
datafusion/datasource/src/instrumentation.rs:
##########
@@ -0,0 +1,104 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF)
under one or more contributor license agreements
+
+//! Application-defined instrumentation for
[`DataSourceExec`](crate::source::DataSourceExec).
+
+use std::fmt::Debug;
+use std::sync::Arc;
+
+use datafusion_common::Result;
+use datafusion_execution::{SendableRecordBatchStream, TaskContext};
+use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
+
+use crate::source::DataSource;
+
+/// Opens one partition and applies `DataSourceExec`'s output batch splitting.
+///
+/// Opening is synchronous; the returned stream performs the asynchronous
reads.
+pub type DataSourceOpen<'a> =
+ Box<dyn FnOnce() -> Result<SendableRecordBatchStream> + Send + 'a>;
+
+/// Instruments the opening and output stream of a
[`DataSourceExec`](crate::source::DataSourceExec).
+///
+/// Implementations must call `open` exactly once and preserve the returned
stream's
+/// schema, batches, errors, and polling behavior. They can establish context
while
+/// opening, observe opening errors, and wrap the stream to record batches,
stream
+/// errors, completion, or cancellation. Stream wrappers should establish any
+/// thread-local context separately on each poll rather than retain a context
guard.
+///
+/// Register an implementation using [`DataSourceInstrumentationExtension`].
+/// This hook covers sources executed through `DataSourceExec`, not arbitrary
+/// execution plans or direct calls to `DataSource::open`.
+pub trait DataSourceInstrumentation: Debug + Send + Sync {
+ /// Instruments one partition using the executing task's context.
+ ///
+ /// `source` allows applications to select which sources to instrument.
+ /// `metrics` is owned by the executor, independently of the source's
metrics,
+ /// and is included in `ExecutionPlan::metrics()`. Use distinct metric
names to
+ /// avoid ambiguity with source metrics. Clones of the plan share these
metrics;
+ /// resetting execution state or replacing the source creates a fresh set.
+ ///
+ /// The callback returns batches after source-level limits and output batch
+ /// splitting, before downstream operators. Stream lifetime includes
consumer
+ /// backpressure; it is not individual storage-request latency.
+ fn instrument(
+ &self,
+ source: &dyn DataSource,
+ partition: usize,
+ context: &TaskContext,
+ metrics: &ExecutionPlanMetricsSet,
+ open: DataSourceOpen<'_>,
Review Comment:
Why does it need to take an intermediate callback definition? wouldn't it be
simpler if it just took the recently created `SendableRecordBatchStream`?
##########
datafusion/datasource/src/instrumentation.rs:
##########
@@ -0,0 +1,104 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF)
under one or more contributor license agreements
+
+//! Application-defined instrumentation for
[`DataSourceExec`](crate::source::DataSourceExec).
+
+use std::fmt::Debug;
+use std::sync::Arc;
+
+use datafusion_common::Result;
+use datafusion_execution::{SendableRecordBatchStream, TaskContext};
+use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
+
+use crate::source::DataSource;
+
+/// Opens one partition and applies `DataSourceExec`'s output batch splitting.
+///
+/// Opening is synchronous; the returned stream performs the asynchronous
reads.
+pub type DataSourceOpen<'a> =
+ Box<dyn FnOnce() -> Result<SendableRecordBatchStream> + Send + 'a>;
+
+/// Instruments the opening and output stream of a
[`DataSourceExec`](crate::source::DataSourceExec).
+///
+/// Implementations must call `open` exactly once and preserve the returned
stream's
+/// schema, batches, errors, and polling behavior. They can establish context
while
+/// opening, observe opening errors, and wrap the stream to record batches,
stream
+/// errors, completion, or cancellation. Stream wrappers should establish any
+/// thread-local context separately on each poll rather than retain a context
guard.
+///
+/// Register an implementation using [`DataSourceInstrumentationExtension`].
+/// This hook covers sources executed through `DataSourceExec`, not arbitrary
+/// execution plans or direct calls to `DataSource::open`.
+pub trait DataSourceInstrumentation: Debug + Send + Sync {
+ /// Instruments one partition using the executing task's context.
+ ///
+ /// `source` allows applications to select which sources to instrument.
+ /// `metrics` is owned by the executor, independently of the source's
metrics,
+ /// and is included in `ExecutionPlan::metrics()`. Use distinct metric
names to
+ /// avoid ambiguity with source metrics. Clones of the plan share these
metrics;
+ /// resetting execution state or replacing the source creates a fresh set.
+ ///
+ /// The callback returns batches after source-level limits and output batch
+ /// splitting, before downstream operators. Stream lifetime includes
consumer
+ /// backpressure; it is not individual storage-request latency.
+ fn instrument(
+ &self,
+ source: &dyn DataSource,
+ partition: usize,
+ context: &TaskContext,
+ metrics: &ExecutionPlanMetricsSet,
+ open: DataSourceOpen<'_>,
+ ) -> Result<SendableRecordBatchStream>;
+}
+
+/// Session extension holding an application-defined data source
instrumentation hook.
+///
+/// Install with
[`SessionConfig::with_extension`](datafusion_execution::config::SessionConfig::with_extension).
+/// The hook is resolved at execution time and is not part of the serialized
plan.
+/// Applications executing plans in other processes must install it in those
sessions too.
+/// Without this extension, `DataSourceExec` opens its source directly.
+///
+/// # Example
+///
+/// ```
+/// use std::sync::Arc;
+/// use datafusion_common::Result;
+/// use datafusion_execution::{SendableRecordBatchStream, TaskContext};
+/// use datafusion_execution::config::SessionConfig;
+/// use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet,
MetricBuilder};
+/// use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
+/// use datafusion_datasource::source::DataSource;
+/// use datafusion_datasource::instrumentation::{
+/// DataSourceInstrumentation, DataSourceInstrumentationExtension,
DataSourceOpen,
+/// };
+/// use futures::TryStreamExt;
+///
+/// #[derive(Debug)]
+/// struct ScanMetrics;
+///
+/// impl DataSourceInstrumentation for ScanMetrics {
+/// fn instrument(
+/// &self,
+/// _source: &dyn DataSource,
+/// partition: usize,
+/// _context: &TaskContext,
+/// metrics: &ExecutionPlanMetricsSet,
+/// open: DataSourceOpen<'_>,
+/// ) -> Result<SendableRecordBatchStream> {
+/// let stream = open()?;
+/// let schema = stream.schema();
+/// let rows = MetricBuilder::new(metrics)
+/// .with_new_label("application", "example")
+/// .counter("observed_rows", partition);
+/// let stream = stream.inspect_ok(move |batch|
rows.add(batch.num_rows()));
+/// Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream)))
+/// }
+/// }
+///
+/// let config = SessionConfig::new().with_extension(Arc::new(
+/// DataSourceInstrumentationExtension(Arc::new(ScanMetrics)),
+/// ));
+/// ```
+#[derive(Debug, Clone)]
+pub struct DataSourceInstrumentationExtension(pub Arc<dyn
DataSourceInstrumentation>);
Review Comment:
Imagine you have multiple teams adding each one their own `DataSource`
implementations, and there are two teams wanting to add observability to their
respective `DataSource`s, for example `DeltaDataSource` and
`IcebergDataSource`, how would they use this instrumentation hook without
stepping in each other toes?
##########
datafusion/datasource/src/source.rs:
##########
@@ -486,6 +506,11 @@ impl ExecutionPlan for DataSourceExec {
metrics.push(output_rows_skew);
}
+ if let Some(instrumentation_metrics) =
self.instrumentation_metrics.get() {
+ for metric in instrumentation_metrics.clone_inner().iter() {
+ metrics.push(Arc::clone(metric));
+ }
+ }
Review Comment:
```suggestion
metrics.extend(instrumentation_metrics.clone_inner());
```
##########
datafusion/datasource/src/instrumentation.rs:
##########
@@ -0,0 +1,104 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF)
under one or more contributor license agreements
+
+//! Application-defined instrumentation for
[`DataSourceExec`](crate::source::DataSourceExec).
+
+use std::fmt::Debug;
+use std::sync::Arc;
+
+use datafusion_common::Result;
+use datafusion_execution::{SendableRecordBatchStream, TaskContext};
+use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
+
+use crate::source::DataSource;
+
+/// Opens one partition and applies `DataSourceExec`'s output batch splitting.
+///
+/// Opening is synchronous; the returned stream performs the asynchronous
reads.
+pub type DataSourceOpen<'a> =
+ Box<dyn FnOnce() -> Result<SendableRecordBatchStream> + Send + 'a>;
+
+/// Instruments the opening and output stream of a
[`DataSourceExec`](crate::source::DataSourceExec).
+///
+/// Implementations must call `open` exactly once and preserve the returned
stream's
+/// schema, batches, errors, and polling behavior. They can establish context
while
+/// opening, observe opening errors, and wrap the stream to record batches,
stream
+/// errors, completion, or cancellation. Stream wrappers should establish any
+/// thread-local context separately on each poll rather than retain a context
guard.
+///
+/// Register an implementation using [`DataSourceInstrumentationExtension`].
+/// This hook covers sources executed through `DataSourceExec`, not arbitrary
+/// execution plans or direct calls to `DataSource::open`.
+pub trait DataSourceInstrumentation: Debug + Send + Sync {
+ /// Instruments one partition using the executing task's context.
+ ///
+ /// `source` allows applications to select which sources to instrument.
+ /// `metrics` is owned by the executor, independently of the source's
metrics,
+ /// and is included in `ExecutionPlan::metrics()`. Use distinct metric
names to
+ /// avoid ambiguity with source metrics. Clones of the plan share these
metrics;
+ /// resetting execution state or replacing the source creates a fresh set.
+ ///
+ /// The callback returns batches after source-level limits and output batch
+ /// splitting, before downstream operators. Stream lifetime includes
consumer
+ /// backpressure; it is not individual storage-request latency.
+ fn instrument(
+ &self,
+ source: &dyn DataSource,
+ partition: usize,
+ context: &TaskContext,
+ metrics: &ExecutionPlanMetricsSet,
+ open: DataSourceOpen<'_>,
+ ) -> Result<SendableRecordBatchStream>;
Review Comment:
By making this return a `SendableRecordBatchStream`, it might no longer be
an instrumentation hook, as people can decide to do whatever they with the data.
I see how the current shape of the API could be abused for reasons other
than observability, but not sure if it's worth doing something about it rather
than just mentioning that in the doc comments.
--
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]
