[beam] 01/01: [BEAM-9027] [SQL] Fix ZetaSQL Byte Literals

[ap...@apache.org]

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

apilloud pushed a commit to branch bytes
in repository https://gitbox.apache.org/repos/asf/beam.git

commit 0ea672005e434e09275727a9175a0c730a6ce878
Author: Andrew Pilloud <ap...@google.com>
AuthorDate: Mon Jan 6 16:43:57 2020 -0800

    [BEAM-9027] [SQL] Fix ZetaSQL Byte Literals
---
 .../meta/provider/bigquery/BeamSqlUnparseContext.java    |  4 +++-
 .../extensions/sql/zetasql/ZetaSQLDialectSpecTest.java   | 16 ++++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BeamSqlUnparseContext.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BeamSqlUnparseContext.java
index 432b3b5..60fd1bc 100644
--- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BeamSqlUnparseContext.java
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/bigquery/BeamSqlUnparseContext.java
@@ -20,6 +20,7 @@ package org.apache.beam.sdk.extensions.sql.meta.provider.bigquery;
 import static org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.rel2sql.SqlImplementor.POS;
 
 import java.util.function.IntFunction;
+import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.avatica.util.ByteString;
 import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.rel2sql.SqlImplementor;
 import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexLiteral;
 import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexNode;
@@ -46,7 +47,8 @@ public class BeamSqlUnparseContext extends SqlImplementor.SimpleContext {
       final RexLiteral literal = (RexLiteral) rex;
       SqlTypeFamily family = literal.getTypeName().getFamily();
       if (SqlTypeFamily.BINARY.equals(family)) {
-        BitString bitString = BitString.createFromBytes(literal.getValueAs(byte[].class));
+        ByteString byteString = literal.getValueAs(ByteString.class);
+        BitString bitString = BitString.createFromHexString(byteString.toString(16));
         return new SqlByteStringLiteral(bitString, POS);
       } else if (SqlTypeFamily.CHARACTER.equals(family)) {
         String escaped = StringEscapeUtils.escapeJava(literal.getValueAs(String.class));
diff --git a/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSQLDialectSpecTest.java b/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSQLDialectSpecTest.java
index e140c02..9eb33b2 100644
--- a/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSQLDialectSpecTest.java
+++ b/sdks/java/extensions/sql/zetasql/src/test/java/org/apache/beam/sdk/extensions/sql/zetasql/ZetaSQLDialectSpecTest.java
@@ -141,6 +141,22 @@ public class ZetaSQLDialectSpecTest {
   }
 
   @Test
+  public void testByteLiterals() {
+    String sql = "SELECT b'abc'";
+
+    byte[] byteString = new byte[] {'a', 'b', 'c'};
+
+    ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
+    BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
+    PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
+
+    final Schema schema = Schema.builder().addNullableField("ColA", FieldType.BYTES).build();
+
+    PAssert.that(stream).containsInAnyOrder(Row.withSchema(schema).addValues(byteString).build());
+
+    pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
+  }
+  @Test
   public void testByteString() {
     String sql = "SELECT @p0 IS NULL AS ColA";