use async_trait::async_trait;
use time::OffsetDateTime;
use time::format_description::FormatItem;
use time::macros::format_description;
use futures::stream::BoxStream;
use crate::arrow::array::RecordBatch;
use crate::error::{Error, Result};
use crate::planner::TimeRange;
use crate::watermark::{ArchivalWindow, TieringWatermark};
pub type BatchStream = BoxStream<'static, Result<RecordBatch>>;
pub fn stream_of(batches: Vec<RecordBatch>) -> BatchStream {
Box::pin(futures::stream::iter(batches.into_iter().map(Ok)))
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ScanSpec {
merge_key: Vec<String>,
extra: Vec<String>,
chunk_rows: Option<usize>,
}
impl ScanSpec {
pub fn new(merge_key: Vec<String>, extra: Vec<String>) -> Self {
Self {
merge_key,
extra,
chunk_rows: None,
}
}
pub fn with_chunk_rows(mut self, rows: usize) -> Self {
self.chunk_rows = Some(rows.max(1));
self
}
pub fn chunk_rows(&self) -> Option<usize> {
self.chunk_rows
}
pub fn core() -> Self {
Self::new(
crate::encode::schema::MERGE_KEY
.iter()
.map(|s| (*s).to_string())
.collect(),
Vec::new(),
)
}
pub fn merge_key(&self) -> &[String] {
&self.merge_key
}
pub fn extra(&self) -> &[String] {
&self.extra
}
pub fn cursor_columns(&self) -> Vec<String> {
let mut columns: Vec<String> = crate::encode::schema::SORT_COLUMNS
.iter()
.map(|s| (*s).to_string())
.collect();
for column in &self.merge_key {
if !columns.contains(column) {
columns.push(column.clone());
}
}
columns.push(crate::encode::schema::col::VERSION.to_string());
columns
}
}
pub fn partitions_ahead(
starts: &[OffsetDateTime],
now: OffsetDateTime,
step: time::Duration,
) -> usize {
let frontier = crate::watermark::align_to_step(now, step);
starts.iter().filter(|start| **start >= frontier).count()
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct WriteHints {
pub distinct_malo_ids: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct PartitionId {
table: String,
start: OffsetDateTime,
}
const PARTITION_SUFFIX: &[FormatItem<'_>] =
format_description!("[year]_[month]_[day]_[hour][minute]");
impl PartitionId {
pub fn new(table: impl Into<String>, start: OffsetDateTime) -> Self {
Self {
table: table.into(),
start,
}
}
pub fn for_window(table: impl Into<String>, window: ArchivalWindow) -> Self {
Self::new(table, window.from())
}
pub fn table(&self) -> &str {
&self.table
}
pub fn start(&self) -> OffsetDateTime {
self.start
}
pub fn from_relation_name(table: &str, relation: &str) -> Result<Self> {
let suffix = relation
.strip_prefix(table)
.and_then(|r| r.strip_prefix('_'))
.ok_or_else(|| {
Error::decode(
"partition name",
format!("{relation:?} is not a partition of {table:?}"),
)
})?;
let start = time::PrimitiveDateTime::parse(suffix, PARTITION_SUFFIX)
.map_err(|e| Error::decode("partition name", format!("{suffix:?}: {e}")))?
.assume_utc();
Ok(Self {
table: table.to_string(),
start,
})
}
pub fn relation_name(&self) -> Result<String> {
let suffix = self
.start
.format(PARTITION_SUFFIX)
.map_err(|e| Error::encode("partition name", e.to_string()))?;
Ok(format!("{}_{suffix}", self.table))
}
}
#[async_trait]
pub trait ArchiveLease: Send + Sync + std::fmt::Debug {
async fn release(self: Box<Self>) -> Result<()>;
}
#[async_trait]
pub trait HotStore: Send + Sync {
async fn try_archive_lease(&self, _table: &str) -> Result<Option<Box<dyn ArchiveLease>>> {
Ok(Some(Box::new(UnenforcedLease)))
}
async fn ensure_partitions(
&self,
table: &str,
from: OffsetDateTime,
until: OffsetDateTime,
step: time::Duration,
) -> Result<Vec<PartitionId>>;
async fn create_tables(
&self,
table: &str,
merge_key: &[String],
extra: &[crate::arrow::datatypes::Field],
) -> Result<()>;
async fn append(
&self,
table: &str,
merge_key: &[String],
batches: &[RecordBatch],
) -> Result<u64>;
async fn scan_range(
&self,
table: &str,
range: TimeRange,
spec: &ScanSpec,
) -> Result<BatchStream>;
async fn partition_exists(&self, partition: &PartitionId) -> Result<bool>;
async fn partition_starts(&self, _table: &str) -> Result<Option<Vec<OffsetDateTime>>> {
Ok(None)
}
async fn detach_partition(&self, partition: &PartitionId) -> Result<()>;
async fn scan_detached(&self, partition: &PartitionId, spec: &ScanSpec) -> Result<BatchStream>;
async fn distinct_malo_ids(&self, _partition: &PartitionId) -> Result<Option<u64>> {
Ok(None)
}
async fn drop_partition(&self, partition: &PartitionId) -> Result<()>;
async fn append_reporting(
&self,
table: &str,
merge_key: &[String],
batches: &[RecordBatch],
) -> Result<Vec<crate::session::Displacement>>;
async fn drop_table(&self, table: &str) -> Result<()>;
async fn orphaned_partitions(&self, table: &str) -> Result<Vec<PartitionId>>;
async fn invariant_violations(&self, table: &str, watermark: TieringWatermark) -> Result<u64>;
}
#[async_trait]
pub trait ColdStore: Send + Sync {
async fn create_tables(
&self,
table: &str,
identity: &[String],
extra: &[crate::arrow::datatypes::Field],
) -> Result<()>;
async fn purge_table(&self, table: &str) -> Result<()>;
async fn watermark(&self, table: &str) -> Result<TieringWatermark>;
async fn append_and_commit(
&self,
table: &str,
batches: BatchStream,
hints: WriteHints,
window: ArchivalWindow,
) -> Result<CommitInfo>;
async fn expire_snapshots(
&self,
table: &str,
retain_for: time::Duration,
retain_last: usize,
now: OffsetDateTime,
) -> Result<usize>;
async fn append_only(
&self,
table: &str,
batches: BatchStream,
hints: WriteHints,
) -> Result<CommitInfo>;
async fn reassert_watermark(&self, _table: &str) -> Result<Option<CommitInfo>> {
Ok(None)
}
async fn version_stats(
&self,
_table: &str,
_range: (OffsetDateTime, OffsetDateTime),
) -> Result<Vec<Option<crate::planner::VersionStats>>> {
Ok(vec![None])
}
async fn snapshot_provider(
&self,
_table: &str,
_at: crate::planner::SnapshotSelector,
) -> Result<std::sync::Arc<dyn datafusion::catalog::TableProvider>> {
Err(Error::config(
"this cold store cannot pin a past snapshot, so as-of reads are unavailable",
))
}
async fn snapshots(&self, _table: &str) -> Result<Vec<SnapshotInfo>> {
Ok(Vec::new())
}
async fn stored_schema(
&self,
_table: &str,
) -> Result<Option<crate::arrow::datatypes::SchemaRef>> {
Ok(None)
}
}
#[derive(Debug)]
pub struct UnenforcedLease;
#[async_trait]
impl ArchiveLease for UnenforcedLease {
async fn release(self: Box<Self>) -> Result<()> {
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotInfo {
pub snapshot_id: i64,
pub committed_at: OffsetDateTime,
pub watermark: Option<TieringWatermark>,
pub rows: Option<u64>,
}
#[async_trait]
impl<T: HotStore + ?Sized> HotStore for std::sync::Arc<T> {
async fn try_archive_lease(&self, table: &str) -> Result<Option<Box<dyn ArchiveLease>>> {
(**self).try_archive_lease(table).await
}
async fn ensure_partitions(
&self,
table: &str,
from: OffsetDateTime,
until: OffsetDateTime,
step: time::Duration,
) -> Result<Vec<PartitionId>> {
(**self).ensure_partitions(table, from, until, step).await
}
async fn create_tables(
&self,
table: &str,
merge_key: &[String],
extra: &[crate::arrow::datatypes::Field],
) -> Result<()> {
(**self).create_tables(table, merge_key, extra).await
}
async fn append(
&self,
table: &str,
merge_key: &[String],
batches: &[RecordBatch],
) -> Result<u64> {
(**self).append(table, merge_key, batches).await
}
async fn scan_range(
&self,
table: &str,
range: TimeRange,
spec: &ScanSpec,
) -> Result<BatchStream> {
(**self).scan_range(table, range, spec).await
}
async fn partition_exists(&self, partition: &PartitionId) -> Result<bool> {
(**self).partition_exists(partition).await
}
async fn partition_starts(&self, table: &str) -> Result<Option<Vec<OffsetDateTime>>> {
(**self).partition_starts(table).await
}
async fn detach_partition(&self, partition: &PartitionId) -> Result<()> {
(**self).detach_partition(partition).await
}
async fn scan_detached(&self, partition: &PartitionId, spec: &ScanSpec) -> Result<BatchStream> {
(**self).scan_detached(partition, spec).await
}
async fn append_reporting(
&self,
table: &str,
merge_key: &[String],
batches: &[RecordBatch],
) -> Result<Vec<crate::session::Displacement>> {
(**self).append_reporting(table, merge_key, batches).await
}
async fn drop_table(&self, table: &str) -> Result<()> {
(**self).drop_table(table).await
}
async fn distinct_malo_ids(&self, partition: &PartitionId) -> Result<Option<u64>> {
(**self).distinct_malo_ids(partition).await
}
async fn drop_partition(&self, partition: &PartitionId) -> Result<()> {
(**self).drop_partition(partition).await
}
async fn orphaned_partitions(&self, table: &str) -> Result<Vec<PartitionId>> {
(**self).orphaned_partitions(table).await
}
async fn invariant_violations(&self, table: &str, watermark: TieringWatermark) -> Result<u64> {
(**self).invariant_violations(table, watermark).await
}
}
#[async_trait]
impl<T: ColdStore + ?Sized> ColdStore for std::sync::Arc<T> {
async fn create_tables(
&self,
table: &str,
identity: &[String],
extra: &[crate::arrow::datatypes::Field],
) -> Result<()> {
(**self).create_tables(table, identity, extra).await
}
async fn purge_table(&self, table: &str) -> Result<()> {
(**self).purge_table(table).await
}
async fn watermark(&self, table: &str) -> Result<TieringWatermark> {
(**self).watermark(table).await
}
async fn append_and_commit(
&self,
table: &str,
batches: BatchStream,
hints: WriteHints,
window: ArchivalWindow,
) -> Result<CommitInfo> {
(**self)
.append_and_commit(table, batches, hints, window)
.await
}
async fn expire_snapshots(
&self,
table: &str,
retain_for: time::Duration,
retain_last: usize,
now: OffsetDateTime,
) -> Result<usize> {
(**self)
.expire_snapshots(table, retain_for, retain_last, now)
.await
}
async fn version_stats(
&self,
table: &str,
range: (OffsetDateTime, OffsetDateTime),
) -> Result<Vec<Option<crate::planner::VersionStats>>> {
(**self).version_stats(table, range).await
}
async fn reassert_watermark(&self, table: &str) -> Result<Option<CommitInfo>> {
(**self).reassert_watermark(table).await
}
async fn append_only(
&self,
table: &str,
batches: BatchStream,
hints: WriteHints,
) -> Result<CommitInfo> {
(**self).append_only(table, batches, hints).await
}
async fn snapshot_provider(
&self,
table: &str,
at: crate::planner::SnapshotSelector,
) -> Result<std::sync::Arc<dyn datafusion::catalog::TableProvider>> {
(**self).snapshot_provider(table, at).await
}
async fn snapshots(&self, table: &str) -> Result<Vec<SnapshotInfo>> {
(**self).snapshots(table).await
}
async fn stored_schema(
&self,
table: &str,
) -> Result<Option<crate::arrow::datatypes::SchemaRef>> {
(**self).stored_schema(table).await
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CommitInfo {
pub snapshot_id: i64,
pub rows: u64,
pub watermark: TieringWatermark,
}
#[cfg(test)]
mod tests {
use super::*;
use time::macros::datetime;
#[test]
fn relation_name_is_derived_from_the_lower_bound() {
let p = PartitionId::new("readings", datetime!(2026-07-20 00:00 UTC));
assert_eq!(p.relation_name().unwrap(), "readings_2026_07_20_0000");
}
#[test]
fn sub_daily_partitions_do_not_collide() {
let a = PartitionId::new("readings", datetime!(2026-07-20 00:00 UTC));
let b = PartitionId::new("readings", datetime!(2026-07-20 06:00 UTC));
assert_ne!(a.relation_name().unwrap(), b.relation_name().unwrap());
}
#[test]
fn partition_for_window_uses_the_windows_lower_bound() {
let w = ArchivalWindow::new(
datetime!(2026-07-20 00:00 UTC),
datetime!(2026-07-21 00:00 UTC),
)
.unwrap();
let p = PartitionId::for_window("readings", w);
assert_eq!(p.start(), w.from());
assert_eq!(p.relation_name().unwrap(), "readings_2026_07_20_0000");
}
#[test]
fn relation_name_round_trips() {
let p = PartitionId::new("readings", datetime!(2026-07-20 06:30 UTC));
let name = p.relation_name().unwrap();
assert_eq!(
PartitionId::from_relation_name("readings", &name).unwrap(),
p
);
}
#[test]
fn from_relation_name_rejects_foreign_relations() {
assert!(PartitionId::from_relation_name("readings", "other_2026_07_20_0000").is_err());
assert!(PartitionId::from_relation_name("readings", "readings").is_err());
assert!(PartitionId::from_relation_name("readings", "readings_garbage").is_err());
}
#[test]
fn partitions_order_by_start() {
let mut v = [
PartitionId::new("readings", datetime!(2026-07-21 00:00 UTC)),
PartitionId::new("readings", datetime!(2026-07-19 00:00 UTC)),
PartitionId::new("readings", datetime!(2026-07-20 00:00 UTC)),
];
v.sort();
assert_eq!(v[0].start(), datetime!(2026-07-19 00:00 UTC));
assert_eq!(v[2].start(), datetime!(2026-07-21 00:00 UTC));
}
}