<commit_before># Licensed 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.

import abc
import logging
import luigi
import luigi.hadoop
import luigi.interface
import os
import subprocess
import tempfile

logger = logging.getLogger('luigi-interface')


def load_hive_cmd():
    return luigi.interface.get_config().get('hive', 'command', 'hive')


def run_hive(args):
    """runs the `hive` from the command line, passing in the given args, and
       returning stdout"""
    cmd = [load_hive_cmd()] + args
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout, _ = p.communicate()
    return stdout


def run_hive_cmd(hivecmd):
    """Runs the given hive query and returns stdout"""
    return run_hive(['-e', hivecmd])


def run_hive_script(script):
    """Runs the contents of the given script in hive and returns stdout"""
    if not os.path.isfile(script):
        raise RuntimeError("Hive script: {0} does not exist.".format(script))
    return run_hive(['-f', script])


def table_location(db, table, partition_spec=None):
    """
    Returns the location of a Hive table/partitions, or None.
    partition_spec should be a string returned by partition_spec()
    """
    cmd = "use {0}; describe formatted {1}".format(db, table)
    if partition_spec:
        cmd += " PARTITION ({0})".format(partition_spec)

    stdout = run_hive_cmd(cmd)

    for line in stdout.split("\n"):
        if "Location:" in line:
            return line.split("\t")[1]


def partition_spec(partition):
    """
    Turns a dict into the a Hive partition specification string
    """
    return ','.join(["{0}='{1}'".format(k, v) for (k, v) in partition.items()])


class HiveQueryTask(luigi.hadoop.BaseHadoopJobTask):
    ''' Task to run a hive query
    '''

    @abc.abstractmethod
    def query(self):
        ''' Text of query to run in hive
        '''
        raise RuntimeError("Must implement query!")

    def hiverc(self):
        ''' Location of an rc file to run before the query
        '''
        return None

    def job_runner(self):
        return HiveQueryRunner()


class HiveQueryRunner(luigi.hadoop.JobRunner):
    ''' Runs a HiveQueryTask by shelling out to hive
    '''

    def run_job(self, job):
        with tempfile.NamedTemporaryFile() as f:
            f.write(job.query())
            f.flush()
            arglist = [load_hive_cmd(), '-f', f.name]
            if job.hiverc():
                arglist += ['-i', job.hiverc()]

            logger.info(arglist)
            luigi.hadoop.run_and_track_hadoop_job(arglist)


class HiveTableTarget(luigi.Target):
    ''' exists returns true if the table exists
    '''

    def __init__(self, table, database='default'):
        self.database = database
        self.table = table
        self.hive_cmd = load_hive_cmd()

    def exists(self):
        stdout = run_hive_cmd('use {0}; describe {1}'.format(self.database, self.table))

        return not "does not exist" in stdout

    @property
    def path(self):
        """Returns the path to this table in HDFS"""
        location = table_location(self.database, self.table)
        if not location:
            raise Exception("Couldn't find location for table: {0}".format(str(self)))
        return location

    def open(self, mode):
        return NotImplementedError("open() is not supported for HiveTableTarget")

<commit_msg>Add support for common hive params.<commit_after># Licensed 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.

import abc
import logging
import luigi
import luigi.hadoop
import luigi.interface
import os
import subprocess
import tempfile

logger = logging.getLogger('luigi-interface')


def load_hive_cmd():
    return luigi.interface.get_config().get('hive', 'command', 'hive')


def run_hive(args):
    """runs the `hive` from the command line, passing in the given args, and
       returning stdout"""
    cmd = [load_hive_cmd()] + args
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout, _ = p.communicate()
    return stdout


def run_hive_cmd(hivecmd):
    """Runs the given hive query and returns stdout"""
    return run_hive(['-e', hivecmd])


def run_hive_script(script):
    """Runs the contents of the given script in hive and returns stdout"""
    if not os.path.isfile(script):
        raise RuntimeError("Hive script: {0} does not exist.".format(script))
    return run_hive(['-f', script])


def table_location(db, table, partition_spec=None):
    """
    Returns the location of a Hive table/partitions, or None.
    partition_spec should be a string returned by partition_spec()
    """
    cmd = "use {0}; describe formatted {1}".format(db, table)
    if partition_spec:
        cmd += " PARTITION ({0})".format(partition_spec)

    stdout = run_hive_cmd(cmd)

    for line in stdout.split("\n"):
        if "Location:" in line:
            return line.split("\t")[1]


def partition_spec(partition):
    """
    Turns a dict into the a Hive partition specification string
    """
    return ','.join(["{0}='{1}'".format(k, v) for (k, v) in partition.items()])


class HiveQueryTask(luigi.hadoop.BaseHadoopJobTask):
    ''' Task to run a hive query
    '''
    # by default, we let hive figure these out.
    n_reduce_tasks = None
    bytes_per_reducer = None
    reducers_max = None

    @abc.abstractmethod
    def query(self):
        ''' Text of query to run in hive
        '''
        raise RuntimeError("Must implement query!")

    def hiverc(self):
        ''' Location of an rc file to run before the query
        '''
        return None

    def hiveconfs(self):
        '''
        Returns an dict of key=value settings to be passed along
        to the hive command line via --hiveconf. By default, sets
        mapred.job.name to task_id and if not None, sets:
        * mapred.reduce.tasks (n_reduce_tasks)
        * mapred.fairscheduler.pool (pool)
        * hive.exec.reducers.bytes.per.reducer (bytes_per_reducer)
        * hive.exec.reducers.max (reducers_max)
        '''
        jcs = {}
        jcs['mapred.job.name'] = self.task_id
        if self.n_reduce_tasks is not None:
            jcs['mapred.reduce.tasks'] = self.n_reduce_tasks
        if self.pool is not None:
            jcs['mapred.fairscheduler.pool'] = self.pool
        if self.bytes_per_reducer is not None:
            jcs['hive.exec.reducers.bytes.per.reducer'] = self.bytes_per_reducer
        if self.reducers_max is not None:
            jcs['hive.exec.reducers.max'] = self.reducers_max
        return jcs

    def job_runner(self):
        return HiveQueryRunner()


class HiveQueryRunner(luigi.hadoop.JobRunner):
    ''' Runs a HiveQueryTask by shelling out to hive
    '''

    def run_job(self, job):
        with tempfile.NamedTemporaryFile() as f:
            f.write(job.query())
            f.flush()
            arglist = [load_hive_cmd(), '-f', f.name]
            if job.hiverc():
                arglist += ['-i', job.hiverc()]
            if job.hiveconfs():
                for k, v in job.hiveconfs().iteritems():
                    arglist += ['--hiveconf', '{0}={1}'.format(k, v)]

            logger.info(arglist)
            luigi.hadoop.run_and_track_hadoop_job(arglist)


class HiveTableTarget(luigi.Target):
    ''' exists returns true if the table exists
    '''

    def __init__(self, table, database='default'):
        self.database = database
        self.table = table
        self.hive_cmd = load_hive_cmd()

    def exists(self):
        stdout = run_hive_cmd('use {0}; describe {1}'.format(self.database, self.table))

        return not "does not exist" in stdout

    @property
    def path(self):
        """Returns the path to this table in HDFS"""
        location = table_location(self.database, self.table)
        if not location:
            raise Exception("Couldn't find location for table: {0}".format(str(self)))
        return location

    def open(self, mode):
        return NotImplementedError("open() is not supported for HiveTableTarget")


{'fraction_non_alphanumeric': 0.11237205162438807, 'fraction_numerical': 0.0033377837116154874, 'mean_word_length': 2.8529789969995716, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '6834897', 'n_tokens_mistral': 2755, 'n_tokens_neox': 2691, 'n_words': 853}