Get all @ManyToOne relations in JpaRepository with Ids

Question

I need to get a List<List<Office>> from a single Repository Method. I ended up using this method:

List<List<Office>> lists = officeDao.findAllByCompanyIdIn(ids);

Which works, except when looking at lists.get(i), there are more elements than needed. This is because it doesn't group the Company's offices together (in the same index at that level) if they have the same company id.

Here are the related entities for some context:

@Setter
@Getter
@Entity
@AllArgsConstructor
public class Company implements Serializable {

  private static final long serialVersionUID = -6007975840330441233L;

  @Id
  @GeneratedValue
  @Column(columnDefinition = "uuid", updatable = false)
  private UUID id;

  private String name;

  private BigDecimal balance;

  private CompanyType type;

  @OneToMany(
      targetEntity = Office.class,
      orphanRemoval = true,
      fetch = FetchType.LAZY,
      mappedBy = "company",
      cascade = CascadeType.ALL)
  private List<Office> offices;

  public enum CompanyType {
    PRIVATE_LIMITED,
    SOLE_TRADER,
    PUBLIC
  }
}

@Setter
@Getter
@Entity
@AllArgsConstructor
public class Office implements Serializable {

  private static final long serialVersionUID = 476906690321337185L;

  @Id
  @GeneratedValue
  @Column(columnDefinition = "uuid", updatable = false)
  private UUID id;

  private String officeLocation;

  private int maximumHeadcount;

  private FinanceType financeType;

  @ManyToOne(targetEntity = Company.class, fetch = FetchType.LAZY, cascade = CascadeType.ALL)
  private Company company;

  public enum FinanceType {
    RENTING,
    OWNED,
    LEASING
  }
}

Apologies for the super-ambiguous question, but please comment if you need more context!

Cheers guys!

Answer 1

It turns out you can achieve the thing I wanted to using a few techniques from the Java Streams API (Java 8+).

I didn't have to change anything on the Entity/Dao sides, I just loaded it using this block of code here:

List<List<Office>> bucketedList = new ArrayList<>();

officeDao.findAllByCompanyIdIn(ids).stream()
    .flatMap(Collection::stream)
    .collect(Collectors.groupingBy(office -> office.getCompany().getId()))
    .forEach((uuid, offices) -> bucketedList.add(offices));

I'm essentially creating a new (nested) list where the each object will be sorted into, grouped by office.getCompany().getId().

If anyone has any suggestions or improvements to this method, feel free to comment below!