Introduction

In this article, we are going to see how we can implement a table partitioning solution when using Spring and Hibernate.

The goal of table partitioning is to split a large table into multiple smaller partition tables so that the associated table and index records can fit into the in-memory Buffer Pool, therefore allowing a more efficient seek or scan.

Table partitioning with PostgreSQL

PostgreSQL provides three strategies for table partitioning:

• List Partitioning
• Range Partitioning
• Hash Partitioning

In our example, we are going to use List Partitioning as we will divide our tables by continents.

For example, the users table is going to be partitioned like this:

CREATE TABLE users (
    id bigint NOT NULL,
    first_name varchar(255),
    last_name varchar(255),
    registered_on timestamp(6),
    partition_key varchar(255),
    PRIMARY KEY (id, partition_key)
) 
PARTITION BY LIST (partition_key)

CREATE TABLE users_asia 
PARTITION OF users FOR VALUES IN ('Asia')

CREATE TABLE users_africa 
PARTITION OF users FOR VALUES IN ('Africa')

CREATE TABLE users_north_america 
PARTITION OF users FOR VALUES IN ('North America')

CREATE TABLE users_south_america 
PARTITION OF users FOR VALUES IN ('South America')

CREATE TABLE users_europe 
PARTITION OF users FOR VALUES IN ('Europe')

CREATE TABLE users_australia 
PARTITION OF users FOR VALUES IN ('Australia')

And, we can also partition the posts table as follows:

CREATE TABLE posts (
    id bigint NOT NULL,
    title varchar(255),
    created_on timestamp(6),
    user_id bigint,
    partition_key varchar(255),
    PRIMARY KEY (id, partition_key)
) 
PARTITION BY LIST (partition_key)

CREATE TABLE posts_asia 
PARTITION OF posts FOR VALUES IN ('Asia')

CREATE TABLE posts_africa 
PARTITION OF posts FOR VALUES IN ('Africa')

CREATE TABLE posts_north_america 
PARTITION OF posts FOR VALUES IN ('North America')

CREATE TABLE posts_south_america 
PARTITION OF posts FOR VALUES IN ('South America')

CREATE TABLE posts_europe 
PARTITION OF posts FOR VALUES IN ('Europe')

CREATE TABLE posts_australia 
PARTITION OF posts FOR VALUES IN ('Australia')

ALTER TABLE IF EXISTS posts
ADD CONSTRAINT fk_posts_user_id 
FOREIGN KEY (user_id, partition_key) 
REFERENCES users

Mapping table partitioning with JPA and Hibernate

Now, implementing table partitioning with Spring and Hibernate requires us to provide the partition key when reading and writing records from and to the partitioned tables.

In our case, we want to map the users and posts tables can be mapped to JPA entities to take advantage of all the benefits offered by Spring Data JPA.

First, we will create a PartitionAware base class that will be extended by every JPA entity that’s mapped to a partitioned table:

@MappedSuperclass
@FilterDef(
    name = PartitionAware.PARTITION_KEY,
    parameters = @ParamDef(
        name = PartitionAware.PARTITION_KEY,
        type = String.class
    )
)
@Filter(
    name = PartitionAware.PARTITION_KEY,
    condition = "partition_key = :partitionKey"
)
public abstract class PartitionAware<T extends PartitionAware> {

    public static final String PARTITION_KEY = "partitionKey";

    @Column(name = "partition_key")
    @PartitionKey
    private String partitionKey;

    public String getPartitionKey() {
        return partitionKey;
    }

    public T setPartitionKey(String partitionKey) {
        this.partitionKey = partitionKey;
        return (T) this;
    }

    public T setPartition(Partition partition) {
        this.partitionKey = partition.getKey();
        return (T) this;
    }
}

• The @Filter annotation allows us to enable the partition filtering dynamically
• The @PartitionKey annotation was introduced in Hibernate 6.2 and allows us to define a partition key

The Partition object is an Enum that defines the partitions that our application is supporting:

public enum Partition {
    ASIA("Asia"),
    AFRICA("Africa"),
    NORTH_AMERICA("North America"),
    SOUTH_AMERICA("South America"),
    EUROPE("Europe"),
    AUSTRALIA("Australia"),
    ;

    private final String key;

    Partition(String key) {
        this.key = key;
    }

    public String getKey() {
        return key;
    }
}

For more details about the @MappedSuperclass annotation, check out this article.

The User entity extends the PartitionAware base class and is mapped like this:

@Entity
@Table(name = "users")
public class User extends PartitionAware<User> {

    @Id
    @GeneratedValue
    private Long id;

    @Column(name = "first_name")
    private String firstName;

    @Column(name = "last_name")
    private String lastName;

    @Column(name = "registered_on")
    @CreationTimestamp
    private LocalDateTime createdOn = LocalDateTime.now();
    
    //Getters and setters omitted for brevity
}

And the Post entity is mapped as follows:

@Entity
@Table(name = "posts")
public class Post extends PartitionAware<Post> {

    @Id
    @GeneratedValue
    private Long id;

    private String title;

    @Column(name = "created_on")
    @CreationTimestamp
    private LocalDateTime createdOn;

    @ManyToOne(fetch = FetchType.LAZY)
    @JoinColumn(name = "user_id")
    private User user;
    
    //Getters and setters omitted for brevity
}

Now, we want the partition filtering to be enabled by default when a JPA EntityManager is created, and this can be achieved via the Spring EntityManager initializer mechanism:

@Bean
public JpaTransactionManager transactionManager(
        EntityManagerFactory entityManagerFactory){
    JpaTransactionManager transactionManager = new JpaTransactionManager();
    transactionManager.setEntityManagerFactory(entityManagerFactory);
    transactionManager.setEntityManagerInitializer(entityManager -> {
        User user = UserContext.getCurrent();
        if (user != null) {
            entityManager.unwrap(Session.class)
                .enableFilter(PartitionAware.PARTITION_KEY)
                .setParameter(
                    PartitionAware.PARTITION_KEY, 
                    user.getPartitionKey()
                );
        }
    });
    return transactionManager;
}

The current partition is being set based on the currently logged User. Since the partition Filter is enabled, every entity query will select the User and Post entities from the partition of the currently logged User.

We also want the partition_key column to be set on every entity that the currently logged User is persisting, and for this reason, we will create the following Hibernate PersistEventListener:

public class PartitionAwareInsertEventListener 
        implements PersistEventListener {

    public static final PartitionAwareInsertEventListener INSTANCE = 
        new PartitionAwareInsertEventListener();

    @Override
    public void onPersist(PersistEvent event) 
            throws HibernateException {
        final Object entity = event.getObject();

        if (entity instanceof PartitionAware partitionAware) {
            if (partitionAware.getPartitionKey() == null) {
                FilterImpl partitionKeyFilter = (FilterImpl) event
                    .getSession()
                    .getEnabledFilter(PartitionAware.PARTITION_KEY);
                partitionAware.setPartitionKey(
                    (String) partitionKeyFilter
                        .getParameter(PartitionAware.PARTITION_KEY)
                );
            }
        }
    }

    @Override
    public void onPersist(PersistEvent event, PersistContext persistContext)
            throws HibernateException {
        onPersist(event);
    }
}

The PartitionAwareInsertEventListener is going to be registered using the following PartitionAwareEventListenerIntegrator:

public class PartitionAwareEventListenerIntegrator 
        implements Integrator {

    public static final PartitionAwareEventListenerIntegrator INSTANCE =
        new PartitionAwareEventListenerIntegrator();

    @Override
    public void integrate(
            Metadata metadata,
            BootstrapContext bootstrapContext,
            SessionFactoryImplementor sessionFactory) {
        sessionFactory
            .getServiceRegistry()
            .getService(EventListenerRegistry.class)
            .prependListeners(
                EventType.PERSIST,
                PartitionAwareInsertEventListener.INSTANCE
            );
    }

    @Override
    public void disintegrate(
        SessionFactoryImplementor sessionFactory,
        SessionFactoryServiceRegistry serviceRegistry) {
    }
}

And the PartitionAwareEventListenerIntegrator is provided to Spring via the Java-based Properties configuration:

properties.put(
    EntityManagerFactoryBuilderImpl.INTEGRATOR_PROVIDER,
    (IntegratorProvider) () -> List.of(
        PartitionAwareEventListenerIntegrator.INSTANCE
    )
);

That’s it!

Testing Time

When persisting a User along with 3 Post entries:

final User vlad = new User()
    .setFirstName("Vlad")
    .setLastName("Mihalcea")
    .setPartition(Partition.EUROPE);

userRepository.persist(vlad);

UserContext.logIn(vlad);

forumService.createPosts(LongStream.rangeClosed(1, POST_COUNT)
    .mapToObj(postId -> new Post()
        .setTitle(
            String.format(
                "High-Performance Java Persistence - Part %d",
                postId
            )
        )
        .setUser(vlad)
    )
    .toList()
);

Hibernate will generate the following SQL INSERT statements:

INSERT INTO users (
    registered_on,
    first_name,
    last_name,
    partition_key,
    id
) 
VALUES (
    '2023-11-09 11:22:55.802704', 
    'Vlad, Mihalcea', 
    'Europe', 
    1
)
INSERT INTO posts (
    created_on,
    partition_key,
    title,
    user_id,
    id
) 
VALUES (
    '2023-11-09 11:19:55.856126', 
    'Europe', 
    'High-Performance Java Persistence - Part 1', 
    1, 
    1
), (
    '2023-11-09 11:19:55.856126', 
    'Europe', 
    'High-Performance Java Persistence - Part 2', 
    1, 
    2
),(
    '2023-11-09 11:19:55.856126', 
    'Europe', 
    'High-Performance Java Persistence - Part 3', 
    1, 
    3
)

And when fetching the Post entities:

List<Post> posts = forumService.findByIds(
    LongStream.rangeClosed(1, POST_COUNT).boxed().toList()
);

Hibernate is going to filter by the partition of the currently logged User:

SELECT 
    p.id,
    p.created_on,
    p.partition_key,
    p.title,
    p.user_id 
FROM 
    posts p 
WHERE 
    p.partition_key = 'Europe' AND 
    p.id in (1, 2, 3)

When running an EXPLAIN ANALYZE on this SQL query, we can see that only the posts_europe table gets scanned:

Seq Scan on posts_europe p  
  (cost=0.00..11.14 rows=1 width=1056) 
  (actual time=0.022..0.023 rows=3 loops=1)
  Filter: (
    ((partition_key)::text = 'Europe'::text) AND 
    (id = ANY ('{1,2,3}'::bigint[]))
  )

Awesome, right?

If you enjoyed this article, I bet you are going to love my Book and Video Courses as well.

Conclusion

Implementing a table partitioning solution with Spring and Hibernate is quite straightforward, as both frameworks are very flexible and can be customized in any possible way.

When operating with very large tables, table partitioning can help you optimize the performance of your SQL queries since the database will prune the partitions that are not needed when scanning the table records.

Insert details about how the information is going to be processed

DOWNLOAD NOW