(iceberg) branch main updated: Spark 3.5: Add utility to load table state reliably (#10984)

[ao...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

aokolnychyi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg.git


The following commit(s) were added to refs/heads/main by this push:
     new 5864850083 Spark 3.5: Add utility to load table state reliably (#10984)
5864850083 is described below

commit 586485008355e2a11ea4c78890964dc88613c54f
Author: Anton Okolnychyi <ao...@apache.org>
AuthorDate: Fri Aug 23 17:31:20 2024 -0700

    Spark 3.5: Add utility to load table state reliably (#10984)
---
 .../java/org/apache/iceberg/spark/SparkTableUtil.java | 19 +++++++++++++++----
 .../apache/iceberg/spark/actions/NDVSketchUtil.java   | 12 ++++--------
 .../spark/actions/TestComputeTableStatsAction.java    | 14 ++++++++++++++
 3 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/SparkTableUtil.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/SparkTableUtil.java
index 7a96e97fb9..e103104171 100644
--- a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/SparkTableUtil.java
+++ b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/SparkTableUtil.java
@@ -851,6 +851,12 @@ public class SparkTableUtil {
         .run(item -> io.deleteFile(item.path()));
   }
 
+  public static Dataset<Row> loadTable(SparkSession spark, Table table, long snapshotId) {
+    SparkTable sparkTable = new SparkTable(table, snapshotId, false);
+    DataSourceV2Relation relation = createRelation(sparkTable, ImmutableMap.of());
+    return Dataset.ofRows(spark, relation);
+  }
+
   public static Dataset<Row> loadMetadataTable(
       SparkSession spark, Table table, MetadataTableType type) {
     return loadMetadataTable(spark, table, type, ImmutableMap.of());
@@ -858,11 +864,16 @@ public class SparkTableUtil {
 
   public static Dataset<Row> loadMetadataTable(
       SparkSession spark, Table table, MetadataTableType type, Map<String, String> extraOptions) {
-    SparkTable metadataTable =
-        new SparkTable(MetadataTableUtils.createMetadataTableInstance(table, type), false);
+    Table metadataTable = MetadataTableUtils.createMetadataTableInstance(table, type);
+    SparkTable sparkMetadataTable = new SparkTable(metadataTable, false);
+    DataSourceV2Relation relation = createRelation(sparkMetadataTable, extraOptions);
+    return Dataset.ofRows(spark, relation);
+  }
+
+  private static DataSourceV2Relation createRelation(
+      SparkTable sparkTable, Map<String, String> extraOptions) {
     CaseInsensitiveStringMap options = new CaseInsensitiveStringMap(extraOptions);
-    return Dataset.ofRows(
-        spark, DataSourceV2Relation.create(metadataTable, Some.empty(), Some.empty(), options));
+    return DataSourceV2Relation.create(sparkTable, Option.empty(), Option.empty(), options);
   }
 
   /**
diff --git a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/NDVSketchUtil.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/NDVSketchUtil.java
index 22055a161e..c8a20d3cca 100644
--- a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/NDVSketchUtil.java
+++ b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/NDVSketchUtil.java
@@ -32,9 +32,10 @@ import org.apache.iceberg.puffin.StandardBlobTypes;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
 import org.apache.iceberg.relocated.com.google.common.collect.Lists;
-import org.apache.iceberg.spark.SparkReadOptions;
+import org.apache.iceberg.spark.SparkTableUtil;
 import org.apache.iceberg.types.Types;
 import org.apache.spark.sql.Column;
+import org.apache.spark.sql.Dataset;
 import org.apache.spark.sql.Row;
 import org.apache.spark.sql.SparkSession;
 import org.apache.spark.sql.stats.ThetaSketchAgg;
@@ -73,13 +74,8 @@ public class NDVSketchUtil {
 
   private static Row computeNDVSketches(
       SparkSession spark, Table table, Snapshot snapshot, List<String> colNames) {
-    return spark
-        .read()
-        .format("iceberg")
-        .option(SparkReadOptions.SNAPSHOT_ID, snapshot.snapshotId())
-        .load(table.name())
-        .select(toAggColumns(colNames))
-        .first();
+    Dataset<Row> inputDF = SparkTableUtil.loadTable(spark, table, snapshot.snapshotId());
+    return inputDF.select(toAggColumns(colNames)).first();
   }
 
   private static Column[] toAggColumns(List<String> colNames) {
diff --git a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/actions/TestComputeTableStatsAction.java b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/actions/TestComputeTableStatsAction.java
index 588bb29f47..4cab0d4281 100644
--- a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/actions/TestComputeTableStatsAction.java
+++ b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/actions/TestComputeTableStatsAction.java
@@ -79,6 +79,20 @@ public class TestComputeTableStatsAction extends CatalogTestBase {
           required(4, "nestedStructCol", NESTED_STRUCT_TYPE),
           required(5, "stringCol", Types.StringType.get()));
 
+  @TestTemplate
+  public void testLoadingTableDirectly() {
+    sql("CREATE TABLE %s (id int, data string) USING iceberg", tableName);
+    sql("INSERT into %s values(1, 'abcd')", tableName);
+
+    Table table = validationCatalog.loadTable(tableIdent);
+
+    SparkActions actions = SparkActions.get();
+    ComputeTableStats.Result results = actions.computeTableStats(table).execute();
+    StatisticsFile statisticsFile = results.statisticsFile();
+    Assertions.assertNotEquals(statisticsFile.fileSizeInBytes(), 0);
+    Assertions.assertEquals(statisticsFile.blobMetadata().size(), 2);
+  }
+
   @TestTemplate
   public void testComputeTableStatsAction() throws NoSuchTableException, ParseException {
     sql("CREATE TABLE %s (id int, data string) USING iceberg", tableName);