org.apache.druid.query.IterableRowsCursorHelper Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of druid-processing Show documentation
A module that is everything required to understands Druid Segments
There is a newer version: 30.0.1
/*
 * 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.druid.query;

import org.apache.druid.java.util.common.Intervals;
import org.apache.druid.java.util.common.Pair;
import org.apache.druid.java.util.common.granularity.Granularities;
import org.apache.druid.java.util.common.guava.Sequence;
import org.apache.druid.java.util.common.guava.Sequences;
import org.apache.druid.java.util.common.guava.Yielder;
import org.apache.druid.segment.Cursor;
import org.apache.druid.segment.RowAdapter;
import org.apache.druid.segment.RowBasedCursor;
import org.apache.druid.segment.RowWalker;
import org.apache.druid.segment.VirtualColumns;
import org.apache.druid.segment.column.RowSignature;

import java.io.Closeable;
import java.util.Iterator;

/**
 * Helper methods to create cursor from iterable of rows
 */
public class IterableRowsCursorHelper
{

  /**
   * Creates a cursor that iterates over all the rows generated by the iterable. Presence of __time column is not a
   * necessity
   */
  public static Pair getCursorFromIterable(Iterable rows, RowSignature rowSignature)
  {
    return getCursorFromSequence(Sequences.simple(rows), rowSignature);
  }

  /**
   * Creates a cursor that iterates over all the rows generated by the sequence. Presence of __time column is not a
   * necessity.
   * 
   * Returns a pair of cursor that iterates over the rows and closeable that cleans up the created rowWalker
   */
  public static Pair getCursorFromSequence(Sequence rows, RowSignature rowSignature)
  {
    RowAdapter rowAdapter = columnName -> {
      if (rowSignature == null) {
        return row -> null;
      }
      final int columnIndex = rowSignature.indexOf(columnName);
      if (columnIndex < 0) {
        return row -> null;
      }
      return row -> row[columnIndex];
    };

    RowWalker rowWalker = new RowWalker<>(rows, rowAdapter);

    Cursor baseCursor = new RowBasedCursor<>(
        rowWalker,
        rowAdapter,
        null,
        Intervals.ETERNITY, // Setting the interval to eternity ensures that we are iterating over all of the rows
        VirtualColumns.EMPTY,
        Granularities.ALL,
        false,
        rowSignature != null ? rowSignature : RowSignature.empty()
    );

    return Pair.of(baseCursor, rowWalker);
  }

  public static Pair getCursorFromYielder(Yielder yielderParam, RowSignature rowSignature)
  {
    return getCursorFromIterable(
        new Iterable()
        {
          Yielder yielder = yielderParam;
          @Override
          public Iterator iterator()
          {
            return new Iterator()
            {
              @Override
              public boolean hasNext()
              {
                return !yielder.isDone();
              }

              @Override
              public Object[] next()
              {
                Object[] retVal = yielder.get();
                yielder = yielder.next(null);
                return retVal;
              }
            };
          }
        },
        rowSignature
    );
  }
}