From 393c38e02b27cae3690b0987021d96e436da05cf Mon Sep 17 00:00:00 2001 From: ohmayr Date: Wed, 16 Sep 2026 22:48:11 +0000 Subject: [PATCH 1/4] feat(firestore): add BSON cross-type query ordering support --- .../google/cloud/firestore_v1/order.py | 140 +++++++++++++++--- .../tests/system/test_system.py | 19 +++ .../tests/unit/v1/test_order.py | 58 ++++++++ 3 files changed, 193 insertions(+), 24 deletions(-) diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/order.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/order.py index a3d65cc5000e..037f447fecca 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/order.py +++ b/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,10 +64,24 @@ 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] @@ -69,16 +89,21 @@ def from_value(value) -> Any: # 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: diff --git a/packages/google-cloud-firestore/tests/system/test_system.py b/packages/google-cloud-firestore/tests/system/test_system.py index f15fb4ec6b14..e9e0f3054c54 100644 --- a/packages/google-cloud-firestore/tests/system/test_system.py +++ b/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(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 diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_order.py b/packages/google-cloud-firestore/tests/unit/v1/test_order.py index 1942a5298438..51659c307f00 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_order.py +++ b/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}) From aa7b44b77008470b5607d31644fde05e6254d0e6 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Mon, 21 Sep 2026 19:30:16 +0000 Subject: [PATCH 2/4] fix(firestore): streamline BSON cross-type ordering and numeric comparisons Use _BSON_KEY_TO_TYPE_ORDER dictionary lookup in order.py for O(1) wire key resolution, decoupling bson.py from query ordering. Consolidate cross-type numeric comparisons in compare_numbers with safe Decimal handling and restore compare_doubles to standard float comparisons. --- .../google/cloud/firestore_v1/order.py | 74 +++++++++++++------ .../tests/system/test_system.py | 2 +- .../tests/unit/v1/test_order.py | 26 +++++++ 3 files changed, 80 insertions(+), 22 deletions(-) diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/order.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/order.py index 037f447fecca..e4d421a77eec 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/order.py +++ b/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 @@ -67,26 +68,33 @@ def from_value(value) -> Any: 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 + 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.BSON_MIN_KEY: 1, @@ -136,8 +144,10 @@ def compare(cls, left, right) -> int: 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) @@ -209,6 +219,7 @@ def compare_bson_regexes(left, right) -> int: @staticmethod def compare_timestamps(left, right) -> Any: + """Compare native Firestore timestamps and BSON timestamps.""" left_pb = getattr(left, "_pb", left) right_pb = getattr(right, "_pb", right) @@ -319,13 +330,34 @@ def compare_objects(left, right) -> int: @staticmethod def compare_numbers(left, right) -> int: - 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)) + """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) + ) + + l = _to_number(left) + r = _to_number(right) + + l_nan = l.is_nan() if hasattr(l, "is_nan") else math.isnan(l) + r_nan = r.is_nan() if hasattr(r, "is_nan") else math.isnan(r) + if l_nan or r_nan: + return 0 if (l_nan and r_nan) else (-1 if l_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(l, decimal.Decimal) and isinstance(r, float): + r = decimal.Decimal(str(r)) + elif isinstance(r, decimal.Decimal) and isinstance(l, float): + l = decimal.Decimal(str(l)) + + return Order._compare_to(l, r) @staticmethod def compare_doubles(left, right) -> int: diff --git a/packages/google-cloud-firestore/tests/system/test_system.py b/packages/google-cloud-firestore/tests/system/test_system.py index e9e0f3054c54..8422ee372a60 100644 --- a/packages/google-cloud-firestore/tests/system/test_system.py +++ b/packages/google-cloud-firestore/tests/system/test_system.py @@ -1325,7 +1325,7 @@ def test_bson_query_ordering(client, cleanup, database): doc3.set({"val": BSONMaxKey()}) query = coll_ref.order_by("val") - results = [doc.to_dict(decode_bson=True)["val"] for doc in query.stream()] + results = [doc.to_dict()["val"] for doc in query.stream()] assert results == [BSONMinKey(), BSONInt32(10), BSONMaxKey()] diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_order.py b/packages/google-cloud-firestore/tests/unit/v1/test_order.py index 51659c307f00..977efd6dd302 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_order.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_order.py @@ -242,6 +242,16 @@ def test_order_bson_type_ordering(): 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 @@ -256,6 +266,22 @@ def test_order_bson_type_ordering(): 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}) From 1273cdbafacced14eddbd2dd3c4330f79058b059 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Mon, 21 Sep 2026 19:47:50 +0000 Subject: [PATCH 3/4] fix(firestore): resolve lint E741 and format in compare_numbers Rename ambiguous single-letter variable l to left_val to satisfy flake8 E741 and align expression formatting with ruff. --- .../google/cloud/firestore_v1/order.py | 34 ++++++++++--------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/order.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/order.py index e4d421a77eec..d5e740c03118 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/order.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/order.py @@ -335,29 +335,31 @@ def compare_numbers(left, right) -> int: 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) - ) + return to_decimal() if callable(to_decimal) else getattr(num, "value", num) - l = _to_number(left) - r = _to_number(right) + left_val = _to_number(left) + right_val = _to_number(right) - l_nan = l.is_nan() if hasattr(l, "is_nan") else math.isnan(l) - r_nan = r.is_nan() if hasattr(r, "is_nan") else math.isnan(r) - if l_nan or r_nan: - return 0 if (l_nan and r_nan) else (-1 if l_nan else 1) + 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(l, decimal.Decimal) and isinstance(r, float): - r = decimal.Decimal(str(r)) - elif isinstance(r, decimal.Decimal) and isinstance(l, float): - l = decimal.Decimal(str(l)) + 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(l, r) + return Order._compare_to(left_val, right_val) @staticmethod def compare_doubles(left, right) -> int: From 309ce0b27f4fb74a91e5d281504e268409e7e273 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Wed, 23 Sep 2026 19:33:43 +0000 Subject: [PATCH 4/4] fix(firestore): separate BSON_TIMESTAMP in TypeOrder and optimize compare_numbers Separate BSON_TIMESTAMP from native TIMESTAMP to conform to the cross-SDK 17-rank TypeOrder specification. Restore compare_timestamps to native Firestore timestamps and introduce compare_bson_timestamps. Extract module-level _to_number and _is_nan helpers, directly inspecting protobuf fields to optimize numeric comparisons. --- .../google/cloud/firestore_v1/order.py | 126 +++++++++++------- .../tests/unit/v1/test_order.py | 12 +- 2 files changed, 90 insertions(+), 48 deletions(-) diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/order.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/order.py index d5e740c03118..dab281d82fcc 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/order.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/order.py @@ -20,6 +20,45 @@ from google.cloud.firestore_v1._helpers import GeoPoint, decode_value +def _to_number(val: Any) -> Any: + """Extract a numeric value (int, float, Decimal) from a Value protobuf or Python value. + + Directly inspects the protobuf value_type without calling decode_value() + for optimal performance. + """ + value_pb = getattr(val, "_pb", val) + which = ( + value_pb.WhichOneof("value_type") if hasattr(value_pb, "WhichOneof") else None + ) + + if which == "integer_value": + return value_pb.integer_value + elif which == "double_value": + return value_pb.double_value + elif which == "map_value": + fields = value_pb.map_value.fields + if "__int__" in fields: + return fields["__int__"].integer_value + elif "__decimal128__" in fields: + return decimal.Decimal(fields["__decimal128__"].string_value) + + num = decode_value(val, None) + to_decimal = getattr(num, "to_decimal", None) + return to_decimal() if callable(to_decimal) else getattr(num, "value", num) + + +def _is_nan(val: Any) -> bool: + """Check if a numeric value is NaN, safely handling OverflowError and non-floats.""" + if hasattr(val, "is_nan"): + return val.is_nan() + if isinstance(val, (int, decimal.Decimal)): + return False + try: + return math.isnan(val) + except (TypeError, OverflowError): + return False + + class TypeOrder(Enum): """The supported Data Type. @@ -42,6 +81,7 @@ class TypeOrder(Enum): BSON_OBJECT_ID = 13 BSON_BINARY = 14 BSON_REGEX = 15 + BSON_TIMESTAMP = 16 @staticmethod def from_value(value) -> Any: @@ -77,8 +117,7 @@ def from_value(value) -> Any: # 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. +# BSONInt32 and BSONDecimal128 map to TypeOrder.NUMBER, enabling cross-type comparisons. _BSON_KEY_TO_TYPE_ORDER = { "__min__": TypeOrder.BSON_MIN_KEY, "__max__": TypeOrder.BSON_MAX_KEY, @@ -86,32 +125,30 @@ def from_value(value) -> Any: "__int__": TypeOrder.NUMBER, "__decimal128__": TypeOrder.NUMBER, "__binary__": TypeOrder.BSON_BINARY, - "__request_timestamp__": TypeOrder.TIMESTAMP, + "__request_timestamp__": TypeOrder.BSON_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.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, + TypeOrder.BSON_TIMESTAMP: 5, + TypeOrder.STRING: 6, + TypeOrder.BLOB: 7, + TypeOrder.BSON_BINARY: 8, + TypeOrder.REF: 9, + TypeOrder.BSON_OBJECT_ID: 10, + TypeOrder.GEO_POINT: 11, + TypeOrder.BSON_REGEX: 12, + TypeOrder.ARRAY: 13, + TypeOrder.VECTOR: 14, + TypeOrder.OBJECT: 15, + TypeOrder.BSON_MAX_KEY: 16, } @@ -147,8 +184,9 @@ def compare(cls, left, right) -> int: # 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.BSON_TIMESTAMP: + return cls.compare_bson_timestamps(left, right) elif leftType == TypeOrder.STRING: return cls._compare_to(left.string_value, right.string_value) elif leftType == TypeOrder.BLOB: @@ -219,25 +257,31 @@ def compare_bson_regexes(left, right) -> int: @staticmethod def compare_timestamps(left, right) -> Any: - """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_pb.timestamp_value.seconds, right_pb.timestamp_value.seconds + ) + if seconds != 0: + return seconds - 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 + return Order._compare_to( + left_pb.timestamp_value.nanos, right_pb.timestamp_value.nanos + ) + + @staticmethod + def compare_bson_timestamps(left, right) -> Any: + left_pb = getattr(left, "_pb", left) + right_pb = getattr(right, "_pb", right) + + 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 + + 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 seconds = Order._compare_to(l_sec, r_sec) if seconds != 0: @@ -331,23 +375,11 @@ def compare_objects(left, right) -> int: @staticmethod def compare_numbers(left, right) -> int: """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) - ) + left_nan = _is_nan(left_val) + right_nan = _is_nan(right_val) if left_nan or right_nan: return 0 if (left_nan and right_nan) else (-1 if left_nan else 1) diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_order.py b/packages/google-cloud-firestore/tests/unit/v1/test_order.py index 977efd6dd302..ff70eb441260 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_order.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_order.py @@ -254,6 +254,16 @@ def test_order_bson_type_ordering(): # Test timestamp comparison (native timestamp < BSON timestamp with increment) assert target.compare(ts_native, ts_bson) == -1 + assert target.compare(ts_bson, ts_native) == 1 + + # Test BSON timestamp comparison + ts_bson2 = encode_value(BSONTimestamp(100, 2)) + ts_bson_later = encode_value(BSONTimestamp(101, 0)) + assert target.compare(ts_bson, ts_bson2) == -1 + assert target.compare(ts_bson2, ts_bson) == 1 + assert target.compare(ts_bson, ts_bson_later) == -1 + assert target.compare(ts_bson_later, ts_bson) == 1 + assert target.compare(ts_bson, ts_bson) == 0 # Test BSON binary > bytes assert target.compare(bytes_native, bin_b) == -1 @@ -277,7 +287,7 @@ def test_order_bson_type_ordering(): "__decimal128__": TypeOrder.NUMBER, "__binary__": TypeOrder.BSON_BINARY, "__regex__": TypeOrder.BSON_REGEX, - "__request_timestamp__": TypeOrder.TIMESTAMP, + "__request_timestamp__": TypeOrder.BSON_TIMESTAMP, } for key, expected_order in expected_orders.items(): assert _BSON_KEY_TO_TYPE_ORDER.get(key) == expected_order