Add support for skipping the float type

This adds support for deserialization of messages that contain the float type (19) from Facebook Thrift.
AntiBro · Oct 22, 2018 · 9c98d1d · 9c98d1d
1 parent e55f4c6
commit 9c98d1d
Show file tree

Hide file tree

Showing 10 changed files with 186 additions and 2 deletions.
diff --git a/drift-protocol/src/main/java/io/airlift/drift/protocol/TBinaryProtocol.java b/drift-protocol/src/main/java/io/airlift/drift/protocol/TBinaryProtocol.java
@@ -24,6 +24,8 @@
 
 import static java.lang.Double.doubleToLongBits;
 import static java.lang.Double.longBitsToDouble;
+import static java.lang.Float.floatToIntBits;
+import static java.lang.Float.intBitsToFloat;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.Objects.requireNonNull;
 
@@ -176,6 +178,13 @@ public void writeI64(long value)
         transport.write(i64out, 0, 8);
     }
 
+    @Override
+    public void writeFloat(float value)
+            throws TException
+    {
+        writeI32(floatToIntBits(value));
+    }
+
     @Override
     public void writeDouble(double value)
             throws TException
@@ -345,6 +354,13 @@ public long readI64()
                 ((long) (buf[off + 7] & 0xff));
     }
 
+    @Override
+    public float readFloat()
+            throws TException
+    {
+        return intBitsToFloat(readI32());
+    }
+
     @Override
     public double readDouble()
             throws TException

diff --git a/drift-protocol/src/main/java/io/airlift/drift/protocol/TCompactProtocol.java b/drift-protocol/src/main/java/io/airlift/drift/protocol/TCompactProtocol.java
@@ -26,6 +26,8 @@
 
 import static java.lang.Double.doubleToLongBits;
 import static java.lang.Double.longBitsToDouble;
+import static java.lang.Float.floatToIntBits;
+import static java.lang.Float.intBitsToFloat;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.Objects.requireNonNull;
 
@@ -46,7 +48,7 @@ public class TCompactProtocol
     private static final TStruct ANONYMOUS_STRUCT = new TStruct("");
     private static final TField TSTOP = new TField("", TType.STOP, (short) 0);
 
-    private static final byte[] TTYPE_TO_COMPACT_TYPE = new byte[16];
+    private static final byte[] TTYPE_TO_COMPACT_TYPE = new byte[20];
 
     static {
         TTYPE_TO_COMPACT_TYPE[TType.STOP] = TType.STOP;
@@ -55,6 +57,7 @@ public class TCompactProtocol
         TTYPE_TO_COMPACT_TYPE[TType.I16] = Types.I16;
         TTYPE_TO_COMPACT_TYPE[TType.I32] = Types.I32;
         TTYPE_TO_COMPACT_TYPE[TType.I64] = Types.I64;
+        TTYPE_TO_COMPACT_TYPE[TType.FLOAT] = Types.FLOAT;
         TTYPE_TO_COMPACT_TYPE[TType.DOUBLE] = Types.DOUBLE;
         TTYPE_TO_COMPACT_TYPE[TType.STRING] = Types.BINARY;
         TTYPE_TO_COMPACT_TYPE[TType.LIST] = Types.LIST;
@@ -87,6 +90,7 @@ private static final class Types
         public static final byte SET = 0x0A;
         public static final byte MAP = 0x0B;
         public static final byte STRUCT = 0x0C;
+        public static final byte FLOAT = 0x0D;
     }
 
     /**
@@ -319,6 +323,18 @@ public void writeI64(long value)
         writeVarint64(longToZigzag(value));
     }
 
+    /**
+     * Write a float to the wire as 4 bytes.
+     */
+    @Override
+    public void writeFloat(float value)
+            throws TException
+    {
+        byte[] data = {0, 0, 0, 0};
+        fixedIntToBytes(floatToIntBits(value), data);
+        transport.write(data);
+    }
+
     /**
      * Write a double to the wire as 8 bytes.
      */
@@ -475,6 +491,17 @@ private static void fixedLongToBytes(long n, byte[] buf)
         buf[7] = (byte) ((n >> 56) & 0xff);
     }
 
+    /**
+     * Convert an int into little-endian bytes in buf
+     */
+    private static void fixedIntToBytes(int n, byte[] buf)
+    {
+        buf[0] = (byte) (n & 0xff);
+        buf[1] = (byte) ((n >> 8) & 0xff);
+        buf[2] = (byte) ((n >> 16) & 0xff);
+        buf[3] = (byte) ((n >> 24) & 0xff);
+    }
+
     private final byte[] byteDirectBuffer = new byte[1];
 
     /**
@@ -693,6 +720,19 @@ public long readI64()
         return zigzagToLong(readVarint64());
     }
 
+    private final byte[] floatBuf = new byte[4];
+
+    /**
+     * No magic here - just read a float off the wire.
+     */
+    @Override
+    public float readFloat()
+            throws TException
+    {
+        transport.read(floatBuf, 0, 4);
+        return intBitsToFloat(bytesToInt(floatBuf));
+    }
+
     private final byte[] doubleBuf = new byte[8];
 
     /**
@@ -861,6 +901,14 @@ private static long bytesToLong(byte[] bytes)
                 ((bytes[0] & 0xffL));
     }
 
+    private static int bytesToInt(byte[] bytes)
+    {
+        return ((bytes[3] & 0xff) << 24) |
+                ((bytes[2] & 0xff) << 16) |
+                ((bytes[1] & 0xff) << 8) |
+                ((bytes[0] & 0xff));
+    }
+
     //
     // type testing and converting
     //
@@ -892,6 +940,8 @@ private static byte getTType(byte type)
                 return TType.I32;
             case Types.I64:
                 return TType.I64;
+            case Types.FLOAT:
+                return TType.FLOAT;
             case Types.DOUBLE:
                 return TType.DOUBLE;
             case Types.BINARY:

diff --git a/drift-protocol/src/main/java/io/airlift/drift/protocol/TFacebookCompactProtocol.java b/drift-protocol/src/main/java/io/airlift/drift/protocol/TFacebookCompactProtocol.java
@@ -26,6 +26,8 @@
 
 import static java.lang.Double.doubleToLongBits;
 import static java.lang.Double.longBitsToDouble;
+import static java.lang.Float.floatToIntBits;
+import static java.lang.Float.intBitsToFloat;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.Objects.requireNonNull;
 
@@ -46,7 +48,7 @@ public class TFacebookCompactProtocol
     private static final TStruct ANONYMOUS_STRUCT = new TStruct("");
     private static final TField TSTOP = new TField("", TType.STOP, (short) 0);
 
-    private static final byte[] TTYPE_TO_COMPACT_TYPE = new byte[16];
+    private static final byte[] TTYPE_TO_COMPACT_TYPE = new byte[20];
 
     static {
         TTYPE_TO_COMPACT_TYPE[TType.STOP] = TType.STOP;
@@ -55,6 +57,7 @@ public class TFacebookCompactProtocol
         TTYPE_TO_COMPACT_TYPE[TType.I16] = Types.I16;
         TTYPE_TO_COMPACT_TYPE[TType.I32] = Types.I32;
         TTYPE_TO_COMPACT_TYPE[TType.I64] = Types.I64;
+        TTYPE_TO_COMPACT_TYPE[TType.FLOAT] = Types.FLOAT;
         TTYPE_TO_COMPACT_TYPE[TType.DOUBLE] = Types.DOUBLE;
         TTYPE_TO_COMPACT_TYPE[TType.STRING] = Types.BINARY;
         TTYPE_TO_COMPACT_TYPE[TType.LIST] = Types.LIST;
@@ -87,6 +90,7 @@ private static final class Types
         public static final byte SET = 0x0A;
         public static final byte MAP = 0x0B;
         public static final byte STRUCT = 0x0C;
+        public static final byte FLOAT = 0x0D;
     }
 
     /**
@@ -319,6 +323,18 @@ public void writeI64(long value)
         writeVarint64(longToZigzag(value));
     }
 
+    /**
+     * Write a float to the wire as 4 bytes.
+     */
+    @Override
+    public void writeFloat(float value)
+            throws TException
+    {
+        byte[] data = {0, 0, 0, 0};
+        fixedIntToBytes(floatToIntBits(value), data);
+        transport.write(data);
+    }
+
     /**
      * Write a double to the wire as 8 bytes.
      */
@@ -475,6 +491,17 @@ private static void fixedLongToBytes(long n, byte[] buf)
         buf[7] = (byte) (n & 0xff);
     }
 
+    /**
+     * Convert an int into big-endian bytes in buf
+     */
+    private static void fixedIntToBytes(int n, byte[] buf)
+    {
+        buf[0] = (byte) ((n >> 24) & 0xff);
+        buf[1] = (byte) ((n >> 16) & 0xff);
+        buf[2] = (byte) ((n >> 8) & 0xff);
+        buf[3] = (byte) (n & 0xff);
+    }
+
     private final byte[] byteDirectBuffer = new byte[1];
 
     /**
@@ -693,6 +720,19 @@ public long readI64()
         return zigzagToLong(readVarint64());
     }
 
+    private final byte[] floatBuf = new byte[4];
+
+    /**
+     * No magic here - just read a float off the wire.
+     */
+    @Override
+    public float readFloat()
+            throws TException
+    {
+        transport.read(floatBuf, 0, 4);
+        return intBitsToFloat(bytesToInt(floatBuf));
+    }
+
     private final byte[] doubleBuf = new byte[8];
 
     /**
@@ -861,6 +901,14 @@ private static long bytesToLong(byte[] bytes)
                 (bytes[7] & 0xffL);
     }
 
+    private static int bytesToInt(byte[] bytes)
+    {
+        return ((bytes[0] & 0xff) << 24) |
+                ((bytes[1] & 0xff) << 16) |
+                ((bytes[2] & 0xff) << 8) |
+                ((bytes[3] & 0xff));
+    }
+
     //
     // type testing and converting
     //
@@ -892,6 +940,8 @@ private static byte getTType(byte type)
                 return TType.I32;
             case Types.I64:
                 return TType.I64;
+            case Types.FLOAT:
+                return TType.FLOAT;
             case Types.DOUBLE:
                 return TType.DOUBLE;
             case Types.BINARY:

diff --git a/drift-protocol/src/main/java/io/airlift/drift/protocol/TProtocolReader.java b/drift-protocol/src/main/java/io/airlift/drift/protocol/TProtocolReader.java
@@ -72,6 +72,9 @@ int readI32()
     long readI64()
             throws TException;
 
+    float readFloat()
+            throws TException;
+
     double readDouble()
             throws TException;
 

diff --git a/drift-protocol/src/main/java/io/airlift/drift/protocol/TProtocolUtil.java b/drift-protocol/src/main/java/io/airlift/drift/protocol/TProtocolUtil.java
@@ -42,6 +42,9 @@ public static void skip(TProtocolReader protocol, byte type)
             case TType.I64:
                 protocol.readI64();
                 return;
+            case TType.FLOAT:
+                protocol.readFloat();
+                return;
             case TType.DOUBLE:
                 protocol.readDouble();
                 return;

diff --git a/drift-protocol/src/main/java/io/airlift/drift/protocol/TProtocolWriter.java b/drift-protocol/src/main/java/io/airlift/drift/protocol/TProtocolWriter.java
@@ -75,6 +75,9 @@ void writeI32(int value)
     void writeI64(long value)
             throws TException;
 
+    void writeFloat(float value)
+            throws TException;
+
     void writeDouble(double value)
             throws TException;
 

diff --git a/drift-protocol/src/main/java/io/airlift/drift/protocol/TType.java b/drift-protocol/src/main/java/io/airlift/drift/protocol/TType.java
@@ -34,6 +34,7 @@ public final class TType
     public static final byte SET = 14;
     public static final byte LIST = 15;
     public static final byte ENUM = 16;
+    public static final byte FLOAT = 19;
 
     private TType() {}
 }
diff --git a/drift-protocol/src/test/java/io/airlift/drift/protocol/TestProtocol.java b/drift-protocol/src/test/java/io/airlift/drift/protocol/TestProtocol.java
@@ -0,0 +1,44 @@
+/*
+ * Copyright (C) 2018 Facebook, Inc.
+ *
+ * Licensed 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.
+ */
+package io.airlift.drift.protocol;
+
+import io.airlift.drift.TException;
+import org.testng.annotations.Test;
+
+import java.util.function.Function;
+
+import static org.testng.Assert.assertEquals;
+
+public class TestProtocol
+{
+    @Test
+    public void testFloat()
+            throws Exception
+    {
+        assertFloat(TBinaryProtocol::new);
+        assertFloat(TCompactProtocol::new);
+        assertFloat(TFacebookCompactProtocol::new);
+    }
+
+    private static void assertFloat(Function<TTransport, TProtocol> factory)
+            throws TException
+    {
+        TTransport transport = new TMemoryBuffer(0);
+        TProtocol protocol = factory.apply(transport);
+        protocol.writeFloat(123.45f);
+        assertEquals(protocol.readFloat(), 123.45f);
+    }
+}
diff --git a/...e/src/main/java/io/airlift/drift/transport/apache/client/ThriftToDriftProtocolReader.java b/...e/src/main/java/io/airlift/drift/transport/apache/client/ThriftToDriftProtocolReader.java
@@ -249,6 +249,13 @@ public long readI64()
         }
     }
 
+    @Override
+    public float readFloat()
+            throws TException
+    {
+        throw new TException("float is not supported");
+    }
+
     @Override
     public double readDouble()
             throws TException

diff --git a/...e/src/main/java/io/airlift/drift/transport/apache/client/ThriftToDriftProtocolWriter.java b/...e/src/main/java/io/airlift/drift/transport/apache/client/ThriftToDriftProtocolWriter.java
@@ -255,6 +255,13 @@ public void writeI64(long value)
         }
     }
 
+    @Override
+    public void writeFloat(float value)
+            throws TException
+    {
+        throw new TException("float is not supported");
+    }
+
     @Override
     public void writeDouble(double value)
             throws TException