全文検索エンジン Solr の検索 plugin 作成 (第1回)

初めまして、ネクストでデータマイニング・レコメンドを担当している古川と申します。本ブログでは、日々の業務でお世話になっている、アルゴリズムやソフトウェアなどを紹介していきたいと思います。

まず第 1 弾として、全文検索エンジン Solr の検索 plugin 作成方法について数回の連載に分けて紹介します。

導入

Solr は、

日本語の全文検索が可能である
RDMBS 同様、フィールド値を使ったソート、複数フィールドを使った高速なソートが可能である
分散処理機能を備え、スケーラビリティに優れている
キャッシュ機能を備え高速なレスポンスが可能である
レプリケーション可能である
オンラインでのデータ更新がサポートされている

などの多くの利点を持つオープンソースの全文検索エンジンで、様々なサイトでの利用が広がっています。

課題

デフォルト機能でも多くのニーズに答えてくれる Solr ですが、サイト固有のニーズを実現するためには、多少のカスタマイズが必要になることがあります。 HOME'S でのニーズを例に説明します。

フィールド値として家賃、面積などを保存しておけば、「家賃が 8 万円以下」という検索条件について、合致する物件一覧を、「家賃が安い順-面積の広い順」、「家賃が高い順 - 面積広い順」といった順序でソートすることは、 RDBMS 同様 Solr でも可能です。しかしながら、これでは「家賃がそこそこの割に、面積もそこそこ広い」という条件の良い物件があったとしても検索結果上位に表示されにくい、という問題があります。もちろん、ユーザにより詳細な検索条件を入力して頂ければ解決する問題ですが、 1 度目の検索でサイトを離脱されてしまうユーザも多く、少ない入力条件でいかに“いい感じ”の検索結果を返すかということは、非常に重要な課題です。

“いい感じ” の結果を返すためには、「家賃、面積」という2つのフィールドに関して、個々に考えるのではなく、2つのフィールドを総合的に加味する必要があります。しかし、これを Solr のデフォルト機能で実現するのは難しく、 plugin を作成ました。

Lucene での実現方法

本連載では、検索クエリと複数フィールド値を加味しながらソート順を決定する単純なサンプルプログラムを例に、作成方法を説明していきたいと思います。

連載 1 回目の今回は、 Solr 用の検索 plugin 作成をする前段階として、 Solr の使っている Lucene ライブラリでの実現方法を紹介します。なお、実現方法に関しては、株式会社ロンウイットの関口さんが、ブログで紹介されている方法のうち、 Collector クラスを継承して実現する方法を参考にさせて頂きました。

以下のようにフィールドが id 、 type 、 x 、 y からなるドキュメントセットについて、検索条件として、 type と x 、 y を指定し、同じ type のドキュメントをスコアの小さい順に取得するという検索を考えます。このとき、

Score = weight(type)(qx‐x)(qx-x) + (1-weight(type)) (qy-y)(qy-y)

のように、typeによって x 、 y の weight が変わる値を検索用スコアとして考えました。

source.tsv

#id	type	x	y
1	a	10	20
2	b	10	20
3	a	11	11
4	b	11	11

次のプログラムは、上記 TSV ファイルを読み込み、検索用のインデックス作成、プログラム内で指定した条件で検索を行うというプログラムです。

BlogTest.java

package blog;

import java.util.HashMap;
import java.util.ArrayList;

import java.io.FileReader;
import java.io.BufferedReader;
import java.io.FileWriter;
import java.io.BufferedWriter;

import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.NumericField;

import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.Term;

import org.apache.lucene.analysis.SimpleAnalyzer;
import org.apache.lucene.util.Version;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.RAMDirectory;

import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.TermQuery;
import org.apache.lucene.search.FieldCache;


public class BlogTest {
    public static void main(String[] args) {
        String srcPath  = args[0];
        Directory directory = new RAMDirectory();
        System.out.println("[create index]");
        int numDocuments = createIndex(directory, srcPath);
        System.out.println("\tloaded " + numDocuments + " documents");
        
        System.out.println("[run test]");
        BlogTestQuery[] testQueries = new BlogTestQuery[2];
        testQueries[0] = new BlogTestQuery("a", 10, 10);
        testQueries[1] = new BlogTestQuery("b", 10, 10);

        searchTest(directory, testQueries, 5);
        System.out.println("[done]");
    }
    
    public static int createIndex(Directory directory, String path) {
        int counter = 0;
        try {
            IndexWriterConfig config = new IndexWriterConfig(
                                                             Version.LUCENE_34, 
                                                             new SimpleAnalyzer(Version.LUCENE_34));
            IndexWriter     iw = new IndexWriter(directory, config);
            FileReader      fr = new FileReader(path);
            BufferedReader  br = new BufferedReader(fr);

            String strLine;
            String[] buf;
            while ( (strLine = br.readLine()) != null ) {
            if (strLine.charAt(0) == '#' ) continue;
                buf = strLine.split("\\t");
                if ( buf.length != 4 ) continue;
                counter += 1;
                String id   = buf[0];
                String type = buf[1];
                int x = Integer.parseInt(buf[2]);
                int y = Integer.parseInt(buf[3]);
                Document doc = new  Document();                
                doc.add(new Field("id", id, Field.Store.YES, Field.Index.NOT_ANALYZED));
                doc.add(new Field("type", type, Field.Store.YES, Field.Index.NOT_ANALYZED));
                doc.add(new NumericField("x",  Field.Store.YES, true).setIntValue(x));
                doc.add(new NumericField("y",  Field.Store.YES, true).setIntValue(y));
                iw.addDocument(doc);
            }
            br.close();
            fr.close();
            iw.close();
        } catch (Exception e) {
            e.printStackTrace();
            System.exit(1);
        }

        return counter;
    }

    public static void searchTest(Directory directory, BlogTestQuery[] testQueries, int numLimit) {
        // searcher作成
        System.out.print("\tinitializing searcher...");
        IndexSearcher is = null;
        IndexReader ir = null;

        HashMap<String, Double> weights = new HashMap<String, Double>(2);
        weights.put("a", 0.5);
        weights.put("b", 0.999);
        int[][] caches = new int[2][];
        try {
            is = new IndexSearcher(directory, true);
            ir = is.getIndexReader();
            caches[0] = FieldCache.DEFAULT.getInts(ir, "x");
            caches[1] = FieldCache.DEFAULT.getInts(ir, "y");
        } catch (Exception e) {
            e.printStackTrace();
            System.exit(1);
        }
        System.out.println(" done.");

        // run search
        for ( int i=0; i<testQueries.length; i++) {
            BlogTestQuery testQuery = testQueries[i];
            Query query = new TermQuery(new Term("type", testQuery.type));
            double weight = weights.get(testQuery.type);
            BlogTestCollector collector = new BlogTestCollector(ir, caches, weight, testQuery.x, testQuery.y);
            try {
                is.search(query, null, collector);
            } catch (Exception e) {
                e.printStackTrace();
                System.exit(1);
            }
            ArrayList<DocIdScore> docIdScores = collector.getDocIdScores(numLimit);
            System.out.print("\ttype=" + testQuery.type + "\tx=" + testQuery.x + "\ty=" + testQuery.y);
            System.out.println("\tweight=" + weight);
            System.out.println("\t\tid\ttype\tx\ty\tscore");
            for(int j=0; j<docIdScores.size(); j++){
                int id            = (docIdScores.get(j)).getDocumentId();
                double score      = (docIdScores.get(j)).getScore();
                try {
                    Document doc = is.doc(id);
                    System.out.print("\t\t");
                    System.out.print(doc.get("id")+"\t");
                    System.out.print(doc.get("type")+"\t");
                    System.out.print(doc.get("x")+"\t");
                    System.out.print(doc.get("y")+"\t");
                    System.out.print(score+"\n");
                } catch (Exception e) {
                    e.printStackTrace();
                    System.exit(1);
                }
            }
        }
    }
}

BlogTestCollector.java

package blog;

import java.lang.Math;
import java.util.Collections;
import java.util.ArrayList;
import java.util.Comparator;

import org.apache.lucene.document.Document;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.FieldCache;

public class BlogTestCollector extends Collector {
    private ArrayList<DocIdScore> docIdScores;
    private IndexReader reader;
    private int docBase;
    private int[][] fieldsCaches;
    private double weight;
    private int x;
    private int y;

    public BlogTestCollector(IndexReader reader, int[][] caches, double weight, int x, int y) {
        this.reader = reader;
        this.docIdScores = new ArrayList<DocIdScore>();
        try {
            this.fieldsCaches = caches;
            this.weight = weight;
            this.x = x;
            this.y = y;
        } catch (Exception e) {
            e.printStackTrace();
            System.exit(1);
        }
    }

    @Override
    public void setNextReader(IndexReader reader , int docBase) {
        this.docBase = docBase;
    }

    @Override
    public void setScorer(Scorer socorer){}

    @Override
    public boolean acceptsDocsOutOfOrder(){
        return true;
    }

    @Override
    public void collect(int id) {
        try {
            int docid = this.docBase + id;
            int x = this.fieldsCaches[0][docid];
            int y = this.fieldsCaches[1][docid];
            double score = 0;
            score = this.weight*(this.x - x)*(this.x - x) + (1.0-this.weight)*(this.y - y)*(this.y - y);
            this.docIdScores.add(new DocIdScore( docid, score ));
        } catch (Exception e) {
            e.printStackTrace();
            System.exit(1);
        }
    }

    public ArrayList<DocIdScore> getDocIdScores(int limit) {
        Collections.sort( this.docIdScores, new ScoreComparator() );
        ArrayList<DocIdScore> retvals = new ArrayList<DocIdScore>(limit);
        int maxSize = this.docIdScores.size();
        if ( limit > maxSize ) {
            limit = maxSize;
        }
        for (int i=0; i<limit; i++) {
            retvals.add(this.docIdScores.get(i));
        }
        return retvals;
    }

    static class ScoreComparator implements Comparator<DocIdScore> {
        public int compare( DocIdScore a, DocIdScore b) {
            double a_score = a.getScore();
            double b_score = b.getScore();
            if ( a_score > b_score ) {
                return 1;
            } else if (a_score < b_score) {
                return -1;
            } else {
                return 0;
            }
        }
    }
}

BlogTestQuery.java

package blog;

class BlogTestQuery {
    public String type;
    public int x;
    public int y;
    public BlogTestQuery(String type, int x, int y) {
        this.type = type;
        this.x    = x;
        this.y    = y;
    }
}

DocidScore.java

package blog;

class DocIdScore {
    private int   id;
    private double score;

    public DocIdScore(int id, double score) {
        this.id    = id;
        this.score = score;
    }

    public int getDocumentId() {
        return this.id;
    }

    public double getScore() {
        return this.score;
    }
}

インデックス作成部分、クエリ指定部分など Lucene の基本的な使い方については、参考文献を参照して頂くとして、今回のポイントは、

BlogTest.java 115 行目の、BlogTestCollector オブジェクトを作成し、 IndexSearcher オブジェクトの search() 関数の引数として渡し検索を行う部分
BlogTest.java 120 行目の、BlogTestCollector オブジェクトの getDocIdScores() 関数で検索結果を取得する部分

になります。

1 のように、 search() 関数の引数として BlogTestCollector オブジェクトを渡すことで、 IndexSearcher の search() 関数が条件に合致するドキュメントを発見するたびに、 BlogTestCollector の collect() 関数が実行され、 BlogTestCollector 内部にスコア計算結果が保存されていきます。そして、 2 でその結果一覧をスコア順にソートして取得しているわけです。 collect() 関数、 getDocIdScores() 関数の詳細は、ソースコードをご覧ください。

プログラムの実行結果は、以下の通りです。検索クエリの type の値により、 weight が変更され、 x 、 y が同じ値であっても、異なるソート結果となっていることが分かります。

[create index]
        loaded 4 documents
[run test]
        initializing searcher... done.
        type=a  x=10    y=10    weight=0.5
                id      type    x       y       score
                3       a       11      11      1.0
                1       a       10      20      50.0
        type=b  x=10    y=10    weight=0.999
                id      type    x       y       score
                2       b       10      20      0.1
                4       b       11      11      1.0

Lucene のみで実現しているシステムの場合、検索部分を今回の例のように変更するだけで、ソート順を思いのままに変更できますが、 Solr へ組み込み時には、更に Solr の検索コンポーネントの作法に合わせる必要があります。次回は、その方法に関して解説します。