Arrow Vector OversizedAllocationException when processing large batches

[XorSum]

Describe the bug

Large batches after join operations exceed Arrow Vector 2GB memory limit, causing OversizedAllocationException.

To Reproduce

./bin/spark-sql \
  --conf spark.sql.autoBroadcastJoinThreshold=-1 \
  --conf spark.sql.adaptive.autoBroadcastJoinThreshold=-1

create table tmp_t1(a int, b int) stored as orc;
with g1 as (select id as a from range(1)),
     g2 as (select id as b from range(10000))
insert overwrite tmp_t1 select g1.a, g2.b from from g1 join g2;

create table tmp_t2(a int, b int) stored as orc;
with g1 as (select id as a from range(1)),
     g2 as (select id as b from range(10000))
insert overwrite tmp_t2 select g1.a, g2.b from from g1 join g2;

select s, count(1) as cnt
from (select concat(
                     cast(date_add('2010-01-01', t1.b) as string),
                     cast(date_add('2010-01-02', t2.b) as string)
             ) as s
      from tmp_t1 t1 join tmp_t2 t2 on t1.a = t2.a)
group by s
order by cnt
limit 100;

Exception in thread "auron native task 0.0 in stage 4.0 (TID 10)" auron.org.apache.arrow.vector.util.OversizedAllocationException: Memory required for vector is (2147483648), which is overflow or more than max allowed (2147483647). You could consider using LargeVarCharVector/LargeVarBinaryVector for large strings/large bytes types
	at auron.org.apache.arrow.vector.BaseVariableWidthVector.checkDataBufferSize(BaseVariableWidthVector.java:465)
	at auron.org.apache.arrow.vector.BaseVariableWidthVector.reallocDataBuffer(BaseVariableWidthVector.java:574)
	at auron.org.apache.arrow.vector.BaseVariableWidthVector.handleSafe(BaseVariableWidthVector.java:1344)
	at auron.org.apache.arrow.vector.BaseVariableWidthVector.setSafe(BaseVariableWidthVector.java:1178)
	at org.apache.spark.sql.execution.auron.arrowio.util.StringWriter.setValue(ArrowWriter.scala:247)
	at org.apache.spark.sql.execution.auron.arrowio.util.ArrowFieldWriter.write(ArrowWriter.scala:126)
	at org.apache.spark.sql.execution.auron.arrowio.util.ArrowWriter.write(ArrowWriter.scala:97)
	at org.apache.auron.spark.sql.SparkAuronUDFWrapperContext.$anonfun$eval$5(SparkAuronUDFWrapperContext.scala:78)
	at org.apache.auron.spark.sql.SparkAuronUDFWrapperContext.$anonfun$eval$5$adapted(SparkAuronUDFWrapperContext.scala:76)
	at scala.collection.Iterator.foreach(Iterator.scala:943)
	at scala.collection.Iterator.foreach$(Iterator.scala:943)
	at scala.collection.AbstractIterator.foreach(Iterator.scala:1431)
	at org.apache.auron.spark.sql.SparkAuronUDFWrapperContext.$anonfun$eval$4(SparkAuronUDFWrapperContext.scala:76)
	at org.apache.auron.spark.sql.SparkAuronUDFWrapperContext.$anonfun$eval$4$adapted(SparkAuronUDFWrapperContext.scala:69)
	at org.apache.spark.sql.auron.util.Using$.$anonfun$resources$9(Using.scala:395)
	at org.apache.spark.sql.auron.util.Using$.resource(Using.scala:273)
	at org.apache.spark.sql.auron.util.Using$.$anonfun$resources$8(Using.scala:394)
	at org.apache.spark.sql.auron.util.Using$.resource(Using.scala:273)
	at org.apache.spark.sql.auron.util.Using$.$anonfun$resources$7(Using.scala:393)
	at org.apache.spark.sql.auron.util.Using$.resource(Using.scala:273)
	at org.apache.spark.sql.auron.util.Using$.$anonfun$resources$6(Using.scala:392)
	at org.apache.spark.sql.auron.util.Using$.resource(Using.scala:273)
	at org.apache.spark.sql.auron.util.Using$.resources(Using.scala:391)
	at org.apache.auron.spark.sql.SparkAuronUDFWrapperContext.eval(SparkAuronUDFWrapperContext.scala:69)

Expected behavior

Query should execute without memory allocation errors.

Screenshots

Additional context

[xumingming]

@XorSum It is indeed a bug.

But your test sql here is not little extreme that tmp_t1.a and tmp_t2.a have only one value: 0, then all the records from both table (10000 x 10000) matches together.

Can you share your real scenario you met in real life query?

[XorSum]

@XorSum It is indeed a bug.

But your test sql here is not little extreme that tmp_t1.a and tmp_t2.a have only one value: 0, then all the records from both table (10000 x 10000) matches together.

Can you share your real scenario you met in real life query?

In our cluster, there are indeed some SQL queries where extreme duplication of join keys leads to a significant data volume. This bug was initially discovered in such SQL queries. These scenarios typically arise from two causes:

• Unintentional – where data engineers accidentally miss certain join conditions during development.
• Intentional – where such a join pattern is designed to meet specific business requirements.

Since it is not easy to distinguish between unintentional and intentional causes, ensuring these SQL queries run successfully whenever possible will enhance Auron’s robustness, resilience, and user experience.

[xumingming]

Since it is not easy to distinguish between unintentional and intentional causes, ensuring these SQL queries run successfully whenever possible will enhance Auron’s robustness, resilience, and user experience.

Totally agree with you. After some digging I find the root cause of this issue is:

auron/native-engine/datafusion-ext-plans/src/joins/smj/full_join.rs

Lines 163 to 166 in 027cffa

	for (&lidx, &ridx) in equal_lindices.iter().cartesian_product(&equal_rindices) {
	self.lindices.push(lidx);
	self.rindices.push(ridx);
	}

It is accumulating huge batches here without checking batchSize, I think we should fix it here.

[XorSum]

It is accumulating huge batches here without checking batchSize, I think we should fix it here.

Thanks for digging out the root cause! I’ve created another PR to fix it. Please take a look.