<commit_before>import argparse
from pyspark.sql import SparkSession, Row
from pyspark.sql.functions import (col, collect_list, lit, sort_array, struct, udf)
import pyspark.sql.functions as f

def catPages(pagea):
    text = ''
    pages = []
    for p in pagea:
        cur = p.page.asDict(True)
        off = len(text)
        if 'regions' not in cur:
            cur['regions'] = []
        i = 0
        while i < len(cur['regions']):
            cur['regions'][i]['start'] += off
            i += 1
        pages.append(cur)
        text += p.text
    return {'text': text, 'pages': pages}

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='IA Periodical',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('-c', '--max-chars', type=int, default=2000000,
                        help='Maximum characters')
    parser.add_argument('-p', '--max-pages', type=int, default=1000,
                        help='Maximum pages')

    parser.add_argument('inputPath', metavar='<path>', help='input path')
    parser.add_argument('mapPath', metavar='<path>', help='map path')
    parser.add_argument('outputPath', metavar='<path>', help='output path')

    config = parser.parse_args()

    spark = SparkSession.builder.appName('IA Periodical').getOrCreate()

    corpus = spark.read.load(config.inputPath).withColumnRenamed('book', 'issue')

    smap = spark.read.json(config.mapPath).withColumnRenamed('identifier', 'issue')

    goods = corpus.groupBy('issue'
                ).agg(f.count('id').alias('pp'), f.sum(f.length('text')).alias('tlen')
                ).filter( (col('pp') <= config.max_pages)
                          & (col('tlen') <= config.max_chars)
                          & ~(col('issue').rlike('_(index|appendix|contents)'))
                          & ~(col('issue').rlike('\d{4}-\d{4}'))
                ).select('issue'
                ).join(smap, 'issue')

    cat_pages = udf(lambda pagea: catPages(pagea),
                    corpus.select('text', 'pages').schema.simpleString())

    spark.conf.set('spark.sql.shuffle.partitions', 5000)

    corpus.join(goods, 'issue'
        ).groupBy('issue', 'series'
        ).agg(
            cat_pages(sort_array(collect_list(struct('seq', 'text',
                                                     col('pages')[0].alias('page'))))).alias('p')
        ).select(col('issue').alias('id'), 'issue', 'series',
                 f.regexp_replace(f.split('issue', '_')[2], r'-\d{4}.*$', ''
                                  ).cast('date').cast('string').alias('date'),
                 col('p.text'), col('p.pages')
        ).filter(col('date').isNotNull()
        ).write.save(config.outputPath)

    spark.stop()
<commit_msg>Use date in series-map.<commit_after>import argparse
from pyspark.sql import SparkSession, Row
from pyspark.sql.functions import (col, collect_list, lit, sort_array, struct, udf)
import pyspark.sql.functions as f

def catPages(pagea):
    text = ''
    pages = []
    for p in pagea:
        cur = p.page.asDict(True)
        off = len(text)
        if 'regions' not in cur:
            cur['regions'] = []
        i = 0
        while i < len(cur['regions']):
            cur['regions'][i]['start'] += off
            i += 1
        pages.append(cur)
        text += p.text
    return {'text': text, 'pages': pages}

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='IA Periodical',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('-c', '--max-chars', type=int, default=2000000,
                        help='Maximum characters')
    parser.add_argument('-p', '--max-pages', type=int, default=1000,
                        help='Maximum pages')

    parser.add_argument('inputPath', metavar='<path>', help='input path')
    parser.add_argument('mapPath', metavar='<path>', help='map path')
    parser.add_argument('outputPath', metavar='<path>', help='output path')

    config = parser.parse_args()

    spark = SparkSession.builder.appName('IA Periodical').getOrCreate()

    corpus = spark.read.load(config.inputPath).withColumnRenamed('book', 'issue')

    smap = spark.read.json(config.mapPath)

    goods = corpus.groupBy('issue'
                ).agg(f.count('id').alias('pp'), f.sum(f.length('text')).alias('tlen')
                ).filter( (col('pp') <= config.max_pages)
                          & (col('tlen') <= config.max_chars)
                          & ~(col('issue').rlike('_(index|appendix|contents)'))
                          & ~(col('issue').rlike('\d{4}-\d{4}'))
                ).select('issue'
                ).join(smap, 'issue')

    cat_pages = udf(lambda pagea: catPages(pagea),
                    corpus.select('text', 'pages').schema.simpleString())

    spark.conf.set('spark.sql.shuffle.partitions', 5000)

    corpus.join(goods, 'issue'
        ).groupBy('issue', 'series', 'date'
        ).agg(
            cat_pages(sort_array(collect_list(struct('seq', 'text',
                                                     col('pages')[0].alias('page'))))).alias('p')
        ).select(col('issue').alias('id'), 'issue', 'series',
                 col('date').cast('date').cast('string').alias('date'),
                 col('p.text'), col('p.pages')
        ).filter(col('date').isNotNull()
        ).write.save(config.outputPath)

    spark.stop()

{'fraction_non_alphanumeric': 0.17603956768638945, 'fraction_numerical': 0.007693716797948343, 'mean_word_length': 2.2851985559566788, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 7, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2066254', 'n_tokens_mistral': 1704, 'n_tokens_neox': 1621, 'n_words': 278}