How to custom sort only specific results in Spring MongoDB?

Is it possible to sort only results that matches a condition in Spring Mongo? Let us say I have this data:

Color	Fruit	Amount
Orange	Orange	23
Red	Apple	4
Red	Strawberry	66
Yellow	Banana	2

I want to sort the list to display the Fruits with color red on Top and the remaining fruits will be sorted by amount. So the resulting table should be.

Color	Fruit	Amount
Red	Apple	4
Red	Strawberry	66
Yellow	Banana	2
Orange	Orange	23

So far here is what I’ve tried using aggregation.

val match1: MatchOperation = Aggregation.match(Criteria("Color").`is`("Red"))
val match2: MatchOperation = Aggregation.match(Criteria("Color").`is`("Red").not())
val sortByAmount=  sort(Sort.Direction.ASC, "Amount")
val aggregation = Aggregation.newAggregation(match1, sortByAmount, match2, sortByAmount)

val output: AggregationResults<Fruits> = mongoTemplate.aggregate(aggregation, "fruits", Fruits::class.java)

But I’m only getting this as a result

Color	Fruit	Amount
Red	Apple	4
Red	Strawberry	66

Answer

One way to do this using the $facet; with two facets one for “red” and the other for the “not-reds”.

Aggregation agg = newAggregation(
    facet(
        match(where("color").is("red")),
        sort(ASC, "amt"))
    .as("reds")
    .and(
        match(where("color").ne("red")),
        sort(ASC, "amt")
    ).as("others"),
    project()
       .and(arrayOf("reds").concat("others"))
       .as("result"),
    unwind("result"),
    replaceRoot("result")
);

AggregationResults<Document> results = mongoTemplate.aggregate(agg, "fruits", Document.class);
results.forEach(doc -> System.out.println(doc.toJson());

I am using input documents as follows, for brevity: { color: 'red', amt: 12 }, { color: 'blue', amt: 2 }, { color: 'green', amt: 4 }, { color: 'red', amt: 3 }, { color: 'yellow', amt: 5 }

Another way is by using the $function operator. This requires Spring Data MongoDB v3.2 and MongoDB v4.4. I didn’t have a chance to actually run the code (I think it should work).

The pipeline is to be built using these four stages:

GroupOperation groupOperation = Aggregation.group().push("$$ROOT").as("docs");
AddFieldsOperation addFieldsOperation = Aggregation.addFields()
                                            .addFieldWithValue("docs",
                                                                ScriptOperators.Function.function(JAVASCRIPT_FUNCTION).args("docs").lang("js"))
                                            .build();
UnwindOperation unwindOperation = Aggregation.unwind("docs");
ReplaceRootOperation replaceRootOperation = Aggregation.replaceRoot("docs");

The string JAVASCRIPT_FUNCTION (used in the AddFieldsOperation) has the following JavaScript function which actually sorts an array of the fruits documents.

function (a, b) {
  if (a.color == 'red' && b.color == 'red') {
      return a.amt - b.amt;
  }
  if (a.color == 'red' || b.color == 'red') {
     if (a.color == 'red') return 0;
     return 1;
  }
  return a.amt - b.amt;
}