pysparkで類似文章を求めてみた

calc_sim.py

import sys,os
from pyspark.sql import SQLContext, Row
import pyspark.sql.functions as func
import logging
from pyspark import SparkContext

logger = logging.getLogger('py4j')

from optparse import OptionParser
parser = OptionParser()
parser.add_option("-s", "--src", dest="src",
                                    help="src file")
parser.add_option("-d", "--dest",dest="dest",
                  help="dest directory")
parser.add_option("-p", "--part", dest="part",type='int',default=4,
                  help="repatition size")
parser.add_option("-g", "--gt", dest="gt",type='float',default=0.75,
                  help="similarity threshold(0.0-1.0)")

(options, args) = parser.parse_args()

def calc_sim():
    sc =SparkContext()
    sqlContext = SQLContext(sc)

    df1 = sqlContext.read.format("com.databricks.spark.csv").option("header", "true").load(options.src)
    df2 = sqlContext.read.format("com.databricks.spark.csv").option("header", "true").load(options.src)
    df1 = df1.repartition(options.part).cache()
    df2 = df2.repartition(options.part).cache()
    test1 = df1.take(1)
    test2 = df2.take(1)
    df1 = df1.withColumnRenamed('cid','cid_1')
    df2 = df2.withColumnRenamed('cid','cid_2')
    df1 = df1.withColumnRenamed('title','title_1')
    df2 = df2.withColumnRenamed('title','title_2')
    dfout=df1.join(df2).\
           filter(df1.cid_1!=df2.cid_2).\
           filter(1.0 - (func.levenshtein(df1.title_1,df2.title_2)  / func.greatest(func.length(df1.title_1),func.length(df2.title_2))) > options.gt)\
           .select(func.col('cid_1'),func.col('title_1'),func.col('cid_2'),func.col('title_2'),\
                   (1.0 - func.levenshtein(df1.title_1,df2.title_2)  / func.greatest(func.length(df1.title_1),func.length(df2.title_2))).alias('simulality'))
    dfout.explain()
    dfout.write.save(options.dest,compression='snappy')
    return 0

def main():
    result = calc_sim()
    if result==0:
        sys.exit(0)
if __name__=='__main__':main()

csv_sample.csv

lunch.sh

#!/bin/bash
DEST=outtest2
rm -rf $DEST
spark-submit --executor-memory 10G  --total-executor-cores 4 --master local[4] calc_sim.py -s src.csv.gz -d $DEST -p 4 -g 0.7

result

+-----+----------------+-----+----------------+-------------------+
|cid_1|         title_1|cid_2|         title_2|         simulality|
+-----+----------------+-----+----------------+-------------------+
|    2|             いいい|    6|             あいう|0.33333333333333337|
|    2|             いいい|   10|    あいあいいいああうあいあ|               0.25|
|    2|             いいい|   11|あいあいうあいあうあいいあうあい|             0.1875|
|    2|             いいい|    4|             ああい|0.33333333333333337|
|    2|             いいい|    8|           ああいああ|0.19999999999999996|
|    2|             いいい|    5|             いいう| 0.6666666666666667|
|    2|             いいい|    9|    あいあいあいあいあいあい|               0.25|
|    6|             あいう|    2|             いいい|0.33333333333333337|
|    6|             あいう|   10|    あいあいいいああうあいあ|               0.25|
|    6|             あいう|    3|             ううう|0.33333333333333337|
|    6|             あいう|    7|          ああああああ|0.16666666666666663|
|    6|             あいう|   11|あいあいうあいあうあいいあうあい|             0.1875|
|    6|             あいう|    4|             ああい|0.33333333333333337|
|    6|             あいう|    8|           ああいああ|                0.4|
|    6|             あいう|    1|             あああ|0.33333333333333337|
|    6|             あいう|    5|             いいう| 0.6666666666666667|
|    6|             あいう|    9|    あいあいあいあいあいあい|0.16666666666666663|
|   10|    あいあいいいああうあいあ|    2|             いいい|               0.25|
|   10|    あいあいいいああうあいあ|    6|             あいう|               0.25|
|   10|    あいあいいいああうあいあ|    7|          ああああああ|                0.5|
|   10|    あいあいいいああうあいあ|   11|あいあいうあいあうあいいあうあい|             0.5625|
|   10|    あいあいいいああうあいあ|    4|             ああい|               0.25|
|   10|    あいあいいいああうあいあ|    8|           ああいああ|0.41666666666666663|
|   10|    あいあいいいああうあいあ|    1|             あああ|               0.25|
|   10|    あいあいいいああうあいあ|    5|             いいう|               0.25|
|   10|    あいあいいいああうあいあ|    9|    あいあいあいあいあいあい| 0.6666666666666667|
|    3|             ううう|    6|             あいう|0.33333333333333337|
|    3|             ううう|   11|あいあいうあいあうあいいあうあい|             0.1875|
|    3|             ううう|    5|             いいう|0.33333333333333337|
|    7|          ああああああ|    6|             あいう|0.16666666666666663|
|    7|          ああああああ|   10|    あいあいいいああうあいあ|                0.5|
|    7|          ああああああ|   11|あいあいうあいあうあいいあうあい|              0.375|
|    7|          ああああああ|    4|             ああい|0.33333333333333337|
|    7|          ああああああ|    8|           ああいああ| 0.6666666666666667|
|    7|          ああああああ|    1|             あああ|                0.5|
|    7|          ああああああ|    9|    あいあいあいあいあいあい|                0.5|
|   11|あいあいうあいあうあいいあうあい|    2|             いいい|             0.1875|
|   11|あいあいうあいあうあいいあうあい|    6|             あいう|             0.1875|
|   11|あいあいうあいあうあいいあうあい|   10|    あいあいいいああうあいあ|             0.5625|
|   11|あいあいうあいあうあいいあうあい|    3|             ううう|             0.1875|
|   11|あいあいうあいあうあいいあうあい|    7|          ああああああ|              0.375|
|   11|あいあいうあいあうあいいあうあい|    4|             ああい|             0.1875|
|   11|あいあいうあいあうあいいあうあい|    8|           ああいああ|             0.3125|
|   11|あいあいうあいあうあいいあうあい|    1|             あああ|             0.1875|
|   11|あいあいうあいあうあいいあうあい|    5|             いいう|             0.1875|
|   11|あいあいうあいあうあいいあうあい|    9|    あいあいあいあいあいあい|             0.6875|
|    4|             ああい|    2|             いいい|0.33333333333333337|
|    4|             ああい|    6|             あいう|0.33333333333333337|
|    4|             ああい|   10|    あいあいいいああうあいあ|               0.25|
|    4|             ああい|    7|          ああああああ|0.33333333333333337|
|    4|             ああい|   11|あいあいうあいあうあいいあうあい|             0.1875|
|    4|             ああい|    8|           ああいああ|                0.6|
|    4|             ああい|    1|             あああ| 0.6666666666666667|
|    4|             ああい|    9|    あいあいあいあいあいあい|               0.25|
|    8|           ああいああ|    2|             いいい|0.19999999999999996|
|    8|           ああいああ|    6|             あいう|                0.4|
|    8|           ああいああ|   10|    あいあいいいああうあいあ|0.41666666666666663|
|    8|           ああいああ|    7|          ああああああ| 0.6666666666666667|
|    8|           ああいああ|   11|あいあいうあいあうあいいあうあい|             0.3125|
|    8|           ああいああ|    4|             ああい|                0.6|
|    8|           ああいああ|    1|             あああ|                0.6|
|    8|           ああいああ|    5|             いいう|0.19999999999999996|
|    8|           ああいああ|    9|    あいあいあいあいあいあい|0.41666666666666663|
|    1|             あああ|    6|             あいう|0.33333333333333337|
|    1|             あああ|   10|    あいあいいいああうあいあ|               0.25|
|    1|             あああ|    7|          ああああああ|                0.5|
|    1|             あああ|   11|あいあいうあいあうあいいあうあい|             0.1875|
|    1|             あああ|    4|             ああい| 0.6666666666666667|
|    1|             あああ|    8|           ああいああ|                0.6|
|    1|             あああ|    9|    あいあいあいあいあいあい|               0.25|
|    5|             いいう|    2|             いいい| 0.6666666666666667|
|    5|             いいう|    6|             あいう| 0.6666666666666667|
|    5|             いいう|   10|    あいあいいいああうあいあ|               0.25|
|    5|             いいう|    3|             ううう|0.33333333333333337|
|    5|             いいう|   11|あいあいうあいあうあいいあうあい|             0.1875|
|    5|             いいう|    8|           ああいああ|0.19999999999999996|
|    5|             いいう|    9|    あいあいあいあいあいあい|0.16666666666666663|
|    9|    あいあいあいあいあいあい|    2|             いいい|               0.25|
|    9|    あいあいあいあいあいあい|    6|             あいう|0.16666666666666663|
|    9|    あいあいあいあいあいあい|   10|    あいあいいいああうあいあ| 0.6666666666666667|
|    9|    あいあいあいあいあいあい|    7|          ああああああ|                0.5|
|    9|    あいあいあいあいあいあい|   11|あいあいうあいあうあいいあうあい|             0.6875|
|    9|    あいあいあいあいあいあい|    4|             ああい|               0.25|
|    9|    あいあいあいあいあいあい|    8|           ああいああ|0.41666666666666663|
|    9|    あいあいあいあいあいあい|    1|             あああ|               0.25|
|    9|    あいあいあいあいあいあい|    5|             いいう|0.16666666666666663|
+-----+----------------+-----+----------------+-------------------+