Add some examples

Author: gabotechs

Cargo.lock

@@ -48,6 +48,7 @@ integration = [
 ]
 
 [dev-dependencies]
+structopt = "0.3"
 insta = { version = "1.43.1", features = ["filters"] }
 tpchgen = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "c8d823432528eed4f70fca5a1296a66c68a389a8" }
 tpchgen-arrow = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "c8d823432528eed4f70fca5a1296a66c68a389a8" }

Cargo.toml

@@ -0,0 +1,106 @@
+use arrow::util::pretty::pretty_format_batches;
+use async_trait::async_trait;
+use dashmap::{DashMap, Entry};
+use datafusion::common::DataFusionError;
+use datafusion::execution::SessionStateBuilder;
+use datafusion::physical_plan::displayable;
+use datafusion::prelude::{ParquetReadOptions, SessionContext};
+use datafusion_distributed::{
+    BoxCloneSyncChannel, ChannelResolver, DistributedExt, DistributedPhysicalOptimizerRule,
+};
+use futures::TryStreamExt;
+use std::error::Error;
+use std::sync::Arc;
+use structopt::StructOpt;
+use tonic::transport::Channel;
+use url::Url;
+
+#[derive(StructOpt)]
+#[structopt(name = "run", about = "A localhost Distributed DataFusion runner")]
+struct Args {
+    #[structopt()]
+    query: String,
+
+    // --cluster-ports 8080,8081,8082
+    #[structopt(long = "cluster-ports", use_delimiter = true)]
+    cluster_ports: Vec<u16>,
+
+    #[structopt(long)]
+    explain: bool,
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn Error>> {
+    let args = Args::from_args();
+
+    let localhost_resolver = LocalhostChannelResolver {
+        ports: args.cluster_ports,
+        cached: DashMap::new(),
+    };
+
+    let state = SessionStateBuilder::new()
+        .with_default_features()
+        .with_distributed_channel_resolver(localhost_resolver)
+        .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule::new()))
+        .build();
+
+    let ctx = SessionContext::from(state);
+
+    ctx.register_parquet(
+        "flights_1m",
+        "testdata/flights-1m.parquet",
+        ParquetReadOptions::default(),
+    )
+    .await?;
+
+    ctx.register_parquet(
+        "weather",
+        "testdata/weather.parquet",
+        ParquetReadOptions::default(),
+    )
+    .await?;
+
+    let df = ctx.sql(&args.query).await?;
+    if args.explain {
+        let plan = df.create_physical_plan().await?;
+        let display = displayable(plan.as_ref()).indent(true).to_string();
+        println!("{display}");
+    } else {
+        let stream = df.execute_stream().await?;
+        let batches = stream.try_collect::<Vec<_>>().await?;
+        let formatted = pretty_format_batches(&batches)?;
+        println!("{formatted}");
+    }
+    Ok(())
+}
+
+#[derive(Clone)]
+struct LocalhostChannelResolver {
+    ports: Vec<u16>,
+    cached: DashMap<Url, BoxCloneSyncChannel>,
+}
+
+#[async_trait]
+impl ChannelResolver for LocalhostChannelResolver {
+    fn get_urls(&self) -> Result<Vec<Url>, DataFusionError> {
+        Ok(self
+            .ports
+            .iter()
+            .map(|port| Url::parse(&format!("http://localhost:{port}")).unwrap())
+            .collect())
+    }
+
+    async fn get_channel_for_url(&self, url: &Url) -> Result<BoxCloneSyncChannel, DataFusionError> {
+        match self.cached.entry(url.clone()) {
+            Entry::Occupied(v) => Ok(v.get().clone()),
+            Entry::Vacant(v) => {
+                let channel = Channel::from_shared(url.to_string())
+                    .unwrap()
+                    .connect_lazy();
+                let channel = BoxCloneSyncChannel::new(channel);
+                v.insert(channel.clone());
+                Ok(channel)
+            }
+        }
+    }
+}

examples/localhost_run.rs

@@ -0,0 +1,84 @@
+use arrow_flight::flight_service_server::FlightServiceServer;
+use async_trait::async_trait;
+use dashmap::{DashMap, Entry};
+use datafusion::common::DataFusionError;
+use datafusion::execution::SessionStateBuilder;
+use datafusion_distributed::{
+    ArrowFlightEndpoint, BoxCloneSyncChannel, ChannelResolver, DistributedExt,
+    DistributedSessionBuilderContext,
+};
+use std::error::Error;
+use std::net::{IpAddr, Ipv4Addr, SocketAddr};
+use structopt::StructOpt;
+use tonic::transport::{Channel, Server};
+use url::Url;
+
+#[derive(StructOpt)]
+#[structopt(name = "localhost_worker", about = "A localhost DataFusion worker")]
+struct Args {
+    #[structopt(default_value = "8080")]
+    port: u16,
+
+    // --cluster-ports 8080,8081,8082
+    #[structopt(long = "cluster-ports", use_delimiter = true)]
+    cluster_ports: Vec<u16>,
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn Error>> {
+    let args = Args::from_args();
+
+    let localhost_resolver = LocalhostChannelResolver {
+        ports: args.cluster_ports,
+        cached: DashMap::new(),
+    };
+
+    let endpoint = ArrowFlightEndpoint::new(move |ctx: DistributedSessionBuilderContext| {
+        let local_host_resolver = localhost_resolver.clone();
+        async move {
+            Ok(SessionStateBuilder::new()
+                .with_runtime_env(ctx.runtime_env)
+                .with_distributed_channel_resolver(local_host_resolver)
+                .with_default_features()
+                .build())
+        }
+    });
+
+    Server::builder()
+        .add_service(FlightServiceServer::new(endpoint))
+        .serve(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), args.port))
+        .await?;
+
+    Ok(())
+}
+
+#[derive(Clone)]
+struct LocalhostChannelResolver {
+    ports: Vec<u16>,
+    cached: DashMap<Url, BoxCloneSyncChannel>,
+}
+
+#[async_trait]
+impl ChannelResolver for LocalhostChannelResolver {
+    fn get_urls(&self) -> Result<Vec<Url>, DataFusionError> {
+        Ok(self
+            .ports
+            .iter()
+            .map(|port| Url::parse(&format!("http://localhost:{port}")).unwrap())
+            .collect())
+    }
+
+    async fn get_channel_for_url(&self, url: &Url) -> Result<BoxCloneSyncChannel, DataFusionError> {
+        match self.cached.entry(url.clone()) {
+            Entry::Occupied(v) => Ok(v.get().clone()),
+            Entry::Vacant(v) => {
+                let channel = Channel::from_shared(url.to_string())
+                    .unwrap()
+                    .connect_lazy();
+                let channel = BoxCloneSyncChannel::new(channel);
+                v.insert(channel.clone());
+                Ok(channel)
+            }
+        }
+    }
+}