Skip to content

Commit 40fd68c

Browse files
zhangyue19921010zhangyue19921010
authored
feat: support time travel (#201)
Support `TIMESTAMP AS OF ` and `VERSION AS OF` ``` spark.sql("SELECT * FROM " + fullName + " TIMESTAMP AS OF " + ts1); spark.sql("SELECT * FROM " + fullName + " VERSION AS OF " + "2"); ```
1 parent 5804975 commit 40fd68c

10 files changed

Lines changed: 415 additions & 80 deletions

File tree

‎lance-spark-base_2.12/src/main/java/org/lance/spark/BaseLanceNamespaceSparkCatalog.java‎

Lines changed: 76 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import org.lance.spark.function.LanceFragmentIdWithDefaultFunction;
2727
import org.lance.spark.utils.Optional;
2828
import org.lance.spark.utils.SchemaConverter;
29+
import org.lance.spark.utils.Utils;
2930

3031
import org.apache.spark.sql.catalyst.analysis.NamespaceAlreadyExistsException;
3132
import org.apache.spark.sql.catalyst.analysis.NoSuchFunctionException;
@@ -57,6 +58,9 @@
5758
import java.util.stream.Collectors;
5859
import java.util.stream.Stream;
5960

61+
import static org.lance.spark.utils.Utils.createReadOptions;
62+
import static org.lance.spark.utils.Utils.getSchema;
63+
6064
public abstract class BaseLanceNamespaceSparkCatalog
6165
implements TableCatalog, SupportsNamespaces, FunctionCatalog {
6266

@@ -411,60 +415,17 @@ public boolean tableExists(Identifier ident) {
411415

412416
@Override
413417
public Table loadTable(Identifier ident) throws NoSuchTableException {
414-
// Transform identifier for API call
415-
Identifier actualIdent = transformIdentifierForApi(ident);
416-
417-
// Build the table ID for credential vending
418-
List<String> tableId = buildTableId(actualIdent);
419-
420-
// Call describeTable to get location and initial storage options
421-
DescribeTableRequest describeRequest = new DescribeTableRequest();
422-
tableId.forEach(describeRequest::addIdItem);
423-
DescribeTableResponse describeResponse;
424-
try {
425-
describeResponse = namespace.describeTable(describeRequest);
426-
} catch (TableNotFoundException e) {
427-
throw new NoSuchTableException(ident);
428-
} catch (RuntimeException e) {
429-
throw new RuntimeException("Failed to describe table: " + ident, e);
430-
}
431-
432-
String location = describeResponse.getLocation();
433-
Map<String, String> initialStorageOptions = describeResponse.getStorageOptions();
434-
435-
// Open dataset to get schema
436-
StructType schema;
437-
try (Dataset dataset =
438-
Dataset.open()
439-
.allocator(LanceRuntime.allocator())
440-
.namespace(namespace)
441-
.tableId(tableId)
442-
.build()) {
443-
schema = LanceArrowUtils.fromArrowSchema(dataset.getSchema());
444-
} catch (TableNotFoundException e) {
445-
throw new NoSuchTableException(ident);
446-
}
418+
return loadTableInternal(ident, Optional.empty(), Optional.empty());
419+
}
447420

448-
// Create read options with namespace support
449-
LanceSparkReadOptions readOptions = createReadOptions(location, tableId);
450-
return createDataset(
451-
readOptions, schema, initialStorageOptions, namespaceImpl, namespaceProperties);
421+
@Override
422+
public Table loadTable(Identifier ident, String version) throws NoSuchTableException {
423+
return loadTableInternal(ident, Optional.empty(), Optional.of(version));
452424
}
453425

454-
/**
455-
* Creates LanceSparkReadOptions with namespace settings for this catalog.
456-
*
457-
* @param location the dataset location URI
458-
* @param tableId the table identifier within the namespace
459-
* @return a new LanceSparkReadOptions with all catalog settings
460-
*/
461-
private LanceSparkReadOptions createReadOptions(String location, List<String> tableId) {
462-
return LanceSparkReadOptions.builder()
463-
.datasetUri(location)
464-
.withCatalogDefaults(catalogConfig)
465-
.namespace(namespace)
466-
.tableId(tableId)
467-
.build();
426+
@Override
427+
public Table loadTable(Identifier ident, long timestamp) throws NoSuchTableException {
428+
return loadTableInternal(ident, Optional.of(timestamp), Optional.empty());
468429
}
469430

470431
@Override
@@ -500,7 +461,13 @@ public Table createTable(
500461
Map<String, String> initialStorageOptions = describeResponse.getStorageOptions();
501462

502463
// Create read options with namespace settings
503-
LanceSparkReadOptions readOptions = createReadOptions(location, tableIdList);
464+
LanceSparkReadOptions readOptions =
465+
createReadOptions(
466+
location,
467+
catalogConfig,
468+
Optional.empty(),
469+
Optional.of(namespace),
470+
Optional.of(tableIdList));
504471
return createDataset(
505472
readOptions, processedSchema, initialStorageOptions, namespaceImpl, namespaceProperties);
506473
}
@@ -651,6 +618,63 @@ private List<String> buildTableId(Identifier ident) {
651618
.collect(Collectors.toList());
652619
}
653620

621+
private Table loadTableInternal(
622+
Identifier ident, Optional<Long> timestamp, Optional<String> version)
623+
throws NoSuchTableException {
624+
625+
// Transform identifier for API call
626+
Identifier actualIdent = transformIdentifierForApi(ident);
627+
628+
// Build the table ID for credential vending
629+
List<String> tableId = buildTableId(actualIdent);
630+
631+
// Call describeTable to get location and initial storage options
632+
DescribeTableRequest describeRequest = new DescribeTableRequest();
633+
tableId.forEach(describeRequest::addIdItem);
634+
DescribeTableResponse describeResponse;
635+
try {
636+
describeResponse = namespace.describeTable(describeRequest);
637+
} catch (TableNotFoundException e) {
638+
throw new NoSuchTableException(ident);
639+
} catch (RuntimeException e) {
640+
throw new RuntimeException("Failed to describe table: " + ident, e);
641+
}
642+
String location = describeResponse.getLocation();
643+
Map<String, String> initialStorageOptions = describeResponse.getStorageOptions();
644+
645+
Optional<Long> versionId = Optional.empty();
646+
if (timestamp.isPresent()) {
647+
try (Dataset dataset =
648+
Dataset.open()
649+
.allocator(LanceRuntime.allocator())
650+
.uri(location)
651+
.readOptions(
652+
createReadOptions(
653+
location,
654+
catalogConfig,
655+
Optional.empty(),
656+
Optional.of(namespace),
657+
Optional.of(tableId))
658+
.toReadOptions())
659+
.build()) {
660+
versionId = Optional.of(Utils.findVersion(dataset.listVersions(), timestamp.get()));
661+
} catch (TableNotFoundException e) {
662+
throw new NoSuchTableException(ident);
663+
}
664+
} else if (version.isPresent()) {
665+
versionId = Optional.of(Utils.parseVersion(version.get()));
666+
}
667+
668+
LanceSparkReadOptions readOptions =
669+
createReadOptions(
670+
location, catalogConfig, versionId, Optional.of(namespace), Optional.of(tableId));
671+
StructType schema = getSchema(ident, location, readOptions, namespace);
672+
673+
// Create read options with namespace support
674+
return createDataset(
675+
readOptions, schema, initialStorageOptions, namespaceImpl, namespaceProperties);
676+
}
677+
654678
public abstract LanceDataset createDataset(
655679
LanceSparkReadOptions readOptions,
656680
StructType sparkSchema,

‎lance-spark-base_2.12/src/main/java/org/lance/spark/LanceCatalog.java‎

Lines changed: 55 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515

1616
import org.lance.Dataset;
1717
import org.lance.WriteParams;
18+
import org.lance.spark.utils.Optional;
19+
import org.lance.spark.utils.Utils;
1820

1921
import org.apache.spark.sql.catalyst.analysis.NoSuchNamespaceException;
2022
import org.apache.spark.sql.catalyst.analysis.NoSuchTableException;
@@ -30,6 +32,9 @@
3032

3133
import java.util.Map;
3234

35+
import static org.lance.spark.utils.Utils.createReadOptions;
36+
import static org.lance.spark.utils.Utils.getSchema;
37+
3338
/**
3439
* A simple Lance catalog that supports both path-based and catalog-based table access.
3540
*
@@ -52,28 +57,27 @@ public Identifier[] listTables(String[] namespace) throws NoSuchNamespaceExcepti
5257

5358
@Override
5459
public Table loadTable(Identifier ident) throws NoSuchTableException {
55-
String datasetUri = getDatasetUri(ident);
56-
LanceSparkReadOptions readOptions = createReadOptions(datasetUri);
57-
StructType schema;
58-
try (Dataset dataset =
59-
Dataset.open()
60-
.allocator(LanceRuntime.allocator())
61-
.uri(datasetUri)
62-
.readOptions(readOptions.toReadOptions())
63-
.build()) {
64-
schema = LanceArrowUtils.fromArrowSchema(dataset.getSchema());
65-
} catch (IllegalArgumentException e) {
66-
throw new NoSuchTableException(ident);
67-
}
68-
return new LanceDataset(readOptions, schema, null, null, null);
60+
return loadTableInternal(ident, Optional.empty(), Optional.empty());
61+
}
62+
63+
@Override
64+
public Table loadTable(Identifier ident, String version) throws NoSuchTableException {
65+
return loadTableInternal(ident, Optional.empty(), Optional.of(version));
66+
}
67+
68+
@Override
69+
public Table loadTable(Identifier ident, long timestamp) throws NoSuchTableException {
70+
return loadTableInternal(ident, Optional.of(timestamp), Optional.empty());
6971
}
7072

7173
@Override
7274
public Table createTable(
7375
Identifier ident, StructType schema, Transform[] partitions, Map<String, String> properties)
7476
throws TableAlreadyExistsException, NoSuchNamespaceException {
7577
String datasetUri = getDatasetUri(ident);
76-
LanceSparkReadOptions readOptions = createReadOptions(datasetUri);
78+
LanceSparkReadOptions readOptions =
79+
createReadOptions(
80+
datasetUri, catalogConfig, Optional.empty(), Optional.empty(), Optional.empty());
7781
try {
7882
Dataset.write()
7983
.allocator(LanceRuntime.allocator())
@@ -115,19 +119,6 @@ public void initialize(String name, CaseInsensitiveStringMap options) {
115119
this.catalogConfig = LanceSparkCatalogConfig.from(options.asCaseSensitiveMap());
116120
}
117121

118-
/**
119-
* Creates LanceSparkReadOptions for this catalog.
120-
*
121-
* @param datasetUri the dataset URI
122-
* @return a new LanceSparkReadOptions with catalog settings
123-
*/
124-
private LanceSparkReadOptions createReadOptions(String datasetUri) {
125-
return LanceSparkReadOptions.builder()
126-
.datasetUri(datasetUri)
127-
.withCatalogDefaults(catalogConfig)
128-
.build();
129-
}
130-
131122
@Override
132123
public String name() {
133124
return catalogName;
@@ -169,4 +160,40 @@ private String getDatasetUri(Identifier ident) {
169160
sb.append(name);
170161
return sb.toString();
171162
}
163+
164+
private Table loadTableInternal(
165+
Identifier ident, Optional<Long> timestamp, Optional<String> version)
166+
throws NoSuchTableException {
167+
String datasetUri = getDatasetUri(ident);
168+
169+
Optional<Long> versionId = Optional.empty();
170+
171+
if (version.isPresent()) {
172+
versionId = Optional.of(Utils.parseVersion(version.get()));
173+
} else if (timestamp.isPresent()) {
174+
try (Dataset dataset =
175+
Dataset.open()
176+
.allocator(LanceRuntime.allocator())
177+
.uri(datasetUri)
178+
.readOptions(
179+
createReadOptions(
180+
datasetUri,
181+
catalogConfig,
182+
Optional.empty(),
183+
Optional.empty(),
184+
Optional.empty())
185+
.toReadOptions())
186+
.build()) {
187+
versionId = Optional.of(Utils.findVersion(dataset.listVersions(), timestamp.get()));
188+
} catch (IllegalArgumentException e) {
189+
throw new NoSuchTableException(ident);
190+
}
191+
}
192+
193+
LanceSparkReadOptions readOptions =
194+
createReadOptions(datasetUri, catalogConfig, versionId, Optional.empty(), Optional.empty());
195+
StructType schema = getSchema(ident, datasetUri, readOptions, null);
196+
197+
return new LanceDataset(readOptions, schema, null, null, null);
198+
}
172199
}

‎lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceCountStarPartitionReader.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,7 @@ private Dataset openDataset(LanceSparkReadOptions readOptions) {
9696
.allocator(allocator)
9797
.namespace(readOptions.getNamespace())
9898
.tableId(readOptions.getTableId())
99+
.readOptions(readOptions.toReadOptions())
99100
.build();
100101
} else {
101102
return Dataset.open()

‎lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScanBuilder.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,7 @@ private Dataset getOrOpenDataset() {
103103
.allocator(LanceRuntime.allocator())
104104
.namespace(readOptions.getNamespace())
105105
.tableId(readOptions.getTableId())
106+
.readOptions(readOptions.toReadOptions())
106107
.build();
107108
} else {
108109
lazyDataset =

‎lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceSplit.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ private static Dataset openDataset(LanceSparkReadOptions readOptions) {
5151
.allocator(LanceRuntime.allocator())
5252
.namespace(readOptions.getNamespace())
5353
.tableId(readOptions.getTableId())
54+
.readOptions(readOptions.toReadOptions())
5455
.build();
5556
} else {
5657
return Dataset.open()

0 commit comments

Comments
 (0)