merge main and resolve conflicts

Author: xx01cyx

optd-cost-model/src/common/mod.rs

@@ -1,4 +1,5 @@
 pub mod nodes;
 pub mod predicates;
+pub mod properties;
 pub mod types;
 pub mod values;

optd-cost-model/src/common/properties/attr_ref.rs

@@ -0,0 +1,245 @@
+use std::collections::HashSet;
+
+use crate::{common::types::TableId, utils::DisjointSets};
+
+pub type AttrRefs = Vec<AttrRef>;
+
+/// [`BaseTableAttrRef`] represents a reference to an attribute in a base table,
+/// i.e. a table existing in the catalog.
+#[derive(Clone, Debug, Default, Eq, Hash, PartialEq)]
+pub struct BaseTableAttrRef {
+    pub table_id: TableId,
+    pub attr_idx: u64,
+}
+
+/// [`AttrRef`] represents a reference to an attribute in a query.
+#[derive(Clone, Debug)]
+pub enum AttrRef {
+    /// Reference to a base table attribute.
+    BaseTableAttrRef(BaseTableAttrRef),
+    /// Reference to a derived attribute (e.g. t.v1 + t.v2).
+    /// TODO: Better representation of derived attributes.
+    Derived,
+}
+
+impl AttrRef {
+    pub fn base_table_attr_ref(table_id: TableId, attr_idx: u64) -> Self {
+        AttrRef::BaseTableAttrRef(BaseTableAttrRef { table_id, attr_idx })
+    }
+}
+
+impl From<BaseTableAttrRef> for AttrRef {
+    fn from(attr: BaseTableAttrRef) -> Self {
+        AttrRef::BaseTableAttrRef(attr)
+    }
+}
+
+/// [`EqPredicate`] represents an equality predicate between two attributes.
+#[derive(Clone, Debug, Eq, Hash, PartialEq)]
+pub struct EqPredicate {
+    pub left: BaseTableAttrRef,
+    pub right: BaseTableAttrRef,
+}
+
+impl EqPredicate {
+    pub fn new(left: BaseTableAttrRef, right: BaseTableAttrRef) -> Self {
+        Self { left, right }
+    }
+}
+
+/// [`SemanticCorrelation`] represents the semantic correlation between attributes in a
+/// query. "Semantic" means that the attributes are correlated based on the
+/// semantics of the query, not the statistics.
+///
+/// [`SemanticCorrelation`] contains equal attributes denoted by disjoint sets of base
+/// table attributes, e.g. {{ t1.c1 = t2.c1 = t3.c1 }, { t1.c2 = t2.c2 }}.
+#[derive(Clone, Debug, Default)]
+pub struct SemanticCorrelation {
+    /// A disjoint set of base table attributes with equal values in the same row.
+    disjoint_eq_attr_sets: DisjointSets<BaseTableAttrRef>,
+    /// The predicates that define the equalities.
+    eq_predicates: HashSet<EqPredicate>,
+}
+
+impl SemanticCorrelation {
+    pub fn new() -> Self {
+        Self {
+            disjoint_eq_attr_sets: DisjointSets::new(),
+            eq_predicates: HashSet::new(),
+        }
+    }
+
+    pub fn add_predicate(&mut self, predicate: EqPredicate) {
+        let left = &predicate.left;
+        let right = &predicate.right;
+
+        // Add the indices to the set if they do not exist.
+        if !self.disjoint_eq_attr_sets.contains(left) {
+            self.disjoint_eq_attr_sets
+                .make_set(left.clone())
+                .expect("just checked left attribute index does not exist");
+        }
+        if !self.disjoint_eq_attr_sets.contains(right) {
+            self.disjoint_eq_attr_sets
+                .make_set(right.clone())
+                .expect("just checked right attribute index does not exist");
+        }
+        // Union the attributes.
+        self.disjoint_eq_attr_sets
+            .union(left, right)
+            .expect("both attribute indices should exist");
+
+        // Keep track of the predicate.
+        self.eq_predicates.insert(predicate);
+    }
+
+    /// Determine if two attributes are in the same set.
+    pub fn is_eq(&mut self, left: &BaseTableAttrRef, right: &BaseTableAttrRef) -> bool {
+        self.disjoint_eq_attr_sets
+            .same_set(left, right)
+            .unwrap_or(false)
+    }
+
+    pub fn contains(&self, base_attr_ref: &BaseTableAttrRef) -> bool {
+        self.disjoint_eq_attr_sets.contains(base_attr_ref)
+    }
+
+    /// Get the number of attributes that are equal to `attr`, including `attr` itself.
+    pub fn num_eq_attributes(&mut self, attr: &BaseTableAttrRef) -> usize {
+        self.disjoint_eq_attr_sets.set_size(attr).unwrap()
+    }
+
+    /// Find the set of predicates that define the equality of the set of attributes `attr` belongs to.
+    pub fn find_predicates_for_eq_attr_set(&mut self, attr: &BaseTableAttrRef) -> Vec<EqPredicate> {
+        let mut predicates = Vec::new();
+        for predicate in &self.eq_predicates {
+            let left = &predicate.left;
+            let right = &predicate.right;
+            if (left != attr && self.disjoint_eq_attr_sets.same_set(attr, left).unwrap())
+                || (right != attr && self.disjoint_eq_attr_sets.same_set(attr, right).unwrap())
+            {
+                predicates.push(predicate.clone());
+            }
+        }
+        predicates
+    }
+
+    /// Find the set of attributes that define the equality of the set of attributes `attr` belongs to.
+    pub fn find_attrs_for_eq_attribute_set(
+        &mut self,
+        attr: &BaseTableAttrRef,
+    ) -> HashSet<BaseTableAttrRef> {
+        let predicates = self.find_predicates_for_eq_attr_set(attr);
+        predicates
+            .into_iter()
+            .flat_map(|predicate| vec![predicate.left, predicate.right])
+            .collect()
+    }
+
+    /// Union two `EqBaseTableattributesets` to produce a new disjoint sets.
+    pub fn union(x: Self, y: Self) -> Self {
+        let mut eq_attr_sets = Self::new();
+        for predicate in x
+            .eq_predicates
+            .into_iter()
+            .chain(y.eq_predicates.into_iter())
+        {
+            eq_attr_sets.add_predicate(predicate);
+        }
+        eq_attr_sets
+    }
+
+    pub fn merge(x: Option<Self>, y: Option<Self>) -> Option<Self> {
+        let eq_attr_sets = match (x, y) {
+            (Some(x), Some(y)) => Self::union(x, y),
+            (Some(x), None) => x.clone(),
+            (None, Some(y)) => y.clone(),
+            _ => return None,
+        };
+        Some(eq_attr_sets)
+    }
+}
+
+/// [`GroupAttrRefs`] represents the attributes of a group in a query.
+#[derive(Clone, Debug)]
+pub struct GroupAttrRefs {
+    attribute_refs: AttrRefs,
+    /// Correlation of the output attributes of the group.
+    output_correlation: Option<SemanticCorrelation>,
+}
+
+impl GroupAttrRefs {
+    pub fn new(attribute_refs: AttrRefs, output_correlation: Option<SemanticCorrelation>) -> Self {
+        Self {
+            attribute_refs,
+            output_correlation,
+        }
+    }
+
+    pub fn base_table_attribute_refs(&self) -> &AttrRefs {
+        &self.attribute_refs
+    }
+
+    pub fn output_correlation(&self) -> Option<&SemanticCorrelation> {
+        self.output_correlation.as_ref()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_eq_base_table_attribute_sets() {
+        let attr1 = BaseTableAttrRef {
+            table_id: TableId(1),
+            attr_idx: 1,
+        };
+        let attr2 = BaseTableAttrRef {
+            table_id: TableId(2),
+            attr_idx: 2,
+        };
+        let attr3 = BaseTableAttrRef {
+            table_id: TableId(3),
+            attr_idx: 3,
+        };
+        let attr4 = BaseTableAttrRef {
+            table_id: TableId(4),
+            attr_idx: 4,
+        };
+        let pred1 = EqPredicate::new(attr1.clone(), attr2.clone());
+        let pred2 = EqPredicate::new(attr3.clone(), attr4.clone());
+        let pred3 = EqPredicate::new(attr1.clone(), attr3.clone());
+
+        let mut eq_attr_sets = SemanticCorrelation::new();
+
+        // (1, 2)
+        eq_attr_sets.add_predicate(pred1.clone());
+        assert!(eq_attr_sets.is_eq(&attr1, &attr2));
+
+        // (1, 2), (3, 4)
+        eq_attr_sets.add_predicate(pred2.clone());
+        assert!(eq_attr_sets.is_eq(&attr3, &attr4));
+        assert!(!eq_attr_sets.is_eq(&attr2, &attr3));
+
+        let predicates = eq_attr_sets.find_predicates_for_eq_attr_set(&attr1);
+        assert_eq!(predicates.len(), 1);
+        assert!(predicates.contains(&pred1));
+
+        let predicates = eq_attr_sets.find_predicates_for_eq_attr_set(&attr3);
+        assert_eq!(predicates.len(), 1);
+        assert!(predicates.contains(&pred2));
+
+        // (1, 2, 3, 4)
+        eq_attr_sets.add_predicate(pred3.clone());
+        assert!(eq_attr_sets.is_eq(&attr1, &attr3));
+        assert!(eq_attr_sets.is_eq(&attr2, &attr4));
+        assert!(eq_attr_sets.is_eq(&attr1, &attr4));
+
+        let predicates = eq_attr_sets.find_predicates_for_eq_attr_set(&attr1);
+        assert_eq!(predicates.len(), 3);
+        assert!(predicates.contains(&pred1));
+        assert!(predicates.contains(&pred2));
+        assert!(predicates.contains(&pred3));
+    }
+}

optd-cost-model/src/common/properties/mod.rs

@@ -0,0 +1,23 @@
+use serde::{Deserialize, Serialize};
+
+use super::predicates::constant_pred::ConstantType;
+
+pub mod attr_ref;
+pub mod schema;
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct Attribute {
+    pub name: String,
+    pub typ: ConstantType,
+    pub nullable: bool,
+}
+
+impl std::fmt::Display for Attribute {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        if self.nullable {
+            write!(f, "{}:{:?}", self.name, self.typ)
+        } else {
+            write!(f, "{}:{:?}(non-null)", self.name, self.typ)
+        }
+    }
+}

optd-cost-model/src/common/properties/schema.rs

@@ -0,0 +1,35 @@
+use itertools::Itertools;
+
+use serde::{Deserialize, Serialize};
+
+use super::Attribute;
+
+/// [`Schema`] represents the schema of a group in the memo. It contains a list of attributes.
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct Schema {
+    pub attributes: Vec<Attribute>,
+}
+
+impl std::fmt::Display for Schema {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "[{}]",
+            self.attributes.iter().map(|x| x.to_string()).join(", ")
+        )
+    }
+}
+
+impl Schema {
+    pub fn new(attributes: Vec<Attribute>) -> Self {
+        Self { attributes }
+    }
+
+    pub fn len(&self) -> usize {
+        self.attributes.len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+}

optd-cost-model/src/cost/agg.rs

@@ -67,13 +67,15 @@ mod tests {
     use std::collections::HashMap;
 
     use crate::{
-        common::{predicates::constant_pred::ConstantType, types::TableId, values::Value},
+        common::{
+            predicates::constant_pred::ConstantType, properties::Attribute, types::TableId,
+            values::Value,
+        },
         cost_model::tests::{
             attr_ref, cnst, create_cost_model_mock_storage, empty_list, empty_per_attr_stats, list,
             TestPerAttributeStats,
         },
         stats::{utilities::simple_map::SimpleMap, MostCommonValues, DEFAULT_NUM_DISTINCT},
-        storage::Attribute,
         EstimatedStatistic,
     };
 

optd-cost-model/src/cost/filter/controller.rs

@@ -98,6 +98,7 @@ mod tests {
                 bin_op_pred::BinOpType, constant_pred::ConstantType, log_op_pred::LogOpType,
                 un_op_pred::UnOpType,
             },
+            properties::Attribute,
             types::TableId,
             values::Value,
         },
@@ -106,7 +107,6 @@ mod tests {
             utilities::{counter::Counter, simple_map::SimpleMap},
             Distribution, MostCommonValues, DEFAULT_EQ_SEL,
         },
-        storage::Attribute,
     };
     use arrow_schema::DataType;