service/cli/ColumnBasedSet.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.cli.thrift.TColumn;
0026 import org.apache.hive.service.cli.thrift.TRow;
0027 import org.apache.hive.service.cli.thrift.TRowSet;
0028
0029 /**
0030  * ColumnBasedSet.
0031  */
0032 public class ColumnBasedSet implements RowSet {
0033
0034   private long startOffset;
0035
0036   private final Type[] types; // non-null only for writing (server-side)
0037   private final List<Column> columns;
0038
0039   public ColumnBasedSet(TableSchema schema) {
0040     types = schema.toTypes();
0041     columns = new ArrayList<Column>();
0042     for (ColumnDescriptor colDesc : schema.getColumnDescriptors()) {
0043       columns.add(new Column(colDesc.getType()));
0044     }
0045   }
0046
0047   public ColumnBasedSet(TRowSet tRowSet) {
0048     types = null;
0049     columns = new ArrayList<Column>();
0050     for (TColumn tvalue : tRowSet.getColumns()) {
0051       columns.add(new Column(tvalue));
0052     }
0053     startOffset = tRowSet.getStartRowOffset();
0054   }
0055
0056   private ColumnBasedSet(Type[] types, List<Column> columns, long startOffset) {
0057     this.types = types;
0058     this.columns = columns;
0059     this.startOffset = startOffset;
0060   }
0061
0062   @Override
0063   public ColumnBasedSet addRow(Object[] fields) {
0064     for (int i = 0; i < fields.length; i++) {
0065       columns.get(i).addValue(types[i], fields[i]);
0066     }
0067     return this;
0068   }
0069
0070   public List<Column> getColumns() {
0071     return columns;
0072   }
0073
0074   @Override
0075   public int numColumns() {
0076     return columns.size();
0077   }
0078
0079   @Override
0080   public int numRows() {
0081     return columns.isEmpty() ? 0 : columns.get(0).size();
0082   }
0083
0084   @Override
0085   public ColumnBasedSet extractSubset(int maxRows) {
0086     int numRows = Math.min(numRows(), maxRows);
0087
0088     List<Column> subset = new ArrayList<Column>();
0089     for (int i = 0; i < columns.size(); i++) {
0090       subset.add(columns.get(i).extractSubset(0, numRows));
0091     }
0092     ColumnBasedSet result = new ColumnBasedSet(types, subset, startOffset);
0093     startOffset += numRows;
0094     return result;
0095   }
0096
0097   @Override
0098   public long getStartOffset() {
0099     return startOffset;
0100   }
0101
0102   @Override
0103   public void setStartOffset(long startOffset) {
0104     this.startOffset = startOffset;
0105   }
0106
0107   public TRowSet toTRowSet() {
0108     TRowSet tRowSet = new TRowSet(startOffset, new ArrayList<TRow>());
0109     for (int i = 0; i < columns.size(); i++) {
0110       tRowSet.addToColumns(columns.get(i).toTColumn());
0111     }
0112     return tRowSet;
0113   }
0114
0115   @Override
0116   public Iterator<Object[]> iterator() {
0117     return new Iterator<Object[]>() {
0118
0119       private int index;
0120       private final Object[] convey = new Object[numColumns()];
0121
0122       @Override
0123       public boolean hasNext() {
0124         return index < numRows();
0125       }
0126
0127       @Override
0128       public Object[] next() {
0129         for (int i = 0; i < columns.size(); i++) {
0130           convey[i] = columns.get(i).get(index);
0131         }
0132         index++;
0133         return convey;
0134       }
0135
0136       @Override
0137       public void remove() {
0138         throw new UnsupportedOperationException("remove");
0139       }
0140     };
0141   }
0142
0143   public Object[] fill(int index, Object[] convey) {
0144     for (int i = 0; i < columns.size(); i++) {
0145       convey[i] = columns.get(i).get(index);
0146     }
0147     return convey;
0148   }
0149 }