Skip to content

Commit 0a650a0

Browse files
r1balamb
andauthored
feat: support temporary views in DataFrameTableProvider (#18158)
## Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123. --> - Closes #18026 ## Rationale for this change <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. --> This makes it possible to support temporary views in datafusion-python without code duplication. Ref: apache/datafusion-python#1267 ## What changes are included in this PR? <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> - Add new public function `DataFrame::into_temporary_view` - Update `DataFrameTableProvider` with a new member that determines the `table_type` - Add a test ## Are these changes tested? <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code If tests are not included in your PR, please explain why (for example, are they covered by existing tests)? --> Yes, see added test `register_temporary_table` ## Are there any user-facing changes? <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. --> <!-- If there are any breaking changes to public APIs, please add the `api change` label. --> Yes, there is a new public function `DataFrame::into_temporary_view` --------- Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
1 parent 606dfd2 commit 0a650a0

File tree

2 files changed

+33
-3
lines changed

2 files changed

+33
-3
lines changed

datafusion/core/src/dataframe/mod.rs

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1654,7 +1654,19 @@ impl DataFrame {
16541654
/// Note: This discards the [`SessionState`] associated with this
16551655
/// [`DataFrame`] in favour of the one passed to [`TableProvider::scan`]
16561656
pub fn into_view(self) -> Arc<dyn TableProvider> {
1657-
Arc::new(DataFrameTableProvider { plan: self.plan })
1657+
Arc::new(DataFrameTableProvider {
1658+
plan: self.plan,
1659+
table_type: TableType::Temporary,
1660+
})
1661+
}
1662+
1663+
/// See [`Self::into_view`]. The returned [`TableProvider`] will
1664+
/// create a transient table.
1665+
pub fn into_temporary_view(self) -> Arc<dyn TableProvider> {
1666+
Arc::new(DataFrameTableProvider {
1667+
plan: self.plan,
1668+
table_type: TableType::Temporary,
1669+
})
16581670
}
16591671

16601672
/// Return a DataFrame with the explanation of its plan so far.
@@ -2524,6 +2536,7 @@ macro_rules! dataframe {
25242536
#[derive(Debug)]
25252537
struct DataFrameTableProvider {
25262538
plan: LogicalPlan,
2539+
table_type: TableType,
25272540
}
25282541

25292542
#[async_trait]
@@ -2549,7 +2562,7 @@ impl TableProvider for DataFrameTableProvider {
25492562
}
25502563

25512564
fn table_type(&self) -> TableType {
2552-
TableType::View
2565+
self.table_type
25532566
}
25542567

25552568
async fn scan(

datafusion/core/tests/dataframe/mod.rs

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ use datafusion_expr::var_provider::{VarProvider, VarType};
7777
use datafusion_expr::{
7878
cast, col, create_udf, exists, in_subquery, lit, out_ref_col, placeholder,
7979
scalar_subquery, when, wildcard, Expr, ExprFunctionExt, ExprSchemable, LogicalPlan,
80-
LogicalPlanBuilder, ScalarFunctionImplementation, SortExpr, WindowFrame,
80+
LogicalPlanBuilder, ScalarFunctionImplementation, SortExpr, TableType, WindowFrame,
8181
WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition,
8282
};
8383
use datafusion_physical_expr::aggregate::AggregateExprBuilder;
@@ -1577,6 +1577,23 @@ async fn register_table() -> Result<()> {
15771577
Ok(())
15781578
}
15791579

1580+
#[tokio::test]
1581+
async fn register_temporary_table() -> Result<()> {
1582+
let df = test_table().await?.select_columns(&["c1", "c12"])?;
1583+
let ctx = SessionContext::new();
1584+
let df_impl = DataFrame::new(ctx.state(), df.logical_plan().clone());
1585+
1586+
let df_table_provider = df_impl.clone().into_temporary_view();
1587+
1588+
// check that we set the correct table_type
1589+
assert_eq!(df_table_provider.table_type(), TableType::Temporary);
1590+
1591+
// check that we can register a dataframe as a temporary table
1592+
ctx.register_table("test_table", df_table_provider)?;
1593+
1594+
Ok(())
1595+
}
1596+
15801597
/// Compare the formatted string representation of two plans for equality
15811598
fn assert_same_plan(plan1: &LogicalPlan, plan2: &LogicalPlan) {
15821599
assert_eq!(format!("{plan1:?}"), format!("{plan2:?}"));

0 commit comments

Comments
 (0)