<commit_before>/*! ******************************************************************************
 *
 * Pentaho Data Integration
 *
 * Copyright (C) 2017 by <NAME> : http://www.pentaho.com
 *
 *******************************************************************************
 *
 * Licensed 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.pentaho.hadoop.shim.common.format.orc;

import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hive.common.type.HiveDecimal;
import org.apache.hadoop.hive.ql.exec.vector.BytesColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.ColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.DecimalColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.DoubleColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.LongColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.TimestampColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.VectorizedRowBatch;
import org.apache.hadoop.hive.serde2.io.HiveDecimalWritable;
import org.apache.log4j.Logger;
import org.apache.orc.OrcFile;
import org.apache.orc.TypeDescription;
import org.apache.orc.Writer;
import org.pentaho.di.core.RowMetaAndData;
import org.apache.hadoop.conf.Configuration;
import org.pentaho.di.core.exception.KettleValueException;
import org.pentaho.di.core.plugins.IValueMetaConverter;
import org.pentaho.di.core.row.RowMeta;
import org.pentaho.di.core.row.ValueMetaInterface;
import org.pentaho.di.core.row.value.ValueMetaBase;
import org.pentaho.di.core.row.value.ValueMetaBigNumber;
import org.pentaho.di.core.row.value.ValueMetaBinary;
import org.pentaho.di.core.row.value.ValueMetaBoolean;
import org.pentaho.di.core.row.value.ValueMetaConversionException;
import org.pentaho.di.core.row.value.ValueMetaConverter;
import org.pentaho.di.core.row.value.ValueMetaDate;
import org.pentaho.di.core.row.value.ValueMetaInteger;
import org.pentaho.di.core.row.value.ValueMetaInternetAddress;
import org.pentaho.di.core.row.value.ValueMetaNumber;
import org.pentaho.di.core.row.value.ValueMetaString;
import org.pentaho.di.core.row.value.ValueMetaTimestamp;
import org.pentaho.hadoop.shim.api.format.IPentahoOutputFormat;
import org.pentaho.hadoop.shim.api.format.SchemaDescription;

import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.math.BigDecimal;
import java.nio.ByteBuffer;
import java.sql.Timestamp;
import java.text.ParseException;
import java.text.SimpleDateFormat;
import java.time.LocalDate;
import java.time.ZoneId;
import java.time.temporal.ChronoUnit;
import java.util.Date;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Created by tkafalas on 11/3/2017.
 */
public class PentahoOrcRecordWriter implements IPentahoOutputFormat.IPentahoRecordWriter {
  private final SchemaDescription schemaDescription;
  private final TypeDescription schema;
  private Configuration conf;
  private VectorizedRowBatch batch;
  private int batchRowNumber;
  private Writer writer;
  private final IValueMetaConverter valueMetaConverter = new ValueMetaConverter();
  private final String dateFormatString = ValueMetaBase.DEFAULT_DATE_FORMAT_MASK; //TODO: Get this from UI
  private SimpleDateFormat datePattern = new SimpleDateFormat( dateFormatString );
  private RowMeta outputRowMeta = new RowMeta();
  private RowMetaAndData outputRowMetaAndData;
  private boolean isFirstRow = true;
  private static final Logger logger = Logger.getLogger( PentahoOrcRecordWriter.class );

  public int batchMaxSize = 10000;

  public PentahoOrcRecordWriter( SchemaDescription schemaDescription, TypeDescription schema, String filePath,
                                 Configuration conf ) {
    this.schemaDescription = schemaDescription;
    this.schema = schema;

    final AtomicInteger fieldNumber = new AtomicInteger();  //Mutable field count
    schemaDescription.forEach( field -> setOutputMeta( fieldNumber, field ) );
    outputRowMetaAndData = new RowMetaAndData( outputRowMeta, new Object[ fieldNumber.get() ] );

    try {
      writer = OrcFile.createWriter( new Path( filePath ),
        OrcFile.writerOptions( conf )
          .setSchema( schema ) );
      batch = schema.createRowBatch();
    } catch ( IOException e ) {
      logger.error( e );
    }

    //Write the addition metadata for the fields
    new OrcMetaDataWriter( writer ).write( schemaDescription );
  }

  private void setOutputMeta( AtomicInteger fieldNumber, SchemaDescription.Field field ) {
    outputRowMeta.addValueMeta( getValueMetaInterface( field.pentahoFieldName, field.pentahoValueMetaType ) );
    fieldNumber.getAndIncrement();
  }

  @Override public void write( RowMetaAndData row ) throws Exception {
    final AtomicInteger fieldNumber = new AtomicInteger();  //Mutable field count
    batchRowNumber = batch.size++;

    schemaDescription.forEach( field -> setFieldValue( fieldNumber, field, row ) );
    if ( batch.size == batchMaxSize ) {
      writer.addRowBatch( batch );
      batch.reset();
    }
  }

  private void setFieldValue( AtomicInteger fieldNumber, SchemaDescription.Field field,
                              RowMetaAndData rowMetaAndData ) {
    int fieldNo = fieldNumber.getAndIncrement();
    ColumnVector columnVector = batch.cols[ fieldNo ];
    int rowMetaIndex = rowMetaAndData.getRowMeta().indexOfValue( field.pentahoFieldName );

    if ( rowMetaAndData.getData()[ rowMetaIndex ] == null ) {
      if ( field.allowNull ) {<commit_msg> [BACKLOG-20030] - ORC does not support user defined batch max size. Using getMaxSize to get the max batch size is supports<commit_after>/*! ******************************************************************************
 *
 * Pentaho Data Integration
 *
 * Copyright (C) 2017 by <NAME> : http://www.pentaho.com
 *
 *******************************************************************************
 *
 * Licensed 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.pentaho.hadoop.shim.common.format.orc;

import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hive.common.type.HiveDecimal;
import org.apache.hadoop.hive.ql.exec.vector.BytesColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.ColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.DecimalColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.DoubleColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.LongColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.TimestampColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.VectorizedRowBatch;
import org.apache.hadoop.hive.serde2.io.HiveDecimalWritable;
import org.apache.log4j.Logger;
import org.apache.orc.OrcFile;
import org.apache.orc.TypeDescription;
import org.apache.orc.Writer;
import org.pentaho.di.core.RowMetaAndData;
import org.apache.hadoop.conf.Configuration;
import org.pentaho.di.core.exception.KettleValueException;
import org.pentaho.di.core.plugins.IValueMetaConverter;
import org.pentaho.di.core.row.RowMeta;
import org.pentaho.di.core.row.ValueMetaInterface;
import org.pentaho.di.core.row.value.ValueMetaBase;
import org.pentaho.di.core.row.value.ValueMetaBigNumber;
import org.pentaho.di.core.row.value.ValueMetaBinary;
import org.pentaho.di.core.row.value.ValueMetaBoolean;
import org.pentaho.di.core.row.value.ValueMetaConversionException;
import org.pentaho.di.core.row.value.ValueMetaConverter;
import org.pentaho.di.core.row.value.ValueMetaDate;
import org.pentaho.di.core.row.value.ValueMetaInteger;
import org.pentaho.di.core.row.value.ValueMetaInternetAddress;
import org.pentaho.di.core.row.value.ValueMetaNumber;
import org.pentaho.di.core.row.value.ValueMetaString;
import org.pentaho.di.core.row.value.ValueMetaTimestamp;
import org.pentaho.hadoop.shim.api.format.IPentahoOutputFormat;
import org.pentaho.hadoop.shim.api.format.SchemaDescription;

import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.math.BigDecimal;
import java.nio.ByteBuffer;
import java.sql.Timestamp;
import java.text.ParseException;
import java.text.SimpleDateFormat;
import java.time.LocalDate;
import java.time.ZoneId;
import java.time.temporal.ChronoUnit;
import java.util.Date;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Created by tkafalas on 11/3/2017.
 */
public class PentahoOrcRecordWriter implements IPentahoOutputFormat.IPentahoRecordWriter {
  private final SchemaDescription schemaDescription;
  private final TypeDescription schema;
  private Configuration conf;
  private VectorizedRowBatch batch;
  private int batchRowNumber;
  private Writer writer;
  private final IValueMetaConverter valueMetaConverter = new ValueMetaConverter();
  private final String dateFormatString = ValueMetaBase.DEFAULT_DATE_FORMAT_MASK; //TODO: Get this from UI
  private SimpleDateFormat datePattern = new SimpleDateFormat( dateFormatString );
  private RowMeta outputRowMeta = new RowMeta();
  private RowMetaAndData outputRowMetaAndData;
  private boolean isFirstRow = true;
  private static final Logger logger = Logger.getLogger( PentahoOrcRecordWriter.class );

  public PentahoOrcRecordWriter( SchemaDescription schemaDescription, TypeDescription schema, String filePath,
                                 Configuration conf ) {
    this.schemaDescription = schemaDescription;
    this.schema = schema;

    final AtomicInteger fieldNumber = new AtomicInteger();  //Mutable field count
    schemaDescription.forEach( field -> setOutputMeta( fieldNumber, field ) );
    outputRowMetaAndData = new RowMetaAndData( outputRowMeta, new Object[ fieldNumber.get() ] );

    try {
      writer = OrcFile.createWriter( new Path( filePath ),
        OrcFile.writerOptions( conf )
          .setSchema( schema ) );
      batch = schema.createRowBatch();
    } catch ( IOException e ) {
      logger.error( e );
    }

    //Write the addition metadata for the fields
    new OrcMetaDataWriter( writer ).write( schemaDescription );
  }

  private void setOutputMeta( AtomicInteger fieldNumber, SchemaDescription.Field field ) {
    outputRowMeta.addValueMeta( getValueMetaInterface( field.pentahoFieldName, field.pentahoValueMetaType ) );
    fieldNumber.getAndIncrement();
  }

  @Override public void write( RowMetaAndData row ) throws Exception {
    final AtomicInteger fieldNumber = new AtomicInteger();  //Mutable field count
    batchRowNumber = batch.size++;

    schemaDescription.forEach( field -> setFieldValue( fieldNumber, field, row ) );
    if ( batch.size == batch.getMaxSize() - 1 ) {
      writer.addRowBatch( batch );
      batch.reset();
    }
  }

  private void setFieldValue( AtomicInteger fieldNumber, SchemaDescription.Field field,
                              RowMetaAndData rowMetaAndData ) {
    int fieldNo = fieldNumber.getAndIncrement();
    ColumnVector columnVector = batch.cols[ fieldNo ];
    int rowMetaIndex = rowMetaAndData.getRowMeta().indexOfValue( field.pentahoFieldName );

    if ( rowMetaAndData.getData()[ rowMetaIndex ] == null ) {
      if ( field.allowNull ) {
{'fraction_non_alphanumeric': 0.13008600727753886, 'fraction_numerical': 0.0037214687396625867, 'mean_word_length': 4.690823529411765, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 4, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6726006', 'n_tokens_mistral': 3426, 'n_tokens_neox': 3331, 'n_words': 793}