wangyong9999 commented on code in PR #315: URL: https://github.com/apache/paimon-cpp/pull/315#discussion_r3996417595
########## src/paimon/format/lance/lance_file_batch_reader.cpp: ########## @@ -0,0 +1,304 @@ +/* + * 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. + */ + +#include "paimon/format/lance/lance_file_batch_reader.h" + +#include <utility> + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "fmt/format.h" +#include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/utils/arrow/arrow_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/core/utils/nested_projection_utils.h" +#include "paimon/format/lance/lance_utils.h" +#include "paimon/fs/file_system.h" + +namespace paimon::lance { + +LanceFileBatchReader::LanceFileBatchReader(const std::shared_ptr<InputStream>& input, + int32_t batch_size, uint32_t batch_readahead, + PaimonLanceReader* reader, + const std::shared_ptr<arrow::Schema>& file_schema, + uint64_t total_rows, + const std::shared_ptr<MemoryPool>& pool, + const std::shared_ptr<arrow::MemoryPool>& arrow_pool) + : input_(input), + batch_size_(batch_size), + batch_readahead_(batch_readahead), + reader_(reader), + file_schema_(file_schema), + total_rows_(total_rows), + pool_(pool), + arrow_pool_(arrow_pool), + metrics_(std::make_shared<MetricsImpl>()) {} + +Result<std::unique_ptr<LanceFileBatchReader>> LanceFileBatchReader::Create( + const std::shared_ptr<InputStream>& input, int32_t batch_size, uint32_t batch_readahead, + const std::map<std::string, std::string>& options, const std::shared_ptr<MemoryPool>& pool, + const std::shared_ptr<arrow::MemoryPool>& arrow_pool) { + if (input == nullptr || pool == nullptr || arrow_pool == nullptr || batch_size <= 0 || + batch_readahead == 0) { + return Status::Invalid( + "Lance reader requires non-null input and memory pools, positive " + "batch size, and positive batch readahead"); + } + PAIMON_ASSIGN_OR_RAISE(std::string uri, input->GetUri()); + if (uri.empty()) { + return Status::Invalid("Lance reader input URI is empty"); + } + LanceStorageOptions storage_options = GetLanceStorageOptions(options, uri); + std::vector<const char*> keys; + std::vector<const char*> values; + keys.reserve(storage_options.size()); + values.reserve(storage_options.size()); + for (const auto& [key, value] : storage_options) { + keys.push_back(key.c_str()); + values.push_back(value.c_str()); + } + PaimonLanceReader* reader = nullptr; + if (paimon_lance_reader_open(uri.c_str(), keys.data(), values.data(), keys.size(), &reader) != + 0) { + return LanceFfiError("open Lance reader"); + } + + ::ArrowSchema ffi_schema = {}; + if (paimon_lance_reader_export_schema(reader, &ffi_schema) != 0) { + paimon_lance_reader_free(reader); + return LanceFfiError("read Lance schema"); + } + arrow::Result<std::shared_ptr<arrow::Schema>> schema_result = arrow::ImportSchema(&ffi_schema); + if (!schema_result.ok()) { + paimon_lance_reader_free(reader); + return ToPaimonStatus(schema_result.status()); + } + uint64_t total_rows = 0; + if (paimon_lance_reader_num_rows(reader, &total_rows) != 0) { + paimon_lance_reader_free(reader); + return LanceFfiError("read Lance row count"); + } + return std::unique_ptr<LanceFileBatchReader>(new LanceFileBatchReader( + input, batch_size, batch_readahead, reader, std::move(schema_result).MoveValueUnsafe(), + total_rows, pool, arrow_pool)); +} + +LanceFileBatchReader::~LanceFileBatchReader() { + CloseInternal(); +} + +Status LanceFileBatchReader::OpenStream() { + if (batch_reader_ != nullptr || (has_selection_ && selection_row_ids_.empty())) { + return Status::OK(); + } + std::vector<const char*> projection_pointers; + projection_pointers.reserve(projection_names_.size()); + for (const std::string& name : projection_names_) { + projection_pointers.push_back(name.c_str()); + } + if (paimon_lance_reader_open_stream(reader_, static_cast<uint32_t>(batch_size_), + batch_readahead_, projection_pointers.data(), + projection_pointers.size(), has_selection_, + selection_starts_.data(), selection_ends_.data(), + selection_starts_.size(), &batch_reader_) != 0) { + return LanceFfiError("open Lance batch reader"); + } + return Status::OK(); +} + +Result<BatchReader::ReadBatch> LanceFileBatchReader::AlignBatch(ReadBatch batch) const { + if (read_schema_ == nullptr) { + return batch; + } + auto& [ffi_array, ffi_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> array, + arrow::ImportArray(ffi_array.get(), ffi_schema.get())); + if (array->type_id() != arrow::Type::STRUCT) { + return Status::Invalid("Lance reader returned a non-struct Arrow batch"); + } + auto struct_array = checked_pointer_cast<arrow::StructArray>(array); + arrow::ArrayVector fields; + fields.reserve(read_schema_->num_fields()); + for (const std::shared_ptr<arrow::Field>& field : read_schema_->fields()) { + std::shared_ptr<arrow::Array> child = struct_array->GetFieldByName(field->name()); + if (child == nullptr) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + child, + arrow::MakeArrayOfNull(field->type(), struct_array->length(), arrow_pool_.get())); + } else if (!child->type()->Equals(field->type()) && Review Comment: Agreed. Removed the format reader's MakeArrayOfNull fallback and stopped filtering missing projection names. Missing table fields remain FieldMappingReader's responsibility. Nested missing-field and data-evolution tests pass. Commit: c45f394c. ########## src/paimon/format/lance/lance_utils.cpp: ########## @@ -0,0 +1,159 @@ +/* + * 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. + */ + +#include "paimon/format/lance/lance_utils.h" + +#include <cstring> +#include <map> +#include <optional> + +#include "paimon/common/utils/path_util.h" +#include "paimon/common/utils/string_utils.h" +#include "paimon/format/lance/lance_ffi.h" + +namespace paimon::lance { +namespace { + +void CopyOption(const std::map<std::string, std::string>& options, const std::string& source, + const std::string& target, std::map<std::string, std::string>* storage_options) { + auto iter = options.find(source); + if (iter != options.end() && !iter->second.empty()) { + storage_options->insert_or_assign(target, iter->second); + } +} + +const char* CanonicalS3OptionName(const std::string& option) { + if (option == "access-key" || option == "access.key" || option == "accessKeyId") { + return "s3.access-key"; + } + if (option == "secret-key" || option == "secret.key" || option == "accessKeySecret") { + return "s3.secret-key"; + } + if (option == "session.token" || option == "session-token" || option == "security.token" || + option == "security-token" || option == "securityToken") { + return "s3.session.token"; + } + if (option == "endpoint") { + return "s3.endpoint"; + } + if (option == "region") { + return "s3.region"; + } + if (option == "path-style-access" || option == "path.style.access") { + return "s3.path-style-access"; + } + return nullptr; +} + +std::map<std::string, std::string> NormalizeS3Options( + const std::map<std::string, std::string>& options) { + std::map<std::string, std::string> normalized = options; + for (const auto& [key, value] : options) { + for (const char* prefix : {"s3a.", "fs.s3.", "fs.s3a.", "s3."}) { + if (!StringUtils::StartsWith(key, prefix)) { + continue; + } + const char* canonical = CanonicalS3OptionName(key.substr(std::strlen(prefix))); + if (canonical != nullptr && normalized.find(canonical) == normalized.end()) { + normalized.emplace(canonical, value); + } + break; + } + } + return normalized; +} + +void AddS3Options(const std::map<std::string, std::string>& options, + std::map<std::string, std::string>* storage_options) { + const std::map<std::string, std::string> normalized = NormalizeS3Options(options); + CopyOption(normalized, "s3.region", "region", storage_options); + CopyOption(normalized, "s3.endpoint", "endpoint", storage_options); + CopyOption(normalized, "s3.access-key", "access_key_id", storage_options); + CopyOption(normalized, "s3.secret-key", "secret_access_key", storage_options); + CopyOption(normalized, "s3.session.token", "session_token", storage_options); + auto endpoint = storage_options->find("endpoint"); + if (endpoint != storage_options->end() && endpoint->second.find("://") == std::string::npos) { + endpoint->second = "https://"; + endpoint->second; + } + auto path_style = normalized.find("s3.path-style-access"); + if (path_style != normalized.end()) { + std::optional<bool> enabled = StringUtils::StringToValue<bool>(path_style->second); + if (enabled) { + storage_options->insert_or_assign("virtual_hosted_style_request", + *enabled ? "false" : "true"); + } + } +} + +void AddOssOptions(const std::map<std::string, std::string>& options, const std::string& bucket, + std::map<std::string, std::string>* storage_options) { + auto copy_oss_option = [&](const std::string& suffix, const std::string& target) { + CopyOption(options, "fs.oss." + suffix, target, storage_options); + if (!bucket.empty()) { + CopyOption(options, "fs.oss.bucket." + bucket + "." + suffix, target, storage_options); + } + }; + copy_oss_option("endpoint", "oss_endpoint"); + copy_oss_option("region", "oss_region"); + copy_oss_option("accessKeyId", "oss_access_key_id"); + copy_oss_option("accessKeySecret", "oss_secret_access_key"); + copy_oss_option("sessionToken", "oss_session_token"); + copy_oss_option("securityToken", "oss_session_token"); + copy_oss_option("accessKeyId", "access_key_id"); + copy_oss_option("accessKeySecret", "secret_access_key"); + copy_oss_option("sessionToken", "session_token"); + copy_oss_option("securityToken", "session_token"); + auto endpoint = storage_options->find("oss_endpoint"); + if (endpoint != storage_options->end() && endpoint->second.find("://") == std::string::npos) { + endpoint->second = "https://"; + endpoint->second; + } +} + +} // namespace + +LanceStorageOptions GetLanceStorageOptions(const std::map<std::string, std::string>& options, + const std::string& uri) { + std::map<std::string, std::string> storage_options; Review Comment: The pinned oss:// backend uses OpenDAL: it consumes oss_endpoint, adds the bucket itself and defaults bare hosts to HTTPS. Java's S3-style endpoint/virtual_hosted_style_request keys are unused here. Kept the equivalent mapping and added comments, endpoint/override tests and documentation. Also documented that this backend does not forward table-option OSS session tokens; mapping tests do not establish STS support. Commit: c45f394c. ########## src/paimon/format/lance/lance_format_writer.cpp: ########## @@ -0,0 +1,139 @@ +/* + * 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. + */ + +#include "paimon/format/lance/lance_format_writer.h" + +#include <cstdint> +#include <utility> +#include <vector> + +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/utils/arrow/arrow_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/math.h" +#include "paimon/common/utils/scope_guard.h" + +namespace paimon::lance { + +LanceFormatWriter::LanceFormatWriter(const std::shared_ptr<arrow::Schema>& schema, + const std::shared_ptr<arrow::MemoryPool>& arrow_pool, + PaimonLanceWriter* writer) + : schema_(schema), + arrow_pool_(arrow_pool), + writer_(writer), + metrics_(std::make_shared<MetricsImpl>()) {} + +Result<std::unique_ptr<LanceFormatWriter>> LanceFormatWriter::Create( + const std::string& path, const std::shared_ptr<arrow::Schema>& schema, + const LanceStorageOptions& storage_options, + const std::shared_ptr<arrow::MemoryPool>& arrow_pool) { + if (path.empty() || schema == nullptr || arrow_pool == nullptr) { + return Status::Invalid("Lance writer requires a path, schema, and memory pool"); + } + std::vector<const char*> keys; + std::vector<const char*> values; + keys.reserve(storage_options.size()); + values.reserve(storage_options.size()); + for (const auto& [key, value] : storage_options) { + keys.push_back(key.c_str()); + values.push_back(value.c_str()); + } + PaimonLanceWriter* writer = nullptr; + if (paimon_lance_writer_open(path.c_str(), keys.data(), values.data(), keys.size(), &writer) != + 0) { + return LanceFfiError("open Lance writer"); + } + return std::unique_ptr<LanceFormatWriter>(new LanceFormatWriter(schema, arrow_pool, writer)); +} + +LanceFormatWriter::~LanceFormatWriter() { + paimon_lance_writer_free(writer_); +} + +Status LanceFormatWriter::AddBatch(::ArrowArray* batch) { + if (batch == nullptr) { + return Status::Invalid("Lance writer batch is nullptr"); + } + if (finished_) { + return Status::Invalid("cannot add a batch after Lance writer is finished"); + } + auto row_count = static_cast<uint64_t>(batch->length); + ::ArrowSchema import_schema = {}; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema_, &import_schema)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> array, + arrow::ImportArray(batch, &import_schema)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> normalized, + ArrowUtils::NormalizeArrayOffsets(array, arrow_pool_.get())); Review Comment: This handles sliced nested arrays crossing the C++/Rust Arrow C Data Interface, rather than rejecting all non-zero offsets. Normalization keeps child buffers and offsets consistent for Lance 0.39. Added a comment; the supported-types test writes sliced batches and the new ROW/ARRAY<ROW> test reads back a non-zero-offset slice. Commit: c45f394c. ########## src/paimon/format/lance/lance_format_writer.cpp: ########## @@ -0,0 +1,139 @@ +/* + * 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. + */ + +#include "paimon/format/lance/lance_format_writer.h" + +#include <cstdint> +#include <utility> +#include <vector> + +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/utils/arrow/arrow_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/math.h" +#include "paimon/common/utils/scope_guard.h" + +namespace paimon::lance { + +LanceFormatWriter::LanceFormatWriter(const std::shared_ptr<arrow::Schema>& schema, + const std::shared_ptr<arrow::MemoryPool>& arrow_pool, + PaimonLanceWriter* writer) + : schema_(schema), + arrow_pool_(arrow_pool), + writer_(writer), + metrics_(std::make_shared<MetricsImpl>()) {} + +Result<std::unique_ptr<LanceFormatWriter>> LanceFormatWriter::Create( + const std::string& path, const std::shared_ptr<arrow::Schema>& schema, + const LanceStorageOptions& storage_options, + const std::shared_ptr<arrow::MemoryPool>& arrow_pool) { + if (path.empty() || schema == nullptr || arrow_pool == nullptr) { + return Status::Invalid("Lance writer requires a path, schema, and memory pool"); + } + std::vector<const char*> keys; + std::vector<const char*> values; + keys.reserve(storage_options.size()); + values.reserve(storage_options.size()); + for (const auto& [key, value] : storage_options) { + keys.push_back(key.c_str()); + values.push_back(value.c_str()); + } + PaimonLanceWriter* writer = nullptr; + if (paimon_lance_writer_open(path.c_str(), keys.data(), values.data(), keys.size(), &writer) != + 0) { + return LanceFfiError("open Lance writer"); + } + return std::unique_ptr<LanceFormatWriter>(new LanceFormatWriter(schema, arrow_pool, writer)); +} + +LanceFormatWriter::~LanceFormatWriter() { + paimon_lance_writer_free(writer_); +} + +Status LanceFormatWriter::AddBatch(::ArrowArray* batch) { + if (batch == nullptr) { + return Status::Invalid("Lance writer batch is nullptr"); + } + if (finished_) { + return Status::Invalid("cannot add a batch after Lance writer is finished"); + } + auto row_count = static_cast<uint64_t>(batch->length); + ::ArrowSchema import_schema = {}; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema_, &import_schema)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> array, + arrow::ImportArray(batch, &import_schema)); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> normalized, + ArrowUtils::NormalizeArrayOffsets(array, arrow_pool_.get())); + ::ArrowArray ffi_array = {}; + ::ArrowSchema ffi_schema = {}; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*normalized, &ffi_array, &ffi_schema)); + ScopeGuard guard([&]() { + ArrowArrayRelease(&ffi_array); + ArrowSchemaRelease(&ffi_schema); + }); + if (paimon_lance_writer_write(writer_, &ffi_array, &ffi_schema) != 0) { + return LanceFfiError("write Lance batch"); + } + rows_written_ += row_count; + return Status::OK(); +} + +Status LanceFormatWriter::Flush() { + return Status::OK(); +} + +Status LanceFormatWriter::Finish() { + if (finished_) { + return Status::OK(); + } + uint64_t row_count = 0; + if (paimon_lance_writer_finish(writer_, &row_count) != 0) { + return LanceFfiError("finish Lance writer"); + } + if (row_count != rows_written_) { + return Status::Invalid("Lance writer row count mismatch"); + } + finished_ = true; + return Status::OK(); +} + +Result<bool> LanceFormatWriter::ReachTargetSize(bool suggested_check, int64_t target_size) const { + if (!suggested_check) { + return false; + } + PAIMON_RETURN_NOT_OK(ValidateValueNonNegative(target_size, "Lance target size")); + uint64_t position = 0; + if (paimon_lance_writer_tell(writer_, &position) != 0) { + return LanceFfiError("get Lance writer position"); + } + return position >= static_cast<uint64_t>(target_size); Review Comment: Added the comment: C++ checks the current encoded file position, excluding buffered data, while Java estimates input Arrow-buffer size. The same target-file-size need not produce identical rolling boundaries. Commit: c45f394c. ########## src/paimon/format/lance/lance_file_format_test.cpp: ########## @@ -0,0 +1,344 @@ +/* + * 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. + */ + +#include <cstdint> +#include <map> +#include <memory> +#include <string> +#include <utility> +#include <vector> + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/ipc/json_simple.h" +#include "gtest/gtest.h" +#include "paimon/common/types/data_field.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/path_util.h" +#include "paimon/format/column_stats.h" +#include "paimon/format/file_format.h" +#include "paimon/format/file_format_factory.h" +#include "paimon/format/format_stats_extractor.h" +#include "paimon/format/format_writer.h" +#include "paimon/format/lance/lance_utils.h" +#include "paimon/format/writer_builder.h" +#include "paimon/fs/file_system.h" +#include "paimon/fs/local/local_file_system.h" +#include "paimon/testing/utils/read_result_collector.h" +#include "paimon/testing/utils/testharness.h" +#include "paimon/utils/roaring_bitmap32.h" + +namespace paimon::lance::test { + +class LanceFileFormatTest : public ::testing::Test { + public: + void SetUp() override { + ASSERT_OK_AND_ASSIGN(format_, FileFormatFactory::Get("lance", {{"file.format", "lance"}})); + file_system_ = std::make_shared<LocalFileSystem>(); + directory_ = paimon::test::UniqueTestDirectory::Create(); + ASSERT_NE(directory_, nullptr); + pool_ = GetDefaultPool(); + arrow_pool_ = GetArrowPool(pool_); + } + + Status WriteFile(const std::string& path, const std::shared_ptr<arrow::Schema>& schema, + const std::shared_ptr<arrow::Array>& array, int32_t batch_size) const { + ::ArrowSchema ffi_schema = {}; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, &ffi_schema)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<WriterBuilder> writer_builder, + format_->CreateWriterBuilder(&ffi_schema, batch_size)); + writer_builder->WithMemoryPool(pool_); + auto* direct_writer_builder = dynamic_cast<DirectWriterBuilder*>(writer_builder.get()); + if (direct_writer_builder == nullptr) { + return Status::Invalid("Lance writer builder is not path based"); + } + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<FormatWriter> writer, + direct_writer_builder->BuildFromPath(path)); + for (int64_t offset = 0; offset < array->length(); offset += batch_size) { + std::shared_ptr<arrow::Array> slice = array->Slice(offset, batch_size); + ::ArrowArray ffi_array = {}; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*slice, &ffi_array)); + PAIMON_RETURN_NOT_OK(writer->AddBatch(&ffi_array)); + } + return writer->Finish(); + } + + Result<std::unique_ptr<FileBatchReader>> OpenReader(const std::string& path, + int32_t batch_size) const { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ReaderBuilder> builder, + format_->CreateReaderBuilder(batch_size)); + builder->WithMemoryPool(pool_); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<InputStream> input, file_system_->Open(path)); + return builder->Build(input); + } + + protected: + std::shared_ptr<FileFormat> format_; + std::shared_ptr<LocalFileSystem> file_system_; + std::unique_ptr<paimon::test::UniqueTestDirectory> directory_; + std::shared_ptr<MemoryPool> pool_; + std::shared_ptr<arrow::MemoryPool> arrow_pool_; +}; + +TEST(LanceUtilsTest, ConvertsObjectStoreOptions) { + auto to_map = [](const LanceStorageOptions& values) { + return std::map<std::string, std::string>(values.begin(), values.end()); + }; + auto s3_options = to_map(GetLanceStorageOptions({{"s3.region", "us-west-2"}, + {"s3.endpoint", "http://s3.example.com"}, + {"s3.access-key", "key"}, + {"s3.secret-key", "secret"}, + {"s3.session.token", "token"}, + {"s3.path-style-access", "true"}, + {"lance.storage.region", "override"}}, + "s3://bucket/data.lance")); + ASSERT_EQ(s3_options.at("region"), "override"); + ASSERT_EQ(s3_options.at("endpoint"), "http://s3.example.com";); + ASSERT_EQ(s3_options.at("access_key_id"), "key"); + ASSERT_EQ(s3_options.at("secret_access_key"), "secret"); + ASSERT_EQ(s3_options.at("session_token"), "token"); + ASSERT_EQ(s3_options.at("virtual_hosted_style_request"), "false"); + ASSERT_EQ(s3_options.at("allow_http"), "true"); + + auto python_s3_options = + to_map(GetLanceStorageOptions({{"fs.s3.region", "us-west-1"}, + {"fs.s3.endpoint", "s3.python.example.com"}, + {"fs.s3.accessKeyId", "python-key"}, + {"fs.s3.accessKeySecret", "python-secret"}, + {"fs.s3.securityToken", "python-token"}, + {"fs.s3.path.style.access", "false"}}, + "s3://bucket/data.lance")); + ASSERT_EQ(python_s3_options.at("region"), "us-west-1"); + ASSERT_EQ(python_s3_options.at("endpoint"), "https://s3.python.example.com";); + ASSERT_EQ(python_s3_options.at("access_key_id"), "python-key"); + ASSERT_EQ(python_s3_options.at("secret_access_key"), "python-secret"); + ASSERT_EQ(python_s3_options.at("session_token"), "python-token"); + ASSERT_EQ(python_s3_options.at("virtual_hosted_style_request"), "true"); + + auto java_s3_options = to_map(GetLanceStorageOptions({{"s3a.region", "eu-west-1"}, + {"s3a.access.key", "java-key"}, + {"s3a.secret.key", "java-secret"}, + {"s3a.session.token", "java-token"}, + {"s3.access-key", "canonical-key"}}, + "s3://bucket/data.lance")); + ASSERT_EQ(java_s3_options.at("region"), "eu-west-1"); + ASSERT_EQ(java_s3_options.at("access_key_id"), "canonical-key"); + ASSERT_EQ(java_s3_options.at("secret_access_key"), "java-secret"); + ASSERT_EQ(java_s3_options.at("session_token"), "java-token"); + + auto oss_options = + to_map(GetLanceStorageOptions({{"fs.oss.endpoint", "global.example.com"}, + {"fs.oss.accessKeyId", "global-key"}, + {"fs.oss.accessKeySecret", "global-secret"}, + {"fs.oss.bucket.bucket.endpoint", "bucket.example.com"}, + {"fs.oss.bucket.bucket.accessKeyId", "bucket-key"}, + {"fs.oss.bucket.bucket.accessKeySecret", "bucket-secret"}, + {"fs.oss.bucket.bucket.securityToken", "bucket-token"}}, + "oss://bucket/data.lance")); + ASSERT_EQ(oss_options.at("oss_endpoint"), "https://bucket.example.com";); + ASSERT_EQ(oss_options.at("oss_access_key_id"), "bucket-key"); + ASSERT_EQ(oss_options.at("oss_secret_access_key"), "bucket-secret"); + ASSERT_EQ(oss_options.at("oss_session_token"), "bucket-token"); + ASSERT_EQ(oss_options.at("access_key_id"), "bucket-key"); + ASSERT_EQ(oss_options.at("secret_access_key"), "bucket-secret"); + ASSERT_EQ(oss_options.at("session_token"), "bucket-token"); +} + +TEST_F(LanceFileFormatTest, WriteThenReadSupportedTypes) { + arrow::FieldVector fields = { + arrow::field("bool_col", arrow::boolean()), + arrow::field("i8_col", arrow::int8()), + arrow::field("i16_col", arrow::int16()), + arrow::field("i32_col", arrow::int32()), + arrow::field("i64_col", arrow::int64()), + arrow::field("float_col", arrow::float32()), + arrow::field("double_col", arrow::float64()), + arrow::field("string_col", arrow::utf8()), + arrow::field("binary_col", arrow::binary()), + arrow::field("date_col", arrow::date32()), + arrow::field("time_col", arrow::time32(arrow::TimeUnit::MILLI)), + arrow::field("timestamp_col", arrow::timestamp(arrow::TimeUnit::MICRO)), + arrow::field("decimal_col", arrow::decimal128(12, 2)), + arrow::field("array_col", arrow::list(arrow::float32())), + arrow::field("row_col", arrow::struct_({arrow::field("value", arrow::int32())})), + arrow::field( + "vector_col", + arrow::fixed_size_list(arrow::field("item", arrow::float32(), /*nullable=*/false), 2)), + }; + std::shared_ptr<arrow::Schema> schema = arrow::schema(fields); + std::shared_ptr<arrow::Array> expected = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([ + [true,-1,-2,-3,-4,1.5,2.5,"one","abc",1,1234, + "1970-01-01 00:00:00.000001","12.34",[1.0,2.0],[7],[3.0,4.0]], + [false,1,2,3,4,3.5,4.5,"two","xyz",2,5678, + "2030-12-31 23:59:59.999999","-12.34",[],[8],[5.0,6.0]], + [null,null,null,null,null,null,null,null,null,null,null,null,null,null,[null],null] + ])") + .ValueOrDie(); + std::string path = PathUtil::JoinPath(directory_->Str(), "supported-types.lance"); + ASSERT_OK(WriteFile(path, schema, expected, /*batch_size=*/2)); + + for (int32_t batch_size : {1, 2, 4}) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileBatchReader> reader, OpenReader(path, batch_size)); + ::ArrowSchema ffi_schema = {}; + ASSERT_TRUE(arrow::ExportSchema(*schema, &ffi_schema).ok()); + ASSERT_OK(reader->SetReadSchema(&ffi_schema, /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::ChunkedArray> actual, + paimon::test::ReadResultCollector::CollectResult(std::move(reader))); + ASSERT_TRUE(actual->Equals(arrow::ChunkedArray(expected))) << actual->ToString() << "\nvs\n" + << expected->ToString(); + } +} + +TEST_F(LanceFileFormatTest, RejectsNullTopLevelRows) { + arrow::FieldVector fields = {arrow::field("id", arrow::int32())}; + std::shared_ptr<arrow::Array> data = + arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([[1],null])") + .ValueOrDie(); + ASSERT_NOK_WITH_MSG(WriteFile(PathUtil::JoinPath(directory_->Str(), "null-row.lance"), + arrow::schema(fields), data, /*batch_size=*/2), + "Lance writer does not accept null top-level rows"); +} + +TEST_F(LanceFileFormatTest, ProjectionSelectionAndFileRowIds) { + arrow::FieldVector fields = {arrow::field("id", arrow::int32()), + arrow::field("name", arrow::utf8())}; + std::shared_ptr<arrow::Schema> schema = arrow::schema(fields); + std::shared_ptr<arrow::Array> data = + arrow::ipc::internal::json::ArrayFromJSON( + arrow::struct_(fields), R"([[0,"zero"],[1,"one"],[2,"two"],[3,"three"],[4,"four"]])") + .ValueOrDie(); + std::string path = PathUtil::JoinPath(directory_->Str(), "selection.lance"); + ASSERT_OK(WriteFile(path, schema, data, /*batch_size=*/2)); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr<FileBatchReader> reader, OpenReader(path, 2)); + std::shared_ptr<arrow::Field> read_field = + arrow::field("name", arrow::utf8(), /*nullable=*/false, + arrow::key_value_metadata({DataField::FIELD_ID}, {"1"})); + std::shared_ptr<arrow::Schema> projected_schema = arrow::schema({read_field}); + ::ArrowSchema ffi_schema = {}; + ASSERT_TRUE(arrow::ExportSchema(*projected_schema, &ffi_schema).ok()); + ASSERT_OK(reader->SetReadSchema(&ffi_schema, /*predicate=*/nullptr, + RoaringBitmap32::From({1, 2, 4}))); + + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch first_batch, reader->NextBatch()); + arrow::Result<std::shared_ptr<arrow::Array>> first_result = + arrow::ImportArray(first_batch.first.get(), first_batch.second.get()); + ASSERT_TRUE(first_result.ok()) << first_result.status().ToString(); + std::shared_ptr<arrow::Array> first = std::move(first_result).MoveValueUnsafe(); + auto first_type = std::static_pointer_cast<arrow::StructType>(first->type()); + ASSERT_TRUE(first_type->field(0)->Equals(*read_field, /*check_metadata=*/true)); + ASSERT_TRUE(first->Equals(arrow::ipc::internal::json::ArrayFromJSON( Review Comment: This test checks field metadata and nullability. CollectResultOneBatch uses DictArrayConverter, which normalizes those attributes away. Kept the direct Arrow import here and added a comment; result-only tests continue using ReadResultCollector. Commit: c45f394c. -- 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]
