You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by bh...@apache.org on 2022/03/04 23:04:34 UTC
[beam] branch master updated: BEAM-14026 - Fixes bug related to Unnesting nested rows in an array (#16988)
This is an automated email from the ASF dual-hosted git repository.
bhulette pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new bb5dd05 BEAM-14026 - Fixes bug related to Unnesting nested rows in an array (#16988)
bb5dd05 is described below
commit bb5dd05053cb8abcc7780a691d21e4b5b2f2f313
Author: abhijeet-lele <56...@users.noreply.github.com>
AuthorDate: Sat Mar 5 04:33:15 2022 +0530
BEAM-14026 - Fixes bug related to Unnesting nested rows in an array (#16988)
* Suggested changes to handle nested row in an array
* Beam-14026 Suggested changes to handle nested row in an array
* Beam-14026 Enhanced by segregating the code from getBaseValues enhanced test case and example.
* Beam-14026 The code is moved from Row to avoid impact to the public interface.
The code is moved to BeamUnnestRel.java since its the caller class.
The Example code was duplicate, hence dropped.
build.gradle updated with the removal of example code.
---
.../sdk/extensions/sql/impl/rel/BeamUnnestRel.java | 35 ++++-
.../extensions/sql/BeamSqlDslUnnestRowsTest.java | 142 +++++++++++++++++++++
2 files changed, 176 insertions(+), 1 deletion(-)
diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamUnnestRel.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamUnnestRel.java
index 655d75a..d454dfe 100644
--- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamUnnestRel.java
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamUnnestRel.java
@@ -17,8 +17,13 @@
*/
package org.apache.beam.sdk.extensions.sql.impl.rel;
+import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
import org.apache.beam.sdk.extensions.sql.impl.planner.BeamCostModel;
import org.apache.beam.sdk.extensions.sql.impl.planner.BeamRelMetadataQuery;
import org.apache.beam.sdk.extensions.sql.impl.planner.NodeStats;
@@ -129,6 +134,34 @@ public class BeamUnnestRel extends Uncollect implements BeamRelNode {
this.outputSchema = outputSchema;
this.unnestIndices = unnestIndices;
}
+ /**
+ * This is recursive call to get all the values of the nested rows. The recusion is bounded by
+ * the amount of nesting with in the data. This mirrors the unnest behavior of calcite towards
+ * schema. *
+ */
+ private List<Object> getNestedRowBaseValues(Row nestedRow) {
+ return IntStream.range(0, nestedRow.getFieldCount())
+ .mapToObj(
+ (i) -> {
+ List<Object> values = new ArrayList<>();
+ Schema.FieldType fieldType = nestedRow.getSchema().getField(i).getType();
+ if (fieldType.getTypeName().equals(Schema.TypeName.ROW)) {
+ @Nullable Row row = nestedRow.getBaseValue(i, Row.class);
+ if (row == null) {
+ return Stream.builder().build();
+ }
+ List<Object> rowValues = getNestedRowBaseValues(row);
+ if (null != rowValues) {
+ values.addAll(rowValues);
+ }
+ } else {
+ values.add(nestedRow.getBaseValue(i));
+ }
+ return values.stream();
+ })
+ .flatMap(Function.identity())
+ .collect(Collectors.toList());
+ }
@ProcessElement
public void process(@Element Row row, OutputReceiver<Row> out) {
@@ -157,7 +190,7 @@ public class BeamUnnestRel extends Uncollect implements BeamRelNode {
out.output(
Row.withSchema(outputSchema)
.addValues(row.getBaseValues())
- .addValues(nestedRow.getBaseValues())
+ .addValues(getNestedRowBaseValues(nestedRow))
.build());
} else {
out.output(
diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslUnnestRowsTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslUnnestRowsTest.java
new file mode 100644
index 0000000..d72bc24
--- /dev/null
+++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslUnnestRowsTest.java
@@ -0,0 +1,142 @@
+/*
+ * 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.beam.sdk.extensions.sql;
+
+import java.util.Arrays;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.testing.PAssert;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.values.PBegin;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TupleTag;
+import org.junit.Rule;
+import org.junit.Test;
+
+/** Tests for nested rows handling. */
+public class BeamSqlDslUnnestRowsTest {
+
+ @Rule public final TestPipeline pipeline = TestPipeline.create();
+
+ /**
+ * TODO([BEAM-14026]): This is a test of the incorrect behavior unnest because calcite flattens
+ * the row.
+ */
+ @Test
+ public void testUnnestArrayWithNestedRows() {
+
+ Schema level3Type =
+ Schema.builder().addInt32Field("c1").addStringField("c2").addDoubleField("c3").build();
+
+ Row level3Row1 = Row.withSchema(level3Type).addValues(1, "row", 1.0).build();
+ Row level3Row2 = Row.withSchema(level3Type).addValues(2, "row", 2.0).build();
+ Row level3Row3 = Row.withSchema(level3Type).addValues(3, "row", 3.0).build();
+
+ // define the input row format level3
+ Schema level2Type =
+ Schema.builder()
+ .addInt32Field("b1")
+ .addStringField("b2")
+ .addRowField("b3", level3Type)
+ .addDoubleField("b4")
+ .build();
+
+ Row level2Row1 = Row.withSchema(level2Type).addValues(1, "row", level3Row1, 1.0).build();
+ Row level2Row2 = Row.withSchema(level2Type).addValues(2, "row", level3Row2, 2.0).build();
+ Row level2Row3 = Row.withSchema(level2Type).addValues(3, "row", level3Row3, 3.0).build();
+
+ // define the input row format level3
+ Schema level1Type =
+ Schema.builder()
+ .addInt32Field("a1")
+ .addStringField("a2")
+ .addDoubleField("a3")
+ .addArrayField("a4", Schema.FieldType.row(level2Type))
+ .build();
+ Row level1Row1 =
+ Row.withSchema(level1Type)
+ .addValues(1, "row", 1.0, Arrays.asList(level2Row1, level2Row2, level2Row3))
+ .build();
+ Row level1Row2 =
+ Row.withSchema(level1Type)
+ .addValues(2, "row", 2.0, Arrays.asList(level2Row1, level2Row2, level2Row3))
+ .build();
+ Row level1Row3 =
+ Row.withSchema(level1Type)
+ .addValues(3, "row", 3.0, Arrays.asList(level2Row1, level2Row2, level2Row3))
+ .build();
+
+ // create a source PCollection with Create.of();
+ PCollection<Row> inputTable =
+ PBegin.in(pipeline)
+ .apply(Create.of(level1Row1, level1Row2, level1Row3).withRowSchema(level1Type));
+
+ String sql =
+ "select t.a1, t.a2, t.a3, d.b1, d.b2, d.b4, "
+ + "d.b3.c1, d.b3.c2, d.b3.c3 from test t cross join unnest(t.a4) d";
+ // Case 1. run a simple SQL query over input PCollection with BeamSql.simpleQuery;
+ PCollection<Row> result =
+ PCollectionTuple.of(new TupleTag<>("test"), inputTable).apply(SqlTransform.query(sql));
+
+ Schema resultSchema =
+ Schema.builder()
+ .addInt32Field("a1")
+ .addStringField("a2")
+ .addDoubleField("a3")
+ .addInt32Field("b1")
+ .addStringField("b2")
+ .addDoubleField("b4")
+ .addInt32Field("c1")
+ .addStringField("c2")
+ .addDoubleField("c3")
+ .build();
+
+ PAssert.that(result)
+ .containsInAnyOrder(
+ Row.withSchema(resultSchema)
+ .addValues(1, "row", 1.0, 1, "row", 1.0, 1, "row", 1.0)
+ .build(),
+ Row.withSchema(resultSchema)
+ .addValues(1, "row", 1.0, 2, "row", 2.0, 2, "row", 2.0)
+ .build(),
+ Row.withSchema(resultSchema)
+ .addValues(1, "row", 1.0, 3, "row", 3.0, 3, "row", 3.0)
+ .build(),
+ Row.withSchema(resultSchema)
+ .addValues(3, "row", 3.0, 1, "row", 1.0, 1, "row", 1.0)
+ .build(),
+ Row.withSchema(resultSchema)
+ .addValues(3, "row", 3.0, 2, "row", 2.0, 2, "row", 2.0)
+ .build(),
+ Row.withSchema(resultSchema)
+ .addValues(3, "row", 3.0, 3, "row", 3.0, 3, "row", 3.0)
+ .build(),
+ Row.withSchema(resultSchema)
+ .addValues(2, "row", 2.0, 1, "row", 1.0, 1, "row", 1.0)
+ .build(),
+ Row.withSchema(resultSchema)
+ .addValues(2, "row", 2.0, 2, "row", 2.0, 2, "row", 2.0)
+ .build(),
+ Row.withSchema(resultSchema)
+ .addValues(2, "row", 2.0, 3, "row", 3.0, 3, "row", 3.0)
+ .build());
+ pipeline.run();
+ }
+}