feat(firestore): add BSON cross-type query ordering support

Author: ohmayr

packages/google-cloud-firestore/google/cloud/firestore_v1/order.py

@@ -36,10 +36,16 @@ class TypeOrder(Enum):
     ARRAY = 8
     OBJECT = 9
     VECTOR = 10
+    BSON_MIN_KEY = 11
+    BSON_MAX_KEY = 12
+    BSON_OBJECT_ID = 13
+    BSON_BINARY = 14
+    BSON_REGEX = 15
 
     @staticmethod
     def from_value(value) -> Any:
-        v = value._pb.WhichOneof("value_type")
+        value_pb = getattr(value, "_pb", value)
+        v = value_pb.WhichOneof("value_type")
         lut = {
             "null_value": TypeOrder.NULL,
             "boolean_value": TypeOrder.BOOLEAN,
@@ -58,27 +64,46 @@ def from_value(value) -> Any:
             raise ValueError(f"Could not detect value type for {v}")
 
         if v == "map_value":
-            if (
-                "__type__" in value.map_value.fields
-                and value.map_value.fields["__type__"].string_value == "__vector__"
-            ):
+            fields = value_pb.map_value.fields
+            if len(fields) == 1:
+                key = next(iter(fields))
+                if key == "__min__":
+                    return TypeOrder.BSON_MIN_KEY
+                if key == "__max__":
+                    return TypeOrder.BSON_MAX_KEY
+                if key == "__oid__":
+                    return TypeOrder.BSON_OBJECT_ID
+                if key in ("__int__", "__decimal128__"):
+                    return TypeOrder.NUMBER
+                if key == "__binary__":
+                    return TypeOrder.BSON_BINARY
+                if key == "__regex__":
+                    return TypeOrder.BSON_REGEX
+                if key == "__request_timestamp__":
+                    return TypeOrder.TIMESTAMP
+            if "__type__" in fields and fields["__type__"].string_value == "__vector__":
                 return TypeOrder.VECTOR
         return lut[v]
 
 
 # NOTE: This order is defined by the backend and cannot be changed.
 _TYPE_ORDER_MAP = {
     TypeOrder.NULL: 0,
-    TypeOrder.BOOLEAN: 1,
-    TypeOrder.NUMBER: 2,
-    TypeOrder.TIMESTAMP: 3,
-    TypeOrder.STRING: 4,
-    TypeOrder.BLOB: 5,
-    TypeOrder.REF: 6,
-    TypeOrder.GEO_POINT: 7,
-    TypeOrder.ARRAY: 8,
-    TypeOrder.VECTOR: 9,
-    TypeOrder.OBJECT: 10,
+    TypeOrder.BSON_MIN_KEY: 1,
+    TypeOrder.BOOLEAN: 2,
+    TypeOrder.NUMBER: 3,
+    TypeOrder.TIMESTAMP: 4,
+    TypeOrder.STRING: 5,
+    TypeOrder.BLOB: 6,
+    TypeOrder.BSON_BINARY: 7,
+    TypeOrder.REF: 8,
+    TypeOrder.BSON_OBJECT_ID: 9,
+    TypeOrder.GEO_POINT: 10,
+    TypeOrder.BSON_REGEX: 11,
+    TypeOrder.ARRAY: 12,
+    TypeOrder.VECTOR: 13,
+    TypeOrder.OBJECT: 14,
+    TypeOrder.BSON_MAX_KEY: 15,
 }
 
 
@@ -102,8 +127,12 @@ def compare(cls, left, right) -> int:
             else:
                 return 1
 
-        if leftType == TypeOrder.NULL:
-            return 0  # nulls are all equal
+        if (
+            leftType == TypeOrder.NULL
+            or leftType == TypeOrder.BSON_MIN_KEY
+            or leftType == TypeOrder.BSON_MAX_KEY
+        ):
+            return 0  # sentinels are equal
         elif leftType == TypeOrder.BOOLEAN:
             return cls._compare_to(left.boolean_value, right.boolean_value)
         elif leftType == TypeOrder.NUMBER:
@@ -114,10 +143,16 @@ def compare(cls, left, right) -> int:
             return cls._compare_to(left.string_value, right.string_value)
         elif leftType == TypeOrder.BLOB:
             return cls.compare_blobs(left, right)
+        elif leftType == TypeOrder.BSON_BINARY:
+            return cls.compare_bson_binaries(left, right)
         elif leftType == TypeOrder.REF:
             return cls.compare_resource_paths(left, right)
+        elif leftType == TypeOrder.BSON_OBJECT_ID:
+            return cls.compare_bson_object_ids(left, right)
         elif leftType == TypeOrder.GEO_POINT:
             return cls.compare_geo_points(left, right)
+        elif leftType == TypeOrder.BSON_REGEX:
+            return cls.compare_bson_regexes(left, right)
         elif leftType == TypeOrder.ARRAY:
             return cls.compare_arrays(left, right)
         elif leftType == TypeOrder.VECTOR:
@@ -135,16 +170,69 @@ def compare_blobs(left, right) -> int:
 
         return Order._compare_to(left_bytes, right_bytes)
 
+    @staticmethod
+    def compare_bson_binaries(left, right) -> int:
+        l_bin = left.map_value.fields["__binary__"].bytes_value
+        r_bin = right.map_value.fields["__binary__"].bytes_value
+
+        l_subtype = l_bin[0] if l_bin else 0
+        r_subtype = r_bin[0] if r_bin else 0
+
+        cmp_subtype = Order._compare_to(l_subtype, r_subtype)
+        if cmp_subtype != 0:
+            return cmp_subtype
+
+        return Order._compare_to(
+            l_bin[1:] if l_bin else b"", r_bin[1:] if r_bin else b""
+        )
+
+    @staticmethod
+    def compare_bson_object_ids(left, right) -> int:
+        l_oid = left.map_value.fields["__oid__"].string_value
+        r_oid = right.map_value.fields["__oid__"].string_value
+        return Order._compare_to(l_oid, r_oid)
+
+    @staticmethod
+    def compare_bson_regexes(left, right) -> int:
+        l_regex = left.map_value.fields["__regex__"].map_value.fields
+        r_regex = right.map_value.fields["__regex__"].map_value.fields
+
+        l_pattern = l_regex["pattern"].string_value if "pattern" in l_regex else ""
+        r_pattern = r_regex["pattern"].string_value if "pattern" in r_regex else ""
+        cmp_pat = Order._compare_to(l_pattern, r_pattern)
+        if cmp_pat != 0:
+            return cmp_pat
+
+        l_options = l_regex["options"].string_value if "options" in l_regex else ""
+        r_options = r_regex["options"].string_value if "options" in r_regex else ""
+        return Order._compare_to(l_options, r_options)
+
     @staticmethod
     def compare_timestamps(left, right) -> Any:
-        left = left._pb.timestamp_value
-        right = right._pb.timestamp_value
+        left_pb = getattr(left, "_pb", left)
+        right_pb = getattr(right, "_pb", right)
+
+        if left_pb.WhichOneof("value_type") == "map_value":
+            l_ts = left_pb.map_value.fields["__request_timestamp__"].map_value.fields
+            l_sec = l_ts["seconds"].integer_value if "seconds" in l_ts else 0
+            l_inc = l_ts["increment"].integer_value if "increment" in l_ts else 0
+        else:
+            l_sec = left_pb.timestamp_value.seconds
+            l_inc = left_pb.timestamp_value.nanos
 
-        seconds = Order._compare_to(left.seconds or 0, right.seconds or 0)
+        if right_pb.WhichOneof("value_type") == "map_value":
+            r_ts = right_pb.map_value.fields["__request_timestamp__"].map_value.fields
+            r_sec = r_ts["seconds"].integer_value if "seconds" in r_ts else 0
+            r_inc = r_ts["increment"].integer_value if "increment" in r_ts else 0
+        else:
+            r_sec = right_pb.timestamp_value.seconds
+            r_inc = right_pb.timestamp_value.nanos
+
+        seconds = Order._compare_to(l_sec, r_sec)
         if seconds != 0:
             return seconds
 
-        return Order._compare_to(left.nanos or 0, right.nanos or 0)
+        return Order._compare_to(l_inc, r_inc)
 
     @staticmethod
     def compare_geo_points(left, right) -> Any:
@@ -231,9 +319,13 @@ def compare_objects(left, right) -> int:
 
     @staticmethod
     def compare_numbers(left, right) -> int:
-        left_value = decode_value(left, None)
-        right_value = decode_value(right, None)
-        return Order.compare_doubles(left_value, right_value)
+        left_val = decode_value(left, None, decode_bson=True)
+        right_val = decode_value(right, None, decode_bson=True)
+        if hasattr(left_val, "value"):
+            left_val = left_val.value
+        if hasattr(right_val, "value"):
+            right_val = right_val.value
+        return Order.compare_doubles(float(left_val), float(right_val))
 
     @staticmethod
     def compare_doubles(left, right) -> int:

packages/google-cloud-firestore/tests/system/test_system.py

@@ -1310,6 +1310,25 @@ def test_bson_document_read_and_write(client, cleanup, database):
     assert snapshot.to_dict(decode_bson=True) == bson_payload
 
 
+def test_bson_query_ordering(client, cleanup, database):
+    """Test server query ordering for BSON types."""
+    collection_id = "bson_ordering_" + UNIQUE_RESOURCE_ID
+    coll_ref = client.collection(collection_id)
+
+    doc1 = coll_ref.document("doc1")
+    doc2 = coll_ref.document("doc2")
+    doc3 = coll_ref.document("doc3")
+    cleanup.extend([doc1.delete, doc2.delete, doc3.delete])
+
+    doc1.set({"val": BSONMinKey()})
+    doc2.set({"val": BSONInt32(10)})
+    doc3.set({"val": BSONMaxKey()})
+
+    query = coll_ref.order_by("val")
+    results = [doc.to_dict(decode_bson=True)["val"] for doc in query.stream()]
+    assert results == [BSONMinKey(), BSONInt32(10), BSONMaxKey()]
+
+
 @pytest.fixture(scope="module")
 def query_docs(client, database):
     collection_id = "qs" + UNIQUE_RESOURCE_ID

packages/google-cloud-firestore/tests/unit/v1/test_order.py

@@ -199,6 +199,64 @@ def test_order_all_value_present():
         assert type_order in _TYPE_ORDER_MAP
 
 
+def test_order_bson_type_ordering():
+    from google.cloud.firestore_v1._helpers import encode_value
+    from google.cloud.firestore_v1.bson import (
+        BSONBinary,
+        BSONDecimal128,
+        BSONInt32,
+        BSONMaxKey,
+        BSONMinKey,
+        BSONObjectId,
+        BSONRegex,
+        BSONTimestamp,
+    )
+    from google.cloud.firestore_v1.order import Order
+
+    min_k = encode_value(BSONMinKey())
+    max_k = encode_value(BSONMaxKey())
+    null_v = nullValue()
+    int32_v = encode_value(BSONInt32(10))
+    int64_v = _int_value(10)
+    dec_v = encode_value(BSONDecimal128("10.0"))
+    ts_bson = encode_value(BSONTimestamp(100, 1))
+    ts_native = _timestamp_value(100, 0)
+    bin_b = encode_value(BSONBinary(b"xyz", subtype=1))
+    bytes_native = _blob_value(b"xyz")
+    ref_v = _reference_value("projects/p1/databases/d1/documents/c1/doc1")
+    oid_v = encode_value(BSONObjectId("507f191e810c19729de860ea"))
+    geo_v = _geoPoint_value(0, 0)
+    regex_v = encode_value(BSONRegex("abc"))
+    arr_v = _array_value()
+    map_v = _object_value({"a": 1})
+
+    # Test 16-rank ordering bounds
+    target = Order()
+    assert target.compare(null_v, min_k) == -1
+    assert target.compare(min_k, null_v) == 1
+
+    assert target.compare(max_k, map_v) == 1
+    assert target.compare(map_v, max_k) == -1
+
+    # Test numbers comparison equality across int32, int64, decimal128
+    assert target.compare(int32_v, int64_v) == 0
+    assert target.compare(int32_v, dec_v) == 0
+
+    # Test timestamp comparison (native timestamp < BSON timestamp with increment)
+    assert target.compare(ts_native, ts_bson) == -1
+
+    # Test BSON binary > bytes
+    assert target.compare(bytes_native, bin_b) == -1
+
+    # Test ObjectId rank (REF < OID < GEO_POINT)
+    assert target.compare(ref_v, oid_v) == -1
+    assert target.compare(oid_v, geo_v) == -1
+
+    # Test Regex rank (GEO_POINT < REGEX < ARRAY)
+    assert target.compare(geo_v, regex_v) == -1
+    assert target.compare(regex_v, arr_v) == -1
+
+
 def test_order_compare_w_objects_different_keys():
     left = _object_value({"foo": 0})
     right = _object_value({"bar": 0})