tustvold commented on code in PR #2226: URL: https://github.com/apache/arrow-datafusion/pull/2226#discussion_r850818283
########## datafusion/scheduler/src/task.rs: ########## @@ -0,0 +1,225 @@ +// 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::{is_worker, spawn_local, spawn_local_fifo, Query}; +use futures::task::ArcWake; +use log::{debug, trace}; +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: Arc<Query>) { + debug!("Spawning query: {:#?}", query); + + let spawner = query.spawner(); + + for (pipeline_idx, query_pipeline) in query.pipelines().iter().enumerate() { + for partition in 0..query_pipeline.pipeline.output_partitions() { + spawner.spawn(Task { + query: query.clone(), + waker: Arc::new(TaskWaker { + query: Arc::downgrade(&query), + wake_count: AtomicUsize::new(1), + pipeline: pipeline_idx, + partition, + }), + }); + } + } +} + +/// 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 [`Query`] this is necessary to be able to + /// route the output of the partition to its destination, and also because + /// when [`Query`] is dropped it signals completion of query execution + query: Arc<Query>, + + /// 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 { + /// 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 + // + // We aren't using the wake count to synchronise other memory, and so can + // use relaxed memory ordering + let wake_count = self.waker.wake_count.load(Ordering::Relaxed); + + 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) { Review Comment: Not sure, my hope is that most operators will use a simpler serial trait, and only a few will need to implement the Pipeline trait. I think some sort of async interface is unavoidable, and using the standard rust types for this seems sensible, but hopefully we can hide most of it away 😃 -- 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]
