feat: Upstream Spice AI changes to `datafusion-federation` #51

phillipleblanc · 2024-08-29T12:27:42Z

This PR contains all of the changes that the spiceai project made to datafusion-federation to fix bugs and add features to make federation work better for our use-cases.

The main two features that were added:

The name of the table registered in DataFusion doesn't need to match the name of the table in the remote database. This was an important requirement for our product, and we did a lot of work to get this working correctly.
The Arrow schema returned by the provider may not exactly match what DataFusion expects it to be, so we add a schema casting layer to make sure. Where possible this leverages the native Arrow cast kernels, but there are a few special cases that are handled.

* testing * testing2 * Rewrite all table scans to use federated table name * Wrap the table scan in a subquery alias

* Rewrite subquery table scans to point to remote table * Add tracing * more tracing * more debug logging * Handle subquery in binary expressions

* Fix the table scan rewrite to properly rewrite column relations * Handle more expressions for table scan rewrites

Update Support for multiple table occurrences in single expression Support for multiple different tables in single aggregation expression Document rewrite_column_name_in_expr logic Update Fix spelling

…uery when executing (#10) * keep original qualified_field when wrap_projections * only mutate the table reference when sending plan to execute * formatting * rename method

* do not over eager rewrite column when col relation is there * add more tests * Update sources/sql/src/lib.rs Co-authored-by: Phillip LeBlanc <[email protected]> --------- Co-authored-by: Phillip LeBlanc <[email protected]>

* Add List types parsing support to schema cast * Update to use arrow-json for lists parsing

* Add List types parsing support to schema cast * Update to use arrow-json for lists parsing * Add interval cast given original schema * Fix test --------- Co-authored-by: sgrebnov <[email protected]> Co-authored-by: Phillip LeBlanc <[email protected]>

* Add Structs parsing support to schema cast * Optimize memory space

Co-authored-by: Suriya Kandaswamy <[email protected]> Co-authored-by: Suriya Kandaswamy <[email protected]>

…r from datafusion-flight-sql-server"

…onOptimizerRule

…erly (datafusion-contrib#49)

…on-migrate-spiceai-41

…ceai-41 rebase spiceai-41 branch

phillipleblanc · 2024-08-29T12:30:52Z

Cargo.toml

 edition = "2021"
-license = "MIT"
+license = "Apache-2.0"


I felt this would be better since DataFusion is an Apache project?

phillipleblanc · 2024-08-29T12:32:35Z

datafusion-federation/src/lib.rs

 pub use table_provider::{FederatedTableProviderAdaptor, FederatedTableSource};

+// TODO clean up this
+// TODO move schema_cast.rs to schema_cast directory
+pub mod schema_cast;


This schema_cast module is useful when the schema registered in DataFusion and the schema returned by the underlying provider don't match exactly. The schema_cast provides a mechanism to keep the schema aligned with what DataFusion expects.

I put TODO there as reminder to clean up the module by moving schema_cast.rs to schema_cast/mod.rs

phillipleblanc · 2024-08-29T12:33:49Z

datafusion-federation/src/sql/mod.rs

@@ -100,6 +113,481 @@ impl OptimizerRule for SQLFederationOptimizerRule {
        false
    }
 }
+
+/// Rewrite table scans to use the original federated table name.
+fn rewrite_table_scans(


This is needed to support the table name in the remote source being different from the table name registered in DataFusion. This required a lot of trial and error to get right, and we have a good integration test coverage of this in spiceai.

datafusion-federation/src/sql/mod.rs

backkem · 2024-09-02T10:36:16Z

The rewriting does seem like good generic functionality so I'm fine to accept it in.

phillipleblanc and others added 30 commits May 9, 2024 14:56

Update to DataFusion 37.1

a25a523

Allow TableScans of non-federated sources

a852cd3

.clone() -> Arc::clone(&

c81a766

Update Cargo.toml license to Apache-2.0

c6c9542

Merge branch 'main' into spiceai

4808d74

Upgrade to DataFusion 38

5d9b98f

Add fallback TableProvider to FederatedTableProviderAdaptor

3de50a9

Make connectorx optional

0fadae1

Rewrite all table scans to use federated table name (#4)

8ba03bd

* testing * testing2 * Rewrite all table scans to use federated table name * Wrap the table scan in a subquery alias

Rewrite subquery table scans to point to remote table (#5)

8fcda49

* Rewrite subquery table scans to point to remote table * Add tracing * more tracing * more debug logging * Handle subquery in binary expressions

Fix the fallback to the table provider for supports_filters_pushdown

1c77d8c

Use SpiceAI datafusion fork

b221f4e

Fix the table scan rewrite to properly rewrite column relations (#3)

f5d4797

* Fix the table scan rewrite to properly rewrite column relations * Handle more expressions for table scan rewrites

Handle table rewrites for aggregation functions

532ecd0

Update datafusion version

f5d20bb

Use Unparser Dialect for plan_to_sql

65d4be2

Use unparser Dialect

6d714cb

bump datafusion fork

9dc7404

Add layer to convert actual schema to expected schema

bf5fac6

Update datafusion-federation to datafusion 39 (#6)

4ea973e

Improve rewrite logic for expressions

e85aa96

Update Support for multiple table occurrences in single expression Support for multiple different tables in single aggregation expression Document rewrite_column_name_in_expr logic Update Fix spelling

update spiceai datafusion to include new unparsers logic (#8)

15d4bd9

Update datafusion to include time predicate fix (#9)

776930f

Remove table reference mutating during logical planning, converting q…

c3f696a

…uery when executing (#10) * keep original qualified_field when wrap_projections * only mutate the table reference when sending plan to execute * formatting * rename method

Update SpiceAI Datafusion fork

1becf0c

Upgrade to DF 40

77548d3

Bump datafusion-federation version to 0.1.5

f430401

Upgrade to DataFusion v41

9b4392f

Use Display for displaying federated plans

ae9da45

sgrebnov and others added 22 commits August 25, 2024 11:07

Add List types parsing support to schema cast (#13)

2cb3330

* Add List types parsing support to schema cast * Update to use arrow-json for lists parsing

Add Structs parsing support to schema cast (#16)

21f07be

* Add Structs parsing support to schema cast * Optimize memory space

feat: initial subquery support (#37)

1033cff

Co-authored-by: Suriya Kandaswamy <[email protected]> Co-authored-by: Suriya Kandaswamy <[email protected]>

delete redundant github workflows & commitlint

7dfb653

update deps && use datafusion 41

8968619

remove connectorx from sources/sql crate

5741fd6

flight-sql: fix up dependencie issues

82f3dfa

move sources/sql code to datafusion-federation crate and add sql feature

3e5996d

move sources/flight-sql to datafusion-flight-sql-server

6e6acd5

"create datafusion-flight-sql-table-provider crate & move the executo…

0cb719d

…r from datafusion-flight-sql-server"

override supports_rewrite for FederationOptimizerRule and SQLFederati…

66fb747

…onOptimizerRule

Improve project overview

5e7bba4

update README.md (datafusion-contrib#48)

b80203a

datafusion-fedeartion: Remove assert_eq macros and handle errors prop…

95d3416

…erly (datafusion-contrib#49)

Add example in datafusion-federation crate (datafusion-contrib#50)

f789187

fix rebase conflicts

49f8be0

Merge remote-tracking branch 'upstream/main' into datafusion-federati…

340a955

…on-migrate-spiceai-41

cargo fmt

ffdbe87

remove datafusion-federation/analyzer.rs

f9abf43

datafusion-federation: export FederatedPlanner

238f836

Merge pull request #18 from hozan23/datafusion-federation-migrate-spi…

ddbc880

…ceai-41 rebase spiceai-41 branch

phillipleblanc changed the title ~~Upstream Spice AI changes to datafusion-federation~~ feat: Upstream Spice AI changes to datafusion-federation Aug 29, 2024

phillipleblanc commented Aug 29, 2024

View reviewed changes

phillipleblanc marked this pull request as ready for review August 29, 2024 12:39

Update datafusion-federation/src/sql/mod.rs

957810c

backkem approved these changes Sep 2, 2024

View reviewed changes

hozan23 merged commit af09abe into datafusion-contrib:main Sep 2, 2024
7 checks passed

github-actions bot mentioned this pull request Nov 15, 2024

chore(datafusion-flight-sql-table-provider): release v0.3.1 #79

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Upstream Spice AI changes to `datafusion-federation` #51

feat: Upstream Spice AI changes to `datafusion-federation` #51

phillipleblanc commented Aug 29, 2024 •

edited

Loading

phillipleblanc Aug 29, 2024

phillipleblanc Aug 29, 2024

hozan23 Aug 29, 2024 •

edited

Loading

phillipleblanc Aug 29, 2024

backkem commented Sep 2, 2024

feat: Upstream Spice AI changes to datafusion-federation #51

feat: Upstream Spice AI changes to datafusion-federation #51

Conversation

phillipleblanc commented Aug 29, 2024 • edited Loading

phillipleblanc Aug 29, 2024

Choose a reason for hiding this comment

phillipleblanc Aug 29, 2024

Choose a reason for hiding this comment

hozan23 Aug 29, 2024 • edited Loading

Choose a reason for hiding this comment

phillipleblanc Aug 29, 2024

Choose a reason for hiding this comment

backkem commented Sep 2, 2024

feat: Upstream Spice AI changes to `datafusion-federation` #51

feat: Upstream Spice AI changes to `datafusion-federation` #51

phillipleblanc commented Aug 29, 2024 •

edited

Loading

hozan23 Aug 29, 2024 •

edited

Loading