MongoDB Aggregation を Java でもう少しスッキリ書きたい

MongoDBでAggregationを使うと

MapReduceより大分処理が早いのがいいのだけど、Javaで書くとどうしても面倒。こんなユーティリティがJavaドライバにあったらなあ、という件。多分どっかにオープンソースみたいなのがあるとお思うけど、とりあえず軽く検索してみて良いのが無かったのでちょっと書いてみた。

例えば、CLIで

db.students.aggregate([
         {$unwind:"$scores"},
         {$group: {_id:"$name", average:{$avg:"$scores.score"}}},
         {$skip:5},
         {$limit:3} ])

みたいな処理をJavaで書くと

DBObject groupFields = new BasicDBObject("_id", "$name");
groupFields.put("average", new BasicDBObject("$avg", "$scores.score"));
AggregationOutput output = collection.aggregate(
        new BasicDBObject("$unwind", "$scores"),
        new BasicDBObject("$group", groupFields),
        new BasicDBObject("$skip", 5),
        new BasicDBObject("$limit", 3));

となってしまう。希望としては

• \$groupとか\$avgとか、Stringで渡すのではなく、開発環境でgrとか入力したらgroup?と言って欲しい

• new BasicDBObject を減らしたい

というわけで、

AggBuilder と AggUtils というクラスを作ってみた。

コンパクトにするためにimport と JavaDoc省略

public class AggBuilder {

    private final DBCollection collection;

    private final List<DBObject> ops = new ArrayList<DBObject>();

    public AggBuilder(DBCollection collection) {
        this.collection = collection;
    }

    public AggregationOutput execute() {
        if (ops.size() > 1) {
            return collection.aggregate(ops.remove(0), ops.toArray(new DBObject[ops.size()]));
        }
        return collection.aggregate(ops.get(0));
    }
    public AggBuilder limit(int param) {
        ops.add(new BasicDBObject("$limit", param));
        return this;
    }

    public AggBuilder skip(int param) {
        ops.add(new BasicDBObject("$skip", param));
        return this;
    }

    public AggBuilder unwind(String param) {
        ops.add(new BasicDBObject("$unwind", param));
        return this;
    }

    public AggBuilder group(String id, DBObject param) {
        DBObject params = new BasicDBObject("_id", id);
        params.putAll(param);
        ops.add(new BasicDBObject("$group", params));
        return this;
    }
}

単純に"$なんちゃら"を"なんちゃらメソッド"にしただけです。あとは、groupとかで使うオペレータをstatic importできるようにUtlisを仕込みます。とりあえずavgだけ。

public class AggUtils {
    /**
     * @param fieldName
     *            name of the average field aggregated
     * @param fieldToAverage
     *            field to calculate average for
     */
    public static DBObject avg(String fieldName, String fieldToAverage) {
        return new BasicDBObject(fieldName, new BasicDBObject("$avg", fieldToAverage));
    }
}

以上仕込んだコードを使ってAggregateすると、CLIにだいぶ近くなる。さっきのJSと比べるとこんな感じ。唯一avgのAggregate後の名前とavgオペレータの順序が逆なのがややこしいかな。

db.students.aggregate([
         {$unwind:"$scores"},
         {$group: {_id:"$name", average:{$avg:"$scores.score"}}},
         {$skip:5},
         {$limit:3} ])

AggregationOutput output = new AggBuilder(collection)
        .unwind("$scores")
        .group("$name", avg("average", "$scores.score"))
        .skip(5)
        .limit(3).execute();

今回はAggregationフレームワークでサポートされている機能の一部を実装したけど、あんまり時間もかからないしgroupをタイプミスしてもコード実行するまで気づかないなんてことがなくなるのは嬉しい。

最後にサンプルコード

package advent;

import static advent.AggUtils.avg;

import java.net.UnknownHostException;

import com.mongodb.AggregationOutput;
import com.mongodb.BasicDBObject;
import com.mongodb.DB;
import com.mongodb.DBCollection;
import com.mongodb.DBObject;
import com.mongodb.MongoClient;

public class AggSample {
    public static void main(String[] args) throws UnknownHostException {
        MongoClient client = new MongoClient("localhost", 27017);
        DB database = client.getDB("school");
        DBCollection collection = database.getCollection("students");

        /*
         * db.students.aggregate([
         * {$unwind:"$scores"},
         * {$group: {_id:"$name", average:{$avg:"$scores.score"}}},
         * {$skip:5},
         * {$limit:3} ])
         */
        // Normal Java code
        DBObject groupFields = new BasicDBObject("_id", "$name");
        groupFields.put("average", new BasicDBObject("$avg", "$scores.score"));
        printResults(collection.aggregate(
                new BasicDBObject("$unwind", "$scores"),
                new BasicDBObject("$group", groupFields),
                new BasicDBObject("$skip", 5),
                new BasicDBObject("$limit", 3)));

        // Use AggBuilder
        printResults(new AggBuilder(collection)
        .unwind("$scores")
        .group("$name", avg("average", "$scores.score"))
        .skip(5)
        .limit(3).execute());
    }

    private static void printResults(AggregationOutput output) {
        for (DBObject obj : output.results()) {
            System.out.println(obj);
        }
    }
}

　まとめ

気付いてしまった人もいると思いますが、\$projectについてはなにもしてません。これは結構自由度が高いので、綺麗に書く方法がすんなり思いつかなかった。多分ちゃんと検索したら誰かがAggregationのオペレータを網羅したツールを作ってるのが見つかると思いますが、とりあえずちゃっちゃと自分で書いてもこれくらいにはできるよね、MongoDBって使いやすいよね、という話でした。