Glue: honor Segment in get_partitions

moto/glue/models.py

@@ -3,7 +3,7 @@
 from collections import OrderedDict
 from collections.abc import Iterator
 from datetime import datetime
-from typing import Any
+from typing import Any, Optional
 
 from moto.core.base_backend import BackendDict, BaseBackend
 from moto.core.common_models import BaseModel
@@ -395,7 +395,11 @@ def get_partition(
         return table.get_partition(values)
 
     def get_partitions(
-        self, database_name: str, table_name: str, expression: str
+        self,
+        database_name: str,
+        table_name: str,
+        expression: str,
+        segment: Optional[dict[str, int]] = None,
     ) -> list["FakePartition"]:
         """
         See https://docs.aws.amazon.com/glue/latest/webapi/API_GetPartitions.html
@@ -409,7 +413,7 @@ def get_partitions(
           Only % and _ wildcards are supported, and SQL escaping using [] does not work.
         """
         table = self.get_table(database_name, table_name)
-        return table.get_partitions(expression)
+        return table.get_partitions(expression, segment)
 
     def update_partition(
         self,
@@ -1769,11 +1773,30 @@ def create_partition(self, partiton_input: dict[str, Any]) -> None:
             raise PartitionAlreadyExistsException()
         self.partitions[str(partition.values)] = partition
 
-    def get_partitions(self, expression: str) -> list["FakePartition"]:
+    def get_partitions(
+        self, expression: str, segment: Optional[dict[str, int]] = None
+    ) -> list["FakePartition"]:
         # Only load pyparsing when necessary
         from .utils import PartitionFilter
 
-        return list(filter(PartitionFilter(expression, self), self.partitions.values()))
+        partitions = list(
+            filter(PartitionFilter(expression, self), self.partitions.values())
+        )
+        if segment is not None:
+            # A segmented GetPartitions (used by the Hive Glue client for parallel
+            # scans) must return a disjoint slice per segment; the union across all
+            # segments equals the full set with no duplicates. Assign each partition
+            # to exactly one segment via a stable hash of its values.
+            total = int(segment["TotalSegments"])
+            number = int(segment["SegmentNumber"])
+            partitions = [
+                p
+                for p in partitions
+                if int(hashlib.md5(str(p.values).encode("utf-8")).hexdigest(), 16)
+                % total
+                == number
+            ]
+        return partitions
 
     def get_partition(self, values: str) -> "FakePartition":
         try:

moto/glue/responses.py

@@ -146,10 +146,12 @@ def get_partitions(self) -> ActionResult:
         database_name = self.parameters.get("DatabaseName")
         table_name = self.parameters.get("TableName")
         expression = self.parameters.get("Expression")
+        segment = self.parameters.get("Segment")
         partitions = self.glue_backend.get_partitions(
             database_name,  # type: ignore[arg-type]
             table_name,  # type: ignore[arg-type]
             expression,  # type: ignore[arg-type]
+            segment,
         )
 
         return ActionResult({"Partitions": [p.as_dict() for p in partitions]})

tests/test_glue/test_datacatalog.py

@@ -642,6 +642,41 @@ def test_batch_create_partition():
             )
 
 
+@mock_aws
+def test_get_partitions_segmented():
+    # A segmented GetPartitions (used by parallel scanners such as the Hive Glue
+    # client) must return disjoint slices whose union is the full set with no
+    # duplicates, not the full set per segment.
+    client = boto3.client("glue", region_name="us-east-1")
+    database_name = "myspecialdatabase"
+    table_name = "myfirsttable"
+    helpers.create_database(client, database_name)
+    helpers.create_table(client, database_name, table_name)
+
+    partition_inputs = [
+        helpers.create_partition_input(database_name, table_name, values=[f"2018-10-{i:02}"])
+        for i in range(20)
+    ]
+    client.batch_create_partition(
+        DatabaseName=database_name,
+        TableName=table_name,
+        PartitionInputList=partition_inputs,
+    )
+
+    total_segments = 5
+    seen = []
+    for segment_number in range(total_segments):
+        response = client.get_partitions(
+            DatabaseName=database_name,
+            TableName=table_name,
+            Segment={"SegmentNumber": segment_number, "TotalSegments": total_segments},
+        )
+        seen.extend(tuple(p["Values"]) for p in response["Partitions"])
+
+    assert len(seen) == 20
+    assert len(set(seen)) == 20
+
+
 @mock_aws
 def test_batch_create_partition_already_exist():
     client = boto3.client("glue", region_name="us-east-1")