You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by di...@apache.org on 2023/01/11 02:46:25 UTC
[flink] 02/03: [FLINK-30607][python] Support BinaryType for Pandas UDF
This is an automated email from the ASF dual-hosted git repository.
dianfu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 72c9cfdb6e8324d9982bf7d05306c9dac98a0bfd
Author: Dian Fu <di...@apache.org>
AuthorDate: Tue Jan 10 16:00:15 2023 +0800
[FLINK-30607][python] Support BinaryType for Pandas UDF
This closes #21639.
---
flink-python/pyflink/fn_execution/coders.py | 5 +-
.../pyflink/table/tests/test_pandas_udf.py | 21 +++--
flink-python/pyflink/table/types.py | 2 +
.../flink/table/runtime/arrow/ArrowUtils.java | 15 ++++
.../arrow/vectors/ArrowBinaryColumnVector.java | 48 +++++++++++
.../table/runtime/arrow/writers/BinaryWriter.java | 95 ++++++++++++++++++++++
6 files changed, 180 insertions(+), 6 deletions(-)
diff --git a/flink-python/pyflink/fn_execution/coders.py b/flink-python/pyflink/fn_execution/coders.py
index 704d278666b..190efb5dea6 100644
--- a/flink-python/pyflink/fn_execution/coders.py
+++ b/flink-python/pyflink/fn_execution/coders.py
@@ -29,7 +29,8 @@ from pyflink.common.typeinfo import TypeInformation, BasicTypeInfo, BasicType, D
ExternalTypeInfo
from pyflink.table.types import TinyIntType, SmallIntType, IntType, BigIntType, BooleanType, \
FloatType, DoubleType, VarCharType, VarBinaryType, DecimalType, DateType, TimeType, \
- LocalZonedTimestampType, RowType, RowField, to_arrow_type, TimestampType, ArrayType, MapType
+ LocalZonedTimestampType, RowType, RowField, to_arrow_type, TimestampType, ArrayType, MapType, \
+ BinaryType
try:
from pyflink.fn_execution import coder_impl_fast as coder_impl
@@ -125,6 +126,8 @@ class LengthPrefixBaseCoder(ABC):
return DoubleType(field_type.nullable)
elif field_type.type_name == flink_fn_execution_pb2.Schema.VARCHAR:
return VarCharType(0x7fffffff, field_type.nullable)
+ elif field_type.type_name == flink_fn_execution_pb2.Schema.BINARY:
+ return BinaryType(field_type.binary_info.length, field_type.nullable)
elif field_type.type_name == flink_fn_execution_pb2.Schema.VARBINARY:
return VarBinaryType(0x7fffffff, field_type.nullable)
elif field_type.type_name == flink_fn_execution_pb2.Schema.DECIMAL:
diff --git a/flink-python/pyflink/table/tests/test_pandas_udf.py b/flink-python/pyflink/table/tests/test_pandas_udf.py
index 46d9560b351..0e4d9d9a8ab 100644
--- a/flink-python/pyflink/table/tests/test_pandas_udf.py
+++ b/flink-python/pyflink/table/tests/test_pandas_udf.py
@@ -220,6 +220,14 @@ class PandasUDFITTests(object):
assert isinstance(map_param, pd.Series)
return map_param
+ @udf(result_type=DataTypes.BINARY(5), func_type="pandas")
+ def binary_func(binary_param):
+ assert isinstance(binary_param, pd.Series)
+ assert isinstance(binary_param[0], bytes), \
+ 'binary_param of wrong type %s !' % type(binary_param[0])
+ assert len(binary_param[0]) == 5
+ return binary_param
+
sink_table_ddl = """
CREATE TABLE Results_test_all_data_types(
a TINYINT,
@@ -243,7 +251,8 @@ class PandasUDFITTests(object):
s ARRAY<INT>,
t ARRAY<STRING>,
u ROW<f1 INT, f2 STRING, f3 TIMESTAMP(3), f4 ARRAY<INT>>,
- v MAP<STRING, STRING>
+ v MAP<STRING, STRING>,
+ w BINARY(5)
) WITH ('connector'='test-sink')
"""
self.t_env.execute_sql(sink_table_ddl)
@@ -255,7 +264,7 @@ class PandasUDFITTests(object):
datetime.date(2014, 9, 13), datetime.time(hour=1, minute=0, second=1),
timestamp_value, ['hello', '中文', None], [timestamp_value], [1, 2],
[['hello', '中文', None]], Row(1, 'hello', timestamp_value, [1, 2]),
- {"1": "hello", "2": "world"})],
+ {"1": "hello", "2": "world"}, bytearray(b'flink'))],
DataTypes.ROW(
[DataTypes.FIELD("a", DataTypes.TINYINT()),
DataTypes.FIELD("b", DataTypes.SMALLINT()),
@@ -278,7 +287,8 @@ class PandasUDFITTests(object):
DataTypes.FIELD("s", DataTypes.ARRAY(DataTypes.INT())),
DataTypes.FIELD("t", DataTypes.ARRAY(DataTypes.ARRAY(DataTypes.STRING()))),
DataTypes.FIELD("u", row_type),
- DataTypes.FIELD("v", map_type)]))
+ DataTypes.FIELD("v", map_type),
+ DataTypes.FIELD("w", DataTypes.BINARY(5))]))
t.select(
tinyint_func(t.a),
@@ -302,7 +312,8 @@ class PandasUDFITTests(object):
array_int_func(t.s),
nested_array_func(t.t),
row_func(t.u),
- map_func(t.v)) \
+ map_func(t.v),
+ binary_func(t.k)) \
.execute_insert("Results_test_all_data_types").wait()
actual = source_sink_utils.results()
self.assert_equals(
@@ -312,7 +323,7 @@ class PandasUDFITTests(object):
"1000000000000000000.059999999999999999, 2014-09-13, 01:00:01, "
"1970-01-02T00:00:00.123, [hello, 中文, null], [1970-01-02T00:00:00.123], "
"[1, 2], [hello, 中文, null], +I[1, hello, 1970-01-02T00:00:00.123, [1, 2]], "
- "{1=hello, 2=world}]"])
+ "{1=hello, 2=world}, [102, 108, 105, 110, 107]]"])
def test_invalid_pandas_udf(self):
diff --git a/flink-python/pyflink/table/types.py b/flink-python/pyflink/table/types.py
index 83d54629dad..c8227f66042 100644
--- a/flink-python/pyflink/table/types.py
+++ b/flink-python/pyflink/table/types.py
@@ -2281,6 +2281,8 @@ def to_arrow_type(data_type: DataType):
return pa.float64()
elif isinstance(data_type, (CharType, VarCharType)):
return pa.utf8()
+ elif isinstance(data_type, BinaryType):
+ return pa.binary(data_type.length)
elif isinstance(data_type, VarBinaryType):
return pa.binary()
elif isinstance(data_type, DecimalType):
diff --git a/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/ArrowUtils.java b/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/ArrowUtils.java
index 2d4e09e4553..bdd691b5c9d 100644
--- a/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/ArrowUtils.java
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/ArrowUtils.java
@@ -35,6 +35,7 @@ import org.apache.flink.table.operations.OutputConversionModifyOperation;
import org.apache.flink.table.runtime.arrow.sources.ArrowTableSource;
import org.apache.flink.table.runtime.arrow.vectors.ArrowArrayColumnVector;
import org.apache.flink.table.runtime.arrow.vectors.ArrowBigIntColumnVector;
+import org.apache.flink.table.runtime.arrow.vectors.ArrowBinaryColumnVector;
import org.apache.flink.table.runtime.arrow.vectors.ArrowBooleanColumnVector;
import org.apache.flink.table.runtime.arrow.vectors.ArrowDateColumnVector;
import org.apache.flink.table.runtime.arrow.vectors.ArrowDecimalColumnVector;
@@ -52,6 +53,7 @@ import org.apache.flink.table.runtime.arrow.vectors.ArrowVarCharColumnVector;
import org.apache.flink.table.runtime.arrow.writers.ArrayWriter;
import org.apache.flink.table.runtime.arrow.writers.ArrowFieldWriter;
import org.apache.flink.table.runtime.arrow.writers.BigIntWriter;
+import org.apache.flink.table.runtime.arrow.writers.BinaryWriter;
import org.apache.flink.table.runtime.arrow.writers.BooleanWriter;
import org.apache.flink.table.runtime.arrow.writers.DateWriter;
import org.apache.flink.table.runtime.arrow.writers.DecimalWriter;
@@ -69,6 +71,7 @@ import org.apache.flink.table.runtime.arrow.writers.VarCharWriter;
import org.apache.flink.table.types.DataType;
import org.apache.flink.table.types.logical.ArrayType;
import org.apache.flink.table.types.logical.BigIntType;
+import org.apache.flink.table.types.logical.BinaryType;
import org.apache.flink.table.types.logical.BooleanType;
import org.apache.flink.table.types.logical.CharType;
import org.apache.flink.table.types.logical.DateType;
@@ -103,6 +106,7 @@ import org.apache.arrow.vector.BitVector;
import org.apache.arrow.vector.DateDayVector;
import org.apache.arrow.vector.DecimalVector;
import org.apache.arrow.vector.FieldVector;
+import org.apache.arrow.vector.FixedSizeBinaryVector;
import org.apache.arrow.vector.Float4Vector;
import org.apache.arrow.vector.Float8Vector;
import org.apache.arrow.vector.IntVector;
@@ -255,6 +259,8 @@ public final class ArrowUtils {
return DoubleWriter.forRow((Float8Vector) vector);
} else if (vector instanceof VarCharVector) {
return VarCharWriter.forRow((VarCharVector) vector);
+ } else if (vector instanceof FixedSizeBinaryVector) {
+ return BinaryWriter.forRow((FixedSizeBinaryVector) vector);
} else if (vector instanceof VarBinaryVector) {
return VarBinaryWriter.forRow((VarBinaryVector) vector);
} else if (vector instanceof DecimalVector) {
@@ -328,6 +334,8 @@ public final class ArrowUtils {
return DoubleWriter.forArray((Float8Vector) vector);
} else if (vector instanceof VarCharVector) {
return VarCharWriter.forArray((VarCharVector) vector);
+ } else if (vector instanceof FixedSizeBinaryVector) {
+ return BinaryWriter.forArray((FixedSizeBinaryVector) vector);
} else if (vector instanceof VarBinaryVector) {
return VarBinaryWriter.forArray((VarBinaryVector) vector);
} else if (vector instanceof DecimalVector) {
@@ -411,6 +419,8 @@ public final class ArrowUtils {
return new ArrowDoubleColumnVector((Float8Vector) vector);
} else if (vector instanceof VarCharVector) {
return new ArrowVarCharColumnVector((VarCharVector) vector);
+ } else if (vector instanceof FixedSizeBinaryVector) {
+ return new ArrowBinaryColumnVector((FixedSizeBinaryVector) vector);
} else if (vector instanceof VarBinaryVector) {
return new ArrowVarBinaryColumnVector((VarBinaryVector) vector);
} else if (vector instanceof DecimalVector) {
@@ -715,6 +725,11 @@ public final class ArrowUtils {
return ArrowType.Utf8.INSTANCE;
}
+ @Override
+ public ArrowType visit(BinaryType varCharType) {
+ return new ArrowType.FixedSizeBinary(varCharType.getLength());
+ }
+
@Override
public ArrowType visit(VarBinaryType varCharType) {
return ArrowType.Binary.INSTANCE;
diff --git a/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/vectors/ArrowBinaryColumnVector.java b/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/vectors/ArrowBinaryColumnVector.java
new file mode 100644
index 00000000000..6a22011739a
--- /dev/null
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/vectors/ArrowBinaryColumnVector.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.flink.table.runtime.arrow.vectors;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.table.data.columnar.vector.BytesColumnVector;
+import org.apache.flink.util.Preconditions;
+
+import org.apache.arrow.vector.FixedSizeBinaryVector;
+
+/** Arrow column vector for Binary. */
+@Internal
+public final class ArrowBinaryColumnVector implements BytesColumnVector {
+
+ /** Container which is used to store the sequence of varbinary values of a column to read. */
+ private final FixedSizeBinaryVector fixedSizeBinaryVector;
+
+ public ArrowBinaryColumnVector(FixedSizeBinaryVector fixedSizeBinaryVector) {
+ this.fixedSizeBinaryVector = Preconditions.checkNotNull(fixedSizeBinaryVector);
+ }
+
+ @Override
+ public Bytes getBytes(int i) {
+ byte[] bytes = fixedSizeBinaryVector.get(i);
+ return new Bytes(bytes, 0, bytes.length);
+ }
+
+ @Override
+ public boolean isNullAt(int i) {
+ return fixedSizeBinaryVector.isNull(i);
+ }
+}
diff --git a/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/writers/BinaryWriter.java b/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/writers/BinaryWriter.java
new file mode 100644
index 00000000000..af9f4724a18
--- /dev/null
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/arrow/writers/BinaryWriter.java
@@ -0,0 +1,95 @@
+/*
+ * 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.table.runtime.arrow.writers;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.table.data.ArrayData;
+import org.apache.flink.table.data.RowData;
+
+import org.apache.arrow.vector.FixedSizeBinaryVector;
+
+/** {@link ArrowFieldWriter} for Binary. */
+@Internal
+public abstract class BinaryWriter<T> extends ArrowFieldWriter<T> {
+
+ public static BinaryWriter<RowData> forRow(FixedSizeBinaryVector fixedSizeBinaryVector) {
+ return new BinaryWriterForRow(fixedSizeBinaryVector);
+ }
+
+ public static BinaryWriter<ArrayData> forArray(FixedSizeBinaryVector fixedSizeBinaryVector) {
+ return new BinaryWriterForArray(fixedSizeBinaryVector);
+ }
+
+ // ------------------------------------------------------------------------------------------
+
+ private BinaryWriter(FixedSizeBinaryVector fixedSizeBinaryVector) {
+ super(fixedSizeBinaryVector);
+ }
+
+ abstract boolean isNullAt(T in, int ordinal);
+
+ abstract byte[] readBinary(T in, int ordinal);
+
+ @Override
+ public void doWrite(T in, int ordinal) {
+ if (isNullAt(in, ordinal)) {
+ ((FixedSizeBinaryVector) getValueVector()).setNull(getCount());
+ } else {
+ ((FixedSizeBinaryVector) getValueVector()).setSafe(getCount(), readBinary(in, ordinal));
+ }
+ }
+
+ // ------------------------------------------------------------------------------------------
+
+ /** {@link BinaryWriter} for {@link RowData} input. */
+ public static final class BinaryWriterForRow extends BinaryWriter<RowData> {
+
+ private BinaryWriterForRow(FixedSizeBinaryVector fixedSizeBinaryVector) {
+ super(fixedSizeBinaryVector);
+ }
+
+ @Override
+ boolean isNullAt(RowData in, int ordinal) {
+ return in.isNullAt(ordinal);
+ }
+
+ @Override
+ byte[] readBinary(RowData in, int ordinal) {
+ return in.getBinary(ordinal);
+ }
+ }
+
+ /** {@link BinaryWriter} for {@link ArrayData} input. */
+ public static final class BinaryWriterForArray extends BinaryWriter<ArrayData> {
+
+ private BinaryWriterForArray(FixedSizeBinaryVector fixedSizeBinaryVector) {
+ super(fixedSizeBinaryVector);
+ }
+
+ @Override
+ boolean isNullAt(ArrayData in, int ordinal) {
+ return in.isNullAt(ordinal);
+ }
+
+ @Override
+ byte[] readBinary(ArrayData in, int ordinal) {
+ return in.getBinary(ordinal);
+ }
+ }
+}