feat(core)!: composite PK plumbing, Key API change, and Arrow filters (#453)

* feat(core)!: composite PK plumbing, Key API change, and Arrow filters

- Key API: replace `Key::to_arrow_datum` with `Key::to_arrow_datums` to supportmulti-component keys; update all builtin `Key` impls (ints/floats/string/date/time/timestamp/dynamic `Value`) to return a single-element Vec for single-PK types.
- Composite key showcase: add `Key2`/`Key2Ref` under `#[cfg(test)]` with `Encode`/`Decode`, `Ord`, `Hash`, `Key`/`KeyRef`; round-trip and ordering tests included.
- Parquet scan filtering: implement lexicographic composite range predicates in`get_range_filter` using `eq/gt/gt_eq/lt/lt_eq` with `and_kleene`/`or_kleene`; retain `_ts` upper-bound predicate.
- Scan plumbing: thread `pk_indices` through `SsTable::{get,scan}`, `Version::{query,streams,table_query}`, and `LevelStream`; update test callers. Build fixed projections as `[0, 1] ∪ pk_indices` in read paths.
- Docs: update composite PK RFC with “Current Status (2025-08-11)”; add draft RFC “Replace Schema Trait with Arrow Schema”.
- Clippy hygiene: remove needless range loops and redundant `.into_iter()`; add a focused `#[allow(clippy::too_many_arguments)]` on `Version::table_query`.

BREAKING CHANGE:
- `Key::to_arrow_datum` was replaced by `Key::to_arrow_datums(Vec<Arc<dyn Datum>>)`; external `Key` implementors must migrate to the new method (return one datum per PK component).

* refactor: remove unsafe code in src/ondisk/arrows.rs

Author: ethe

guide/src/contribution/composite_primary_keys.md

@@ -2,6 +2,15 @@
 
 This document outlines a practical, incremental plan to add composite (multi-column) primary key support to Tonbo while maintaining backward compatibility. It explains design goals, changes required across the codebase, and a step-by-step implementation and validation plan.
 
+## Current Status (2025-08-11)
+
+- Phase 1: Completed. Plural `Schema` APIs, fixed projections, and Parquet writer configuration are implemented for single-PK schemas.
+  - `Schema` now exposes `primary_key_indices()` and `primary_key_paths_and_sorting()` (src/record/mod.rs). Macro-generated single-PK schemas return one-element slices.
+  - Read paths build fixed projections as `[0, 1] ∪ PKs` using `primary_key_indices()` (src/lib.rs, src/transaction.rs).
+  - `DbOption::new` configures sorting columns (`_ts` then PKs) and enables stats + bloom filters for each PK column path (src/option.rs).
+- Phase 2: Not implemented. Composite key types under `src/record/key/composite/` are placeholders; derive macro still accepts only a single `#[record(primary_key)]` and generates a single-field key. No multi-PK trybuild/integration tests.
+- Phase 3: Not implemented. `DynSchema` remains single-PK (stores one `primary_index_arrow`, one `pk_path`, and sorting with a single PK column).
+
 ## Goals
 
 - Support multi-column primary keys with lexicographic ordering of PK components.

src/compaction/leveled.rs

@@ -289,6 +289,7 @@ where
                                 None,
                                 ProjectionMask::all(),
                                 None, // Default order for compaction
+                                instance.primary_key_indices(),
                             )
                             .await?,
                     });
@@ -307,6 +308,7 @@ where
                     level_fs.clone(),
                     ctx.parquet_lru.clone(),
                     None, // Default order for compaction
+                    instance.primary_key_indices(),
                 )
                 .ok_or(CompactionError::EmptyLevel)?;
 
@@ -333,6 +335,7 @@ where
                     level_l_fs.clone(),
                     ctx.parquet_lru.clone(),
                     None, // Default order for compaction
+                    instance.primary_key_indices(),
                 )
                 .ok_or(CompactionError::EmptyLevel)?;
 

src/lib.rs

@@ -749,6 +749,7 @@ where
                 TsRef::new(key, ts),
                 projection,
                 ctx.cache().clone(),
+                self.record_schema.primary_key_indices(),
             )
             .await?
             .map(|entry| Entry::RecordBatch(entry)))
@@ -1009,6 +1010,7 @@ where
                 self.limit,
                 self.projection,
                 self.order,
+                self.mem_storage.record_schema.primary_key_indices(),
             )
             .await?;
 
@@ -1070,6 +1072,7 @@ where
                 self.limit,
                 self.projection,
                 self.order,
+                self.mem_storage.record_schema.primary_key_indices(),
             )
             .await?;
         let merge_stream = MergeStream::from_vec(streams, self.ts, self.order).await?;

src/ondisk/arrows.rs

@@ -2,9 +2,10 @@ use std::ops::Bound;
 
 use arrow::{
     array::{BooleanArray, Datum},
-    buffer::BooleanBuffer,
-    compute::kernels::cmp::{gt, gt_eq, lt_eq},
-    error::ArrowError,
+    compute::kernels::{
+        boolean::{and_kleene, or_kleene},
+        cmp::{eq, gt, gt_eq, lt, lt_eq},
+    },
 };
 use parquet::{
     arrow::{
@@ -19,70 +20,130 @@ use crate::{
     version::timestamp::Timestamp,
 };
 
-unsafe fn get_range_bound_fn<R>(
-    range: Bound<&<R::Schema as Schema>::Key>,
-) -> (
-    Option<&'static <R::Schema as Schema>::Key>,
-    &'static (dyn Fn(&dyn Datum, &dyn Datum) -> Result<BooleanArray, ArrowError> + Sync),
-)
+enum BoundKind {
+    Lower { inclusive: bool },
+    Upper { inclusive: bool },
+}
+
+fn lower_bound_owned<R>(
+    b: Bound<&<R::Schema as Schema>::Key>,
+) -> (Option<<R::Schema as Schema>::Key>, BoundKind)
+where
+    R: Record,
+{
+    match b {
+        Bound::Included(k) => (Some(k.clone()), BoundKind::Lower { inclusive: true }),
+        Bound::Excluded(k) => (Some(k.clone()), BoundKind::Lower { inclusive: false }),
+        Bound::Unbounded => (None, BoundKind::Lower { inclusive: true }),
+    }
+}
+
+fn upper_bound_owned<R>(
+    b: Bound<&<R::Schema as Schema>::Key>,
+) -> (Option<<R::Schema as Schema>::Key>, BoundKind)
 where
     R: Record,
 {
-    let cmp: &'static (dyn Fn(&dyn Datum, &dyn Datum) -> Result<BooleanArray, ArrowError> + Sync);
-    let key = match range {
-        Bound::Included(key) => {
-            cmp = &gt_eq;
-            Some(&*(key as *const _))
-        }
-        Bound::Excluded(key) => {
-            cmp = &gt;
-            Some(&*(key as *const _))
-        }
-        Bound::Unbounded => {
-            cmp = &|this, _| {
-                let len = this.get().0.len();
-                Ok(BooleanArray::new(
-                    BooleanBuffer::collect_bool(len, |_| true),
-                    None,
-                ))
-            };
-            None
-        }
-    };
-    (key, cmp)
+    match b {
+        Bound::Included(k) => (Some(k.clone()), BoundKind::Upper { inclusive: true }),
+        Bound::Excluded(k) => (Some(k.clone()), BoundKind::Upper { inclusive: false }),
+        Bound::Unbounded => (None, BoundKind::Upper { inclusive: true }),
+    }
 }
 
-pub(crate) unsafe fn get_range_filter<R>(
+pub(crate) fn get_range_filter<R>(
     schema_descriptor: &SchemaDescriptor,
     range: (
         Bound<&<R::Schema as Schema>::Key>,
         Bound<&<R::Schema as Schema>::Key>,
     ),
     ts: Timestamp,
+    pk_indices: &[usize],
 ) -> RowFilter
 where
     R: Record,
 {
-    let (lower_key, lower_cmp) = get_range_bound_fn::<R>(range.0);
-    let (upper_key, upper_cmp) = get_range_bound_fn::<R>(range.1);
+    let (lower_key, lower_kind) = lower_bound_owned::<R>(range.0);
+    let (upper_key, upper_kind) = upper_bound_owned::<R>(range.1);
 
+    let ts_scalar = ts.to_arrow_scalar();
     let mut predictions: Vec<Box<dyn ArrowPredicate>> = vec![Box::new(ArrowPredicateFn::new(
         ProjectionMask::roots(schema_descriptor, [1]),
-        move |record_batch| lt_eq(record_batch.column(0), &ts.to_arrow_scalar() as &dyn Datum),
+        move |record_batch| lt_eq(record_batch.column(0), &ts_scalar as &dyn Datum),
     ))];
+
     if let Some(lower_key) = lower_key {
+        let pk_len = pk_indices.len();
         predictions.push(Box::new(ArrowPredicateFn::new(
-            ProjectionMask::roots(schema_descriptor, [2]),
+            ProjectionMask::roots(schema_descriptor, pk_indices.to_vec()),
             move |record_batch| {
-                lower_cmp(record_batch.column(0), lower_key.to_arrow_datum().as_ref())
+                let datums = lower_key.to_arrow_datums();
+                debug_assert_eq!(datums.len(), pk_len);
+                let n = datums.len();
+                let mut acc: Option<BooleanArray> = None;
+                for i in 0..n {
+                    let cmp_i = if i == n - 1 {
+                        match lower_kind {
+                            BoundKind::Lower { inclusive: true } => {
+                                gt_eq(record_batch.column(i), datums[i].as_ref())?
+                            }
+                            BoundKind::Lower { inclusive: false } => {
+                                gt(record_batch.column(i), datums[i].as_ref())?
+                            }
+                            _ => unreachable!(),
+                        }
+                    } else {
+                        gt(record_batch.column(i), datums[i].as_ref())?
+                    };
+                    let mut term = cmp_i;
+                    for (j, d) in datums.iter().enumerate().take(i) {
+                        let eq_j = eq(record_batch.column(j), d.as_ref())?;
+                        term = and_kleene(&term, &eq_j)?;
+                    }
+                    acc = Some(match acc {
+                        None => term,
+                        Some(prev) => or_kleene(&prev, &term)?,
+                    });
+                }
+                Ok(acc.expect("at least one key component"))
             },
         )));
     }
+
     if let Some(upper_key) = upper_key {
+        let pk_len = pk_indices.len();
         predictions.push(Box::new(ArrowPredicateFn::new(
-            ProjectionMask::roots(schema_descriptor, [2]),
+            ProjectionMask::roots(schema_descriptor, pk_indices.to_vec()),
             move |record_batch| {
-                upper_cmp(upper_key.to_arrow_datum().as_ref(), record_batch.column(0))
+                let datums = upper_key.to_arrow_datums();
+                debug_assert_eq!(datums.len(), pk_len);
+                let n = datums.len();
+                let mut acc: Option<BooleanArray> = None;
+                for i in 0..n {
+                    let cmp_i = if i == n - 1 {
+                        match upper_kind {
+                            BoundKind::Upper { inclusive: true } => {
+                                lt_eq(record_batch.column(i), datums[i].as_ref())?
+                            }
+                            BoundKind::Upper { inclusive: false } => {
+                                lt(record_batch.column(i), datums[i].as_ref())?
+                            }
+                            _ => unreachable!(),
+                        }
+                    } else {
+                        lt(record_batch.column(i), datums[i].as_ref())?
+                    };
+                    let mut term = cmp_i;
+                    for (j, d) in datums.iter().enumerate().take(i) {
+                        let eq_j = eq(record_batch.column(j), d.as_ref())?;
+                        term = and_kleene(&term, &eq_j)?;
+                    }
+                    acc = Some(match acc {
+                        None => term,
+                        Some(prev) => or_kleene(&prev, &term)?,
+                    });
+                }
+                Ok(acc.expect("at least one key component"))
             },
         )));
     }

src/ondisk/sstable.rs

@@ -96,13 +96,15 @@ where
         self,
         key: &TsRef<<R::Schema as Schema>::Key>,
         projection_mask: ProjectionMask,
+        pk_indices: &[usize],
     ) -> ParquetResult<Option<RecordBatchEntry<R>>> {
         self.scan(
             (Bound::Included(key.value()), Bound::Included(key.value())),
             key.ts(),
             Some(1),
             projection_mask,
             None, // Order doesn't matter for single-key get
+            pk_indices,
         )
         .await?
         .next()
@@ -120,6 +122,7 @@ where
         limit: Option<usize>,
         projection_mask: ProjectionMask,
         order: Option<Order>,
+        pk_indices: &[usize],
     ) -> Result<SsTableScan<'scan, R>, parquet::errors::ParquetError> {
         let builder = self
             .into_parquet_builder(limit, projection_mask.clone())
@@ -128,9 +131,8 @@ where
         let schema_descriptor = builder.metadata().file_metadata().schema_descr();
         let full_schema = builder.schema().clone();
 
-        // Safety: filter's lifetime relies on range's lifetime, sstable must not live longer than
-        // it
-        let filter = unsafe { get_range_filter::<R>(schema_descriptor, range, ts) };
+        // Build a row filter for ts and primary key range
+        let filter = get_range_filter::<R>(schema_descriptor, range, ts, pk_indices);
 
         Ok(SsTableScan::new(
             builder.with_row_filter(filter).build()?,
@@ -252,6 +254,7 @@ pub(crate) mod tests {
                             .unwrap(),
                         [0, 1, 2, 3],
                     ),
+                    TestSchema {}.primary_key_indices(),
                 )
                 .await
                 .unwrap()
@@ -271,6 +274,7 @@ pub(crate) mod tests {
                             .unwrap(),
                         [0, 1, 2, 4],
                     ),
+                    TestSchema {}.primary_key_indices(),
                 )
                 .await
                 .unwrap()
@@ -290,6 +294,7 @@ pub(crate) mod tests {
                             .unwrap(),
                         [0, 1, 2],
                     ),
+                    TestSchema {}.primary_key_indices(),
                 )
                 .await
                 .unwrap()
@@ -337,6 +342,7 @@ pub(crate) mod tests {
                         [0, 1, 2, 3],
                     ),
                     None,
+                    TestSchema {}.primary_key_indices(),
                 )
                 .await
                 .unwrap();
@@ -365,6 +371,7 @@ pub(crate) mod tests {
                         [0, 1, 2, 4],
                     ),
                     None,
+                    TestSchema {}.primary_key_indices(),
                 )
                 .await
                 .unwrap();
@@ -393,6 +400,7 @@ pub(crate) mod tests {
                         [0, 1, 2],
                     ),
                     None,
+                    TestSchema {}.primary_key_indices(),
                 )
                 .await
                 .unwrap();
@@ -442,6 +450,7 @@ pub(crate) mod tests {
                     None,
                     ProjectionMask::all(),
                     None, // Forward order (default)
+                    TestSchema {}.primary_key_indices(),
                 )
                 .await
                 .unwrap();
@@ -466,6 +475,7 @@ pub(crate) mod tests {
                     None,
                     ProjectionMask::all(),
                     Some(Order::Desc), // Reverse order
+                    TestSchema {}.primary_key_indices(),
                 )
                 .await
                 .unwrap();
@@ -520,6 +530,7 @@ pub(crate) mod tests {
                         [0, 1, 2, 3], // Include vu32, exclude vbool
                     ),
                     Some(Order::Desc),
+                    TestSchema {}.primary_key_indices(),
                 )
                 .await
                 .unwrap();
@@ -552,6 +563,7 @@ pub(crate) mod tests {
                         [0, 1, 2, 4], // Include vbool, exclude vu32
                     ),
                     Some(Order::Desc),
+                    TestSchema {}.primary_key_indices(),
                 )
                 .await
                 .unwrap();
@@ -608,6 +620,7 @@ pub(crate) mod tests {
                     None,
                     ProjectionMask::all(),
                     Some(Order::Desc),
+                    TestSchema {}.primary_key_indices(),
                 )
                 .await
                 .unwrap();
@@ -631,6 +644,7 @@ pub(crate) mod tests {
                     None,
                     ProjectionMask::all(),
                     Some(Order::Desc),
+                    TestSchema {}.primary_key_indices(),
                 )
                 .await
                 .unwrap();
@@ -678,6 +692,7 @@ pub(crate) mod tests {
                     None,
                     ProjectionMask::all(),
                     None, // Forward order
+                    TestSchema {}.primary_key_indices(),
                 )
                 .await
                 .unwrap();
@@ -698,6 +713,7 @@ pub(crate) mod tests {
                     None,
                     ProjectionMask::all(),
                     Some(Order::Desc), // Reverse order
+                    TestSchema {}.primary_key_indices(),
                 )
                 .await
                 .unwrap();