add unsafe map

Author: cloud-fan

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeArrayData.java

@@ -277,17 +277,17 @@ public ArrayData getArray(int ordinal) {
     assertIndexIsValid(ordinal);
     final int offset = getElementOffset(ordinal);
     if (offset < 0) return null;
-    final int numElements = PlatformDependent.UNSAFE.getInt(baseObject, baseOffset + offset);
     final int size = getElementSize(offset, ordinal);
-    final UnsafeArrayData array = new UnsafeArrayData();
-    // Skip the first 4 bytes.
-    array.pointTo(baseObject, baseOffset + offset + 4, numElements, size - 4);
-    return array;
+    return UnsafeReaders.readArray(baseObject, baseOffset + offset, size);
   }
 
   @Override
   public MapData getMap(int ordinal) {
-    return null;
+    assertIndexIsValid(ordinal);
+    final int offset = getElementOffset(ordinal);
+    if (offset < 0) return null;
+    final int size = getElementSize(offset, ordinal);
+    return UnsafeReaders.readMap(baseObject, baseOffset + offset, size);
   }
 
   @Override

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeMapData.java

@@ -0,0 +1,56 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.sql.catalyst.expressions;
+
+import org.apache.spark.sql.types.ArrayData;
+import org.apache.spark.sql.types.MapData;
+
+/**
+ * An Unsafe implementation of Map which is backed by raw memory instead of Java objects.
+ *
+ * Currently we just use 2 UnsafeArrayData to represent UnsafeMapData.
+ */
+public class UnsafeMapData extends MapData {
+
+  public final UnsafeArrayData keys;
+  public final UnsafeArrayData values;
+  // The number of elements in this array
+  private int numElements;
+
+  public UnsafeMapData(UnsafeArrayData keys, UnsafeArrayData values) {
+    assert keys.numElements() == values.numElements();
+    this.numElements = keys.numElements();
+    this.keys = keys;
+    this.values = values;
+  }
+
+  @Override
+  public int numElements() {
+    return numElements;
+  }
+
+  @Override
+  public ArrayData keyArray() {
+    return keys;
+  }
+
+  @Override
+  public ArrayData valueArray() {
+    return values;
+  }
+}

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeReaders.java

@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.sql.catalyst.expressions;
+
+import org.apache.spark.unsafe.PlatformDependent;
+
+public class UnsafeReaders {
+
+  public static UnsafeArrayData readArray(Object baseObject, long baseOffset, int numBytes) {
+    // Read the number of elements from first 4 bytes.
+    final int numElements = PlatformDependent.UNSAFE.getInt(baseObject, baseOffset);
+    final UnsafeArrayData array = new UnsafeArrayData();
+    // Skip the first 4 bytes.
+    array.pointTo(baseObject, baseOffset + 4, numElements, numBytes - 4);
+    return array;
+  }
+
+  public static UnsafeMapData readMap(Object baseObject, long baseOffset, int numBytes) {
+    // Read the number of elements from first 4 bytes.
+    final int numElements = PlatformDependent.UNSAFE.getInt(baseObject, baseOffset);
+    // Read the numBytes of key array in second 4 bytes.
+    final int keyArraySize = PlatformDependent.UNSAFE.getInt(baseObject, baseOffset + 4);
+    final int valueArraySize = numBytes - 8 - keyArraySize;
+
+    final UnsafeArrayData keyArray = new UnsafeArrayData();
+    keyArray.pointTo(baseObject, baseOffset + 8, numElements, keyArraySize);
+
+    final UnsafeArrayData valueArray = new UnsafeArrayData();
+    valueArray.pointTo(baseObject, baseOffset + 8 + keyArraySize, numElements, valueArraySize);
+
+    return new UnsafeMapData(keyArray, valueArray);
+  }
+}

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java

@@ -365,7 +365,6 @@ public Decimal getDecimal(int ordinal, int precision, int scale) {
 
   @Override
   public UTF8String getUTF8String(int ordinal) {
-    assertIndexIsValid(ordinal);
     if (isNullAt(ordinal)) return null;
     final long offsetAndSize = getLong(ordinal);
     final int offset = (int) (offsetAndSize >> 32);
@@ -375,7 +374,6 @@ public UTF8String getUTF8String(int ordinal) {
 
   @Override
   public byte[] getBinary(int ordinal) {
-    assertIndexIsValid(ordinal);
     if (isNullAt(ordinal)) {
       return null;
     } else {
@@ -430,19 +428,25 @@ public ArrayData getArray(int ordinal) {
       final long offsetAndSize = getLong(ordinal);
       final int offset = (int) (offsetAndSize >> 32);
       final int size = (int) (offsetAndSize & ((1L << 32) - 1));
-      final int numElements = PlatformDependent.UNSAFE.getInt(baseObject, baseOffset + offset);
-      final UnsafeArrayData array = new UnsafeArrayData();
-      // Skip the first 4 bytes.
-      array.pointTo(baseObject, baseOffset + offset + 4, numElements, size - 4);
-      return array;
+      return UnsafeReaders.readArray(baseObject, baseOffset + offset, size);
+    }
+  }
+
+  @Override
+  public MapData getMap(int ordinal) {
+    if (isNullAt(ordinal)) {
+      return null;
+    } else {
+      final long offsetAndSize = getLong(ordinal);
+      final int offset = (int) (offsetAndSize >> 32);
+      final int size = (int) (offsetAndSize & ((1L << 32) - 1));
+      return UnsafeReaders.readMap(baseObject, baseOffset + offset, size);
     }
   }
 
   /**
    * Copies this row, returning a self-contained UnsafeRow that stores its data in an internal
    * byte array rather than referencing data stored in a data page.
-   * <p>
-   * This method is only supported on UnsafeRows that do not use ObjectPools.
    */
   @Override
   public UnsafeRow copy() {

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRowWriters.java

@@ -19,6 +19,7 @@
 
 import org.apache.spark.sql.catalyst.InternalRow;
 import org.apache.spark.sql.types.Decimal;
+import org.apache.spark.sql.types.MapData;
 import org.apache.spark.unsafe.PlatformDependent;
 import org.apache.spark.unsafe.array.ByteArrayMethods;
 import org.apache.spark.unsafe.types.ByteArray;
@@ -215,4 +216,44 @@ public static int write(UnsafeRow target, int ordinal, int cursor, UnsafeArrayDa
       return ByteArrayMethods.roundNumberOfBytesToNearestWord(numBytes);
     }
   }
+
+  public static class MapWriter {
+
+    public static int getSize(UnsafeMapData input) {
+      // we need extra 8 bytes to store number of elements and numBytes of key array.
+      final int sizeInBytes = 4 + 4 + input.keys.getSizeInBytes() + input.values.getSizeInBytes();
+      return ByteArrayMethods.roundNumberOfBytesToNearestWord(sizeInBytes);
+    }
+
+    public static int write(UnsafeRow target, int ordinal, int cursor, UnsafeMapData input) {
+      final long offset = target.getBaseOffset() + cursor;
+      final UnsafeArrayData keyArray = input.keys;
+      final UnsafeArrayData valueArray = input.values;
+      final int keysNumBytes = keyArray.getSizeInBytes();
+      final int valuesNumBytes = valueArray.getSizeInBytes();
+      final int numBytes = 4 + 4 + keysNumBytes + valuesNumBytes;
+
+      // write the number of elements into first 4 bytes.
+      PlatformDependent.UNSAFE.putInt(target.getBaseObject(), offset, input.numElements());
+      // write the numBytes of key array into second 4 bytes.
+      PlatformDependent.UNSAFE.putInt(target.getBaseObject(), offset + 4, keysNumBytes);
+
+      // zero-out the padding bytes
+      if ((numBytes & 0x07) > 0) {
+        PlatformDependent.UNSAFE.putLong(
+          target.getBaseObject(), offset + ((numBytes >> 3) << 3), 0L);
+      }
+
+      // Write the bytes of key array to the variable length portion.
+      keyArray.writeToMemory(target.getBaseObject(), offset + 8);
+
+      // Write the bytes of value array to the variable length portion.
+      valueArray.writeToMemory(target.getBaseObject(), offset + 8 + keysNumBytes);
+
+      // Set the fixed length portion.
+      target.setLong(ordinal, (((long) cursor) << 32) | ((long) numBytes));
+
+      return ByteArrayMethods.roundNumberOfBytesToNearestWord(numBytes);
+    }
+  }
 }

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeWriters.java

@@ -35,9 +35,9 @@ public static void writeToMemory(
       int numBytes) {
 
     // zero-out the padding bytes
-    if ((numBytes & 0x07) > 0) {
-      PlatformDependent.UNSAFE.putLong(targetObject, targetOffset + ((numBytes >> 3) << 3), 0L);
-    }
+//    if ((numBytes & 0x07) > 0) {
+//      PlatformDependent.UNSAFE.putLong(targetObject, targetOffset + ((numBytes >> 3) << 3), 0L);
+//    }
 
     // Write the UnsafeData to the target memory.
     PlatformDependent.copyMemory(
@@ -171,7 +171,39 @@ public static int write(Object targetObject, long targetOffset, UnsafeArrayData
       writeToMemory(input.getBaseObject(), input.getBaseOffset(),
         targetObject, targetOffset + 4, numBytes);
 
-      return getRoundedSize(numBytes) + 4;
+      return getRoundedSize(numBytes + 4);
+    }
+  }
+
+  public static class MapWriter {
+
+    public static int getSize(UnsafeMapData input) {
+      // we need extra 8 bytes to store number of elements and numBytes of key array.
+      final int sizeInBytes = 4 + 4 + input.keys.getSizeInBytes() + input.values.getSizeInBytes();
+      return getRoundedSize(sizeInBytes);
+    }
+
+    public static int write(Object targetObject, long targetOffset, UnsafeMapData input) {
+      final UnsafeArrayData keyArray = input.keys;
+      final UnsafeArrayData valueArray = input.values;
+      final int keysNumBytes = keyArray.getSizeInBytes();
+      final int valuesNumBytes = valueArray.getSizeInBytes();
+      final int numBytes = 4 + 4 + keysNumBytes + valuesNumBytes;
+
+      // write the number of elements into first 4 bytes.
+      PlatformDependent.UNSAFE.putInt(targetObject, targetOffset, input.numElements());
+      // write the numBytes of key array into second 4 bytes.
+      PlatformDependent.UNSAFE.putInt(targetObject, targetOffset + 4, keysNumBytes);
+
+      // Write the bytes of key array to the variable length portion.
+      writeToMemory(keyArray.getBaseObject(), keyArray.getBaseOffset(),
+        targetObject, targetOffset + 8, keysNumBytes);
+
+      // Write the bytes of value array to the variable length portion.
+      writeToMemory(valueArray.getBaseObject(), valueArray.getBaseOffset(),
+        targetObject, targetOffset + 8 + keysNumBytes, valuesNumBytes);
+
+      return getRoundedSize(numBytes);
     }
   }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/FromUnsafe.scala

@@ -23,7 +23,8 @@ import org.apache.spark.sql.types._
 case class FromUnsafe(child: Expression) extends UnaryExpression
   with ExpectsInputTypes with CodegenFallback {
 
-  override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(ArrayType, StructType))
+  override def inputTypes: Seq[AbstractDataType] =
+    Seq(TypeCollection(ArrayType, StructType, MapType))
 
   override def dataType: DataType = child.dataType
 
@@ -38,19 +39,25 @@ case class FromUnsafe(child: Expression) extends UnaryExpression
       }
       new GenericInternalRow(result)
 
-    case ArrayType(elemnentType, _) =>
+    case ArrayType(elementType, _) =>
       val array = value.asInstanceOf[UnsafeArrayData]
       val length = array.numElements()
       val result = new Array[Any](length)
       var i = 0
       while (i < length) {
         if (!array.isNullAt(i)) {
-          result(i) = convert(array.get(i, elemnentType), elemnentType)
+          result(i) = convert(array.get(i, elementType), elementType)
         }
         i += 1
       }
       new GenericArrayData(result)
 
+    case MapType(kt, vt, _) =>
+      val map = value.asInstanceOf[UnsafeMapData]
+      val safeKeyArray = convert(map.keys, ArrayType(kt)).asInstanceOf[GenericArrayData]
+      val safeValueArray = convert(map.values, ArrayType(vt)).asInstanceOf[GenericArrayData]
+      new ArrayBasedMapData(safeKeyArray, safeValueArray)
+
     case _ => value
   }
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Projection.scala

@@ -156,7 +156,7 @@ object FromUnsafeProjection {
       // todo: this is quite slow, maybe remove this whole projection after remove generic getter of
       // InternalRow?
       b.dataType match {
-        case _: StructType | _: ArrayType => FromUnsafe(b)
+        case _: StructType | _: ArrayType | _: MapType => FromUnsafe(b)
         case _ => b
       }
     }))

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala

@@ -306,7 +306,8 @@ abstract class CodeGenerator[InType <: AnyRef, OutType <: AnyRef] extends Loggin
       classOf[CalendarInterval].getName,
       classOf[ArrayData].getName,
       classOf[UnsafeArrayData].getName,
-      classOf[MapData].getName
+      classOf[MapData].getName,
+      classOf[UnsafeMapData].getName
     ))
     evaluator.setExtendedClass(classOf[GeneratedClass])
     try {