<commit_before>#!/usr/bin/env python
# vim: set fileencoding=utf-8 :
#
# MovieLens indexing tool for generating Elasticsearch indices.
#
# Re-write of an original idea/script by <NAME>.
# Copyright 2014 © <NAME>
#
# This file is licensed to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#  http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

from pprint import pprint
from pyes import *

import argparse
import string
import sys
import os
import json

from threading import Thread
from Queue import Queue


def parse(fname, field_types, custom_append=None):
    """Parse a data file from the Movie Lens data set.

       This function parses a file from the Movie Lens data set.
       The function returns a generator for JSON dict strings.

       Optionally, a callback for extending the JSON dict string can be passed.
       The callback function will be called for each line parsed. Fields parsed
       will be provided via an array passed to the callback.
       The callback signature thusly is:

        custom_append(fields)

       Arguments:
       fname -- file name of the data file.
       field_types -- array of field identifiers. Must match the number of
                       fields per line.
                       NOTE: The "Genres" field entries will be put in an
                       array in the return dict.

       Example Usage:
        for line in parse('ratings.dat',
                            ("UserID", "MovieID", "Rating", "Timestamp"):
            print line

       Example output:
        {"UserID":"1","MovieID":"122","Rating":"5","Timestamp":"838985046"}
        {"UserID":"1","MovieID":"362","Rating":"5","Timestamp":"838984885"}
        {"UserID":"1","MovieID":"364","Rating":"5","Timestamp":"838983707"}
        ...

       Exaample for "movies.dat":
        for line in parse('movies.dat', ("MovieID", "Title", "Genres")):
            print line

       Output:
        {"MovieID":"65130","Title":"Revolutionary Road (2008)","Genres":["Drama","Romance"]}
        ...
    """
    if not custom_append:
        custom_append = lambda x: ""
    sz = os.stat(fname).st_size
    rd = 0
    with open(fname) as f:
        num_fields = len(field_types)
        for line in f: 
            line = line.strip()
            fields = line.split("::")
            ret = "{"
            for i in range(num_fields):
                val = fields[i]
                if field_types[i] == "Genres":
                    g   = fields[i].split("|")
                    ret = ret + '"Genres":["%s"],' % '","'.join(g)
                else:
                    ret = ret + '"%s":%s,' % (field_types[i], json.dumps(val))
            rd = rd + len(line)
            ret += custom_append(fields)
            yield ret[:-1] + '}', rd, sz


def index_writer(con, q, index, doctype):
    """Reads data tuple from queue
        (start-document-num, end-document-num, documents-buf,
            bytes-read, bytes-total),
       writes documents to elasticsearch, and prints progress. Function is
       intended to run in a separate thread.

       Arguments:<commit_msg>Fix timestamp in ratings index.<commit_after>#!/usr/bin/env python
# vim: set fileencoding=utf-8 :
#
# MovieLens indexing tool for generating Elasticsearch indices.
#
# Re-write of an original idea/script by <NAME>.
# Copyright 2014 © <NAME>
#
# This file is licensed to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#  http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

from pprint import pprint
from pyes import *

import argparse
import string
import sys
import os
import json

from threading import Thread
from Queue import Queue


def parse(fname, field_types, custom_append=None):
    """Parse a data file from the Movie Lens data set.

       This function parses a file from the Movie Lens data set.
       The function returns a generator for JSON dict strings.

       Optionally, a callback for extending the JSON dict string can be passed.
       The callback function will be called for each line parsed. Fields parsed
       will be provided via an array passed to the callback.
       The callback signature thusly is:

        custom_append(fields)

       Arguments:
       fname -- file name of the data file.
       field_types -- array of field identifiers. Must match the number of
                       fields per line.
                       NOTE: The "Genres" field entries will be put in an
                       array in the return dict.

       Example Usage:
        for line in parse('ratings.dat',
                            ("UserID", "MovieID", "Rating", "Timestamp"):
            print line

       Example output:
        {"UserID":"1","MovieID":"122","Rating":"5","Timestamp":"838985046"}
        {"UserID":"1","MovieID":"362","Rating":"5","Timestamp":"838984885"}
        {"UserID":"1","MovieID":"364","Rating":"5","Timestamp":"838983707"}
        ...

       Exaample for "movies.dat":
        for line in parse('movies.dat', ("MovieID", "Title", "Genres")):
            print line

       Output:
        {"MovieID":"65130","Title":"Revolutionary Road (2008)","Genres":["Drama","Romance"]}
        ...
    """
    if not custom_append:
        custom_append = lambda x: ""
    sz = os.stat(fname).st_size
    rd = 0
    with open(fname) as f:
        num_fields = len(field_types)
        for line in f: 
            line = line.strip()
            fields = line.split("::")
            ret = "{"
            for i in range(num_fields):
                val = fields[i]
                if field_types[i] == "Genres":
                    g   = fields[i].split("|")
                    ret = ret + '"Genres":["%s"],' % '","'.join(g)
                elif field_types[i] == "Timestamp":
                    ret = ret + '"%s":%s000,' % (field_types[i], val)
                else:
                    ret = ret + '"%s":%s,' % (field_types[i], json.dumps(val))
            rd = rd + len(line)
            ret += custom_append(fields)
            yield ret[:-1] + '}', rd, sz


def index_writer(con, q, index, doctype):
    """Reads data tuple from queue
        (start-document-num, end-document-num, documents-buf,
            bytes-read, bytes-total),
       writes documents to elasticsearch, and prints progress. Function is
       intended to run in a separate thread.

       Arguments:
{'fraction_non_alphanumeric': 0.112409571508069, 'fraction_numerical': 0.01766833611574847, 'mean_word_length': 2.3068077276908925, 'pattern_counts': {'":': 42, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1301482', 'n_tokens_mistral': 2088, 'n_tokens_neox': 1937, 'n_words': 733}