support float type in data frame (#129)

Signed-off-by: Yaliang Wu <ylwu@amazon.com>
opensearch-project · Jan 26, 2022 · 82177a5 · 82177a5
1 parent 5a0e40e
commit 82177a5
Show file tree

Hide file tree

Showing 9 changed files with 138 additions and 0 deletions.
diff --git a/common/src/main/java/org/opensearch/ml/common/dataframe/ColumnType.java b/common/src/main/java/org/opensearch/ml/common/dataframe/ColumnType.java
@@ -17,6 +17,7 @@ public enum ColumnType {
     INTEGER,
     DOUBLE,
     BOOLEAN,
+    FLOAT,
     NULL;
 
     public static ColumnType from(Object object) {
@@ -36,6 +37,10 @@ public static ColumnType from(Object object) {
             return BOOLEAN;
         }
 
+        if(object instanceof Float) {
+            return FLOAT;
+        }
+
         throw new IllegalArgumentException("unsupported type:" + object.getClass().getName());
     }
 }
diff --git a/common/src/main/java/org/opensearch/ml/common/dataframe/ColumnValue.java b/common/src/main/java/org/opensearch/ml/common/dataframe/ColumnValue.java
@@ -37,6 +37,10 @@ default double doubleValue() {
         throw new RuntimeException("the value isn't Double type");
     }
 
+    default float floatValue() {
+        throw new RuntimeException("the value isn't Float type");
+    }
+
     default boolean booleanValue() {
         throw new RuntimeException("the value isn't Boolean type");
     }

diff --git a/common/src/main/java/org/opensearch/ml/common/dataframe/ColumnValueBuilder.java b/common/src/main/java/org/opensearch/ml/common/dataframe/ColumnValueBuilder.java
@@ -45,6 +45,10 @@ public ColumnValue build(Object object) {
             return new BooleanValue((Boolean)object);
         }
 
+        if(object instanceof Float) {
+            return new FloatValue((Float)object);
+        }
+
         throw new IllegalArgumentException("unsupported type:" + object.getClass().getName());
     }
 }
diff --git a/common/src/main/java/org/opensearch/ml/common/dataframe/ColumnValueReader.java b/common/src/main/java/org/opensearch/ml/common/dataframe/ColumnValueReader.java
@@ -30,6 +30,8 @@ public ColumnValue read(StreamInput in) throws IOException {
                 return new StringValue(in.readString());
             case BOOLEAN:
                 return new BooleanValue(in.readBoolean());
+            case FLOAT:
+                return new FloatValue(in.readFloat());
             case NULL:
                 return new NullValue();
             default:

diff --git a/common/src/main/java/org/opensearch/ml/common/dataframe/FloatValue.java b/common/src/main/java/org/opensearch/ml/common/dataframe/FloatValue.java
@@ -0,0 +1,54 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ *
+ */
+
+package org.opensearch.ml.common.dataframe;
+
+import lombok.AccessLevel;
+import lombok.RequiredArgsConstructor;
+import lombok.ToString;
+import lombok.experimental.FieldDefaults;
+import org.opensearch.common.io.stream.StreamOutput;
+
+import java.io.IOException;
+
+@FieldDefaults(makeFinal = true, level = AccessLevel.PRIVATE)
+@RequiredArgsConstructor
+@ToString
+public class FloatValue implements ColumnValue {
+    float value;
+
+    @Override
+    public ColumnType columnType() {
+        return ColumnType.FLOAT;
+    }
+
+    @Override
+    public Object getValue() {
+        return value;
+    }
+
+    @Override
+    public double doubleValue() {
+        return Float.valueOf(value).doubleValue();
+    }
+
+    @Override
+    public float floatValue() {
+        return value;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeEnum(columnType());
+        out.writeFloat(value);
+    }
+}
diff --git a/common/src/test/java/org/opensearch/ml/common/dataframe/ColumnValueBuilderTest.java b/common/src/test/java/org/opensearch/ml/common/dataframe/ColumnValueBuilderTest.java
@@ -45,6 +45,10 @@ public void build() {
         value = ColumnValueBuilder.build(true);
         assertEquals(ColumnType.BOOLEAN, value.columnType());
         assertEquals(true, value.booleanValue());
+
+        value = ColumnValueBuilder.build(2.1f);
+        assertEquals(ColumnType.FLOAT, value.columnType());
+        assertEquals(2.1f, value.floatValue(), 1e-5);
     }
 
     @Test

diff --git a/common/src/test/java/org/opensearch/ml/common/dataframe/ColumnValueReaderTest.java b/common/src/test/java/org/opensearch/ml/common/dataframe/ColumnValueReaderTest.java
@@ -80,4 +80,14 @@ public void read_BooleanValue() throws IOException {
         assertEquals(ColumnType.BOOLEAN, value.columnType());
         assertEquals(true, value.booleanValue());
     }
+
+    @Test
+    public void read_FloatValue() throws IOException {
+        ColumnValue value = new FloatValue(2.1f);
+        BytesStreamOutput bytesStreamOutput = new BytesStreamOutput();
+        value.writeTo(bytesStreamOutput);
+        value = reader.read(bytesStreamOutput.bytes().streamInput());
+        assertEquals(ColumnType.FLOAT, value.columnType());
+        assertEquals(2.1f, value.floatValue(), 1e-5);
+    }
 }
diff --git a/common/src/test/java/org/opensearch/ml/common/dataframe/ColumnValueTest.java b/common/src/test/java/org/opensearch/ml/common/dataframe/ColumnValueTest.java
@@ -60,4 +60,12 @@ public void wrongDoubleValue() {
         ColumnValue value = new IntValue(1);
         value.doubleValue();
     }
+
+    @Test
+    public void wrongFloatValue() {
+        exceptionRule.expect(RuntimeException.class);
+        exceptionRule.expectMessage("the value isn't Float type");
+        ColumnValue value = new IntValue(1);
+        value.floatValue();
+    }
 }
diff --git a/common/src/test/java/org/opensearch/ml/common/dataframe/FloatValueTest.java b/common/src/test/java/org/opensearch/ml/common/dataframe/FloatValueTest.java
@@ -0,0 +1,47 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ *
+ */
+
+package org.opensearch.ml.common.dataframe;
+
+import org.junit.Test;
+import org.opensearch.common.Strings;
+import org.opensearch.common.xcontent.XContentBuilder;
+import org.opensearch.common.xcontent.XContentFactory;
+import org.opensearch.common.xcontent.XContentType;
+
+import java.io.IOException;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+
+public class FloatValueTest {
+
+    @Test
+    public void floatValue() {
+        FloatValue floatValue = new FloatValue(2.1f);
+        assertEquals(ColumnType.FLOAT, floatValue.columnType());
+        assertEquals(2.1f, floatValue.getValue());
+        assertEquals(2.1f, floatValue.floatValue(), 1e-5);
+        assertEquals(2.1d, floatValue.doubleValue(), 1e-5);
+    }
+
+    @Test
+    public void testToXContent() throws IOException {
+        FloatValue floatValue = new FloatValue(2.1f);
+        XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON);
+        floatValue.toXContent(builder);
+
+        assertNotNull(builder);
+        String jsonStr = Strings.toString(builder);
+        assertEquals("{\"column_type\":\"FLOAT\",\"value\":2.1}", jsonStr);
+    }
+}