Skip to content

Commit

Permalink
fix: unit tests
Browse files Browse the repository at this point in the history
  • Loading branch information
discord9 committed Dec 19, 2024
1 parent 2cb3dd0 commit e002a63
Show file tree
Hide file tree
Showing 3 changed files with 61 additions and 72 deletions.
6 changes: 5 additions & 1 deletion src/flow/src/compute/render/reduce.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1585,7 +1585,11 @@ mod test {

if let Some(expected) = expected.get(&now) {
let batch = expected.iter().map(|v| Value::from(*v)).collect_vec();
let batch = Batch::try_from_rows(vec![batch.into()]).unwrap();
let batch = Batch::try_from_rows_with_types(
vec![batch.into()],
&[CDT::int64_datatype()],
)
.unwrap();
assert_eq!(res.first(), Some(&batch));
}
});
Expand Down
57 changes: 0 additions & 57 deletions src/flow/src/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -123,63 +123,6 @@ impl Batch {
Ok(batch)
}

/// Get batch from rows, will try best to determine data type
///
/// for test purposes only
#[cfg(test)]
pub fn try_from_rows(rows: Vec<crate::repr::Row>) -> Result<Self, EvalError> {
if rows.is_empty() {
return Ok(Self::empty());
}
let len = rows.len();
let mut builder = rows
.first()
.unwrap()
.iter()
.enumerate()
.map(|(i, v)| {
let mut ty = None;
if v.data_type().is_null() {
for row in rows.iter() {
if let Some(t) = row.get(i)
&& !t.data_type().is_null()
{
ty = Some(t.data_type().clone());
break;
}
}
}
// if all rows are null, use null type
let ty = ty.unwrap_or(datatypes::prelude::ConcreteDataType::null_datatype());

ty.create_mutable_vector(len)
})
.collect_vec();
for row in rows {
ensure!(
row.len() == builder.len(),
InvalidArgumentSnafu {
reason: format!(
"row length not match, expect {}, found {}",
builder.len(),
row.len()
)
}
);
for (idx, value) in row.iter().enumerate() {
builder[idx]
.try_push_value_ref(value.as_value_ref())
.context(DataTypeSnafu {
msg: "Failed to convert rows to columns",
})?;
}
}

let columns = builder.into_iter().map(|mut b| b.to_vector()).collect_vec();
let batch = Self::try_new(columns, len)?;
Ok(batch)
}

pub fn empty() -> Self {
Self {
batch: vec![],
Expand Down
70 changes: 56 additions & 14 deletions src/flow/src/expr/linear.rs
Original file line number Diff line number Diff line change
Expand Up @@ -936,20 +936,33 @@ mod test {
.unwrap()
.unwrap();
assert_eq!(ret, Row::pack(vec![Value::from(false), Value::from(true)]));

let ty = [
ConcreteDataType::int32_datatype(),
ConcreteDataType::int32_datatype(),
ConcreteDataType::int32_datatype(),
];
// batch mode
let mut batch = Batch::try_from_rows(vec![Row::from(vec![
Value::from(4),
Value::from(2),
Value::from(3),
])])
let mut batch = Batch::try_from_rows_with_types(
vec![Row::from(vec![
Value::from(4),
Value::from(2),
Value::from(3),
])],
&ty,
)
.unwrap();
let ret = safe_mfp.eval_batch_into(&mut batch).unwrap();

assert_eq!(
ret,
Batch::try_from_rows(vec![Row::from(vec![Value::from(false), Value::from(true)])])
.unwrap()
Batch::try_from_rows_with_types(
vec![Row::from(vec![Value::from(false), Value::from(true)])],
&[
ConcreteDataType::boolean_datatype(),
ConcreteDataType::boolean_datatype(),
],
)
.unwrap()
);
}

Expand Down Expand Up @@ -984,7 +997,15 @@ mod test {
.unwrap();
assert_eq!(ret, None);

let mut input1_batch = Batch::try_from_rows(vec![Row::new(input1)]).unwrap();
let input_type = [
ConcreteDataType::int32_datatype(),
ConcreteDataType::int32_datatype(),
ConcreteDataType::int32_datatype(),
ConcreteDataType::string_datatype(),
];

let mut input1_batch =
Batch::try_from_rows_with_types(vec![Row::new(input1)], &input_type).unwrap();
let ret_batch = safe_mfp.eval_batch_into(&mut input1_batch).unwrap();
assert_eq!(
ret_batch,
Expand All @@ -1002,7 +1023,8 @@ mod test {
.unwrap();
assert_eq!(ret, Some(Row::pack(vec![Value::from(11)])));

let mut input2_batch = Batch::try_from_rows(vec![Row::new(input2)]).unwrap();
let mut input2_batch =
Batch::try_from_rows_with_types(vec![Row::new(input2)], &input_type).unwrap();
let ret_batch = safe_mfp.eval_batch_into(&mut input2_batch).unwrap();
assert_eq!(
ret_batch,
Expand Down Expand Up @@ -1055,7 +1077,14 @@ mod test {
let ret = safe_mfp.evaluate_into(&mut input1.clone(), &mut Row::empty());
assert!(matches!(ret, Err(EvalError::InvalidArgument { .. })));

let mut input1_batch = Batch::try_from_rows(vec![Row::new(input1)]).unwrap();
let input_type = [
ConcreteDataType::int64_datatype(),
ConcreteDataType::int32_datatype(),
ConcreteDataType::int32_datatype(),
ConcreteDataType::int32_datatype(),
];
let mut input1_batch =
Batch::try_from_rows_with_types(vec![Row::new(input1)], &input_type).unwrap();
let ret_batch = safe_mfp.eval_batch_into(&mut input1_batch);
assert!(matches!(ret_batch, Err(EvalError::InvalidArgument { .. })));

Expand All @@ -1065,7 +1094,13 @@ mod test {
.unwrap();
assert_eq!(ret, Some(Row::new(input2.clone())));

let input2_batch = Batch::try_from_rows(vec![Row::new(input2)]).unwrap();
let input_type = [
ConcreteDataType::int64_datatype(),
ConcreteDataType::int32_datatype(),
ConcreteDataType::int32_datatype(),
];
let input2_batch =
Batch::try_from_rows_with_types(vec![Row::new(input2)], &input_type).unwrap();
let ret_batch = safe_mfp.eval_batch_into(&mut input2_batch.clone()).unwrap();
assert_eq!(ret_batch, input2_batch);

Expand All @@ -1075,7 +1110,8 @@ mod test {
.unwrap();
assert_eq!(ret, None);

let input3_batch = Batch::try_from_rows(vec![Row::new(input3)]).unwrap();
let input3_batch =
Batch::try_from_rows_with_types(vec![Row::new(input3)], &input_type).unwrap();
let ret_batch = safe_mfp.eval_batch_into(&mut input3_batch.clone()).unwrap();
assert_eq!(
ret_batch,
Expand Down Expand Up @@ -1111,7 +1147,13 @@ mod test {
let ret = safe_mfp.evaluate_into(&mut input1.clone(), &mut Row::empty());
assert_eq!(ret.unwrap(), Some(Row::new(vec![Value::from(false)])));

let mut input1_batch = Batch::try_from_rows(vec![Row::new(input1)]).unwrap();
let input_type = [
ConcreteDataType::int32_datatype(),
ConcreteDataType::int32_datatype(),
ConcreteDataType::int32_datatype(),
];
let mut input1_batch =
Batch::try_from_rows_with_types(vec![Row::new(input1)], &input_type).unwrap();
let ret_batch = safe_mfp.eval_batch_into(&mut input1_batch).unwrap();

assert_eq!(
Expand Down

0 comments on commit e002a63

Please sign in to comment.