support short and long type in data frame

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

Hide file tree

Showing 12 changed files with 257 additions and 2 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
@@ -14,17 +14,27 @@
 
 public enum ColumnType {
     STRING,
+    SHORT,
     INTEGER,
+    LONG,
     DOUBLE,
     BOOLEAN,
     FLOAT,
     NULL;
 
     public static ColumnType from(Object object) {
+        if(object instanceof Short) {
+            return SHORT;
+        }
+
         if(object instanceof Integer) {
             return INTEGER;
         }
 
+        if(object instanceof Long) {
+            return LONG;
+        }
+
         if(object instanceof String) {
             return STRING;
         }

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
@@ -25,10 +25,18 @@ public interface ColumnValue extends Writeable, ToXContentObject {
 
     Object getValue();
 
+    default short shortValue() {
+        throw new RuntimeException("the value isn't Short type");
+    }
+
     default int intValue() {
         throw new RuntimeException("the value isn't Integer type");
     }
 
+    default long longValue() {
+        throw new RuntimeException("the value isn't Long type");
+    }
+
     default String stringValue() {
         throw new RuntimeException("the value isn't String 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
@@ -29,10 +29,18 @@ public ColumnValue build(Object object) {
             return new NullValue();
         }
 
+        if(object instanceof Short) {
+            return new ShortValue((Short)object);
+        }
+
         if(object instanceof Integer) {
             return new IntValue((Integer)object);
         }
 
+        if(object instanceof Long) {
+            return new LongValue((Long)object);
+        }
+
         if(object instanceof String) {
             return new StringValue((String)object);
         }

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
@@ -22,8 +22,12 @@ public class ColumnValueReader implements Writeable.Reader<ColumnValue> {
     public ColumnValue read(StreamInput in) throws IOException {
         ColumnType columnType = in.readEnum(ColumnType.class);
         switch (columnType){
+            case SHORT:
+                return new ShortValue(in.readShort());
             case INTEGER:
                 return new IntValue(in.readInt());
+            case LONG:
+                return new LongValue(in.readLong());
             case DOUBLE:
                 return new DoubleValue(in.readDouble());
             case STRING:

diff --git a/common/src/main/java/org/opensearch/ml/common/dataframe/IntValue.java b/common/src/main/java/org/opensearch/ml/common/dataframe/IntValue.java
@@ -42,6 +42,11 @@ public int intValue() {
         return value;
     }
 
+    @Override
+    public double doubleValue() {
+        return Integer.valueOf(value).doubleValue();
+    }
+
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeEnum(columnType());

diff --git a/common/src/main/java/org/opensearch/ml/common/dataframe/LongValue.java b/common/src/main/java/org/opensearch/ml/common/dataframe/LongValue.java
@@ -0,0 +1,47 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+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 LongValue implements ColumnValue {
+    long value;
+
+    @Override
+    public ColumnType columnType() {
+        return ColumnType.LONG;
+    }
+
+    @Override
+    public Object getValue() {
+        return value;
+    }
+
+    @Override
+    public long longValue() {
+        return value;
+    }
+
+    @Override
+    public double doubleValue() {
+        return Long.valueOf(value).doubleValue();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeEnum(columnType());
+        out.writeLong(value);
+    }
+}
diff --git a/common/src/main/java/org/opensearch/ml/common/dataframe/ShortValue.java b/common/src/main/java/org/opensearch/ml/common/dataframe/ShortValue.java
@@ -0,0 +1,47 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+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 ShortValue implements ColumnValue {
+    short value;
+
+    @Override
+    public ColumnType columnType() {
+        return ColumnType.SHORT;
+    }
+
+    @Override
+    public Object getValue() {
+        return value;
+    }
+
+    @Override
+    public short shortValue() {
+        return value;
+    }
+
+    @Override
+    public double doubleValue() {
+        return Short.valueOf(value).doubleValue();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeEnum(columnType());
+        out.writeShort(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
@@ -33,6 +33,7 @@ public void build() {
         value = ColumnValueBuilder.build(2);
         assertEquals(ColumnType.INTEGER, value.columnType());
         assertEquals(2, value.intValue());
+        assertEquals(2.0d, value.doubleValue(), 1e-5);
 
         value = ColumnValueBuilder.build("string");
         assertEquals(ColumnType.STRING, value.columnType());
@@ -49,6 +50,17 @@ public void build() {
         value = ColumnValueBuilder.build(2.1f);
         assertEquals(ColumnType.FLOAT, value.columnType());
         assertEquals(2.1f, value.floatValue(), 1e-5);
+        assertEquals(2.1d, value.doubleValue(), 1e-5);
+
+        value = ColumnValueBuilder.build((short)2);
+        assertEquals(ColumnType.SHORT, value.columnType());
+        assertEquals(2, value.shortValue());
+        assertEquals(2.0d, value.doubleValue(), 1e-5);
+
+        value = ColumnValueBuilder.build((long)2);
+        assertEquals(ColumnType.LONG, value.columnType());
+        assertEquals(2, value.longValue());
+        assertEquals(2.0d, value.doubleValue(), 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
@@ -90,4 +90,24 @@ public void read_FloatValue() throws IOException {
         assertEquals(ColumnType.FLOAT, value.columnType());
         assertEquals(2.1f, value.floatValue(), 1e-5);
     }
-}
+
+    @Test
+    public void read_ShortValue() throws IOException {
+        ColumnValue value = new ShortValue((short)2);
+        BytesStreamOutput bytesStreamOutput = new BytesStreamOutput();
+        value.writeTo(bytesStreamOutput);
+        value = reader.read(bytesStreamOutput.bytes().streamInput());
+        assertEquals(ColumnType.SHORT, value.columnType());
+        assertEquals(2, value.shortValue());
+    }
+
+    @Test
+    public void read_LongValue() throws IOException {
+        ColumnValue value = new LongValue((long)2);
+        BytesStreamOutput bytesStreamOutput = new BytesStreamOutput();
+        value.writeTo(bytesStreamOutput);
+        value = reader.read(bytesStreamOutput.bytes().streamInput());
+        assertEquals(ColumnType.LONG, value.columnType());
+        assertEquals(2, value.longValue());
+    }
+}
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
@@ -57,7 +57,7 @@ public void wrongStringValue() {
     public void wrongDoubleValue() {
         exceptionRule.expect(RuntimeException.class);
         exceptionRule.expectMessage("the value isn't Double type");
-        ColumnValue value = new IntValue(1);
+        ColumnValue value = new BooleanValue(true);
         value.doubleValue();
     }
 
@@ -68,4 +68,20 @@ public void wrongFloatValue() {
         ColumnValue value = new IntValue(1);
         value.floatValue();
     }
+
+    @Test
+    public void wrongShortValue() {
+        exceptionRule.expect(RuntimeException.class);
+        exceptionRule.expectMessage("the value isn't Short type");
+        ColumnValue value = new IntValue(1);
+        value.shortValue();
+    }
+
+    @Test
+    public void wrongLongValue() {
+        exceptionRule.expect(RuntimeException.class);
+        exceptionRule.expectMessage("the value isn't Long type");
+        ColumnValue value = new IntValue(1);
+        value.longValue();
+    }
 }
diff --git a/common/src/test/java/org/opensearch/ml/common/dataframe/LongValueTest.java b/common/src/test/java/org/opensearch/ml/common/dataframe/LongValueTest.java
@@ -0,0 +1,39 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+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 LongValueTest {
+
+    @Test
+    public void longValue() {
+        LongValue longValue = new LongValue((long)2);
+        assertEquals(ColumnType.LONG, longValue.columnType());
+        assertEquals(2L, longValue.getValue());
+        assertEquals(2.0d, longValue.doubleValue(), 1e-5);
+    }
+
+    @Test
+    public void testToXContent() throws IOException {
+        LongValue longValue = new LongValue((long)2);
+        XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON);
+        longValue.toXContent(builder);
+
+        assertNotNull(builder);
+        String jsonStr = Strings.toString(builder);
+        assertEquals("{\"column_type\":\"LONG\",\"value\":2}", jsonStr);
+    }
+}
diff --git a/common/src/test/java/org/opensearch/ml/common/dataframe/ShortValueTest.java b/common/src/test/java/org/opensearch/ml/common/dataframe/ShortValueTest.java
@@ -0,0 +1,39 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+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 ShortValueTest {
+
+    @Test
+    public void shortValue() {
+        ShortValue shortValue = new ShortValue((short)2);
+        assertEquals(ColumnType.SHORT, shortValue.columnType());
+        assertEquals(2, shortValue.getValue());
+        assertEquals(2.0d, shortValue.doubleValue(), 1e-5);
+    }
+
+    @Test
+    public void testToXContent() throws IOException {
+        ShortValue shortValue = new ShortValue((short)2);
+        XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON);
+        shortValue.toXContent(builder);
+
+        assertNotNull(builder);
+        String jsonStr = Strings.toString(builder);
+        assertEquals("{\"column_type\":\"SHORT\",\"value\":2}", jsonStr);
+    }
+}