Skip to content

Commit

Permalink
DeltaObjectStore: move additional functionality into LogStore
Browse files Browse the repository at this point in the history
  • Loading branch information
dispanser committed Nov 6, 2023
1 parent 4493f56 commit 69f9f3d
Show file tree
Hide file tree
Showing 13 changed files with 118 additions and 126 deletions.
11 changes: 5 additions & 6 deletions crates/deltalake-core/src/delta_datafusion/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,6 @@ use crate::errors::{DeltaResult, DeltaTableError};
use crate::kernel::{Add, DataType as DeltaDataType, Invariant, PrimitiveType};
use crate::logstore::LogStoreRef;
use crate::protocol::{self};
use crate::storage::ObjectStoreRef;
use crate::table::builder::ensure_table_uri;
use crate::table::state::DeltaTableState;
use crate::{open_table, open_table_with_storage_options, DeltaTable};
Expand Down Expand Up @@ -358,10 +357,10 @@ impl PruningStatistics for DeltaTable {

// each delta table must register a specific object store, since paths are internally
// handled relative to the table root.
pub(crate) fn register_store(store: ObjectStoreRef, env: Arc<RuntimeEnv>) {
pub(crate) fn register_store(store: LogStoreRef, env: Arc<RuntimeEnv>) {
let object_store_url = store.object_store_url();
let url: &Url = object_store_url.as_ref();
env.register_object_store(url, store);
env.register_object_store(url, store.object_store());
}

pub(crate) fn logical_schema(
Expand Down Expand Up @@ -633,7 +632,7 @@ impl<'a> DeltaScanBuilder<'a> {
.create_physical_plan(
self.state,
FileScanConfig {
object_store_url: self.log_store.object_store().object_store_url(),
object_store_url: self.log_store.object_store_url(),
file_schema,
file_groups: file_groups.into_values().collect(),
statistics: self.snapshot.datafusion_table_statistics(),
Expand Down Expand Up @@ -685,7 +684,7 @@ impl TableProvider for DeltaTable {
filters: &[Expr],
limit: Option<usize>,
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
register_store(self.object_store(), session.runtime_env().clone());
register_store(self.log_store(), session.runtime_env().clone());
let filter_expr = conjunction(filters.iter().cloned());

let scan = DeltaScanBuilder::new(&self.state, self.log_store(), session)
Expand Down Expand Up @@ -763,7 +762,7 @@ impl TableProvider for DeltaTableProvider {
filters: &[Expr],
limit: Option<usize>,
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
register_store(self.log_store.object_store(), session.runtime_env().clone());
register_store(self.log_store.clone(), session.runtime_env().clone());
let filter_expr = conjunction(filters.iter().cloned());

let scan = DeltaScanBuilder::new(&self.snapshot, self.log_store.clone(), session)
Expand Down
5 changes: 5 additions & 0 deletions crates/deltalake-core/src/logstore/default_logstore.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
use std::sync::Arc;

use bytes::Bytes;
use datafusion::execution::object_store::ObjectStoreUrl;
use object_store::path::Path;
use url::Url;

Expand Down Expand Up @@ -71,4 +72,8 @@ impl LogStore for DefaultLogStore {
fn to_uri(&self, location: &Path) -> String {
super::to_uri(&self.location, location)
}

fn object_store_url(&self) -> ObjectStoreUrl {
super::object_store_url(&self.location)
}
}
76 changes: 73 additions & 3 deletions crates/deltalake-core/src/logstore/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,12 @@ use crate::{
};
use bytes::Bytes;
use log::debug;
use object_store::{path::Path, Error as ObjectStoreError, ObjectStore};
use object_store::{
path::Path, Error as ObjectStoreError, ObjectStore, Result as ObjectStoreResult,
};

#[cfg(feature = "datafusion")]
use datafusion::datasource::object_store::ObjectStoreUrl;

pub mod default_logstore;

Expand Down Expand Up @@ -69,6 +74,30 @@ pub trait LogStore: Sync + Send {
fn log_path(&self) -> &Path {
&DELTA_LOG_PATH
}

/// Check if the location is a delta table location
async fn is_delta_table_location(&self) -> ObjectStoreResult<bool> {
// TODO We should really be using HEAD here, but this fails in windows tests
let object_store = self.object_store();
let mut stream = object_store.list(Some(self.log_path())).await?;
if let Some(res) = stream.next().await {
match res {
Ok(_) => Ok(true),
Err(ObjectStoreError::NotFound { .. }) => Ok(false),
Err(err) => Err(err),
}
} else {
Ok(false)
}
}

#[cfg(feature = "datafusion")]
/// Generate a unique enough url to identify the store in datafusion.
/// The DF object store registry only cares about the scheme and the host of the url for
/// registering/fetching. In our case the scheme is hard-coded to "delta-rs", so to get a unique
/// host we convert the location from this `DeltaObjectStore` to a valid name, combining the
/// original scheme, host and path with invalid characters replaced.
fn object_store_url(&self) -> ObjectStoreUrl;
}

// TODO: maybe a bit of a hack, required to `#[derive(Debug)]` for the operation builders
Expand Down Expand Up @@ -111,6 +140,20 @@ fn to_uri(root: &Url, location: &Path) -> String {
}
}

fn object_store_url(location: &Url) -> ObjectStoreUrl {
// we are certain, that the URL can be parsed, since
// we make sure when we are parsing the table uri

use object_store::path::DELIMITER;
ObjectStoreUrl::parse(format!(
"delta-rs://{}-{}{}",
location.scheme(),
location.host_str().unwrap_or("-"),
location.path().replace(DELIMITER, "-").replace(':', "-")
))
.expect("Invalid object store url.")
}

/// Extract version from a file name in the delta log
pub fn extract_version_from_filename(name: &str) -> Option<i64> {
DELTA_LOG_REGEX
Expand All @@ -119,8 +162,7 @@ pub fn extract_version_from_filename(name: &str) -> Option<i64> {
}

async fn get_latest_version(log_store: &dyn LogStore, current_version: i64) -> DeltaResult<i64> {
let object_store = log_store.object_store();
let version_start = match get_last_checkpoint(&object_store).await {
let version_start = match get_last_checkpoint(log_store).await {
Ok(last_check_point) => last_check_point.version,
Err(ProtocolError::CheckpointNotFound) => {
// no checkpoint
Expand All @@ -140,6 +182,8 @@ async fn get_latest_version(log_store: &dyn LogStore, current_version: i64) -> D
let mut max_version: i64 = version_start;
let prefix = Some(log_store.log_path());
let offset_path = commit_uri_from_version(max_version);
// let log_store = log_store.clone();
let object_store = log_store.object_store();
let mut files = object_store.list_with_offset(prefix, &offset_path).await?;

while let Some(obj_meta) = files.next().await {
Expand Down Expand Up @@ -189,3 +233,29 @@ async fn write_commit_entry(
})?;
Ok(())
}

#[cfg(feature = "datafusion")]
#[cfg(test)]
mod tests {
use url::Url;

#[tokio::test]
async fn test_unique_object_store_url() {
for (location_1, location_2) in [
// Same scheme, no host, different path
("file:///path/to/table_1", "file:///path/to/table_2"),
// Different scheme/host, same path
("s3://my_bucket/path/to/table_1", "file:///path/to/table_1"),
// Same scheme, different host, same path
("s3://bucket_1/table_1", "s3://bucket_2/table_1"),
] {
let url_1 = Url::parse(location_1).unwrap();
let url_2 = Url::parse(location_2).unwrap();

assert_ne!(
super::object_store_url(&url_1).as_str(),
super::object_store_url(&url_2).as_str(),
);
}
}
}
3 changes: 2 additions & 1 deletion crates/deltalake-core/src/operations/create.rs
Original file line number Diff line number Diff line change
Expand Up @@ -291,7 +291,8 @@ impl std::future::IntoFuture for CreateBuilder {
Box::pin(async move {
let mode = this.mode.clone();
let (mut table, actions, operation) = this.into_table_and_actions()?;
let table_state = if table.object_store().is_delta_table_location().await? {
let log_store = table.log_store();
let table_state = if log_store.is_delta_table_location().await? {
match mode {
SaveMode::ErrorIfExists => return Err(CreateError::TableAlreadyExists.into()),
SaveMode::Append => return Err(CreateError::AppendNotAllowed.into()),
Expand Down
2 changes: 1 addition & 1 deletion crates/deltalake-core/src/operations/delete.rs
Original file line number Diff line number Diff line change
Expand Up @@ -278,7 +278,7 @@ impl std::future::IntoFuture for DeleteBuilder {
let session = SessionContext::new();

// If a user provides their own their DF state then they must register the store themselves
register_store(this.log_store.object_store().clone(), session.runtime_env());
register_store(this.log_store.clone(), session.runtime_env());

session.state()
});
Expand Down
2 changes: 1 addition & 1 deletion crates/deltalake-core/src/operations/merge.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1212,7 +1212,7 @@ impl std::future::IntoFuture for MergeBuilder {
let session = SessionContext::new();

// If a user provides their own their DF state then they must register the store themselves
register_store(this.log_store.object_store().clone(), session.runtime_env());
register_store(this.log_store.clone(), session.runtime_env());

session.state()
});
Expand Down
2 changes: 1 addition & 1 deletion crates/deltalake-core/src/operations/update.rs
Original file line number Diff line number Diff line change
Expand Up @@ -430,7 +430,7 @@ impl std::future::IntoFuture for UpdateBuilder {
let session = SessionContext::new();

// If a user provides their own their DF state then they must register the store themselves
register_store(this.log_store.object_store().clone(), session.runtime_env());
register_store(this.log_store.clone(), session.runtime_env());

session.state()
});
Expand Down
3 changes: 1 addition & 2 deletions crates/deltalake-core/src/operations/write.rs
Original file line number Diff line number Diff line change
Expand Up @@ -211,8 +211,7 @@ impl WriteBuilder {
}

async fn check_preconditions(&self) -> DeltaResult<Vec<Action>> {
let object_store = self.log_store.object_store();
match object_store.is_delta_table_location().await? {
match self.log_store.is_delta_table_location().await? {
true => {
let min_writer = self.snapshot.min_writer_version();
if min_writer > MAX_SUPPORTED_WRITER_VERSION {
Expand Down
33 changes: 15 additions & 18 deletions crates/deltalake-core/src/protocol/checkpoints.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ use crate::kernel::{
Action, Add as AddAction, DataType, Metadata, PrimitiveType, Protocol, StructField, StructType,
Txn,
};
use crate::storage::DeltaObjectStore;
use crate::logstore::LogStore;
use crate::table::state::DeltaTableState;
use crate::table::{CheckPoint, CheckPointBuilder};
use crate::{open_table_with_version, DeltaTable};
Expand Down Expand Up @@ -70,12 +70,7 @@ pub const CHECKPOINT_RECORD_BATCH_SIZE: usize = 5000;

/// Creates checkpoint at current table version
pub async fn create_checkpoint(table: &DeltaTable) -> Result<(), ProtocolError> {
create_checkpoint_for(
table.version(),
table.get_state(),
table.object_store().as_ref(),
)
.await?;
create_checkpoint_for(table.version(), table.get_state(), table.log_store.as_ref()).await?;
Ok(())
}

Expand All @@ -86,7 +81,7 @@ pub async fn cleanup_metadata(table: &DeltaTable) -> Result<usize, ProtocolError
Utc::now().timestamp_millis() - table.get_state().log_retention_millis();
cleanup_expired_logs_for(
table.version(),
table.object_store().as_ref(),
table.log_store.as_ref(),
log_retention_timestamp,
)
.await
Expand All @@ -103,7 +98,7 @@ pub async fn create_checkpoint_from_table_uri_and_cleanup(
let table = open_table_with_version(table_uri, version)
.await
.map_err(|err| ProtocolError::Generic(err.to_string()))?;
create_checkpoint_for(version, table.get_state(), table.object_store().as_ref()).await?;
create_checkpoint_for(version, table.get_state(), table.log_store.as_ref()).await?;

let enable_expired_log_cleanup =
cleanup.unwrap_or_else(|| table.get_state().enable_expired_log_cleanup());
Expand All @@ -120,27 +115,28 @@ pub async fn create_checkpoint_from_table_uri_and_cleanup(
pub async fn create_checkpoint_for(
version: i64,
state: &DeltaTableState,
storage: &DeltaObjectStore,
log_store: &dyn LogStore,
) -> Result<(), ProtocolError> {
// TODO: checkpoints _can_ be multi-part... haven't actually found a good reference for
// an appropriate split point yet though so only writing a single part currently.
// See https://github.com/delta-io/delta-rs/issues/288
let last_checkpoint_path = storage.log_path().child("_last_checkpoint");
let last_checkpoint_path = log_store.log_path().child("_last_checkpoint");

debug!("Writing parquet bytes to checkpoint buffer.");
let (checkpoint, parquet_bytes) = parquet_bytes_from_state(state)?;

let file_name = format!("{version:020}.checkpoint.parquet");
let checkpoint_path = storage.log_path().child(file_name);
let checkpoint_path = log_store.log_path().child(file_name);

let object_store = log_store.object_store();
debug!("Writing checkpoint to {:?}.", checkpoint_path);
storage.put(&checkpoint_path, parquet_bytes).await?;
object_store.put(&checkpoint_path, parquet_bytes).await?;

let last_checkpoint_content: Value = serde_json::to_value(checkpoint)?;
let last_checkpoint_content = bytes::Bytes::from(serde_json::to_vec(&last_checkpoint_content)?);

debug!("Writing _last_checkpoint to {:?}.", last_checkpoint_path);
storage
object_store
.put(&last_checkpoint_path, last_checkpoint_content)
.await?;

Expand All @@ -151,7 +147,7 @@ pub async fn create_checkpoint_for(
/// and less than the specified version.
pub async fn cleanup_expired_logs_for(
until_version: i64,
storage: &DeltaObjectStore,
log_store: &dyn LogStore,
cutoff_timestamp: i64,
) -> Result<usize, ProtocolError> {
lazy_static! {
Expand All @@ -162,10 +158,11 @@ pub async fn cleanup_expired_logs_for(
// Feed a stream of candidate deletion files directly into the delete_stream
// function to try to improve the speed of cleanup and reduce the need for
// intermediate memory.
let deleted = storage
let object_store = log_store.object_store();
let deleted = object_store
.delete_stream(
storage
.list(Some(storage.log_path()))
object_store
.list(Some(log_store.log_path()))
.await?
// This predicate function will filter out any locations that don't
// match the given timestamp range
Expand Down
12 changes: 7 additions & 5 deletions crates/deltalake-core/src/protocol/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ use std::mem::take;

use crate::errors::DeltaResult;
use crate::kernel::{Add, CommitInfo, Metadata, Protocol, Remove};
use crate::storage::ObjectStoreRef;
use crate::logstore::LogStore;
use crate::table::CheckPoint;
use crate::table::DeltaTableMetaData;

Expand Down Expand Up @@ -601,14 +601,15 @@ pub enum OutputMode {
}

pub(crate) async fn get_last_checkpoint(
object_store: &ObjectStoreRef,
log_store: &dyn LogStore,
) -> Result<CheckPoint, ProtocolError> {
let last_checkpoint_path = Path::from_iter(["_delta_log", "_last_checkpoint"]);
debug!("loading checkpoint from {last_checkpoint_path}");
let object_store = log_store.object_store();
match object_store.get(&last_checkpoint_path).await {
Ok(data) => Ok(serde_json::from_slice(&data.bytes().await?)?),
Err(ObjectStoreError::NotFound { .. }) => {
match find_latest_check_point_for_version(object_store, i64::MAX).await {
match find_latest_check_point_for_version(log_store, i64::MAX).await {
Ok(Some(cp)) => Ok(cp),
_ => Err(ProtocolError::CheckpointNotFound),
}
Expand All @@ -618,7 +619,7 @@ pub(crate) async fn get_last_checkpoint(
}

pub(crate) async fn find_latest_check_point_for_version(
object_store: &ObjectStoreRef,
log_store: &dyn LogStore,
version: i64,
) -> Result<Option<CheckPoint>, ProtocolError> {
lazy_static! {
Expand All @@ -629,7 +630,8 @@ pub(crate) async fn find_latest_check_point_for_version(
}

let mut cp: Option<CheckPoint> = None;
let mut stream = object_store.list(Some(object_store.log_path())).await?;
let object_store = log_store.object_store();
let mut stream = object_store.list(Some(log_store.log_path())).await?;

while let Some(obj_meta) = stream.next().await {
// Exit early if any objects can't be listed.
Expand Down
6 changes: 2 additions & 4 deletions crates/deltalake-core/src/storage/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -232,10 +232,8 @@ pub(crate) fn configure_log_store(
options: impl Into<StorageOptions> + Clone,
) -> DeltaResult<LogStoreRef> {
let mut options = options.into();
let (scheme, _prefix) = ObjectStoreScheme::parse(&url, &mut options)?;
match scheme {
_ => Ok(Arc::new(DefaultLogStore::try_new(url, options)?)),
}
let (_scheme, _prefix) = ObjectStoreScheme::parse(&url, &mut options)?;
Ok(Arc::new(DefaultLogStore::try_new(url, options)?))
}

pub(crate) fn configure_store(
Expand Down
Loading

0 comments on commit 69f9f3d

Please sign in to comment.