Java tutorial
/** * 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.hive.hplsql; import java.sql.Connection; import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.util.HashMap; import java.util.Map; import java.util.List; import java.io.FileOutputStream; import java.io.OutputStream; import java.io.IOException; import org.apache.hadoop.fs.FileSystem; import org.apache.hadoop.fs.Path; import org.apache.hive.hplsql.Var; import org.antlr.v4.runtime.ParserRuleContext; import org.apache.commons.lang.StringEscapeUtils; import org.apache.commons.lang3.tuple.Pair; public class Copy { Exec exec; Timer timer = new Timer(); boolean trace = false; boolean info = false; long srcSizeInBytes = 0; String delimiter = "\t"; boolean sqlInsert = false; String sqlInsertName; String targetConn; int batchSize = 1000; boolean overwrite = false; boolean delete = false; boolean ignore = false; Copy(Exec e) { exec = e; trace = exec.getTrace(); info = exec.getInfo(); } /** * Run COPY command */ Integer run(HplsqlParser.Copy_stmtContext ctx) { trace(ctx, "COPY"); initOptions(ctx); StringBuilder sql = new StringBuilder(); String conn = null; if (ctx.table_name() != null) { String table = evalPop(ctx.table_name()).toString(); conn = exec.getObjectConnection(ctx.table_name().getText()); sql.append("SELECT * FROM "); sql.append(table); } else { sql.append(evalPop(ctx.select_stmt()).toString()); conn = exec.getStatementConnection(); if (trace) { trace(ctx, "Statement:\n" + sql); } } Query query = exec.executeQuery(ctx, sql.toString(), conn); if (query.error()) { exec.signal(query); return 1; } exec.setSqlSuccess(); try { if (targetConn != null) { copyToTable(ctx, query); } else { copyToFile(ctx, query); } } catch (Exception e) { exec.signal(e); return 1; } finally { exec.closeQuery(query, conn); } return 0; } /** * Copy the query results to another table * @throws Exception */ void copyToTable(HplsqlParser.Copy_stmtContext ctx, Query query) throws Exception { ResultSet rs = query.getResultSet(); if (rs == null) { return; } ResultSetMetaData rm = rs.getMetaData(); int cols = rm.getColumnCount(); int rows = 0; if (trace) { trace(ctx, "SELECT executed: " + cols + " columns"); } Connection conn = exec.getConnection(targetConn); StringBuilder sql = new StringBuilder(); sql.append("INSERT INTO " + sqlInsertName + " VALUES ("); for (int i = 0; i < cols; i++) { sql.append("?"); if (i + 1 < cols) { sql.append(","); } } sql.append(")"); PreparedStatement ps = conn.prepareStatement(sql.toString()); long start = timer.start(); long prev = start; boolean batchOpen = false; while (rs.next()) { for (int i = 1; i <= cols; i++) { ps.setObject(i, rs.getObject(i)); } rows++; if (batchSize > 1) { ps.addBatch(); batchOpen = true; if (rows % batchSize == 0) { ps.executeBatch(); batchOpen = false; } } else { ps.executeUpdate(); } if (trace && rows % 100 == 0) { long cur = timer.current(); if (cur - prev > 10000) { trace(ctx, "Copying rows: " + rows + " (" + rows / ((cur - start) / 1000) + " rows/sec)"); prev = cur; } } } if (batchOpen) { ps.executeBatch(); } ps.close(); exec.returnConnection(targetConn, conn); exec.setRowCount(rows); long elapsed = timer.stop(); if (info) { info(ctx, "COPY completed: " + rows + " row(s), " + timer.format() + ", " + rows / (elapsed / 1000) + " rows/sec"); } } /** * Copy the query results to a file * @throws Exception */ void copyToFile(HplsqlParser.Copy_stmtContext ctx, Query query) throws Exception { ResultSet rs = query.getResultSet(); if (rs == null) { return; } ResultSetMetaData rm = rs.getMetaData(); String filename = null; if (ctx.copy_target().expr() != null) { filename = evalPop(ctx.copy_target().expr()).toString(); } else { filename = ctx.copy_target().getText(); } byte[] del = delimiter.getBytes(); byte[] rowdel = "\n".getBytes(); byte[] nullstr = "NULL".getBytes(); int cols = rm.getColumnCount(); int rows = 0; long bytes = 0; if (trace || info) { String mes = "Query executed: " + cols + " columns, output file: " + filename; if (trace) { trace(ctx, mes); } else { info(ctx, mes); } } java.io.File file = null; File hdfsFile = null; if (ctx.T_HDFS() == null) { file = new java.io.File(filename); } else { hdfsFile = new File(); } OutputStream out = null; timer.start(); try { if (file != null) { if (!file.exists()) { file.createNewFile(); } out = new FileOutputStream(file, false /*append*/); } else { out = hdfsFile.create(filename, true /*overwrite*/); } String col; String sql = ""; if (sqlInsert) { sql = "INSERT INTO " + sqlInsertName + " VALUES ("; rowdel = ");\n".getBytes(); } while (rs.next()) { if (sqlInsert) { out.write(sql.getBytes()); } for (int i = 1; i <= cols; i++) { if (i > 1) { out.write(del); bytes += del.length; } col = rs.getString(i); if (col != null) { if (sqlInsert) { col = Utils.quoteString(col); } byte[] b = col.getBytes(); out.write(b); bytes += b.length; } else if (sqlInsert) { out.write(nullstr); } } out.write(rowdel); bytes += rowdel.length; rows++; } exec.setRowCount(rows); } finally { if (out != null) { out.close(); } } long elapsed = timer.stop(); if (info) { info(ctx, "COPY completed: " + rows + " row(s), " + Utils.formatSizeInBytes(bytes) + ", " + timer.format() + ", " + rows / elapsed / 1000 + " rows/sec"); } } /** * Run COPY FROM LOCAL statement */ public Integer runFromLocal(HplsqlParser.Copy_from_local_stmtContext ctx) { trace(ctx, "COPY FROM LOCAL"); initFileOptions(ctx.copy_file_option()); HashMap<String, Pair<String, Long>> srcFiles = new HashMap<String, Pair<String, Long>>(); String src = evalPop(ctx.copy_source(0)).toString(); String dest = evalPop(ctx.copy_target()).toString(); int srcItems = ctx.copy_source().size(); for (int i = 0; i < srcItems; i++) { createLocalFileList(srcFiles, evalPop(ctx.copy_source(i)).toString(), null); } if (info) { info(ctx, "Files to copy: " + srcFiles.size() + " (" + Utils.formatSizeInBytes(srcSizeInBytes) + ")"); } if (srcFiles.size() == 0) { exec.setHostCode(2); return 2; } timer.start(); File file = new File(); FileSystem fs = null; int succeed = 0; int failed = 0; long copiedSize = 0; try { fs = file.createFs(); boolean multi = false; if (srcFiles.size() > 1) { multi = true; } for (Map.Entry<String, Pair<String, Long>> i : srcFiles.entrySet()) { try { Path s = new Path(i.getKey()); Path d = null; if (multi) { String relativePath = i.getValue().getLeft(); if (relativePath == null) { d = new Path(dest, s.getName()); } else { d = new Path(dest, relativePath + Path.SEPARATOR + s.getName()); } } else { // Path to file is specified (can be relative), so treat target as a file name (hadoop fs -put behavior) if (srcItems == 1 && i.getKey().endsWith(src)) { d = new Path(dest); } // Source directory is specified, so treat the target as a directory else { d = new Path(dest + Path.SEPARATOR + s.getName()); } } fs.copyFromLocalFile(delete, overwrite, s, d); succeed++; long size = i.getValue().getRight(); copiedSize += size; if (info) { info(ctx, "Copied: " + file.resolvePath(d) + " (" + Utils.formatSizeInBytes(size) + ")"); } } catch (IOException e) { failed++; if (!ignore) { throw e; } } } } catch (IOException e) { exec.signal(e); exec.setHostCode(1); return 1; } finally { long elapsed = timer.stop(); if (info) { info(ctx, "COPY completed: " + succeed + " succeed, " + failed + " failed, " + timer.format() + ", " + Utils.formatSizeInBytes(copiedSize) + ", " + Utils.formatBytesPerSec(copiedSize, elapsed)); } if (failed == 0) { exec.setHostCode(0); } else { exec.setHostCode(1); } file.close(); } return 0; } /** * Create the list of local files for the specified path (including subdirectories) */ void createLocalFileList(HashMap<String, Pair<String, Long>> list, String path, String relativePath) { java.io.File file = new java.io.File(path); if (file.exists()) { if (file.isDirectory()) { for (java.io.File i : file.listFiles()) { if (i.isDirectory()) { String rel = null; if (relativePath == null) { rel = i.getName(); } else { rel = relativePath + java.io.File.separator + i.getName(); } createLocalFileList(list, i.getAbsolutePath(), rel); } else { long size = i.length(); list.put(i.getAbsolutePath(), Pair.of(relativePath, size)); srcSizeInBytes += size; } } } else { long size = file.length(); list.put(file.getAbsolutePath(), Pair.of(relativePath, size)); srcSizeInBytes += size; } } } /** * Initialize COPY command options */ void initOptions(HplsqlParser.Copy_stmtContext ctx) { int cnt = ctx.copy_option().size(); for (int i = 0; i < cnt; i++) { HplsqlParser.Copy_optionContext option = ctx.copy_option(i); if (option.T_DELIMITER() != null) { delimiter = StringEscapeUtils.unescapeJava(evalPop(option.expr()).toString()); } else if (option.T_SQLINSERT() != null) { sqlInsert = true; delimiter = ", "; if (option.ident() != null) { sqlInsertName = option.ident().getText(); } } else if (option.T_AT() != null) { targetConn = option.ident().getText(); if (ctx.copy_target().expr() != null) { sqlInsertName = evalPop(ctx.copy_target().expr()).toString(); } else { sqlInsertName = ctx.copy_target().getText(); } } else if (option.T_BATCHSIZE() != null) { batchSize = evalPop(option.expr()).intValue(); } } } /** * Initialize COPY FILE options */ void initFileOptions(List<HplsqlParser.Copy_file_optionContext> options) { srcSizeInBytes = 0; for (HplsqlParser.Copy_file_optionContext i : options) { if (i.T_OVERWRITE() != null) { overwrite = true; } else if (i.T_DELETE() != null) { delete = true; } else if (i.T_IGNORE() != null) { ignore = true; } } } /** * Evaluate the expression and pop value from the stack */ Var evalPop(ParserRuleContext ctx) { exec.visit(ctx); if (!exec.stack.isEmpty()) { return exec.stackPop(); } return Var.Empty; } /** * Trace and information */ public void trace(ParserRuleContext ctx, String message) { exec.trace(ctx, message); } public void info(ParserRuleContext ctx, String message) { exec.info(ctx, message); } }