cgivre commented on code in PR #3067: URL: https://github.com/apache/drill/pull/3067#discussion_r3958982286
########## contrib/storage-accumulo/src/main/java/org/apache/drill/exec/store/accumulo/AccumuloScanBatchCreator.java: ########## @@ -0,0 +1,132 @@ +/* + * 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.drill.exec.store.accumulo; + +import java.util.LinkedList; +import java.util.List; + +import org.apache.accumulo.core.client.AccumuloClient; +import org.apache.drill.common.exceptions.ExecutionSetupException; +import org.apache.drill.common.expression.SchemaPath; +import org.apache.drill.exec.ops.ExecutorFragmentContext; +import org.apache.drill.exec.physical.base.GroupScan; +import org.apache.drill.exec.physical.impl.BatchCreator; +import org.apache.drill.exec.physical.impl.ScanBatch; +import org.apache.drill.exec.record.RecordBatch; +import org.apache.drill.exec.store.RecordReader; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import com.google.common.base.Preconditions; + +/** + * BatchCreator for Accumulo scan operations. + * + * <p>This class creates the execution pipeline for Accumulo scans by wiring + * together the AccumuloSubScan with AccumuloRecordReaders.</p> + * + * <p>For user impersonation mode, this class creates clients using the + * delegation token passed from the SubScan. When a delegation token is + * present, the reader owns the client and is responsible for closing it.</p> + */ +public class AccumuloScanBatchCreator implements BatchCreator<AccumuloSubScan> { + private static final Logger logger = LoggerFactory.getLogger(AccumuloScanBatchCreator.class); + + @Override + public ScanBatch getBatch( + ExecutorFragmentContext context, + AccumuloSubScan subScan, + List<RecordBatch> children) throws ExecutionSetupException { + + Preconditions.checkArgument(children.isEmpty(), "AccumuloSubScan should have no children"); + + List<RecordReader> readers = new LinkedList<>(); + List<SchemaPath> columns = subScan.getColumns(); + + if (columns == null) { + columns = GroupScan.ALL_COLUMNS; + } + + try { + // Determine if we need to create a new client from delegation token + // or use the shared service client + AccumuloClient client; + boolean ownsClient; + + if (subScan.hasDelegationToken()) { + // User impersonation mode: create a new client from the delegation token + // The reader will own this client and close it when done + DelegationTokenInfo tokenInfo = subScan.getDelegationTokenInfo(); + logger.debug("Creating Accumulo client from delegation token for user: {}", + tokenInfo.getUserName()); + + client = subScan.getStoragePlugin().getConnectionManager() + .createClientWithDelegationToken(tokenInfo); + ownsClient = true; + + logger.info("Created impersonated Accumulo client for user '{}' to scan table '{}'", + tokenInfo.getUserName(), subScan.getScanSpec().getTableName()); + } else { + // Shared user mode: use the service client + // The reader does not own this client and should not close it + client = subScan.getStoragePlugin().getClient(); + ownsClient = false; + } + + // Create a record reader for this sub-scan + // In the future, we may have multiple readers for different tablet ranges + AccumuloRecordReader reader = new AccumuloRecordReader( + client, + subScan.getScanSpec(), + columns, + getMaxRecords(subScan), + ownsClient); + + readers.add(reader); + + } catch (Exception e) { + throw new ExecutionSetupException( Review Comment: Fixed ########## contrib/storage-accumulo/src/main/java/org/apache/drill/exec/store/accumulo/AccumuloRecordReader.java: ########## @@ -0,0 +1,473 @@ +/* + * 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.drill.exec.store.accumulo; + +import java.nio.charset.StandardCharsets; +import java.util.Collection; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.TimeUnit; + +import org.apache.accumulo.core.client.AccumuloClient; +import org.apache.accumulo.core.client.Scanner; +import org.apache.accumulo.core.client.TableNotFoundException; +import org.apache.accumulo.core.data.Key; +import org.apache.accumulo.core.data.Range; +import org.apache.accumulo.core.data.Value; +import org.apache.accumulo.core.security.Authorizations; +import org.apache.drill.common.exceptions.DrillRuntimeException; +import org.apache.drill.common.exceptions.ExecutionSetupException; +import org.apache.drill.common.expression.PathSegment; +import org.apache.drill.common.expression.PathSegment.NameSegment; +import org.apache.drill.common.expression.SchemaPath; +import org.apache.drill.exec.exception.SchemaChangeException; +import org.apache.drill.exec.ops.OperatorContext; +import org.apache.drill.exec.ops.OperatorStats; +import org.apache.drill.exec.physical.impl.OutputMutator; +import org.apache.drill.exec.record.MaterializedField; +import org.apache.drill.exec.store.AbstractRecordReader; +import org.apache.drill.exec.vector.NullableVarBinaryVector; +import org.apache.drill.exec.vector.ValueVector; +import org.apache.drill.exec.vector.VarBinaryVector; +import org.apache.drill.exec.vector.complex.MapVector; +import org.apache.hadoop.io.Text; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import com.google.common.base.Preconditions; +import com.google.common.base.Stopwatch; +import com.google.common.collect.Sets; + +/** + * RecordReader for Accumulo storage plugin. + * + * <p>This reader scans Accumulo tables and populates Drill value vectors. + * It uses the dynamic schema approach similar to HBase, where column families + * are represented as maps containing their qualifiers as fields.</p> + * + * <p>Row structure:</p> + * <ul> + * <li>row_key: VARBINARY - the Accumulo row key</li> + * <li>Each column family becomes a MAP with qualifier names as keys</li> + * </ul> + * + * <p>For user impersonation mode, the reader may own the AccumuloClient + * (created from a delegation token) and is responsible for closing it. + * For shared user mode, the reader uses a shared client and should not close it.</p> + */ +public class AccumuloRecordReader extends AbstractRecordReader implements DrillAccumuloConstants { + private static final Logger logger = LoggerFactory.getLogger(AccumuloRecordReader.class); + + // Batch constraints to avoid OOM + private static final int MAX_ALLOCATED_MEMORY_PER_BATCH = 64 * 1024 * 1024; // 64 MB + private static final int TARGET_RECORD_COUNT = DEFAULT_BATCH_SIZE; + + private final AccumuloClient client; + private final AccumuloScanSpec scanSpec; + private final int maxRecords; + + /** + * Whether this reader owns the client and should close it. + * True for user impersonation mode (client created from delegation token), + * false for shared user mode (client is shared/pooled). + */ + private final boolean ownsClient; + + private OutputMutator outputMutator; + private OperatorContext operatorContext; + + private Scanner scanner; + private Iterator<Map.Entry<Key, Value>> scanIterator; + + private Map<String, MapVector> familyVectorMap; + private VarBinaryVector rowKeyVector; + + private Set<String> requestedFamilies; + private Map<String, Set<String>> requestedColumns; // family -> set of qualifiers + private boolean rowKeyOnly; + private int recordsRead; + + /** + * Creates an AccumuloRecordReader with a shared client (does not own the client). + */ + public AccumuloRecordReader( + AccumuloClient client, + AccumuloScanSpec scanSpec, + List<SchemaPath> projectedColumns, + int maxRecords) { + this(client, scanSpec, projectedColumns, maxRecords, false); + } + + /** + * Creates an AccumuloRecordReader with explicit client ownership. + * + * @param client the Accumulo client to use + * @param scanSpec the scan specification + * @param projectedColumns columns to project + * @param maxRecords maximum records to read (-1 for unlimited) + * @param ownsClient true if this reader owns the client and should close it + */ + public AccumuloRecordReader( + AccumuloClient client, + AccumuloScanSpec scanSpec, + List<SchemaPath> projectedColumns, + int maxRecords, + boolean ownsClient) { + this.client = Preconditions.checkNotNull(client, "AccumuloClient required"); + this.scanSpec = Preconditions.checkNotNull(scanSpec, "AccumuloScanSpec required"); + this.maxRecords = maxRecords > 0 ? maxRecords : Integer.MAX_VALUE; + this.recordsRead = 0; + this.ownsClient = ownsClient; + + setColumns(projectedColumns); + + if (ownsClient) { + logger.debug("RecordReader owns the AccumuloClient and will close it when done"); + } + } + + /** + * Transforms projected columns and determines which Accumulo columns to fetch. + */ + @Override + protected Collection<SchemaPath> transformColumns(Collection<SchemaPath> columns) { + Set<SchemaPath> transformed = Sets.newLinkedHashSet(); + requestedFamilies = Sets.newHashSet(); + requestedColumns = new HashMap<>(); + + rowKeyOnly = true; + + if (!isStarQuery()) { + for (SchemaPath column : columns) { + if (column.getRootSegment().getPath().equalsIgnoreCase(ROW_KEY)) { + transformed.add(ROW_KEY_PATH); + continue; + } + + rowKeyOnly = false; + NameSegment root = column.getRootSegment(); + String family = root.getPath(); + transformed.add(SchemaPath.getSimplePath(family)); + + PathSegment child = root.getChild(); + if (child != null && child.isNamed()) { + // Specific column within family: cf.qualifier + String qualifier = child.getNameSegment().getPath(); + requestedColumns.computeIfAbsent(family, k -> Sets.newHashSet()).add(qualifier); + } else { + // Entire column family requested + requestedFamilies.add(family); + } + } + } else { + rowKeyOnly = false; + transformed.add(ROW_KEY_PATH); + } + + return transformed; + } + + @Override + public void setup(OperatorContext context, OutputMutator output) throws ExecutionSetupException { + this.operatorContext = context; + this.outputMutator = output; + familyVectorMap = new HashMap<>(); + + try { + // Create scanner + scanner = client.createScanner(scanSpec.getTableName(), Authorizations.EMPTY); Review Comment: Fixed -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
