jiacai2050 commented on code in PR #1594:
URL: https://github.com/apache/horaedb/pull/1594#discussion_r1849875235
##########
src/partition_table_engine/src/scan_builder.rs:
##########
@@ -89,18 +92,186 @@ impl TableScanBuilder for PartitionedTableScanBuilder {
DataFusionError::Internal(format!("failed to build
partition rule, err:{e}"))
})?;
+ let mut partitioned_key_indices = PartitionedFilterKeyIndex::new();
// Evaluate expr and locate partition.
let partitions = df_partition_rule
- .locate_partitions_for_read(request.predicate.exprs())
+ .locate_partitions_for_read(request.predicate.exprs(), &mut
partitioned_key_indices)
.map_err(|e| {
DataFusionError::Internal(format!("failed to locate partition
for read, err:{e}"))
})?;
+
let sub_tables =
- self.get_sub_table_idents(&self.table_name, &self.partition_info,
partitions);
+ self.get_sub_table_idents(&self.table_name, &self.partition_info,
&partitions);
+
+ let predicates = if partitioned_key_indices.len() == partitions.len() {
+ Some(
+ partitioned_predicates(
+ request.predicate.clone(),
+ &partitions,
+ &mut partitioned_key_indices,
+ )
+ .map_err(|e| DataFusionError::Internal(format!("err:{e}")))?,
+ )
+ } else {
+ // since FilterExtractor.extract only cover some specific expr
+ // cases, partitioned_key_indices.len() could be 0.
+ // All partition requests will have the same predicate.
+ None
+ };
// Build plan.
- let plan = UnresolvedPartitionedScan::new(&self.table_name,
sub_tables, request);
+ let plan =
+ UnresolvedPartitionedScan::new(&self.table_name, sub_tables,
request, predicates);
Ok(Arc::new(plan))
}
}
+
+#[cfg(test)]
+mod tests {
+ use common_types::{column_schema::Builder as ColBuilder, datum::DatumKind,
schema::Builder};
+ use datafusion::logical_expr::{
+ expr::{BinaryExpr, InList},
+ Expr, Operator,
+ };
+ use table_engine::{
+ partition::{
+ rule::df_adapter::{DfPartitionRuleAdapter,
PartitionedFilterKeyIndex},
+ KeyPartitionInfo, PartitionDefinition, PartitionInfo,
+ },
+ predicate::PredicateBuilder,
+ };
+
+ use super::partitioned_predicates;
+
+ #[test]
+ fn test_partitioned_predicate() {
+ // conditions:
+ // 1) table schema: col_ts, col1, col2, in which col1 and col2 are
both keys,
+ // and with two partitions
+ // 2) sql: select * from table where col1 = '33' and col2 in ("aa",
"bb",
+ // "cc", "dd")
+ // partition expectations:
+ // 1) query fit in two partitions
+ // 2) yield two predicates, p0: col1 = '33' and col2 in ("aa",
"bb", "cc");
+ // p1: col1 = '33' and col2 in ("dd")
+ let definitions = vec![
+ PartitionDefinition {
+ name: "p1".to_string(),
+ origin_name: None,
+ },
+ PartitionDefinition {
+ name: "p2".to_string(),
+ origin_name: None,
+ },
+ ];
+
+ let partition_info = PartitionInfo::Key(KeyPartitionInfo {
+ version: 0,
+ definitions,
+ partition_key: vec!["col1".to_string(), "col2".to_string()],
+ linear: false,
+ });
+
+ let schema = {
+ let builder = Builder::new();
+ let col_ts = ColBuilder::new("col_ts".to_string(),
DatumKind::Timestamp)
+ .build()
+ .expect("ts");
+ let col1 = ColBuilder::new("col1".to_string(), DatumKind::String)
+ .build()
+ .expect("should succeed to build column schema");
+ let col2 = ColBuilder::new("col2".to_string(), DatumKind::String)
+ .build()
+ .expect("should succeed to build column schema");
+ builder
+ .auto_increment_column_id(true)
+ .add_key_column(col_ts)
+ .unwrap()
+ .add_key_column(col1)
+ .unwrap()
+ .add_key_column(col2)
+ .unwrap()
+ .primary_key_indexes(vec![1, 2])
+ .build()
+ .unwrap()
+ };
+
+ let df_partition_rule = DfPartitionRuleAdapter::new(partition_info,
&schema).unwrap();
+
+ let exprs = vec![
+ Expr::BinaryExpr(BinaryExpr {
Review Comment:
DataFusion provider some helper functions to define expr:
-
https://docs.rs/datafusion/latest/datafusion/logical_expr/index.html#functions
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
