Support BigQuery dialect in Beam SQL (#35472)

* Support BigQuery dialect in Beam SQL

* fixed the tests

* switched to setCalciteConnectionProperties

Author: liferoad

sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamSqlPipelineOptions.java

@@ -17,6 +17,7 @@
  */
 package org.apache.beam.sdk.extensions.sql.impl;
 
+import java.util.Map;
 import org.apache.beam.sdk.options.Default;
 import org.apache.beam.sdk.options.Description;
 import org.apache.beam.sdk.options.PipelineOptions;
@@ -30,6 +31,12 @@ public interface BeamSqlPipelineOptions extends PipelineOptions {
 
   void setPlannerName(String className);
 
+  @Description(
+      "The properties of calcite connection. For example <('fun','bigquery'), ('lex','big_query')>")
+  Map<String, String> getCalciteConnectionProperties();
+
+  void setCalciteConnectionProperties(Map<String, String> properties);
+
   @Description(
       "Default timezone for ZetaSQL analyzer; has no effect on CalciteSQL. Allow tz database "
           + "timezone names in https://en.wikipedia.org/wiki/List_of_tz_database_time_zones")

sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/JdbcDriver.java

@@ -25,6 +25,7 @@
 import java.sql.SQLException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Properties;
 import java.util.function.Consumer;
 import org.apache.beam.sdk.extensions.sql.SqlTransform;
@@ -165,6 +166,13 @@ private static JdbcConnection getConnection(PipelineOptions options) {
     Properties properties = new Properties();
     properties.setProperty(
         SCHEMA_FACTORY.camelName(), BeamCalciteSchemaFactory.Empty.class.getName());
+    BeamSqlPipelineOptions sqlOptions = options.as(BeamSqlPipelineOptions.class);
+    if (sqlOptions != null) {
+      Map<String, String> calciteConnectionProperties = sqlOptions.getCalciteConnectionProperties();
+      if (calciteConnectionProperties != null) {
+        properties.putAll(calciteConnectionProperties);
+      }
+    }
     JdbcConnection connection;
     try {
       connection = (JdbcConnection) INSTANCE.connect(CONNECT_STRING_PREFIX, properties);

sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/JdbcDriverTest.java

@@ -472,6 +472,18 @@ private Row row(Schema schema, Object... values) {
     return Row.withSchema(schema).addValues(values).build();
   }
 
+  @Test
+  public void testCalciteConnectionProperties_BigQuery() throws Exception {
+    BeamSqlPipelineOptions options =
+        PipelineOptionsFactory.create().as(BeamSqlPipelineOptions.class);
+    options.setCalciteConnectionProperties(ImmutableMap.of("fun", "BIGQUERY", "lex", "BIG_QUERY"));
+    CalciteConnection connection = JdbcDriver.connect(BOUNDED_TABLE, options);
+    assertEquals(
+        "BIGQUERY", connection.getProperties().getProperty("fun", "CALCITE").toUpperCase());
+    assertEquals(
+        "BIG_QUERY", connection.getProperties().getProperty("lex", "CALCITE").toUpperCase());
+  }
+
   @Test
   public void testInternalConnect_setDirectRunner() throws Exception {
     CalciteConnection connection =
@@ -509,4 +521,53 @@ public void testInternalConnect_driverManagerDifferentProtocol() throws Exceptio
 
     DriverManager.getConnection("jdbc:baaaaaad");
   }
+
+  @Test
+  public void testBigQueryDialect() throws Exception {
+    ReadOnlyTableProvider tableProvider =
+        new ReadOnlyTableProvider(
+            "test",
+            ImmutableMap.of(
+                "test",
+                TestUnboundedTable.of(
+                        Schema.FieldType.INT32, "order_id",
+                        Schema.FieldType.INT32, "site_id",
+                        Schema.FieldType.INT32, "price",
+                        Schema.FieldType.STRING, "label",
+                        Schema.FieldType.STRING, "date",
+                        Schema.FieldType.DATETIME, "order_time")
+                    .timestampColumnIndex(5)
+                    .addRows(
+                        Duration.ZERO,
+                        1,
+                        1,
+                        1,
+                        "test-1",
+                        "2025-06-30",
+                        FIRST_DATE,
+                        1,
+                        2,
+                        6,
+                        "test-2",
+                        "2025-06-30",
+                        FIRST_DATE)));
+    BeamSqlPipelineOptions options =
+        PipelineOptionsFactory.create().as(BeamSqlPipelineOptions.class);
+    options.setCalciteConnectionProperties(ImmutableMap.of("fun", "bigquery"));
+
+    CalciteConnection connection = JdbcDriver.connect(tableProvider, options);
+    assertEquals("bigquery", connection.getProperties().getProperty("fun", "CALCITE"));
+    Statement statement = connection.createStatement();
+
+    ResultSet resultSet1 =
+        statement.executeQuery(
+            "SELECT SUBSTRING(`label`,2), SUBSTR(`label`, 2,4) FROM test LIMIT 1");
+    assertTrue(resultSet1.next());
+    assertFalse(resultSet1.next());
+
+    ResultSet resultSet2 = statement.executeQuery("SELECT * FROM test LIMIT 2");
+    assertTrue(resultSet2.next());
+    assertTrue(resultSet2.next());
+    assertFalse(resultSet2.next());
+  }
 }