[SPARK-36670][SQL][TEST] Add FileSourceCodecSuite

core/src/main/java/org/apache/hadoop/shaded/net/jpountz/lz4/LZ4Compressor.java

@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.hadoop.shaded.net.jpountz.lz4;
+
+/**
+ * A temporary workaround for SPARK-36669. We should remove this after Hadoop 3.3.2 release
+ * which fixes the LZ4 relocation in shaded Hadoop client libraries. This does not need
+ * implement all net.jpountz.lz4.LZ4Compressor API, just the ones used by Hadoop Lz4Compressor.
+ */
+public final class LZ4Compressor {
+
+  private net.jpountz.lz4.LZ4Compressor lz4Compressor;
+
+  public LZ4Compressor(net.jpountz.lz4.LZ4Compressor lz4Compressor) {
+    this.lz4Compressor = lz4Compressor;
+  }
+
+  public void compress(java.nio.ByteBuffer src, java.nio.ByteBuffer dest) {
+    lz4Compressor.compress(src, dest);
+  }
+}
+

core/src/main/java/org/apache/hadoop/shaded/net/jpountz/lz4/LZ4Factory.java

@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.hadoop.shaded.net.jpountz.lz4;
+
+/**
+ * A temporary workaround for SPARK-36669. We should remove this after Hadoop 3.3.2 release
+ * which fixes the LZ4 relocation in shaded Hadoop client libraries. This does not need
+ * implement all net.jpountz.lz4.LZ4Factory API, just the ones used by Hadoop Lz4Compressor.
+ */
+public final class LZ4Factory {
+
+  private net.jpountz.lz4.LZ4Factory lz4Factory;
+
+  public LZ4Factory(net.jpountz.lz4.LZ4Factory lz4Factory) {
+    this.lz4Factory = lz4Factory;
+  }
+
+  public static LZ4Factory fastestInstance() {
+	return new LZ4Factory(net.jpountz.lz4.LZ4Factory.fastestInstance());
+  }
+
+  public LZ4Compressor highCompressor() {
+	return new LZ4Compressor(lz4Factory.highCompressor());
+  }
+
+  public LZ4Compressor fastCompressor() {
+	return new LZ4Compressor(lz4Factory.fastCompressor());
+  }
+
+  public LZ4SafeDecompressor safeDecompressor() {
+    return new LZ4SafeDecompressor(lz4Factory.safeDecompressor());
+  }
+}

core/src/main/java/org/apache/hadoop/shaded/net/jpountz/lz4/LZ4SafeDecompressor.java

@@ -0,0 +1,36 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.hadoop.shaded.net.jpountz.lz4;
+
+/**
+ * A temporary workaround for SPARK-36669. We should remove this after Hadoop 3.3.2 release
+ * which fixes the LZ4 relocation in shaded Hadoop client libraries. This does not need
+ * implement all net.jpountz.lz4.LZ4SafeDecompressor API, just the ones used by Hadoop
+ * Lz4Decompressor.
+ */
+public final class LZ4SafeDecompressor {
+  private net.jpountz.lz4.LZ4SafeDecompressor lz4Decompressor;
+
+  public LZ4SafeDecompressor(net.jpountz.lz4.LZ4SafeDecompressor lz4Decompressor) {
+	this.lz4Decompressor = lz4Decompressor;
+  }
+
+  public void decompress(java.nio.ByteBuffer src, java.nio.ByteBuffer dest) {
+    lz4Decompressor.decompress(src, dest);
+  }
+}

pom.xml

@@ -300,6 +300,13 @@
         <enabled>false</enabled>
       </snapshots>
     </repository>
+
+    <repository>
+      <id>jitpack.io</id>
+      <url>https://jitpack.io</url>
+      <name>Jitpack.io repository</name>
+      <!-- needed for brotli-codec -->
+    </repository>
   </repositories>
   <pluginRepositories>
     <pluginRepository>

project/SparkBuild.scala

@@ -274,7 +274,9 @@ object SparkBuild extends PomBuild {
       "gcs-maven-central-mirror" at "https://maven-central.storage-download.googleapis.com/maven2/",
       DefaultMavenRepository,
       Resolver.mavenLocal,
-      Resolver.file("ivyLocal", file(Path.userHome.absolutePath + "/.ivy2/local"))(Resolver.ivyStylePatterns)
+      Resolver.file("ivyLocal", file(Path.userHome.absolutePath + "/.ivy2/local"))(Resolver.ivyStylePatterns),
+      // needed for brotli-codec
+      "jitpack.io" at "https://jitpack.io"
     ),
     externalResolvers := resolvers.value,
     otherResolvers := SbtPomKeys.mvnLocalRepository(dotM2 => Seq(Resolver.file("dotM2", dotM2))).value,

sql/core/pom.xml

@@ -178,6 +178,17 @@
       <artifactId>htmlunit-driver</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.lz4</groupId>
+      <artifactId>lz4-java</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>com.github.rdblue</groupId>
+      <artifactId>brotli-codec</artifactId>
+      <version>0.1.1</version>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
   <build>
     <outputDirectory>target/scala-${scala.binary.version}/classes</outputDirectory>

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/DataSourceCodecTest.scala

@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.execution.datasources
+
+import org.apache.spark.sql.QueryTest
+import org.apache.spark.sql.test.SQLTestUtils
+
+abstract class DataSourceCodecTest extends QueryTest with SQLTestUtils {
+
+  protected def dataSourceName: String
+  protected val codecConfigName: String
+  protected def availableCodecs: Seq[String]
+
+  def testWithAllCodecs(name: String)(f: => Unit): Unit = {
+    for (codec <- availableCodecs) {
+      test(s"$name - data source $dataSourceName - codec: $codec") {
+        withSQLConf(codecConfigName -> codec) {
+          f
+        }
+      }
+    }
+  }
+
+  testWithAllCodecs("write and read - single partition") {
+    withTempPath { dir =>
+      testData
+        .repartition(1)
+        .write
+        .format(dataSourceName)
+        .save(dir.getCanonicalPath)
+
+      val df = spark.read.format(dataSourceName).load(dir.getCanonicalPath)
+      checkAnswer(df, testData)
+    }
+  }
+
+  testWithAllCodecs("write and read") {
+    withTempPath { dir =>
+      testData
+        .repartition(5)
+        .write
+        .format(dataSourceName)
+        .save(dir.getCanonicalPath)
+
+      val df = spark.read.format(dataSourceName).load(dir.getCanonicalPath)
+      checkAnswer(df, testData)
+    }
+  }
+}
+
+

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcCodecTestSuite.scala

@@ -0,0 +1,31 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.execution.datasources.orc
+
+import org.apache.spark.sql.execution.datasources.DataSourceCodecTest
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+
+class OrcCodecTestSuite extends DataSourceCodecTest with SharedSparkSession{
+
+  override def dataSourceName: String = "orc"
+  override val codecConfigName = SQLConf.ORC_COMPRESSION.key
+  override protected def availableCodecs = Seq("none", "uncompressed", "snappy",
+    "zlib", "zstd", "lz4", "lzo")
+}
+

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetCodecTestSuite.scala

@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.execution.datasources.parquet
+
+import org.apache.spark.sql.execution.datasources.DataSourceCodecTest
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+
+class ParquetCodecTestSuite extends DataSourceCodecTest with SharedSparkSession {
+
+  override def dataSourceName: String = "parquet"
+  override val codecConfigName = SQLConf.PARQUET_COMPRESSION.key
+  // Exclude "lzo" because it is GPL-licenced so not included in Hadoop.
+  override protected def availableCodecs: Seq[String] = Seq("none", "uncompressed", "snappy",
+    "gzip", "brotli", "zstd", "lz4")
+}
+