This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 32b23306da Minor: Move test with code (#7392)
32b23306da is described below
commit 32b23306da35eafdf6081df021aa3af6a1e19c3e
Author: Andrew Lamb <[email protected]>
AuthorDate: Thu Aug 24 06:25:19 2023 -0400
Minor: Move test with code (#7392)
---
datafusion/core/src/physical_plan/mod.rs | 78 ----------------------------
datafusion/physical-expr/src/partitioning.rs | 77 +++++++++++++++++++++++++++
2 files changed, 77 insertions(+), 78 deletions(-)
diff --git a/datafusion/core/src/physical_plan/mod.rs
b/datafusion/core/src/physical_plan/mod.rs
index 1aa6baaf9f..43ce14a1e8 100644
--- a/datafusion/core/src/physical_plan/mod.rs
+++ b/datafusion/core/src/physical_plan/mod.rs
@@ -410,81 +410,3 @@ use datafusion_execution::TaskContext;
pub use datafusion_physical_expr::{
expressions, functions, hash_utils,
ordering_equivalence_properties_helper, udf,
};
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use arrow::datatypes::DataType;
- use arrow::datatypes::Schema;
-
- use crate::physical_plan::Distribution;
- use crate::physical_plan::Partitioning;
- use crate::physical_plan::PhysicalExpr;
- use datafusion_physical_expr::expressions::Column;
-
- use std::sync::Arc;
-
- #[tokio::test]
- async fn partitioning_satisfy_distribution() -> Result<()> {
- let schema = Arc::new(Schema::new(vec![
- arrow::datatypes::Field::new("column_1", DataType::Int64, false),
- arrow::datatypes::Field::new("column_2", DataType::Utf8, false),
- ]));
-
- let partition_exprs1: Vec<Arc<dyn PhysicalExpr>> = vec![
- Arc::new(Column::new_with_schema("column_1", &schema).unwrap()),
- Arc::new(Column::new_with_schema("column_2", &schema).unwrap()),
- ];
-
- let partition_exprs2: Vec<Arc<dyn PhysicalExpr>> = vec![
- Arc::new(Column::new_with_schema("column_2", &schema).unwrap()),
- Arc::new(Column::new_with_schema("column_1", &schema).unwrap()),
- ];
-
- let distribution_types = vec![
- Distribution::UnspecifiedDistribution,
- Distribution::SinglePartition,
- Distribution::HashPartitioned(partition_exprs1.clone()),
- ];
-
- let single_partition = Partitioning::UnknownPartitioning(1);
- let unspecified_partition = Partitioning::UnknownPartitioning(10);
- let round_robin_partition = Partitioning::RoundRobinBatch(10);
- let hash_partition1 = Partitioning::Hash(partition_exprs1, 10);
- let hash_partition2 = Partitioning::Hash(partition_exprs2, 10);
-
- for distribution in distribution_types {
- let result = (
- single_partition.satisfy(distribution.clone(), || {
- EquivalenceProperties::new(schema.clone())
- }),
- unspecified_partition.satisfy(distribution.clone(), || {
- EquivalenceProperties::new(schema.clone())
- }),
- round_robin_partition.satisfy(distribution.clone(), || {
- EquivalenceProperties::new(schema.clone())
- }),
- hash_partition1.satisfy(distribution.clone(), || {
- EquivalenceProperties::new(schema.clone())
- }),
- hash_partition2.satisfy(distribution.clone(), || {
- EquivalenceProperties::new(schema.clone())
- }),
- );
-
- match distribution {
- Distribution::UnspecifiedDistribution => {
- assert_eq!(result, (true, true, true, true, true))
- }
- Distribution::SinglePartition => {
- assert_eq!(result, (true, false, false, false, false))
- }
- Distribution::HashPartitioned(_) => {
- assert_eq!(result, (false, false, false, true, false))
- }
- }
- }
-
- Ok(())
- }
-}
diff --git a/datafusion/physical-expr/src/partitioning.rs
b/datafusion/physical-expr/src/partitioning.rs
index 991d6ba982..76567c8050 100644
--- a/datafusion/physical-expr/src/partitioning.rs
+++ b/datafusion/physical-expr/src/partitioning.rs
@@ -168,3 +168,80 @@ impl Distribution {
}
}
}
+
+#[cfg(test)]
+mod tests {
+ use crate::expressions::Column;
+
+ use super::*;
+ use arrow::datatypes::DataType;
+ use arrow::datatypes::Field;
+ use arrow::datatypes::Schema;
+ use datafusion_common::Result;
+
+ use std::sync::Arc;
+
+ #[test]
+ fn partitioning_satisfy_distribution() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("column_1", DataType::Int64, false),
+ Field::new("column_2", DataType::Utf8, false),
+ ]));
+
+ let partition_exprs1: Vec<Arc<dyn PhysicalExpr>> = vec![
+ Arc::new(Column::new_with_schema("column_1", &schema).unwrap()),
+ Arc::new(Column::new_with_schema("column_2", &schema).unwrap()),
+ ];
+
+ let partition_exprs2: Vec<Arc<dyn PhysicalExpr>> = vec![
+ Arc::new(Column::new_with_schema("column_2", &schema).unwrap()),
+ Arc::new(Column::new_with_schema("column_1", &schema).unwrap()),
+ ];
+
+ let distribution_types = vec![
+ Distribution::UnspecifiedDistribution,
+ Distribution::SinglePartition,
+ Distribution::HashPartitioned(partition_exprs1.clone()),
+ ];
+
+ let single_partition = Partitioning::UnknownPartitioning(1);
+ let unspecified_partition = Partitioning::UnknownPartitioning(10);
+ let round_robin_partition = Partitioning::RoundRobinBatch(10);
+ let hash_partition1 = Partitioning::Hash(partition_exprs1, 10);
+ let hash_partition2 = Partitioning::Hash(partition_exprs2, 10);
+
+ for distribution in distribution_types {
+ let result = (
+ single_partition.satisfy(distribution.clone(), || {
+ EquivalenceProperties::new(schema.clone())
+ }),
+ unspecified_partition.satisfy(distribution.clone(), || {
+ EquivalenceProperties::new(schema.clone())
+ }),
+ round_robin_partition.satisfy(distribution.clone(), || {
+ EquivalenceProperties::new(schema.clone())
+ }),
+ hash_partition1.satisfy(distribution.clone(), || {
+ EquivalenceProperties::new(schema.clone())
+ }),
+ hash_partition2.satisfy(distribution.clone(), || {
+ EquivalenceProperties::new(schema.clone())
+ }),
+ );
+
+ match distribution {
+ Distribution::UnspecifiedDistribution => {
+ assert_eq!(result, (true, true, true, true, true))
+ }
+ Distribution::SinglePartition => {
+ assert_eq!(result, (true, false, false, false, false))
+ }
+ Distribution::HashPartitioned(_) => {
+ assert_eq!(result, (false, false, false, true, false))
+ }
+ }
+ }
+
+ Ok(())
+ }
+}
