support _id filters without typescript (#24298)

support index filters on `_id` fields. don't change typescript yet so we can do a soft launch, testing in prod a bit before exposing it to customers.

the tricky part here is when it's a virtual table. The `_id` strings use the virtual table number but the underlying Intervals use the physical table number. Luckily, base32 preserves order, so ordering is all the same. When the `_id` filter is a valid virtual ID, it's easy to map it to the corresponding physical ID. But what if the `_id` filter isn't a virtual ID? Then we need to find a string in the physical ID space that corresponds to the same place in the order. Added a bunch of proptests so I think it's working.

GitOrigin-RevId: 6677664011760c805d32d87851c470c9d84876f9

Author: ldanilek

crates/common/src/bootstrap_model/index/database_index/indexed_fields.rs

@@ -2,6 +2,7 @@ use std::{
     collections::HashSet,
     convert::TryFrom,
     fmt::Display,
+    iter,
     ops::Deref,
 };
 
@@ -44,6 +45,10 @@ impl IndexedFields {
             .expect("Invalid _creationTime field path");
         IndexedFields(vec![field_path].into())
     }
+
+    pub fn iter_with_id(&self) -> impl Iterator<Item = &FieldPath> {
+        self.iter().chain(iter::once(&*ID_FIELD_PATH))
+    }
 }
 
 impl HeapSize for IndexedFields {

crates/common/src/query.rs

@@ -25,16 +25,19 @@ use sha2::{
 };
 use value::{
     heap_size::HeapSize,
+    id_v6::VirtualTableNumberMap,
     utils::display_sequence,
     val,
     ConvexObject,
     ConvexValue,
+    DeveloperDocumentId,
     TableId,
     TableIdAndTableNumber,
 };
 
 use crate::{
     bootstrap_model::index::database_index::IndexedFields,
+    document::ID_FIELD_PATH,
     index::IndexKeyBytes,
     interval::{
         BinaryKey,
@@ -176,26 +179,36 @@ pub struct IndexRange {
 }
 
 impl IndexRange {
-    pub fn compile(self, indexed_fields: IndexedFields) -> anyhow::Result<Interval> {
+    pub fn compile(
+        self,
+        indexed_fields: IndexedFields,
+        virtual_table_number_map: Option<VirtualTableNumberMap>,
+    ) -> anyhow::Result<Interval> {
         let index_name = self.index_name.clone();
         let SplitIndexRange {
             equalities,
             inequality,
-        } = self.split()?;
+        } = self.split()?.map_values(|field, v| {
+            if field == &*ID_FIELD_PATH {
+                map_id_value_to_tablet(v, virtual_table_number_map)
+            } else {
+                Ok(v)
+            }
+        })?;
 
         // Check that some permutation of the equality field paths + the (optional)
         // inequality field path is a prefix of the indexed paths.
         //
         // NB: `indexed_fields` does not include the implicit `_id` field at the end of
         // every index, so this omission prevents the user from using it in an
         // index expression.
-        let index_rank: BTreeMap<_, _> = indexed_fields[..]
-            .iter()
+        let index_rank: BTreeMap<_, _> = indexed_fields
+            .iter_with_id()
             .enumerate()
             .map(|(i, field_name)| (field_name, i))
             .collect();
         anyhow::ensure!(
-            index_rank.len() == indexed_fields.len(),
+            index_rank.len() == indexed_fields.iter_with_id().count(),
             "{index_name} has duplicate fields?"
         );
 
@@ -237,7 +250,7 @@ impl IndexRange {
 
         let query_fields = QueryFields(used_paths.clone());
 
-        let mut fields_iter = indexed_fields.iter();
+        let mut fields_iter = indexed_fields.iter_with_id();
         for field_path in used_paths {
             let matching_field = fields_iter.next().ok_or_else(|| {
                 invalid_index_range(&index_name, &indexed_fields, &query_fields, &field_path)
@@ -387,6 +400,49 @@ struct SplitIndexRange {
     inequality: Option<IndexInequality>,
 }
 
+impl SplitIndexRange {
+    pub fn map_values(
+        self,
+        f: impl Fn(&FieldPath, ConvexValue) -> anyhow::Result<ConvexValue>,
+    ) -> anyhow::Result<SplitIndexRange> {
+        let equalities = self
+            .equalities
+            .into_iter()
+            .map(|(field, value)| {
+                let new_value = match value.0 {
+                    Some(value) => MaybeValue(Some(f(&field, value)?)),
+                    None => MaybeValue(None),
+                };
+                anyhow::Ok((field, new_value))
+            })
+            .try_collect()?;
+        let inequality = self
+            .inequality
+            .map(|inequality| {
+                let start = match inequality.start {
+                    Bound::Unbounded => Bound::Unbounded,
+                    Bound::Included(value) => Bound::Included(f(&inequality.field_path, value)?),
+                    Bound::Excluded(value) => Bound::Excluded(f(&inequality.field_path, value)?),
+                };
+                let end = match inequality.end {
+                    Bound::Unbounded => Bound::Unbounded,
+                    Bound::Included(value) => Bound::Included(f(&inequality.field_path, value)?),
+                    Bound::Excluded(value) => Bound::Excluded(f(&inequality.field_path, value)?),
+                };
+                anyhow::Ok(IndexInequality {
+                    field_path: inequality.field_path,
+                    start,
+                    end,
+                })
+            })
+            .transpose()?;
+        Ok(SplitIndexRange {
+            equalities,
+            inequality,
+        })
+    }
+}
+
 struct IndexInequality {
     field_path: FieldPath,
     start: Bound<ConvexValue>,
@@ -403,6 +459,21 @@ impl Display for QueryFields {
     }
 }
 
+fn map_id_value_to_tablet(
+    value: ConvexValue,
+    virtual_table_number_map: Option<VirtualTableNumberMap>,
+) -> anyhow::Result<ConvexValue> {
+    let val = match (&value, virtual_table_number_map) {
+        (ConvexValue::String(id), Some(virtual_table_number_map)) => {
+            let mapped =
+                DeveloperDocumentId::map_string_between_table_numbers(id, virtual_table_number_map);
+            val!(mapped)
+        },
+        _ => value,
+    };
+    Ok(val)
+}
+
 fn already_defined_bound_error(
     bound_type: &str,
     field_path: &FieldPath,

crates/common/src/types/index.rs

@@ -11,13 +11,16 @@ use sync_types::identifier::{
 };
 use value::{
     heap_size::HeapSize,
+    id_v6::VirtualTableNumberMap,
     FieldName,
     InternalId,
     ResolvedDocumentId,
     TableId,
     TableIdAndTableNumber,
     TableIdentifier,
+    TableMapping,
     TableName,
+    VirtualTableMapping,
 };
 
 use crate::{
@@ -119,6 +122,35 @@ pub enum StableIndexName {
     Missing,
 }
 
+impl StableIndexName {
+    pub fn virtual_table_number_map(
+        &self,
+        table_mapping: &TableMapping,
+        virtual_table_mapping: &VirtualTableMapping,
+    ) -> anyhow::Result<Option<VirtualTableNumberMap>> {
+        match self {
+            StableIndexName::Physical(index_name) => {
+                let table_number =
+                    table_mapping.inject_table_number()(*index_name.table())?.table_number;
+                Ok(Some(VirtualTableNumberMap {
+                    virtual_table_number: table_number,
+                    physical_table_number: table_number,
+                }))
+            },
+            StableIndexName::Virtual(index_name, tablet_index_name) => {
+                Ok(Some(VirtualTableNumberMap {
+                    virtual_table_number: virtual_table_mapping.number(index_name.table())?,
+                    physical_table_number: table_mapping.inject_table_number()(
+                        *tablet_index_name.table(),
+                    )?
+                    .table_number,
+                }))
+            },
+            StableIndexName::Missing => Ok(None),
+        }
+    }
+}
+
 impl HeapSize for TabletIndexName {
     fn heap_size(&self) -> usize {
         self.descriptor.heap_size()

crates/database/src/query/mod.rs

@@ -375,7 +375,10 @@ impl<RT: Runtime, T: QueryType> CompiledQuery<RT, T> {
             )),
             QuerySource::IndexRange(index_range) => {
                 let order = index_range.order;
-                let interval = index_range.compile(indexed_fields)?;
+                let virtual_table_mapping = tx.virtual_table_mapping().clone();
+                let virtual_table_number_map = stable_index_name
+                    .virtual_table_number_map(tx.table_mapping(), &virtual_table_mapping)?;
+                let interval = index_range.compile(indexed_fields, virtual_table_number_map)?;
                 QueryNode::IndexRange(IndexRange::new(
                     stable_index_name,
                     index_name,

crates/isolate/src/tests/id_strings.rs

@@ -196,6 +196,20 @@ async fn test_system_normalize_id(rt: TestRuntime) -> anyhow::Result<()> {
     Ok(())
 }
 
+#[convex_macro::test_runtime]
+async fn test_virtual_id_query(rt: TestRuntime) -> anyhow::Result<()> {
+    let t = UdfTest::default(rt).await?;
+    let scheduled_id = t.mutation("idStrings:schedule", assert_obj!()).await?;
+
+    t.query(
+        "idStrings:queryVirtualId",
+        assert_obj!("id" => scheduled_id),
+    )
+    .await?;
+
+    Ok(())
+}
+
 proptest! {
     #![proptest_config(ProptestConfig { cases: 32 * env_config("CONVEX_PROPTEST_MULTIPLIER", 1), failure_persistence: None, .. ProptestConfig::default() })]
 

crates/value/src/base32.rs

@@ -1,3 +1,5 @@
+use std::cmp;
+
 // Forked from https://github.com/andreasots/base32 @ 58909ac.
 //
 // Copyright (c) 2015 The base32 Developers - MIT License
@@ -121,6 +123,72 @@ pub fn decode(data: &str) -> Result<Vec<u8>, InvalidBase32Error> {
     Ok(out)
 }
 
+fn clamp_char_to_alphabet(c: char) -> char {
+    match c {
+        ..'0' => '0',
+        '0'..='9' => c,
+        ':'..'a' => 'a',
+        'a'..='h' => c,
+        'i' => 'j',
+        'j'..='k' => c,
+        'l' => 'm',
+        'm'..='n' => c,
+        'o' => 'p',
+        'p'..='t' => c,
+        'u' => 'v',
+        'v'..='z' => c,
+        '{'.. => 'z',
+    }
+}
+
+/// Returns a string that can be base32 decoded, and is the closest such string
+/// to the input string in lexicographic order.
+/// i.e. for every `t` which is valid base32 of length <= `target_len`,
+/// if s = t, then clamp_to_alphabet(s) = t,
+/// if s < t, then clamp_to_alphabet(s) <= t.
+/// if s > t, then clamp_to_alphabet(s) >= t.
+///
+/// How does it work?
+/// Each character is clamped to the closest character in the base32 alphabet.
+/// If a character has to be rounded up, following characters all become '0'.
+/// If a character has to be rounded down, following characters all become 'z'.
+/// And then we pad to a multiple of 5 characters to avoid dropping chars when
+/// decoding.
+///
+/// e.g. "azi" is between "azhzz" and "azj00", so it is clamped to "azj00".
+/// e.g. "abcd~" is between "abcdz" and "abce0", so it is clamped to "abcdz".
+pub fn clamp_to_alphabet(s: &str, target_len: usize) -> String {
+    let mut out = String::with_capacity(s.len());
+    let mut order = cmp::Ordering::Equal;
+    for c in s.chars() {
+        match order {
+            cmp::Ordering::Equal => {
+                let clamped = clamp_char_to_alphabet(c);
+                out.push(clamped);
+                order = c.cmp(&clamped);
+            },
+            cmp::Ordering::Less => {
+                out.push('0');
+            },
+            cmp::Ordering::Greater => {
+                out.push('z');
+            },
+        }
+    }
+    // Pad the output to a multiple of 5 with 0s so no characters are lost.
+    while out.len() < target_len || out.len() % 5 != 0 {
+        match order {
+            cmp::Ordering::Equal | cmp::Ordering::Less => {
+                out.push('0');
+            },
+            cmp::Ordering::Greater => {
+                out.push('z');
+            },
+        }
+    }
+    out
+}
+
 #[cfg(test)]
 mod tests {
     use proptest::prelude::*;