[FEAT]: expr simplifier

daft/delta_lake/delta_lake_scan.py

@@ -111,6 +111,9 @@ def __init__(
     def schema(self) -> Schema:
         return self._schema
 
+    def name(self) -> str:
+        return "DeltaLakeScanOperator"
+
     def display_name(self) -> str:
         return f"DeltaLakeScanOperator({self._table.metadata().name})"
 

daft/hudi/hudi_scan.py

@@ -37,6 +37,9 @@ def __init__(self, table_uri: str, storage_config: StorageConfig) -> None:
     def schema(self) -> Schema:
         return self._schema
 
+    def name(self) -> str:
+        return "HudiScanOperator"
+
     def display_name(self) -> str:
         return f"HudiScanOperator({self._table.props.name})"
 

daft/iceberg/iceberg_scan.py

@@ -115,6 +115,9 @@ def __init__(self, iceberg_table: Table, snapshot_id: int | None, storage_config
     def schema(self) -> Schema:
         return self._schema
 
+    def name(self) -> str:
+        return "IcebergScanOperator"
+
     def display_name(self) -> str:
         return f"IcebergScanOperator({'.'.join(self._table.name())})"
 

daft/io/_generator.py

@@ -80,6 +80,9 @@ def __init__(
         self._generators = generators
         self._schema = schema
 
+    def name(self) -> str:
+        return self.display_name()
+
     def display_name(self) -> str:
         return "GeneratorScanOperator"
 

daft/io/_lance.py

@@ -69,6 +69,9 @@ class LanceDBScanOperator(ScanOperator):
     def __init__(self, ds: "lance.LanceDataset"):
         self._ds = ds
 
+    def name(self) -> str:
+        return "LanceDBScanOperator"
+
     def display_name(self) -> str:
         return f"LanceDBScanOperator({self._ds.uri})"
 

daft/sql/sql_scan.py

@@ -69,6 +69,9 @@ def __init__(
     def schema(self) -> Schema:
         return self._schema
 
+    def name(self) -> str:
+        return "SQLScanOperator"
+
     def display_name(self) -> str:
         return f"SQLScanOperator(sql={self.sql}, conn={self.conn})"
 

src/common/scan-info/src/scan_operator.rs

@@ -10,6 +10,8 @@ use daft_schema::schema::SchemaRef;
 use crate::{PartitionField, Pushdowns, ScanTaskLikeRef};
 
 pub trait ScanOperator: Send + Sync + Debug {
+    fn name(&self) -> &str;
+
     fn schema(&self) -> SchemaRef;
     fn partitioning_keys(&self) -> &[PartitionField];
     fn file_path_column(&self) -> Option<&str>;

src/common/scan-info/src/test/mod.rs

@@ -97,6 +97,9 @@ Pushdowns: {pushdowns}
 }
 
 impl ScanOperator for DummyScanOperator {
+    fn name(&self) -> &'static str {
+        "dummy"
+    }
     fn schema(&self) -> SchemaRef {
         self.schema.clone()
     }

src/daft-logical-plan/src/optimization/optimizer.rs

@@ -7,7 +7,8 @@ use super::{
     logical_plan_tracker::LogicalPlanTracker,
     rules::{
         DropRepartition, EliminateCrossJoin, EnrichWithStats, LiftProjectFromAgg, MaterializeScans,
-        OptimizerRule, PushDownFilter, PushDownLimit, PushDownProjection, SplitActorPoolProjects,
+        OptimizerRule, PushDownFilter, PushDownLimit, PushDownProjection, SimplifyExpressionsRule,
+        SplitActorPoolProjects,
     },
 };
 use crate::LogicalPlan;
@@ -97,6 +98,11 @@ impl Optimizer {
                 ],
                 RuleExecutionStrategy::Once,
             ),
+            // we want to simplify expressions first to make the rest of the rules easier
+            RuleBatch::new(
+                vec![Box::new(SimplifyExpressionsRule::new())],
+                RuleExecutionStrategy::FixedPoint(Some(3)),
+            ),
             // --- Bulk of our rules ---
             RuleBatch::new(
                 vec![
@@ -129,6 +135,11 @@ impl Optimizer {
                 vec![Box::new(EnrichWithStats::new())],
                 RuleExecutionStrategy::Once,
             ),
+            // try to simplify expressions again as other rules could introduce new exprs
+            RuleBatch::new(
+                vec![Box::new(SimplifyExpressionsRule::new())],
+                RuleExecutionStrategy::FixedPoint(Some(3)),
+            ),
         ];
 
         Self::with_rule_batches(rule_batches, config)

src/daft-logical-plan/src/optimization/rules/mod.rs

@@ -7,6 +7,7 @@ mod push_down_filter;
 mod push_down_limit;
 mod push_down_projection;
 mod rule;
+mod simplify_expressions;
 mod split_actor_pool_projects;
 
 pub use drop_repartition::DropRepartition;
@@ -18,4 +19,5 @@ pub use push_down_filter::PushDownFilter;
 pub use push_down_limit::PushDownLimit;
 pub use push_down_projection::PushDownProjection;
 pub use rule::OptimizerRule;
+pub use simplify_expressions::SimplifyExpressionsRule;
 pub use split_actor_pool_projects::SplitActorPoolProjects;