if (queryObject.getModel() instanceof IdentifiableType && !((IdentifiableType) queryObject.getModel()).hasSingleIdAttribute()) {
  IdentifiableType entityType = (IdentifiableType) queryObject.getModel();
  ManagedType idType = (ManagedType<?>) entityType.getIdType();
  List<Expression<?>> expressions = new ArrayList<>();
  for (Attribute<?, ?> attr : idType.getAttributes()) {
    expressions.add(queryObject.get(queryObject.get(entityType.getId(idType)).get(attr));
  }
  query.select(criteriaBuilder.function("my_count", Long.class, expressions.toArray(new Expression[expressions.size()])));
} else {
  query.select(criteriaBuilder.countDistinct(queryObject));
}