resolve conflict

Author: xudong963

datafusion/datasource/benches/split_groups_by_statistics.rs

@@ -18,8 +18,7 @@
 use arrow::datatypes::{DataType, Field, Schema};
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
 use datafusion_datasource::file_scan_config::FileScanConfig;
-use datafusion_datasource::test_util::generate_test_files;
-use datafusion_datasource::test_util::verify_sort_integrity;
+use datafusion_datasource::{generate_test_files, verify_sort_integrity};
 use datafusion_physical_expr::PhysicalSortExpr;
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
 use std::sync::Arc;

datafusion/datasource/src/file_scan_config.rs

@@ -1467,7 +1467,10 @@ pub fn wrap_partition_value_in_dict(val: ScalarValue) -> ScalarValue {
 
 #[cfg(test)]
 mod tests {
-    use crate::{test_util::MockSource, tests::aggr_test_schema};
+    use crate::{
+        generate_test_files, test_util::MockSource, tests::aggr_test_schema,
+        verify_sort_integrity,
+    };
 
     use super::*;
     use arrow::{
@@ -2311,9 +2314,10 @@ mod tests {
         );
         assert_eq!(new_config.constraints, Constraints::default());
         assert!(new_config.new_lines_in_values);
+    }
+
+    #[test]
     fn test_split_groups_by_statistics_with_target_partitions() -> Result<()> {
-        use crate::test_util::generate_test_files;
-        use crate::test_util::verify_sort_integrity;
         use datafusion_common::DFSchema;
         use datafusion_expr::{col, execution_props::ExecutionProps};
 

datafusion/datasource/src/mod.rs

@@ -43,26 +43,27 @@ pub mod sink;
 pub mod source;
 mod statistics;
 
-#[cfg(any(test))]
+#[cfg(test)]
 pub mod test_util;
 
 pub mod url;
 pub mod write;
+pub use self::url::ListingTableUrl;
+use crate::file_groups::FileGroup;
 use chrono::TimeZone;
-use datafusion_common::Result;
+use datafusion_common::stats::Precision;
+use datafusion_common::{ColumnStatistics, Result};
 use datafusion_common::{ScalarValue, Statistics};
 use file_meta::FileMeta;
 use futures::{Stream, StreamExt};
 use object_store::{path::Path, ObjectMeta};
 use object_store::{GetOptions, GetRange, ObjectStore};
+pub use statistics::add_row_stats;
+pub use statistics::compute_all_files_statistics;
 use std::ops::Range;
 use std::pin::Pin;
 use std::sync::Arc;
 
-pub use self::url::ListingTableUrl;
-pub use statistics::add_row_stats;
-pub use statistics::compute_all_files_statistics;
-
 /// Stream of files get listed from object store
 pub type PartitionedFileStream =
     Pin<Box<dyn Stream<Item = Result<PartitionedFile>> + Send + Sync + 'static>>;
@@ -313,6 +314,78 @@ async fn find_first_newline(
     Ok(index)
 }
 
+/// Generates test files with min-max statistics in different overlap patterns
+/// Used by tests and benchmarks
+pub fn generate_test_files(num_files: usize, overlap_factor: f64) -> Vec<FileGroup> {
+    let mut files = Vec::with_capacity(num_files);
+    if num_files == 0 {
+        return vec![];
+    }
+    let range_size = if overlap_factor == 0.0 {
+        100 / num_files as i64
+    } else {
+        (100.0 / (overlap_factor * num_files as f64)).max(1.0) as i64
+    };
+
+    for i in 0..num_files {
+        let base = (i as f64 * range_size as f64 * (1.0 - overlap_factor)) as i64;
+        let min = base as f64;
+        let max = (base + range_size) as f64;
+
+        let file = PartitionedFile {
+            object_meta: ObjectMeta {
+                location: Path::from(format!("file_{}.parquet", i)),
+                last_modified: chrono::Utc::now(),
+                size: 1000,
+                e_tag: None,
+                version: None,
+            },
+            partition_values: vec![],
+            range: None,
+            statistics: Some(Arc::new(Statistics {
+                num_rows: Precision::Exact(100),
+                total_byte_size: Precision::Exact(1000),
+                column_statistics: vec![ColumnStatistics {
+                    null_count: Precision::Exact(0),
+                    max_value: Precision::Exact(ScalarValue::Float64(Some(max))),
+                    min_value: Precision::Exact(ScalarValue::Float64(Some(min))),
+                    sum_value: Precision::Absent,
+                    distinct_count: Precision::Absent,
+                }],
+            })),
+            extensions: None,
+            metadata_size_hint: None,
+        };
+        files.push(file);
+    }
+
+    vec![FileGroup::new(files)]
+}
+
+// Helper function to verify that files within each group maintain sort order
+/// Used by tests and benchmarks
+pub fn verify_sort_integrity(file_groups: &[FileGroup]) -> bool {
+    for group in file_groups {
+        let files = group.iter().collect::<Vec<_>>();
+        for i in 1..files.len() {
+            let prev_file = files[i - 1];
+            let curr_file = files[i];
+
+            // Check if the min value of current file is greater than max value of previous file
+            if let (Some(prev_stats), Some(curr_stats)) =
+                (&prev_file.statistics, &curr_file.statistics)
+            {
+                let prev_max = &prev_stats.column_statistics[0].max_value;
+                let curr_min = &curr_stats.column_statistics[0].min_value;
+                if curr_min.get_value().unwrap() <= prev_max.get_value().unwrap() {
+                    return false;
+                }
+            }
+        }
+    }
+    true
+}
+
 #[cfg(test)]
 mod tests {
     use super::ListingTableUrl;

datafusion/datasource/src/test_util.rs

@@ -17,18 +17,14 @@
 
 use crate::{
     file::FileSource, file_scan_config::FileScanConfig, file_stream::FileOpener,
-    PartitionedFile,
 };
 
 use std::sync::Arc;
 
-use crate::file_groups::FileGroup;
 use arrow::datatypes::SchemaRef;
-use datafusion_common::stats::Precision;
-use datafusion_common::{ColumnStatistics, Result, ScalarValue, Statistics};
+use datafusion_common::{Result, Statistics};
 use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
-use object_store::path::Path;
-use object_store::{ObjectMeta, ObjectStore};
+use object_store::ObjectStore;
 
 /// Minimal [`crate::file::FileSource`] implementation for use in tests.
 #[derive(Clone, Default)]
@@ -85,73 +81,3 @@ impl FileSource for MockSource {
         "mock"
     }
 }
-
-/// Generates test files with min-max statistics in different overlap patterns
-pub fn generate_test_files(num_files: usize, overlap_factor: f64) -> Vec<FileGroup> {
-    let mut files = Vec::with_capacity(num_files);
-    if num_files == 0 {
-        return vec![];
-    }
-    let range_size = if overlap_factor == 0.0 {
-        100 / num_files as i64
-    } else {
-        (100.0 / (overlap_factor * num_files as f64)).max(1.0) as i64
-    };
-
-    for i in 0..num_files {
-        let base = (i as f64 * range_size as f64 * (1.0 - overlap_factor)) as i64;
-        let min = base as f64;
-        let max = (base + range_size) as f64;
-
-        let file = PartitionedFile {
-            object_meta: ObjectMeta {
-                location: Path::from(format!("file_{}.parquet", i)),
-                last_modified: chrono::Utc::now(),
-                size: 1000,
-                e_tag: None,
-                version: None,
-            },
-            partition_values: vec![],
-            range: None,
-            statistics: Some(Statistics {
-                num_rows: Precision::Exact(100),
-                total_byte_size: Precision::Exact(1000),
-                column_statistics: vec![ColumnStatistics {
-                    null_count: Precision::Exact(0),
-                    max_value: Precision::Exact(ScalarValue::Float64(Some(max))),
-                    min_value: Precision::Exact(ScalarValue::Float64(Some(min))),
-                    sum_value: Precision::Absent,
-                    distinct_count: Precision::Absent,
-                }],
-            }),
-            extensions: None,
-            metadata_size_hint: None,
-        };
-        files.push(file);
-    }
-
-    vec![FileGroup::new(files)]
-}
-
-// Helper function to verify that files within each group maintain sort order
-pub fn verify_sort_integrity(file_groups: &[FileGroup]) -> bool {
-    for group in file_groups {
-        let files = group.iter().collect::<Vec<_>>();
-        for i in 1..files.len() {
-            let prev_file = files[i - 1];
-            let curr_file = files[i];
-
-            // Check if the min value of current file is greater than max value of previous file
-            if let (Some(prev_stats), Some(curr_stats)) =
-                (&prev_file.statistics, &curr_file.statistics)
-            {
-                let prev_max = &prev_stats.column_statistics[0].max_value;
-                let curr_min = &curr_stats.column_statistics[0].min_value;
-                if curr_min.get_value().unwrap() <= prev_max.get_value().unwrap() {
-                    return false;
-                }
-            }
-        }
-    }
-    true
-}