From 630b2fc93be2301bcf80b66db7b385aaf729c156 Mon Sep 17 00:00:00 2001 From: Nicolas Paris Date: Mon, 20 Jul 2026 16:36:36 +0200 Subject: [PATCH] Glue: honor Segment in get_partitions A segmented GetPartitions (used by the SDK-v1 Hive metastore Glue client for parallel partition scans) must return a disjoint slice per segment whose union is the full set with no duplicates. moto ignored the Segment parameter and returned every partition for each segment, so a client issuing TotalSegments=N saw each partition N times. This made Spark read every partitioned table row N times when reading a Delta/Hudi table synced into the moto Glue catalog. Assign each partition to exactly one segment via a stable md5 hash of its values, mirroring real AWS Glue segmentation semantics. --- moto/glue/models.py | 33 ++++++++++++++++++++++----- moto/glue/responses.py | 2 ++ tests/test_glue/test_datacatalog.py | 35 +++++++++++++++++++++++++++++ 3 files changed, 65 insertions(+), 5 deletions(-) diff --git a/moto/glue/models.py b/moto/glue/models.py index 0f8ae774fba2..f2dc1fb243f9 100644 --- a/moto/glue/models.py +++ b/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: diff --git a/moto/glue/responses.py b/moto/glue/responses.py index ece081e8fc82..e7b4c0a66cd6 100644 --- a/moto/glue/responses.py +++ b/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]}) diff --git a/tests/test_glue/test_datacatalog.py b/tests/test_glue/test_datacatalog.py index 9d4858011f47..34f0fee18d71 100644 --- a/tests/test_glue/test_datacatalog.py +++ b/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")