Boost PostgreSQL IN Query Performance with Spring AOP SplitWork Annotation

Introduction

Large IN queries with hundreds or thousands of parameters can severely degrade PostgreSQL performance and cause slow API responses. This article presents an optimization approach using multithreading and Spring AOP.

Problem Statement

Typical query: SELECT * FROM device WHERE id IN (1, 2, 3, 4) When the parameter list becomes very large, the query slows down. The idea is to split the list into smaller chunks and execute them in parallel, then merge the results.

Solution Overview

A custom annotation @SplitWorkAnnotation is defined to configure thread pool, split limits, and result‑handling logic. By annotating a method, the framework automatically performs the split‑execute‑merge process.

Define AOP Annotation

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface SplitWorkAnnotation {
    ThreadPoolEnum setThreadPool();
    Class<? extends HandleReturn> handlerReturnClass() default MergeFunction.class;
    int splitLimit() default 1000;
    int splitGroupNum() default 100;
}

Parameters:

setThreadPool: The thread pool to use (avoid shared pools).

handlerReturnClass: Class that processes merged results (e.g., sum, count, top‑N).

splitLimit: Minimum size to trigger splitting.

splitGroupNum: Number of items per batch.

AOP Aspect Implementation

@Aspect
@Component
@Slf4j
public class SplitWorkAspect {
    @Pointcut("@annotation(com.demo.SplitWorkAnnotation)")
    public void needSplit() {}

    @Around("needSplit()")
    public Object around(ProceedingJoinPoint pjp) throws Throwable {
        Method targetMethod = ((MethodSignature) pjp.getSignature()).getMethod();
        SplitWorkAnnotation anno = targetMethod.getAnnotation(SplitWorkAnnotation.class);
        Object[] args = pjp.getArgs();
        int splitLimit = anno.splitLimit();
        int splitGroupNum = anno.splitGroupNum();
        if (args == null || args.length == 0 || splitLimit <= splitGroupNum) {
            return pjp.proceed();
        }
        int splitParamIndex = -1;
        for (int i = 0; i < targetMethod.getParameters().length; i++) {
            if (targetMethod.getParameters()[i].isAnnotationPresent(NeedSplitParam.class)) {
                splitParamIndex = i;
                break;
            }
        }
        if (splitParamIndex == -1) {
            return pjp.proceed();
        }
        Object splitParam = args[splitParamIndex];
        if (!(splitParam instanceof Object[]) && !(splitParam instanceof List) && !(splitParam instanceof Set)) {
            return pjp.proceed();
        }
        boolean notMeet = (splitParam instanceof Object[] && ((Object[]) splitParam).length <= splitLimit)
                || (splitParam instanceof List && ((List<?>) splitParam).size() <= splitLimit)
                || (splitParam instanceof Set && ((Set<?>) splitParam).size() <= splitLimit);
        if (notMeet) {
            return pjp.proceed();
        }
        // Optional deduplication for List
        if (splitParam instanceof List) {
            List<?> list = (List<?>) splitParam;
            if (list.size() > 1) {
                splitParam = new ArrayList<>(new HashSet<>(list));
            }
        }
        int batchNum = getBatchNum(splitParam, splitGroupNum);
        if (batchNum == 1) {
            return pjp.proceed();
        }
        CompletableFuture<?>[] futures = new CompletableFuture[batchNum];
        ThreadPoolEnum threadPool = anno.setThreadPool();
        if (threadPool == null) {
            return pjp.proceed();
        }
        for (int batch = 0; batch < batchNum; batch++) {
            final int currentBatch = batch;
            futures[batch] = CompletableFuture.supplyAsync(() -> {
                Object[] newArgs = Arrays.copyOf(args, args.length);
                try {
                    newArgs[splitParamIndex] = getPartParam(splitParam, splitGroupNum, currentBatch);
                    return pjp.proceed(newArgs);
                } catch (Throwable e) {
                    throw new RuntimeException(e);
                }
            }, threadPool.getThreadPoolExecutor());
        }
        CompletableFuture.allOf(futures).get();
        Class<? extends HandleReturn> handlerCls = anno.handlerReturnClass();
        List<Object> results = new ArrayList<>(futures.length);
        for (CompletableFuture<?> f : futures) {
            results.add(f.get());
        }
        return handlerCls.getDeclaredMethods()[0]
                .invoke(handlerCls.getDeclaredConstructor().newInstance(), results);
    }

    public Integer getBatchNum(Object param, Integer groupSize) {
        if (param instanceof Object[]) {
            return ((Object[]) param).length % groupSize == 0 ? ((Object[]) param).length / groupSize
                    : ((Object[]) param).length / groupSize + 1;
        } else if (param instanceof Collection) {
            return ((Collection<?>) param).size() % groupSize == 0 ? ((Collection<?>) param).size() / groupSize
                    : ((Collection<?>) param).size() / groupSize + 1;
        }
        return 1;
    }

    public Object getPartParam(Object param, Integer groupSize, Integer batch)
            throws NoSuchMethodException, InvocationTargetException, InstantiationException, IllegalAccessException {
        if (param instanceof Object[]) {
            Object[] arr = (Object[]) param;
            int end = Math.min((batch + 1) * groupSize, arr.length);
            return Arrays.copyOfRange(arr, batch * groupSize, end);
        } else if (param instanceof List) {
            List<?> list = (List<?>) param;
            int end = Math.min((batch + 1) * groupSize, list.size());
            return list.subList(batch * groupSize, end);
        } else if (param instanceof Set) {
            List<?> list = new ArrayList<>((Set<?>) param);
            int end = Math.min((batch + 1) * groupSize, list.size());
            Set<Object> set = param.getClass().getDeclaredConstructor().newInstance();
            set.addAll(list.subList(batch * groupSize, end));
            return set;
        }
        return null;
    }
}

Define Parameter Marker

@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.PARAMETER)
public @interface NeedSplitParam {}

Return Handling Interface

public interface HandleReturn {
    Object handleReturn(List<?> results);
}

Merge Function Example

public class MergeFunction implements HandleReturn {
    @Override
    public Object handleReturn(List<?> results) {
        if (results == null) return null;
        if (results.size() <= 1) return results.get(0);
        List<Object> merged = new ArrayList<>((List<?>) results.get(0));
        for (int i = 1; i < results.size(); i++) {
            merged.addAll((List<?>) results.get(i));
        }
        return merged;
    }
}

Apply the annotation to a service method, e.g.:

@SplitWorkAnnotation(setThreadPool = LIST_DEVICE_EXECUTOR, splitLimit = 20, splitGroupNum = 10)
public List<DeviceDetail> listDeviceDetail(Long projectId, @NeedSplitParam List<Long> deviceId) {
    // original logic
}

This solution is suitable for large‑batch IN queries where the result can be simply merged (e.g., SUM, COUNT, TOP‑N). It is not appropriate for pagination or cases that do not fit the split‑merge formula.