| # coding=utf-8 |
| # |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # 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. |
| |
| import plpy |
| from utilities.control import MinWarning |
| from utilities.utilities import _assert |
| from utilities.utilities import _check_groups |
| from utilities.utilities import get_table_qualified_col_str |
| from utilities.utilities import extract_keyvalue_params |
| from utilities.utilities import add_postfix |
| from utilities.utilities import unique_string |
| from utilities.utilities import split_quoted_delimited_str |
| from utilities.validate_args import table_exists |
| from utilities.validate_args import columns_exist_in_table |
| from utilities.validate_args import table_is_empty |
| from utilities.validate_args import get_expr_type |
| from utilities.validate_args import get_cols |
| |
| m4_changequote(` <!', `!>') |
| |
| |
| def _get_sql_string(str): |
| if str: |
| return "'" + str + "'" |
| return "NULL" |
| |
| |
| def train_test_split(schema_madlib, source_table, output_table, train_proportion, |
| test_proportion, grouping_cols, target_cols, with_replacement, |
| separate_output_tables, **kwargs): |
| """ |
| test train split function |
| Args: |
| @param source_table Input table name. |
| @param output_table Output table name. |
| @param train_proportion The ratio of training data to the entire |
| input table |
| @param test_proportion The ratio of test data to the entire |
| input table |
| @param grouping_cols (Default: NULL) The columns to distinguish |
| each strata. |
| @param target_cols (Default: NULL) The columns to include in |
| the output. |
| @param with_replacement (Default: FALSE) The sampling method. |
| @param separate_output_tables (Default: FALSE) Create two output tables, |
| <output_table>_train and <output_table>_test. |
| Otherwise one output table is created with |
| and additional column 'split' which takes the |
| value 0 for test and 1 for training. |
| |
| """ |
| with MinWarning("warning"): |
| if test_proportion is None: |
| test_proportion = 1 - train_proportion |
| validate_strs(source_table, output_table, train_proportion, test_proportion, |
| split_quoted_delimited_str(grouping_cols), target_cols, |
| with_replacement) |
| grouping_cols = _get_sql_string(grouping_cols) |
| target_cols = _get_sql_string(target_cols) |
| with_replacement = with_replacement or "False" |
| strat_query = """ |
| SELECT {schema_madlib}.stratified_sample( |
| '{strat_source_table}', |
| '{strat_out_table}', |
| '{strat_proportion}', |
| {strat_grouping_cols}, |
| {strat_target_cols}, |
| {strat_with_replacement} |
| ) |
| """ |
| strat_out_table = unique_string() |
| q = strat_query.format( |
| schema_madlib=schema_madlib, |
| strat_source_table=source_table, |
| strat_out_table=strat_out_table, |
| strat_proportion=train_proportion + test_proportion, |
| strat_grouping_cols=grouping_cols, |
| strat_with_replacement=with_replacement, |
| strat_target_cols=target_cols |
| ) |
| plpy.execute(q) |
| test_table = add_postfix(output_table, "_test") |
| train_table = add_postfix(output_table, "_train") |
| if not separate_output_tables: |
| test_table = unique_string() |
| train_table = unique_string() |
| test_query = strat_query.format( |
| schema_madlib=schema_madlib, |
| strat_source_table=strat_out_table, |
| strat_out_table=test_table, |
| strat_proportion=(test_proportion / |
| (train_proportion + test_proportion)), |
| strat_grouping_cols=grouping_cols, |
| strat_with_replacement=False, |
| strat_target_cols=target_cols |
| ) |
| plpy.execute(test_query) |
| train_query = """ |
| CREATE TABLE {train_table} AS |
| SELECT * FROM {strat_out_table} |
| EXCEPT ALL |
| SELECT * FROM {test_table} |
| """.format(train_table=train_table, |
| strat_out_table=strat_out_table, |
| test_table=test_table) |
| plpy.execute(train_query) |
| clean_up_tables = [strat_out_table] |
| if not separate_output_tables: |
| union_query = """ |
| CREATE TABLE {output_table} AS |
| SELECT *,0 AS split FROM {test_table} |
| UNION ALL |
| SELECT *,1 AS split FROM {train_table} |
| """.format(output_table=output_table, |
| test_table=test_table, |
| train_table=train_table) |
| plpy.execute(union_query) |
| clean_up_tables += [train_table, test_table] |
| clean_up_query = """ |
| DROP TABLE IF EXISTS {clean_up_tables} |
| """.format(clean_up_tables=",".join(clean_up_tables)) |
| plpy.execute(clean_up_query) |
| return |
| |
| |
| def validate_strs(source_table, output_table, train_proportion, test_proportion, glist, target_cols, with_replacement): |
| |
| _assert(output_table and output_table.strip().lower() not in ('null', ''), |
| "Sample: Invalid output table name {output_table}!".format(**locals())) |
| _assert(not table_exists(output_table), |
| "Sample: Output table already exists!".format(**locals())) |
| |
| _assert(source_table and source_table.strip().lower() not in ('null', ''), |
| "Sample: Invalid Source table name!".format(**locals())) |
| _assert(table_exists(source_table), |
| "Sample: Source table ({source_table}) is missing!".format(**locals())) |
| _assert(not table_is_empty(source_table), |
| "Sample: Source table ({source_table}) is empty!".format(**locals())) |
| |
| for proportion in [train_proportion, test_proportion]: |
| _assert(proportion > 0 and proportion < 1, |
| "Sample: Proportions aren't in the range (0,1)!") |
| if not with_replacement: |
| _assert(train_proportion + test_proportion <= 1, |
| "Sample: Proportions add up to greater than 1!") |
| |
| if glist is not None: |
| _assert(columns_exist_in_table(source_table, glist), |
| ("""Sample: Not all columns from {glist} are present in source""" + |
| """ table ({source_table}).""").format(**locals())) |
| |
| if not (target_cols is None or target_cols is '*'): |
| tlist = split_quoted_delimited_str(target_cols) |
| _assert(columns_exist_in_table(source_table, tlist), |
| ("""Sample: Not all columns from {target_cols} are present in""" + |
| """ edge table ({source_table})""").format(**locals())) |
| return |
| |
| |
| def train_test_split_help(schema_madlib, message, **kwargs): |
| """ |
| Help function for train_test_split |
| |
| Args: |
| @param schema_madlib |
| @param message: string, Help message string |
| @param kwargs |
| |
| Returns: |
| String. Help/usage information |
| """ |
| if not message: |
| help_string = """ |
| ----------------------------------------------------------------------- |
| SUMMARY |
| ----------------------------------------------------------------------- |
| |
| Given a table, train_test_split returns a random sample of the |
| table for testing and training. It is possible to use with or without |
| replacement sampling methods, specify a set of target columns, and a |
| set of grouping columns, in which case, stratified sampling will be |
| performed. |
| |
| For more details on function usage: |
| SELECT {schema_madlib}.train_test_split('usage'); |
| """ |
| elif message.lower() in ['usage', 'help', '?']: |
| help_string = """ |
| |
| Given a table, test train split returns a proportion of records for |
| each group (strata). It is possible to use with or without replacement |
| sampling methods, specify a set of target columns, and assume the |
| whole table is a single strata. |
| |
| ---------------------------------------------------------------------------- |
| USAGE |
| ---------------------------------------------------------------------------- |
| |
| SELECT {schema_madlib}.train_test_split( |
| source_table TEXT, -- Name of the table containing the input data. |
| output_table TEXT, -- Output table name. |
| train_proportion FLOAT8, -- The ratio of train sample size to the |
| -- number of records. |
| test_proportion FLOAT8, -- The ratio of test sample size to the |
| -- number of records. |
| grouping_cols TEXT -- (Default: NULL) The columns to distinguish |
| -- each strata. |
| target_cols TEXT, -- (Default: NULL) The columns to include in |
| -- the output. |
| with_replacement BOOLEAN -- (Default: FALSE) The sampling method. |
| separate_output_tables |
| BOOLEAN -- (Default: FALSE) Separate the output table |
| -- into $output_table$_train and |
| -- $output_table$_test, otherwise, the split |
| -- column in output_table will identify 1 for |
| -- train set and 0 for test set. |
| |
| If grouping_cols is NULL, the whole table is treated as a single group and |
| sampled accordingly. |
| |
| If target_cols is NULL or '*', all of the columns will be included in the |
| output table. |
| |
| If with_replacement is TRUE, each sample is independent (the same row may |
| be selected in the sample set more than once). Else (if with_replacement |
| is FALSE), a row can be selected at most once. |
| ); |
| """ |
| else: |
| help_string = "No such option. Use {schema_madlib}.train_test_split()" |
| |
| return help_string.format(schema_madlib=schema_madlib) |
