feat: Added IcebergArrowInputSourceReader

extensions-contrib/druid-iceberg-extensions/pom.xml

@@ -760,6 +760,23 @@
       <scope>provided</scope>
     </dependency>
 
+    <dependency>
+      <groupId>org.apache.iceberg</groupId>
+      <artifactId>iceberg-arrow</artifactId>
+      <version>${iceberg.core.version}</version>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.arrow</groupId>
+      <artifactId>arrow-vector</artifactId>
+      <version>${arrow.version}</version>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.arrow</groupId>
+      <artifactId>arrow-memory-unsafe</artifactId>
+      <version>${arrow.version}</version>
+      <scope>runtime</scope>
+    </dependency>
+
     <dependency>
       <groupId>org.apache.iceberg</groupId>
       <artifactId>iceberg-parquet</artifactId>

extensions-contrib/druid-iceberg-extensions/src/main/java/org/apache/druid/iceberg/common/IcebergDruidModule.java

@@ -27,6 +27,7 @@
 import org.apache.druid.iceberg.guice.HiveConf;
 import org.apache.druid.iceberg.input.GlueIcebergCatalog;
 import org.apache.druid.iceberg.input.HiveIcebergCatalog;
+import org.apache.druid.iceberg.input.IcebergArrowInputSource;
 import org.apache.druid.iceberg.input.IcebergInputSource;
 import org.apache.druid.iceberg.input.LocalCatalog;
 import org.apache.druid.iceberg.input.RestIcebergCatalog;
@@ -49,6 +50,7 @@ public List<? extends Module> getJacksonModules()
                 new NamedType(LocalCatalog.class, LocalCatalog.TYPE_KEY),
                 new NamedType(RestIcebergCatalog.class, RestIcebergCatalog.TYPE_KEY),
                 new NamedType(IcebergInputSource.class, IcebergInputSource.TYPE_KEY),
+                new NamedType(IcebergArrowInputSource.class, IcebergArrowInputSource.TYPE_KEY),
                 new NamedType(GlueIcebergCatalog.class, GlueIcebergCatalog.TYPE_KEY)
             )
     );

extensions-contrib/druid-iceberg-extensions/src/main/java/org/apache/druid/iceberg/input/AbstractIcebergInputSource.java

@@ -0,0 +1,98 @@
+/*
+ * 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.iceberg.input;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.base.Preconditions;
+import org.apache.druid.common.config.Configs;
+import org.apache.druid.iceberg.filter.IcebergFilter;
+import org.apache.iceberg.Table;
+import org.joda.time.DateTime;
+
+import javax.annotation.Nullable;
+
+public abstract class AbstractIcebergInputSource
+{
+  protected final String tableName;
+  protected final String namespace;
+  protected final IcebergCatalog icebergCatalog;
+  protected final IcebergFilter icebergFilter;
+  protected final DateTime snapshotTime;
+  protected final ResidualFilterMode residualFilterMode;
+
+  protected AbstractIcebergInputSource(
+      final String tableName,
+      final String namespace,
+      @Nullable final IcebergFilter icebergFilter,
+      final IcebergCatalog icebergCatalog,
+      @Nullable final DateTime snapshotTime,
+      @Nullable final ResidualFilterMode residualFilterMode
+  )
+  {
+    this.tableName = Preconditions.checkNotNull(tableName, "tableName cannot be null");
+    this.namespace = Preconditions.checkNotNull(namespace, "namespace cannot be null");
+    this.icebergCatalog = Preconditions.checkNotNull(icebergCatalog, "icebergCatalog cannot be null");
+    this.icebergFilter = icebergFilter;
+    this.snapshotTime = snapshotTime;
+    this.residualFilterMode = Configs.valueOrDefault(residualFilterMode, ResidualFilterMode.IGNORE);
+  }
+
+  @JsonProperty
+  public String getTableName()
+  {
+    return tableName;
+  }
+
+  @JsonProperty
+  public String getNamespace()
+  {
+    return namespace;
+  }
+
+  @JsonProperty
+  public IcebergCatalog getIcebergCatalog()
+  {
+    return icebergCatalog;
+  }
+
+  @JsonProperty
+  public IcebergFilter getIcebergFilter()
+  {
+    return icebergFilter;
+  }
+
+  @Nullable
+  @JsonProperty
+  public DateTime getSnapshotTime()
+  {
+    return snapshotTime;
+  }
+
+  @JsonProperty
+  public ResidualFilterMode getResidualFilterMode()
+  {
+    return residualFilterMode;
+  }
+
+  protected Table retrieveTable()
+  {
+    return icebergCatalog.retrieveTable(namespace, tableName);
+  }
+}

extensions-contrib/druid-iceberg-extensions/src/main/java/org/apache/druid/iceberg/input/IcebergArrowInputSource.java

@@ -0,0 +1,101 @@
+/*
+ * 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.iceberg.input;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import org.apache.druid.data.input.InputFormat;
+import org.apache.druid.data.input.InputRowSchema;
+import org.apache.druid.data.input.InputSource;
+import org.apache.druid.data.input.InputSourceReader;
+import org.apache.druid.iceberg.filter.IcebergFilter;
+import org.apache.iceberg.Table;
+import org.apache.iceberg.TableScan;
+import org.joda.time.DateTime;
+
+import javax.annotation.Nullable;
+import java.io.File;
+
+public class IcebergArrowInputSource extends AbstractIcebergInputSource implements InputSource
+{
+  public static final String TYPE_KEY = "iceberg_arrow";
+
+  @JsonProperty
+  private final int arrowBatchSize;
+
+  @JsonCreator
+  public IcebergArrowInputSource(
+      @JsonProperty("tableName") String tableName,
+      @JsonProperty("namespace") String namespace,
+      @JsonProperty("icebergFilter") @Nullable IcebergFilter icebergFilter,
+      @JsonProperty("icebergCatalog") IcebergCatalog icebergCatalog,
+      @JsonProperty("snapshotTime") @Nullable DateTime snapshotTime,
+      @JsonProperty("residualFilterMode") @Nullable ResidualFilterMode residualFilterMode,
+      @JsonProperty("arrowBatchSize") @Nullable Integer arrowBatchSize
+  )
+  {
+    super(tableName, namespace, icebergFilter, icebergCatalog, snapshotTime, residualFilterMode);
+    this.arrowBatchSize = arrowBatchSize != null && arrowBatchSize > 0
+                          ? arrowBatchSize
+                          : IcebergArrowInputSourceReader.DEFAULT_BATCH_SIZE;
+  }
+
+  @JsonProperty
+  public int getArrowBatchSize()
+  {
+    return arrowBatchSize;
+  }
+
+  @Override
+  public boolean needsFormat()
+  {
+    return false;
+  }
+
+  @Override
+  public boolean isSplittable()
+  {
+    return false;
+  }
+
+  @Override
+  public InputSourceReader reader(
+      InputRowSchema inputRowSchema,
+      @Nullable InputFormat inputFormat,
+      File temporaryDirectory
+  )
+  {
+    final Table table = retrieveTable();
+    if (icebergFilter != null) {
+      final TableScan filteredScan = icebergFilter.filter(
+          table.newScan().caseSensitive(icebergCatalog.isCaseSensitive())
+      );
+      icebergCatalog.enforceResidualMode(filteredScan, getResidualFilterMode());
+    }
+    return new IcebergArrowInputSourceReader(
+        table,
+        icebergFilter,
+        snapshotTime,
+        icebergCatalog.isCaseSensitive(),
+        inputRowSchema,
+        arrowBatchSize
+    );
+  }
+}