Skip to content

Commit 29a0147

Browse files
danipozodanipozo
committed
feat(storage): Support composite primary keys
Signed-off-by: Dani Pozo <[email protected]>
1 parent a85b544 commit 29a0147

File tree

9 files changed

+74
-31
lines changed

9 files changed

+74
-31
lines changed

‎src/catalog/column.rs

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -126,17 +126,14 @@ impl ColumnCatalog {
126126
}
127127

128128
/// Find the id of the sort key among column catalogs
129-
pub fn find_sort_key_id(column_infos: &[ColumnCatalog]) -> Option<usize> {
130-
let mut key = None;
129+
pub fn find_sort_key_id(column_infos: &[ColumnCatalog]) -> Vec<usize> {
130+
let mut keys = vec![];
131131
for (id, column_info) in column_infos.iter().enumerate() {
132132
if column_info.is_primary() {
133-
if key.is_some() {
134-
panic!("only one primary key is supported");
135-
}
136-
key = Some(id);
133+
keys.push(id);
137134
}
138135
}
139-
key
136+
keys
140137
}
141138

142139
#[cfg(test)]

‎src/storage/memory/table.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ pub struct InMemoryTable {
1616
pub(super) table_ref_id: TableRefId,
1717
pub(super) columns: Arc<[ColumnCatalog]>,
1818
pub(super) inner: InMemoryTableInnerRef,
19+
pub(super) ordered_pk_ids: Vec<ColumnId>,
1920
}
2021

2122
pub(super) struct InMemoryTableInner {
@@ -58,6 +59,7 @@ impl InMemoryTable {
5859
table_ref_id,
5960
columns: columns.into(),
6061
inner: Arc::new(RwLock::new(InMemoryTableInner::new())),
62+
ordered_pk_ids: Vec::new(),
6163
}
6264
}
6365
}
@@ -84,4 +86,8 @@ impl Table for InMemoryTable {
8486
async fn update(&self) -> StorageResult<InMemoryTransaction> {
8587
InMemoryTransaction::start(self)
8688
}
89+
90+
fn ordered_pk_ids(&self) -> Vec<ColumnId> {
91+
self.ordered_pk_ids.clone()
92+
}
8793
}

‎src/storage/memory/transaction.rs

Lines changed: 28 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,8 @@ use itertools::Itertools;
77

88
use super::table::InMemoryTableInnerRef;
99
use super::{InMemoryRowHandler, InMemoryTable, InMemoryTxnIterator};
10-
use crate::array::{ArrayBuilderImpl, ArrayImplBuilderPickExt, ArrayImplSortExt, DataChunk};
11-
use crate::catalog::{find_sort_key_id, ColumnCatalog};
12-
use crate::storage::{ScanOptions, StorageColumnRef, StorageResult, Transaction};
10+
use crate::array::{ArrayBuilderImpl, ArrayImplBuilderPickExt, DataChunk};
11+
use crate::storage::{ScanOptions, StorageColumnRef, StorageResult, Table, Transaction};
1312

1413
/// A transaction running on `InMemoryStorage`.
1514
pub struct InMemoryTransaction {
@@ -34,31 +33,42 @@ pub struct InMemoryTransaction {
3433
/// Snapshot of all deleted rows
3534
deleted_rows: Arc<HashSet<usize>>,
3635

37-
/// All information about columns
38-
column_infos: Arc<[ColumnCatalog]>,
36+
/// Ordered primary key indexes in `column_infos`
37+
ordered_pk_idx: Vec<usize>,
3938
}
4039

4140
impl InMemoryTransaction {
4241
pub(super) fn start(table: &InMemoryTable) -> StorageResult<Self> {
4342
let inner = table.inner.read().unwrap();
43+
let ordered_pk_idx = table
44+
.ordered_pk_ids()
45+
.iter()
46+
.map(|id| {
47+
table
48+
.columns
49+
.iter()
50+
.position(|c| c.id() == *id)
51+
.expect("Malformed table object")
52+
})
53+
.collect_vec();
4454
Ok(Self {
4555
finished: false,
4656
buffer: vec![],
4757
delete_buffer: vec![],
4858
table: table.inner.clone(),
4959
snapshot: Arc::new(inner.get_all_chunks()),
5060
deleted_rows: Arc::new(inner.get_all_deleted_rows()),
51-
column_infos: table.columns.clone(),
61+
ordered_pk_idx,
5262
})
5363
}
5464
}
5565

5666
/// If primary key is found in [`ColumnCatalog`], sort all in-memory data using that key.
5767
fn sort_datachunk_by_pk(
5868
chunks: &Arc<Vec<DataChunk>>,
59-
column_infos: &[ColumnCatalog],
69+
ordered_pk_idx: &[usize],
6070
) -> Arc<Vec<DataChunk>> {
61-
if let Some(sort_key_id) = find_sort_key_id(column_infos) {
71+
if !ordered_pk_idx.is_empty() {
6272
if chunks.is_empty() {
6373
return chunks.clone();
6474
}
@@ -78,7 +88,15 @@ fn sort_datachunk_by_pk(
7888
.into_iter()
7989
.map(|builder| builder.finish())
8090
.collect_vec();
81-
let sorted_index = arrays[sort_key_id].get_sorted_indices();
91+
92+
let pk_arrays = Vec::from(ordered_pk_idx)
93+
.iter()
94+
.map(|idx| &arrays[*idx])
95+
.collect_vec();
96+
let pk_array = itertools::izip!(pk_arrays).collect_vec();
97+
let sorted_index = (0..pk_array.len())
98+
.sorted_by_key(|idx| pk_array[*idx])
99+
.collect_vec();
82100

83101
let chunk = arrays
84102
.into_iter()
@@ -109,7 +127,7 @@ impl Transaction for InMemoryTransaction {
109127
assert!(!opts.reversed, "reverse iterator is not supported for now");
110128

111129
let snapshot = if opts.is_sorted {
112-
sort_datachunk_by_pk(&self.snapshot, &self.column_infos)
130+
sort_datachunk_by_pk(&self.snapshot, &self.ordered_pk_idx)
113131
} else {
114132
self.snapshot.clone()
115133
};

‎src/storage/mod.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,9 @@ pub trait Table: Sync + Send + Clone + 'static {
106106

107107
/// Get table id
108108
fn table_id(&self) -> TableRefId;
109+
110+
/// Get primary key
111+
fn ordered_pk_ids(&self) -> Vec<ColumnId>;
109112
}
110113

111114
/// Reference to a column.

‎src/storage/secondary/compactor.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -82,11 +82,11 @@ impl Compactor {
8282
);
8383
}
8484

85-
let sort_key = find_sort_key_id(&table.columns);
86-
let mut iter: SecondaryIterator = if let Some(sort_key) = sort_key {
85+
let sort_keys = find_sort_key_id(&table.columns);
86+
let mut iter: SecondaryIterator = if !sort_keys.is_empty() {
8787
MergeIterator::new(
8888
iters.into_iter().map(|iter| iter.into()).collect_vec(),
89-
vec![sort_key],
89+
sort_keys,
9090
)
9191
.into()
9292
} else {

‎src/storage/secondary/manifest.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -206,7 +206,7 @@ impl SecondaryStorage {
206206
table_name.clone(),
207207
column_descs.to_vec(),
208208
false,
209-
ordered_pk_ids,
209+
ordered_pk_ids.clone(),
210210
)
211211
.map_err(|_| TracedStorageError::duplicated("table", table_name))?;
212212

@@ -222,6 +222,7 @@ impl SecondaryStorage {
222222
self.version.clone(),
223223
self.block_cache.clone(),
224224
self.txn_mgr.clone(),
225+
ordered_pk_ids,
225226
);
226227
self.tables.write().insert(id, table);
227228

‎src/storage/secondary/rowset/mem_rowset.rs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -174,9 +174,10 @@ impl SecondaryMemRowsetImpl {
174174
column_options: ColumnBuilderOptions,
175175
rowset_id: u32,
176176
) -> Self {
177-
if let Some(sort_key_idx) = find_sort_key_id(&columns) {
177+
let sort_keys = find_sort_key_id(&columns);
178+
if !sort_keys.is_empty() {
178179
Self::BTree(SecondaryMemRowset::<BTreeMapMemTable> {
179-
mem_table: BTreeMapMemTable::new(columns.clone(), sort_key_idx),
180+
mem_table: BTreeMapMemTable::new(columns.clone(), sort_keys[0]),
180181
rowset_builder: RowsetBuilder::new(columns, column_options),
181182
rowset_id,
182183
})

‎src/storage/secondary/table.rs

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ use crate::storage::Table;
1414
/// A table in Secondary engine.
1515
///
1616
/// As `SecondaryStorage` holds the reference to `SecondaryTable`, we cannot store
17-
/// `Arc<SecondaryStorage>` inside `SecondaryTable`. This sturct only contains necessary information
17+
/// `Arc<SecondaryStorage>` inside `SecondaryTable`. This struct only contains necessary information
1818
/// to decode the columns of the table.
1919
#[derive(Clone)]
2020
pub struct SecondaryTable {
@@ -27,6 +27,9 @@ pub struct SecondaryTable {
2727
/// Mapping from [`ColumnId`] to column index in `columns`.
2828
pub column_map: HashMap<ColumnId, usize>,
2929

30+
/// Ordered list of primary key columns.
31+
pub ordered_pk_ids: Vec<ColumnId>,
32+
3033
/// Root directory of the storage
3134
pub storage_options: Arc<StorageOptions>,
3235

@@ -54,6 +57,7 @@ impl SecondaryTable {
5457
version: Arc<VersionManager>,
5558
block_cache: Cache<BlockCacheKey, Block>,
5659
txn_mgr: Arc<TransactionManager>,
60+
ordered_pk_ids: Vec<ColumnId>,
5761
) -> Self {
5862
Self {
5963
columns: columns.into(),
@@ -68,6 +72,7 @@ impl SecondaryTable {
6872
version,
6973
block_cache,
7074
txn_mgr,
75+
ordered_pk_ids,
7176
}
7277
}
7378

@@ -126,4 +131,8 @@ impl Table for SecondaryTable {
126131
async fn update(&self) -> StorageResult<SecondaryTransaction> {
127132
SecondaryTransaction::start(self, false, true).await
128133
}
134+
135+
fn ordered_pk_ids(&self) -> Vec<ColumnId> {
136+
self.ordered_pk_ids.clone()
137+
}
129138
}

‎src/storage/secondary/transaction.rs

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -259,15 +259,23 @@ impl SecondaryTransaction {
259259
let final_iter = if iters.len() == 1 {
260260
iters.pop().unwrap().into()
261261
} else if opts.is_sorted {
262-
let sort_key = find_sort_key_id(&self.table.columns);
263-
if let Some(sort_key) = sort_key {
264-
let real_col_idx = col_idx.iter().position(|x| match x {
265-
StorageColumnRef::Idx(y) => *y as usize == sort_key,
266-
_ => false,
267-
});
262+
let sort_keys = find_sort_key_id(&self.table.columns);
263+
if !sort_keys.is_empty() {
264+
let real_col_idx = sort_keys
265+
.iter()
266+
.map(|id| {
267+
col_idx
268+
.iter()
269+
.position(|x| match x {
270+
StorageColumnRef::Idx(y) => *y as usize == *id,
271+
_ => false,
272+
})
273+
.expect("sorting key not in column list")
274+
})
275+
.collect_vec();
268276
MergeIterator::new(
269277
iters.into_iter().map(|iter| iter.into()).collect_vec(),
270-
vec![real_col_idx.expect("sort key not in column list")],
278+
real_col_idx,
271279
)
272280
.into()
273281
} else {

0 commit comments

Comments
 (0)