Boost PostgreSQL IN Query Performance with Spring AOP Parallel Splitting

This guide addresses the performance degradation caused by very large

IN

clauses in PostgreSQL queries within Spring Boot projects and shows how to optimize them by splitting the parameter list, executing the sub‑queries in parallel, and merging the results.

<code>SELECT * FROM device WHERE id IN (1, 2, 3, 4)</code>

The query can be divided into smaller batches:

<code>SELECT * FROM device WHERE id IN (1, 2)
SELECT * FROM device WHERE id IN (3, 4)</code>

These batches are executed concurrently, and the results are combined.

Define AOP Annotation

<code>@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;
}</code>

Define Parameter Annotation

<code>@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.PARAMETER)
public @interface NeedSplitParam {
}
</code>

Implement AOP Aspect for Parallel Splitting

<code>@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].getAnnotation(NeedSplitParam.class) != null) {
                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 tooSmall = (splitParam instanceof Object[] && ((Object[]) splitParam).length <= splitLimit)
                || (splitParam instanceof List && ((List<?>) splitParam).size() <= splitLimit)
                || (splitParam instanceof Set && ((Set<?>) splitParam).size() <= splitLimit);
        if (tooSmall) {
            return pjp.proceed();
        }
        // deduplicate if 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 i = 0; i < batchNum; i++) {
            final int currentBatch = i;
            futures[i] = 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[]) {
            int len = ((Object[]) param).length;
            return (len + groupSize - 1) / groupSize;
        } else if (param instanceof Collection) {
            int len = ((Collection<?>) param).size();
            return (len + groupSize - 1) / groupSize;
        }
        return 1;
    }

    public Object getPartParam(Object param, Integer groupSize, Integer batch) throws Exception {
        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 = (Set<Object>) param.getClass().getDeclaredConstructor().newInstance();
            set.addAll(list.subList(batch * groupSize, end));
            return set;
        }
        return null;
    }
}
</code>

HandleReturn Interface

<code>public interface HandleReturn {
    Object handleReturn(List t);
}
</code>

MergeFunction Implementation

<code>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 first = (List) results.get(0);
        for (int i = 1; i < results.size(); i++) {
            first.addAll((List) results.get(i));
        }
        return first;
    }
}
</code>

By annotating a service method with

@SplitWorkAnnotation

and marking the large collection parameter with

@NeedSplitParam

, developers can automatically split, parallelize, and merge IN‑query results without changing business logic.