Using row_number().over(Window) leads to Arrow exception in 4.0 preview (0.43.1-preview)

[mariussoutier]

My code is working fine with Spark 3.5.2 and 0.43.1 connector. I'm trying to upgrade to Spark 4.0 and I'm hitting this error:

WARN org.apache.spark.scheduler.TaskSetManager Lost task 0.0 in stage 0.0 (TID 0) (executor driver): java.lang.IllegalArgumentException: should have as many children as in the schema: found 0 expected 15
	at com.google.cloud.spark.bigquery.repackaged.org.apache.arrow.util.Preconditions.checkArgument(Preconditions.java:255)
	at com.google.cloud.spark.bigquery.repackaged.org.apache.arrow.vector.VectorLoader.loadBuffers(VectorLoader.java:153)
	at com.google.cloud.spark.bigquery.repackaged.org.apache.arrow.vector.VectorLoader.load(VectorLoader.java:88)
	at com.google.cloud.spark.bigquery.repackaged.org.apache.arrow.vector.ipc.ArrowReader.loadRecordBatch(ArrowReader.java:213)
	at com.google.cloud.spark.bigquery.repackaged.org.apache.arrow.vector.ipc.ArrowStreamReader.loadNextBatch(ArrowStreamReader.java:161)
	at com.google.cloud.spark.bigquery.v2.context.ArrowColumnBatchPartitionReaderContext$SimpleAdapter.loadNextBatch(ArrowColumnBatchPartitionReaderContext.java:77)
	at com.google.cloud.spark.bigquery.v2.context.ArrowColumnBatchPartitionReaderContext.next(ArrowColumnBatchPartitionReaderContext.java:252)
	at com.google.cloud.spark.bigquery.v2.BigQueryPartitionReader.next(BigQueryPartitionReader.java:32)

I've narrowed it down to this statement:

.withColumn("row_num", row_number().over(
  Window
    .partitionBy(concat(col("user_pseudo_id"), col("event_timestamp"), col("event_name")))
    .orderBy(lit("window_ordering"))
))

Without it, the code runs fine.

[mariussoutier]

Hm, but the 0.43.1 connector seems to work with Spark 4 as well. I thought I had to explicitly use the preview connector.

[davidrabinowitz]

Hi @mariussoutier , what do you mean by the 0.43.1 connector seems to work with Spark 4 as well ? which connector is that?

[mariussoutier]

I'm including spark-bigquery-with-dependencies_2.13, version 0.43.1, and it seems to work fine. My assumption was that I have to use the preview connector (spark-4.0-bigquery:0.43.1-preview).

[mariussoutier]

@davidrabinowitz ping

[yalimu-g]

Hi @mariussoutier,
Can you provide more information? To clarify, you are using com.google.cloud.spark:spark-bigquery-with-dependencies_2.13:0.43.1 and it works. When you try to use gs://spark-lib/bigquery/spark-4.0-bigquery-0.43.1-preview.jar you are running into the IllegalArgumentException. Are you running on Scala or Python and what kind of dataset that you suspect would reproduce the issue?

[davidrabinowitz]

Hi @mariussoutier please have a look at PR #1439 that tries to reproduce the issue. Is it possible your issue is schema related?

[mariussoutier]

you are using com.google.cloud.spark:spark-bigquery-with-dependencies_2.13:0.43.1 and it works.

Yes

When you try to use gs://spark-lib/bigquery/spark-4.0-bigquery-0.43.1-preview.jar you are running into the IllegalArgumentException.

Yes

Are you running on Scala or Python and what kind of dataset that you suspect would reproduce the issue?

Scala, we are processing GA4 raw data

[mariussoutier]

Hi @mariussoutier please have a look at PR #1439 that tries to reproduce the issue. Is it possible your issue is schema related?

I think that's unlikely, as the schema always contains the same number of columns. I don't understand why it should suddently drop to 0 columns?

[davidrabinowitz]

Hi @mariussoutier , we are trying to understand the difference between your schema and the schema that we used in PR #1439, as the PR is not able to reproduce the issue.

You are also welcome to open an issue via the Google Cloud console or your account manager in order to securely share the schema

[yalimu-g]

Hi @mariussoutier , I am able to reproduce the exception in the new PR #1447.

I noticed the exception only happens when I select the column device which has deeply nested struct type. Notice in the two tests(testWindowFunctionPartitionBy and testWindowFunctionPartitionBy_withAVRO): without device column using the same window function or with data format AVRO and device column, both works.

This is because the Arrow vector loader didn't handle the struct type correctly when mapping from Bigquery to Spark. Please let me know if use AVRO format will help you mitigate the issue.

[mariussoutier]

We are not fully extracting device, only col("device").getItem("category").

Edit: You are right, even col("device").getItem("category") triggers this. Accessing any complex object (like geo, traffic_source, collected_traffic_source, ...) triggers this error.

event_params on the other hand doesn't seem affected. Example:
.withColumn("campaign_id", col("event_params").getItem("campaign_id").getItem("int_value"))

When switching to Avro, I'm encountering this error:
java.lang.ClassCastException: class org.apache.spark.sql.types.MapType cannot be cast to class org.apache.spark.sql.types.ArrayType

Arrow is also more efficient than Avro, and since it works in the previous connector, this should be a regression somewhere?