Skip to content

Add test window function to ReadByFormatIntegrationTestBase #1447

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yalimu-g merged 10 commits into from
Jan 6, 2026
Merged

Add test window function to ReadByFormatIntegrationTestBase #1447

Changes from 1 commit
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
c6c7307
Add testWindowFunctionPartitionBy in ReadByFormatIntegrationTestBase …
yalimu-g Dec 1, 2025
8a116ed
Change format to AVRO to bypass the exception.
yalimu-g Dec 1, 2025
42c4211
Change format to AVRO to bypass the exception.
yalimu-g Dec 1, 2025
9d636fe
merge upstream
yalimu-g Dec 9, 2025
d0d6d61
Fix the schema converter for array type.
yalimu-g Dec 10, 2025
7960dc9
Revert the local changes
yalimu-g Dec 10, 2025
a75dea7
Format and fix tests
yalimu-g Dec 10, 2025
7d67a4a
Refactor according to gemini-code-assist.
yalimu-g Dec 16, 2025
b77b6bb
Add Parameterized tests for ARROW and AVRO format in Spark35ReadByFor…
yalimu-g Dec 18, 2025
e32af7e
Rename test for AVRO format.
yalimu-g Jan 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 34 additions & 0 deletions ...est/java/com/google/cloud/spark/bigquery/integration/ReadByFormatIntegrationTestBase.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,10 @@
package com.google.cloud.spark.bigquery.integration;

import static com.google.common.truth.Truth.assertThat;
import static org.apache.spark.sql.functions.col;
import static org.apache.spark.sql.functions.concat;
import static org.apache.spark.sql.functions.lit;
import static org.apache.spark.sql.functions.row_number;
import static org.hamcrest.CoreMatchers.is;
import static org.junit.Assume.assumeThat;
import static org.junit.Assume.assumeTrue;
Expand All @@ -42,6 +46,8 @@
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.expressions.Window;
import org.apache.spark.sql.expressions.WindowSpec;
import org.apache.spark.sql.types.DataType;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.StructField;
Expand Down Expand Up @@ -316,6 +322,34 @@ public void testTimestampNTZReadFromBigQuery() {
assertThat(row.get(0)).isEqualTo(dateTime);
}

@Test
public void testWindowFunctionPartitionBy() {
String tableName = "bigquery-public-data.ga4_obfuscated_sample_ecommerce.events_20210131";
WindowSpec windowSpec =
Window.partitionBy(concat(col("user_pseudo_id"), col("event_timestamp"), col("event_name")))
.orderBy(lit("window_ordering"));

Dataset<Row> df =
spark
.read()
.format("bigquery")
.option("table", tableName)
.option("readDataFormat", dataFormat)
.load()
.withColumn("row_num", row_number().over(windowSpec));

Dataset<Row> selectedDF =
df.select("user_pseudo_id", "event_name", "event_timestamp", "row_num");

assertThat(selectedDF.columns().length).isEqualTo(4);
assertThat(
Arrays.stream(df.schema().fields())
.filter(field -> field.name().equals("row_num"))
.count())
.isEqualTo(1);
assertThat(selectedDF.head().get(3)).isEqualTo(1);
}

static <K, V> Map<K, V> scalaMapToJavaMap(scala.collection.Map<K, V> map) {
ImmutableMap.Builder<K, V> result = ImmutableMap.<K, V>builder();
map.foreach(entry -> result.put(entry._1(), entry._2()));
Expand Down
Loading