In SMB and ParquetAvroIOTap set GenericDataSupplier and read schemas #5121

shnapz · 2023-12-14T23:46:35Z

Addressing:

shnapz · 2023-12-14T23:47:21Z

scio-parquet/src/main/scala/com/spotify/scio/parquet/avro/ParquetAvroSink.scala

@@ -0,0 +1,38 @@
+package com.spotify.scio.parquet.avro


moved this class out from package object

shnapz · 2023-12-14T23:48:13Z

scio-parquet/src/main/scala/com/spotify/scio/parquet/avro/ParquetAvroTap.scala

+    val filePattern = ScioUtil.filePattern(path, params.suffix)
+    val conf = ParquetConfiguration
+      .ofNullable(params.conf)
+    conf.setReadSchemas(params)


Except of these two lines, the rest of the class remains the same, just moved

shnapz · 2023-12-14T23:49:37Z

scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/AvroSortedBucketIO.java

@@ -241,21 +241,24 @@ public Read<T> withPredicate(Predicate<T> predicate) {
      return toBuilder().setPredicate(predicate).build();
    }

+    @SuppressWarnings("unchecked")
+    @Override
+    FileOperations<T> getFileOperations() {


Added this to reflect the same style as in Write and TransformOutput

shnapz · 2023-12-14T23:51:24Z

scio-smb/src/test/java/org/apache/beam/sdk/extensions/smb/ParquetAvroFileOperationsTest.java

  @Test
  public void testSpecificRecord() throws Exception {
    final ParquetAvroFileOperations<AvroGeneratedUser> fileOperations =
-        ParquetAvroFileOperations.of(AvroGeneratedUser.getClassSchema());
+        ParquetAvroFileOperations.of(AvroGeneratedUser.class);


In Prod we use the overload with Schema for GenericRecords only

codecov · 2023-12-16T00:17:37Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (41c8df7) 62.99% compared to head (1ba9610) 63.06%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #5121      +/-   ##
==========================================
+ Coverage   62.99%   63.06%   +0.06%     
==========================================
  Files         298      300       +2     
  Lines       10946    10945       -1     
  Branches      754      760       +6     
==========================================
+ Hits         6895     6902       +7     
+ Misses       4051     4043       -8

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

shnapz · 2023-12-18T15:26:01Z

scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/ParquetAvroSortedBucketIO.java

-      final Schema schema =
-          getRecordClass() == null
-              ? getSchema()
-              : new ReflectData(getRecordClass().getClassLoader()).getSchema(getRecordClass());


there is identical resolution of schema through ReflectData in ParquetAvroFileOperations

clairemcginty · 2023-12-18T18:58:56Z

scio-parquet/src/main/scala/com/spotify/scio/parquet/avro/package.scala

+
+  implicit class ConfigurationSyntax(conf: Configuration) {
+
+    def setReadSchemas[A, T](params: ParquetAvroIO.ReadParam[A, T]): Unit = {


What do you think of broadening the scope of this method and moving it into a class method of ParquetAvroIO.ReadParam? Like:

final case class ReadParam[A: ClassTag, T: ClassTag] private ( projectionFn: A => T, projection: Schema = ReadParam.DefaultProjection, predicate: FilterPredicate = ReadParam.DefaultPredicate, conf: Configuration = ReadParam.DefaultConfiguration, suffix: String = null ) { val avroClass: Class[A] = ScioUtil.classOf[A] val isSpecific = classOf[SpecificRecord] isAssignableFrom ScioUtil.classOf[T] private[avro] def getReadConfiguration(): Configuration = { val c = ParquetConfiguration.ofNullable(conf) // Set schema/projection val readSchema: Schema = if (isSpecific) ReflectData.get().getSchema(avroClass) else params.projection AvroReadSupport.setAvroReadSchema(readSchema) if (params.projection != null) { /* set projection... */ } if (params.predicate != null) { /* set predicate */ } // Set data supplier if (!isSpecific && conf.get(AvroReadSupport.AVRO_DATA_SUPPLIER) == null) { .... } }

That way it's easy to access from both read and tap methods 👍

This is great! I missed it, the great place for common functions on input conf!

RustedBones · 2023-12-19T14:40:12Z

scio-parquet/src/main/scala/com/spotify/scio/parquet/avro/ParquetAvroIO.scala

@@ -219,13 +212,27 @@ object ParquetAvroIO {
      }
    }

+    def setReadSchemas[A, T](params: ParquetAvroIO.ReadParam[A, T]): Unit = {


ReadParam is already parameterized with [A, T]. I don't think we want type shadowing here

Removing this function

RustedBones · 2023-12-19T14:41:06Z

scio-parquet/src/main/scala/com/spotify/scio/parquet/avro/ParquetAvroTap.scala

+    val filePattern = ScioUtil.filePattern(path, params.suffix)
+    val conf = ParquetConfiguration
+      .ofNullable(params.conf)
+    params.setReadSchemas(params)


do we want to do this ?

probably was tired when committed this, I am changing the approach

RustedBones

I see that using an empty config as parameter allows some simplification, but I'd still prefer to pass a null there for consistency.

This is a personal preference, not a blocker. @clairemcginty might also have an opinion about it.

RustedBones · 2024-01-17T08:18:03Z

scio-parquet/src/main/scala/com/spotify/scio/parquet/avro/ParquetAvroIO.scala

@@ -159,7 +145,7 @@ object ParquetAvroIO {
  object ReadParam {
    val DefaultProjection: Schema = null
    val DefaultPredicate: FilterPredicate = null
-    val DefaultConfiguration: Configuration = null
+    def DefaultConfiguration: Configuration = ParquetConfiguration.empty()


TBH, I'd prefer keeping this as null and handle it internally

Agree, I think, partly because it matches other ScioIO param conventions and partly because I'm worried that over time we'll forget why this was made a def and refactor it back into a val :)

I wanted to challenge this approach :) but fully agree, it is so easy to introduce a bug. As an alternative would be full clone of ParquetConfiguration after used passed it, because we modify this instance and we don't know what else a user does

RustedBones · 2024-01-17T08:21:07Z

scio-parquet/src/main/scala/com/spotify/scio/parquet/avro/ParquetAvroIO.scala

+          Functions.serializableFn { x =>
+            cleanedProjectionFn(x)
+          }


Nit. can we revert to previous syntax?

clairemcginty · 2024-01-17T13:35:45Z

scio-parquet/src/main/scala/com/spotify/scio/parquet/avro/ParquetAvroIO.scala

-          current = reader.read()
-          r
-        }
+    private[avro] def setupConfigAndGetSchema[T: ClassTag](): Schema = {


I think I would prefer getting rid of this helper method an just inlining it in ParquetAvroIO::write as before. Especially since once we merge the Parquet 1.13.1 upgrade, we'll be able to remove all the logical type checks, so there won't be much logic left :D

Alternately, we could convert it it into a method that just gets the Schema for a type T:

private[avro] def getSchemaForClass[T: ClassTag](): Schema = { val avroClass = ScioUtil.classOf[T] val isSpecific: Boolean = classOf[SpecificRecord] isAssignableFrom avroClass if (isSpecific) ReflectData.get().getSchema(avroClass) else schema }

since that logic is reused a bit in this file :)

Agreed. In case of write it is used only once, but I aimed to have similarity with read param handling. I will inline

scio-parquet/src/main/scala/com/spotify/scio/parquet/avro/ParquetAvroTap.scala

scio-parquet/src/main/scala/com/spotify/scio/parquet/avro/ParquetAvroSink.scala

clairemcginty · 2024-01-17T13:39:39Z

scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/ParquetAvroFileOperations.java

+  public static <V extends IndexedRecord> ParquetAvroFileOperations<V> of(
+      Schema schema, CompressionCodecName compression, FilterPredicate predicate,
+      Configuration conf) {
+    if (conf.get(AvroReadSupport.AVRO_DATA_SUPPLIER) == null) {


Prefer setting this in ParquetAvroReader::prepareReader, similar to how we set model in ParquetAvroSink::open -- that way this logic doesn't depend on the user choosing this constructor specifically :)

Oh right, with the help of newly introduced private final Class<ValueT> recordClass; we can do that

shnapz · 2024-01-18T03:55:04Z

scio-parquet/src/test/scala/com/spotify/scio/parquet/avro/ParquetAvroIOTest.scala

-      _.parquetAvroFile[TestRecord](_).map(identity)
-    )(_.saveAsParquetAvroFile(_))
+
+    forAllCases(readConfigs) { case (c, _) =>


the main idea of changes in this file is to test both SDR and Legacy in the same test cases

shnapz commented Dec 14, 2023

View reviewed changes

shnapz commented Dec 18, 2023

View reviewed changes

shnapz marked this pull request as ready for review December 18, 2023 15:27

shnapz requested a review from clairemcginty December 18, 2023 15:27

clairemcginty reviewed Dec 18, 2023

View reviewed changes

shnapz requested a review from clairemcginty December 18, 2023 21:40

shnapz force-pushed the akabas/ParquetAvro-dataSupplier branch from 1673bf2 to b952934 Compare December 18, 2023 23:01

RustedBones reviewed Dec 19, 2023

View reviewed changes

shnapz requested a review from RustedBones January 3, 2024 23:56

shnapz force-pushed the akabas/ParquetAvro-dataSupplier branch from aa95561 to 1b8ddac Compare January 17, 2024 01:24

RustedBones reviewed Jan 17, 2024

View reviewed changes

clairemcginty reviewed Jan 17, 2024

View reviewed changes

RustedBones added this to the 0.14.0 milestone Jan 17, 2024

shnapz added 11 commits January 17, 2024 17:24

In SMB and ParquetAvroIOTap set GenericDataSupplier and read schemas

8dc1a5a

fixing build

d82e1b4

fix build

0e24e56

Address PR comments

f66a2cb

fix

7aed7a9

Refactored AvroParquet ReadParams and improved tests

d35fd13

fix

6b09630

fix PR comments

70ccfa9

fix 3

f6c2b51

Added more test cases

7ad2840

fix tests

338dc55

shnapz force-pushed the akabas/ParquetAvro-dataSupplier branch from 0cbfe14 to 338dc55 Compare January 18, 2024 03:51

shnapz commented Jan 18, 2024

View reviewed changes

fix cleanup

1ba9610

shnapz requested review from RustedBones and clairemcginty January 18, 2024 04:06

RustedBones approved these changes Jan 18, 2024

View reviewed changes

RustedBones merged commit b48d9d6 into main Jan 18, 2024
11 checks passed

RustedBones deleted the akabas/ParquetAvro-dataSupplier branch January 18, 2024 08:04

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

In SMB and ParquetAvroIOTap set GenericDataSupplier and read schemas #5121

In SMB and ParquetAvroIOTap set GenericDataSupplier and read schemas #5121

shnapz commented Dec 14, 2023

shnapz Dec 14, 2023

shnapz Dec 14, 2023

shnapz Dec 14, 2023

shnapz Dec 14, 2023

codecov bot commented Dec 16, 2023 •

edited

Loading

shnapz Dec 18, 2023

clairemcginty Dec 18, 2023 •

edited

Loading

shnapz Dec 18, 2023

RustedBones Dec 19, 2023

shnapz Dec 21, 2023

RustedBones Dec 19, 2023

shnapz Dec 21, 2023

RustedBones left a comment

RustedBones Jan 17, 2024

clairemcginty Jan 17, 2024

shnapz Jan 17, 2024 •

edited

Loading

RustedBones Jan 17, 2024

clairemcginty Jan 17, 2024

shnapz Jan 17, 2024

clairemcginty Jan 17, 2024

shnapz Jan 17, 2024

shnapz Jan 18, 2024


		implicit class ConfigurationSyntax(conf: Configuration) {

		def setReadSchemas[A, T](params: ParquetAvroIO.ReadParam[A, T]): Unit = {

In SMB and ParquetAvroIOTap set GenericDataSupplier and read schemas #5121

In SMB and ParquetAvroIOTap set GenericDataSupplier and read schemas #5121

Conversation

shnapz commented Dec 14, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented Dec 16, 2023 • edited Loading

Codecov Report

Choose a reason for hiding this comment

clairemcginty Dec 18, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

RustedBones left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shnapz Jan 17, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented Dec 16, 2023 •

edited

Loading

clairemcginty Dec 18, 2023 •

edited

Loading

shnapz Jan 17, 2024 •

edited

Loading