| # 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 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 stratified_sample(schema_madlib, source_table, output_table, proportion, |
| grouping_cols, target_cols, with_replacement, **kwargs): |
| |
| """ |
| Stratified sampling function |
| Args: |
| @param source_table Input table name. |
| @param output_table Output table name. |
| @param proportion The ratio of sample size to the number of |
| records. |
| @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. |
| |
| """ |
| with MinWarning("warning"): |
| label = unique_string(desp='label') |
| perc = unique_string(desp='perc') |
| |
| checkg_lp = "" |
| window = "" |
| grp_by = "" |
| grp_from_perc = "" |
| grp_comma = "" |
| glist = None |
| if grouping_cols is not None: |
| glist = split_quoted_delimited_str(grouping_cols) |
| checkg_lp = " AND " + _check_groups(label,perc,glist) |
| window = "PARTITION BY {0}".format(grouping_cols) |
| grp_by = "GROUP BY {0}".format(grouping_cols) |
| grp_from_perc = get_table_qualified_col_str(perc,glist) + " , " |
| grp_comma = grouping_cols + " , " |
| |
| validate_strs(source_table, output_table, proportion, glist, target_cols) |
| |
| if target_cols is None or target_cols is '*': |
| cols = get_cols(source_table) |
| if grouping_cols is not None: |
| cols = [item for item in cols if item not in glist] |
| target_cols = " , ".join(cols) |
| |
| plpy.execute("DROP TABLE IF EXISTS {0},{1}".format(label,perc)) |
| if not with_replacement : |
| if grouping_cols: |
| |
| # Create a random label for each record |
| sql1 = """ CREATE TEMP TABLE {label} AS ( |
| SELECT {target_cols},{grouping_cols},random() AS __samp_out_label |
| FROM {source_table})""".format(**locals()) |
| plpy.execute(sql1) |
| |
| # Find the cut-off label for the given proportion |
| sql2 = """ CREATE TEMP TABLE {perc} AS ( |
| SELECT {grouping_cols}, percentile_disc({proportion}) |
| WITHIN GROUP (ORDER BY __samp_out_label) AS __samp_out_label |
| FROM {label} GROUP BY {grouping_cols})""".format(**locals()) |
| plpy.execute(sql2) |
| |
| # Select every record that has a label under the threshold |
| sql3 = """ CREATE TABLE {output_table} AS ( |
| SELECT {grp_from_perc} {target_cols} |
| FROM {label} INNER JOIN {perc} ON ( |
| {label}.__samp_out_label <= {perc}.__samp_out_label |
| {checkg_lp}) )""".format(**locals()) |
| plpy.execute(sql3) |
| else: |
| |
| # Find the number of records to select |
| count = plpy.execute("SELECT count(*) AS count FROM {0}". |
| format(source_table))[0]['count'] |
| count = count * proportion |
| |
| # Order randomly and select the required number of records |
| sql1 = """ CREATE TABLE {output_table} AS ( |
| SELECT {target_cols} |
| FROM {source_table} |
| ORDER BY random() |
| LIMIT {count})""".format(**locals()) |
| plpy.execute(sql1) |
| else: |
| |
| # Set the row number as the label for each record |
| # OVER clause ensures that different groups have independent |
| # row_numbers |
| sql1 = """ CREATE TEMP TABLE {label} AS ( |
| SELECT {grp_comma} {target_cols}, |
| row_number() OVER ({window}) AS __samp_out_label |
| FROM {source_table})""".format(**locals()) |
| plpy.execute(sql1) |
| |
| # Generate a series of random values for each group based on their |
| # individual row counts. |
| # These random values are independent from each other and may have |
| # the same value. |
| |
| sql2 = """ CREATE TEMP TABLE {perc} AS ( |
| SELECT {grp_comma} |
| GENERATE_SERIES(0,(count*{proportion}-1)::int) AS __i, |
| ((random()*(count-1)+1)::int) AS __samp_out_label |
| FROM ( |
| SELECT {grp_comma} count(*) AS count |
| FROM {source_table} {grp_by}) AS sub) |
| """.format(**locals()) |
| plpy.execute(sql2) |
| |
| # Join the two tables to get the selected samples. |
| # If a random value is generated twice, the join will ensure that |
| # the record is selected twice |
| sql3 = """ CREATE TABLE {output_table} AS ( |
| SELECT {grp_from_perc} {target_cols} |
| FROM {label} INNER JOIN {perc} ON ( |
| {label}.__samp_out_label = {perc}.__samp_out_label |
| {checkg_lp}) )""".format(**locals()) |
| plpy.execute(sql3) |
| plpy.execute("DROP TABLE IF EXISTS {0},{1}".format(label,perc)) |
| |
| return |
| |
| def validate_strs (source_table, output_table, proportion, glist, target_cols): |
| |
| _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())) |
| |
| _assert(proportion > 0 and proportion <= 1, |
| "Sample: Proportion isn't in the range (0,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 stratified_sample_help(schema_madlib, message, **kwargs): |
| """ |
| Help function for stratified_sample |
| |
| 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, stratified sampling 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. |
| |
| For more details on function usage: |
| SELECT {schema_madlib}.stratified_sample('usage'); |
| """ |
| elif message.lower() in ['usage', 'help', '?']: |
| help_string = """ |
| |
| Given a table, stratified sampling 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}.stratified_sample( |
| source_table TEXT, -- Input table name. |
| output_table TEXT, -- Output table name. |
| proportion FLOAT8, -- The ratio of 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. |
| |
| 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}.stratified_sample()" |
| |
| return help_string.format(schema_madlib=schema_madlib) |
