| /* |
| * 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.flink.ml.python; |
| |
| import org.apache.flink.api.common.typeinfo.BasicArrayTypeInfo; |
| import org.apache.flink.api.common.typeinfo.BasicTypeInfo; |
| import org.apache.flink.api.common.typeinfo.PrimitiveArrayTypeInfo; |
| import org.apache.flink.api.common.typeinfo.SqlTimeTypeInfo; |
| import org.apache.flink.api.common.typeinfo.TypeInformation; |
| import org.apache.flink.api.common.typeutils.TypeSerializer; |
| import org.apache.flink.api.java.tuple.Tuple; |
| import org.apache.flink.api.java.typeutils.ListTypeInfo; |
| import org.apache.flink.api.java.typeutils.MapTypeInfo; |
| import org.apache.flink.api.java.typeutils.ObjectArrayTypeInfo; |
| import org.apache.flink.api.java.typeutils.RowTypeInfo; |
| import org.apache.flink.api.java.typeutils.TupleTypeInfo; |
| import org.apache.flink.api.java.typeutils.TupleTypeInfoBase; |
| import org.apache.flink.api.python.shaded.net.razorvine.pickle.Pickler; |
| import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos; |
| import org.apache.flink.core.memory.DataOutputViewStreamWrapper; |
| import org.apache.flink.streaming.api.typeinfo.python.PickledByteArrayTypeInfo; |
| import org.apache.flink.types.Row; |
| import org.apache.flink.util.Preconditions; |
| |
| import java.io.IOException; |
| import java.sql.Date; |
| import java.sql.Time; |
| import java.util.ArrayList; |
| import java.util.List; |
| import java.util.Map; |
| |
| import static org.apache.flink.api.common.typeinfo.BasicTypeInfo.FLOAT_TYPE_INFO; |
| import static org.apache.flink.api.common.typeinfo.SqlTimeTypeInfo.DATE; |
| import static org.apache.flink.api.common.typeinfo.SqlTimeTypeInfo.TIME; |
| |
| /** |
| * Utility functions used to override PyFlink methods to provide a temporary solution for certain |
| * bugs. |
| */ |
| // TODO: Remove this class after Flink ML depends on a Flink version with FLINK-30168 and |
| // FLINK-29477 fixed. |
| public class PythonBridgeUtils { |
| public static Object getPickledBytesFromJavaObject(Object obj, TypeInformation<?> dataType) |
| throws IOException { |
| Pickler pickler = new Pickler(); |
| |
| // triggers the initialization process |
| org.apache.flink.api.common.python.PythonBridgeUtils.getPickledBytesFromJavaObject( |
| null, null); |
| |
| if (obj == null) { |
| return pickler.dumps(null); |
| } else { |
| if (dataType instanceof SqlTimeTypeInfo) { |
| SqlTimeTypeInfo<?> sqlTimeTypeInfo = |
| SqlTimeTypeInfo.getInfoFor(dataType.getTypeClass()); |
| if (sqlTimeTypeInfo == DATE) { |
| return pickler.dumps(((Date) obj).toLocalDate().toEpochDay()); |
| } else if (sqlTimeTypeInfo == TIME) { |
| return pickler.dumps(((Time) obj).toLocalTime().toNanoOfDay() / 1000); |
| } |
| } else if (dataType instanceof RowTypeInfo || dataType instanceof TupleTypeInfo) { |
| TypeInformation<?>[] fieldTypes = ((TupleTypeInfoBase<?>) dataType).getFieldTypes(); |
| int arity = |
| dataType instanceof RowTypeInfo |
| ? ((Row) obj).getArity() |
| : ((Tuple) obj).getArity(); |
| |
| List<Object> fieldBytes = new ArrayList<>(arity + 1); |
| if (dataType instanceof RowTypeInfo) { |
| fieldBytes.add(new byte[] {((Row) obj).getKind().toByteValue()}); |
| } |
| for (int i = 0; i < arity; i++) { |
| Object field = |
| dataType instanceof RowTypeInfo |
| ? ((Row) obj).getField(i) |
| : ((Tuple) obj).getField(i); |
| fieldBytes.add(getPickledBytesFromJavaObject(field, fieldTypes[i])); |
| } |
| return fieldBytes; |
| } else if (dataType instanceof BasicArrayTypeInfo |
| || dataType instanceof PrimitiveArrayTypeInfo |
| || dataType instanceof ObjectArrayTypeInfo) { |
| Object[] objects; |
| TypeInformation<?> elementType; |
| if (dataType instanceof BasicArrayTypeInfo) { |
| objects = (Object[]) obj; |
| elementType = ((BasicArrayTypeInfo<?, ?>) dataType).getComponentInfo(); |
| } else if (dataType instanceof PrimitiveArrayTypeInfo) { |
| objects = primitiveArrayConverter(obj, dataType); |
| elementType = ((PrimitiveArrayTypeInfo<?>) dataType).getComponentType(); |
| } else { |
| objects = (Object[]) obj; |
| elementType = ((ObjectArrayTypeInfo<?, ?>) dataType).getComponentInfo(); |
| } |
| |
| List<Object> serializedElements = new ArrayList<>(objects.length); |
| |
| for (Object object : objects) { |
| serializedElements.add(getPickledBytesFromJavaObject(object, elementType)); |
| } |
| return pickler.dumps(serializedElements); |
| } else if (dataType instanceof MapTypeInfo) { |
| List<List<Object>> serializedMapKV = new ArrayList<>(2); |
| Map<Object, Object> mapObj = (Map) obj; |
| List<Object> keyBytesList = new ArrayList<>(mapObj.size()); |
| List<Object> valueBytesList = new ArrayList<>(mapObj.size()); |
| for (Map.Entry entry : mapObj.entrySet()) { |
| keyBytesList.add( |
| getPickledBytesFromJavaObject( |
| entry.getKey(), ((MapTypeInfo) dataType).getKeyTypeInfo())); |
| valueBytesList.add( |
| getPickledBytesFromJavaObject( |
| entry.getValue(), ((MapTypeInfo) dataType).getValueTypeInfo())); |
| } |
| serializedMapKV.add(keyBytesList); |
| serializedMapKV.add(valueBytesList); |
| return pickler.dumps(serializedMapKV); |
| } else if (dataType instanceof ListTypeInfo) { |
| List objects = (List) obj; |
| List<Object> serializedElements = new ArrayList<>(objects.size()); |
| TypeInformation elementType = ((ListTypeInfo) dataType).getElementTypeInfo(); |
| for (Object object : objects) { |
| serializedElements.add(getPickledBytesFromJavaObject(object, elementType)); |
| } |
| return pickler.dumps(serializedElements); |
| } |
| if (dataType instanceof BasicTypeInfo |
| && BasicTypeInfo.getInfoFor(dataType.getTypeClass()) == FLOAT_TYPE_INFO) { |
| // Serialization of float type with pickler loses precision. |
| return pickler.dumps(String.valueOf(obj)); |
| } else if (dataType instanceof PickledByteArrayTypeInfo |
| || dataType instanceof BasicTypeInfo) { |
| return pickler.dumps(obj); |
| } else { |
| // other typeinfos will use the corresponding serializer to serialize data. |
| TypeSerializer serializer = dataType.createSerializer(null); |
| ByteArrayOutputStreamWithPos baos = new ByteArrayOutputStreamWithPos(); |
| DataOutputViewStreamWrapper baosWrapper = new DataOutputViewStreamWrapper(baos); |
| serializer.serialize(obj, baosWrapper); |
| return pickler.dumps(baos.toByteArray()); |
| } |
| } |
| } |
| |
| private static Object[] primitiveArrayConverter( |
| Object array, TypeInformation<?> arrayTypeInfo) { |
| Preconditions.checkArgument(arrayTypeInfo instanceof PrimitiveArrayTypeInfo); |
| Preconditions.checkArgument(array.getClass().isArray()); |
| Object[] objects; |
| if (PrimitiveArrayTypeInfo.BOOLEAN_PRIMITIVE_ARRAY_TYPE_INFO.equals(arrayTypeInfo)) { |
| boolean[] booleans = (boolean[]) array; |
| objects = new Object[booleans.length]; |
| for (int i = 0; i < booleans.length; i++) { |
| objects[i] = booleans[i]; |
| } |
| } else if (PrimitiveArrayTypeInfo.BYTE_PRIMITIVE_ARRAY_TYPE_INFO.equals(arrayTypeInfo)) { |
| byte[] bytes = (byte[]) array; |
| objects = new Object[bytes.length]; |
| for (int i = 0; i < bytes.length; i++) { |
| objects[i] = bytes[i]; |
| } |
| } else if (PrimitiveArrayTypeInfo.SHORT_PRIMITIVE_ARRAY_TYPE_INFO.equals(arrayTypeInfo)) { |
| short[] shorts = (short[]) array; |
| objects = new Object[shorts.length]; |
| for (int i = 0; i < shorts.length; i++) { |
| objects[i] = shorts[i]; |
| } |
| } else if (PrimitiveArrayTypeInfo.INT_PRIMITIVE_ARRAY_TYPE_INFO.equals(arrayTypeInfo)) { |
| int[] ints = (int[]) array; |
| objects = new Object[ints.length]; |
| for (int i = 0; i < ints.length; i++) { |
| objects[i] = ints[i]; |
| } |
| } else if (PrimitiveArrayTypeInfo.LONG_PRIMITIVE_ARRAY_TYPE_INFO.equals(arrayTypeInfo)) { |
| long[] longs = (long[]) array; |
| objects = new Object[longs.length]; |
| for (int i = 0; i < longs.length; i++) { |
| objects[i] = longs[i]; |
| } |
| } else if (PrimitiveArrayTypeInfo.FLOAT_PRIMITIVE_ARRAY_TYPE_INFO.equals(arrayTypeInfo)) { |
| float[] floats = (float[]) array; |
| objects = new Object[floats.length]; |
| for (int i = 0; i < floats.length; i++) { |
| objects[i] = floats[i]; |
| } |
| } else if (PrimitiveArrayTypeInfo.DOUBLE_PRIMITIVE_ARRAY_TYPE_INFO.equals(arrayTypeInfo)) { |
| double[] doubles = (double[]) array; |
| objects = new Object[doubles.length]; |
| for (int i = 0; i < doubles.length; i++) { |
| objects[i] = doubles[i]; |
| } |
| } else if (PrimitiveArrayTypeInfo.CHAR_PRIMITIVE_ARRAY_TYPE_INFO.equals(arrayTypeInfo)) { |
| char[] chars = (char[]) array; |
| objects = new Object[chars.length]; |
| for (int i = 0; i < chars.length; i++) { |
| objects[i] = chars[i]; |
| } |
| } else { |
| throw new UnsupportedOperationException( |
| String.format( |
| "Primitive array of %s is not supported in PyFlink yet", |
| ((PrimitiveArrayTypeInfo<?>) arrayTypeInfo) |
| .getComponentType() |
| .getTypeClass() |
| .getSimpleName())); |
| } |
| return objects; |
| } |
| } |
