/*
 * 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.runners.spark.structuredstreaming.translation.batch;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import org.apache.beam.runners.core.construction.PipelineOptionsSerializationUtils;
import org.apache.beam.runners.core.construction.ReadTranslation;
import org.apache.beam.runners.core.serialization.Base64Serializer;
import org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator;
import org.apache.beam.runners.spark.structuredstreaming.translation.TranslationContext;
import org.apache.beam.sdk.io.BoundedSource;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.spark.api.java.function.MapFunction;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;

class ReadSourceTranslatorBatch<T>
    implements TransformTranslator<PTransform<PBegin, PCollection<T>>> {

  private static String sourceProviderClass = DatasetSourceBatch.class.getCanonicalName();

  @SuppressWarnings("unchecked")
  @Override
  public void translateTransform(
      PTransform<PBegin, PCollection<T>> transform, TranslationContext context) {
    AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>> rootTransform =
        (AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>>)
            context.getCurrentTransform();

        BoundedSource<T> source;
    try {
      source = ReadTranslation.boundedSourceFromTransform(rootTransform);
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
    SparkSession sparkSession = context.getSparkSession();

    String serializedSource = Base64Serializer.serializeUnchecked(source);
    Dataset<Row> rowDataset = sparkSession.read().format(sourceProviderClass)
        .option(DatasetSourceBatch.BEAM_SOURCE_OPTION, serializedSource)
        .option(DatasetSourceBatch.DEFAULT_PARALLELISM,
            String.valueOf(context.getSparkSession().sparkContext().defaultParallelism()))
        .option(DatasetSourceBatch.PIPELINE_OPTIONS,
            PipelineOptionsSerializationUtils.serializeToJson(context.getOptions())).load();

    // extract windowedValue from Row
    MapFunction<Row, WindowedValue> func = new MapFunction<Row, WindowedValue>() {
      @Override public WindowedValue call(Row value) throws Exception {
        //there is only one value put in each Row by the InputPartitionReader
        byte[] bytes = (byte[]) value.get(0);
        WindowedValue.FullWindowedValueCoder<T> windowedValueCoder = WindowedValue.FullWindowedValueCoder
            .of(source.getOutputCoder(), GlobalWindow.Coder.INSTANCE);
        WindowedValue<T> windowedValue = windowedValueCoder.decode(new ByteArrayInputStream(bytes));
        return windowedValue;
      }
    };
    //TODO: is there a better way than using the raw WindowedValue? Can an Encoder<WindowedValue<T>>
    // be created ?
    Dataset<WindowedValue> dataset = rowDataset.map(func, Encoders.kryo(WindowedValue.class));

    PCollection<T> output = (PCollection<T>) context.getOutput();
    context.putDatasetRaw(output, dataset);
  }
}

@@ -1160,16 +1160,101 @@
 alizer;%0A
+import org.apache.beam.runners.spark.structuredstreaming.translation.EncoderHelpers;%0A
 import o
@@ -3411,16 +3411,19 @@
 wedValue
+%3CT%3E
 %3E func =
@@ -3457,16 +3457,19 @@
 wedValue
+%3CT%3E
 %3E() %7B%0A  
@@ -3502,16 +3502,19 @@
 wedValue
+%3CT%3E
  call(Ro
@@ -3996,129 +3996,8 @@
     
-//TODO: is there a better way than using the raw WindowedValue? Can an Encoder%3CWindowedValue%3CT%3E%3E%0A    // be created ?%0A    
 Data
@@ -4013,16 +4013,19 @@
 wedValue
+%3CT%3E
 %3E datase
@@ -4060,34 +4060,24 @@
 oder
-s.kryo(WindowedValue.class
+Helpers.encoder(
 ));%0A
@@ -4143,16 +4143,16 @@
 tput();%0A
+
     cont
@@ -4169,11 +4169,8 @@
 aset
-Raw
 (out

Fix type checking with Encoder of WindowedValue<T>