Skip to content

Commit

Permalink
Add Dataset::maybe_batch for faster database write
Browse files Browse the repository at this point in the history
  • Loading branch information
Atreyagaurav committed Nov 28, 2024
1 parent b448b08 commit b3e28f5
Show file tree
Hide file tree
Showing 3 changed files with 83 additions and 1 deletion.
1 change: 1 addition & 0 deletions CHANGES.md
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
- Add methods `alternative_name`, `is_nullable`, `is_unique`, `default_value` to `Field` ([#561](https://github.com/georust/gdal/pull/561))
- Add `Defn::geometry_type` ([#562](https://github.com/georust/gdal/pull/562))
- Add `Defn::field_index` and `Feature::field_index` ([#581](https://github.com/georust/gdal/pull/581))
- Add `Dataset::maybe_batch` which always succeeds to use transaction when possible for speed reasons -- no rollback ([#584](https://github.com/georust/gdal/pull/584))

### Fixed

Expand Down
35 changes: 35 additions & 0 deletions src/test_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,41 @@ pub fn open_gpkg_for_update(path: &Path) -> (TempPath, Dataset) {
(temp_path, ds)
}

/// Copies the given file to a temporary file and opens it for writing. When the returned
/// `TempPath` is dropped, the file is deleted.
pub fn open_dataset_for_update(path: &Path) -> (TempPath, Dataset) {
use std::fs;
use std::io::Write;

let input_data = fs::read(path).unwrap();
let (mut file, temp_path) = tempfile::Builder::new()
// using the whole filename as suffix should be fine (can't
// use .extension() for .shp.zip and such)
.suffix(
path.file_name()
.unwrap_or_default()
.to_string_lossy()
.as_ref(),
)
.tempfile()
.unwrap()
.into_parts();
file.write_all(&input_data).unwrap();
// Close the temporary file so that Dataset can open it safely even if the filesystem uses
// exclusive locking (Windows?).
drop(file);

let ds = Dataset::open_ex(
&temp_path,
DatasetOptions {
open_flags: GDALAccess::GA_Update.into(),
..DatasetOptions::default()
},
)
.unwrap();
(temp_path, ds)
}

/// Assert numerical difference between two expressions is less than
/// 64-bit machine epsilon or a specified epsilon.
///
Expand Down
48 changes: 47 additions & 1 deletion src/vector/transaction.rs
Original file line number Diff line number Diff line change
Expand Up @@ -177,11 +177,31 @@ impl Dataset {
}
Ok(Transaction::new(self))
}

/// Optionally start a transaction before running `func` for performance
///
/// This uses transaction if the dataset supports it, otherwise it
/// runs the `func` function as it is on the dataset.
pub fn maybe_batch(&mut self, func: impl Fn(&Dataset) -> Result<()>) -> Result<()> {
let force = 0; // since this is for speed
let rv = unsafe { gdal_sys::GDALDatasetStartTransaction(self.c_dataset(), force) };
let res = func(self);
if rv == OGRErr::OGRERR_NONE {
let rv = unsafe { gdal_sys::GDALDatasetCommitTransaction(self.c_dataset()) };
if rv != OGRErr::OGRERR_NONE {
return Err(GdalError::OgrError {
err: rv,
method_name: "GDALDatasetCommitTransaction",
});
}
}
res
}
}

#[cfg(test)]
mod tests {
use crate::test_utils::{fixture, open_gpkg_for_update};
use crate::test_utils::{fixture, open_dataset_for_update, open_gpkg_for_update};
use crate::vector::{Geometry, LayerAccess};
use crate::Dataset;

Expand Down Expand Up @@ -241,4 +261,30 @@ mod tests {
let mut ds = Dataset::open(fixture("roads.geojson")).unwrap();
assert!(ds.start_transaction().is_err());
}

#[test]
fn test_maybe_transaction() {
let (_temp_path, mut ds) = open_gpkg_for_update(&fixture("poly.gpkg"));
let orig_feature_count = ds.layer(0).unwrap().feature_count();

let res = ds.maybe_transaction(|d| {
let mut layer = d.layer(0).unwrap();
layer.create_feature(polygon())
});
assert!(res.is_ok());
assert_eq!(ds.layer(0).unwrap().feature_count(), orig_feature_count + 1);
}

#[test]
fn test_maybe_transaction_unsupported() {
let (_temp_path, mut ds) = open_dataset_for_update(&fixture("roads.geojson"));
let orig_feature_count = ds.layer(0).unwrap().feature_count();

let res = ds.maybe_transaction(|d| {
let mut layer = d.layer(0).unwrap();
layer.create_feature(polygon())
});
assert!(res.is_ok());
assert_eq!(ds.layer(0).unwrap().feature_count(), orig_feature_count + 1);
}
}

0 comments on commit b3e28f5

Please sign in to comment.