Query condition missed key schema element -> trying to read from dynamodb table

I have this dynamodb table

import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.*;

@DynamoDbBean public class ImageDetail {

private String documentId;
private String id;
private int trackPosition;
private String url;
private String name;
private double width;
private double height;

public ImageDetail() {
    // Default constructor is required for DynamoDBBean
}

public ImageDetail(String name) {
    this.name = name;
}

public ImageDetail(String name, String url, double width, double height) {
    this.name = name;
    this.url = url;
    this.width = width;
    this.height = height;
}

public ImageDetail(String name, double width, double height) {
    this.name = name;
    this.width = width;
    this.height = height;
}

public ImageDetail(String name, String url) {
    this.name = name;
    this.url = url;
}




@DynamoDbAttribute("Name")
@DynamoDbIgnoreNulls
public String getName() {
    return name;
}

public void setName(String name) {
    this.name = name;
}

@DynamoDbAttribute("Url")
@DynamoDbIgnoreNulls
public String getUrl() {
    return url;
}

public void setUrl(String url) {
    this.url = url;
}

@DynamoDbAttribute("Width")
@DynamoDbIgnoreNulls
public double getWidth() {
    return width;
}

public void setWidth(double width) {
    this.width = width;
}

@DynamoDbAttribute("Height")
@DynamoDbIgnoreNulls
public double getHeight() {
    return height;
}

public void setHeight(double height) {
    this.height = height;
}

// Partition key
@DynamoDbPartitionKey
public String getDocumentId() {
    return documentId;
}

public void setDocumentId(String documentId) {
    this.documentId = documentId;
}

@DynamoDbAttribute("Id")
@DynamoDbIgnoreNulls
public String getId() {
    return id;
}

public void setId(String id) {
    this.id = id;
}

@DynamoDbSortKey
@DynamoDbIgnoreNulls
public int getTrackPosition() {
    return trackPosition;
}

public void setTrackPosition(int trackPosition) {
    this.trackPosition = trackPosition;
}

@Override
public String toString() {
    return "ImageDetail{" +
            "documentId='" + documentId + '\'' +
            ", id='" + id + '\'' +
            ", trackPosition=" + trackPosition +
            ", url='" + url + '\'' +
            ", name='" + name + '\'' +
            ", width=" + width +
            ", height=" + height +
            '}';
}

} And then this ImageService class:

import com.amazonaws.services.lambda.runtime.LambdaLogger; import com.exiq.columbia.pdftools.model.image.ImageDetail; import software.amazon.awssdk.regions.Region; import software.amazon.awssdk.services.dynamodb.DynamoDbClient; import software.amazon.awssdk.services.dynamodb.model.DynamoDbException;

import java.net.URI; import java.text.SimpleDateFormat; import java.util.TimeZone;

import software.amazon.awssdk.enhanced.dynamodb.; import software.amazon.awssdk.enhanced.dynamodb.model.;

import java.util.*; import java.util.UUID;

public class ImageService {

private DynamoDbClient client;
private DynamoDbEnhancedClient enhancedClient;
private DynamoDbTable<ImageDetail> imagesTable;

private SimpleDateFormat sdf;

public ImageService local() {
    client = DynamoDbClient.builder()
            .endpointOverride(URI.create("http://localhost:8000"))
            .region(Region.of(System.getenv("AWS_REGION")))
            .build();
    enhancedClient = DynamoDbEnhancedClient.builder()
            .dynamoDbClient(client)
            .build();
    return this;
}

public ImageService cloud() {
    client = DynamoDbClient.builder()
            .region(Region.of(System.getenv("AWS_REGION")))
            .build();
    enhancedClient = DynamoDbEnhancedClient.builder()
            .dynamoDbClient(client)
            .build();
    return this;
}

public ImageService init() {
    sdf = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSXXX");
    sdf.setTimeZone(TimeZone.getTimeZone("UTC"));
    return this;
}

public ImageService withDefaultTable() {
    if (client == null) {
        cloud();
    }
    imagesTable = enhancedClient.table(System.getenv("COLUMBIA_IMAGES_TABLE"), TableSchema.fromBean(ImageDetail.class)); // Reflect the change here as well
    return this;
}

public ImageService withImagesTable(String tableName) {
    this.imagesTable = enhancedClient.table(tableName, TableSchema.fromBean(ImageDetail.class));
    return this;
}


public List<ImageDetail> getImagesByDocumentId(String documentId, LambdaLogger logger) {
    List<ImageDetail> imageDetails = new ArrayList<>();

    try {
        // Query condition for the partition key only
        QueryConditional queryConditional = QueryConditional
                .keyEqualTo(k -> k.partitionValue(documentId));

        logger.log("Querying for DocumentId: " + documentId);

        // Enhanced client operation to query the images
        Iterator<Page<ImageDetail>> results = imagesTable.query(r -> r.queryConditional(queryConditional)).iterator();
        while (results.hasNext()) {
            Page<ImageDetail> page = results.next();
            imageDetails.addAll(page.items());
        }
    } catch (DynamoDbException e) {
        logger.log("Error retrieving images for documentId " + documentId + ": " + e.getMessage());
        // Consider rethrowing the exception or handling it as necessary
    }

    return imageDetails;
}



public void saveOrUpdateImageRecords(String documentId, List<ImageDetail> imageDetails, LambdaLogger logger) {
    DynamoDbTable<ImageDetail> table = enhancedClient.table(System.getenv("COLUMBIA_IMAGES_TABLE"), TableSchema.fromBean(ImageDetail.class));

    for (int i = 0; i < imageDetails.size(); i++) {
        ImageDetail imageDetail = imageDetails.get(i);
        // Set DocumentId and TrackPosition as they are your keys.
        logger.log("image detail " + i);
        logger.log(imageDetail.toString());
        imageDetail.setDocumentId(documentId);
        imageDetail.setTrackPosition(i);
        imageDetail.setHeight(imageDetail.getHeight());
        imageDetail.setWidth(imageDetail.getWidth());
        imageDetail.setUrl(imageDetail.getUrl());
        logger.log("Processing image: " + imageDetail.getName());

        try {
            // Check if the record already exists
            ImageDetail itemToCheck = table.getItem(Key.builder()
                    .partitionValue(documentId)
                    .sortValue(i)
                    .build());

            if (itemToCheck != null) {
                // Update the existing record
                table.updateItem(imageDetail);
            } else {
                // Put a new item if it doesn't exist
                table.putItem(imageDetail);
            }
            logger.log("Operation succeeded for image: " + imageDetail.getName());
        } catch (DynamoDbException e) {
            logger.log("Unable to process image: " + imageDetail.getName());
            logger.log(e.getMessage());
            // Handle the exception based on your use case
        }
    }
}

} When I perform a get query for a document id value that corresponds to existing entries in the dynamodb table I am getting this error:

Error retrieving images for documentId 2b820dad-4cf4-4a21-b243-8c0b840f56a2: Query condition missed key schema element: DocumentId (Service: DynamoDb, Status Code: 400, Request ID: B1MN9OTQVQ275DC8BGHTJVEFONVV4KQNSO5AEMVJF66Q9ASUAAJG)

The name of the table is set correctly in the environment variables of my lambda. Attached a screenshot of the table

Haven't managed to figure out what I am doing wrong yet. Any recommendations or help, very welcome