diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs index 287a133273d8..98345d3f0325 100644 --- a/datafusion/core/src/dataframe/mod.rs +++ b/datafusion/core/src/dataframe/mod.rs @@ -1598,7 +1598,19 @@ impl DataFrame { /// Note: This discards the [`SessionState`] associated with this /// [`DataFrame`] in favour of the one passed to [`TableProvider::scan`] pub fn into_view(self) -> Arc { - Arc::new(DataFrameTableProvider { plan: self.plan }) + Arc::new(DataFrameTableProvider { + plan: self.plan, + df_table_type: DataFrameTableType::View, + }) + } + + /// See [`Self::into_view`]. The returned [`TableProvider`] will + /// create a transient table. + pub fn into_temporary_view(self) -> Arc { + Arc::new(DataFrameTableProvider { + plan: self.plan, + df_table_type: DataFrameTableType::TemporaryView, + }) } /// Return a DataFrame with the explanation of its plan so far. @@ -2423,9 +2435,25 @@ macro_rules! dataframe { }}; } +#[derive(Debug, Clone, Copy)] +enum DataFrameTableType { + View, + TemporaryView, +} + +impl From for TableType { + fn from(value: DataFrameTableType) -> TableType { + match value { + DataFrameTableType::View => TableType::View, + DataFrameTableType::TemporaryView => TableType::Temporary, + } + } +} + #[derive(Debug)] struct DataFrameTableProvider { plan: LogicalPlan, + df_table_type: DataFrameTableType, } #[async_trait] @@ -2451,7 +2479,7 @@ impl TableProvider for DataFrameTableProvider { } fn table_type(&self) -> TableType { - TableType::View + self.df_table_type.into() } async fn scan( diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index aa538f6dee81..d32d45d618c1 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -78,7 +78,7 @@ use datafusion_expr::var_provider::{VarProvider, VarType}; use datafusion_expr::{ cast, col, create_udf, exists, in_subquery, lit, out_ref_col, placeholder, scalar_subquery, when, wildcard, Expr, ExprFunctionExt, ExprSchemable, LogicalPlan, - LogicalPlanBuilder, ScalarFunctionImplementation, SortExpr, WindowFrame, + LogicalPlanBuilder, ScalarFunctionImplementation, SortExpr, TableType, WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, }; use datafusion_physical_expr::aggregate::AggregateExprBuilder; @@ -1578,6 +1578,23 @@ async fn register_table() -> Result<()> { Ok(()) } +#[tokio::test] +async fn register_temporary_table() -> Result<()> { + let df = test_table().await?.select_columns(&["c1", "c12"])?; + let ctx = SessionContext::new(); + let df_impl = DataFrame::new(ctx.state(), df.logical_plan().clone()); + + let df_table_provider = df_impl.clone().into_temporary_view(); + + // check that we set the correct table_type + assert_eq!(df_table_provider.table_type(), TableType::Temporary); + + // check that we can register a dataframe as a temporary table + ctx.register_table("test_table", df_table_provider)?; + + Ok(()) +} + /// Compare the formatted string representation of two plans for equality fn assert_same_plan(plan1: &LogicalPlan, plan2: &LogicalPlan) { assert_eq!(format!("{plan1:?}"), format!("{plan2:?}"));