alamb commented on code in PR #2226: URL: https://github.com/apache/arrow-datafusion/pull/2226#discussion_r856456282
########## datafusion/scheduler/src/task.rs: ########## @@ -0,0 +1,439 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::query::Query; +use crate::{is_worker, spawn_local, spawn_local_fifo, RoutablePipeline, Spawner}; +use arrow::record_batch::RecordBatch; +use datafusion::error::{DataFusionError, Result}; +use futures::channel::mpsc; +use futures::task::ArcWake; +use futures::{Stream, StreamExt}; +use log::{debug, trace}; +use std::pin::Pin; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Weak}; +use std::task::{Context, Poll}; + +/// Spawns a query using the provided [`Spawner`] +pub fn spawn_query(query: Query, spawner: Spawner) -> QueryResults { + debug!("Spawning query: {:#?}", query); + + let (sender, receiver) = mpsc::unbounded(); + let query = Arc::new(QueryTask { + spawner, + pipelines: query.pipelines, + output: sender, + }); + + for (pipeline_idx, query_pipeline) in query.pipelines.iter().enumerate() { + for partition in 0..query_pipeline.pipeline.output_partitions() { + query.spawner.spawn(Task { + query: query.clone(), + waker: Arc::new(TaskWaker { + query: Arc::downgrade(&query), + wake_count: AtomicUsize::new(1), + pipeline: pipeline_idx, + partition, + }), + }); + } + } + + QueryResults { + query, + inner: receiver, + } +} + +/// A [`Task`] identifies an output partition within a given pipeline that may be able to +/// make progress. The [`Scheduler`][super::Scheduler] maintains a list of outstanding +/// [`Task`] and distributes them amongst its worker threads. +/// +/// A [`Query`] is considered completed when it has no outstanding [`Task`] +pub struct Task { + /// Maintain a link to the [`QueryTask`] this is necessary to be able to + /// route the output of the partition to its destination, and also because + /// when [`QueryTask`] is dropped it signals completion of query execution + query: Arc<QueryTask>, + + /// A [`ArcWake`] that can be used to re-schedule this [`Task`] for execution + waker: Arc<TaskWaker>, +} + +impl std::fmt::Debug for Task { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let output = &self.query.pipelines[self.waker.pipeline].output; + + f.debug_struct("Task") + .field("pipeline", &self.waker.pipeline) + .field("partition", &self.waker.partition) + .field("output", &output) + .finish() + } +} + +impl Task { + fn handle_error(&self, routable: &RoutablePipeline, error: DataFusionError) { + self.query.send_query_output(Err(error)); + if let Some(link) = routable.output { + trace!( + "Closing pipeline: {:?}, partition: {}, due to error", + link, + self.waker.partition, + ); + + self.query.pipelines[link.pipeline] + .pipeline + .close(link.child, self.waker.partition); + } + } + /// Call [`Pipeline::poll_partition`] attempting to make progress on query execution + pub fn do_work(self) { + assert!(is_worker(), "Task::do_work called outside of worker pool"); + if self.query.is_cancelled() { + return; + } + + // Capture the wake count prior to calling [`Pipeline::poll_partition`] + // this allows us to detect concurrent wake ups and handle them correctly + let wake_count = self.waker.wake_count.load(Ordering::SeqCst); + + let node = self.waker.pipeline; + let partition = self.waker.partition; + + let waker = futures::task::waker_ref(&self.waker); + let mut cx = Context::from_waker(&*waker); + + let pipelines = &self.query.pipelines; + let routable = &pipelines[node]; + match routable.pipeline.poll_partition(&mut cx, partition) { + Poll::Ready(Some(Ok(batch))) => { + trace!("Poll {:?}: Ok: {}", self, batch.num_rows()); + match routable.output { + Some(link) => { + trace!( + "Publishing batch to pipeline {:?} partition {}", + link, + partition + ); + + let r = pipelines[link.pipeline] + .pipeline + .push(batch, link.child, partition); + + if let Err(e) = r { + self.handle_error(routable, e); + + // Return without rescheduling this output again + return; + } + } + None => { + trace!("Publishing batch to output"); + self.query.send_query_output(Ok(batch)) + } + } + + // Reschedule this pipeline again + // + // We want to prioritise running tasks triggered by the most recent + // batch, so reschedule with FIFO ordering + // + // Note: We must schedule after we have routed the batch, otherwise + // we introduce a potential ordering race where the newly scheduled + // task runs before this task finishes routing the output + spawn_local_fifo(self); + } + Poll::Ready(Some(Err(e))) => { + trace!("Poll {:?}: Error: {:?}", self, e); + self.handle_error(routable, e) + } + Poll::Ready(None) => { + trace!("Poll {:?}: None", self); + match routable.output { + Some(link) => { + trace!("Closing pipeline: {:?}, partition: {}", link, partition); + pipelines[link.pipeline] + .pipeline + .close(link.child, partition) + } + None => self.query.finish(), + } + } + Poll::Pending => { + trace!("Poll {:?}: Pending", self); + // Attempt to reset the wake count with the value obtained prior + // to calling [`Pipeline::poll_partition`]. + // + // If this fails it indicates a wakeup was received whilst executing + // [`Pipeline::poll_partition`] and we should reschedule the task + let reset = self.waker.wake_count.compare_exchange( + wake_count, + 0, + Ordering::SeqCst, + Ordering::SeqCst, + ); + + if reset.is_err() { + trace!("Wakeup triggered whilst polling: {:?}", self); + spawn_local(self); + } + } + } + } +} + +/// The result stream for a query +/// +/// # Cancellation +/// +/// Dropping this will cancel the inflight query +pub struct QueryResults { + inner: mpsc::UnboundedReceiver<Option<Result<RecordBatch>>>, + + /// Keep a reference to the [`QueryTask`] so it isn't dropped early + #[allow(unused)] + query: Arc<QueryTask>, +} + +impl Stream for QueryResults { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + self.inner.poll_next_unpin(cx).map(Option::flatten) + } +} + +/// The shared state of all [`Task`] created from the same [`Query`] +#[derive(Debug)] +struct QueryTask { + /// Spawner for this query + spawner: Spawner, + + /// List of pipelines that belong to this query, pipelines are addressed + /// based on their index within this list + pipelines: Vec<RoutablePipeline>, + + /// The output stream for this query's execution + output: mpsc::UnboundedSender<Option<Result<RecordBatch>>>, +} + +impl Drop for QueryTask { + fn drop(&mut self) { + debug!("Query dropped"); + } +} + +impl QueryTask { + /// Returns `true` if this query has been dropped, specifically if the + /// stream returned by [`super::Scheduler::schedule`] has been dropped + fn is_cancelled(&self) -> bool { + self.output.is_closed() + } + + /// Sends `output` to this query's output stream + fn send_query_output(&self, output: Result<RecordBatch>) { + let _ = self.output.unbounded_send(Some(output)); + } + + /// Mark this query as finished + fn finish(&self) { + let _ = self.output.unbounded_send(None); + } +} + +struct TaskWaker { Review Comment: lets not do that lol -- 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]
