Skip to content

Commit

Permalink
feat: implement Display for PartitionExpr
Browse files Browse the repository at this point in the history
  • Loading branch information
Kelvinyu1117 committed Jun 1, 2024
1 parent 45fee94 commit 53d81cf
Showing 1 changed file with 75 additions and 0 deletions.
75 changes: 75 additions & 0 deletions src/partition/src/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fmt::{Debug, Display, Formatter};

use datatypes::value::Value;
use serde::{Deserialize, Serialize};
use sql::statements::value_to_sql_value;
Expand All @@ -35,6 +37,16 @@ pub enum Operand {
Expr(PartitionExpr),
}

impl Display for Operand {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Self::Column(v) => write!(f, "{v}"),
Self::Value(v) => write!(f, "{v}"),
Self::Expr(v) => write!(f, "{v}"),
}
}
}

/// A restricted set of [Operator](datafusion_expr::Operator) that can be used in
/// partition expressions.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
Expand Down Expand Up @@ -80,6 +92,20 @@ impl RestrictedOp {
}
}
}
impl Display for RestrictedOp {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Self::Eq => write!(f, "="),
Self::NotEq => write!(f, "<>"),
Self::Lt => write!(f, "<"),
Self::LtEq => write!(f, "<="),
Self::Gt => write!(f, ">"),
Self::GtEq => write!(f, ">="),
Self::And => write!(f, "AND"),
Self::Or => write!(f, "OR"),
}
}
}

impl PartitionExpr {
pub fn new(lhs: Operand, op: RestrictedOp, rhs: Operand) -> Self {
Expand Down Expand Up @@ -115,3 +141,52 @@ impl PartitionExpr {
}
}
}

impl Display for PartitionExpr {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "{} {} {}", self.lhs, self.op, self.rhs)
}
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn test_partition_expr() {
let partition_expr_s = [
"a = 10", "a <> 10", "a < 10", "a <= 10", "a > 10", "a >= 10", "a AND b", "a OR b",
];

let ops = [
RestrictedOp::Eq,
RestrictedOp::NotEq,
RestrictedOp::Lt,
RestrictedOp::LtEq,
RestrictedOp::Gt,
RestrictedOp::GtEq,
RestrictedOp::And,
RestrictedOp::Or,
];

assert_eq!(partition_expr_s.len(), ops.len());

for i in 0..partition_expr_s.len() - 2 {
let expr = PartitionExpr::new(
Operand::Column("a".to_string()),
ops[i].clone(),
Operand::Value(Value::UInt32(10)),
);
assert_eq!(partition_expr_s[i], expr.to_string());
}

for i in partition_expr_s.len() - 2..partition_expr_s.len() {
let expr = PartitionExpr::new(
Operand::Column("a".to_string()),
ops[i].clone(),
Operand::Column("b".to_string()),
);
assert_eq!(partition_expr_s[i], expr.to_string());
}
}
}

0 comments on commit 53d81cf

Please sign in to comment.