fix: reject cross-segment ExclusiveStartKey in parallel Scan

crates/engine/src/scan.rs

@@ -414,6 +414,34 @@ pub async fn handle_scan(
     // Validate ExclusiveStartKey matches the key schema
     if let Some(ref start_key) = input.exclusive_start_key {
         validate_scan_exclusive_start_key(start_key, &key_info, index_info.as_ref())?;
+
+        // Parallel scan: the ExclusiveStartKey must belong to the segment being
+        // scanned. Segment assignment is backend-private (each backend hashes or
+        // partitions differently), so the backend answers membership and the
+        // engine owns the refusal, uniformly across backends. Without this, a
+        // LastEvaluatedKey from one segment replayed against another returns a
+        // silently wrong page instead of an error. Message measured against the
+        // service (us-east-1).
+        if let (Some(seg), Some(total)) = (input.segment, input.total_segments) {
+            let in_segment = ctx
+                .storage
+                .scan_key_in_segment(
+                    &scan_key_info,
+                    start_key,
+                    seg,
+                    total,
+                    input.index_name.as_deref(),
+                )
+                .await
+                .map_err(storage_err_to_dynamo)?;
+            if !in_segment {
+                return Err(DynamoDbError::ValidationException(format!(
+                    "The provided starting key is invalid: Invalid ExclusiveStartKey. \
+                     Please use ExclusiveStartKey with correct Segment. \
+                     TotalSegments: {total} Segment: {seg}"
+                )));
+            }
+        }
     }
 
     // Validate begins_with operand types upfront (before any rows are scanned).

crates/storage-mongodb/src/data_engine.rs

@@ -226,6 +226,32 @@ impl DataEngine for MongoEngine {
         })
     }
 
+    fn scan_key_in_segment(
+        &self,
+        key_info: &TableKeyInfo,
+        key: &Item,
+        segment: i64,
+        total_segments: i64,
+        _index_name: Option<&str>,
+    ) -> BoxFuture<'_, Result<bool, StorageError>> {
+        // Same CRC32-of-pk assignment the scan's per-item segment filter uses,
+        // evaluated for the candidate key. `key_info` is scan-shaped, so
+        // `key_schema` names the pk of whichever table (data or index) the
+        // scan targets, matching what the filter hashes. Pure computation, no
+        // round trip.
+        let key_schema = key_info.key_schema.clone();
+        let key = key.clone();
+        Box::pin(async move {
+            let pk_text = composite_pk_to_text(&key, &key_schema)?;
+            let hash = crc32fast::hash(pk_text.as_bytes());
+            #[allow(clippy::cast_sign_loss)]
+            let total_u = total_segments as u32;
+            #[allow(clippy::cast_sign_loss)]
+            let seg_u = segment as u32;
+            Ok(hash % total_u == seg_u)
+        })
+    }
+
     fn transact_get_items(
         &self,
         ops: &[TransactGetOp<'_>],

crates/storage-postgres/src/data/data_engine.rs

@@ -7,6 +7,7 @@
 use extenddb_core::expression::{Expr, ExpressionMaps, KeyCondition, UpdateAction};
 use extenddb_core::types::{Item, TableKeyInfo};
 use extenddb_storage::error::StorageError;
+use extenddb_storage::util::pk_to_text;
 use extenddb_storage::{DataEngine, IdempotencyKey, StreamCapture, TransactGetOp, TransactWriteOp};
 use futures::future::BoxFuture;
 
@@ -179,6 +180,38 @@ impl DataEngine for PostgresEngine {
         })
     }
 
+    fn scan_key_in_segment(
+        &self,
+        key_info: &TableKeyInfo,
+        key: &Item,
+        segment: i64,
+        total_segments: i64,
+        _index_name: Option<&str>,
+    ) -> BoxFuture<'_, Result<bool, StorageError>> {
+        // Same hashtext-based assignment the scan predicate uses
+        // ((hashtext(pk)::bigint & 2147483647) % total = seg), evaluated for
+        // the candidate key. `key_info` is scan-shaped, so `key_schema[0]` is
+        // the pk of whichever table (data or index) the scan targets, matching
+        // the `pk` column the predicate hashes. Bigint bitmask instead of
+        // abs() per CB-20 / SP-SCN-002 (hashtext() == i32::MIN).
+        let pk_name = key_info.key_schema[0].attribute_name.clone();
+        let key = key.clone();
+        Box::pin(async move {
+            let Some(pk_value) = key.get(&pk_name) else {
+                // Schema validation upstream already rejects this shape.
+                return Ok(false);
+            };
+            let pk_text = pk_to_text(pk_value)?;
+            sqlx::query_scalar("SELECT (hashtext($1)::bigint & 2147483647) % $2 = $3")
+                .bind(pk_text.as_ref())
+                .bind(total_segments)
+                .bind(segment)
+                .fetch_one(&self.data_pool)
+                .await
+                .map_err(|e| StorageError::Internal(e.to_string()))
+        })
+    }
+
     fn transact_get_items(
         &self,
         ops: &[TransactGetOp<'_>],

crates/storage-sqlite/src/data/data_engine.rs

@@ -173,6 +173,29 @@ impl DataEngine for SqliteEngine {
         })
     }
 
+    fn scan_key_in_segment(
+        &self,
+        key_info: &TableKeyInfo,
+        key: &Item,
+        segment: i64,
+        total_segments: i64,
+        index_name: Option<&str>,
+    ) -> BoxFuture<'_, Result<bool, StorageError>> {
+        let key_info = key_info.clone();
+        let key = key.clone();
+        let index_name = index_name.map(str::to_owned);
+        Box::pin(async move {
+            self.scan_key_in_segment_impl(
+                &key_info,
+                &key,
+                segment,
+                total_segments,
+                index_name.as_deref(),
+            )
+            .await
+        })
+    }
+
     fn transact_get_items(
         &self,
         ops: &[TransactGetOp<'_>],

crates/storage-sqlite/src/data/query_scan.rs

@@ -311,6 +311,84 @@ fn resolve_attr_name(path: &[PathElement], maps: &ExpressionMaps) -> Option<Stri
 }
 
 /// Extract the base-table partition key text from a (combined) start key.
+impl SqliteEngine {
+    /// Whether `key` belongs to `segment` under the rowid partitioning the
+    /// scan predicate uses (`rowid % total_segments = segment`).
+    ///
+    /// rowid is storage identity, not key content, so membership is only
+    /// decidable for a key whose row still exists. A key that resolves to no
+    /// row returns `true` (indeterminate, permitted): DynamoDB accepts an
+    /// `ExclusiveStartKey` that has been deleted since the previous page, and
+    /// refusing it would reject legitimate resumptions. A row that does exist
+    /// answers definitively, which is what makes the engine's cross-segment
+    /// refusal discriminating on this backend.
+    pub(crate) async fn scan_key_in_segment_impl(
+        &self,
+        key_info: &TableKeyInfo,
+        key: &Item,
+        segment: i64,
+        total_segments: i64,
+        index_name: Option<&str>,
+    ) -> Result<bool, StorageError> {
+        let ddb_table = if let Some(idx_name) = index_name {
+            let info = self
+                .fetch_index_info_by_table_id(&key_info.table_id, idx_name)
+                .await?;
+            index_table_name(&info.index_id)
+        } else {
+            data_table_name(&key_info.table_id)
+        };
+
+        let pk_name = &key_info.key_schema[0].attribute_name;
+        let Some(pk_value) = key.get(pk_name) else {
+            // Schema validation upstream already rejects this shape.
+            return Ok(false);
+        };
+        let pk_text = pk_to_text(pk_value)?.into_owned();
+
+        let mut conditions: Vec<String> = vec!["pk = ?".to_owned()];
+        let mut binds: Vec<BoundValue> = vec![BoundValue::Text(pk_text)];
+
+        let sk_info_val = sk_info(&key_info.key_schema, &key_info.attribute_definitions);
+        if let Some((sk_name, sk_type)) = sk_info_val
+            && let Some(v) = key.get(sk_name)
+        {
+            conditions.push(format!("{} = ?", sk_column(sk_type)));
+            binds.push(sk_bound(&parse_sk(v, sk_type)?));
+        }
+
+        if index_name.is_some() {
+            conditions.push("base_pk = ?".to_owned());
+            binds.push(BoundValue::Text(base_pk_from_start_key(key, key_info)?));
+            if let Some((base_name, base_type)) =
+                sk_info(&key_info.base_key_schema, &key_info.attribute_definitions)
+                && let Some(v) = key.get(base_name)
+            {
+                conditions.push(format!("base_{} = ?", sk_column(base_type)));
+                binds.push(sk_bound(&parse_sk(v, base_type)?));
+            }
+        }
+
+        let sql = format!(
+            "SELECT rowid FROM {ddb_table} WHERE {}",
+            conditions.join(" AND ")
+        );
+        let mut query = sqlx::query_as::<_, (i64,)>(&sql);
+        for bound in binds {
+            query = super::bind_bound!(query, bound);
+        }
+        let row: Option<(i64,)> = query
+            .fetch_optional(&self.pool)
+            .await
+            .map_err(|e| StorageError::Internal(e.to_string()))?;
+
+        Ok(match row {
+            Some((rowid,)) => rowid.rem_euclid(total_segments) == segment,
+            None => true,
+        })
+    }
+}
+
 fn base_pk_from_start_key(
     start_key: &Item,
     key_info: &TableKeyInfo,

crates/storage/src/lib.rs

@@ -376,6 +376,35 @@ pub trait DataEngine: Send + Sync {
         index_name: Option<&str>,
     ) -> BoxFuture<'_, QueryResult>;
 
+    /// Whether `key` belongs to `segment` of a `total_segments`-way parallel
+    /// scan, under this backend's own segment-assignment function.
+    ///
+    /// Segment assignment is backend-private (Postgres hashes the partition
+    /// key with `hashtext`, MongoDB with CRC32, SQLite partitions by rowid),
+    /// so only the backend can answer this. The engine owns the rule and the
+    /// error: a `Scan` whose `ExclusiveStartKey` reports `false` here is
+    /// refused before storage is asked to scan, because resuming a segment
+    /// with another segment's key would otherwise return a silently wrong
+    /// (truncated or fabricated) page.
+    ///
+    /// `Ok(true)` means "cannot be proven foreign" rather than "proven local":
+    /// a backend whose assignment depends on storage identity rather than key
+    /// content (SQLite's rowid) MUST return `true` for a key it cannot resolve,
+    /// because DynamoDB permits an `ExclusiveStartKey` that no longer exists
+    /// and refusing it would reject legitimate resumptions after a delete.
+    ///
+    /// Required, not defaulted: a new backend must decide its answer
+    /// explicitly, since returning a constant `true` silently disables the
+    /// engine's cross-segment refusal for that backend.
+    fn scan_key_in_segment(
+        &self,
+        key_info: &TableKeyInfo,
+        key: &Item,
+        segment: i64,
+        total_segments: i64,
+        index_name: Option<&str>,
+    ) -> BoxFuture<'_, Result<bool, StorageError>>;
+
     /// The vector-search implementation, if this backend has one.
     ///
     /// Defaults to `None`, so a backend that has never heard of vector search is
@@ -816,6 +845,17 @@ mod vector_opt_out_tests {
     struct MinimalDataEngine;
 
     impl DataEngine for MinimalDataEngine {
+        fn scan_key_in_segment(
+            &self,
+            _key_info: &TableKeyInfo,
+            _key: &Item,
+            _segment: i64,
+            _total_segments: i64,
+            _index_name: Option<&str>,
+        ) -> BoxFuture<'_, Result<bool, StorageError>> {
+            Box::pin(async { Ok(true) })
+        }
+
         fn put_item(
             &self,
             _key_info: &TableKeyInfo,

tests/rust/src/scan.rs

@@ -324,3 +324,98 @@ async fn scan_filter_contains_binary() {
         .collect();
     assert_eq!(keys, vec!["i2"], "NOT contains b\"ell\" should match only i2");
 }
+
+/// Parallel scan validates that an `ExclusiveStartKey` belongs to the segment
+/// being scanned: a key is accepted as the start key for its own segment and
+/// rejected (ValidationException) for any other segment. Regression for the
+/// previously-missing cross-segment check, which silently returned a truncated
+/// or empty page instead of erroring.
+#[tokio::test]
+async fn parallel_scan_rejects_cross_segment_exclusive_start_key() {
+    use aws_sdk_dynamodb::types::{
+        AttributeDefinition, BillingMode, KeySchemaElement, KeyType, ScalarAttributeType,
+    };
+    let c = client();
+    let table = format!("ScanEskSeg_{}", ts());
+    c.create_table()
+        .table_name(&table)
+        .billing_mode(BillingMode::PayPerRequest)
+        .key_schema(
+            KeySchemaElement::builder()
+                .attribute_name(HASH_KEY_S)
+                .key_type(KeyType::Hash)
+                .build()
+                .unwrap(),
+        )
+        .attribute_definitions(
+            AttributeDefinition::builder()
+                .attribute_name(HASH_KEY_S)
+                .attribute_type(ScalarAttributeType::S)
+                .build()
+                .unwrap(),
+        )
+        .send()
+        .await
+        .unwrap();
+    wait_for_active(c, &table).await;
+
+    for i in 0..20 {
+        let mut item = std::collections::HashMap::new();
+        item.insert(HASH_KEY_S.into(), s(&i.to_string()));
+        c.put_item()
+            .table_name(&table)
+            .set_item(Some(item))
+            .send()
+            .await
+            .unwrap();
+    }
+
+    const TOTAL: i32 = 4;
+    // Find a segment that owns at least one item, and grab one of its keys.
+    let mut owning_seg = None;
+    let mut key = None;
+    for seg in 0..TOTAL {
+        let r = c
+            .scan()
+            .table_name(&table)
+            .total_segments(TOTAL)
+            .segment(seg)
+            .send()
+            .await
+            .unwrap();
+        if let Some(first) = r.items().first() {
+            owning_seg = Some(seg);
+            key = Some(first.clone());
+            break;
+        }
+    }
+    let owning_seg = owning_seg.expect("some segment must own an item");
+    let key = key.unwrap();
+
+    // Same-segment ESK is accepted.
+    c.scan()
+        .table_name(&table)
+        .total_segments(TOTAL)
+        .segment(owning_seg)
+        .set_exclusive_start_key(Some(key.clone()))
+        .send()
+        .await
+        .expect("ESK for its own segment should be accepted");
+
+    // Cross-segment ESK is rejected with a ValidationException.
+    let other_seg = (owning_seg + 1) % TOTAL;
+    let err = c
+        .scan()
+        .table_name(&table)
+        .total_segments(TOTAL)
+        .segment(other_seg)
+        .set_exclusive_start_key(Some(key))
+        .send()
+        .await
+        .expect_err("ESK from a different segment must be rejected");
+    let msg = format!("{err:?}");
+    assert!(
+        msg.contains("Invalid ExclusiveStartKey"),
+        "expected cross-segment ExclusiveStartKey rejection, got: {msg}"
+    );
+}