Fix Row equality/hashing, RowUtils validation and ByteBuffer handling, and LogicalTypeCoder DATETIME conversion

sdks/java/core/src/main/java/org/apache/beam/sdk/schemas/SchemaCoderHelpers.java

@@ -80,13 +80,20 @@ private static class LogicalTypeCoder<InputT, BaseT> extends Coder<InputT> {
       this.isDateTime = logicalType.getBaseType().equals(FieldType.DATETIME);
     }
 
-    @Override
-    public void encode(InputT value, OutputStream outStream) throws CoderException, IOException {
+    private BaseT toBaseType(InputT value) {
+      if (value == null) {
+        return null;
+      }
       BaseT baseType = logicalType.toBaseType(value);
-      if (isDateTime) {
+      if (isDateTime && baseType != null) {
         baseType = (BaseT) ((ReadableInstant) baseType).toInstant();
       }
-      baseTypeCoder.encode(baseType, outStream);
+      return baseType;
+    }
+
+    @Override
+    public void encode(InputT value, OutputStream outStream) throws CoderException, IOException {
+      baseTypeCoder.encode(toBaseType(value), outStream);
     }
 
     @Override
@@ -114,22 +121,23 @@ public boolean consistentWithEquals() {
 
     @Override
     public Object structuralValue(InputT value) {
+      BaseT baseValue = toBaseType(value);
       if (baseTypeCoder.consistentWithEquals()) {
-        return logicalType.toBaseType(value);
+        return baseValue;
       } else {
-        return baseTypeCoder.structuralValue(logicalType.toBaseType(value));
+        return baseTypeCoder.structuralValue(baseValue);
       }
     }
 
     @Override
     public boolean isRegisterByteSizeObserverCheap(InputT value) {
-      return baseTypeCoder.isRegisterByteSizeObserverCheap(logicalType.toBaseType(value));
+      return baseTypeCoder.isRegisterByteSizeObserverCheap(toBaseType(value));
     }
 
     @Override
     public void registerByteSizeObserver(InputT value, ElementByteSizeObserver observer)
         throws Exception {
-      baseTypeCoder.registerByteSizeObserver(logicalType.toBaseType(value), observer);
+      baseTypeCoder.registerByteSizeObserver(toBaseType(value), observer);
     }
   }
 

sdks/java/core/src/main/java/org/apache/beam/sdk/values/Row.java

@@ -486,7 +486,10 @@ public static int deepHashCode(Object a, Schema.FieldType fieldType) {
       if (a == null) {
         return 0;
       } else if (fieldType.getTypeName() == TypeName.LOGICAL_TYPE) {
-        return deepHashCode(a, fieldType.getLogicalType().getBaseType());
+        Schema.LogicalType<Object, Object> logicalType =
+            (Schema.LogicalType<Object, Object>) fieldType.getLogicalType();
+        return deepHashCode(
+            SchemaUtils.toLogicalBaseType(logicalType, a), logicalType.getBaseType());
       } else if (fieldType.getTypeName() == Schema.TypeName.BYTES) {
         return Arrays.hashCode((byte[]) a);
       } else if (fieldType.getTypeName().isCollectionType()) {

sdks/java/core/src/main/java/org/apache/beam/sdk/values/RowUtils.java

@@ -173,11 +173,7 @@ public Object match(
           processedValue = cases.processByte(rowPosition, (Byte) value, this);
           break;
         case BYTES:
-          processedValue =
-              cases.processBytes(
-                  rowPosition,
-                  (byte[]) (value instanceof ByteBuffer ? ((ByteBuffer) value).array() : value),
-                  this);
+          processedValue = cases.processBytes(rowPosition, toByteArray(value), this);
           break;
         case INT16:
           processedValue = cases.processInt16(rowPosition, (Short) value, this);
@@ -423,6 +419,12 @@ public Row processRow(
         }
         retValue = new RowWithStorage(schema, values);
       }
+      if (retValue != null && !schema.equals(retValue.getSchema())) {
+        throw new IllegalArgumentException(
+            String.format(
+                "Expected schema %s for field %s, but got %s",
+                schema, rowPosition.descriptor, retValue.getSchema()));
+      }
       return retValue;
     }
 
@@ -605,7 +607,26 @@ public Boolean processBoolean(RowPosition rowPosition, Boolean value, RowFieldMa
     @Override
     public byte[] processBytes(RowPosition rowPosition, byte[] value, RowFieldMatcher matcher) {
       Object retValue = overrideOrReturn(rowPosition, value);
-      return (retValue instanceof ByteBuffer) ? ((ByteBuffer) retValue).array() : (byte[]) retValue;
+      return toByteArray(retValue);
+    }
+  }
+
+  static byte @Nullable [] toByteArray(@Nullable Object value) {
+    if (value == null) {
+      return null;
+    }
+    if (value instanceof ByteBuffer) {
+      ByteBuffer buf = (ByteBuffer) value;
+      if (buf.hasArray()
+          && buf.arrayOffset() == 0
+          && buf.position() == 0
+          && buf.limit() == buf.capacity()) {
+        return buf.array();
+      }
+      byte[] bytes = new byte[buf.remaining()];
+      buf.duplicate().get(bytes);
+      return bytes;
     }
+    return (byte[]) value;
   }
 }

sdks/java/core/src/main/java/org/apache/beam/sdk/values/RowWithGetters.java

@@ -140,16 +140,19 @@ public boolean equals(@Nullable Object o) {
     }
     if (o instanceof RowWithGetters) {
       RowWithGetters other = (RowWithGetters) o;
-      return Objects.equals(getSchema(), other.getSchema())
-          && Objects.equals(getterTarget, other.getterTarget);
-    } else if (o instanceof Row) {
+      if (Objects.equals(getSchema(), other.getSchema())
+          && Objects.equals(getterTarget, other.getterTarget)) {
+        return true;
+      }
+    }
+    if (o instanceof Row) {
       return super.equals(o);
     }
     return false;
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(getSchema(), getterTarget);
+    return super.hashCode();
   }
 }

sdks/java/core/src/test/java/org/apache/beam/sdk/coders/RowCoderTest.java

@@ -645,4 +645,47 @@ public void testStaticEncoding() throws Exception {
 
     CoderProperties.coderDecodeEncodeEqual(RowCoder.of(schema), row);
   }
+
+  private static class MillisDateTimeLogicalType
+      implements LogicalType<Long, org.joda.time.ReadableInstant> {
+    @Override
+    public String getIdentifier() {
+      return "millis_datetime";
+    }
+
+    @Override
+    public FieldType getArgumentType() {
+      return FieldType.STRING;
+    }
+
+    @Override
+    public String getArgument() {
+      return "";
+    }
+
+    @Override
+    public FieldType getBaseType() {
+      return FieldType.DATETIME;
+    }
+
+    @Override
+    public org.joda.time.ReadableInstant toBaseType(Long input) {
+      return new DateTime(input, DateTimeZone.UTC);
+    }
+
+    @Override
+    public Long toInputType(org.joda.time.ReadableInstant base) {
+      return base.getMillis();
+    }
+  }
+
+  @Test
+  public void testDateTimeLogicalTypeCoderMethods() throws Exception {
+    Coder<Long> coder =
+        org.apache.beam.sdk.schemas.SchemaCoder.coderForFieldType(
+            FieldType.logicalType(new MillisDateTimeLogicalType()));
+    CoderProperties.coderDecodeEncodeEqual(coder, 123456789L);
+    CoderProperties.coderConsistentWithEquals(coder, 123456789L, 123456789L);
+    CoderProperties.structuralValueConsistentWithEquals(coder, 123456789L, 123456789L);
+  }
 }

sdks/java/core/src/test/java/org/apache/beam/sdk/schemas/JavaFieldSchemaTest.java

@@ -1071,4 +1071,35 @@ public void testGenericBeamSchemaArrayFromRow() throws Exception {
     assertEquals("v1", pojos[0].t);
     assertEquals("v2", pojos[1].t);
   }
+
+  @Test
+  public void testRowWithGettersEqualsAndHashCodeContract() throws NoSuchSchemaException {
+    SchemaRegistry registry = SchemaRegistry.createDefault();
+    SimplePOJO pojo1 = createSimple("string");
+    SimplePOJO pojo2 = createSimple("string");
+    Row rowWithGetters1 = registry.getToRowFunction(SimplePOJO.class).apply(pojo1);
+    Row rowWithGetters2 = registry.getToRowFunction(SimplePOJO.class).apply(pojo2);
+    Row rowWithStorage =
+        Row.withSchema(rowWithGetters1.getSchema())
+            .addValues(
+                "string",
+                (byte) 1,
+                (short) 2,
+                3,
+                4L,
+                true,
+                DATE,
+                INSTANT,
+                BYTE_ARRAY,
+                BYTE_BUFFER.array(),
+                BigDecimal.ONE,
+                "stringbuilder")
+            .build();
+
+    assertEquals(rowWithGetters1, rowWithGetters2);
+    assertEquals(rowWithGetters1.hashCode(), rowWithGetters2.hashCode());
+    assertEquals(rowWithStorage, rowWithGetters1);
+    assertEquals(rowWithGetters1, rowWithStorage);
+    assertEquals(rowWithStorage.hashCode(), rowWithGetters1.hashCode());
+  }
 }

sdks/java/core/src/test/java/org/apache/beam/sdk/util/RowStringInterpolatorTest.java

@@ -102,7 +102,7 @@ public void testInvalidRowThrowsHelpfulErrorForDoublyNestedFields() {
     Schema nestedSchema =
         Schema.builder().addNullableRowField("nested_row", doublyNestedSchema).build();
     Row invalidRow =
-        Row.withSchema(Schema.builder().addNullableRowField("row", doublyNestedSchema).build())
+        Row.withSchema(Schema.builder().addNullableRowField("row", nestedSchema).build())
             .addValue(
                 Row.withSchema(nestedSchema).addValue(Row.nullRow(doublyNestedSchema)).build())
             .build();

sdks/java/core/src/test/java/org/apache/beam/sdk/values/RowTest.java

@@ -896,4 +896,90 @@ public void testToCamelCase() {
     assertEquals(expectedCamelCaseInnerRow, row.toCamelCase().getRow("myRowField"));
     assertEquals(expectedCamelCaseRow, row.toCamelCase());
   }
+
+  @Test
+  public void testLogicalTypeBytesEqualsAndHashCode() {
+    Schema.LogicalType<ByteBuffer, byte[]> wrappedBytesType =
+        new Schema.LogicalType<ByteBuffer, byte[]>() {
+          @Override
+          public String getIdentifier() {
+            return "wrapped_bytes";
+          }
+
+          @Override
+          public FieldType getArgumentType() {
+            return FieldType.STRING;
+          }
+
+          @Override
+          public String getArgument() {
+            return "";
+          }
+
+          @Override
+          public FieldType getBaseType() {
+            return FieldType.BYTES;
+          }
+
+          @Override
+          public byte[] toBaseType(ByteBuffer input) {
+            byte[] bytes = new byte[input.remaining()];
+            input.duplicate().get(bytes);
+            return bytes;
+          }
+
+          @Override
+          public ByteBuffer toInputType(byte[] base) {
+            return ByteBuffer.wrap(base);
+          }
+        };
+
+    Schema schema = Schema.builder().addLogicalTypeField("f_bytes", wrappedBytesType).build();
+    Row row1 = Row.withSchema(schema).addValue(ByteBuffer.wrap(new byte[] {1, 2, 3, 4})).build();
+    Row row2 = Row.withSchema(schema).addValue(ByteBuffer.wrap(new byte[] {1, 2, 3, 4})).build();
+
+    assertEquals(row1, row2);
+    assertEquals(row1.hashCode(), row2.hashCode());
+  }
+
+  @Test
+  public void testRejectsMismatchedNestedRowSchema() {
+    Schema expectedNestedSchema = Schema.builder().addStringField("f_str").build();
+    Schema wrongNestedSchema = Schema.builder().addInt32Field("f_int").build();
+    Schema parentSchema = Schema.builder().addRowField("nested", expectedNestedSchema).build();
+
+    Row wrongNestedRow = Row.withSchema(wrongNestedSchema).addValue(42).build();
+    thrown.expect(IllegalArgumentException.class);
+    Row.withSchema(parentSchema).addValue(wrongNestedRow).build();
+  }
+
+  @Test
+  public void testRejectsMismatchedRowInArraySchema() {
+    Schema expectedNestedSchema = Schema.builder().addStringField("f_str").build();
+    Schema wrongNestedSchema = Schema.builder().addInt32Field("f_int").build();
+    Schema parentSchema =
+        Schema.builder().addArrayField("items", FieldType.row(expectedNestedSchema)).build();
+
+    Row wrongNestedRow = Row.withSchema(wrongNestedSchema).addValue(42).build();
+    thrown.expect(IllegalArgumentException.class);
+    Row.withSchema(parentSchema).addArray(ImmutableList.of(wrongNestedRow)).build();
+  }
+
+  @Test
+  public void testSlicedAndReadOnlyByteBufferEquality() {
+    byte[] expected = new byte[] {2, 3, 4};
+    ByteBuffer sliced = ByteBuffer.wrap(new byte[] {99, 2, 3, 4, 99}, 1, 3).slice();
+    ByteBuffer readOnly = ByteBuffer.wrap(expected).asReadOnlyBuffer();
+
+    Schema schema = Schema.of(Schema.Field.of("bytes", Schema.FieldType.BYTES));
+
+    Row rowFromBytes = Row.withSchema(schema).addValue(expected).build();
+    Row rowFromSlice = Row.withSchema(schema).addValue(sliced).build();
+    Row rowFromReadOnly = Row.withSchema(schema).withFieldValue("bytes", readOnly).build();
+
+    assertEquals(rowFromBytes, rowFromSlice);
+    assertEquals(rowFromBytes, rowFromReadOnly);
+    assertTrue(Arrays.equals(expected, rowFromSlice.getBytes(0)));
+    assertTrue(Arrays.equals(expected, rowFromReadOnly.getBytes(0)));
+  }
 }