Add Default impl for LanceScanConfig to simplify tests as per PR sugg…

…estions
lancedb · Oct 15, 2024 · accb77c · accb77c
1 parent 2681afb
commit accb77c
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 6 deletions.
diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs
@@ -68,7 +68,7 @@ use snafu::{location, Location};
 #[cfg(feature = "substrait")]
 use lance_datafusion::substrait::parse_substrait;
 
-const BATCH_SIZE_FALLBACK: usize = 8192;
+pub(crate) const BATCH_SIZE_FALLBACK: usize = 8192;
 // For backwards compatibility / historical reasons we re-calculate the default batch size
 // on each call
 pub fn get_default_batch_size() -> Option<usize> {

diff --git a/rust/lance/src/io/exec/scan.rs b/rust/lance/src/io/exec/scan.rs
@@ -30,7 +30,10 @@ use log::debug;
 use snafu::{location, Location};
 
 use crate::dataset::fragment::{FileFragment, FragReadConfig, FragmentReader};
-use crate::dataset::scanner::{DEFAULT_FRAGMENT_READAHEAD, LEGACY_DEFAULT_FRAGMENT_READAHEAD};
+use crate::dataset::scanner::{
+    BATCH_SIZE_FALLBACK, DEFAULT_FRAGMENT_READAHEAD, DEFAULT_IO_BUFFER_SIZE,
+    LEGACY_DEFAULT_FRAGMENT_READAHEAD,
+};
 use crate::dataset::Dataset;
 use crate::datatypes::Schema;
 
@@ -395,6 +398,23 @@ pub struct LanceScanConfig {
     pub ordered_output: bool,
 }
 
+// This is mostly for testing purposes, end users are unlikely to create this
+// on their own.
+impl Default for LanceScanConfig {
+    fn default() -> Self {
+        Self {
+            batch_size: BATCH_SIZE_FALLBACK,
+            batch_readahead: get_num_compute_intensive_cpus(),
+            fragment_readahead: None,
+            io_buffer_size: *DEFAULT_IO_BUFFER_SIZE,
+            with_row_id: false,
+            with_row_address: false,
+            with_make_deletions_null: false,
+            ordered_output: false,
+        }
+    }
+}
+
 /// DataFusion [ExecutionPlan] for scanning one Lance dataset
 #[derive(Debug)]
 pub struct LanceScanExec {

diff --git a/rust/lance/src/io/exec/take.rs b/rust/lance/src/io/exec/take.rs
@@ -341,7 +341,7 @@ mod tests {
     use tempfile::tempdir;
 
     use crate::{
-        dataset::{scanner::DEFAULT_IO_BUFFER_SIZE, WriteParams},
+        dataset::WriteParams,
         io::exec::{LanceScanConfig, LanceScanExec},
     };
 
@@ -430,7 +430,7 @@ mod tests {
 
         let config = LanceScanConfig {
             with_row_id: true,
-            ..Default::default(),
+            ..Default::default()
         };
         let input = Arc::new(LanceScanExec::new(
             dataset.clone(),
@@ -466,7 +466,7 @@ mod tests {
             dataset.fragments().clone(),
             None,
             scan_schema,
-            ..Default::default(),
+            LanceScanConfig::default(),
         ));
         assert!(TakeExec::try_new(dataset, input, extra_schema, 10).is_err());
     }
@@ -477,7 +477,7 @@ mod tests {
 
         let config = LanceScanConfig {
             with_row_id: true,
-            ..Default::default(),
+            ..Default::default()
         };
         let input = Arc::new(LanceScanExec::new(
             dataset.clone(),