Clean up

Author: deliahu

pkg/consts/consts.go

@@ -46,7 +46,6 @@ var (
 	AggregatesDir         = "aggregates"
 	TransformersDir       = "transformers"
 	EstimatorsDir         = "estimators"
-	ModelImplsDir         = "model_implementations"
 	PythonPackagesDir     = "python_packages"
 	ModelsDir             = "models"
 	ConstantsDir          = "constants"

pkg/workloads/lib/context.py

@@ -415,14 +415,15 @@ def get_inferred_column_type(self, column_name):
 
         return column_type
 
-    # replaces column references with column names (unless preserve_column_refs = true, then leaves them untouched)
+    # Replace aggregates and constants with their values, and columns with their names (unless preserve_column_refs == False)
+    # Also validate against input_schema (if not None)
     def populate_values(self, input, input_schema, preserve_column_refs):
         if input is None:
             if input_schema is None:
                 return None
-            if input_schema["_allow_null"]:
+            if input_schema.get("_allow_null") == True:
                 return None
-            raise UserException("Null is not allowed")
+            raise UserException("Null value is not allowed")
 
         if util.is_resource_ref(input):
             res_name = util.get_resource_ref(input)
@@ -447,8 +448,10 @@ def populate_values(self, input, input_schema, preserve_column_refs):
                     col_type = self.get_inferred_column_type(res_name)
                     if col_type not in input_schema["_type"]:
                         raise UserException(
-                            "column {}: column type mismatch: got {}, expected {}".format(
-                                res_name, col_type, input_schema["_type"]
+                            "column {}: unsupported input type (expected type {}, got type {})".format(
+                                res_name,
+                                util.data_type_str(input_schema["_type"]),
+                                util.data_type_str(col_type),
                             )
                         )
                 if preserve_column_refs:
@@ -460,21 +463,25 @@ def populate_values(self, input, input_schema, preserve_column_refs):
             elem_schema = None
             if input_schema is not None:
                 if not util.is_list(input_schema["_type"]):
-                    raise UserException("unexpected type (list)")
+                    raise UserException(
+                        "unsupported input type (expected type {}, got {})".format(
+                            util.data_type_str(input_schema["_type"]), util.pp_str_flat(input)
+                        )
+                    )
                 elem_schema = input_schema["_type"][0]
 
                 min_count = input_schema.get("_min_count")
                 if min_count is not None and len(input) < min_count:
                     raise UserException(
-                        "list has length {}, but the minimum length is {}".format(
+                        "list has length {}, but the minimum allowed length is {}".format(
                             len(input), min_count
                         )
                     )
 
                 max_count = input_schema.get("_max_count")
                 if max_count is not None and len(input) > max_count:
                     raise UserException(
-                        "list has length {}, but the maximum length is {}".format(
+                        "list has length {}, but the maximum allowed length is {}".format(
                             len(input), max_count
                         )
                     )
@@ -496,24 +503,32 @@ def populate_values(self, input, input_schema, preserve_column_refs):
                     try:
                         val_casted = self.populate_values(val, None, preserve_column_refs)
                     except CortexException as e:
-                        e.wrap(util.pp_str_flat(key_casted))
+                        e.wrap(util.pp_str_flat(key))
                         raise
                     casted[key_casted] = val_casted
                 return casted
 
             if not util.is_dict(input_schema["_type"]):
-                raise UserException("unexpected type (map)")
+                raise UserException(
+                    "unsupported input type (expected type {}, got {})".format(
+                        util.data_type_str(input_schema["_type"]), util.pp_str_flat(input)
+                    )
+                )
 
             min_count = input_schema.get("_min_count")
             if min_count is not None and len(input) < min_count:
                 raise UserException(
-                    "map has length {}, but the minimum length is {}".format(len(input), min_count)
+                    "map has length {}, but the minimum allowed length is {}".format(
+                        len(input), min_count
+                    )
                 )
 
             max_count = input_schema.get("_max_count")
             if max_count is not None and len(input) > max_count:
                 raise UserException(
-                    "map has length {}, but the maximum length is {}".format(len(input), max_count)
+                    "map has length {}, but the maximum allowed length is {}".format(
+                        len(input), max_count
+                    )
                 )
 
             is_generic_map = False
@@ -535,23 +550,23 @@ def populate_values(self, input, input_schema, preserve_column_refs):
                             val, generic_map_value, preserve_column_refs
                         )
                     except CortexException as e:
-                        e.wrap(util.pp_str_flat(key_casted))
+                        e.wrap(util.pp_str_flat(key))
                         raise
                     casted[key_casted] = val_casted
                 return casted
 
             # fixed map
             casted = {}
             for key, val_schema in input_schema["_type"].items():
-                default = None
-                if key not in input:
+                if key in input:
+                    val = input[key]
+                else:
                     if val_schema.get("_optional") is not True:
                         raise UserException("missing key: " + util.pp_str_flat(key))
                     if val_schema.get("_default") is None:
                         continue
-                    default = val_schema["_default"]
+                    val = val_schema["_default"]
 
-                val = input.get(key, default)
                 try:
                     val_casted = self.populate_values(val, val_schema, preserve_column_refs)
                 except CortexException as e:
@@ -562,8 +577,12 @@ def populate_values(self, input, input_schema, preserve_column_refs):
 
         if input_schema is None:
             return input
-        if util.is_list(input_schema["_type"]) or util.is_dict(input_schema["_type"]):
-            raise UserException("unexpected type (scalar)")
+        if not util.is_str(input_schema["_type"]):
+            raise UserException(
+                "unsupported input type (expected type {}, got {})".format(
+                    util.data_type_str(input_schema["_type"]), util.pp_str_flat(input)
+                )
+            )
         return cast_compound_type(input, input_schema["_type"])
 
 
@@ -605,8 +624,8 @@ def cast_compound_type(value, type_str):
             return value
 
     raise UserException(
-        "input value's type is not supported by the schema (got {}, expected input with type {})".format(
-            util.pp_str_flat(value), type_str
+        "unsupported input type (expected type {}, got {})".format(
+            util.data_type_str(type_str), util.pp_str_flat(value)
         )
     )
 
@@ -689,7 +708,7 @@ def _deserialize_raw_ctx(raw_ctx):
 def create_transformer_inputs_from_map(input, col_value_map):
     if util.is_str(input):
         res_name = util.get_resource_ref(input)
-        if res_name in col_value_map:
+        if res_name is not None and res_name in col_value_map:
             return col_value_map[res_name]
         return input
 

pkg/workloads/lib/util.py

@@ -66,6 +66,11 @@ def pp_str_flat(obj, indent=0):
     return indent_str(out, indent)
 
 
+def data_type_str(obj):
+    # TODO. Also call this method with output types?
+    return pp_str_flat(obj)
+
+
 def log_indent(obj, indent=0, logging_func=logger.info):
     if not is_str(obj):
         text = repr(obj)
@@ -748,8 +753,6 @@ def validate_output_type(value, output_type):
         return False
 
     if is_list(output_type):
-        if not (len(output_type) == 1 and is_str(output_type[0])):
-            return False
         if not is_list(value):
             return False
         for value_item in value:
@@ -760,8 +763,6 @@ def validate_output_type(value, output_type):
     if is_dict(output_type):
         if not is_dict(value):
             return False
-        if len(output_type) == 0:
-            return False
 
         is_generic_map = False
         if len(output_type) == 1:
@@ -787,10 +788,10 @@ def validate_output_type(value, output_type):
                 return False
         return True
 
-    return False
+    return False  # unexpected
 
 
-# Casts int -> float. Input is assumed to be already validated
+# value is assumed to be already validated against output_type
 def cast_output_type(value, output_type):
     if is_str(output_type):
         if (
@@ -858,17 +859,17 @@ def extract_resource_refs(input):
             return {res}
         return set()
 
+    if is_list(input):
+        resources = set()
+        for item in input:
+            resources = resources.union(extract_resource_refs(item))
+        return resources
+
     if is_dict(input):
         resources = set()
         for key, val in input.items():
             resources = resources.union(extract_resource_refs(key))
             resources = resources.union(extract_resource_refs(val))
         return resources
 
-    if is_list(input):
-        resources = set()
-        for item in input:
-            resources = resources.union(extract_resource_refs(item))
-        return resources
-
     return set()

pkg/workloads/spark_job/spark_util.py

@@ -397,12 +397,12 @@ def read_parquet(ctx, spark):
 
 
 def split_aggregators(aggregate_names, ctx):
-    aggregate_resources = [ctx.aggregates[agg_name] for agg_name in aggregate_names]
+    aggregates = [ctx.aggregates[agg_name] for agg_name in aggregate_names]
 
     builtin_aggregates = []
     custom_aggregates = []
 
-    for agg in aggregate_resources:
+    for agg in aggregates:
         aggregator = ctx.aggregators[agg["aggregator"]]
         if aggregator.get("namespace", None) == "cortex" and aggregator["name"] in AGG_SPARK_LIST:
             builtin_aggregates.append(agg)
@@ -416,52 +416,50 @@ def run_builtin_aggregators(builtin_aggregates, df, ctx, spark):
     agg_cols = []
     for agg in builtin_aggregates:
         aggregator = ctx.aggregators[agg["aggregator"]]
-        input_repl = ctx.populate_values(
-            agg["input"], aggregator["input"], preserve_column_refs=False
-        )
+        input = ctx.populate_values(agg["input"], aggregator["input"], preserve_column_refs=False)
 
         if aggregator["name"] == "approx_count_distinct":
             agg_cols.append(
-                F.approxCountDistinct(input_repl["col"], input_repl.get("rsd")).alias(agg["name"])
+                F.approxCountDistinct(input["col"], input.get("rsd")).alias(agg["name"])
             )
         if aggregator["name"] == "avg":
-            agg_cols.append(F.avg(input_repl).alias(agg["name"]))
+            agg_cols.append(F.avg(input).alias(agg["name"]))
         if aggregator["name"] in {"collect_set_int", "collect_set_float", "collect_set_string"}:
-            agg_cols.append(F.collect_set(input_repl).alias(agg["name"]))
+            agg_cols.append(F.collect_set(input).alias(agg["name"]))
         if aggregator["name"] == "count":
-            agg_cols.append(F.count(input_repl).alias(agg["name"]))
+            agg_cols.append(F.count(input).alias(agg["name"]))
         if aggregator["name"] == "count_distinct":
-            agg_cols.append(F.countDistinct(*input_repl).alias(agg["name"]))
+            agg_cols.append(F.countDistinct(*input).alias(agg["name"]))
         if aggregator["name"] == "covar_pop":
-            agg_cols.append(F.covar_pop(input_repl["col1"], input_repl["col2"]).alias(agg["name"]))
+            agg_cols.append(F.covar_pop(input["col1"], input["col2"]).alias(agg["name"]))
         if aggregator["name"] == "covar_samp":
-            agg_cols.append(F.covar_samp(input_repl["col1"], input_repl["col2"]).alias(agg["name"]))
+            agg_cols.append(F.covar_samp(input["col1"], input["col2"]).alias(agg["name"]))
         if aggregator["name"] == "kurtosis":
-            agg_cols.append(F.kurtosis(input_repl).alias(agg["name"]))
+            agg_cols.append(F.kurtosis(input).alias(agg["name"]))
         if aggregator["name"] in {"max_int", "max_float", "max_string"}:
-            agg_cols.append(F.max(input_repl).alias(agg["name"]))
+            agg_cols.append(F.max(input).alias(agg["name"]))
         if aggregator["name"] == "mean":
-            agg_cols.append(F.mean(input_repl).alias(agg["name"]))
+            agg_cols.append(F.mean(input).alias(agg["name"]))
         if aggregator["name"] in {"min_int", "min_float", "min_string"}:
-            agg_cols.append(F.min(input_repl).alias(agg["name"]))
+            agg_cols.append(F.min(input).alias(agg["name"]))
         if aggregator["name"] == "skewness":
-            agg_cols.append(F.skewness(input_repl).alias(agg["name"]))
+            agg_cols.append(F.skewness(input).alias(agg["name"]))
         if aggregator["name"] == "stddev":
-            agg_cols.append(F.stddev(input_repl).alias(agg["name"]))
+            agg_cols.append(F.stddev(input).alias(agg["name"]))
         if aggregator["name"] == "stddev_pop":
-            agg_cols.append(F.stddev_pop(input_repl).alias(agg["name"]))
+            agg_cols.append(F.stddev_pop(input).alias(agg["name"]))
         if aggregator["name"] == "stddev_samp":
-            agg_cols.append(F.stddev_samp(input_repl).alias(agg["name"]))
+            agg_cols.append(F.stddev_samp(input).alias(agg["name"]))
         if aggregator["name"] in {"sum_int", "sum_float"}:
-            agg_cols.append(F.sum(input_repl).alias(agg["name"]))
+            agg_cols.append(F.sum(input).alias(agg["name"]))
         if aggregator["name"] in {"sum_distinct_int", "sum_distinct_float"}:
-            agg_cols.append(F.sumDistinct(input_repl).alias(agg["name"]))
+            agg_cols.append(F.sumDistinct(input).alias(agg["name"]))
         if aggregator["name"] == "var_pop":
-            agg_cols.append(F.var_pop(input_repl).alias(agg["name"]))
+            agg_cols.append(F.var_pop(input).alias(agg["name"]))
         if aggregator["name"] == "var_samp":
-            agg_cols.append(F.var_samp(input_repl).alias(agg["name"]))
+            agg_cols.append(F.var_samp(input).alias(agg["name"]))
         if aggregator["name"] == "variance":
-            agg_cols.append(F.variance(input_repl).alias(agg["name"]))
+            agg_cols.append(F.variance(input).alias(agg["name"]))
 
     results = df.agg(*agg_cols).collect()[0].asDict()
 
@@ -479,12 +477,10 @@ def run_builtin_aggregators(builtin_aggregates, df, ctx, spark):
 def run_custom_aggregator(aggregate, df, ctx, spark):
     aggregator = ctx.aggregators[aggregate["aggregator"]]
     aggregator_impl, _ = ctx.get_aggregator_impl(aggregate["name"])
-    input_repl = ctx.populate_values(
-        aggregate["input"], aggregator["input"], preserve_column_refs=False
-    )
+    input = ctx.populate_values(aggregate["input"], aggregator["input"], preserve_column_refs=False)
 
     try:
-        result = aggregator_impl.aggregate_spark(df, input_repl)
+        result = aggregator_impl.aggregate_spark(df, input)
     except Exception as e:
         raise UserRuntimeException(
             "aggregate " + aggregate["name"],
@@ -517,11 +513,11 @@ def execute_transform_spark(column_name, df, ctx, spark):
         spark.sparkContext.addPyFile(trans_impl_path)  # Executor pods need this because of the UDF
         ctx.spark_uploaded_impls[trans_impl_path] = True
 
-    input_repl = ctx.populate_values(
+    input = ctx.populate_values(
         transformed_column["input"], transformer["input"], preserve_column_refs=False
     )
     try:
-        return trans_impl.transform_spark(df, input_repl, column_name)
+        return trans_impl.transform_spark(df, input, column_name)
     except Exception as e:
         raise UserRuntimeException("function transform_spark") from e
 
@@ -532,7 +528,7 @@ def execute_transform_python(column_name, df, ctx, spark, validate=False):
     transformer = ctx.transformers[transformed_column["transformer"]]
 
     input_cols_sorted = sorted(ctx.extract_column_names(transformed_column["input"]))
-    input_repl = ctx.populate_values(
+    input = ctx.populate_values(
         transformed_column["input"], transformer["input"], preserve_column_refs=True
     )
 
@@ -541,9 +537,7 @@ def execute_transform_python(column_name, df, ctx, spark, validate=False):
         ctx.spark_uploaded_impls[trans_impl_path] = True
 
     def _transform(*values):
-        transformer_input = create_transformer_inputs_from_lists(
-            input_repl, input_cols_sorted, values
-        )
+        transformer_input = create_transformer_inputs_from_lists(input, input_cols_sorted, values)
         return trans_impl.transform_python(transformer_input)
 
     transform_python_func = _transform
@@ -593,15 +587,15 @@ def validate_transformer(column_name, test_df, ctx, spark):
             if transformer["output_type"] == consts.COLUMN_TYPE_INFERRED:
                 sample_df = test_df.collect()
                 sample = sample_df[0]
-                input_repl = ctx.populate_values(
+                input = ctx.populate_values(
                     transformed_column["input"], transformer["input"], preserve_column_refs=True
                 )
-                transformer_input = create_transformer_inputs_from_map(input_repl, sample)
+                transformer_input = create_transformer_inputs_from_map(input, sample)
                 initial_transformed_sample = trans_impl.transform_python(transformer_input)
                 inferred_python_type = infer_type(initial_transformed_sample)
 
                 for row in sample_df:
-                    transformer_input = create_transformer_inputs_from_map(input_repl, row)
+                    transformer_input = create_transformer_inputs_from_map(input, row)
                     transformed_sample = trans_impl.transform_python(transformer_input)
                     if inferred_python_type != infer_type(transformed_sample):
                         raise UserRuntimeException(