Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance.
Project price only 1 $
You can buy this project and download/modify it how often you want.
io.trino.plugin.tpch.TpchRecordSet Maven / Gradle / Ivy
/*
* 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 io.trino.plugin.tpch;
import com.google.common.collect.ImmutableList;
import io.airlift.slice.Slice;
import io.airlift.slice.Slices;
import io.trino.spi.TrinoException;
import io.trino.spi.connector.ColumnHandle;
import io.trino.spi.connector.RecordCursor;
import io.trino.spi.connector.RecordSet;
import io.trino.spi.predicate.NullableValue;
import io.trino.spi.predicate.TupleDomain;
import io.trino.spi.type.Type;
import io.trino.tpch.TpchColumn;
import io.trino.tpch.TpchColumnType;
import io.trino.tpch.TpchEntity;
import io.trino.tpch.TpchTable;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.ImmutableMap.toImmutableMap;
import static io.trino.plugin.tpch.TpchMetadata.getTrinoType;
import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED;
import static java.lang.String.format;
import static java.util.Objects.requireNonNull;
public class TpchRecordSet
implements RecordSet
{
public static TpchRecordSet createTpchRecordSet(TpchTable table, double scaleFactor)
{
return createTpchRecordSet(table, table.getColumns(), DecimalTypeMapping.DOUBLE, scaleFactor, 1, 1, TupleDomain.all());
}
public static TpchRecordSet createTpchRecordSet(
TpchTable table,
DecimalTypeMapping decimalTypeMapping,
double scaleFactor,
int part,
int partCount,
TupleDomain predicate)
{
return createTpchRecordSet(table, table.getColumns(), decimalTypeMapping, scaleFactor, part, partCount, predicate);
}
public static TpchRecordSet createTpchRecordSet(
TpchTable table,
List> columns,
DecimalTypeMapping decimalTypeMapping,
double scaleFactor,
int part,
int partCount,
TupleDomain predicate)
{
return new TpchRecordSet<>(table.createGenerator(scaleFactor, part, partCount), table, columns, decimalTypeMapping, predicate);
}
private final Iterable rows;
private final TpchTable table;
private final List> columns;
private final List columnTypes;
private final TupleDomain predicate;
public TpchRecordSet(Iterable rows,
TpchTable table,
List> columns,
DecimalTypeMapping decimalTypeMapping,
TupleDomain predicate)
{
this.rows = requireNonNull(rows, "rows is null");
this.table = requireNonNull(table, "table is null");
this.columns = ImmutableList.copyOf(requireNonNull(columns, "columns is null"));
this.columnTypes = columns.stream()
.map(column -> getTrinoType(column, decimalTypeMapping))
.collect(toImmutableList());
this.predicate = requireNonNull(predicate, "predicate is null");
}
@Override
public List getColumnTypes()
{
return columnTypes;
}
@Override
public RecordCursor cursor()
{
return new TpchRecordCursor<>(rows.iterator(), table, columns, columnTypes, predicate);
}
public static final class TpchRecordCursor
implements RecordCursor
{
private final Iterator rows;
private final TpchTable table;
private final List> columns;
private final List columnTypes;
private final TupleDomain predicate;
private E row;
private boolean closed;
public TpchRecordCursor(Iterator rows, TpchTable table, List> columns, List columnTypes, TupleDomain predicate)
{
this.rows = requireNonNull(rows, "rows is null");
this.table = requireNonNull(table, "table is null");
this.columns = requireNonNull(columns, "columns is null");
this.columnTypes = requireNonNull(columnTypes, "columnTypes is null");
this.predicate = requireNonNull(predicate, "predicate is null");
}
@Override
public long getCompletedBytes()
{
return 0;
}
@Override
public long getReadTimeNanos()
{
return 0;
}
@Override
public Type getType(int field)
{
return columnTypes.get(field);
}
@Override
public boolean advanceNextPosition()
{
while (!closed && rows.hasNext()) {
row = rows.next();
if (rowMatchesPredicate()) {
return true;
}
}
closed = true;
row = null;
return false;
}
@Override
public boolean getBoolean(int field)
{
throw new UnsupportedOperationException();
}
@Override
public long getLong(int field)
{
checkState(row != null, "No current row");
return getLong(getTpchColumn(field));
}
private long getLong(TpchColumn tpchColumn)
{
if (tpchColumn.getType().getBase() == TpchColumnType.Base.DATE) {
return tpchColumn.getDate(row);
}
if (tpchColumn.getType().getBase() == TpchColumnType.Base.INTEGER) {
return tpchColumn.getInteger(row);
}
return tpchColumn.getIdentifier(row);
}
@Override
public double getDouble(int field)
{
checkState(row != null, "No current row");
return getDouble(getTpchColumn(field));
}
private double getDouble(TpchColumn tpchColumn)
{
return tpchColumn.getDouble(row);
}
@Override
public Slice getSlice(int field)
{
checkState(row != null, "No current row");
return getSlice(getTpchColumn(field));
}
private Slice getSlice(TpchColumn tpchColumn)
{
return Slices.utf8Slice(tpchColumn.getString(row));
}
@Override
public Object getObject(int field)
{
throw new UnsupportedOperationException();
}
@Override
public boolean isNull(int field)
{
return false;
}
@Override
public void close()
{
row = null;
closed = true;
}
private boolean rowMatchesPredicate()
{
if (predicate.isAll()) {
return true;
}
if (predicate.isNone()) {
return false;
}
Map rowMap = predicate.getDomains().get().keySet().stream()
.collect(toImmutableMap(
column -> column,
column -> {
TpchColumnHandle tpchColumnHandle = (TpchColumnHandle) column;
Type type = tpchColumnHandle.getType();
TpchColumn tpchColumn = table.getColumn(tpchColumnHandle.getColumnName());
return NullableValue.of(type, getTrinoObject(tpchColumn, type));
}));
TupleDomain rowTupleDomain = TupleDomain.fromFixedValues(rowMap);
return predicate.contains(rowTupleDomain);
}
private Object getTrinoObject(TpchColumn column, Type type)
{
if (type.getJavaType() == long.class) {
return getLong(column);
}
if (type.getJavaType() == double.class) {
return getDouble(column);
}
if (type.getJavaType() == Slice.class) {
return getSlice(column);
}
throw new TrinoException(NOT_SUPPORTED, format("Unsupported column type %s", type.getDisplayName()));
}
private TpchColumn getTpchColumn(int field)
{
return columns.get(field);
}
}
}
