com.google.cloud.bigquery.BigQueryResultImpl Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of google-cloud-bigquery Show documentation
Show all versions of google-cloud-bigquery Show documentation
Java idiomatic client for Google Cloud BigQuery.
/*
* Copyright 2021 Google LLC
*
* 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 com.google.cloud.bigquery;
import java.math.BigDecimal;
import java.sql.Date;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Time;
import java.sql.Timestamp;
import java.time.LocalDateTime;
import java.time.LocalTime;
import java.time.ZoneId;
import java.util.Map;
import java.util.TimeZone;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.TimeUnit;
import org.apache.arrow.vector.util.JsonStringArrayList;
import org.apache.arrow.vector.util.Text;
public class BigQueryResultImpl implements BigQueryResult {
private static final String NULL_CURSOR_MSG =
"Error occurred while reading the cursor. This could happen if getters are called after we are done reading all the records";
// This class represents a row of records, the columns are represented as a map
// (columnName:columnValue pair)
static class Row {
private Map value;
private boolean isLast;
public Row(Map value) {
this.value = value;
}
public Row(Map value, boolean isLast) {
this.value = value;
this.isLast = isLast;
}
public Map getValue() {
return value;
}
public boolean isLast() {
return isLast;
}
public boolean hasField(String fieldName) {
return this.value.containsKey(fieldName);
}
public Object get(String fieldName) {
return this.value.get(fieldName);
}
}
private final Schema schema;
private final long totalRows;
private final BlockingQueue buffer;
private T cursor;
private final BigQueryResultSet underlyingResultSet;
private final BigQueryResultStats bigQueryResultStats;
private final FieldList schemaFieldList;
public BigQueryResultImpl(
Schema schema,
long totalRows,
BlockingQueue buffer,
BigQueryResultStats bigQueryResultStats) {
this.schema = schema;
this.totalRows = totalRows;
this.buffer = buffer;
this.underlyingResultSet = new BigQueryResultSet();
this.bigQueryResultStats = bigQueryResultStats;
this.schemaFieldList = schema.getFields();
}
@Override
public Schema getSchema() {
return schema;
}
@Override
public long getTotalRows() {
return totalRows;
}
@Override
public ResultSet getResultSet() {
return underlyingResultSet;
}
private class BigQueryResultSet extends AbstractJdbcResultSet {
private boolean hasReachedEnd =
false; // flag which will be set to true when we have encountered a EndOfStream or when
// curTup.isLast(). Ref: https://github.com/googleapis/java-bigquery/issues/2033
@Override
/*Advances the result set to the next row, returning false if no such row exists. Potentially blocking operation*/
public boolean next() throws SQLException {
if (hasReachedEnd) { // if end of stream is reached then we can simply return false
return false;
}
try {
cursor = buffer.take(); // advance the cursor,Potentially blocking operation
if (isEndOfStream(cursor)) { // check for end of stream
cursor = null;
hasReachedEnd = true;
return false;
} else if (cursor instanceof Row) {
Row curTup = (Row) cursor;
if (curTup.isLast()) { // last Tuple
cursor = null;
hasReachedEnd = true;
return false;
}
return true;
} else if (cursor instanceof FieldValueList) { // cursor is advanced, we can return true now
return true;
} else { // this case should never occur as the cursor will either be a Row of EoS
throw new BigQuerySQLException("Could not process the current row");
}
} catch (InterruptedException e) {
throw new SQLException(
"Error occurred while advancing the cursor. This could happen when connection is closed while we call the next method");
}
}
private boolean isEndOfStream(T cursor) {
return cursor instanceof ConnectionImpl.EndOfFieldValueList;
}
@Override
public Object getObject(String fieldName) throws SQLException {
if (fieldName == null) {
throw new SQLException("fieldName can't be null");
}
if (cursor == null) {
throw new BigQuerySQLException(NULL_CURSOR_MSG);
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(fieldName);
return (fieldValue == null || fieldValue.getValue() == null) ? null : fieldValue.getValue();
} else { // Data received from Read API (Arrow)
Row curRow = (Row) cursor;
if (!curRow.hasField(fieldName)) {
throw new SQLException(String.format("Field %s not found", fieldName));
}
return curRow.get(fieldName);
}
}
@Override
public Object getObject(int columnIndex) throws SQLException {
if (cursor == null) {
return null;
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(columnIndex);
return (fieldValue == null || fieldValue.getValue() == null) ? null : fieldValue.getValue();
} else { // Data received from Read API (Arrow)
return getObject(schemaFieldList.get(columnIndex).getName());
}
}
@Override
public String getString(String fieldName) throws SQLException {
if (fieldName == null) {
throw new SQLException("fieldName can't be null");
}
if (cursor == null) {
throw new BigQuerySQLException(NULL_CURSOR_MSG);
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(fieldName);
if ((fieldValue == null || fieldValue.getValue() == null)) {
return null;
} else if (fieldValue
.getAttribute()
.equals(FieldValue.Attribute.REPEATED)) { // Case for Arrays
return fieldValue.getValue().toString();
} else {
return fieldValue.getStringValue();
}
} else { // Data received from Read API (Arrow)
Row curRow = (Row) cursor;
if (!curRow.hasField(fieldName)) {
throw new SQLException(String.format("Field %s not found", fieldName));
}
Object currentVal = curRow.get(fieldName);
if (currentVal == null) {
return null;
} else if (currentVal instanceof JsonStringArrayList) { // arrays
JsonStringArrayList jsnAry = (JsonStringArrayList) currentVal;
return jsnAry.toString();
} else if (currentVal instanceof LocalDateTime) {
LocalDateTime dateTime = (LocalDateTime) currentVal;
return dateTime.toString();
} else {
Text textVal = (Text) currentVal;
return textVal.toString();
}
}
}
@Override
public String getString(int columnIndex) throws SQLException {
if (cursor == null) {
return null;
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(columnIndex);
return (fieldValue == null || fieldValue.getValue() == null)
? null
: fieldValue.getStringValue();
} else { // Data received from Read API (Arrow)
return getString(schemaFieldList.get(columnIndex).getName());
}
}
@Override
public int getInt(String fieldName) throws SQLException {
if (fieldName == null) {
throw new SQLException("fieldName can't be null");
}
if (cursor == null) {
return 0; // the column value; if the value is SQL NULL, the value returned is 0 as per
// java.sql.ResultSet definition
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(fieldName);
return (fieldValue == null || fieldValue.getValue() == null)
? 0
: fieldValue.getNumericValue().intValue();
} else { // Data received from Read API (Arrow)
Row curRow = (Row) cursor;
if (!curRow.hasField(fieldName)) {
throw new SQLException(String.format("Field %s not found", fieldName));
}
Object curVal = curRow.get(fieldName);
if (curVal == null) {
return 0;
}
if (curVal instanceof Text) { // parse from text to int
return Integer.parseInt(((Text) curVal).toString());
} else if (curVal
instanceof
Long) { // incase getInt is called for a Long value. Loss of precision might occur
return ((Long) curVal).intValue();
}
return ((BigDecimal) curVal).intValue();
}
}
@Override
public int getInt(int columnIndex) throws SQLException {
if (cursor == null) {
return 0; // the column value; if the value is SQL NULL, the value returned is 0 as per
// java.sql.ResultSet definition
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(columnIndex);
return (fieldValue == null || fieldValue.getValue() == null)
? 0
: fieldValue.getNumericValue().intValue();
} else { // Data received from Read API (Arrow)
return getInt(schemaFieldList.get(columnIndex).getName());
}
}
@Override
public long getLong(String fieldName) throws SQLException {
if (fieldName == null) {
throw new SQLException("fieldName can't be null");
}
if (cursor == null) {
throw new BigQuerySQLException(NULL_CURSOR_MSG);
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(fieldName);
return (fieldValue == null || fieldValue.getValue() == null)
? 0L
: fieldValue.getNumericValue().longValue();
} else { // Data received from Read API (Arrow)
Row curRow = (Row) cursor;
if (!curRow.hasField(fieldName)) {
throw new SQLException(String.format("Field %s not found", fieldName));
}
Object curVal = curRow.get(fieldName);
if (curVal == null) {
return 0L;
} else { // value will be Long or BigDecimal, but are Number
return ((Number) curVal).longValue();
}
}
}
@Override
public long getLong(int columnIndex) throws SQLException {
if (cursor == null) {
return 0L; // the column value; if the value is SQL NULL, the value returned is 0 as per
// java.sql.ResultSet definition
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(columnIndex);
return (fieldValue == null || fieldValue.getValue() == null)
? 0L
: fieldValue.getNumericValue().longValue();
} else { // Data received from Read API (Arrow)
return getInt(schemaFieldList.get(columnIndex).getName());
}
}
@Override
public double getDouble(String fieldName) throws SQLException {
if (fieldName == null) {
throw new SQLException("fieldName can't be null");
}
if (cursor == null) {
throw new BigQuerySQLException(NULL_CURSOR_MSG);
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(fieldName);
return (fieldValue == null || fieldValue.getValue() == null)
? 0d
: fieldValue.getNumericValue().doubleValue();
} else { // Data received from Read API (Arrow)
Row curRow = (Row) cursor;
if (!curRow.hasField(fieldName)) {
throw new SQLException(String.format("Field %s not found", fieldName));
}
Object curVal = curRow.get(fieldName);
return curVal == null ? 0.0d : new BigDecimal(curVal.toString()).doubleValue();
}
}
@Override
public double getDouble(int columnIndex) throws SQLException {
if (cursor == null) {
return 0d; // the column value; if the value is SQL NULL, the value returned is 0 as per
// java.sql.ResultSet definition
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(columnIndex);
return (fieldValue == null || fieldValue.getValue() == null)
? 0d
: fieldValue.getNumericValue().doubleValue();
} else { // Data received from Read API (Arrow)
return getDouble(schemaFieldList.get(columnIndex).getName());
}
}
@Override
public BigDecimal getBigDecimal(String fieldName) throws SQLException {
if (fieldName == null) {
throw new SQLException("fieldName can't be null");
}
if (cursor == null) {
throw new BigQuerySQLException(NULL_CURSOR_MSG);
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(fieldName);
return (fieldValue == null || fieldValue.getValue() == null)
? null
: BigDecimal.valueOf(fieldValue.getNumericValue().doubleValue());
} else { // Data received from Read API (Arrow)
return BigDecimal.valueOf(getDouble(fieldName));
}
}
@Override
public BigDecimal getBigDecimal(int columnIndex) throws SQLException {
if (cursor == null) {
throw new BigQuerySQLException(NULL_CURSOR_MSG);
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(columnIndex);
return (fieldValue == null || fieldValue.getValue() == null)
? null
: BigDecimal.valueOf(fieldValue.getNumericValue().doubleValue());
} else { // Data received from Read API (Arrow)
return getBigDecimal(schemaFieldList.get(columnIndex).getName());
}
}
@Override
public boolean getBoolean(String fieldName) throws SQLException {
if (fieldName == null) {
throw new SQLException("fieldName can't be null");
}
if (cursor == null) {
throw new BigQuerySQLException(NULL_CURSOR_MSG);
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(fieldName);
return fieldValue.getValue() != null && fieldValue.getBooleanValue();
} else { // Data received from Read API (Arrow)
Row curRow = (Row) cursor;
if (!curRow.hasField(fieldName)) {
throw new SQLException(String.format("Field %s not found", fieldName));
}
Object curVal = curRow.get(fieldName);
return curVal != null && (Boolean) curVal;
}
}
@Override
public boolean getBoolean(int columnIndex) throws SQLException {
if (cursor == null) {
throw new BigQuerySQLException(NULL_CURSOR_MSG);
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(columnIndex);
return fieldValue.getValue() != null && fieldValue.getBooleanValue();
} else { // Data received from Read API (Arrow)
return getBoolean(schemaFieldList.get(columnIndex).getName());
}
}
@Override
public byte[] getBytes(String fieldName) throws SQLException {
if (fieldName == null) {
throw new SQLException("fieldName can't be null");
}
if (cursor == null) {
throw new BigQuerySQLException(NULL_CURSOR_MSG);
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(fieldName);
return (fieldValue == null || fieldValue.getValue() == null)
? null
: fieldValue.getBytesValue();
} else { // Data received from Read API (Arrow)
Row curRow = (Row) cursor;
if (!curRow.hasField(fieldName)) {
throw new SQLException(String.format("Field %s not found", fieldName));
}
Object curVal = curRow.get(fieldName);
return curVal == null ? null : (byte[]) curVal;
}
}
@Override
public byte[] getBytes(int columnIndex) throws SQLException {
if (cursor == null) {
return null; // if the value is SQL NULL, the value returned is null
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(columnIndex);
return (fieldValue == null || fieldValue.getValue() == null)
? null
: fieldValue.getBytesValue();
} else { // Data received from Read API (Arrow)
return getBytes(schemaFieldList.get(columnIndex).getName());
}
}
@Override
public Timestamp getTimestamp(String fieldName) throws SQLException {
if (fieldName == null) {
throw new BigQuerySQLException(NULL_CURSOR_MSG);
}
if (cursor == null) {
return null; // if the value is SQL NULL, the value returned is null
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(fieldName);
return (fieldValue == null || fieldValue.getValue() == null)
? null
: new Timestamp(
fieldValue.getTimestampValue()
/ 1000); // getTimestampValue returns time in microseconds, and TimeStamp
// expects it in millis
} else {
Row curRow = (Row) cursor;
if (!curRow.hasField(fieldName)) {
throw new SQLException(String.format("Field %s not found", fieldName));
}
Object timeStampVal = curRow.get(fieldName);
return timeStampVal == null
? null
: new Timestamp((Long) timeStampVal / 1000); // Timestamp is represented as a Long
}
}
@Override
public Timestamp getTimestamp(int columnIndex) throws SQLException {
if (cursor == null) {
throw new BigQuerySQLException(NULL_CURSOR_MSG);
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(columnIndex);
return (fieldValue == null || fieldValue.getValue() == null)
? null
: new Timestamp(
fieldValue.getTimestampValue()
/ 1000); // getTimestampValue returns time in microseconds, and TimeStamp
// expects it in millis
} else { // Data received from Read API (Arrow)
return getTimestamp(schemaFieldList.get(columnIndex).getName());
}
}
@Override
public Time getTime(String fieldName) throws SQLException {
if (fieldName == null) {
throw new BigQuerySQLException(NULL_CURSOR_MSG);
}
if (cursor == null) {
return null; // if the value is SQL NULL, the value returned is null
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(fieldName);
return getTimeFromFieldVal(fieldValue);
} else { // Data received from Read API (Arrow)
Row curRow = (Row) cursor;
if (!curRow.hasField(fieldName)) {
throw new SQLException(String.format("Field %s not found", fieldName));
}
Object timeStampObj = curRow.get(fieldName);
return timeStampObj == null
? null
: new Time(
((Long) timeStampObj)
/ 1000); // Time.toString() will return 12:11:35 in GMT as 17:41:35 in
// (GMT+5:30). This can be offset using getTimeZoneOffset
}
}
private int getTimeZoneOffset() {
TimeZone timeZone = TimeZone.getTimeZone(ZoneId.systemDefault());
return timeZone.getOffset(new java.util.Date().getTime()); // offset in seconds
}
@Override
public Time getTime(int columnIndex) throws SQLException {
if (cursor == null) {
throw new BigQuerySQLException(NULL_CURSOR_MSG);
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(columnIndex);
return getTimeFromFieldVal(fieldValue);
} else { // Data received from Read API (Arrow)
return getTime(schemaFieldList.get(columnIndex).getName());
}
}
private Time getTimeFromFieldVal(FieldValue fieldValue) throws SQLException {
if (fieldValue.getValue() != null) {
// Time ranges from 00:00:00 to 23:59:59.99999. in BigQuery. Parsing it to java.sql.Time
String strTime = fieldValue.getStringValue();
String[] timeSplt = strTime.split(":");
if (timeSplt.length != 3) {
throw new SQLException("Can not parse the value " + strTime + " to java.sql.Time");
}
int hr = Integer.parseInt(timeSplt[0]);
int min = Integer.parseInt(timeSplt[1]);
int sec = 0, nanoSec = 0;
if (timeSplt[2].contains(".")) {
String[] secSplt = timeSplt[2].split("\\.");
sec = Integer.parseInt(secSplt[0]);
nanoSec = Integer.parseInt(secSplt[1]);
} else {
sec = Integer.parseInt(timeSplt[2]);
}
return Time.valueOf(LocalTime.of(hr, min, sec, nanoSec));
} else {
return null;
}
}
@Override
public Date getDate(String fieldName) throws SQLException {
if (fieldName == null) {
throw new SQLException("fieldName can't be null");
}
if (cursor == null) {
throw new BigQuerySQLException(NULL_CURSOR_MSG);
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(fieldName);
return (fieldValue == null || fieldValue.getValue() == null)
? null
: Date.valueOf(fieldValue.getStringValue());
} else { // Data received from Read API (Arrow)
Row curRow = (Row) cursor;
if (!curRow.hasField(fieldName)) {
throw new SQLException(String.format("Field %s not found", fieldName));
}
Object dateObj = curRow.get(fieldName);
if (dateObj == null) {
return null;
} else {
Integer dateInt = (Integer) dateObj;
long dateInMillis =
TimeUnit.DAYS.toMillis(
Long.valueOf(
dateInt)); // For example int 18993 represents 2022-01-01, converting time to
// milli seconds
return new Date(dateInMillis);
}
}
}
@Override
public Date getDate(int columnIndex) throws SQLException {
if (cursor == null) {
throw new BigQuerySQLException(NULL_CURSOR_MSG);
} else if (cursor instanceof FieldValueList) {
FieldValue fieldValue = ((FieldValueList) cursor).get(columnIndex);
return (fieldValue == null || fieldValue.getValue() == null)
? null
: Date.valueOf(fieldValue.getStringValue());
} else { // Data received from Read API (Arrow)
return getDate(schemaFieldList.get(columnIndex).getName());
}
}
}
@Override
public BigQueryResultStats getBigQueryResultStats() {
return bigQueryResultStats;
}
}