Hibernate - Using @PartitionKey Annotation Examples

This tutorial shows you how to use Hibernate's @PartitionKey annotation.

Partitioning a database table can be a solution to improve query performance if the system only needs to access certain rows. After creating the partitions, you have to make sure to use queries that utilize the partition key. For example, if a table is partitioned by a column, usually you need to include the column as a criteria in the queries. Otherwise, the database may need to perform scans on all partitions, making the partitioning useless.

Insert queries usually include the partition key. Therefore, the database should know to which partitions a row belongs to. However, for non-insert queries, the partition key should be added as a criteria, so that the database knows that it should only find matching rows on certain partition(s). If you use Hibernate version 6.2 or above, you can use the @PartitionKey annotation.

Using @PartitionKey Annotation

For example, we have a table named orders and we want to partition it by the value of the date column using a range partition. Below are the queries for creating the table along with the partitions. I use PostgreSQL for this tutorial. You may need to adjust the queries if you use another database.

  CREATE TABLE orders (
      id          uuid        NOT NULL,
      amount      numeric     NOT NULL,
      date        date        NOT NULL,
      CONSTRAINT "pk_orders" PRIMARY KEY (id, created_at)
  ) PARTITION BY RANGE (created_at);

  CREATE TABLE orders_2023_01 PARTITION OF orders
      FOR VALUES FROM ('2023-01-01') TO ('2023-01-31');
  CREATE TABLE orders_2023_02 PARTITION OF orders
      FOR VALUES FROM ('2023-02-01') TO ('2023-02-28');
  CREATE TABLE orders_2023_03 PARTITION OF orders
      FOR VALUES FROM ('2023-03-01') TO ('2023-03-31');

Below is the entity class.

  @Entity
  @Table(name = "orders")
  public class Order {
  
    @Id
    @UuidGenerator(style = UuidGenerator.Style.TIME)
    private UUID id;
  
    private BigDecimal amount;
  
    private LocalDate date;
  
    // Constructor, getters, setters
  }

Let's say we want to fetch a row of the orders table, increment the price by one, and update it to the database. For fetching, we can explicitly include the date column as a criteria. However, for persisting the change to the database, usually we call the persist or merge method of Hibernate's EntityManager.

  String sql = "SELECT o FROM Order o WHERE o.id = :id AND o.date = :date";
  Order order = (Order) this.entityManager.createQuery(sql)
      .setParameter("id", id)
      .setParameter("date", date)
      .getSingleResult();

  order.setAmount(order.getAmount().add(BigDecimal.ONE));
  this.entityManager.merge(order);

Below is the generated query. Hibernate doesn't know the partition key. As a result, it doesn't include the date column in the criteria. That causes the database to scan on all partitions.

  update orders set amount=?,date=? where id=?

The solution is to add @PartitionKey annotation to the date field.

  @Entity
  @Table(name = "orders")
  public class Order {
  
    @Id
    @UuidGenerator(style = UuidGenerator.Style.TIME)
    private UUID id;
  
    private BigDecimal amount;
  
    @PartitionKey
    private LocalDate date;
  
    // Constructor, getters, setters
  }

The @PartitionKey annotation works by adding the annotated field as a criteria clause for UPDATE and DELETE queries generated by Hibernate. Below is the generated query after adding the annotation.

  update orders set amount=?,date=? where id=? and date=?

The same also applies for delete queries. Here is the generated query for deleting a row.

  delete from orders where id=? and date=?

You have to know that the annotation doesn't work for the following cases.

  • SELECT queries, since we usually use explicit criteria for fetching rows.
  • UPDATE/DELETE queries where the criteria is not generated by Hibernate or has to be set explicitly. For example, using a native query, Spring Data JPA's Query annotation, or CriteriaBuilder.

Summary

The @PartitionKey annotation can be useful for automatically adding the partition key as a criteria in UPDATE and DELETE queries. It may reduce the chance of missing partition key in the query criteria. Just make sure you have handled other queries not affected by the annotation.

You can also read about: