[ENH]: implement Display for Where (#5755)

Author: codetheweb

Cargo.lock

@@ -410,6 +410,18 @@ pub enum MetadataValue {
     SparseVector(SparseVector),
 }
 
+impl std::fmt::Display for MetadataValue {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            MetadataValue::Bool(v) => write!(f, "{}", v),
+            MetadataValue::Int(v) => write!(f, "{}", v),
+            MetadataValue::Float(v) => write!(f, "{}", v),
+            MetadataValue::Str(v) => write!(f, "\"{}\"", v),
+            MetadataValue::SparseVector(v) => write!(f, "SparseVector(len={})", v.values.len()),
+        }
+    }
+}
+
 impl Eq for MetadataValue {}
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
@@ -924,6 +936,27 @@ pub enum Where {
     Metadata(MetadataExpression),
 }
 
+impl std::fmt::Display for Where {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Where::Composite(composite) => {
+                let fragment = composite
+                    .children
+                    .iter()
+                    .map(|child| format!("{}", child))
+                    .collect::<Vec<_>>()
+                    .join(match composite.operator {
+                        BooleanOperator::And => " & ",
+                        BooleanOperator::Or => " | ",
+                    });
+                write!(f, "({})", fragment)
+            }
+            Where::Metadata(expr) => write!(f, "{}", expr),
+            Where::Document(expr) => write!(f, "{}", expr),
+        }
+    }
+}
+
 impl serde::Serialize for Where {
     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
     where
@@ -1230,6 +1263,18 @@ pub struct DocumentExpression {
     pub pattern: String,
 }
 
+impl std::fmt::Display for DocumentExpression {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let op_str = match self.operator {
+            DocumentOperator::Contains => "CONTAINS",
+            DocumentOperator::NotContains => "NOT CONTAINS",
+            DocumentOperator::Regex => "REGEX",
+            DocumentOperator::NotRegex => "NOT REGEX",
+        };
+        write!(f, "#document {} \"{}\"", op_str, self.pattern)
+    }
+}
+
 impl From<chroma_proto::DirectWhereDocument> for DocumentExpression {
     fn from(value: chroma_proto::DirectWhereDocument) -> Self {
         Self {
@@ -1285,6 +1330,19 @@ pub struct MetadataExpression {
     pub comparison: MetadataComparison,
 }
 
+impl std::fmt::Display for MetadataExpression {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match &self.comparison {
+            MetadataComparison::Primitive(op, value) => {
+                write!(f, "{} {} {}", self.key, op, value)
+            }
+            MetadataComparison::Set(op, set_value) => {
+                write!(f, "{} {} {}", self.key, op, set_value)
+            }
+        }
+    }
+}
+
 impl TryFrom<chroma_proto::DirectComparison> for MetadataExpression {
     type Error = WhereConversionError;
 
@@ -1428,6 +1486,20 @@ pub enum PrimitiveOperator {
     LessThanOrEqual,
 }
 
+impl std::fmt::Display for PrimitiveOperator {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let op_str = match self {
+            PrimitiveOperator::Equal => "=",
+            PrimitiveOperator::NotEqual => "≠",
+            PrimitiveOperator::GreaterThan => ">",
+            PrimitiveOperator::GreaterThanOrEqual => "≥",
+            PrimitiveOperator::LessThan => "<",
+            PrimitiveOperator::LessThanOrEqual => "≤",
+        };
+        write!(f, "{}", op_str)
+    }
+}
+
 impl From<chroma_proto::GenericComparator> for PrimitiveOperator {
     fn from(value: chroma_proto::GenericComparator) -> Self {
         match value {
@@ -1476,13 +1548,23 @@ impl TryFrom<PrimitiveOperator> for chroma_proto::NumberComparator {
     }
 }
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 #[cfg_attr(feature = "testing", derive(proptest_derive::Arbitrary))]
 pub enum SetOperator {
     In,
     NotIn,
 }
 
+impl std::fmt::Display for SetOperator {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let op_str = match self {
+            SetOperator::In => "∈",
+            SetOperator::NotIn => "∉",
+        };
+        write!(f, "{}", op_str)
+    }
+}
+
 impl From<chroma_proto::ListOperator> for SetOperator {
     fn from(value: chroma_proto::ListOperator) -> Self {
         match value {
@@ -1510,6 +1592,45 @@ pub enum MetadataSetValue {
     Str(Vec<String>),
 }
 
+impl std::fmt::Display for MetadataSetValue {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            MetadataSetValue::Bool(values) => {
+                let values_str = values
+                    .iter()
+                    .map(|v| format!("\"{}\"", v))
+                    .collect::<Vec<_>>()
+                    .join(", ");
+                write!(f, "[{}]", values_str)
+            }
+            MetadataSetValue::Int(values) => {
+                let values_str = values
+                    .iter()
+                    .map(|v| v.to_string())
+                    .collect::<Vec<_>>()
+                    .join(", ");
+                write!(f, "[{}]", values_str)
+            }
+            MetadataSetValue::Float(values) => {
+                let values_str = values
+                    .iter()
+                    .map(|v| v.to_string())
+                    .collect::<Vec<_>>()
+                    .join(", ");
+                write!(f, "[{}]", values_str)
+            }
+            MetadataSetValue::Str(values) => {
+                let values_str = values
+                    .iter()
+                    .map(|v| format!("\"{}\"", v))
+                    .collect::<Vec<_>>()
+                    .join(", ");
+                write!(f, "[{}]", values_str)
+            }
+        }
+    }
+}
+
 impl MetadataSetValue {
     pub fn value_type(&self) -> MetadataValueType {
         match self {