[BEAM-7860] Python Datastore: fix key sort order

This is a regression from the v1 client.

Author: udim

sdks/python/apache_beam/io/gcp/datastore/v1/query_splitter_test.py

@@ -171,11 +171,16 @@ def check_get_splits(self, query, num_splits, num_entities, batch_size):
       batch_size: the number of entities returned by fake datastore in one req.
     """
 
-    # Test for both random long ids and string ids.
-    id_or_name = [True, False]
+    # Test for random long ids, string ids, and a mix of both.
+    id_or_name = [True, False, None]
 
     for id_type in id_or_name:
-      entities = fake_datastore.create_entities(num_entities, id_type)
+      if id_type is None:
+        entities = fake_datastore.create_entities(num_entities, False)
+        entities.extend(fake_datastore.create_entities(num_entities, True))
+        num_entities *= 2
+      else:
+        entities = fake_datastore.create_entities(num_entities, id_type)
       mock_datastore = MagicMock()
       # Assign a fake run_query method as a side_effect to the mock.
       mock_datastore.run_query.side_effect = \

sdks/python/apache_beam/io/gcp/datastore/v1new/query_splitter.py

@@ -26,6 +26,9 @@
 from builtins import range
 from builtins import round
 
+from past.builtins import long
+from past.builtins import unicode
+
 from apache_beam.io.gcp.datastore.v1new import types
 
 __all__ = ['QuerySplitterError', 'SplitNotPossibleError', 'get_splits']
@@ -123,10 +126,59 @@ def _create_scatter_query(query, num_splits):
   return scatter_query
 
 
+class IdOrName(object):
+  """Represents an ID or name of a Datastore key,
+
+   Implements sort ordering: by ID, then by name, keys with IDs before those
+   with names.
+   """
+  def __init__(self, id_or_name):
+    self.id_or_name = id_or_name
+    if isinstance(id_or_name, (str, unicode)):
+      self.id = None
+      self.name = id_or_name
+    elif isinstance(id_or_name, (int, long)):
+      self.id = id_or_name
+      self.name = None
+    else:
+      raise TypeError('Unexpected type of id_or_name: %s' % id_or_name)
+
+  def __lt__(self, other):
+    if not isinstance(other, IdOrName):
+      return super(IdOrName, self).__lt__(other)
+
+    if self.id is not None:
+      if other.id is None:
+        return True
+      else:
+        return self.id < other.id
+
+    if other.id is not None:
+      return False
+
+    return self.name < other.name
+
+  def __eq__(self, other):
+    if not isinstance(other, IdOrName):
+      return super(IdOrName, self).__eq__(other)
+    return self.id == other.id and self.name == other.name
+
+  def __hash__(self):
+    return hash((self.id, self.other))
+
+
 def client_key_sort_key(client_key):
   """Key function for sorting lists of ``google.cloud.datastore.key.Key``."""
-  return [client_key.project, client_key.namespace or ''] + [
-      str(element) for element in client_key.flat_path]
+  sort_key = [client_key.project, client_key.namespace or '']
+  # A key path is made up of (kind, id_or_name) pairs. The last pair might be
+  # missing an id_or_name.
+  flat_path = list(client_key.flat_path)
+  while flat_path:
+    sort_key.append(flat_path.pop(0))  # kind
+    if flat_path:
+      sort_key.append(IdOrName(flat_path.pop(0)))
+
+  return sort_key
 
 
 def _get_scatter_keys(client, query, num_splits):

sdks/python/apache_beam/io/gcp/datastore/v1new/query_splitter_test.py

@@ -103,9 +103,16 @@ def check_get_splits(self, query, num_splits, num_entities,
       unused_batch_size: ignored in v1new since query results are entirely
         handled by the Datastore client.
     """
-    # Test for both random long ids and string ids.
-    for id_or_name in [True, False]:
-      client_entities = helper.create_client_entities(num_entities, id_or_name)
+    # Test for random long ids, string ids, and a mix of both.
+    for id_or_name in [True, False, None]:
+      if id_or_name is None:
+        client_entities = helper.create_client_entities(num_entities, False)
+        client_entities.extend(helper.create_client_entities(num_entities,
+                                                             True))
+        num_entities *= 2
+      else:
+        client_entities = helper.create_client_entities(num_entities,
+                                                        id_or_name)
 
       mock_client = mock.MagicMock()
       mock_client_query = mock.MagicMock()
@@ -154,6 +161,19 @@ def check_get_splits(self, query, num_splits, num_entities,
           if lt_key is None:
             last_query_seen = True
 
+  def test_id_or_name(self):
+    id_ = query_splitter.IdOrName(1)
+    self.assertEqual(1, id_.id)
+    self.assertIsNone(id_.name)
+    name = query_splitter.IdOrName('1')
+    self.assertIsNone(name.id)
+    self.assertEqual('1', name.name)
+    self.assertEqual(query_splitter.IdOrName(1), query_splitter.IdOrName(1))
+    self.assertEqual(query_splitter.IdOrName('1'), query_splitter.IdOrName('1'))
+    self.assertLess(query_splitter.IdOrName(2), query_splitter.IdOrName('1'))
+    self.assertLess(query_splitter.IdOrName(1), query_splitter.IdOrName(2))
+    self.assertLess(query_splitter.IdOrName('1'), query_splitter.IdOrName('2'))
+
   def test_client_key_sort_key(self):
     k = key.Key('kind1', 1, project=self._PROJECT, namespace=self._NAMESPACE)
     k2 = key.Key('kind2', 'a', parent=k)
@@ -165,6 +185,31 @@ def test_client_key_sort_key(self):
     keys.sort(key=query_splitter.client_key_sort_key)
     self.assertEqual(expected_sort, keys)
 
+  def test_client_key_sort_key_ids(self):
+    k1 = key.Key('kind', 2, project=self._PROJECT)
+    k2 = key.Key('kind', 1, project=self._PROJECT)
+    keys = [k1, k2]
+    expected_sort = [k2, k1]
+    keys.sort(key=query_splitter.client_key_sort_key)
+    self.assertEqual(expected_sort, keys)
+
+  def test_client_key_sort_key_names(self):
+    k1 = key.Key('kind', '2', project=self._PROJECT)
+    k2 = key.Key('kind', '1', project=self._PROJECT)
+    keys = [k1, k2]
+    expected_sort = [k2, k1]
+    keys.sort(key=query_splitter.client_key_sort_key)
+    self.assertEqual(expected_sort, keys)
+
+  def test_client_key_sort_key_ids_vs_names(self):
+    # Keys with IDs always come before keys with names.
+    k1 = key.Key('kind', '1', project=self._PROJECT)
+    k2 = key.Key('kind', 2, project=self._PROJECT)
+    keys = [k1, k2]
+    expected_sort = [k2, k1]
+    keys.sort(key=query_splitter.client_key_sort_key)
+    self.assertEqual(expected_sort, keys)
+
 
 # Hide base class from collection by nose.
 del QuerySplitterTestBase