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

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

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import decimal
 import math
 from enum import Enum
 from typing import Any
@@ -36,10 +37,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 +65,53 @@ 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))
+                bson_order = _BSON_KEY_TO_TYPE_ORDER.get(key)
+                if bson_order is not None:
+                    return bson_order
+            if "__type__" in fields and fields["__type__"].string_value == "__vector__":
                 return TypeOrder.VECTOR
         return lut[v]
 
 
+# Maps BSON wire map keys directly to their corresponding TypeOrder.
+# BSONTimestamp maps to TypeOrder.TIMESTAMP, and BSONInt32 / BSONDecimal128
+# map to TypeOrder.NUMBER, enabling cross-type comparisons.
+_BSON_KEY_TO_TYPE_ORDER = {
+    "__min__": TypeOrder.BSON_MIN_KEY,
+    "__max__": TypeOrder.BSON_MAX_KEY,
+    "__oid__": TypeOrder.BSON_OBJECT_ID,
+    "__int__": TypeOrder.NUMBER,
+    "__decimal128__": TypeOrder.NUMBER,
+    "__binary__": TypeOrder.BSON_BINARY,
+    "__request_timestamp__": TypeOrder.TIMESTAMP,
+    "__regex__": TypeOrder.BSON_REGEX,
+}
+
+
 # NOTE: This order is defined by the backend and cannot be changed.
+# BSONTimestamp shares TypeOrder.TIMESTAMP with native timestamps, and
+# BSONInt32 / BSONDecimal128 share TypeOrder.NUMBER, enabling direct cross-type
+# value comparison within those categories.
 _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,22 +135,34 @@ 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:
+            # Handles int64, double, BSONInt32, and BSONDecimal128.
             return cls.compare_numbers(left, right)
         elif leftType == TypeOrder.TIMESTAMP:
+            # Handles native Firestore timestamps and BSONTimestamp.
             return cls.compare_timestamps(left, right)
         elif leftType == TypeOrder.STRING:
             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 +180,70 @@ 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
+        """Compare native Firestore timestamps and BSON timestamps."""
+        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 +330,36 @@ 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)
+        """Compare numeric values across int, float, BSONInt32, and BSONDecimal128."""
+
+        def _to_number(val):
+            num = decode_value(val, None)
+            to_decimal = getattr(num, "to_decimal", None)
+            return to_decimal() if callable(to_decimal) else getattr(num, "value", num)
+
+        left_val = _to_number(left)
+        right_val = _to_number(right)
+
+        left_nan = (
+            left_val.is_nan() if hasattr(left_val, "is_nan") else math.isnan(left_val)
+        )
+        right_nan = (
+            right_val.is_nan()
+            if hasattr(right_val, "is_nan")
+            else math.isnan(right_val)
+        )
+        if left_nan or right_nan:
+            return 0 if (left_nan and right_nan) else (-1 if left_nan else 1)
+
+        # Python raises TypeError when comparing Decimal with float directly,
+        # but allows comparing Decimal with int. Convert float to Decimal
+        # to ensure safe cross-type comparison without float overflow.
+        if isinstance(left_val, decimal.Decimal) and isinstance(right_val, float):
+            right_val = decimal.Decimal(str(right_val))
+        elif isinstance(right_val, decimal.Decimal) and isinstance(left_val, float):
+            left_val = decimal.Decimal(str(left_val))
+
+        return Order._compare_to(left_val, 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() == 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()["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,90 @@ 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 large decimal comparison exceeding float limit
+    large_dec = encode_value(BSONDecimal128("1e1000"))
+    assert target.compare(large_dec, _double_value(1e300)) == 1
+    assert target.compare(_double_value(1e300), large_dec) == -1
+
+    # Test decimal NaN comparison
+    nan_dec = encode_value(BSONDecimal128("NaN"))
+    assert target.compare(nan_dec, int32_v) == -1
+    assert target.compare(int32_v, nan_dec) == 1
+
+    # 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
+
+    # Verify _BSON_KEY_TO_TYPE_ORDER mapping directly
+    from google.cloud.firestore_v1.order import _BSON_KEY_TO_TYPE_ORDER, TypeOrder
+
+    expected_orders = {
+        "__min__": TypeOrder.BSON_MIN_KEY,
+        "__max__": TypeOrder.BSON_MAX_KEY,
+        "__oid__": TypeOrder.BSON_OBJECT_ID,
+        "__int__": TypeOrder.NUMBER,
+        "__decimal128__": TypeOrder.NUMBER,
+        "__binary__": TypeOrder.BSON_BINARY,
+        "__regex__": TypeOrder.BSON_REGEX,
+        "__request_timestamp__": TypeOrder.TIMESTAMP,
+    }
+    for key, expected_order in expected_orders.items():
+        assert _BSON_KEY_TO_TYPE_ORDER.get(key) == expected_order
+
+
 def test_order_compare_w_objects_different_keys():
     left = _object_value({"foo": 0})
     right = _object_value({"bar": 0})