CODEWITHSUNDEEP
pythonmongodbpymongo
Sreeragh A R
Sreeragh A R

Reputation: 3021

Grouping documents in MongoDB on special condition

My collection contains

{name:'p1', age: 20}
{name: 'p2', age: 21}
{name: 'p3', age: 23}
{name: 'p4', ag:41 }

I want to group persons such that for any person in the group there exist another person int the group such that difference between their ages is at most 2. Here resulting group will contain

expected result 

[{name:'p1' ...}, {name:'p2' ...}, {name: 'p3'}]

since ages of p2 -p1 = 1 and p3-p2 = 2

p1,p2,p3 form a group

Upvotes: 0

Views: 102

Answers (1)

Alex Blex
Alex Blex

Reputation: 37018

Disclaimer

Before reading the rest of the answer, please read https://docs.mongodb.com/manual/core/aggregation-pipeline-limits/ The resulting document in the question is expected to have an array of all documents that belong to particular age group. Size of that array cannot exceed 16MB, so the code below will work only for very small collections of tiny documents.

The code:

db.collection.aggregate([
    { $sort: { age: 1 } },
    { $group: {
            _id: null,
            ages: { $push: "$age" }
    } },
    { $addFields: {
        ranges: { $reduce: { 
            input: { $range: [ 1, { $size: "$ages" }, 1 ] }, 
            initialValue: [ [ { $arrayElemAt: [ "$ages", 0 ] } ] ], 
            in: { $cond: { 
                if:  { $gt: [
                    { $subtract: [ { $arrayElemAt: [ "$ages", "$$this" ] }, { $arrayElemAt: [ "$ages", { $subtract: [ "$$this", 1 ] } ] } ] },
                    2
                    ] }, 
                then: { $concatArrays: [ "$$value",  [ [ { $arrayElemAt: [ "$ages", "$$this" ] } ] ] ] }, 
                else: { $concatArrays: [ 
                    { $slice: [ "$$value" , { $subtract: [ { $size: "$$value" }, 1 ] } ] },
                    [ { $concatArrays: [ 
                        { $arrayElemAt: [ { $slice: [ "$$value" , -1 ] }, 0 ] }  ,  
                        [ { $arrayElemAt: [ "$ages", "$$this" ] } ]
                    ]  } ]
                ] }
            } }
        } } 
    } },
    { $unwind: "$ranges" }, 
    { $lookup: {
       from: "collection",
       localField: "ranges",
       foreignField: "age",
       as: "group"
     } },
     { $project: { _id: 0, group: 1 } }
])

The part that may require a bit of explanation is how to calculate age groups.

For that, we get all ages using $group into a single array and then $addFields "ranges" - a 2D array of age groups with gaps between oldest person in a younger group and a youngest person in the older group is greater than 2 years.

The array is calculated using $reduce of a $range array of indexes of all ages but first, which goes to initial value.

The reduce expression is a $cond which calculates difference between current and previous ($subtract) element of the array of all ages.

If it is greater than 2, a new age group is added using $concatArrays. Otherwise the age is added to the oldest group using $slice to push to the last group in the ranges array and $setUnion to eliminate duplicates.

When the age groups are calculated, we $lookup the same collection by age to group them in the "group" array.

Upvotes: 2

Related Questions