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+
19+ package org .apache .wayang .java .operators ;
20+
21+ import org .apache .avro .Schema ;
22+ import org .apache .avro .SchemaBuilder ;
23+ import org .apache .avro .generic .GenericData ;
24+ import org .apache .avro .generic .GenericRecord ;
25+ import org .apache .hadoop .conf .Configuration ;
26+ import org .apache .hadoop .fs .FileSystem ;
27+ import org .apache .hadoop .fs .Path ;
28+ import org .apache .parquet .avro .AvroParquetWriter ;
29+ import org .apache .parquet .hadoop .ParquetWriter ;
30+ import org .apache .parquet .hadoop .metadata .CompressionCodecName ;
31+
32+ import org .apache .wayang .basic .data .Record ;
33+ import org .apache .wayang .basic .operators .ParquetSink ;
34+ import org .apache .wayang .basic .types .RecordType ;
35+ import org .apache .wayang .core .optimizer .OptimizationContext ;
36+ import org .apache .wayang .core .plan .wayangplan .ExecutionOperator ;
37+ import org .apache .wayang .core .platform .ChannelDescriptor ;
38+ import org .apache .wayang .core .platform .ChannelInstance ;
39+ import org .apache .wayang .core .platform .lineage .ExecutionLineageNode ;
40+ import org .apache .wayang .core .types .DataSetType ;
41+ import org .apache .wayang .core .util .Tuple ;
42+ import org .apache .wayang .java .channels .CollectionChannel ;
43+ import org .apache .wayang .java .channels .StreamChannel ;
44+ import org .apache .wayang .java .execution .JavaExecutor ;
45+ import org .apache .wayang .java .platform .JavaPlatform ;
46+
47+ import java .io .IOException ;
48+ import java .math .BigDecimal ;
49+ import java .sql .Timestamp ;
50+ import java .util .Arrays ;
51+ import java .util .Collection ;
52+ import java .util .List ;
53+ import java .util .stream .Collectors ;
54+
55+ /**
56+ * Writes {@link Record}s to a Parquet file using the Java platform.
57+ */
58+ public class JavaParquetSink extends ParquetSink implements JavaExecutionOperator {
59+
60+ private static final int SCHEMA_SAMPLE_SIZE = 50 ;
61+
62+ public JavaParquetSink (ParquetSink that ) {
63+ super (that .getOutputUrl (), that .isOverwrite (), that .prefersDataset (), that .getType ());
64+ }
65+
66+ @ Override
67+ public Tuple <Collection <ExecutionLineageNode >, Collection <ChannelInstance >> evaluate (
68+ ChannelInstance [] inputs ,
69+ ChannelInstance [] outputs ,
70+ JavaExecutor javaExecutor ,
71+ OptimizationContext .OperatorContext operatorContext ) {
72+
73+ assert inputs .length == 1 ;
74+ assert outputs .length == 0 ;
75+
76+ // Get the input stream and collect all records into a list
77+ final List <Record > records = this .getRecords (inputs [0 ]);
78+
79+ if (records .isEmpty ()) {
80+ return ExecutionOperator .modelEagerExecution (inputs , outputs , operatorContext );
81+ }
82+
83+ try {
84+ // Handle overwrite — delete existing file if needed
85+ Path outputPath = new Path (this .getOutputUrl ());
86+ Configuration conf = new Configuration ();
87+ if (this .isOverwrite ()) {
88+ FileSystem fs = outputPath .getFileSystem (conf );
89+ fs .delete (outputPath , true );
90+ }
91+
92+ // Infer schema from RecordType + sampled records
93+ Schema schema = this .inferSchema (records );
94+
95+ // Write records as Parquet
96+ try (ParquetWriter <GenericRecord > writer = AvroParquetWriter .<GenericRecord >builder (outputPath )
97+ .withSchema (schema )
98+ .withConf (conf )
99+ .withCompressionCodec (CompressionCodecName .SNAPPY )
100+ .build ()) {
101+
102+ for (Record record : records ) {
103+ writer .write (this .convertToGenericRecord (record , schema ));
104+ }
105+ }
106+
107+ } catch (IOException e ) {
108+ throw new RuntimeException ("Failed to write Parquet file: " + this .getOutputUrl (), e );
109+ }
110+
111+ return ExecutionOperator .modelEagerExecution (inputs , outputs , operatorContext );
112+ }
113+
114+ /**
115+ * Extracts records from the input channel, handling both Stream and Collection channels.
116+ */
117+ private List <Record > getRecords (ChannelInstance input ) {
118+ if (input instanceof CollectionChannel .Instance ) {
119+ return ((CollectionChannel .Instance ) input ).<Record >provideCollection ()
120+ .stream ().collect (Collectors .toList ());
121+ }
122+ return ((StreamChannel .Instance ) input ).<Record >provideStream ()
123+ .collect (Collectors .toList ());
124+ }
125+
126+ /**
127+ * Infers an Avro schema from the RecordType (if available) and sampled record values.
128+ */
129+ private Schema inferSchema (List <Record > records ) {
130+ String [] fieldNames = this .resolveFieldNames (records );
131+ List <Record > samples = records .subList (0 , Math .min (SCHEMA_SAMPLE_SIZE , records .size ()));
132+
133+ SchemaBuilder .FieldAssembler <Schema > fields = SchemaBuilder
134+ .record ("WayangRecord" )
135+ .namespace ("org.apache.wayang" )
136+ .fields ();
137+
138+ for (int i = 0 ; i < fieldNames .length ; i ++) {
139+ Schema .Type avroType = this .inferColumnType (samples , i );
140+ // Make fields nullable — union of [null, type]
141+ Schema fieldSchema = Schema .createUnion (
142+ Schema .create (Schema .Type .NULL ),
143+ Schema .create (avroType )
144+ );
145+ fields .name (fieldNames [i ]).type (fieldSchema ).noDefault ();
146+ }
147+
148+ return fields .endRecord ();
149+ }
150+
151+ /**
152+ * Resolves field names from RecordType if available, otherwise generates field0, field1, etc.
153+ */
154+ private String [] resolveFieldNames (List <Record > records ) {
155+ DataSetType <Record > dataSetType = this .getType ();
156+ if (dataSetType != null && dataSetType .getDataUnitType () instanceof RecordType ) {
157+ RecordType recordType = (RecordType ) dataSetType .getDataUnitType ();
158+ if (recordType .getFieldNames () != null && recordType .getFieldNames ().length > 0 ) {
159+ return recordType .getFieldNames ();
160+ }
161+ }
162+
163+ // Fallback: generate generic field names
164+ int numFields = records .get (0 ).size ();
165+ String [] names = new String [numFields ];
166+ for (int i = 0 ; i < numFields ; i ++) {
167+ names [i ] = "field" + i ;
168+ }
169+ return names ;
170+ }
171+
172+ /**
173+ * Infers the Avro type for a column by sampling record values.
174+ */
175+ private Schema .Type inferColumnType (List <Record > samples , int columnIndex ) {
176+ for (Record sample : samples ) {
177+ if (sample == null || columnIndex >= sample .size ()) {
178+ continue ;
179+ }
180+ Object value = sample .getField (columnIndex );
181+ if (value == null ) {
182+ continue ;
183+ }
184+ return this .toAvroType (value );
185+ }
186+ // Default to string if all values are null
187+ return Schema .Type .STRING ;
188+ }
189+
190+ /**
191+ * Maps a Java value to an Avro schema type.
192+ */
193+ private Schema .Type toAvroType (Object value ) {
194+ if (value instanceof String || value instanceof Character ) {
195+ return Schema .Type .STRING ;
196+ } else if (value instanceof Integer ) {
197+ return Schema .Type .INT ;
198+ } else if (value instanceof Long || value instanceof Timestamp ) {
199+ return Schema .Type .LONG ;
200+ } else if (value instanceof Float ) {
201+ return Schema .Type .FLOAT ;
202+ } else if (value instanceof Double || value instanceof BigDecimal ) {
203+ return Schema .Type .DOUBLE ;
204+ } else if (value instanceof Boolean ) {
205+ return Schema .Type .BOOLEAN ;
206+ } else if (value instanceof byte []) {
207+ return Schema .Type .BYTES ;
208+ }
209+ return Schema .Type .STRING ;
210+ }
211+
212+ /**
213+ * Converts a Wayang Record to an Avro GenericRecord using the given schema.
214+ */
215+ private GenericRecord convertToGenericRecord (Record record , Schema schema ) {
216+ GenericRecord genericRecord = new GenericData .Record (schema );
217+ List <Schema .Field > fields = schema .getFields ();
218+ for (int i = 0 ; i < fields .size (); i ++) {
219+ Object value = i < record .size () ? record .getField (i ) : null ;
220+ // Convert value to match the Avro type if needed
221+ if (value != null ) {
222+ value = this .convertValue (value , fields .get (i ).schema ());
223+ }
224+ genericRecord .put (fields .get (i ).name (), value );
225+ }
226+ return genericRecord ;
227+ }
228+
229+ /**
230+ * Converts a value to match the expected Avro schema type.
231+ */
232+ private Object convertValue (Object value , Schema fieldSchema ) {
233+ // Handle nullable union types — extract the actual type
234+ Schema actualSchema = fieldSchema ;
235+ if (fieldSchema .getType () == Schema .Type .UNION ) {
236+ for (Schema s : fieldSchema .getTypes ()) {
237+ if (s .getType () != Schema .Type .NULL ) {
238+ actualSchema = s ;
239+ break ;
240+ }
241+ }
242+ }
243+
244+ switch (actualSchema .getType ()) {
245+ case STRING :
246+ return value .toString ();
247+ case INT :
248+ return value instanceof Number ? ((Number ) value ).intValue () : Integer .parseInt (value .toString ());
249+ case LONG :
250+ if (value instanceof Timestamp ) return ((Timestamp ) value ).getTime ();
251+ return value instanceof Number ? ((Number ) value ).longValue () : Long .parseLong (value .toString ());
252+ case FLOAT :
253+ return value instanceof Number ? ((Number ) value ).floatValue () : Float .parseFloat (value .toString ());
254+ case DOUBLE :
255+ return value instanceof Number ? ((Number ) value ).doubleValue () : Double .parseDouble (value .toString ());
256+ case BOOLEAN :
257+ return value instanceof Boolean ? value : Boolean .parseBoolean (value .toString ());
258+ default :
259+ return value ;
260+ }
261+ }
262+
263+ @ Override
264+ public List <ChannelDescriptor > getSupportedInputChannels (int index ) {
265+ return Arrays .asList (CollectionChannel .DESCRIPTOR , StreamChannel .DESCRIPTOR );
266+ }
267+
268+ @ Override
269+ public List <ChannelDescriptor > getSupportedOutputChannels (int index ) {
270+ throw new UnsupportedOperationException ("This operator has no outputs." );
271+ }
272+
273+ @ Override
274+ public JavaPlatform getPlatform () {
275+ return JavaPlatform .getInstance ();
276+ }
277+ }
0 commit comments