mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-26 07:58:31 +00:00
fix(python): expose optimize compaction options
This commit is contained in:
+60
-7
@@ -20,8 +20,9 @@ use arrow::{
|
||||
use lancedb::blob::{BlobFile, BlobRangeRequest};
|
||||
use lancedb::index::scalar::FtsIndexBuilder;
|
||||
use lancedb::table::{
|
||||
AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken,
|
||||
NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
|
||||
AddDataMode, ColumnAlteration, CompactionMode, CompactionOptions, Duration,
|
||||
FieldMetadataUpdate, FtsToken as LanceDbFtsToken, IndexRemapMode, NewColumnTransform,
|
||||
OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
|
||||
};
|
||||
use lancedb::tokenize as lancedb_tokenize;
|
||||
use pyo3::{
|
||||
@@ -39,6 +40,56 @@ enum PredicateArg {
|
||||
Sql(String),
|
||||
}
|
||||
|
||||
fn parse_compaction_options(options: Option<&Bound<'_, PyDict>>) -> PyResult<CompactionOptions> {
|
||||
let mut parsed = CompactionOptions::default();
|
||||
let Some(options) = options else {
|
||||
return Ok(parsed);
|
||||
};
|
||||
|
||||
for (key, value) in options.iter() {
|
||||
let key: String = key.extract()?;
|
||||
match key.as_str() {
|
||||
"target_rows_per_fragment" => parsed.target_rows_per_fragment = value.extract()?,
|
||||
"max_rows_per_group" => parsed.max_rows_per_group = value.extract()?,
|
||||
"max_bytes_per_file" => parsed.max_bytes_per_file = value.extract()?,
|
||||
"materialize_deletions" => parsed.materialize_deletions = value.extract()?,
|
||||
"materialize_deletions_threshold" => {
|
||||
parsed.materialize_deletions_threshold = value.extract()?
|
||||
}
|
||||
"num_threads" => parsed.num_threads = value.extract()?,
|
||||
"batch_size" => parsed.batch_size = value.extract()?,
|
||||
"io_buffer_size" => parsed.io_buffer_size = value.extract()?,
|
||||
"defer_index_remap" => parsed.defer_index_remap = value.extract()?,
|
||||
"index_remap_mode" => {
|
||||
let mode: String = value.extract()?;
|
||||
parsed.index_remap_mode = IndexRemapMode::try_from(mode.as_str())
|
||||
.map_err(|err| PyValueError::new_err(err.to_string()))?;
|
||||
}
|
||||
"compaction_mode" => {
|
||||
let mode: Option<String> = value.extract()?;
|
||||
parsed.compaction_mode = mode
|
||||
.map(|mode| {
|
||||
CompactionMode::try_from(mode.as_str())
|
||||
.map_err(|err| PyValueError::new_err(err.to_string()))
|
||||
})
|
||||
.transpose()?;
|
||||
}
|
||||
"binary_copy_read_batch_bytes" => {
|
||||
parsed.binary_copy_read_batch_bytes = value.extract()?
|
||||
}
|
||||
"max_source_fragments" => parsed.max_source_fragments = value.extract()?,
|
||||
"max_overlays_per_fragment" => parsed.max_overlays_per_fragment = value.extract()?,
|
||||
_ => {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"Invalid compaction option: {key}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(parsed)
|
||||
}
|
||||
|
||||
/// Statistics about a compaction operation.
|
||||
#[pyclass(get_all, from_py_object)]
|
||||
#[derive(Clone, Debug)]
|
||||
@@ -1195,13 +1246,15 @@ impl Table {
|
||||
}
|
||||
|
||||
/// Optimize the on-disk data by compacting and pruning old data, for better performance.
|
||||
#[pyo3(signature = (cleanup_since_ms=None, delete_unverified=None))]
|
||||
pub fn optimize(
|
||||
self_: PyRef<'_, Self>,
|
||||
#[pyo3(signature = (cleanup_since_ms=None, delete_unverified=None, compaction_options=None))]
|
||||
pub fn optimize<'py>(
|
||||
self_: PyRef<'py, Self>,
|
||||
cleanup_since_ms: Option<u64>,
|
||||
delete_unverified: Option<bool>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
compaction_options: Option<&Bound<'py, PyDict>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
let compaction_options = parse_compaction_options(compaction_options)?;
|
||||
let older_than = if let Some(ms) = cleanup_since_ms {
|
||||
if ms > i64::MAX as u64 {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
@@ -1217,7 +1270,7 @@ impl Table {
|
||||
future_into_py(self_.py(), async move {
|
||||
let compaction_stats = inner
|
||||
.optimize(OptimizeAction::Compact {
|
||||
options: lancedb::table::CompactionOptions::default(),
|
||||
options: compaction_options,
|
||||
remap_options: None,
|
||||
})
|
||||
.await
|
||||
|
||||
Reference in New Issue
Block a user