Handle NULL sentinel for nullable partition expressions (e.g. 'bucket(N, col)') (spiceai#10880)

* handle NULL sentinel for nullable partition expressions (e.g. 'bucket(N, col)')

* handle nullability in PartitionValue

* Update 'PartitionValue'

* Fix test compile errors after PartitionValue -> HashMap<String, Option<String>> change

* linting

* change test to handle additional partition

* fix bucket count

* fix: wrap PartitionValue test fixtures in Some() to match HashMap<String, Option<String>> type

* Fix 'test_distributed_acceleration_order_by_limit_pushdown'

* fix trunk?

* missing colon

* fix merge

* fmt

---------

Co-authored-by: Jeadie <jeadie@users.noreply.github.com>
Co-authored-by: Sergei Grebnov <sergei.grebnov@gmail.com>

Author: Jeadie

crates/runtime-cluster/src/executor_registry.rs

@@ -497,7 +497,7 @@ pub(crate) fn get_partitions_from_store(
     };
 
     // All required partitions (future: filter by query predicates)
-    let required_partitions: Vec<HashMap<String, String>> = table_metadata
+    let required_partitions: Vec<HashMap<String, Option<String>>> = table_metadata
         .partitions
         .iter()
         .map(|p| p.partition_value.clone())

crates/runtime-cluster/src/executor_selection.rs

@@ -143,7 +143,7 @@ mod tests {
 
     fn make_partition(key: &str, value: &str) -> PartitionValue {
         let mut p = HashMap::new();
-        p.insert(key.to_string(), value.to_string());
+        p.insert(key.to_string(), Some(value.to_string()));
         p
     }
 
@@ -279,12 +279,12 @@ mod tests {
     #[test]
     fn test_select_executors_composite_partitions() {
         let mut part_a = HashMap::new();
-        part_a.insert("region".to_string(), "us-east".to_string());
-        part_a.insert("date".to_string(), "2024-01-01".to_string());
+        part_a.insert("region".to_string(), Some("us-east".to_string()));
+        part_a.insert("date".to_string(), Some("2024-01-01".to_string()));
 
         let mut part_b = HashMap::new();
-        part_b.insert("region".to_string(), "us-west".to_string());
-        part_b.insert("date".to_string(), "2024-01-01".to_string());
+        part_b.insert("region".to_string(), Some("us-west".to_string()));
+        part_b.insert("date".to_string(), Some("2024-01-01".to_string()));
 
         let required = vec![part_a.clone(), part_b.clone()];
 

crates/runtime-cluster/src/metadata.rs

@@ -39,7 +39,7 @@ use crate::context::PartitionExprResolver;
 /// {"date": "2024-01-01", "region": "us-west"}
 /// {"date": "2024-01-02", "region": "us-east"}
 /// ```
-pub type PartitionValue = HashMap<String, String>;
+pub type PartitionValue = HashMap<String, Option<String>>;
 
 /// Metadata for a single partition of an accelerated table
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
@@ -56,7 +56,7 @@ pub struct PartitionMetadata {
 
 impl PartitionMetadata {
     #[must_use]
-    pub fn new(partition_value: HashMap<String, String>) -> Self {
+    pub fn new(partition_value: HashMap<String, Option<String>>) -> Self {
         Self {
             partition_value,
             assigned_executors: Vec::new(),
@@ -99,7 +99,10 @@ pub async fn partition_value_to_bytes(
     let mut expr: Option<Expr> = None;
     for (partition_expr, val) in p {
         let partition_by = resolver.try_parse_expr(tbl, &partition_expr).await?;
-        let e = partition_by.eq(lit(val));
+        let e = match val {
+            None => partition_by.is_null(),
+            Some(v) => partition_by.eq(lit(v)),
+        };
         expr = match expr {
             Some(existing) => Some(existing.and(e)),
             None => Some(e),
@@ -179,17 +182,21 @@ impl TablePartitionMetadata {
             //   key1 = val1 AND key2 = val2 AND ...
             let partition_predicate = partition_value
                 .iter()
-                .map(|(proj, lit)| {
-                    // Ensure lit is same type as proj
+                .map(|(proj, val)| {
                     let col = ctx.parse_sql_expr(proj, &df_schema)?;
+                    let Some(val) = val else {
+                        // NULL partition values need IS NULL, not = NULL
+                        // (SQL: `col = NULL` is always UNKNOWN, never TRUE)
+                        return Ok(col.is_null());
+                    };
                     let col_type = col.get_type(&df_schema)?;
-                    let mut lit = ctx.parse_sql_expr(lit, &df_schema)?;
-                    if let Expr::Literal(ref s, None) = lit
+                    let mut lit_expr = ctx.parse_sql_expr(val, &df_schema)?;
+                    if let Expr::Literal(ref s, None) = lit_expr
                         && s.data_type() != col_type
                     {
-                        lit = lit.cast_to(&col_type, &df_schema)?;
+                        lit_expr = lit_expr.cast_to(&col_type, &df_schema)?;
                     }
-                    Ok(col.eq(lit))
+                    Ok(col.eq(lit_expr))
                 })
                 .collect::<Result<Vec<Expr>, DataFusionError>>()?
                 .into_iter()

crates/runtime-cluster/src/service.rs

@@ -455,12 +455,12 @@ impl PartitionService {
                 }
             };
 
-        let existing: HashSet<Vec<(String, String)>> = existing_partitions
+        let existing: HashSet<Vec<(String, Option<String>)>> = existing_partitions
             .iter()
             .map(|p| sorted_kv(&p.partition_value))
             .collect();
 
-        let source_set: HashSet<Vec<(String, String)>> =
+        let source_set: HashSet<Vec<(String, Option<String>)>> =
             source_partitions.iter().map(sorted_kv).collect();
 
         let new: Vec<PartitionValue> = source_partitions
@@ -585,7 +585,7 @@ fn resolved_equality(a: &TableReference, b: &TableReference) -> bool {
 }
 
 /// Sort a `PartitionValue` into a deterministic `Vec<(k, v)>` for equality comparisons.
-fn sorted_kv(p: &PartitionValue) -> Vec<(String, String)> {
+fn sorted_kv(p: &PartitionValue) -> Vec<(String, Option<String>)> {
     let mut v: Vec<_> = p.clone().into_iter().collect();
     v.sort();
     v
@@ -1176,7 +1176,7 @@ mod tests {
     }
 
     fn pv(key: &str, val: &str) -> PartitionValue {
-        HashMap::from([(key.to_string(), val.to_string())])
+        HashMap::from([(key.to_string(), Some(val.to_string()))])
     }
 
     async fn setup_table(store: &PartitionStore, table: &str, partitions: Vec<PartitionMetadata>) {

crates/runtime-cluster/src/store.rs

@@ -225,7 +225,7 @@ impl PartitionStore {
     pub async fn set_unassigned_partitions(
         &self,
         table: &TableReference,
-        partition_values: Vec<HashMap<String, String>>,
+        partition_values: Vec<HashMap<String, Option<String>>>,
         partition_expressions: Vec<String>,
     ) -> Result<()> {
         let key = normalized_table_name(table);
@@ -664,7 +664,7 @@ mod tests {
     }
 
     fn partition_value(key: &str, val: &str) -> PartitionValue {
-        HashMap::from([(key.to_string(), val.to_string())])
+        HashMap::from([(key.to_string(), Some(val.to_string()))])
     }
 
     #[tokio::test]
@@ -702,7 +702,7 @@ mod tests {
             metadata
                 .partitions
                 .iter()
-                .find(|p| p.partition_value.get("bucket(3, id)") == Some(&val.to_string()))
+                .find(|p| p.partition_value.get("bucket(3, id)") == Some(&Some(val.to_string())))
                 .expect("partition not found")
         };
 
@@ -786,14 +786,14 @@ mod tests {
         let east = metadata
             .partitions
             .iter()
-            .find(|p| p.partition_value.get("region") == Some(&"us-east".to_string()))
+            .find(|p| p.partition_value.get("region") == Some(&Some("us-east".to_string())))
             .expect("us-east");
         assert!(east.is_assigned_to("executor-1"));
 
         let west = metadata
             .partitions
             .iter()
-            .find(|p| p.partition_value.get("region") == Some(&"us-west".to_string()))
+            .find(|p| p.partition_value.get("region") == Some(&Some("us-west".to_string())))
             .expect("us-west");
         assert!(west.is_assigned_to("executor-2"));
     }
@@ -830,12 +830,12 @@ mod tests {
             .await
             .expect("should initialize");
 
-        let pv = HashMap::from([("region".to_string(), "us-east-1".to_string())]);
+        let pv = HashMap::from([("region".to_string(), Some("us-east-1".to_string()))]);
         pm.set_unassigned_partitions(&source, vec![pv], vec![])
             .await
             .expect("should set partitions");
         let partition_value: PartitionValue =
-            HashMap::from([("region".to_string(), "us-east-1".to_string())]);
+            HashMap::from([("region".to_string(), Some("us-east-1".to_string()))]);
         pm.assign_partition(&source, &partition_value, "executor-1")
             .await
             .expect("should assign");
@@ -890,7 +890,7 @@ mod tests {
             .await
             .expect("should initialize target");
 
-        let pv = HashMap::from([("region".to_string(), "eu-west-1".to_string())]);
+        let pv = HashMap::from([("region".to_string(), Some("eu-west-1".to_string()))]);
         pm.set_unassigned_partitions(&source, vec![pv], vec![])
             .await
             .expect("should set partitions");
@@ -959,7 +959,7 @@ mod tests {
             .await
             .expect("should initialize");
 
-        let pv = HashMap::from([("org_id".to_string(), "test_org_name".to_string())]);
+        let pv = HashMap::from([("org_id".to_string(), Some("test_org_name".to_string()))]);
         pm.set_unassigned_partitions(&bare, vec![pv], vec![])
             .await
             .expect("should set partitions");
@@ -973,7 +973,7 @@ mod tests {
         assert_eq!(meta_via_full.partitions.len(), 1);
 
         let partition_value: PartitionValue =
-            HashMap::from([("org_id".to_string(), "test_org_name".to_string())]);
+            HashMap::from([("org_id".to_string(), Some("test_org_name".to_string()))]);
         pm.assign_partition(&full, &partition_value, "executor-1")
             .await
             .expect("should assign via fully qualified ref");
@@ -999,7 +999,7 @@ mod tests {
         acc.initialize_metadata(&table, vec!["region".to_string()])
             .await
             .expect("init");
-        let pv = HashMap::from([("region".to_string(), "us-east-1".to_string())]);
+        let pv = HashMap::from([("region".to_string(), Some("us-east-1".to_string()))]);
         acc.set_unassigned_partitions(&table, vec![pv.clone()], vec![])
             .await
             .expect("set");
@@ -1018,8 +1018,8 @@ mod tests {
         pm.initialize_metadata(&table, vec!["region".to_string()])
             .await
             .expect("init");
-        let values: Vec<HashMap<String, String>> = (0..50)
-            .map(|i| HashMap::from([("region".to_string(), format!("r-{i}"))]))
+        let values: Vec<HashMap<String, Option<String>>> = (0..50)
+            .map(|i| HashMap::from([("region".to_string(), Some(format!("r-{i}")))]))
             .collect();
         pm.set_unassigned_partitions(&table, values.clone(), vec![])
             .await

crates/runtime-cluster/src/write_through.rs

@@ -454,7 +454,14 @@ async fn route_batch_and_assign_unseen(
             let partition_value: PartitionValue = partition_expr_keys
                 .iter()
                 .zip(scalar_values.iter())
-                .map(|(expr_key, scalar)| (expr_key.clone(), scalar_to_sql_literal(scalar)))
+                .map(|(expr_key, scalar)| {
+                    let val = if scalar.is_null() {
+                        None
+                    } else {
+                        Some(scalar_to_sql_literal(scalar))
+                    };
+                    (expr_key.clone(), val)
+                })
                 .collect();
             Some((scalar_values, partition_value, sub_batch))
         })
@@ -510,7 +517,13 @@ async fn route_batch_and_assign_unseen(
             let new_pred = partition_phys_exprs
                 .iter()
                 .zip(scalar_values.iter())
-                .map(|((logical_expr, _), scalar)| logical_expr.clone().eq(lit(scalar.clone())))
+                .map(|((logical_expr, _), scalar)| {
+                    if scalar.is_null() {
+                        logical_expr.clone().is_null()
+                    } else {
+                        logical_expr.clone().eq(lit(scalar.clone()))
+                    }
+                })
                 .reduce(Expr::and);
 
             if let Some(new_pred) = new_pred {

crates/runtime-datafusion/src/analyzer_rule/partitioned_table_scan_rewrite.rs

@@ -50,7 +50,7 @@ use parking_lot::RwLock;
 /// {"date": "2024-01-01", "region": "us-west"}
 /// {"date": "2024-01-02", "region": "us-east"}
 /// ```
-pub type PartitionValue = HashMap<String, String>;
+pub type PartitionValue = HashMap<String, Option<String>>;
 
 /// Define how to get partitions for a given table, and how they are partitioned.
 pub trait TablePartitionProvider: Send + Sync + Debug {
@@ -378,16 +378,19 @@ fn partition_value_to_expr(
 ) -> Result<Option<Expr>, DataFusionError> {
     let mut expr: Option<Expr> = None;
     for (partition_expr_str, val) in pv {
-        let new_expr = state
+        let col_expr = state
             .upgrade()
             .ok_or_else(|| {
                 DataFusionError::Plan(
                     "SessionState has been dropped, cannot parse partition expression".to_string(),
                 )
             })?
             .read()
-            .create_logical_expr(partition_expr_str, df_schema)?
-            .eq(lit(val.clone()));
+            .create_logical_expr(partition_expr_str, df_schema)?;
+        let new_expr = match val {
+            None => col_expr.is_null(),
+            Some(v) => col_expr.eq(lit(v.clone())),
+        };
         expr = match expr {
             Some(existing) => Some(existing.and(new_expr)),
             None => Some(new_expr),
@@ -426,14 +429,14 @@ mod tests {
                     p1,
                     vec![HashMap::from([(
                         "partition_id".to_string(),
-                        "0".to_string(),
+                        Some("0".to_string()),
                     )])],
                 ),
                 (
                     p2,
                     vec![HashMap::from([(
                         "partition_id".to_string(),
-                        "1".to_string(),
+                        Some("1".to_string()),
                     )])],
                 ),
             ]