diff --git a/integration/src/test/scala/com/spotify/scio/iceberg/IcebergIOIT.scala b/integration/src/test/scala/com/spotify/scio/iceberg/IcebergIOIT.scala index 0cb451ec71..fa6d33a01a 100644 --- a/integration/src/test/scala/com/spotify/scio/iceberg/IcebergIOIT.scala +++ b/integration/src/test/scala/com/spotify/scio/iceberg/IcebergIOIT.scala @@ -17,8 +17,10 @@ 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.{ @@ -26,18 +28,28 @@ import org.apache.iceberg.types.Types.{ 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() + } + } } diff --git a/scio-managed/src/main/scala/com/spotify/scio/iceberg/IcebergIO.scala b/scio-managed/src/main/scala/com/spotify/scio/iceberg/IcebergIO.scala index daa616bfad..bb1ec3ffe8 100644 --- a/scio-managed/src/main/scala/com/spotify/scio/iceberg/IcebergIO.scala +++ b/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] } diff --git a/scio-managed/src/main/scala/com/spotify/scio/iceberg/syntax/IcebergSCollectionSyntax.scala b/scio-managed/src/main/scala/com/spotify/scio/iceberg/syntax/IcebergSCollectionSyntax.scala index 88d1d0efce..2add91d0c7 100644 --- a/scio-managed/src/main/scala/com/spotify/scio/iceberg/syntax/IcebergSCollectionSyntax.scala +++ b/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) } diff --git a/scio-managed/src/test/scala/com/spotify/scio/iceberg/IcebergIOTest.scala b/scio-managed/src/test/scala/com/spotify/scio/iceberg/IcebergIOTest.scala index 37f0305f12..44bc1e0210 100644 --- a/scio-managed/src/test/scala/com/spotify/scio/iceberg/IcebergIOTest.scala +++ b/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" )