Expanding pyarrow buffer size more than 2GB for pandas_udf functions

[asfimport]

There is 2GB limit for data that can be passed to any pandas_udf function and the aim of this issue is to expand this limit. It's very small buffer size if we use pyspark and our goal is fitting machine learning models.

Steps to reproduce - just use following spark-submit for executing following after python function.

%sh
cd /home/zeppelin/code && \
export PYSPARK_DRIVER_PYTHON=/home/zeppelin/envs/env3/bin/python && \
export PYSPARK_PYTHON=./env3/bin/python && \
export ARROW_PRE_0_15_IPC_FORMAT=1 && \
spark-submit \
--master yarn \
--deploy-mode client \
--num-executors 5 \
--executor-cores 5 \
--driver-memory 8G \
--executor-memory 8G \
--conf spark.executor.memoryOverhead=4G \
--conf spark.driver.memoryOverhead=4G \
--archives /home/zeppelin/env3.tar.gz#env3 \
--jars "/opt/deltalake/delta-core_2.11-0.5.0.jar" \
--py-files jobs.zip,"/opt/deltalake/delta-core_2.11-0.5.0.jar" main.py \
--job temp

 

import pyspark
from pyspark.sql import functions as F, types as T
import pandas as pd

def analyze(spark):

    pdf1 = pd.DataFrame(
        [[1234567, 0.0, "abcdefghij", "2000-01-01T00:00:00.000Z"]],
        columns=['df1_c1', 'df1_c2', 'df1_c3', 'df1_c4']
    )
    df1 = spark.createDataFrame(pd.concat([pdf1 for i in range(429)]).reset_index()).drop('index')

    pdf2 = pd.DataFrame(
        [[1234567, 0.0, "abcdefghijklmno", "2000-01-01", "abcdefghijklmno", "abcdefghijklmno"]],
        columns=['df2_c1', 'df2_c2', 'df2_c3', 'df2_c4', 'df2_c5', 'df2_c6']
    )
    df2 = spark.createDataFrame(pd.concat([pdf2 for i in range(48993)]).reset_index()).drop('index')
    df3 = df1.join(df2, df1['df1_c1'] == df2['df2_c1'], how='inner')

    def myudf(df):
        import os
        os.environ["ARROW_PRE_0_15_IPC_FORMAT"] = "1"
        return df

    df4 = df3 \
        .withColumn('df1_c1', F.col('df1_c1').cast(T.IntegerType())) \
        .withColumn('df1_c2', F.col('df1_c2').cast(T.DoubleType())) \
        .withColumn('df1_c3', F.col('df1_c3').cast(T.StringType())) \
        .withColumn('df1_c4', F.col('df1_c4').cast(T.StringType())) \
        .withColumn('df2_c1', F.col('df2_c1').cast(T.IntegerType())) \
        .withColumn('df2_c2', F.col('df2_c2').cast(T.DoubleType())) \
        .withColumn('df2_c3', F.col('df2_c3').cast(T.StringType())) \
        .withColumn('df2_c4', F.col('df2_c4').cast(T.StringType())) \
        .withColumn('df2_c5', F.col('df2_c5').cast(T.StringType())) \
        .withColumn('df2_c6', F.col('df2_c6').cast(T.StringType()))
    print(df4.printSchema())

    udf = F.pandas_udf(df4.schema, F.PandasUDFType.GROUPED_MAP)(myudf)

    df5 = df4.groupBy('df1_c1').apply(udf)
    print('df5.count()', df5.count())

If you need more details please let me know.

Environment: Spark: 2.4.4

Python:
Dcycler (0.10.0)
glmnet-py (0.1.0b2)
joblib (1.0.0)
kiwisolver (1.3.1)
lightgbm (3.1.1) EPRECATION
matplotlib (3.0.3)
numpy (1.19.4)
pandas (1.1.5)
pip (9.0.3: The default format will switch to columns in the future. You can)
pyarrow 2.0.0
pyparsing (2.4.7) use --format=(legacy|columns) (or define a format=(python-dateutil (2.8.1)
pytz (202legacy|columns) in yo0.4)
scikit-learn (0.23.2)
scipy (1.5.4)
setuptools (51.0.0) ur pip.conf under the [list] section) to disable this warnsix (1.15.0)
sklearn (0.0)
threadpoolctl (2.1.0)
venv-paing. ck (0.2.0)
wheel (0.36.2)
Reporter: Dmitry Kravchuk

Original Issue Attachments:

• python env.png
• spark3 env.png
• spark3 error.png

Externally tracked issue: https://issues.apache.org/jira/browse/ARROW-4890#

_{Note: This issue was originally created as ARROW-10957. Please see the migration documentation for further details.}

[asfimport]

Dmitry Kravchuk:
[~fan_li_ya] here we go.

[asfimport]

Dmitry Kravchuk:
I've also tried this code with env spark 3.0.1/scala 2.12 and everyting went okay.

It this issue related just for spark with version lover than 3.0.0?

[asfimport]

Liya Fan / @liyafan82:
[~dishka_krauch] I guess the reason is that the metadata size for Spark 3.0.1 is smaller than 2GB?

And this is the root cause of the problem. Currently, Arrow supports data buffers larger than 2GB, but not metadata.

Could you please help investigate why is the metadata size so large?

[asfimport]

Dmitry Kravchuk:
[~fan_li_ya] I've tested pandas dataframe with size 1.72 gb with spark 2.4.4 and it fails.

Also I've tested padnas datafraem withsize 1.72 gb and 2.5 gb with spark 3.0.1 and it works.

What does actually you mean as metadata in process of serialization data between scala and python using pandas_udf? Is it spark dataframe schema?

[asfimport]

Liya Fan / @liyafan82:
Yes. I have copied the stack trace of the original exception:

java.lang.IllegalArgumentException
	at java.nio.ByteBuffer.allocate(ByteBuffer.java:334)
	at org.apache.arrow.vector.ipc.message.MessageSerializer.readMessage(MessageSerializer.java:543)
	at org.apache.arrow.vector.ipc.message.MessageChannelReader.readNext(MessageChannelReader.java:58)
	at org.apache.arrow.vector.ipc.ArrowStreamReader.readSchema(ArrowStreamReader.java:132)
	at org.apache.arrow.vector.ipc.ArrowReader.initialize(ArrowReader.java:181)
	at org.apache.arrow.vector.ipc.ArrowReader.ensureInitialized(ArrowReader.java:172)
	at org.apache.arrow.vector.ipc.ArrowReader.getVectorSchemaRoot(ArrowReader.java:65)

It can be seen that the exception was thrown when reading the schema. So I am wondering if the schema size was too large?

BTW, are you running the program in a BigEndian platform? Byte order may also casue the negative message length.

[asfimport]

Dmitry Kravchuk:
[~fan_li_ya] sorry, but I can't catch the main meaning of your question.

I have a spark dataset with size 1.72 gb. When I pass this data through pandas_udf I get this exception.

How can I give you information about schema size?

I'm using Hadoop from Hortonworks Cloudera repository.

[asfimport]

Liya Fan / @liyafan82:
[~dishka_krauch] Sorry I don't know how to help. It seems difficult to share a data set with 1.72 gb size.

Currently, Arrow supports buffer sizes larger than 2GB. So it is OK to transfer large vectors/arrays/record batches. However, for metadata (e.g. schema), the 2gb limit still exists. This is because in most applications, the metadata size is relatively small, so 2gb should be more than enough.

So my suggestion is to debug your code to see why the metadata size is so large. Is it plausible, or there are some bugs?

[asfimport]

Dmitry Kravchuk:
[~fan_li_ya] I can't understand how it is possible that metadata size (information about columns, types) can be more than 2 gb in my case. I can see than data size is 1.72 gb after convertion to pandas.

I gave you code that generates dataset above.

import pyspark
from pyspark.sql import functions as F, types as T
import pandas as pd

def analyze(spark):

    pdf1 = pd.DataFrame(
        [[1234567, 0.0, "abcdefghij", "2000-01-01T00:00:00.000Z"]],
        columns=['df1_c1', 'df1_c2', 'df1_c3', 'df1_c4']
    )
    df1 = spark.createDataFrame(pd.concat([pdf1 for i in range(429)]).reset_index()).drop('index')

    pdf2 = pd.DataFrame(
        [[1234567, 0.0, "abcdefghijklmno", "2000-01-01", "abcdefghijklmno", "abcdefghijklmno"]],
        columns=['df2_c1', 'df2_c2', 'df2_c3', 'df2_c4', 'df2_c5', 'df2_c6']
    )
    df2 = spark.createDataFrame(pd.concat([pdf2 for i in range(48993)]).reset_index()).drop('index')
    df3 = df1.join(df2, df1['df1_c1'] == df2['df2_c1'], how='inner')

    def myudf(df):
        import os
        os.environ["ARROW_PRE_0_15_IPC_FORMAT"] = "1"
        return df

    df4 = df3 \
        .withColumn('df1_c1', F.col('df1_c1').cast(T.IntegerType())) \
        .withColumn('df1_c2', F.col('df1_c2').cast(T.DoubleType())) \
        .withColumn('df1_c3', F.col('df1_c3').cast(T.StringType())) \
        .withColumn('df1_c4', F.col('df1_c4').cast(T.StringType())) \
        .withColumn('df2_c1', F.col('df2_c1').cast(T.IntegerType())) \
        .withColumn('df2_c2', F.col('df2_c2').cast(T.DoubleType())) \
        .withColumn('df2_c3', F.col('df2_c3').cast(T.StringType())) \
        .withColumn('df2_c4', F.col('df2_c4').cast(T.StringType())) \
        .withColumn('df2_c5', F.col('df2_c5').cast(T.StringType())) \
        .withColumn('df2_c6', F.col('df2_c6').cast(T.StringType()))
    print(df4.printSchema())

    udf = F.pandas_udf(df4.schema, F.PandasUDFType.GROUPED_MAP)(myudf)

    df5 = df4.groupBy('df1_c1').apply(udf)
    print('df5.count()', df5.count())

I am sure that size of df4.schema output is less than 2 gb, but size of df4.toPandas() is approximately 2 gb (1.72 gb).
May be there is some misunterstanding between us.

[asfimport]

Liya Fan / @liyafan82:
[~dishka_krauch] Thanks for the information.
I can take a look with the program you provided. However, I do not have a python or Spark environment, so it may take me quite a while to reproduce the problem.

I think the most efficient approach is that you can debug the program to see why the metadata is so large?

[asfimport]

Dmitry Kravchuk:
[~fan_li_ya] Okay, it will be easy. Just help me to find out how to get info about metadata size. Can you share some code examples?

[asfimport]

Liya Fan / @liyafan82:
[~dishka_krauch] Sorry for my late reply.
I am not farmiliar with python code base.
For Java code, you can look at TestRoundTrip#testMetadata, which contains a complete example of serializing/deserializing data schema.

[asfimport]

Dmitry Kravchuk:
[~fan_li_ya] Where should I search Java code example of TestRoundTrip#testMetadata? Can you share link?

[asfimport]

Liya Fan / @liyafan82:
Sure. Please find the test case here: https://github.com/apache/arrow/blob/master/java/vector/src/test/java/org/apache/arrow/vector/ipc/TestRoundTrip.java#L254

[asfimport]

Dmitry Kravchuk:
[~liyafan] sorry for late reply...
I have no any environment or any possible way to use Java code from your example.
Can you help me with metadata size calculation?
May be it can be done with known schema if my data?
I mean it should possible to calculate metadata size using schema columns and theirs types, am I right?

[asfimport]

Micah Kornfield / @emkornfield:
Quick question: I was under the impression that "
export ARROW_PRE_0_15_IPC_FORMAT=1
was insufficient to make this work properly.  Per [instructions|conf/spark-env.sh] it needs to be added to you spark-conf?