Expand IcebergIO write API to cover all write options

integration/src/test/scala/com/spotify/scio/iceberg/IcebergIOIT.scala

@@ -17,27 +17,39 @@
 package com.spotify.scio.iceberg
 
 import com.dimafeng.testcontainers.{ForAllTestContainer, GenericContainer}
+import com.spotify.scio.parquet.BeamInputFile
 import com.spotify.scio.testing.PipelineSpec
 import magnolify.beam._
+import magnolify.beam.logical.millis._
 import org.apache.iceberg.catalog.{Namespace, TableIdentifier}
 import org.apache.iceberg.rest.RESTCatalog
 import org.apache.iceberg.types.Types.{
   BooleanType,
   IntegerType,
   NestedField,
   StringType,
-  StructType
+  StructType,
+  TimestampType
 }
-import org.apache.iceberg.{CatalogProperties, CatalogUtil, PartitionSpec, Schema}
+import org.apache.iceberg.{
+  CatalogProperties,
+  CatalogUtil,
+  NullOrder,
+  PartitionSpec,
+  Schema,
+  SortOrder
+}
+import org.apache.parquet.hadoop.ParquetFileReader
 import org.testcontainers.containers.wait.strategy.HostPortWaitStrategy
 
-import java.time.Duration
+import java.time.{Duration, Instant}
 import java.io.File
 import java.nio.file.Files
+import java.time.temporal.ChronoUnit
 import scala.jdk.CollectionConverters._
 
 case class Nested(d: Boolean)
-case class IcebergIOITRecord(a: Int, b: String, c: Nested)
+case class IcebergIOITRecord(ts: Instant, a: Int, b: String, c: Nested)
 object IcebergIOITRecord {
   implicit val icebergIOITRecordRowType: RowType[IcebergIOITRecord] = RowType[IcebergIOITRecord]
 }
@@ -65,32 +77,41 @@ class IcebergIOIT extends PipelineSpec with ForAllTestContainer {
 
   lazy val uri = s"http://${container.containerIpAddress}:${container.mappedPort(ContainerPort)}"
 
-  override def afterStart(): Unit = {
+  lazy val tableSchema = new Schema(
+    NestedField.required(1, "ts", TimestampType.withZone()),
+    NestedField.required(2, "a", IntegerType.get()),
+    NestedField.required(3, "b", StringType.get()),
+    NestedField.required(
+      4,
+      "c",
+      StructType.of(NestedField.required(5, "d", BooleanType.get()))
+    )
+  )
+
+  lazy val catalog: RESTCatalog = {
     val cat = new RESTCatalog()
     cat.initialize(CatalogName, Map("uri" -> uri).asJava)
+    cat
+  }
 
-    cat.createNamespace(Namespace.of(NamespaceName))
-    cat.createTable(
+  override def afterStart(): Unit = {
+    catalog.createNamespace(Namespace.of(NamespaceName))
+    catalog.createTable(
       TableIdentifier.parse(TableName),
-      new Schema(
-        NestedField.required(0, "a", IntegerType.get()),
-        NestedField.required(1, "b", StringType.get()),
-        NestedField.required(
-          2,
-          "c",
-          StructType.of(NestedField.required(3, "d", BooleanType.get()))
-        )
-      ),
+      tableSchema,
       PartitionSpec.unpartitioned()
     )
   }
 
+  override def beforeStop(): Unit = catalog.close()
+
   "IcebergIO" should "work" in {
     val catalogProperties = Map(
       CatalogUtil.ICEBERG_CATALOG_TYPE -> CatalogUtil.ICEBERG_CATALOG_TYPE_REST,
       CatalogProperties.URI -> uri
     )
-    val elements = 1.to(10).map(i => IcebergIOITRecord(i, s"$i", Nested(i % 2 == 0)))
+    val ts = Instant.now().truncatedTo(ChronoUnit.DAYS)
+    val elements = 1.to(10).map(i => IcebergIOITRecord(ts, i, s"$i", Nested(i % 2 == 0)))
 
     runWithRealContext() { sc =>
       sc.parallelize(elements)
@@ -104,4 +125,71 @@ class IcebergIOIT extends PipelineSpec with ForAllTestContainer {
       ) should containInAnyOrder(elements)
     }
   }
+
+  it should "propagate Iceberg dynamic table creation properties" in {
+    val tableName = s"${NamespaceName}.dynamic_table_creation"
+    val catalogProperties = Map(
+      CatalogUtil.ICEBERG_CATALOG_TYPE -> CatalogUtil.ICEBERG_CATALOG_TYPE_REST,
+      CatalogProperties.URI -> uri
+    )
+    val elements =
+      1.to(100).map(i => IcebergIOITRecord(Instant.now(), i, s"value_$i", Nested(i % 2 == 0)))
+
+    val customWriteDataPath = s"$tempDir/custom_path"
+
+    runWithRealContext() { sc =>
+      sc.parallelize(elements)
+        .saveAsIceberg(
+          tableName,
+          catalogProperties = catalogProperties,
+          tableProperties = Map(
+            "write.data.path" -> customWriteDataPath,
+            "write.parquet.bloom-filter-enabled.column.b" -> "true"
+          ),
+          partitionFields = List("day(ts)"),
+          sortFields = List("a asc nulls first")
+        )
+    }
+
+    val table = catalog.loadTable(TableIdentifier.parse(tableName))
+    table.schema().sameSchema(tableSchema) shouldBe true
+
+    // Validate PartitionSpec and SortOrder
+    table.spec() shouldEqual PartitionSpec.builderFor(table.schema()).day("ts").build()
+    table.sortOrder() shouldEqual SortOrder
+      .builderFor(table.schema())
+      .asc("a", NullOrder.NULLS_FIRST)
+      .build()
+
+    // Validate table properties
+    table.properties().get("write.data.path") shouldBe customWriteDataPath
+
+    val tasks = table.newScan().planFiles()
+    try {
+      val dataFiles = tasks.iterator().asScala.map(_.file().location()).toSeq
+      dataFiles should not be empty
+
+      dataFiles.foreach { path =>
+        path should startWith(s"$customWriteDataPath/ts_day=")
+        val reader = ParquetFileReader.open(BeamInputFile.of(path))
+        try {
+          reader.getFooter.getBlocks.asScala.foreach { block =>
+            block.getColumns.asScala.foreach { col =>
+              val hasBloom = col.getBloomFilterOffset > 0
+              col.getPath.toDotString match {
+                case "b" =>
+                  hasBloom shouldBe true
+                case _ =>
+                  hasBloom shouldBe false
+              }
+            }
+          }
+        } finally {
+          reader.close()
+        }
+      }
+    } finally {
+      tasks.close()
+    }
+  }
 }

scio-managed/src/main/scala/com/spotify/scio/iceberg/IcebergIO.scala

@@ -41,6 +41,7 @@ private[scio] object ConfigMap {
   implicit val stringToMap: ConfigMapType[String] = _ => Map.empty
   implicit val intToMap: ConfigMapType[Int] = _ => Map.empty
   implicit val mapToMap: ConfigMapType[Map[String, String]] = _ => Map.empty
+  implicit val mapAnyRefToMap: ConfigMapType[Map[String, AnyRef]] = _ => Map.empty
   implicit val listToMap: ConfigMapType[List[String]] = _ => Map.empty
   implicit def optionToMap[T]: ConfigMapType[Option[T]] = _ => Map.empty
 
@@ -86,8 +87,10 @@ final case class IcebergIO[T: RowType: Coder](table: String, catalogName: Option
 
   private[scio] def config(params: WriteP)(implicit
     mapper: ConfigMap.ConfigMapType[WriteP]
-  ): Map[String, AnyRef] =
-    baseConfig(params)
+  ): Map[String, AnyRef] = {
+    val extra = Option(params.extraConfigProperties).getOrElse(Map.empty)
+    baseConfig(params) - "extra_config_properties" ++ extra
+  }
 
   private[scio] def config(
     params: ReadP
@@ -133,15 +136,21 @@ object IcebergIO {
   }
   case class WriteParam private (
     catalogProperties: Map[String, String] = WriteParam.DefaultCatalogProperties,
-    configProperties: Map[String, String] = WriteParam.DefaultHadoopConfigProperties,
+    tableProperties: Map[String, String] = WriteParam.DefaultTableProperties,
+    sortFields: List[String] = WriteParam.DefaultSortFields,
+    partitionFields: List[String] = WriteParam.DefaultPartitionFields,
     triggeringFrequencySeconds: Option[Int] = None,
-    directWriteByteLimit: Option[Int] = None
+    directWriteByteLimit: Option[Int] = None,
+    extraConfigProperties: Map[String, AnyRef] = WriteParam.DefaultExtraConfigProperties
   )
   object WriteParam {
     val DefaultCatalogProperties: Map[String, String] = null
-    val DefaultHadoopConfigProperties: Map[String, String] = null
+    val DefaultTableProperties: Map[String, String] = null
+    val DefaultSortFields: List[String] = null
+    val DefaultPartitionFields: List[String] = null
     val DefaultTriggeringFrequencySeconds: Int = -1
     val DefaultDirectWriteByteLimit: Int = -1
+    val DefaultExtraConfigProperties: Map[String, AnyRef] = null
 
     implicit val configMap: ConfigMap.ConfigMapType[WriteParam] = ConfigMap.gen[WriteParam]
   }

scio-managed/src/main/scala/com/spotify/scio/iceberg/syntax/IcebergSCollectionSyntax.scala

@@ -34,12 +34,20 @@ class IcebergSCollectionSyntax[T: RowType: Coder](self: SCollection[T]) {
    * @param catalogProperties
    *   any additional properties required by the Iceberg catalog; see:
    *   https://iceberg.apache.org/docs/latest/catalog-properties
-   * @param hadoopConfigProperties
-   *   any additional Hadoop configuration properties
+   * @param tableProperties
+   *   any additional Iceberg table properties to set during dynamic table creation; see:
+   *   https://iceberg.apache.org/docs/latest/configuration/#write-properties
+   * @param sortFields
+   *   list of field names defining the sort order for written files
+   * @param partitionFields
+   *   list of field names defining the partition spec for the table
    * @param triggeringFrequencySeconds
    *   (streaming only) frequency at which snapshots are produced
    * @param directWriteByteLimit
    *   (streaming only) limit for lifting bundles into the direct write path.
+   * @param extraConfigProperties
+   *   additional properties to pass to the Managed IO config, i.e. `distribution_mode: hash` or
+   *   `autosharding: true`
    *
    * For a complete reference, see:
    * https://docs.cloud.google.com/dataflow/docs/guides/managed-io-iceberg
@@ -48,19 +56,24 @@ class IcebergSCollectionSyntax[T: RowType: Coder](self: SCollection[T]) {
     table: String,
     catalogName: String = null,
     catalogProperties: Map[String, String] = IcebergIO.WriteParam.DefaultCatalogProperties,
-    hadoopConfigProperties: Map[String, String] =
-      IcebergIO.WriteParam.DefaultHadoopConfigProperties,
+    tableProperties: Map[String, String] = IcebergIO.WriteParam.DefaultTableProperties,
+    sortFields: List[String] = IcebergIO.WriteParam.DefaultSortFields,
+    partitionFields: List[String] = IcebergIO.WriteParam.DefaultPartitionFields,
+    extraConfigProperties: Map[String, AnyRef] = IcebergIO.WriteParam.DefaultExtraConfigProperties,
     triggeringFrequencySeconds: Int = IcebergIO.WriteParam.DefaultTriggeringFrequencySeconds,
     directWriteByteLimit: Int = IcebergIO.WriteParam.DefaultDirectWriteByteLimit
   ): ClosedTap[Nothing] = {
 
     val params = IcebergIO.WriteParam(
       catalogProperties,
-      hadoopConfigProperties,
+      tableProperties,
+      sortFields,
+      partitionFields,
       Option(triggeringFrequencySeconds).filter(
         _ != IcebergIO.WriteParam.DefaultTriggeringFrequencySeconds
       ),
-      Option(directWriteByteLimit).filter(_ != IcebergIO.WriteParam.DefaultDirectWriteByteLimit)
+      Option(directWriteByteLimit).filter(_ != IcebergIO.WriteParam.DefaultDirectWriteByteLimit),
+      extraConfigProperties
     )
     self.write(IcebergIO(table, Option(catalogName)))(params)
   }

scio-managed/src/test/scala/com/spotify/scio/iceberg/IcebergIOTest.scala

@@ -48,12 +48,16 @@ class IcebergIOTest extends ScioIOSpec {
       IcebergIO.WriteParam(
         Map.empty,
         Map.empty,
+        Nil,
+        Nil,
         None,
         None
       ),
       IcebergIO.WriteParam(
         Map("catalogProp1" -> "catalogProp1Value"),
         Map("configProp1" -> "configProp1Value", "configProp2" -> "configProp2Value"),
+        List("sortField1"),
+        List("partField1"),
         Some(10),
         Some(100)
       )
@@ -70,6 +74,9 @@ class IcebergIOTest extends ScioIOSpec {
       // reads
       "filter",
       // writes
+      "table_properties",
+      "sort_fields",
+      "partition_fields",
       "triggering_frequency_seconds",
       "direct_write_byte_limit"
     )
@@ -129,14 +136,24 @@ class IcebergIOTest extends ScioIOSpec {
 
     val writeParam = IcebergIO.WriteParam(
       Map("a" -> "b"),
-      Map("c" -> "d", "e" -> "f")
+      Map("c" -> "d", "e" -> "f"),
+      List("col1", "col2"),
+      List("partCol1"),
+      extraConfigProperties = Map(
+        "distribution_mode" -> "hash",
+        "autosharding" -> (true: java.lang.Boolean)
+      )
     )
 
     val managedConfig: Map[String, AnyRef] = io.config(writeParam)
 
     managedConfig should contain only (
-      "config_properties" -> Map("c" -> "d", "e" -> "f"),
+      "table_properties" -> Map("c" -> "d", "e" -> "f"),
+      "sort_fields" -> List("col1", "col2"),
+      "partition_fields" -> List("partCol1"),
       "catalog_properties" -> Map("a" -> "b"),
+      "distribution_mode" -> "hash",
+      "autosharding" -> true,
       "table" -> "tableName",
       "catalog_name" -> "catalogName"
     )