| #! /usr/bin/env python |
| import sys |
| import os |
| import logging |
| |
| from datetime import datetime |
| try: |
| from dateutil import parser |
| USE_DATEUTIL = True |
| except ImportError: |
| USE_DATEUTIL = False |
| |
| from pig_util import write_user_exception, udf_logging |
| |
| FIELD_DELIMITER = b',' |
| TUPLE_START = b'(' |
| TUPLE_END = b')' |
| BAG_START = b'{' |
| BAG_END = b'}' |
| MAP_START = b'[' |
| MAP_END = b']' |
| MAP_KEY = b'#' |
| PARAMETER_DELIMITER = b'\t' |
| PRE_WRAP_DELIM = b'|' |
| POST_WRAP_DELIM = b'_' |
| NULL_BYTE = b"-" |
| END_RECORD_DELIM = b'|_\n' |
| END_RECORD_DELIM_LENGTH = len(END_RECORD_DELIM) |
| |
| WRAPPED_FIELD_DELIMITER = PRE_WRAP_DELIM + FIELD_DELIMITER + POST_WRAP_DELIM |
| WRAPPED_TUPLE_START = PRE_WRAP_DELIM + TUPLE_START + POST_WRAP_DELIM |
| WRAPPED_TUPLE_END = PRE_WRAP_DELIM + TUPLE_END + POST_WRAP_DELIM |
| WRAPPED_BAG_START = PRE_WRAP_DELIM + BAG_START + POST_WRAP_DELIM |
| WRAPPED_BAG_END = PRE_WRAP_DELIM + BAG_END + POST_WRAP_DELIM |
| WRAPPED_MAP_START = PRE_WRAP_DELIM + MAP_START + POST_WRAP_DELIM |
| WRAPPED_MAP_END = PRE_WRAP_DELIM + MAP_END + POST_WRAP_DELIM |
| WRAPPED_PARAMETER_DELIMITER = PRE_WRAP_DELIM + PARAMETER_DELIMITER + POST_WRAP_DELIM |
| WRAPPED_NULL_BYTE = PRE_WRAP_DELIM + NULL_BYTE + POST_WRAP_DELIM |
| |
| TYPE_TUPLE = TUPLE_START |
| TYPE_BAG = BAG_START |
| TYPE_MAP = MAP_START |
| |
| TYPE_BOOLEAN = b"B" |
| TYPE_INTEGER = b"I" |
| TYPE_LONG = b"L" |
| TYPE_FLOAT = b"F" |
| TYPE_DOUBLE = b"D" |
| TYPE_BYTEARRAY = b"A" |
| TYPE_CHARARRAY = b"C" |
| TYPE_DATETIME = b"T" |
| TYPE_BIGINTEGER = b"N" |
| TYPE_BIGDECIMAL = b"E" |
| END_OF_STREAM = TYPE_CHARARRAY + b"\x04" + END_RECORD_DELIM |
| TURN_ON_OUTPUT_CAPTURING = TYPE_CHARARRAY + b"TURN_ON_OUTPUT_CAPTURING" + END_RECORD_DELIM |
| |
| NUM_LINES_OFFSET_TRACE = int(os.environ.get('PYTHON_TRACE_OFFSET', 0)) |
| |
| # Deal with Python 2.x to 3.x incompatibilities |
| if sys.version_info[0] >= 3: |
| PYTHON3 = True |
| long = int |
| #https://bugs.python.org/issue15610 |
| IMPORT_LEVEL = 0 |
| # Python 3 handles str as unicode by default. unicode type is dropped |
| def tounicode(input, encoding="utf-8"): |
| if type(input) == str: |
| return input |
| else: |
| return str(input) |
| else: |
| PYTHON3 = False |
| IMPORT_LEVEL = -1 |
| def tounicode(input, encoding=None): |
| return unicode(input) if encoding is None else unicode(input, encoding) |
| |
| class PythonStreamingController: |
| def __init__(self, profiling_mode=False): |
| self.profiling_mode = profiling_mode |
| |
| self.input_count = 0 |
| self.next_input_count_to_log = 1 |
| |
| def main(self, |
| module_name, file_path, func_name, cache_path, |
| output_stream_path, error_stream_path, log_file_name, is_illustrate_str): |
| sys.stdin = os.fdopen(sys.stdin.fileno(), 'rb', 0) |
| |
| #Need to ensure that user functions can't write to the streams we use to |
| #communicate with pig. |
| self.stream_output = os.fdopen(sys.stdout.fileno(), 'wb', 0) |
| self.stream_error = os.fdopen(sys.stderr.fileno(), 'wb', 0) |
| |
| self.input_stream = sys.stdin |
| self.output_stream = open(output_stream_path, 'a') |
| sys.stderr = open(error_stream_path, 'w') |
| is_illustrate = is_illustrate_str == "true" |
| |
| sys.path.append(file_path) |
| sys.path.append(cache_path) |
| sys.path.append('.') |
| |
| logging.basicConfig(filename=log_file_name, format="%(asctime)s %(levelname)s %(message)s", level=udf_logging.udf_log_level) |
| logging.info("To reduce the amount of information being logged only a small subset of rows are logged at the INFO level. Call udf_logging.set_log_level_debug in pig_util to see all rows being processed.") |
| |
| input_bytes = self.get_next_input() |
| |
| try: |
| func = __import__(module_name, globals(), locals(), [func_name], IMPORT_LEVEL).__dict__[func_name] |
| except: |
| #These errors should always be caused by user code. |
| write_user_exception(module_name, self.stream_error, NUM_LINES_OFFSET_TRACE) |
| self.close_controller(-1) |
| |
| if is_illustrate or udf_logging.udf_log_level != logging.DEBUG: |
| #Only log output for illustrate after we get the flag to capture output. |
| sys.stdout = open(os.devnull, 'w') |
| else: |
| sys.stdout = self.output_stream |
| |
| while input_bytes != END_OF_STREAM: |
| should_log = False |
| if self.input_count == self.next_input_count_to_log: |
| should_log = True |
| log_message = logging.info |
| self.update_next_input_count_to_log() |
| elif udf_logging.udf_log_level == logging.DEBUG: |
| should_log = True |
| log_message = logging.debug |
| |
| try: |
| try: |
| if should_log: |
| log_message("Row %s: Serialized Input: %s" % (self.input_count, input_bytes)) |
| inputs = deserialize_input(input_bytes) |
| if should_log: |
| log_message("Row %s: Deserialized Input: %s" % (self.input_count, tounicode(inputs))) |
| except: |
| #Capture errors where the user passes in bad data. |
| write_user_exception(module_name, self.stream_error, NUM_LINES_OFFSET_TRACE) |
| self.close_controller(-3) |
| |
| try: |
| func_output = func(*inputs) |
| if should_log: |
| try: |
| log_message("Row %s: UDF Output: %s" % (self.input_count, tounicode(func_output))) |
| except: |
| #This is probably an error with unicoding the output. Calling unicode on bytearray will |
| #throw an exception. Since its just a log statement, just skip and carry on. |
| logging.exception("Couldn't log output. Try to continue.") |
| except: |
| #These errors should always be caused by user code. |
| write_user_exception(module_name, self.stream_error, NUM_LINES_OFFSET_TRACE) |
| self.close_controller(-2) |
| |
| output = serialize_output(func_output) |
| if should_log: |
| log_message("Row %s: Serialized Output: %s" % (self.input_count, output)) |
| |
| self.stream_output.write( b"%s%s" % (output, END_RECORD_DELIM) ) |
| except Exception as e: |
| #This should only catch internal exceptions with the controller |
| #and pig- not with user code. |
| import traceback |
| traceback.print_exc(file=self.stream_error) |
| sys.exit(-3) |
| |
| sys.stdout.flush() |
| sys.stderr.flush() |
| self.stream_output.flush() |
| self.stream_error.flush() |
| |
| input_bytes = self.get_next_input() |
| |
| def get_next_input(self): |
| input_stream = self.input_stream |
| output_stream = self.output_stream |
| |
| input_bytes = input_stream.readline() |
| |
| while input_bytes.endswith(END_RECORD_DELIM) == False: |
| line = input_stream.readline() |
| if line == '': |
| input_bytes = '' |
| break |
| input_bytes += line |
| |
| if input_bytes == '': |
| return END_OF_STREAM |
| |
| if input_bytes == TURN_ON_OUTPUT_CAPTURING: |
| logging.debug("Turned on Output Capturing") |
| sys.stdout = output_stream |
| return self.get_next_input() |
| |
| if input_bytes == END_OF_STREAM: |
| return input_bytes |
| |
| self.input_count += 1 |
| |
| return input_bytes[:-END_RECORD_DELIM_LENGTH] |
| |
| def update_next_input_count_to_log(self): |
| """ |
| Want to log enough rows that you can see progress being made and see timings without wasting time logging thousands of rows. |
| Show first 10 rows, and then the first 5 rows of every order of magnitude (10-15, 100-105, 1000-1005, ...) |
| """ |
| if self.next_input_count_to_log < 10: |
| self.next_input_count_to_log = self.next_input_count_to_log + 1 |
| elif self.next_input_count_to_log % 10 == 5: |
| self.next_input_count_to_log = (self.next_input_count_to_log - 5) * 10 |
| else: |
| self.next_input_count_to_log = self.next_input_count_to_log + 1 |
| |
| def close_controller(self, exit_code): |
| sys.stderr.close() |
| self.stream_error.write(b"\n") |
| self.stream_error.close() |
| sys.stdout.close() |
| self.stream_output.write(b"\n") |
| self.stream_output.close() |
| sys.exit(exit_code) |
| |
| def deserialize_input(input_bytes): |
| if len(input_bytes) == 0: |
| return [] |
| |
| return [_deserialize_input(param, 0, len(param)-1) for param in input_bytes.split(WRAPPED_PARAMETER_DELIMITER)] |
| |
| def _deserialize_input(input_bytes, si, ei): |
| if ei - si < 1: |
| #Handle all of the cases where you can have valid empty input. |
| if ei == si: |
| if input_bytes[si:si+1] == TYPE_CHARARRAY: |
| if PYTHON3 == True: |
| return "" |
| else: |
| return u"" |
| elif input_bytes[si:si+1] == TYPE_BYTEARRAY: |
| return bytearray(b"") |
| else: |
| raise Exception("Got input type flag %s, but no data to go with it.\nInput string: %s\nSlice: %s" % (input_bytes[si], input_bytes, input_bytes[si:ei+1])) |
| else: |
| raise Exception("Start index %d greater than end index %d.\nInput string: %s\n, Slice: %s" % (si, ei, input_bytes[si:ei+1])) |
| |
| first = input_bytes[si:si+1] |
| schema = input_bytes[si+1:si+2] if first == PRE_WRAP_DELIM else first |
| |
| # Pig to streaming input is serialized in PigStreaming.serializeToBytes() via StorageUtil.putField() |
| if schema == NULL_BYTE: |
| return None |
| elif schema == TYPE_TUPLE or schema == TYPE_MAP or schema == TYPE_BAG: |
| return _deserialize_collection(input_bytes, schema, si+3, ei-3) |
| elif schema == TYPE_CHARARRAY: |
| if PYTHON3 == True: |
| return input_bytes[si+1:ei+1].decode("utf-8") |
| else: |
| return tounicode(input_bytes[si+1:ei+1], 'utf-8') |
| elif schema == TYPE_BYTEARRAY: |
| return bytearray(input_bytes[si+1:ei+1]) |
| elif schema == TYPE_INTEGER: |
| return int(input_bytes[si+1:ei+1]) |
| elif schema == TYPE_LONG or schema == TYPE_BIGINTEGER: |
| return long(input_bytes[si + 1:ei + 1]) |
| elif schema == TYPE_FLOAT or schema == TYPE_DOUBLE or schema == TYPE_BIGDECIMAL: |
| return float(input_bytes[si+1:ei+1]) |
| elif schema == TYPE_BOOLEAN: |
| return input_bytes[si+1:ei+1] == b"true" |
| elif schema == TYPE_DATETIME: |
| #Format is "yyyy-MM-ddTHH:mm:ss.SSS+00:00" or "2013-08-23T18:14:03.123+ZZ" |
| if USE_DATEUTIL: |
| if PYTHON3 == True: |
| return parser.parse(input_bytes[si+1:ei+1].decode('utf-8')) |
| else: |
| return parser.parse(input_bytes[si+1:ei+1]) |
| else: |
| #Try to use datetime even though it doesn't handle time zones properly, |
| #We only use the first 3 microsecond digits and drop time zone (first 23 characters) |
| if PYTHON3 == True: |
| return datetime.strptime(input_bytes[si+1:si+24].decode('utf-8'), "%Y-%m-%dT%H:%M:%S.%f") |
| else: |
| return datetime.strptime(input_bytes[si+1:si+24], "%Y-%m-%dT%H:%M:%S.%f") |
| else: |
| raise Exception("Can't determine type of input: %s" % input_bytes[si:ei+1]) |
| |
| def _deserialize_collection(input_bytes, return_type, si, ei): |
| list_result = [] |
| append_to_list_result = list_result.append |
| dict_result = {} |
| |
| index = si |
| field_start = si |
| depth = 0 |
| |
| key = None |
| |
| # recurse to deserialize elements if the collection is not empty |
| if ei-si+1 > 0: |
| while True: |
| if index >= ei - 2: |
| if return_type == TYPE_MAP: |
| dict_result[key] = _deserialize_input(input_bytes, value_start, ei) |
| else: |
| append_to_list_result(_deserialize_input(input_bytes, field_start, ei)) |
| break |
| |
| if return_type == TYPE_MAP and not key: |
| key_index = input_bytes.find(MAP_KEY, index) |
| if PYTHON3 == True: |
| key = input_bytes[index+1:key_index].decode("utf-8") |
| else: |
| key = tounicode(input_bytes[index+1:key_index], 'utf-8') |
| index = key_index + 1 |
| value_start = key_index + 1 |
| continue |
| |
| if not (input_bytes[index:index+1] == PRE_WRAP_DELIM and input_bytes[index+2:index+3] == POST_WRAP_DELIM): |
| prewrap_index = input_bytes.find(PRE_WRAP_DELIM, index+1) |
| index = (prewrap_index if prewrap_index != -1 else end_index) |
| continue |
| |
| mid = input_bytes[index+1:index+2] |
| |
| if mid == BAG_START or mid == TUPLE_START or mid == MAP_START: |
| depth += 1 |
| elif mid == BAG_END or mid == TUPLE_END or mid == MAP_END: |
| depth -= 1 |
| elif depth == 0 and mid == FIELD_DELIMITER: |
| if return_type == TYPE_MAP: |
| dict_result[key] = _deserialize_input(input_bytes, value_start, index - 1) |
| key = None |
| else: |
| append_to_list_result(_deserialize_input(input_bytes, field_start, index - 1)) |
| field_start = index + 3 |
| |
| index += 3 |
| |
| if return_type == TYPE_MAP: |
| return dict_result |
| elif return_type == TYPE_TUPLE: |
| return tuple(list_result) |
| else: |
| return list_result |
| |
| def wrap_tuple(o, serialized_item): |
| if type(o) != tuple: |
| return WRAPPED_TUPLE_START + serialized_item + WRAPPED_TUPLE_END |
| else: |
| return serialized_item |
| |
| def encode_map_key(key): |
| if PYTHON3 == True and type(key) == bytes: |
| return bytes |
| else: |
| return key.encode('utf-8') |
| |
| def serialize_output(output, utfEncodeAllFields=False): |
| """ |
| @param utfEncodeStrings - Generally we want to utf encode only strings. But for |
| Maps we utf encode everything because on the Java side we don't know the schema |
| for maps so we wouldn't be able to tell which fields were encoded or not. |
| """ |
| |
| output_type = type(output) |
| |
| if output is None: |
| return WRAPPED_NULL_BYTE |
| elif output_type == tuple: |
| return (WRAPPED_TUPLE_START + |
| WRAPPED_FIELD_DELIMITER.join([serialize_output(o, utfEncodeAllFields) for o in output]) + |
| WRAPPED_TUPLE_END) |
| elif output_type == list: |
| return (WRAPPED_BAG_START + |
| WRAPPED_FIELD_DELIMITER.join([wrap_tuple(o, serialize_output(o, utfEncodeAllFields)) for o in output]) + |
| WRAPPED_BAG_END) |
| elif output_type == dict: |
| if PYTHON3 == True: |
| items = output.items() |
| else: |
| items = output.iteritems() |
| return (WRAPPED_MAP_START + |
| WRAPPED_FIELD_DELIMITER.join([b'%s%s%s' % (encode_map_key(k), MAP_KEY, serialize_output(v, True)) for k, v in items]) |
| + WRAPPED_MAP_END) |
| elif output_type == bool: |
| return (b"true" if output else b"false") |
| elif output_type == bytearray or (PYTHON3 == True and output_type == bytes): |
| if PYTHON3 == True: |
| return bytes(output) if output_type == bytearray else output |
| else: |
| return str(output) |
| elif output_type == datetime: |
| if PYTHON3 == True: |
| return bytes(output.isoformat(), 'utf-8') |
| else: |
| return output.isoformat() |
| elif PYTHON3 == True and output_type == str: |
| return output.encode('utf-8') |
| elif PYTHON3 == False and (utfEncodeAllFields or output_type == str or output_type == unicode): |
| #unicode is necessary in cases where we're encoding non-strings. |
| return tounicode(output).encode('utf-8') |
| else: |
| if PYTHON3 == True: |
| return str(output).encode("utf-8") |
| else: |
| return str(output) |
| |
| if __name__ == '__main__': |
| controller = PythonStreamingController() |
| controller.main(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4], |
| sys.argv[5], sys.argv[6], sys.argv[7], sys.argv[8]) |
