Rebase upstream

.gitignore

@@ -23,3 +23,6 @@ Cargo.lock
 # Added by cargo
 
 /target
+
+.idea
+.DS_Store

Cargo.toml

@@ -28,19 +28,20 @@ keywords = ["arrow", "arrow-rs", "datafusion"]
 rust-version = "1.80"
 
 [dependencies]
-arrow = "54"
-arrow-schema = "54"
+arrow = { version = "54.1.0" }
+arrow-schema = { version = "54.1.0" }
 async-trait = "0.1"
+chrono = "= 0.4.39"
 dashmap = "6"
-datafusion = "45"
-datafusion-common = "45"
-datafusion-expr = "45"
-datafusion-functions = "45"
-datafusion-functions-aggregate = "45"
-datafusion-optimizer = "45"
-datafusion-physical-expr = "45"
-datafusion-physical-plan = "45"
-datafusion-sql = "45"
+datafusion = { git = "https://github.com/polygon-io/arrow-datafusion", rev = "1c92803" }
+datafusion-common = { git = "https://github.com/polygon-io/arrow-datafusion", rev = "1c92803" }
+datafusion-expr = { git = "https://github.com/polygon-io/arrow-datafusion", rev = "1c92803" }
+datafusion-functions = { git = "https://github.com/polygon-io/arrow-datafusion", rev = "1c92803" }
+datafusion-functions-aggregate = { git = "https://github.com/polygon-io/arrow-datafusion", rev = "1c92803" }
+datafusion-optimizer = { git = "https://github.com/polygon-io/arrow-datafusion", rev = "1c92803" }
+datafusion-physical-expr = { git = "https://github.com/polygon-io/arrow-datafusion", rev = "1c92803" }
+datafusion-physical-plan = { git = "https://github.com/polygon-io/arrow-datafusion", rev = "1c92803" }
+datafusion-sql = { git = "https://github.com/polygon-io/arrow-datafusion", rev = "1c92803" }
 futures = "0.3"
 itertools = "0.13"
 log = "0.4"

deny.toml

@@ -23,6 +23,7 @@ allow = [
   "BSD-2-Clause",
   "BSD-3-Clause",
   "CC0-1.0",
-  "Unicode-3.0"
+  "Unicode-3.0",
+  "Zlib"
 ]
 version = 2

src/lib.rs

@@ -43,3 +43,18 @@ pub mod materialized;
 
 /// An implementation of Query Rewriting, an optimization that rewrites queries to make use of materialized views.
 pub mod rewrite;
+
+/// Configuration options for materialized view related features.
+#[derive(Debug, Clone)]
+pub struct MaterializedConfig {
+    /// Whether or not query rewriting should exploit this materialized view.
+    pub use_in_query_rewrite: bool,
+}
+
+impl Default for MaterializedConfig {
+    fn default() -> Self {
+        Self {
+            use_in_query_rewrite: true,
+        }
+    }
+}

src/materialized.rs

@@ -44,6 +44,8 @@ use datafusion::{
 use datafusion_expr::LogicalPlan;
 use itertools::Itertools;
 
+use crate::MaterializedConfig;
+
 /// The identifier of the column that [`RowMetadataSource`](row_metadata::RowMetadataSource) implementations should store row metadata in.
 pub const META_COLUMN: &str = "__meta";
 
@@ -102,6 +104,12 @@ pub fn cast_to_listing_table(table: &dyn TableProvider) -> Option<&dyn ListingTa
 pub trait Materialized: ListingTableLike {
     /// The query that defines this materialized view.
     fn query(&self) -> LogicalPlan;
+
+    /// Configuration to control materialized view related features.
+    /// By default, returns the default value for [`MaterializedConfig`]
+    fn config(&self) -> MaterializedConfig {
+        MaterializedConfig::default()
+    }
 }
 
 /// Register a [`Materialized`] implementation in this registry.

src/materialized/dependencies.rs

@@ -31,7 +31,6 @@ use datafusion_expr::{
     col, lit, utils::split_conjunction, Expr, LogicalPlan, LogicalPlanBuilder, TableScan,
 };
 use datafusion_functions::string::expr_fn::{concat, concat_ws};
-use datafusion_optimizer::{analyzer::expand_wildcard_rule::ExpandWildcardRule, AnalyzerRule};
 use datafusion_sql::TableReference;
 use itertools::{Either, Itertools};
 use std::{collections::HashSet, sync::Arc};
@@ -188,9 +187,18 @@ impl TableFunctionImpl for StaleFilesUdtf {
                     )?
                     .aggregate(
                         vec![
-                            // Trim the final path element
-                            regexp_replace(col("file_path"), lit(r"/[^/]*$"), lit("/"), None)
-                                .alias("existing_target"),
+                            // We want to omit the file name along with any "special" partitions
+                            // from the path before comparing it to the target partition. Special
+                            // partitions must be leaf most nodes and are designated by a leading
+                            // underscore. These are useful for adding additional information to a
+                            // filename without affecting partitioning or staleness checks.
+                            regexp_replace(
+                                col("file_path"),
+                                lit(r"(/_[^/=]+=[^/]+)*/[^/]*$"),
+                                lit("/"),
+                                None,
+                            )
+                            .alias("existing_target"),
                         ],
                         vec![max(col("last_modified")).alias("target_last_modified")],
                     )?
@@ -250,9 +258,6 @@ pub fn mv_dependencies_plan(
         .filter_map(|(i, f)| partition_cols.contains(f.name()).then_some(i))
         .collect();
 
-    // First expand all wildcards
-    let plan = ExpandWildcardRule {}.analyze(plan, config_options)?;
-
     let pruned_plan_with_source_files = if partition_cols.is_empty() {
         get_source_files_all_partitions(
             materialized_view,
@@ -1173,8 +1178,7 @@ mod test {
         }
 
         let cases = &[
-            TestCase {
-                name: "un-transformed partition column",
+            TestCase { name: "un-transformed partition column",
                 query_to_analyze:
                     "SELECT column1 AS partition_column, concat(column2, column3) AS some_value FROM t1",
                 table_name: "m1",
@@ -1206,6 +1210,38 @@ mod test {
                     "+--------------------------------+----------------------+-----------------------+----------+",
                 ],
             },
+            TestCase { name: "omit 'special' partition columns",
+                query_to_analyze:
+                    "SELECT column1 AS partition_column, concat(column2, column3) AS some_value FROM t1",
+                table_name: "m1",
+                table_path: ListingTableUrl::parse("s3://m1/").unwrap(),
+                partition_cols: vec!["partition_column"],
+                file_extension: ".parquet",
+                expected_output: vec![
+                    "+--------------------------------+----------------------+---------------------+-------------------+--------------------------------------+----------------------+",
+                    "| target                         | source_table_catalog | source_table_schema | source_table_name | source_uri                           | source_last_modified |",
+                    "+--------------------------------+----------------------+---------------------+-------------------+--------------------------------------+----------------------+",
+                    "| s3://m1/partition_column=2021/ | datafusion           | test                | t1                | s3://t1/column1=2021/data.01.parquet | 2023-07-11T16:29:26  |",
+                    "| s3://m1/partition_column=2022/ | datafusion           | test                | t1                | s3://t1/column1=2022/data.01.parquet | 2023-07-11T16:45:22  |",
+                    "| s3://m1/partition_column=2023/ | datafusion           | test                | t1                | s3://t1/column1=2023/data.01.parquet | 2023-07-11T16:45:44  |",
+                    "+--------------------------------+----------------------+---------------------+-------------------+--------------------------------------+----------------------+",
+                ],
+                // second file is old
+                file_metadata: "
+                    ('datafusion', 'test', 'm1', 's3://m1/partition_column=2021/_v=123/data.01.parquet', '2023-07-12T16:00:00Z', 0),
+                    ('datafusion', 'test', 'm1', 's3://m1/partition_column=2022/_v=123/data.01.parquet', '2023-07-10T16:00:00Z', 0),
+                    ('datafusion', 'test', 'm1', 's3://m1/partition_column=2023/_v=123/data.01.parquet', '2023-07-12T16:00:00Z', 0)
+                ",
+                expected_stale_files_output: vec![
+                    "+--------------------------------+----------------------+-----------------------+----------+",
+                    "| target                         | target_last_modified | sources_last_modified | is_stale |",
+                    "+--------------------------------+----------------------+-----------------------+----------+",
+                    "| s3://m1/partition_column=2021/ | 2023-07-12T16:00:00  | 2023-07-11T16:29:26   | false    |",
+                    "| s3://m1/partition_column=2022/ | 2023-07-10T16:00:00  | 2023-07-11T16:45:22   | true     |",
+                    "| s3://m1/partition_column=2023/ | 2023-07-12T16:00:00  | 2023-07-11T16:45:44   | false    |",
+                    "+--------------------------------+----------------------+-----------------------+----------+",
+                ],
+            },
             TestCase {
                 name: "transform year/month/day partition into timestamp partition",
                 query_to_analyze: "