service/cli/RowBasedSet.java

0001 /**
0002  * Licensed to the Apache Software Foundation (ASF) under one
0003  * or more contributor license agreements.  See the NOTICE file
0004  * distributed with this work for additional information
0005  * regarding copyright ownership.  The ASF licenses this file
0006  * to you under the Apache License, Version 2.0 (the
0007  * "License"); you may not use this file except in compliance
0008  * with the License.  You may obtain a copy of the License at
0009  *
0010  *     http://www.apache.org/licenses/LICENSE-2.0
0011  *
0012  * Unless required by applicable law or agreed to in writing, software
0013  * distributed under the License is distributed on an "AS IS" BASIS,
0014  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
0015  * See the License for the specific language governing permissions and
0016  * limitations under the License.
0017  */
0018
0019 package org.apache.hive.service.cli;
0020
0021 import java.util.ArrayList;
0022 import java.util.Iterator;
0023 import java.util.List;
0024
0025 import org.apache.hive.service.rpc.thrift.TColumnValue;
0026 import org.apache.hive.service.rpc.thrift.TRow;
0027 import org.apache.hive.service.rpc.thrift.TRowSet;
0028
0029 /**
0030  * RowBasedSet
0031  */
0032 public class RowBasedSet implements RowSet {
0033
0034   private long startOffset;
0035
0036   private final TypeDescriptor[] descriptors; // non-null only for writing (server-side)
0037   private final RemovableList<TRow> rows;
0038
0039   public RowBasedSet(TableSchema schema) {
0040     descriptors = schema.toTypeDescriptors();
0041     rows = new RemovableList<TRow>();
0042   }
0043
0044   public RowBasedSet(TRowSet tRowSet) {
0045     descriptors = null;
0046     rows = new RemovableList<TRow>(tRowSet.getRows());
0047     startOffset = tRowSet.getStartRowOffset();
0048   }
0049
0050   private RowBasedSet(TypeDescriptor[] descriptors, List<TRow> rows, long startOffset) {
0051     this.descriptors = descriptors;
0052     this.rows = new RemovableList<TRow>(rows);
0053     this.startOffset = startOffset;
0054   }
0055
0056   @Override
0057   public RowBasedSet addRow(Object[] fields) {
0058     TRow tRow = new TRow();
0059     for (int i = 0; i < fields.length; i++) {
0060       tRow.addToColVals(ColumnValue.toTColumnValue(descriptors[i], fields[i]));
0061     }
0062     rows.add(tRow);
0063     return this;
0064   }
0065
0066   @Override
0067   public int numColumns() {
0068     return rows.isEmpty() ? 0 : rows.get(0).getColVals().size();
0069   }
0070
0071   @Override
0072   public int numRows() {
0073     return rows.size();
0074   }
0075
0076   public RowBasedSet extractSubset(int maxRows) {
0077     int numRows = Math.min(numRows(), maxRows);
0078     RowBasedSet result = new RowBasedSet(descriptors, rows.subList(0, numRows), startOffset);
0079     rows.removeRange(0, numRows);
0080     startOffset += numRows;
0081     return result;
0082   }
0083
0084   public long getStartOffset() {
0085     return startOffset;
0086   }
0087
0088   public void setStartOffset(long startOffset) {
0089     this.startOffset = startOffset;
0090   }
0091
0092   public int getSize() {
0093     return rows.size();
0094   }
0095
0096   public TRowSet toTRowSet() {
0097     TRowSet tRowSet = new TRowSet();
0098     tRowSet.setStartRowOffset(startOffset);
0099     tRowSet.setRows(new ArrayList<TRow>(rows));
0100     return tRowSet;
0101   }
0102
0103   @Override
0104   public Iterator<Object[]> iterator() {
0105     return new Iterator<Object[]>() {
0106
0107       final Iterator<TRow> iterator = rows.iterator();
0108       final Object[] convey = new Object[numColumns()];
0109
0110       @Override
0111       public boolean hasNext() {
0112         return iterator.hasNext();
0113       }
0114
0115       @Override
0116       public Object[] next() {
0117         TRow row = iterator.next();
0118         List<TColumnValue> values = row.getColVals();
0119         for (int i = 0; i < values.size(); i++) {
0120           convey[i] = ColumnValue.toColumnValue(values.get(i));
0121         }
0122         return convey;
0123       }
0124
0125       @Override
0126       public void remove() {
0127         throw new UnsupportedOperationException("remove");
0128       }
0129     };
0130   }
0131
0132   private static class RemovableList<E> extends ArrayList<E> {
0133     RemovableList() { super(); }
0134     RemovableList(List<E> rows) { super(rows); }
0135     @Override
0136     public void removeRange(int fromIndex, int toIndex) {
0137       super.removeRange(fromIndex, toIndex);
0138     }
0139   }
0140 }