ml/attribute/JavaAttributeGroupSuite.java

0001 /*
0002  * Licensed to the Apache Software Foundation (ASF) under one or more
0003  * contributor license agreements.  See the NOTICE file distributed with
0004  * this work for additional information regarding copyright ownership.
0005  * The ASF licenses this file to You under the Apache License, Version 2.0
0006  * (the "License"); you may not use this file except in compliance with
0007  * the License.  You may obtain a copy of the License at
0008  *
0009  *    http://www.apache.org/licenses/LICENSE-2.0
0010  *
0011  * Unless required by applicable law or agreed to in writing, software
0012  * distributed under the License is distributed on an "AS IS" BASIS,
0013  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
0014  * See the License for the specific language governing permissions and
0015  * limitations under the License.
0016  */
0017
0018 package org.apache.spark.ml.attribute;
0019
0020 import org.junit.Assert;
0021 import org.junit.Test;
0022
0023 public class JavaAttributeGroupSuite {
0024
0025   @Test
0026   public void testAttributeGroup() {
0027     Attribute[] attrs = new Attribute[]{
0028       NumericAttribute.defaultAttr(),
0029       NominalAttribute.defaultAttr(),
0030       BinaryAttribute.defaultAttr().withIndex(0),
0031       NumericAttribute.defaultAttr().withName("age").withSparsity(0.8),
0032       NominalAttribute.defaultAttr().withName("size").withValues("small", "medium", "large"),
0033       BinaryAttribute.defaultAttr().withName("clicked").withValues("no", "yes"),
0034       NumericAttribute.defaultAttr(),
0035       NumericAttribute.defaultAttr()
0036     };
0037     AttributeGroup group = new AttributeGroup("user", attrs);
0038     Assert.assertEquals(8, group.size());
0039     Assert.assertEquals("user", group.name());
0040     Assert.assertEquals(NumericAttribute.defaultAttr().withIndex(0), group.getAttr(0));
0041     Assert.assertEquals(3, group.indexOf("age"));
0042     Assert.assertFalse(group.hasAttr("abc"));
0043     Assert.assertEquals(group, AttributeGroup.fromStructField(group.toStructField()));
0044   }
0045 }