feat: add file writer option for python sdk

Author: SaintBacchus

python/python/lance/dataset.py

@@ -3667,6 +3667,7 @@ def write_dataset(
     use_legacy_format: Optional[bool] = None,
     enable_v2_manifest_paths: bool = False,
     enable_move_stable_row_ids: bool = False,
+    file_writer_options: Optional[Dict[str, str]] = None,
 ) -> LanceDataset:
     """Write a given data_obj to the given uri
 
@@ -3725,6 +3726,9 @@ def write_dataset(
         These row ids are stable after compaction operations, but not after updates.
         This makes compaction more efficient, since with stable row ids no
         secondary indices need to be updated to point to new row ids.
+    file_writer_options : optional, dict
+        Extra file write options that make sense to control lance data file. This is
+        used to set page size for file in disk.
     """
     if use_legacy_format is not None:
         warnings.warn(
@@ -3759,6 +3763,7 @@ def write_dataset(
         "data_storage_version": data_storage_version,
         "enable_v2_manifest_paths": enable_v2_manifest_paths,
         "enable_move_stable_row_ids": enable_move_stable_row_ids,
+        "file_writer_options": file_writer_options,
     }
 
     if commit_lock:

python/python/lance/fragment.py

@@ -265,6 +265,7 @@ def create(
         data_storage_version: Optional[str] = None,
         use_legacy_format: Optional[bool] = None,
         storage_options: Optional[Dict[str, str]] = None,
+        file_writer_options: Optional[Dict[str, str]] = None,
     ) -> FragmentMetadata:
         """Create a :class:`FragmentMetadata` from the given data.
 
@@ -304,7 +305,9 @@ def create(
         storage_options : optional, dict
             Extra options that make sense for a particular storage connection. This is
             used to store connection parameters like credentials, endpoint, etc.
-
+        file_writer_options : optional, dict
+            Extra file write options that make sense to control lance data file. This is
+            used to set page size for file in disk.
         See Also
         --------
         lance.dataset.LanceOperation.Overwrite :
@@ -345,6 +348,7 @@ def create(
             mode=mode,
             data_storage_version=data_storage_version,
             storage_options=storage_options,
+            file_writer_options=file_writer_options
         )
 
     @property

python/src/dataset.rs

@@ -75,7 +75,7 @@ use pyo3::{
     PyObject, PyResult,
 };
 use snafu::location;
-
+use lance_file::v2::writer::FileWriterOptions;
 use crate::error::PythonErrorExt;
 use crate::file::object_store_from_uri_or_path;
 use crate::fragment::FileFragment;
@@ -1715,6 +1715,34 @@ pub fn get_write_params(options: &PyDict) -> PyResult<Option<WriteParams>> {
             });
         }
 
+        if let Some(writer_options) =
+            get_dict_opt::<HashMap<String, String>>(options, "file_writer_options")?
+        {
+            let mut file_writer_options = FileWriterOptions::default();
+            if let Some(max_page_bytes) = writer_options.get("max_page_bytes")
+            {
+                file_writer_options.max_page_bytes = match max_page_bytes.parse::<u64>() {
+                    Ok(n) => Some(n),
+                    Err(_e) => None
+                }
+            }
+            if let Some(data_cache_bytes) = writer_options.get("data_cache_bytes")
+            {
+                file_writer_options.data_cache_bytes = match data_cache_bytes.parse::<u64>() {
+                    Ok(n) => Some(n),
+                    Err(_e) => None
+                }
+            }
+            if let Some(keep_original_array) = writer_options.get("keep_original_array")
+            {
+                file_writer_options.keep_original_array = match keep_original_array.parse::<bool>() {
+                    Ok(n) => Some(n),
+                    Err(_e) => None
+                }
+            }
+            p.file_writer_options = Some(file_writer_options);
+        }
+
         if let Some(enable_move_stable_row_ids) =
             get_dict_opt::<bool>(options, "enable_move_stable_row_ids")?
         {

rust/lance-encoding/src/encodings/physical/value.rs

@@ -140,8 +140,9 @@ impl ValuePageDecoder {
         let bytes_u8: Vec<u8> = self.data[0].to_vec();
         let buffer_compressor = GeneralBufferCompressor::get_compressor(self.compression_config);
         let mut uncompressed_bytes: Vec<u8> = Vec::new();
+        println!("before decompress {}", bytes_u8.len());
         buffer_compressor.decompress(&bytes_u8, &mut uncompressed_bytes)?;
-
+        println!("after  decompress {}", uncompressed_bytes.len());
         let mut bytes_in_ranges: Vec<Bytes> =
             Vec::with_capacity(self.uncompressed_range_offsets.len());
         for range in &self.uncompressed_range_offsets {

rust/lance/src/dataset/fragment/write.rs

@@ -99,11 +99,12 @@ impl<'a> FragmentCreateBuilder<'a> {
         let mut fragment = Fragment::new(id);
         let full_path = base_path.child(DATA_DIR).child(filename.clone());
         let obj_writer = object_store.create(&full_path).await?;
-        let mut writer = lance_file::v2::writer::FileWriter::try_new(
-            obj_writer,
-            schema,
-            FileWriterOptions::default(),
-        )?;
+        let file_write_options = params
+            .file_writer_options
+            .clone()
+            .unwrap_or(FileWriterOptions::default());
+        let mut writer =
+            lance_file::v2::writer::FileWriter::try_new(obj_writer, schema, file_write_options)?;
 
         let (major, minor) = writer.version().to_numbers();
 

rust/lance/src/dataset/updater.rs

@@ -148,6 +148,7 @@ impl Updater {
             &schema,
             &self.fragment.dataset().base,
             data_storage_version,
+            None
         )
         .await
     }

rust/lance/src/dataset/write.rs

@@ -172,6 +172,8 @@ pub struct WriteParams {
     pub object_store_registry: Arc<ObjectStoreRegistry>,
 
     pub session: Option<Arc<Session>>,
+
+    pub file_writer_options: Option<FileWriterOptions>,
 }
 
 impl Default for WriteParams {
@@ -191,6 +193,7 @@ impl Default for WriteParams {
             enable_v2_manifest_paths: false,
             object_store_registry: Arc::new(ObjectStoreRegistry::default()),
             session: None,
+            file_writer_options: None,
         }
     }
 }
@@ -249,7 +252,7 @@ pub async fn do_write_fragments(
             .boxed()
     };
 
-    let writer_generator = WriterGenerator::new(object_store, base_dir, schema, storage_version);
+    let writer_generator = WriterGenerator::new(object_store, base_dir, schema, storage_version, params.clone());
     let mut writer: Option<Box<dyn GenericWriter>> = None;
     let mut num_rows_in_current_file = 0;
     let mut fragments = Vec::new();
@@ -506,6 +509,7 @@ pub async fn open_writer(
     schema: &Schema,
     base_dir: &Path,
     storage_version: LanceFileVersion,
+    params: Option<&WriteParams>,
 ) -> Result<Box<dyn GenericWriter>> {
     let filename = format!("{}.lance", Uuid::new_v4());
 
@@ -523,14 +527,17 @@ pub async fn open_writer(
             filename,
         ))
     } else {
+        let mut file_write_options = params
+            .as_ref()
+            .and_then(|wp| wp.file_writer_options.as_ref())
+            .cloned()
+            .unwrap_or_default();
+        file_write_options.format_version = Some(storage_version);
         let writer = object_store.create(&full_path).await?;
         let file_writer = v2::writer::FileWriter::try_new(
             writer,
             schema.clone(),
-            FileWriterOptions {
-                format_version: Some(storage_version),
-                ..Default::default()
-            },
+            file_write_options,
         )?;
         let writer_adapter = V2WriterAdapter {
             writer: file_writer,
@@ -547,6 +554,7 @@ struct WriterGenerator {
     base_dir: Path,
     schema: Schema,
     storage_version: LanceFileVersion,
+    params: WriteParams,
 }
 
 impl WriterGenerator {
@@ -555,12 +563,14 @@ impl WriterGenerator {
         base_dir: &Path,
         schema: &Schema,
         storage_version: LanceFileVersion,
+        params: WriteParams,
     ) -> Self {
         Self {
             object_store,
             base_dir: base_dir.clone(),
             schema: schema.clone(),
             storage_version,
+            params,
         }
     }
 
@@ -573,6 +583,7 @@ impl WriterGenerator {
             &self.schema,
             &self.base_dir,
             self.storage_version,
+            Some(&self.params)
         )
         .await?;
 

rust/lance/src/dataset/write/merge_insert.rs

@@ -721,6 +721,7 @@ impl MergeInsertJob {
                         &write_schema,
                         &dataset.base,
                         data_storage_version,
+                        None,
                     )
                     .await?;