Spark-3.0: support CALL procedure for rewrite_data_files

spark/v3.0/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRewriteDataFilesProcedure.java

@@ -0,0 +1,274 @@
+/*
+ * 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.iceberg.spark.extensions;
+
+import java.util.List;
+import java.util.Map;
+import java.util.stream.IntStream;
+import org.apache.iceberg.AssertHelpers;
+import org.apache.iceberg.exceptions.ValidationException;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.apache.iceberg.spark.source.ThreeColumnRecord;
+import org.apache.spark.sql.AnalysisException;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.Row;
+import org.apache.spark.sql.catalyst.analysis.NoSuchProcedureException;
+import org.junit.After;
+import org.junit.Test;
+
+
+public class TestRewriteDataFilesProcedure extends SparkExtensionsTestBase {
+
+  public TestRewriteDataFilesProcedure(String catalogName, String implementation, Map<String, String> config) {
+    super(catalogName, implementation, config);
+  }
+
+  @After
+  public void removeTable() {
+    sql("DROP TABLE IF EXISTS %s", tableName);
+  }
+
+  @Test
+  public void testRewriteDataFilesInEmptyTable() {
+    createTable();
+    List<Object[]> output = sql(
+        "CALL %s.system.rewrite_data_files('%s')", catalogName, tableIdent);
+    assertEquals("Procedure output must match",
+        ImmutableList.of(row(0, 0)),
+        output);
+  }
+
+  @Test
+  public void testRewriteDataFilesOnPartitionTable() {
+    createPartitionTable();
+    // create 5 files for each partition (c2 = 'foo' and c2 = 'bar')
+    insertData(10);
+    List<Object[]> expectedRecords = currentData();
+
+    List<Object[]> output = sql(
+            "CALL %s.system.rewrite_data_files(table => '%s')", catalogName, tableIdent);
+
+    assertEquals("Action should rewrite 10 data files and add 2 data files (one per partition) ",
+            ImmutableList.of(row(10, 2)),
+            output);
+
+    List<Object[]> actualRecords = currentData();
+    assertEquals("Data after compaction should not change", expectedRecords, actualRecords);
+  }
+
+  @Test
+  public void testRewriteDataFilesOnNonPartitionTable() {
+    createTable();
+    // create 10 files under non-partitioned table
+    insertData(10);
+    List<Object[]> expectedRecords = currentData();
+
+    List<Object[]> output = sql(
+            "CALL %s.system.rewrite_data_files(table => '%s')", catalogName, tableIdent);
+
+    assertEquals("Action should rewrite 10 data files and add 1 data files",
+            ImmutableList.of(row(10, 1)),
+            output);
+
+    List<Object[]> actualRecords = currentData();
+    assertEquals("Data after compaction should not change", expectedRecords, actualRecords);
+  }
+
+  @Test
+  public void testRewriteDataFilesWithOptions() {
+    createTable();
+    // create 10 files under non-partitioned table
+    insertData(10);
+    List<Object[]> expectedRecords = currentData();
+
+    // set the min-input-files = 12, instead of default 5 to skip compacting the files.
+    List<Object[]> output = sql(
+            "CALL %s.system.rewrite_data_files(table => '%s', options => map('min-input-files','12'))",
+            catalogName, tableIdent);
+
+    assertEquals("Action should rewrite 0 data files and add 0 data files",
+            ImmutableList.of(row(0, 0)),
+            output);
+
+    List<Object[]> actualRecords = currentData();
+    assertEquals("Data should not change", expectedRecords, actualRecords);
+  }
+
+  @Test
+  public void testRewriteDataFilesWithSortStrategy() {
+    createTable();
+    // create 10 files under non-partitioned table
+    insertData(10);
+    List<Object[]> expectedRecords = currentData();
+
+    // set sort_order = c1 DESC LAST
+    List<Object[]> output = sql(
+            "CALL %s.system.rewrite_data_files(table => '%s', " +
+                    "strategy => 'sort', sort_order => 'c1 DESC NULLS LAST')",
+            catalogName, tableIdent);
+
+    assertEquals("Action should rewrite 10 data files and add 1 data files",
+            ImmutableList.of(row(10, 1)),
+            output);
+
+    List<Object[]> actualRecords = currentData();
+    assertEquals("Data after compaction should not change", expectedRecords, actualRecords);
+  }
+
+  @Test
+  public void testRewriteDataFilesWithFilter() {
+    createTable();
+    // create 10 files under non-partitioned table
+    insertData(10);
+    List<Object[]> expectedRecords = currentData();
+
+    // select only 5 files for compaction (files that may have c1 = 1)
+    List<Object[]> output = sql(
+            "CALL %s.system.rewrite_data_files(table => '%s'," +
+                    " where => 'c1 = 1 and c2 is not null')", catalogName, tableIdent);
+
+    assertEquals("Action should rewrite 5 data files (containing c1 = 1) and add 1 data files",
+            ImmutableList.of(row(5, 1)),
+            output);
+
+    List<Object[]> actualRecords = currentData();
+    assertEquals("Data after compaction should not change", expectedRecords, actualRecords);
+  }
+
+  @Test
+  public void testRewriteDataFilesWithFilterOnPartitionTable() {
+    createPartitionTable();
+    // create 5 files for each partition (c2 = 'foo' and c2 = 'bar')
+    insertData(10);
+    List<Object[]> expectedRecords = currentData();
+
+    // select only 5 files for compaction (files in the partition c2 = 'bar')
+    List<Object[]> output = sql(
+        "CALL %s.system.rewrite_data_files(table => '%s'," +
+            " where => 'c2 = \"bar\"')", catalogName, tableIdent);
+
+    assertEquals("Action should rewrite 5 data files from single matching partition" +
+            "(containing c2 = bar) and add 1 data files",
+        ImmutableList.of(row(5, 1)),
+        output);
+
+    List<Object[]> actualRecords = currentData();
+    assertEquals("Data after compaction should not change", expectedRecords, actualRecords);
+  }
+
+  @Test
+  public void testRewriteDataFilesWithInvalidInputs() {
+    createTable();
+    // create 2 files under non-partitioned table
+    insertData(2);
+
+    // Test for invalid strategy
+    AssertHelpers.assertThrows("Should reject calls with unsupported strategy error message",
+            IllegalArgumentException.class, "unsupported strategy: temp. Only binpack,sort is supported",
+            () -> sql("CALL %s.system.rewrite_data_files(table => '%s', options => map('min-input-files','2'), " +
+                            "strategy => 'temp')", catalogName, tableIdent));
+
+    // Test for sort_order with binpack strategy
+    AssertHelpers.assertThrows("Should reject calls with error message",
+            IllegalArgumentException.class, "Cannot set strategy to sort, it has already been set",
+            () -> sql("CALL %s.system.rewrite_data_files(table => '%s', strategy => 'binpack', " +
+                    "sort_order => 'c1 ASC NULLS FIRST')", catalogName, tableIdent));
+
+    // Test for sort_order with invalid null order
+    AssertHelpers.assertThrows("Should reject calls with error message",
+            IllegalArgumentException.class, "Unable to parse sortOrder:",
+            () -> sql("CALL %s.system.rewrite_data_files(table => '%s', strategy => 'sort', " +
+                    "sort_order => 'c1 ASC none')", catalogName, tableIdent));
+
+    // Test for sort_order with invalid sort direction
+    AssertHelpers.assertThrows("Should reject calls with error message",
+            IllegalArgumentException.class, "Unable to parse sortOrder:",
+            () -> sql("CALL %s.system.rewrite_data_files(table => '%s', strategy => 'sort', " +
+                    "sort_order => 'c1 none NULLS FIRST')", catalogName, tableIdent));
+
+    // Test for sort_order with invalid column name
+    AssertHelpers.assertThrows("Should reject calls with error message",
+            ValidationException.class, "Cannot find field 'col1' in struct:" +
+                    " struct<1: c1: optional int, 2: c2: optional string, 3: c3: optional string>",
+            () -> sql("CALL %s.system.rewrite_data_files(table => '%s', strategy => 'sort', " +
+                    "sort_order => 'col1 DESC NULLS FIRST')", catalogName, tableIdent));
+
+    // Test for sort_order with invalid filter column col1
+    AssertHelpers.assertThrows("Should reject calls with error message",
+            ValidationException.class, "Cannot find field 'col1' in struct:" +
+                    " struct<1: c1: optional int, 2: c2: optional string, 3: c3: optional string>",
+            () -> sql("CALL %s.system.rewrite_data_files(table => '%s', " +
+                    "where => 'col1 = 3')", catalogName, tableIdent));
+  }
+
+  @Test
+  public void testInvalidCasesForRewriteDataFiles() {
+    AssertHelpers.assertThrows("Should not allow mixed args",
+        AnalysisException.class, "Named and positional arguments cannot be mixed",
+        () -> sql("CALL %s.system.rewrite_data_files('n', table => 't')", catalogName));
+
+    AssertHelpers.assertThrows("Should not resolve procedures in arbitrary namespaces",
+        NoSuchProcedureException.class, "not found",
+        () -> sql("CALL %s.custom.rewrite_data_files('n', 't')", catalogName));
+
+    AssertHelpers.assertThrows("Should reject calls without all required args",
+        AnalysisException.class, "Missing required parameters",
+        () -> sql("CALL %s.system.rewrite_data_files()", catalogName));
+
+    AssertHelpers.assertThrows("Should reject duplicate arg names name",
+        AnalysisException.class, "Duplicate procedure argument: table",
+        () -> sql("CALL %s.system.rewrite_data_files(table => 't', table => 't')", catalogName));
+
+    AssertHelpers.assertThrows("Should reject calls with empty table identifier",
+        IllegalArgumentException.class, "Cannot handle an empty identifier",
+        () -> sql("CALL %s.system.rewrite_data_files('')", catalogName));
+  }
+
+  private void createTable() {
+    sql("CREATE TABLE %s (c1 int, c2 string, c3 string) USING iceberg", tableName);
+  }
+
+  private void createPartitionTable() {
+    sql("CREATE TABLE %s (c1 int, c2 string, c3 string) USING iceberg PARTITIONED BY (c2)", tableName);
+  }
+
+  private void insertData(int filesCount) {
+    ThreeColumnRecord record1 = new ThreeColumnRecord(1, "foo", "detail1");
+    ThreeColumnRecord record2 = new ThreeColumnRecord(2, "bar", "detail2");
+
+    List<ThreeColumnRecord> records = Lists.newArrayList();
+    IntStream.range(0, filesCount / 2).forEach(i -> {
+      records.add(record1);
+      records.add(record2);
+    });
+
+    Dataset<Row> df = spark.createDataFrame(records, ThreeColumnRecord.class).repartition(filesCount);
+    try {
+      df.writeTo(tableName).append();
+    } catch (org.apache.spark.sql.catalyst.analysis.NoSuchTableException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  private List<Object[]> currentData() {
+    return rowsToJava(spark.sql("SELECT * FROM " + tableName + " order by c1, c2, c3").collectAsList());
+  }
+}