moving all Seq references to the Scala specific SparkSqlUtils implemntaiton so it wouldn't leak

Author: davidrabinowitz

spark-bigquery-connector-common/src/main/java/com/google/cloud/spark/bigquery/ProtobufUtils.java

@@ -22,7 +22,6 @@
 import com.google.cloud.bigquery.storage.v1.ProtoRows;
 import com.google.cloud.bigquery.storage.v1.ProtoSchema;
 import com.google.cloud.bigquery.storage.v1.ProtoSchemaConverter;
-import com.google.cloud.spark.bigquery.util.ScalaUtils;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
@@ -35,9 +34,28 @@
 import org.apache.spark.bigquery.BigNumericUDT;
 import org.apache.spark.bigquery.BigQueryDataTypes;
 import org.apache.spark.sql.Row;
+import org.apache.spark.sql.SparkSqlUtils;
 import org.apache.spark.sql.catalyst.InternalRow;
 import org.apache.spark.sql.catalyst.util.ArrayData;
-import org.apache.spark.sql.types.*;
+import org.apache.spark.sql.types.ArrayType;
+import org.apache.spark.sql.types.BinaryType;
+import org.apache.spark.sql.types.BooleanType;
+import org.apache.spark.sql.types.ByteType;
+import org.apache.spark.sql.types.DataType;
+import org.apache.spark.sql.types.DataTypes;
+import org.apache.spark.sql.types.DateType;
+import org.apache.spark.sql.types.Decimal;
+import org.apache.spark.sql.types.DecimalType;
+import org.apache.spark.sql.types.DoubleType;
+import org.apache.spark.sql.types.FloatType;
+import org.apache.spark.sql.types.IntegerType;
+import org.apache.spark.sql.types.LongType;
+import org.apache.spark.sql.types.MapType;
+import org.apache.spark.sql.types.ShortType;
+import org.apache.spark.sql.types.StringType;
+import org.apache.spark.sql.types.StructField;
+import org.apache.spark.sql.types.StructType;
+import org.apache.spark.sql.types.TimestampType;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import scala.collection.mutable.IndexedSeq;
@@ -354,7 +372,7 @@ private static Object convertSparkValueToProtoRowValue(
     if (sparkType instanceof StructType) {
       InternalRow internalRow = null;
       if (sparkValue instanceof Row) {
-        internalRow = ScalaUtils.getInstance().rowToInternalRow((Row) sparkValue);
+        internalRow = SparkSqlUtils.getInstance().rowToInternalRow((Row) sparkValue);
       } else {
         internalRow = (InternalRow) sparkValue;
       }

spark-bigquery-connector-common/src/main/java/com/google/cloud/spark/bigquery/util/Scala213Utils.java

@@ -0,0 +1,3 @@
+package com.google.cloud.spark.bigquery.util;
+
+public class ScalaVersionSupport {}

spark-bigquery-connector-common/src/main/java/com/google/cloud/spark/bigquery/util/ScalaVersionSupport.java

@@ -1,6 +1,5 @@
 package com.google.cloud.spark.bigquery.write;
 
-import com.google.cloud.spark.bigquery.util.ScalaUtils;
 import com.google.cloud.spark.bigquery.write.context.DataWriterContext;
 import com.google.cloud.spark.bigquery.write.context.DataWriterContextFactory;
 import com.google.cloud.spark.bigquery.write.context.WriterCommitMessageContext;
@@ -11,6 +10,7 @@
 import org.apache.spark.TaskContext;
 import org.apache.spark.api.java.function.Function2;
 import org.apache.spark.sql.Row;
+import org.apache.spark.sql.SparkSqlUtils;
 import org.apache.spark.sql.catalyst.InternalRow;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -42,7 +42,7 @@ public Iterator<WriterCommitMessageContext> call(Integer partitionId, Iterator<R
     try {
       while (rowIterator.hasNext()) {
         Row row = rowIterator.next();
-        InternalRow internalRow = ScalaUtils.getInstance().rowToInternalRow(row);
+        InternalRow internalRow = SparkSqlUtils.getInstance().rowToInternalRow(row);
         dataWriterContext.write(internalRow);
       }
       return Iterators.forArray(dataWriterContext.commit());

spark-bigquery-connector-common/src/main/java/com/google/cloud/spark/bigquery/write/DataSourceWriterContextPartitionHandler.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2022 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql;
+
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import org.apache.spark.sql.catalyst.InternalRow;
+import org.apache.spark.sql.catalyst.analysis.SimpleAnalyzer$;
+import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder;
+import org.apache.spark.sql.catalyst.encoders.RowEncoder;
+import org.apache.spark.sql.catalyst.expressions.Attribute;
+import org.apache.spark.sql.catalyst.expressions.AttributeReference;
+import org.apache.spark.sql.catalyst.expressions.NamedExpression;
+import org.apache.spark.sql.types.StructType;
+import scala.collection.JavaConverters;
+import scala.collection.mutable.ListBuffer;
+
+public class Scala213SparkSqlUtils extends SparkSqlUtils {
+
+  @Override
+  public boolean supportsScalaVersion(String scalaVersion) {
+    return scalaVersion.compareTo("2.13") > 0;
+  }
+
+  @Override
+  public InternalRow rowToInternalRow(Row row) {
+    return InternalRow.fromSeq(row.toSeq());
+  }
+
+  @Override
+  public ExpressionEncoder<Row> createExpressionEncoder(StructType schema) {
+    List<Attribute> attributes =
+        JavaConverters.asJavaCollection(toAttributes(schema)).stream()
+            .map(Attribute::toAttribute)
+            .collect(Collectors.toList());
+    ExpressionEncoder<Row> expressionEncoder =
+        RowEncoder.apply(schema)
+            .resolveAndBind(
+                JavaConverters.asScalaIteratorConverter(attributes.iterator()).asScala().toSeq(),
+                SimpleAnalyzer$.MODULE$);
+    return expressionEncoder;
+  }
+
+  // `toAttributes` is protected[sql] starting spark 3.2.0, so we need this call to be in the same
+  // package. Since Scala 2.13/Spark 3.3 forbids it, the implementation has been ported to Java
+  public static scala.collection.immutable.Seq<AttributeReference> toAttributes(StructType schema) {
+    List<AttributeReference> result =
+        Stream.of(schema.fields())
+            .map(
+                field ->
+                    new AttributeReference(
+                        field.name(),
+                        field.dataType(),
+                        field.nullable(),
+                        field.metadata(),
+                        NamedExpression.newExprId(),
+                        new ListBuffer<String>().toSeq()))
+            .collect(Collectors.toList());
+    return JavaConverters.asScalaBuffer(result).toSeq();
+  }
+}

spark-bigquery-connector-common/src/main/java/org/apache/spark/sql/Scala213SparkSqlUtils.java

@@ -13,21 +13,21 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.google.cloud.spark.bigquery.util;
+package org.apache.spark.sql;
 
 import com.google.common.collect.Streams;
 import java.util.ServiceLoader;
-import org.apache.spark.sql.Row;
 import org.apache.spark.sql.catalyst.InternalRow;
-import scala.collection.immutable.Seq;
+import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder;
+import org.apache.spark.sql.types.StructType;
 
-public abstract class ScalaUtils {
-  private static ScalaUtils instance;
+public abstract class SparkSqlUtils {
+  private static SparkSqlUtils instance;
 
-  public static ScalaUtils getInstance() {
+  public static SparkSqlUtils getInstance() {
     String scalaVersion = scala.util.Properties.versionNumberString();
     if (instance == null) {
-      ServiceLoader<ScalaUtils> serviceLoader = ServiceLoader.load(ScalaUtils.class);
+      ServiceLoader<SparkSqlUtils> serviceLoader = ServiceLoader.load(SparkSqlUtils.class);
       instance =
           Streams.stream(serviceLoader.iterator())
               .filter(s -> s.supportsScalaVersion(scalaVersion))
@@ -37,14 +37,14 @@ public static ScalaUtils getInstance() {
                       new IllegalArgumentException(
                           String.format(
                               "Could not load instance of [%], please check the META-INF/services directory in the connector's jar",
-                              ScalaUtils.class.getCanonicalName())));
+                              SparkSqlUtils.class.getCanonicalName())));
     }
     return instance;
   }
 
   public abstract boolean supportsScalaVersion(String scalaVersion);
 
-  public abstract <T> Seq<T> emptySeq(Class<T> clazz);
-
   public abstract InternalRow rowToInternalRow(Row row);
+
+  public abstract ExpressionEncoder<Row> createExpressionEncoder(StructType schema);
 }

spark-bigquery-connector-common/src/main/java/com/google/cloud/spark/bigquery/util/ScalaUtils.java renamed to spark-bigquery-connector-common/src/main/java/org/apache/spark/sql/SparkSqlUtils.java

@@ -0,0 +1 @@
+org.apache.spark.sql.Scala213SparkSqlUtils

spark-bigquery-connector-common/src/main/resources/META-INF/services/com.google.cloud.spark.bigquery.util.ScalaUtils

@@ -13,33 +13,23 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.google.cloud.spark.bigquery.util;
+package org.apache.spark.sql;
 
 import static com.google.common.truth.Truth.assertThat;
 
-import org.apache.spark.sql.Row;
 import org.apache.spark.sql.catalyst.InternalRow;
 import org.apache.spark.sql.catalyst.expressions.GenericRow;
 import org.apache.spark.unsafe.types.UTF8String;
 import org.junit.Test;
-import scala.collection.immutable.Seq;
 
-public class Scala213UtilsTest {
-
-  @Test
-  public void testRowToSeq() throws Exception {
-    ScalaUtils su = ScalaUtils.getInstance();
-    assertThat(su).isInstanceOf(Scala213Utils.class);
-    Seq<Object> seq = su.emptySeq(Object.class);
-    assertThat(seq.size()).isEqualTo(0);
-  }
+public class Scala213SparkSqlUtilsTest {
 
   @Test
   public void testRowToInternalRow() throws Exception {
-    ScalaUtils su = ScalaUtils.getInstance();
-    assertThat(su).isInstanceOf(Scala213Utils.class);
+    SparkSqlUtils ssu = SparkSqlUtils.getInstance();
+    assertThat(ssu).isInstanceOf(Scala213SparkSqlUtils.class);
     Row row = new GenericRow(new Object[] {UTF8String.fromString("a"), 1});
-    InternalRow internalRow = su.rowToInternalRow(row);
+    InternalRow internalRow = ssu.rowToInternalRow(row);
     assertThat(internalRow.numFields()).isEqualTo(2);
     assertThat(internalRow.getString(0).toString()).isEqualTo("a");
     assertThat(internalRow.getInt(1)).isEqualTo(1);

spark-bigquery-connector-common/src/main/resources/META-INF/services/org.apache.spark.sql.SparkSqlUtils

@@ -15,7 +15,6 @@
  */
 package com.google.cloud.spark.bigquery;
 
-import com.google.cloud.spark.bigquery.util.ScalaUtils;
 import com.google.common.base.Preconditions;
 import org.apache.avro.Conversions;
 import org.apache.avro.LogicalTypes;
@@ -24,6 +23,7 @@
 import org.apache.avro.generic.GenericData;
 import org.apache.avro.util.Utf8;
 import org.apache.spark.sql.Row;
+import org.apache.spark.sql.SparkSqlUtils;
 import org.apache.spark.sql.catalyst.InternalRow;
 import org.apache.spark.sql.catalyst.expressions.SpecializedGetters;
 import org.apache.spark.sql.catalyst.expressions.UnsafeArrayData;
@@ -306,7 +306,7 @@ static Converter createConverterFor(DataType sparkType, Schema avroType) {
         } else {
           Object obj = getter.get(ordinal, /* unused */ null);
           if (obj instanceof Row) {
-            internalRow = ScalaUtils.getInstance().rowToInternalRow((Row) obj);
+            internalRow = SparkSqlUtils.getInstance().rowToInternalRow((Row) obj);
           } else {
             internalRow = (InternalRow) obj;
           }