Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Count rows as a metadata only operation #1388

Open
wants to merge 18 commits into
base: main
Choose a base branch
from
+781 −2
Open

Count rows as a metadata only operation #1388

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
731542e
Create test_scan_count.py
tusharchou Nov 28, 2024
c6c971e
moved test_scan_count.py to tests
tusharchou Nov 28, 2024
da18837
implemented count in data scan
tusharchou Nov 28, 2024
3104a2f
tested table scan count in test_sql catalog
tusharchou Nov 28, 2024
c2740ea
refactoring
tusharchou Nov 28, 2024
90bca84
make lint
tusharchou Nov 28, 2024
f7202b9
Merge pull request #1 from tusharchou/gh-1223-count-rows-metadata-onl…
tusharchou Nov 28, 2024
c7205b3
Merge branch 'apache:main' into main
tusharchou Dec 2, 2024
09f9c10
Merge branch 'apache:main' into main
tusharchou Dec 11, 2024
1e9da22
Merge branch 'apache:main' into main
tusharchou Dec 18, 2024
3ab20d4
Merge branch 'apache:main' into main
tusharchou Dec 21, 2024
091c0af
implemeted residual_evaluator.py with tests
tusharchou Dec 24, 2024
3cd797d
added license
tusharchou Dec 24, 2024
6b0924e
fixed lint
tusharchou Dec 24, 2024
96cb4e9
fixed lint errors
tusharchou Dec 24, 2024
212c83b
Merge pull request #3 from tusharchou/gh-1223-metadata-only-row-count
tusharchou Dec 24, 2024
8bc65fa
Merge branch 'apache:main' into main
tusharchou Dec 30, 2024
8bb039f
Gh 1223 metadata only row count (#4)
tusharchou Dec 31, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions pyiceberg/table/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1227,6 +1227,9 @@ def filter(self: S, expr: Union[str, BooleanExpression]) -> S:
def with_case_sensitive(self: S, case_sensitive: bool = True) -> S:
return self.update(case_sensitive=case_sensitive)

@abstractmethod
def count(self) -> int: ...


class ScanTask(ABC):
pass
Expand Down Expand Up @@ -1493,6 +1496,13 @@ def to_ray(self) -> ray.data.dataset.Dataset:

return ray.data.from_arrow(self.to_arrow())

def count(self) -> int:
res = 0
tasks = self.plan_files()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this count will not be accurate when there are deletes files

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @kevinjqliu thank you for the review. I am trying to account for positional deletes, do you have a suggestion on how that can be achieved?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, this can be widely off, not just because the merge-on-read deletes, but because plan_files returns all the files that (might) contain relevant rows. For example, if it cannot be determined if has relevant data, it will be returned by plan_files.

I think there are two ways forward:

  • One is similar to how we handle deletes. For deletes, we check if the whole file matches, if this is the case, then we can simply drop the file from the metadata. You can find the code here. If a file fully matches, is is valid to use task.file.record_count. We would need to extend this to also see if there are also merge-on-read deletes as Kevin already mentioned, or just fail when there are positional deletes.
  • A cleaner option, but this is a bit more work, but pretty exciting, would be to include the residual-predicate in the FileScanTask. When we run a query, like day_added = 2024-12-01 and user_id = 10, then the day_added = 2024-12-01 might be satisfied with the partitioning already. This is the case when the table is partitioned by day, and we know that all the data in the file evaluates true for day_added = 2024-12-01, then we need to open the file, and filter for user_id = 10. If we would leave out the user_id = 10, then it would be ALWAYS_TRUE, and then we know that we can just use task.file.record_count. This way we could very easily loop over the .plan_files():
def count(self) -> int:
    res = 0
    tasks = self.plan_files()
    for task in tasks:
        if task.residual == ALWAYS_TRUE and len(task.delete_files):
            res += task.file.record_count
            else:
            # pseudocode, open the table, and apply the filter and deletes
            res += len(_table_from_scan_task(task))
    return res

To get to the second step, we first have to port the ResidualEvaluator. The java code can be found here, including some excellent tests.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @Fokko I have added Residual Evaluator with Tests.
Now I am trying to create the breaking tests for count where delete has occurred and the counts should differ

@pytest.mark.integration
@pytest.mark.filterwarnings("ignore:Merge on read is not yet supported, falling back to copy-on-write")
def test_delete_partitioned_table_positional_deletes(spark: SparkSession, session_catalog: RestCatalog) -> None:
    identifier = "default.table_partitioned_delete"

    run_spark_commands(
        spark,
        [
            f"DROP TABLE IF EXISTS {identifier}",
            f"""
            CREATE TABLE {identifier} (
                number_partitioned  int,
                number              int
            )
            USING iceberg
            PARTITIONED BY (number_partitioned)
            TBLPROPERTIES(
                'format-version' = 2,
                'write.delete.mode'='merge-on-read',
                'write.update.mode'='merge-on-read',
                'write.merge.mode'='merge-on-read'
            )
        """,
            f"""
            INSERT INTO {identifier} VALUES (10, 20), (10, 30), (10, 40)
        """,
            # Generate a positional delete
            f"""
            DELETE FROM {identifier} WHERE number = 30
        """,
        ],
    )

    tbl = session_catalog.load_table(identifier)

    # Assert that there is just a single Parquet file, that has one merge on read file
    files = list(tbl.scan().plan_files())
    assert len(files) == 1
    assert len(files[0].delete_files) == 1
    # Will rewrite a data file without the positional delete
    assert tbl.scan().to_arrow().to_pydict() == {"number_partitioned": [10, 10], "number": [20, 40]}
    assert tbl.scan().count() == 2
    assert tbl.scan().count() == 1

    tbl.delete(EqualTo("number", 40))

    # One positional delete has been added, but an OVERWRITE status is set
    # https://github.com/apache/iceberg/issues/10122
    assert [snapshot.summary.operation.value for snapshot in tbl.snapshots()] == ["append", "overwrite", "overwrite"]
    assert tbl.scan().to_arrow().to_pydict() == {"number_partitioned": [10], "number": [20]}
    assert tbl.scan().count() == 1

for task in tasks:
res += task.file.record_count
return res
Copy link

@gli-chris-hao gli-chris-hao Dec 31, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I love this approach! My only concern is about loading too much data into memory at once, although this is loading just one file at a time, in the worst case some file could potentially be very large? Shall we define a threshold and check, for example, if file size < XXX, load entire file, otherwise turn it into pa.RecordBatchReader and read stream of record batches for counting.

target_schema = schema_to_pyarrow(self.projection())

batches = ArrowScan(
    self.table_metadata, self.io, self.projection(), self.row_filter, self.case_sensitive, self.limit
).to_record_batches([task])

reader = pa.RecordBatchReader.from_batches(
    target_schema,
    batches,
)

count = 0
for batch in reader:
    count += batch.num_rows
return count

https://github.com/apache/iceberg-python/blob/main/pyiceberg/table/__init__.py#L1541-L1564



@dataclass(frozen=True)
class WriteTask:
Expand Down
52 changes: 52 additions & 0 deletions tests/catalog/test_sql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1431,6 +1431,58 @@ def test_append_table(catalog: SqlCatalog, table_schema_simple: Schema, table_id
assert df == table.scan().to_arrow()


@pytest.mark.parametrize(
"catalog",
[
lazy_fixture("catalog_memory"),
lazy_fixture("catalog_sqlite"),
lazy_fixture("catalog_sqlite_without_rowcount"),
lazy_fixture("catalog_sqlite_fsspec"),
],
)
@pytest.mark.parametrize(
"table_identifier",
[
lazy_fixture("random_table_identifier"),
lazy_fixture("random_hierarchical_identifier"),
lazy_fixture("random_table_identifier_with_catalog"),
],
)
def test_count_table(catalog: SqlCatalog, table_schema_simple: Schema, table_identifier: Identifier) -> None:
table_identifier_nocatalog = catalog._identifier_to_tuple_without_catalog(table_identifier)
namespace = Catalog.namespace_from(table_identifier_nocatalog)
catalog.create_namespace(namespace)
table = catalog.create_table(table_identifier, table_schema_simple)

df = pa.Table.from_pydict(
{
"foo": ["a"],
"bar": [1],
"baz": [True],
},
schema=schema_to_pyarrow(table_schema_simple),
)

table.append(df)

# new snapshot is written in APPEND mode
assert len(table.metadata.snapshots) == 1
assert table.metadata.snapshots[0].snapshot_id == table.metadata.current_snapshot_id
assert table.metadata.snapshots[0].parent_snapshot_id is None
assert table.metadata.snapshots[0].sequence_number == 1
assert table.metadata.snapshots[0].summary is not None
assert table.metadata.snapshots[0].summary.operation == Operation.APPEND
assert table.metadata.snapshots[0].summary["added-data-files"] == "1"
assert table.metadata.snapshots[0].summary["added-records"] == "1"
assert table.metadata.snapshots[0].summary["total-data-files"] == "1"
assert table.metadata.snapshots[0].summary["total-records"] == "1"
assert len(table.metadata.metadata_log) == 1

# read back the data
assert df == table.scan().to_arrow()
assert len(table.scan().to_arrow()) == table.scan().count()


@pytest.mark.parametrize(
"catalog",
[
Expand Down