|
| 1 | +/* |
| 2 | + * Licensed to the Apache Software Foundation (ASF) under one |
| 3 | + * or more contributor license agreements. See the NOTICE file |
| 4 | + * distributed with this work for additional information |
| 5 | + * regarding copyright ownership. The ASF licenses this file |
| 6 | + * to you under the Apache License, Version 2.0 (the |
| 7 | + * "License"); you may not use this file except in compliance |
| 8 | + * with the License. You may obtain a copy of the License at |
| 9 | + * |
| 10 | + * http://www.apache.org/licenses/LICENSE-2.0 |
| 11 | + * |
| 12 | + * Unless required by applicable law or agreed to in writing, software |
| 13 | + * distributed under the License is distributed on an "AS IS" BASIS, |
| 14 | + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 15 | + * See the License for the specific language governing permissions and |
| 16 | + * limitations under the License. |
| 17 | + */ |
| 18 | +package org.apache.beam.sdk.io.delta; |
| 19 | + |
| 20 | +import static org.apache.beam.sdk.io.delta.DeltaReadSchemaTransformProvider.Configuration; |
| 21 | +import static org.apache.beam.sdk.io.delta.DeltaReadSchemaTransformProvider.OUTPUT_TAG; |
| 22 | + |
| 23 | +import java.io.File; |
| 24 | +import java.nio.charset.StandardCharsets; |
| 25 | +import java.nio.file.Files; |
| 26 | +import org.apache.avro.generic.GenericRecord; |
| 27 | +import org.apache.beam.sdk.extensions.avro.coders.AvroCoder; |
| 28 | +import org.apache.beam.sdk.extensions.avro.schemas.utils.AvroUtils; |
| 29 | +import org.apache.beam.sdk.io.Compression; |
| 30 | +import org.apache.beam.sdk.io.FileIO; |
| 31 | +import org.apache.beam.sdk.io.parquet.ParquetIO; |
| 32 | +import org.apache.beam.sdk.schemas.Schema; |
| 33 | +import org.apache.beam.sdk.testing.PAssert; |
| 34 | +import org.apache.beam.sdk.testing.TestPipeline; |
| 35 | +import org.apache.beam.sdk.transforms.Create; |
| 36 | +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; |
| 37 | +import org.apache.beam.sdk.transforms.windowing.PaneInfo; |
| 38 | +import org.apache.beam.sdk.values.PCollection; |
| 39 | +import org.apache.beam.sdk.values.PCollectionRowTuple; |
| 40 | +import org.apache.beam.sdk.values.Row; |
| 41 | +import org.junit.Rule; |
| 42 | +import org.junit.Test; |
| 43 | +import org.junit.rules.TemporaryFolder; |
| 44 | +import org.junit.runner.RunWith; |
| 45 | +import org.junit.runners.JUnit4; |
| 46 | + |
| 47 | +/** Tests for {@link DeltaReadSchemaTransformProvider}. */ |
| 48 | +@RunWith(JUnit4.class) |
| 49 | +public class DeltaReadSchemaTransformProviderTest { |
| 50 | + |
| 51 | + @Rule public TestPipeline writePipeline = TestPipeline.create(); |
| 52 | + @Rule public TestPipeline readPipeline = TestPipeline.create(); |
| 53 | + @Rule public TemporaryFolder tempFolder = new TemporaryFolder(); |
| 54 | + |
| 55 | + @Test |
| 56 | + public void testBuildTransformWithRow() { |
| 57 | + java.util.Map<String, String> hadoopConfig = new java.util.HashMap<>(); |
| 58 | + hadoopConfig.put("fs.gs.project.id", "test-project"); |
| 59 | + |
| 60 | + Row config = |
| 61 | + Row.withSchema(new DeltaReadSchemaTransformProvider().configurationSchema()) |
| 62 | + .withFieldValue("table", "/path/to/table") |
| 63 | + .withFieldValue("version", 5L) |
| 64 | + .withFieldValue("timestamp", "2026-06-04T12:00:00Z") |
| 65 | + .withFieldValue("hadoop_config", hadoopConfig) |
| 66 | + .build(); |
| 67 | + |
| 68 | + new DeltaReadSchemaTransformProvider().from(config); |
| 69 | + } |
| 70 | + |
| 71 | + @Test |
| 72 | + public void testSimpleScan() throws Exception { |
| 73 | + File tableDir = tempFolder.newFolder("delta-table-simple"); |
| 74 | + |
| 75 | + // 1. Write a Parquet file using Beam |
| 76 | + Schema schema = Schema.builder().addField("name", Schema.FieldType.STRING).build(); |
| 77 | + Row row = Row.withSchema(schema).addValues("test-name").build(); |
| 78 | + |
| 79 | + org.apache.avro.Schema avroSchema = AvroUtils.toAvroSchema(schema); |
| 80 | + GenericRecord record = AvroUtils.toGenericRecord(row, avroSchema); |
| 81 | + |
| 82 | + writePipeline |
| 83 | + .apply("Create Input", Create.of(record).withCoder(AvroCoder.of(avroSchema))) |
| 84 | + .apply( |
| 85 | + "Write Parquet", |
| 86 | + FileIO.<GenericRecord>write() |
| 87 | + .via(ParquetIO.sink(avroSchema)) |
| 88 | + .to(tableDir.getAbsolutePath() + "/") |
| 89 | + .withNaming( |
| 90 | + (BoundedWindow window, |
| 91 | + PaneInfo paneInfo, |
| 92 | + int numShards, |
| 93 | + int shardIndex, |
| 94 | + Compression compression) -> "part-00000.parquet")); |
| 95 | + |
| 96 | + writePipeline.run().waitUntilFinish(); |
| 97 | + |
| 98 | + File parquetFile = new File(tableDir, "part-00000.parquet"); |
| 99 | + byte[] fileBytes = Files.readAllBytes(parquetFile.toPath()); |
| 100 | + |
| 101 | + // 2. Create the Delta log |
| 102 | + File logDir = new File(tableDir, "_delta_log"); |
| 103 | + logDir.mkdirs(); |
| 104 | + File commitFile = new File(logDir, "00000000000000000000.json"); |
| 105 | + |
| 106 | + String commitContent = |
| 107 | + "{\"protocol\":{\"minReaderVersion\":1,\"minWriterVersion\":2}}\n" |
| 108 | + + "{\"metaData\":{\"id\":\"test-id\",\"format\":{\"provider\":\"parquet\",\"options\":{}},\"schemaString\":\"{\\\"type\\\":\\\"struct\\\",\\\"fields\\\":[{\\\"name\\\":\\\"name\\\",\\\"type\\\":\\\"string\\\",\\\"nullable\\\":true,\\\"metadata\\\":{}}]}\",\"partitionColumns\":[],\"configuration\":{},\"createdAt\":123456789}}\n" |
| 109 | + + "{\"add\":{\"path\":\"part-00000.parquet\",\"partitionValues\":{},\"size\":" |
| 110 | + + fileBytes.length |
| 111 | + + ",\"modificationTime\":123456789,\"dataChange\":true}}"; |
| 112 | + |
| 113 | + Files.write(commitFile.toPath(), commitContent.getBytes(StandardCharsets.UTF_8)); |
| 114 | + |
| 115 | + // 3. Read it using DeltaReadSchemaTransformProvider |
| 116 | + Configuration readConfig = Configuration.builder().setTable(tableDir.getAbsolutePath()).build(); |
| 117 | + |
| 118 | + PCollection<Row> output = |
| 119 | + PCollectionRowTuple.empty(readPipeline) |
| 120 | + .apply(new DeltaReadSchemaTransformProvider().from(readConfig)) |
| 121 | + .get(OUTPUT_TAG); |
| 122 | + |
| 123 | + PAssert.that(output).containsInAnyOrder(row); |
| 124 | + |
| 125 | + readPipeline.run().waitUntilFinish(); |
| 126 | + } |
| 127 | +} |
0 commit comments