def test_to_feature_columns_and_input_fn(self): df = setup_test_df_3layer() feature_columns, input_fn = ( estimator_utils.to_feature_columns_and_input_fn( df, base_input_keys_with_defaults={ "a": 1, "b": 2, "c": 3, "d": 4 }, label_keys=["g"], feature_keys=["a", "b", "f"])) expected_feature_column_a = feature_column.DataFrameColumn( "a", learn.PredefinedSeries( "a", parsing_ops.FixedLenFeature(tensor_shape.unknown_shape(), dtypes.int32, 1))) expected_feature_column_b = feature_column.DataFrameColumn( "b", learn.PredefinedSeries("b", parsing_ops.VarLenFeature(dtypes.int32))) expected_feature_column_f = feature_column.DataFrameColumn( "f", learn.TransformedSeries([ learn.PredefinedSeries( "c", parsing_ops.FixedLenFeature(tensor_shape.unknown_shape(), dtypes.int32, 3)), learn.PredefinedSeries("d", parsing_ops.VarLenFeature(dtypes.int32)) ], mocks.Mock2x2Transform("iue", "eui", "snt"), "out2")) expected_feature_columns = [ expected_feature_column_a, expected_feature_column_b, expected_feature_column_f ] self.assertEqual(sorted(expected_feature_columns), sorted(feature_columns)) base_features, labels = input_fn() expected_base_features = { "a": mocks.MockTensor("Tensor a", dtypes.int32), "b": mocks.MockSparseTensor("SparseTensor b", dtypes.int32), "c": mocks.MockTensor("Tensor c", dtypes.int32), "d": mocks.MockSparseTensor("SparseTensor d", dtypes.int32) } self.assertEqual(expected_base_features, base_features) expected_labels = mocks.MockTensor("Out iue", dtypes.int32) self.assertEqual(expected_labels, labels) self.assertEqual(3, len(feature_columns))
def setup_test_df_3layer(): """Create a dataframe populated with some test columns.""" df = learn.DataFrame() df["a"] = mocks.MockSeries("a_series", mocks.MockTensor("Tensor a", tf.int32)) df["b"] = mocks.MockSeries("b_series", mocks.MockSparseTensor("SparseTensor b", tf.int32)) df["c"] = mocks.MockSeries("c_series", mocks.MockTensor("Tensor c", tf.int32)) df["d"] = mocks.MockSeries("d_series", mocks.MockSparseTensor("SparseTensor d", tf.int32)) df["e"] = learn.TransformedSeries([df["a"], df["b"]], mocks.Mock2x2Transform("iue", "eui", "snt"), "out1") df["f"] = learn.TransformedSeries([df["c"], df["d"]], mocks.Mock2x2Transform("iue", "eui", "snt"), "out2") df["g"] = learn.TransformedSeries([df["e"], df["f"]], mocks.Mock2x2Transform("iue", "eui", "snt"), "out1") return df