Search code examples
mongodbgoaggregate

Mongo median operation in Golang


I have collection in mongo (go) which type is:

type CreateFeedbackRequest struct {
    UserID     string    `json:"user_id" validate:"required"`
    WaybillID  uint64    `json:"waybill_id" validate:"required"`
    Rating     int       `json:"rating" validate:"required"`
    Comment    string    `json:"comment"`
    ReceivedAt time.Time `json:"received_at" validate:"required"`
}

I need to evaluate median value of rating for last 5 records (by receivedAt time field) for certain user (by his user_id). I have already got this:

matchStage := bson.D{{"$match", bson.D{{"_id", userID}}}}
sortStage := bson.D{{"$sort", bson.D{{"created_at", 1}}}}
limitStage := bson.D{{"$limit", tripsCount}}

cursor, err := r.c.Aggregate(ctx, mongo.Pipeline{matchStage, sortStage, limitStage})

But I don't know then how to got rating median for this 5 rows. And I am not sure about correct way I'm doing this. Help, thank you


Solution

  • After the $limit stage, one option since mongodb version 7.0 is to $group with $median accumulator

    groupgStage := bson.D{{"$group", bson.D{
      {"_id", 0}, 
      {"median", bson.D{{"$median", 
        bson.D{{"$input", "$rating"}, {"method", "approximate"}}
      }}}
    }}}
    

    For older versions, you can

    1. $sort by rating
    2. $group and $push all ratings to an array (all 5 of them after you limit)
    3. $project the item in the middle of the array

    It will look like that:

    sortRatingStage := bson.D{{"$sort", bson.D{{"rating", 1}}}}
    groupStage := bson.D{{"$group", bson.D{{"_id", 0}, {"ratings", bson.D{{"$push", "ratings"}}}}}}
    projectStage := bson.D{{"$project", bson.D{
      {"_id", 0}, 
      {median, bson.D{{"$arrayElemAt", bson.D{
        {"$ratings", bson.D{{"$floor", bson.D{
          {"$divide", bson.A{{bson.D{{"$size", "$ratings"}}, 2}}}
        }}}}
      }}}}
    }}}}