Skip to content

Commit

Permalink
[datafusion-cli] support for flight tables
Browse files Browse the repository at this point in the history
  • Loading branch information
ccciudatu committed Aug 12, 2024
1 parent 4fe546d commit 820021b
Show file tree
Hide file tree
Showing 4 changed files with 67 additions and 0 deletions.
1 change: 1 addition & 0 deletions datafusion-cli/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ datafusion = { path = "../datafusion/core", version = "41.0.0", features = [
"regex_expressions",
"unicode_expressions",
"compression",
"flight",
] }
dirs = "4.0.0"
env_logger = "0.9"
Expand Down
3 changes: 3 additions & 0 deletions datafusion-cli/src/exec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ use datafusion::sql::sqlparser;
use rustyline::error::ReadlineError;
use rustyline::Editor;
use tokio::signal;
use datafusion::datasource::flight::config::FlightOptions;

/// run and execute SQL statements and commands, against a context with the given print options
pub async fn exec_from_commands(
Expand Down Expand Up @@ -383,6 +384,8 @@ pub(crate) async fn register_object_store_and_config_extensions(
let mut table_options = ctx.session_state().default_table_options().clone();
if let Some(format) = format {
table_options.set_config_format(format);
} else {
table_options.extensions.insert(FlightOptions::default())
}
table_options.alter_with_string_hash_map(options)?;

Expand Down
62 changes: 62 additions & 0 deletions datafusion/core/src/datasource/flight/config.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
// 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.

//! Only meant for datafusion-cli to accept the `flight` namespace

use datafusion_common::config::{ConfigEntry, ConfigExtension, ExtensionOptions};
use std::any::Any;
use std::collections::HashMap;

/// Collects the config entries
#[derive(Default, Debug, Clone)]
pub struct FlightOptions {
inner: HashMap<String, String>,
}

impl ConfigExtension for FlightOptions {
const PREFIX: &'static str = "flight";
}

impl ExtensionOptions for FlightOptions {
fn as_any(&self) -> &dyn Any {
self
}

fn as_any_mut(&mut self) -> &mut dyn Any {
self
}

fn cloned(&self) -> Box<dyn ExtensionOptions> {
Box::new(self.clone())
}

fn set(&mut self, key: &str, value: &str) -> datafusion_common::Result<()> {
self.inner.insert(key.into(), value.into());
Ok(())
}

fn entries(&self) -> Vec<ConfigEntry> {
self.inner
.iter()
.map(|(key, value)| ConfigEntry {
key: key.to_owned(),
value: Some(value.to_owned()).filter(|s| !s.is_empty()),
description: "",
})
.collect()
}
}
1 change: 1 addition & 0 deletions datafusion/core/src/datasource/flight/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ use datafusion_physical_plan::{ExecutionMode, ExecutionPlan, PlanProperties};

use crate::datasource::physical_plan::FlightExec;

pub mod config;
pub mod sql;

/// Generic Arrow Flight data source. Requires a [FlightDriver] that allows implementors
Expand Down

0 comments on commit 820021b

Please sign in to comment.