org.apache.spark.sql.sources.v2.writer.DataWriter Java Examples

@Override
public DataWriter<Row> createDataWriter(int partitionId, int attemptNumber)
{
    try {
        log.info("creating a data writer for table " + _table +
                " partition " + partitionId + " attempt " + attemptNumber);
        return new TaskDataWriter(_host, _port, _table, _schema, partitionId, attemptNumber);
    } catch (UnknownTableException ute) {
        throw new RuntimeException(ute);
    }
}
 

@Override
public DataWriter<InternalRow> createDataWriter(int partitionId, long taskId, long epochId) {
  OutputFileFactory fileFactory = new OutputFileFactory(
      spec, format, locations, io.value(), encryptionManager.value(), partitionId, taskId);
  SparkAppenderFactory appenderFactory = new SparkAppenderFactory(properties, writeSchema, dsSchema);

  if (spec.fields().isEmpty()) {
    return new Unpartitioned24Writer(spec, format, appenderFactory, fileFactory, io.value(), targetFileSize);
  } else {
    return new Partitioned24Writer(
        spec, format, appenderFactory, fileFactory, io.value(), targetFileSize, writeSchema);
  }
}
 

@Override
public DataWriter<InternalRow> createDataWriter(int partitionId, int attemptNumber) {
  String filename = format.addExtension(String.format("%05d-%d-%s",
      partitionId, attemptNumber, uuid));
  AppenderFactory<InternalRow> factory = new SparkAppenderFactory();
  if (spec.fields().isEmpty()) {
    return new UnpartitionedWriter(lazyDataPath(), filename, format, conf.value(), factory);

  } else {
    Path baseDataPath = lazyDataPath(); // avoid calling this in the output path function
    Function<PartitionKey, Path> outputPathFunc = key ->
        new Path(new Path(baseDataPath, key.toPath()), filename);

    boolean useObjectStorage = (
        Boolean.parseBoolean(properties.get(OBJECT_STORE_ENABLED)) ||
        OBJECT_STORE_ENABLED_DEFAULT
    );

    if (useObjectStorage) {
      // try to get db and table portions of the path for context in the object store
      String context = pathContext(baseDataPath);
      String objectStore = properties.get(OBJECT_STORE_PATH);
      Preconditions.checkNotNull(objectStore,
          "Cannot use object storage, missing location: " + OBJECT_STORE_PATH);
      Path objectStorePath = new Path(objectStore);

      outputPathFunc = key -> {
        String partitionAndFilename = key.toPath() + "/" + filename;
        int hash = HASH_FUNC.apply(partitionAndFilename);
        return new Path(objectStorePath,
            String.format("%08x/%s/%s", hash, context, partitionAndFilename));
      };
    }

    return new PartitionedWriter(spec, format, conf.value(), factory, outputPathFunc);
  }
}
 

@Override public DataWriter<InternalRow> createDataWriter(int partitionId, int attemptNumber) {
  Path filePath = new Path(this.path, String.format("%s_%s_%s", jobId, partitionId, attemptNumber));
  FileSystem fs = null;
  try {
    fs = filePath.getFileSystem(conf.get());
  } catch (IOException e) {
    e.printStackTrace();
  }
  return getDataWriter(conf.get(), jobId, schema, partitionId, attemptNumber,
      fs, filePath);
}
 

@Override
public DataWriter<InternalRow> createDataWriter(int partitionId, int attemptNumber) {
  ClassLoader restoredClassloader = Thread.currentThread().getContextClassLoader();
  ClassLoader isolatedClassloader = HiveIsolatedClassLoader.isolatedClassLoader();
  try {
    Thread.currentThread().setContextClassLoader(isolatedClassloader);
    return new HiveStreamingDataWriter(jobId, schema, commitIntervalRows, partitionId, attemptNumber, db,
      table, partition, metastoreUri, metastoreKrbPrincipal);
  } finally {
    Thread.currentThread().setContextClassLoader(restoredClassloader);
  }
}
 

protected DataWriter<InternalRow> getDataWriter(Configuration conf, String jobId,
    StructType schema, int partitionId, int attemptNumber,
    FileSystem fs, Path filePath) {
  return new HiveWarehouseDataWriter(conf, jobId, schema, partitionId, attemptNumber, fs, filePath);
}
 

protected DataWriter<InternalRow> getDataWriter(Configuration conf, String jobId,
    StructType schema, int partitionId, int attemptNumber,
    FileSystem fs, Path filePath) {
  return new MockHiveWarehouseDataWriter(conf, jobId, schema, partitionId, attemptNumber, fs, filePath);
}