Skip to content

Commit a57845b

Browse files
bibith4bibith4
committed
Changes make similarity search behind a feature flag
1 parent 95fa012 commit a57845b

File tree

11 files changed

+68
-43
lines changed

11 files changed

+68
-43
lines changed

‎presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergAbstractMetadata.java‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import com.facebook.presto.iceberg.changelog.ChangelogOperation;
3333
import com.facebook.presto.iceberg.changelog.ChangelogUtil;
3434
import com.facebook.presto.iceberg.statistics.StatisticsFileCache;
35+
import com.facebook.presto.iceberg.tvf.ApproxNearestNeighborsFunction;
3536
import com.facebook.presto.spi.ColumnHandle;
3637
import com.facebook.presto.spi.ColumnMetadata;
3738
import com.facebook.presto.spi.ConnectorDeleteTableHandle;
@@ -261,6 +262,7 @@ public abstract class IcebergAbstractMetadata
261262
protected Transaction transaction;
262263
protected final StatisticsFileCache statisticsFileCache;
263264
protected final IcebergTableProperties tableProperties;
265+
protected final IcebergConfig icebergConfig;
264266

265267
private final StandardFunctionResolution functionResolution;
266268
private final ConcurrentMap<SchemaTableName, Table> icebergTables = new ConcurrentHashMap<>();
@@ -274,7 +276,8 @@ public IcebergAbstractMetadata(
274276
NodeVersion nodeVersion,
275277
FilterStatsCalculatorService filterStatsCalculatorService,
276278
StatisticsFileCache statisticsFileCache,
277-
IcebergTableProperties tableProperties)
279+
IcebergTableProperties tableProperties,
280+
IcebergConfig icebergConfig)
278281
{
279282
this.typeManager = requireNonNull(typeManager, "typeManager is null");
280283
this.commitTaskCodec = requireNonNull(commitTaskCodec, "commitTaskCodec is null");
@@ -285,6 +288,7 @@ public IcebergAbstractMetadata(
285288
this.filterStatsCalculatorService = requireNonNull(filterStatsCalculatorService, "filterStatsCalculatorService is null");
286289
this.statisticsFileCache = requireNonNull(statisticsFileCache, "statisticsFileCache is null");
287290
this.tableProperties = requireNonNull(tableProperties, "tableProperties is null");
291+
this.icebergConfig = requireNonNull(icebergConfig, "icebergConfig is null");
288292
}
289293

290294
protected final Table getIcebergTable(ConnectorSession session, SchemaTableName schemaTableName)
@@ -1774,5 +1778,16 @@ private boolean viewExists(ConnectorSession session, ConnectorTableMetadata view
17741778
}
17751779

17761780
@Override
1777-
public abstract Optional<TableFunctionApplicationResult<ConnectorTableHandle>> applyTableFunction(ConnectorSession session, ConnectorTableFunctionHandle handle);
1781+
public Optional<TableFunctionApplicationResult<ConnectorTableHandle>> applyTableFunction(ConnectorSession session, ConnectorTableFunctionHandle handle)
1782+
{
1783+
if (!icebergConfig.isSimilaritySearchEnabled()) {
1784+
return Optional.empty();
1785+
}
1786+
if (handle instanceof ApproxNearestNeighborsFunction.IcebergAnnTableFunctionHandle) {
1787+
ApproxNearestNeighborsFunction.IcebergAnnTableFunctionHandle annTableFunctionHandle = (ApproxNearestNeighborsFunction.IcebergAnnTableFunctionHandle) handle;
1788+
return Optional.of(new TableFunctionApplicationResult<>(annTableFunctionHandle.getTableHandle(), annTableFunctionHandle.getColumnHandles()));
1789+
}
1790+
1791+
throw new IllegalArgumentException("Unsupported function handle: " + handle);
1792+
}
17781793
}

‎presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergCommonModule.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,9 @@ protected void setup(Binder binder)
191191
procedures.addBinding().toProvider(SetTablePropertyProcedure.class).in(Scopes.SINGLETON);
192192
procedures.addBinding().toProvider(StatisticsFileCacheInvalidationProcedure.class).in(Scopes.SINGLETON);
193193
procedures.addBinding().toProvider(ManifestFileCacheInvalidationProcedure.class).in(Scopes.SINGLETON);
194-
procedures.addBinding().toProvider(BuildVectorIndexProcedure.class).in(Scopes.SINGLETON);
194+
if (icebergConfig.isSimilaritySearchEnabled()) {
195+
procedures.addBinding().toProvider(BuildVectorIndexProcedure.class).in(Scopes.SINGLETON);
196+
}
195197

196198
// for orc
197199
binder.bind(EncryptionLibrary.class).annotatedWith(HiveDwrfEncryptionProvider.ForCryptoService.class).to(UnsupportedEncryptionLibrary.class).in(Scopes.SINGLETON);
@@ -202,7 +204,9 @@ protected void setup(Binder binder)
202204
configBinder(binder).bindConfig(OrcFileWriterConfig.class);
203205

204206
configBinder(binder).bindConfig(ParquetCacheConfig.class, connectorId);
205-
newSetBinder(binder, ConnectorTableFunction.class).addBinding().toProvider(ApproxNearestNeighborsFunction.class).in(Scopes.SINGLETON);
207+
if (icebergConfig.isSimilaritySearchEnabled()) {
208+
newSetBinder(binder, ConnectorTableFunction.class).addBinding().toProvider(ApproxNearestNeighborsFunction.class).in(Scopes.SINGLETON);
209+
}
206210
binder.bind(ConnectorPlanOptimizerProvider.class).to(IcebergPlanOptimizerProvider.class).in(Scopes.SINGLETON);
207211
}
208212

‎presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergConfig.java‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ public class IcebergConfig
5959
private boolean mergeOnReadModeEnabled = true;
6060
private double statisticSnapshotRecordDifferenceWeight;
6161
private boolean pushdownFilterEnabled;
62+
private boolean similaritySearchEnabled;
6263
private boolean deleteAsJoinRewriteEnabled = true;
6364
private int deleteAsJoinRewriteMaxDeleteColumns = 400;
6465
private int rowsForMetadataOptimizationThreshold = 1000;
@@ -496,4 +497,17 @@ public IcebergConfig setMaterializedViewStoragePrefix(String materializedViewSto
496497
this.materializedViewStoragePrefix = materializedViewStoragePrefix;
497498
return this;
498499
}
500+
501+
public boolean isSimilaritySearchEnabled()
502+
{
503+
return similaritySearchEnabled;
504+
}
505+
506+
@Config("iceberg.similarity-search-enabled")
507+
@ConfigDescription("Enable filter for similarity search")
508+
public IcebergConfig setSimilaritySearchEnabled(boolean similaritySearchEnabled)
509+
{
510+
this.similaritySearchEnabled = similaritySearchEnabled;
511+
return this;
512+
}
499513
}

‎presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergHiveMetadata.java‎

Lines changed: 3 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,6 @@
3636
import com.facebook.presto.hive.metastore.PrincipalPrivileges;
3737
import com.facebook.presto.hive.metastore.Table;
3838
import com.facebook.presto.iceberg.statistics.StatisticsFileCache;
39-
import com.facebook.presto.iceberg.tvf.ApproxNearestNeighborsFunction;
4039
import com.facebook.presto.spi.ColumnHandle;
4140
import com.facebook.presto.spi.ColumnMetadata;
4241
import com.facebook.presto.spi.ConnectorNewTableLayout;
@@ -55,9 +54,7 @@
5554
import com.facebook.presto.spi.SchemaTablePrefix;
5655
import com.facebook.presto.spi.TableNotFoundException;
5756
import com.facebook.presto.spi.ViewNotFoundException;
58-
import com.facebook.presto.spi.connector.TableFunctionApplicationResult;
5957
import com.facebook.presto.spi.function.StandardFunctionResolution;
60-
import com.facebook.presto.spi.function.table.ConnectorTableFunctionHandle;
6158
import com.facebook.presto.spi.plan.FilterStatsCalculatorService;
6259
import com.facebook.presto.spi.relation.RowExpressionService;
6360
import com.facebook.presto.spi.security.PrestoPrincipal;
@@ -190,9 +187,10 @@ public IcebergHiveMetadata(
190187
StatisticsFileCache statisticsFileCache,
191188
ManifestFileCache manifestFileCache,
192189
IcebergTableProperties tableProperties,
193-
ConnectorSystemConfig connectorSystemConfig)
190+
ConnectorSystemConfig connectorSystemConfig,
191+
IcebergConfig icebergConfig)
194192
{
195-
super(typeManager, functionResolution, rowExpressionService, commitTaskCodec, columnMappingsCodec, nodeVersion, filterStatsCalculatorService, statisticsFileCache, tableProperties);
193+
super(typeManager, functionResolution, rowExpressionService, commitTaskCodec, columnMappingsCodec, nodeVersion, filterStatsCalculatorService, statisticsFileCache, tableProperties, icebergConfig);
196194
this.catalogName = requireNonNull(catalogName, "catalogName is null");
197195
this.metastore = requireNonNull(metastore, "metastore is null");
198196
this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null");
@@ -725,14 +723,4 @@ protected void updateIcebergViewProperties(
725723
throw new PrestoException(NOT_SUPPORTED, "Iceberg Hive catalog does not support native Iceberg views for materialized views.");
726724
}
727725

728-
@Override
729-
public Optional<TableFunctionApplicationResult<ConnectorTableHandle>> applyTableFunction(ConnectorSession session, ConnectorTableFunctionHandle handle)
730-
{
731-
if (handle instanceof ApproxNearestNeighborsFunction.IcebergAnnTableFunctionHandle) {
732-
ApproxNearestNeighborsFunction.IcebergAnnTableFunctionHandle annTableFunctionHandle = (ApproxNearestNeighborsFunction.IcebergAnnTableFunctionHandle) handle;
733-
return Optional.of(new TableFunctionApplicationResult<>(annTableFunctionHandle.getTableHandle(), annTableFunctionHandle.getColumnHandles()));
734-
}
735-
736-
throw new IllegalArgumentException("Unsupported function handle: " + handle);
737-
}
738726
}

‎presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergHiveMetadataFactory.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ public class IcebergHiveMetadataFactory
4949
final ManifestFileCache manifestFileCache;
5050
final IcebergTableProperties tableProperties;
5151
final ConnectorSystemConfig connectorSystemConfig;
52+
final IcebergConfig icebergConfig;
5253

5354
@Inject
5455
public IcebergHiveMetadataFactory(
@@ -66,7 +67,8 @@ public IcebergHiveMetadataFactory(
6667
StatisticsFileCache statisticsFileCache,
6768
ManifestFileCache manifestFileCache,
6869
IcebergTableProperties tableProperties,
69-
ConnectorSystemConfig connectorSystemConfig)
70+
ConnectorSystemConfig connectorSystemConfig,
71+
IcebergConfig icebergConfig)
7072
{
7173
this.catalogName = requireNonNull(catalogName, "catalogName is null");
7274
this.metastore = requireNonNull(metastore, "metastore is null");
@@ -83,6 +85,7 @@ public IcebergHiveMetadataFactory(
8385
this.manifestFileCache = requireNonNull(manifestFileCache, "manifestFileCache is null");
8486
this.tableProperties = requireNonNull(tableProperties, "icebergTableProperties is null");
8587
this.connectorSystemConfig = requireNonNull(connectorSystemConfig, "connectorSystemConfig is null");
88+
this.icebergConfig = requireNonNull(icebergConfig, "icebergConfig is null");
8689
}
8790

8891
public ConnectorMetadata create()
@@ -102,6 +105,7 @@ public ConnectorMetadata create()
102105
statisticsFileCache,
103106
manifestFileCache,
104107
tableProperties,
105-
connectorSystemConfig);
108+
connectorSystemConfig,
109+
icebergConfig);
106110
}
107111
}

‎presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergNativeMetadata.java‎

Lines changed: 3 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
import com.facebook.presto.hive.NodeVersion;
1919
import com.facebook.presto.hive.TableAlreadyExistsException;
2020
import com.facebook.presto.iceberg.statistics.StatisticsFileCache;
21-
import com.facebook.presto.iceberg.tvf.ApproxNearestNeighborsFunction;
2221
import com.facebook.presto.iceberg.util.IcebergPrestoModelConverters;
2322
import com.facebook.presto.spi.ColumnMetadata;
2423
import com.facebook.presto.spi.ConnectorNewTableLayout;
@@ -31,9 +30,7 @@
3130
import com.facebook.presto.spi.PrestoException;
3231
import com.facebook.presto.spi.SchemaTableName;
3332
import com.facebook.presto.spi.SchemaTablePrefix;
34-
import com.facebook.presto.spi.connector.TableFunctionApplicationResult;
3533
import com.facebook.presto.spi.function.StandardFunctionResolution;
36-
import com.facebook.presto.spi.function.table.ConnectorTableFunctionHandle;
3734
import com.facebook.presto.spi.plan.FilterStatsCalculatorService;
3835
import com.facebook.presto.spi.relation.RowExpressionService;
3936
import com.google.common.collect.ImmutableList;
@@ -120,9 +117,10 @@ public IcebergNativeMetadata(
120117
NodeVersion nodeVersion,
121118
FilterStatsCalculatorService filterStatsCalculatorService,
122119
StatisticsFileCache statisticsFileCache,
123-
IcebergTableProperties tableProperties)
120+
IcebergTableProperties tableProperties,
121+
IcebergConfig icebergConfig)
124122
{
125-
super(typeManager, functionResolution, rowExpressionService, commitTaskCodec, columnMappingsCodec, nodeVersion, filterStatsCalculatorService, statisticsFileCache, tableProperties);
123+
super(typeManager, functionResolution, rowExpressionService, commitTaskCodec, columnMappingsCodec, nodeVersion, filterStatsCalculatorService, statisticsFileCache, tableProperties, icebergConfig);
126124
this.catalogFactory = requireNonNull(catalogFactory, "catalogFactory is null");
127125
this.catalogType = requireNonNull(catalogType, "catalogType is null");
128126
this.warehouseDataDir = Optional.ofNullable(catalogFactory.getCatalogWarehouseDataDir());
@@ -532,14 +530,4 @@ protected void updateIcebergViewProperties(
532530
icebergViews.remove(viewName);
533531
}
534532

535-
@Override
536-
public Optional<TableFunctionApplicationResult<ConnectorTableHandle>> applyTableFunction(ConnectorSession session, ConnectorTableFunctionHandle handle)
537-
{
538-
if (handle instanceof ApproxNearestNeighborsFunction.IcebergAnnTableFunctionHandle) {
539-
ApproxNearestNeighborsFunction.IcebergAnnTableFunctionHandle annTableFunctionHandle = (ApproxNearestNeighborsFunction.IcebergAnnTableFunctionHandle) handle;
540-
return Optional.of(new TableFunctionApplicationResult<>(annTableFunctionHandle.getTableHandle(), annTableFunctionHandle.getColumnHandles()));
541-
}
542-
543-
throw new IllegalArgumentException("Unsupported function handle: " + handle);
544-
}
545533
}

‎presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergNativeMetadataFactory.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ public class IcebergNativeMetadataFactory
4242
final FilterStatsCalculatorService filterStatsCalculatorService;
4343
final StatisticsFileCache statisticsFileCache;
4444
final IcebergTableProperties tableProperties;
45+
final IcebergConfig icebergConfig;
4546

4647
@Inject
4748
public IcebergNativeMetadataFactory(
@@ -55,7 +56,8 @@ public IcebergNativeMetadataFactory(
5556
NodeVersion nodeVersion,
5657
FilterStatsCalculatorService filterStatsCalculatorService,
5758
StatisticsFileCache statisticsFileCache,
58-
IcebergTableProperties tableProperties)
59+
IcebergTableProperties tableProperties,
60+
IcebergConfig icebergConfig)
5961
{
6062
this.catalogFactory = requireNonNull(catalogFactory, "catalogFactory is null");
6163
this.typeManager = requireNonNull(typeManager, "typeManager is null");
@@ -68,10 +70,11 @@ public IcebergNativeMetadataFactory(
6870
this.filterStatsCalculatorService = requireNonNull(filterStatsCalculatorService, "filterStatsCalculatorService is null");
6971
this.statisticsFileCache = requireNonNull(statisticsFileCache, "statisticsFileCache is null");
7072
this.tableProperties = requireNonNull(tableProperties, "tableProperties is null");
73+
this.icebergConfig = requireNonNull(icebergConfig, "icebergConfig is null");
7174
}
7275

7376
public ConnectorMetadata create()
7477
{
75-
return new IcebergNativeMetadata(catalogFactory, typeManager, functionResolution, rowExpressionService, commitTaskCodec, columnMappingsCodec, catalogType, nodeVersion, filterStatsCalculatorService, statisticsFileCache, tableProperties);
78+
return new IcebergNativeMetadata(catalogFactory, typeManager, functionResolution, rowExpressionService, commitTaskCodec, columnMappingsCodec, catalogType, nodeVersion, filterStatsCalculatorService, statisticsFileCache, tableProperties, icebergConfig);
7679
}
7780
}

‎presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergPageSourceProvider.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -216,6 +216,7 @@ public class IcebergPageSourceProvider
216216
private final IcebergHiveTableOperationsConfig tableOperationsConfig;
217217
private final ManifestFileCache manifestFileCache;
218218
private final IcebergCatalogName catalogName;
219+
private final boolean similaritySearchEnabled;
219220

220221
@Inject
221222
public IcebergPageSourceProvider(
@@ -255,6 +256,7 @@ public IcebergPageSourceProvider(
255256
this.tableOperationsConfig = requireNonNull(tableOperationsConfig, "tableOperationsConfig is null");
256257
this.manifestFileCache = requireNonNull(manifestFileCache, "manifestFileCache is null");
257258
this.catalogName = requireNonNull(catalogName, "catalogName is null");
259+
this.similaritySearchEnabled = icebergConfig.isSimilaritySearchEnabled();
258260
}
259261

260262
private static ConnectorPageSourceWithRowPositions createParquetPageSource(
@@ -750,7 +752,7 @@ public ConnectorPageSource createPageSource(
750752

751753
IcebergSplit split = (IcebergSplit) connectorSplit;
752754
IcebergTableHandle table = icebergLayout.getTable();
753-
if (split.isAnn()) {
755+
if (similaritySearchEnabled && split.isAnn()) {
754756
SchemaTableName schemaTableName = table.getSchemaTableName();
755757
Table icebergTable = IcebergUtil.getHiveIcebergTable(
756758
metastore,

‎presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergSplitManager.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -58,17 +58,20 @@ public class IcebergSplitManager
5858
private final TypeManager typeManager;
5959
private final ExecutorService executor;
6060
private final ThreadPoolExecutorMBean executorServiceMBean;
61+
private final IcebergConfig icebergConfig;
6162

6263
@Inject
6364
public IcebergSplitManager(
6465
IcebergTransactionManager transactionManager,
6566
TypeManager typeManager,
66-
@ForIcebergSplitManager ExecutorService executor)
67+
@ForIcebergSplitManager ExecutorService executor,
68+
IcebergConfig icebergConfig)
6769
{
6870
this.transactionManager = requireNonNull(transactionManager, "transactionManager is null");
6971
this.typeManager = requireNonNull(typeManager, "typeManager is null");
7072
this.executor = requireNonNull(executor, "executor is null");
7173
this.executorServiceMBean = new ThreadPoolExecutorMBean((ThreadPoolExecutor) executor);
74+
this.icebergConfig = requireNonNull(icebergConfig, "icebergConfig is null");
7275
}
7376

7477
@Override
@@ -81,7 +84,7 @@ public ConnectorSplitSource getSplits(
8184
IcebergTableLayoutHandle layoutHandle = (IcebergTableLayoutHandle) layout;
8285
IcebergTableHandle table = layoutHandle.getTable();
8386

84-
if (table instanceof ApproxNearestNeighborsFunction.IcebergAnnTableHandle) {
87+
if (icebergConfig.isSimilaritySearchEnabled() && table instanceof ApproxNearestNeighborsFunction.IcebergAnnTableHandle) {
8588
ApproxNearestNeighborsFunction.IcebergAnnTableHandle annHandle = (ApproxNearestNeighborsFunction.IcebergAnnTableHandle) table;
8689
IcebergSplit annSplit = new IcebergSplit(
8790
/* path */ "", // non-null dummy string

‎presto-iceberg/src/main/java/com/facebook/presto/iceberg/InternalIcebergConnectorFactory.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -122,13 +122,16 @@ public static Connector createConnector(
122122
IcebergSessionProperties icebergSessionProperties = injector.getInstance(IcebergSessionProperties.class);
123123
HiveCommonSessionProperties hiveCommonSessionProperties = injector.getInstance(HiveCommonSessionProperties.class);
124124
IcebergTableProperties icebergTableProperties = injector.getInstance(IcebergTableProperties.class);
125+
IcebergConfig icebergConfig = injector.getInstance(IcebergConfig.class);
125126
Set<BaseProcedure<?>> procedures =
126127
injector.getInstance(Key.get(new TypeLiteral<Set<BaseProcedure<?>>>() {}));
127128
ConnectorPlanOptimizerProvider planOptimizerProvider = injector.getInstance(ConnectorPlanOptimizerProvider.class);
128129

129130
List<PropertyMetadata<?>> allSessionProperties = new ArrayList<>(icebergSessionProperties.getSessionProperties());
130131
allSessionProperties.addAll(hiveCommonSessionProperties.getSessionProperties());
131-
Set<ConnectorTableFunction> connectorTableFunctions = new HashSet<>(Collections.singleton(new ApproxNearestNeighborsFunction().get())); // TODO update accordingly
132+
Set<ConnectorTableFunction> connectorTableFunctions = icebergConfig.isSimilaritySearchEnabled()
133+
? new HashSet<>(Collections.singleton(new ApproxNearestNeighborsFunction().get()))
134+
: new HashSet<>();
132135
return new IcebergConnector(
133136
lifeCycleManager,
134137
transactionManager,

0 commit comments

Comments
 (0)