fix: take benchmark can not read dataset (#4065)

Author: yanghua

rust/lance/benches/take.rs

@@ -7,10 +7,8 @@ use arrow_array::{
 use arrow_schema::{DataType, Field, FieldRef, Schema as ArrowSchema};
 use criterion::{criterion_group, criterion_main, Criterion};
 
-use lance::{
-    arrow::FixedSizeListArrayExt,
-    dataset::{builder::DatasetBuilder, ProjectionRequest},
-};
+use lance::dataset::{Dataset, WriteMode, WriteParams};
+use lance::{arrow::FixedSizeListArrayExt, dataset::ProjectionRequest};
 use lance_file::version::LanceFileVersion;
 #[cfg(target_os = "linux")]
 use pprof::criterion::{Output, PProfProfiler};
@@ -19,8 +17,6 @@ use std::sync::Arc;
 #[cfg(target_os = "linux")]
 use std::time::Duration;
 
-use lance::dataset::{Dataset, WriteMode, WriteParams};
-
 const BATCH_SIZE: u64 = 1024;
 
 fn gen_ranges(num_rows: u64, file_size: u64, n: usize) -> Vec<u64> {
@@ -92,25 +88,6 @@ async fn create_dataset(
     num_batches: i32,
     file_size: i32,
 ) -> Dataset {
-    create_file(
-        std::path::Path::new(path),
-        WriteMode::Create,
-        data_storage_version,
-        num_batches,
-        file_size,
-    )
-    .await;
-
-    DatasetBuilder::from_uri(path).load().await.unwrap()
-}
-
-async fn create_file(
-    path: &std::path::Path,
-    mode: WriteMode,
-    data_storage_version: LanceFileVersion,
-    num_batches: i32,
-    file_size: i32,
-) {
     let schema = Arc::new(ArrowSchema::new(vec![
         Field::new("i", DataType::Int32, false),
         Field::new("f", DataType::Float32, false),
@@ -163,18 +140,17 @@ async fn create_file(
         })
         .collect();
 
-    let test_uri = path.to_str().unwrap();
     let write_params = WriteParams {
         max_rows_per_file: file_size as usize,
         max_rows_per_group: batch_size as usize,
-        mode,
+        mode: WriteMode::Create,
         data_storage_version: Some(data_storage_version),
         ..Default::default()
     };
     let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone());
-    Dataset::write(reader, test_uri, Some(write_params))
+    Dataset::write(reader, path, Some(write_params))
         .await
-        .unwrap();
+        .unwrap()
 }
 
 #[cfg(target_os = "linux")]